Skip to main content

reddb_engine/storage/
tokenization.rs

1//! Shared tokenization utilities for indexing and search.
2//!
3//! Used by the `ContextIndex` and search paths to ensure consistent
4//! token generation across all search operations.
5
6use std::collections::BTreeSet;
7
8use reddb_types::Value;
9
10pub const MAX_INDEX_TOKENS: usize = 256;
11pub const MAX_TOKEN_LEN: usize = 160;
12pub const MAX_JSON_TOKEN_BUDGET: usize = 128;
13
14pub fn push_value_tokens(tokens: &mut BTreeSet<String>, value: &Value) {
15    match value {
16        Value::Null => {}
17        Value::Integer(v) => {
18            push_text_tokens(tokens, &v.to_string(), false);
19            if *v >= 0 {
20                push_text_tokens(tokens, &format!("e{v}"), false);
21            }
22        }
23        Value::UnsignedInteger(v) => {
24            push_text_tokens(tokens, &v.to_string(), false);
25            push_text_tokens(tokens, &format!("e{v}"), false);
26        }
27        Value::Float(v) => {
28            if v.is_finite() {
29                push_text_tokens(tokens, &v.to_string(), false);
30            }
31        }
32        Value::Text(v) => push_text_tokens(tokens, v, true),
33        Value::Boolean(v) => push_text_tokens(tokens, if *v { "true" } else { "false" }, false),
34        Value::Timestamp(v)
35        | Value::Duration(v)
36        | Value::TimestampMs(v)
37        | Value::BigInt(v)
38        | Value::Decimal(v) => push_text_tokens(tokens, &v.to_string(), false),
39        Value::Phone(v) => push_text_tokens(tokens, &v.to_string(), false),
40        Value::Port(v) => push_text_tokens(tokens, &v.to_string(), false),
41        Value::NodeRef(v) | Value::EdgeRef(v) | Value::Email(v) | Value::Url(v) => {
42            push_text_tokens(tokens, v, true)
43        }
44        Value::RowRef(collection, id)
45        | Value::VectorRef(collection, id)
46        | Value::DocRef(collection, id) => {
47            push_text_tokens(tokens, collection, false);
48            push_text_tokens(tokens, &id.to_string(), false);
49            push_text_tokens(tokens, &format!("e{id}"), false);
50            push_text_tokens(tokens, &format!("{collection}:{id}"), false);
51        }
52        Value::KeyRef(collection, key) => {
53            push_text_tokens(tokens, collection, false);
54            push_text_tokens(tokens, key, true);
55            push_text_tokens(tokens, &format!("{collection}:{key}"), true);
56        }
57        Value::TableRef(table) => push_text_tokens(tokens, table, false),
58        Value::Json(bytes) => push_json_tokens(tokens, bytes),
59        Value::Array(values) => {
60            for item in values {
61                push_value_tokens(tokens, item);
62                if tokens.len() >= MAX_INDEX_TOKENS {
63                    break;
64                }
65            }
66        }
67        _ => {
68            push_text_tokens(tokens, &value.to_string(), false);
69        }
70    }
71}
72
73pub fn push_json_tokens(tokens: &mut BTreeSet<String>, bytes: &[u8]) {
74    fn collect(
75        value: &crate::serde_json::Value,
76        tokens: &mut BTreeSet<String>,
77        budget: &mut usize,
78    ) {
79        if *budget == 0 || tokens.len() >= MAX_INDEX_TOKENS {
80            return;
81        }
82        match value {
83            crate::serde_json::Value::Null => {}
84            crate::serde_json::Value::Bool(v) => {
85                push_text_tokens(tokens, if *v { "true" } else { "false" }, false);
86                *budget = budget.saturating_sub(1);
87            }
88            crate::serde_json::Value::Integer(v) => {
89                push_text_tokens(tokens, &v.to_string(), false);
90                *budget = budget.saturating_sub(1);
91            }
92            crate::serde_json::Value::Number(v) => {
93                push_text_tokens(tokens, &v.to_string(), false);
94                *budget = budget.saturating_sub(1);
95            }
96            crate::serde_json::Value::Decimal(v) => {
97                push_text_tokens(tokens, v, false);
98                *budget = budget.saturating_sub(1);
99            }
100            crate::serde_json::Value::String(v) => {
101                push_text_tokens(tokens, v, true);
102                *budget = budget.saturating_sub(1);
103            }
104            crate::serde_json::Value::Array(values) => {
105                for item in values {
106                    collect(item, tokens, budget);
107                    if *budget == 0 || tokens.len() >= MAX_INDEX_TOKENS {
108                        break;
109                    }
110                }
111            }
112            crate::serde_json::Value::Object(fields) => {
113                for (key, item) in fields {
114                    push_text_tokens(tokens, key, false);
115                    collect(item, tokens, budget);
116                    if *budget == 0 || tokens.len() >= MAX_INDEX_TOKENS {
117                        break;
118                    }
119                }
120            }
121        }
122    }
123
124    // A document body may be the native binary container (PRD-1398); decode it
125    // to JSON so its index tokens match those of an equivalent plain JSON body.
126    if let Some(value) = reddb_types::document_body_codec::decode_container_to_json(bytes) {
127        let mut budget = MAX_JSON_TOKEN_BUDGET;
128        collect(&value, tokens, &mut budget);
129        return;
130    }
131    if let Ok(value) = crate::serde_json::from_slice::<crate::serde_json::Value>(bytes) {
132        let mut budget = MAX_JSON_TOKEN_BUDGET;
133        collect(&value, tokens, &mut budget);
134    }
135}
136
137pub fn push_text_tokens(tokens: &mut BTreeSet<String>, text: &str, split_words: bool) {
138    let trimmed = text.trim();
139    if trimmed.is_empty() {
140        return;
141    }
142
143    push_token_variant(tokens, trimmed);
144
145    if let Some((_, rhs)) = trimmed.split_once(':') {
146        let rhs = rhs.trim();
147        if !rhs.is_empty() {
148            push_token_variant(tokens, rhs);
149        }
150    }
151
152    if split_words {
153        for word in trimmed
154            .split(|ch: char| ch.is_ascii_whitespace() || [',', ';', '|'].contains(&ch))
155            .map(str::trim)
156            .filter(|word| !word.is_empty())
157        {
158            push_token_variant(tokens, word);
159        }
160    }
161}
162
163fn push_token_variant(tokens: &mut BTreeSet<String>, token: &str) {
164    let normalized = normalize_token(token);
165    if !normalized.is_empty() {
166        tokens.insert(normalized);
167    }
168
169    let canonical = canonical_token(token);
170    if !canonical.is_empty() {
171        tokens.insert(canonical);
172    }
173}
174
175fn normalize_token(token: &str) -> String {
176    token
177        .trim()
178        .to_ascii_lowercase()
179        .replace(['-', ' '], "_")
180        .chars()
181        .take(MAX_TOKEN_LEN)
182        .collect()
183}
184
185fn canonical_token(token: &str) -> String {
186    token
187        .chars()
188        .filter(|ch| ch.is_ascii_alphanumeric())
189        .flat_map(|ch| ch.to_lowercase())
190        .take(MAX_TOKEN_LEN)
191        .collect()
192}