hydro_lang/builder/
compiled.rs

1use std::collections::BTreeMap;
2#[cfg(feature = "staged_macro")]
3use std::marker::PhantomData;
4
5use dfir_lang::graph::DfirGraph;
6#[cfg(feature = "staged_macro")]
7use dfir_rs::scheduled::graph::Dfir;
8#[cfg(feature = "staged_macro")]
9use proc_macro2::TokenStream;
10#[cfg(feature = "staged_macro")]
11use quote::quote;
12#[cfg(feature = "staged_macro")]
13use stageleft::QuotedWithContext;
14#[cfg(feature = "staged_macro")]
15use stageleft::runtime_support::{FreeVariableWithContext, QuoteTokens};
16
17use crate::Location;
18use crate::staging_util::Invariant;
19
20pub struct CompiledFlow<'a, ID> {
21    pub(super) dfir: BTreeMap<usize, DfirGraph>,
22    #[cfg(feature = "staged_macro")]
23    pub(super) extra_stmts: BTreeMap<usize, Vec<syn::Stmt>>,
24    pub(super) _phantom: Invariant<'a, ID>,
25}
26
27impl<'a, ID> CompiledFlow<'a, ID> {
28    pub fn dfir_for(&self, location: &impl Location<'a>) -> &DfirGraph {
29        self.dfir.get(&location.id().raw_id()).unwrap()
30    }
31
32    pub fn all_dfir(&self) -> &BTreeMap<usize, DfirGraph> {
33        &self.dfir
34    }
35}
36
37impl<'a> CompiledFlow<'a, usize> {
38    #[cfg(feature = "staged_macro")]
39    pub fn with_dynamic_id(
40        self,
41        id: impl QuotedWithContext<'a, usize, ()>,
42    ) -> CompiledFlowWithId<'a> {
43        let hydro_lang_crate = proc_macro_crate::crate_name("hydro_lang")
44            .expect("hydro_lang should be present in `Cargo.toml`");
45        let root = match hydro_lang_crate {
46            proc_macro_crate::FoundCrate::Itself => quote! { hydro_lang::runtime_support::dfir_rs },
47            proc_macro_crate::FoundCrate::Name(name) => {
48                let ident = syn::Ident::new(&name, proc_macro2::Span::call_site());
49                quote! { #ident::runtime_support::dfir_rs }
50            }
51        };
52
53        let mut conditioned_tokens = None;
54        for (subgraph_id, partitioned_graph) in self.dfir {
55            let mut diagnostics = Vec::new();
56            let tokens = partitioned_graph.as_code(&root, true, quote::quote!(), &mut diagnostics);
57            let my_extra_stmts = self
58                .extra_stmts
59                .get(&subgraph_id)
60                .cloned()
61                .unwrap_or_default();
62
63            if let Some(conditioned_tokens) = conditioned_tokens.as_mut() {
64                *conditioned_tokens = syn::parse_quote! {
65                    #conditioned_tokens else if __given_id == #subgraph_id {
66                        #(#my_extra_stmts)*
67                        #tokens
68                    }
69                };
70            } else {
71                conditioned_tokens = Some(syn::parse_quote! {
72                    if __given_id == #subgraph_id {
73                        #(#my_extra_stmts)*
74                        #tokens
75                    }
76                });
77            }
78        }
79
80        let conditioned_tokens: TokenStream = conditioned_tokens.unwrap();
81        let id = id.splice_untyped();
82        CompiledFlowWithId {
83            tokens: syn::parse_quote!({
84                let __given_id = #id;
85                #conditioned_tokens else {
86                    panic!("Invalid node id: {}", __given_id);
87                }
88            }),
89            _phantom: PhantomData,
90        }
91    }
92}
93
94#[cfg(feature = "staged_macro")]
95impl<'a, Ctx> QuotedWithContext<'a, Dfir<'a>, Ctx> for CompiledFlow<'a, ()> {}
96
97#[cfg(feature = "staged_macro")]
98impl<'a, Ctx> FreeVariableWithContext<Ctx> for CompiledFlow<'a, ()> {
99    type O = Dfir<'a>;
100
101    fn to_tokens(mut self, _ctx: &Ctx) -> QuoteTokens {
102        let hydro_lang_crate = proc_macro_crate::crate_name("hydro_lang")
103            .expect("hydro_lang should be present in `Cargo.toml`");
104        let root = match hydro_lang_crate {
105            proc_macro_crate::FoundCrate::Itself => quote! { hydro_lang::runtime_support::dfir_rs },
106            proc_macro_crate::FoundCrate::Name(name) => {
107                let ident = syn::Ident::new(&name, proc_macro2::Span::call_site());
108                quote! { #ident::runtime_support::dfir_rs }
109            }
110        };
111
112        if self.dfir.len() != 1 {
113            panic!("Expected exactly one subgraph in the DFIR.");
114        }
115
116        let partitioned_graph = self.dfir.remove(&0).unwrap();
117
118        let mut diagnostics = Vec::new();
119        let tokens = partitioned_graph.as_code(&root, true, quote::quote!(), &mut diagnostics);
120
121        QuoteTokens {
122            prelude: None,
123            expr: Some(tokens),
124        }
125    }
126}
127
128#[cfg(feature = "staged_macro")]
129pub struct CompiledFlowWithId<'a> {
130    tokens: TokenStream,
131    _phantom: Invariant<'a>,
132}
133
134#[cfg(feature = "staged_macro")]
135impl<'a, Ctx> QuotedWithContext<'a, Dfir<'a>, Ctx> for CompiledFlowWithId<'a> {}
136
137#[cfg(feature = "staged_macro")]
138impl<'a, Ctx> FreeVariableWithContext<Ctx> for CompiledFlowWithId<'a> {
139    type O = Dfir<'a>;
140
141    fn to_tokens(self, _ctx: &Ctx) -> QuoteTokens {
142        QuoteTokens {
143            prelude: None,
144            expr: Some(self.tokens),
145        }
146    }
147}
hydro_lang/builder/compiled.rs

hydro_lang/builder/
compiled.rs
