Allow structural recursion in type aliases

engine/baml-lib/baml-core/src/ir/ir_helpers/to_baml_arg.rs

@@ -263,8 +263,24 @@ impl ArgCoercer {
                     Err(())
                 }
             },
-            (FieldType::Alias { resolution, .. }, _) => {
-                self.coerce_arg(ir, &resolution, value, scope)
+            // TODO: Is this even possible?
+            (FieldType::RecursiveTypeAlias(name), _) => {
+                let mut maybe_coerced = None;
+                // TODO: Fix this O(n)
+                for cycle in ir.structural_recursive_alias_cycles().iter() {
+                    if let Some(target) = cycle.get(name) {
+                        maybe_coerced = Some(self.coerce_arg(ir, target, value, scope)?);
+                        break;
+                    }
+                }
+
+                match maybe_coerced {
+                    Some(coerced) => Ok(coerced),
+                    None => {
+                        scope.push_error(format!("Recursive type alias {} not found", name));
+                        Err(())
+                    }
+                }
             }
             (FieldType::List(item), _) => match value {
                 BamlValue::List(arr) => {

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

@@ -156,10 +156,12 @@ impl WithJsonSchema for FieldType {
             FieldType::Class(name) | FieldType::Enum(name) => json!({
                 "$ref": format!("#/definitions/{}", name),
             }),
-            FieldType::Alias { resolution, .. } => resolution.json_schema(),
             FieldType::Literal(v) => json!({
                 "const": v.to_string(),
             }),
+            FieldType::RecursiveTypeAlias(_) => json!({
+                "type": ["number", "string", "boolean", "object", "array", "null"]
+            }),
             FieldType::Primitive(t) => match t {
                 TypeValue::String => json!({
                     "type": "string",

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<IndexMap<String, FieldType>>,
+
     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![],
@@ -98,6 +102,10 @@ impl IntermediateRepr {
         &self.finite_recursive_cycles
     }
 
+    pub fn structural_recursive_alias_cycles(&self) -> &[IndexMap<String, FieldType>] {
+        &self.structural_recursive_alias_cycles
+    }
+
     pub fn walk_enums(&self) -> impl ExactSizeIterator<Item = Walker<'_, &Node<Enum>>> {
         self.enums.iter().map(|e| Walker { db: self, item: e })
     }
@@ -106,6 +114,14 @@ impl IntermediateRepr {
         self.classes.iter().map(|e| Walker { db: self, item: e })
     }
 
+    // TODO: Exact size Iterator + Node<>?
+    pub fn walk_alias_cycles(&self) -> impl Iterator<Item = Walker<'_, (&String, &FieldType)>> {
+        self.structural_recursive_alias_cycles
+            .iter()
+            .flatten()
+            .map(|e| Walker { db: self, item: e })
+    }
+
     pub fn function_names(&self) -> impl ExactSizeIterator<Item = &str> {
         self.functions.iter().map(|f| f.elem.name())
     }
@@ -168,6 +184,18 @@ impl IntermediateRepr {
                         .collect()
                 })
                 .collect(),
+            structural_recursive_alias_cycles: {
+                let mut recursive_aliases = vec![];
+                for cycle in db.structural_recursive_alias_cycles() {
+                    let mut component = IndexMap::new();
+                    for id in cycle {
+                        let alias = &db.ast()[*id];
+                        component.insert(alias.name().to_string(), alias.value.repr(db)?);
+                    }
+                    recursive_aliases.push(component);
+                }
+                recursive_aliases
+            },
             functions: db
                 .walk_functions()
                 .map(|e| e.node(db))
@@ -419,11 +447,18 @@ impl WithRepr<FieldType> for ast::FieldType {
                             _ => base_type,
                         }
                     }
-                    Some(TypeWalker::TypeAlias(alias_walker)) => FieldType::Alias {
-                        name: alias_walker.name().to_owned(),
-                        target: Box::new(alias_walker.target().repr(db)?),
-                        resolution: Box::new(alias_walker.resolved().repr(db)?),
-                    },
+                    Some(TypeWalker::TypeAlias(alias_walker)) => {
+                        if db
+                            .structural_recursive_alias_cycles()
+                            .iter()
+                            .any(|cycle| cycle.contains(&alias_walker.id))
+                        {
+                            FieldType::RecursiveTypeAlias(alias_walker.name().to_string())
+                        } else {
+                            alias_walker.resolved().to_owned().repr(db)?
+                        }
+                    }
+
                     None => return Err(anyhow!("Field type uses unresolvable local identifier")),
                 },
                 arity,
@@ -1224,11 +1259,7 @@ mod tests {
         let class = ir.find_class("Test").unwrap();
         let alias = class.find_field("field").unwrap();
 
-        let FieldType::Alias { resolution, .. } = alias.r#type() else {
-            panic!("expected alias type, found {:?}", alias.r#type());
-        };
-
-        assert_eq!(**resolution, FieldType::Primitive(TypeValue::Int));
+        assert_eq!(*alias.r#type(), FieldType::Primitive(TypeValue::Int));
     }
 
     #[test]
@@ -1249,12 +1280,11 @@ mod tests {
         let class = ir.find_class("Test").unwrap();
         let alias = class.find_field("field").unwrap();
 
-        let FieldType::Alias { resolution, .. } = alias.r#type() else {
-            panic!("expected alias type, found {:?}", alias.r#type());
-        };
-
-        let FieldType::Constrained { base, constraints } = &**resolution else {
-            panic!("expected resolved constrained type, found {:?}", resolution);
+        let FieldType::Constrained { base, constraints } = alias.r#type() else {
+            panic!(
+                "expected resolved constrained type, found {:?}",
+                alias.r#type()
+            );
         };
 
         assert_eq!(constraints.len(), 3);

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-types/src/field_type/mod.rs

@@ -85,14 +85,7 @@ pub enum FieldType {
     Union(Vec<FieldType>),
     Tuple(Vec<FieldType>),
     Optional(Box<FieldType>),
-    Alias {
-        /// Name of the alias.
-        name: String,
-        /// Type that the alias points to.
-        target: Box<FieldType>,
-        /// Final resolved type (an alias can point to other aliases).
-        resolution: Box<FieldType>,
-    },
+    RecursiveTypeAlias(String),
     Constrained {
         base: Box<FieldType>,
         constraints: Vec<Constraint>,
@@ -103,8 +96,9 @@ pub enum FieldType {
 impl std::fmt::Display for FieldType {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            FieldType::Enum(name) | FieldType::Class(name) => write!(f, "{name}"),
-            FieldType::Alias { name, .. } => write!(f, "{name}"),
+            FieldType::Enum(name)
+            | FieldType::Class(name)
+            | FieldType::RecursiveTypeAlias(name) => write!(f, "{name}"),
             FieldType::Primitive(t) => write!(f, "{t}"),
             FieldType::Literal(v) => write!(f, "{v}"),
             FieldType::Union(choices) => {
@@ -187,8 +181,6 @@ impl FieldType {
         }
 
         match (self, other) {
-            (FieldType::Alias { resolution, .. }, _) => resolution.is_subtype_of(other),
-            (_, FieldType::Alias { resolution, .. }) => self.is_subtype_of(resolution),
             (FieldType::Primitive(TypeValue::Null), FieldType::Optional(_)) => true,
             (FieldType::Optional(self_item), FieldType::Optional(other_item)) => {
                 self_item.is_subtype_of(other_item)
@@ -207,6 +199,12 @@ impl FieldType {
             }
             (FieldType::Map(_, _), _) => false,
 
+            // TODO: is it necessary to check if the alias is part of the same
+            // cycle?
+            (FieldType::RecursiveTypeAlias(_), _) | (_, FieldType::RecursiveTypeAlias(_)) => {
+                self == other
+            }
+
             (
                 FieldType::Constrained {
                     base: self_base,

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
 //    |