Skip to main content

somatize_core/
keys.rs

1//! Reserved keys in a run's output store.
2//!
3//! A run threads node outputs through one `HashMap<String, Value>`, and a
4//! few entries in it are not node outputs: the graph's input, a specific
5//! node's input, and the state a trainable node learned. They are
6//! distinguished by a `__` prefix.
7//!
8//! The prefixes used to be spelled inline — written with `format!` in the
9//! runner and read back with `strip_prefix` in three different crates. This
10//! module is the only place that knows how the key is spelled, so a change
11//! to it cannot leave one reader behind.
12//!
13//! The prefix is a convention, not a guarantee: a node whose id is literally
14//! `__state_x` would collide. [`is_reserved`] is what a caller uses to keep
15//! these out of a list of node outputs.
16
17/// Prefix for a state a trainable node learned during a fit.
18const STATE: &str = "__state_";
19/// Prefix for the input handed to a specific node.
20const INPUT: &str = "__input_";
21/// The graph's own input, available to every root node.
22pub const GRAPH_INPUT: &str = "__input__";
23
24/// Where the state fitted for `node_id` is stored.
25pub fn state_key(node_id: &str) -> String {
26    format!("{STATE}{node_id}")
27}
28
29/// Where the input handed to `node_id` is stored.
30pub fn input_key(node_id: &str) -> String {
31    format!("{INPUT}{node_id}")
32}
33
34/// The node whose state this key holds, or `None` if it is not a state key.
35pub fn node_of_state_key(key: &str) -> Option<&str> {
36    key.strip_prefix(STATE)
37}
38
39/// Does this key hold an input rather than an output or a state?
40pub fn is_input_key(key: &str) -> bool {
41    key == GRAPH_INPUT || key.starts_with(INPUT)
42}
43
44/// Is this key one of the run's own entries rather than a node's output?
45///
46/// Callers that report "what did each node produce" filter with this;
47/// without it a fit's answer included its own inputs and states as though
48/// nodes had produced them.
49pub fn is_reserved(key: &str) -> bool {
50    key == GRAPH_INPUT || key.starts_with(STATE) || key.starts_with(INPUT)
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56
57    #[test]
58    fn state_keys_round_trip() {
59        let key = state_key("scaler");
60        assert_eq!(node_of_state_key(&key), Some("scaler"));
61        assert!(is_reserved(&key));
62    }
63
64    #[test]
65    fn an_ordinary_node_id_is_not_a_state_key() {
66        assert_eq!(node_of_state_key("scaler"), None);
67        assert!(!is_reserved("scaler"));
68    }
69
70    #[test]
71    fn every_reserved_shape_is_recognised() {
72        assert!(is_reserved(GRAPH_INPUT));
73        assert!(is_reserved(&input_key("model")));
74        assert!(is_reserved(&state_key("model")));
75    }
76}