1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
//! Node representation and execution
use crate::distribution::DistTransferFn;
use crate::graph_data::GraphData;
use std::collections::HashMap;
use std::sync::Arc;
/// Unique identifier for a node
pub type NodeId = usize;
/// Type alias for node execution functions using GraphData
/// Takes GraphData ports as input, returns output ports
pub type NodeFunction = Arc<
dyn Fn(&HashMap<String, GraphData>) -> HashMap<String, GraphData>
+ Send
+ Sync,
>;
/// Represents a node in the graph
#[derive(Clone)]
pub struct Node {
/// Unique identifier
pub id: NodeId,
/// Optional label for visualization
pub label: Option<String>,
/// Function to execute
pub function: NodeFunction,
/// Explicit node implementation/version token used for execution-result caching
pub code_fingerprint: String,
/// Whether the version token was explicitly provided by the caller
pub has_explicit_cache_version: bool,
/// Whether this node is eligible for execution-result caching
pub cacheable: bool,
/// Optional subset of impl-var inputs to include in the cache key
pub cache_key_inputs: Option<Vec<String>>,
/// Input mapping: broadcast_var -> impl_var (what the function sees)
pub input_mapping: HashMap<String, String>,
/// Output mapping: impl_var -> broadcast_var (where function output goes in context)
pub output_mapping: HashMap<String, String>,
/// Branch ID for branch-specific variable resolution (None for main graph nodes)
pub branch_id: Option<usize>,
/// Nodes that this node depends on (connected from)
pub dependencies: Vec<NodeId>,
/// Whether this node is part of a branch
pub is_branch: bool,
/// Variant index if this is part of a variant sweep
pub variant_index: Option<usize>,
/// Variant parameters for this node (param_name -> value)
pub variant_params: HashMap<String, GraphData>,
/// Optional analytical distribution transfer.
///
/// Receives input distributions keyed by **impl_var** names (same keys the function sees)
/// and returns output distributions keyed by **impl_var** output names, or `None` to
/// signal that Monte Carlo fallback should be used for this node.
pub dist_transfer: Option<DistTransferFn>,
}
impl Node {
/// Create a new node
pub fn new(
id: NodeId,
function: NodeFunction,
code_fingerprint: String,
label: Option<String>,
input_mapping: HashMap<String, String>,
output_mapping: HashMap<String, String>,
) -> Self {
Self {
id,
label,
function,
code_fingerprint,
has_explicit_cache_version: false,
cacheable: true,
cache_key_inputs: None,
input_mapping,
output_mapping,
branch_id: None,
dependencies: Vec::new(),
is_branch: false,
variant_index: None,
variant_params: HashMap::new(),
dist_transfer: None,
}
}
/// Gather this node's inputs using impl-var names.
pub fn gather_inputs(&self, context: &HashMap<String, GraphData>) -> HashMap<String, GraphData> {
// Map broadcast context vars to impl vars using input_mapping
// input_mapping: broadcast_var -> impl_var
// Special case: For merge nodes, broadcast_var may be "branch_id:var_name"
self
.input_mapping
.iter()
.filter_map(|(broadcast_key, impl_var)| {
// Handle merge node special format: "branch_id:broadcast_var"
if broadcast_key.contains(':') {
// Parse "branch_id:var_name" and look for "__branch_{id}__{var}"
let parts: Vec<&str> = broadcast_key.split(':').collect();
if parts.len() == 2 {
let prefixed_key = format!("__branch_{}__{}", parts[0], parts[1]);
context
.get(&prefixed_key)
.map(|val| (impl_var.clone(), val.clone()))
} else {
None
}
} else {
// Normal case: direct lookup
context
.get(broadcast_key)
.map(|val| (impl_var.clone(), val.clone()))
}
})
.collect()
}
/// Execute this node with pre-resolved impl-var inputs.
pub fn execute_with_inputs(
&self,
inputs: &HashMap<String, GraphData>,
) -> HashMap<String, GraphData> {
// Execute function with inputs
let func_outputs = (self.function)(inputs);
// Map function outputs to broadcast vars using output_mapping
// output_mapping: impl_var -> broadcast_var
let mut context_outputs = HashMap::new();
for (impl_var, broadcast_var) in &self.output_mapping {
if let Some(value) = func_outputs.get(impl_var) {
context_outputs.insert(broadcast_var.clone(), value.clone());
}
}
context_outputs
}
/// Execute this node with the given context
pub fn execute(&self, context: &HashMap<String, GraphData>) -> HashMap<String, GraphData> {
let inputs = self.gather_inputs(context);
self.execute_with_inputs(&inputs)
}
/// Restrict cache-key inputs to the configured impl-var subset when present.
pub fn cache_key_inputs(
&self,
inputs: &HashMap<String, GraphData>,
) -> HashMap<String, GraphData> {
match &self.cache_key_inputs {
Some(keys) => keys
.iter()
.filter_map(|key| inputs.get(key).cloned().map(|value| (key.clone(), value)))
.collect(),
None => inputs.clone(),
}
}
/// Get display name for this node
pub fn display_name(&self) -> String {
self.label
.clone()
.unwrap_or_else(|| format!("Node {}", self.id))
}
}