context_forge/analysis/
frequency.rs

1use std::collections::HashMap;
2
3use crate::analysis::ngrams::{bigrams, trigrams};
4
5/// Compute term counts for a list of tokens.
6///
7/// Returns a map from term to count of occurrences.
8#[must_use]
9pub fn term_counts(tokens: &[String]) -> HashMap<String, usize> {
10    let mut counts: HashMap<String, usize> = HashMap::new();
11
12    for token in tokens {
13        *counts.entry(token.clone()).or_insert(0) += 1;
14    }
15
16    counts
17}
18
19/// Compute term counts including n-grams.
20///
21/// Computes occurrence counts for unigrams, bigrams, and trigrams combined.
22/// N-grams are space-separated strings (e.g., "system openssl").
23#[must_use]
24pub fn term_counts_with_ngrams(tokens: &[String]) -> HashMap<String, usize> {
25    let mut counts: HashMap<String, usize> = HashMap::new();
26
27    for token in tokens {
28        *counts.entry(token.clone()).or_insert(0) += 1;
29    }
30
31    for bigram in bigrams(tokens) {
32        *counts.entry(bigram).or_insert(0) += 1;
33    }
34
35    for trigram in trigrams(tokens) {
36        *counts.entry(trigram).or_insert(0) += 1;
37    }
38
39    counts
40}
41
42#[cfg(test)]
43mod tests {
44    use std::collections::HashMap;
45
46    use super::{term_counts, term_counts_with_ngrams};
47
48    #[test]
49    fn test_term_counts() {
50        let tokens = vec![
51            "hello".to_string(),
52            "world".to_string(),
53            "hello".to_string(),
54        ];
55        let counts = term_counts(&tokens);
56
57        assert_eq!(counts.get("hello"), Some(&2));
58        assert_eq!(counts.get("world"), Some(&1));
59        assert_eq!(counts.len(), 2);
60    }
61
62    #[test]
63    fn test_term_counts_empty() {
64        let tokens: Vec<String> = Vec::new();
65        let counts = term_counts(&tokens);
66        assert!(counts.is_empty());
67    }
68
69    #[test]
70    fn test_term_counts_with_ngrams() {
71        let tokens = vec![
72            "system".to_string(),
73            "openssl".to_string(),
74            "upgrade".to_string(),
75        ];
76        let counts = term_counts_with_ngrams(&tokens);
77
78        assert_eq!(counts.get("system"), Some(&1));
79        assert_eq!(counts.get("openssl"), Some(&1));
80        assert_eq!(counts.get("upgrade"), Some(&1));
81        assert_eq!(counts.get("system openssl"), Some(&1));
82        assert_eq!(counts.get("openssl upgrade"), Some(&1));
83        assert_eq!(counts.get("system openssl upgrade"), Some(&1));
84        assert_eq!(counts.len(), 6);
85    }
86
87    #[test]
88    fn test_single_token_with_ngrams() {
89        let tokens = vec!["solo".to_string()];
90        let counts = term_counts_with_ngrams(&tokens);
91
92        let mut expected = HashMap::new();
93        expected.insert("solo".to_string(), 1_usize);
94        assert_eq!(counts, expected);
95    }
96}
context_forge/analysis/frequency.rs

context_forge/analysis/
frequency.rs
