use super::{extract_doc_comment, extract_signature, field_text, node_text};
use crate::types::*;
pub fn extract(tree: &tree_sitter::Tree, source: &[u8]) -> (Vec<Symbol>, Vec<Import>) {
let root = tree.root_node();
let mut symbols = Vec::new();
let mut imports = Vec::new();
let mut cursor = root.walk();
for child in root.children(&mut cursor) {
extract_node(&child, source, &mut symbols, &mut imports, None);
}
(symbols, imports)
}
fn extract_node(
node: &tree_sitter::Node,
source: &[u8],
symbols: &mut Vec<Symbol>,
imports: &mut Vec<Import>,
parent_name: Option<&str>,
) {
match node.kind() {
"function_signature" => {
if let Some(sym) = extract_function(node, source, parent_name) {
symbols.push(sym);
}
}
"method_signature" => {
if let Some(sym) = extract_function(node, source, parent_name) {
symbols.push(sym);
}
}
"class_definition" | "class_declaration" => {
if let Some(sym) = extract_class_def(node, source, SymbolKind::Class, parent_name) {
symbols.push(sym);
}
}
"enum_declaration" => {
if let Some(sym) = extract_class_def(node, source, SymbolKind::Enum, parent_name) {
symbols.push(sym);
}
}
"mixin_declaration" => {
if let Some(sym) = extract_class_def(node, source, SymbolKind::Trait, parent_name) {
symbols.push(sym);
}
}
"type_alias" => {
if let Some(sym) = extract_type_alias(node, source, parent_name) {
symbols.push(sym);
}
}
"import_or_export" => {
let text = node_text(node, source);
let path = text
.trim_start_matches("import")
.trim_start_matches("export")
.trim()
.trim_end_matches(';')
.trim();
let (clean_path, alias) = if let Some(as_idx) = path.find(" as ") {
let p = path[..as_idx].trim().trim_matches('\'').trim_matches('"');
let a = path[as_idx + 4..].trim().trim_end_matches(';').trim();
(p.to_string(), Some(a.to_string()))
} else {
(path.trim_matches('\'').trim_matches('"').to_string(), None)
};
imports.push(Import {
path: clean_path,
alias,
span: Span::from_node(node),
});
}
"class_body" | "declaration" | "class_member" => {
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
extract_node(&child, source, symbols, imports, parent_name);
}
}
_ => {}
}
}
fn first_identifier<'a>(node: &tree_sitter::Node, source: &'a [u8]) -> Option<&'a str> {
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
if child.kind() == "identifier" {
return Some(node_text(&child, source));
}
}
None
}
fn extract_function(
node: &tree_sitter::Node,
source: &[u8],
parent_name: Option<&str>,
) -> Option<Symbol> {
let func_node = if node.kind() == "method_signature" {
let mut cursor = node.walk();
let mut found = None;
for child in node.children(&mut cursor) {
if child.kind() == "function_signature" {
found = Some(child);
break;
}
}
found
} else {
None
};
let target = func_node.as_ref().unwrap_or(node);
let name = field_text(target, "name", source).or_else(|| first_identifier(target, source))?;
if name.is_empty() {
return None;
}
let kind = if parent_name.is_some() {
SymbolKind::Method
} else {
SymbolKind::Function
};
Some(Symbol {
name: name.to_string(),
kind,
span: Span::from_node(node),
signature: extract_signature(node, "body", source),
doc_comment: extract_doc_comment(node, source),
parent: parent_name.map(|s| s.to_string()),
children: Vec::new(),
})
}
fn extract_class_def(
node: &tree_sitter::Node,
source: &[u8],
kind: SymbolKind,
parent_name: Option<&str>,
) -> Option<Symbol> {
let name = field_text(node, "name", source).or_else(|| first_identifier(node, source))?;
let signature = {
let mut cursor = node.walk();
let mut sig = node_text(node, source).to_string();
for child in node.children(&mut cursor) {
if child.kind() == "class_body" || child.kind() == "enum_body" {
let s = &source[node.start_byte()..child.start_byte()];
sig = std::str::from_utf8(s).unwrap_or("").trim().to_string();
break;
}
}
sig
};
let mut children = Vec::new();
let mut child_imports = Vec::new();
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
if child.kind() == "class_body" || child.kind() == "enum_body" {
let mut inner = child.walk();
for body_child in child.children(&mut inner) {
extract_node(
&body_child,
source,
&mut children,
&mut child_imports,
Some(name),
);
}
}
}
Some(Symbol {
name: name.to_string(),
kind,
span: Span::from_node(node),
signature,
doc_comment: extract_doc_comment(node, source),
parent: parent_name.map(|s| s.to_string()),
children,
})
}
fn extract_type_alias(
node: &tree_sitter::Node,
source: &[u8],
parent_name: Option<&str>,
) -> Option<Symbol> {
let name = field_text(node, "name", source).or_else(|| {
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
if child.kind() == "type_identifier" || child.kind() == "identifier" {
return Some(node_text(&child, source));
}
}
None
})?;
Some(Symbol {
name: name.to_string(),
kind: SymbolKind::TypeAlias,
span: Span::from_node(node),
signature: node_text(node, source).trim().to_string(),
doc_comment: extract_doc_comment(node, source),
parent: parent_name.map(|s| s.to_string()),
children: Vec::new(),
})
}
#[cfg(test)]
mod tests {
use super::*;
fn parse_dart(source: &str) -> (Vec<Symbol>, Vec<Import>) {
let mut parser = tree_sitter::Parser::new();
parser
.set_language(&tree_sitter_dart::LANGUAGE.into())
.unwrap();
let tree = parser.parse(source, None).unwrap();
extract(&tree, source.as_bytes())
}
#[test]
fn test_dart_class_with_methods() {
let source = r#"
import 'dart:math';
import 'package:flutter/material.dart';
class Calculator {
int add(int x) {
return x + 1;
}
int multiply(int x, int y) {
return x * y;
}
}
"#;
let (symbols, imports) = parse_dart(source);
assert_eq!(imports.len(), 2);
assert!(imports[0].path.contains("dart:math"));
assert!(imports[1].path.contains("flutter/material.dart"));
let calc = symbols.iter().find(|s| s.name == "Calculator").unwrap();
assert_eq!(calc.kind, SymbolKind::Class);
let child_names: Vec<&str> = calc.children.iter().map(|s| s.name.as_str()).collect();
assert!(
child_names.contains(&"add"),
"expected 'add' in {:?}",
child_names
);
assert!(
child_names.contains(&"multiply"),
"expected 'multiply' in {:?}",
child_names
);
}
#[test]
fn test_dart_enum_and_mixin() {
let source = r#"
enum Status {
active,
inactive,
pending
}
mixin Loggable {
void log(String message) {
print(message);
}
}
"#;
let (symbols, _imports) = parse_dart(source);
let status = symbols.iter().find(|s| s.name == "Status");
assert!(
status.is_some(),
"expected Status in {:?}",
symbols
.iter()
.map(|s| (&s.name, &s.kind))
.collect::<Vec<_>>()
);
assert_eq!(status.unwrap().kind, SymbolKind::Enum);
let loggable = symbols.iter().find(|s| s.name == "Loggable");
assert!(
loggable.is_some(),
"expected Loggable in {:?}",
symbols
.iter()
.map(|s| (&s.name, &s.kind))
.collect::<Vec<_>>()
);
assert_eq!(loggable.unwrap().kind, SymbolKind::Trait);
let log_method = loggable.unwrap().children.iter().find(|s| s.name == "log");
assert!(
log_method.is_some(),
"expected 'log' method in Loggable children"
);
assert_eq!(log_method.unwrap().kind, SymbolKind::Method);
}
#[test]
fn test_dart_type_alias_and_import() {
let source = r#"
import 'dart:async' as async_lib;
typedef StringCallback = void Function(String);
void greet(String name) {
print('Hello, $name');
}
"#;
let (symbols, imports) = parse_dart(source);
assert_eq!(imports.len(), 1);
assert!(imports[0].path.contains("dart:async"));
let type_alias = symbols.iter().find(|s| s.kind == SymbolKind::TypeAlias);
assert!(
type_alias.is_some(),
"expected TypeAlias in {:?}",
symbols
.iter()
.map(|s| (&s.name, &s.kind))
.collect::<Vec<_>>()
);
assert_eq!(type_alias.unwrap().name, "StringCallback");
}
}