dfir_lang/graph/
mod.rs

1//! Graph representation stages for DFIR graphs.
2
3use std::borrow::Cow;
4use std::hash::Hash;
5
6use proc_macro2::{Ident, Span, TokenStream};
7use quote::ToTokens;
8use serde::{Deserialize, Serialize};
9use syn::punctuated::Punctuated;
10use syn::spanned::Spanned;
11use syn::{Expr, ExprPath, GenericArgument, Token, Type};
12
13use self::ops::{OperatorConstraints, Persistence};
14use crate::diagnostic::{Diagnostic, Diagnostics, Level};
15use crate::parse::{DfirCode, IndexInt, Operator, PortIndex, Ported, SingletonRef};
16use crate::pretty_span::PrettySpan;
17
18mod di_mul_graph;
19mod eliminate_extra_unions_tees;
20mod flat_graph_builder;
21mod flat_to_partitioned;
22mod graph_write;
23mod meta_graph;
24mod meta_graph_debugging;
25
26use std::fmt::Display;
27
28pub use di_mul_graph::DiMulGraph;
29pub use eliminate_extra_unions_tees::eliminate_extra_unions_tees;
30pub use flat_graph_builder::{FlatGraphBuilder, FlatGraphBuilderOutput};
31pub use flat_to_partitioned::partition_graph;
32pub use meta_graph::{DfirGraph, WriteConfig, WriteGraphType};
33
34pub use crate::graph_ids::{GraphEdgeId, GraphLoopId, GraphNodeId, GraphSubgraphId};
35
36pub mod graph_algorithms;
37pub mod ops;
38
39impl GraphSubgraphId {
40    /// Generate a deterministic `Ident` for the given subgraph ID.
41    pub fn as_ident(self, span: Span) -> Ident {
42        use slotmap::Key;
43        Ident::new(&format!("sgid_{:?}", self.data()), span)
44    }
45}
46
47impl GraphLoopId {
48    /// Generate a deterministic `Ident` for the given loop ID.
49    pub fn as_ident(self, span: Span) -> Ident {
50        use slotmap::Key;
51        Ident::new(&format!("loop_{:?}", self.data()), span)
52    }
53}
54
55/// Context identifier as a string.
56const CONTEXT: &str = "context";
57/// Runnable DFIR graph object identifier as a string.
58const GRAPH: &str = "df";
59
60const HANDOFF_NODE_STR: &str = "handoff";
61const SINGLETON_SLOT_NODE_STR: &str = "singleton";
62const MODULE_BOUNDARY_NODE_STR: &str = "module_boundary";
63
64mod serde_syn {
65    use serde::{Deserialize, Deserializer, Serializer};
66
67    pub fn serialize<S, T>(value: &T, serializer: S) -> Result<S::Ok, S::Error>
68    where
69        S: Serializer,
70        T: quote::ToTokens,
71    {
72        serializer.serialize_str(&value.to_token_stream().to_string())
73    }
74
75    pub fn deserialize<'de, D, T>(deserializer: D) -> Result<T, D::Error>
76    where
77        D: Deserializer<'de>,
78        T: syn::parse::Parse,
79    {
80        let s = String::deserialize(deserializer)?;
81        syn::parse_str(&s).map_err(<D::Error as serde::de::Error>::custom)
82    }
83}
84
85/// A variable name assigned to a pipeline in DFIR syntax.
86///
87/// Fundamentally a serializable/deserializable wrapper around [`syn::Ident`].
88#[derive(Clone, Debug, Serialize, Deserialize, PartialOrd, Ord, PartialEq, Eq, Hash)]
89pub struct Varname(#[serde(with = "serde_syn")] pub Ident);
90
91/// The kind of inter-subgraph handoff.
92#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
93pub enum HandoffKind {
94    /// A `Vec<T>` buffer for streams (zero or more items).
95    Vec,
96    /// An `Option<T>` slot for singletons (exactly one item expected).
97    /// `#varname` gives `&T` (panics if empty).
98    Singleton,
99    /// An `Option<T>` slot for optionals (zero or one item).
100    /// `#varname` gives `&Option<T>`.
101    Optional,
102}
103
104/// A node, corresponding to an operator or a handoff.
105#[derive(Clone, Serialize, Deserialize)]
106pub enum GraphNode {
107    /// An operator.
108    Operator(#[serde(with = "serde_syn")] Operator),
109    /// An inter-subgraph handoff point for buffering data between subgraphs.
110    Handoff {
111        /// What kind of storage this handoff uses.
112        kind: HandoffKind,
113        /// The span of the input into the handoff.
114        #[serde(skip, default = "Span::call_site")]
115        src_span: Span,
116        /// The span of the output out of the handoff.
117        #[serde(skip, default = "Span::call_site")]
118        dst_span: Span,
119    },
120
121    /// Module Boundary, used for importing modules. Only exists prior to partitioning.
122    ModuleBoundary {
123        /// If this module is an input or output boundary.
124        input: bool,
125
126        /// The span of the import!() expression that imported this module.
127        /// The value of this span when the ModuleBoundary node is still inside the module is Span::call_site()
128        /// TODO: This could one day reference into the module file itself?
129        #[serde(skip, default = "Span::call_site")]
130        import_expr: Span,
131    },
132}
133impl GraphNode {
134    /// Return the node as a human-readable string.
135    pub fn to_pretty_string(&self) -> Cow<'static, str> {
136        match self {
137            GraphNode::Operator(op) => op.to_pretty_string().into(),
138            GraphNode::Handoff {
139                kind: HandoffKind::Vec,
140                ..
141            } => HANDOFF_NODE_STR.into(),
142            GraphNode::Handoff {
143                kind: HandoffKind::Singleton | HandoffKind::Optional,
144                ..
145            } => SINGLETON_SLOT_NODE_STR.into(),
146            GraphNode::ModuleBoundary { .. } => MODULE_BOUNDARY_NODE_STR.into(),
147        }
148    }
149
150    /// Return the name of the node as a string, excluding parenthesis and op source code.
151    pub fn to_name_string(&self) -> Cow<'static, str> {
152        match self {
153            GraphNode::Operator(op) => op.name_string().into(),
154            GraphNode::Handoff {
155                kind: HandoffKind::Vec,
156                ..
157            } => HANDOFF_NODE_STR.into(),
158            GraphNode::Handoff {
159                kind: HandoffKind::Singleton | HandoffKind::Optional,
160                ..
161            } => SINGLETON_SLOT_NODE_STR.into(),
162            GraphNode::ModuleBoundary { .. } => MODULE_BOUNDARY_NODE_STR.into(),
163        }
164    }
165
166    /// Return the source code span of the node.
167    pub fn span(&self) -> Span {
168        match self {
169            Self::Operator(op) => op.span(),
170            &Self::Handoff {
171                src_span, dst_span, ..
172            } => src_span.join(dst_span).unwrap_or(src_span),
173            Self::ModuleBoundary { import_expr, .. } => *import_expr,
174        }
175    }
176}
177impl std::fmt::Debug for GraphNode {
178    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
179        match self {
180            Self::Operator(operator) => {
181                write!(f, "Node::Operator({} span)", PrettySpan(operator.span()))
182            }
183            Self::Handoff { kind, .. } => write!(f, "Node::Handoff({kind:?})"),
184            Self::ModuleBoundary { input, .. } => {
185                write!(f, "Node::ModuleBoundary{{input: {}}}", input)
186            }
187        }
188    }
189}
190
191/// Meta-data relating to operators which may be useful throughout the compilation process.
192///
193/// This data can be generated from the graph, but it is useful to have it readily available
194/// pre-computed as many algorithms use the same info. Stuff like port names, arguments, and the
195/// [`OperatorConstraints`] for the operator.
196///
197/// Because it is derived from the graph itself, there can be "cache invalidation"-esque issues
198/// if this data is not kept in sync with the graph.
199#[derive(Clone, Debug)]
200pub struct OperatorInstance {
201    /// Name of the operator (will match [`OperatorConstraints::name`]).
202    pub op_constraints: &'static OperatorConstraints,
203    /// Port values used as this operator's input.
204    pub input_ports: Vec<PortIndexValue>,
205    /// Port values used as this operator's output.
206    pub output_ports: Vec<PortIndexValue>,
207    /// Singleton references within the operator arguments.
208    pub singletons_referenced: Vec<SingletonRef>,
209
210    /// Generic arguments.
211    pub generics: OpInstGenerics,
212    /// Arguments provided by the user into the operator as arguments.
213    /// I.e. the `a, b, c` in `-> my_op(a, b, c) -> `.
214    ///
215    /// These arguments do not include singleton postprocessing codegen. Instead use
216    /// [`ops::WriteContextArgs::arguments`].
217    pub arguments_pre: Punctuated<Expr, Token![,]>,
218    /// Unparsed arguments, for singleton parsing.
219    pub arguments_raw: TokenStream,
220}
221
222/// Operator generic arguments, split into specific categories.
223#[derive(Clone, Debug)]
224pub struct OpInstGenerics {
225    /// Operator generic (type or lifetime) arguments.
226    pub generic_args: Option<Punctuated<GenericArgument, Token![,]>>,
227    /// Lifetime persistence arguments. Corresponds to a prefix of [`Self::generic_args`].
228    pub persistence_args: Vec<Persistence>,
229    /// Type persistence arguments. Corersponds to a (suffix) of [`Self::generic_args`].
230    pub type_args: Vec<Type>,
231}
232
233impl OpInstGenerics {
234    /// Helper to join a sequence of spans into a single span, if possible.
235    ///
236    /// Returns `None` if there are no spans or if any `Span::join` call fails
237    /// (for example, when spans are not contiguous).
238    fn join_spans<I>(mut spans: I) -> Option<Span>
239    where
240        I: Iterator<Item = Span>,
241    {
242        let mut span = spans.next()?;
243        for s in spans {
244            span = span.join(s)?;
245        }
246        Some(span)
247    }
248
249    /// Returns a [`Span`] containing all persistence (lifetime) args if possible.
250    pub fn persistence_args_span(&self) -> Option<Span> {
251        self.generic_args.as_ref().and_then(|args| {
252            Self::join_spans(
253                args.iter()
254                    .filter(|a| matches!(a, GenericArgument::Lifetime(_)))
255                    .map(|a| a.span()),
256            )
257        })
258    }
259
260    /// Returns a [`Span`] containing all type args if possible.
261    pub fn type_args_span(&self) -> Option<Span> {
262        self.generic_args.as_ref().and_then(|args| {
263            Self::join_spans(
264                args.iter()
265                    .filter(|a| matches!(a, GenericArgument::Type(_)))
266                    .map(|a| a.span()),
267            )
268        })
269    }
270}
271
272/// Gets the generic arguments for the operator.
273///
274/// This helper method is useful due to the special handling of persistence lifetimes (`'static`,
275/// `'tick`) which must come before other generic parameters.
276pub fn get_operator_generics(diagnostics: &mut Diagnostics, operator: &Operator) -> OpInstGenerics {
277    // Generic arguments.
278    let generic_args = operator.type_arguments().cloned();
279    let persistence_args = generic_args.iter().flatten().map_while(|generic_arg| match generic_arg {
280            GenericArgument::Lifetime(lifetime) => {
281                match &*lifetime.ident.to_string() {
282                    "none" => Some(Persistence::None),
283                    "loop" => Some(Persistence::Loop),
284                    "tick" => Some(Persistence::Tick),
285                    "static" => Some(Persistence::Static),
286                    _ => {
287                        diagnostics.push(Diagnostic::spanned(
288                            generic_arg.span(),
289                            Level::Error,
290                            format!("Unknown lifetime generic argument `'{}`, expected `'none`, `'loop`, `'tick`, or `'static`.", lifetime.ident),
291                        ));
292                        // TODO(mingwei): should really keep going and not short circuit?
293                        None
294                    }
295                }
296            },
297            _ => None,
298        }).collect::<Vec<_>>();
299    let type_args = generic_args
300        .iter()
301        .flatten()
302        .skip(persistence_args.len())
303        .map_while(|generic_arg| match generic_arg {
304            GenericArgument::Type(typ) => Some(typ),
305            _ => None,
306        })
307        .cloned()
308        .collect::<Vec<_>>();
309
310    OpInstGenerics {
311        generic_args,
312        persistence_args,
313        type_args,
314    }
315}
316
317/// Push, Pull, Comp, or Hoff polarity.
318#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
319pub enum Color {
320    /// Pull (green)
321    Pull,
322    /// Push (blue)
323    Push,
324    /// Computation (yellow)
325    Comp,
326    /// Handoff (grey) -- not a color for operators, inserted between subgraphs.
327    Hoff,
328}
329
330/// Helper struct for [`PortIndex`] which keeps span information for elided ports.
331#[derive(Clone, Debug, Serialize, Deserialize)]
332pub enum PortIndexValue {
333    /// An integer value: `[0]`, `[1]`, etc. Can be negative although we don't use that (2023-08-16).
334    Int(#[serde(with = "serde_syn")] IndexInt),
335    /// A name or path. `[pos]`, `[neg]`, etc. Can use `::` separators but we don't use that (2023-08-16).
336    Path(#[serde(with = "serde_syn")] ExprPath),
337    /// Elided, unspecified port. We have this variant, rather than wrapping in `Option`, in order
338    /// to preserve the `Span` information.
339    Elided(#[serde(skip)] Option<Span>),
340}
341impl PortIndexValue {
342    /// For a [`Ported`] value like `[port_in]name[port_out]`, get the `port_in` and `port_out` as
343    /// [`PortIndexValue`]s.
344    pub fn from_ported<Inner>(ported: Ported<Inner>) -> (Self, Inner, Self)
345    where
346        Inner: Spanned,
347    {
348        let ported_span = Some(ported.inner.span());
349        let port_inn = ported
350            .inn
351            .map(|idx| idx.index.into())
352            .unwrap_or_else(|| Self::Elided(ported_span));
353        let inner = ported.inner;
354        let port_out = ported
355            .out
356            .map(|idx| idx.index.into())
357            .unwrap_or_else(|| Self::Elided(ported_span));
358        (port_inn, inner, port_out)
359    }
360
361    /// Returns `true` if `self` is not [`PortIndexValue::Elided`].
362    pub fn is_specified(&self) -> bool {
363        !matches!(self, Self::Elided(_))
364    }
365
366    /// Returns whichever of the two ports are specified.
367    /// If both are [`Self::Elided`], returns [`Self::Elided`].
368    /// If both are specified, returns `Err(self)`.
369    #[allow(clippy::allow_attributes, reason = "Only triggered on nightly.")]
370    #[allow(
371        clippy::result_large_err,
372        reason = "variants are same size, error isn't to be propagated."
373    )]
374    pub fn combine(self, other: Self) -> Result<Self, Self> {
375        match (self.is_specified(), other.is_specified()) {
376            (false, _other) => Ok(other),
377            (true, false) => Ok(self),
378            (true, true) => Err(self),
379        }
380    }
381
382    /// Formats self as a human-readable string for error messages.
383    pub fn as_error_message_string(&self) -> String {
384        match self {
385            PortIndexValue::Int(n) => format!("`{}`", n.value),
386            PortIndexValue::Path(path) => format!("`{}`", path.to_token_stream()),
387            PortIndexValue::Elided(_) => "<elided>".to_owned(),
388        }
389    }
390
391    /// Returns the span of this port value.
392    pub fn span(&self) -> Span {
393        match self {
394            PortIndexValue::Int(x) => x.span(),
395            PortIndexValue::Path(x) => x.span(),
396            PortIndexValue::Elided(span) => span.unwrap_or_else(Span::call_site),
397        }
398    }
399}
400impl From<PortIndex> for PortIndexValue {
401    fn from(value: PortIndex) -> Self {
402        match value {
403            PortIndex::Int(x) => Self::Int(x),
404            PortIndex::Path(x) => Self::Path(x),
405        }
406    }
407}
408impl PartialEq for PortIndexValue {
409    fn eq(&self, other: &Self) -> bool {
410        match (self, other) {
411            (Self::Int(l0), Self::Int(r0)) => l0 == r0,
412            (Self::Path(l0), Self::Path(r0)) => l0 == r0,
413            (Self::Elided(_), Self::Elided(_)) => true,
414            _else => false,
415        }
416    }
417}
418impl Eq for PortIndexValue {}
419impl PartialOrd for PortIndexValue {
420    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
421        Some(self.cmp(other))
422    }
423}
424impl Ord for PortIndexValue {
425    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
426        match (self, other) {
427            (Self::Int(s), Self::Int(o)) => s.cmp(o),
428            (Self::Path(s), Self::Path(o)) => s
429                .to_token_stream()
430                .to_string()
431                .cmp(&o.to_token_stream().to_string()),
432            (Self::Elided(_), Self::Elided(_)) => std::cmp::Ordering::Equal,
433            (Self::Int(_), Self::Path(_)) => std::cmp::Ordering::Less,
434            (Self::Path(_), Self::Int(_)) => std::cmp::Ordering::Greater,
435            (_, Self::Elided(_)) => std::cmp::Ordering::Less,
436            (Self::Elided(_), _) => std::cmp::Ordering::Greater,
437        }
438    }
439}
440
441impl Display for PortIndexValue {
442    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
443        match self {
444            PortIndexValue::Int(x) => write!(f, "{}", x.to_token_stream()),
445            PortIndexValue::Path(x) => write!(f, "{}", x.to_token_stream()),
446            PortIndexValue::Elided(_) => write!(f, "[]"),
447        }
448    }
449}
450
451/// Output of [`build_dfir_code`].
452pub struct BuildDfirCodeOutput {
453    /// The now-partitioned graph.
454    pub partitioned_graph: DfirGraph,
455    /// The Rust source code tokens for the DFIR.
456    pub code: TokenStream,
457    /// Any (non-error) diagnostics emitted.
458    pub diagnostics: Diagnostics,
459}
460
461/// Compiles a [`DfirCode`] AST into inline source code that runs the dataflow.
462pub fn build_dfir_code(
463    dfir_code: DfirCode,
464    root: &TokenStream,
465) -> Result<BuildDfirCodeOutput, Diagnostics> {
466    let flat_graph_builder = FlatGraphBuilder::from_dfir(dfir_code);
467
468    let FlatGraphBuilderOutput {
469        mut flat_graph,
470        uses,
471        mut diagnostics,
472    } = flat_graph_builder.build()?;
473
474    let () = match flat_graph.merge_modules() {
475        Ok(()) => (),
476        Err(d) => {
477            diagnostics.push(d);
478            return Err(diagnostics);
479        }
480    };
481
482    eliminate_extra_unions_tees(&mut flat_graph);
483
484    // Detect adjacent handoffs (e.g. `handoff() -> handoff()` or `singleton() -> singleton()`),
485    // which can arise after unary tee/union elimination.
486    for (edge_id, (src, dst)) in flat_graph.edges() {
487        let _ = edge_id;
488        if matches!(flat_graph.node(src), GraphNode::Handoff { .. })
489            && matches!(flat_graph.node(dst), GraphNode::Handoff { .. })
490        {
491            let span = flat_graph.node(dst).span();
492            diagnostics.push(Diagnostic::spanned(
493                span,
494                Level::Error,
495                "Adjacent handoff/singleton operators are not allowed. \
496                 Remove one or insert an operator between them.",
497            ));
498        }
499    }
500
501    if diagnostics.has_error() {
502        return Err(diagnostics);
503    }
504
505    let partitioned_graph = match partition_graph(flat_graph) {
506        Ok(partitioned_graph) => partitioned_graph,
507        Err(d) => {
508            diagnostics.push(d);
509            return Err(diagnostics);
510        }
511    };
512
513    let code =
514        partitioned_graph.as_code(root, true, quote::quote! { #( #uses )* }, &mut diagnostics)?;
515
516    Ok(BuildDfirCodeOutput {
517        partitioned_graph,
518        code,
519        diagnostics,
520    })
521}
522
523/// Changes all of token's spans to `span`, recursing into groups.
524fn change_spans(tokens: TokenStream, span: Span) -> TokenStream {
525    use proc_macro2::{Group, TokenTree};
526    tokens
527        .into_iter()
528        .map(|token| match token {
529            TokenTree::Group(mut group) => {
530                group.set_span(span);
531                TokenTree::Group(Group::new(
532                    group.delimiter(),
533                    change_spans(group.stream(), span),
534                ))
535            }
536            TokenTree::Ident(mut ident) => {
537                ident.set_span(span.resolved_at(ident.span()));
538                TokenTree::Ident(ident)
539            }
540            TokenTree::Punct(mut punct) => {
541                punct.set_span(span);
542                TokenTree::Punct(punct)
543            }
544            TokenTree::Literal(mut literal) => {
545                literal.set_span(span);
546                TokenTree::Literal(literal)
547            }
548        })
549        .collect()
550}
dfir_lang/graph/mod.rs

dfir_lang/graph/
mod.rs
