Skip to main content

sift_core/index/trigram/
lifecycle.rs

1use std::path::{Path, PathBuf};
2
3use crate::index::snapshot::ArtifactData;
4use crate::index::{CorpusKind, IndexConfig, IndexDestination, IndexSource};
5
6use super::TrigramIndex;
7use super::TrigramIndexError;
8use super::builder::IndexTables;
9use super::file_table::{FileFingerprint, FileTable};
10use super::storage;
11use super::storage::lexicon::Lexicon;
12use super::storage::postings::Postings;
13use super::storage::trigram_sets::TrigramSets;
14
15impl TrigramIndex {
16    /// Write tables to `dir` as persistence files and return an mmap-backed index.
17    fn create_in_dir(
18        tables: &IndexTables,
19        root: &Path,
20        corpus_kind: CorpusKind,
21        dir: &Path,
22    ) -> crate::Result<Self> {
23        std::fs::create_dir_all(dir)?;
24
25        let files_path = dir.join(crate::FILES_BIN);
26        let lexicon_path = dir.join(crate::LEXICON_BIN);
27        let postings_path = dir.join(crate::POSTINGS_BIN);
28        let trigrams_path = dir.join(crate::TRIGRAMS_BIN);
29
30        let ((fr, lr), (pr, tr)) = rayon::join(
31            || {
32                rayon::join(
33                    || FileTable::create(&files_path, &tables.fingerprints),
34                    || storage::lexicon::Lexicon::create(&lexicon_path, &tables.lexicon),
35                )
36            },
37            || {
38                rayon::join(
39                    || storage::postings::Postings::create(&postings_path, &tables.postings),
40                    || {
41                        storage::trigram_sets::TrigramSets::create(
42                            &trigrams_path,
43                            &tables.file_trigrams,
44                        )
45                    },
46                )
47            },
48        );
49
50        let files = fr.map_err(crate::Error::Io)?;
51        let lexicon = lr.map_err(crate::Error::Io)?;
52        let postings = pr.map_err(crate::Error::Io)?;
53        let trigram_sets = tr.map_err(crate::Error::Io)?;
54
55        let root = root.to_path_buf();
56        let fingerprints = files.to_fingerprints().map_err(crate::Error::Io)?;
57        Self::validate_file_paths(&fingerprints)?;
58        Self::validate_lexicon_postings(&lexicon, &postings)?;
59
60        Ok(Self {
61            root,
62            fingerprints,
63            trigram_sets,
64            lexicon,
65            postings,
66            corpus_kind,
67        })
68    }
69
70    /// Build a trigram index from an explicit path list (or full walk when `paths` is empty).
71    ///
72    /// # Errors
73    ///
74    /// Returns an error if corpus walking, extraction, or encoding fails.
75    pub fn build(
76        config: &IndexConfig<'_>,
77        output_dir: &Path,
78        paths: &[PathBuf],
79    ) -> crate::Result<Self> {
80        let tables = IndexTables::build(config, paths)?;
81        let root = config.corpus.root.canonicalize()?;
82        Self::persist_tables(
83            &tables,
84            &root,
85            config.corpus.kind,
86            IndexDestination::Directory(output_dir),
87        )
88    }
89
90    /// Encode and store tables at the given destination, returning a live index.
91    pub(crate) fn persist_tables(
92        tables: &IndexTables,
93        root: &Path,
94        corpus_kind: CorpusKind,
95        dest: IndexDestination,
96    ) -> crate::Result<Self> {
97        match dest {
98            IndexDestination::Directory(dir) => Self::create_in_dir(tables, root, corpus_kind, dir),
99            IndexDestination::Snapshot { writer, namespace } => {
100                let ((fr, lr), (pr, tr)) = rayon::join(
101                    || {
102                        rayon::join(
103                            || FileTable::encode(&tables.fingerprints),
104                            || Lexicon::encode(&tables.lexicon),
105                        )
106                    },
107                    || {
108                        rayon::join(
109                            || Postings::encode(&tables.postings),
110                            || TrigramSets::encode(&tables.file_trigrams),
111                        )
112                    },
113                );
114
115                let files_bytes = fr.map_err(crate::Error::Io)?;
116                let lexicon_bytes = lr.map_err(crate::Error::Io)?;
117                let postings_bytes = pr.map_err(crate::Error::Io)?;
118                let trigram_sets_bytes = tr.map_err(crate::Error::Io)?;
119
120                let files =
121                    FileTable::from_artifact(ArtifactData::Memory(files_bytes.clone().into()))?;
122                let lexicon =
123                    Lexicon::from_artifact(ArtifactData::Memory(lexicon_bytes.clone().into()))?;
124                let postings =
125                    Postings::from_artifact(ArtifactData::Memory(postings_bytes.clone().into()))?;
126                let trigram_sets = TrigramSets::from_artifact(ArtifactData::Memory(
127                    trigram_sets_bytes.clone().into(),
128                ))?;
129
130                writer.put_artifact(namespace, crate::FILES_BIN, files_bytes)?;
131                writer.put_artifact(namespace, crate::LEXICON_BIN, lexicon_bytes)?;
132                writer.put_artifact(namespace, crate::POSTINGS_BIN, postings_bytes)?;
133                writer.put_artifact(namespace, crate::TRIGRAMS_BIN, trigram_sets_bytes)?;
134
135                let fingerprints = files.to_fingerprints().map_err(crate::Error::Io)?;
136                Self::validate_file_paths(&fingerprints)?;
137                Self::validate_lexicon_postings(&lexicon, &postings)?;
138
139                Ok(Self {
140                    root: root.to_path_buf(),
141                    fingerprints,
142                    trigram_sets,
143                    lexicon,
144                    postings,
145                    corpus_kind,
146                })
147            }
148        }
149    }
150
151    /// Open a previously persisted trigram index from `index_dir`.
152    ///
153    /// # Errors
154    ///
155    /// Returns an error if persistence files are missing or malformed.
156    pub fn open(index_dir: &Path, root: &Path, corpus_kind: CorpusKind) -> crate::Result<Self> {
157        Self::open_tables(IndexSource::Directory(index_dir), root, corpus_kind)
158    }
159
160    /// Update the index from the current corpus, writing artifact files
161    /// into `output_dir`.
162    ///
163    /// Returns `Ok(Some(index))` if a new index was written, or `Ok(None)`
164    /// if no files changed.
165    ///
166    /// # Errors
167    ///
168    /// Returns an error if corpus walking, extraction, or encoding fails.
169    pub fn update(
170        &self,
171        config: &IndexConfig<'_>,
172        output_dir: &Path,
173        paths: &[PathBuf],
174    ) -> crate::Result<Option<Self>> {
175        self.rebuild(config, IndexDestination::Directory(output_dir), paths)
176    }
177
178    /// Rebuild index tables for changed files and persist to `dest`.
179    pub(crate) fn rebuild(
180        &self,
181        config: &IndexConfig<'_>,
182        dest: IndexDestination,
183        paths: &[PathBuf],
184    ) -> crate::Result<Option<Self>> {
185        use rayon::prelude::*;
186        use std::collections::HashMap;
187
188        let fingerprints = if paths.is_empty() {
189            let corpus_paths =
190                crate::index::trigram::builder::CorpusWalker::new(config).collect()?;
191            crate::index::trigram::builder::FingerprintCollector::new(
192                config.corpus.root,
193                &corpus_paths,
194            )
195            .collect()?
196        } else {
197            Self::merge_partial_fingerprints(&self.fingerprints, config.corpus.root, paths)?
198        };
199
200        if fingerprints == self.fingerprints {
201            return Ok(None);
202        }
203
204        let prev_id_by_fp: HashMap<(&Path, i64, u64), usize> = self
205            .fingerprints
206            .iter()
207            .enumerate()
208            .map(|(id, fp)| ((fp.path.as_path(), fp.mtime_secs, fp.size), id))
209            .collect();
210
211        let file_trigrams: Vec<storage::trigram_sets::TrigramSet> = fingerprints
212            .par_iter()
213            .map(|fp| {
214                if let Some(&prev_id) =
215                    prev_id_by_fp.get(&(fp.path.as_path(), fp.mtime_secs, fp.size))
216                {
217                    return self.trigram_sets.get(prev_id).map_err(crate::Error::Io);
218                }
219                let abs = config.corpus.root.join(&fp.path);
220                storage::trigram_sets::TrigramSet::from_file(&abs).map_err(crate::Error::Io)
221            })
222            .collect::<crate::Result<_>>()?;
223
224        let (lexicon, postings) =
225            crate::index::trigram::builder::PostingAssembler::new(&file_trigrams).assemble()?;
226
227        let tables = IndexTables {
228            fingerprints,
229            file_trigrams,
230            lexicon,
231            postings,
232        };
233
234        let root = config.corpus.root.canonicalize()?;
235        let index = Self::persist_tables(&tables, &root, config.corpus.kind, dest)?;
236        Ok(Some(index))
237    }
238
239    fn merge_partial_fingerprints(
240        existing: &[super::file_table::FileFingerprint],
241        root: &Path,
242        paths: &[PathBuf],
243    ) -> crate::Result<Vec<super::file_table::FileFingerprint>> {
244        use std::collections::HashMap;
245
246        let mut by_path: HashMap<PathBuf, super::file_table::FileFingerprint> = existing
247            .iter()
248            .map(|fp| (fp.path.clone(), fp.clone()))
249            .collect();
250        for rel in paths {
251            let abs = root.join(rel);
252            let meta = std::fs::metadata(&abs).map_err(crate::Error::Io)?;
253            let mtime_secs = meta
254                .modified()
255                .ok()
256                .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
257                .map_or(0, |d| i64::try_from(d.as_secs()).unwrap_or(0));
258            let fp = super::file_table::FileFingerprint {
259                path: rel.clone(),
260                mtime_secs,
261                size: meta.len(),
262            };
263            by_path.insert(rel.clone(), fp);
264        }
265        let mut merged: Vec<_> = by_path.into_values().collect();
266        merged.sort_by(|a, b| a.path.cmp(&b.path));
267        Ok(merged)
268    }
269
270    /// Open index tables from a storage source (directory or snapshot).
271    pub(crate) fn open_tables(
272        source: IndexSource,
273        root: &Path,
274        corpus_kind: CorpusKind,
275    ) -> crate::Result<Self> {
276        match source {
277            IndexSource::Directory(dir) => {
278                let files_path = dir.join(crate::FILES_BIN);
279                let lexicon_path = dir.join(crate::LEXICON_BIN);
280                let postings_path = dir.join(crate::POSTINGS_BIN);
281                let trigrams_path = dir.join(crate::TRIGRAMS_BIN);
282
283                for p in [&files_path, &lexicon_path, &postings_path, &trigrams_path] {
284                    if !p.is_file() {
285                        return Err(TrigramIndexError::MissingComponent(p.clone()).into());
286                    }
287                }
288
289                let files = FileTable::open(&files_path).map_err(TrigramIndexError::Io)?;
290                let fingerprints = files.to_fingerprints().map_err(TrigramIndexError::Io)?;
291                Self::validate_file_paths(&fingerprints)?;
292
293                let lexicon = storage::lexicon::Lexicon::open(&lexicon_path)
294                    .map_err(TrigramIndexError::Io)?;
295                let postings = storage::postings::Postings::open(&postings_path)
296                    .map_err(TrigramIndexError::Io)?;
297
298                let trigram_sets = storage::trigram_sets::TrigramSets::open(&trigrams_path)
299                    .map_err(TrigramIndexError::Io)?;
300
301                Ok(Self {
302                    root: root.to_path_buf(),
303                    fingerprints,
304                    trigram_sets,
305                    lexicon,
306                    postings,
307                    corpus_kind,
308                })
309            }
310            IndexSource::Snapshot { reader, namespace } => {
311                let files_data = reader.artifact(namespace, crate::FILES_BIN)?;
312                let files = FileTable::from_artifact(files_data).map_err(TrigramIndexError::Io)?;
313                let fingerprints = files.to_fingerprints().map_err(TrigramIndexError::Io)?;
314                Self::validate_file_paths(&fingerprints)?;
315
316                let lexicon_data = reader.artifact(namespace, crate::LEXICON_BIN)?;
317                let lexicon =
318                    Lexicon::from_artifact(lexicon_data).map_err(TrigramIndexError::Io)?;
319
320                let postings_data = reader.artifact(namespace, crate::POSTINGS_BIN)?;
321                let postings =
322                    Postings::from_artifact(postings_data).map_err(TrigramIndexError::Io)?;
323
324                let trigram_sets_data = reader.artifact(namespace, crate::TRIGRAMS_BIN)?;
325                let trigram_sets =
326                    TrigramSets::from_artifact(trigram_sets_data).map_err(TrigramIndexError::Io)?;
327
328                Ok(Self {
329                    root: root.to_path_buf(),
330                    fingerprints,
331                    trigram_sets,
332                    lexicon,
333                    postings,
334                    corpus_kind,
335                })
336            }
337        }
338    }
339
340    fn validate_lexicon_postings(
341        lexicon: &storage::lexicon::Lexicon,
342        postings: &storage::postings::Postings,
343    ) -> Result<(), TrigramIndexError> {
344        let payload_len = postings.payload_len();
345        for entry in lexicon {
346            let start = usize::try_from(entry.offset).map_err(|_| {
347                TrigramIndexError::Io(std::io::Error::new(
348                    std::io::ErrorKind::InvalidData,
349                    format!(
350                        "lexicon entry {:?} offset {} exceeds usize",
351                        entry.trigram, entry.offset,
352                    ),
353                ))
354            })?;
355            let end = lexicon.posting_byte_end(entry.offset, payload_len);
356            if start > end || end > payload_len {
357                return Err(TrigramIndexError::Io(std::io::Error::new(
358                    std::io::ErrorKind::InvalidData,
359                    format!(
360                        "lexicon entry {:?} posting range [{start},{end}) exceeds payload_len {payload_len}",
361                        entry.trigram,
362                    ),
363                )));
364            }
365            let slice = postings.slice(start, end.saturating_sub(start));
366            let decoded_count = storage::postings::Postings::validate_list(slice).map_err(|e| {
367                TrigramIndexError::Io(std::io::Error::new(
368                    std::io::ErrorKind::InvalidData,
369                    format!("posting list for trigram {:?}: {e}", entry.trigram),
370                ))
371            })?;
372            if decoded_count != entry.len as usize {
373                return Err(TrigramIndexError::Io(std::io::Error::new(
374                    std::io::ErrorKind::InvalidData,
375                    format!(
376                        "lexicon entry {:?} claims len {} but posting list has {decoded_count} entries",
377                        entry.trigram, entry.len,
378                    ),
379                )));
380            }
381        }
382        Ok(())
383    }
384
385    fn validate_file_paths(fingerprints: &[FileFingerprint]) -> Result<(), TrigramIndexError> {
386        for fp in fingerprints {
387            if fp.path.as_os_str().is_empty()
388                || fp.path.is_absolute()
389                || fp
390                    .path
391                    .components()
392                    .any(|c| matches!(c, std::path::Component::ParentDir))
393            {
394                return Err(TrigramIndexError::Io(std::io::Error::new(
395                    std::io::ErrorKind::InvalidData,
396                    format!("invalid file path in index: {}", fp.path.display()),
397                )));
398            }
399        }
400        Ok(())
401    }
402}
403
404#[cfg(test)]
405mod tests {
406    use super::*;
407    use std::path::PathBuf;
408
409    #[test]
410    fn validate_file_paths_accepts_normal_relative_paths() {
411        let fps = vec![
412            FileFingerprint {
413                path: PathBuf::from("a.txt"),
414                mtime_secs: 0,
415                size: 0,
416            },
417            FileFingerprint {
418                path: PathBuf::from("sub/b.txt"),
419                mtime_secs: 0,
420                size: 0,
421            },
422        ];
423        let result = TrigramIndex::validate_file_paths(&fps);
424        assert!(result.is_ok());
425    }
426
427    #[test]
428    fn validate_file_paths_rejects_absolute_paths() {
429        let abs = std::env::current_dir().unwrap().join("a.txt");
430        let fps = vec![FileFingerprint {
431            path: abs,
432            mtime_secs: 0,
433            size: 0,
434        }];
435        let result = TrigramIndex::validate_file_paths(&fps);
436        assert!(result.is_err());
437    }
438
439    #[test]
440    fn validate_file_paths_rejects_empty_paths() {
441        let fps = vec![FileFingerprint {
442            path: PathBuf::from(""),
443            mtime_secs: 0,
444            size: 0,
445        }];
446        let result = TrigramIndex::validate_file_paths(&fps);
447        assert!(result.is_err());
448    }
449
450    #[test]
451    fn validate_file_paths_rejects_parent_dir_paths() {
452        let fps = vec![FileFingerprint {
453            path: PathBuf::from("../escape.txt"),
454            mtime_secs: 0,
455            size: 0,
456        }];
457        let result = TrigramIndex::validate_file_paths(&fps);
458        assert!(result.is_err());
459    }
460}