Skip to main content

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
87pub(super) struct PreparedSimilarity {
88    method: SimilarityMethod,
89    lines: Option<HashSet<String>>,
90    tokens: HashSet<String>,
91    ast_counts: HashMap<Language, Option<HashMap<String, usize>>>,
92}
93
94impl PreparedSimilarity {
95    pub(super) fn new(
96        content: &str,
97        method: SimilarityMethod,
98        languages: impl IntoIterator<Item = Language>,
99    ) -> Self {
100        let tokens = content
101            .split_whitespace()
102            .map(String::from)
103            .collect::<HashSet<_>>();
104        let lines = (method == SimilarityMethod::Lines).then(|| {
105            content
106                .lines()
107                .filter(|line| !line.trim().is_empty())
108                .map(String::from)
109                .collect()
110        });
111        let ast_counts = if method == SimilarityMethod::Ast {
112            languages
113                .into_iter()
114                .map(|language| (language, ast_node_counts(content, language)))
115                .collect()
116        } else {
117            HashMap::new()
118        };
119        Self {
120            method,
121            lines,
122            tokens,
123            ast_counts,
124        }
125    }
126
127    pub(super) fn similarity(&self, other: &Self, language: Language) -> f64 {
128        debug_assert_eq!(self.method, other.method);
129        match self.method {
130            SimilarityMethod::Lines => {
131                let score = set_similarity(
132                    self.lines.as_ref().expect("lines prepared"),
133                    other.lines.as_ref().expect("lines prepared"),
134                );
135                if score == 0.0 {
136                    set_similarity(&self.tokens, &other.tokens)
137                } else {
138                    score
139                }
140            }
141            SimilarityMethod::Tokens => set_similarity(&self.tokens, &other.tokens),
142            SimilarityMethod::Ast => match (
143                self.ast_counts.get(&language).and_then(Option::as_ref),
144                other.ast_counts.get(&language).and_then(Option::as_ref),
145            ) {
146                (Some(left), Some(right)) => count_similarity(left, right),
147                _ => set_similarity(&self.tokens, &other.tokens),
148            },
149        }
150    }
151}
152
153fn set_similarity(left: &HashSet<String>, right: &HashSet<String>) -> f64 {
154    if left.is_empty() && right.is_empty() {
155        return 1.0;
156    }
157    if left.is_empty() || right.is_empty() {
158        return 0.0;
159    }
160    left.intersection(right).count() as f64 / left.union(right).count() as f64
161}
162
163fn count_similarity(left: &HashMap<String, usize>, right: &HashMap<String, usize>) -> f64 {
164    if left.is_empty() && right.is_empty() {
165        return 1.0;
166    }
167    if left.is_empty() || right.is_empty() {
168        return 0.0;
169    }
170    let keys = left.keys().chain(right.keys()).collect::<HashSet<_>>();
171    let (intersection, union) = keys.into_iter().fold((0usize, 0usize), |totals, key| {
172        let left_count = left.get(key).copied().unwrap_or(0);
173        let right_count = right.get(key).copied().unwrap_or(0);
174        (
175            totals.0 + left_count.min(right_count),
176            totals.1 + left_count.max(right_count),
177        )
178    });
179    if union == 0 {
180        0.0
181    } else {
182        intersection as f64 / union as f64
183    }
184}
185
186fn ast_node_counts(content: &str, language: Language) -> Option<HashMap<String, usize>> {
187    let parsed = ParsedFile::parse(content, language)?;
188    let mut counts = HashMap::new();
189    collect_node_kinds(parsed.root_node(), &mut counts);
190    Some(counts)
191}
192
193fn compute_ast_similarity(a: &str, b: &str, language: Language) -> Option<f64> {
194    let counts_a = ast_node_counts(a, language)?;
195    let counts_b = ast_node_counts(b, language)?;
196    Some(count_similarity(&counts_a, &counts_b))
197}
198
199fn collect_node_kinds(node: tree_sitter::Node<'_>, counts: &mut HashMap<String, usize>) {
200    let mut stack = vec![node];
201
202    while let Some(current) = stack.pop() {
203        let kind = current.kind();
204        let entry = counts.entry(kind.to_string()).or_insert(0);
205        *entry += 1;
206
207        let child_count = current.child_count();
208        for index in (0..child_count).rev() {
209            if let Some(child) = current.child(index as u32) {
210                stack.push(child);
211            }
212        }
213    }
214}