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}