Allow structural recursion in type aliases

engine/baml-lib/baml-core/src/ir/repr.rs

@@ -35,6 +35,9 @@ pub struct IntermediateRepr {
     /// Strongly connected components of the dependency graph (finite cycles).
     finite_recursive_cycles: Vec<IndexSet<String>>,
 
+    /// Type alias cycles introduced by lists and maps.
+    structural_recursive_alias_cycles: Vec<IndexSet<String>>,
+
     configuration: Configuration,
 }
 
@@ -53,6 +56,7 @@ impl IntermediateRepr {
             enums: vec![],
             classes: vec![],
             finite_recursive_cycles: vec![],
+            structural_recursive_alias_cycles: vec![],
             functions: vec![],
             clients: vec![],
             retry_policies: vec![],
@@ -168,6 +172,15 @@ impl IntermediateRepr {
                         .collect()
                 })
                 .collect(),
+            structural_recursive_alias_cycles: db
+                .structural_recursive_alias_cycles()
+                .iter()
+                .map(|ids| {
+                    ids.iter()
+                        .map(|id| db.ast()[*id].name().to_string())
+                        .collect()
+                })
+                .collect(),
             functions: db
                 .walk_functions()
                 .map(|e| e.node(db))

engine/baml-lib/baml-core/src/validate/validation_pipeline/validations/cycle.rs

