Correctly implement NotifyBlock instruction (#2678)

Remove the hack used to implement `Instruction::NotifyBlock` and
properly store metadata somewhere else and index into it.

Author: antoniosarosi

engine/Cargo.lock

@@ -50,7 +50,6 @@ default-members = [
 
 [workspace.dependencies]
 anyhow = "1.0"
-arraystring = "0.3.0"
 askama = { version = "0.14.0", features = ["code-in-doc"] }
 axum = "0.7.5"
 baml-cli = { path = "cli" }

engine/Cargo.toml

@@ -8,7 +8,6 @@ license-file.workspace = true
 
 [dependencies]
 anyhow = { workspace = true }
-arraystring.workspace = true
 baml-types = { path = "../baml-lib/baml-types" }
 baml-vm = { path = "../baml-vm" }
 internal-baml-ast = { path = "../baml-lib/ast" }

engine/baml-compiler/Cargo.toml

@@ -2,7 +2,6 @@
 
 use std::collections::{HashMap, HashSet};
 
-use arraystring::{typenum::U255, ArrayString};
 use baml_types::{ir_type::TypeIR, BamlMap, BamlMediaType, BamlValueWithMeta, TypeValue};
 use baml_vm::{
     BamlVmProgram, BinOp, Bytecode, Class, CmpOp, Enum, Function, FunctionKind, GlobalIndex,
@@ -152,6 +151,7 @@ fn compile_thir_to_bytecode(
             kind: FunctionKind::Llm,
             locals_in_scope: vec![func.parameters.iter().map(|p| p.name.clone()).collect()],
             span: func.span.clone(),
+            block_notifications: Vec::new(),
         });
 
         let object_index = objects.insert(bytecode_llm_function);
@@ -235,6 +235,7 @@ fn compile_thir_to_bytecode(
             kind: FunctionKind::Native(func),
             locals_in_scope: vec![], // TODO.
             span: Span::fake_builtin_baml(),
+            block_notifications: Vec::new(),
         });
 
         let object_index = objects.insert(native_function);
@@ -247,6 +248,7 @@ fn compile_thir_to_bytecode(
         kind: FunctionKind::Future,
         locals_in_scope: vec![],
         span: Span::fake_builtin_baml(),
+        block_notifications: Vec::new(),
     }))));
 
     let mut resolved_class_names = HashMap::new();
@@ -431,6 +433,9 @@ struct HirCompiler<'g> {
     /// `AllocInstance` instructions that have a placeholder, which must be resolved when location
     /// of the class object is resolved.
     class_alloc_patch_list: &'g mut Vec<AllocInstancePatch>,
+
+    /// Block notifications for the current function being compiled.
+    block_notifications: Vec<baml_vm::bytecode::BlockNotification>,
 }
 
 #[derive(Debug)]
@@ -470,6 +475,7 @@ impl<'g> HirCompiler<'g> {
             scopes: Vec::new(),
             current_source_line: 0,
             locals_in_scope: Vec::new(),
+            block_notifications: Vec::new(),
         }
     }
 
@@ -508,6 +514,7 @@ impl<'g> HirCompiler<'g> {
             })),
 
             span: func.span.clone(),
+            block_notifications: self.block_notifications.clone(),
         })
     }
 
@@ -1585,18 +1592,24 @@ impl<'g> HirCompiler<'g> {
         self.emit(Instruction::LoadConst(const_index));
     }
 
