feat(typecheck): enhance instanceof type checking with support for primitive types and enums

Author: AntennaeVY

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

@@ -17,12 +17,13 @@
 ///
 /// However, the current implementation is simple and ad-hoc, likely wrong
 /// in several places. Bidirectional typing is the target.
-use std::{borrow::Cow, sync::Arc};
+use std::{borrow::Cow, str::FromStr, sync::Arc};
 
 use baml_types::{
     ir_type::{ArrowGeneric, TypeIR},
     BamlMap, BamlMediaType, BamlValueWithMeta, TypeValue,
 };
+use baml_vm::types::Type;
 use internal_baml_ast::ast::WithSpan;
 use internal_baml_diagnostics::{DatamodelError, Diagnostics, Span};
 
@@ -1356,6 +1357,7 @@ pub fn typecheck_expression(
                     "image" | "audio" | "video" | "pdf" | "baml" => {}
 
                     cls if context.classes.contains_key(cls) => {}
+                    p if TypeValue::from_str(p).is_ok() => {}
 
                     _ => {
                         diagnostics.push_error(DatamodelError::new_validation_error(
@@ -2344,17 +2346,22 @@ pub fn typecheck_expression(
                     thir::Expr::Var(name, _) => {
                         if context.classes.get(name).is_some() {
                             Some(TypeIR::bool())
+                        } else if context.enums.get(name).is_some() {
+                            Some(TypeIR::bool())
+                        } else if TypeValue::from_str(name).is_ok() {
+                            Some(TypeIR::bool())
                         } else {
+                            // TODO: Check type aliases (may be recursive)
                             diagnostics.push_error(DatamodelError::new_validation_error(
-                                &format!("Class {name} not found"),
+                                &format!("Type {name} not found"),
                                 span.clone(),
                             ));
                             None
                         }
                     }
                     _ => {
                         diagnostics.push_error(DatamodelError::new_validation_error(
-                            "Invalid binary operation (instanceof): right operand must be a class",
+                            "Invalid binary operation (instanceof): right operand must be a type",
                             span.clone(),
                         ));
                         None

integ-tests/baml_src/test-files/vm/expr_funcs.baml

@@ -154,8 +154,26 @@ class DummyJsonTodo {
     userId int
 }
 
+enum DummyEnum {
+   POSITIVE
+   NEGATIVE
+   NEUTRAL
+}
+
 function ExecFetchAs(url: string) -> DummyJsonTodo {
     let todo = baml.fetch_as<DummyJsonTodo>(url);
 
     todo
 }
+
+function InstanceofWithPrimitives(n: int) -> bool {
+    n instanceof int
+}
+
+function InstanceofWithClasses(dummy: DummyJsonTodo) -> bool {
+    dummy instanceof DummyJsonTodo
+}
+
+function InstanceofWithEnums(e: DummyEnum) -> bool {
+    e instanceof DummyEnum
+}