separate adding graph edges from recursive visit

Author: Jesus Lapastora Núñez

engine/baml-lib/ast/src/ast/baml_vis/graph.rs

@@ -224,13 +224,41 @@ impl<'index, 'pre> GraphBuilder<'index, 'pre> {
             self.build_scope_sequence(scope, None);
         }
 
+        self.add_scope_edges();
+
         if self.cfg.show_call_nodes {
             self.add_header_calls();
         }
 
         (self.graph, self.span_map)
     }
 
+    /// Links the items in each scope in source order, showing execution order.
+    fn add_scope_edges(&mut self) {
+        for scope in self.index.scopes() {
+            // TODO: I think we can remove collect()?
+            let items: Vec<Hid> = self
+                .index
+                .headers_in_scope_iter(scope)
+                .filter(|h| !self.has_md_parent.contains(&h.hid))
+                .map(|h| h.hid)
+                .collect();
+
+            for (hid, prev_hid) in items[1..].iter().zip(&items) {
+                let entry = self.header_entry[hid];
+                let prev_exits = &self.header_exits[prev_hid];
+                // NOTE: `extend` on a loop. TL;DR: each exits vec is pretty small.
+                // More thorough explanation at the end of `build_header`.
+                self.graph.edges.extend(
+                    prev_exits
+                        .iter()
+                        .copied()
+                        .map(|e| Edge { from: e, to: entry }),
+                );
+            }
+        }
+    }
+
     fn build_scope_sequence(&mut self, scope: ScopeId, parent_cluster: Option<ClusterId>) {
         let items: Vec<Hid> = self
             .index
@@ -243,20 +271,6 @@ impl<'index, 'pre> GraphBuilder<'index, 'pre> {
         for &hid in &items {
             self.build_header(hid, parent_cluster);
         }
-
-        // unordered: add edges from scope <- header_entry, header_exits
-        for (hid, prev_hid) in items[1..].iter().zip(&items) {
-            let entry = self.header_entry[hid];
-            let prev_exits = &self.header_exits[prev_hid];
-            // NOTE: `extend` on a loop. TL;DR: each exits vec is pretty small.
-            // More thorough explanation at the end of `build_header`.
-            self.graph.edges.extend(
-                prev_exits
-                    .iter()
-                    .copied()
-                    .map(|e| Edge { from: e, to: entry }),
-            );
-        }
     }
 
     fn build_header(&mut self, hid: Hid, parent_cluster: Option<ClusterId>) {

engine/baml-lib/ast/src/ast/header_collector.rs

@@ -14,6 +14,7 @@
 
 use std::{collections::HashMap, sync::Arc};
 
+use baml_types::BamlMap;
 use internal_baml_diagnostics::Span;
 
 use super::{Ast, Expression, ExpressionBlock, Field, Header, Stmt, Top, WithName, WithSpan};
@@ -70,8 +71,10 @@ pub struct RenderableHeader {
 pub struct HeaderIndex {
     /// All headers in source order per scope, flattened
     pub headers: Vec<RenderableHeader>,
-    // header indexes in source order per scope
-    by_scope: HashMap<ScopeId, Vec<usize>>,
+
+    /// Header indexes in source order per scope.
+    /// Since it uses `BamlMap` it will maintain source order in snapshot testing.
+    by_scope: BamlMap<ScopeId, Vec<usize>>,
     /// Mapping of internal Hid -> names of functions called by the labeled expression
     pub header_calls: HashMap<Hid, Vec<String>>, // hid -> [callee_name]
     hid_to_idx: Vec<usize>,
@@ -91,6 +94,12 @@ impl HeaderIndex {
             .flat_map(|idxs| idxs.iter().map(|i| &self.headers[*i]))
     }
 
+    /// Iterates all scopes. Under snapshot testing, order is guaranteed to be
+    /// consistent with source order.
+    pub fn scopes<'iter>(&'iter self) -> impl Iterator<Item = ScopeId> + 'iter {
+        self.by_scope.keys().copied()
+    }
+
     /// O(1) access to a header by its Hid via internal index
     pub fn get_by_hid(&self, hid: Hid) -> Option<&RenderableHeader> {
         let idx = *self.hid_to_idx.get(hid.0 as usize)?;
@@ -107,12 +116,13 @@ impl HeaderIndex {
 }
 
 /// Internal collector to walk AST and build a HeaderIndex
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub struct HeaderCollector {
     scope_counter: u32,
     scope_stack: Vec<ScopeId>,
-    // Raw headers by scope before normalization and in-scope parenting
-    raw_by_scope: HashMap<ScopeId, Vec<RawHeader>>, // source order
+    /// Raw headers by scope before normalization and in-scope parenting.
+    /// Will maintain source order under snapshot testing.
+    raw_by_scope: BamlMap<ScopeId, Vec<RawHeader>>, // source order
     // Accumulated nested edges (Hid -> Hid)
     nested_edges_hid: Vec<(Hid, Hid)>,
     // Mapping during collection: header (by Hid) -> function names called by the labeled expression
@@ -134,15 +144,7 @@ struct RawHeader {
 
 impl HeaderCollector {
     pub fn collect(ast: &Ast) -> HeaderIndex {
-        let mut c = Self {
-            scope_counter: 0,
-            scope_stack: Vec::new(),
-            raw_by_scope: HashMap::new(),
-            nested_edges_hid: Vec::new(),
-            header_fn_calls: HashMap::new(),
-            next_hid: 0,
-            last_hdr_stack: Vec::new(),
-        };
+        let mut c = Self::default();
         c.visit_ast(ast);
         c.build_index()
     }
@@ -464,7 +466,7 @@ impl HeaderCollector {
     fn build_index(self) -> HeaderIndex {
         let mut index = HeaderIndex {
             headers: Vec::new(),
-            by_scope: HashMap::new(),
+            by_scope: BamlMap::new(),
             header_calls: HashMap::new(),
             hid_to_idx: Vec::new(),
             nested_edges_hid: Vec::new(),

engine/baml-lib/baml/tests/validation_files/headers/ai_content_pipeline.mmd

@@ -37,10 +37,8 @@ flowchart TD
   end
   n0 --> n1
   n3 --> n4
-  n4 --> n5
   n2 --> n3
   n2 --> n6
-  n7 --> n8
   n4 --> n7
   n6 --> n7
   n9 --> n10
@@ -49,6 +47,8 @@ flowchart TD
   n8 --> n9
   n10 --> n12
   n11 --> n12
+  n4 --> n5
+  n7 --> n8
 %%__BAML_SPANMAP__={"n0":{"file_path":"tests/validation_files/headers/ai_content_pipeline.baml","start_line":4,"start":4,"end_line":5,"end":0},"n1":{"file_path":"tests/validation_files/headers/ai_content_pipeline.baml","start_line":8,"start":4,"end_line":9,"end":0},"n2":{"file_path":"tests/validation_files/headers/ai_content_pipeline.baml","start_line":15,"start":4,"end_line":16,"end":0},"n3":{"file_path":"tests/validation_files/headers/ai_content_pipeline.baml","start_line":19,"start":8,"end_line":20,"end":0},"n4":{"file_path":"tests/validation_files/headers/ai_content_pipeline.baml","start_line":22,"start":8,"end_line":23,"end":0},"n5":{"file_path":"tests/validation_files/headers/ai_content_pipeline.baml","start_line":25,"start":8,"end_line":26,"end":0},"n6":{"file_path":"tests/validation_files/headers/ai_content_pipeline.baml","start_line":30,"start":8,"end_line":31,"end":0},"n7":{"file_path":"tests/validation_files/headers/ai_content_pipeline.baml","start_line":38,"start":8,"end_line":39,"end":0},"n8":{"file_path":"tests/validation_files/headers/ai_content_pipeline.baml","start_line":42,"start":8,"end_line":43,"end":0},"n9":{"file_path":"tests/validation_files/headers/ai_content_pipeline.baml","start_line":48,"start":4,"end_line":49,"end":0},"n10":{"file_path":"tests/validation_files/headers/ai_content_pipeline.baml","start_line":50,"start":8,"end_line":51,"end":0},"n11":{"file_path":"tests/validation_files/headers/ai_content_pipeline.baml","start_line":54,"start":8,"end_line":55,"end":0},"n12":{"file_path":"tests/validation_files/headers/ai_content_pipeline.baml","start_line":58,"start":4,"end_line":59,"end":0}}
   click n0 call bamlMermaidNodeClick() "Go to source"
   click n1 call bamlMermaidNodeClick() "Go to source"

engine/baml-lib/baml/tests/validation_files/headers/complex_headers_test.mmd

@@ -44,8 +44,8 @@ flowchart TD
   n0 --> n1
   n7 --> n8
   n8 --> n9
-  n10 --> n11
   n11 --> n12
+  n10 --> n11
 %%__BAML_SPANMAP__={"n0":{"file_path":"tests/validation_files/headers/complex_headers_test.baml","start_line":3,"start":4,"end_line":4,"end":0},"n1":{"file_path":"tests/validation_files/headers/complex_headers_test.baml","start_line":13,"start":8,"end_line":14,"end":0},"n2":{"file_path":"tests/validation_files/headers/complex_headers_test.baml","start_line":16,"start":8,"end_line":17,"end":0},"n3":{"file_path":"tests/validation_files/headers/complex_headers_test.baml","start_line":22,"start":12,"end_line":23,"end":0},"n4":{"file_path":"tests/validation_files/headers/complex_headers_test.baml","start_line":25,"start":12,"end_line":26,"end":0},"n5":{"file_path":"tests/validation_files/headers/complex_headers_test.baml","start_line":28,"start":12,"end_line":29,"end":0},"n6":{"file_path":"tests/validation_files/headers/complex_headers_test.baml","start_line":32,"start":8,"end_line":33,"end":0},"n7":{"file_path":"tests/validation_files/headers/complex_headers_test.baml","start_line":35,"start":8,"end_line":36,"end":0},"n8":{"file_path":"tests/validation_files/headers/complex_headers_test.baml","start_line":39,"start":4,"end_line":40,"end":0},"n9":{"file_path":"tests/validation_files/headers/complex_headers_test.baml","start_line":42,"start":4,"end_line":43,"end":0},"n10":{"file_path":"tests/validation_files/headers/complex_headers_test.baml","start_line":53,"start":8,"end_line":54,"end":0},"n11":{"file_path":"tests/validation_files/headers/complex_headers_test.baml","start_line":56,"start":8,"end_line":57,"end":0},"n12":{"file_path":"tests/validation_files/headers/complex_headers_test.baml","start_line":62,"start":4,"end_line":63,"end":0}}
   click n0 call bamlMermaidNodeClick() "Go to source"
   click n1 call bamlMermaidNodeClick() "Go to source"

engine/baml-lib/baml/tests/validation_files/headers/edge_cases.mmd

@@ -86,39 +86,39 @@ flowchart TD
   n13 --> n14
   n12 --> n13
   n15 --> n16
-  n14 --> n15
   n18 --> n19
   n17 --> n18
   n21 --> n22
-  n22 --> n23
   n20 --> n21
   n20 --> n24
-  n22 --> n25
-  n24 --> n25
-  n25 --> n26
   n31 --> n32
   n30 --> n31
   n29 --> n30
   n28 --> n29
   n27 --> n28
   n35 --> n36
   n36 --> n37
-  n33 --> n34
-  n34 --> n35
   n38 --> n39
   n38 --> n40
   n39 --> n41
   n40 --> n41
   n42 --> n43
   n41 --> n42
   n43 --> n44
-  n44 --> n45
   n46 --> n47
   n52 --> n53
   n51 --> n52
   n50 --> n51
   n49 --> n50
   n48 --> n49
+  n14 --> n15
+  n25 --> n26
+  n22 --> n25
+  n24 --> n25
+  n22 --> n23
+  n33 --> n34
+  n34 --> n35
+  n44 --> n45
 %%__BAML_SPANMAP__={"n0":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":3,"start":4,"end_line":4,"end":0},"n1":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":6,"start":4,"end_line":7,"end":0},"n2":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":9,"start":4,"end_line":10,"end":0},"n3":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":15,"start":4,"end_line":16,"end":0},"n4":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":16,"start":4,"end_line":17,"end":0},"n5":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":17,"start":4,"end_line":18,"end":0},"n6":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":18,"start":4,"end_line":19,"end":0},"n7":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":19,"start":4,"end_line":20,"end":0},"n8":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":20,"start":4,"end_line":21,"end":0},"n9":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":21,"start":4,"end_line":22,"end":0},"n10":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":22,"start":4,"end_line":23,"end":0},"n11":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":23,"start":4,"end_line":24,"end":0},"n12":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":29,"start":4,"end_line":30,"end":0},"n13":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":30,"start":4,"end_line":31,"end":0},"n14":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":31,"start":4,"end_line":32,"end":0},"n15":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":35,"start":4,"end_line":36,"end":0},"n16":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":36,"start":4,"end_line":37,"end":0},"n17":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":42,"start":4,"end_line":43,"end":0},"n18":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":48,"start":4,"end_line":49,"end":0},"n19":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":51,"start":4,"end_line":52,"end":0},"n20":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":60,"start":8,"end_line":61,"end":0},"n21":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":63,"start":12,"end_line":64,"end":0},"n22":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":66,"start":16,"end_line":67,"end":0},"n23":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":70,"start":12,"end_line":71,"end":0},"n24":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":73,"start":12,"end_line":74,"end":0},"n25":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":77,"start":8,"end_line":78,"end":0},"n26":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":81,"start":4,"end_line":82,"end":0},"n27":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":86,"start":4,"end_line":87,"end":0},"n28":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":87,"start":4,"end_line":88,"end":0},"n29":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":88,"start":4,"end_line":89,"end":0},"n30":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":89,"start":4,"end_line":90,"end":0},"n31":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":90,"start":4,"end_line":91,"end":0},"n32":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":91,"start":4,"end_line":92,"end":0},"n33":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":97,"start":4,"end_line":98,"end":0},"n34":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":98,"start":4,"end_line":99,"end":0},"n35":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":101,"start":4,"end_line":102,"end":0},"n36":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":102,"start":4,"end_line":103,"end":0},"n37":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":103,"start":4,"end_line":104,"end":0},"n38":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":113,"start":4,"end_line":114,"end":0},"n39":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":115,"start":8,"end_line":116,"end":0},"n40":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":118,"start":8,"end_line":119,"end":0},"n41":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":121,"start":4,"end_line":122,"end":0},"n42":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":123,"start":4,"end_line":124,"end":0},"n43":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":125,"start":8,"end_line":126,"end":0},"n44":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":128,"start":4,"end_line":129,"end":0},"n45":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":130,"start":4,"end_line":131,"end":0},"n46":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":135,"start":4,"end_line":136,"end":0},"n47":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":137,"start":4,"end_line":138,"end":0},"n48":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":143,"start":4,"end_line":144,"end":0},"n49":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":144,"start":4,"end_line":145,"end":0},"n50":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":145,"start":4,"end_line":146,"end":0},"n51":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":146,"start":4,"end_line":147,"end":0},"n52":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":147,"start":4,"end_line":148,"end":0},"n53":{"file_path":"tests/validation_files/headers/edge_cases.baml","start_line":148,"start":4,"end_line":149,"end":0}}
   click n0 call bamlMermaidNodeClick() "Go to source"
   click n1 call bamlMermaidNodeClick() "Go to source"

