baml_language(feat): implement generic parameter declarations

baml_language/crates/baml_compiler_hir/src/generics.rs

@@ -4,8 +4,17 @@
 //! from the `ItemTree` to maintain the invalidation barrier. Changes to generic
 //! parameters don't invalidate the `ItemTree`.
 
-use baml_base::Name;
+use std::sync::Arc;
+
+use baml_base::{Name, SourceFile};
+use baml_compiler_parser::syntax_tree;
+use baml_compiler_syntax::ast;
 use la_arena::{Arena, Idx};
+use rowan::ast::AstNode;
+
+use crate::fqn::QualifiedName;
+use crate::ids::{ItemKind, LocalIdAllocator};
+use crate::{ClassId, Db, EnumId, FunctionId, TypeAliasId};
 
 /// Type parameter in a generic definition.
 ///
@@ -53,3 +62,233 @@ impl std::ops::Index<LocalTypeParamId> for GenericParams {
         &self.type_params[index]
     }
 }
+
+fn empty_generic_params() -> Arc<GenericParams> {
+    Arc::new(GenericParams::new())
+}
+
+fn generic_params_from_list(list: Option<ast::GenericParamList>) -> Arc<GenericParams> {
+    let Some(list) = list else {
+        return empty_generic_params();
+    };
+    let mut params = GenericParams::new();
+    for param_token in list.params() {
+        params.type_params.alloc(TypeParam {
+            name: Name::new(param_token.text()),
+        });
+    }
+    Arc::new(params)
+}
+
+fn with_source_file<T>(
+    db: &dyn Db,
+    file: SourceFile,
+    f: impl FnOnce(ast::SourceFile) -> T,
+) -> Option<T> {
+    let tree = syntax_tree(db, file);
+    ast::SourceFile::cast(tree).map(f)
+}
+
+fn find_top_level_generic_params(
+    source_file: &ast::SourceFile,
+    allocator: &mut LocalIdAllocator,
+    target_id: u32,
+    item_kind: ItemKind,
+    mut extract: impl FnMut(ast::Item) -> Option<(Name, Option<ast::GenericParamList>)>,
+) -> Option<Arc<GenericParams>> {
+    for item in source_file.items() {
+        if let Some((name, list)) = extract(item) {
+            let id = allocator.alloc_id::<()>(item_kind, &name);
+            if id.as_u32() == target_id {
+                return Some(generic_params_from_list(list));
+            }
+        }
+    }
+    None
+}
+
+fn scan_function_item_for_generics(
+    allocator: &mut LocalIdAllocator,
+    func_node: &ast::FunctionDef,
+    target_id: u32,
+) -> Option<Arc<GenericParams>> {
+    let name_token = func_node.name()?;
+    let base_name = Name::new(name_token.text());
+    let base_id = allocator.alloc_id::<()>(ItemKind::Function, &base_name);
+    if base_id.as_u32() == target_id {
+        return Some(generic_params_from_list(func_node.generic_param_list()));
+    }
+
+    func_node.llm_body()?;
+
+    let render_name = Name::new(format!("{base_name}.render_prompt"));
+    let render_id = allocator.alloc_id::<()>(ItemKind::Function, &render_name);
+    if render_id.as_u32() == target_id {
+        return Some(generic_params_from_list(func_node.generic_param_list()));
+    }
+
+    let build_name = Name::new(format!("{base_name}.build_request"));
+    let build_id = allocator.alloc_id::<()>(ItemKind::Function, &build_name);
+    if build_id.as_u32() == target_id {
+        return Some(generic_params_from_list(func_node.generic_param_list()));
+    }
+
+    None
+}
+
+fn scan_class_methods_for_generics(
+    allocator: &mut LocalIdAllocator,
+    class_node: &ast::ClassDef,
+    target_id: u32,
+) -> Option<Arc<GenericParams>> {
+    let class_name = class_node
+        .name()
+        .map(|token| token.text().to_string())
+        .unwrap_or_else(|| "UnnamedClass".to_string());
+    for method in class_node.methods() {
+        let Some(method_name) = method.name() else {
+            continue;
+        };
+        let qualified_method_name =
+            QualifiedName::local_method_from_str(&class_name, method_name.text());
+        let id = allocator.alloc_id::<()>(ItemKind::Function, &qualified_method_name);
+        if id.as_u32() == target_id {
+            return Some(generic_params_from_list(method.generic_param_list()));
+        }
+    }
+    None
+}
+
+pub(crate) fn function_generic_params_from_cst(
+    db: &dyn Db,
+    func: FunctionId<'_>,
+) -> Arc<GenericParams> {
+    let file = func.file(db);
+    let target_id = func.id(db).as_u32();
+
+    let result = with_source_file(db, file, |source_file| {
+        let mut allocator = LocalIdAllocator::new();
+        for item in source_file.items() {
+            match item {
+                ast::Item::Function(func_node) => {
+                    if let Some(params) =
+                        scan_function_item_for_generics(&mut allocator, &func_node, target_id)
+                    {
+                        return params;
+                    }
+                }
+                ast::Item::Class(class_node) => {
+                    if let Some(params) =
+                        scan_class_methods_for_generics(&mut allocator, &class_node, target_id)
+                    {
+                        return params;
+                    }
+                }
+                ast::Item::Client(client_node) => {
+                    if let Some(name_token) = client_node.name() {
+                        let resolve_name = Name::new(format!("{}.resolve", name_token.text()));
+                        let resolve_id =
+                            allocator.alloc_id::<()>(ItemKind::Function, &resolve_name);
+                        if resolve_id.as_u32() == target_id {
+                            return empty_generic_params();
+                        }
+                    }
+                }
+                _ => {}
+            }
+        }
+        empty_generic_params()
+    });
+
+    result.unwrap_or_else(empty_generic_params)
+}
+
+pub(crate) fn class_generic_params_from_cst(db: &dyn Db, class: ClassId<'_>) -> Arc<GenericParams> {
+    let file = class.file(db);
+    let target_id = class.id(db).as_u32();
+
+    let result = with_source_file(db, file, |source_file| {
+        let mut allocator = LocalIdAllocator::new();
+        find_top_level_generic_params(
+            &source_file,
+            &mut allocator,
+            target_id,
+            ItemKind::Class,
+            |item| {
+                if let ast::Item::Class(class_node) = item {
+                    class_node.name().map(|name_token| {
+                        (
+                            Name::new(name_token.text()),
+                            class_node.generic_param_list(),
+                        )
+                    })
+                } else {
+                    None
+                }
+            },
+        )
+    });
+
+    result.flatten().unwrap_or_else(empty_generic_params)
+}
+
+pub(crate) fn enum_generic_params_from_cst(
+    db: &dyn Db,
+    enum_def: EnumId<'_>,
+) -> Arc<GenericParams> {
+    let file = enum_def.file(db);
+    let target_id = enum_def.id(db).as_u32();
+
+    let result = with_source_file(db, file, |source_file| {
+        let mut allocator = LocalIdAllocator::new();
+        find_top_level_generic_params(
+            &source_file,
+            &mut allocator,
+            target_id,
+            ItemKind::Enum,
+            |item| {
+                if let ast::Item::Enum(enum_node) = item {
+                    enum_node.name().map(|name_token| {
+                        (Name::new(name_token.text()), enum_node.generic_param_list())
+                    })
+                } else {
+                    None
+                }
+            },
+        )
+    });
+
+    result.flatten().unwrap_or_else(empty_generic_params)
+}
+
+pub(crate) fn type_alias_generic_params_from_cst(
+    db: &dyn Db,
+    alias: TypeAliasId<'_>,
+) -> Arc<GenericParams> {
+    let file = alias.file(db);
+    let target_id = alias.id(db).as_u32();
+
+    let result = with_source_file(db, file, |source_file| {
+        let mut allocator = LocalIdAllocator::new();
+        find_top_level_generic_params(
+            &source_file,
+            &mut allocator,
+            target_id,
+            ItemKind::TypeAlias,
+            |item| {
+                if let ast::Item::TypeAlias(alias_node) = item {
+                    alias_node.name().map(|name_token| {
+                        (
+                            Name::new(name_token.text()),
+                            alias_node.generic_param_list(),
+                        )
+                    })
+                } else {
+                    None
+                }
+            },
+        )
+    });
+
+    result.flatten().unwrap_or_else(empty_generic_params)
+}

