Skip to main content

weir/fixed_point_graph/
graph.rs

1#![allow(clippy::too_many_arguments)]
2
3use super::FixedPointAnalysisKind;
4
5/// Packed forward CSR graph buffers ready for Vyre dispatch.
6#[derive(Clone, Debug, PartialEq)]
7pub struct FixedPointForwardGraph {
8    pub(crate) kind: FixedPointAnalysisKind,
9    pub(crate) node_count: u32,
10    pub(crate) edge_count: u32,
11    pub(crate) pg_nodes_bytes: Vec<u8>,
12    pub(crate) edge_offsets_bytes: Vec<u8>,
13    pub(crate) edge_targets_bytes: Vec<u8>,
14    pub(crate) edge_kind_mask_bytes: Vec<u8>,
15    pub(crate) pg_node_tags_bytes: Vec<u8>,
16    pub(crate) stable_layout_hash: u64,
17    pub(crate) normalized_graph: crate::graph_layout::NormalizedCsrGraph,
18    pub(crate) normalization_scratch: crate::graph_layout::CsrGraphNormalizationScratch,
19}
20
21impl FixedPointForwardGraph {
22    /// Validate and pack invariant forward-CSR buffers once.
23    pub fn new(
24        stage: &str,
25        node_count: u32,
26        edge_offsets: &[u32],
27        edge_targets: &[u32],
28        edge_kind_mask: &[u32],
29    ) -> Result<Self, String> {
30        Self::new_for_kind(
31            FixedPointAnalysisKind::Generic,
32            stage,
33            node_count,
34            edge_offsets,
35            edge_targets,
36            edge_kind_mask,
37        )
38    }
39
40    /// Validate and pack invariant forward-CSR buffers for one analysis family.
41    pub fn new_for_kind(
42        kind: FixedPointAnalysisKind,
43        stage: &str,
44        node_count: u32,
45        edge_offsets: &[u32],
46        edge_targets: &[u32],
47        edge_kind_mask: &[u32],
48    ) -> Result<Self, String> {
49        let mut normalized_graph = crate::graph_layout::NormalizedCsrGraph::empty();
50        let mut normalization_scratch = crate::graph_layout::CsrGraphNormalizationScratch::new();
51        crate::graph_layout::CsrGraph::new(node_count, edge_offsets, edge_targets, edge_kind_mask)
52            .normalize_into(stage, &mut normalized_graph, &mut normalization_scratch)?;
53        Self::from_normalized(
54            kind,
55            stage,
56            node_count,
57            normalized_graph,
58            normalization_scratch,
59        )
60    }
61
62    /// Validate and pack invariant CSR buffers once after masking edge kinds.
63    pub fn new_with_edge_kind_mask(
64        stage: &str,
65        node_count: u32,
66        edge_offsets: &[u32],
67        edge_targets: &[u32],
68        edge_kind_mask: &[u32],
69        allowed_mask: u32,
70    ) -> Result<Self, String> {
71        Self::new_with_edge_kind_mask_for_kind(
72            FixedPointAnalysisKind::Generic,
73            stage,
74            node_count,
75            edge_offsets,
76            edge_targets,
77            edge_kind_mask,
78            allowed_mask,
79        )
80    }
81
82    /// Validate and pack masked CSR buffers for one analysis family.
83    pub fn new_with_edge_kind_mask_for_kind(
84        kind: FixedPointAnalysisKind,
85        stage: &str,
86        node_count: u32,
87        edge_offsets: &[u32],
88        edge_targets: &[u32],
89        edge_kind_mask: &[u32],
90        allowed_mask: u32,
91    ) -> Result<Self, String> {
92        let mut normalized_graph = crate::graph_layout::NormalizedCsrGraph::empty();
93        let mut normalization_scratch = crate::graph_layout::CsrGraphNormalizationScratch::new();
94        crate::graph_layout::CsrGraph::new(node_count, edge_offsets, edge_targets, edge_kind_mask)
95            .normalize_with_edge_kind_mask_into(
96                stage,
97                allowed_mask,
98                &mut normalized_graph,
99                &mut normalization_scratch,
100            )?;
101        Self::from_normalized(
102            kind,
103            stage,
104            node_count,
105            normalized_graph,
106            normalization_scratch,
107        )
108    }
109
110    fn from_normalized(
111        kind: FixedPointAnalysisKind,
112        stage: &str,
113        node_count: u32,
114        normalized_graph: crate::graph_layout::NormalizedCsrGraph,
115        normalization_scratch: crate::graph_layout::CsrGraphNormalizationScratch,
116    ) -> Result<Self, String> {
117        let mut graph = Self {
118            kind,
119            node_count: 0,
120            edge_count: 0,
121            pg_nodes_bytes: Vec::new(),
122            edge_offsets_bytes: Vec::new(),
123            edge_targets_bytes: Vec::new(),
124            edge_kind_mask_bytes: Vec::new(),
125            pg_node_tags_bytes: Vec::new(),
126            stable_layout_hash: 0,
127            normalized_graph,
128            normalization_scratch,
129        };
130        graph.replace_from_current_normalized(kind, stage, node_count)?;
131        Ok(graph)
132    }
133
134    /// Validate and repack this graph in place while retaining buffer capacity.
135    pub fn replace(
136        &mut self,
137        stage: &str,
138        node_count: u32,
139        edge_offsets: &[u32],
140        edge_targets: &[u32],
141        edge_kind_mask: &[u32],
142    ) -> Result<(), String> {
143        self.replace_for_kind(
144            self.kind,
145            stage,
146            node_count,
147            edge_offsets,
148            edge_targets,
149            edge_kind_mask,
150        )
151    }
152
153    /// Validate and repack this graph for the same analysis family.
154    pub fn replace_for_kind(
155        &mut self,
156        kind: FixedPointAnalysisKind,
157        stage: &str,
158        node_count: u32,
159        edge_offsets: &[u32],
160        edge_targets: &[u32],
161        edge_kind_mask: &[u32],
162    ) -> Result<(), String> {
163        self.require_kind(kind, stage)?;
164        crate::graph_layout::CsrGraph::new(node_count, edge_offsets, edge_targets, edge_kind_mask)
165            .normalize_into(
166                stage,
167                &mut self.normalized_graph,
168                &mut self.normalization_scratch,
169            )?;
170        self.replace_from_current_normalized(kind, stage, node_count)
171    }
172
173    /// Validate and repack this graph after masking edge kinds in place while
174    /// retaining buffer capacity.
175    pub fn replace_with_edge_kind_mask(
176        &mut self,
177        stage: &str,
178        node_count: u32,
179        edge_offsets: &[u32],
180        edge_targets: &[u32],
181        edge_kind_mask: &[u32],
182        allowed_mask: u32,
183    ) -> Result<(), String> {
184        self.replace_with_edge_kind_mask_for_kind(
185            self.kind,
186            stage,
187            node_count,
188            edge_offsets,
189            edge_targets,
190            edge_kind_mask,
191            allowed_mask,
192        )
193    }
194
195    /// Validate and repack masked CSR buffers for the same analysis family.
196    pub fn replace_with_edge_kind_mask_for_kind(
197        &mut self,
198        kind: FixedPointAnalysisKind,
199        stage: &str,
200        node_count: u32,
201        edge_offsets: &[u32],
202        edge_targets: &[u32],
203        edge_kind_mask: &[u32],
204        allowed_mask: u32,
205    ) -> Result<(), String> {
206        self.require_kind(kind, stage)?;
207        crate::graph_layout::CsrGraph::new(node_count, edge_offsets, edge_targets, edge_kind_mask)
208            .normalize_with_edge_kind_mask_into(
209                stage,
210                allowed_mask,
211                &mut self.normalized_graph,
212                &mut self.normalization_scratch,
213            )?;
214        self.replace_from_current_normalized(kind, stage, node_count)
215    }
216
217    fn replace_from_current_normalized(
218        &mut self,
219        kind: FixedPointAnalysisKind,
220        stage: &str,
221        node_count: u32,
222    ) -> Result<(), String> {
223        let normalized = &self.normalized_graph;
224        let nodes = usize::try_from(normalized.node_count()).map_err(|_| {
225            format!("{stage} node_count={node_count} cannot fit usize. Fix: shard the graph before dispatch.")
226        })?;
227        self.kind = kind;
228        self.node_count = normalized.node_count();
229        self.edge_count = normalized.edge_count();
230        self.stable_layout_hash = normalized.stable_layout_hash();
231        crate::dispatch_decode::pack_repeated_u32_into(0, nodes, &mut self.pg_nodes_bytes)?;
232        crate::dispatch_decode::try_pack_u32_into(
233            normalized.edge_offsets(),
234            &mut self.edge_offsets_bytes,
235            "fixed-point graph edge offsets bytes",
236        )?;
237        crate::dispatch_decode::try_pack_u32_into(
238            normalized.edge_targets(),
239            &mut self.edge_targets_bytes,
240            "fixed-point graph edge targets bytes",
241        )?;
242        crate::dispatch_decode::try_pack_u32_into(
243            normalized.edge_kind_mask(),
244            &mut self.edge_kind_mask_bytes,
245            "fixed-point graph edge kind mask bytes",
246        )?;
247        crate::dispatch_decode::pack_repeated_u32_into(0, nodes, &mut self.pg_node_tags_bytes)?;
248        Ok(())
249    }
250
251    /// Analysis family this prepared graph layout belongs to.
252    #[must_use]
253    pub fn kind(&self) -> FixedPointAnalysisKind {
254        self.kind
255    }
256
257    /// Verify this graph layout is consumed by the analysis family that
258    /// prepared it.
259    pub fn require_kind(
260        &self,
261        expected: FixedPointAnalysisKind,
262        consumer: &str,
263    ) -> Result<(), String> {
264        if self.kind == expected {
265            return Ok(());
266        }
267        Err(prepared_kind_mismatch(consumer, self.kind, expected))
268    }
269
270    /// Number of graph nodes in the prepared dispatch domain.
271    #[must_use]
272    pub fn node_count(&self) -> u32 {
273        self.node_count
274    }
275
276    /// Number of graph edges in the prepared dispatch domain.
277    #[must_use]
278    pub fn edge_count(&self) -> u32 {
279        self.edge_count
280    }
281
282    /// Total retained bytes across invariant packed graph buffers.
283    #[must_use]
284    pub fn retained_bytes(&self) -> usize {
285        self.pg_nodes_bytes.capacity()
286            + self.edge_offsets_bytes.capacity()
287            + self.edge_targets_bytes.capacity()
288            + self.edge_kind_mask_bytes.capacity()
289            + self.pg_node_tags_bytes.capacity()
290    }
291
292    /// Deterministic hash of the normalized graph layout.
293    ///
294    /// This is intentionally not `HashMap`'s randomized hasher. Plan and
295    /// resident-resource caches need a stable key for equivalent CSR layouts
296    /// across runs, devices, and processes.
297    #[must_use]
298    pub fn stable_layout_hash(&self) -> u64 {
299        self.stable_layout_hash
300    }
301}
302
303#[cold]
304fn prepared_kind_mismatch(
305    consumer: &str,
306    got: FixedPointAnalysisKind,
307    expected: FixedPointAnalysisKind,
308) -> String {
309    let mut scratch = crate::error_format::ErrorFormatScratch::default();
310    let _ = std::fmt::Write::write_fmt(
311        &mut scratch.buf,
312        format_args!(
313            "{consumer} received a {got:?} fixed-point graph, expected {expected:?}. Fix: prepare the graph with the matching Weir analysis constructor."
314        ),
315    );
316    scratch.finish()
317}