entrenar/monitor/inference/provenance/
graph.rs1use super::edge::ProvenanceEdge;
4use super::node::{NodeId, ProvenanceNode};
5use crate::monitor::inference::path::DecisionPath;
6use crate::monitor::inference::trace::DecisionTrace;
7use std::collections::HashMap;
8
9pub struct ProvenanceGraph {
16 nodes: HashMap<NodeId, ProvenanceNode>,
18 edges: Vec<ProvenanceEdge>,
20 forward: HashMap<NodeId, Vec<usize>>,
22 backward: HashMap<NodeId, Vec<usize>>,
24 next_id: NodeId,
26}
27
28impl ProvenanceGraph {
29 pub fn new() -> Self {
31 Self {
32 nodes: HashMap::new(),
33 edges: Vec::new(),
34 forward: HashMap::new(),
35 backward: HashMap::new(),
36 next_id: 0,
37 }
38 }
39
40 pub fn add_node(&mut self, node: ProvenanceNode) -> NodeId {
42 let id = self.next_id;
43 self.next_id += 1;
44 self.nodes.insert(id, node);
45 id
46 }
47
48 pub fn add_edge(&mut self, edge: ProvenanceEdge) {
50 let edge_idx = self.edges.len();
51 self.forward.entry(edge.from).or_default().push(edge_idx);
52 self.backward.entry(edge.to).or_default().push(edge_idx);
53 self.edges.push(edge);
54 }
55
56 pub fn get_node(&self, id: NodeId) -> Option<&ProvenanceNode> {
58 self.nodes.get(&id)
59 }
60
61 pub fn nodes(&self) -> &HashMap<NodeId, ProvenanceNode> {
63 &self.nodes
64 }
65
66 pub fn edges(&self) -> &[ProvenanceEdge] {
68 &self.edges
69 }
70
71 pub fn incoming_edges(&self, id: NodeId) -> Vec<&ProvenanceEdge> {
73 self.backward
74 .get(&id)
75 .map(|indices| indices.iter().map(|&i| &self.edges[i]).collect())
76 .unwrap_or_default()
77 }
78
79 pub fn outgoing_edges(&self, id: NodeId) -> Vec<&ProvenanceEdge> {
81 self.forward
82 .get(&id)
83 .map(|indices| indices.iter().map(|&i| &self.edges[i]).collect())
84 .unwrap_or_default()
85 }
86
87 pub fn predecessors(&self, id: NodeId) -> Vec<NodeId> {
89 self.incoming_edges(id).into_iter().map(|e| e.from).collect()
90 }
91
92 pub fn successors(&self, id: NodeId) -> Vec<NodeId> {
94 self.outgoing_edges(id).into_iter().map(|e| e.to).collect()
95 }
96
97 pub fn node_count(&self) -> usize {
99 self.nodes.len()
100 }
101
102 pub fn edge_count(&self) -> usize {
104 self.edges.len()
105 }
106
107 pub fn add_inference<P: DecisionPath>(
109 &mut self,
110 trace: &DecisionTrace<P>,
111 model_id: &str,
112 model_version: &str,
113 ) -> NodeId {
114 self.add_node(ProvenanceNode::Inference {
115 model_id: model_id.to_string(),
116 model_version: model_version.to_string(),
117 confidence: trace.confidence(),
118 output: trace.output,
119 })
120 }
121}
122
123impl Default for ProvenanceGraph {
124 fn default() -> Self {
125 Self::new()
126 }
127}