Baml salsa HIR

baml_language/Cargo.toml

@@ -36,9 +36,11 @@ divan = { package = "codspeed-divan-compat", version = "*" }
 glob = { version = "0.3" }
 insta = { version = "1.34", features = ["yaml", "redactions"] }
 insta-cmd = { version = "0.6.0" }
+la-arena = { version = "0.3" }
 logos = { version = "0.15" }
 pretty_assertions = { version = "1.4" }
 rowan = { version = "0.16" }
+rustc-hash = { version = "2.1" }
 # When updating salsa, make sure to also update the revision in `fuzz/Cargo.toml`
 salsa = { git = "https://github.com/salsa-rs/salsa.git", rev = "cdd0b85516a52c18b8a6d17a2279a96ed6c3e198", default-features = false, features = [
     "compact_str",

baml_language/crates/baml_db/Cargo.toml

@@ -18,6 +18,7 @@ workspace = true
 
 [dependencies]
 salsa = { workspace = true }
+rowan = { workspace = true }
 
 # All compiler crates
 baml_base = { workspace = true }

baml_language/crates/baml_db/src/lib.rs

@@ -18,6 +18,7 @@ pub use baml_parser;
 pub use baml_syntax;
 pub use baml_thir;
 pub use baml_workspace;
+use rowan::ast::AstNode;
 use salsa::Storage;
 
 /// Type alias for Salsa event callbacks
@@ -35,6 +36,9 @@ pub struct RootDatabase {
 #[salsa::db]
 impl salsa::Database for RootDatabase {}
 
+#[salsa::db]
+impl baml_hir::Db for RootDatabase {}
+
 impl RootDatabase {
     /// Create a new empty database.
     pub fn new() -> Self {
@@ -80,3 +84,73 @@ impl Default for RootDatabase {
         Self::new()
     }
 }
+
+//
+// ────────────────────────────────────────────────── FUNCTION QUERIES ─────
+//
+
+/// Returns the signature of a function (params, return type, generics).
+///
+/// This is separate from the `ItemTree` to provide fine-grained incrementality.
+/// Changing a function body does NOT invalidate this query.
+#[salsa::tracked]
+pub fn function_signature<'db>(
+    db: &'db dyn baml_hir::Db,
+    file: SourceFile,
+    function: baml_hir::FunctionLoc<'db>,
+) -> Arc<baml_hir::FunctionSignature> {
+    let tree = baml_parser::syntax_tree(db, file);
+    let source_file = baml_syntax::ast::SourceFile::cast(tree).unwrap();
+
+    // Find the function node by name
+    let item_tree = baml_hir::file_item_tree(db, file);
+    let func = &item_tree[function.id(db)];
+
+    for item in source_file.items() {
+        if let baml_syntax::ast::Item::Function(func_node) = item {
+            if let Some(name_token) = func_node.name() {
+                if name_token.text() == func.name.as_str() {
+                    return baml_hir::FunctionSignature::lower(&func_node);
+                }
+            }
+        }
+    }
+
+    // Function not found - return minimal signature
+    Arc::new(baml_hir::FunctionSignature {
+        name: func.name.clone(),
+        params: vec![],
+        return_type: baml_hir::TypeRef::Unknown,
+        attrs: baml_hir::FunctionAttributes::default(),
+    })
+}
+
+/// Returns the body of a function (LLM prompt or expression IR).
+///
+/// This is the most frequently invalidated query - it changes whenever
+/// the function body is edited.
+#[salsa::tracked]
+pub fn function_body<'db>(
+    db: &'db dyn baml_hir::Db,
+    file: SourceFile,
+    function: baml_hir::FunctionLoc<'db>,
+) -> Arc<baml_hir::FunctionBody> {
+    let tree = baml_parser::syntax_tree(db, file);
+    let source_file = baml_syntax::ast::SourceFile::cast(tree).unwrap();
+
+    let item_tree = baml_hir::file_item_tree(db, file);
+    let func = &item_tree[function.id(db)];
+
+    for item in source_file.items() {
+        if let baml_syntax::ast::Item::Function(func_node) = item {
+            if let Some(name_token) = func_node.name() {
+                if name_token.text() == func.name.as_str() {
+                    return baml_hir::FunctionBody::lower(&func_node);
+                }
+            }
+        }
+    }
+
+    // No body found
+    Arc::new(baml_hir::FunctionBody::Missing)
+}

baml_language/crates/baml_hir/Cargo.toml

@@ -22,4 +22,7 @@ baml_parser = { workspace = true }
 baml_syntax = { workspace = true }
 baml_workspace = { workspace = true }
 
+la-arena = { workspace = true }
+rowan = { workspace = true }
+rustc-hash = { workspace = true }
 salsa = { workspace = true }