Skip to main content

codesynapse_core/ts_extract/
objc.rs

1use super::make_file_node;
2use crate::error::{CodeSynapseError, Result};
3use crate::extract::{make_id, ImportNode, LanguageExtractor};
4use crate::types::{Edge, ExtractionFragment, Node};
5use std::collections::HashMap;
6use std::path::Path;
7use tree_sitter::{Node as TsNode, Parser};
8
9pub struct ObjCExtractor;
10
11fn oc_text(source: &[u8], node: &TsNode<'_>) -> String {
12    std::str::from_utf8(&source[node.start_byte()..node.end_byte()])
13        .unwrap_or("")
14        .trim()
15        .to_string()
16}
17
18fn oc_node(id: String, label: String, file_type: &str, path: &Path) -> Node {
19    Node {
20        id,
21        label,
22        file_type: file_type.to_string(),
23        source_file: path.to_string_lossy().to_string(),
24        source_location: None,
25        community: None,
26        rationale: None,
27        docstring: None,
28        metadata: HashMap::new(),
29    }
30}
31
32fn oc_edge(fragment: &mut ExtractionFragment, src: &str, tgt: &str, rel: &str, path: &Path) {
33    fragment.edges.push(Edge {
34        source: src.to_string(),
35        target: tgt.to_string(),
36        relation: rel.to_string(),
37        confidence: "EXTRACTED".to_string(),
38        source_file: Some(path.to_string_lossy().to_string()),
39        weight: 1.0,
40        context: None,
41    });
42}
43
44fn add_node_if_missing(fragment: &mut ExtractionFragment, node: Node) {
45    if !fragment.nodes.iter().any(|n| n.id == node.id) {
46        fragment.nodes.push(node);
47    }
48}
49
50fn first_identifier(source: &[u8], node: &TsNode<'_>) -> Option<String> {
51    for i in 0..node.child_count() {
52        if let Some(child) = node.child(i) {
53            if child.kind() == "identifier" {
54                let t = oc_text(source, &child);
55                if !t.is_empty() {
56                    return Some(t);
57                }
58            }
59        }
60    }
61    None
62}
63
64/// Build ObjC method selector: join all identifier children (selector parts).
65fn method_selector(source: &[u8], node: &TsNode<'_>) -> Option<String> {
66    let parts: Vec<String> = (0..node.child_count())
67        .filter_map(|i| node.child(i))
68        .filter(|c| c.kind() == "identifier")
69        .map(|c| oc_text(source, &c))
70        .filter(|s| !s.is_empty())
71        .collect();
72    if parts.is_empty() {
73        None
74    } else {
75        Some(parts.join(""))
76    }
77}
78
79fn walk_oc<'t>(
80    node: TsNode<'t>,
81    source: &[u8],
82    file_id: &str,
83    stem: &str,
84    path: &Path,
85    fragment: &mut ExtractionFragment,
86    parent_nid: Option<String>,
87) {
88    match node.kind() {
89        "preproc_include" => {
90            for i in 0..node.child_count() {
91                if let Some(child) = node.child(i) {
92                    match child.kind() {
93                        "system_lib_string" => {
94                            let raw = oc_text(source, &child);
95                            let module = raw.trim_matches(|c| c == '<' || c == '>');
96                            let module = module.rsplit('/').next().unwrap_or(module);
97                            let module = module.trim_end_matches(".h");
98                            if !module.is_empty() {
99                                let tgt = make_id(&[module]);
100                                fragment.edges.push(Edge {
101                                    source: file_id.to_string(),
102                                    target: tgt,
103                                    relation: "imports".to_string(),
104                                    confidence: "EXTRACTED".to_string(),
105                                    source_file: Some(path.to_string_lossy().to_string()),
106                                    weight: 1.0,
107                                    context: Some("import".to_string()),
108                                });
109                            }
110                        }
111                        "string_literal" => {
112                            // string_content child
113                            for j in 0..child.child_count() {
114                                if let Some(sub) = child.child(j) {
115                                    if sub.kind() == "string_content" {
116                                        let raw = oc_text(source, &sub);
117                                        let module = raw.rsplit('/').next().unwrap_or(&raw);
118                                        let module = module.trim_end_matches(".h");
119                                        if !module.is_empty() {
120                                            let tgt = make_id(&[module]);
121                                            fragment.edges.push(Edge {
122                                                source: file_id.to_string(),
123                                                target: tgt,
124                                                relation: "imports".to_string(),
125                                                confidence: "EXTRACTED".to_string(),
126                                                source_file: Some(
127                                                    path.to_string_lossy().to_string(),
128                                                ),
129                                                weight: 1.0,
130                                                context: Some("import".to_string()),
131                                            });
132                                        }
133                                    }
134                                }
135                            }
136                        }
137                        _ => {}
138                    }
139                }
140            }
141        }
142        "class_interface" => {
143            if let Some(name) = first_identifier(source, &node) {
144                let cls_nid = make_id(&[stem, &name]);
145                fragment
146                    .nodes
147                    .push(oc_node(cls_nid.clone(), name.clone(), "class", path));
148                oc_edge(fragment, file_id, &cls_nid, "contains", path);
149
150                // Second identifier after ':' is superclass
151                let mut colon_seen = false;
152                for i in 0..node.child_count() {
153                    if let Some(child) = node.child(i) {
154                        match child.kind() {
155                            ":" => colon_seen = true,
156                            "identifier" if colon_seen => {
157                                let super_name = oc_text(source, &child);
158                                if !super_name.is_empty() && super_name != name {
159                                    let super_nid = make_id(&[&super_name]);
160                                    oc_edge(fragment, &cls_nid, &super_nid, "inherits", path);
161                                }
162                                colon_seen = false;
163                            }
164                            "method_declaration" => {
165                                walk_oc(
166                                    child,
167                                    source,
168                                    file_id,
169                                    stem,
170                                    path,
171                                    fragment,
172                                    Some(cls_nid.clone()),
173                                );
174                            }
175                            _ => {}
176                        }
177                    }
178                }
179            }
180        }
181        "class_implementation" => {
182            if let Some(name) = first_identifier(source, &node) {
183                let impl_nid = make_id(&[stem, &name]);
184                if !fragment.nodes.iter().any(|n| n.id == impl_nid) {
185                    fragment
186                        .nodes
187                        .push(oc_node(impl_nid.clone(), name, "class", path));
188                    oc_edge(fragment, file_id, &impl_nid, "contains", path);
189                }
190                for i in 0..node.child_count() {
191                    if let Some(child) = node.child(i) {
192                        if child.kind() == "implementation_definition" {
193                            for j in 0..child.child_count() {
194                                if let Some(sub) = child.child(j) {
195                                    walk_oc(
196                                        sub,
197                                        source,
198                                        file_id,
199                                        stem,
200                                        path,
201                                        fragment,
202                                        Some(impl_nid.clone()),
203                                    );
204                                }
205                            }
206                        }
207                    }
208                }
209            }
210        }
211        "protocol_declaration" => {
212            if let Some(name) = first_identifier(source, &node) {
213                let proto_nid = make_id(&[stem, &name]);
214                fragment.nodes.push(oc_node(
215                    proto_nid.clone(),
216                    format!("<{}>", name),
217                    "code",
218                    path,
219                ));
220                oc_edge(fragment, file_id, &proto_nid, "contains", path);
221                for i in 0..node.child_count() {
222                    if let Some(child) = node.child(i) {
223                        walk_oc(
224                            child,
225                            source,
226                            file_id,
227                            stem,
228                            path,
229                            fragment,
230                            Some(proto_nid.clone()),
231                        );
232                    }
233                }
234            }
235        }
236        "method_declaration" | "method_definition" => {
237            let container = parent_nid.as_deref().unwrap_or(file_id);
238            if let Some(sel) = method_selector(source, &node) {
239                let method_nid = make_id(&[container, &sel]);
240                let label = format!("-{}", sel);
241                add_node_if_missing(
242                    fragment,
243                    oc_node(method_nid.clone(), label, "function", path),
244                );
245                oc_edge(fragment, container, &method_nid, "method", path);
246            }
247        }
248        _ => {
249            for i in 0..node.child_count() {
250                if let Some(child) = node.child(i) {
251                    walk_oc(
252                        child,
253                        source,
254                        file_id,
255                        stem,
256                        path,
257                        fragment,
258                        parent_nid.clone(),
259                    );
260                }
261            }
262        }
263    }
264}
265
266impl ObjCExtractor {
267    pub fn extract(source: &[u8], path: &Path) -> Result<ExtractionFragment> {
268        let (file_id, _, file_node) = make_file_node(path);
269        let stem = file_id.clone();
270        let mut fragment = ExtractionFragment {
271            nodes: vec![file_node],
272            edges: vec![],
273        };
274
275        let lang: tree_sitter::Language = tree_sitter_objc::LANGUAGE.into();
276        let mut parser = Parser::new();
277        parser
278            .set_language(&lang)
279            .map_err(|e| CodeSynapseError::Parse(format!("objc set_language: {e}")))?;
280        let tree = parser
281            .parse(source, None)
282            .ok_or_else(|| CodeSynapseError::Parse("objc parse failed".to_string()))?;
283
284        walk_oc(
285            tree.root_node(),
286            source,
287            &file_id,
288            &stem,
289            path,
290            &mut fragment,
291            None,
292        );
293
294        Ok(fragment)
295    }
296}
297
298impl LanguageExtractor for ObjCExtractor {
299    fn file_extensions(&self) -> Vec<&'static str> {
300        vec!["m", "mm"]
301    }
302    fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment> {
303        Self::extract(source, path)
304    }
305    fn resolve_imports(&self, _imports: &[ImportNode]) -> Vec<Edge> {
306        vec![]
307    }
308    fn collect_type_refs(&self, _fragment: &mut ExtractionFragment) {}
309}