Merge remote-tracking branch 'origin/canary' into sam/jetbrains

Author: sxlijin

engine/baml-compiler/src/codegen.rs

@@ -2,7 +2,10 @@
 //!
 //! This files contains the convertions between Baml AST nodes to HIR nodes.
 
+use std::collections::HashSet;
+
 use baml_types::{
+    ir_type::TypeGeneric,
     type_meta::{self, base::StreamingBehavior},
     Constraint, ConstraintLevel, TypeIR, TypeValue,
 };
@@ -48,6 +51,38 @@ impl Hir {
             }
         }
 
+        let enums = HashSet::<&str>::from_iter(hir.enums.iter().map(|e| e.name.as_str()));
+
+        let param_type: fn(&mut Parameter) -> &mut TypeIR = |p| &mut p.r#type;
+
+        // Patch return types because only here in the code we have the full
+        // context for enums.
+        hir.expr_functions
+            .iter_mut()
+            .map(|f| (f.parameters.iter_mut().map(param_type), &mut f.return_type))
+            .chain(
+                hir.llm_functions
+                    .iter_mut()
+                    .map(|f| (f.parameters.iter_mut().map(param_type), &mut f.return_type)),
+            )
+            .chain(hir.classes.iter_mut().flat_map(|c| {
+                c.methods
+                    .iter_mut()
+                    .map(|f| (f.parameters.iter_mut().map(param_type), &mut f.return_type))
+            }))
+            .flat_map(|(parameters, return_type)| parameters.chain(std::iter::once(return_type)))
+            .for_each(|ty| match ty {
+                TypeIR::Class { name, meta, .. } if enums.contains(name.as_str()) => {
+                    *ty = TypeIR::Enum {
+                        name: name.to_owned(),
+                        dynamic: false, // TODO: How to know if it's dynamic.
+                        meta: meta.clone(),
+                    }
+                }
+
+                _ => {}
+            });
+
         hir
     }
 }
