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