baml_language/crates/baml_compiler_hir/src/ids.rs

@@ -8,6 +8,9 @@
 
 use std::marker::PhantomData;
 
+use baml_base::Name;
+use rustc_hash::FxHashMap;
+
 /// Identifier for a class definition.
 pub use crate::loc::ClassLoc as ClassId;
 /// Identifier for a client configuration.
@@ -188,3 +191,43 @@ pub enum ItemKind {
     TemplateString,
     RetryPolicy,
 }
+
+pub(crate) fn allocate_local_id<T>(
+    next_index: &mut FxHashMap<(ItemKind, u16), u16>,
+    kind: ItemKind,
+    name: &Name,
+) -> LocalItemId<T> {
+    let hash = hash_name(name);
+    let index = next_index.entry((kind, hash)).or_insert(0);
+    let id = LocalItemId::new(hash, *index);
+    *index += 1;
+    id
+}
+
+/// Allocator for `LocalItemId`s with collision handling.
+///
+/// Replays the same hashing and collision-indexing logic as `ItemTree`.
+/// This allows queries to reproduce the same local IDs when scanning CST.
+pub struct LocalIdAllocator {
+    next_index: FxHashMap<(ItemKind, u16), u16>,
+}
+
+impl LocalIdAllocator {
+    /// Create a new allocator with empty collision state.
+    pub fn new() -> Self {
+        Self {
+            next_index: FxHashMap::default(),
+        }
+    }
+
+    /// Allocate a `LocalItemId` for a named item, updating collision state.
+    pub fn alloc_id<T>(&mut self, kind: ItemKind, name: &Name) -> LocalItemId<T> {
+        allocate_local_id(&mut self.next_index, kind, name)
+    }
+}
+
+impl Default for LocalIdAllocator {
+    fn default() -> Self {
+        Self::new()
+    }
+}

