diff --git a/dfir_lang/src/graph/hydroflow_graph.rs b/dfir_lang/src/graph/hydroflow_graph.rs index f61ab3e6a9c..9e8792e8d00 100644 --- a/dfir_lang/src/graph/hydroflow_graph.rs +++ b/dfir_lang/src/graph/hydroflow_graph.rs @@ -2,13 +2,13 @@ extern crate proc_macro; -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::{BTreeMap, BTreeSet, VecDeque}; use std::fmt::Debug; use std::iter::FusedIterator; use itertools::Itertools; use proc_macro2::{Ident, Literal, Span, TokenStream}; -use quote::{format_ident, quote, quote_spanned, ToTokens}; +use quote::{format_ident, quote, quote_spanned, ToTokens, TokenStreamExt}; use serde::{Deserialize, Serialize}; use slotmap::{Key, SecondaryMap, SlotMap, SparseSecondaryMap}; use syn::spanned::Spanned; @@ -52,10 +52,13 @@ pub struct DfirGraph { /// Which loop a node belongs to (or none for top-level). node_loops: SecondaryMap, + /// Which nodes belong to each loop. loop_nodes: SlotMap>, - /// For the key loop, what is its parent (`None` for top-level). + /// For the loop, what is its parent (`None` for top-level). loop_parent: SparseSecondaryMap, - /// For the key loop, what are its child loops. + /// What loops are at the root. + root_loops: Vec, + /// For the loop, what are its child loops. loop_children: SecondaryMap>, /// Which subgraph each node belongs to. @@ -802,6 +805,32 @@ impl DfirGraph { subgraph_handoffs } + /// Generate a deterministic `Ident` for the given loop ID. + fn loop_as_ident(loop_id: GraphLoopId) -> Ident { + Ident::new(&format!("loop_{:?}", loop_id.data()), Span::call_site()) + } + + /// Code for adding all nested loops. + fn helper_loop_code(&self, hf: &Ident) -> TokenStream { + // Breadth-first iteration from outermost (root) loops to deepest nested loops. + let mut out = TokenStream::new(); + let mut queue = VecDeque::from_iter(self.root_loops.iter().copied()); + while let Some(loop_id) = queue.pop_front() { + let parent_code = if let Some(&parent_id) = self.loop_parent.get(loop_id) { + let parent_ident = Self::loop_as_ident(parent_id); + quote! { Some(#parent_ident) } + } else { + quote! { None } + }; + let loop_name = Self::loop_as_ident(loop_id); + out.append_all(quote! { + let #loop_name = #hf.add_loop(#parent_code); + }); + queue.extend(self.loop_children.get(loop_id).into_iter().flatten()); + } + out + } + /// Emit this `HydroflowGraph` as runnable Rust source code tokens. pub fn as_code( &self, @@ -813,7 +842,8 @@ impl DfirGraph { let hf = Ident::new(HYDROFLOW, Span::call_site()); let context = Ident::new(CONTEXT, Span::call_site()); - let handoffs = self + // Code for adding handoffs. + let handoff_code = self .nodes .iter() .filter_map(|(node_id, node)| match node { @@ -1161,6 +1191,14 @@ impl DfirGraph { self.subgraph_stratum.get(subgraph_id).cloned().unwrap_or(0), ); let laziness = self.subgraph_laziness(subgraph_id); + + let loop_id_code = + self.node_loop(subgraph_nodes[0]) + .map_or(quote! { None }, |loop_id| { + let loop_ident = Self::loop_as_ident(loop_id); + quote! { Some(#loop_ident) } + }); + subgraphs.push(quote! { #hf.add_subgraph_stratified( #hoff_name, @@ -1168,7 +1206,7 @@ impl DfirGraph { var_expr!( #( #recv_ports ),* ), var_expr!( #( #send_ports ),* ), #laziness, - None, // `LoopId` + #loop_id_code, move |#context, var_args!( #( #recv_ports ),* ), var_args!( #( #send_ports ),* )| { #( #recv_port_code )* #( #send_port_code )* @@ -1180,12 +1218,15 @@ impl DfirGraph { } } + let loop_code = self.helper_loop_code(&hf); + // These two are quoted separately here because iterators are lazily evaluated, so this // forces them to do their work. This work includes populating some data, namely // `diagonstics`, which we need to determine if it compilation was actually successful. // -Mingwei let code = quote! { - #( #handoffs )* + #( #handoff_code )* + #loop_code #( #op_prologue_code )* #( #subgraphs )* }; @@ -1557,6 +1598,8 @@ impl DfirGraph { .get_mut(parent_loop) .unwrap() .push(loop_id); + } else { + self.root_loops.push(loop_id); } loop_id } diff --git a/dfir_rs/src/scheduled/graph.rs b/dfir_rs/src/scheduled/graph.rs index 2152710bbb6..a7283457951 100644 --- a/dfir_rs/src/scheduled/graph.rs +++ b/dfir_rs/src/scheduled/graph.rs @@ -23,7 +23,7 @@ use super::state::StateHandle; use super::subgraph::Subgraph; use super::{HandoffId, HandoffTag, LoopId, LoopTag, SubgraphId, SubgraphTag}; use crate::scheduled::ticks::{TickDuration, TickInstant}; -use crate::util::slot_vec::SlotVec; +use crate::util::slot_vec::{SecondarySlotVec, SlotVec}; use crate::Never; /// A DFIR graph. Owns, schedules, and runs the compiled subgraphs. @@ -31,8 +31,11 @@ use crate::Never; pub struct Dfir<'a> { pub(super) subgraphs: SlotVec>, pub(super) context: Context, + + // Depth of loop (zero for top-level). + loop_depth: SlotVec, // Map from `LoopId` to parent `LoopId` (or `None` for top-level). - pub(super) loop_parent: SlotVec>, + loop_parent: SecondarySlotVec>, handoffs: SlotVec, @@ -790,8 +793,13 @@ impl<'a> Dfir<'a> { /// Adds a new loop with the given parent (or `None` for top-level). Returns a loop ID which /// is used in [`Self::add_subgraph_stratified`] or for nested loops. + /// + /// TODO(mingwei): add loop names to ensure traceability while debugging? pub fn add_loop(&mut self, parent: Option) -> LoopId { - self.loop_parent.insert(parent) + let depth = parent.map_or(0, |p| self.loop_depth[p] + 1); + let loop_id = self.loop_depth.insert(depth); + self.loop_parent.insert(loop_id, parent); + loop_id } }