pelicode_core/
extract.rs

1//! Function extraction and feature computation for a single source file.
2
3use std::collections::BTreeMap;
4use std::path::{Path, PathBuf};
5
6use serde::Serialize;
7use sha2::{Digest, Sha256};
8use tree_sitter::{Node, Parser};
9
10use crate::Error;
11use crate::config::{Config, Level};
12use crate::hash::{self, Fingerprint, Hash128};
13use crate::language::{Language, LanguageSpec};
14use crate::normalize::normalize;
15
16/// Features of one source file.
17#[derive(Debug, Clone, Serialize)]
18pub struct FileFeatures {
19    /// Path of the file as recorded, with `/` as the separator.
20    pub path: PathBuf,
21    /// Language of the file.
22    pub language: Language,
23    /// SHA-256 of the file content (hex).
24    pub content_sha256: String,
25    /// Functions of the file, in the order they start; a function precedes the functions nested in it.
26    pub functions: Vec<FunctionFeatures>,
27}
28
29/// Features of one function unit.
30#[derive(Debug, Clone, Serialize)]
31pub struct FunctionFeatures {
32    /// Function name for display; `<anonymous>` for lambdas and blocks. Never used for matching.
33    pub name: String,
34    /// Tree-sitter node kind of the function.
35    pub kind: String,
36    /// 1-based, inclusive.
37    pub start_line: usize,
38    /// Last line of the function (1-based, inclusive).
39    pub end_line: usize,
40    /// 1-based byte columns of the first and last character.
41    pub start_column: usize,
42    /// 1-based byte column of the last character.
43    pub end_column: usize,
44    /// Byte range in the file (0-based, end exclusive).
45    pub start_byte: usize,
46    /// End of the byte range in the file (exclusive).
47    pub end_byte: usize,
48    /// Index of the enclosing function in the same file.
49    pub parent: Option<usize>,
50    /// Number of N0 tokens.
51    pub token_count: usize,
52    /// Number of named AST nodes (comments excluded).
53    pub node_count: usize,
54    /// Depth of the syntax tree of the function.
55    pub depth: usize,
56    /// Hash of the whole function at every normalization level.
57    pub exact_hashes: BTreeMap<Level, Hash128>,
58    /// Present only for functions with at least `min_tokens` tokens.
59    pub features: BTreeMap<Level, LevelFeatures>,
60    /// Normalized tokens at every level, kept only with [`Config::keep_tokens`].
61    #[serde(skip_serializing_if = "Option::is_none")]
62    pub tokens: Option<BTreeMap<Level, Vec<String>>>,
63}
64
65impl FunctionFeatures {
66    /// Whether this function is large enough to carry fingerprints.
67    ///
68    /// # Examples
69    ///
70    /// ```
71    /// use pelicode_core::{Config, Language, extract_source};
72    ///
73    /// let source = b"def area(w, h):\n    return w * h\n";
74    /// let small = Config { min_tokens: 100, ..Config::default() };
75    /// assert!(!extract_source(Language::Python, source, &small).unwrap()[0].is_indexed());
76    /// let all = Config { min_tokens: 0, ..Config::default() };
77    /// assert!(extract_source(Language::Python, source, &all).unwrap()[0].is_indexed());
78    /// ```
79    pub fn is_indexed(&self) -> bool {
80        !self.features.is_empty()
81    }
82}
83
84/// Partial-match features of one normalization level.
85#[derive(Debug, Clone, Serialize)]
86pub struct LevelFeatures {
87    /// Fingerprints selected by winnowing, in token order.
88    pub fingerprints: Vec<Fingerprint>,
89    /// MinHash signature of the fingerprint set.
90    #[serde(serialize_with = "hash::serialize_u64_hex_vec")]
91    pub minhash: Vec<u64>,
92    /// LSH bucket of each band of the MinHash signature.
93    #[serde(serialize_with = "hash::serialize_u64_hex_vec")]
94    pub lsh_buckets: Vec<u64>,
95}
96
97/// SHA-256 of `content` as 64 lower-case hexadecimal digits, the identity of a file.
98///
99/// # Examples
100///
101/// ```
102/// use pelicode_core::extract::sha256_hex;
103///
104/// assert_eq!(
105///     sha256_hex(b"abc"),
106///     "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
107/// );
108/// ```
109pub fn sha256_hex(content: &[u8]) -> String {
110    hash::to_hex(&Sha256::digest(content))
111}
112
113/// Reads and extracts a file; `display_path` is the path recorded in the result.
114///
115/// # Errors
116///
117/// [`Error::Io`] when the file cannot be read, and the errors of [`extract_source`].
118///
119/// # Examples
120///
121/// ```
122/// use std::path::Path;
123/// use pelicode_core::{Config, Language, extract_file};
124///
125/// let dir = tempfile::tempdir().unwrap();
126/// let path = dir.path().join("greet.rb");
127/// std::fs::write(&path, "def greet(name)\n  \"Hello, #{name}\"\nend\n").unwrap();
128///
129/// let file = extract_file(&path, Path::new("lib/greet.rb"), Language::Ruby, &Config::default()).unwrap();
130/// assert_eq!(file.path, Path::new("lib/greet.rb"));
131/// assert_eq!(file.functions[0].name, "greet");
132/// ```
133pub fn extract_file(
134    path: &Path,
135    display_path: &Path,
136    language: Language,
137    config: &Config,
138) -> Result<FileFeatures, Error> {
139    let content = std::fs::read(path).map_err(|e| Error::Io(path.to_path_buf(), e))?;
140    extract_bytes(display_path, language, &content, config)
141}
142
143/// Extracts a file whose content is already in memory (e.g. read from a git revision).
144///
145/// # Errors
146///
147/// The errors of [`extract_source`].
148///
149/// # Examples
150///
151/// ```
152/// use std::path::Path;
153/// use pelicode_core::{Config, Language, extract_bytes};
154///
155/// let source = b"package main\nfunc Add(a, b int) int { return a + b }\n";
156/// let file = extract_bytes(Path::new("add.go"), Language::Go, source, &Config::default()).unwrap();
157/// assert_eq!(file.functions[0].name, "Add");
158/// assert_eq!(file.content_sha256.len(), 64);
159/// ```
160pub fn extract_bytes(
161    path: &Path,
162    language: Language,
163    content: &[u8],
164    config: &Config,
165) -> Result<FileFeatures, Error> {
166    Ok(FileFeatures {
167        path: path.to_path_buf(),
168        language,
169        content_sha256: sha256_hex(content),
170        functions: extract_source(language, content, config)?,
171    })
172}
173
174/// Extracts the features of every function in `source`.
175///
176/// Every function, method, constructor, lambda or block of the language is a
177/// function unit, nested ones included. Functions with syntax errors are skipped.
178/// Fingerprints, MinHash signatures and LSH buckets are computed only for
179/// functions with at least [`Config::min_tokens`] tokens.
180///
181/// # Errors
182///
183/// [`Error::InvalidConfig`] when `config` is invalid, and [`Error::Parse`] when
184/// the source cannot be parsed at all.
185///
186/// # Examples
187///
188/// Renaming identifiers changes the N0 hash but not the N1 hash:
189///
190/// ```
191/// use pelicode_core::{Config, Language, Level, extract_source};
192///
193/// let a = extract_source(Language::Java, b"class A { int twice(int x) { return x * 2; } }", &Config::default()).unwrap();
194/// let b = extract_source(Language::Java, b"class B { int doubled(int v) { return v * 2; } }", &Config::default()).unwrap();
195/// assert_eq!(a[0].name, "twice");
196/// assert_ne!(a[0].exact_hashes[&Level::N0], b[0].exact_hashes[&Level::N0]);
197/// assert_eq!(a[0].exact_hashes[&Level::N1], b[0].exact_hashes[&Level::N1]);
198/// ```
199pub fn extract_source(
200    language: Language,
201    source: &[u8],
202    config: &Config,
203) -> Result<Vec<FunctionFeatures>, Error> {
204    config.validate()?;
205    let mut parser = Parser::new();
206    parser
207        .set_language(&language.ts_language())
208        .map_err(|e| Error::Parse(e.to_string()))?;
209    let tree = parser
210        .parse(source, None)
211        .ok_or_else(|| Error::Parse("parser returned no tree".into()))?;
212
213    let spec = language.spec();
214    let mut functions = Vec::new();
215    collect(
216        tree.root_node(),
217        None,
218        None,
219        source,
220        spec,
221        config,
222        &mut functions,
223    );
224    Ok(functions)
225}
226
227fn collect(
228    node: Node<'_>,
229    parent_kind: Option<&str>,
230    parent_index: Option<usize>,
231    source: &[u8],
232    spec: &'static LanguageSpec,
233    config: &Config,
234    out: &mut Vec<FunctionFeatures>,
235) {
236    let kind = node.kind();
237    let mut current = parent_index;
238    // A function directly wrapping another (e.g. Ruby `-> { }` with its block) is one unit.
239    let wrapped = parent_kind.is_some_and(|p| spec.functions.contains(&p));
240    if spec.functions.contains(&kind) && !wrapped && !node.has_error() {
241        out.push(features(node, parent_index, source, spec, config));
242        current = Some(out.len() - 1);
243    }
244    let mut cursor = node.walk();
245    for child in node.children(&mut cursor) {
246        collect(child, Some(kind), current, source, spec, config, out);
247    }
248}
249
250fn features(
251    node: Node<'_>,
252    parent: Option<usize>,
253    source: &[u8],
254    spec: &'static LanguageSpec,
255    config: &Config,
256) -> FunctionFeatures {
257    let normalized = normalize(node, source, spec, config.ident_mode);
258    let token_count = normalized.tokens[&Level::N0].len();
259    let exact_hashes = normalized
260        .tokens
261        .iter()
262        .map(|(level, tokens)| (*level, Hash128::of_tokens(tokens)))
263        .collect();
264
265    let mut features = BTreeMap::new();
266    if token_count >= config.min_tokens {
267        for level in &config.fingerprint_levels {
268            let tokens = &normalized.tokens[level];
269            let rows = &normalized.rows[level];
270            let k = config.kgram.min(tokens.len());
271            let fingerprints: Vec<Fingerprint> =
272                hash::winnow(&hash::kgram_hashes(tokens, config.kgram), config.window)
273                    .into_iter()
274                    .map(|(position, hash)| Fingerprint {
275                        hash,
276                        position: position as u32,
277                        start_line: rows[position].0 as u32 + 1,
278                        end_line: rows[position + k - 1].1 as u32 + 1,
279                    })
280                    .collect();
281            let Some(minhash) =
282                hash::minhash(fingerprints.iter().map(|f| f.hash), config.minhash_size)
283            else {
284                continue;
285            };
286            let lsh_buckets = hash::lsh_buckets(&minhash, config.bands);
287            features.insert(
288                *level,
289                LevelFeatures {
290                    fingerprints,
291                    minhash,
292                    lsh_buckets,
293                },
294            );
295        }
296    }
297
298    FunctionFeatures {
299        name: function_name(node, source),
300        kind: node.kind().to_string(),
301        start_line: node.start_position().row + 1,
302        end_line: node.end_position().row + 1,
303        start_column: node.start_position().column + 1,
304        // The end position points just past the last character.
305        end_column: node.end_position().column.max(1),
306        start_byte: node.start_byte(),
307        end_byte: node.end_byte(),
308        parent,
309        token_count,
310        node_count: normalized.node_count,
311        depth: normalized.depth,
312        exact_hashes,
313        features,
314        tokens: config.keep_tokens.then_some(normalized.tokens),
315    }
316}
317
318fn function_name(node: Node<'_>, source: &[u8]) -> String {
319    if let Some(name) = node.child_by_field_name("name")
320        && let Ok(text) = name.utf8_text(source)
321    {
322        return text.to_string();
323    }
324    match node.kind() {
325        "getter" => "get".into(),
326        "setter" => "set".into(),
327        "secondary_constructor" => "constructor".into(),
328        _ => "<anonymous>".into(),
329    }
330}
pelicode_core/extract.rs

pelicode_core/
extract.rs
