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        // Fewer documents under Miri. What comes back is the stored bytes over
1397        // the plain bytes summed across the documents, and neither half of that
1398        // counts the name table, so the ratio is a per document number and the
1399        // count only decides how many times it is averaged.
1400        let n = if cfg!(miri) { 16i64 } else { 256 };
1401        for i in 0..n {
1402            let bytes = shape(i);
1403            plain += bytes.len();
1404            docs.put_bytes(format!("d:{i}").as_bytes(), &bytes)
1405                .expect("put");
1406        }
1407        let stored: usize = (0..n)
1408            .map(|i| {
1409                docs.bytes(format!("d:{i}").as_bytes())
1410                    .expect("stored")
1411                    .len()
1412            })
1413            .sum();
1414        stored as f64 / plain as f64
1415    }
1416
1417    #[test]
1418    fn interning_makes_a_collection_of_the_same_shape_smaller() {
1419        // The claim in `09` section 4 is that the same field names on every
1420        // document are most of what a document collection costs, and that
1421        // interning them is worth about forty percent. How much it is actually
1422        // worth depends on how much of a document is names, so both ends are
1423        // measured here rather than one number being asserted twice.
1424        //
1425        // A document that is mostly names, which is what a typed collection of
1426        // small records looks like, keeps a little over half its bytes.
1427        let names = shrinkage(|i| {
1428            let mut b = Builder::new();
1429            b.begin_object().expect("open");
1430            for f in 0..20 {
1431                b.key(format!("some_field_name_{f:02}").as_bytes())
1432                    .expect("key");
1433                b.int(i + f).expect("value");
1434            }
1435            b.end_object().expect("close");
1436            b.finish().expect("finished").to_vec()
1437        });
1438        assert!(names < 0.60, "a document of names kept {names}");
1439
1440        // An order, which carries real payload as well, keeps about three
1441        // quarters. That is the honest floor for the claim and it is still a
1442        // fifth of the collection gone for nothing but a table of twenty
1443        // strings.
1444        let orders = shrinkage(|i| order(i, "open", 2));
1445        assert!(orders < 0.80, "an order collection kept {orders}");
1446    }
1447
1448    #[test]
1449    fn a_document_whose_keys_are_already_ids_is_refused() {
1450        let mut b = Builder::new();
1451        b.begin_object_interned().expect("open");
1452        b.key_id(0).expect("key");
1453        b.int(1).expect("value");
1454        b.end_object().expect("close");
1455        let bytes = b.finish().expect("finished").to_vec();
1456
1457        let mut docs = Docs::new();
1458        let err = docs.put_bytes(b"x", &bytes).expect_err("refused");
1459        assert_eq!(err.code(), Code::Invalid);
1460    }
1461
1462    #[test]
1463    fn a_document_that_is_not_readable_is_refused() {
1464        let mut docs = Docs::new();
1465        let err = docs.put_bytes(b"x", &[2, 0, 0, 0]).expect_err("refused");
1466        assert_eq!(err.code(), Code::Corrupt);
1467        assert!(docs.is_empty());
1468    }
1469
1470    #[test]
1471    fn a_removal_leaves_every_other_document_where_it_was() {
1472        let mut docs = Docs::new();
1473        // Fewer documents under Miri. Every count below comes off this one so
1474        // that a smaller run still removes a third of them and still looks at
1475        // all of the rest.
1476        let n = if cfg!(miri) { 24i64 } else { 64 };
1477        for i in 0..n {
1478            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1479                .expect("put");
1480        }
1481        let gone = (0..n).step_by(3).count();
1482        for i in (0..n).step_by(3) {
1483            assert!(docs.remove(format!("order:{i}").as_bytes()));
1484        }
1485        assert_eq!(docs.len(), n as usize - gone);
1486        for i in 0..n {
1487            let id = format!("order:{i}");
1488            match docs.get(id.as_bytes()) {
1489                Some(d) => {
1490                    assert!(i % 3 != 0, "{id} was removed");
1491                    assert_eq!(d.get(b"id").and_then(|v| v.as_int()), Some(i));
1492                }
1493                None => assert!(i % 3 == 0, "{id} was not removed"),
1494            }
1495        }
1496        assert_eq!(docs.keys().len(), 6, "a removal does not un-intern a name");
1497    }
1498
1499    #[test]
1500    fn a_walk_sees_every_document_once() {
1501        // Fewer documents under Miri, but still several batches of the scan,
1502        // which is what the cursor is being checked over. Every count below
1503        // comes from this one.
1504        let n = if cfg!(miri) { 48i64 } else { 200 };
1505        let mut docs = Docs::new();
1506        for i in 0..n {
1507            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1508                .expect("put");
1509        }
1510
1511        let mut seen: Vec<i64> = docs
1512            .iter()
1513            .map(|(_, d)| d.get(b"id").and_then(|v| v.as_int()).expect("an id"))
1514            .collect();
1515        seen.sort_unstable();
1516        assert_eq!(seen, (0..n).collect::<Vec<i64>>());
1517
1518        let mut scanned = Vec::new();
1519        let mut cursor = Cursor::START;
1520        loop {
1521            cursor = docs.scan(cursor, 16, |id, _| scanned.push(id.to_vec()));
1522            if cursor.is_end() {
1523                break;
1524            }
1525        }
1526        scanned.sort_unstable();
1527        scanned.dedup();
1528        assert_eq!(scanned.len(), n as usize);
1529    }
1530
1531    #[test]
1532    fn an_empty_collection_answers_nothing_rather_than_failing() {
1533        let docs = Docs::new();
1534        assert!(docs.is_empty());
1535        assert!(docs.get(b"nothing").is_none());
1536        assert!(docs.bytes(b"nothing").is_none());
1537        assert!(!docs.contains(b"nothing"));
1538        assert_eq!(docs.iter().count(), 0);
1539    }
1540
1541    #[test]
1542    fn a_document_prints_with_its_names_back_on() {
1543        let mut docs = Docs::new();
1544        docs.put_bytes(b"order:1", &order(1, "open", 1))
1545            .expect("put");
1546        let text = format!("{:?}", docs.get(b"order:1").expect("stored"));
1547        assert!(text.contains("\"status\": \"open\""), "{text}");
1548        assert!(text.contains("\"sku\": \"sku-0\""), "{text}");
1549    }
1550
1551    /// The ids `find` answers for one key, sorted so a test can compare them.
1552    fn found(docs: &Docs, path: &str, key: &Key) -> Vec<String> {
1553        let mut out = Vec::new();
1554        let n = docs
1555            .find(path, key, |id, d| {
1556                assert!(!d.is_empty(), "the document came back whole");
1557                out.push(String::from_utf8_lossy(id).into_owned());
1558            })
1559            .expect("indexed");
1560        assert_eq!(n, out.len(), "the count is what the callback saw");
1561        out.sort();
1562        out
1563    }
1564
1565    #[test]
1566    fn an_index_declared_after_the_documents_finds_them() {
1567        // A quarter of them shut, whatever the count is. Fewer under Miri, and
1568        // the two counts below come from it rather than being written out
1569        // again, which is what would go quietly wrong here.
1570        let n = if cfg!(miri) { 24i64 } else { 64 };
1571        let (shut, open) = ((n / 4) as usize, (n - n / 4) as usize);
1572        let mut docs = Docs::new();
1573        for i in 0..n {
1574            let status = if i % 4 == 0 { "shut" } else { "open" };
1575            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, status, 1))
1576                .expect("put");
1577        }
1578        docs.create_index("$.status").expect("indexed");
1579        assert_eq!(docs.index("$.status").expect("there").len(), 2);
1580        assert_eq!(docs.count("$.status", &Key::text("shut")).expect("i"), shut);
1581        assert_eq!(docs.count("$.status", &Key::text("open")).expect("i"), open);
1582        assert_eq!(found(&docs, "$.status", &Key::text("shut")).len(), shut);
1583        assert!(found(&docs, "$.status", &Key::text("gone")).is_empty());
1584
1585        // A document written after the index exists is filed by the write.
1586        docs.put_bytes(format!("order:{n}").as_bytes(), &order(n, "shut", 1))
1587            .expect("put");
1588        assert_eq!(
1589            docs.count("$.status", &Key::text("shut")).expect("i"),
1590            shut + 1
1591        );
1592    }
1593
1594    #[test]
1595    fn an_overwrite_moves_a_document_from_one_key_to_the_other() {
1596        let mut docs = Docs::new();
1597        docs.create_index("$.status").expect("indexed");
1598        docs.put_bytes(b"order:1", &order(1, "open", 1))
1599            .expect("put");
1600        assert_eq!(found(&docs, "$.status", &Key::text("open")), ["order:1"]);
1601
1602        docs.put_bytes(b"order:1", &order(1, "shut", 1))
1603            .expect("put");
1604        assert!(
1605            found(&docs, "$.status", &Key::text("open")).is_empty(),
1606            "the old key kept it"
1607        );
1608        assert_eq!(found(&docs, "$.status", &Key::text("shut")), ["order:1"]);
1609        assert_eq!(docs.index("$.status").expect("there").postings(), 1);
1610    }
1611
1612    #[test]
1613    fn a_removal_takes_a_document_out_of_every_index() {
1614        let mut docs = Docs::new();
1615        docs.create_index("$.status").expect("indexed");
1616        docs.create_index("$.customer").expect("indexed");
1617        for i in 0..8i64 {
1618            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1619                .expect("put");
1620        }
1621        assert!(docs.remove(b"order:3"));
1622        assert_eq!(found(&docs, "$.status", &Key::text("open")).len(), 7);
1623        assert_eq!(docs.count("$.customer", &Key::int(21)).expect("i"), 0);
1624        assert_eq!(docs.count("$.customer", &Key::int(28)).expect("i"), 1);
1625        for index in docs.indexes() {
1626            assert_eq!(index.postings(), 7);
1627        }
1628
1629        assert!(!docs.remove(b"order:3"), "it is already gone");
1630        assert_eq!(docs.index("$.status").expect("there").postings(), 7);
1631    }
1632
1633    #[test]
1634    fn a_path_that_names_a_container_or_nothing_is_simply_not_filed() {
1635        let mut docs = Docs::new();
1636        docs.create_index("$.lines").expect("indexed");
1637        docs.create_index("$.shipped").expect("indexed");
1638        docs.create_index("$.lines[0].qty").expect("indexed");
1639        for i in 0..4i64 {
1640            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 2))
1641                .expect("put");
1642        }
1643        assert_eq!(docs.len(), 4);
1644        assert!(
1645            docs.index("$.lines").expect("there").is_empty(),
1646            "an array has no equality key"
1647        );
1648        assert!(
1649            docs.index("$.shipped").expect("there").is_empty(),
1650            "no document has that path"
1651        );
1652        assert_eq!(
1653            docs.count("$.lines[0].qty", &Key::int(1)).expect("i"),
1654            4,
1655            "a path through an array reaches a scalar"
1656        );
1657    }
1658
1659    #[test]
1660    fn a_value_too_long_to_index_fails_the_write_and_stores_nothing() {
1661        let mut b = Builder::new();
1662        b.begin_object().expect("open");
1663        b.key(b"status").expect("key");
1664        b.text(&"x".repeat(crate::KEY_MAX)).expect("value");
1665        b.end_object().expect("close");
1666        let huge = b.finish().expect("finished").to_vec();
1667
1668        let mut docs = Docs::new();
1669        docs.create_index("$.status").expect("indexed");
1670        let err = docs.put_bytes(b"order:1", &huge).expect_err("refused");
1671        assert_eq!(err.code(), Code::Full);
1672        assert!(
1673            docs.is_empty(),
1674            "a write that cannot be indexed leaves nothing behind"
1675        );
1676
1677        // Without the index it is an ordinary document and goes in fine.
1678        assert!(docs.drop_index("$.status"));
1679        docs.put_bytes(b"order:1", &huge).expect("put");
1680        assert_eq!(docs.len(), 1);
1681    }
1682
1683    #[test]
1684    fn a_query_on_a_path_with_no_index_says_so_rather_than_scanning() {
1685        let mut docs = Docs::new();
1686        docs.put_bytes(b"order:1", &order(1, "open", 1))
1687            .expect("put");
1688        let err = docs
1689            .find("$.status", &Key::text("open"), |_, _| ())
1690            .expect_err("refused");
1691        assert_eq!(err.code(), Code::Invalid);
1692        assert_eq!(
1693            docs.count("$.status", &Key::text("open"))
1694                .expect_err("refused")
1695                .code(),
1696            Code::Invalid
1697        );
1698        assert!(docs.index("$.status").is_none());
1699        assert!(!docs.drop_index("$.status"));
1700    }
1701
1702    #[test]
1703    fn declaring_the_same_index_twice_leaves_the_first_one_alone() {
1704        let mut docs = Docs::new();
1705        docs.create_index("$.status").expect("indexed");
1706        docs.put_bytes(b"order:1", &order(1, "open", 1))
1707            .expect("put");
1708        docs.create_index("$.status").expect("indexed again");
1709        assert_eq!(docs.indexes().len(), 1);
1710        assert_eq!(
1711            docs.index("$.status").expect("there").postings(),
1712            1,
1713            "a redeclaration did not double file anything"
1714        );
1715        assert!(docs.create_index("$.[").is_err(), "the path has to parse");
1716    }
1717
1718    #[test]
1719    fn clearing_a_collection_empties_its_indexes_and_keeps_them() {
1720        let mut docs = Docs::new();
1721        docs.create_index("$.status").expect("indexed");
1722        for i in 0..8i64 {
1723            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1724                .expect("put");
1725        }
1726        docs.clear();
1727        assert!(docs.is_empty());
1728        assert!(docs.index("$.status").expect("still declared").is_empty());
1729        assert_eq!(docs.count("$.status", &Key::text("open")).expect("i"), 0);
1730
1731        docs.put_bytes(b"order:9", &order(9, "open", 1))
1732            .expect("put");
1733        assert_eq!(found(&docs, "$.status", &Key::text("open")), ["order:9"]);
1734    }
1735
1736    #[test]
1737    fn two_indexes_intersect_as_the_sets_they_are() {
1738        let mut docs = Docs::new();
1739        docs.create_index("$.status").expect("indexed");
1740        docs.create_index("$.customer").expect("indexed");
1741        for i in 0..32i64 {
1742            let status = if i % 2 == 0 { "open" } else { "shut" };
1743            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i % 4, status, 1))
1744                .expect("put");
1745        }
1746
1747        // What a planner does: probe both, walk the smaller, ask the larger.
1748        // That is `SINTER` and there is no code here that is not already the
1749        // set's.
1750        let open = Key::text("open");
1751        let customer = Key::int(14);
1752        let small = docs.count("$.customer", &customer).expect("indexed");
1753        let large = docs.count("$.status", &open).expect("indexed");
1754        assert_eq!((small, large), (8, 16));
1755
1756        let small = docs.index("$.customer").expect("there").get(&customer);
1757        let large = docs.index("$.status").expect("there").get(&open);
1758        let (Some(small), Some(large)) = (small, large) else {
1759            panic!("both keys are filed");
1760        };
1761        let mut both = Vec::new();
1762        index::each_id(small, |id| {
1763            if large.contains(id) {
1764                both.push(String::from_utf8_lossy(id).into_owned());
1765            }
1766        });
1767        both.sort();
1768        assert_eq!(
1769            both,
1770            [
1771                "order:10", "order:14", "order:18", "order:2", "order:22", "order:26", "order:30",
1772                "order:6"
1773            ]
1774        );
1775    }
1776
1777    /// The customer numbers a range answers, in the order it answered them.
1778    fn ranged(docs: &Docs, lo: Bound<&Key>, hi: Bound<&Key>) -> Vec<i64> {
1779        let mut out = Vec::new();
1780        let n = docs
1781            .range("$.customer", lo, hi, |_, d| {
1782                out.push(d.get(b"customer").and_then(|v| v.as_int()).expect("there"));
1783            })
1784            .expect("ordered");
1785        assert_eq!(n, out.len());
1786
1787        let mut back = Vec::new();
1788        docs.range_rev("$.customer", lo, hi, |_, d| {
1789            back.push(d.get(b"customer").and_then(|v| v.as_int()).expect("there"));
1790        })
1791        .expect("ordered");
1792        back.reverse();
1793        assert_eq!(out, back, "backwards is forwards read the other way");
1794        assert_eq!(
1795            docs.count_range("$.customer", lo, hi).expect("ordered"),
1796            out.len()
1797        );
1798        out
1799    }
1800
1801    #[test]
1802    fn an_ordered_index_answers_a_range_in_order() {
1803        let mut docs = Docs::new();
1804        docs.create_ordered_index("$.customer").expect("ordered");
1805        // Customer is seven times the id, so the values are 0, 7, 14 and on.
1806        // Fewer under Miri. The bounds below are all well inside a collection
1807        // this size and the empty one is well past the end of either, so they
1808        // ask the same questions of a shorter run of values.
1809        let n = if cfg!(miri) { 24i64 } else { 64 };
1810        for i in 0..n {
1811            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1812                .expect("put");
1813        }
1814
1815        assert_eq!(
1816            ranged(&docs, Bound::Unbounded, Bound::Unbounded),
1817            (0..n).map(|i| i * 7).collect::<Vec<i64>>()
1818        );
1819        let (lo, hi) = (Key::int(70), Key::int(105));
1820        assert_eq!(
1821            ranged(&docs, Bound::Included(&lo), Bound::Included(&hi)),
1822            [70, 77, 84, 91, 98, 105]
1823        );
1824        assert_eq!(
1825            ranged(&docs, Bound::Excluded(&lo), Bound::Excluded(&hi)),
1826            [77, 84, 91, 98]
1827        );
1828        // Bounds that fall between two values, which is the ordinary case.
1829        assert_eq!(
1830            ranged(
1831                &docs,
1832                Bound::Included(&Key::int(71)),
1833                Bound::Excluded(&Key::int(90))
1834            ),
1835            [77, 84]
1836        );
1837        assert!(ranged(&docs, Bound::Included(&Key::int(442)), Bound::Unbounded).is_empty());
1838
1839        // Equality still works on the same index.
1840        assert_eq!(docs.count("$.customer", &Key::int(70)).expect("i"), 1);
1841        assert_eq!(
1842            docs.index("$.customer").expect("there").kind(),
1843            IndexKind::Ordered
1844        );
1845    }
1846
1847    #[test]
1848    fn a_range_stays_right_through_writes_and_removals() {
1849        // Fewer documents under Miri. Half of them are removed either way, so
1850        // the renumbering still runs through the whole collection, and every
1851        // count below comes from this one rather than being written out again.
1852        let n = if cfg!(miri) { 32i64 } else { 128 };
1853        let mut docs = Docs::new();
1854        for i in 0..n {
1855            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1856                .expect("put");
1857        }
1858        // Declared after the fact, so this is the backfill and not the write
1859        // path putting the tree together.
1860        docs.create_ordered_index("$.customer").expect("ordered");
1861        assert_eq!(
1862            ranged(&docs, Bound::Unbounded, Bound::Unbounded).len(),
1863            n as usize
1864        );
1865
1866        // Every removal moves the key table's last row into the hole, so this is
1867        // the renumbering going through the whole collection.
1868        for i in (0..n).step_by(2) {
1869            assert!(docs.remove(format!("order:{i}").as_bytes()));
1870        }
1871        assert_eq!(
1872            ranged(&docs, Bound::Unbounded, Bound::Unbounded),
1873            (0..n)
1874                .filter(|i| i % 2 == 1)
1875                .map(|i| i * 7)
1876                .collect::<Vec<i64>>()
1877        );
1878
1879        // And an overwrite that moves a document from one key to another.
1880        docs.put_bytes(b"order:1", &order(200, "open", 1))
1881            .expect("put");
1882        let after = ranged(&docs, Bound::Unbounded, Bound::Unbounded);
1883        assert_eq!(after.first(), Some(&21), "seven is gone");
1884        assert_eq!(after.last(), Some(&1400), "and it came back at the top");
1885    }
1886
1887    #[test]
1888    fn an_equality_index_refuses_a_range_rather_than_answering_nothing() {
1889        let mut docs = Docs::new();
1890        docs.create_index("$.customer").expect("indexed");
1891        docs.put_bytes(b"order:1", &order(1, "open", 1))
1892            .expect("put");
1893        let err = docs
1894            .range("$.customer", Bound::Unbounded, Bound::Unbounded, |_, _| ())
1895            .expect_err("refused");
1896        assert_eq!(err.code(), Code::Invalid);
1897        assert!(err.to_string().contains("equality"), "{err}");
1898        assert_eq!(
1899            docs.range("$.status", Bound::Unbounded, Bound::Unbounded, |_, _| ())
1900                .expect_err("refused")
1901                .code(),
1902            Code::Invalid
1903        );
1904    }
1905
1906    #[test]
1907    fn asking_for_an_order_on_an_equality_index_upgrades_it() {
1908        let mut docs = Docs::new();
1909        docs.create_index("$.customer").expect("indexed");
1910        for i in 0..8i64 {
1911            docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1912                .expect("put");
1913        }
1914        assert_eq!(
1915            docs.index("$.customer").expect("there").kind(),
1916            IndexKind::Equality
1917        );
1918
1919        docs.create_ordered_index("$.customer").expect("upgraded");
1920        assert_eq!(docs.indexes().len(), 1, "it replaced rather than added");
1921        assert_eq!(ranged(&docs, Bound::Unbounded, Bound::Unbounded).len(), 8);
1922
1923        // And going the other way leaves the order alone, because an ordered
1924        // index answers equality too.
1925        docs.create_index("$.customer").expect("already there");
1926        assert_eq!(
1927            docs.index("$.customer").expect("there").kind(),
1928            IndexKind::Ordered
1929        );
1930        assert_eq!(docs.indexes().len(), 1);
1931    }
1932
1933    /// A document with a list of tags at `$.tags` and a title at `$.title`.
1934    fn tagged(title: &str, tags: &[&str]) -> Vec<u8> {
1935        let mut b = Builder::new();
1936        b.begin_object().expect("open");
1937        b.key(b"title").expect("key");
1938        b.text(title).expect("value");
1939        b.key(b"tags").expect("key");
1940        b.begin_array().expect("open");
1941        for tag in tags {
1942            b.text(tag).expect("value");
1943        }
1944        b.end_array().expect("close");
1945        b.end_object().expect("close");
1946        b.finish().expect("finished").to_vec()
1947    }
1948
1949    #[test]
1950    fn an_array_index_files_a_document_under_every_element() {
1951        let mut docs = Docs::new();
1952        docs.create_array_index("$.tags").expect("indexed");
1953        docs.put_bytes(b"a", &tagged("one", &["red", "blue"]))
1954            .expect("put");
1955        docs.put_bytes(b"b", &tagged("two", &["blue", "green"]))
1956            .expect("put");
1957        docs.put_bytes(b"c", &tagged("three", &[])).expect("put");
1958
1959        assert_eq!(found(&docs, "$.tags", &Key::text("red")), ["a"]);
1960        assert_eq!(found(&docs, "$.tags", &Key::text("blue")), ["a", "b"]);
1961        assert_eq!(found(&docs, "$.tags", &Key::text("green")), ["b"]);
1962        assert!(found(&docs, "$.tags", &Key::text("puce")).is_empty());
1963        assert_eq!(
1964            docs.index("$.tags").expect("there").len(),
1965            3,
1966            "three distinct tags over two documents"
1967        );
1968    }
1969
1970    #[test]
1971    fn an_array_index_takes_every_element_back_out_again() {
1972        let mut docs = Docs::new();
1973        docs.create_array_index("$.tags").expect("indexed");
1974        docs.put_bytes(b"a", &tagged("one", &["red", "blue"]))
1975            .expect("put");
1976        docs.put_bytes(b"b", &tagged("two", &["blue"]))
1977            .expect("put");
1978
1979        // An overwrite drops one tag and gains another.
1980        docs.put_bytes(b"a", &tagged("one", &["blue", "green"]))
1981            .expect("put");
1982        assert!(found(&docs, "$.tags", &Key::text("red")).is_empty());
1983        assert_eq!(found(&docs, "$.tags", &Key::text("blue")), ["a", "b"]);
1984        assert_eq!(found(&docs, "$.tags", &Key::text("green")), ["a"]);
1985
1986        assert!(docs.remove(b"a"));
1987        assert_eq!(found(&docs, "$.tags", &Key::text("blue")), ["b"]);
1988        assert!(found(&docs, "$.tags", &Key::text("green")).is_empty());
1989        assert_eq!(
1990            docs.index("$.tags").expect("there").len(),
1991            1,
1992            "a tag nobody has left is not a key any more"
1993        );
1994    }
1995
1996    #[test]
1997    fn an_array_index_treats_one_value_as_a_list_of_one() {
1998        let mut docs = Docs::new();
1999        docs.create_array_index("$.status").expect("indexed");
2000        docs.put_bytes(b"order:1", &order(1, "open", 1))
2001            .expect("put");
2002        assert_eq!(found(&docs, "$.status", &Key::text("open")), ["order:1"]);
2003    }
2004
2005    #[test]
2006    fn the_same_element_twice_is_one_posting() {
2007        let mut docs = Docs::new();
2008        docs.create_array_index("$.tags").expect("indexed");
2009        docs.put_bytes(b"a", &tagged("one", &["red", "red", "red"]))
2010            .expect("put");
2011        assert_eq!(found(&docs, "$.tags", &Key::text("red")), ["a"]);
2012        assert_eq!(docs.index("$.tags").expect("there").postings(), 1);
2013
2014        // And taking it out once takes it out, rather than three times over.
2015        assert!(docs.remove(b"a"));
2016        assert_eq!(docs.index("$.tags").expect("there").postings(), 0);
2017        assert!(docs.index("$.tags").expect("there").is_empty());
2018    }
2019
2020    #[test]
2021    fn a_text_index_files_a_document_under_every_word() {
2022        let mut docs = Docs::new();
2023        docs.create_text_index("$.title").expect("indexed");
2024        docs.put_bytes(b"a", &tagged("A red bicycle", &[]))
2025            .expect("put");
2026        docs.put_bytes(b"b", &tagged("The red car, and a bicycle!", &[]))
2027            .expect("put");
2028
2029        assert_eq!(found(&docs, "$.title", &word("bicycle")), ["a", "b"]);
2030        assert_eq!(found(&docs, "$.title", &word("car")), ["b"]);
2031        assert_eq!(
2032            found(&docs, "$.title", &word("RED")),
2033            ["a", "b"],
2034            "a search folds case the same way the write did"
2035        );
2036        assert!(found(&docs, "$.title", &word("lorry")).is_empty());
2037    }
2038
2039    #[test]
2040    fn a_text_index_follows_the_words_through_a_rewrite() {
2041        let mut docs = Docs::new();
2042        docs.create_text_index("$.title").expect("indexed");
2043        docs.put_bytes(b"a", &tagged("a red bicycle", &[]))
2044            .expect("put");
2045        docs.put_bytes(b"a", &tagged("a blue bicycle", &[]))
2046            .expect("put");
2047        assert!(found(&docs, "$.title", &word("red")).is_empty());
2048        assert_eq!(found(&docs, "$.title", &word("blue")), ["a"]);
2049        assert_eq!(found(&docs, "$.title", &word("bicycle")), ["a"]);
2050
2051        assert!(docs.remove(b"a"));
2052        assert!(docs.index("$.title").expect("there").is_empty());
2053    }
2054
2055    #[test]
2056    fn a_text_index_declared_after_the_documents_finds_them() {
2057        let mut docs = Docs::new();
2058        for i in 0..16i64 {
2059            let title = if i % 2 == 0 {
2060                "a red one"
2061            } else {
2062                "a blue one"
2063            };
2064            docs.put_bytes(format!("t:{i}").as_bytes(), &tagged(title, &[]))
2065                .expect("put");
2066        }
2067        docs.create_text_index("$.title").expect("indexed");
2068        assert_eq!(docs.count("$.title", &word("red")).expect("i"), 8);
2069        assert_eq!(docs.count("$.title", &word("one")).expect("i"), 16);
2070        assert_eq!(
2071            docs.index("$.title").expect("there").len(),
2072            4,
2073            "a, red, blue and one"
2074        );
2075    }
2076
2077    #[test]
2078    fn changing_what_an_index_is_asked_rebuilds_it() {
2079        let mut docs = Docs::new();
2080        docs.create_index("$.tags").expect("indexed");
2081        docs.put_bytes(b"a", &tagged("one", &["red", "blue"]))
2082            .expect("put");
2083        assert!(
2084            found(&docs, "$.tags", &Key::text("red")).is_empty(),
2085            "an equality index over an array files nothing"
2086        );
2087
2088        docs.create_array_index("$.tags").expect("rebuilt");
2089        assert_eq!(docs.indexes().len(), 1, "it replaced rather than added");
2090        assert_eq!(found(&docs, "$.tags", &Key::text("red")), ["a"]);
2091
2092        docs.create_array_index("$.tags").expect("already there");
2093        assert_eq!(docs.indexes().len(), 1);
2094    }
2095
2096    /// The key a text index files one word under.
2097    fn word(w: &str) -> Key {
2098        Key::word(w).expect("one word")
2099    }
2100
2101    #[test]
2102    #[cfg_attr(miri, ignore = "a full key table is the claim and it is 65536 names")]
2103    fn a_collection_whose_key_table_is_full_stores_the_rest_with_names() {
2104        // Fill the table with names no document below uses, then write one and
2105        // check it is stored whole rather than refused.
2106        let mut docs = Docs::new();
2107        for i in 0..crate::KEYS_MAX {
2108            let name = format!("filler{i}");
2109            assert!(docs.keys.intern(name.as_bytes()).is_some());
2110        }
2111        assert!(docs.keys().is_full());
2112
2113        docs.put_bytes(b"order:1", &order(1, "open", 1))
2114            .expect("put");
2115        let d = docs.get(b"order:1").expect("stored");
2116        assert!(!d.value().is_interned(), "there were no ids left to use");
2117        assert_eq!(d.get(b"status").and_then(|v| v.as_text()), Some("open"));
2118        assert_eq!(
2119            d.path("$.lines[0].sku")
2120                .expect("a path")
2121                .and_then(|v| v.as_text()),
2122            Some("sku-0")
2123        );
2124    }
2125
2126    // ---- the vector index
2127
2128    /// A document with a language and an embedding, which is the shape `10`
2129    /// section 3 uses as its example.
2130    fn item(lang: &str, v: &[f32]) -> Vec<u8> {
2131        let mut b = Builder::new();
2132        b.begin_object().expect("open");
2133        b.key(b"lang").expect("key");
2134        b.text(lang).expect("value");
2135        b.key(b"embedding").expect("key");
2136        b.begin_array().expect("open");
2137        for x in v {
2138            b.float(f64::from(*x)).expect("value");
2139        }
2140        b.end_array().expect("close");
2141        b.end_object().expect("close");
2142        b.finish().expect("finished").to_vec()
2143    }
2144
2145    /// The same document with no embedding in it at all.
2146    fn bare(lang: &str) -> Vec<u8> {
2147        let mut b = Builder::new();
2148        b.begin_object().expect("open");
2149        b.key(b"lang").expect("key");
2150        b.text(lang).expect("value");
2151        b.end_object().expect("close");
2152        b.finish().expect("finished").to_vec()
2153    }
2154
2155    /// Eight coordinates that depend only on `n`, so a test that fails fails
2156    /// the same way twice.
2157    fn spread(n: u64) -> [f32; 8] {
2158        let mut s = n.wrapping_mul(0x9e37_79b9_7f4a_7c15) | 1;
2159        let mut v = [0.0f32; 8];
2160        for x in &mut v {
2161            s ^= s << 13;
2162            s ^= s >> 7;
2163            s ^= s << 17;
2164            *x = (s >> 40) as f32 / 4096.0 - 1.0;
2165        }
2166        v
2167    }
2168
2169    #[test]
2170    fn a_document_and_its_embedding_are_one_write() {
2171        let mut docs = Docs::new();
2172        docs.create_vector_index("$.embedding", 3).expect("index");
2173        for (id, v) in [
2174            ("a", [1.0, 0.0, 0.0]),
2175            ("b", [0.0, 1.0, 0.0]),
2176            ("c", [0.0, 0.0, 1.0]),
2177        ] {
2178            docs.put_bytes(id.as_bytes(), &item("en", &v)).expect("put");
2179        }
2180        assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 3);
2181
2182        // The answer is documents and not ids to go and look up somewhere else,
2183        // and it comes back nearest first.
2184        let mut got = Vec::new();
2185        let n = docs
2186            .nearest("$.embedding", &[0.9, 0.1, 0.0], 3, |id, doc, d| {
2187                let lang = doc
2188                    .get(b"lang")
2189                    .and_then(|v| v.as_text())
2190                    .map(str::to_owned);
2191                got.push((id.to_vec(), lang, d));
2192            })
2193            .expect("nearest");
2194        assert_eq!(n, 3);
2195        assert_eq!(got[0].0, b"a".to_vec());
2196        assert_eq!(got[0].1.as_deref(), Some("en"));
2197        assert!(got[0].2 <= got[1].2 && got[1].2 <= got[2].2);
2198
2199        // A path with no vector index on it says so rather than scanning.
2200        assert!(
2201            docs.nearest("$.lang", &[1.0, 0.0, 0.0], 1, |_, _, _| {})
2202                .is_err()
2203        );
2204    }
2205
2206    #[test]
2207    fn an_embedding_of_the_wrong_shape_fails_the_write_and_stores_nothing() {
2208        let mut docs = Docs::new();
2209        docs.create_vector_index("$.embedding", 3).expect("index");
2210        docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2211            .expect("put");
2212
2213        for wrong in [vec![1.0, 0.0], vec![1.0, 0.0, 0.0, 0.0]] {
2214            assert!(docs.put_bytes(b"b", &item("en", &wrong)).is_err());
2215        }
2216        assert!(docs.get(b"b").is_none(), "the write left nothing behind");
2217        assert_eq!(docs.len(), 1);
2218        assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 1);
2219    }
2220
2221    #[test]
2222    fn a_document_with_nothing_at_the_path_is_not_in_the_index() {
2223        let mut docs = Docs::new();
2224        docs.create_vector_index("$.embedding", 3).expect("index");
2225        docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2226            .expect("put");
2227        docs.put_bytes(b"b", &bare("en")).expect("put");
2228        assert_eq!(docs.len(), 2);
2229        assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 1);
2230
2231        // Rewriting a document without its embedding takes the old vector out
2232        // rather than leaving the one the last version had.
2233        docs.put_bytes(b"a", &bare("en")).expect("put");
2234        assert!(
2235            docs.vector_index("$.embedding")
2236                .expect("declared")
2237                .is_empty()
2238        );
2239        assert!(docs.get(b"a").is_some());
2240
2241        // And a removal takes it with it.
2242        docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2243            .expect("put");
2244        assert!(docs.remove(b"a"));
2245        assert!(
2246            docs.vector_index("$.embedding")
2247                .expect("declared")
2248                .is_empty()
2249        );
2250    }
2251
2252    #[test]
2253    fn a_filter_finds_the_nearest_match_and_not_the_nearest_that_matches() {
2254        let mut docs = Docs::new();
2255        docs.create_index("$.lang").expect("index");
2256        docs.create_vector_index("$.embedding", 8).expect("index");
2257
2258        // English documents spread about, and one French one that sits nowhere
2259        // near the query. Fewer under Miri: what the test needs is more English
2260        // documents near the query than the twenty the first search asks for,
2261        // so that the French one is not in the answer by accident.
2262        let count = if cfg!(miri) { 30u64 } else { 400 };
2263        for n in 0..count {
2264            let id = format!("en:{n}");
2265            docs.put_bytes(id.as_bytes(), &item("en", &spread(n)))
2266                .expect("put");
2267        }
2268        docs.put_bytes(b"fr", &item("fr", &spread(9_999)))
2269            .expect("put");
2270
2271        let q = spread(3);
2272        let mut top = Vec::new();
2273        docs.nearest("$.embedding", &q, 20, |id, _, _| top.push(id.to_vec()))
2274            .expect("nearest");
2275        assert_eq!(top[0], b"en:3".to_vec());
2276        assert!(
2277            !top.iter().any(|id| id == b"fr"),
2278            "searching and then filtering would have answered nothing"
2279        );
2280
2281        // Filtering inside the scan finds it anyway.
2282        let french = [("$.lang", Key::text("fr"))];
2283        let mut found = Vec::new();
2284        docs.nearest_where("$.embedding", &q, 5, &french, |id, _, _| {
2285            found.push(id.to_vec())
2286        })
2287        .expect("nearest");
2288        assert_eq!(found, [b"fr".to_vec()]);
2289
2290        // A path with no index on it cannot be filtered on, and says so.
2291        let nothing = [("$.topic", Key::text("finance"))];
2292        assert!(
2293            docs.nearest_where("$.embedding", &q, 5, &nothing, |_, _, _| {})
2294                .is_err()
2295        );
2296    }
2297
2298    #[test]
2299    fn declaring_either_index_last_gives_the_same_answers() {
2300        let q = spread(11);
2301        let french = [("$.lang", Key::text("fr"))];
2302
2303        // Vectors first, then the field the filter reads, so every tag was
2304        // written before there was anything to put in it.
2305        // Fewer documents under Miri, with the French ones proportionally as
2306        // often, so there are still the four the asserts below ask for. The two
2307        // collections are built the same way twice, so this test costs double
2308        // whatever the count is.
2309        let (count, every) = if cfg!(miri) { (20u64, 5) } else { (200, 50) };
2310        let mut late = Docs::new();
2311        late.create_vector_index("$.embedding", 8).expect("index");
2312        for n in 0..count {
2313            let lang = if n % every == 0 { "fr" } else { "en" };
2314            let id = format!("{n}");
2315            late.put_bytes(id.as_bytes(), &item(lang, &spread(n)))
2316                .expect("put");
2317        }
2318        late.create_index("$.lang").expect("index");
2319
2320        // The other way round, where every write already knew.
2321        let mut early = Docs::new();
2322        early.create_index("$.lang").expect("index");
2323        for n in 0..count {
2324            let lang = if n % every == 0 { "fr" } else { "en" };
2325            let id = format!("{n}");
2326            early
2327                .put_bytes(id.as_bytes(), &item(lang, &spread(n)))
2328                .expect("put");
2329        }
2330        early.create_vector_index("$.embedding", 8).expect("index");
2331
2332        let mut a = Vec::new();
2333        late.nearest_where("$.embedding", &q, 4, &french, |id, _, _| {
2334            a.push(id.to_vec())
2335        })
2336        .expect("nearest");
2337        let mut b = Vec::new();
2338        early
2339            .nearest_where("$.embedding", &q, 4, &french, |id, _, _| {
2340                b.push(id.to_vec())
2341            })
2342            .expect("nearest");
2343        assert_eq!(a.len(), 4, "there are four French documents to find");
2344        assert_eq!(a, b);
2345
2346        // Dropping the field index and declaring it again leaves the tags right.
2347        assert!(late.drop_index("$.lang"));
2348        late.create_index("$.lang").expect("index");
2349        let mut again = Vec::new();
2350        late.nearest_where("$.embedding", &q, 4, &french, |id, _, _| {
2351            again.push(id.to_vec());
2352        })
2353        .expect("nearest");
2354        assert_eq!(again, a);
2355    }
2356
2357    #[test]
2358    fn nearest_to_leaves_the_document_itself_out() {
2359        let mut docs = Docs::new();
2360        docs.create_vector_index("$.embedding", 3).expect("index");
2361        for (id, v) in [
2362            ("a", [1.0, 0.0, 0.0]),
2363            ("b", [0.9, 0.1, 0.0]),
2364            ("c", [0.0, 0.0, 1.0]),
2365        ] {
2366            docs.put_bytes(id.as_bytes(), &item("en", &v)).expect("put");
2367        }
2368
2369        let mut like = Vec::new();
2370        docs.nearest_to("$.embedding", b"a", 2, |id, _, _| like.push(id.to_vec()))
2371            .expect("nearest");
2372        assert_eq!(like, [b"b".to_vec(), b"c".to_vec()]);
2373
2374        // A document with no embedding has nothing to be like.
2375        docs.put_bytes(b"d", &bare("en")).expect("put");
2376        let mut none = 0;
2377        assert_eq!(
2378            docs.nearest_to("$.embedding", b"d", 2, |_, _, _| none += 1)
2379                .expect("nearest"),
2380            0
2381        );
2382    }
2383
2384    #[test]
2385    fn declaring_the_same_vector_index_again_rebuilds_nothing() {
2386        let mut docs = Docs::new();
2387        for n in 0..8u64 {
2388            let id = format!("{n}");
2389            docs.put_bytes(id.as_bytes(), &item("en", &spread(n)))
2390                .expect("put");
2391        }
2392        docs.create_vector_index("$.embedding", 8).expect("index");
2393        assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 8);
2394
2395        // The same declaration is nothing at all.
2396        docs.create_vector_index("$.embedding", 8).expect("again");
2397        assert_eq!(docs.vector_indexes().len(), 1);
2398
2399        // A different width is a different question, so it is rebuilt, and
2400        // documents whose embedding is not that wide fail the declaration. The
2401        // index that was already there is the one that is still there.
2402        assert!(docs.create_vector_index("$.embedding", 4).is_err());
2403        let still = docs.vector_index("$.embedding").expect("still declared");
2404        assert_eq!(still.dim(), 8);
2405        assert_eq!(still.len(), 8);
2406
2407        assert!(docs.drop_vector_index("$.embedding"));
2408        assert!(!docs.drop_vector_index("$.embedding"));
2409        assert!(docs.vector_indexes().is_empty());
2410        assert_eq!(docs.len(), 8, "the documents are untouched");
2411    }
2412
2413    #[test]
2414    fn clearing_a_collection_empties_the_vector_index_and_keeps_it_declared() {
2415        let mut docs = Docs::new();
2416        docs.create_vector_index("$.embedding", 3).expect("index");
2417        docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2418            .expect("put");
2419        let full = docs.memory_bytes();
2420
2421        docs.clear();
2422        assert!(docs.is_empty());
2423        assert!(
2424            docs.vector_index("$.embedding")
2425                .expect("declared")
2426                .is_empty()
2427        );
2428        assert!(docs.memory_bytes() < full);
2429
2430        docs.put_bytes(b"b", &item("en", &[0.0, 1.0, 0.0]))
2431            .expect("put");
2432        let mut got = Vec::new();
2433        docs.nearest("$.embedding", &[0.0, 1.0, 0.0], 1, |id, _, _| {
2434            got.push(id.to_vec())
2435        })
2436        .expect("nearest");
2437        assert_eq!(got, [b"b".to_vec()]);
2438    }
2439}