-    fn emit_annotated_block(&mut self, v: &str) {
-        self.emit_string_literal(v);
+    fn emit_annotated_block(&mut self, annotation: &str) {
+        // Create the notification metadata
+        let notification = baml_vm::bytecode::BlockNotification {
+            function_name: String::new(), // Will be populated at runtime from Function::name
+            block_name: annotation.to_string(),
+            level: self.scopes.len(), // Current scope depth (1-based)
+            block_type: baml_vm::bytecode::BlockNotificationType::Statement,
+            is_enter: true,
+        };
 
-        self.emit(Instruction::NotifyBlock(
-            baml_vm::bytecode::BlockNotification {
-                function_name: ArrayString::<U255>::from_str_truncate(v),
-                block_name: ArrayString::<U255>::from_str_truncate(v),
-                level: 1,
-                block_type: baml_vm::bytecode::BlockNotificationType::Statement,
-                is_enter: true,
-            },
-        ));
+        // Add to the function's notification list
+        let notification_index = self.block_notifications.len();
+        self.block_notifications.push(notification);
+
+        // Emit instruction with just the index
+        self.emit(Instruction::NotifyBlock(notification_index));
+
+        // TODO: Emit exit notification when leaving the block
     }
 
     /// Emits a single instruction and returns the index of the instruction.

engine/baml-compiler/src/codegen.rs

@@ -7,7 +7,6 @@ description = "BAML Virtual Machine"
 license-file.workspace = true
 
 [dependencies]
-arraystring.workspace = true
 colored.workspace = true
 tokio = { workspace = true }
 baml-types = { path = "../baml-lib/baml-types" }

engine/baml-vm/Cargo.toml

@@ -1,7 +1,5 @@
 //! Instruction set and bytecode representation.
 
-use arraystring::{typenum::U255, ArrayString};
-
 use crate::{types::Value, GlobalIndex, ObjectIndex};
 
 /// Individual bytecode instruction.
@@ -242,15 +240,17 @@ pub enum Instruction {
 
     /// Notifies about entering or exiting a block.
     ///
-    /// Format: `NOTIFY_BLOCK function_name block_name`
-    NotifyBlock(BlockNotification),
+    /// Format: `NOTIFY_BLOCK block_index` where `block_index` is the index
+    /// into the current function's block_notifications array.
+    NotifyBlock(usize),
 }
 
-#[derive(Clone, Copy, Debug, PartialEq)]
+/// Block notification metadata stored in the Function struct.
+/// The function_name field is populated at runtime from the Function containing this notification.
+#[derive(Clone, Debug, PartialEq)]
 pub struct BlockNotification {
-    // This is a hack cause i don't wanna deal with implementing Copy by allocating this on the heap + doing an object pointer.
-    pub function_name: ArrayString<U255>,
-    pub block_name: ArrayString<U255>,
+    pub function_name: String, // Populated at runtime from Function::name
+    pub block_name: String,
     pub level: usize,
     pub block_type: BlockNotificationType,
     pub is_enter: bool,
@@ -368,18 +368,8 @@ impl std::fmt::Display for Instruction {
             Instruction::Assert => f.write_str("ASSERT"),
             Instruction::AllocMap(n) => write!(f, "ALLOC_MAP {n}"),
             Instruction::Watch(i) => write!(f, "WATCH {i}"),
-            Instruction::NotifyBlock(notification) => {
-                write!(
-                    f,
-                    "{}_BLOCK {function_name}.{block_name}",
-                    if notification.is_enter {
-                        "ENTER"
-                    } else {
-                        "EXIT"
-                    },
-                    function_name = &notification.function_name,
-                    block_name = &notification.block_name,
-                )
+            Instruction::NotifyBlock(block_index) => {
+                write!(f, "NOTIFY_BLOCK {block_index}")
             }
             Instruction::Notify(i) => write!(f, "NOTIFY {i}"),
         }

engine/baml-vm/src/bytecode.rs

@@ -57,8 +57,12 @@ pub fn display_instruction(
     let instruction = &function.bytecode.instructions[instruction_ptr as usize];
 
     let metadata = match instruction {
-        Instruction::NotifyBlock(notification) => {
-            format!("({})", &notification.block_name)
+        Instruction::NotifyBlock(block_index) => {
+            if let Some(notification) = function.block_notifications.get(*block_index) {
+                format!("({})", &notification.block_name)
+            } else {
+                format!("(invalid block index: {})", block_index)
+            }
         }
         Instruction::LoadConst(index) => format!(
             "({})",

engine/baml-vm/src/debug.rs

@@ -53,6 +53,12 @@ pub struct Function {
 
     /// Span of the function as computed by the parser.
     pub span: internal_baml_diagnostics::Span,
+
+    /// Block notifications for this function.
+    ///
+    /// Stores metadata about annotated blocks (//# annotations) in this function.
+    /// Instructions reference these by index.
+    pub block_notifications: Vec<crate::bytecode::BlockNotification>,
 }
 
 impl std::fmt::Display for Function {

engine/baml-vm/src/types.rs

@@ -639,8 +639,22 @@ impl Vm {
             // }
 
             match function.bytecode.instructions[instruction_ptr as usize] {
-                Instruction::NotifyBlock(notification) => {
-                    return Ok(VmExecState::Notify(WatchNotification::Block(notification)));
+                Instruction::NotifyBlock(block_index) => {
+                    // Get the notification from the function's storage
+                    let notification = &function.block_notifications[block_index];
+
+                    // Create a copy with the function name populated
+                    let full_notification = crate::bytecode::BlockNotification {
+                        function_name: function.name.clone(),
+                        block_name: notification.block_name.clone(),
+                        level: notification.level,
+                        block_type: notification.block_type,
+                        is_enter: notification.is_enter,
+                    };
+
+                    return Ok(VmExecState::Notify(WatchNotification::Block(
+                        full_notification,
+                    )));
                 }
                 Instruction::LoadConst(index) => {
                     let value = &function.bytecode.constants[index];

engine/baml-vm/src/vm.rs

@@ -462,6 +462,7 @@ pub fn assert_vm_executes_bytecode_with_inspection(
             vec![names]
         },
         span: internal_baml_diagnostics::Span::fake(),
+        block_notifications: Vec::new(),
     };
 
     let objects = vec![VmObject::Function(function)];