Skip to main content

codesynapse_core/ts_extract/
verilog.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 VerilogExtractor;
10
11fn vl_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 vl_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 vl_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 walk_vl<'t>(
51    node: TsNode<'t>,
52    source: &[u8],
53    file_id: &str,
54    stem: &str,
55    path: &Path,
56    fragment: &mut ExtractionFragment,
57    module_nid: Option<String>,
58) {
59    match node.kind() {
60        "module_declaration" => {
61            let name_node = node.child_by_field_name("name");
62            if let Some(n) = name_node {
63                let name = vl_text(source, &n);
64                if !name.is_empty() {
65                    let nid = make_id(&[stem, &name]);
66                    fragment
67                        .nodes
68                        .push(vl_node(nid.clone(), name, "module", path));
69                    vl_edge(fragment, file_id, &nid, "defines", path);
70                    for i in 0..node.child_count() {
71                        if let Some(child) = node.child(i) {
72                            walk_vl(
73                                child,
74                                source,
75                                file_id,
76                                stem,
77                                path,
78                                fragment,
79                                Some(nid.clone()),
80                            );
81                        }
82                    }
83                    return;
84                }
85            }
86            for i in 0..node.child_count() {
87                if let Some(child) = node.child(i) {
88                    walk_vl(
89                        child,
90                        source,
91                        file_id,
92                        stem,
93                        path,
94                        fragment,
95                        module_nid.clone(),
96                    );
97                }
98            }
99        }
100        "function_declaration" | "function_prototype" => {
101            let name_node = node.child_by_field_name("name");
102            if let Some(n) = name_node {
103                let name = vl_text(source, &n);
104                if !name.is_empty() {
105                    let parent = module_nid.as_deref().unwrap_or(file_id);
106                    let nid = make_id(&[parent, &name]);
107                    fragment.nodes.push(vl_node(
108                        nid.clone(),
109                        format!("{}()", name),
110                        "function",
111                        path,
112                    ));
113                    vl_edge(fragment, parent, &nid, "contains", path);
114                }
115            }
116        }
117        "task_declaration" => {
118            let name_node = node.child_by_field_name("name");
119            if let Some(n) = name_node {
120                let name = vl_text(source, &n);
121                if !name.is_empty() {
122                    let parent = module_nid.as_deref().unwrap_or(file_id);
123                    let nid = make_id(&[parent, &name]);
124                    fragment
125                        .nodes
126                        .push(vl_node(nid.clone(), name, "code", path));
127                    vl_edge(fragment, parent, &nid, "contains", path);
128                }
129            }
130        }
131        "package_import_declaration" => {
132            for i in 0..node.child_count() {
133                if let Some(child) = node.child(i) {
134                    if child.kind() == "package_import_item" {
135                        let text = vl_text(source, &child);
136                        let pkg_name = text.split("::").next().unwrap_or("").trim().to_string();
137                        if !pkg_name.is_empty() {
138                            let tgt_nid = make_id(&[&pkg_name]);
139                            let tgt_node = vl_node(tgt_nid.clone(), pkg_name, "module", path);
140                            add_node_if_missing(fragment, tgt_node);
141                            let src = module_nid.as_deref().unwrap_or(file_id);
142                            vl_edge(fragment, src, &tgt_nid, "imports_from", path);
143                        }
144                    }
145                }
146            }
147        }
148        "module_instantiation" => {
149            let type_node = node.child_by_field_name("module_type");
150            if let Some(n) = type_node {
151                if let Some(mod_nid) = &module_nid {
152                    let inst_type = vl_text(source, &n);
153                    if !inst_type.is_empty() {
154                        let tgt_nid = make_id(&[&inst_type]);
155                        let tgt_node = vl_node(tgt_nid.clone(), inst_type, "module", path);
156                        add_node_if_missing(fragment, tgt_node);
157                        vl_edge(fragment, mod_nid, &tgt_nid, "instantiates", path);
158                    }
159                }
160            }
161        }
162        _ => {
163            for i in 0..node.child_count() {
164                if let Some(child) = node.child(i) {
165                    walk_vl(
166                        child,
167                        source,
168                        file_id,
169                        stem,
170                        path,
171                        fragment,
172                        module_nid.clone(),
173                    );
174                }
175            }
176        }
177    }
178}
179
180impl VerilogExtractor {
181    pub fn extract(source: &[u8], path: &Path) -> Result<ExtractionFragment> {
182        let (file_id, _, file_node) = make_file_node(path);
183        let stem = file_id.clone();
184        let mut fragment = ExtractionFragment {
185            nodes: vec![file_node],
186            edges: vec![],
187        };
188
189        let lang: tree_sitter::Language = tree_sitter_verilog::LANGUAGE.into();
190        let mut parser = Parser::new();
191        parser
192            .set_language(&lang)
193            .map_err(|e| CodeSynapseError::Parse(format!("verilog set_language: {e}")))?;
194        let tree = parser
195            .parse(source, None)
196            .ok_or_else(|| CodeSynapseError::Parse("verilog parse failed".to_string()))?;
197
198        walk_vl(
199            tree.root_node(),
200            source,
201            &file_id,
202            &stem,
203            path,
204            &mut fragment,
205            None,
206        );
207
208        Ok(fragment)
209    }
210}
211
212impl LanguageExtractor for VerilogExtractor {
213    fn file_extensions(&self) -> Vec<&'static str> {
214        vec!["v", "sv", "svh"]
215    }
216    fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment> {
217        Self::extract(source, path)
218    }
219    fn resolve_imports(&self, _imports: &[ImportNode]) -> Vec<Edge> {
220        vec![]
221    }
222    fn collect_type_refs(&self, _fragment: &mut ExtractionFragment) {}
223}