1use serde::{Deserialize, Serialize};
4use sha2::{Digest, Sha256};
5use std::path::Path;
6use tree_sitter::{Language, Node, Parser};
7
8const SNAPSHOT_SCHEMA: &str = "xbp-code-snapshot-v1";
9
10#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
11#[serde(rename_all = "camelCase")]
12pub struct CodeSnapshot {
13 pub schema: String,
14 pub language: String,
15 pub parser_version: String,
16 pub path: String,
17 pub start_line: usize,
18 pub end_line: usize,
19 pub symbol_path: String,
20 pub name: String,
21 pub kind: String,
22 pub parameters: Vec<String>,
23 pub parameter_types: Vec<String>,
24 pub return_type: Option<String>,
25 pub local_binding_count: usize,
26 pub call_names: Vec<String>,
27 pub import_names: Vec<String>,
28 pub line_count: usize,
29 pub symbol_id: String,
30 pub signature_hash: String,
31 pub body_hash: String,
32}
33
34pub fn snapshot_for_marker(path: &str, source: &str, marker_line: usize) -> Option<CodeSnapshot> {
35 let snapshots = snapshots_for_file(path, source);
36 snapshots
37 .into_iter()
38 .filter(|snapshot| {
39 (snapshot.start_line..=snapshot.end_line).contains(&marker_line)
40 || (snapshot.start_line > marker_line && snapshot.start_line <= marker_line + 12)
41 })
42 .min_by_key(|snapshot| snapshot.end_line - snapshot.start_line)
43}
44
45pub fn snapshots_for_file(path: &str, source: &str) -> Vec<CodeSnapshot> {
46 let Some((language, language_name)) = language_for(path) else {
47 return Vec::new();
48 };
49 let mut parser = Parser::new();
50 if parser.set_language(&language).is_err() {
51 return Vec::new();
52 }
53 let Some(tree) = parser.parse(source, None) else {
54 return Vec::new();
55 };
56 let root = tree.root_node();
57 let mut functions = Vec::new();
58 collect_functions(root, &mut functions);
59 functions
60 .into_iter()
61 .filter_map(|node| build_snapshot(path, source, node, language_name))
62 .collect()
63}
64
65fn language_for(path: &str) -> Option<(Language, &'static str)> {
66 match Path::new(path).extension()?.to_str()?.to_ascii_lowercase().as_str() {
67 "rs" => Some((tree_sitter_rust::LANGUAGE.into(), "rust")),
68 "js" | "jsx" | "mjs" | "cjs" => {
69 Some((tree_sitter_javascript::LANGUAGE.into(), "javascript"))
70 }
71 "ts" => Some((tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(), "typescript")),
72 "tsx" => Some((tree_sitter_typescript::LANGUAGE_TSX.into(), "tsx")),
73 _ => None,
74 }
75}
76
77fn collect_functions<'a>(node: Node<'a>, output: &mut Vec<Node<'a>>) {
78 if is_function(node.kind()) {
79 output.push(node);
80 }
81 let mut cursor = node.walk();
82 for child in node.children(&mut cursor) {
83 collect_functions(child, output);
84 }
85}
86
87fn is_function(kind: &str) -> bool {
88 matches!(
89 kind,
90 "function_item"
91 | "function_declaration"
92 | "function_expression"
93 | "arrow_function"
94 | "method_definition"
95 | "method_declaration"
96 )
97}
98
99fn build_snapshot(
100 path: &str,
101 source: &str,
102 node: Node<'_>,
103 language: &str,
104) -> Option<CodeSnapshot> {
105 let body = node.child_by_field_name("body");
106 let name = node
107 .child_by_field_name("name")
108 .map(|n| text(n, source))
109 .or_else(|| node.child_by_field_name("function").map(|n| text(n, source)))
110 .unwrap_or_else(|| "<anonymous>".into());
111 let parameters = node
112 .child_by_field_name("parameters")
113 .map(|n| named_children_text(n, source))
114 .unwrap_or_default();
115 let parameter_types = parameters
116 .iter()
117 .filter_map(|parameter| parameter.split_once(':').map(|(_, ty)| ty.trim().to_string()))
118 .collect::<Vec<_>>();
119 let return_type = node
120 .child_by_field_name("return_type")
121 .map(|n| text(n, source));
122 let signature_end = body.map(|n| n.start_byte()).unwrap_or(node.end_byte());
123 let signature = normalize(&source[node.start_byte()..signature_end]);
124 let body_text = body.map(|n| text(n, source)).unwrap_or_default();
125 let mut calls = Vec::new();
126 let mut imports = Vec::new();
127 let mut locals = 0;
128 collect_metrics(node, source, &mut calls, &mut imports, &mut locals);
129 calls.sort();
130 calls.dedup();
131 imports.sort();
132 imports.dedup();
133 let mut parents = Vec::new();
134 let mut parent = node.parent();
135 while let Some(container) = parent {
136 let container_name = container
137 .child_by_field_name("name")
138 .or_else(|| container.child_by_field_name("type"))
139 .map(|child| text(child, source));
140 if matches!(
141 container.kind(),
142 "mod_item" | "impl_item" | "trait_item" | "class_declaration" | "class"
143 ) {
144 if let Some(container_name) = container_name {
145 parents.push(container_name);
146 }
147 }
148 parent = container.parent();
149 }
150 parents.reverse();
151 let qualified_name = if parents.is_empty() {
152 name.clone()
153 } else {
154 format!("{}::{}", parents.join("::"), name)
155 };
156 let symbol_path = format!("{}::{}", path.replace('\\', "/"), qualified_name);
157 Some(CodeSnapshot {
158 schema: SNAPSHOT_SCHEMA.into(),
159 language: language.into(),
160 parser_version: "tree-sitter-0.25".into(),
161 path: path.replace('\\', "/"),
162 start_line: node.start_position().row + 1,
163 end_line: node.end_position().row + 1,
164 line_count: node.end_position().row + 1 - node.start_position().row,
165 symbol_id: digest(&symbol_path),
166 symbol_path,
167 name,
168 kind: node.kind().into(),
169 parameters,
170 parameter_types,
171 return_type,
172 local_binding_count: locals,
173 call_names: calls,
174 import_names: imports,
175 signature_hash: digest(&signature),
176 body_hash: digest(&normalize(&body_text)),
177 })
178}
179
180fn collect_metrics(node: Node<'_>, source: &str, calls: &mut Vec<String>, imports: &mut Vec<String>, locals: &mut usize) {
181 match node.kind() {
182 "call_expression" => {
183 if let Some(function) = node.child_by_field_name("function") {
184 calls.push(text(function, source));
185 }
186 }
187 "macro_invocation" => {
188 if let Some(name) = node.child_by_field_name("macro") {
189 calls.push(format!("{}!", text(name, source)));
190 }
191 }
192 "use_declaration" | "import_statement" | "import_clause" => {
193 imports.push(normalize(&text(node, source)));
194 }
195 "let_declaration" | "lexical_declaration" | "variable_declarator" => *locals += 1,
196 _ => {}
197 }
198 let mut cursor = node.walk();
199 for child in node.children(&mut cursor) {
200 collect_metrics(child, source, calls, imports, locals);
201 }
202}
203
204fn named_children_text(node: Node<'_>, source: &str) -> Vec<String> {
205 let mut cursor = node.walk();
206 node.named_children(&mut cursor)
207 .map(|child| normalize(&text(child, source)))
208 .collect()
209}
210
211fn text(node: Node<'_>, source: &str) -> String {
212 source
213 .get(node.byte_range())
214 .unwrap_or_default()
215 .to_string()
216}
217
218fn normalize(value: &str) -> String {
219 value.split_whitespace().collect::<Vec<_>>().join(" ")
220}
221
222fn digest(value: &str) -> String {
223 let mut hasher = Sha256::new();
224 hasher.update(value.as_bytes());
225 format!("{:x}", hasher.finalize())
226}
227
228#[cfg(test)]
229mod tests {
230 use super::*;
231
232 #[test]
233 fn snapshots_rust_function_metadata_deterministically() {
234 let source = "use std::fmt;\n\nfn greet(name: &str) -> String { let value = format!(\"hi {name}\"); value }";
235 let snapshot = snapshot_for_marker("src/lib.rs", source, 2).expect("snapshot");
236 assert_eq!(snapshot.name, "greet");
237 assert_eq!(snapshot.start_line, 3);
238 assert_eq!(snapshot.return_type.as_deref(), Some("String"));
239 assert_eq!(snapshot.local_binding_count, 1);
240 assert!(snapshot.call_names.iter().any(|call| call == "format!"));
241 assert_eq!(snapshot, snapshot_for_marker("src/lib.rs", source, 2).unwrap());
242 }
243
244 #[test]
245 fn snapshots_javascript_and_typescript_functions() {
246 let js = snapshot_for_marker("src/a.js", "function add(a, b) { return helper(a) + b; }", 1).unwrap();
247 let ts = snapshot_for_marker("src/a.ts", "function add(a: number, b: number): number { return a + b; }", 1).unwrap();
248 assert_eq!(js.name, "add");
249 assert_eq!(ts.parameter_types, vec!["number", "number"]);
250 assert_eq!(ts.return_type.as_deref(), Some(": number"));
251 assert_ne!(js.body_hash, ts.body_hash);
252 }
253}