use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::path::Path;
use tree_sitter::{Language, Node, Parser};
const SNAPSHOT_SCHEMA: &str = "xbp-code-snapshot-v1";
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CodeSnapshot {
pub schema: String,
pub language: String,
pub parser_version: String,
pub path: String,
pub start_line: usize,
pub end_line: usize,
pub symbol_path: String,
pub name: String,
pub kind: String,
pub parameters: Vec<String>,
pub parameter_types: Vec<String>,
pub return_type: Option<String>,
pub local_binding_count: usize,
pub call_names: Vec<String>,
pub import_names: Vec<String>,
pub line_count: usize,
pub symbol_id: String,
pub signature_hash: String,
pub body_hash: String,
}
pub fn snapshot_for_marker(path: &str, source: &str, marker_line: usize) -> Option<CodeSnapshot> {
let snapshots = snapshots_for_file(path, source);
snapshots
.into_iter()
.filter(|snapshot| {
(snapshot.start_line..=snapshot.end_line).contains(&marker_line)
|| (snapshot.start_line > marker_line && snapshot.start_line <= marker_line + 12)
})
.min_by_key(|snapshot| snapshot.end_line - snapshot.start_line)
}
pub fn snapshots_for_file(path: &str, source: &str) -> Vec<CodeSnapshot> {
let Some((language, language_name)) = language_for(path) else {
return Vec::new();
};
let mut parser = Parser::new();
if parser.set_language(&language).is_err() {
return Vec::new();
}
let Some(tree) = parser.parse(source, None) else {
return Vec::new();
};
let root = tree.root_node();
let mut functions = Vec::new();
collect_functions(root, &mut functions);
functions
.into_iter()
.filter_map(|node| build_snapshot(path, source, node, language_name))
.collect()
}
fn language_for(path: &str) -> Option<(Language, &'static str)> {
match Path::new(path).extension()?.to_str()?.to_ascii_lowercase().as_str() {
"rs" => Some((tree_sitter_rust::LANGUAGE.into(), "rust")),
"js" | "jsx" | "mjs" | "cjs" => {
Some((tree_sitter_javascript::LANGUAGE.into(), "javascript"))
}
"ts" => Some((tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(), "typescript")),
"tsx" => Some((tree_sitter_typescript::LANGUAGE_TSX.into(), "tsx")),
_ => None,
}
}
fn collect_functions<'a>(node: Node<'a>, output: &mut Vec<Node<'a>>) {
if is_function(node.kind()) {
output.push(node);
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
collect_functions(child, output);
}
}
fn is_function(kind: &str) -> bool {
matches!(
kind,
"function_item"
| "function_declaration"
| "function_expression"
| "arrow_function"
| "method_definition"
| "method_declaration"
)
}
fn build_snapshot(
path: &str,
source: &str,
node: Node<'_>,
language: &str,
) -> Option<CodeSnapshot> {
let body = node.child_by_field_name("body");
let name = node
.child_by_field_name("name")
.map(|n| text(n, source))
.or_else(|| node.child_by_field_name("function").map(|n| text(n, source)))
.unwrap_or_else(|| "<anonymous>".into());
let parameters = node
.child_by_field_name("parameters")
.map(|n| named_children_text(n, source))
.unwrap_or_default();
let parameter_types = parameters
.iter()
.filter_map(|parameter| parameter.split_once(':').map(|(_, ty)| ty.trim().to_string()))
.collect::<Vec<_>>();
let return_type = node
.child_by_field_name("return_type")
.map(|n| text(n, source));
let signature_end = body.map(|n| n.start_byte()).unwrap_or(node.end_byte());
let signature = normalize(&source[node.start_byte()..signature_end]);
let body_text = body.map(|n| text(n, source)).unwrap_or_default();
let mut calls = Vec::new();
let mut imports = Vec::new();
let mut locals = 0;
collect_metrics(node, source, &mut calls, &mut imports, &mut locals);
calls.sort();
calls.dedup();
imports.sort();
imports.dedup();
let mut parents = Vec::new();
let mut parent = node.parent();
while let Some(container) = parent {
let container_name = container
.child_by_field_name("name")
.or_else(|| container.child_by_field_name("type"))
.map(|child| text(child, source));
if matches!(
container.kind(),
"mod_item" | "impl_item" | "trait_item" | "class_declaration" | "class"
) {
if let Some(container_name) = container_name {
parents.push(container_name);
}
}
parent = container.parent();
}
parents.reverse();
let qualified_name = if parents.is_empty() {
name.clone()
} else {
format!("{}::{}", parents.join("::"), name)
};
let symbol_path = format!("{}::{}", path.replace('\\', "/"), qualified_name);
Some(CodeSnapshot {
schema: SNAPSHOT_SCHEMA.into(),
language: language.into(),
parser_version: "tree-sitter-0.25".into(),
path: path.replace('\\', "/"),
start_line: node.start_position().row + 1,
end_line: node.end_position().row + 1,
line_count: node.end_position().row + 1 - node.start_position().row,
symbol_id: digest(&symbol_path),
symbol_path,
name,
kind: node.kind().into(),
parameters,
parameter_types,
return_type,
local_binding_count: locals,
call_names: calls,
import_names: imports,
signature_hash: digest(&signature),
body_hash: digest(&normalize(&body_text)),
})
}
fn collect_metrics(node: Node<'_>, source: &str, calls: &mut Vec<String>, imports: &mut Vec<String>, locals: &mut usize) {
match node.kind() {
"call_expression" => {
if let Some(function) = node.child_by_field_name("function") {
calls.push(text(function, source));
}
}
"macro_invocation" => {
if let Some(name) = node.child_by_field_name("macro") {
calls.push(format!("{}!", text(name, source)));
}
}
"use_declaration" | "import_statement" | "import_clause" => {
imports.push(normalize(&text(node, source)));
}
"let_declaration" | "lexical_declaration" | "variable_declarator" => *locals += 1,
_ => {}
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
collect_metrics(child, source, calls, imports, locals);
}
}
fn named_children_text(node: Node<'_>, source: &str) -> Vec<String> {
let mut cursor = node.walk();
node.named_children(&mut cursor)
.map(|child| normalize(&text(child, source)))
.collect()
}
fn text(node: Node<'_>, source: &str) -> String {
source
.get(node.byte_range())
.unwrap_or_default()
.to_string()
}
fn normalize(value: &str) -> String {
value.split_whitespace().collect::<Vec<_>>().join(" ")
}
fn digest(value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn snapshots_rust_function_metadata_deterministically() {
let source = "use std::fmt;\n\nfn greet(name: &str) -> String { let value = format!(\"hi {name}\"); value }";
let snapshot = snapshot_for_marker("src/lib.rs", source, 2).expect("snapshot");
assert_eq!(snapshot.name, "greet");
assert_eq!(snapshot.start_line, 3);
assert_eq!(snapshot.return_type.as_deref(), Some("String"));
assert_eq!(snapshot.local_binding_count, 1);
assert!(snapshot.call_names.iter().any(|call| call == "format!"));
assert_eq!(snapshot, snapshot_for_marker("src/lib.rs", source, 2).unwrap());
}
#[test]
fn snapshots_javascript_and_typescript_functions() {
let js = snapshot_for_marker("src/a.js", "function add(a, b) { return helper(a) + b; }", 1).unwrap();
let ts = snapshot_for_marker("src/a.ts", "function add(a: number, b: number): number { return a + b; }", 1).unwrap();
assert_eq!(js.name, "add");
assert_eq!(ts.parameter_types, vec!["number", "number"]);
assert_eq!(ts.return_type.as_deref(), Some(": number"));
assert_ne!(js.body_hash, ts.body_hash);
}
}