move baml_vis stuff to baml_vis; make scope traversal always source order

Author: Jesus Lapastora Núñez

engine/Cargo.lock

@@ -26,6 +26,7 @@ either.workspace = true
 test-log.workspace = true
 pretty = { workspace = true }
 regex.workspace = true
+indexmap.workspace = true
 
 [dev-dependencies]
 pretty_assertions.workspace = true

engine/baml-lib/ast/Cargo.toml

@@ -23,25 +23,21 @@ mod type_expression_block;
 mod value_expression_block;
 
 mod baml_vis;
-pub mod header_collector;
-pub mod mermaid_debug;
 pub use app::App;
 pub use argument::{Argument, ArgumentId, ArgumentsList};
 pub use assignment::Assignment;
 pub use attribute::{Attribute, AttributeContainer, AttributeId};
-pub use baml_vis::diagram_generator;
+pub use baml_vis::{diagram_generator, mermaid_debug::MermaidDiagramGenerator};
 pub use config::ConfigBlockProperty;
 pub use expr::{ExprFn, TopLevelAssignment};
 pub use expression::{
     BinaryOperator, ClassConstructor, ClassConstructorField, Expression, ExpressionBlock,
     RawString, UnaryOperator,
 };
 pub use field::{Field, FieldArity, FieldType};
-pub use header_collector::{HeaderCollector, HeaderIndex, RenderableHeader, ScopeId};
 pub use identifier::{Identifier, RefIdentifier};
 pub use indentation_type::IndentationType;
 pub use internal_baml_diagnostics::Span;