@@ -125,22 +160,12 @@ pub fn type_ir_from_ast(type_: &ast::FieldType) -> TypeIR {
     };
 
     match type_ {
-        ast::FieldType::Symbol(_, name, _) => {
-            if name.name().starts_with("Enum") {
-                TypeIR::Enum {
-                    name: name.name().to_string(),
-                    dynamic: false,
-                    meta,
-                }
-            } else {
-                TypeIR::Class {
-                    name: name.name().to_string(),
-                    mode: baml_types::ir_type::StreamingMode::NonStreaming,
-                    dynamic: false,
-                    meta,
-                }
-            }
-        }
+        ast::FieldType::Symbol(_, name, _) => TypeIR::Class {
+            name: name.name().to_string(),
+            mode: baml_types::ir_type::StreamingMode::NonStreaming,
+            dynamic: false,
+            meta,
+        },
         ast::FieldType::Primitive(_, prim, _, _) => TypeIR::Primitive(*prim, meta),
         ast::FieldType::List(_, inner, dims, _, _) => {
             // Respect multi-dimensional arrays (e.g., int[][] has dims=2)

engine/baml-compiler/src/hir/lowering.rs

@@ -9,10 +9,20 @@ pub mod test {
     use internal_baml_diagnostics::Diagnostics;
     use internal_baml_parser_database::{parse_and_diagnostics, ParserDatabase};
 
+    use crate::{hir, thir};
+
     /// Shim helper function for testing.
     pub fn ast(source: &'static str) -> anyhow::Result<ParserDatabase> {
-        let (parser_db, diagnostics) = parse_and_diagnostics(source)?;
+        let (parser_db, mut diagnostics) = parse_and_diagnostics(source)?;
+
+        if diagnostics.has_errors() {
+            let errors = diagnostics.to_pretty_string();
+            anyhow::bail!("{errors}");
+        }
 
+        // Here because of cycle dependencies between crates and shit.
+        // TODO: We're building this like 3 different times, needs refactoring.
+        thir::typecheck::typecheck(&hir::Hir::from_ast(&parser_db.ast), &mut diagnostics);
         if diagnostics.has_errors() {
             let errors = diagnostics.to_pretty_string();
             anyhow::bail!("{errors}");

engine/baml-compiler/src/lib.rs

@@ -2,7 +2,7 @@
 ///
 use baml_types::ir_type::TypeIR;
 
-use crate::hir::{self, AssignOp, BinaryOperator, Enum, LlmFunction, UnaryOperator};
+use crate::hir::{self, AssignOp, BinaryOperator, LlmFunction, UnaryOperator};
 
 pub mod interpret;
 pub mod typecheck;
@@ -54,6 +54,14 @@ pub struct Class<T> {
     pub span: Span,
 }
 
+#[derive(Clone, Debug)]
+pub struct Enum {
+    pub name: String,
+    pub variants: Vec<hir::EnumVariant>,
+    pub span: Span,
+    pub ty: TypeIR, // TODO: Used for type checking, but do we need this?
+}
+
 /// A BAML expression term.
 /// T is the type of the metadata.
 #[derive(Debug, Clone)]

engine/baml-compiler/src/thir.rs

@@ -44,7 +44,7 @@ pub fn typecheck_returning_context<'a>(
         .map(|c| (c.name.clone(), c))
         .collect();
 
-    let enums = hir
+    let enums: BamlMap<String, hir::Enum> = hir
         .enums
         .clone()
         .into_iter()
@@ -54,6 +54,7 @@ pub fn typecheck_returning_context<'a>(
     // Create typing context with all functions
     let mut typing_context = TypeContext::new();
     typing_context.classes.extend(classes.clone());
+    typing_context.enums.extend(enums.clone());
 
     // Add expr functions to typing context
     for func in &hir.expr_functions {
@@ -271,11 +272,32 @@ pub fn typecheck_returning_context<'a>(
         })
         .collect();
 
+    // TODO: Those are HIR enums, figure out if there's something different we
+    // would need in a THIR enum? Does it need a "type"?.
+    let thir_enums = enums
+        .iter()
+        .map(|(name, enum_def)| {
+            (
+                name.clone(),
+                thir::Enum {
+                    name: enum_def.name.clone(),
+                    variants: enum_def.variants.clone(),
+                    span: enum_def.span.clone(),
+                    ty: TypeIR::Enum {
+                        name: enum_def.name.clone(),
+                        dynamic: false,
+                        meta: Default::default(),
+                    },
+                },
+            )
+        })
+        .collect();
+
     (
         THir {
             llm_functions: hir.llm_functions.clone(),
             classes: thir_classes,
-            enums,
+            enums: thir_enums,
             expr_functions,
             global_assignments: BamlMap::new(),
         },
@@ -302,6 +324,7 @@ pub struct TypeContext<'func> {
     // Variables in scope with mutability info
     pub vars: BamlMap<String, VarInfo>,
     pub classes: BamlMap<String, hir::Class>,
+    pub enums: BamlMap<String, hir::Enum>,
     // Used for knowing whether `break` and `continue` are inside a loop or not.
     pub is_inside_loop: bool,
 
@@ -336,6 +359,7 @@ impl TypeContext<'_> {
             symbols: BamlMap::new(),
             vars,
             classes: BamlMap::new(),
+            enums: BamlMap::new(),
             is_inside_loop: false,
             function_return_type: None,
         }
@@ -428,7 +452,9 @@ impl TypeContext<'_> {
             }
             hir::Expression::Identifier(name, _) => {
                 // Look up type in context
-                self.get_type(name).cloned()
+                self.get_type(name)
+                    .cloned()
+                    .or_else(|| self.enums.get(name).map(|e| TypeIR::r#enum(&e.name)))
             }
             hir::Expression::Array(items, _) => {
                 // Infer array type from first item
@@ -575,6 +601,14 @@ impl TypeContext<'_> {
                                 None
                             }
                         }
+                        TypeIR::Enum {
+                            name: enum_name, ..
+                        } => {
+                            // Look up field in enum definition
+                            self.enums
+                                .get(&enum_name)
+                                .map(|enum_def| TypeIR::r#enum(&enum_def.name))
+                        }
                         _ => None, // Not a class
                     }
                 } else {
@@ -1162,6 +1196,14 @@ pub fn typecheck_expression(
             BamlValueWithMeta::String(value.clone(), (span.clone(), Some(TypeIR::string()))),
         ),
         hir::Expression::Identifier(name, span) => {
+            // Enum access: let x = Shape.Rectangle
+            if let Some(enum_def) = context.enums.get(name) {
+                return thir::Expr::Var(
+                    name.clone(),
+                    (span.clone(), Some(TypeIR::r#enum(&enum_def.name))),
+                );
+            }
+
             // Look up type in context
             let var_type = context.get_type(name).cloned();
             if var_type.is_none() {
@@ -1740,6 +1782,20 @@ pub fn typecheck_expression(
                         None
                     }
                 }
+                Some(TypeIR::Enum {
+                    name: enum_name, ..
+                }) => {
+                    // Look up field in enum definition
+                    if let Some(enum_def) = context.enums.get(enum_name) {
+                        Some(TypeIR::r#enum(&enum_def.name))
+                    } else {
+                        diagnostics.push_error(DatamodelError::new_validation_error(
+                            &format!("Enum {enum_name} not found"),
+                            span.clone(),
+                        ));
+                        None
+                    }
+                }
                 _ => {
                     diagnostics.push_error(DatamodelError::new_validation_error(
                         "Can only access fields on class instances",

engine/baml-compiler/src/thir/typecheck.rs

@@ -0,0 +1,29 @@
+//! Compiler tests for array construction.
+
+use baml_vm::Instruction;
+
+mod common;
+use common::{assert_compiles, Program};
+
+#[test]
+fn array_constructor() -> anyhow::Result<()> {
+    assert_compiles(Program {
+        source: "
+            fn main() -> int[] {
+                let a = [1, 2, 3];
+                a
+            }
+        ",
+        expected: vec![(
+            "main",
+            vec![
+                Instruction::LoadConst(0),
+                Instruction::LoadConst(1),
+                Instruction::LoadConst(2),
+                Instruction::AllocArray(3),
+                Instruction::LoadVar(1),
+                Instruction::Return,
+            ],
+        )],
+    })
+}

engine/baml-compiler/tests/arrays.rs

@@ -0,0 +1,54 @@
+//! Compiler tests for assert statements.
+
+use baml_vm::{BinOp, CmpOp, Instruction};
+
+mod common;
+use common::{assert_compiles, Program};
+
+#[test]
+fn assert_statement_ok() -> anyhow::Result<()> {
+    assert_compiles(Program {
+        source: "
+            fn assertOk() -> int {
+                assert 2 + 2 == 4;
+                3
+            }
+        ",
+        expected: vec![(
+            "assertOk",
+            vec![
+                Instruction::LoadConst(0), // 2
+                Instruction::LoadConst(1), // 2
+                Instruction::BinOp(BinOp::Add),
+                Instruction::LoadConst(2), // 4
+                Instruction::CmpOp(CmpOp::Eq),
+                Instruction::Assert,
+                Instruction::LoadConst(3), // 3
+                Instruction::Return,
+            ],
+        )],
+    })
+}
+
+#[test]
+fn assert_statement_not_ok() -> anyhow::Result<()> {
+    assert_compiles(Program {
+        source: "
+            fn assertNotOk() -> int {
+                assert 3 == 1;
+                2
+            }
+        ",
+        expected: vec![(
+            "assertNotOk",
+            vec![
+                Instruction::LoadConst(0), // 3
+                Instruction::LoadConst(1), // 1
+                Instruction::CmpOp(CmpOp::Eq),
+                Instruction::Assert,
+                Instruction::LoadConst(2), // 2
+                Instruction::Return,
+            ],
+        )],
+    })
+}

engine/baml-compiler/tests/assertions.rs

@@ -0,0 +1,32 @@
+//! Compiler tests for built-in method calls.
+
+use baml_vm::{GlobalIndex, Instruction};
+
+mod common;
+use common::{assert_compiles, Program};
+
+#[test]
+fn builtin_method_call() -> anyhow::Result<()> {
+    assert_compiles(Program {
+        source: r#"
+            fn main() -> int {
+                let arr = [1, 2, 3];
+                arr.len()
+            }
+        "#,
+        expected: vec![(
+            "main",
+            vec![
+                Instruction::LoadConst(0),
+                Instruction::LoadConst(1),
+                Instruction::LoadConst(2),
+                Instruction::AllocArray(3),
+                Instruction::LoadGlobal(GlobalIndex::from_raw(3)),
+                Instruction::LoadVar(1),
+                // call with one argument (self)
+                Instruction::Call(1),
+                Instruction::Return,
+            ],
+        )],
+    })
+}