Skip to main content

yo_doc/
docs.rs

1//! A document collection: the primary table and the key table that goes with
2//! it (`09` section 4).
3//!
4//! The primary table is an element table keyed by document id with the
5//! document's bytes stored behind its id in the same blob. That is not a family
6//! resemblance to a hash's field table, it is the same code: [`Elements`] in
7//! tailed mode, which is what `HSET` writes into. A document collection and a
8//! hash differ in what the bytes behind the name mean and in nothing else, and
9//! that is the point of R25.
10//!
11//! ```
12//! use yo_doc::{Builder, Docs};
13//!
14//! let mut b = Builder::new();
15//! b.begin_object()?;
16//! b.key(b"customer")?;
17//! b.int(7)?;
18//! b.key(b"status")?;
19//! b.text("open")?;
20//! b.end_object()?;
21//! let order = b.finish()?.to_vec();
22//!
23//! let mut docs = Docs::new();
24//! assert!(docs.put_bytes(b"order:1", &order)?);
25//! let got = docs.get(b"order:1").expect("stored");
26//! assert_eq!(got.get(b"status").and_then(|v| v.as_text()), Some("open"));
27//! # Ok::<(), yo_common::Error>(())
28//! ```
29//!
30//! # What a write does
31//!
32//! [`Docs::put`] does not store the value it is given. It walks it once and
33//! writes it again with every object key replaced by its id from the
34//! collection's [`Keys`], which is where the forty percent that interning is
35//! worth actually gets saved.
36//!
37//! If the key table fills part way through, the document is stored as it
38//! arrived with its keys as bytes. Nothing about that is a fallback mode: the
39//! interned flag sits in each container's header, so a collection holds both
40//! kinds at once, a reader tells them apart per container, and documents
41//! written before the table filled stay exactly as they were.
42//!
43//! # What a write does about the indexes
44//!
45//! A write takes one path lookup per declared index, not a comparison against
46//! every index path at every node of the document. I had it the other way round
47//! at first, on the argument that the interning walk is already touching every
48//! node so the extraction may as well ride along on it. That is worse: with N
49//! indexes it costs N comparisons at every node, where a lookup per index costs
50//! N times the two or three binary searches a shallow path takes, and index
51//! paths are shallow. It is also much simpler, and it is the same code the
52//! backfill in [`Docs::create_index`] runs.
53//!
54//! The keys are worked out before anything is stored, so a value that is too
55//! long to be an index key fails the write rather than leaving behind a document
56//! that is in the collection and in none of its indexes. Then the old document
57//! under that id is taken out of the indexes, then the new one is stored, then
58//! it is filed. An overwrite and a removal both un-index through the same code,
59//! because both make the old entries wrong.
60//!
61//! A vector index at a path is read in the same pass and for the same reason: a
62//! document whose embedding is the wrong shape fails the write rather than
63//! landing in the collection with no vector in it. It is filed after the
64//! document is stored, along with the tag that a filtered search tests, which is
65//! the keys the other indexes just filed this document under. See
66//! [`vector`](crate::vector).
67//!
68//! # Reading one back
69//!
70//! [`Docs::get`] answers a [`Doc`], which is a [`Value`] with the collection's
71//! key table beside it. Everything that needs a name rather than an id goes
72//! through the table: `get(b"status")` resolves the name to an id once and then
73//! searches the document by id, which is a binary search over integers.
74
75use core::ops::Bound;
76
77use yo_common::{Code, Error, Result};
78use yo_kv::{Cursor, Elements, Full};
79use yo_shape::Metric;
80use yo_vector::{Match, Signature};
81
82use crate::head::{DEPTH_MAX, Kind};
83use crate::index::{self, IndexKind, Key, PathIndex};
84use crate::path::{Step, Steps};
85use crate::vector::{self, VectorIndex};
86use crate::{Builder, Keys, Value};
87
88/// Documents by id, with the key table their keys are interned against.
89#[derive(Debug)]
90pub struct Docs {
91    /// Document id to the document's bytes, the bytes behind the id.
92    rows: Elements<()>,
93    /// The names every interned object in this collection uses.
94    keys: Keys,
95    /// The buffer a write is re-encoded into, kept so a write does not allocate.
96    build: Builder,
97    /// One per indexed path, in the order they were declared.
98    indexes: Vec<PathIndex>,
99    /// The key each index takes from the document being written, one slot per
100    /// index and empty where the document has nothing to file.
101    ///
102    /// Worked out before anything is stored, so a value that cannot be indexed
103    /// fails the write rather than leaving a document behind that no query will
104    /// ever find. Kept on the collection so a write allocates nothing.
105    taken: Vec<Vec<u8>>,
106    /// One per path holding an embedding, in the order they were declared.
107    ///
108    /// Beside the path indexes rather than among them, because nearness has no
109    /// key to file under. See [`vector`](crate::vector).
110    vectors: Vec<VectorIndex>,
111    /// The vector each of those takes from the document being written, one slot
112    /// per index and empty where the document has nothing at the path.
113    ///
114    /// Read before anything is stored, for the same reason the index keys are:
115    /// a document whose embedding is the wrong shape fails the write rather than
116    /// landing in the collection and in none of its vector indexes.
117    drawn: Vec<Vec<f32>>,
118}
119
120impl Default for Docs {
121    /// Not derived, because the primary table has to be the kind that keeps a
122    /// document behind its id and an empty [`Elements`] is not.
123    fn default() -> Docs {
124        Docs::new()
125    }
126}
127
128impl Docs {
129    /// An empty collection that has not allocated anything yet.
130    #[must_use]
131    pub fn new() -> Docs {
132        Docs {
133            rows: Elements::tailed(0, 0),
134            keys: Keys::new(),
135            build: Builder::new(),
136            indexes: Vec::new(),
137            taken: Vec::new(),
138            vectors: Vec::new(),
139            drawn: Vec::new(),
140        }
141    }
142
143    /// An empty collection with room for `n` documents of about `each` bytes.
144    ///
145    /// The ids and the documents share one blob, so the size asked for is the
146    /// two of them together. Getting it wrong costs a growth, not a rewrite.
147    #[must_use]
148    pub fn with_capacity(n: usize, each: usize) -> Docs {
149        Docs {
150            rows: Elements::tailed(n, n.saturating_mul(each)),
151            keys: Keys::new(),
152            build: Builder::with_capacity(each),
153            indexes: Vec::new(),
154            taken: Vec::new(),
155            vectors: Vec::new(),
156            drawn: Vec::new(),
157        }
158    }
159
160    /// Store `value` under `id`, and say whether the id is new.
161    ///
162    /// The value is re-encoded with this collection's interned keys on the way
163    /// in. It may not already be interned: a document whose keys are ids
164    /// belongs to whichever collection handed those ids out, and moving it to
165    /// another one without the names is how a collection ends up reading the
166    /// wrong field.
167    pub fn put(&mut self, id: &[u8], value: Value<'_>) -> Result<bool> {
168        self.write(id, value, None)
169    }
170
171    /// Store the document `doc` encodes under `id`, and say whether the id is
172    /// new.
173    ///
174    /// The bytes are checked far enough to be readable and no further, the same
175    /// as [`Value::new`]. A caller holding bytes it did not write should run
176    /// [`Value::validate`] first.
177    pub fn put_bytes(&mut self, id: &[u8], doc: &[u8]) -> Result<bool> {
178        let value = Value::new(doc)
179            .ok_or_else(|| Error::new(Code::Corrupt, "the document is not a readable value"))?;
180        self.write(id, value, Some(doc))
181    }
182
183    /// The write both forms of put go through.
184    ///
185    /// `raw` is the caller's bytes when it had some, so that the path where the
186    /// key table is full stores them directly instead of copying them through
187    /// the builder to get back what it was already holding.
188    ///
189    /// The order is: work out the index keys, un-index whatever was under this
190    /// id, store, index. Working the keys out first is what makes a write that
191    /// cannot be indexed leave the collection exactly as it was, rather than
192    /// storing a document that no query will find.
193    fn write(&mut self, id: &[u8], value: Value<'_>, raw: Option<&[u8]>) -> Result<bool> {
194        let Docs {
195            rows,
196            keys,
197            build,
198            indexes,
199            taken,
200            vectors,
201            drawn,
202        } = self;
203
204        taken.resize(indexes.len(), Vec::new());
205        for (slot, index) in taken.iter_mut().zip(indexes.iter()) {
206            slot.clear();
207            // The incoming value has its keys as bytes, since put refuses one
208            // that does not, so its paths resolve without the key table.
209            let Some(at) = value.path_bytes(index.path())? else {
210                continue;
211            };
212            if index.keys_at(at, slot).is_err() {
213                return Err(Error::fmt(
214                    Code::Full,
215                    format_args!(
216                        "a value at {} is longer than {} bytes and cannot be indexed",
217                        String::from_utf8_lossy(index.path()),
218                        index::KEY_MAX
219                    ),
220                ));
221            }
222        }
223
224        drawn.resize(vectors.len(), Vec::new());
225        for (slot, index) in drawn.iter_mut().zip(vectors.iter()) {
226            slot.clear();
227            let Some(at) = value.path_bytes(index.path())? else {
228                continue;
229            };
230            vector::coordinates(at, index.dim(), index.path(), slot)?;
231        }
232        // What a filtered search will meet in the posting scan: one bit per key
233        // the other indexes file this document under, worked out here because
234        // `taken` already holds exactly those keys.
235        let tag = if vectors.is_empty() {
236            0
237        } else {
238            vector::tag_of(
239                indexes
240                    .iter()
241                    .map(PathIndex::path)
242                    .zip(taken.iter().map(Vec::as_slice)),
243            )
244        };
245
246        unindex(rows, keys, indexes, id);
247
248        build.clear();
249        let fresh = if intern_into(keys, build, value, 0)? {
250            store(rows, id, build.finish()?)?
251        } else if let Some(raw) = raw {
252            // The key table filled part way through, and the caller is holding
253            // exactly what should be stored.
254            store(rows, id, raw)?
255        } else {
256            build.clear();
257            build.embed(&value)?;
258            store(rows, id, build.finish()?)?
259        };
260
261        for (slot, index) in taken.iter().zip(indexes.iter_mut()) {
262            let mut filed = Ok(());
263            index::each_key(slot, |key| {
264                if filed.is_ok() {
265                    filed = index.add(key, id);
266                }
267            });
268            filed?;
269        }
270        // A document that no longer has anything at the path leaves the vector
271        // index, rather than keeping the vector the last version of it had.
272        for (slot, index) in drawn.iter().zip(vectors.iter_mut()) {
273            if slot.is_empty() {
274                index.collection_mut().remove(id);
275            } else {
276                index.collection_mut().put_tagged(id, slot, tag)?;
277            }
278        }
279        Ok(fresh)
280    }
281
282    /// Start indexing `path` for equality, and file every document already here
283    /// under it.
284    ///
285    /// Declaring the same path twice is not an error and does not rebuild
286    /// anything, because a caller that opens a collection and declares its
287    /// indexes on the way in should be able to do that every time it opens it.
288    /// An ordered index that is already there stays ordered, since it answers
289    /// equality as well.
290    ///
291    /// The backfill is a path lookup per document, so it costs the collection
292    /// once. There is no background indexer and no window in which the index is
293    /// declared and not yet true, which is Y3.
294    pub fn create_index(&mut self, path: &str) -> Result<()> {
295        self.create_index_bytes(path.as_bytes(), IndexKind::Equality)
296    }
297
298    /// Start indexing `path` for equality and for ranges.
299    ///
300    /// An ordered index is an equality index with a counted B+ tree over the
301    /// rows of its key table, which is the same tree a sorted set ranks with.
302    /// It costs about three bytes per distinct value on top of the equality
303    /// index and a logarithmic search per new value, and it is what
304    /// [`Docs::range`] needs.
305    ///
306    /// A path that is already indexed for equality is upgraded and rebuilt. The
307    /// alternative is answering `Ok` and then having every range on it come back
308    /// empty, which is a query that lies.
309    pub fn create_ordered_index(&mut self, path: &str) -> Result<()> {
310        self.create_index_bytes(path.as_bytes(), IndexKind::Ordered)
311    }
312
313    /// Start indexing every element of the array at `path`.
314    ///
315    /// A document with `["red", "blue"]` there is filed under both, so a search
316    /// for either finds it. A scalar at the path is an array of one, so a
317    /// collection where some documents have a list of tags and some have a
318    /// single tag works without the caller having to normalise it first.
319    ///
320    /// The lookup is [`Docs::find`] with the element as the key, unchanged. An
321    /// array index costs what the document has at the path, so a document with
322    /// ten elements costs ten postings and a document with none costs nothing.
323    pub fn create_array_index(&mut self, path: &str) -> Result<()> {
324        self.create_index_bytes(path.as_bytes(), IndexKind::Array)
325    }
326
327    /// Start indexing every word of the string at `path`.
328    ///
329    /// A document with `"A red bicycle"` there is filed under `a`, `red` and
330    /// `bicycle`, and the lookup is [`Docs::find`] with [`Key::word`] as the
331    /// key. Case is folded on both sides, so a search does not have to know how
332    /// the document was written.
333    ///
334    /// This is a word index and not a search engine. There is no ranking, no
335    /// stemming and no phrase matching, and a path that holds something other
336    /// than a string files nothing. What it answers is which documents contain
337    /// a word, which is a filter, and the ranking that belongs on top of it is
338    /// `10`.
339    pub fn create_text_index(&mut self, path: &str) -> Result<()> {
340        self.create_index_bytes(path.as_bytes(), IndexKind::Text)
341    }
342
343    /// [`Docs::create_index`] and [`Docs::create_ordered_index`] for a path that
344    /// is already bytes.
345    pub fn create_index_bytes(&mut self, path: &[u8], kind: IndexKind) -> Result<()> {
346        for step in Steps::new(path) {
347            step?;
348        }
349        // The same kind again is nothing at all, and equality on top of ordered
350        // is already answered. Every other pair means the path is being asked a
351        // different question, so it gets rebuilt. The old one stays in place
352        // until the new one is filled, so a backfill that fails leaves the
353        // collection with the index it already had rather than with none.
354        let old = match self.indexes.iter().position(|i| i.path() == path) {
355            Some(at) if self.indexes[at].kind() == kind => return Ok(()),
356            Some(at)
357                if kind == IndexKind::Equality && self.indexes[at].kind() == IndexKind::Ordered =>
358            {
359                return Ok(());
360            }
361            found => found,
362        };
363        let mut index = PathIndex::new(path, kind);
364        let mut list = Vec::new();
365        for (id, bytes) in self.rows.pairs() {
366            let Some(value) = Value::new(bytes) else {
367                continue;
368            };
369            let doc = Doc {
370                value,
371                keys: &self.keys,
372            };
373            let Some(at) = doc.path_bytes(path)? else {
374                continue;
375            };
376            list.clear();
377            if index.keys_at(at.value(), &mut list).is_err() {
378                return Err(Error::fmt(
379                    Code::Full,
380                    format_args!(
381                        "a value at {} in {} is longer than {} bytes and cannot be indexed",
382                        String::from_utf8_lossy(path),
383                        String::from_utf8_lossy(id),
384                        index::KEY_MAX
385                    ),
386                ));
387            }
388            let mut filed = Ok(());
389            index::each_key(&list, |key| {
390                if filed.is_ok() {
391                    filed = index.add(key, id);
392                }
393            });
394            filed?;
395        }
396        match old {
397            Some(at) => self.indexes[at] = index,
398            None => {
399                self.indexes.push(index);
400                self.taken.push(Vec::new());
401            }
402        }
403        self.retag();
404        Ok(())
405    }
406
407    /// Stop indexing `path`, and say whether it was indexed.
408    pub fn drop_index(&mut self, path: &str) -> bool {
409        self.drop_index_bytes(path.as_bytes())
410    }
411
412    /// [`Docs::drop_index`] for a path that is already bytes.
413    pub fn drop_index_bytes(&mut self, path: &[u8]) -> bool {
414        let Some(at) = self.indexes.iter().position(|i| i.path() == path) else {
415            return false;
416        };
417        self.indexes.remove(at);
418        self.taken.truncate(self.indexes.len());
419        self.retag();
420        true
421    }
422
423    /// The indexes this collection keeps, in the order they were declared.
424    #[must_use]
425    pub fn indexes(&self) -> &[PathIndex] {
426        &self.indexes
427    }
428
429    /// The index on `path`, if there is one.
430    #[must_use]
431    pub fn index(&self, path: &str) -> Option<&PathIndex> {
432        self.indexes.iter().find(|i| i.path() == path.as_bytes())
433    }
434
435    /// Start indexing the `dim` wide embedding at `path` for nearness, by
436    /// cosine, and file every document already here.
437    ///
438    /// Cosine because that is what a text or image embedding is compared by, and
439    /// a collection built for one measure and searched as if it were another
440    /// gives wrong answers quietly. [`Docs::create_vector_index_with`] takes the
441    /// other one.
442    ///
443    /// Declaring the same path at the same width and measure twice is not an
444    /// error and rebuilds nothing, the same as [`Docs::create_index`], so a
445    /// caller that declares its indexes every time it opens a collection can.
446    /// Changing the width or the measure rebuilds, because there is nothing in
447    /// the old collection that answers the new question.
448    pub fn create_vector_index(&mut self, path: &str, dim: usize) -> Result<()> {
449        self.create_vector_index_bytes(path.as_bytes(), dim, Metric::Cosine)
450    }
451
452    /// The same, saying what nearness means.
453    pub fn create_vector_index_with(
454        &mut self,
455        path: &str,
456        dim: usize,
457        metric: Metric,
458    ) -> Result<()> {
459        self.create_vector_index_bytes(path.as_bytes(), dim, metric)
460    }
461
462    /// [`Docs::create_vector_index_with`] for a path that is already bytes.
463    ///
464    /// # Errors
465    ///
466    /// [`Code::Invalid`] for a path that does not parse, a width of zero or past
467    /// the format's limit, and for a document already here whose value at the
468    /// path is not an array of `dim` numbers. [`Code::Unsupported`] for a
469    /// measure this build does not do.
470    pub fn create_vector_index_bytes(
471        &mut self,
472        path: &[u8],
473        dim: usize,
474        metric: Metric,
475    ) -> Result<()> {
476        for step in Steps::new(path) {
477            step?;
478        }
479        // As with a path index, the one that is already here stays until the new
480        // one is filled, so a document whose embedding is not the new width
481        // leaves the collection with the index it had rather than with none.
482        let old = match self.vectors.iter().position(|v| v.path() == path) {
483            Some(at) if self.vectors[at].dim() == dim && self.vectors[at].metric() == metric => {
484                return Ok(());
485            }
486            found => found,
487        };
488        let mut index = VectorIndex::new(path, dim, metric)?;
489        let mut list = Vec::new();
490        let mut v = Vec::new();
491        for (id, bytes) in self.rows.pairs() {
492            let Some(value) = Value::new(bytes) else {
493                continue;
494            };
495            let doc = Doc {
496                value,
497                keys: &self.keys,
498            };
499            let Some(at) = doc.path_bytes(path)? else {
500                continue;
501            };
502            vector::coordinates(at.value(), dim, path, &mut v)?;
503            let tag = tag_for(&doc, &self.indexes, &mut list);
504            index.collection_mut().put_tagged(id, &v, tag)?;
505        }
506        match old {
507            Some(at) => self.vectors[at] = index,
508            None => {
509                self.vectors.push(index);
510                self.drawn.push(Vec::new());
511            }
512        }
513        Ok(())
514    }
515
516    /// Stop indexing the embedding at `path`, and say whether it was indexed.
517    pub fn drop_vector_index(&mut self, path: &str) -> bool {
518        self.drop_vector_index_bytes(path.as_bytes())
519    }
520
521    /// [`Docs::drop_vector_index`] for a path that is already bytes.
522    pub fn drop_vector_index_bytes(&mut self, path: &[u8]) -> bool {
523        let Some(at) = self.vectors.iter().position(|v| v.path() == path) else {
524            return false;
525        };
526        self.vectors.remove(at);
527        self.drawn.truncate(self.vectors.len());
528        true
529    }
530
531    /// The vector indexes this collection keeps, in the order they were
532    /// declared.
533    #[must_use]
534    pub fn vector_indexes(&self) -> &[VectorIndex] {
535        &self.vectors
536    }
537
538    /// The vector index on `path`, if there is one.
539    #[must_use]
540    pub fn vector_index(&self, path: &str) -> Option<&VectorIndex> {
541        self.vectors.iter().find(|v| v.path() == path.as_bytes())
542    }
543
544    /// The vector stored for `id` at `path`, as the index holds it.
545    ///
546    /// For a cosine index that is the normalised vector and not the one the
547    /// document carries, because normalising once at write time is what makes
548    /// every later comparison a dot product. The document itself still has what
549    /// was written.
550    #[must_use]
551    pub fn embedding(&self, path: &str, id: &[u8]) -> Option<&[f32]> {
552        self.vector_index(path)?.collection().get(id)
553    }
554
555    /// Hand the `k` documents nearest to `q` at `path` to `f`, nearest first,
556    /// and say how many there were.
557    ///
558    /// The third argument to `f` is the distance, which for a cosine index is
559    /// one minus the cosine so that nearer is smaller, and it is measured
560    /// against the full precision vector rather than against the code.
561    ///
562    /// A path with no vector index on it is an error and not a scan, the same
563    /// rule [`Docs::find`] follows.
564    pub fn nearest(
565        &self,
566        path: &str,
567        q: &[f32],
568        k: usize,
569        f: impl FnMut(&[u8], Doc<'_>, f32),
570    ) -> Result<usize> {
571        let hits = self.vector(path)?.collection().search(q, k, None)?;
572        Ok(self.answer(&hits, f))
573    }
574
575    /// The same, over only the documents whose indexed fields hold every value
576    /// in `want`.
577    ///
578    /// The filter runs inside the posting scan rather than over the answers, so
579    /// a selective filter returns the nearest documents that match instead of
580    /// whichever of the nearest happened to match. See
581    /// [`vector`](crate::vector) for the encoding and for the one direction it
582    /// is not exact in: a document can pass a filter it does not really match,
583    /// never fail one it does, so a caller with a predicate of its own still
584    /// gets every answer to check.
585    ///
586    /// Every path in `want` has to carry an ordinary index, because the bits the
587    /// filter tests are the keys those indexes filed the document under. A path
588    /// with no index is an error rather than a filter that matches nothing.
589    pub fn nearest_where(
590        &self,
591        path: &str,
592        q: &[f32],
593        k: usize,
594        want: &[(&str, Key)],
595        f: impl FnMut(&[u8], Doc<'_>, f32),
596    ) -> Result<usize> {
597        let filter = self.wanted(want)?;
598        let hits = self
599            .vector(path)?
600            .collection()
601            .search_where(q, k, None, &filter)?;
602        Ok(self.answer(&hits, f))
603    }
604
605    /// Hand the `k` documents most like the one under `id` to `f`, `id` itself
606    /// left out.
607    ///
608    /// More like this, which is the query a document collection with embeddings
609    /// in it is really for. A document with nothing at the path has nothing to
610    /// be like, so it answers zero rather than an error.
611    pub fn nearest_to(
612        &self,
613        path: &str,
614        id: &[u8],
615        k: usize,
616        f: impl FnMut(&[u8], Doc<'_>, f32),
617    ) -> Result<usize> {
618        let index = self.vector(path)?;
619        let Some(q) = index.collection().get(id) else {
620            return Ok(0);
621        };
622        let hits = index.collection().search(q, k, Some(id))?;
623        Ok(self.answer(&hits, f))
624    }
625
626    /// The vector index on `path`, or the error that says why there is not one.
627    fn vector(&self, path: &str) -> Result<&VectorIndex> {
628        self.vector_index(path).ok_or_else(|| {
629            Error::fmt(
630                Code::Invalid,
631                format_args!("there is no vector index on {path}, so this would be a scan"),
632            )
633        })
634    }
635
636    /// Turn what a query requires into the signature the scan tests.
637    fn wanted(&self, want: &[(&str, Key)]) -> Result<Signature> {
638        let mut sig = Signature::default();
639        for (path, key) in want {
640            if self.index(path).is_none() {
641                return Err(Error::fmt(
642                    Code::Invalid,
643                    format_args!("there is no index on {path}, so a search cannot filter on it"),
644                ));
645            }
646            sig.insert(path, key.as_bytes());
647        }
648        Ok(sig)
649    }
650
651    /// Read the documents a search answered with, skipping any that have gone.
652    fn answer(&self, hits: &[Match], mut f: impl FnMut(&[u8], Doc<'_>, f32)) -> usize {
653        let mut n = 0usize;
654        for hit in hits {
655            if let Some(doc) = self.get(&hit.key) {
656                f(&hit.key, doc, hit.distance);
657                n += 1;
658            }
659        }
660        n
661    }
662
663    /// Work out every document's tag again and write it back.
664    ///
665    /// A tag summarises the keys the other indexes filed the document under, so
666    /// declaring or dropping one makes every tag wrong. This is one store per
667    /// document per vector index, with no requantising and no maintenance, which
668    /// is cheap enough to do on the spot and is the only alternative to a filter
669    /// that used to work quietly answering nothing.
670    fn retag(&mut self) {
671        let Docs {
672            rows,
673            keys,
674            indexes,
675            vectors,
676            ..
677        } = self;
678        if vectors.is_empty() {
679            return;
680        }
681        let mut list = Vec::new();
682        for (id, bytes) in rows.pairs() {
683            let Some(value) = Value::new(bytes) else {
684                continue;
685            };
686            let doc = Doc { value, keys };
687            let tag = tag_for(&doc, indexes, &mut list);
688            for index in vectors.iter_mut() {
689                index.collection_mut().retag(id, tag);
690            }
691        }
692    }
693
694    /// Hand every document whose value at `path` is `key` to `f`, and say how
695    /// many there were.
696    ///
697    /// One probe of the index and one probe of the primary table per document,
698    /// which is the cost model `09` section 5 states rather than hides. A path
699    /// with no index on it is an error and not a scan: a query that silently
700    /// turns into a walk of the collection is the thing this API exists not to
701    /// do.
702    pub fn find(&self, path: &str, key: &Key, mut f: impl FnMut(&[u8], Doc<'_>)) -> Result<usize> {
703        let index = self.index(path).ok_or_else(|| {
704            Error::fmt(
705                Code::Invalid,
706                format_args!("there is no index on {path}, so this would be a scan"),
707            )
708        })?;
709        let Some(set) = index.get(key) else {
710            return Ok(0);
711        };
712        let mut n = 0usize;
713        index::each_id(set, |id| {
714            if let Some(doc) = self.get(id) {
715                f(id, doc);
716                n += 1;
717            }
718        });
719        Ok(n)
720    }
721
722    /// How many documents have `key` at `path`, without reading any of them.
723    ///
724    /// The number a caller sorts its filters by before it intersects them, and
725    /// it is a probe rather than a walk.
726    pub fn count(&self, path: &str, key: &Key) -> Result<usize> {
727        let index = self.index(path).ok_or_else(|| {
728            Error::fmt(
729                Code::Invalid,
730                format_args!("there is no index on {path}, so this would be a scan"),
731            )
732        })?;
733        Ok(index.count(key))
734    }
735
736    /// Hand every document whose value at `path` falls between `lo` and `hi` to
737    /// `f`, smallest first, and say how many there were.
738    ///
739    /// One search of the tree and then a walk, so the cost is the size of the
740    /// answer and not the size of the collection. The bounds are the ordinary
741    /// [`Bound`], so a half open range, a range open at one end and a range open
742    /// at both are all the same call.
743    ///
744    /// The path has to carry an ordered index. An equality index has no order to
745    /// walk, and answering nothing would be a query that lies rather than a
746    /// query that says no.
747    pub fn range(
748        &self,
749        path: &str,
750        lo: Bound<&Key>,
751        hi: Bound<&Key>,
752        mut f: impl FnMut(&[u8], Doc<'_>),
753    ) -> Result<usize> {
754        let index = self.ordered(path)?;
755        let mut n = 0usize;
756        for (_, set) in index.range(lo, hi) {
757            index::each_id(set, |id| {
758                if let Some(doc) = self.get(id) {
759                    f(id, doc);
760                    n += 1;
761                }
762            });
763        }
764        Ok(n)
765    }
766
767    /// [`Docs::range`] backwards, largest value first.
768    pub fn range_rev(
769        &self,
770        path: &str,
771        lo: Bound<&Key>,
772        hi: Bound<&Key>,
773        mut f: impl FnMut(&[u8], Doc<'_>),
774    ) -> Result<usize> {
775        let index = self.ordered(path)?;
776        let mut n = 0usize;
777        for (_, set) in index.range_rev(lo, hi) {
778            index::each_id(set, |id| {
779                if let Some(doc) = self.get(id) {
780                    f(id, doc);
781                    n += 1;
782                }
783            });
784        }
785        Ok(n)
786    }
787
788    /// How many documents fall between `lo` and `hi` at `path`, without reading
789    /// any of them.
790    ///
791    /// This reads the distinct values in the range rather than the documents, so
792    /// a range covering a million documents under a hundred values costs a
793    /// hundred.
794    pub fn count_range(&self, path: &str, lo: Bound<&Key>, hi: Bound<&Key>) -> Result<usize> {
795        Ok(self.ordered(path)?.count_in(lo, hi))
796    }
797
798    /// The ordered index on `path`, or the error that says why there is not one.
799    fn ordered(&self, path: &str) -> Result<&PathIndex> {
800        match self.index(path) {
801            Some(index) if index.kind() == IndexKind::Ordered => Ok(index),
802            Some(_) => Err(Error::fmt(
803                Code::Invalid,
804                format_args!("the index on {path} answers equality and not ranges"),
805            )),
806            None => Err(Error::fmt(
807                Code::Invalid,
808                format_args!("there is no index on {path}, so this would be a scan"),
809            )),
810        }
811    }
812
813    /// The document stored under `id`.
814    #[must_use]
815    pub fn get(&self, id: &[u8]) -> Option<Doc<'_>> {
816        let value = Value::new(self.rows.tail(id)?)?;
817        Some(Doc {
818            value,
819            keys: &self.keys,
820        })
821    }
822
823    /// The stored bytes of the document under `id`, as they sit in the blob.
824    ///
825    /// For a caller that is going to write them somewhere else rather than read
826    /// them, which is `DUMP`, replication and the record plane.
827    #[must_use]
828    pub fn bytes(&self, id: &[u8]) -> Option<&[u8]> {
829        self.rows.tail(id)
830    }
831
832    /// Whether there is a document under `id`.
833    #[must_use]
834    pub fn contains(&self, id: &[u8]) -> bool {
835        self.rows.contains(id)
836    }
837
838    /// Take the document under `id` out, and say whether there was one.
839    ///
840    /// Every index the document was filed in loses it first, so a removal costs
841    /// a path lookup per index on the way out.
842    ///
843    /// The key table is left alone. A name it interned stays interned even if
844    /// this was the last document using it, which is [`Keys`]'s rule and the
845    /// reason an id is a row index.
846    pub fn remove(&mut self, id: &[u8]) -> bool {
847        let Docs {
848            rows,
849            keys,
850            indexes,
851            vectors,
852            ..
853        } = self;
854        unindex(rows, keys, indexes, id);
855        for index in vectors.iter_mut() {
856            index.collection_mut().remove(id);
857        }
858        rows.remove(id).is_some()
859    }
860
861    /// How many documents there are.
862    #[must_use]
863    pub fn len(&self) -> usize {
864        self.rows.len()
865    }
866
867    /// Whether the collection holds nothing.
868    #[must_use]
869    pub fn is_empty(&self) -> bool {
870        self.rows.is_empty()
871    }
872
873    /// The names this collection has interned.
874    #[must_use]
875    pub fn keys(&self) -> &Keys {
876        &self.keys
877    }
878
879    /// Every document, in insertion order.
880    pub fn iter(&self) -> impl Iterator<Item = (&[u8], Doc<'_>)> {
881        let keys = &self.keys;
882        self.rows.pairs().filter_map(move |(id, bytes)| {
883            let value = Value::new(bytes)?;
884            Some((id, Doc { value, keys }))
885        })
886    }
887
888    /// Walk part of the collection and say where to resume, the same contract
889    /// [`Elements::scan`] has.
890    pub fn scan<F>(&self, cursor: Cursor, count: usize, mut f: F) -> Cursor
891    where
892        F: FnMut(&[u8], Doc<'_>),
893    {
894        let keys = &self.keys;
895        self.rows.scan_pairs(cursor, count, |id, bytes| {
896            if let Some(value) = Value::new(bytes) {
897                f(id, Doc { value, keys });
898            }
899        })
900    }
901
902    /// Throw every document away and keep the key table and the allocations.
903    ///
904    /// The indexes stay declared and go empty, for the same reason the key table
905    /// stays: a caller that empties a collection is refilling it, and an index
906    /// that quietly disappeared when the last document did would turn the next
907    /// query into an error.
908    ///
909    /// The key table stays because a collection that is emptied is usually a
910    /// collection that is about to be refilled with the same shape of document,
911    /// and relearning twenty names is work with nothing to show for it.
912    pub fn clear(&mut self) {
913        self.rows.clear();
914        self.build.clear();
915        for index in &mut self.indexes {
916            index.clear();
917        }
918        for index in &mut self.vectors {
919            index.clear();
920        }
921    }
922
923    /// What the collection costs, the key table and the indexes included.
924    #[must_use]
925    pub fn memory_bytes(&self) -> usize {
926        self.rows.memory_bytes()
927            + self.keys.memory_bytes()
928            + self
929                .indexes
930                .iter()
931                .map(PathIndex::memory_bytes)
932                .sum::<usize>()
933            + self
934                .vectors
935                .iter()
936                .map(VectorIndex::memory_bytes)
937                .sum::<usize>()
938    }
939}
940
941/// The tag a stored document carries, worked out from the indexes it is filed
942/// in.
943///
944/// Nothing here can fail. A path that no longer resolves or a value that is too
945/// long to be an index key contributes no bit, which is the same absence the
946/// document has in that index.
947fn tag_for(doc: &Doc<'_>, indexes: &[PathIndex], list: &mut Vec<u8>) -> u64 {
948    let mut sig = Signature::default();
949    for index in indexes {
950        let Ok(Some(at)) = doc.path_bytes(index.path()) else {
951            continue;
952        };
953        list.clear();
954        let _ = index.keys_at(at.value(), list);
955        vector::add_keys(&mut sig, index.path(), list);
956    }
957    sig.bits()
958}
959
960/// Take whatever is stored under `id` out of every index, leaving the primary
961/// table alone.
962///
963/// Both an overwrite and a removal go through here, because both of them make
964/// the old document's index entries wrong and neither of them can work out what
965/// those entries were once the bytes are gone. Nothing here can fail: a document
966/// that is no longer readable, or a path that no longer resolves, simply has
967/// nothing filed under it, and refusing a removal because the thing being
968/// removed is damaged is the wrong answer.
969fn unindex(rows: &Elements<()>, keys: &Keys, indexes: &mut [PathIndex], id: &[u8]) {
970    if indexes.is_empty() {
971        return;
972    }
973    let Some(bytes) = rows.tail(id) else {
974        return;
975    };
976    let Some(value) = Value::new(bytes) else {
977        return;
978    };
979    let doc = Doc { value, keys };
980    let mut list = Vec::new();
981    for index in indexes {
982        let Ok(Some(at)) = doc.path_bytes(index.path()) else {
983            continue;
984        };
985        list.clear();
986        // The keys came out of this same code on the way in, so a key that was
987        // refused then is not filed now and there is nothing to take out.
988        let _ = index.keys_at(at.value(), &mut list);
989        index::each_key(&list, |key| index.take(key, id));
990    }
991}
992
993/// Put `bytes` in the primary table under `id`, turning a refusal into the
994/// error the layer above would have written anyway.
995fn store(rows: &mut Elements<()>, id: &[u8], bytes: &[u8]) -> Result<bool> {
996    match rows.set_tailed(id, bytes, ()) {
997        Ok((_, fresh)) => Ok(fresh),
998        Err(Full::Name) => Err(Error::fmt(
999            Code::Full,
1000            format_args!("a document id is at most {} bytes", yo_kv::NAME_MAX),
1001        )),
1002        Err(Full::Rows) => Err(Error::fmt(
1003            Code::Full,
1004            format_args!("a collection holds at most {} documents", yo_kv::MAX_ROWS),
1005        )),
1006    }
1007}
1008
1009/// Write `value` into `b` with every object key replaced by its id.
1010///
1011/// `Ok(false)` means the key table ran out of ids part way through, and the
1012/// caller stores the document with its keys as bytes instead. The builder is
1013/// left half open in that case, so the caller clears it.
1014///
1015/// The recursion is bounded by the builder: it refuses to open a container more
1016/// than [`DEPTH_MAX`] deep, so a document nested deeper than that, which only a
1017/// damaged one can be, stops with an error rather than with the stack.
1018fn intern_into(keys: &mut Keys, b: &mut Builder, value: Value<'_>, depth: usize) -> Result<bool> {
1019    let corrupt = || Error::new(Code::Corrupt, "the document is not readable at that point");
1020    match value.kind() {
1021        Kind::Null => b.null()?,
1022        Kind::Bool => b.bool(value.as_bool().ok_or_else(corrupt)?)?,
1023        Kind::Int => b.int(value.as_int().ok_or_else(corrupt)?)?,
1024        Kind::Float => b.float(value.as_float().ok_or_else(corrupt)?)?,
1025        Kind::Text => b.text_bytes(value.text_bytes().ok_or_else(corrupt)?)?,
1026        Kind::Array => {
1027            b.begin_array()?;
1028            for i in 0..value.len() {
1029                let child = value.at(i).ok_or_else(corrupt)?;
1030                if !intern_into(keys, b, child, depth + 1)? {
1031                    return Ok(false);
1032                }
1033            }
1034            b.end_array()?;
1035        }
1036        Kind::Object => {
1037            if value.is_interned() {
1038                return Err(Error::new(
1039                    Code::Invalid,
1040                    "this document's keys are ids from another collection's key table",
1041                ));
1042            }
1043            b.begin_object_interned()?;
1044            for i in 0..value.len() {
1045                let name = value.key_at(i).ok_or_else(corrupt)?;
1046                let Some(id) = keys.intern(name) else {
1047                    return Ok(false);
1048                };
1049                b.key_id(id)?;
1050                let child = value.at(i).ok_or_else(corrupt)?;
1051                if !intern_into(keys, b, child, depth + 1)? {
1052                    return Ok(false);
1053                }
1054            }
1055            b.end_object()?;
1056        }
1057    }
1058    debug_assert!(depth <= DEPTH_MAX, "the builder caps the depth");
1059    Ok(true)
1060}
1061
1062/// A value with the key table its keys are interned against.
1063///
1064/// Everything a [`Value`] offers is here too, and the things that need a name
1065/// rather than an id, which are a lookup, a walk over the members and printing
1066/// the thing, go through the table. A document whose keys are bytes works the
1067/// same way and simply never asks the table anything, so a caller does not have
1068/// to know which kind it is holding.
1069#[derive(Clone, Copy)]
1070pub struct Doc<'a> {
1071    value: Value<'a>,
1072    keys: &'a Keys,
1073}
1074
1075impl<'a> Doc<'a> {
1076    /// A view of `value` against `keys`.
1077    #[must_use]
1078    pub fn new(value: Value<'a>, keys: &'a Keys) -> Doc<'a> {
1079        Doc { value, keys }
1080    }
1081
1082    /// The value underneath, for the accessors that never need a name.
1083    #[must_use]
1084    pub fn value(&self) -> Value<'a> {
1085        self.value
1086    }
1087
1088    /// The key table this reads names out of.
1089    #[must_use]
1090    pub fn keys(&self) -> &'a Keys {
1091        self.keys
1092    }
1093
1094    /// What this value is.
1095    #[must_use]
1096    pub fn kind(&self) -> Kind {
1097        self.value.kind()
1098    }
1099
1100    /// Whether this is `null`.
1101    #[must_use]
1102    pub fn is_null(&self) -> bool {
1103        self.value.is_null()
1104    }
1105
1106    /// The boolean this holds, if it holds one.
1107    #[must_use]
1108    pub fn as_bool(&self) -> Option<bool> {
1109        self.value.as_bool()
1110    }
1111
1112    /// The integer this holds, if it holds one.
1113    #[must_use]
1114    pub fn as_int(&self) -> Option<i64> {
1115        self.value.as_int()
1116    }
1117
1118    /// The float this holds, if it holds one.
1119    #[must_use]
1120    pub fn as_float(&self) -> Option<f64> {
1121        self.value.as_float()
1122    }
1123
1124    /// The string this holds, if it holds one and it is UTF-8.
1125    #[must_use]
1126    pub fn as_text(&self) -> Option<&'a str> {
1127        self.value.as_text()
1128    }
1129
1130    /// The string this holds as it is stored, without the UTF-8 check.
1131    #[must_use]
1132    pub fn text_bytes(&self) -> Option<&'a [u8]> {
1133        self.value.text_bytes()
1134    }
1135
1136    /// How many elements a container holds. Zero for anything else.
1137    #[must_use]
1138    pub fn len(&self) -> usize {
1139        self.value.len()
1140    }
1141
1142    /// Whether this is a container with nothing in it.
1143    #[must_use]
1144    pub fn is_empty(&self) -> bool {
1145        self.value.is_empty()
1146    }
1147
1148    /// The value stored under `key`.
1149    ///
1150    /// For an interned object this is a name to id lookup in the table and then
1151    /// a binary search over integers. A name the table has never seen cannot be
1152    /// in the document, so it answers `None` without touching the document at
1153    /// all.
1154    #[must_use]
1155    pub fn get(&self, key: &[u8]) -> Option<Doc<'a>> {
1156        let value = if self.value.is_interned() {
1157            self.value.get_id(self.keys.id(key)?)?
1158        } else {
1159            self.value.get(key)?
1160        };
1161        Some(Doc {
1162            value,
1163            keys: self.keys,
1164        })
1165    }
1166
1167    /// Element `i` of a container, in the container's own order.
1168    #[must_use]
1169    pub fn at(&self, i: usize) -> Option<Doc<'a>> {
1170        Some(Doc {
1171            value: self.value.at(i)?,
1172            keys: self.keys,
1173        })
1174    }
1175
1176    /// The name of member `i` of an object, whichever way the keys are stored.
1177    #[must_use]
1178    pub fn key_at(&self, i: usize) -> Option<&'a [u8]> {
1179        if self.value.is_interned() {
1180            self.keys.name(self.value.key_id_at(i)?)
1181        } else {
1182            self.value.key_at(i)
1183        }
1184    }
1185
1186    /// Every member of an object, name first, in the order the document stores
1187    /// them.
1188    ///
1189    /// That order is by key id for an interned object and by key bytes for one
1190    /// whose keys are bytes, so it is stable for a given collection and it is
1191    /// not alphabetical. Sort it if the order is part of the answer.
1192    #[must_use]
1193    pub fn members(&self) -> DocMembers<'a> {
1194        DocMembers { d: *self, i: 0 }
1195    }
1196
1197    /// Every element of a container, in the container's own order.
1198    #[must_use]
1199    pub fn iter(&self) -> DocElems<'a> {
1200        DocElems { d: *self, i: 0 }
1201    }
1202
1203    /// The value at `path`, where a path names exactly one place.
1204    ///
1205    /// The same grammar [`Value::path`] takes, with the names resolved through
1206    /// the key table on the way down.
1207    pub fn path(&self, path: &str) -> Result<Option<Doc<'a>>> {
1208        self.path_bytes(path.as_bytes())
1209    }
1210
1211    /// [`Doc::path`] for a path that is already bytes.
1212    pub fn path_bytes(&self, path: &[u8]) -> Result<Option<Doc<'a>>> {
1213        let mut at = *self;
1214        for step in Steps::new(path) {
1215            let next = match step? {
1216                Step::Key(k) => at.get(k),
1217                Step::Index(_) if at.kind() != Kind::Array => None,
1218                Step::Index(i) => {
1219                    let n = at.len();
1220                    let i = if i < 0 {
1221                        match n.checked_sub(i.unsigned_abs() as usize) {
1222                            Some(i) => i,
1223                            None => return Ok(None),
1224                        }
1225                    } else {
1226                        i as usize
1227                    };
1228                    at.at(i)
1229                }
1230            };
1231            let Some(next) = next else {
1232                return Ok(None);
1233            };
1234            at = next;
1235        }
1236        Ok(Some(at))
1237    }
1238}
1239
1240impl core::fmt::Debug for Doc<'_> {
1241    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1242        match self.kind() {
1243            Kind::Object => {
1244                let mut m = f.debug_map();
1245                for (k, v) in self.members() {
1246                    m.entry(&String::from_utf8_lossy(k), &v);
1247                }
1248                m.finish()
1249            }
1250            Kind::Array => f.debug_list().entries(self.iter()).finish(),
1251            _ => self.value.fmt(f),
1252        }
1253    }
1254}
1255
1256/// Every member of an object, from [`Doc::members`].
1257#[derive(Clone)]
1258pub struct DocMembers<'a> {
1259    d: Doc<'a>,
1260    i: usize,
1261}
1262
1263impl<'a> Iterator for DocMembers<'a> {
1264    type Item = (&'a [u8], Doc<'a>);
1265
1266    fn next(&mut self) -> Option<(&'a [u8], Doc<'a>)> {
1267        let key = self.d.key_at(self.i)?;
1268        let val = self.d.at(self.i)?;
1269        self.i += 1;
1270        Some((key, val))
1271    }
1272
1273    fn size_hint(&self) -> (usize, Option<usize>) {
1274        let left = self.d.len().saturating_sub(self.i);
1275        (left, Some(left))
1276    }
1277}
1278
1279/// Every element of a container, from [`Doc::iter`].
1280#[derive(Clone)]
1281pub struct DocElems<'a> {
1282    d: Doc<'a>,
1283    i: usize,
1284}
1285
1286impl<'a> Iterator for DocElems<'a> {
1287    type Item = Doc<'a>;
1288
1289    fn next(&mut self) -> Option<Doc<'a>> {
1290        let out = self.d.at(self.i)?;
1291        self.i += 1;
1292        Some(out)
1293    }
1294
1295    fn size_hint(&self) -> (usize, Option<usize>) {
1296        let left = self.d.len().saturating_sub(self.i);
1297        (left, Some(left))
1298    }
1299}
1300
1301#[cfg(test)]
1302mod tests {
1303    use super::*;
1304
1305    /// An order, the shape `09` section 5 uses as its example.
1306    fn order(id: i64, status: &str, lines: usize) -> Vec<u8> {
1307        let mut b = Builder::new();
1308        b.begin_object().expect("open");
1309        b.key(b"id").expect("key");
1310        b.int(id).expect("value");
1311        b.key(b"customer").expect("key");
1312        b.int(id * 7).expect("value");
1313        b.key(b"status").expect("key");
1314        b.text(status).expect("value");
1315        b.key(b"lines").expect("key");
1316        b.begin_array().expect("open");
1317        for i in 0..lines {
1318            b.begin_object().expect("open");
1319            b.key(b"sku").expect("key");
1320            b.text(&format!("sku-{i}")).expect("value");
1321            b.key(b"qty").expect("key");
1322            b.int(i as i64 + 1).expect("value");
1323            b.end_object().expect("close");
1324        }
1325        b.end_array().expect("close");
1326        b.end_object().expect("close");
1327        b.finish().expect("finished").to_vec()
1328    }
1329
1330    #[test]
1331    fn a_document_reads_back_the_way_it_went_in() {
1332        let mut docs = Docs::new();
1333        assert!(
1334            docs.put_bytes(b"order:1", &order(1, "open", 3))
1335                .expect("put")
1336        );
1337        assert!(
1338            !docs
1339                .put_bytes(b"order:1", &order(1, "shut", 3))
1340                .expect("put")
1341        );
1342        assert_eq!(docs.len(), 1);
1343
1344        let d = docs.get(b"order:1").expect("stored");
1345        assert_eq!(d.get(b"id").and_then(|v| v.as_int()), Some(1));
1346        assert_eq!(d.get(b"status").and_then(|v| v.as_text()), Some("shut"));
1347        assert_eq!(d.get(b"lines").map(|v| v.len()), Some(3));
1348        assert_eq!(
1349            d.path("$.lines[1].sku")
1350                .expect("a path")
1351                .and_then(|v| v.as_text()),
1352            Some("sku-1")
1353        );
1354        assert_eq!(
1355            d.path("$.lines[-1].qty")
1356                .expect("a path")
1357                .and_then(|v| v.as_int()),
1358            Some(3)
1359        );
1360        assert!(d.get(b"missing").is_none());
1361    }
1362
1363    #[test]
1364    fn the_keys_are_interned_and_the_names_come_back() {
1365        let mut docs = Docs::new();
1366        docs.put_bytes(b"order:1", &order(1, "open", 2))
1367            .expect("put");
1368        let names: Vec<String> = docs
1369            .keys()
1370            .iter()
1371            .map(|(n, _)| String::from_utf8_lossy(n).into_owned())
1372            .collect();
1373        names.iter().for_each(|n| assert!(!n.is_empty()));
1374        assert_eq!(
1375            docs.keys().len(),
1376            6,
1377            "id customer status lines sku qty: {names:?}"
1378        );
1379
1380        let d = docs.get(b"order:1").expect("stored");
1381        assert!(d.value().is_interned());
1382        let mut got: Vec<&[u8]> = d.members().map(|(k, _)| k).collect();
1383        got.sort_unstable();
1384        assert_eq!(got, [&b"customer"[..], b"id", b"lines", b"status"]);
1385        let line = d.path("$.lines[0]").expect("a path").expect("there");
1386        assert!(line.value().is_interned());
1387        let mut inner: Vec<&[u8]> = line.members().map(|(k, _)| k).collect();
1388        inner.sort_unstable();
1389        assert_eq!(inner, [&b"qty"[..], b"sku"]);
1390    }
1391
1392    /// Store 256 copies of a shape and say what fraction of the bytes survived.
1393    fn shrinkage(shape: impl Fn(i64) -> Vec<u8>) -> f64 {
1394        let mut docs = Docs::new();
1395        let mut plain = 0usize;
1396        for i in 0..256i64 {
1397            let bytes = shape(i);
1398            plain += bytes.len();
1399            docs.put_bytes(format!("d:{i}").as_bytes(), &bytes)
1400                .expect("put");
1401        }
1402        let stored: usize = (0..256i64)
1403            .map(|i| {
1404                docs.bytes(format!("d:{i}").as_bytes())
1405                    .expect("stored")
1406                    .len()
1407            })
1408            .sum();
1409        stored as f64 / plain as f64
1410    }
1411
1412    #[test]
1413    fn interning_makes_a_collection_of_the_same_shape_smaller() {
1414        // The claim in `09` section 4 is that the same field names on every
1415        // document are most of what a document collection costs, and that
1416        // interning them is worth about forty percent. How much it is actually
1417        // worth depends on how much of a document is names, so both ends are
1418        // measured here rather than one number being asserted twice.
1419        //
1420        // A document that is mostly names, which is what a typed collection of
1421        // small records looks like, keeps a little over half its bytes.
1422        let names = shrinkage(|i| {
1423            let mut b = Builder::new();
1424            b.begin_object().expect("open");
1425            for f in 0..20 {
1426                b.key(format!("some_field_name_{f:02}").as_bytes())
1427                    .expect("key");
1428                b.int(i + f).expect("value");
1429            }
1430            b.end_object().expect("close");
1431            b.finish().expect("finished").to_vec()
1432        });
1433        assert!(names < 0.60, "a document of names kept {names}");
1434
1435        // An order, which carries real payload as well, keeps about three
1436        // quarters. That is the honest floor for the claim and it is still a
1437        // fifth of the collection gone for nothing but a table of twenty
1438        // strings.
1439        let orders = shrinkage(|i| order(i, "open", 2));
1440        assert!(orders < 0.80, "an order collection kept {orders}");
1441    }
1442
1443    #[test]
1444    fn a_document_whose_keys_are_already_ids_is_refused() {
1445        let mut b = Builder::new();
1446        b.begin_object_interned().expect("open");
1447        b.key_id(0).expect("key");
1448        b.int(1).expect("value");
1449        b.end_object().expect("close");
1450        let bytes = b.finish().expect("finished").to_vec();
1451
1452        let mut docs = Docs::new();
1453        let err = docs.put_bytes(b"x", &bytes).expect_err("refused");
1454        assert_eq!(err.code(), Code::Invalid);
1455    }
1456
1457    #[test]
1458    fn a_document_that_is_not_readable_is_refused() {
1459        let mut docs = Docs::new();
1460        let err = docs.put_bytes(b"x", &[2, 0, 0, 0]).expect_err("refused");
1461        assert_eq!(err.code(), Code::Corrupt);
1462        assert!(docs.is_empty());
1463    }
1464
1465    #[test]
1466    fn a_removal_leaves_every_other_document_where_it_was() {
1467        let mut docs = Docs::new();
1468        for i in 0..64i64 {
1469            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1470                .expect("put");
1471        }
1472        for i in (0..64i64).step_by(3) {
1473            assert!(docs.remove(format!("order:{i}").as_bytes()));
1474        }
1475        assert_eq!(docs.len(), 64 - 22);
1476        for i in 0..64i64 {
1477            let id = format!("order:{i}");
1478            match docs.get(id.as_bytes()) {
1479                Some(d) => {
1480                    assert!(i % 3 != 0, "{id} was removed");
1481                    assert_eq!(d.get(b"id").and_then(|v| v.as_int()), Some(i));
1482                }
1483                None => assert!(i % 3 == 0, "{id} was not removed"),
1484            }
1485        }
1486        assert_eq!(docs.keys().len(), 6, "a removal does not un-intern a name");
1487    }
1488
1489    #[test]
1490    fn a_walk_sees_every_document_once() {
1491        let mut docs = Docs::new();
1492        for i in 0..200i64 {
1493            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1494                .expect("put");
1495        }
1496
1497        let mut seen: Vec<i64> = docs
1498            .iter()
1499            .map(|(_, d)| d.get(b"id").and_then(|v| v.as_int()).expect("an id"))
1500            .collect();
1501        seen.sort_unstable();
1502        assert_eq!(seen, (0..200).collect::<Vec<i64>>());
1503
1504        let mut scanned = Vec::new();
1505        let mut cursor = Cursor::START;
1506        loop {
1507            cursor = docs.scan(cursor, 16, |id, _| scanned.push(id.to_vec()));
1508            if cursor.is_end() {
1509                break;
1510            }
1511        }
1512        scanned.sort_unstable();
1513        scanned.dedup();
1514        assert_eq!(scanned.len(), 200);
1515    }
1516
1517    #[test]
1518    fn an_empty_collection_answers_nothing_rather_than_failing() {
1519        let docs = Docs::new();
1520        assert!(docs.is_empty());
1521        assert!(docs.get(b"nothing").is_none());
1522        assert!(docs.bytes(b"nothing").is_none());
1523        assert!(!docs.contains(b"nothing"));
1524        assert_eq!(docs.iter().count(), 0);
1525    }
1526
1527    #[test]
1528    fn a_document_prints_with_its_names_back_on() {
1529        let mut docs = Docs::new();
1530        docs.put_bytes(b"order:1", &order(1, "open", 1))
1531            .expect("put");
1532        let text = format!("{:?}", docs.get(b"order:1").expect("stored"));
1533        assert!(text.contains("\"status\": \"open\""), "{text}");
1534        assert!(text.contains("\"sku\": \"sku-0\""), "{text}");
1535    }
1536
1537    /// The ids `find` answers for one key, sorted so a test can compare them.
1538    fn found(docs: &Docs, path: &str, key: &Key) -> Vec<String> {
1539        let mut out = Vec::new();
1540        let n = docs
1541            .find(path, key, |id, d| {
1542                assert!(!d.is_empty(), "the document came back whole");
1543                out.push(String::from_utf8_lossy(id).into_owned());
1544            })
1545            .expect("indexed");
1546        assert_eq!(n, out.len(), "the count is what the callback saw");
1547        out.sort();
1548        out
1549    }
1550
1551    #[test]
1552    fn an_index_declared_after_the_documents_finds_them() {
1553        let mut docs = Docs::new();
1554        for i in 0..64i64 {
1555            let status = if i % 4 == 0 { "shut" } else { "open" };
1556            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, status, 1))
1557                .expect("put");
1558        }
1559        docs.create_index("$.status").expect("indexed");
1560        assert_eq!(docs.index("$.status").expect("there").len(), 2);
1561        assert_eq!(docs.count("$.status", &Key::text("shut")).expect("i"), 16);
1562        assert_eq!(docs.count("$.status", &Key::text("open")).expect("i"), 48);
1563        assert_eq!(found(&docs, "$.status", &Key::text("shut")).len(), 16);
1564        assert!(found(&docs, "$.status", &Key::text("gone")).is_empty());
1565
1566        // A document written after the index exists is filed by the write.
1567        docs.put_bytes(b"order:64", &order(64, "shut", 1))
1568            .expect("put");
1569        assert_eq!(docs.count("$.status", &Key::text("shut")).expect("i"), 17);
1570    }
1571
1572    #[test]
1573    fn an_overwrite_moves_a_document_from_one_key_to_the_other() {
1574        let mut docs = Docs::new();
1575        docs.create_index("$.status").expect("indexed");
1576        docs.put_bytes(b"order:1", &order(1, "open", 1))
1577            .expect("put");
1578        assert_eq!(found(&docs, "$.status", &Key::text("open")), ["order:1"]);
1579
1580        docs.put_bytes(b"order:1", &order(1, "shut", 1))
1581            .expect("put");
1582        assert!(
1583            found(&docs, "$.status", &Key::text("open")).is_empty(),
1584            "the old key kept it"
1585        );
1586        assert_eq!(found(&docs, "$.status", &Key::text("shut")), ["order:1"]);
1587        assert_eq!(docs.index("$.status").expect("there").postings(), 1);
1588    }
1589
1590    #[test]
1591    fn a_removal_takes_a_document_out_of_every_index() {
1592        let mut docs = Docs::new();
1593        docs.create_index("$.status").expect("indexed");
1594        docs.create_index("$.customer").expect("indexed");
1595        for i in 0..8i64 {
1596            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1597                .expect("put");
1598        }
1599        assert!(docs.remove(b"order:3"));
1600        assert_eq!(found(&docs, "$.status", &Key::text("open")).len(), 7);
1601        assert_eq!(docs.count("$.customer", &Key::int(21)).expect("i"), 0);
1602        assert_eq!(docs.count("$.customer", &Key::int(28)).expect("i"), 1);
1603        for index in docs.indexes() {
1604            assert_eq!(index.postings(), 7);
1605        }
1606
1607        assert!(!docs.remove(b"order:3"), "it is already gone");
1608        assert_eq!(docs.index("$.status").expect("there").postings(), 7);
1609    }
1610
1611    #[test]
1612    fn a_path_that_names_a_container_or_nothing_is_simply_not_filed() {
1613        let mut docs = Docs::new();
1614        docs.create_index("$.lines").expect("indexed");
1615        docs.create_index("$.shipped").expect("indexed");
1616        docs.create_index("$.lines[0].qty").expect("indexed");
1617        for i in 0..4i64 {
1618            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 2))
1619                .expect("put");
1620        }
1621        assert_eq!(docs.len(), 4);
1622        assert!(
1623            docs.index("$.lines").expect("there").is_empty(),
1624            "an array has no equality key"
1625        );
1626        assert!(
1627            docs.index("$.shipped").expect("there").is_empty(),
1628            "no document has that path"
1629        );
1630        assert_eq!(
1631            docs.count("$.lines[0].qty", &Key::int(1)).expect("i"),
1632            4,
1633            "a path through an array reaches a scalar"
1634        );
1635    }
1636
1637    #[test]
1638    fn a_value_too_long_to_index_fails_the_write_and_stores_nothing() {
1639        let mut b = Builder::new();
1640        b.begin_object().expect("open");
1641        b.key(b"status").expect("key");
1642        b.text(&"x".repeat(crate::KEY_MAX)).expect("value");
1643        b.end_object().expect("close");
1644        let huge = b.finish().expect("finished").to_vec();
1645
1646        let mut docs = Docs::new();
1647        docs.create_index("$.status").expect("indexed");
1648        let err = docs.put_bytes(b"order:1", &huge).expect_err("refused");
1649        assert_eq!(err.code(), Code::Full);
1650        assert!(
1651            docs.is_empty(),
1652            "a write that cannot be indexed leaves nothing behind"
1653        );
1654
1655        // Without the index it is an ordinary document and goes in fine.
1656        assert!(docs.drop_index("$.status"));
1657        docs.put_bytes(b"order:1", &huge).expect("put");
1658        assert_eq!(docs.len(), 1);
1659    }
1660
1661    #[test]
1662    fn a_query_on_a_path_with_no_index_says_so_rather_than_scanning() {
1663        let mut docs = Docs::new();
1664        docs.put_bytes(b"order:1", &order(1, "open", 1))
1665            .expect("put");
1666        let err = docs
1667            .find("$.status", &Key::text("open"), |_, _| ())
1668            .expect_err("refused");
1669        assert_eq!(err.code(), Code::Invalid);
1670        assert_eq!(
1671            docs.count("$.status", &Key::text("open"))
1672                .expect_err("refused")
1673                .code(),
1674            Code::Invalid
1675        );
1676        assert!(docs.index("$.status").is_none());
1677        assert!(!docs.drop_index("$.status"));
1678    }
1679
1680    #[test]
1681    fn declaring_the_same_index_twice_leaves_the_first_one_alone() {
1682        let mut docs = Docs::new();
1683        docs.create_index("$.status").expect("indexed");
1684        docs.put_bytes(b"order:1", &order(1, "open", 1))
1685            .expect("put");
1686        docs.create_index("$.status").expect("indexed again");
1687        assert_eq!(docs.indexes().len(), 1);
1688        assert_eq!(
1689            docs.index("$.status").expect("there").postings(),
1690            1,
1691            "a redeclaration did not double file anything"
1692        );
1693        assert!(docs.create_index("$.[").is_err(), "the path has to parse");
1694    }
1695
1696    #[test]
1697    fn clearing_a_collection_empties_its_indexes_and_keeps_them() {
1698        let mut docs = Docs::new();
1699        docs.create_index("$.status").expect("indexed");
1700        for i in 0..8i64 {
1701            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1702                .expect("put");
1703        }
1704        docs.clear();
1705        assert!(docs.is_empty());
1706        assert!(docs.index("$.status").expect("still declared").is_empty());
1707        assert_eq!(docs.count("$.status", &Key::text("open")).expect("i"), 0);
1708
1709        docs.put_bytes(b"order:9", &order(9, "open", 1))
1710            .expect("put");
1711        assert_eq!(found(&docs, "$.status", &Key::text("open")), ["order:9"]);
1712    }
1713
1714    #[test]
1715    fn two_indexes_intersect_as_the_sets_they_are() {
1716        let mut docs = Docs::new();
1717        docs.create_index("$.status").expect("indexed");
1718        docs.create_index("$.customer").expect("indexed");
1719        for i in 0..32i64 {
1720            let status = if i % 2 == 0 { "open" } else { "shut" };
1721            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i % 4, status, 1))
1722                .expect("put");
1723        }
1724
1725        // What a planner does: probe both, walk the smaller, ask the larger.
1726        // That is `SINTER` and there is no code here that is not already the
1727        // set's.
1728        let open = Key::text("open");
1729        let customer = Key::int(14);
1730        let small = docs.count("$.customer", &customer).expect("indexed");
1731        let large = docs.count("$.status", &open).expect("indexed");
1732        assert_eq!((small, large), (8, 16));
1733
1734        let small = docs.index("$.customer").expect("there").get(&customer);
1735        let large = docs.index("$.status").expect("there").get(&open);
1736        let (Some(small), Some(large)) = (small, large) else {
1737            panic!("both keys are filed");
1738        };
1739        let mut both = Vec::new();
1740        index::each_id(small, |id| {
1741            if large.contains(id) {
1742                both.push(String::from_utf8_lossy(id).into_owned());
1743            }
1744        });
1745        both.sort();
1746        assert_eq!(
1747            both,
1748            [
1749                "order:10", "order:14", "order:18", "order:2", "order:22", "order:26", "order:30",
1750                "order:6"
1751            ]
1752        );
1753    }
1754
1755    /// The customer numbers a range answers, in the order it answered them.
1756    fn ranged(docs: &Docs, lo: Bound<&Key>, hi: Bound<&Key>) -> Vec<i64> {
1757        let mut out = Vec::new();
1758        let n = docs
1759            .range("$.customer", lo, hi, |_, d| {
1760                out.push(d.get(b"customer").and_then(|v| v.as_int()).expect("there"));
1761            })
1762            .expect("ordered");
1763        assert_eq!(n, out.len());
1764
1765        let mut back = Vec::new();
1766        docs.range_rev("$.customer", lo, hi, |_, d| {
1767            back.push(d.get(b"customer").and_then(|v| v.as_int()).expect("there"));
1768        })
1769        .expect("ordered");
1770        back.reverse();
1771        assert_eq!(out, back, "backwards is forwards read the other way");
1772        assert_eq!(
1773            docs.count_range("$.customer", lo, hi).expect("ordered"),
1774            out.len()
1775        );
1776        out
1777    }
1778
1779    #[test]
1780    fn an_ordered_index_answers_a_range_in_order() {
1781        let mut docs = Docs::new();
1782        docs.create_ordered_index("$.customer").expect("ordered");
1783        // Customer is seven times the id, so the values are 0, 7, 14 and on.
1784        for i in 0..64i64 {
1785            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1786                .expect("put");
1787        }
1788
1789        assert_eq!(
1790            ranged(&docs, Bound::Unbounded, Bound::Unbounded),
1791            (0..64i64).map(|i| i * 7).collect::<Vec<i64>>()
1792        );
1793        let (lo, hi) = (Key::int(70), Key::int(105));
1794        assert_eq!(
1795            ranged(&docs, Bound::Included(&lo), Bound::Included(&hi)),
1796            [70, 77, 84, 91, 98, 105]
1797        );
1798        assert_eq!(
1799            ranged(&docs, Bound::Excluded(&lo), Bound::Excluded(&hi)),
1800            [77, 84, 91, 98]
1801        );
1802        // Bounds that fall between two values, which is the ordinary case.
1803        assert_eq!(
1804            ranged(
1805                &docs,
1806                Bound::Included(&Key::int(71)),
1807                Bound::Excluded(&Key::int(90))
1808            ),
1809            [77, 84]
1810        );
1811        assert!(ranged(&docs, Bound::Included(&Key::int(442)), Bound::Unbounded).is_empty());
1812
1813        // Equality still works on the same index.
1814        assert_eq!(docs.count("$.customer", &Key::int(70)).expect("i"), 1);
1815        assert_eq!(
1816            docs.index("$.customer").expect("there").kind(),
1817            IndexKind::Ordered
1818        );
1819    }
1820
1821    #[test]
1822    fn a_range_stays_right_through_writes_and_removals() {
1823        let mut docs = Docs::new();
1824        for i in 0..128i64 {
1825            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1826                .expect("put");
1827        }
1828        // Declared after the fact, so this is the backfill and not the write
1829        // path putting the tree together.
1830        docs.create_ordered_index("$.customer").expect("ordered");
1831        assert_eq!(ranged(&docs, Bound::Unbounded, Bound::Unbounded).len(), 128);
1832
1833        // Every removal moves the key table's last row into the hole, so this is
1834        // the renumbering going through the whole collection.
1835        for i in (0..128i64).step_by(2) {
1836            assert!(docs.remove(format!("order:{i}").as_bytes()));
1837        }
1838        assert_eq!(
1839            ranged(&docs, Bound::Unbounded, Bound::Unbounded),
1840            (0..128i64)
1841                .filter(|i| i % 2 == 1)
1842                .map(|i| i * 7)
1843                .collect::<Vec<i64>>()
1844        );
1845
1846        // And an overwrite that moves a document from one key to another.
1847        docs.put_bytes(b"order:1", &order(200, "open", 1))
1848            .expect("put");
1849        let after = ranged(&docs, Bound::Unbounded, Bound::Unbounded);
1850        assert_eq!(after.first(), Some(&21), "seven is gone");
1851        assert_eq!(after.last(), Some(&1400), "and it came back at the top");
1852    }
1853
1854    #[test]
1855    fn an_equality_index_refuses_a_range_rather_than_answering_nothing() {
1856        let mut docs = Docs::new();
1857        docs.create_index("$.customer").expect("indexed");
1858        docs.put_bytes(b"order:1", &order(1, "open", 1))
1859            .expect("put");
1860        let err = docs
1861            .range("$.customer", Bound::Unbounded, Bound::Unbounded, |_, _| ())
1862            .expect_err("refused");
1863        assert_eq!(err.code(), Code::Invalid);
1864        assert!(err.to_string().contains("equality"), "{err}");
1865        assert_eq!(
1866            docs.range("$.status", Bound::Unbounded, Bound::Unbounded, |_, _| ())
1867                .expect_err("refused")
1868                .code(),
1869            Code::Invalid
1870        );
1871    }
1872
1873    #[test]
1874    fn asking_for_an_order_on_an_equality_index_upgrades_it() {
1875        let mut docs = Docs::new();
1876        docs.create_index("$.customer").expect("indexed");
1877        for i in 0..8i64 {
1878            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1879                .expect("put");
1880        }
1881        assert_eq!(
1882            docs.index("$.customer").expect("there").kind(),
1883            IndexKind::Equality
1884        );
1885
1886        docs.create_ordered_index("$.customer").expect("upgraded");
1887        assert_eq!(docs.indexes().len(), 1, "it replaced rather than added");
1888        assert_eq!(ranged(&docs, Bound::Unbounded, Bound::Unbounded).len(), 8);
1889
1890        // And going the other way leaves the order alone, because an ordered
1891        // index answers equality too.
1892        docs.create_index("$.customer").expect("already there");
1893        assert_eq!(
1894            docs.index("$.customer").expect("there").kind(),
1895            IndexKind::Ordered
1896        );
1897        assert_eq!(docs.indexes().len(), 1);
1898    }
1899
1900    /// A document with a list of tags at `$.tags` and a title at `$.title`.
1901    fn tagged(title: &str, tags: &[&str]) -> Vec<u8> {
1902        let mut b = Builder::new();
1903        b.begin_object().expect("open");
1904        b.key(b"title").expect("key");
1905        b.text(title).expect("value");
1906        b.key(b"tags").expect("key");
1907        b.begin_array().expect("open");
1908        for tag in tags {
1909            b.text(tag).expect("value");
1910        }
1911        b.end_array().expect("close");
1912        b.end_object().expect("close");
1913        b.finish().expect("finished").to_vec()
1914    }
1915
1916    #[test]
1917    fn an_array_index_files_a_document_under_every_element() {
1918        let mut docs = Docs::new();
1919        docs.create_array_index("$.tags").expect("indexed");
1920        docs.put_bytes(b"a", &tagged("one", &["red", "blue"]))
1921            .expect("put");
1922        docs.put_bytes(b"b", &tagged("two", &["blue", "green"]))
1923            .expect("put");
1924        docs.put_bytes(b"c", &tagged("three", &[])).expect("put");
1925
1926        assert_eq!(found(&docs, "$.tags", &Key::text("red")), ["a"]);
1927        assert_eq!(found(&docs, "$.tags", &Key::text("blue")), ["a", "b"]);
1928        assert_eq!(found(&docs, "$.tags", &Key::text("green")), ["b"]);
1929        assert!(found(&docs, "$.tags", &Key::text("puce")).is_empty());
1930        assert_eq!(
1931            docs.index("$.tags").expect("there").len(),
1932            3,
1933            "three distinct tags over two documents"
1934        );
1935    }
1936
1937    #[test]
1938    fn an_array_index_takes_every_element_back_out_again() {
1939        let mut docs = Docs::new();
1940        docs.create_array_index("$.tags").expect("indexed");
1941        docs.put_bytes(b"a", &tagged("one", &["red", "blue"]))
1942            .expect("put");
1943        docs.put_bytes(b"b", &tagged("two", &["blue"]))
1944            .expect("put");
1945
1946        // An overwrite drops one tag and gains another.
1947        docs.put_bytes(b"a", &tagged("one", &["blue", "green"]))
1948            .expect("put");
1949        assert!(found(&docs, "$.tags", &Key::text("red")).is_empty());
1950        assert_eq!(found(&docs, "$.tags", &Key::text("blue")), ["a", "b"]);
1951        assert_eq!(found(&docs, "$.tags", &Key::text("green")), ["a"]);
1952
1953        assert!(docs.remove(b"a"));
1954        assert_eq!(found(&docs, "$.tags", &Key::text("blue")), ["b"]);
1955        assert!(found(&docs, "$.tags", &Key::text("green")).is_empty());
1956        assert_eq!(
1957            docs.index("$.tags").expect("there").len(),
1958            1,
1959            "a tag nobody has left is not a key any more"
1960        );
1961    }
1962
1963    #[test]
1964    fn an_array_index_treats_one_value_as_a_list_of_one() {
1965        let mut docs = Docs::new();
1966        docs.create_array_index("$.status").expect("indexed");
1967        docs.put_bytes(b"order:1", &order(1, "open", 1))
1968            .expect("put");
1969        assert_eq!(found(&docs, "$.status", &Key::text("open")), ["order:1"]);
1970    }
1971
1972    #[test]
1973    fn the_same_element_twice_is_one_posting() {
1974        let mut docs = Docs::new();
1975        docs.create_array_index("$.tags").expect("indexed");
1976        docs.put_bytes(b"a", &tagged("one", &["red", "red", "red"]))
1977            .expect("put");
1978        assert_eq!(found(&docs, "$.tags", &Key::text("red")), ["a"]);
1979        assert_eq!(docs.index("$.tags").expect("there").postings(), 1);
1980
1981        // And taking it out once takes it out, rather than three times over.
1982        assert!(docs.remove(b"a"));
1983        assert_eq!(docs.index("$.tags").expect("there").postings(), 0);
1984        assert!(docs.index("$.tags").expect("there").is_empty());
1985    }
1986
1987    #[test]
1988    fn a_text_index_files_a_document_under_every_word() {
1989        let mut docs = Docs::new();
1990        docs.create_text_index("$.title").expect("indexed");
1991        docs.put_bytes(b"a", &tagged("A red bicycle", &[]))
1992            .expect("put");
1993        docs.put_bytes(b"b", &tagged("The red car, and a bicycle!", &[]))
1994            .expect("put");
1995
1996        assert_eq!(found(&docs, "$.title", &word("bicycle")), ["a", "b"]);
1997        assert_eq!(found(&docs, "$.title", &word("car")), ["b"]);
1998        assert_eq!(
1999            found(&docs, "$.title", &word("RED")),
2000            ["a", "b"],
2001            "a search folds case the same way the write did"
2002        );
2003        assert!(found(&docs, "$.title", &word("lorry")).is_empty());
2004    }
2005
2006    #[test]
2007    fn a_text_index_follows_the_words_through_a_rewrite() {
2008        let mut docs = Docs::new();
2009        docs.create_text_index("$.title").expect("indexed");
2010        docs.put_bytes(b"a", &tagged("a red bicycle", &[]))
2011            .expect("put");
2012        docs.put_bytes(b"a", &tagged("a blue bicycle", &[]))
2013            .expect("put");
2014        assert!(found(&docs, "$.title", &word("red")).is_empty());
2015        assert_eq!(found(&docs, "$.title", &word("blue")), ["a"]);
2016        assert_eq!(found(&docs, "$.title", &word("bicycle")), ["a"]);
2017
2018        assert!(docs.remove(b"a"));
2019        assert!(docs.index("$.title").expect("there").is_empty());
2020    }
2021
2022    #[test]
2023    fn a_text_index_declared_after_the_documents_finds_them() {
2024        let mut docs = Docs::new();
2025        for i in 0..16i64 {
2026            let title = if i % 2 == 0 {
2027                "a red one"
2028            } else {
2029                "a blue one"
2030            };
2031            docs.put_bytes(format!("t:{i}").as_bytes(), &tagged(title, &[]))
2032                .expect("put");
2033        }
2034        docs.create_text_index("$.title").expect("indexed");
2035        assert_eq!(docs.count("$.title", &word("red")).expect("i"), 8);
2036        assert_eq!(docs.count("$.title", &word("one")).expect("i"), 16);
2037        assert_eq!(
2038            docs.index("$.title").expect("there").len(),
2039            4,
2040            "a, red, blue and one"
2041        );
2042    }
2043
2044    #[test]
2045    fn changing_what_an_index_is_asked_rebuilds_it() {
2046        let mut docs = Docs::new();
2047        docs.create_index("$.tags").expect("indexed");
2048        docs.put_bytes(b"a", &tagged("one", &["red", "blue"]))
2049            .expect("put");
2050        assert!(
2051            found(&docs, "$.tags", &Key::text("red")).is_empty(),
2052            "an equality index over an array files nothing"
2053        );
2054
2055        docs.create_array_index("$.tags").expect("rebuilt");
2056        assert_eq!(docs.indexes().len(), 1, "it replaced rather than added");
2057        assert_eq!(found(&docs, "$.tags", &Key::text("red")), ["a"]);
2058
2059        docs.create_array_index("$.tags").expect("already there");
2060        assert_eq!(docs.indexes().len(), 1);
2061    }
2062
2063    /// The key a text index files one word under.
2064    fn word(w: &str) -> Key {
2065        Key::word(w).expect("one word")
2066    }
2067
2068    #[test]
2069    fn a_collection_whose_key_table_is_full_stores_the_rest_with_names() {
2070        // Fill the table with names no document below uses, then write one and
2071        // check it is stored whole rather than refused.
2072        let mut docs = Docs::new();
2073        for i in 0..crate::KEYS_MAX {
2074            let name = format!("filler{i}");
2075            assert!(docs.keys.intern(name.as_bytes()).is_some());
2076        }
2077        assert!(docs.keys().is_full());
2078
2079        docs.put_bytes(b"order:1", &order(1, "open", 1))
2080            .expect("put");
2081        let d = docs.get(b"order:1").expect("stored");
2082        assert!(!d.value().is_interned(), "there were no ids left to use");
2083        assert_eq!(d.get(b"status").and_then(|v| v.as_text()), Some("open"));
2084        assert_eq!(
2085            d.path("$.lines[0].sku")
2086                .expect("a path")
2087                .and_then(|v| v.as_text()),
2088            Some("sku-0")
2089        );
2090    }
2091
2092    // ---- the vector index
2093
2094    /// A document with a language and an embedding, which is the shape `10`
2095    /// section 3 uses as its example.
2096    fn item(lang: &str, v: &[f32]) -> Vec<u8> {
2097        let mut b = Builder::new();
2098        b.begin_object().expect("open");
2099        b.key(b"lang").expect("key");
2100        b.text(lang).expect("value");
2101        b.key(b"embedding").expect("key");
2102        b.begin_array().expect("open");
2103        for x in v {
2104            b.float(f64::from(*x)).expect("value");
2105        }
2106        b.end_array().expect("close");
2107        b.end_object().expect("close");
2108        b.finish().expect("finished").to_vec()
2109    }
2110
2111    /// The same document with no embedding in it at all.
2112    fn bare(lang: &str) -> Vec<u8> {
2113        let mut b = Builder::new();
2114        b.begin_object().expect("open");
2115        b.key(b"lang").expect("key");
2116        b.text(lang).expect("value");
2117        b.end_object().expect("close");
2118        b.finish().expect("finished").to_vec()
2119    }
2120
2121    /// Eight coordinates that depend only on `n`, so a test that fails fails
2122    /// the same way twice.
2123    fn spread(n: u64) -> [f32; 8] {
2124        let mut s = n.wrapping_mul(0x9e37_79b9_7f4a_7c15) | 1;
2125        let mut v = [0.0f32; 8];
2126        for x in &mut v {
2127            s ^= s << 13;
2128            s ^= s >> 7;
2129            s ^= s << 17;
2130            *x = (s >> 40) as f32 / 4096.0 - 1.0;
2131        }
2132        v
2133    }
2134
2135    #[test]
2136    fn a_document_and_its_embedding_are_one_write() {
2137        let mut docs = Docs::new();
2138        docs.create_vector_index("$.embedding", 3).expect("index");
2139        for (id, v) in [
2140            ("a", [1.0, 0.0, 0.0]),
2141            ("b", [0.0, 1.0, 0.0]),
2142            ("c", [0.0, 0.0, 1.0]),
2143        ] {
2144            docs.put_bytes(id.as_bytes(), &item("en", &v)).expect("put");
2145        }
2146        assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 3);
2147
2148        // The answer is documents and not ids to go and look up somewhere else,
2149        // and it comes back nearest first.
2150        let mut got = Vec::new();
2151        let n = docs
2152            .nearest("$.embedding", &[0.9, 0.1, 0.0], 3, |id, doc, d| {
2153                let lang = doc
2154                    .get(b"lang")
2155                    .and_then(|v| v.as_text())
2156                    .map(str::to_owned);
2157                got.push((id.to_vec(), lang, d));
2158            })
2159            .expect("nearest");
2160        assert_eq!(n, 3);
2161        assert_eq!(got[0].0, b"a".to_vec());
2162        assert_eq!(got[0].1.as_deref(), Some("en"));
2163        assert!(got[0].2 <= got[1].2 && got[1].2 <= got[2].2);
2164
2165        // A path with no vector index on it says so rather than scanning.
2166        assert!(
2167            docs.nearest("$.lang", &[1.0, 0.0, 0.0], 1, |_, _, _| {})
2168                .is_err()
2169        );
2170    }
2171
2172    #[test]
2173    fn an_embedding_of_the_wrong_shape_fails_the_write_and_stores_nothing() {
2174        let mut docs = Docs::new();
2175        docs.create_vector_index("$.embedding", 3).expect("index");
2176        docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2177            .expect("put");
2178
2179        for wrong in [vec![1.0, 0.0], vec![1.0, 0.0, 0.0, 0.0]] {
2180            assert!(docs.put_bytes(b"b", &item("en", &wrong)).is_err());
2181        }
2182        assert!(docs.get(b"b").is_none(), "the write left nothing behind");
2183        assert_eq!(docs.len(), 1);
2184        assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 1);
2185    }
2186
2187    #[test]
2188    fn a_document_with_nothing_at_the_path_is_not_in_the_index() {
2189        let mut docs = Docs::new();
2190        docs.create_vector_index("$.embedding", 3).expect("index");
2191        docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2192            .expect("put");
2193        docs.put_bytes(b"b", &bare("en")).expect("put");
2194        assert_eq!(docs.len(), 2);
2195        assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 1);
2196
2197        // Rewriting a document without its embedding takes the old vector out
2198        // rather than leaving the one the last version had.
2199        docs.put_bytes(b"a", &bare("en")).expect("put");
2200        assert!(
2201            docs.vector_index("$.embedding")
2202                .expect("declared")
2203                .is_empty()
2204        );
2205        assert!(docs.get(b"a").is_some());
2206
2207        // And a removal takes it with it.
2208        docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2209            .expect("put");
2210        assert!(docs.remove(b"a"));
2211        assert!(
2212            docs.vector_index("$.embedding")
2213                .expect("declared")
2214                .is_empty()
2215        );
2216    }
2217
2218    #[test]
2219    fn a_filter_finds_the_nearest_match_and_not_the_nearest_that_matches() {
2220        let mut docs = Docs::new();
2221        docs.create_index("$.lang").expect("index");
2222        docs.create_vector_index("$.embedding", 8).expect("index");
2223
2224        // Four hundred English documents spread about, and one French one that
2225        // sits nowhere near the query.
2226        for n in 0..400u64 {
2227            let id = format!("en:{n}");
2228            docs.put_bytes(id.as_bytes(), &item("en", &spread(n)))
2229                .expect("put");
2230        }
2231        docs.put_bytes(b"fr", &item("fr", &spread(9_999)))
2232            .expect("put");
2233
2234        let q = spread(3);
2235        let mut top = Vec::new();
2236        docs.nearest("$.embedding", &q, 20, |id, _, _| top.push(id.to_vec()))
2237            .expect("nearest");
2238        assert_eq!(top[0], b"en:3".to_vec());
2239        assert!(
2240            !top.iter().any(|id| id == b"fr"),
2241            "searching and then filtering would have answered nothing"
2242        );
2243
2244        // Filtering inside the scan finds it anyway.
2245        let french = [("$.lang", Key::text("fr"))];
2246        let mut found = Vec::new();
2247        docs.nearest_where("$.embedding", &q, 5, &french, |id, _, _| {
2248            found.push(id.to_vec())
2249        })
2250        .expect("nearest");
2251        assert_eq!(found, [b"fr".to_vec()]);
2252
2253        // A path with no index on it cannot be filtered on, and says so.
2254        let nothing = [("$.topic", Key::text("finance"))];
2255        assert!(
2256            docs.nearest_where("$.embedding", &q, 5, &nothing, |_, _, _| {})
2257                .is_err()
2258        );
2259    }
2260
2261    #[test]
2262    fn declaring_either_index_last_gives_the_same_answers() {
2263        let q = spread(11);
2264        let french = [("$.lang", Key::text("fr"))];
2265
2266        // Vectors first, then the field the filter reads, so every tag was
2267        // written before there was anything to put in it.
2268        let mut late = Docs::new();
2269        late.create_vector_index("$.embedding", 8).expect("index");
2270        for n in 0..200u64 {
2271            let lang = if n % 50 == 0 { "fr" } else { "en" };
2272            let id = format!("{n}");
2273            late.put_bytes(id.as_bytes(), &item(lang, &spread(n)))
2274                .expect("put");
2275        }
2276        late.create_index("$.lang").expect("index");
2277
2278        // The other way round, where every write already knew.
2279        let mut early = Docs::new();
2280        early.create_index("$.lang").expect("index");
2281        for n in 0..200u64 {
2282            let lang = if n % 50 == 0 { "fr" } else { "en" };
2283            let id = format!("{n}");
2284            early
2285                .put_bytes(id.as_bytes(), &item(lang, &spread(n)))
2286                .expect("put");
2287        }
2288        early.create_vector_index("$.embedding", 8).expect("index");
2289
2290        let mut a = Vec::new();
2291        late.nearest_where("$.embedding", &q, 4, &french, |id, _, _| {
2292            a.push(id.to_vec())
2293        })
2294        .expect("nearest");
2295        let mut b = Vec::new();
2296        early
2297            .nearest_where("$.embedding", &q, 4, &french, |id, _, _| {
2298                b.push(id.to_vec())
2299            })
2300            .expect("nearest");
2301        assert_eq!(a.len(), 4, "there are four French documents to find");
2302        assert_eq!(a, b);
2303
2304        // Dropping the field index and declaring it again leaves the tags right.
2305        assert!(late.drop_index("$.lang"));
2306        late.create_index("$.lang").expect("index");
2307        let mut again = Vec::new();
2308        late.nearest_where("$.embedding", &q, 4, &french, |id, _, _| {
2309            again.push(id.to_vec());
2310        })
2311        .expect("nearest");
2312        assert_eq!(again, a);
2313    }
2314
2315    #[test]
2316    fn nearest_to_leaves_the_document_itself_out() {
2317        let mut docs = Docs::new();
2318        docs.create_vector_index("$.embedding", 3).expect("index");
2319        for (id, v) in [
2320            ("a", [1.0, 0.0, 0.0]),
2321            ("b", [0.9, 0.1, 0.0]),
2322            ("c", [0.0, 0.0, 1.0]),
2323        ] {
2324            docs.put_bytes(id.as_bytes(), &item("en", &v)).expect("put");
2325        }
2326
2327        let mut like = Vec::new();
2328        docs.nearest_to("$.embedding", b"a", 2, |id, _, _| like.push(id.to_vec()))
2329            .expect("nearest");
2330        assert_eq!(like, [b"b".to_vec(), b"c".to_vec()]);
2331
2332        // A document with no embedding has nothing to be like.
2333        docs.put_bytes(b"d", &bare("en")).expect("put");
2334        let mut none = 0;
2335        assert_eq!(
2336            docs.nearest_to("$.embedding", b"d", 2, |_, _, _| none += 1)
2337                .expect("nearest"),
2338            0
2339        );
2340    }
2341
2342    #[test]
2343    fn declaring_the_same_vector_index_again_rebuilds_nothing() {
2344        let mut docs = Docs::new();
2345        for n in 0..8u64 {
2346            let id = format!("{n}");
2347            docs.put_bytes(id.as_bytes(), &item("en", &spread(n)))
2348                .expect("put");
2349        }
2350        docs.create_vector_index("$.embedding", 8).expect("index");
2351        assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 8);
2352
2353        // The same declaration is nothing at all.
2354        docs.create_vector_index("$.embedding", 8).expect("again");
2355        assert_eq!(docs.vector_indexes().len(), 1);
2356
2357        // A different width is a different question, so it is rebuilt, and
2358        // documents whose embedding is not that wide fail the declaration. The
2359        // index that was already there is the one that is still there.
2360        assert!(docs.create_vector_index("$.embedding", 4).is_err());
2361        let still = docs.vector_index("$.embedding").expect("still declared");
2362        assert_eq!(still.dim(), 8);
2363        assert_eq!(still.len(), 8);
2364
2365        assert!(docs.drop_vector_index("$.embedding"));
2366        assert!(!docs.drop_vector_index("$.embedding"));
2367        assert!(docs.vector_indexes().is_empty());
2368        assert_eq!(docs.len(), 8, "the documents are untouched");
2369    }
2370
2371    #[test]
2372    fn clearing_a_collection_empties_the_vector_index_and_keeps_it_declared() {
2373        let mut docs = Docs::new();
2374        docs.create_vector_index("$.embedding", 3).expect("index");
2375        docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2376            .expect("put");
2377        let full = docs.memory_bytes();
2378
2379        docs.clear();
2380        assert!(docs.is_empty());
2381        assert!(
2382            docs.vector_index("$.embedding")
2383                .expect("declared")
2384                .is_empty()
2385        );
2386        assert!(docs.memory_bytes() < full);
2387
2388        docs.put_bytes(b"b", &item("en", &[0.0, 1.0, 0.0]))
2389            .expect("put");
2390        let mut got = Vec::new();
2391        docs.nearest("$.embedding", &[0.0, 1.0, 0.0], 1, |id, _, _| {
2392            got.push(id.to_vec())
2393        })
2394        .expect("nearest");
2395        assert_eq!(got, [b"b".to_vec()]);
2396    }
2397}