codesynapse_core/ts_extract/
rust.rs

1use super::{
2    add_contains_edge, add_node_if_missing, make_file_node, run_query_matches_ranged,
3    run_query_named,
4};
5use crate::error::Result;
6use crate::extract::{make_id, ImportNode, LanguageExtractor};
7use crate::types::{Edge, ExtractionFragment, Node};
8use std::collections::HashMap;
9use std::path::Path;
10
11pub struct TsRustExtractor;
12
13impl TsRustExtractor {
14    pub fn extract(source: &[u8], path: &Path) -> Result<ExtractionFragment> {
15        let (file_id, _, file_node) = make_file_node(path);
16        let mut fragment = ExtractionFragment {
17            nodes: vec![file_node],
18            edges: vec![],
19        };
20
21        let lang = tree_sitter_rust::LANGUAGE.into();
22
23        // Struct definitions with generics
24        let struct_query = r#"
25            (struct_item
26                name: (type_identifier) @struct.name
27                (type_parameters
28                    (type_parameter
29                        (type_identifier) @struct.generic
30                    )
31                )?
32            )
33        "#;
34        if let Ok(mut captures) = run_query_named(source, &lang, struct_query) {
35            let names = captures.remove("struct.name").unwrap_or_default();
36            let _generics = captures.remove("struct.generic").unwrap_or_default();
37
38            for name in &names {
39                let struct_id = make_id(&[&file_id, name]);
40                fragment.nodes.push(Node {
41                    id: struct_id.clone(),
42                    label: name.to_string(),
43                    file_type: "struct".to_string(),
44                    source_file: path.to_string_lossy().to_string(),
45                    source_location: None,
46                    community: None,
47                    rationale: None,
48                    docstring: None,
49                    metadata: HashMap::new(),
50                });
51                add_contains_edge(&mut fragment, &file_id, struct_id.clone(), path);
52
53                // Find generic params for this struct
54                let content = std::str::from_utf8(source);
55                if let Ok(content) = content {
56                    if let Some(line) = content
57                        .lines()
58                        .find(|l| l.contains(&format!("struct {}", name)))
59                    {
60                        if let Some(gen_start) = line.find('<') {
61                            if let Some(gen_end) = line.find('>') {
62                                let gen_str = &line[gen_start + 1..gen_end].trim();
63                                for gen_name in gen_str.split(',').map(|s| s.trim()) {
64                                    if !gen_name.is_empty() {
65                                        let gen_id = make_id(&[&file_id, gen_name, "ty"]);
66                                        if !fragment.nodes.iter().any(|n| n.id == gen_id) {
67                                            fragment.nodes.push(Node {
68                                                id: gen_id.clone(),
69                                                label: gen_name.to_string(),
70                                                file_type: "code".to_string(),
71                                                source_file: path.to_string_lossy().to_string(),
72                                                source_location: None,
73                                                community: None,
74                                                rationale: None,
75                                                docstring: None,
76                                                metadata: {
77                                                    let mut m = HashMap::new();
78                                                    m.insert(
79                                                        "kind".to_string(),
80                                                        "generic_param".to_string(),
81                                                    );
82                                                    m
83                                                },
84                                            });
85                                        }
86                                        fragment.edges.push(Edge {
87                                            source: struct_id.clone(),
88                                            target: gen_id,
89                                            relation: "generic".to_string(),
90                                            confidence: "EXTRACTED".to_string(),
91                                            source_file: Some(path.to_string_lossy().to_string()),
92                                            weight: 1.0,
93                                            context: None,
94                                        });
95                                    }
96                                }
97                            }
98                        }
99                    }
100                }
101            }
102        }
103
104        // Use declarations
105        let use_query = r#"
106            (use_declaration
107                argument: (scoped_identifier) @use.path
108            )
109        "#;
110        if let Ok(mut captures) = run_query_named(source, &lang, use_query) {
111            let paths = captures.remove("use.path").unwrap_or_default();
112            for p in &paths {
113                let parts: Vec<&str> = p.split("::").collect();
114                for (i, part) in parts.iter().enumerate() {
115                    let seg_id = if i == 0 {
116                        make_id(&[part])
117                    } else {
118                        let parent = make_id(parts[..i].as_ref());
119                        make_id(&[&parent, part])
120                    };
121                    let seg_node = Node {
122                        id: seg_id.clone(),
123                        label: part.to_string(),
124                        file_type: "code".to_string(),
125                        source_file: path.to_string_lossy().to_string(),
126                        source_location: None,
127                        community: None,
128                        rationale: None,
129                        docstring: None,
130                        metadata: HashMap::new(),
131                    };
132                    add_node_if_missing(&mut fragment, seg_node);
133                    fragment.edges.push(Edge {
134                        source: file_id.clone(),
135                        target: seg_id,
136                        relation: "imports".to_string(),
137                        confidence: "EXTRACTED".to_string(),
138                        source_file: Some(path.to_string_lossy().to_string()),
139                        weight: 1.0,
140                        context: None,
141                    });
142                }
143            }
144        }
145
146        // Function items (top-level and impl methods) with source_location for BM25 body indexing
147        let fn_query = r#"
148            (function_item
149                name: (identifier) @fn.name
150            ) @fn.node
151        "#;
152        if let Ok(matches) = run_query_matches_ranged(source, &lang, fn_query) {
153            for (text_map, range_map) in &matches {
154                let Some(name) = text_map.get("fn.name") else {
155                    continue;
156                };
157                if name.is_empty() {
158                    continue;
159                }
160                let label = format!("{}()", name);
161                let fn_id = make_id(&[&file_id, &label]);
162                let (start, end) = range_map.get("fn.node").copied().unwrap_or((0, 0));
163                let source_location = if start < end {
164                    Some(format!("{}:{}", start, end))
165                } else {
166                    None
167                };
168                let fn_node = Node {
169                    id: fn_id.clone(),
170                    label,
171                    file_type: "function".to_string(),
172                    source_file: path.to_string_lossy().to_string(),
173                    source_location,
174                    community: None,
175                    rationale: None,
176                    docstring: None,
177                    metadata: {
178                        let mut m = HashMap::new();
179                        m.insert("kind".to_string(), "function".to_string());
180                        m
181                    },
182                };
183                add_node_if_missing(&mut fragment, fn_node);
184                add_contains_edge(&mut fragment, &file_id, fn_id, path);
185            }
186        }
187
188        // Impl method → struct contains edges for simple impl types: impl Foo { fn m() {} }
189        let impl_simple_query = r#"
190            (impl_item
191                type: (type_identifier) @impl.type
192                body: (declaration_list
193                    (function_item
194                        name: (identifier) @method.name
195                    )
196                )
197            )
198        "#;
199        Self::add_impl_edges(
200            &mut fragment,
201            source,
202            &lang,
203            impl_simple_query,
204            &file_id,
205            path,
206        );
207
208        // Impl method → struct contains edges for generic impl types: impl<T> Foo<T> { fn m() {} }
209        let impl_generic_query = r#"
210            (impl_item
211                type: (generic_type
212                    type: (type_identifier) @impl.type
213                )
214                body: (declaration_list
215                    (function_item
216                        name: (identifier) @method.name
217                    )
218                )
219            )
220        "#;
221        Self::add_impl_edges(
222            &mut fragment,
223            source,
224            &lang,
225            impl_generic_query,
226            &file_id,
227            path,
228        );
229
230        Ok(fragment)
231    }
232
233    fn add_impl_edges(
234        fragment: &mut ExtractionFragment,
235        source: &[u8],
236        lang: &tree_sitter::Language,
237        query_str: &str,
238        file_id: &str,
239        path: &Path,
240    ) {
241        let Ok(matches) = run_query_matches_ranged(source, lang, query_str) else {
242            return;
243        };
244        for (text_map, _) in &matches {
245            let Some(impl_type) = text_map.get("impl.type") else {
246                continue;
247            };
248            let Some(method_name) = text_map.get("method.name") else {
249                continue;
250            };
251            let impl_type_id = make_id(&[file_id, impl_type]);
252            let method_label = format!("{}()", method_name);
253            let method_id = make_id(&[file_id, &method_label]);
254            if fragment.nodes.iter().any(|n| n.id == impl_type_id) {
255                fragment.edges.push(Edge {
256                    source: impl_type_id,
257                    target: method_id,
258                    relation: "contains".to_string(),
259                    confidence: "EXTRACTED".to_string(),
260                    source_file: Some(path.to_string_lossy().to_string()),
261                    weight: 1.0,
262                    context: None,
263                });
264            }
265        }
266    }
267}
268
269/// Tree-sitter based Rust extractor
270impl LanguageExtractor for TsRustExtractor {
271    fn file_extensions(&self) -> Vec<&'static str> {
272        vec!["rs"]
273    }
274    fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment> {
275        Self::extract(source, path)
276    }
277    fn resolve_imports(&self, _imports: &[ImportNode]) -> Vec<Edge> {
278        vec![]
279    }
280    fn collect_type_refs(&self, _fragment: &mut ExtractionFragment) {}
281}
codesynapse_core/ts_extract/rust.rs

codesynapse_core/ts_extract/
rust.rs
