From 612b844dbf86041959ddf27ce44241d1a87bc7dd Mon Sep 17 00:00:00 2001 From: Aaron Villalpando Date: Fri, 2 Oct 2026 15:46:11 -0700 Subject: [PATCH 1/8] Build each parse model from only the types its target reaches Every LLM parse built a schema-aligned-parsing model of every class, enum, and type alias in the program, the standard library included: about 4 MB per model and two models per call, held by a VM object whose native size the GC budget does not count. A busy engine never collected it, so RSS grew with the call count. TypeCtx::for_target now keeps only the definitions the parse target reaches, walked with RuntimeTy::visit_heads (the head walk bex_vm::reachable uses). A small LLM call drops from about 13.3 MB to about 3.4 MB allocated, and 400 unrelated classes add about 22 KB per call instead of 6.3 MB. baml_tests/tests/llm_call_memory.rs counts every allocation with a global allocator and bounds both, so the runtime cannot regress to per-program parse cost again. Co-Authored-By: Claude Opus 5.5 --- .../baml_tests/tests/llm_call_memory.rs | 218 +++++++++++++++ baml_language/crates/bex_sap/src/lib.rs | 2 +- .../crates/bex_sap/src/sap_model/convert.rs | 255 ++++++++++++++++-- 3 files changed, 457 insertions(+), 18 deletions(-) create mode 100644 baml_language/crates/baml_tests/tests/llm_call_memory.rs diff --git a/baml_language/crates/baml_tests/tests/llm_call_memory.rs b/baml_language/crates/baml_tests/tests/llm_call_memory.rs new file mode 100644 index 00000000000..0c618258d7b --- /dev/null +++ b/baml_language/crates/baml_tests/tests/llm_call_memory.rs @@ -0,0 +1,218 @@ +//! The bytes the runtime allocates for one LLM call stay small. +//! +//! A Rust test because BAML cannot observe allocations. Each parse used to build +//! a schema-aligned-parsing model of every type in the program, the standard +//! library included: about 5 MB per call for a one-function program, kept alive +//! until the VM collected it. A counting global allocator measures every byte +//! allocated while the calls run, on every thread, so the bound does not depend +//! on when the GC runs. + +#![allow( + unsafe_code, + reason = "a counting global allocator forwards every call to `System`" +)] + +use std::{ + alloc::{GlobalAlloc, Layout, System}, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, +}; + +use baml_tests::engine::compile_source; +use bex_engine::{BexEngine, BexExternalValue, FunctionCallContextBuilder}; +use sys_native::SysOpsExt; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{method, path}, +}; + +/// Every byte handed out since the process started. +static ALLOCATED: AtomicUsize = AtomicUsize::new(0); + +struct CountingAllocator; + +// SAFETY: every method forwards its arguments unchanged to `System`, which +// upholds the `GlobalAlloc` contract; the counter has no effect on memory. +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + ALLOCATED.fetch_add(layout.size(), Ordering::Relaxed); + // SAFETY: the caller's contract for `alloc` is forwarded unchanged. + unsafe { System.alloc(layout) } + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + ALLOCATED.fetch_add(layout.size(), Ordering::Relaxed); + // SAFETY: the caller's contract for `alloc_zeroed` is forwarded unchanged. + unsafe { System.alloc_zeroed(layout) } + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + ALLOCATED.fetch_add(new_size.saturating_sub(layout.size()), Ordering::Relaxed); + // SAFETY: the caller's contract for `realloc` is forwarded unchanged. + unsafe { System.realloc(ptr, layout, new_size) } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + // SAFETY: the caller's contract for `dealloc` is forwarded unchanged. + unsafe { System.dealloc(ptr, layout) } + } +} + +#[global_allocator] +static GLOBAL: CountingAllocator = CountingAllocator; + +/// Calls before measuring: the first ones pay one-time setup (lazy statics, +/// the HTTP connection pool). +const WARM_UP_CALLS: usize = 3; +/// Calls measured. +const MEASURED_CALLS: usize = 20; +/// The ceiling for one small LLM call. A debug build allocated about 3.4 MB per +/// call with the fix and about 13 MB before it. +const MAX_BYTES_PER_CALL: usize = 8 * 1024 * 1024; +/// The extra bytes per call that 400 unrelated classes may cost: about 22 KB with +/// the fix, about 6.3 MB before it (each call converted every class twice). +const MAX_UNRELATED_GROWTH_PER_CALL: usize = 256 * 1024; + +fn completed_json(text: &str) -> String { + let text = serde_json::to_string(text).expect("response text is JSON-serializable"); + format!( + "{{\"status\":\"completed\",\"output\":[{{\"type\":\"message\",\"role\":\"assistant\",\ + \"content\":[{{\"type\":\"output_text\",\"text\":{text}}}]}}],\ + \"usage\":{{\"input_tokens\":1,\"output_tokens\":1}}}}" + ) +} + +/// A program with `functions` and a Responses client named `Fake` that a mock +/// server answers with `reply`. Keep the server alive while calling. +async fn engine_for(functions: &str, reply: &str) -> (MockServer, Arc) { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/json") + .set_body_string(completed_json(reply)), + ) + .mount(&server) + .await; + let source = format!( + r#" +client Fake = openai.ResponsesClient.new(model = "fake", api_key = "test-key", base_url = "{base_url}"); +{functions} +"#, + base_url = server.uri() + ); + let engine = BexEngine::new_with_runtime_compiler( + compile_source(&source), + Arc::new(sys_ops::SysOps::native()), + Vec::new(), + bex_project::runtime_compiler(), + ) + .expect("engine builds"); + (server, Arc::new(engine)) +} + +async fn call(engine: &Arc, function: &str) -> BexExternalValue { + engine + .call_function( + function, + Vec::new(), + FunctionCallContextBuilder::new(sys_types::CallId::next()).build(), + true, + ) + .await + .expect("the LLM call succeeds") +} + +/// The average bytes allocated per call of `function` (a plain function that +/// makes one LLM call: an LLM function's own hidden parameters are not passed +/// here), after warm-up. +async fn bytes_per_call(engine: &Arc, function: &str) -> usize { + for _ in 0..WARM_UP_CALLS { + call(engine, function).await; + } + let before = ALLOCATED.load(Ordering::Relaxed); + for _ in 0..MEASURED_CALLS { + call(engine, function).await; + } + (ALLOCATED.load(Ordering::Relaxed) - before) / MEASURED_CALLS +} + +/// `count` classes that nothing in the parse target refers to, each with a few +/// fields and a reference to the one before it. +fn unrelated_classes(count: usize) -> String { + (0..count) + .map(|index| { + let previous = if index == 0 { + String::new() + } else { + format!("\n previous Unrelated{}?", index - 1) + }; + format!( + "class Unrelated{index} {{\n name string\n count int\n tags string[]{previous}\n}}\n" + ) + }) + .collect() +} + +const PICK_TOOL: &str = r#" +class Choice { + tool string +} + +function PickTool() -> Choice { + client: Fake + prompt: ` + Pick the tool to run. + ${ctx.output_format()} + ` +} + +function main() -> string { + PickTool().tool +} +"#; + +#[tokio::test] +async fn call_cost_does_not_grow_with_unrelated_types() { + let (_small_server, small) = engine_for(PICK_TOOL, r#"{"tool": "search"}"#).await; + let large_program = format!("{PICK_TOOL}\n{}", unrelated_classes(400)); + let (_large_server, large) = engine_for(&large_program, r#"{"tool": "search"}"#).await; + let small_bytes = bytes_per_call(&small, "main").await; + let large_bytes = bytes_per_call(&large, "main").await; + let growth = large_bytes.saturating_sub(small_bytes); + assert!( + growth <= MAX_UNRELATED_GROWTH_PER_CALL, + "400 unrelated classes added {growth} bytes per call ({small_bytes} -> {large_bytes}); \ + the parse must convert only the types its target reaches" + ); +} + +#[tokio::test] +async fn string_output_call_allocates_little() { + let (_server, engine) = engine_for( + r#" +function SayHi() -> string { + client: Fake + prompt: `Say hi.` +} + +function main() -> string { + SayHi() +} +"#, + "hi", + ) + .await; + assert_eq!( + call(&engine, "main").await, + BexExternalValue::String("hi".to_string().into()) + ); + let bytes = bytes_per_call(&engine, "main").await; + assert!( + bytes <= MAX_BYTES_PER_CALL, + "one LLM call allocated {bytes} bytes (limit {MAX_BYTES_PER_CALL})" + ); +} diff --git a/baml_language/crates/bex_sap/src/lib.rs b/baml_language/crates/bex_sap/src/lib.rs index 025c040aaa5..d3dc4f5cd05 100644 --- a/baml_language/crates/bex_sap/src/lib.rs +++ b/baml_language/crates/bex_sap/src/lib.rs @@ -47,7 +47,7 @@ impl CompiledSapModel { ctx: &::sys_types::SysOpContext, target: sys_types::SapTy, ) -> Result { - let type_ctx = sap_model::TypeCtx::from_sys_op_context(ctx); + let type_ctx = sap_model::TypeCtx::for_target(ctx, &target); Self::from_type_ctx(type_ctx, target) } diff --git a/baml_language/crates/bex_sap/src/sap_model/convert.rs b/baml_language/crates/bex_sap/src/sap_model/convert.rs index 33166f7542d..ba51552ffe9 100644 --- a/baml_language/crates/bex_sap/src/sap_model/convert.rs +++ b/baml_language/crates/bex_sap/src/sap_model/convert.rs @@ -2,10 +2,7 @@ use std::borrow::Cow; -use ::std::{ - collections::{HashMap, HashSet}, - sync::Arc, -}; +use ::std::collections::{HashMap, HashSet}; use ::sys_types::{ClassDefinition, DefKey, EnumDefinition, SapTy}; use indexmap::IndexMap; @@ -50,7 +47,8 @@ pub enum ConvertError { const MAX_RECURSION_DEPTH: usize = 16; -/// Contains stuff from [`sys_types::SysOpContext`] that we need for converting to the sap model. +/// Contains stuff from [`sys_types::SysOpContext`] that we need for converting to the sap model: +/// the definitions one parse target reaches ([`TypeCtx::for_target`]), not the whole program's. /// /// ## Representation /// - Unions should be flattened: @@ -67,17 +65,18 @@ const MAX_RECURSION_DEPTH: usize = 16; #[allow(clippy::struct_field_names)] pub struct TypeCtx { class_definitions: IndexMap, - enum_definitions: Arc>, + enum_definitions: IndexMap, type_alias_definitions: HashMap, sap_parseable: HashMap, } impl TypeCtx { - /// The reason `enum_definitions` is an `Arc` while the others aren't is that - /// we do transformations on the others (so we don't need arc) but we don't - /// need to transform `enum_definitions` so we can just share it. - pub fn new( - class_definitions: &IndexMap, - enum_definitions: Arc>, + /// Builds the context from definitions that are closed under reference: + /// every name a class field or an alias body refers to is present too + /// ([`TypeCtx::for_target`] collects such a set). A missing name makes the + /// types that refer to it unparsable. + pub fn new<'d>( + class_definitions: impl IntoIterator, + enum_definitions: IndexMap, type_alias_definitions: &HashMap, ) -> Self { // todo: we can hold more of this by reference probably @@ -94,7 +93,7 @@ impl TypeCtx { }) .collect(); let class_definitions: IndexMap = class_definitions - .iter() + .into_iter() .map(|(k, v)| { let fields = v.fields.iter().map(|field| { let mut field = field.clone(); @@ -150,18 +149,34 @@ impl TypeCtx { } } - pub fn from_sys_op_context( + /// The context for parsing into `target`: only the classes, enums, and type + /// aliases that `target` reaches ([`reachable_definitions`]), in program + /// order. Converting every definition in the program on each parse cost + /// several MB per LLM call; the reachable set is all the converter looks up. + pub fn for_target( ctx: &::sys_types::SysOpContext, + target: &SapTy, ) -> Self { + let reachable = + reachable_definitions(target, &ctx.class_definitions, &ctx.type_alias_definitions); + let enum_definitions = ctx + .enum_definitions + .iter() + .filter(|(name, _)| reachable.contains(*name)) + .map(|(name, definition)| (name.clone(), definition.clone())) + .collect(); let type_alias_definitions = ctx .type_alias_definitions .iter() + .filter(|(name, _)| reachable.contains(*name)) .map(|(name, ty)| (name.clone(), ty.clone())) .collect(); Self::new( - &ctx.class_definitions, - ctx.enum_definitions.clone(), + ctx.class_definitions + .iter() + .filter(|(name, _)| reachable.contains(*name)), + enum_definitions, &type_alias_definitions, ) } @@ -208,7 +223,7 @@ impl TypeCtx { TyResolved::Class(self.convert_class(name, cls)?), )?; } - for (name, enum_def) in &*self.enum_definitions { + for (name, enum_def) in &self.enum_definitions { add( &mut types, name, @@ -483,6 +498,41 @@ impl TypeCtx { } } +/// The definitions `target` can reach: the classes, interfaces, enums, and type +/// aliases it names, then, transitively, the ones named by a reached class's +/// field types or a reached alias's body. +/// +/// The walk is [`RuntimeTy::visit_heads`](baml_type::RuntimeTy::visit_heads), +/// the head walk `bex_vm::reachable` uses for heap declarations, run over the +/// context's definition tables (which the engine extracts from those same +/// declarations). Heads the converter never resolves, such as type arguments +/// and function types, are kept too. That only adds definitions, so every name +/// the converter looks up is present. +fn reachable_definitions( + target: &SapTy, + class_definitions: &IndexMap, + type_alias_definitions: &IndexMap, +) -> HashSet { + let mut reachable = HashSet::new(); + let mut pending = Vec::new(); + target.visit_heads(&mut |name: &DefKey| pending.push(name.clone())); + while let Some(name) = pending.pop() { + if !reachable.insert(name.clone()) { + continue; + } + if let Some(class) = class_definitions.get(&name) { + for field in &class.fields { + field + .field_type + .visit_heads(&mut |head: &DefKey| pending.push(head.clone())); + } + } else if let Some(alias) = type_alias_definitions.get(&name) { + alias.visit_heads(&mut |head: &DefKey| pending.push(head.clone())); + } + } + reachable +} + fn check_parseable( name: &DefKey, class_definitions: &IndexMap, @@ -582,3 +632,174 @@ fn is_sap_parseable(ty: &SapTy) -> Result, ()> { | SapTy::Type => Err(()), } } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use baml_type::{DeclarationName, Name, typetag::TypeTag}; + use sys_types::{ClassFieldDefinition, EnumVariantDefinition, SysOpContext}; + + use super::*; + + fn key(name: &str) -> DefKey { + DefKey::new( + TypeTag::fresh_dynamic(), + DeclarationName::Anonymous(Name::new(name)), + ) + } + + fn class(key: &DefKey, fields: Vec<(&str, SapTy)>) -> (DefKey, ClassDefinition) { + let fields = fields + .into_iter() + .map(|(name, field_type)| ClassFieldDefinition { + name: name.to_string(), + field_type, + field_template: None, + description: None, + docstring: None, + alias: None, + skip: false, + stream_done: false, + must_exist: false, + }) + .collect(); + let definition = ClassDefinition { + name: key.to_string(), + description: None, + docstring: None, + alias: None, + fields, + stream_done: false, + }; + (key.clone(), definition) + } + + fn enumeration(key: &DefKey) -> (DefKey, EnumDefinition) { + let definition = EnumDefinition { + name: key.to_string(), + description: None, + docstring: None, + alias: None, + variants: vec![EnumVariantDefinition { + name: "Only".to_string(), + description: None, + docstring: None, + alias: None, + }], + }; + (key.clone(), definition) + } + + fn class_ty(key: &DefKey) -> SapTy { + SapTy::Class(key.clone(), Box::new([])) + } + + /// A program with a reachable cluster around `Root` and types nothing in it + /// refers to. `Root -> Child -> Root` is a cycle; `Child` also names an enum + /// and a recursive alias whose body names `Leaf`. + struct Program { + root: DefKey, + child: DefKey, + leaf: DefKey, + color: DefKey, + tree: DefKey, + unrelated: DefKey, + unrelated_enum: DefKey, + ctx: SysOpContext, + } + + fn program() -> Program { + let [root, child, leaf, color, tree, unrelated, unrelated_enum] = [ + "Root", + "Child", + "Leaf", + "Color", + "Tree", + "Unrelated", + "UnrelatedEnum", + ] + .map(key); + let tree_body = + SapTy::union([SapTy::list(SapTy::TypeAlias(tree.clone())), class_ty(&leaf)]); + let mut ctx = SysOpContext::empty(); + ctx.class_definitions = Arc::new(IndexMap::from([ + class(&unrelated, vec![("value", SapTy::string())]), + class(&root, vec![("child", class_ty(&child))]), + class( + &child, + vec![ + ("parent", SapTy::optional(class_ty(&root))), + ("color", SapTy::Enum(color.clone())), + ("tree", SapTy::TypeAlias(tree.clone())), + ], + ), + class(&leaf, vec![("value", SapTy::int())]), + ])); + ctx.enum_definitions = Arc::new(IndexMap::from([ + enumeration(&unrelated_enum), + enumeration(&color), + ])); + ctx.type_alias_definitions = Arc::new(IndexMap::from([(tree.clone(), tree_body)])); + Program { + root, + child, + leaf, + color, + tree, + unrelated, + unrelated_enum, + ctx, + } + } + + fn names<'d, V: 'd>( + definitions: impl IntoIterator, + ) -> Vec<&'d DefKey> { + definitions.into_iter().map(|(name, _)| name).collect() + } + + #[test] + fn primitive_target_keeps_no_definitions() { + let program = program(); + let ctx = TypeCtx::for_target(&program.ctx, &SapTy::string()); + assert!(ctx.class_definitions.is_empty()); + assert!(ctx.enum_definitions.is_empty()); + assert!(ctx.type_alias_definitions.is_empty()); + } + + #[test] + fn class_target_keeps_what_its_fields_reach_in_program_order() { + let program = program(); + let ctx = TypeCtx::for_target(&program.ctx, &class_ty(&program.root)); + + assert_eq!( + names(&ctx.class_definitions), + vec![&program.root, &program.child, &program.leaf], + ); + assert_eq!(names(&ctx.enum_definitions), vec![&program.color]); + assert_eq!(names(&ctx.type_alias_definitions), vec![&program.tree]); + assert!(!ctx.class_definitions.contains_key(&program.unrelated)); + assert!(!ctx.enum_definitions.contains_key(&program.unrelated_enum)); + } + + #[test] + fn alias_target_keeps_its_body() { + let program = program(); + let ctx = TypeCtx::for_target(&program.ctx, &SapTy::TypeAlias(program.tree.clone())); + + assert_eq!(names(&ctx.class_definitions), vec![&program.leaf]); + assert!(ctx.enum_definitions.is_empty()); + assert_eq!(names(&ctx.type_alias_definitions), vec![&program.tree]); + } + + #[test] + fn reachable_definitions_convert_without_unknown_names() { + let program = program(); + let root = class_ty(&program.root); + let ctx = TypeCtx::for_target(&program.ctx, &root); + let db = ctx.build_db().expect("the reachable set converts"); + let target = ctx.convert_ty(&root).expect("the target converts"); + assert!(db.resolve(&target).is_ok()); + } +} From bfc6255a1fd860a2486a018bb18a93bf875711ab Mon Sep 17 00:00:00 2001 From: Aaron Villalpando Date: Fri, 2 Oct 2026 17:03:48 -0700 Subject: [PATCH 2/8] Refute cross-kind type tests before canonicalizing, and tag-test json match arms A runtime match type test that missed canonicalized both sides: heads_definitely_differ decided only same-variant nominal pairs, so a string or a list tested against map expanded the recursive json alias. openai.internal._prune_nulls tests every node of the request JSON that way, so each prompt message cost about 1.56 MB per LLM call. heads_definitely_differ and its interned twin now also refute two heads of different categories. Heads that canonicalization can rewrite (alias, union, projection) or that a context-driven rule relates across kinds (interface, type variable) have no category, nor do the sentinels; a literal and its base, and a variant and its enum, share one. One macro table serves the plain and interned kinds. In MIR, a let x: T arm tests a snapshot of the match scrutinee, so the coarse LIST/MAP tag shortcut never fired for narrowing arms. MatchScrutinee now records the snapshot copies, so json[] and map arms on a json scrutinee test only the runtime tag. A small LLM call drops from about 3.4 MB to 0.7 MB allocated, and a prompt message from 1.56 MB to 35 KB. llm_call_memory.rs bounds both. Co-Authored-By: Claude Opus 5.5 --- .../crates/baml_compiler2_mir/src/lower.rs | 65 +++-- .../baml_src/ns_json_alias/json_alias.baml | 67 ++++++ .../baml_tests/src/compiler2_mir/mod.rs | 85 ++++++- .../baml_tests/tests/llm_call_memory.rs | 59 ++++- .../crates/baml_type/src/normalize.rs | 119 ++++++++-- .../crates/baml_type/src/normalize/tests.rs | 222 +++++++++++++++++- 6 files changed, 565 insertions(+), 52 deletions(-) diff --git a/baml_language/crates/baml_compiler2_mir/src/lower.rs b/baml_language/crates/baml_compiler2_mir/src/lower.rs index 535a9d23390..f15dbba8250 100644 --- a/baml_language/crates/baml_compiler2_mir/src/lower.rs +++ b/baml_language/crates/baml_compiler2_mir/src/lower.rs @@ -805,6 +805,37 @@ pub(crate) fn ty_to_pattern_template_from_resolved_ty(ty: &RuntimeTy) -> Option< } } +/// The value an enclosing `match` scrutinizes, for the coarse-tag proofs. +struct MatchScrutinee { + /// The local the match evaluated its scrutinee into, then each snapshot a + /// narrowing arm (`let x: T`) copies it into and tests in its place. + locals: Vec, + /// The scrutinee's resolved static type. + ty: RuntimeTy, +} + +impl MatchScrutinee { + fn new(local: Local, ty: RuntimeTy) -> Self { + Self { + locals: vec![local], + ty, + } + } + + /// The scrutinee's static type when `local` holds the scrutinee. + fn ty_held_by(&self, local: Local) -> Option<&RuntimeTy> { + self.locals.contains(&local).then_some(&self.ty) + } + + /// Records that `copy` was assigned a copy of `source`: if `source` holds + /// the scrutinee, so does `copy`. + fn record_copy(&mut self, source: Local, copy: Local) { + if self.locals.contains(&source) { + self.locals.push(copy); + } + } +} + /// Flatten a `RuntimeTy` union into its leaf members (recursively), pushing each /// non-union leaf into `out`. A recursive alias is kept as a single `TypeAlias` /// leaf — it is not expanded — matching `ResolvedAliases::convert`. @@ -1693,14 +1724,14 @@ struct LoweringContext<'db> { // baml.ops.Compare`). Saved/restored across nested lambdas. lambda_param_tir_types: FxHashMap, - /// The `(local, resolved static type)` of the value the enclosing `match` is - /// scrutinizing, set for the duration of its arm lowering (save/restored - /// across nested matches). A container arm's runtime type test consults this - /// — guarded on the local matching — to decide whether a coarse `LIST`/`MAP` - /// tag suffices in place of a structural element check (see + /// The value the enclosing `match` is scrutinizing, set for the duration of + /// its arm lowering (save/restored across nested matches). A container + /// arm's runtime type test consults this — guarded on the tested local + /// holding the scrutinee — to decide whether a coarse `LIST`/`MAP` tag + /// suffices in place of a structural element check (see /// [`parametric_arm_tag_sufficient`]). `None` outside a match, and the local /// guard keeps `is` / standalone tests element-precise even inside one. - match_scrutinee: Option<(Local, RuntimeTy)>, + match_scrutinee: Option, /// Stable values read while testing patterns, keyed by pattern identity. /// Binding lowering reuses these exact values instead of rereading mutable @@ -12507,9 +12538,10 @@ impl<'db> LoweringContext<'db> { // Expose the scrutinee's static type to each arm's container type test so // it can decide whether a coarse `LIST`/`MAP` tag suffices; restore the // enclosing match's scrutinee (if any) once the arms are lowered. - let saved_scrutinee = self - .match_scrutinee - .replace((scrutinee_local, self.expr_ty(scrutinee))); + let saved_scrutinee = self.match_scrutinee.replace(MatchScrutinee::new( + scrutinee_local, + self.expr_ty(scrutinee), + )); let switched = self.try_lower_as_switch( scrutinee_local, @@ -12585,8 +12617,8 @@ impl<'db> LoweringContext<'db> { let scrutinee_static_ty: Option = self .match_scrutinee .as_ref() - .filter(|(local, _)| *local == scrutinee) - .map(|(_, ty)| ty.clone()); + .and_then(|match_scrutinee| match_scrutinee.ty_held_by(scrutinee)) + .cloned(); // Classify arms: collect (key, arm_index) for int literal, enum // variant, or type patterns, and check for a trailing wildcard/binding. @@ -13208,14 +13240,14 @@ impl<'db> LoweringContext<'db> { // If this is a parametric arm whose type arguments the emitter would // otherwise check structurally, but the enclosing match's scrutinee // statically proves a coarse `LIST`/`MAP`/`FUTURE` tag suffices, - // emit the tag test directly. Guarded on the scrutinee local so an + // emit the tag test directly. Guarded on the tested local holding + // the scrutinee (itself, or a narrowing arm's snapshot of it) so an // `is` / nested test against a different value stays arg-precise. if self .match_scrutinee .as_ref() - .is_some_and(|(local, scrutinee_ty)| { - *local == scrutinee && parametric_arm_tag_sufficient(&ty, scrutinee_ty) - }) + .and_then(|match_scrutinee| match_scrutinee.ty_held_by(scrutinee)) + .is_some_and(|scrutinee_ty| parametric_arm_tag_sufficient(&ty, scrutinee_ty)) { let tag = match &ty { RuntimeTy::List(..) => baml_type::typetag::LIST, @@ -13770,6 +13802,9 @@ impl<'db> LoweringContext<'db> { Place::local(snapshot), Rvalue::Use(Operand::Copy(Place::Local(scrutinee))), ); + if let Some(match_scrutinee) = self.match_scrutinee.as_mut() { + match_scrutinee.record_copy(scrutinee, snapshot); + } self.tested_pattern_values .insert(self.pat_metadata_key(pat_id), snapshot); snapshot diff --git a/baml_language/crates/baml_tests/baml_src/ns_json_alias/json_alias.baml b/baml_language/crates/baml_tests/baml_src/ns_json_alias/json_alias.baml index 5144d1e6726..64b5f55a508 100644 --- a/baml_language/crates/baml_tests/baml_src/ns_json_alias/json_alias.baml +++ b/baml_language/crates/baml_tests/baml_src/ns_json_alias/json_alias.baml @@ -137,3 +137,70 @@ function json_wrap_tag() -> string { test "impl_on_recursive_alias_generic_arg_dispatches" { assert.equal(json_wrap_tag(), "json-wrap") } + +// ─── 1.9 Container arms on a json scrutinee ────────────────────────────────── + +// `json[]` and `map` are json's only list and map members, so a +// match on a json value tests these arms by the runtime tag alone. Every kind +// of json value must still reach its own arm, also through recursion (the +// request builders prune nulls this way). + +function json_kind(j: json) -> string { + match (j) { + let object: map => "map", + let items: json[] => "list", + _ => "scalar", + } +} + +function json_kinds() -> string[] { + [ + json_kind({"a": 1}), + json_kind({}), + json_kind([1, "two"]), + json_kind([]), + json_kind("map"), + json_kind(1), + json_kind(true), + json_kind(null), + ] +} + +test "json_container_arms_match_each_kind" { + assert.equal(json_kinds(), ["map", "map", "list", "list", "scalar", "scalar", "scalar", "scalar"]) +} + +function drop_nulls(j: json) -> json { + match (j) { + let object: map => { + let out: map = {}; + for (let key in object.keys()) { + let entry = object.get(key); + match (entry) { + null => null, + _ => { + let _ = out.set(key, drop_nulls(entry)); + null + }, + }; + } + out + }, + let items: json[] => { + let out: json[] = []; + for (let item in items) { + out.push(drop_nulls(item)); + } + out + }, + _ => j, + } +} + +function drop_nulls_nested() -> string { + baml.json.stringify(drop_nulls(baml.json.parse("{\"a\":null,\"b\":[1,{\"c\":null,\"d\":\"x\"}],\"e\":{\"f\":null}}"))) +} + +test "json_container_arms_recurse_through_nested_values" { + assert.equal(drop_nulls_nested(), "{\"b\":[1,{\"d\":\"x\"}],\"e\":{}}") +} diff --git a/baml_language/crates/baml_tests/src/compiler2_mir/mod.rs b/baml_language/crates/baml_tests/src/compiler2_mir/mod.rs index e8d7dc6cbcb..3740955607f 100644 --- a/baml_language/crates/baml_tests/src/compiler2_mir/mod.rs +++ b/baml_language/crates/baml_tests/src/compiler2_mir/mod.rs @@ -23,6 +23,25 @@ fn make_db() -> ProjectDatabase { db } +/// The rendered MIR of the user function `name`, out of [`render_mir`]'s output. +fn function_mir<'a>(output: &'a str, name: &str) -> &'a str { + output + .split(&format!("fn user.{name}(")) + .nth(1) + .and_then(|tail| tail.split("\nfn ").next()) + .unwrap_or_else(|| panic!("no MIR for `user.{name}`:\n{output}")) +} + +/// Whether a line of MIR is a structural type test against a list or map type: +/// an `is_type` (which renders the type, `int[]`) or a `narrow_bind` (which +/// renders the type test, `List(Int)`). +fn is_structural_container_test(line: &str) -> bool { + (line.contains("is_type(") || line.contains("narrow_bind")) + && ["[]", "map<", "List(", "Map {"] + .iter() + .any(|spelling| line.contains(spelling)) +} + /// Lower all functions in a file to MIR and pretty-print them. fn render_mir(db: &ProjectDatabase, file: baml_base::SourceFile) -> String { // Dump in source order (by declaration span) — intrinsic and @@ -425,11 +444,7 @@ fn generic_class_destructure_field_projection_uses_instantiated_type() { // Scope the check to `user.f`'s body — auto-derived `to_json` / // `from_json` methods on `Box` legitimately have `void` locals // because T isn't instantiated in the auto-derive body. - let f_body = output - .split("fn user.f(") - .nth(1) - .and_then(|tail| tail.split("\nfn ").next()) - .unwrap_or(&output); + let f_body = function_mir(&output, "f"); assert!( !f_body .lines() @@ -438,6 +453,66 @@ fn generic_class_destructure_field_projection_uses_instantiated_type() { ); } +/// `json[]` and `map` are a `json` scrutinee's only list and map +/// members, so arms of exactly those types test only the runtime tag, also +/// when they bind the value (a narrowing arm tests a snapshot of the +/// scrutinee). +#[test] +fn json_container_arms_test_only_the_runtime_tag() { + let mut db = make_db(); + let file = db.file( + "test.baml", + r#" + function size(value: json) -> int { + match (value) { + let object: map => object.length(), + let items: json[] => items.length(), + _ => 0, + } + } + "#, + ); + let output = render_mir(&db, file); + let body = function_mir(&output, "size"); + assert!( + body.contains(", MAP);") && body.contains(", LIST);"), + "the json container arms do not test the runtime tag:\n{body}" + ); + assert!( + !body.lines().any(is_structural_container_test), + "a json container arm kept its structural test:\n{body}" + ); +} + +/// An element-specific arm stays element-precise on a `json` scrutinee: its +/// list member is `json[]`, so the `LIST` tag alone cannot prove `int[]`. +#[test] +fn json_scrutinee_keeps_element_specific_arms_structural() { + let mut db = make_db(); + let file = db.file( + "test.baml", + r#" + function ints(value: json) -> int { + match (value) { + let xs: int[] => xs.length(), + _ => 0, + } + } + "#, + ); + let output = render_mir(&db, file); + let body = function_mir(&output, "ints"); + assert!( + body.lines().any(|line| is_structural_container_test(line) + && (line.contains("int[]") || line.contains("List(Int)"))), + "the `int[]` arm lost its structural test:\n{body}" + ); + assert!( + !body.contains(", LIST);"), + "the `int[]` arm tests only the LIST tag:\n{body}" + ); +} + #[test] fn match_or_mixed_array_class_binding_uses_branch_local_rest_type() { let mut db = make_db(); diff --git a/baml_language/crates/baml_tests/tests/llm_call_memory.rs b/baml_language/crates/baml_tests/tests/llm_call_memory.rs index 0c618258d7b..3943c1d986b 100644 --- a/baml_language/crates/baml_tests/tests/llm_call_memory.rs +++ b/baml_language/crates/baml_tests/tests/llm_call_memory.rs @@ -3,9 +3,11 @@ //! A Rust test because BAML cannot observe allocations. Each parse used to build //! a schema-aligned-parsing model of every type in the program, the standard //! library included: about 5 MB per call for a one-function program, kept alive -//! until the VM collected it. A counting global allocator measures every byte -//! allocated while the calls run, on every thread, so the bound does not depend -//! on when the GC runs. +//! until the VM collected it. And every `match` type test that failed +//! canonicalized the recursive `json` alias: the request builder tests each +//! node of the request JSON that way, so the cost grew with the conversation. +//! A counting global allocator measures every byte allocated while the calls +//! run, on every thread, so the bounds do not depend on when the GC runs. #![allow( unsafe_code, @@ -68,12 +70,19 @@ static GLOBAL: CountingAllocator = CountingAllocator; const WARM_UP_CALLS: usize = 3; /// Calls measured. const MEASURED_CALLS: usize = 20; -/// The ceiling for one small LLM call. A debug build allocated about 3.4 MB per -/// call with the fix and about 13 MB before it. -const MAX_BYTES_PER_CALL: usize = 8 * 1024 * 1024; +/// The ceiling for one small LLM call. A debug build allocates about 0.7 MB; it +/// allocated about 3.4 MB while failed type tests expanded `json`, and about +/// 13 MB while each parse converted the whole program. +const MAX_BYTES_PER_CALL: usize = 2 * 1024 * 1024; /// The extra bytes per call that 400 unrelated classes may cost: about 22 KB with /// the fix, about 6.3 MB before it (each call converted every class twice). const MAX_UNRELATED_GROWTH_PER_CALL: usize = 256 * 1024; +/// User/assistant exchanges in the long conversation. +const LONG_CONVERSATION_TURNS: usize = 10; +/// The extra bytes per call that one more prompt message may cost: about 35 KB, +/// and about 1.56 MB while every failed type test on the request JSON expanded +/// `json`. +const MAX_BYTES_PER_MESSAGE: usize = 128 * 1024; fn completed_json(text: &str) -> String { let text = serde_json::to_string(text).expect("response text is JSON-serializable"); @@ -157,6 +166,28 @@ fn unrelated_classes(count: usize) -> String { .collect() } +/// A chat program whose prompt holds `turns` user/assistant exchanges before +/// the last question. +fn chat_program(turns: usize) -> String { + let history: String = (0..turns) + .map(|turn| { + format!("${{role(\"user\")}}Question {turn}?${{role(\"assistant\")}}Answer {turn}.") + }) + .collect(); + format!( + r#" +function Chat() -> string {{ + client: Fake + prompt: `${{role("system")}}Be concise.{history}${{role("user")}}Say hi.` +}} + +function main() -> string {{ + Chat() +}} +"# + ) +} + const PICK_TOOL: &str = r#" class Choice { tool string @@ -216,3 +247,19 @@ function main() -> string { "one LLM call allocated {bytes} bytes (limit {MAX_BYTES_PER_CALL})" ); } + +#[tokio::test] +async fn call_cost_does_not_grow_with_conversation_length() { + let (_short_server, short) = engine_for(&chat_program(0), "hi").await; + let (_long_server, long) = engine_for(&chat_program(LONG_CONVERSATION_TURNS), "hi").await; + let short_bytes = bytes_per_call(&short, "main").await; + let long_bytes = bytes_per_call(&long, "main").await; + let extra_messages = 2 * LONG_CONVERSATION_TURNS; + let per_message = long_bytes.saturating_sub(short_bytes) / extra_messages; + assert!( + per_message <= MAX_BYTES_PER_MESSAGE, + "each prompt message added {per_message} bytes per call ({short_bytes} -> {long_bytes} \ + for {extra_messages} more messages); building the request must stay linear in its size \ + with a small constant" + ); +} diff --git a/baml_language/crates/baml_type/src/normalize.rs b/baml_language/crates/baml_type/src/normalize.rs index 9a82584a24e..6b545217525 100644 --- a/baml_language/crates/baml_type/src/normalize.rs +++ b/baml_language/crates/baml_type/src/normalize.rs @@ -414,17 +414,18 @@ pub trait TypeContext { if sub == sup { return true; } - // Narrow definite-mismatch filter. Unlike `equivalent`, differing heads - // do NOT generally refute subtyping (a literal is a subtype of its base, - // a member of its union, a class of an interface it implements) — but - // those pairs are different `Ty` *variants*, for which - // `heads_definitely_differ` already answers `false`. Of the same-variant - // nominal pairs it does claim, only interface-to-interface can still be - // a subtype under different names (`A <: B` iff `A` requires `B`), so it - // is excluded. For the rest the canonical walk below is a foregone - // `false`: class heads are preserved by canonicalization and the class - // subtype rule requires equal names (generic args are invariant), while - // enums and enum variants are nominal with no cross-name rule at all. + // Definite-mismatch filter. Unlike `equivalent`, differing heads do NOT + // generally refute subtyping (a literal is a subtype of its base, a + // member of its union, a class of an interface it implements) — but + // `heads_definitely_differ` claims none of those pairs: a literal and + // its base share a category, and unions and interfaces have none. Of + // the same-variant nominal pairs it does claim, only + // interface-to-interface can still be a subtype under different names + // (`A <: B` iff `A` requires `B`), so it is excluded. For the rest the + // canonical walk below is a foregone `false`: class heads are preserved + // by canonicalization and the class subtype rule requires equal names + // (generic args are invariant), enums and enum variants are nominal + // with no cross-name rule at all, and no rule relates two categories. if !matches!((sub, sup), (Ty::Interface(..), Ty::Interface(..))) && heads_definitely_differ(sub, sup) { @@ -624,24 +625,34 @@ pub fn is_subtype>(sub: &Ty, sup: &Ty, ctx: &C) /// are rewritten), and equality is nominal, so two of the same kind with /// different names can never share a canonical form. /// -/// This is deliberately *conservative* — it decides only same-kind pairs whose -/// heads are stable and never uses a cross-kind "different discriminant ⇒ -/// differ" catch-all. That catch-all would be unsound under the current -/// normalization: `TypeAlias` / `Union` / `AssociatedTypeProjection` / generic `Media` / -/// `Infer` heads can be rewritten into any other shape. Leaving every such case -/// undecided (`false`) preserves correctness; only the unambiguous nominal -/// misses are fast-rejected. Context-independent: canonicalization preserves -/// these nominal heads regardless of the `TypeContext`. +/// Across kinds it decides only pairs whose heads both have a [`Category`] +/// (see [`Category::of_head`]), never with a "different discriminant ⇒ differ" +/// catch-all. That catch-all would be unsound: `TypeAlias` / `Union` / +/// `AssociatedTypeProjection` heads can be rewritten into any other shape, so +/// they have no category and stay undecided (`false`). Canonicalization keeps +/// every head that has a category, so two different categories never share a +/// canonical form — and no subtyping rule relates them either: no concrete type +/// is a subtype of another (TYPE_SYSTEM.md §Concrete Types), and a literal and +/// its base, or a variant and its enum, share a category. This is the common +/// miss of a runtime `match` on a `json` value: a string or a list tested +/// against `map` would otherwise expand the recursive alias on +/// every test. Context-independent: canonicalization preserves these heads +/// regardless of the `TypeContext`. fn heads_definitely_differ(a: &Ty, b: &Ty) -> bool { match (a, b) { (Ty::Class(q1, ..), Ty::Class(q2, ..)) | (Ty::Interface(q1, ..), Ty::Interface(q2, ..)) | (Ty::Enum(q1, ..), Ty::Enum(q2, ..)) => q1 != q2, (Ty::EnumVariant(q1, v1, ..), Ty::EnumVariant(q2, v2, ..)) => q1 != q2 || v1 != v2, - _ => false, + _ => categories_differ(Category::of_head(a), Category::of_head(b)), } } +/// Whether two heads are of different, known categories. +fn categories_differ(a: Option, b: Option) -> bool { + matches!((a, b), (Some(a), Some(b)) if a != b) +} + /// Free-function form of [`TypeContext::definitely_disjoint`], for a context held /// by value. Pending removal once every caller uses the method form. pub fn definitely_disjoint>(a: &Ty, b: &Ty, ctx: &C) -> bool { @@ -773,10 +784,11 @@ impl NormalTy { // CONCRETE-TYPE DISJOINTNESS // ═══════════════════════════════════════════════════════════════════════════ -/// Top-level concrete category of a ground-headed type, used only for the -/// cross-category check (a value of one category is never a value of another). -/// Same-category pairs are decided by the structural arms of -/// [`NormalTy::is_disjoint_from`], not by this. +/// Top-level concrete category of a ground-headed type, for the cross-category +/// checks: a value of one category is never a value of another +/// ([`NormalTy::is_disjoint_from`]), and no subtyping rule or canonical form +/// relates two categories ([`heads_definitely_differ`]). Same-category pairs +/// are decided by the structural arms, not by this. #[derive(PartialEq, Eq)] enum Category { Int, @@ -1017,6 +1029,47 @@ impl NormalTy { } } +/// The one table behind [`Category::of_head`] and +/// [`Category::of_interned_head`]. Both are `ty_family!` members with the same +/// variant names, so `$member` is the enum's path and `$none` adds the member's own +/// variants that have no category. Exhaustive on purpose: a new variant must +/// be classified here. +macro_rules! family_head_category { + ($ty:expr, $($member:ident)::+ $(, $none:pat)*) => { + Some(match $ty { + $($member)::+::Int => Category::Int, + $($member)::+::Bigint => Category::Bigint, + $($member)::+::Float => Category::Float, + $($member)::+::String => Category::String, + $($member)::+::Bool => Category::Bool, + $($member)::+::Literal(lit, _) => Category::of_literal(lit), + $($member)::+::Null => Category::Null, + $($member)::+::Uint8Array => Category::Uint8Array, + $($member)::+::Media(kind) => Category::Media(*kind), + $($member)::+::Void => Category::Void, + $($member)::+::RustType => Category::RustType, + $($member)::+::Type => Category::Type, + $($member)::+::Resource => Category::Resource, + $($member)::+::PromptAst => Category::PromptAst, + $($member)::+::Class(..) => Category::Class, + $($member)::+::List(_) => Category::List, + $($member)::+::Map { .. } => Category::Map, + $($member)::+::Enum(_) | $($member)::+::EnumVariant(..) => Category::Enum, + $($member)::+::Function { .. } => Category::Function, + $($member)::+::Future(..) => Category::Future, + $($member)::+::Interface(..) + | $($member)::+::Union(_) + | $($member)::+::TypeAlias(_) + | $($member)::+::TypeVar(_) + | $($member)::+::AssociatedTypeProjection { .. } + | $($member)::+::Unknown + | $($member)::+::Never + | $($member)::+::Error + $(| $none)* => return None, + }) + }; +} + impl Category { fn of_literal(lit: &Literal) -> Category { match lit { @@ -1027,6 +1080,22 @@ impl Category { Literal::Bool(_) => Category::Bool, } } + + /// The category of a head that canonicalization keeps, read before + /// canonicalizing: the [`Ty`] twin of [`NormalTy::head_category`]. `None` + /// for a head canonicalization can rewrite into another kind (an alias, a + /// union, a projection), one a context-driven rule relates across kinds + /// (an interface, a type variable), and the sentinels (`unknown`, `never`, + /// the error type). + fn of_head(ty: &Ty) -> Option { + family_head_category!(ty, Ty) + } + + /// [`Category::of_head`] for the interned kind. A live inference variable + /// can still become any type, so it has no category. + fn of_interned_head(ty: &interned::InferTy) -> Option { + family_head_category!(ty, interned::InferTy, interned::InferTy::InferVar { .. }) + } } // ═══════════════════════════════════════════════════════════════════════════ @@ -3166,7 +3235,7 @@ fn interned_heads_definitely_differ(a: &interned::Ty, b: &interned::Ty) -> bool | (K::Interface(q1, ..), K::Interface(q2, ..)) | (K::Enum(q1, ..), K::Enum(q2, ..)) => q1 != q2, (K::EnumVariant(q1, v1, ..), K::EnumVariant(q2, v2, ..)) => q1 != q2 || v1 != v2, - _ => false, + (a, b) => categories_differ(Category::of_interned_head(a), Category::of_interned_head(b)), } } diff --git a/baml_language/crates/baml_type/src/normalize/tests.rs b/baml_language/crates/baml_type/src/normalize/tests.rs index 09a3f303a75..b69f65ff0c0 100644 --- a/baml_language/crates/baml_type/src/normalize/tests.rs +++ b/baml_language/crates/baml_type/src/normalize/tests.rs @@ -1,7 +1,7 @@ //! Unit tests for the canonical type algebra, driven by a hand-built //! [`TypeContext`] stub that records nominal facts directly. -use std::collections::HashMap; +use std::{cell::Cell, collections::HashMap}; use super::*; use crate::{DeclName, Freshness, FunctionParamTy, Literal, Name, Ty}; @@ -25,6 +25,9 @@ struct Ctx { /// `(base, interface head, member) → reduced type` projection facts, for the /// `project` oracle. A `Vec` (not a map) because `Ty` is not `Hash`. projections: Vec<(Ty, DeclName, Name, Ty)>, + /// How many times the algebra asked for an alias definition: a refutation + /// that skips canonicalization never asks. + alias_lookups: Cell, } fn nominal_head(ty: &Ty) -> Option { @@ -55,6 +58,7 @@ impl TypeContext for Ctx { } fn alias_def(&self, name: &DeclName) -> Option { + self.alias_lookups.set(self.alias_lookups.get() + 1); self.aliases.get(name).cloned() } @@ -1457,6 +1461,166 @@ fn any_function_existentials_are_covariant_in_their_pins() { assert!(!is_subtype(&other, &any_function(vec![]), &ctx)); } +// ── head-category refutation ─────────────────────────────────────────────-- + +/// `type Json = null | bool | int | float | string | Json[] | map` +/// registered under the given name: the recursive alias whose canonical form +/// is expensive to build. +fn json_alias(ctx: &mut Ctx, name: &str) { + ctx.aliases.insert( + qtn(name), + union(vec![ + Ty::null(), + Ty::bool(), + Ty::int(), + Ty::float(), + Ty::string(), + list(alias(name)), + map_ty(Ty::string(), alias(name)), + ]), + ); +} + +/// One type of every `Ty` variant, plus the facts that make the context-driven +/// rules fire: `Box` and `int` implement `Displayable`, `T: Displayable`, +/// `(Box as Displayable).Item` reduces to `int`, and `Json` is recursive. +fn kind_corpus() -> (Ctx, Vec) { + let mut ctx = Ctx::default(); + json_alias(&mut ctx, "Json"); + ctx.enums + .insert(qtn("Side"), vec![Name::new("L"), Name::new("R")]); + ctx.impls.push((qtn("Box"), qtn("Displayable"))); + ctx.prim_impls.push(("int", qtn("Displayable"))); + ctx.var_bounds + .insert(param(0, "T"), vec![iface("Displayable")]); + ctx.projections.push(( + class("Box"), + qtn("Displayable"), + Name::new("Item"), + Ty::int(), + )); + let corpus = vec![ + Ty::int(), + bigint(), + Ty::float(), + Ty::string(), + Ty::bool(), + Ty::null(), + Ty::Uint8Array, + Ty::Media(MediaKind::Image), + Ty::Media(MediaKind::Generic), + Ty::Void, + Ty::RustType, + Ty::Type, + Ty::Resource, + Ty::PromptAst, + lit_int(1), + lit_float("1.5"), + lit_str("user"), + lit_bool(true), + class("Box"), + class1("Gen", Ty::int()), + enum_ty("Side"), + variant("Side", "L"), + list(Ty::int()), + list(alias("Json")), + map_ty(Ty::string(), alias("Json")), + map_ty(Ty::string(), Ty::int()), + Ty::Function { + params: Box::new([FunctionParamTy::required(None, Ty::int())]), + ret: Box::new(Ty::bool()), + throws: Box::new(Ty::Never), + }, + Ty::Future(Box::new(Ty::int()), Box::new(Ty::Never)), + union(vec![Ty::int(), Ty::string()]), + Ty::optional(Ty::string()), + alias("Json"), + iface("Displayable"), + typevar(0, "T"), + projection(class("Box"), "Displayable", "Item"), + Ty::Unknown, + Ty::Never, + Ty::Error, + ]; + (ctx, corpus) +} + +/// The head-category refutation is only a shortcut: over every pair of +/// variants, `is_subtype` and `equivalent` give the verdict the canonical +/// relation gives. +#[test] +fn subtype_and_equivalence_agree_with_the_canonical_relation_for_every_kind_pair() { + let (ctx, corpus) = kind_corpus(); + for sub in &corpus { + for sup in &corpus { + let canonical_sub = NormalTy::canonical(sub, &ctx); + let canonical_sup = NormalTy::canonical(sup, &ctx); + assert_eq!( + ctx.is_subtype(sub, sup), + canonical_sub.is_subtype_of(&canonical_sup, &ctx, &mut HashSet::new()), + "subtype verdict diverged for {sub:?} <: {sup:?}" + ); + assert_eq!( + ctx.equivalent(sub, sup), + canonical_sub == canonical_sup, + "equivalence verdict diverged for {sub:?} == {sup:?}" + ); + } + } +} + +/// Pairs whose top-level kinds no rule relates, the ones a runtime `match` on +/// a `json` value asks about for every node: none of them may canonicalize +/// (and so expand) the recursive alias. +fn cross_category_pairs() -> Vec<(Ty, Ty)> { + let json_map = map_ty(Ty::string(), alias("Json")); + let json_list = list(alias("Json")); + vec![ + (lit_str("user"), json_map.clone()), + (lit_str("user"), json_list.clone()), + (json_list.clone(), json_map.clone()), + (json_map.clone(), json_list.clone()), + (Ty::int(), json_map.clone()), + (lit_int(1), Ty::float()), + (Ty::null(), json_list), + (class("Box"), json_map.clone()), + (enum_ty("Side"), json_map), + ] +} + +#[test] +fn cross_category_pairs_are_refuted_without_canonicalizing() { + let mut ctx = Ctx::default(); + json_alias(&mut ctx, "Json"); + for (sub, sup) in cross_category_pairs() { + assert!(!ctx.is_subtype(&sub, &sup), "{sub:?} <: {sup:?}"); + assert!(!ctx.equivalent(&sub, &sup), "{sub:?} == {sup:?}"); + assert_eq!( + ctx.alias_lookups.get(), + 0, + "{sub:?} vs {sup:?} canonicalized an alias" + ); + } +} + +/// Kinds a canonical form can rewrite, or that a context-driven rule relates +/// across kinds, are never refuted by their head alone. +#[test] +fn rewritable_and_context_driven_kinds_are_not_refuted() { + let (ctx, _) = kind_corpus(); + let json_map = map_ty(Ty::string(), alias("Json")); + assert!(ctx.is_subtype(&lit_str("user"), &alias("Json"))); + assert!(ctx.is_subtype(&json_map, &alias("Json"))); + assert!(ctx.is_subtype(&Ty::int(), &iface("Displayable"))); + assert!(ctx.is_subtype(&class("Box"), &iface("Displayable"))); + assert!(ctx.is_subtype(&typevar(0, "T"), &iface("Displayable"))); + assert!(ctx.is_subtype(&projection(class("Box"), "Displayable", "Item"), &Ty::int())); + assert!(ctx.is_subtype(&Ty::Never, &json_map)); + assert!(ctx.is_subtype(&json_map, &Ty::Unknown)); + assert!(ctx.is_subtype(&Ty::Error, &json_map)); + assert!(ctx.is_subtype(&lit_str("user"), &Ty::optional(Ty::string()))); +} + // ── interned entry (S4b) ─────────────────────────────────────────────────── mod interned_entry { @@ -1524,6 +1688,62 @@ mod interned_entry { )); } + /// The interned entries give the canonical verdict for every kind pair too. + #[test] + fn interned_verdicts_agree_with_the_canonical_relation_for_every_kind_pair() { + let (ctx, corpus) = kind_corpus(); + let corpus: Vec<_> = corpus.iter().map(it).collect(); + let cache = InternedCanonicalCache::default(); + for sub in &corpus { + for sup in &corpus { + let canonical_sub = NormalTy::canonical_interned(sub, &ctx); + let canonical_sup = NormalTy::canonical_interned(sup, &ctx); + let subtype = + canonical_sub.is_subtype_of(&canonical_sup, &ctx, &mut HashSet::new()); + let equivalent = canonical_sub == canonical_sup; + assert_eq!( + is_subtype_interned(sub, sup, &ctx), + subtype, + "{sub:?} <: {sup:?}" + ); + assert_eq!( + cache.is_subtype(sub, sup, &ctx), + subtype, + "{sub:?} <: {sup:?}" + ); + assert_eq!( + equivalent_interned(sub, sup, &ctx), + equivalent, + "{sub:?} == {sup:?}" + ); + assert_eq!( + cache.equivalent(sub, sup, &ctx), + equivalent, + "{sub:?} == {sup:?}" + ); + } + } + } + + #[test] + fn interned_cross_category_pairs_are_refuted_without_canonicalizing() { + let mut ctx = Ctx::default(); + json_alias(&mut ctx, "Json"); + let cache = InternedCanonicalCache::default(); + for (sub, sup) in cross_category_pairs() { + let (sub, sup) = (it(&sub), it(&sup)); + assert!(!is_subtype_interned(&sub, &sup, &ctx), "{sub:?} <: {sup:?}"); + assert!(!equivalent_interned(&sub, &sup, &ctx), "{sub:?} == {sup:?}"); + assert!(!cache.is_subtype(&sub, &sup, &ctx), "{sub:?} <: {sup:?}"); + assert!(!cache.equivalent(&sub, &sup, &ctx), "{sub:?} == {sup:?}"); + assert_eq!( + ctx.alias_lookups.get(), + 0, + "{sub:?} vs {sup:?} canonicalized an alias" + ); + } + } + #[test] fn canonical_cache_matches_reject_free_canonical_relations() { let mut ctx = Ctx::default(); From c2064d1b66b6c762bd421980df3266e115f51c83 Mon Sep 17 00:00:00 2001 From: Aaron Villalpando Date: Fri, 2 Oct 2026 17:12:25 -0700 Subject: [PATCH 3/8] Serialize the allocation-counting memory tests ALLOCATED counts every thread, and cargo test runs a binary's tests on parallel threads, so one test's bytes could land in another test's measurement and fail it at random. Each test now holds a process-wide lock for its whole run. nextest gives each test its own process, so CI never showed it. Co-Authored-By: Claude Opus 5.5 --- baml_language/crates/baml_tests/Cargo.toml | 2 +- baml_language/crates/baml_tests/tests/llm_call_memory.rs | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/baml_language/crates/baml_tests/Cargo.toml b/baml_language/crates/baml_tests/Cargo.toml index 13eee8dc04d..bac066ff219 100644 --- a/baml_language/crates/baml_tests/Cargo.toml +++ b/baml_language/crates/baml_tests/Cargo.toml @@ -64,7 +64,7 @@ num-bigint = { workspace = true } rayon = { workspace = true } serde_json = { workspace = true } tempfile = { workspace = true } -tokio = { workspace = true, features = [ "rt-multi-thread", "macros", "net" ] } +tokio = { workspace = true, features = [ "rt-multi-thread", "macros", "net", "sync" ] } walkdir = { workspace = true } wiremock = { workspace = true } diff --git a/baml_language/crates/baml_tests/tests/llm_call_memory.rs b/baml_language/crates/baml_tests/tests/llm_call_memory.rs index 0c618258d7b..35fe12c4b49 100644 --- a/baml_language/crates/baml_tests/tests/llm_call_memory.rs +++ b/baml_language/crates/baml_tests/tests/llm_call_memory.rs @@ -31,6 +31,12 @@ use wiremock::{ /// Every byte handed out since the process started. static ALLOCATED: AtomicUsize = AtomicUsize::new(0); +/// Held for the whole of each test. `ALLOCATED` counts every thread, so a test +/// running beside a measurement (`cargo test` runs this binary's tests on +/// parallel threads; nextest gives each its own process) would add its bytes +/// to that measurement. +static MEASURING: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + struct CountingAllocator; // SAFETY: every method forwards its arguments unchanged to `System`, which @@ -177,6 +183,7 @@ function main() -> string { #[tokio::test] async fn call_cost_does_not_grow_with_unrelated_types() { + let _measuring = MEASURING.lock().await; let (_small_server, small) = engine_for(PICK_TOOL, r#"{"tool": "search"}"#).await; let large_program = format!("{PICK_TOOL}\n{}", unrelated_classes(400)); let (_large_server, large) = engine_for(&large_program, r#"{"tool": "search"}"#).await; @@ -192,6 +199,7 @@ async fn call_cost_does_not_grow_with_unrelated_types() { #[tokio::test] async fn string_output_call_allocates_little() { + let _measuring = MEASURING.lock().await; let (_server, engine) = engine_for( r#" function SayHi() -> string { From 5132085943935a0961643c9fcbd5ca635dcc4b63 Mon Sep 17 00:00:00 2001 From: Aaron Villalpando Date: Fri, 2 Oct 2026 17:13:31 -0700 Subject: [PATCH 4/8] Serialize the conversation-length memory test too It measures the shared ALLOCATED counter like the other two tests in the file, so it takes the same lock. Co-Authored-By: Claude Opus 5.5 --- baml_language/crates/baml_tests/tests/llm_call_memory.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/baml_language/crates/baml_tests/tests/llm_call_memory.rs b/baml_language/crates/baml_tests/tests/llm_call_memory.rs index 5a25bd512e6..955f1a52e83 100644 --- a/baml_language/crates/baml_tests/tests/llm_call_memory.rs +++ b/baml_language/crates/baml_tests/tests/llm_call_memory.rs @@ -258,6 +258,7 @@ function main() -> string { #[tokio::test] async fn call_cost_does_not_grow_with_conversation_length() { + let _measuring = MEASURING.lock().await; let (_short_server, short) = engine_for(&chat_program(0), "hi").await; let (_long_server, long) = engine_for(&chat_program(LONG_CONVERSATION_TURNS), "hi").await; let short_bytes = bytes_per_call(&short, "main").await; From f1532838ce078ebd9074e95a993136f2bef74c8b Mon Sep 17 00:00:00 2001 From: Aaron Villalpando Date: Fri, 2 Oct 2026 22:59:14 -0700 Subject: [PATCH 5/8] Keep a public doc from linking to a private item cargo doc with -D warnings rejects the intra-doc link from TypeCtx::for_target to the private reachable_definitions. Co-Authored-By: Claude Opus 5.5 --- baml_language/crates/bex_sap/src/sap_model/convert.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/baml_language/crates/bex_sap/src/sap_model/convert.rs b/baml_language/crates/bex_sap/src/sap_model/convert.rs index ba51552ffe9..ace0511150c 100644 --- a/baml_language/crates/bex_sap/src/sap_model/convert.rs +++ b/baml_language/crates/bex_sap/src/sap_model/convert.rs @@ -150,7 +150,7 @@ impl TypeCtx { } /// The context for parsing into `target`: only the classes, enums, and type - /// aliases that `target` reaches ([`reachable_definitions`]), in program + /// aliases that `target` reaches (`reachable_definitions`), in program /// order. Converting every definition in the program on each parse cost /// several MB per LLM call; the reachable set is all the converter looks up. pub fn for_target( From 0154a9827653a654e803952d651818dc7c3b1a31 Mon Sep 17 00:00:00 2001 From: Aaron Villalpando Date: Fri, 2 Oct 2026 23:11:35 -0700 Subject: [PATCH 6/8] Leave skipped fields out of the parse model's reachable set The parser drops a @skip field, so a type only a skipped field names never needs converting. Reported by CodeRabbit. Co-Authored-By: Claude Opus 5.5 --- .../crates/bex_sap/src/sap_model/convert.rs | 26 +++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/baml_language/crates/bex_sap/src/sap_model/convert.rs b/baml_language/crates/bex_sap/src/sap_model/convert.rs index ace0511150c..d9ef52c1aed 100644 --- a/baml_language/crates/bex_sap/src/sap_model/convert.rs +++ b/baml_language/crates/bex_sap/src/sap_model/convert.rs @@ -500,7 +500,8 @@ impl TypeCtx { /// The definitions `target` can reach: the classes, interfaces, enums, and type /// aliases it names, then, transitively, the ones named by a reached class's -/// field types or a reached alias's body. +/// field types (except `@skip` fields, which the parser drops) or a reached +/// alias's body. /// /// The walk is [`RuntimeTy::visit_heads`](baml_type::RuntimeTy::visit_heads), /// the head walk `bex_vm::reachable` uses for heap declarations, run over the @@ -521,7 +522,7 @@ fn reachable_definitions( continue; } if let Some(class) = class_definitions.get(&name) { - for field in &class.fields { + for field in class.fields.iter().filter(|field| !field.skip) { field .field_type .visit_heads(&mut |head: &DefKey| pending.push(head.clone())); @@ -759,6 +760,27 @@ mod tests { definitions.into_iter().map(|(name, _)| name).collect() } + /// The parser drops a `@skip` field, so a type only a skipped field names + /// is not part of the model. + #[test] + fn skipped_fields_reach_nothing() { + let (root, hidden) = (key("Root"), key("Hidden")); + let (root_key, mut root_definition) = class( + &root, + vec![("name", SapTy::string()), ("secret", class_ty(&hidden))], + ); + root_definition.fields[1].skip = true; + let mut ctx = SysOpContext::empty(); + ctx.class_definitions = Arc::new(IndexMap::from([ + (root_key, root_definition), + class(&hidden, vec![("value", SapTy::int())]), + ])); + + let type_ctx = TypeCtx::for_target(&ctx, &class_ty(&root)); + + assert_eq!(names(&type_ctx.class_definitions), vec![&root]); + } + #[test] fn primitive_target_keeps_no_definitions() { let program = program(); From b5ba31c5684f3e40c6edcf06e50b5acaaefe741b Mon Sep 17 00:00:00 2001 From: aaronvg Date: Fri, 2 Oct 2026 23:46:44 -0700 Subject: [PATCH 7/8] Run schema-aligned parsing as VM natives (#5119) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Issue Reference No issue filed. Stacked on #5113 (which is stacked on #5111): it removes `TypeCtx::for_target` from #5111 and extends `llm_call_memory.rs`. Retarget to `canary` once those merge. ## Changes Schema-aligned parsing (SAP) was three sys ops (`$rust_io_function`) although it does no IO: before `spawn`, sys ops were the only way to get concurrency. As sys ops, each parse: - received `T` as an owned copy (`SapTy`), and the engine copied the class and enum tables into the per-call `SysOpContext` (twice more for runtime-built types); - built the parse model from those copies (`TypeCtx::for_target`); - returned a `BexExternalValue` tree that the engine landed on the heap by looking every class and enum up again by name. The SAP ops already returned `SysOpOutput::Ready`, so they ran inline on the VM task under the heap permit. Moving them into the VM changes nothing about what they block. This PR makes them VM natives (`$rust_function`, `bex_vm/src/package_baml/sap.rs`): - **`_new_parse_cache`** reads `T` from the call's type arguments, walks the declarations it reaches (`reachable::all_declarations`), reads each with the shared definition reader, and builds the same `CompiledSapModel` through the unchanged `TypeCtx::new` / `build_db` / `convert_ty`. The model stays in the `_ParseCache` `RustData` (one model per cache, only the types `T` reaches); it holds tagged names, never heap pointers, because `RustData` is a GC leaf. The instance now carries its real `T`. - **`_parse_final` / `_parse_partial`** run jsonish and the coercer as before and allocate the result directly on the heap, mapping the model's tagged names back to the declarations the target reaches. Error behavior is unchanged: a final failure is `baml.errors.LlmClient` with the same messages, a partial that does not parse yet is `_NoYield`. Supporting changes: - **One definition reader** (`bex_vm::definitions`): `class_definition` / `enum_definition` replace the engine's three class readers and two enum readers, which disagreed on `@skip` (the runtime-type overlay dropped skipped fields; the program table and the runtime-package overlay kept them flagged). Every consumer (output-format rendering, `json.schema`, SAP) already filters on the flag, so the reader keeps skipped fields flagged and drops skipped enum variants. - **`reachable::all_declarations`**: every declaration a type reaches, aliases included; `all_nominals` now filters it, and the walk uses a `HashSet` instead of `Vec::contains`. - **`BexVm::rust_data_field`**: one helper for "clone the `Arc` out of a `$rust_type` field", replacing the copies in `regex.rs` and `csv.rs` (which also reported a misleading `MissingNativeFunction`). - **Deleted:** `sys_ops/src/sap.rs` and the SAP sys-op impls, `CompiledSapModel::from_sys_op_context`, `TypeCtx::for_target` / `reachable_definitions`, and the value half of `bex_sap::to_external` (the type half remains, renamed `to_baml_ty`). `bex_sap` no longer depends on `bex_external_types`; `sys_ops` no longer depends on `bex_sap`. ### Bug fixed: `@skip` fields Parsing into a class with a `@skip` field failed: SAP drops the field (the prompt never shows it), and the heap landing required every field (`Missing field \`age\` in external Instance for class \`user.Person\``). The native fills a skipped field with its type's empty value: `null` for nullable types, `0` / `0.0` / `0n`, `false`, `""`, an empty list or map, a literal's own value, an enum's first variant, and a class of its fields' empty values. A skipped field whose type has none (media, functions, a self-containing class) panics with a message naming the field. ### Memory Measured with the counting allocator in `llm_call_memory.rs` (debug build). The gain grows with the size of the parsed output, because the parse no longer builds an external copy of the value and lands it by name. | Parse of a 13 KB JSON reply with 200 items (`Item { id, name, tags }[]`) | #5113 | this PR | | --- | --- | --- | | Final parse (`baml.sap.parse`) | 1,149,933 bytes | 664,622 bytes (−42%) | | Same reply streamed in 52 batches of 256 characters, each reparsed through one `_ParseCache` (as `ai.stream.Stream` does) | 45,427,554 bytes | 32,301,944 bytes (−29%) | (SAP's own bytes: the totals minus a baseline function doing the same work without parsing, about 4-6 KB.) | Small LLM call (a one-line reply) | #5113 | this PR | | --- | --- | --- | | Bytes allocated per call | 699,060 | 687,584 | | Extra bytes per call from 400 unrelated classes | 21,918 | 0 | | Extra bytes per call per prompt message | 35,023 | 34,967 | The streamed case is still large because every batch reparses the whole text so far; incremental parsing is not in this PR. ## Testing - [x] `bex_sap` unit tests unchanged and passing (they build `TypeRefDb` directly). - [x] New `bex_vm` unit tests: `definitions` keeps skipped fields flagged and drops skipped variants. - [x] New corpus test `sap_parse_fills_skipped_fields_with_empty_values` (failed before with the missing-field error). - [x] `baml_builtins2_codegen`: `test_sap_parse_is_a_vm_builtin` replaces the IO-builtin assertion for `_parse_final`. - [x] New memory guards in `llm_call_memory.rs`: `large_parse_allocates_little` (limit 900 KiB) and `streamed_parse_allocates_little` (limit 40 MiB). Both limits sit between the numbers above, so a return to the external-copy path fails them. - [x] Locally, the tests this touches: `bex_sap`, `bex_vm`, `baml_builtins2_codegen`, `sys_ops`, and the `baml_tests` library plus `llm_call_memory` and `reflect_call_any` (1616 passed); the corpus namespaces that parse (`parse_companions`, `generic_union_returns`, `streaming_partial_parse`, `stdlib_runner_stream_batching`, `runtime_type_bindings_phase1_wrapped`, `runtime_identity_seams`, `provider_stdlib`, and every `llm_*` namespace: 498 passed). CI runs the rest. - [x] `prek` passes (cargo fmt, cargo clippy, cargo stow, baml fmt). ## Additional Notes - Not in this PR: yielding to BAML code for custom coercion (the `YieldToCall` continuation path `baml.json.to` uses), and moving `ctx.output_format()` / `baml.json.schema`, the last users of the `SysOpContext` tables, to natives. - `StreamStateTy` still has no production constructor; left as is. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- > [!NOTE] > **Medium Risk** > Touches core LLM output parsing and heap allocation paths; behavior is intended to match prior SAP errors but execution moved from sys ops to VM with definition-walk semantics changes. > > **Overview** > **Schema-aligned parsing (SAP)** moves from async sys ops to **VM natives** (`baml.sap` uses `$rust_function` with `//baml:mut_vm`). Parsing no longer copies definition tables into `SysOpContext`, builds models via `TypeCtx::for_target`, or returns `BexExternalValue` for a second heap landing—`_new_parse_cache` walks only declarations `T` reaches on the heap, and `_parse_final` / `_parse_partial` allocate results directly. > > Supporting cleanup: shared **`bex_vm::definitions`** for class/enum tables (skipped fields stay flagged), **`reachable::declarations_reached`**, **`BexVm::rust_data_field`**, removal of **`sys_ops` SAP** and **`bex_sap::to_external`** (types live in **`to_baml_ty`**). > > **`@skip` fields** are filled with type empty values instead of failing parse/landing. New corpus and **allocation guards** in `llm_call_memory.rs` lock in lower parse memory (~42% on large final parses). > > Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit e82af144200370b05e402ff4c5fbf28385f1d000. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot). --------- Co-authored-by: Claude Opus 5.5 --- baml_language/Cargo.lock | 3 +- .../baml_std/baml/ns_sap/sap.baml | 10 +- .../baml_builtins2_codegen/src/extract.rs | 9 +- ...tests__render_builtin_package_listing.snap | 4 +- .../ns_parse_companions/parse_companions.baml | 46 ++ .../baml_tests/tests/llm_call_memory.rs | 118 +++- baml_language/crates/bex_engine/src/lib.rs | 181 +------ baml_language/crates/bex_sap/Cargo.toml | 1 - baml_language/crates/bex_sap/src/lib.rs | 27 +- .../crates/bex_sap/src/sap_model/convert.rs | 271 +--------- .../crates/bex_sap/src/to_baml_ty.rs | 140 +++++ .../crates/bex_sap/src/to_external.rs | 275 ---------- baml_language/crates/bex_vm/Cargo.toml | 1 + .../crates/bex_vm/src/definitions.rs | 181 +++++++ baml_language/crates/bex_vm/src/lib.rs | 1 + .../crates/bex_vm/src/package_baml/csv.rs | 52 +- .../crates/bex_vm/src/package_baml/mod.rs | 1 + .../crates/bex_vm/src/package_baml/regex.rs | 19 +- .../crates/bex_vm/src/package_baml/sap.rs | 503 ++++++++++++++++++ baml_language/crates/bex_vm/src/reachable.rs | 41 +- baml_language/crates/bex_vm/src/vm.rs | 32 ++ .../crates/bex_vm_types/src/errors.rs | 4 + baml_language/crates/sys_ops/Cargo.toml | 1 - baml_language/crates/sys_ops/src/lib.rs | 85 +-- baml_language/crates/sys_ops/src/sap.rs | 173 ------ 25 files changed, 1120 insertions(+), 1059 deletions(-) create mode 100644 baml_language/crates/bex_sap/src/to_baml_ty.rs delete mode 100644 baml_language/crates/bex_sap/src/to_external.rs create mode 100644 baml_language/crates/bex_vm/src/definitions.rs create mode 100644 baml_language/crates/bex_vm/src/package_baml/sap.rs delete mode 100644 baml_language/crates/sys_ops/src/sap.rs diff --git a/baml_language/Cargo.lock b/baml_language/Cargo.lock index 242d6fc41b1..e730fbb4cd7 100644 --- a/baml_language/Cargo.lock +++ b/baml_language/Cargo.lock @@ -1476,7 +1476,6 @@ name = "bex_sap" version = "0.0.0-beta" dependencies = [ "baml_type", - "bex_external_types", "bex_str", "bex_vm_types", "derive_more", @@ -1516,6 +1515,7 @@ dependencies = [ "baml_type", "base64", "bex_heap", + "bex_sap", "bex_str", "bex_vm_types", "btel_clock", @@ -7424,7 +7424,6 @@ dependencies = [ "bex_external_types", "bex_heap", "bex_resource_types", - "bex_sap", "bex_vm_types", "indexmap", "num-bigint", diff --git a/baml_language/crates/baml_builtins2/baml_std/baml/ns_sap/sap.baml b/baml_language/crates/baml_builtins2/baml_std/baml/ns_sap/sap.baml index 31c3375bf60..0c437523e66 100644 --- a/baml_language/crates/baml_builtins2/baml_std/baml/ns_sap/sap.baml +++ b/baml_language/crates/baml_builtins2/baml_std/baml/ns_sap/sap.baml @@ -8,14 +8,16 @@ class _ParseCache { _data: $rust_type, + //baml:mut_vm /// (internal) Final SAP operation shared by streams and public SAP. function _parse_final(self, json: string) -> T throws root.errors.LlmClient { - $rust_io_function + $rust_function } + //baml:mut_vm /// (internal) Partial SAP operation used by `ai.stream.Stream`. function _parse_partial(self, json: string) -> T | _NoYield throws root.errors.LlmClient { - $rust_io_function + $rust_function } } @@ -23,6 +25,8 @@ class _ParseCache { class _NoYield { } +//baml:mut_vm +//baml:fallible /// (internal) Builds a parse cache for `T`. /// /// A free function rather than a `_ParseCache` static method: a *static* @@ -31,7 +35,7 @@ class _NoYield { /// operand slots the call site never fills. Instance methods on a generic /// class (`_parse_final`, `_parse_partial`) are threaded correctly. function _new_parse_cache() -> _ParseCache throws never { - $rust_io_function + $rust_function } /// Parse an LLM's raw text output into `T` with Schema-Aligned Parsing. diff --git a/baml_language/crates/baml_builtins2_codegen/src/extract.rs b/baml_language/crates/baml_builtins2_codegen/src/extract.rs index 08723fb013a..43425fffe30 100644 --- a/baml_language/crates/baml_builtins2_codegen/src/extract.rs +++ b/baml_language/crates/baml_builtins2_codegen/src/extract.rs @@ -1579,8 +1579,15 @@ implements Modifier for Root { .find(|b| b.path == "baml.http._fetch") .unwrap(); assert_eq!(http_fetch.throws, throws(&["Io", "Timeout"])); + } + + /// Schema-aligned parsing does no IO, so its entry points are VM natives. + #[test] + fn test_sap_parse_is_a_vm_builtin() { + let (vm_builtins, io_builtins, _class_defs) = extract_native_builtins().unwrap(); - let sap_final = io_builtins + assert!(!io_builtins.iter().any(|b| b.path.starts_with("baml.sap."))); + let sap_final = vm_builtins .iter() .find(|b| b.path == "baml.sap._ParseCache._parse_final") .unwrap(); diff --git a/baml_language/crates/baml_cli/src/snapshots/baml_cli__describe_command_tests__render_builtin_package_listing.snap b/baml_language/crates/baml_cli/src/snapshots/baml_cli__describe_command_tests__render_builtin_package_listing.snap index f8e38098e22..0105f657791 100644 --- a/baml_language/crates/baml_cli/src/snapshots/baml_cli__describe_command_tests__render_builtin_package_listing.snap +++ b/baml_language/crates/baml_cli/src/snapshots/baml_cli__describe_command_tests__render_builtin_package_listing.snap @@ -216,8 +216,8 @@ class baml.regex.Error /baml/ns_regex/regex. class baml.regex.Regex /baml/ns_regex/regex.baml:97 function baml.regex.new /baml/ns_regex/regex.baml:244 function baml.regex.escape /baml/ns_regex/regex.baml:268 -function baml.sap.parse /baml/ns_sap/sap.baml:41 -function baml.sap.parse_type /baml/ns_sap/sap.baml:54 +function baml.sap.parse /baml/ns_sap/sap.baml:45 +function baml.sap.parse_type /baml/ns_sap/sap.baml:58 class baml.spawn.CancelToken /baml/ns_spawn/spawn.baml:23 interface baml.spawn.Modifier /baml/ns_spawn/spawn.baml:60 class baml.spawn.Plan /baml/ns_spawn/spawn.baml:78 diff --git a/baml_language/crates/baml_tests/baml_src/ns_parse_companions/parse_companions.baml b/baml_language/crates/baml_tests/baml_src/ns_parse_companions/parse_companions.baml index fad5822eeb1..133789c83c1 100644 --- a/baml_language/crates/baml_tests/baml_src/ns_parse_companions/parse_companions.baml +++ b/baml_language/crates/baml_tests/baml_src/ns_parse_companions/parse_companions.baml @@ -172,3 +172,49 @@ test "sap_parse_matches_the_enum_lines_the_prompt_renders" { assert.is_true(prompt.includes("- A: eaten in the morning")); assert.is_true(prompt.includes("- B: LUNCH eaten at noon")) } + +// ─── @skip fields ──────────────────────────────────────────────────────────── +// +// The parser never reads a `@skip` field (the prompt does not show it), so a +// parsed instance holds its type's empty value: "" for a string, the zero of a +// number, false, an empty list, null for an optional field, an enum's first +// variant, and a class of its fields' empty values. + +enum ParseCompanionsSkipColor { + Red, + Green, +} + +class ParseCompanionsSkipInner { + label: string, + count: int, +} + +class ParseCompanionsSkipPerson { + name: string, + age: int @skip, + score: float @skip, + active: bool @skip, + tags: string[] @skip, + nick: string? @skip, + color: ParseCompanionsSkipColor @skip, + inner: ParseCompanionsSkipInner @skip, +} + +function parse_companions_skipped_fields_fn() -> bool { + let parsed = baml.sap.parse("{\"name\": \"Ada\", \"age\": 37}"); + parsed == ParseCompanionsSkipPerson { + name: "Ada", + age: 0, + score: 0.0, + active: false, + tags: [], + nick: null, + color: ParseCompanionsSkipColor.Red, + inner: ParseCompanionsSkipInner { label: "", count: 0 }, + } +} + +test "sap_parse_fills_skipped_fields_with_empty_values" { + assert.is_true(parse_companions_skipped_fields_fn()) +} diff --git a/baml_language/crates/baml_tests/tests/llm_call_memory.rs b/baml_language/crates/baml_tests/tests/llm_call_memory.rs index 955f1a52e83..5401c58b043 100644 --- a/baml_language/crates/baml_tests/tests/llm_call_memory.rs +++ b/baml_language/crates/baml_tests/tests/llm_call_memory.rs @@ -6,6 +6,9 @@ //! until the VM collected it. And every `match` type test that failed //! canonicalized the recursive `json` alias: the request builder tests each //! node of the request JSON that way, so the cost grew with the conversation. +//! Parsing a reply also used to build a second, external copy of the parsed +//! value and land it on the heap by name, so the parse cost of a large or +//! streamed output was larger than it needs to be. //! A counting global allocator measures every byte allocated while the calls //! run, on every thread, so the bounds do not depend on when the GC runs. @@ -85,6 +88,15 @@ const MAX_BYTES_PER_CALL: usize = 2 * 1024 * 1024; const MAX_UNRELATED_GROWTH_PER_CALL: usize = 256 * 1024; /// User/assistant exchanges in the long conversation. const LONG_CONVERSATION_TURNS: usize = 10; +/// Items in the large parsed output, and the characters a stream adds per batch. +const PARSED_ITEMS: i64 = 200; +const STREAM_STEP: usize = 256; +/// The ceiling for parsing a 13 KB reply of 200 items: about 665 KB; about +/// 1.15 MB while each parse built an external copy of the value. +const MAX_FINAL_PARSE_BYTES: usize = 900 * 1024; +/// The ceiling for the same reply streamed in 52 batches, each reparsed from +/// the start: about 32.3 MB; about 45.4 MB with the external copy. +const MAX_STREAMED_PARSE_BYTES: usize = 40 * 1024 * 1024; /// The extra bytes per call that one more prompt message may cost: about 35 KB, /// and about 1.56 MB while every failed type test on the request JSON expanded /// `json`. @@ -119,14 +131,19 @@ client Fake = openai.ResponsesClient.new(model = "fake", api_key = "test-key", b "#, base_url = server.uri() ); + (server, engine(&source)) +} + +/// An engine running `source` with the native sys ops. +fn engine(source: &str) -> Arc { let engine = BexEngine::new_with_runtime_compiler( - compile_source(&source), + compile_source(source), Arc::new(sys_ops::SysOps::native()), Vec::new(), bex_project::runtime_compiler(), ) .expect("engine builds"); - (server, Arc::new(engine)) + Arc::new(engine) } async fn call(engine: &Arc, function: &str) -> BexExternalValue { @@ -194,6 +211,47 @@ function main() -> string {{ ) } +/// A program that parses a 13 KB JSON reply of `PARSED_ITEMS` items: whole in +/// `final_parse`, and in `stream_parse` the way `ai.stream.Stream` does, one +/// growing prefix per batch through one parse cache. +fn parse_program() -> String { + let items: Vec = (0..PARSED_ITEMS) + .map(|i| format!(r#"{{"id": {i}, "name": "item number {i}", "tags": ["alpha", "beta"]}}"#)) + .collect(); + let text = serde_json::to_string(&format!("[{}]", items.join(", "))) + .expect("the reply is JSON-serializable"); + format!( + r#" +class Item {{ + id: int, + name: string, + tags: string[], +}} + +function final_parse() -> int {{ + baml.sap.parse({text}).length() +}} + +function stream_parse() -> int {{ + let text = {text}; + let cache = baml.sap._new_parse_cache(); + let yielded = 0; + let end = {STREAM_STEP}; + while (end < text.length() + {STREAM_STEP}) {{ + let upto = if (end > text.length()) {{ text.length() }} else {{ end }}; + let parsed: Item[] | baml.sap._NoYield = cache._parse_partial(text.slice(0, upto)); + match (parsed) {{ + let items: Item[] => {{ yielded = yielded + 1; }}, + _ => {{}}, + }}; + end = end + {STREAM_STEP}; + }} + yielded +}} +"# + ) +} + const PICK_TOOL: &str = r#" class Choice { tool string @@ -272,3 +330,59 @@ async fn call_cost_does_not_grow_with_conversation_length() { with a small constant" ); } + +#[tokio::test] +async fn large_parse_allocates_little() { + let _measuring = MEASURING.lock().await; + let engine = engine(&parse_program()); + assert_eq!( + call(&engine, "final_parse").await, + BexExternalValue::Int(PARSED_ITEMS) + ); + let bytes = bytes_per_call(&engine, "final_parse").await; + assert!( + bytes <= MAX_FINAL_PARSE_BYTES, + "parsing {PARSED_ITEMS} items allocated {bytes} bytes (limit {MAX_FINAL_PARSE_BYTES})" + ); +} + +#[tokio::test] +async fn streamed_parse_allocates_little() { + let _measuring = MEASURING.lock().await; + let engine = engine(&parse_program()); + assert!( + matches!(call(&engine, "stream_parse").await, BexExternalValue::Int(yields) if yields > 0), + "the stream never yielded a partial value" + ); + let bytes = bytes_per_call(&engine, "stream_parse").await; + assert!( + bytes <= MAX_STREAMED_PARSE_BYTES, + "streaming {PARSED_ITEMS} items allocated {bytes} bytes (limit {MAX_STREAMED_PARSE_BYTES})" + ); +} + +/// The parser drops a `@skip` field, so the types only a skipped field names +/// cost a parse nothing. +#[tokio::test] +async fn call_cost_does_not_grow_with_types_only_skipped_fields_name() { + let _measuring = MEASURING.lock().await; + let (_small_server, small) = engine_for(PICK_TOOL, r#"{"tool": "search"}"#).await; + let skipping_program = format!( + "{}\n{}", + unrelated_classes(400), + PICK_TOOL.replace( + "class Choice {\n tool string\n}", + "class Choice {\n tool string\n extra Unrelated399 @skip\n}", + ) + ); + assert!(skipping_program.contains("extra Unrelated399 @skip")); + let (_skipping_server, skipping) = engine_for(&skipping_program, r#"{"tool": "search"}"#).await; + let small_bytes = bytes_per_call(&small, "main").await; + let skipping_bytes = bytes_per_call(&skipping, "main").await; + let growth = skipping_bytes.saturating_sub(small_bytes); + assert!( + growth <= MAX_UNRELATED_GROWTH_PER_CALL, + "400 classes a skipped field names added {growth} bytes per call \ + ({small_bytes} -> {skipping_bytes}); the parse must not convert them" + ); +} diff --git a/baml_language/crates/bex_engine/src/lib.rs b/baml_language/crates/bex_engine/src/lib.rs index 646b9de35a9..b7218484b42 100644 --- a/baml_language/crates/bex_engine/src/lib.rs +++ b/baml_language/crates/bex_engine/src/lib.rs @@ -2248,43 +2248,7 @@ impl BexEngine { if let Object::Class(cls) = obj { defs.insert( ::sys_types::DefKey::new(cls.type_tag, cls.name.clone()), - sys_types::ClassDefinition { - name: cls.name.display_name().to_string(), - description: cls.description.clone(), - docstring: cls.docstring.clone(), - alias: cls.alias.clone(), - stream_done: cls.stream_done, - fields: cls - .fields - .iter() - .map(|f| sys_types::ClassFieldDefinition { - name: f.name.clone(), - field_type: f - .field_type - .try_map_heads(&mut bex_vm_types::TypeHead::to_tagged_name) - .unwrap_or_else(|_| { - unreachable!( - "compiled class fields name compiled declarations" - ) - }), - field_template: Some( - f.field_template - .try_map_heads(&mut bex_vm_types::TypeHead::to_tagged_name) - .unwrap_or_else(|_| { - unreachable!( - "compiled class fields name compiled declarations" - ) - }), - ), - description: f.description.clone(), - docstring: f.docstring.clone(), - alias: f.alias.clone(), - skip: f.skip, - stream_done: f.stream_done, - must_exist: f.must_exist, - }) - .collect(), - }, + bex_vm::definitions::class_definition(cls), ); } } @@ -2302,109 +2266,18 @@ impl BexEngine { if let Object::Enum(enm) = obj { defs.insert( ::sys_types::DefKey::new(enm.type_tag, enm.name.clone()), - sys_types::EnumDefinition { - name: enm.name.display_name().to_string(), - description: enm.description.clone(), - docstring: enm.docstring.clone(), - alias: enm.alias.clone(), - variants: enm - .variants - .iter() - .filter(|v| !v.skip) - .map(|v| sys_types::EnumVariantDefinition { - name: v.name.clone(), - description: v.description.clone(), - docstring: v.docstring.clone(), - alias: v.alias.clone(), - }) - .collect(), - }, + bex_vm::definitions::enum_definition(enm), ); } } defs } - /// Carry a declaration's field type onto the sys-op lane. - /// - /// Total by invariant: a live declaration's heads are resolved pointers to - /// declarations, so each yields its identity and its own name. - fn lane_ty(ty: &bex_vm_types::RuntimeTy) -> ::sys_types::SapTy { - ty.try_map_heads(&mut bex_vm_types::TypeHead::to_tagged_name) - .unwrap_or_else(|head| { - unreachable!("a live declaration's field names a declaration: {head}") - }) - } - - /// [`Self::lane_ty`] for a field's symbolic template, same totality - /// argument: every head a live declaration's template names is resolved. - fn lane_template(template: &bex_vm_types::TyTemplate) -> ::sys_types::SapTyTemplate { - template - .try_map_heads(&mut bex_vm_types::TypeHead::to_tagged_name) - .unwrap_or_else(|head| { - unreachable!("a live declaration's field template names a declaration: {head}") - }) - } - - fn enum_definition(enm: &bex_vm_types::Enum) -> sys_types::EnumDefinition { - sys_types::EnumDefinition { - name: enm.name.display_name().to_string(), - description: enm.description.clone(), - docstring: enm.docstring.clone(), - alias: enm.alias.clone(), - variants: enm - .variants - .iter() - .filter(|variant| !variant.skip) - .map(|variant| sys_types::EnumVariantDefinition { - name: variant.name.clone(), - description: variant.description.clone(), - docstring: variant.docstring.clone(), - alias: variant.alias.clone(), - }) - .collect(), - } - } - - fn class_definition( - class: &bex_vm_types::Class, - _permit: bex_heap::PermitProof<'_>, - ) -> sys_types::ClassDefinition { - sys_types::ClassDefinition { - name: class.name.display_name().to_string(), - description: class.description.clone(), - docstring: class.docstring.clone(), - alias: class.alias.clone(), - stream_done: class.stream_done, - fields: class - .fields - .iter() - .filter(|field| !field.skip) - .map(|field| sys_types::ClassFieldDefinition { - name: field.name.clone(), - field_type: Self::lane_ty(&field.field_type), - field_template: Some(Self::lane_template(&field.field_template)), - description: field.description.clone(), - docstring: field.docstring.clone(), - alias: field.alias.clone(), - skip: field.skip, - stream_done: field.stream_done, - must_exist: field.must_exist, - }) - .collect(), - } - } - /// Gather runtime definitions from the type descriptors passed directly /// to a sys-op. Type arguments are lowered as ordinary `Object::Type` /// arguments, so this is the last synchronous chokepoint before the permit /// is released for async work. - fn runtime_type_overlay( - &self, - vm: &BexVm, - args: &[Value], - permit: bex_heap::PermitProof<'_>, - ) -> RuntimeTypeOverlay { + fn runtime_type_overlay(&self, vm: &BexVm, args: &[Value]) -> RuntimeTypeOverlay { let mut overlay = RuntimeTypeOverlay::default(); for value in args { let Some(type_ptr) = value.as_object_ptr() else { @@ -2425,7 +2298,7 @@ impl BexEngine { overlay .class_definitions .entry(head.clone()) - .or_insert_with(|| Self::class_definition(class, permit)); + .or_insert_with(|| bex_vm::definitions::class_definition(class)); overlay .class_handles .entry(head) @@ -2439,7 +2312,7 @@ impl BexEngine { overlay .enum_definitions .entry(head.clone()) - .or_insert_with(|| Self::enum_definition(enm)); + .or_insert_with(|| bex_vm::definitions::enum_definition(enm)); overlay .enum_handles .entry(head) @@ -2499,28 +2372,7 @@ impl BexEngine { }; classes.insert( ::sys_types::DefKey::new(class.type_tag, class.name.clone()), - sys_types::ClassDefinition { - name: class.name.display_name().to_string(), - description: class.description.clone(), - docstring: class.docstring.clone(), - alias: class.alias.clone(), - stream_done: class.stream_done, - fields: class - .fields - .iter() - .map(|field| sys_types::ClassFieldDefinition { - name: field.name.clone(), - field_type: Self::lane_ty(&field.field_type), - field_template: Some(Self::lane_template(&field.field_template)), - description: field.description.clone(), - docstring: field.docstring.clone(), - alias: field.alias.clone(), - skip: field.skip, - stream_done: field.stream_done, - must_exist: field.must_exist, - }) - .collect(), - }, + bex_vm::definitions::class_definition(class), ); named_owners.insert(class.name.to_string(), handle.clone()); } @@ -2530,23 +2382,7 @@ impl BexEngine { }; enums.insert( ::sys_types::DefKey::new(enm.type_tag, enm.name.clone()), - sys_types::EnumDefinition { - name: enm.name.display_name().to_string(), - description: enm.description.clone(), - docstring: enm.docstring.clone(), - alias: enm.alias.clone(), - variants: enm - .variants - .iter() - .filter(|variant| !variant.skip) - .map(|variant| sys_types::EnumVariantDefinition { - name: variant.name.clone(), - description: variant.description.clone(), - docstring: variant.docstring.clone(), - alias: variant.alias.clone(), - }) - .collect(), - }, + bex_vm::definitions::enum_definition(enm), ); named_owners.insert(enm.name.to_string(), handle.clone()); } @@ -5853,8 +5689,7 @@ impl BexEngine { // runs its sys-ops to completion. let op_cancel = thread.sysop_cancel.clone(); - let runtime_type_overlay = - self.runtime_type_overlay(&thread.vm, &args, thread.proof()); + let runtime_type_overlay = self.runtime_type_overlay(&thread.vm, &args); let runtime_compile_request = match operation { SysOp::ReflectPackageCompile => { Some(Ok(Self::runtime_compile_request(&thread.vm, &args)?)) diff --git a/baml_language/crates/bex_sap/Cargo.toml b/baml_language/crates/bex_sap/Cargo.toml index d9085f053eb..f392936d57b 100644 --- a/baml_language/crates/bex_sap/Cargo.toml +++ b/baml_language/crates/bex_sap/Cargo.toml @@ -18,7 +18,6 @@ workspace = true [dependencies] baml_type = { workspace = true } -bex_external_types = { workspace = true } bex_str = { workspace = true } bex_vm_types = { workspace = true } sys_types = { workspace = true } diff --git a/baml_language/crates/bex_sap/src/lib.rs b/baml_language/crates/bex_sap/src/lib.rs index d3dc4f5cd05..2db7436a6f4 100644 --- a/baml_language/crates/bex_sap/src/lib.rs +++ b/baml_language/crates/bex_sap/src/lib.rs @@ -2,7 +2,10 @@ use ouroboros::self_referencing; use sys_types::DefKey; pub use crate::jsonish::parse; -use crate::sap_model::{Ty, TypeRefDb}; +use crate::{ + deserializer::coercer::{ParsingContext, ParsingError}, + sap_model::{Ty, TyResolvedRef, TypeRefDb}, +}; pub mod baml_value; pub mod deserializer; @@ -10,7 +13,7 @@ pub mod jsonish; pub mod sap_model; #[cfg(test)] mod tests; -pub mod to_external; +pub mod to_baml_ty; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum StreamingMode { @@ -43,14 +46,6 @@ impl CompiledSapModel { )?; Ok(Self { inner }) } - pub fn from_sys_op_context( - ctx: &::sys_types::SysOpContext, - target: sys_types::SapTy, - ) -> Result { - let type_ctx = sap_model::TypeCtx::for_target(ctx, &target); - Self::from_type_ctx(type_ctx, target) - } - pub fn db(&self) -> &TypeRefDb<'_, DefKey> { self.inner.borrow_db() } @@ -58,6 +53,18 @@ impl CompiledSapModel { pub fn ty(&self) -> &Ty<'_, DefKey> { self.inner.borrow_ty() } + + /// The parse target, resolved in [`Self::db`]. + /// + /// # Errors + /// + /// The target names a type the database does not hold: a schema the + /// converter accepted but cannot resolve, which no input text can fix. + pub fn resolved_target(&self) -> Result, ParsingError> { + self.db() + .resolve(self.ty()) + .map_err(|name| ParsingContext::new(self.db()).error_type_resolution(name)) + } } #[self_referencing] diff --git a/baml_language/crates/bex_sap/src/sap_model/convert.rs b/baml_language/crates/bex_sap/src/sap_model/convert.rs index d9ef52c1aed..d66aca35502 100644 --- a/baml_language/crates/bex_sap/src/sap_model/convert.rs +++ b/baml_language/crates/bex_sap/src/sap_model/convert.rs @@ -47,8 +47,9 @@ pub enum ConvertError { const MAX_RECURSION_DEPTH: usize = 16; -/// Contains stuff from [`sys_types::SysOpContext`] that we need for converting to the sap model: -/// the definitions one parse target reaches ([`TypeCtx::for_target`]), not the whole program's. +/// The class, enum, and type-alias definitions one parse target reaches, which the converter +/// needs to build the sap model. The caller reads them off the target's heap declarations, so a +/// parse pays only for the types its target uses, not for the whole program's. /// /// ## Representation /// - Unions should be flattened: @@ -71,9 +72,9 @@ pub struct TypeCtx { } impl TypeCtx { /// Builds the context from definitions that are closed under reference: - /// every name a class field or an alias body refers to is present too - /// ([`TypeCtx::for_target`] collects such a set). A missing name makes the - /// types that refer to it unparsable. + /// every name a class field or an alias body refers to is present too (the + /// declarations a type reaches on the heap are such a set). A missing name + /// makes the types that refer to it unparsable. pub fn new<'d>( class_definitions: impl IntoIterator, enum_definitions: IndexMap, @@ -149,38 +150,6 @@ impl TypeCtx { } } - /// The context for parsing into `target`: only the classes, enums, and type - /// aliases that `target` reaches (`reachable_definitions`), in program - /// order. Converting every definition in the program on each parse cost - /// several MB per LLM call; the reachable set is all the converter looks up. - pub fn for_target( - ctx: &::sys_types::SysOpContext, - target: &SapTy, - ) -> Self { - let reachable = - reachable_definitions(target, &ctx.class_definitions, &ctx.type_alias_definitions); - let enum_definitions = ctx - .enum_definitions - .iter() - .filter(|(name, _)| reachable.contains(*name)) - .map(|(name, definition)| (name.clone(), definition.clone())) - .collect(); - let type_alias_definitions = ctx - .type_alias_definitions - .iter() - .filter(|(name, _)| reachable.contains(*name)) - .map(|(name, ty)| (name.clone(), ty.clone())) - .collect(); - - Self::new( - ctx.class_definitions - .iter() - .filter(|(name, _)| reachable.contains(*name)), - enum_definitions, - &type_alias_definitions, - ) - } - /// Normalize a runtime-materialized parse target before converting it to /// the SAP model. Generic substitution happens after [`TypeCtx::new`] has /// simplified the declared class and alias definitions, so it can create a @@ -498,42 +467,6 @@ impl TypeCtx { } } -/// The definitions `target` can reach: the classes, interfaces, enums, and type -/// aliases it names, then, transitively, the ones named by a reached class's -/// field types (except `@skip` fields, which the parser drops) or a reached -/// alias's body. -/// -/// The walk is [`RuntimeTy::visit_heads`](baml_type::RuntimeTy::visit_heads), -/// the head walk `bex_vm::reachable` uses for heap declarations, run over the -/// context's definition tables (which the engine extracts from those same -/// declarations). Heads the converter never resolves, such as type arguments -/// and function types, are kept too. That only adds definitions, so every name -/// the converter looks up is present. -fn reachable_definitions( - target: &SapTy, - class_definitions: &IndexMap, - type_alias_definitions: &IndexMap, -) -> HashSet { - let mut reachable = HashSet::new(); - let mut pending = Vec::new(); - target.visit_heads(&mut |name: &DefKey| pending.push(name.clone())); - while let Some(name) = pending.pop() { - if !reachable.insert(name.clone()) { - continue; - } - if let Some(class) = class_definitions.get(&name) { - for field in class.fields.iter().filter(|field| !field.skip) { - field - .field_type - .visit_heads(&mut |head: &DefKey| pending.push(head.clone())); - } - } else if let Some(alias) = type_alias_definitions.get(&name) { - alias.visit_heads(&mut |head: &DefKey| pending.push(head.clone())); - } - } - reachable -} - fn check_parseable( name: &DefKey, class_definitions: &IndexMap, @@ -633,195 +566,3 @@ fn is_sap_parseable(ty: &SapTy) -> Result, ()> { | SapTy::Type => Err(()), } } - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use baml_type::{DeclarationName, Name, typetag::TypeTag}; - use sys_types::{ClassFieldDefinition, EnumVariantDefinition, SysOpContext}; - - use super::*; - - fn key(name: &str) -> DefKey { - DefKey::new( - TypeTag::fresh_dynamic(), - DeclarationName::Anonymous(Name::new(name)), - ) - } - - fn class(key: &DefKey, fields: Vec<(&str, SapTy)>) -> (DefKey, ClassDefinition) { - let fields = fields - .into_iter() - .map(|(name, field_type)| ClassFieldDefinition { - name: name.to_string(), - field_type, - field_template: None, - description: None, - docstring: None, - alias: None, - skip: false, - stream_done: false, - must_exist: false, - }) - .collect(); - let definition = ClassDefinition { - name: key.to_string(), - description: None, - docstring: None, - alias: None, - fields, - stream_done: false, - }; - (key.clone(), definition) - } - - fn enumeration(key: &DefKey) -> (DefKey, EnumDefinition) { - let definition = EnumDefinition { - name: key.to_string(), - description: None, - docstring: None, - alias: None, - variants: vec![EnumVariantDefinition { - name: "Only".to_string(), - description: None, - docstring: None, - alias: None, - }], - }; - (key.clone(), definition) - } - - fn class_ty(key: &DefKey) -> SapTy { - SapTy::Class(key.clone(), Box::new([])) - } - - /// A program with a reachable cluster around `Root` and types nothing in it - /// refers to. `Root -> Child -> Root` is a cycle; `Child` also names an enum - /// and a recursive alias whose body names `Leaf`. - struct Program { - root: DefKey, - child: DefKey, - leaf: DefKey, - color: DefKey, - tree: DefKey, - unrelated: DefKey, - unrelated_enum: DefKey, - ctx: SysOpContext, - } - - fn program() -> Program { - let [root, child, leaf, color, tree, unrelated, unrelated_enum] = [ - "Root", - "Child", - "Leaf", - "Color", - "Tree", - "Unrelated", - "UnrelatedEnum", - ] - .map(key); - let tree_body = - SapTy::union([SapTy::list(SapTy::TypeAlias(tree.clone())), class_ty(&leaf)]); - let mut ctx = SysOpContext::empty(); - ctx.class_definitions = Arc::new(IndexMap::from([ - class(&unrelated, vec![("value", SapTy::string())]), - class(&root, vec![("child", class_ty(&child))]), - class( - &child, - vec![ - ("parent", SapTy::optional(class_ty(&root))), - ("color", SapTy::Enum(color.clone())), - ("tree", SapTy::TypeAlias(tree.clone())), - ], - ), - class(&leaf, vec![("value", SapTy::int())]), - ])); - ctx.enum_definitions = Arc::new(IndexMap::from([ - enumeration(&unrelated_enum), - enumeration(&color), - ])); - ctx.type_alias_definitions = Arc::new(IndexMap::from([(tree.clone(), tree_body)])); - Program { - root, - child, - leaf, - color, - tree, - unrelated, - unrelated_enum, - ctx, - } - } - - fn names<'d, V: 'd>( - definitions: impl IntoIterator, - ) -> Vec<&'d DefKey> { - definitions.into_iter().map(|(name, _)| name).collect() - } - - /// The parser drops a `@skip` field, so a type only a skipped field names - /// is not part of the model. - #[test] - fn skipped_fields_reach_nothing() { - let (root, hidden) = (key("Root"), key("Hidden")); - let (root_key, mut root_definition) = class( - &root, - vec![("name", SapTy::string()), ("secret", class_ty(&hidden))], - ); - root_definition.fields[1].skip = true; - let mut ctx = SysOpContext::empty(); - ctx.class_definitions = Arc::new(IndexMap::from([ - (root_key, root_definition), - class(&hidden, vec![("value", SapTy::int())]), - ])); - - let type_ctx = TypeCtx::for_target(&ctx, &class_ty(&root)); - - assert_eq!(names(&type_ctx.class_definitions), vec![&root]); - } - - #[test] - fn primitive_target_keeps_no_definitions() { - let program = program(); - let ctx = TypeCtx::for_target(&program.ctx, &SapTy::string()); - assert!(ctx.class_definitions.is_empty()); - assert!(ctx.enum_definitions.is_empty()); - assert!(ctx.type_alias_definitions.is_empty()); - } - - #[test] - fn class_target_keeps_what_its_fields_reach_in_program_order() { - let program = program(); - let ctx = TypeCtx::for_target(&program.ctx, &class_ty(&program.root)); - - assert_eq!( - names(&ctx.class_definitions), - vec![&program.root, &program.child, &program.leaf], - ); - assert_eq!(names(&ctx.enum_definitions), vec![&program.color]); - assert_eq!(names(&ctx.type_alias_definitions), vec![&program.tree]); - assert!(!ctx.class_definitions.contains_key(&program.unrelated)); - assert!(!ctx.enum_definitions.contains_key(&program.unrelated_enum)); - } - - #[test] - fn alias_target_keeps_its_body() { - let program = program(); - let ctx = TypeCtx::for_target(&program.ctx, &SapTy::TypeAlias(program.tree.clone())); - - assert_eq!(names(&ctx.class_definitions), vec![&program.leaf]); - assert!(ctx.enum_definitions.is_empty()); - assert_eq!(names(&ctx.type_alias_definitions), vec![&program.tree]); - } - - #[test] - fn reachable_definitions_convert_without_unknown_names() { - let program = program(); - let root = class_ty(&program.root); - let ctx = TypeCtx::for_target(&program.ctx, &root); - let db = ctx.build_db().expect("the reachable set converts"); - let target = ctx.convert_ty(&root).expect("the target converts"); - assert!(db.resolve(&target).is_ok()); - } -} diff --git a/baml_language/crates/bex_sap/src/to_baml_ty.rs b/baml_language/crates/bex_sap/src/to_baml_ty.rs new file mode 100644 index 00000000000..a084b9b0aae --- /dev/null +++ b/baml_language/crates/bex_sap/src/to_baml_ty.rs @@ -0,0 +1,140 @@ +//! Conversion from SAP model types back to `SapTy`. +//! +//! A parsed value's metadata carries SAP model types; the VM, which builds the +//! value on the heap, turns them into the element, key, and value types of the +//! lists and maps it allocates. + +use sys_types::{DefKey, SapTy}; + +use crate::sap_model::{ + ArrayTy, ClassTy, EnumTy, EnumVariantTy, MapTy, MediaTy, StreamStateTy, Ty, TyResolvedRef, + TypeRefDb, UnionTy, +}; + +// ============================================================================ +// Type conversion: SAP types → SapTy +// ============================================================================ + +/// Convert a SAP type back to a `SapTy`. +/// +/// # Known simplifications (not round-trip safe) +/// +/// - `Optional(T)` was converted to `Union([Null, T])` on the way in; it stays as `Union` here. +/// - Aliases already resolved into the target type have no nominal name to recover; named alias +/// references inside containers are preserved. +/// - `EnumVariant` becomes `Enum` (variant specificity lost at the type level). +pub trait ToBamlTy { + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy; +} + +impl ToBamlTy for TyResolvedRef<'_, DefKey> { + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { + match self { + TyResolvedRef::Int(_) => SapTy::Int, + TyResolvedRef::Bigint(_) => SapTy::Bigint, + TyResolvedRef::Float(_) => SapTy::Float, + TyResolvedRef::String(_) => SapTy::String, + TyResolvedRef::Bool(_) => SapTy::Bool, + TyResolvedRef::Null(_) => SapTy::Null, + TyResolvedRef::Media(media) => SapTy::Media(media.to_baml_media_kind()), + TyResolvedRef::LiteralInt(v) => { + SapTy::Literal(baml_type::Literal::Int(v.0), baml_type::Freshness::Regular) + } + TyResolvedRef::LiteralBigint(v) => SapTy::Literal( + baml_type::Literal::Bigint(v.0.clone()), + baml_type::Freshness::Regular, + ), + TyResolvedRef::LiteralString(v) => SapTy::Literal( + baml_type::Literal::String(v.0.to_string()), + baml_type::Freshness::Regular, + ), + TyResolvedRef::LiteralBool(v) => { + SapTy::Literal(baml_type::Literal::Bool(v.0), baml_type::Freshness::Regular) + } + TyResolvedRef::Array(a) => a.to_baml_ty(db), + TyResolvedRef::Map(m) => m.to_baml_ty(db), + TyResolvedRef::Class(c) => c.to_baml_ty(db), + TyResolvedRef::Enum(e) => e.to_baml_ty(db), + TyResolvedRef::EnumVariant(ev) => ev.to_baml_ty(db), + TyResolvedRef::Union(u) => u.to_baml_ty(db), + TyResolvedRef::StreamState(s) => s.to_baml_ty(db), + } + } +} + +impl ToBamlTy for Ty<'_, DefKey> { + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { + match self { + Ty::Resolved(resolved) => resolved.as_ref().to_baml_ty(db), + Ty::ResolvedRef(resolved_ref) => resolved_ref.to_baml_ty(db), + Ty::Unresolved(name) => match db.resolve_name(name) { + Some(TyResolvedRef::Class(class)) if class.name == *name => class.to_baml_ty(db), + Some(TyResolvedRef::Enum(enm)) if enm.name == *name => enm.to_baml_ty(db), + // Any other named entry is an alias. Keep it nominal here + // instead of recursively expanding aliases such as `json`. + Some(_) => SapTy::TypeAlias(name.clone()), + None => SapTy::unknown(), + }, + } + } +} + +impl ToBamlTy for ArrayTy<'_, DefKey> { + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { + let inner = self.ty.to_baml_ty(db); + SapTy::List(Box::new(inner)) + } +} + +impl ToBamlTy for MapTy<'_, DefKey> { + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { + SapTy::Map { + key: Box::new(self.key.to_baml_ty(db)), + value: Box::new(self.value.to_baml_ty(db)), + } + } +} + +impl ToBamlTy for ClassTy<'_, DefKey> { + fn to_baml_ty(&self, _db: &TypeRefDb<'_, DefKey>) -> SapTy { + SapTy::Class(self.name.clone(), Box::new([])) + } +} + +impl ToBamlTy for EnumTy<'_, DefKey> { + fn to_baml_ty(&self, _db: &TypeRefDb<'_, DefKey>) -> SapTy { + SapTy::Enum(self.name.clone()) + } +} + +impl ToBamlTy for EnumVariantTy<'_, DefKey> { + /// Loses variant specificity — maps back to the parent enum type. + fn to_baml_ty(&self, _db: &TypeRefDb<'_, DefKey>) -> SapTy { + SapTy::Enum(self.name.clone()) + } +} + +impl ToBamlTy for UnionTy<'_, DefKey> { + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { + let members: Vec = self.variants.iter().map(|v| v.to_baml_ty(db)).collect(); + SapTy::Union(members.into()) + } +} + +impl ToBamlTy for StreamStateTy<'_, DefKey> { + /// `StreamState` is a value-level concept; at the type level we return the inner type. + fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { + self.value.to_baml_ty(db) + } +} + +impl MediaTy { + fn to_baml_media_kind(self) -> baml_type::MediaKind { + match self { + MediaTy::Image => baml_type::MediaKind::Image, + MediaTy::Audio => baml_type::MediaKind::Audio, + MediaTy::Pdf => baml_type::MediaKind::Pdf, + MediaTy::Video => baml_type::MediaKind::Video, + } + } +} diff --git a/baml_language/crates/bex_sap/src/to_external.rs b/baml_language/crates/bex_sap/src/to_external.rs deleted file mode 100644 index e914efc9a35..00000000000 --- a/baml_language/crates/bex_sap/src/to_external.rs +++ /dev/null @@ -1,275 +0,0 @@ -//! Conversion from SAP types and values to `BexExternalValue` and `SapTy`. -//! -//! This module provides the bridge between SAP's internal type/value representation -//! and the external value tree used by the engine to push results back onto the VM heap. - -use bex_external_types::BexExternalValue; -use indexmap::IndexMap; -use sys_types::{DefKey, SapTy}; - -use crate::{ - baml_value::{BamlStreamState, BamlValue}, - deserializer::types::BamlValueWithFlags, - sap_model::{ - ArrayTy, ClassTy, EnumTy, EnumVariantTy, MapTy, MediaTy, StreamStateTy, Ty, TyResolvedRef, - TypeRefDb, UnionTy, - }, -}; - -// ============================================================================ -// Type conversion: SAP types → SapTy -// ============================================================================ - -/// Convert a SAP type back to a `SapTy`. -/// -/// # Known simplifications (not round-trip safe) -/// -/// - `Optional(T)` was converted to `Union([Null, T])` on the way in; it stays as `Union` here. -/// - Aliases already resolved into the target type have no nominal name to recover; named alias -/// references inside containers are preserved. -/// - `EnumVariant` becomes `Enum` (variant specificity lost at the type level). -pub trait ToBamlTy { - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy; -} - -impl ToBamlTy for TyResolvedRef<'_, DefKey> { - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { - match self { - TyResolvedRef::Int(_) => SapTy::Int, - TyResolvedRef::Bigint(_) => SapTy::Bigint, - TyResolvedRef::Float(_) => SapTy::Float, - TyResolvedRef::String(_) => SapTy::String, - TyResolvedRef::Bool(_) => SapTy::Bool, - TyResolvedRef::Null(_) => SapTy::Null, - TyResolvedRef::Media(media) => SapTy::Media(media.to_baml_media_kind()), - TyResolvedRef::LiteralInt(v) => { - SapTy::Literal(baml_type::Literal::Int(v.0), baml_type::Freshness::Regular) - } - TyResolvedRef::LiteralBigint(v) => SapTy::Literal( - baml_type::Literal::Bigint(v.0.clone()), - baml_type::Freshness::Regular, - ), - TyResolvedRef::LiteralString(v) => SapTy::Literal( - baml_type::Literal::String(v.0.to_string()), - baml_type::Freshness::Regular, - ), - TyResolvedRef::LiteralBool(v) => { - SapTy::Literal(baml_type::Literal::Bool(v.0), baml_type::Freshness::Regular) - } - TyResolvedRef::Array(a) => a.to_baml_ty(db), - TyResolvedRef::Map(m) => m.to_baml_ty(db), - TyResolvedRef::Class(c) => c.to_baml_ty(db), - TyResolvedRef::Enum(e) => e.to_baml_ty(db), - TyResolvedRef::EnumVariant(ev) => ev.to_baml_ty(db), - TyResolvedRef::Union(u) => u.to_baml_ty(db), - TyResolvedRef::StreamState(s) => s.to_baml_ty(db), - } - } -} - -impl ToBamlTy for Ty<'_, DefKey> { - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { - match self { - Ty::Resolved(resolved) => resolved.as_ref().to_baml_ty(db), - Ty::ResolvedRef(resolved_ref) => resolved_ref.to_baml_ty(db), - Ty::Unresolved(name) => match db.resolve_name(name) { - Some(TyResolvedRef::Class(class)) if class.name == *name => class.to_baml_ty(db), - Some(TyResolvedRef::Enum(enm)) if enm.name == *name => enm.to_baml_ty(db), - // Any other named entry is an alias. Keep it nominal here - // instead of recursively expanding aliases such as `json`. - Some(_) => SapTy::TypeAlias(name.clone()), - None => SapTy::unknown(), - }, - } - } -} - -impl ToBamlTy for ArrayTy<'_, DefKey> { - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { - let inner = self.ty.to_baml_ty(db); - SapTy::List(Box::new(inner)) - } -} - -impl ToBamlTy for MapTy<'_, DefKey> { - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { - SapTy::Map { - key: Box::new(self.key.to_baml_ty(db)), - value: Box::new(self.value.to_baml_ty(db)), - } - } -} - -impl ToBamlTy for ClassTy<'_, DefKey> { - fn to_baml_ty(&self, _db: &TypeRefDb<'_, DefKey>) -> SapTy { - SapTy::Class(self.name.clone(), Box::new([])) - } -} - -impl ToBamlTy for EnumTy<'_, DefKey> { - fn to_baml_ty(&self, _db: &TypeRefDb<'_, DefKey>) -> SapTy { - SapTy::Enum(self.name.clone()) - } -} - -impl ToBamlTy for EnumVariantTy<'_, DefKey> { - /// Loses variant specificity — maps back to the parent enum type. - fn to_baml_ty(&self, _db: &TypeRefDb<'_, DefKey>) -> SapTy { - SapTy::Enum(self.name.clone()) - } -} - -impl ToBamlTy for UnionTy<'_, DefKey> { - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { - let members: Vec = self.variants.iter().map(|v| v.to_baml_ty(db)).collect(); - SapTy::Union(members.into()) - } -} - -impl ToBamlTy for StreamStateTy<'_, DefKey> { - /// `StreamState` is a value-level concept; at the type level we return the inner type. - fn to_baml_ty(&self, db: &TypeRefDb<'_, DefKey>) -> SapTy { - self.value.to_baml_ty(db) - } -} - -impl MediaTy { - fn to_baml_media_kind(self) -> baml_type::MediaKind { - match self { - MediaTy::Image => baml_type::MediaKind::Image, - MediaTy::Audio => baml_type::MediaKind::Audio, - MediaTy::Pdf => baml_type::MediaKind::Pdf, - MediaTy::Video => baml_type::MediaKind::Video, - } - } -} - -// ============================================================================ -// Value conversion: SAP BamlValueWithFlags → BexExternalValue -// ============================================================================ - -/// Convert a lane type to the name-headed form `BexExternalValue` carries. -/// -/// The value metadata a parsed result carries has nowhere to put a declaration -/// identity, so it converts here — through the same per-call overlay spelling -/// the definition tables for this call were built with, so a returned value's -/// type names agree with the ones that call rendered. An anonymous declaration -/// has no resolvable name; its overlay spelling is meaningful only inside this -/// call and must not be treated as a lookup key beyond it. -fn to_value_metadata_ty(ty: &SapTy) -> baml_type::RuntimeTy { - match ty.try_map_heads(&mut |head: &DefKey| { - Ok::<_, std::convert::Infallible>(head.name().overlay_name()) - }) { - Ok(converted) => converted, - Err(never) => match never {}, - } -} - -/// Convert a SAP `BamlValueWithFlags` to a `BexExternalValue`, using the -/// type metadata from the deserializer to populate type information on -/// composite values (arrays, maps). -/// -/// Deserializer metadata (flags, scores) is discarded. -pub fn baml_value_to_external( - value: &BamlValueWithFlags<'_, '_, '_, DefKey>, - db: &TypeRefDb<'_, DefKey>, -) -> BexExternalValue { - baml_value_inner_to_external(&value.value, value.meta.ty, db) -} - -fn baml_value_inner_to_external( - value: &BamlValue<'_, '_, '_, DefKey>, - ty: TyResolvedRef<'_, DefKey>, - db: &TypeRefDb<'_, DefKey>, -) -> BexExternalValue { - match value { - BamlValue::String(s) => BexExternalValue::String(bex_str::BexStr::from(s.value.as_ref())), - BamlValue::Int(i) => BexExternalValue::Int(i.value), - BamlValue::Bigint(b) => BexExternalValue::Bigint(b.value.clone()), - BamlValue::Float(f) => BexExternalValue::Float(f.value), - BamlValue::Bool(b) => BexExternalValue::Bool(b.value), - BamlValue::Null(_) => BexExternalValue::Null, - BamlValue::Media(_) => { - unimplemented!("Media value conversion to BexExternalValue is not yet implemented") - } - BamlValue::Array(arr) => { - let element_type = to_value_metadata_ty(&match ty { - TyResolvedRef::Array(a) => a.ty.to_baml_ty(db), - _ => SapTy::unknown(), - }); - let items: Vec = arr - .value - .iter() - .map(|item| baml_value_to_external(item, db)) - .collect(); - BexExternalValue::Array { - element_type, - items, - } - } - BamlValue::Map(map) => { - let (key_type, value_type) = match ty { - TyResolvedRef::Map(m) => (m.key.to_baml_ty(db), m.value.to_baml_ty(db)), - _ => (SapTy::string(), SapTy::unknown()), - }; - let (key_type, value_type) = ( - to_value_metadata_ty(&key_type), - to_value_metadata_ty(&value_type), - ); - let entries: IndexMap = map - .value - .iter() - .map(|(k, v)| (k.to_string(), baml_value_to_external(v, db))) - .collect(); - BexExternalValue::Map { - key_type, - value_type, - entries, - } - } - BamlValue::Enum(e) => BexExternalValue::Variant { - // As for classes above: the qualified overlay spelling. - enum_name: e.name.name().overlay_name().to_string(), - variant_name: e.value.to_string(), - }, - BamlValue::Class(c) => { - let fields: IndexMap = c - .value - .iter() - .map(|(k, v)| (k.to_string(), baml_value_to_external(v, db))) - .collect(); - BexExternalValue::Instance { - // The landing resolves this against the per-call handle view, - // which is keyed by the *qualified* overlay spelling — so emit - // that, not the display name. One declaration has several name - // forms; the two sides of this hand-off must pick the same one. - class_name: c.name.name().overlay_name().to_string(), - type_args: vec![], - fields, - } - } - BamlValue::StreamState(state) => stream_state_to_external(state, db), - } -} - -fn stream_state_to_external( - state: &BamlStreamState<'_, '_, '_, DefKey>, - db: &TypeRefDb<'_, DefKey>, -) -> BexExternalValue { - let (inner, state_name) = match state { - BamlStreamState::Pending(v) => (v, "Pending"), - BamlStreamState::Incomplete(v) => (v, "Incomplete"), - BamlStreamState::Complete(v) => (v, "Complete"), - }; - let mut fields = IndexMap::new(); - fields.insert("value".to_string(), baml_value_to_external(inner, db)); - fields.insert( - "state".to_string(), - BexExternalValue::String(bex_str::BexStr::from(state_name)), - ); - BexExternalValue::Instance { - class_name: "StreamState".to_string(), - type_args: vec![], - fields, - } -} diff --git a/baml_language/crates/bex_vm/Cargo.toml b/baml_language/crates/bex_vm/Cargo.toml index 578f63f1d02..3195c19c78a 100644 --- a/baml_language/crates/bex_vm/Cargo.toml +++ b/baml_language/crates/bex_vm/Cargo.toml @@ -24,6 +24,7 @@ baml_compiler_diagnostics = { workspace = true } baml_linker_types = { workspace = true } baml_type = { workspace = true } bex_heap = { workspace = true } +bex_sap = { workspace = true } bex_str = { workspace = true } bex_vm_types = { workspace = true } btel_clock = { workspace = true } diff --git a/baml_language/crates/bex_vm/src/definitions.rs b/baml_language/crates/bex_vm/src/definitions.rs new file mode 100644 index 00000000000..2636b9af0ab --- /dev/null +++ b/baml_language/crates/bex_vm/src/definitions.rs @@ -0,0 +1,181 @@ +//! Class and enum definitions ([`sys_types::ClassDefinition`], +//! [`sys_types::EnumDefinition`]) read out of heap declarations. +//! +//! The one reader behind every definition table: the engine's program-wide +//! tables, its per-call runtime-type and runtime-package overlays, and +//! schema-aligned parsing. A class keeps every field with its `skip` flag, and +//! each consumer (output-format rendering, the JSON schema, the parser) leaves +//! skipped fields out itself. An enum leaves out skipped variants, which no +//! consumer reads. + +use bex_vm_types::{Class, Enum, RuntimeTy, TyTemplate, TypeHead}; +use sys_types::{ + ClassDefinition, ClassFieldDefinition, EnumDefinition, EnumVariantDefinition, SapTy, + SapTyTemplate, +}; + +/// The definition of `class`, every field included. +#[must_use] +pub fn class_definition(class: &Class) -> ClassDefinition { + ClassDefinition { + name: class.name.display_name().to_string(), + description: class.description.clone(), + docstring: class.docstring.clone(), + alias: class.alias.clone(), + stream_done: class.stream_done, + fields: class + .fields + .iter() + .map(|field| ClassFieldDefinition { + name: field.name.clone(), + field_type: lane_ty(&field.field_type), + field_template: Some(lane_template(&field.field_template)), + description: field.description.clone(), + docstring: field.docstring.clone(), + alias: field.alias.clone(), + skip: field.skip, + stream_done: field.stream_done, + must_exist: field.must_exist, + }) + .collect(), + } +} + +/// The definition of `enm`, without its skipped variants. +#[must_use] +pub fn enum_definition(enm: &Enum) -> EnumDefinition { + EnumDefinition { + name: enm.name.display_name().to_string(), + description: enm.description.clone(), + docstring: enm.docstring.clone(), + alias: enm.alias.clone(), + variants: enm + .variants + .iter() + .filter(|variant| !variant.skip) + .map(|variant| EnumVariantDefinition { + name: variant.name.clone(), + description: variant.description.clone(), + docstring: variant.docstring.clone(), + alias: variant.alias.clone(), + }) + .collect(), + } +} + +/// A declaration's type with each head as the tagged name it identifies. +/// +/// Total by invariant: a live declaration's heads are resolved pointers to +/// declarations, so each yields its identity and its own name. +#[must_use] +pub fn lane_ty(ty: &RuntimeTy) -> SapTy { + ty.try_map_heads(&mut TypeHead::to_tagged_name) + .unwrap_or_else(|head| { + unreachable!("a live declaration's field names a declaration: {head}") + }) +} + +/// [`lane_ty`] for a field's symbolic template, by the same invariant. +#[must_use] +pub fn lane_template(template: &TyTemplate) -> SapTyTemplate { + template + .try_map_heads(&mut TypeHead::to_tagged_name) + .unwrap_or_else(|head| { + unreachable!("a live declaration's field template names a declaration: {head}") + }) +} + +#[cfg(test)] +mod tests { + use baml_type::{DeclarationName, Name, typetag::TypeTag}; + use bex_vm_types::{ClassField, EnumVariant, types::Owner}; + use indexmap::IndexMap; + + use super::*; + + fn field( + name: &str, + field_type: RuntimeTy, + field_template: TyTemplate, + skip: bool, + ) -> ClassField { + ClassField { + name: name.to_string(), + field_type, + field_template, + description: None, + alias: None, + docstring: None, + other: IndexMap::new(), + skip, + stream_done: false, + must_exist: false, + runtime_type: None, + } + } + + fn variant(name: &str, skip: bool) -> EnumVariant { + EnumVariant { + name: name.to_string(), + description: None, + alias: None, + docstring: None, + other: IndexMap::new(), + skip, + } + } + + #[test] + fn class_definition_keeps_skipped_fields_flagged() { + let class = Class { + name: DeclarationName::Anonymous(Name::new("Person")), + fields: vec![ + field("name", RuntimeTy::String, TyTemplate::String, false), + field("age", RuntimeTy::Int, TyTemplate::Int, true), + ], + description: None, + alias: None, + docstring: None, + other: IndexMap::new(), + stream_done: false, + type_tag: TypeTag::fresh_dynamic(), + has_cleanup: false, + generic_param_count: 0, + owner: Owner::anonymous(), + methods: IndexMap::new(), + }; + + let definition = class_definition(&class); + + let fields: Vec<(&str, bool)> = definition + .fields + .iter() + .map(|field| (field.name.as_str(), field.skip)) + .collect(); + assert_eq!(fields, [("name", false), ("age", true)]); + assert_eq!(definition.fields[1].field_type, SapTy::Int); + } + + #[test] + fn enum_definition_leaves_out_skipped_variants() { + let enm = Enum { + name: DeclarationName::Anonymous(Name::new("Color")), + type_tag: TypeTag::fresh_dynamic(), + variants: vec![variant("Red", false), variant("Hidden", true)], + description: None, + alias: None, + docstring: None, + other: IndexMap::new(), + owner: Owner::anonymous(), + }; + + let definition = enum_definition(&enm); + + let names: Vec<&str> = definition + .variants + .iter() + .map(|variant| variant.name.as_str()) + .collect(); + assert_eq!(names, ["Red"]); + } +} diff --git a/baml_language/crates/bex_vm/src/lib.rs b/baml_language/crates/bex_vm/src/lib.rs index b14ed31d89d..6602b630454 100644 --- a/baml_language/crates/bex_vm/src/lib.rs +++ b/baml_language/crates/bex_vm/src/lib.rs @@ -14,6 +14,7 @@ pub(crate) mod array_index; mod call_specialize; pub mod compile_artifact; pub mod debug; +pub mod definitions; pub mod errors; pub mod indexable; pub mod kperf; diff --git a/baml_language/crates/bex_vm/src/package_baml/csv.rs b/baml_language/crates/bex_vm/src/package_baml/csv.rs index 2b31a2dbbb8..89ab0a6d292 100644 --- a/baml_language/crates/bex_vm/src/package_baml/csv.rs +++ b/baml_language/crates/bex_vm/src/package_baml/csv.rs @@ -190,34 +190,6 @@ fn done_value(vm: &mut BexVm) -> Result { // Handle-state access // ============================================================================= -/// Clone the `Arc` out of a `$rust_type` field so the caller can lock it while -/// still holding `&mut BexVm` for allocation. -fn state_arc( - vm: &BexVm, - holder: Value, - field_idx: usize, -) -> Result, VmRustFnError> { - let inst = vm.as_instance(&holder)?; - let handle = inst.load_field(field_idx); - let ptr = handle - .as_object_ptr() - .ok_or_else(|| VmInternalError::MissingNativeFunction { - name: "csv handle field is not an object".to_string(), - })?; - match vm.get_object(ptr) { - Object::RustData(arc) => arc.clone().downcast::().map_err(|_| { - VmRustFnError::InternalError(VmInternalError::MissingNativeFunction { - name: "csv handle holds an unexpected Rust type".to_string(), - }) - }), - _ => Err(VmRustFnError::InternalError( - VmInternalError::MissingNativeFunction { - name: "csv handle field is not RustData".to_string(), - }, - )), - } -} - fn lock(m: &Mutex) -> MutexGuard<'_, T> { m.lock().unwrap_or_else(std::sync::PoisonError::into_inner) } @@ -1490,7 +1462,7 @@ impl From for DecodeFail { } fn record_arc(vm: &BexVm, rec: Value) -> Result, VmRustFnError> { - state_arc::(vm, rec, 0) + Ok(vm.rust_data_field::(&rec, 0)?) } /// Decode a whole record into an instance of class `ty`. @@ -2088,7 +2060,7 @@ fn render_markdown(headers: &[String], rows: &[Vec], total_rows: usize) impl BamlClassCsvReader for PackageBamlImpl { fn skipped(vm: &mut BexVm, csvreader: &Value) -> Vec { - let Ok(st) = state_arc::>(vm, *csvreader, 0) else { + let Ok(st) = vm.rust_data_field::>(csvreader, 0) else { return Vec::new(); }; let infos: Vec = lock(&st).skipped.clone(); @@ -2103,7 +2075,7 @@ impl BamlClassCsvReader for PackageBamlImpl { } fn position(vm: &mut BexVm, csvreader: &Value) -> Value { - match state_arc::>(vm, *csvreader, 0) { + match vm.rust_data_field::>(csvreader, 0) { Ok(st) => { let (byte, line, record) = { let s = lock(&st); @@ -2121,7 +2093,7 @@ impl BamlClassCsvReader for PackageBamlImpl { } fn _poll(vm: &mut BexVm, csvreader: &Value) -> Result { - let st = state_arc::>(vm, *csvreader, 0)?; + let st = vm.rust_data_field::>(csvreader, 0)?; let polled = { let mut s = lock(&st); poll_record(&mut s) @@ -2142,7 +2114,7 @@ impl BamlClassCsvReader for PackageBamlImpl { } fn _poll_headers(vm: &mut BexVm, csvreader: &Value) -> Result { - let st = state_arc::>(vm, *csvreader, 0)?; + let st = vm.rust_data_field::>(csvreader, 0)?; let polled = { let mut s = lock(&st); poll_headers(&mut s) @@ -2257,7 +2229,7 @@ impl BamlClassCsvWriter for PackageBamlImpl { } fn text(vm: &mut BexVm, csvwriter: &Value) -> Result { - let st = state_arc::>(vm, *csvwriter, 0)?; + let st = vm.rust_data_field::>(csvwriter, 0)?; let text = { let s = lock(&st); s.buffer.clone() @@ -2279,7 +2251,7 @@ impl BamlClassCsvWriter for PackageBamlImpl { csvwriter: &Value, record: &[Value], ) -> Result { - let st = state_arc::>(vm, *csvwriter, 0)?; + let st = vm.rust_data_field::>(csvwriter, 0)?; let mut s = lock(&st); if s.closed { drop(s); @@ -2314,7 +2286,7 @@ impl BamlClassCsvWriter for PackageBamlImpl { csvwriter: &Value, names: &[Value], ) -> Result { - let st = state_arc::>(vm, *csvwriter, 0)?; + let st = vm.rust_data_field::>(csvwriter, 0)?; let mut s = lock(&st); if s.closed { drop(s); @@ -2432,7 +2404,7 @@ impl BamlNamespaceCsv for PackageBamlImpl { } }; - let st = state_arc::>(vm, *r, 0)?; + let st = vm.rust_data_field::>(r, 0)?; let header = lock(&st).header.clone(); for cf in &class_fields { @@ -2480,7 +2452,7 @@ impl BamlNamespaceCsv for PackageBamlImpl { Err(DecodeFail::Fatal(e)) => Err(e), Err(DecodeFail::Info(info)) => { let skip = info.kind == Kind::Decode && { - let st = state_arc::>(vm, *r, 0)?; + let st = vm.rust_data_field::>(r, 0)?; let mut s = lock(&st); if s.opts.skip_on_error { s.register_skip(&info); @@ -2557,7 +2529,7 @@ impl BamlNamespaceCsv for PackageBamlImpl { w: &Value, row: &Value, ) -> Result { - let st = state_arc::>(vm, *w, 0)?; + let st = vm.rust_data_field::>(w, 0)?; let mut s = lock(&st); if s.closed { drop(s); @@ -2584,7 +2556,7 @@ impl BamlNamespaceCsv for PackageBamlImpl { w: &Value, rows: &[Value], ) -> Result { - let st = state_arc::>(vm, *w, 0)?; + let st = vm.rust_data_field::>(w, 0)?; let mut s = lock(&st); if s.closed { drop(s); diff --git a/baml_language/crates/bex_vm/src/package_baml/mod.rs b/baml_language/crates/bex_vm/src/package_baml/mod.rs index d003fffa7fe..f7bcdb3399d 100644 --- a/baml_language/crates/bex_vm/src/package_baml/mod.rs +++ b/baml_language/crates/bex_vm/src/package_baml/mod.rs @@ -44,6 +44,7 @@ mod regex; pub(crate) mod resolve; pub(crate) use resolve::ImplResolver; pub(crate) mod root; +mod sap; mod spawn; pub use spawn::{ SpawnLaunch, alloc_projected_cancel_token, cancel_token_members, plan_body, diff --git a/baml_language/crates/bex_vm/src/package_baml/regex.rs b/baml_language/crates/bex_vm/src/package_baml/regex.rs index 08bae5a897c..8729f467930 100644 --- a/baml_language/crates/bex_vm/src/package_baml/regex.rs +++ b/baml_language/crates/bex_vm/src/package_baml/regex.rs @@ -215,24 +215,7 @@ fn regex_value(vm: &mut BexVm, prog: Program) -> Value { /// Clone the `Arc` out of `Regex._handle` so the program stays alive while /// `&mut BexVm` is borrowed for allocation. fn program_of(vm: &BexVm, regex: Value) -> Result, VmRustFnError> { - let instance = vm.as_instance(®ex)?; - let ptr = instance.load_field(0).as_object_ptr().ok_or_else(|| { - VmInternalError::MissingNativeFunction { - name: "baml.regex.Regex._handle is not an object".to_string(), - } - })?; - match vm.get_object(ptr) { - Object::RustData(data) => data.clone().downcast::().map_err(|_| { - VmRustFnError::InternalError(VmInternalError::MissingNativeFunction { - name: "baml.regex.Regex._handle holds an unexpected Rust type".to_string(), - }) - }), - _ => Err(VmRustFnError::InternalError( - VmInternalError::MissingNativeFunction { - name: "baml.regex.Regex._handle is not RustData".to_string(), - }, - )), - } + Ok(vm.rust_data_field::(®ex, 0)?) } /// A search that ran out of budget is a panic, not a "no match". diff --git a/baml_language/crates/bex_vm/src/package_baml/sap.rs b/baml_language/crates/bex_vm/src/package_baml/sap.rs new file mode 100644 index 00000000000..172a307c53f --- /dev/null +++ b/baml_language/crates/bex_vm/src/package_baml/sap.rs @@ -0,0 +1,503 @@ +//! `baml.sap`: schema-aligned parsing (SAP) of an LLM's text into a declared +//! type. +//! +//! VM natives, not sys ops: parsing does no IO, so it runs on the calling VM +//! with the heap at hand. `_new_parse_cache` builds the parse model from the +//! declarations `T` reaches, and the parse methods allocate their result +//! directly on the heap, so no definition table is copied and no value tree is +//! landed by name. + +use std::{any::Any, collections::HashMap, sync::Arc}; + +use baml_type::{DeclarationName, Literal, typetag::TypeTag}; +use bex_heap::TlabHolder; +use bex_sap::{ + CompiledSapModel, + baml_value::BamlValue, + deserializer::{coercer::ParsingContext, types::BamlValueWithFlags}, + jsonish, + sap_model::{TyResolvedRef, TypeCtx, TypeRefDb}, + to_baml_ty::ToBamlTy, +}; +use bex_str::BexStr; +use bex_vm_types::{ + HeapPtr, RealizedTy, RuntimeTy, TypeHead, + types::{Instance, Object, Value}, +}; +use indexmap::IndexMap; +use sys_types::{DefKey, SapTy}; + +use super::{BamlNamespaceSap, BamlNamespaceSap_ParseCache, PackageBamlImpl}; +use crate::{ + BexVm, definitions, + errors::{VmBamlError, VmInternalError, VmPanic, VmRustFnError}, + reachable, +}; + +/// The `_data` of a `baml.sap._ParseCache`: the parse model for `T`. +/// +/// A `RustData` payload is a GC leaf, so the model holds no heap pointers: it +/// names declarations by tagged name, and each parse maps those names back to +/// the declarations through [`Declarations`]. +struct ParseCache { + model: CompiledSapModel, +} + +impl BamlNamespaceSap for PackageBamlImpl { + fn _new_parse_cache(vm: &mut BexVm) -> Result { + let target = parse_target(vm)?; + let model = build_model(vm, &target).map_err(|e| { + // `throws never`: a `T` the parser cannot model comes from the + // caller's own `parse`, a program bug rather than bad input. + VmPanic::UserPanic { + message: format!("schema-aligned parsing cannot model this type: {e}"), + } + })?; + let data: Arc = Arc::new(ParseCache { model }); + let data = Value::object(vm.alloc_rust_data(data)); + let class = vm.resolve_class("baml.sap._ParseCache"); + let cache = vm.tlab.alloc(Object::Instance(Instance::new( + class, + Box::new([target]), + vec![data], + ))); + Ok(Value::object(cache)) + } +} + +impl BamlNamespaceSap_ParseCache for PackageBamlImpl { + fn _parse_final(vm: &mut BexVm, cache: &Value, json: &BexStr) -> Result { + let parsed = parse(vm, *cache, json, Completion::Final)?; + parsed.ok_or_else(|| llm_client("SAP parse returned no value when complete")) + } + + fn _parse_partial( + vm: &mut BexVm, + cache: &Value, + json: &BexStr, + ) -> Result { + match parse(vm, *cache, json, Completion::Partial)? { + Some(value) => Ok(value), + None => { + let no_yield = vm.resolve_class("baml.sap._NoYield"); + Ok(Value::object(vm.alloc_instance(no_yield, Vec::new()))) + } + } + } +} + +/// Whether the text is the whole reply, or the part a stream has received. +#[derive(Clone, Copy, PartialEq, Eq)] +enum Completion { + Final, + Partial, +} + +/// The `T` of the current `_new_parse_cache` or `_ParseCache` method +/// call: the first type argument either way. +fn parse_target(vm: &BexVm) -> Result { + vm.current_call_type_args().first().cloned().ok_or_else(|| { + VmInternalError::SapValue { + message: "a `baml.sap` native was called without its `T`".to_string(), + } + .into() + }) +} + +/// The declarations the parse model of `target` needs: those `target` names, +/// then those a reached class's field types name (except `@skip` fields, which +/// the parser drops) and those a reached alias's body names. +fn parsed_declarations(vm: &BexVm, target: &RealizedTy) -> Vec { + reachable::declarations_reached(vm, target, |object, visit| match object { + Object::Class(class) => { + for field in class.fields.iter().filter(|field| !field.skip) { + field.field_type.visit_heads(&mut |head| visit(head)); + } + } + Object::TypeAlias(alias) => alias.definition.visit_heads(&mut |head| visit(head)), + _ => {} + }) +} + +/// The parse model for `target`, built from the declarations it reaches. +fn build_model( + vm: &BexVm, + target: &RealizedTy, +) -> Result { + let mut classes = IndexMap::new(); + let mut enums = IndexMap::new(); + let mut aliases = HashMap::new(); + for ptr in parsed_declarations(vm, target) { + match vm.get_object(ptr) { + Object::Class(class) => { + classes.insert( + DefKey::new(class.type_tag, class.name.clone()), + definitions::class_definition(class), + ); + } + Object::Enum(enm) => { + enums.insert( + DefKey::new(enm.type_tag, enm.name.clone()), + definitions::enum_definition(enm), + ); + } + Object::TypeAlias(alias) => { + aliases.insert( + DefKey::new( + alias.type_tag, + DeclarationName::Declared(alias.name.clone()), + ), + definitions::lane_ty(&RuntimeTy::from(alias.definition.clone())), + ); + } + _ => {} + } + } + let target = definitions::lane_ty(&RuntimeTy::from(target.clone())); + CompiledSapModel::from_type_ctx(TypeCtx::new(&classes, enums, &aliases), target) +} + +/// Parse `text` with the model in `cache` and allocate the result, or `None` +/// for a partial reply that does not parse yet. +fn parse( + vm: &mut BexVm, + cache: Value, + text: &BexStr, + completion: Completion, +) -> Result, VmRustFnError> { + let cache = vm.rust_data_field::(&cache, 0)?; + let target_type = parse_target(vm)?; + let model = &cache.model; + // An unresolvable target is a schema bug, not an incomplete reply. + let target = model + .resolved_target() + .map_err(|e| llm_client(e.to_string()))?; + let is_done = completion == Completion::Final; + let value = match jsonish::parse(text.as_str(), jsonish::ParseOptions::default(), is_done) { + Ok(value) => value, + Err(e) if is_done => return Err(llm_client(e.to_string())), + // A reply that opens with a fence or with prose before the JSON cannot + // be coerced until the JSON starts: skip the partial, wait for text. + Err(_) => return Ok(None), + }; + let context = ParsingContext::new(model.db()); + let parsed = match TyResolvedRef::coerce(&context, target, &value) { + Ok(parsed) => parsed, + Err(e) if is_done => return Err(llm_client(e.to_string())), + Err(_) => return Ok(None), + }; + // Only a value to allocate needs the declarations: a partial that does not + // parse yet skips the walk. + parsed + .map(|parsed| { + let declarations = Declarations::reached_by(vm, &target_type); + to_heap(vm, &parsed, model.db(), &declarations) + }) + .transpose() +} + +/// A parse failure, as `baml.errors.LlmClient`. +fn llm_client(message: impl Into) -> VmRustFnError { + VmBamlError::LlmClient { + message: message.into(), + } + .into() +} + +/// The declarations a parse target reaches, by tag, for mapping the model's +/// tagged names back to the heap. Built for one call: a parse never yields, so +/// no collection can move them meanwhile. +struct Declarations(HashMap); + +impl Declarations { + fn reached_by(vm: &BexVm, target: &RealizedTy) -> Self { + Self( + parsed_declarations(vm, target) + .into_iter() + .filter_map(|ptr| { + let tag = match vm.get_object(ptr) { + Object::Class(class) => class.type_tag, + Object::Enum(enm) => enm.type_tag, + Object::TypeAlias(alias) => alias.type_tag, + _ => return None, + }; + Some((tag, ptr)) + }) + .collect(), + ) + } + + fn get(&self, name: &DefKey) -> Result { + self.0.get(&name.tag()).copied().ok_or_else(|| { + invariant(format!( + "the parsed value names `{name}`, which its target does not reach" + )) + }) + } + + /// `ty` with each tagged name replaced by the head of its declaration. + fn realize(&self, ty: &SapTy) -> Result { + let ty = ty.try_map_heads(&mut |name: &DefKey| { + self.get(name).map(|ptr| TypeHead::new(ptr, name.tag())) + })?; + RealizedTy::try_from(ty) + .map_err(|e| invariant(format!("a parsed value's type is not realized: {e}"))) + } +} + +fn invariant(message: String) -> VmRustFnError { + VmInternalError::SapValue { message }.into() +} + +/// Allocate `value`: the same value the parser produced, on the heap. +fn to_heap( + vm: &mut BexVm, + value: &BamlValueWithFlags<'_, '_, '_, DefKey>, + db: &TypeRefDb<'_, DefKey>, + declarations: &Declarations, +) -> Result { + Ok(match &value.value { + BamlValue::String(s) => Value::object(vm.alloc_string(s.value.as_ref())), + BamlValue::Int(i) => Value::try_int(i.value).ok_or_else(|| { + llm_client(format!( + "integer {} is outside the BAML integer range [{}, {}]", + i.value, + Value::INT_MIN, + Value::INT_MAX + )) + })?, + BamlValue::Bigint(b) => vm.try_alloc_bigint(Arc::new(b.value.clone()))?, + BamlValue::Float(f) => Value::object(vm.alloc_float(f.value)), + BamlValue::Bool(b) => Value::bool(b.value), + BamlValue::Null(_) => Value::NULL, + BamlValue::Media(_) => { + return Err(invariant( + "the parser produced a media value, which it cannot parse".to_string(), + )); + } + BamlValue::Array(array) => { + let element = match value.meta.ty { + TyResolvedRef::Array(array) => array.ty.to_baml_ty(db), + _ => SapTy::unknown(), + }; + let element = declarations.realize(&element)?; + let items = array + .value + .iter() + .map(|item| to_heap(vm, item, db, declarations)) + .collect::, _>>()?; + Value::object(vm.alloc_array(element, items)) + } + BamlValue::Map(map) => { + let (key, value_ty) = match value.meta.ty { + TyResolvedRef::Map(map) => (map.key.to_baml_ty(db), map.value.to_baml_ty(db)), + _ => (SapTy::string(), SapTy::unknown()), + }; + let (key, value_ty) = ( + declarations.realize(&key)?, + declarations.realize(&value_ty)?, + ); + let mut entries = IndexMap::with_capacity(map.value.len()); + for (k, v) in &map.value { + entries.insert(BexStr::from(k.as_ref()), to_heap(vm, v, db, declarations)?); + } + Value::object(vm.alloc_map(key, value_ty, entries)) + } + BamlValue::Enum(e) => { + let enm = declarations.get(e.name)?; + let Object::Enum(declaration) = vm.get_object(enm) else { + return Err(invariant(format!("`{}` is not an enum", e.name))); + }; + let index = declaration + .variants + .iter() + .position(|variant| variant.name == e.value) + .ok_or_else(|| { + invariant(format!("enum `{}` has no variant `{}`", e.name, e.value)) + })?; + Value::object(vm.alloc_variant(enm, index)) + } + BamlValue::Class(c) => { + let class = declarations.get(c.name)?; + let Object::Class(declaration) = vm.get_object(class) else { + return Err(invariant(format!("`{}` is not a class", c.name))); + }; + let fields: Vec<(String, bool, RuntimeTy)> = declaration + .fields + .iter() + .map(|field| (field.name.clone(), field.skip, field.field_type.clone())) + .collect(); + let mut values = Vec::with_capacity(fields.len()); + for (name, skip, field_type) in fields { + let field_value = match c.value.get(name.as_str()) { + Some(parsed) => to_heap(vm, parsed, db, declarations)?, + None if skip => skipped_field_default(vm, c.name, &name, &field_type)?, + None => { + return Err(invariant(format!( + "the parsed `{}` has no field `{name}`", + c.name + ))); + } + }; + values.push(field_value); + } + Value::object(vm.alloc_instance(class, values)) + } + BamlValue::StreamState(_) => { + return Err(invariant( + "the parser produced a stream state, which no target declares".to_string(), + )); + } + }) +} + +/// The value a `@skip` field holds in a parsed instance. The parser never reads +/// the field, so it gets its type's empty value. +fn skipped_field_default( + vm: &mut BexVm, + class: &DefKey, + field: &str, + field_type: &RuntimeTy, +) -> Result { + let no_default = |ty: &str| -> VmRustFnError { + VmPanic::UserPanic { + message: format!( + "the `@skip` field `{class}.{field}` has type {ty}, which has no default value \ + for schema-aligned parsing to fill in" + ), + } + .into() + }; + let realized = RealizedTy::try_from(field_type.clone()) + .map_err(|_| no_default("of a generic parameter"))?; + empty_value(vm, &realized, &mut Vec::new())?.ok_or_else(|| no_default(&realized.to_string())) +} + +/// The empty value of `ty`: `null` for a nullable type, the zero of a number, +/// `false`, `""`, an empty list or map, a literal's own value, an enum's first +/// variant, and a class of its fields' empty values. `None` when `ty` has none +/// (media, functions, a class that contains itself, ...). +/// +/// `within` holds the classes and aliases being built, to stop a type that +/// needs its own empty value. +fn empty_value( + vm: &mut BexVm, + ty: &RealizedTy, + within: &mut Vec, +) -> Result, VmRustFnError> { + Ok(Some(match ty { + RealizedTy::Null | RealizedTy::Unknown => Value::NULL, + RealizedTy::Int => Value::int(0), + RealizedTy::Bigint => vm.try_alloc_bigint(Arc::new(num_bigint::BigInt::default()))?, + RealizedTy::Float => Value::object(vm.alloc_float(0.0)), + RealizedTy::Bool => Value::bool(false), + RealizedTy::String => Value::object(vm.alloc_string("")), + RealizedTy::Uint8Array => Value::object(vm.alloc_uint8array(Vec::new())), + RealizedTy::Literal(literal, _) => match literal { + Literal::Int(n) => Value::int(*n), + Literal::Bigint(n) => vm.try_alloc_bigint(Arc::new(n.clone()))?, + Literal::Float(text) => match text.parse::() { + Ok(f) => Value::object(vm.alloc_float(f)), + Err(_) => return Ok(None), + }, + Literal::String(s) => Value::object(vm.alloc_string(s.as_str())), + Literal::Bool(b) => Value::bool(*b), + }, + RealizedTy::List(element) => Value::object(vm.alloc_array((**element).clone(), Vec::new())), + RealizedTy::Map { key, value } => { + Value::object(vm.alloc_map((**key).clone(), (**value).clone(), IndexMap::new())) + } + RealizedTy::Enum(head) => { + let Object::Enum(enm) = vm.get_object(head.ptr()) else { + return Ok(None); + }; + let Some(index) = enm.variants.iter().position(|variant| !variant.skip) else { + return Ok(None); + }; + Value::object(vm.alloc_variant(head.ptr(), index)) + } + RealizedTy::EnumVariant(head, name) => { + let Object::Enum(enm) = vm.get_object(head.ptr()) else { + return Ok(None); + }; + let Some(index) = enm + .variants + .iter() + .position(|variant| variant.name == name.as_str()) + else { + return Ok(None); + }; + Value::object(vm.alloc_variant(head.ptr(), index)) + } + RealizedTy::Union(members) => { + if members + .iter() + .any(|member| matches!(member, RealizedTy::Null)) + { + Value::NULL + } else { + let mut found = None; + for member in members { + if let Some(value) = empty_value(vm, member, within)? { + found = Some(value); + break; + } + } + return Ok(found); + } + } + RealizedTy::Class(head, type_args) => { + let class = head.ptr(); + if within.contains(&class) { + return Ok(None); + } + let Object::Class(declaration) = vm.get_object(class) else { + return Ok(None); + }; + let templates: Vec<_> = declaration + .fields + .iter() + .map(|field| field.field_template.clone()) + .collect(); + within.push(class); + let mut values = Vec::with_capacity(templates.len()); + for template in &templates { + let field_ty = vm.realize_field_ty(template, type_args); + let Some(value) = empty_value(vm, &field_ty, within)? else { + within.pop(); + return Ok(None); + }; + values.push(value); + } + within.pop(); + Value::object( + vm.tlab + .alloc_instance_with_type_args(class, type_args.clone(), values), + ) + } + RealizedTy::TypeAlias(head) => { + let alias = head.ptr(); + if within.contains(&alias) { + return Ok(None); + } + let Object::TypeAlias(declaration) = vm.get_object(alias) else { + return Ok(None); + }; + let definition = declaration.definition.clone(); + within.push(alias); + let value = empty_value(vm, &definition, within)?; + within.pop(); + return Ok(value); + } + RealizedTy::Media(_) + | RealizedTy::Function { .. } + | RealizedTy::Future(..) + | RealizedTy::Interface(..) + | RealizedTy::RustType + | RealizedTy::Type + | RealizedTy::Resource + | RealizedTy::PromptAst + | RealizedTy::Void + | RealizedTy::Never => return Ok(None), + })) +} diff --git a/baml_language/crates/bex_vm/src/reachable.rs b/baml_language/crates/bex_vm/src/reachable.rs index 3f7752881f6..66430e59c1e 100644 --- a/baml_language/crates/bex_vm/src/reachable.rs +++ b/baml_language/crates/bex_vm/src/reachable.rs @@ -121,14 +121,17 @@ pub fn runtime_nominals_under_permit( (classes, enums) } -/// Every class and enum `ty` reaches, including compile-time declarations. -/// -/// Most runtime consumers intentionally stop at the immutable program image, -/// but source rendering must emit a standalone graph: a runtime class field -/// that names a static class needs that static declaration in the output too. +/// Every declaration `ty` reaches when each reached declaration leads on to the +/// heads `follow` reports for it, in the order a depth-first walk meets them: +/// the heads of `ty`, then each declaration's followed heads. #[must_use] -pub fn all_nominals(vm: &BexVm, ty: &bex_vm_types::RealizedTy) -> (Vec, Vec) { +pub fn declarations_reached( + vm: &BexVm, + ty: &bex_vm_types::RealizedTy, + follow: impl Fn(&Object, &mut dyn FnMut(&bex_vm_types::TypeHead)), +) -> Vec { let mut found = Vec::new(); + let mut seen = std::collections::HashSet::new(); let mut pending = Vec::new(); ty.visit_heads(&mut |head| { if head.is_resolved() { @@ -137,13 +140,12 @@ pub fn all_nominals(vm: &BexVm, ty: &bex_vm_types::RealizedTy) -> (Vec, }); pending.reverse(); while let Some(ptr) = pending.pop() { - if found.contains(&ptr) { + if !seen.insert(ptr) { continue; } found.push(ptr); let mut next = Vec::new(); - let object = vm.get_object(ptr); - bex_vm_types::head_walk::visit_object_heads(object, &mut |head| { + follow(vm.get_object(ptr), &mut |head| { if head.is_resolved() { next.push(head.ptr()); } @@ -151,10 +153,29 @@ pub fn all_nominals(vm: &BexVm, ty: &bex_vm_types::RealizedTy) -> (Vec, next.reverse(); pending.extend(next); } + found +} +/// Every declaration `ty` reaches (classes, enums, interfaces, type aliases), +/// including compile-time ones, through every head a declaration holds (field +/// types, templates, alias bodies, ...). +#[must_use] +pub fn all_declarations(vm: &BexVm, ty: &bex_vm_types::RealizedTy) -> Vec { + declarations_reached(vm, ty, |object, visit| { + bex_vm_types::head_walk::visit_object_heads(object, &mut |head| visit(head)); + }) +} + +/// Every class and enum `ty` reaches, including compile-time declarations. +/// +/// Most runtime consumers intentionally stop at the immutable program image, +/// but source rendering must emit a standalone graph: a runtime class field +/// that names a static class needs that static declaration in the output too. +#[must_use] +pub fn all_nominals(vm: &BexVm, ty: &bex_vm_types::RealizedTy) -> (Vec, Vec) { let mut classes = Vec::new(); let mut enums = Vec::new(); - for ptr in found { + for ptr in all_declarations(vm, ty) { match vm.get_object(ptr) { Object::Class(_) => classes.push(ptr), Object::Enum(_) => enums.push(ptr), diff --git a/baml_language/crates/bex_vm/src/vm.rs b/baml_language/crates/bex_vm/src/vm.rs index cb6003d1e82..8a02525584b 100644 --- a/baml_language/crates/bex_vm/src/vm.rs +++ b/baml_language/crates/bex_vm/src/vm.rs @@ -4940,6 +4940,38 @@ impl BexVm { } } + /// Clone the `Arc` out of the `$rust_type` field `field` of the instance + /// `holder`, so a native can keep using the payload while it borrows + /// `&mut BexVm` to allocate. + pub fn rust_data_field( + &self, + holder: &Value, + field: usize, + ) -> Result, VmInternalError> { + let handle = self.as_instance(holder)?.load_field(field); + let Some(ptr) = handle.as_object_ptr() else { + return Err(VmInternalError::TypeError { + expected: Type::Object(ObjectType::RustData), + got: self.type_of(&handle), + }); + }; + match self.get_object(ptr) { + Object::RustData(arc) => { + let got = arc.as_ref().type_id(); + arc.clone() + .downcast::() + .map_err(|_| VmInternalError::RustTypeError { + expected: TypeId::of::(), + got, + }) + } + _ => Err(VmInternalError::TypeError { + expected: Type::Object(ObjectType::RustData), + got: self.type_of(&handle), + }), + } + } + /// Extract an `&Instance` from a `Value` carrying a heap-object pointer. /// /// Used by generated glue code to construct `view::` structs. diff --git a/baml_language/crates/bex_vm_types/src/errors.rs b/baml_language/crates/bex_vm_types/src/errors.rs index 452c13a106d..b04b02188a2 100644 --- a/baml_language/crates/bex_vm_types/src/errors.rs +++ b/baml_language/crates/bex_vm_types/src/errors.rs @@ -244,6 +244,10 @@ pub enum VmInternalError { #[error("missing native function: {name}")] MissingNativeFunction { name: String }, + /// Schema-aligned parsing produced a value its target's declarations + /// cannot hold: the parse model and the heap disagree. + #[error("schema-aligned parsing produced a value its target cannot hold: {message}")] + SapValue { message: String }, #[error("unexpected constant kind: expected a TyTemplate constant at this index")] UnexpectedConstantKind, diff --git a/baml_language/crates/sys_ops/Cargo.toml b/baml_language/crates/sys_ops/Cargo.toml index 93e2d6f7c8a..a2e5668183f 100644 --- a/baml_language/crates/sys_ops/Cargo.toml +++ b/baml_language/crates/sys_ops/Cargo.toml @@ -25,7 +25,6 @@ baml_type = { workspace = true } bex_external_types = { workspace = true } bex_heap = { workspace = true } bex_resource_types = { workspace = true } -bex_sap = { workspace = true } bex_vm_types = { workspace = true } sys_types = { workspace = true } indexmap = { workspace = true } diff --git a/baml_language/crates/sys_ops/src/lib.rs b/baml_language/crates/sys_ops/src/lib.rs index 022c7cc3d50..553839ed03c 100644 --- a/baml_language/crates/sys_ops/src/lib.rs +++ b/baml_language/crates/sys_ops/src/lib.rs @@ -94,7 +94,6 @@ pub use io_adapter::build_runtime_io; // ============================================================================ pub mod output_format; -pub mod sap; // ============================================================================ // Blanket IO LLM implementation @@ -150,9 +149,8 @@ fn llm_function_lookup_error( } } -/// Blanket impl — schema-aligned parsing, backing both public `baml.sap.parse` -/// and incremental `ai.stream.Stream` parsing. -/// +/// Blanket impl — `ai.OutputFormat`: rendering the output-format schema with +/// the caller's options. impl io::IoClassAiOutputFormat for T { #[allow(clippy::too_many_arguments)] fn _render( @@ -899,85 +897,6 @@ pub fn get_return_type_op( SysOpOutput::ok(info.return_type.clone()) } -/// Blanket impl — schema-aligned parsing, backing both public `baml.sap.parse` -/// and incremental `ai.stream.Stream` parsing. All three are free functions -/// (see `ns_sap/sap.baml`), so each carries its own `T` type-arg operand; the -/// cache already holds the compiled model, so the two parse entry points -/// ignore theirs. -impl io::IoClassSapParseCache for T { - fn _parse_final( - &self, - _heap: &std::sync::Arc, - _call_id: CallId, - cache: io::owned::sap::ParseCache, - json: String, - _type_arg_0: ::sys_types::SapTy, - ctx: &SysOpContext, - ) -> SysOpOutput { - let Ok(sap) = cache._data.clone().downcast::() else { - return SysOpOutput::err(VmInternalError::RustTypeError { - expected: std::any::TypeId::of::(), - got: cache._data.type_id(), - }); - }; - SysOpOutput::Ready( - crate::sap::execute_sap_parse_final(&json, &sap, ctx).map_err(VmRustFnError::from), - ) - } - - fn _parse_partial( - &self, - _heap: &std::sync::Arc, - _call_id: CallId, - cache: io::owned::sap::ParseCache, - json: String, - _type_arg_0: ::sys_types::SapTy, - ctx: &SysOpContext, - ) -> SysOpOutput { - let Ok(sap) = cache._data.clone().downcast::() else { - return SysOpOutput::err(VmInternalError::RustTypeError { - expected: std::any::TypeId::of::(), - got: cache._data.type_id(), - }); - }; - let result = match crate::sap::execute_sap_parse_partial(&json, &sap, ctx) { - Ok(Some(value)) => Ok(value), - Ok(None) => Ok(BexExternalValue::instance( - "baml.sap._NoYield", - ::indexmap::IndexMap::new(), - )), - Err(e) => Err(VmRustFnError::from(e)), - }; - SysOpOutput::Ready(result) - } -} - -impl io::IoNamespaceSap for T { - fn _new_parse_cache( - &self, - _heap: &std::sync::Arc, - _call_id: CallId, - target: ::sys_types::SapTy, - ctx: &SysOpContext, - ) -> SysOpOutput { - let compiled = match ::bex_sap::CompiledSapModel::from_sys_op_context(ctx, target) { - Ok(compiled) => compiled, - Err(e) => { - // `_new_parse_cache` declares `throws never`, and the type - // arguments that reach it come from the caller's own - // `parse` — a `T` schema-aligned parsing cannot model is - // a program bug, not a recoverable condition, so it panics. - return SysOpOutput::err(VmPanic::UserPanic { - message: format!("schema-aligned parsing cannot model this type: {e}"), - }); - } - }; - let sap = crate::sap::SapParseCache::new(compiled); - let data: std::sync::Arc = std::sync::Arc::new(sap); - SysOpOutput::ok(io::owned::sap::ParseCache { _data: data }) - } -} - /// Wrap an `OutputFormatContent` into the generated `owned::ai::OutputFormat` handle. fn wrap_output_format( content: std::sync::Arc, diff --git a/baml_language/crates/sys_ops/src/sap.rs b/baml_language/crates/sys_ops/src/sap.rs deleted file mode 100644 index 03255b042c0..00000000000 --- a/baml_language/crates/sys_ops/src/sap.rs +++ /dev/null @@ -1,173 +0,0 @@ -use ::bex_sap::sap_model; -use ::sys_types::DefKey; - -/// Cached schema information for incremental or one-shot SAP parsing. -pub struct SapParseCache { - types: bex_sap::CompiledSapModel, -} - -impl SapParseCache { - pub fn new(types: bex_sap::CompiledSapModel) -> Self { - Self { types } - } - - pub fn db(&self) -> &sap_model::TypeRefDb<'_, DefKey> { - self.types.db() - } - - pub fn ty(&self) -> &sap_model::Ty<'_, DefKey> { - self.types.ty() - } - - pub fn ty_resolved(&self) -> Result, &DefKey> { - self.db().resolve(self.ty()) - } -} - -/// Errors that can occur during LLM operations. Relocated verbatim from -/// `sys_llm`; only `ParseResponseError`, `JsonishError` and `SapError` are -/// still constructible now that the SAP parse entry points are the sole users. -#[derive(Debug, thiserror::Error)] -pub enum LlmOpError { - #[error("Expected {expected}, got {actual}")] - TypeError { - expected: &'static str, - actual: String, - }, - - #[error("Parse response error: {0}")] - ParseResponseError(String), - - #[error("Jsonish error: {0}")] - JsonishError(::bex_sap::jsonish::JsonishError), - - #[error("SAP error: {0}")] - SapError(::bex_sap::deserializer::coercer::ParsingError), -} - -impl From for ::sys_types::VmRustFnError { - fn from(e: LlmOpError) -> Self { - let baml: ::sys_types::VmBamlError = match e { - LlmOpError::TypeError { expected, actual } => { - ::sys_types::VmBamlError::InvalidArgument { - message: format!("expected {expected}, got {actual}"), - } - } - LlmOpError::ParseResponseError(e) => ::sys_types::VmBamlError::LlmClient { message: e }, - LlmOpError::JsonishError(e) => ::sys_types::VmBamlError::LlmClient { - message: e.to_string(), - }, - LlmOpError::SapError(e) => ::sys_types::VmBamlError::LlmClient { - message: e.to_string(), - }, - }; - baml.into() - } -} - -pub fn execute_sap_parse_final( - json: &str, - sap: &SapParseCache, - _ctx: &::sys_types::SysOpContext, -) -> Result { - // === Jsonish === - let jsonish_options = ::bex_sap::jsonish::ParseOptions::default(); - let jsonish = - ::bex_sap::jsonish::parse(json, jsonish_options, true).map_err(LlmOpError::JsonishError)?; - - let parse_ctx = ::bex_sap::deserializer::coercer::ParsingContext::new(sap.db()); - let target = sap - .ty_resolved() - .map_err(|err| parse_ctx.error_type_resolution(err)) - .map_err(LlmOpError::SapError)?; - let parsed = ::bex_sap::sap_model::TyResolvedRef::coerce(&parse_ctx, target, &jsonish) - .map_err(LlmOpError::SapError)? - .ok_or_else(|| { - LlmOpError::ParseResponseError("SAP parse returned no value when complete".to_string()) - })?; - - // === Convert back to baml === - Ok(::bex_sap::to_external::baml_value_to_external( - &parsed, - sap.db(), - )) -} - -/// Parses the text received so far as a partial value, or `None` for no yield. -/// -/// Text that doesn't parse as the target type yet is not an error here: a reply -/// that opens with a ```` ```json ```` fence or with prose before the JSON can't -/// be coerced until the JSON starts. As in BAML v0, the stream skips that -/// partial and waits for more text; [`execute_sap_parse_final`] still reports a -/// reply that never parses. -pub fn execute_sap_parse_partial( - json: &str, - sap: &SapParseCache, - _ctx: &::sys_types::SysOpContext, -) -> Result, LlmOpError> { - // An unresolvable target type is a schema bug, not an incomplete reply. - let parse_ctx = ::bex_sap::deserializer::coercer::ParsingContext::new(sap.db()); - let target = sap - .ty_resolved() - .map_err(|err| parse_ctx.error_type_resolution(err)) - .map_err(LlmOpError::SapError)?; - - // === Jsonish === - let jsonish_options = ::bex_sap::jsonish::ParseOptions::default(); - let Ok(jsonish) = ::bex_sap::jsonish::parse(json, jsonish_options, false) else { - return Ok(None); - }; - - // === SAP parsing (a partial is the same type, parsed from the text so far) === - let Ok(parsed) = ::bex_sap::sap_model::TyResolvedRef::coerce(&parse_ctx, target, &jsonish) - else { - return Ok(None); - }; - // === Convert back to baml === - match parsed { - Some(parsed) => { - let converted = ::bex_sap::to_external::baml_value_to_external(&parsed, sap.db()); - Ok(Some(converted)) - } - None => Ok(None), - } -} - -#[cfg(test)] -mod tests { - use super::LlmOpError; - - /// A parse failure must reach BAML as `LlmClient`, which is what - /// `ai.errors.normalize` classifies into `ai.errors.ParseFailed`. Mapping - /// it to a generic internal error instead would make a provider's bad - /// answer look like an engine defect and escape the client error taxonomy. - #[test] - fn parse_response_error_maps_to_llm_client() { - let mapped: ::sys_types::VmRustFnError = - LlmOpError::ParseResponseError("bad json".to_string()).into(); - assert_eq!( - mapped, - ::sys_types::VmRustFnError::BamlError(::sys_types::VmBamlError::LlmClient { - message: "bad json".to_string(), - }) - ); - } - - /// A shape mismatch is the CALLER's problem, so it maps to - /// `InvalidArgument` — a different `baml.errors.*` class, and catchable on - /// its own. - #[test] - fn type_error_maps_to_invalid_argument() { - let mapped: ::sys_types::VmRustFnError = LlmOpError::TypeError { - expected: "string", - actual: "int".to_string(), - } - .into(); - assert_eq!( - mapped, - ::sys_types::VmRustFnError::BamlError(::sys_types::VmBamlError::InvalidArgument { - message: "expected string, got int".to_string(), - }) - ); - } -} From 3691a2f2d248a1ba15b9a49d2453e434529206e9 Mon Sep 17 00:00:00 2001 From: Aaron Villalpando Date: Sat, 3 Oct 2026 00:12:08 -0700 Subject: [PATCH 8/8] Count only partial yields in the streamed parse test The loop ends with the whole reply, so a parser that yields only then would still have passed. Suggested by CodeRabbit. Co-Authored-By: Claude Opus 5.5 --- baml_language/crates/baml_tests/tests/llm_call_memory.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/baml_language/crates/baml_tests/tests/llm_call_memory.rs b/baml_language/crates/baml_tests/tests/llm_call_memory.rs index 5401c58b043..54dd491693e 100644 --- a/baml_language/crates/baml_tests/tests/llm_call_memory.rs +++ b/baml_language/crates/baml_tests/tests/llm_call_memory.rs @@ -241,7 +241,10 @@ function stream_parse() -> int {{ let upto = if (end > text.length()) {{ text.length() }} else {{ end }}; let parsed: Item[] | baml.sap._NoYield = cache._parse_partial(text.slice(0, upto)); match (parsed) {{ - let items: Item[] => {{ yielded = yielded + 1; }}, + let items: Item[] => {{ + // Only a prefix shorter than the reply shows a partial value. + if (upto < text.length()) {{ yielded = yielded + 1; }} + }}, _ => {{}}, }}; end = end + {STREAM_STEP};