Skip to main content

llmcc_rust/
token.rs

1use llmcc_core::LanguageTraitImpl;
2use llmcc_core::graph_builder::BlockKind;
3use llmcc_core::ir::{HirKind, HirNode};
4use llmcc_core::lang_def::{LanguageTrait, ParseNode, ParseTree, TreeSitterParseTree};
5use llmcc_core::scope::{Scope, ScopeStack};
6use llmcc_core::symbol::{SymKind, Symbol};
7use llmcc_core::{CompileCtxt, CompileUnit};
8use llmcc_resolver::ResolverOption;
9
10#[allow(clippy::single_component_path_imports)]
11use tree_sitter_rust;
12
13// Include the auto-generated language definition from build script
14// The generated file contains a define_lang! call that expands to LangRust
15include!(concat!(env!("OUT_DIR"), "/rust_tokens.rs"));
16
17impl LanguageTraitImpl for LangRust {
18    /// Block kind with parent context - handles tuple struct fields
19    fn block_kind_with_parent_impl(kind_id: u16, field_id: u16, parent_kind_id: u16) -> BlockKind {
20        // Tuple struct fields: types inside ordered_field_declaration_list with field "type"
21        if parent_kind_id == LangRust::ordered_field_declaration_list
22            && field_id == LangRust::field_type
23        {
24            return BlockKind::Field;
25        }
26        // Don't create return blocks inside function_type (type annotations, not function definitions)
27        // e.g., `type F = impl FnOnce() -> T;` should not create a return block for T
28        if parent_kind_id == LangRust::function_type && field_id == LangRust::field_return_type {
29            return BlockKind::Undefined;
30        }
31        // Default behavior: check field kind first, then node kind
32        let field_kind = <Self as LanguageTrait>::block_kind(field_id);
33        if field_kind != BlockKind::Undefined {
34            field_kind
35        } else {
36            <Self as LanguageTrait>::block_kind(kind_id)
37        }
38    }
39
40    #[rustfmt::skip]
41    fn collect_init_impl<'tcx>(cc: &'tcx CompileCtxt<'tcx>) -> ScopeStack<'tcx> {
42        let stack = ScopeStack::new(cc.arena(), &cc.interner);
43        let globals = cc.create_globals();
44        stack.push(globals);
45        debug_assert!(stack.depth() == 1);
46
47        for prim in crate::RUST_PRIMITIVES {
48            let name = cc.interner.intern(prim);
49            let symbol_val = Symbol::new(CompileCtxt::GLOBAL_SCOPE_OWNER, name);
50            let sym_id = symbol_val.id().0;
51            let symbol = cc.arena().alloc_with_id(sym_id, symbol_val);
52            symbol.set_kind(SymKind::Primitive);
53            symbol.set_is_global(true);
54            globals.insert(symbol);
55        }
56
57        stack
58    }
59
60    fn parse_impl(text: impl AsRef<[u8]>) -> Option<Box<dyn ParseTree>> {
61        use std::cell::RefCell;
62
63        // Thread-local parser reuse to avoid contention from Parser::new()
64        thread_local! {
65            static PARSER: RefCell<tree_sitter::Parser> = {
66                let mut parser = tree_sitter::Parser::new();
67                parser.set_language(&tree_sitter_rust::LANGUAGE.into()).unwrap();
68                RefCell::new(parser)
69            };
70        }
71
72        PARSER.with(|parser| {
73            let mut parser = parser.borrow_mut();
74            let bytes = text.as_ref();
75            let tree = parser.parse(bytes, None)?;
76            Some(Box::new(TreeSitterParseTree { tree }) as Box<dyn ParseTree>)
77        })
78    }
79
80    fn supported_extensions_impl() -> &'static [&'static str] {
81        &["rs"]
82    }
83
84    fn manifest_name_impl() -> &'static str {
85        "Cargo.toml"
86    }
87
88    fn container_dirs_impl() -> &'static [&'static str] {
89        &["src"]
90    }
91
92    /// Check if the given parse node is a Rust test attribute.
93    /// Detects: #[test], #[cfg(test)], #[tokio::test], #[async_std::test], etc.
94    fn is_test_attribute_impl(node: &dyn ParseNode, source: &[u8]) -> bool {
95        // First check if this is an attribute_item or inner_attribute_item
96        let kind_id = node.kind_id();
97        if kind_id != LangRust::attribute_item && kind_id != LangRust::inner_attribute_item {
98            return false;
99        }
100
101        // Extract the text of the attribute
102        let start = node.start_byte();
103        let end = node.end_byte();
104        if end <= start || end > source.len() {
105            return false;
106        }
107
108        let attr_text = match std::str::from_utf8(&source[start..end]) {
109            Ok(text) => text,
110            Err(_) => return false,
111        };
112
113        // Check for test-related attributes
114        attr_text.contains("#[test]")
115            || attr_text.contains("#[cfg(test)]")
116            || attr_text.contains("::test]") // catches #[tokio::test], #[async_std::test], etc.
117    }
118
119    fn collect_symbols_impl<'tcx, C>(
120        unit: CompileUnit<'tcx>,
121        node: HirNode<'tcx>,
122        scope_stack: ScopeStack<'tcx>,
123        config: &C,
124    ) -> &'tcx Scope<'tcx> {
125        unsafe {
126            let config_ref = config as *const C as *const ResolverOption;
127            crate::collect::collect_symbols(unit, &node, scope_stack, &*config_ref)
128        }
129    }
130
131    fn bind_symbols_impl<'tcx, C>(
132        unit: CompileUnit<'tcx>,
133        node: HirNode<'tcx>,
134        globals: &'tcx Scope<'tcx>,
135        config: &C,
136    ) {
137        unsafe {
138            let config = config as *const C as *const ResolverOption;
139            crate::bind::bind_symbols(unit, &node, globals, &*config);
140        }
141    }
142}