engine/baml-lib/baml/tests/validation_files/headers/for_loop_mermaid.mmd

@@ -8,8 +8,8 @@ flowchart TD
       n1["Accumulation"]
     end
   end
-  n0 --> n1
   n1 --> n2
+  n0 --> n1
 %%__BAML_SPANMAP__={"n0":{"file_path":"tests/validation_files/headers/for_loop_mermaid.baml","start_line":7,"start":8,"end_line":8,"end":0},"n1":{"file_path":"tests/validation_files/headers/for_loop_mermaid.baml","start_line":10,"start":8,"end_line":11,"end":0},"n2":{"file_path":"tests/validation_files/headers/for_loop_mermaid.baml","start_line":16,"start":4,"end_line":17,"end":0}}
   click n0 call bamlMermaidNodeClick() "Go to source"
   click n1 call bamlMermaidNodeClick() "Go to source"

engine/baml-lib/baml/tests/validation_files/headers/mixed_children_not_flattened.mmd

@@ -5,8 +5,8 @@ flowchart TD
     n1["Loop Child"]
     n2["And stuff"]
   end
-  n1 --> n2
   n0 --> n1
+  n1 --> n2
 %%__BAML_SPANMAP__={"n0":{"file_path":"tests/validation_files/headers/mixed_children_not_flattened.baml","start_line":0,"start":0,"end_line":1,"end":0},"n1":{"file_path":"tests/validation_files/headers/mixed_children_not_flattened.baml","start_line":5,"start":8,"end_line":6,"end":0},"n2":{"file_path":"tests/validation_files/headers/mixed_children_not_flattened.baml","start_line":6,"start":8,"end_line":7,"end":0}}
   click n0 call bamlMermaidNodeClick() "Go to source"
   click n1 call bamlMermaidNodeClick() "Go to source"