semantic/analysis/
analysis_similarity.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Similarity computation utilities.
3
4use std::collections::{HashMap, HashSet};
5
6use crate::parser::{Language, ParsedFile};
7
8/// Method for computing content similarity.
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10pub enum SimilarityMethod {
11    /// Simple line-by-line comparison.
12    Lines,
13    /// Token-based comparison (ignores whitespace).
14    Tokens,
15    /// AST-based comparison (structure only).
16    Ast,
17}
18
19/// Compute similarity between two strings (0.0 to 1.0).
20pub fn compute_similarity(a: &str, b: &str, method: SimilarityMethod) -> f64 {
21    match method {
22        SimilarityMethod::Lines => {
23            let lines_a: HashSet<&str> = a.lines().filter(|l| !l.trim().is_empty()).collect();
24            let lines_b: HashSet<&str> = b.lines().filter(|l| !l.trim().is_empty()).collect();
25
26            if lines_a.is_empty() && lines_b.is_empty() {
27                return 1.0;
28            }
29            if lines_a.is_empty() || lines_b.is_empty() {
30                return 0.0;
31            }
32
33            let intersection: HashSet<_> = lines_a.intersection(&lines_b).collect();
34            let union: HashSet<_> = lines_a.union(&lines_b).collect();
35
36            let line_similarity = intersection.len() as f64 / union.len() as f64;
37            if line_similarity == 0.0 {
38                return compute_similarity(a, b, SimilarityMethod::Tokens);
39            }
40
41            line_similarity
42        }
43        SimilarityMethod::Tokens => {
44            let tokens_a: HashSet<&str> = a.split_whitespace().collect();
45            let tokens_b: HashSet<&str> = b.split_whitespace().collect();
46
47            if tokens_a.is_empty() && tokens_b.is_empty() {
48                return 1.0;
49            }
50            if tokens_a.is_empty() || tokens_b.is_empty() {
51                return 0.0;
52            }
53
54            let intersection: HashSet<_> = tokens_a.intersection(&tokens_b).collect();
55            let union: HashSet<_> = tokens_a.union(&tokens_b).collect();
56
57            intersection.len() as f64 / union.len() as f64
58        }
59        // AST similarity is language-dependent: without a grammar there is no
60        // tree to compare. The language-free entry point therefore cannot
61        // honor `Ast` itself — it forwards to the one sanctioned AST path,
62        // which degrades to token similarity for `Language::Unknown` rather
63        // than silently masquerading token similarity as an AST result.
64        SimilarityMethod::Ast => {
65            compute_similarity_with_language(a, b, SimilarityMethod::Ast, Language::Unknown)
66        }
67    }
68}
69
70pub fn compute_similarity_with_language(
71    a: &str,
72    b: &str,
73    method: SimilarityMethod,
74    language: Language,
75) -> f64 {
76    match method {
77        SimilarityMethod::Ast => {
78            if let Some(score) = compute_ast_similarity(a, b, language) {
79                return score;
80            }
81            compute_similarity(a, b, SimilarityMethod::Tokens)
82        }
83        _ => compute_similarity(a, b, method),
84    }
85}
86
87/// AST kind-bag similarity with no token fallback.
88///
89/// `None` means a language has no grammar or either side failed to
90/// parse. Callers that must not invent a novelty/uniqueness signal
91/// from identifier tokens should treat that as fail-closed.
92pub fn try_compute_ast_similarity(a: &str, b: &str, language: Language) -> Option<f64> {
93    try_compute_ast_similarity_for_languages(a, language, b, language)
94}
95
96/// Like [`try_compute_ast_similarity`], but each side uses its own grammar.
97pub fn try_compute_ast_similarity_for_languages(
98    a: &str,
99    a_language: Language,
100    b: &str,
101    b_language: Language,
102) -> Option<f64> {
103    if a_language == Language::Unknown || b_language == Language::Unknown {
104        return None;
105    }
106    let counts_a = ast_node_counts(a, a_language)?;
107    let counts_b = ast_node_counts(b, b_language)?;
108    Some(count_similarity(&counts_a, &counts_b))
109}
110
111pub(super) struct PreparedSimilarity {
112    method: SimilarityMethod,
113    lines: Option<HashSet<String>>,
114    tokens: HashSet<String>,
115    ast_counts: HashMap<Language, Option<HashMap<String, usize>>>,
116}
117
118impl PreparedSimilarity {
119    pub(super) fn new(
120        content: &str,
121        method: SimilarityMethod,
122        languages: impl IntoIterator<Item = Language>,
123    ) -> Self {
124        let tokens = content
125            .split_whitespace()
126            .map(String::from)
127            .collect::<HashSet<_>>();
128        let lines = (method == SimilarityMethod::Lines).then(|| {
129            content
130                .lines()
131                .filter(|line| !line.trim().is_empty())
132                .map(String::from)
133                .collect()
134        });
135        let ast_counts = if method == SimilarityMethod::Ast {
136            languages
137                .into_iter()
138                .map(|language| (language, ast_node_counts(content, language)))
139                .collect()
140        } else {
141            HashMap::new()
142        };
143        Self {
144            method,
145            lines,
146            tokens,
147            ast_counts,
148        }
149    }
150
151    pub(super) fn similarity(&self, other: &Self, language: Language) -> f64 {
152        debug_assert_eq!(self.method, other.method);
153        match self.method {
154            SimilarityMethod::Lines => {
155                let score = set_similarity(
156                    self.lines.as_ref().expect("lines prepared"),
157                    other.lines.as_ref().expect("lines prepared"),
158                );
159                if score == 0.0 {
160                    set_similarity(&self.tokens, &other.tokens)
161                } else {
162                    score
163                }
164            }
165            SimilarityMethod::Tokens => set_similarity(&self.tokens, &other.tokens),
166            SimilarityMethod::Ast => match (
167                self.ast_counts.get(&language).and_then(Option::as_ref),
168                other.ast_counts.get(&language).and_then(Option::as_ref),
169            ) {
170                (Some(left), Some(right)) => count_similarity(left, right),
171                _ => set_similarity(&self.tokens, &other.tokens),
172            },
173        }
174    }
175}
176
177fn set_similarity(left: &HashSet<String>, right: &HashSet<String>) -> f64 {
178    if left.is_empty() && right.is_empty() {
179        return 1.0;
180    }
181    if left.is_empty() || right.is_empty() {
182        return 0.0;
183    }
184    left.intersection(right).count() as f64 / left.union(right).count() as f64
185}
186
187fn count_similarity(left: &HashMap<String, usize>, right: &HashMap<String, usize>) -> f64 {
188    if left.is_empty() && right.is_empty() {
189        return 1.0;
190    }
191    if left.is_empty() || right.is_empty() {
192        return 0.0;
193    }
194    let keys = left.keys().chain(right.keys()).collect::<HashSet<_>>();
195    let (intersection, union) = keys.into_iter().fold((0usize, 0usize), |totals, key| {
196        let left_count = left.get(key).copied().unwrap_or(0);
197        let right_count = right.get(key).copied().unwrap_or(0);
198        (
199            totals.0 + left_count.min(right_count),
200            totals.1 + left_count.max(right_count),
201        )
202    });
203    if union == 0 {
204        0.0
205    } else {
206        intersection as f64 / union as f64
207    }
208}
209
210fn ast_node_counts(content: &str, language: Language) -> Option<HashMap<String, usize>> {
211    let parsed = ParsedFile::parse(content, language)?;
212    let mut counts = HashMap::new();
213    collect_node_kinds(parsed.root_node(), &mut counts);
214    Some(counts)
215}
216
217fn compute_ast_similarity(a: &str, b: &str, language: Language) -> Option<f64> {
218    let counts_a = ast_node_counts(a, language)?;
219    let counts_b = ast_node_counts(b, language)?;
220    Some(count_similarity(&counts_a, &counts_b))
221}
222
223fn collect_node_kinds(node: tree_sitter::Node<'_>, counts: &mut HashMap<String, usize>) {
224    let mut stack = vec![node];
225
226    while let Some(current) = stack.pop() {
227        let kind = current.kind();
228        let entry = counts.entry(kind.to_string()).or_insert(0);
229        *entry += 1;
230
231        let child_count = current.child_count();
232        for index in (0..child_count).rev() {
233            if let Some(child) = current.child(index as u32) {
234                stack.push(child);
235            }
236        }
237    }
238}
semantic/analysis/analysis_similarity.rs

semantic/analysis/
analysis_similarity.rs
