1#![expect(clippy::let_underscore_must_use, reason = "removing what is already meant to be gone")]
25
26use std::collections::{HashMap, HashSet};
27use std::path::Path;
28
29use kevy_index::ColdBloom;
30use kevy_text::TextSegment;
31use kevy_text::cold::decode_fwd;
32
33const TXT_TAG: &[u8] = b"txtcold:";
36
37#[derive(Debug)]
40pub(super) struct ColdSeg {
41 pub(super) seg: kevy_seg::Seg,
42 pub(super) seq: u32,
43 pub(super) n_docs: u64,
44 pub(super) total_len: u64,
45}
46
47#[path = "text_query.rs"]
48mod query;
49pub use query::{ColdHit, ColdPage, ColdPageQuery};
50
51#[derive(Debug)]
55pub struct TextColdDir {
56 pub(super) segs: Vec<ColdSeg>,
57 seq: u32,
58 cleaned: bool,
59 bloom: ColdBloom,
60 pub(super) tombs: HashMap<Vec<u8>, HashSet<u32>>,
62 pub(super) df_dead: HashMap<Vec<u8>, u32>,
65}
66
67impl Default for TextColdDir {
68 fn default() -> Self {
69 Self::new()
70 }
71}
72
73impl TextColdDir {
74 pub fn new() -> Self {
76 Self {
77 segs: Vec::new(),
78 seq: 0,
79 cleaned: false,
80 bloom: ColdBloom::new(4096),
81 tombs: HashMap::new(),
82 df_dead: HashMap::new(),
83 }
84 }
85
86 pub fn has_cold(&self) -> bool {
89 !self.segs.is_empty()
90 }
91
92 pub fn on_row_write(&mut self, row_key: &[u8]) {
96 if !self.bloom.contains(row_key) {
97 return;
98 }
99 let mut fwd_key = vec![0u8];
100 fwd_key.extend_from_slice(row_key);
101 for cs in &mut self.segs {
102 let shadowed = self.tombs.get(row_key).is_some_and(|s| s.contains(&cs.seq));
103 if shadowed {
104 continue;
105 }
106 let Ok(Some(payload)) = cs.seg.get(&fwd_key) else { continue };
107 let Some(rec) = decode_fwd(&payload) else { continue };
108 cs.n_docs = cs.n_docs.saturating_sub(1);
109 cs.total_len = cs.total_len.saturating_sub(u64::from(rec.dl));
110 for t in rec.terms {
111 *self.df_dead.entry(t).or_insert(0) += 1;
112 }
113 self.tombs.entry(row_key.to_vec()).or_default().insert(cs.seq);
114 }
115 }
116
117 pub fn freeze_batch(
122 &mut self,
123 ts: &mut TextSegment,
124 index_name: &[u8],
125 keys: &[Vec<u8>],
126 segs_dir: &Path,
127 ) -> Result<bool, String> {
128 if !self.cleaned {
129 clean_stale(index_name, segs_dir)?;
130 self.cleaned = true;
131 }
132 let Some(bucket) = ts.freeze_docs(keys) else { return Ok(false) };
133 std::fs::create_dir_all(segs_dir).map_err(|e| e.to_string())?;
134 let file = format!("txt-{}-{}.seg", hex_stem(index_name), self.seq);
135 let seq = self.seq;
136 self.seq += 1;
137 let path = segs_dir.join(&file);
138 write_seg_file(&path, &bucket).inspect_err(|_| {
139 let _ = std::fs::remove_file(&path);
140 })?;
141 let mut m = kevy_seg::Manifest::open(segs_dir).map_err(|e| e.to_string())?;
142 let mut meta = TXT_TAG.to_vec();
143 meta.extend_from_slice(index_name);
144 meta.extend_from_slice(format!(":{}:{}", bucket.n_docs, bucket.total_len).as_bytes());
145 m.add(kevy_seg::ManifestEntry {
146 file: file.clone(),
147 meta,
148 min_key: bucket.fwd.keys().next().cloned().unwrap_or_default(),
149 max_key: bucket.terms.keys().next_back().cloned().unwrap_or_default(),
150 records: (bucket.fwd.len() + bucket.terms.len()) as u64,
151 })
152 .map_err(|e| e.to_string())?;
153 let seg = kevy_seg::Seg::open(&path).map_err(|e| format!("reopen {file}: {e}"))?;
154 self.segs.push(ColdSeg { seg, seq, n_docs: bucket.n_docs, total_len: bucket.total_len });
155 for k in keys {
156 self.bloom.insert(k);
157 }
158 Ok(true)
159 }
160}
161
162fn write_seg_file(path: &Path, bucket: &kevy_text::cold::FrozenBucket) -> Result<(), String> {
166 let mut b = kevy_seg::SegBuilder::create(path).map_err(|e| e.to_string())?;
167 for (row_key, payload) in &bucket.fwd {
168 let mut k = vec![0u8];
169 k.extend_from_slice(row_key);
170 b.push(&k, payload).map_err(|e| e.to_string())?;
171 }
172 for (term, payload) in &bucket.terms {
173 b.push(term, payload).map_err(|e| e.to_string())?;
174 }
175 b.finish().map(|_| ()).map_err(|e| e.to_string())
176}
177
178fn clean_stale(index_name: &[u8], segs_dir: &Path) -> Result<(), String> {
181 if !segs_dir.exists() {
182 return Ok(());
183 }
184 let mut m = kevy_seg::Manifest::open(segs_dir).map_err(|e| e.to_string())?;
185 let mut tag = TXT_TAG.to_vec();
186 tag.extend_from_slice(index_name);
187 tag.push(b':');
188 let stale: Vec<String> =
189 m.live().filter(|e| e.meta.starts_with(&tag)).map(|e| e.file.clone()).collect();
190 for f in stale {
191 m.drop_seg(&f).map_err(|e| e.to_string())?;
192 let _ = std::fs::remove_file(segs_dir.join(&f));
193 }
194 Ok(())
195}
196
197fn hex_stem(name: &[u8]) -> String {
198 name.iter().map(|b| format!("{b:02x}")).collect()
199}