add loops to codegen

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<GraphNodeId, GraphLoopId>,
+    /// Which nodes belong to each loop.
     loop_nodes: SlotMap<GraphLoopId, Vec<GraphNodeId>>,
-    /// 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<GraphLoopId, GraphLoopId>,
-    /// For the key loop, what are its child loops.
+    /// What loops are at the root.
+    root_loops: Vec<GraphLoopId>,
+    /// For the loop, what are its child loops.
     loop_children: SecondaryMap<GraphLoopId, Vec<GraphLoopId>>,
 
     /// 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,14 +1191,22 @@ 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,
                         #stratum,
                         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
     }

dfir_rs/src/scheduled/graph.rs

@@ -23,16 +23,19 @@ 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.
 #[derive(Default)]
 pub struct Dfir<'a> {
     pub(super) subgraphs: SlotVec<SubgraphTag, SubgraphData<'a>>,
     pub(super) context: Context,
+
+    // Depth of loop (zero for top-level).
+    loop_depth: SlotVec<LoopTag, usize>,
     // Map from `LoopId` to parent `LoopId` (or `None` for top-level).
-    pub(super) loop_parent: SlotVec<LoopTag, Option<LoopId>>,
+    loop_parent: SecondarySlotVec<LoopTag, Option<LoopId>>,
 
     handoffs: SlotVec<HandoffTag, HandoffData>,
 
@@ -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>) -> 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
     }
 }