Skip to main content

sift_core/index/trigram/
lifecycle.rs

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