@@ -14,24 +14,40 @@ use crate::validate::validation_pipeline::context::Context;
 
 /// Validates if the dependency graph contains one or more infinite cycles.
 pub(super) fn validate(ctx: &mut Context<'_>) {
-    // Solve cycles first. We need that information in case a class points to
-    // an unresolveble type alias.
-    let type_aliases_components = report_infinite_cycles(
-        &ctx.db.type_alias_dependencies(),
+    // We'll check type alias cycles first. Just like Typescript, cycles are
+    // allowed only for maps and lists. We'll call such cycles "structural
+    // recursion". Anything else like nulls or unions won't terminate a cycle.
+    let structural_type_aliases = HashMap::from_iter(ctx.db.walk_type_aliases().map(|alias| {
+        let mut dependencies = HashSet::new();
+        insert_required_alias_deps(alias.target(), ctx, &mut dependencies);
+
+        (alias.id, dependencies)
+    }));
+
+    // Based on the graph we've built with does not include the edges created
+    // by maps and lists, check the cycles and report them.
+    report_infinite_cycles(
+        &structural_type_aliases,
         ctx,
         "These aliases form a dependency cycle",
     );
 
-    // Store this locally to pass refs to the insert_required_deps function.
-    let alias_cycles = type_aliases_components.iter().flatten().copied().collect();
+    // In order to avoid infinite recursion when resolving types for class
+    // dependencies below, we'll compute the cycles of aliases including maps
+    // and lists so that the recursion can be stopped before entering a cycle.
+    let complete_alias_cycles = Tarjan::components(ctx.db.type_alias_dependencies())
+        .iter()
+        .flatten()
+        .copied()
+        .collect();
 
-    // First, build a graph of all the "required" dependencies represented as an
+    // Now build a graph of all the "required" dependencies represented as an
     // adjacency list. We're only going to consider type dependencies that can
     // actually cause infinite recursion. Unions and optionals can stop the
     // recursion at any point, so they don't have to be part of the "dependency"
     // graph because technically an optional field doesn't "depend" on anything,
     // it can just be null.
-    let dependency_graph = HashMap::from_iter(ctx.db.walk_classes().map(|class| {
+    let class_dependency_graph = HashMap::from_iter(ctx.db.walk_classes().map(|class| {
         let expr_block = &ctx.db.ast()[class.id];
 
         // TODO: There's already a hash set that returns "dependencies" in
@@ -47,15 +63,21 @@ pub(super) fn validate(ctx: &mut Context<'_>) {
 
         for field in &expr_block.fields {
             if let Some(field_type) = &field.expr {
-                insert_required_deps(class.id, field_type, ctx, &mut dependencies, &alias_cycles);
+                insert_required_class_deps(
+                    class.id,
+                    field_type,
+                    ctx,
+                    &mut dependencies,
+                    &complete_alias_cycles,
+                );
             }
         }
 
         (class.id, dependencies)
     }));
 
     report_infinite_cycles(
-        &dependency_graph,
+        &class_dependency_graph,
         ctx,
         "These classes form a dependency cycle",
     );
@@ -103,7 +125,7 @@ where
 /// it reaches stack overflows with large inputs.
 ///
 /// TODO: Use a struct to keep all this state. Too many parameters already.
-fn insert_required_deps(
+fn insert_required_class_deps(
     id: TypeExpId,
     field: &FieldType,
     ctx: &Context<'_>,
@@ -139,7 +161,7 @@ fn insert_required_deps(
                     // We also have to stop recursion if we know the alias is
                     // part of a cycle.
                     if !alias_cycles.contains(&alias.id) {
-                        insert_required_deps(id, alias.target(), ctx, deps, alias_cycles)
+                        insert_required_class_deps(id, alias.target(), ctx, deps, alias_cycles)
                     }
                 }
                 _ => {}
@@ -156,7 +178,7 @@ fn insert_required_deps(
             let mut nested_deps = HashSet::new();
 
             for f in field_types {
-                insert_required_deps(id, f, ctx, &mut nested_deps, alias_cycles);
+                insert_required_class_deps(id, f, ctx, &mut nested_deps, alias_cycles);
 
                 // No nested deps found on this component, this makes the
                 // union finite, so no need to go deeper.
@@ -186,3 +208,26 @@ fn insert_required_deps(
         _ => {}
     }
 }
+
+/// Implemented a la TS, maps and lists are not included as edges.
+fn insert_required_alias_deps(
+    field_type: &FieldType,
+    ctx: &Context<'_>,
+    required: &mut HashSet<TypeAliasId>,
+) {
+    match field_type {
+        FieldType::Symbol(_, ident, _) => {
+            if let Some(TypeWalker::TypeAlias(alias)) = ctx.db.find_type_by_str(ident.name()) {
+                required.insert(alias.id);
+            }
+        }
+
+        FieldType::Union(_, field_types, ..) | FieldType::Tuple(_, field_types, ..) => {
+            for f in field_types {
+                insert_required_alias_deps(f, ctx, required);
+            }
+        }
+
+        _ => {}
+    }
+}

engine/baml-lib/baml/tests/validation_files/class/recursive_type_aliases.baml

@@ -73,12 +73,6 @@ type Map = map<string, Map>
 // 50 | 
 // 51 | type EnterCycle = NoStop
 //    | 
-// error: Error validating: These aliases form a dependency cycle: Map
-//   -->  class/recursive_type_aliases.baml:56
-//    | 
-// 55 | // RecursiveMap
-// 56 | type Map = map<string, Map>
-//    | 
 // error: Error validating: These classes form a dependency cycle: Recursive
 //   -->  class/recursive_type_aliases.baml:22
 //    | 

engine/baml-lib/parser-database/src/lib.rs

@@ -138,6 +138,11 @@ impl ParserDatabase {
             self.types.resolved_type_aliases.insert(*alias_id, resolved);
         }
 
+        // Cycles left here after cycle validation are allowed. Basically lists
+        // and maps can introduce cycles.
+        self.types.structural_recursive_alias_cycles =
+            Tarjan::components(&self.types.type_alias_dependencies);
+
         // NOTE: Class dependency cycles are already checked at
         // baml-lib/baml-core/src/validate/validation_pipeline/validations/cycle.rs
         //
@@ -326,7 +331,7 @@ mod test {
     }
 
     fn assert_finite_cycles(baml: &'static str, expected: &[&[&str]]) -> Result<(), Diagnostics> {
-        let mut db = parse(baml)?;
+        let db = parse(baml)?;
 
         assert_eq!(
             db.finite_recursive_cycles()
@@ -342,6 +347,26 @@ mod test {
         Ok(())
     }
 
+    fn assert_structural_alias_cycles(
+        baml: &'static str,
+        expected: &[&[&str]],
+    ) -> Result<(), Diagnostics> {
+        let db = parse(baml)?;
+
+        assert_eq!(
+            db.structural_recursive_alias_cycles()
+                .iter()
+                .map(|ids| Vec::from_iter(ids.iter().map(|id| db.ast()[*id].name().to_string())))
+                .collect::<Vec<_>>(),
+            expected
+                .iter()
+                .map(|cycle| Vec::from_iter(cycle.iter().map(ToString::to_string)))
+                .collect::<Vec<_>>()
+        );
+
+        Ok(())
+    }
+
     #[test]
     fn find_simple_recursive_class() -> Result<(), Diagnostics> {
         assert_finite_cycles(
@@ -600,4 +625,41 @@ mod test {
 
         Ok(())
     }
+
+    #[test]
+    fn find_basic_map_structural_cycle() -> Result<(), Diagnostics> {
+        assert_structural_alias_cycles(
+            "type RecursiveMap = map<string, RecursiveMap>",
+            &[&["RecursiveMap"]],
+        )
+    }
+
+    #[test]
+    fn find_basic_list_structural_cycle() -> Result<(), Diagnostics> {
+        assert_structural_alias_cycles("type A = A[]", &[&["A"]])
+    }
+
+    #[test]
+    fn find_long_list_structural_cycle() -> Result<(), Diagnostics> {
+        assert_structural_alias_cycles(
+            r#"
+                type A = B
+                type B = C
+                type C = A[]
+            "#,
+            &[&["A", "B", "C"]],
+        )
+    }
+
+    #[test]
+    fn find_intricate_structural_cycle() -> Result<(), Diagnostics> {
+        assert_structural_alias_cycles(
+            r#"
+                type JsonValue = string | int | float | bool | null | JsonArray | JsonObject
+                type JsonArray = JsonValue[]
+                type JsonObject = map<string, JsonValue>
+            "#,
+            &[&["JsonValue", "JsonArray", "JsonObject"]],
+        )
+    }
 }

engine/baml-lib/parser-database/src/tarjan.rs

@@ -6,6 +6,7 @@
 use std::{
     cmp,
     collections::{HashMap, HashSet},
+    fmt::Debug,
     hash::Hash,
 };
 
@@ -126,8 +127,14 @@ impl<'g, V: Eq + Ord + Hash + Copy> Tarjan<'g, V> {
         self.index += 1;
         self.stack.push(node_id);
 
+        // TODO: @antoniosarosi: HashSet is random, won't always iterate in the
+        // same order. Fix this with IndexSet or something, we really don't want
+        // to sort this every single time.
+        let mut successors = Vec::from_iter(&self.graph[&node_id]);
+        successors.sort();
+
         // Visit neighbors to find strongly connected components.
-        for successor_id in &self.graph[&node_id] {
+        for successor_id in successors {
             // Grab owned state to circumvent borrow checker.
             let mut successor = self.state[successor_id];
             if successor.index == Self::UNVISITED {

engine/baml-lib/parser-database/src/types/mod.rs

@@ -276,6 +276,13 @@ pub(super) struct Types {
     /// Merge-Find Set or something like that.
     pub(super) finite_recursive_cycles: Vec<Vec<ast::TypeExpId>>,
 
+    /// Contains recursive type aliases.
+    ///
+    /// Recursive type aliases are a little bit trickier than recursive classes
+    /// because the termination condition is tied to lists and maps only. Nulls
+    /// and unions won't allow type alias cycles to be resolved.
+    pub(super) structural_recursive_alias_cycles: Vec<Vec<ast::TypeAliasId>>,
+
     pub(super) function: HashMap<ast::ValExpId, FunctionType>,
 
     pub(super) client_properties: HashMap<ast::ValExpId, ClientProperties>,

engine/baml-lib/parser-database/src/walkers/mod.rs

@@ -22,7 +22,9 @@ pub use configuration::*;
 use either::Either;
 pub use field::*;
 pub use function::FunctionWalker;
-use internal_baml_schema_ast::ast::{FieldType, Identifier, TopId, TypeExpId, WithName};
+use internal_baml_schema_ast::ast::{
+    FieldType, Identifier, TopId, TypeAliasId, TypeExpId, WithName,
+};
 pub use r#class::*;
 pub use r#enum::*;
 pub use template_string::TemplateStringWalker;
@@ -142,6 +144,14 @@ impl<'db> crate::ParserDatabase {
         &self.types.finite_recursive_cycles
     }
 
+    /// Set of all aliases that are part of a structural cycle.
+    ///
+    /// A structural cycle is created through a map or list, which introduce one
+    /// level of indirection.
+    pub fn structural_recursive_alias_cycles(&self) -> &[Vec<TypeAliasId>] {
+        &self.types.structural_recursive_alias_cycles
+    }
+
     /// Returns the resolved aliases map.
     pub fn resolved_type_alias_by_name(&self, alias: &str) -> Option<&FieldType> {
         match self.find_type_by_str(alias) {
@@ -209,6 +219,17 @@ impl<'db> crate::ParserDatabase {
             })
     }
 
+    /// Walk all the type aliases in the AST.
+    pub fn walk_type_aliases(&self) -> impl Iterator<Item = TypeAliasWalker<'_>> {
+        self.ast()
+            .iter_tops()
+            .filter_map(|(top_id, _)| top_id.as_type_alias_id())
+            .map(move |top_id| Walker {
+                db: self,
+                id: top_id,
+            })
+    }
+
     /// Walk all template strings in the schema.
     pub fn walk_templates(&self) -> impl Iterator<Item = TemplateStringWalker<'_>> {
         self.ast()

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

@@ -191,6 +191,14 @@ impl TopId {
         }
     }
 
+    /// Try to interpret the top as a type alias.
+    pub fn as_type_alias_id(self) -> Option<TypeAliasId> {
+        match self {
+            TopId::TypeAlias(id) => Some(id),
+            _ => None,
+        }
+    }
+
     /// Try to interpret the top as a function.
     pub fn as_function_id(self) -> Option<ValExpId> {
         match self {