@@ -29,8 +29,8 @@ impl TreeSitterTestDetector {
2929 Ok ( Self { parser } )
3030 }
3131
32- /// Detect all test-related entry points in the file
33- pub fn detect_test_entry_points (
32+ /// Detect all entry points in the file (tests, main, benchmarks, etc.)
33+ pub fn detect_entry_points (
3434 & mut self ,
3535 source : & str ,
3636 cursor : Option < Position > ,
@@ -339,12 +339,11 @@ impl TreeSitterTestDetector {
339339 }
340340 "function_item" => {
341341 let fn_name = self . get_function_name ( & node, source) ?;
342+ let start_line = node. start_position ( ) . row as u32 + 1 ;
343+ let end_line = node. end_position ( ) . row as u32 + 1 ;
342344
343345 // Check if this is a test function
344346 if self . is_test_function ( & node, source) {
345- let start_line = node. start_position ( ) . row as u32 + 1 ;
346- let end_line = node. end_position ( ) . row as u32 + 1 ;
347-
348347 // Build full path including module hierarchy
349348 let full_path = if module_stack. is_empty ( ) {
350349 fn_name. clone ( )
@@ -368,6 +367,26 @@ impl TreeSitterTestDetector {
368367 line_range : ( start_line, end_line) ,
369368 full_path : Some ( full_path) ,
370369 } ) ;
370+ } else if fn_name == "main" {
371+ // Detect main function with proper line range
372+ entry_points. push ( EntryPoint {
373+ name : "main" . to_string ( ) ,
374+ entry_type : EntryPointType :: Main ,
375+ line : start_line,
376+ column : node. start_position ( ) . column as u32 ,
377+ line_range : ( start_line, end_line) ,
378+ full_path : None ,
379+ } ) ;
380+ } else if self . is_bench_function ( & node, source) {
381+ // Detect benchmark function with proper line range
382+ entry_points. push ( EntryPoint {
383+ name : fn_name,
384+ entry_type : EntryPointType :: Benchmark ,
385+ line : start_line,
386+ column : node. start_position ( ) . column as u32 ,
387+ line_range : ( start_line, end_line) ,
388+ full_path : None ,
389+ } ) ;
371390 }
372391 }
373392 _ => { }
@@ -419,7 +438,24 @@ impl TreeSitterTestDetector {
419438 let attr_text = prev_sibling. utf8_text ( source. as_bytes ( ) ) ;
420439 if let Ok ( text) = attr_text {
421440 // Match various test attributes
422- if text. contains ( "test" ) {
441+ if text. contains ( "test" ) && !text. contains ( "bench" ) {
442+ return true ;
443+ }
444+ }
445+ }
446+ }
447+
448+ false
449+ }
450+
451+ /// Check if a function node is a benchmark function (#[bench])
452+ fn is_bench_function ( & self , node : & Node , source : & str ) -> bool {
453+ // Check previous sibling for bench attributes
454+ if let Some ( prev_sibling) = node. prev_sibling ( ) {
455+ if prev_sibling. kind ( ) == "attribute_item" {
456+ let attr_text = prev_sibling. utf8_text ( source. as_bytes ( ) ) ;
457+ if let Ok ( text) = attr_text {
458+ if text. contains ( "bench" ) {
423459 return true ;
424460 }
425461 }
0 commit comments