semantic/analysis/
analysis_similarity.rs1use std::collections::{HashMap, HashSet};
5
6use crate::parser::{Language, ParsedFile};
7
8#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10pub enum SimilarityMethod {
11 Lines,
13 Tokens,
15 Ast,
17}
18
19pub 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 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}