1use std::collections::{HashMap, HashSet};
2use std::error::Error;
3use std::fmt::{Display, Formatter};
4
5#[derive(Debug, Clone)]
6pub enum KnoloError {
7 InvalidPack(String),
8}
9
10impl Display for KnoloError {
11 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
12 match self {
13 KnoloError::InvalidPack(msg) => write!(f, "invalid pack: {msg}"),
14 }
15 }
16}
17
18impl Error for KnoloError {}
19
20#[derive(Debug, Clone)]
21pub struct PackMeta {
22 pub version: u32,
23 pub stats: PackStats,
24}
25
26#[derive(Debug, Clone)]
27pub struct PackStats {
28 pub docs: usize,
29 pub blocks: usize,
30 pub terms: usize,
31 pub avg_block_len: Option<f64>,
32}
33
34#[derive(Debug, Clone)]
35pub struct Pack {
36 pub meta: PackMeta,
37 pub lexicon: HashMap<String, u32>,
38 pub postings: Vec<u32>,
39 pub blocks: Vec<String>,
40 pub headings: Vec<Option<String>>,
41 pub doc_ids: Vec<Option<String>>,
42 pub namespaces: Vec<Option<String>>,
43 pub block_token_lens: Vec<usize>,
44}
45
46#[derive(Debug, Clone)]
47pub struct QueryOptions {
48 pub top_k: usize,
49 pub min_score: f64,
50 pub namespace: Option<Vec<String>>,
51 pub source: Option<Vec<String>>,
52}
53
54impl Default for QueryOptions {
55 fn default() -> Self {
56 Self {
57 top_k: 10,
58 min_score: 0.0,
59 namespace: None,
60 source: None,
61 }
62 }
63}
64
65#[derive(Debug, Clone, PartialEq)]
66pub struct Hit {
67 pub block_id: usize,
68 pub score: f64,
69 pub text: String,
70 pub source: Option<String>,
71 pub namespace: Option<String>,
72}
73
74pub fn mount_pack_from_bytes(bytes: &[u8]) -> Result<Pack, KnoloError> {
75 let mut cursor = 0usize;
76
77 let meta_len = read_u32(bytes, &mut cursor)? as usize;
78 let meta_json = read_slice(bytes, &mut cursor, meta_len)?;
79 let meta = parse_meta(std::str::from_utf8(meta_json).map_err(|_| KnoloError::InvalidPack("meta utf8".into()))?)?;
80
81 let lex_len = read_u32(bytes, &mut cursor)? as usize;
82 let lex_json = read_slice(bytes, &mut cursor, lex_len)?;
83 let lexicon = parse_lexicon(std::str::from_utf8(lex_json).map_err(|_| KnoloError::InvalidPack("lexicon utf8".into()))?)?;
84
85 let post_count = read_u32(bytes, &mut cursor)? as usize;
86 let postings = read_u32_array(bytes, &mut cursor, post_count)?;
87
88 let blocks_len = read_u32(bytes, &mut cursor)? as usize;
89 let blocks_json = read_slice(bytes, &mut cursor, blocks_len)?;
90 let blocks_str = std::str::from_utf8(blocks_json).map_err(|_| KnoloError::InvalidPack("blocks utf8".into()))?;
91 let parsed_blocks = parse_blocks(blocks_str)?;
92
93 Ok(Pack {
94 meta,
95 lexicon,
96 postings,
97 blocks: parsed_blocks.texts,
98 headings: parsed_blocks.headings,
99 doc_ids: parsed_blocks.doc_ids,
100 namespaces: parsed_blocks.namespaces,
101 block_token_lens: parsed_blocks.lens,
102 })
103}
104
105pub fn query(pack: &Pack, q: &str, opts: QueryOptions) -> Vec<Hit> {
106 if q.trim().is_empty() {
107 return vec![];
108 }
109 let tokens = tokenize(q);
110 if tokens.is_empty() {
111 return vec![];
112 }
113
114 let term_ids = tokens
115 .iter()
116 .filter_map(|t| pack.lexicon.get(t).copied())
117 .collect::<HashSet<_>>();
118 if term_ids.is_empty() {
119 return vec![];
120 }
121
122 let namespace_filter = normalize_filter(opts.namespace.as_ref());
123 let source_filter = normalize_filter(opts.source.as_ref());
124
125 let mut candidates: HashMap<usize, HashMap<u32, f64>> = HashMap::new();
126 let mut dfs: HashMap<u32, usize> = HashMap::new();
127 let uses_offset_block_ids = pack.meta.version >= 3;
128
129 let mut i = 0usize;
130 while i < pack.postings.len() {
131 let tid = pack.postings[i];
132 i += 1;
133 if tid == 0 {
134 continue;
135 }
136 let relevant = term_ids.contains(&tid);
137 let mut term_df = 0usize;
138
139 if i >= pack.postings.len() { break; }
140 let mut encoded_bid = pack.postings[i];
141 i += 1;
142
143 while encoded_bid != 0 && i < pack.postings.len() {
144 let bid = if uses_offset_block_ids {
145 encoded_bid.saturating_sub(1) as usize
146 } else {
147 encoded_bid as usize
148 };
149
150 let mut tf = 0usize;
151 while i < pack.postings.len() {
152 let pos = pack.postings[i];
153 i += 1;
154 if pos == 0 {
155 break;
156 }
157 tf += 1;
158 }
159
160 term_df += 1;
161 if relevant && bid < pack.blocks.len() {
162 let entry = candidates.entry(bid).or_default();
163 *entry.entry(tid).or_insert(0.0) += tf as f64;
164 }
165
166 if i >= pack.postings.len() { break; }
167 encoded_bid = pack.postings[i];
168 i += 1;
169 }
170
171 if relevant {
172 dfs.insert(tid, term_df);
173 }
174 }
175
176 if !namespace_filter.is_empty() {
177 candidates.retain(|bid, _| {
178 pack.namespaces
179 .get(*bid)
180 .and_then(|n| n.clone())
181 .map(|n| namespace_filter.contains(&normalize(&n)))
182 .unwrap_or(false)
183 });
184 }
185
186 if !source_filter.is_empty() {
187 candidates.retain(|bid, _| {
188 pack.doc_ids
189 .get(*bid)
190 .and_then(|n| n.clone())
191 .map(|n| source_filter.contains(&normalize(&n)))
192 .unwrap_or(false)
193 });
194 }
195
196 let doc_count = pack.meta.stats.blocks.max(1) as f64;
197 let avg_len = pack
198 .meta
199 .stats
200 .avg_block_len
201 .unwrap_or_else(|| {
202 if pack.block_token_lens.is_empty() {
203 1.0
204 } else {
205 pack.block_token_lens.iter().sum::<usize>() as f64 / pack.block_token_lens.len() as f64
206 }
207 })
208 .max(1.0);
209
210 let mut scored = candidates
211 .into_iter()
212 .map(|(bid, tf_map)| {
213 let mut score = 0.0;
214 let len = *pack.block_token_lens.get(bid).unwrap_or(&1) as f64;
215 for (tid, tf) in tf_map {
216 let df = *dfs.get(&tid).unwrap_or(&0) as f64;
217 let idf = (1.0 + (doc_count - df + 0.5) / (df + 0.5)).ln();
218 let k1 = 1.5;
219 let b = 0.75;
220 let numer = tf * (k1 + 1.0);
221 let denom = tf + k1 * (1.0 - b + b * (len / avg_len));
222 score += idf * (numer / denom);
223 }
224 (bid, score)
225 })
226 .filter(|(_, score)| *score >= opts.min_score)
227 .collect::<Vec<_>>();
228
229 scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
230
231 scored
232 .into_iter()
233 .take(opts.top_k.max(1))
234 .map(|(bid, score)| Hit {
235 block_id: bid,
236 score,
237 text: pack.blocks.get(bid).cloned().unwrap_or_default(),
238 source: pack.doc_ids.get(bid).and_then(|s| s.clone()),
239 namespace: pack.namespaces.get(bid).and_then(|s| s.clone()),
240 })
241 .collect()
242}
243
244struct ParsedBlocks {
245 texts: Vec<String>,
246 headings: Vec<Option<String>>,
247 doc_ids: Vec<Option<String>>,
248 namespaces: Vec<Option<String>>,
249 lens: Vec<usize>,
250}
251
252fn parse_meta(json: &str) -> Result<PackMeta, KnoloError> {
253 Ok(PackMeta {
254 version: parse_u32_field(json, "version")?,
255 stats: PackStats {
256 docs: parse_u32_field(json, "docs")? as usize,
257 blocks: parse_u32_field(json, "blocks")? as usize,
258 terms: parse_u32_field(json, "terms")? as usize,
259 avg_block_len: parse_f64_field(json, "avgBlockLen"),
260 },
261 })
262}
263
264fn parse_lexicon(json: &str) -> Result<HashMap<String, u32>, KnoloError> {
265 let mut map = HashMap::new();
266 let s = compact(json);
267 let mut i = 0usize;
268 while let Some(start) = s[i..].find("[\"") {
269 let abs = i + start + 2;
270 let rest = &s[abs..];
271 let end = rest.find('"').ok_or_else(|| KnoloError::InvalidPack("lexicon key".into()))?;
272 let key = rest[..end].to_string();
273 let rest2 = &rest[end + 1..];
274 let comma = rest2.find(',').ok_or_else(|| KnoloError::InvalidPack("lexicon comma".into()))?;
275 let rest3 = &rest2[comma + 1..];
276 let mut n = String::new();
277 for ch in rest3.chars() {
278 if ch.is_ascii_digit() {
279 n.push(ch);
280 } else {
281 break;
282 }
283 }
284 if !n.is_empty() {
285 map.insert(key, n.parse::<u32>().map_err(|_| KnoloError::InvalidPack("lexicon tid".into()))?);
286 }
287 i = abs + end + 1;
288 }
289 Ok(map)
290}
291
292fn parse_blocks(json: &str) -> Result<ParsedBlocks, KnoloError> {
293 let s = compact(json);
294 if s.starts_with("[\"") {
295 let mut texts = Vec::new();
296 let mut i = 2usize;
297 while i < s.len() {
298 if let Some(end) = s[i..].find('"') {
299 let piece = &s[i..i + end];
300 texts.push(unescape(piece));
301 i += end + 1;
302 if let Some(next) = s[i..].find('"') {
303 i += next + 1;
304 } else {
305 break;
306 }
307 } else {
308 break;
309 }
310 }
311 let lens = texts.iter().map(|t| tokenize(t).len()).collect::<Vec<_>>();
312 return Ok(ParsedBlocks {
313 headings: vec![None; texts.len()],
314 doc_ids: vec![None; texts.len()],
315 namespaces: vec![None; texts.len()],
316 lens,
317 texts,
318 });
319 }
320
321 let objects = split_top_level_objects(&s)?;
322 let mut texts = Vec::new();
323 let mut headings = Vec::new();
324 let mut doc_ids = Vec::new();
325 let mut namespaces = Vec::new();
326 let mut lens = Vec::new();
327
328 for obj in objects {
329 let text = parse_string_or_null(&obj, "text").unwrap_or_default();
330 let len = parse_u32_field_optional(&obj, "len").map(|v| v as usize).unwrap_or_else(|| tokenize(&text).len());
331 texts.push(text);
332 headings.push(parse_string_or_null(&obj, "heading"));
333 doc_ids.push(parse_string_or_null(&obj, "docId"));
334 namespaces.push(parse_string_or_null(&obj, "namespace"));
335 lens.push(len);
336 }
337
338 Ok(ParsedBlocks { texts, headings, doc_ids, namespaces, lens })
339}
340
341fn split_top_level_objects(s: &str) -> Result<Vec<String>, KnoloError> {
342 let mut out = Vec::new();
343 let mut depth = 0i32;
344 let mut start = None;
345 let chars: Vec<char> = s.chars().collect();
346 for (i, ch) in chars.iter().enumerate() {
347 if *ch == '{' {
348 if depth == 0 {
349 start = Some(i);
350 }
351 depth += 1;
352 } else if *ch == '}' {
353 depth -= 1;
354 if depth == 0 {
355 if let Some(st) = start {
356 out.push(chars[st..=i].iter().collect());
357 }
358 start = None;
359 }
360 }
361 }
362 if out.is_empty() {
363 return Err(KnoloError::InvalidPack("blocks objects".into()));
364 }
365 Ok(out)
366}
367
368fn parse_string_or_null(obj: &str, key: &str) -> Option<String> {
369 let needle = format!("\"{}\":", key);
370 let idx = obj.find(&needle)? + needle.len();
371 let tail = &obj[idx..];
372 if tail.starts_with("null") {
373 return None;
374 }
375 if !tail.starts_with('"') {
376 return None;
377 }
378 let rest = &tail[1..];
379 let end = rest.find('"')?;
380 Some(unescape(&rest[..end]))
381}
382
383fn parse_u32_field(json: &str, key: &str) -> Result<u32, KnoloError> {
384 parse_u32_field_optional(json, key).ok_or_else(|| KnoloError::InvalidPack(format!("missing {key}")))
385}
386
387fn parse_u32_field_optional(json: &str, key: &str) -> Option<u32> {
388 let needle = format!("\"{}\":", key);
389 let idx = json.find(&needle)? + needle.len();
390 let tail = &json[idx..];
391 let mut n = String::new();
392 for ch in tail.chars() {
393 if ch.is_ascii_digit() {
394 n.push(ch);
395 } else if !n.is_empty() {
396 break;
397 }
398 }
399 n.parse().ok()
400}
401
402fn parse_f64_field(json: &str, key: &str) -> Option<f64> {
403 let needle = format!("\"{}\":", key);
404 let idx = json.find(&needle)? + needle.len();
405 let tail = &json[idx..];
406 let mut n = String::new();
407 for ch in tail.chars() {
408 if ch.is_ascii_digit() || ch == '.' {
409 n.push(ch);
410 } else if !n.is_empty() {
411 break;
412 }
413 }
414 n.parse().ok()
415}
416
417fn normalize_filter(values: Option<&Vec<String>>) -> HashSet<String> {
418 values
419 .map(|arr| arr.iter().map(|s| normalize(s)).collect::<HashSet<_>>())
420 .unwrap_or_default()
421}
422
423fn normalize(s: &str) -> String {
424 s.to_lowercase().trim().to_string()
425}
426
427fn tokenize(text: &str) -> Vec<String> {
428 let mut out = Vec::new();
429 let mut cur = String::new();
430 for ch in text.chars() {
431 if ch.is_alphanumeric() {
432 cur.push(ch.to_ascii_lowercase());
433 } else if !cur.is_empty() {
434 out.push(std::mem::take(&mut cur));
435 }
436 }
437 if !cur.is_empty() {
438 out.push(cur);
439 }
440 out
441}
442
443fn compact(s: &str) -> String {
444 let mut out = String::with_capacity(s.len());
445 let mut in_string = false;
446 let mut escaped = false;
447
448 for ch in s.chars() {
449 if in_string {
450 out.push(ch);
451 if escaped {
452 escaped = false;
453 } else if ch == '\\' {
454 escaped = true;
455 } else if ch == '"' {
456 in_string = false;
457 }
458 continue;
459 }
460
461 if ch.is_whitespace() {
462 continue;
463 }
464
465 out.push(ch);
466 if ch == '"' {
467 in_string = true;
468 }
469 }
470
471 out
472}
473
474fn unescape(s: &str) -> String {
475 s.replace("\\\"", "\"")
476}
477
478fn read_u32(bytes: &[u8], cursor: &mut usize) -> Result<u32, KnoloError> {
479 let chunk = read_slice(bytes, cursor, 4)?;
480 Ok(u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
481}
482
483fn read_u32_array(bytes: &[u8], cursor: &mut usize, len: usize) -> Result<Vec<u32>, KnoloError> {
484 let mut out = Vec::with_capacity(len);
485 for _ in 0..len {
486 out.push(read_u32(bytes, cursor)?);
487 }
488 Ok(out)
489}
490
491fn read_slice<'a>(bytes: &'a [u8], cursor: &mut usize, len: usize) -> Result<&'a [u8], KnoloError> {
492 let end = cursor.saturating_add(len);
493 if end > bytes.len() {
494 return Err(KnoloError::InvalidPack("unexpected end-of-buffer".into()));
495 }
496 let slice = &bytes[*cursor..end];
497 *cursor = end;
498 Ok(slice)
499}