provenant/license_detection/index/
dictionary.rs

1// SPDX-FileCopyrightText: nexB Inc. and others
2// ScanCode is a trademark of nexB Inc.
3// SPDX-FileCopyrightText: Provenant contributors
4// SPDX-License-Identifier: Apache-2.0
5// Derived from ScanCode Toolkit (Apache-2.0); modified. See NOTICE.
6
7//! Token string to integer ID mapping.
8//!
9//! TokenDictionary maps token strings to unique integer IDs. This enables
10//! efficient token-based matching and indexing.
11
12use std::collections::{BTreeMap, HashMap};
13
14use rkyv::Archive;
15use rkyv::Archived;
16use serde::{Deserialize, Serialize};
17
18#[derive(
19    Debug,
20    Clone,
21    Copy,
22    PartialEq,
23    Eq,
24    Hash,
25    PartialOrd,
26    Ord,
27    Serialize,
28    Deserialize,
29    Archive,
30    rkyv::Serialize,
31    rkyv::Deserialize,
32)]
33#[rkyv(derive(Hash, Eq, PartialEq, PartialOrd, Ord))]
34pub struct TokenId(u16);
35
36impl TokenId {
37    pub const fn new(raw: u16) -> Self {
38        Self(raw)
39    }
40
41    pub const fn raw(self) -> u16 {
42        self.0
43    }
44
45    pub const fn as_usize(self) -> usize {
46        self.0 as usize
47    }
48
49    pub const fn to_le_bytes(self) -> [u8; 2] {
50        self.0.to_le_bytes()
51    }
52}
53
54#[cfg(test)]
55pub const fn tid(raw: u16) -> TokenId {
56    TokenId::new(raw)
57}
58
59impl From<u16> for TokenId {
60    fn from(value: u16) -> Self {
61        Self(value)
62    }
63}
64
65impl From<TokenId> for u16 {
66    fn from(value: TokenId) -> Self {
67        value.0
68    }
69}
70
71impl PartialEq<u16> for TokenId {
72    fn eq(&self, other: &u16) -> bool {
73        self.0 == *other
74    }
75}
76
77impl PartialOrd<u16> for TokenId {
78    fn partial_cmp(&self, other: &u16) -> Option<std::cmp::Ordering> {
79        self.0.partial_cmp(other)
80    }
81}
82
83impl PartialEq<TokenId> for u16 {
84    fn eq(&self, other: &TokenId) -> bool {
85        *self == other.0
86    }
87}
88
89impl PartialOrd<TokenId> for u16 {
90    fn partial_cmp(&self, other: &TokenId) -> Option<std::cmp::Ordering> {
91        self.partial_cmp(&other.0)
92    }
93}
94
95#[derive(
96    Debug,
97    Clone,
98    Copy,
99    PartialEq,
100    Eq,
101    Hash,
102    Serialize,
103    Deserialize,
104    Archive,
105    rkyv::Serialize,
106    rkyv::Deserialize,
107)]
108pub enum TokenKind {
109    Legalese,
110    Regular,
111}
112
113#[derive(
114    Debug,
115    Clone,
116    Copy,
117    PartialEq,
118    Eq,
119    Hash,
120    Serialize,
121    Deserialize,
122    Archive,
123    rkyv::Serialize,
124    rkyv::Deserialize,
125)]
126pub struct KnownToken {
127    pub id: TokenId,
128    pub kind: TokenKind,
129    pub is_digit_only: bool,
130    pub is_short_or_digit: bool,
131}
132
133#[derive(
134    Debug,
135    Clone,
136    Copy,
137    PartialEq,
138    Eq,
139    Hash,
140    Serialize,
141    Deserialize,
142    Archive,
143    rkyv::Serialize,
144    rkyv::Deserialize,
145)]
146pub enum QueryToken {
147    Known(KnownToken),
148    Unknown,
149    Stopword,
150}
151
152#[derive(
153    Debug, Clone, Copy, Serialize, Deserialize, Archive, rkyv::Serialize, rkyv::Deserialize,
154)]
155pub struct TokenMetadata {
156    pub kind: TokenKind,
157    pub is_digit_only: bool,
158    pub is_short_or_digit: bool,
159}
160
161/// Token dictionary mapping token strings to unique integer IDs.
162///
163/// Token IDs are assigned as follows:
164/// - IDs 0 to len_legalese-1: Reserved for legalese tokens (high-value words)
165/// - IDs len_legalese and above: Assigned to other tokens as encountered
166///
167/// The `len_legalese` delimiter allows the matching engine to distinguish
168/// between high-value (legalese) tokens and regular tokens.
169///
170/// Based on the Python ScanCode Toolkit implementation at:
171/// reference/scancode-toolkit/src/licensedcode/index.py
172#[derive(Debug, Clone, Archive, rkyv::Serialize, rkyv::Deserialize)]
173pub struct TokenDictionary {
174    /// Mapping from token string to token ID
175    tokens_to_ids: HashMap<String, TokenId>,
176
177    token_metadata: Vec<Option<TokenMetadata>>,
178
179    /// Number of legalese tokens (lower IDs = higher value)
180    len_legalese: usize,
181
182    /// Next token ID to assign (for non-legalese tokens)
183    next_id: TokenId,
184}
185
186impl TokenDictionary {
187    const DEFAULT_METADATA: TokenMetadata = TokenMetadata {
188        kind: TokenKind::Regular,
189        is_digit_only: false,
190        is_short_or_digit: false,
191    };
192
193    /// Create a new token dictionary initialized with legalese tokens.
194    ///
195    /// This follows the Python ScanCode Toolkit pattern where the dictionary
196    /// starts with pre-defined legalese words that get low IDs (high value).
197    ///
198    /// # Arguments
199    /// * `legalese` - Archived BTreeMap of word → u16 pairs for legalese words.
200    ///   Values are bare u16 (not TokenId) because the rkyv artifact is built
201    ///   by `build.rs` which cannot depend on this crate's types.
202    ///
203    /// # Returns
204    /// A new TokenDictionary instance with legalese tokens pre-populated
205    pub fn new_with_legalese(legalese: &Archived<BTreeMap<String, u16>>) -> Self {
206        let mut tokens_to_ids = HashMap::new();
207        let max_existing_id = legalese
208            .iter()
209            .map(|(_, id)| id.to_native() as usize)
210            .max()
211            .unwrap_or(0);
212        let mut token_metadata = vec![None; max_existing_id.saturating_add(1)];
213
214        for (word, id) in legalese.iter() {
215            let native_id = TokenId::new(id.to_native());
216            tokens_to_ids.insert(word.to_string(), native_id);
217            token_metadata[native_id.as_usize()] = Some(TokenMetadata {
218                kind: TokenKind::Legalese,
219                is_digit_only: word.chars().all(|c: char| c.is_ascii_digit()),
220                is_short_or_digit: word.len() == 1
221                    || word.chars().all(|c: char| c.is_ascii_digit()),
222            });
223        }
224
225        let len_legalese = legalese.len();
226        let next_id = TokenId::new((max_existing_id + 1).max(len_legalese) as u16);
227
228        Self {
229            tokens_to_ids,
230            token_metadata,
231            len_legalese,
232            next_id,
233        }
234    }
235
236    /// Create a new token dictionary initialized with legalese token pairs.
237    ///
238    /// Convenience constructor for tests that don't use the rkyv-archived legalese data.
239    pub fn new_with_legalese_pairs(legalese_entries: &[(&str, u16)]) -> Self {
240        let mut tokens_to_ids = HashMap::new();
241        let max_existing_id = legalese_entries
242            .iter()
243            .map(|(_, token_id)| *token_id as usize)
244            .max()
245            .unwrap_or(0);
246        let mut token_metadata = vec![None; max_existing_id.saturating_add(1)];
247
248        for (word, token_id) in legalese_entries {
249            let id = TokenId::from(*token_id);
250            tokens_to_ids.insert(word.to_string(), id);
251            token_metadata[id.as_usize()] = Some(TokenMetadata {
252                kind: TokenKind::Legalese,
253                is_digit_only: word.chars().all(|c| c.is_ascii_digit()),
254                is_short_or_digit: word.len() == 1 || word.chars().all(|c| c.is_ascii_digit()),
255            });
256        }
257
258        let len_legalese = legalese_entries.len();
259        let next_id = TokenId::new((max_existing_id + 1).max(len_legalese) as u16);
260
261        Self {
262            tokens_to_ids,
263            token_metadata,
264            len_legalese,
265            next_id,
266        }
267    }
268
269    /// Create a new empty token dictionary (for testing).
270    ///
271    /// # Arguments
272    /// * `legalese_count` - Number of reserved legalese token IDs
273    ///
274    /// # Returns
275    /// A new TokenDictionary instance
276    pub fn new(legalese_count: usize) -> Self {
277        Self {
278            tokens_to_ids: HashMap::new(),
279            token_metadata: Vec::new(),
280            len_legalese: legalese_count,
281            next_id: TokenId::new(legalese_count as u16),
282        }
283    }
284
285    fn metadata_for(&self, id: TokenId) -> TokenMetadata {
286        self.token_metadata
287            .get(id.as_usize())
288            .and_then(|meta| *meta)
289            .unwrap_or(Self::DEFAULT_METADATA)
290    }
291
292    fn build_known_token(&self, id: TokenId) -> KnownToken {
293        let metadata = self.metadata_for(id);
294        KnownToken {
295            id,
296            kind: metadata.kind,
297            is_digit_only: metadata.is_digit_only,
298            is_short_or_digit: metadata.is_short_or_digit,
299        }
300    }
301
302    fn insert_metadata(&mut self, id: TokenId, kind: TokenKind, token: &str) {
303        let raw = id.as_usize();
304        if self.token_metadata.len() <= raw {
305            self.token_metadata.resize(raw + 1, None);
306        }
307        self.token_metadata[raw] = Some(TokenMetadata {
308            kind,
309            is_digit_only: token.chars().all(|c| c.is_ascii_digit()),
310            is_short_or_digit: token.len() == 1 || token.chars().all(|c| c.is_ascii_digit()),
311        });
312    }
313
314    pub fn intern(&mut self, token: &str) -> KnownToken {
315        if let Some(&id) = self.tokens_to_ids.get(token) {
316            return self.build_known_token(id);
317        }
318
319        let id = self.next_id;
320        self.next_id = TokenId::new(self.next_id.raw() + 1);
321        self.tokens_to_ids.insert(token.to_string(), id);
322        self.insert_metadata(id, TokenKind::Regular, token);
323        self.build_known_token(id)
324    }
325
326    pub fn lookup(&self, token: &str) -> Option<KnownToken> {
327        self.tokens_to_ids
328            .get(token)
329            .copied()
330            .map(|id| self.build_known_token(id))
331    }
332
333    pub fn classify_query_token(&self, token: &str) -> QueryToken {
334        self.lookup(token)
335            .map_or(QueryToken::Unknown, QueryToken::Known)
336    }
337
338    pub fn token_kind(&self, token_id: TokenId) -> TokenKind {
339        self.metadata_for(token_id).kind
340    }
341
342    pub fn is_digit_only_token(&self, token_id: TokenId) -> bool {
343        self.metadata_for(token_id).is_digit_only
344    }
345
346    #[cfg(test)]
347    pub fn get_or_assign(&mut self, token: &str) -> TokenId {
348        self.intern(token).id
349    }
350
351    /// Get the token ID for a token string if it exists.
352    ///
353    /// # Arguments
354    /// * `token` - The token string
355    ///
356    /// # Returns
357    /// Some(token_id) if the token exists, None otherwise
358    pub fn get_token_id(&self, token: &str) -> Option<TokenId> {
359        self.lookup(token).map(|token| token.id)
360    }
361
362    /// Get the token ID (alias for backward compatibility).
363    #[inline]
364    pub fn get(&self, token: &str) -> Option<TokenId> {
365        self.get_token_id(token)
366    }
367
368    /// Get the number of legalese tokens.
369    pub const fn legalese_count(&self) -> usize {
370        self.len_legalese
371    }
372
373    /// Get an iterator over all token string and ID pairs.
374    #[cfg(test)]
375    pub fn tokens_to_ids(&self) -> impl Iterator<Item = (&String, &TokenId)> {
376        self.tokens_to_ids.iter()
377    }
378
379    /// Get the number of tokens in the dictionary.
380    // This method will be used by the embedded index roundtrip tests in upcoming phases.
381    #[allow(dead_code)]
382    pub fn tokens_to_ids_len(&self) -> usize {
383        self.tokens_to_ids.len()
384    }
385}
386
387impl Default for TokenDictionary {
388    fn default() -> Self {
389        Self::new(0)
390    }
391}
392
393#[cfg(test)]
394mod tests {
395    use super::*;
396
397    #[test]
398    fn test_token_dictionary_new() {
399        let dict = TokenDictionary::new(10);
400        assert_eq!(dict.legalese_count(), 10);
401        assert_eq!(dict.tokens_to_ids.len(), 0);
402        assert!(dict.tokens_to_ids.is_empty());
403    }
404
405    #[test]
406    fn test_new_with_legalese() {
407        let legalese = [
408            ("license".to_string(), 0u16),
409            ("copyright".to_string(), 1u16),
410            ("permission".to_string(), 2u16),
411        ];
412
413        let mut dict = TokenDictionary::new_with_legalese_pairs(
414            &legalese
415                .iter()
416                .map(|(s, i)| (s.as_str(), *i))
417                .collect::<Vec<_>>(),
418        );
419
420        assert_eq!(dict.legalese_count(), 3);
421        assert_eq!(dict.tokens_to_ids.len(), 3);
422        assert!(!dict.tokens_to_ids.is_empty());
423
424        // Check that legalese tokens are registered
425        assert_eq!(dict.get_token_id("license"), Some(tid(0)));
426        assert_eq!(dict.get_token_id("copyright"), Some(tid(1)));
427        assert_eq!(dict.get_token_id("permission"), Some(tid(2)));
428
429        // Check that new tokens get IDs starting after legalese
430        let test_id = dict.get_or_assign("test");
431        assert_eq!(test_id, 3);
432    }
433
434    #[test]
435    fn test_new_with_legalese_sorted() {
436        let legalese = [
437            ("copyright".to_string(), 5u16),
438            ("license".to_string(), 0u16),
439            ("permission".to_string(), 10u16),
440        ];
441
442        let mut dict = TokenDictionary::new_with_legalese_pairs(
443            &legalese
444                .iter()
445                .map(|(s, i)| (s.as_str(), *i))
446                .collect::<Vec<_>>(),
447        );
448
449        assert_eq!(dict.legalese_count(), 3);
450        assert_eq!(dict.tokens_to_ids.len(), 3);
451
452        // Check legalese IDs are correct regardless of input order
453        assert_eq!(dict.get_token_id("copyright"), Some(tid(5)));
454        assert_eq!(dict.get_token_id("license"), Some(tid(0)));
455        assert_eq!(dict.get_token_id("permission"), Some(tid(10)));
456
457        // Next ID should advance past the highest explicit legalese token ID.
458        let test_id = dict.get_or_assign("test");
459        assert_eq!(test_id, tid(11));
460    }
461
462    #[test]
463    fn test_get_or_assign_new_token() {
464        let mut dict = TokenDictionary::new(5);
465
466        let id1 = dict.get_or_assign("hello");
467        let id2 = dict.get_or_assign("world");
468
469        // Should assign IDs starting at legalese_count (5)
470        assert_eq!(id1, 5);
471        assert_eq!(id2, 6);
472        assert_eq!(dict.tokens_to_ids.len(), 2);
473    }
474
475    #[test]
476    fn test_get_or_assign_existing_token() {
477        let mut dict = TokenDictionary::new(5);
478
479        let id1 = dict.get_or_assign("hello");
480        let id2 = dict.get_or_assign("hello");
481
482        // Should return the same ID for the same token
483        assert_eq!(id1, id2);
484        assert_eq!(dict.tokens_to_ids.len(), 1);
485    }
486
487    #[test]
488    fn test_get_or_assign_with_preexisting_legalese() {
489        let legalese = [("license".to_string(), 0u16)];
490        let mut dict = TokenDictionary::new_with_legalese_pairs(
491            &legalese
492                .iter()
493                .map(|(s, i)| (s.as_str(), *i))
494                .collect::<Vec<_>>(),
495        );
496
497        // Legalese tokens should already exist
498        let id = dict.get_or_assign("license");
499        assert_eq!(id, 0);
500        assert_eq!(dict.tokens_to_ids.len(), 1);
501
502        // New tokens should get IDs after legalese
503        let new_id = dict.get_or_assign("new");
504        assert_eq!(new_id, 1);
505        assert_eq!(dict.tokens_to_ids.len(), 2);
506    }
507
508    #[test]
509    fn test_get_existing_token() {
510        let mut dict = TokenDictionary::new(5);
511
512        dict.get_or_assign("hello");
513        assert_eq!(dict.get_token_id("hello"), Some(tid(5)));
514    }
515
516    #[test]
517    fn test_get_nonexistent_token() {
518        let dict = TokenDictionary::new(5);
519        assert_eq!(dict.get_token_id("hello"), None);
520    }
521
522    #[test]
523    fn test_legalese_range() {
524        let dict = TokenDictionary::new(10);
525
526        // IDs 0-9 are legalese
527        assert!(0 < dict.legalese_count() as u16);
528        assert!(5 < dict.legalese_count() as u16);
529        assert!(9 < dict.legalese_count() as u16);
530
531        // ID 10+ are not legalese
532        assert!(10 >= dict.legalese_count() as u16);
533        assert!(100 >= dict.legalese_count() as u16);
534    }
535
536    #[test]
537    fn test_legalese_range_with_actual_legalese() {
538        let legalese = [
539            ("license".to_string(), 0u16),
540            ("copyright".to_string(), 1u16),
541        ];
542
543        let mut dict = TokenDictionary::new_with_legalese_pairs(
544            &legalese
545                .iter()
546                .map(|(s, i)| (s.as_str(), *i))
547                .collect::<Vec<_>>(),
548        );
549
550        // Legalese tokens should have IDs in the legalese range
551        assert!(dict.get_token_id("license").unwrap() < dict.legalese_count() as u16);
552        assert!(dict.get_token_id("copyright").unwrap() < dict.legalese_count() as u16);
553
554        // Regular tokens should not be legalese
555        let regular_id = dict.get_or_assign("regular");
556        assert!(regular_id >= dict.legalese_count() as u16);
557    }
558
559    #[test]
560    fn test_token_dictionary_default() {
561        let dict = TokenDictionary::default();
562        assert_eq!(dict.legalese_count(), 0);
563        assert!(dict.tokens_to_ids.is_empty());
564    }
565
566    #[test]
567    fn test_get_alias() {
568        let mut dict = TokenDictionary::new(5);
569        dict.get_or_assign("hello");
570
571        // get() should be an alias for get_token_id()
572        assert_eq!(dict.get("hello"), dict.get_token_id("hello"));
573    }
574
575    #[test]
576    fn test_with_actual_legalese_module() {
577        use crate::license_detection::rules::legalese;
578
579        let legalese = legalese::archived_legalese();
580        assert!(!legalese.is_empty(), "Should have legalese words");
581
582        let mut dict = TokenDictionary::new_with_legalese(legalese);
583
584        // Verify dictionary has the right structure
585        assert_eq!(dict.legalese_count(), legalese.len());
586        assert_eq!(dict.tokens_to_ids.len(), legalese.len());
587
588        // Verify some legalese words are correctly registered
589        let license_id = dict.get_token_id("license");
590        assert!(license_id.is_some(), "License should be in dictionary");
591        assert!(
592            license_id.unwrap() < dict.legalese_count() as u16,
593            "License should be a legalese token"
594        );
595
596        // Note: Standalone "copyright" is NOT in the Python reference dictionary
597        // Only compound words like "copyrighted", "copyrights" are present
598        let copyrighted_id = dict.get_token_id("copyrighted");
599        assert!(
600            copyrighted_id.is_some(),
601            "Copyrighted should be in dictionary"
602        );
603        assert!(
604            copyrighted_id.unwrap() < dict.legalese_count() as u16,
605            "Copyrighted should be a legalese token"
606        );
607
608        // New tokens should get IDs after legalese
609        let hello_id = dict.get_or_assign("hello");
610        assert!(hello_id >= dict.legalese_count() as u16);
611        assert!(hello_id >= dict.legalese_count() as u16);
612    }
613}
provenant/license_detection/index/dictionary.rs

provenant/license_detection/index/
dictionary.rs