-pub use mermaid_debug::MermaidDiagramGenerator;
 pub use newline_type::NewlineType;
 pub use stmt::{
     AssertStmt, AssignOp, AssignOpStmt, AssignStmt, CForLoopStmt, ExprStmt, ForLoopStmt, Header,

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

@@ -1,4 +1,7 @@
+pub mod diagram_generator;
 mod graph;
+mod header_collector;
+pub(super) mod mermaid_debug;
 
 use std::collections::{HashMap, HashSet};
 
@@ -9,10 +12,3 @@ use baml_types::BamlMap;
 use graph::{Cluster, ClusterId, Direction, Graph, Node, NodeId, NodeKind};
 use internal_baml_diagnostics::SerializedSpan;
 use serde_json;
-
-use super::{
-    header_collector::{HeaderLabelKind, Hid},
-    Ast, HeaderCollector, HeaderIndex, RenderableHeader, ScopeId,
-};
-
-pub mod diagram_generator;

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

@@ -12,12 +12,16 @@ use internal_baml_diagnostics::SerializedSpan;
 
 use crate::ast::{
     baml_vis::{
-        graph::{self, BuilderConfig, Cluster, ClusterId, Direction, Graph, Node, NodeId, NodeKind},
+        graph::{
+            self, BuilderConfig, Cluster, ClusterId, Direction, Graph, Node, NodeId, NodeKind,
+        },
         SHOW_CALL_NODES,
     },
-    Ast, HeaderCollector,
+    Ast,
 };
 
+use super::header_collector::HeaderCollector;
+
 /// Generate a Mermaid flowchart (LR) showing headers as linear steps and
 /// nested scopes as subgraphs.
 pub fn generate_headers_flowchart(ast: &Ast) -> String {
@@ -64,17 +68,11 @@ fn render_mermaid_graph(
 
     let mut children_by_parent: BamlMap<_, Vec<_>> = BamlMap::new();
     for c in &graph.clusters {
-        children_by_parent
-            .entry(c.parent)
-            .or_default()
-            .push(c);
+        children_by_parent.entry(c.parent).or_default().push(c);
     }
     let mut nodes_by_cluster: BamlMap<_, Vec<_>> = BamlMap::new();
     for n in &graph.nodes {
-        nodes_by_cluster
-            .entry(n.cluster)
-            .or_default()
-            .push(n);
+        nodes_by_cluster.entry(n.cluster).or_default().push(n);
     }
 
     fn emit<'index>(

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

@@ -7,10 +7,7 @@ use std::{
 use baml_types::BamlMap;
 use internal_baml_diagnostics::SerializedSpan;
 
-use crate::ast::{
-    header_collector::{HeaderLabelKind, Hid},
-    HeaderIndex, RenderableHeader, ScopeId,
-};
+use super::header_collector::{HeaderIndex, HeaderLabelKind, Hid, RenderableHeader, ScopeId};
 
 /// Config: maximum number of direct children (markdown + nested) a non-branching
 /// container may have to be flattened into a linear sequence instead of a subgraph.

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

@@ -14,10 +14,12 @@
 
 use std::{collections::HashMap, sync::Arc};
 
-use baml_types::BamlMap;
+use indexmap::IndexMap;
 use internal_baml_diagnostics::Span;
 
-use super::{Ast, Expression, ExpressionBlock, Field, Header, Stmt, Top, WithName, WithSpan};
+use crate::ast::{
+    Ast, ClassConstructorField, Expression, ExpressionBlock, Field, Header, Stmt, Top,
+};
 
 /// Alias for external header identifiers for public consumption
 type HeaderId = String;
@@ -30,16 +32,6 @@ pub struct Hid(pub u32);
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct ScopeId(pub u32);
 
-/// Classification of scope kinds for visualization semantics
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub enum ScopeKind {
-    TopLevel,
-    ForBody,
-    IfThen,
-    IfElse,
-    Generic,
-}
-
 /// Classification of what kind of statement a header labels
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum HeaderLabelKind {
@@ -73,8 +65,10 @@ pub struct HeaderIndex {
     pub headers: Vec<RenderableHeader>,
 
     /// 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>>,
+    /// Uses [`IndexMap`] instead of BamlMap to always guarantee iteration order = insertion order
+    /// = source order.
+    /// See [`graph`](crate::ast::baml_vis::graph) module, where we add
+    by_scope: IndexMap<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>,
@@ -94,7 +88,7 @@ impl HeaderIndex {
             .flat_map(|idxs| idxs.iter().map(|i| &self.headers[*i]))
     }
 
-    /// Iterates all scopes. Under snapshot testing, order is guaranteed to be
+    /// Iterates all scopes. Order of iteration is guaranteed to be
     /// consistent with source order.
     pub fn scopes<'iter>(&'iter self) -> impl Iterator<Item = ScopeId> + 'iter {
         self.by_scope.keys().copied()
@@ -121,8 +115,8 @@ pub struct HeaderCollector {
     scope_counter: u32,
     scope_stack: Vec<ScopeId>,
     /// 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
+    /// [`IndexMap`] to maintain source order.
+    raw_by_scope: IndexMap<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
@@ -353,8 +347,8 @@ impl HeaderCollector {
             Expression::ClassConstructor(cons, _) => {
                 for f in &cons.fields {
                     match f {
-                        super::ClassConstructorField::Named(_, e) => self.visit_expression(e),
-                        super::ClassConstructorField::Spread(e) => self.visit_expression(e),
+                        ClassConstructorField::Named(_, e) => self.visit_expression(e),
+                        ClassConstructorField::Spread(e) => self.visit_expression(e),
                     }
                 }
             }
@@ -466,7 +460,7 @@ impl HeaderCollector {
     fn build_index(self) -> HeaderIndex {
         let mut index = HeaderIndex {
             headers: Vec::new(),
-            by_scope: BamlMap::new(),
+            by_scope: IndexMap::new(),
             header_calls: HashMap::new(),
             hid_to_idx: Vec::new(),
             nested_edges_hid: Vec::new(),
@@ -545,6 +539,7 @@ impl HeaderCollector {
 /// Only captures the outermost call(s) that structurally represent the expression,
 /// ignoring nested calls within arguments or sub-expressions.
 fn collect_top_level_calls(expr: &Expression) -> Vec<String> {
+    use crate::ast::WithName;
     match expr {
         // If the expression is a block, the top-level expression is inside it
         Expression::ExprBlock(block, _span) => {

engine/baml-lib/ast/src/ast/header_collector.rs renamed to engine/baml-lib/ast/src/ast/baml_vis/header_collector.rs

@@ -1,12 +1,14 @@
 use std::collections::HashMap;
 
-use super::{
-    App, Argument, ArgumentsList, Assignment, Ast, Attribute, BlockArgs, ClassConstructor,
-    ClassConstructorField, ExprFn, Expression, ExpressionBlock, Field, FieldType, Header,
-    HeaderCollector, HeaderIndex, Identifier, RawString, Stmt, TemplateString, Top,
-    TopLevelAssignment, TypeExpressionBlock, ValueExprBlock, WithIdentifier, WithName, WithSpan,
+use crate::ast::{
+    traits::{WithIdentifier, WithName},
+    Argument, ArgumentsList, Assignment, Ast, Attribute, BlockArgs, ClassConstructor,
+    ClassConstructorField, ExprFn, Expression, ExpressionBlock, Field, FieldType, Header, Stmt,
+    TemplateString, Top, TopLevelAssignment, TypeExpressionBlock, ValueExprBlock,
 };
 
+use super::header_collector::{HeaderCollector, HeaderIndex};
+
 /// A debug utility for converting AST structures to Mermaid diagrams
 #[derive(Debug)]
 pub struct MermaidDiagramGenerator {