[wip] annotate & arrange processing order (dependency lists)

Only done for if branches. Should be done for the rest, slowly.

Author: Jesus Lapastora Núñez

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

@@ -239,9 +239,12 @@ impl<'index, 'pre> GraphBuilder<'index, 'pre> {
 
     // NOTE: we may wont to make the caches have more lifetime
     // than build so that we can return &'cache [NodeId] after writing.
+    // TODO: in process of changing signature so that it does not return anything & we just get
+    // from the built cache.
     fn build_header(
         &mut self,
         hid: Hid,
+        // TODO: understand  why this exists!
         visited_scopes: &mut HashSet<ScopeId>,
         parent_cluster: Option<ClusterId>,
     ) -> (NodeId, Vec<NodeId>) {
@@ -260,11 +263,8 @@ impl<'index, 'pre> GraphBuilder<'index, 'pre> {
             .get(&hid)
             .map(Vec::as_slice)
             .unwrap_or(&[]);
-        let has_md = !md_children.is_empty();
-        let has_nested = !nested_children.is_empty();
-        let is_branching = header.label_kind == HeaderLabelKind::If;
 
-        if !has_md && !has_nested {
+        if md_children.is_empty() && nested_children.is_empty() {
             let node_id = self.new_node_id();
             let span = SerializedSpan::serialize(&header.span);
             self.span_map.insert(node_id, span.clone());
@@ -282,65 +282,19 @@ impl<'index, 'pre> GraphBuilder<'index, 'pre> {
             return (node_id, vec![node_id]);
         }
 
-        let total_children = md_children.len() + nested_children.len();
-        let mut single_nested_child_has_multiple_items = false;
-        if nested_children.len() == 1 {
-            let child_root = self.by_hid[&nested_children[0]];
-            let count = self.index.headers_in_scope_iter(child_root.scope).count();
-            single_nested_child_has_multiple_items = count > 1;
-        }
-        let should_flatten = !is_branching
-            && total_children <= MAX_CHILDREN_TO_FLATTEN
-            && !(nested_children.len() == 1 && single_nested_child_has_multiple_items);
-
-        if should_flatten {
-            let node_id = self.new_node_id();
-            let span = SerializedSpan::serialize(&header.span);
-            self.span_map.insert(node_id, span.clone());
-            let kind = if is_branching {
-                NodeKind::Decision(hid, Some(span.clone()))
-            } else {
-                NodeKind::Header(hid, Some(span.clone()))
-            };
-            self.graph.nodes.push(Node {
-                id: node_id,
-                label: header.title.as_ref(),
-                kind,
-                cluster: parent_cluster,
-            });
-            self.header_entry.insert(hid, node_id);
-            let exits = if md_children.len() == 1 {
-                let (c_entry, c_exits) =
-                    self.build_header(md_children[0], visited_scopes, parent_cluster);
-                self.graph.edges.push(Edge {
-                    from: node_id,
-                    to: c_entry,
-                });
-                c_exits
-            } else if nested_children.len() == 1 {
-                let child_root_hid = nested_children[0];
-                let child_scope = self.by_hid[&child_root_hid].scope;
-                self.build_scope_sequence(child_scope, visited_scopes, parent_cluster);
-                let (c_entry, c_exits) =
-                    self.build_header(child_root_hid, visited_scopes, parent_cluster);
-                self.graph.edges.push(Edge {
-                    from: node_id,
-                    to: c_entry,
-                });
-                c_exits
-            };
-            self.header_exits.insert(hid, exits);
-            return (node_id, exits);
-        }
+        if header.label_kind == HeaderLabelKind::If {
+            // pre-order/post-order notation:
+            // <pre/post>-order: <name> [[output] <- <dependency list>]
 
-        if is_branching {
+            // pre-order: assign cluster ids: cluster_id <- header, parent_cluster
             let cluster_id = self.new_cluster_id();
             self.graph.clusters.push(Cluster {
                 id: cluster_id,
                 label: header.title.as_ref(),
                 parent: parent_cluster,
             });
 
+            // pre-order: insert entry <- header, cluster_id
             let decision_id = self.new_node_id();
             let span = SerializedSpan::serialize(&header.span);
             self.span_map.insert(decision_id, span.clone());
@@ -351,33 +305,56 @@ impl<'index, 'pre> GraphBuilder<'index, 'pre> {
                 cluster: Some(cluster_id),
             });
             self.header_entry.insert(hid, decision_id);
+
+            // unordered: render_calls_for_header <- header_entry, cluster_id
+
             if self.cfg.show_call_nodes {
                 self.render_calls_for_header(hid, decision_id, Some(cluster_id));
             }
 
-            let mut branch_exits = Vec::new();
-            for child_root_hid in nested_children.iter() {
-                let child_scope = self.by_hid[child_root_hid].scope;
+            // post-order: build scope sequence for scope & build header
+            // header_entry, header_exit <- cluster_id.
+            // For some reason it can't work without a visited_scopes? That's pre-order data.
+            for child_root_id in nested_children {
+                let child_scope = self.by_hid[child_root_id].scope;
                 self.build_scope_sequence(child_scope, visited_scopes, Some(cluster_id));
-                let (entry, exits) =
-                    self.build_header(*child_root_hid, visited_scopes, Some(cluster_id));
-                self.graph.edges.push(Edge {
+                self.build_header(*child_root_id, visited_scopes, Some(cluster_id));
+            }
+
+            // post-order: build header for each of the markdown children <- cluster_id.
+            // Same doubt wrt visited_scopes.
+            for child in md_children {
+                self.build_header(*child, visited_scopes, Some(cluster_id));
+            }
+
+            // unordered: add edges from decision id to children <- decision_id
+            self.graph
+                .edges
+                .extend(nested_children.iter().map(|child_root_id| Edge {
                     from: decision_id,
-                    to: entry,
-                });
-                if self.cfg.show_call_nodes {
-                    self.render_calls_for_header(*child_root_hid, entry, Some(cluster_id));
+                    to: self.header_entry[child_root_id],
+                }));
+
+            // unordered: render_calls_for_header <- header_entry, cluster_id
+            if self.cfg.show_call_nodes {
+                for child_root_hid in nested_children {
+                    self.render_calls_for_header(
+                        *child_root_hid,
+                        self.header_entry[child_root_hid],
+                        Some(cluster_id),
+                    );
                 }
-                branch_exits.extend(exits);
             }
 
-            // NOTE: here we discard the Vec<>  from `build_header`!
-            let md_ids_only: Vec<_> = md_children
+            // post-order: collect branch exits <- header_exits
+            let branch_exits: Vec<_> = nested_children
                 .iter()
+                .flat_map(|child_hid| &self.header_exits[child_hid])
                 .copied()
-                .map(|ch| self.build_header(ch, visited_scopes, Some(cluster_id)).0)
                 .collect();
 
+            let md_ids_only: Vec<_> = md_children.iter().map(|ch| self.header_entry[ch]).collect();
+
             if let Some(first_md) = md_ids_only.first().copied() {
                 if !branch_exits.is_empty() {
                     for e in branch_exits.iter() {
@@ -393,10 +370,11 @@ impl<'index, 'pre> GraphBuilder<'index, 'pre> {
                     });
                 }
             }
-            for win in md_ids_only.windows(2) {
-                let (a, b) = (&win[0], &win[1]);
-                self.graph.edges.push(Edge { from: *a, to: *b });
-            }
+            self.graph.edges.extend(md_ids_only.windows(2).map(|win| {
+                let from = win[0];
+                let to = win[1];
+                Edge { from, to }
+            }));
             let outward = if let Some(last) = md_ids_only.last().copied() {
                 vec![last]
             } else {
@@ -406,6 +384,54 @@ impl<'index, 'pre> GraphBuilder<'index, 'pre> {
             return (decision_id, outward);
         }
 
+        let single_nested_child_has_multiple_items = nested_children.len() == 1 && {
+            let child_root = self.by_hid[&nested_children[0]];
+            let count = self.index.headers_in_scope_iter(child_root.scope).count();
+            count > 1
+        };
+
+        let total_children = md_children.len() + nested_children.len();
+
+        let should_flatten =
+            total_children <= MAX_CHILDREN_TO_FLATTEN && !single_nested_child_has_multiple_items;
+
+        if should_flatten {
+            let node_id = self.new_node_id();
+            let span = SerializedSpan::serialize(&header.span);
+            self.span_map.insert(node_id, span.clone());
+            self.graph.nodes.push(Node {
+                id: node_id,
+                label: header.title.as_ref(),
+                kind: NodeKind::Header(hid, Some(span.clone())),
+                cluster: parent_cluster,
+            });
+            self.header_entry.insert(hid, node_id);
+            let exits = if md_children.len() == 1 {
+                let (c_entry, c_exits) =
+                    self.build_header(md_children[0], visited_scopes, parent_cluster);
+                self.graph.edges.push(Edge {
+                    from: node_id,
+                    to: c_entry,
+                });
+                c_exits
+            } else if nested_children.len() == 1 {
+                let child_root_hid = nested_children[0];
+                let child_scope = self.by_hid[&child_root_hid].scope;
+                self.build_scope_sequence(child_scope, visited_scopes, parent_cluster);
+                let (c_entry, c_exits) =
+                    self.build_header(child_root_hid, visited_scopes, parent_cluster);
+                self.graph.edges.push(Edge {
+                    from: node_id,
+                    to: c_entry,
+                });
+                c_exits
+            } else {
+                vec![node_id]
+            };
+            self.header_exits.insert(hid, exits.clone());
+            return (node_id, exits);
+        }
+
         let cluster_id = self.new_cluster_id();
         self.graph.clusters.push(Cluster {
             id: cluster_id,
@@ -415,6 +441,7 @@ impl<'index, 'pre> GraphBuilder<'index, 'pre> {
         });
 
         // Merge markdown children and direct nested roots, preserving each list's internal order
+        // NOTE: we could do this merge beforehand! & get both nested & md_children.
         let items_merged = merge_by_pos(self.by_hid, &md_children, &nested_children);
 
         let mut first_rep: Option<_> = None;
@@ -472,6 +499,8 @@ impl<'index, 'pre> GraphBuilder<'index, 'pre> {
         (entry, exits)
     }
 
+    /// If header is in [`HeaderIndex::header_calls`], inserts & links a call node to the
+    /// given header node.
     fn render_calls_for_header(
         &mut self,
         hid: Hid,

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

@@ -45,17 +45,17 @@ flowchart TD
   end
   n3 --> n4
   n3 --> n5
-  n2 --> n3
   n6 --> n7
   n6 --> n8
+  n2 --> n3
   n2 --> n6
-  n1 --> n2
   n10 --> n11
   n10 --> n12
-  n9 --> n10
   n13 --> n14
   n13 --> n15
+  n9 --> n10
   n9 --> n13
+  n1 --> n2
   n1 --> n9
   n0 --> n1
 %%__BAML_SPANMAP__={"n0":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":0,"start":0,"end_line":1,"end":0},"n1":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":2,"start":4,"end_line":3,"end":0},"n2":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":5,"start":8,"end_line":6,"end":0},"n3":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":7,"start":12,"end_line":8,"end":0},"n4":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":9,"start":16,"end_line":10,"end":0},"n5":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":12,"start":16,"end_line":13,"end":0},"n6":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":18,"start":12,"end_line":19,"end":0},"n7":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":20,"start":16,"end_line":21,"end":0},"n8":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":23,"start":16,"end_line":24,"end":0},"n9":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":31,"start":8,"end_line":32,"end":0},"n10":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":33,"start":12,"end_line":34,"end":0},"n11":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":35,"start":16,"end_line":36,"end":0},"n12":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":38,"start":16,"end_line":39,"end":0},"n13":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":44,"start":12,"end_line":45,"end":0},"n14":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":46,"start":16,"end_line":47,"end":0},"n15":{"file_path":"tests/validation_files/headers/nested_if_statements.baml","start_line":49,"start":16,"end_line":50,"end":0}}