codesynapse_core/ts_extract/
fortran.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 FortranExtractor;
10
11fn ft_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_lowercase()
16}
17
18fn ft_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 ft_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
50/// Extract the name from a Fortran statement node (program_statement,
51/// module_statement, subroutine_statement, function_statement, or use_statement).
52/// Fortran is case-insensitive — returns lowercase.
53fn fortran_name(source: &[u8], stmt: &TsNode<'_>) -> Option<String> {
54    for i in 0..stmt.child_count() {
55        if let Some(child) = stmt.child(i) {
56            match child.kind() {
57                "name" | "identifier" | "module_name" => {
58                    let t = ft_text(source, &child);
59                    if !t.is_empty() {
60                        return Some(t);
61                    }
62                }
63                _ => {}
64            }
65        }
66    }
67    None
68}
69
70fn walk_ft<'t>(
71    node: TsNode<'t>,
72    source: &[u8],
73    file_id: &str,
74    stem: &str,
75    path: &Path,
76    fragment: &mut ExtractionFragment,
77    scope_nid: String,
78) {
79    match node.kind() {
80        "program" => {
81            let stmt = (0..node.child_count())
82                .filter_map(|i| node.child(i))
83                .find(|c| c.kind() == "program_statement");
84            let name = stmt.and_then(|s| fortran_name(source, &s));
85            if let Some(name) = name {
86                let nid = make_id(&[stem, &name]);
87                fragment
88                    .nodes
89                    .push(ft_node(nid.clone(), name, "code", path));
90                ft_edge(fragment, file_id, &nid, "defines", path);
91                for i in 0..node.child_count() {
92                    if let Some(child) = node.child(i) {
93                        walk_ft(child, source, file_id, stem, path, fragment, nid.clone());
94                    }
95                }
96            }
97        }
98        "module" => {
99            let stmt = (0..node.child_count())
100                .filter_map(|i| node.child(i))
101                .find(|c| c.kind() == "module_statement");
102            let name = stmt.and_then(|s| fortran_name(source, &s));
103            if let Some(name) = name {
104                let nid = make_id(&[stem, &name]);
105                fragment
106                    .nodes
107                    .push(ft_node(nid.clone(), name, "module", path));
108                ft_edge(fragment, file_id, &nid, "defines", path);
109                for i in 0..node.child_count() {
110                    if let Some(child) = node.child(i) {
111                        walk_ft(child, source, file_id, stem, path, fragment, nid.clone());
112                    }
113                }
114            }
115        }
116        "internal_procedures" => {
117            for i in 0..node.child_count() {
118                if let Some(child) = node.child(i) {
119                    walk_ft(
120                        child,
121                        source,
122                        file_id,
123                        stem,
124                        path,
125                        fragment,
126                        scope_nid.clone(),
127                    );
128                }
129            }
130        }
131        "subroutine" => {
132            let stmt = (0..node.child_count())
133                .filter_map(|i| node.child(i))
134                .find(|c| c.kind() == "subroutine_statement");
135            let name = stmt.and_then(|s| fortran_name(source, &s));
136            if let Some(name) = name {
137                let nid = make_id(&[stem, &name]);
138                fragment.nodes.push(ft_node(
139                    nid.clone(),
140                    format!("{}()", name),
141                    "function",
142                    path,
143                ));
144                ft_edge(fragment, &scope_nid, &nid, "defines", path);
145                for i in 0..node.child_count() {
146                    if let Some(child) = node.child(i) {
147                        walk_ft(child, source, file_id, stem, path, fragment, nid.clone());
148                    }
149                }
150            }
151        }
152        "function" => {
153            let stmt = (0..node.child_count())
154                .filter_map(|i| node.child(i))
155                .find(|c| c.kind() == "function_statement");
156            let name = stmt.and_then(|s| fortran_name(source, &s));
157            if let Some(name) = name {
158                let nid = make_id(&[stem, &name]);
159                fragment.nodes.push(ft_node(
160                    nid.clone(),
161                    format!("{}()", name),
162                    "function",
163                    path,
164                ));
165                ft_edge(fragment, &scope_nid, &nid, "defines", path);
166                for i in 0..node.child_count() {
167                    if let Some(child) = node.child(i) {
168                        walk_ft(child, source, file_id, stem, path, fragment, nid.clone());
169                    }
170                }
171            }
172        }
173        "use_statement" => {
174            let name = fortran_name(source, &node);
175            if let Some(mod_name) = name {
176                let imp_nid = make_id(&[&mod_name]);
177                let imp_node = ft_node(imp_nid.clone(), mod_name, "module", path);
178                add_node_if_missing(fragment, imp_node);
179                fragment.edges.push(Edge {
180                    source: scope_nid.to_string(),
181                    target: imp_nid,
182                    relation: "imports".to_string(),
183                    confidence: "EXTRACTED".to_string(),
184                    source_file: Some(path.to_string_lossy().to_string()),
185                    weight: 1.0,
186                    context: Some("use".to_string()),
187                });
188            }
189        }
190        _ => {
191            for i in 0..node.child_count() {
192                if let Some(child) = node.child(i) {
193                    walk_ft(
194                        child,
195                        source,
196                        file_id,
197                        stem,
198                        path,
199                        fragment,
200                        scope_nid.clone(),
201                    );
202                }
203            }
204        }
205    }
206}
207
208impl FortranExtractor {
209    pub fn extract(source: &[u8], path: &Path) -> Result<ExtractionFragment> {
210        let (file_id, _, file_node) = make_file_node(path);
211        let stem = file_id.clone();
212        let mut fragment = ExtractionFragment {
213            nodes: vec![file_node],
214            edges: vec![],
215        };
216
217        let lang: tree_sitter::Language = tree_sitter_fortran::LANGUAGE.into();
218        let mut parser = Parser::new();
219        parser
220            .set_language(&lang)
221            .map_err(|e| CodeSynapseError::Parse(format!("fortran set_language: {e}")))?;
222        let tree = parser
223            .parse(source, None)
224            .ok_or_else(|| CodeSynapseError::Parse("fortran parse failed".to_string()))?;
225
226        let scope = file_id.clone();
227        walk_ft(
228            tree.root_node(),
229            source,
230            &file_id,
231            &stem,
232            path,
233            &mut fragment,
234            scope,
235        );
236
237        Ok(fragment)
238    }
239}
240
241impl LanguageExtractor for FortranExtractor {
242    fn file_extensions(&self) -> Vec<&'static str> {
243        vec![
244            "f", "f90", "f95", "f03", "f08", "F", "F90", "F95", "F03", "F08",
245        ]
246    }
247    fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment> {
248        Self::extract(source, path)
249    }
250    fn resolve_imports(&self, _imports: &[ImportNode]) -> Vec<Edge> {
251        vec![]
252    }
253    fn collect_type_refs(&self, _fragment: &mut ExtractionFragment) {}
254}
codesynapse_core/ts_extract/fortran.rs

codesynapse_core/ts_extract/
fortran.rs