baml_language/crates/baml_compiler_hir/src/item_tree.rs

@@ -12,7 +12,7 @@ use rowan::TextRange;
 use rustc_hash::FxHashMap;
 
 use crate::{
-    ids::{ItemKind, LocalItemId, hash_name},
+    ids::{ItemKind, LocalItemId, allocate_local_id},
     loc::{
         ClassMarker, ClientMarker, EnumMarker, FunctionMarker, GeneratorMarker, RetryPolicyMarker,
         TemplateStringMarker, TestMarker, TypeAliasMarker,
@@ -111,11 +111,7 @@ impl ItemTree {
     /// Allocate a collision-resistant ID for an item.
     /// Returns a `LocalItemId` with the name's hash and a unique collision index.
     fn alloc_id<T>(&mut self, kind: ItemKind, name: &Name) -> LocalItemId<T> {
-        let hash = hash_name(name);
-        let index = self.next_index.entry((kind, hash)).or_insert(0);
-        let id = LocalItemId::new(hash, *index);
-        *index += 1;
-        id
+        allocate_local_id(&mut self.next_index, kind, name)
     }
 
     /// Add a function and return its local ID.

baml_language/crates/baml_compiler_hir/src/lib.rs

@@ -20,8 +20,8 @@ use std::sync::Arc;
 use baml_base::{FileId, Name, SourceFile, Span, TyAttr};
 use baml_compiler_diagnostics::{HirDiagnostic, NameError};
 use baml_compiler_parser::syntax_tree;
-use baml_compiler_syntax::SyntaxNode;
-use rowan::{SyntaxToken, TextRange, ast::AstNode};
+use baml_compiler_syntax::{SyntaxNode, SyntaxToken};
+use rowan::{TextRange, ast::AstNode};
 
 // Module declarations
 mod body;
@@ -210,34 +210,27 @@ pub fn project_items(db: &dyn Db, root: baml_workspace::Project) -> ProjectItems
 ///
 /// This is queried separately from `ItemTree` for incrementality - changes to
 /// generic parameters don't invalidate the `ItemTree`.
-///
-/// For now, this returns empty generic parameters since BAML doesn't currently
-/// parse generic syntax. Future work will extract `<T>` from the CST.
 #[salsa::tracked]
-pub fn function_generic_params(_db: &dyn Db, _func: FunctionId<'_>) -> Arc<GenericParams> {
-    // TODO: Extract generic parameters from CST when BAML adds generic syntax
-    Arc::new(GenericParams::new())
+pub fn function_generic_params(db: &dyn Db, func: FunctionId<'_>) -> Arc<GenericParams> {
+    generics::function_generic_params_from_cst(db, func)
 }
 
 /// Tracked: Get generic parameters for a class.
 #[salsa::tracked]
-pub fn class_generic_params(_db: &dyn Db, _class: ClassId<'_>) -> Arc<GenericParams> {
-    // TODO: Extract generic parameters from CST when BAML adds generic syntax
-    Arc::new(GenericParams::new())
+pub fn class_generic_params(db: &dyn Db, class: ClassId<'_>) -> Arc<GenericParams> {
+    generics::class_generic_params_from_cst(db, class)
 }
 
 /// Tracked: Get generic parameters for an enum.
 #[salsa::tracked]
-pub fn enum_generic_params(_db: &dyn Db, _enum: EnumId<'_>) -> Arc<GenericParams> {
-    // TODO: Extract generic parameters from CST when BAML adds generic syntax
-    Arc::new(GenericParams::new())
+pub fn enum_generic_params(db: &dyn Db, enum_def: EnumId<'_>) -> Arc<GenericParams> {
+    generics::enum_generic_params_from_cst(db, enum_def)
 }
 
 /// Tracked: Get generic parameters for a type alias.
 #[salsa::tracked]
-pub fn type_alias_generic_params(_db: &dyn Db, _alias: TypeAliasId<'_>) -> Arc<GenericParams> {
-    // TODO: Extract generic parameters from CST when BAML adds generic syntax
-    Arc::new(GenericParams::new())
+pub fn type_alias_generic_params(db: &dyn Db, alias: TypeAliasId<'_>) -> Arc<GenericParams> {
+    generics::type_alias_generic_params_from_cst(db, alias)
 }
 
 //