provenant/license_detection/index/
dictionary.rs

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

provenant/license_detection/index/
dictionary.rs
