hydro_lang/builder/
compiled.rs

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

hydro_lang/builder/
compiled.rs
