Skip to main content

yo/
doc.rs

1//! `Docs<T>`, the typed document handle, and the traits `#[derive(Yo)]` writes
2//! (`15` sections 2 and 4).
3//!
4//! A document collection is your own struct, stored as your own struct. There
5//! is no schema to declare, no JSON text to parse on either side, and no query
6//! language: a struct goes in, the same struct comes out, and the fields worth
7//! looking documents up by say so with an attribute.
8//!
9//! ```
10//! use yo::Yo;
11//!
12//! #[derive(Yo, Debug, PartialEq)]
13//! struct Order {
14//!     #[yo(id)]
15//!     id: u64,
16//!     #[yo(index)]
17//!     status: String,
18//!     #[yo(ordered)]
19//!     total: f64,
20//! }
21//!
22//! let db = yo::open(yo::MEMORY)?;
23//! let orders = db.docs::<Order>("orders")?;
24//!
25//! orders.put(&Order { id: 1, status: "open".to_owned(), total: 12.5 })?;
26//! orders.put(&Order { id: 2, status: "shipped".to_owned(), total: 99.0 })?;
27//!
28//! assert_eq!(orders.get(&1)?.unwrap().total, 12.5);
29//! assert_eq!(orders.find(Order::STATUS, "open")?.len(), 1);
30//! assert_eq!(orders.count(Order::STATUS, "shipped")?, 1);
31//! # Ok::<(), yo::Error>(())
32//! ```
33//!
34//! # The query is a constant, not a string
35//!
36//! `Order::STATUS` is a [`Path`] the derive wrote, and `Order::TOTAL` is an
37//! [`Ordered`], which is what a `#[yo(ordered)]` field gets. A field that is not
38//! indexed has no constant at all, so asking for one is a name that does not
39//! exist rather than a query that quietly turns into a scan. [`Docs::range`]
40//! takes an `Ordered` and nothing else, so asking an equality index for a range
41//! is a type error at the call site.
42//!
43//! ```compile_fail
44//! # use yo::Yo;
45//! # #[derive(Yo)]
46//! # struct Order { #[yo(id)] id: u64, #[yo(index)] status: String }
47//! # let db = yo::open(yo::MEMORY).unwrap();
48//! # let orders = db.docs::<Order>("orders").unwrap();
49//! // The index on status answers equality, so there is no range to walk.
50//! orders.range(Order::STATUS, "a".."z").unwrap();
51//! ```
52//!
53//! The value side is typed too, so comparing a number field against a string is
54//! the same kind of mistake and gets the same answer.
55//!
56//! ```compile_fail
57//! # use yo::Yo;
58//! # #[derive(Yo)]
59//! # struct Order { #[yo(id)] id: u64, #[yo(ordered)] total: f64 }
60//! # let db = yo::open(yo::MEMORY).unwrap();
61//! # let orders = db.docs::<Order>("orders").unwrap();
62//! orders.find(Order::TOTAL, "twelve").unwrap();
63//! ```
64//!
65//! # A range over a string field takes a pair of bounds
66//!
67//! `orders.range(Order::NAME, "a".."m")` does not compile, and the reason is not
68//! this crate. `Range<&str>` only implements `RangeBounds<str>` when `str` is
69//! sized, which it is not, so the standard library's own
70//! `BTreeMap<String, u8>::range("a".."m")` is rejected the same way. Writing the
71//! two ends out is what works there and it is what works here.
72//!
73//! ```
74//! # use std::ops::Bound;
75//! # use yo::Yo;
76//! # #[derive(Yo, Debug)]
77//! # struct Order { #[yo(id)] id: u64, #[yo(ordered)] name: String }
78//! # let db = yo::open(yo::MEMORY)?;
79//! # let orders = db.docs::<Order>("orders")?;
80//! # orders.put(&Order { id: 1, name: "banana".to_owned() })?;
81//! # orders.put(&Order { id: 2, name: "quince".to_owned() })?;
82//! let early = orders.range(Order::NAME, (Bound::Included("a"), Bound::Excluded("m")))?;
83//! assert_eq!(early.len(), 1);
84//! # Ok::<(), yo::Error>(())
85//! ```
86//!
87//! A range over a number field is written the way anyone would write it, because
88//! the numbers are sized and `0.0..50.0` is a `RangeBounds<f64>` already.
89//!
90//! # What a field can be
91//!
92//! [`Field`] is the list, and it is the JSON types rather than the Rust ones,
93//! because a document is JSON shaped whatever it was written from. The integers
94//! and floats, `bool`, `String`, `Option<T>` for a field that may be absent,
95//! `Vec<T>` for a list, and any other type that derives `Yo`, which nests.
96//!
97//! An integer is stored as an `i64`, which is the one number type JSON has, so
98//! a `u64` above `i64::MAX` is refused on the way in rather than silently
99//! rounded through a float.
100//!
101//! # An embedding is a field
102//!
103//! `#[yo(vector = 384)]` on a `Vec<f32>` gives that path a vector index, and the
104//! field is still an ordinary field: it is written with the document, it comes
105//! back with the document, and there is no second collection to keep in step.
106//! The constant the derive writes is a [`Vector`], so [`Docs::near`] takes it
107//! and [`Docs::find`] does not.
108//!
109//! ```
110//! use yo::Yo;
111//!
112//! #[derive(Yo, Debug)]
113//! struct Note {
114//!     #[yo(id)]
115//!     id: u64,
116//!     #[yo(index)]
117//!     lang: String,
118//!     #[yo(vector = 3)]
119//!     embedding: Vec<f32>,
120//! }
121//!
122//! let db = yo::open(yo::MEMORY)?;
123//! let notes = db.docs::<Note>("notes")?;
124//! for (id, lang, v) in [
125//!     (1u64, "en", [1.0, 0.0, 0.0]),
126//!     (2, "fr", [0.9, 0.1, 0.0]),
127//!     (3, "en", [0.0, 0.0, 1.0]),
128//! ] {
129//!     notes.put(&Note { id, lang: lang.to_owned(), embedding: v.to_vec() })?;
130//! }
131//!
132//! let close = notes.nearest(Note::EMBEDDING, &[1.0, 0.05, 0.0], 2)?;
133//! assert_eq!(close.iter().map(|n| n.id).collect::<Vec<_>>(), [1, 2]);
134//!
135//! // The same search, narrowed by another indexed field.
136//! let english = notes
137//!     .near(Note::EMBEDDING, &[1.0, 0.05, 0.0])
138//!     .filter(Note::LANG, "en")
139//!     .take(2)?;
140//! assert_eq!(english.iter().map(|n| n.id).collect::<Vec<_>>(), [1, 3]);
141//! # Ok::<(), yo::Error>(())
142//! ```
143//!
144//! The filter is decided inside the scan and not over the answers, so asking for
145//! two English notes gives the two nearest English notes rather than whichever
146//! of the nearest few happened to be English. That distinction is the whole
147//! reason the two live in one collection, and [`yo_doc::vector`] has the rest of
148//! it.
149
150use core::marker::PhantomData;
151use core::ops::{Bound, RangeBounds};
152
153use yo_common::{Code, Error, Result};
154use yo_shape::{Shape, Tag};
155
156use crate::db::Handle;
157
158pub use yo_doc::{Builder, Doc, IndexKind, Key};
159
160/// A type that can be a field of a document.
161///
162/// The encoding is YOJB and not a Rust layout, so what goes in the store is
163/// what a document is: an object with named fields, readable by the RESP
164/// surface and by another language's binding without either of them knowing
165/// what Rust is.
166pub trait Field: Shape + Sized {
167    /// Write this value into the document being built.
168    ///
169    /// # Errors
170    ///
171    /// [`Code::Invalid`] for a value the document encoding cannot hold, which
172    /// is a `u64` past `i64::MAX` and nothing else so far.
173    fn write(&self, b: &mut Builder) -> Result<()>;
174
175    /// Read this value back out.
176    ///
177    /// # Errors
178    ///
179    /// [`Code::Corrupt`] when the stored value is not this type, which means
180    /// the collection disagrees with its own shape.
181    fn read(d: Doc<'_>) -> Result<Self>;
182
183    /// What to do when the field is not in the document at all.
184    ///
185    /// An error for everything except [`Option`], which is the whole point of
186    /// having an `Option`: a field that may be absent says so in the type, and
187    /// every other field being absent is a document that does not match the
188    /// shape it was stored under.
189    ///
190    /// # Errors
191    ///
192    /// [`Code::Corrupt`], unless the type is an `Option`.
193    fn missing(name: &str) -> Result<Self> {
194        Err(Error::fmt(
195            Code::Corrupt,
196            format_args!(
197                "this document has no {name}, and the field is not an Option. Either the collection holds something written under another shape, or the field was added without a default"
198            ),
199        ))
200    }
201}
202
203/// A value that can be an index key.
204///
205/// Separate from [`Field`] because a lookup takes the borrowed form, the same
206/// way `HashMap::get` does, so a `String` field is searched with `&str` and not
207/// with a `String` built for the length of one call.
208pub trait Query {
209    /// The key this value is filed under in an index of `kind`, or `None` if an
210    /// index of that kind does not file this type at all.
211    fn key(&self, kind: IndexKind) -> Option<Key>;
212}
213
214/// How a field's type is written in a query.
215///
216/// An associated type rather than a `Borrow` bound on the call, because a
217/// `Borrow` bound leaves the compiler two ways to read `"a".."z"` and it picks
218/// the wrong one. This way the borrowed form follows from the field's type and
219/// there is nothing to infer.
220pub trait Asked: Query {
221    /// The borrowed form, which is `str` for a `String` and the type itself for
222    /// everything else.
223    type Ask: Query + ?Sized;
224}
225
226macro_rules! asks_for_itself {
227    ($($t:ty),* $(,)?) => {
228        $(impl Asked for $t {
229            type Ask = $t;
230        })*
231    };
232}
233
234asks_for_itself!(i8, i16, i32, i64, u8, u16, u32, u64, f32, f64, bool);
235
236impl Asked for String {
237    type Ask = str;
238}
239
240/// A type that is a whole document: a [`Field`] with an id and its indexes.
241///
242/// Written by `#[derive(Yo)]` from the field marked `#[yo(id)]`.
243#[diagnostic::on_unimplemented(
244    message = "`{Self}` is not a document",
245    label = "this type has no id",
246    note = "add `#[derive(Yo)]` to it and mark one field `#[yo(id)]`, which is what a document is stored under"
247)]
248pub trait Document: Field + Indexed {
249    /// The type of the field marked `#[yo(id)]`.
250    type Id: Field + Asked;
251
252    /// This document's id.
253    fn id(&self) -> &Self::Id;
254}
255
256/// The indexes a type declares.
257///
258/// Written by `#[derive(Yo)]` for every type it is put on, whether or not the
259/// type has an id, because an edge type declares indexes and has no id. That is
260/// the whole reason this is a trait of its own rather than a constant on
261/// [`Document`].
262pub trait Indexed {
263    /// The paths this type asks to be indexed, and how.
264    const INDEXES: &'static [(&'static str, IndexKind)];
265
266    /// The paths that hold an embedding, and how wide it is.
267    ///
268    /// Defaulted to nothing, so a type written before vector indexes existed
269    /// and a type that has no embedding both say the same thing without saying
270    /// anything.
271    const VECTORS: &'static [(&'static str, usize)] = &[];
272}
273
274/// A path into a document, what its index can be asked, and the type of the
275/// value that lives there.
276///
277/// Written by `#[derive(Yo)]` as a constant per indexed field, so a query names
278/// the field rather than spelling a string the compiler cannot check. A field
279/// marked `#[yo(ordered)]` gets an [`Ordered`] instead, which is the same thing
280/// with ranges on it.
281pub struct Path<T, V> {
282    path: &'static str,
283    kind: IndexKind,
284    /// `fn() -> (T, V)` so the constant's auto traits do not come from what it
285    /// points at, which lets it be a `const` in any type.
286    marker: PhantomData<fn() -> (T, V)>,
287}
288
289impl<T, V> Clone for Path<T, V> {
290    fn clone(&self) -> Path<T, V> {
291        *self
292    }
293}
294
295impl<T, V> Copy for Path<T, V> {}
296
297impl<T, V> core::fmt::Debug for Path<T, V> {
298    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
299        f.debug_struct("Path")
300            .field("path", &self.path)
301            .field("kind", &self.kind)
302            .finish()
303    }
304}
305
306impl<T, V> Path<T, V> {
307    /// A path and what its index answers.
308    ///
309    /// The derive calls this. Calling it by hand is allowed and is how a path
310    /// into a nested object is reached until the derive learns to follow one,
311    /// but nothing checks that the collection has the index you named until the
312    /// query runs.
313    #[must_use]
314    pub const fn new(path: &'static str, kind: IndexKind) -> Path<T, V> {
315        Path {
316            path,
317            kind,
318            marker: PhantomData,
319        }
320    }
321
322    /// The path, as `$.status`.
323    #[must_use]
324    pub const fn path(&self) -> &'static str {
325        self.path
326    }
327
328    /// What the index on this path can be asked.
329    #[must_use]
330    pub const fn kind(&self) -> IndexKind {
331        self.kind
332    }
333}
334
335/// A path whose index keeps its keys in order, so it answers ranges as well as
336/// equality.
337///
338/// A separate type rather than a flag on [`Path`], because which questions a
339/// path can answer is decided when the type is written and there is no reason
340/// for the compiler not to know it. [`Docs::range`] takes one of these and
341/// nothing else, so asking an equality index for a range is a type error at the
342/// call site rather than a message at run time.
343pub struct Ordered<T, V> {
344    path: Path<T, V>,
345}
346
347impl<T, V> Clone for Ordered<T, V> {
348    fn clone(&self) -> Ordered<T, V> {
349        *self
350    }
351}
352
353impl<T, V> Copy for Ordered<T, V> {}
354
355impl<T, V> core::fmt::Debug for Ordered<T, V> {
356    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
357        f.debug_struct("Ordered")
358            .field("path", &self.path.path)
359            .finish()
360    }
361}
362
363impl<T, V> Ordered<T, V> {
364    /// A path whose index is ordered.
365    ///
366    /// The derive calls this.
367    #[must_use]
368    pub const fn new(path: &'static str) -> Ordered<T, V> {
369        Ordered {
370            path: Path::new(path, IndexKind::Ordered),
371        }
372    }
373
374    /// The path, as `$.total`.
375    #[must_use]
376    pub const fn path(&self) -> &'static str {
377        self.path.path
378    }
379}
380
381/// An ordered path answers equality too, so everything that takes a [`Path`]
382/// takes one of these.
383impl<T, V> From<Ordered<T, V>> for Path<T, V> {
384    fn from(o: Ordered<T, V>) -> Path<T, V> {
385        o.path
386    }
387}
388
389/// A path that holds an embedding, and how wide it is.
390///
391/// A third type rather than another kind on [`Path`], for the same reason
392/// [`Ordered`] is a second one: what a path can be asked is decided when the
393/// type is written. [`Docs::near`] takes one of these and nothing else, so
394/// asking an equality index for the nearest anything is a type error at the call
395/// site, and so is handing a vector path to [`Docs::find`].
396pub struct Vector<T> {
397    path: &'static str,
398    dim: usize,
399    marker: PhantomData<fn() -> T>,
400}
401
402impl<T> Clone for Vector<T> {
403    fn clone(&self) -> Vector<T> {
404        *self
405    }
406}
407
408impl<T> Copy for Vector<T> {}
409
410impl<T> core::fmt::Debug for Vector<T> {
411    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
412        f.debug_struct("Vector")
413            .field("path", &self.path)
414            .field("dim", &self.dim)
415            .finish()
416    }
417}
418
419impl<T> Vector<T> {
420    /// A path holding a `dim` wide embedding.
421    ///
422    /// The derive calls this.
423    #[must_use]
424    pub const fn new(path: &'static str, dim: usize) -> Vector<T> {
425        Vector {
426            path,
427            dim,
428            marker: PhantomData,
429        }
430    }
431
432    /// The path, as `$.embedding`.
433    #[must_use]
434    pub const fn path(&self) -> &'static str {
435        self.path
436    }
437
438    /// How many coordinates the embedding there has.
439    #[must_use]
440    pub const fn dim(&self) -> usize {
441        self.dim
442    }
443}
444
445/// A collection of `T`.
446///
447/// Cheap to clone and cheap to keep around, the same way [`crate::Map`] is: the
448/// handle is a pointer and an index, and every clone is the same collection.
449pub struct Docs<T> {
450    db: Handle,
451    at: usize,
452    tag: Tag,
453    marker: PhantomData<fn() -> T>,
454}
455
456impl<T> Clone for Docs<T> {
457    fn clone(&self) -> Docs<T> {
458        Docs {
459            db: self.db.clone(),
460            at: self.at,
461            tag: self.tag,
462            marker: PhantomData,
463        }
464    }
465}
466
467impl<T> core::fmt::Debug for Docs<T> {
468    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
469        let name = self
470            .db
471            .read(|inner| Ok(inner.collections[self.at].name.clone()))
472            .unwrap_or_else(|_| "?".to_owned());
473        f.debug_struct("Docs").field("name", &name).finish()
474    }
475}
476
477impl<T: Document> Docs<T> {
478    pub(crate) fn new(db: Handle, at: usize, tag: Tag) -> Docs<T> {
479        Docs {
480            db,
481            at,
482            tag,
483            marker: PhantomData,
484        }
485    }
486
487    /// The name this collection was opened under.
488    ///
489    /// # Errors
490    ///
491    /// [`Code::Invalid`] if called from inside a callback that is already
492    /// holding this database.
493    pub fn name(&self) -> Result<String> {
494        self.db
495            .read(|inner| Ok(inner.collections[self.at].name.clone()))
496    }
497
498    /// This collection's shape tag.
499    #[must_use]
500    pub fn tag(&self) -> Tag {
501        self.tag
502    }
503
504    /// Store a document, replacing whatever was under its id.
505    ///
506    /// Answers whether the id was new. Every index the type declares is brought
507    /// up to date in the same call, and the old document is taken back out of
508    /// them first, so an overwrite cannot leave a stale posting behind.
509    ///
510    /// # Errors
511    ///
512    /// [`Code::Invalid`] for an id that cannot be a key, and [`Code::Full`] for
513    /// a value at an indexed path that is too long to be one.
514    pub fn put(&self, doc: &T) -> Result<bool> {
515        let id = key_of(doc.id(), IndexKind::Equality, "the id")?;
516        self.write(|c| {
517            c.scratch.clear();
518            Field::write(doc, &mut c.scratch)?;
519            let bytes = c.scratch.finish()?;
520            c.docs.put_bytes(id.as_bytes(), bytes)
521        })
522    }
523
524    /// Read a document by its id.
525    ///
526    /// # Errors
527    ///
528    /// [`Code::Corrupt`] if the stored document is not a `T`.
529    pub fn get(&self, id: &<T::Id as Asked>::Ask) -> Result<Option<T>> {
530        let id = key_of(id, IndexKind::Equality, "the id")?;
531        self.read(|docs| match docs.get(id.as_bytes()) {
532            Some(doc) => T::read(doc).map(Some),
533            None => Ok(None),
534        })
535    }
536
537    /// Whether an id is in the collection, without reading the document.
538    ///
539    /// # Errors
540    ///
541    /// [`Code::Invalid`] for an id that cannot be a key.
542    pub fn contains(&self, id: &<T::Id as Asked>::Ask) -> Result<bool> {
543        let id = key_of(id, IndexKind::Equality, "the id")?;
544        self.read(|docs| Ok(docs.contains(id.as_bytes())))
545    }
546
547    /// Take a document out, answering whether it was there.
548    ///
549    /// # Errors
550    ///
551    /// [`Code::Invalid`] for an id that cannot be a key.
552    pub fn remove(&self, id: &<T::Id as Asked>::Ask) -> Result<bool> {
553        let id = key_of(id, IndexKind::Equality, "the id")?;
554        self.write(|c| Ok(c.docs.remove(id.as_bytes())))
555    }
556
557    /// How many documents there are.
558    ///
559    /// # Errors
560    ///
561    /// [`Code::Invalid`] if called from inside a callback that is already
562    /// holding this database.
563    pub fn len(&self) -> Result<usize> {
564        self.read(|docs| Ok(docs.len()))
565    }
566
567    /// Whether the collection is empty.
568    ///
569    /// # Errors
570    ///
571    /// The same as [`Docs::len`].
572    pub fn is_empty(&self) -> Result<bool> {
573        self.read(|docs| Ok(docs.is_empty()))
574    }
575
576    /// Every document, in no particular order.
577    ///
578    /// A walk of the whole collection, which is what it says it is. The indexed
579    /// calls are the ones with a cost model.
580    ///
581    /// # Errors
582    ///
583    /// [`Code::Corrupt`] if any stored document is not a `T`.
584    pub fn all(&self) -> Result<Vec<T>> {
585        self.read(|docs| {
586            let mut out = Vec::with_capacity(docs.len());
587            for (_, doc) in docs.iter() {
588                out.push(T::read(doc)?);
589            }
590            Ok(out)
591        })
592    }
593
594    /// Every document whose value at `path` is `value`.
595    ///
596    /// One probe of the index and one probe of the primary table per document
597    /// in the answer, so the cost is the size of the answer rather than the
598    /// size of the collection.
599    ///
600    /// # Errors
601    ///
602    /// [`Code::Invalid`] if the collection has no index on that path, because a
603    /// query that quietly turns into a scan is the thing this API exists not to
604    /// do.
605    pub fn find<V: Asked>(&self, path: impl Into<Path<T, V>>, value: &V::Ask) -> Result<Vec<T>> {
606        let path = path.into();
607        let key = key_of(value, path.kind, path.path)?;
608        self.read(|docs| {
609            let mut out = Vec::new();
610            let mut bad = Ok(());
611            docs.find(path.path, &key, |_, doc| {
612                if bad.is_ok() {
613                    match T::read(doc) {
614                        Ok(v) => out.push(v),
615                        Err(e) => bad = Err(e),
616                    }
617                }
618            })?;
619            bad?;
620            Ok(out)
621        })
622    }
623
624    /// How many documents have `value` at `path`, without reading any of them.
625    ///
626    /// The number to sort filters by before intersecting them, and it is a
627    /// probe rather than a walk.
628    ///
629    /// # Errors
630    ///
631    /// The same as [`Docs::find`].
632    pub fn count<V: Asked>(&self, path: impl Into<Path<T, V>>, value: &V::Ask) -> Result<usize> {
633        let path = path.into();
634        let key = key_of(value, path.kind, path.path)?;
635        self.read(|docs| docs.count(path.path, &key))
636    }
637
638    /// Every document whose value at `path` falls in `range`, smallest first.
639    ///
640    /// The bounds are the ordinary Rust range syntax, so `..`, `a..b`, `a..=b`
641    /// and `..b` all work and mean what they say.
642    ///
643    /// # Errors
644    ///
645    /// [`Code::Invalid`] if the collection has no index on that path. An index
646    /// that answers equality only cannot get here at all, because [`Ordered`] is
647    /// a different type from [`Path`] and this takes one of them.
648    pub fn range<V: Asked, R: RangeBounds<V::Ask>>(
649        &self,
650        path: Ordered<T, V>,
651        range: R,
652    ) -> Result<Vec<T>> {
653        let path = path.path();
654        let (lo, hi) = bounds(&range, path)?;
655        self.read(|docs| {
656            let mut out = Vec::new();
657            let mut bad = Ok(());
658            docs.range(path, as_ref(&lo), as_ref(&hi), |_, doc| {
659                if bad.is_ok() {
660                    match T::read(doc) {
661                        Ok(v) => out.push(v),
662                        Err(e) => bad = Err(e),
663                    }
664                }
665            })?;
666            bad?;
667            Ok(out)
668        })
669    }
670
671    /// [`Docs::range`] backwards, largest value first.
672    ///
673    /// # Errors
674    ///
675    /// The same as [`Docs::range`].
676    pub fn range_rev<V: Asked, R: RangeBounds<V::Ask>>(
677        &self,
678        path: Ordered<T, V>,
679        range: R,
680    ) -> Result<Vec<T>> {
681        let path = path.path();
682        let (lo, hi) = bounds(&range, path)?;
683        self.read(|docs| {
684            let mut out = Vec::new();
685            let mut bad = Ok(());
686            docs.range_rev(path, as_ref(&lo), as_ref(&hi), |_, doc| {
687                if bad.is_ok() {
688                    match T::read(doc) {
689                        Ok(v) => out.push(v),
690                        Err(e) => bad = Err(e),
691                    }
692                }
693            })?;
694            bad?;
695            Ok(out)
696        })
697    }
698
699    /// How many documents fall in `range` at `path`, without reading any.
700    ///
701    /// This reads the distinct values in the range rather than the documents,
702    /// so a range covering a million documents under a hundred values costs a
703    /// hundred.
704    ///
705    /// # Errors
706    ///
707    /// The same as [`Docs::range`].
708    pub fn count_range<V: Asked, R: RangeBounds<V::Ask>>(
709        &self,
710        path: Ordered<T, V>,
711        range: R,
712    ) -> Result<usize> {
713        let path = path.path();
714        let (lo, hi) = bounds(&range, path)?;
715        self.read(|docs| docs.count_range(path, as_ref(&lo), as_ref(&hi)))
716    }
717
718    /// The `k` documents whose embedding at `path` is nearest to `q`, nearest
719    /// first.
720    ///
721    /// # Errors
722    ///
723    /// [`Code::Invalid`] if `q` is not as wide as the path says, and
724    /// [`Code::Corrupt`] if a stored document is not a `T`.
725    pub fn nearest(&self, path: Vector<T>, q: &[f32], k: usize) -> Result<Vec<T>> {
726        self.near(path, q).take(k)
727    }
728
729    /// The `k` documents most like the one under `id`, that one left out.
730    ///
731    /// More like this, which is the question a collection with embeddings in it
732    /// is really for, and it does not make the caller read the document back
733    /// out to get its vector first. A document with no embedding has nothing to
734    /// be like, so this answers nothing rather than an error.
735    ///
736    /// # Errors
737    ///
738    /// [`Code::Invalid`] for an id that cannot be a key, and [`Code::Corrupt`]
739    /// if a stored document is not a `T`.
740    pub fn nearest_to(
741        &self,
742        path: Vector<T>,
743        id: &<T::Id as Asked>::Ask,
744        k: usize,
745    ) -> Result<Vec<T>> {
746        let id = key_of(id, IndexKind::Equality, "the id")?;
747        self.read(|docs| {
748            let mut out = Vec::new();
749            let mut bad = Ok(());
750            docs.nearest_to(path.path(), id.as_bytes(), k, |_, doc, _| {
751                collect::<T>(&mut out, &mut bad, doc);
752            })?;
753            bad?;
754            Ok(out)
755        })
756    }
757
758    /// A nearest neighbour search that other indexed fields can narrow.
759    ///
760    /// ```
761    /// # use yo::Yo;
762    /// #[derive(Yo, Debug)]
763    /// struct Note {
764    ///     #[yo(id)]
765    ///     id: u64,
766    ///     #[yo(index)]
767    ///     lang: String,
768    ///     #[yo(vector = 3)]
769    ///     embedding: Vec<f32>,
770    /// }
771    ///
772    /// # let db = yo::open(yo::MEMORY)?;
773    /// # let notes = db.docs::<Note>("notes")?;
774    /// # for (id, lang, v) in [
775    /// #     (1u64, "en", [1.0, 0.0, 0.0]),
776    /// #     (2, "fr", [0.9, 0.1, 0.0]),
777    /// #     (3, "en", [0.0, 0.0, 1.0]),
778    /// # ] {
779    /// #     notes.put(&Note { id, lang: lang.to_owned(), embedding: v.to_vec() })?;
780    /// # }
781    /// let close = notes
782    ///     .near(Note::EMBEDDING, &[1.0, 0.05, 0.0])
783    ///     .filter(Note::LANG, "en")
784    ///     .take(2)?;
785    /// assert_eq!(close.iter().map(|n| n.id).collect::<Vec<_>>(), [1, 3]);
786    /// # Ok::<(), yo::Error>(())
787    /// ```
788    ///
789    /// Every filter is decided inside the scan rather than over the answers, so
790    /// asking for two English notes gives the two nearest English notes and not
791    /// whichever of the nearest few happened to be English. See
792    /// [`yo_doc::vector`] for the encoding and for the one direction it is not
793    /// exact in.
794    pub fn near<'a>(&'a self, path: Vector<T>, q: &'a [f32]) -> Near<'a, T> {
795        Near {
796            docs: self,
797            path,
798            q,
799            want: Vec::new(),
800            bad: None,
801        }
802    }
803
804    /// What this collection is holding, documents and indexes together.
805    ///
806    /// # Errors
807    ///
808    /// [`Code::Invalid`] if called from inside a callback that is already
809    /// holding this database.
810    pub fn memory_bytes(&self) -> Result<usize> {
811        self.read(|docs| Ok(docs.memory_bytes()))
812    }
813
814    fn read<R>(&self, f: impl FnOnce(&yo_doc::Docs) -> Result<R>) -> Result<R> {
815        self.db
816            .read(|inner| f(inner.collections[self.at].data.docs()))
817    }
818
819    fn write<R>(&self, f: impl FnOnce(&mut Documents) -> Result<R>) -> Result<R> {
820        self.db
821            .write(|inner| f(inner.collections[self.at].data.docs_mut()))
822    }
823}
824
825/// A nearest neighbour search being put together, from [`Docs::near`].
826///
827/// A builder rather than a method with a list of filters, because the filters
828/// are over different fields with different types and a slice of them would
829/// have to give that up. Turning a value into an index key can fail, so a
830/// filter that cannot be one is kept here and handed over at the end rather
831/// than making every step return a `Result`.
832pub struct Near<'a, T> {
833    docs: &'a Docs<T>,
834    path: Vector<T>,
835    q: &'a [f32],
836    want: Vec<(&'static str, Key)>,
837    bad: Option<Error>,
838}
839
840impl<T> core::fmt::Debug for Near<'_, T> {
841    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
842        f.debug_struct("Near")
843            .field("path", &self.path.path())
844            .field("filters", &self.want.len())
845            .finish()
846    }
847}
848
849impl<'a, T: Document> Near<'a, T> {
850    /// Only documents whose value at `path` is `value`.
851    ///
852    /// Two of these means both, so it is a conjunction and not a choice. The
853    /// path has to carry an ordinary index, because what the scan tests is the
854    /// keys that index filed the document under.
855    #[must_use]
856    pub fn filter<V: Asked>(mut self, path: impl Into<Path<T, V>>, value: &V::Ask) -> Near<'a, T> {
857        let path = path.into();
858        match key_of(value, path.kind, path.path) {
859            Ok(key) => self.want.push((path.path, key)),
860            Err(e) => self.bad = self.bad.or(Some(e)),
861        }
862        self
863    }
864
865    /// The `k` nearest that pass every filter, nearest first.
866    ///
867    /// # Errors
868    ///
869    /// [`Code::Invalid`] if a filter names a path with no index on it, if a
870    /// filter value cannot be an index key, or if the query vector is not as
871    /// wide as the path says. [`Code::Corrupt`] if a stored document is not a
872    /// `T`.
873    pub fn take(self, k: usize) -> Result<Vec<T>> {
874        Ok(self.scored(k)?.into_iter().map(|(doc, _)| doc).collect())
875    }
876
877    /// The same, each document with how far it is.
878    ///
879    /// The distance is what the collection measures, so for the cosine default
880    /// it is one minus the cosine and nearer is smaller. It is measured against
881    /// the full precision vector rather than against the code.
882    ///
883    /// # Errors
884    ///
885    /// The same as [`Near::take`].
886    pub fn scored(self, k: usize) -> Result<Vec<(T, f32)>> {
887        if let Some(e) = self.bad {
888            return Err(e);
889        }
890        let (path, q, want) = (self.path.path(), self.q, self.want);
891        self.docs.read(|docs| {
892            let mut out = Vec::new();
893            let mut bad = Ok(());
894            docs.nearest_where(path, q, k, &want, |_, doc, at| {
895                if bad.is_ok() {
896                    match T::read(doc) {
897                        Ok(v) => out.push((v, at)),
898                        Err(e) => bad = Err(e),
899                    }
900                }
901            })?;
902            bad?;
903            Ok(out)
904        })
905    }
906}
907
908/// Read one answer into `out`, keeping the first failure rather than the last.
909fn collect<T: Document>(out: &mut Vec<T>, bad: &mut Result<()>, doc: Doc<'_>) {
910    if bad.is_ok() {
911        match T::read(doc) {
912            Ok(v) => out.push(v),
913            Err(e) => *bad = Err(e),
914        }
915    }
916}
917
918/// The documents of one collection, and the builder a write goes through.
919///
920/// The builder lives here rather than on the stack of [`Docs::put`] so that a
921/// write reuses the buffer it filled last time and allocates nothing.
922pub(crate) struct Documents {
923    pub(crate) docs: yo_doc::Docs,
924    pub(crate) scratch: Builder,
925}
926
927impl Documents {
928    pub(crate) fn new() -> Documents {
929        Documents {
930            docs: yo_doc::Docs::new(),
931            scratch: Builder::new(),
932        }
933    }
934}
935
936/// The key a value is filed under, or the sentence saying why it has none.
937pub(crate) fn key_of<Q: Query + ?Sized>(value: &Q, kind: IndexKind, what: &str) -> Result<Key> {
938    // The kind matters because a text index files words rather than whole
939    // strings, so a query against one has to be folded the same way the write
940    // was. Everything else asks its value for the key it always gives.
941    let key = value.key(kind).ok_or_else(|| {
942        let why = if kind == IndexKind::Text {
943            "a text index holds one word at a time, and this is not one word"
944        } else {
945            "an index does not file this type, so it cannot be looked up"
946        };
947        Error::fmt(Code::Invalid, format_args!("{what}: {why}"))
948    })?;
949    if key.is_too_long() {
950        return Err(Error::fmt(
951            Code::Full,
952            format_args!(
953                "{what} is longer than {} bytes, which is as long as a key can be",
954                yo_doc::KEY_MAX
955            ),
956        ));
957    }
958    Ok(key)
959}
960
961/// Turn a Rust range over the query type into the pair of key bounds the index
962/// walks between.
963fn bounds<Q, R>(range: &R, path: &str) -> Result<(Bound<Key>, Bound<Key>)>
964where
965    Q: Query + ?Sized,
966    R: RangeBounds<Q>,
967{
968    Ok((
969        one(range.start_bound(), path)?,
970        one(range.end_bound(), path)?,
971    ))
972}
973
974fn one<Q: Query + ?Sized>(b: Bound<&Q>, path: &str) -> Result<Bound<Key>> {
975    Ok(match b {
976        Bound::Included(v) => Bound::Included(key_of(v, IndexKind::Ordered, path)?),
977        Bound::Excluded(v) => Bound::Excluded(key_of(v, IndexKind::Ordered, path)?),
978        Bound::Unbounded => Bound::Unbounded,
979    })
980}
981
982fn as_ref(b: &Bound<Key>) -> Bound<&Key> {
983    match b {
984        Bound::Included(k) => Bound::Included(k),
985        Bound::Excluded(k) => Bound::Excluded(k),
986        Bound::Unbounded => Bound::Unbounded,
987    }
988}
989
990/// Read one field out of a document, which is what the derive calls per field.
991///
992/// # Errors
993///
994/// [`Code::Corrupt`] if the field is missing and its type is not an `Option`,
995/// or if it is there and is the wrong type.
996pub fn at<V: Field>(d: Doc<'_>, name: &str) -> Result<V> {
997    match d.get(name.as_bytes()) {
998        Some(at) => V::read(at),
999        None => V::missing(name),
1000    }
1001}
1002
1003/// Check that what is stored under this collection is an object at all, which
1004/// is what the derive calls before it reads the fields.
1005///
1006/// # Errors
1007///
1008/// [`Code::Corrupt`] for anything that is not an object.
1009pub fn expect_object(d: Doc<'_>, name: &str) -> Result<()> {
1010    if d.kind() == yo_doc::Kind::Object {
1011        return Ok(());
1012    }
1013    Err(Error::fmt(
1014        Code::Corrupt,
1015        format_args!("a {name} in this collection is stored as {:?}", d.kind()),
1016    ))
1017}
1018
1019fn not_a(want: &str, d: Doc<'_>) -> Error {
1020    Error::fmt(
1021        Code::Corrupt,
1022        format_args!(
1023            "this field should be a {want} and is stored as {:?}",
1024            d.kind()
1025        ),
1026    )
1027}
1028
1029macro_rules! ints {
1030    ($($t:ty),* $(,)?) => {
1031        $(
1032            impl Field for $t {
1033                fn write(&self, b: &mut Builder) -> Result<()> {
1034                    b.int(i64::from(*self))
1035                }
1036
1037                fn read(d: Doc<'_>) -> Result<$t> {
1038                    let n = d.as_int().ok_or_else(|| not_a(stringify!($t), d))?;
1039                    <$t>::try_from(n).map_err(|_| {
1040                        Error::fmt(
1041                            Code::Corrupt,
1042                            format_args!("{n} does not fit in a {}", stringify!($t)),
1043                        )
1044                    })
1045                }
1046            }
1047
1048            impl Query for $t {
1049                fn key(&self, _kind: IndexKind) -> Option<Key> {
1050                    Some(Key::int(i64::from(*self)))
1051                }
1052            }
1053        )*
1054    };
1055}
1056
1057ints!(i8, i16, i32, i64, u8, u16, u32);
1058
1059/// A `u64` is the one integer that does not fit, because JSON has one number
1060/// type and it is signed. Anything past `i64::MAX` is refused on the way in
1061/// rather than rounded through a float on the way out.
1062impl Field for u64 {
1063    fn write(&self, b: &mut Builder) -> Result<()> {
1064        match i64::try_from(*self) {
1065            Ok(n) => b.int(n),
1066            Err(_) => Err(Error::fmt(
1067                Code::Invalid,
1068                format_args!(
1069                    "{self} is past i64::MAX, and a document holds one number type, which is signed"
1070                ),
1071            )),
1072        }
1073    }
1074
1075    fn read(d: Doc<'_>) -> Result<u64> {
1076        let n = d.as_int().ok_or_else(|| not_a("u64", d))?;
1077        u64::try_from(n).map_err(|_| {
1078            Error::fmt(
1079                Code::Corrupt,
1080                format_args!("{n} is negative and this field is a u64"),
1081            )
1082        })
1083    }
1084}
1085
1086impl Query for u64 {
1087    fn key(&self, _kind: IndexKind) -> Option<Key> {
1088        i64::try_from(*self).ok().map(Key::int)
1089    }
1090}
1091
1092macro_rules! floats {
1093    ($($t:ty),* $(,)?) => {
1094        $(
1095            impl Field for $t {
1096                fn write(&self, b: &mut Builder) -> Result<()> {
1097                    b.float(f64::from(*self))
1098                }
1099
1100                fn read(d: Doc<'_>) -> Result<$t> {
1101                    // An integer reads back as a float, because a whole number
1102                    // written as a float is stored as an integer and refusing
1103                    // it here would make a round trip fail on 12.0.
1104                    match (d.as_float(), d.as_int()) {
1105                        (Some(v), _) => Ok(v as $t),
1106                        (None, Some(n)) => Ok(n as $t),
1107                        (None, None) => Err(not_a(stringify!($t), d)),
1108                    }
1109                }
1110            }
1111
1112            impl Query for $t {
1113                fn key(&self, _kind: IndexKind) -> Option<Key> {
1114                    Some(Key::float(f64::from(*self)))
1115                }
1116            }
1117        )*
1118    };
1119}
1120
1121floats!(f32, f64);
1122
1123impl Field for bool {
1124    fn write(&self, b: &mut Builder) -> Result<()> {
1125        b.bool(*self)
1126    }
1127
1128    fn read(d: Doc<'_>) -> Result<bool> {
1129        d.as_bool().ok_or_else(|| not_a("bool", d))
1130    }
1131}
1132
1133impl Query for bool {
1134    fn key(&self, _kind: IndexKind) -> Option<Key> {
1135        Some(Key::bool(*self))
1136    }
1137}
1138
1139impl Field for String {
1140    fn write(&self, b: &mut Builder) -> Result<()> {
1141        b.text(self)
1142    }
1143
1144    fn read(d: Doc<'_>) -> Result<String> {
1145        d.as_text()
1146            .map(str::to_owned)
1147            .ok_or_else(|| not_a("string", d))
1148    }
1149}
1150
1151impl Query for String {
1152    fn key(&self, kind: IndexKind) -> Option<Key> {
1153        self.as_str().key(kind)
1154    }
1155}
1156
1157/// The borrowed form, so a `String` field is searched with a `&str`.
1158impl Query for str {
1159    fn key(&self, kind: IndexKind) -> Option<Key> {
1160        match kind {
1161            // A text index filed the words of the string, folded, so one word
1162            // is what can be asked for and a phrase is not a key at all.
1163            IndexKind::Text => Key::word(self),
1164            _ => Some(Key::text(self)),
1165        }
1166    }
1167}
1168
1169/// A field that may be absent, which is the only type whose absence is not an
1170/// error. `None` is stored as null rather than left out, so a document always
1171/// has the fields its shape says it has.
1172impl<T: Field> Field for Option<T> {
1173    fn write(&self, b: &mut Builder) -> Result<()> {
1174        match self {
1175            Some(v) => v.write(b),
1176            None => b.null(),
1177        }
1178    }
1179
1180    fn read(d: Doc<'_>) -> Result<Option<T>> {
1181        if d.is_null() {
1182            return Ok(None);
1183        }
1184        T::read(d).map(Some)
1185    }
1186
1187    fn missing(_name: &str) -> Result<Option<T>> {
1188        Ok(None)
1189    }
1190}
1191
1192impl<T: Field> Field for Vec<T> {
1193    fn write(&self, b: &mut Builder) -> Result<()> {
1194        b.begin_array()?;
1195        for v in self {
1196            v.write(b)?;
1197        }
1198        b.end_array()
1199    }
1200
1201    fn read(d: Doc<'_>) -> Result<Vec<T>> {
1202        if d.kind() != yo_doc::Kind::Array {
1203            return Err(not_a("list", d));
1204        }
1205        let mut out = Vec::with_capacity(d.len());
1206        for elem in d.iter() {
1207            out.push(T::read(elem)?);
1208        }
1209        Ok(out)
1210    }
1211}
1212
1213#[cfg(test)]
1214mod tests {
1215    use super::*;
1216    use crate::{Yo, open};
1217
1218    #[derive(Yo, Debug, Clone, PartialEq)]
1219    struct Order {
1220        #[yo(id)]
1221        id: u64,
1222        #[yo(index)]
1223        status: String,
1224        #[yo(ordered)]
1225        total: f64,
1226        #[yo(array)]
1227        tags: Vec<String>,
1228        #[yo(text)]
1229        note: String,
1230        sent: Option<String>,
1231    }
1232
1233    fn order(id: u64, status: &str, total: f64) -> Order {
1234        Order {
1235            id,
1236            status: status.to_owned(),
1237            total,
1238            tags: Vec::new(),
1239            note: String::new(),
1240            sent: None,
1241        }
1242    }
1243
1244    /// A collection holding the three orders most of these tests want.
1245    fn three() -> (crate::Db, Docs<Order>) {
1246        let db = open(crate::MEMORY).expect("a database in memory");
1247        let orders = db.docs::<Order>("orders").expect("a new collection");
1248        for o in [
1249            order(1, "open", 12.5),
1250            order(2, "shipped", 99.0),
1251            order(3, "open", 40.0),
1252        ] {
1253            orders.put(&o).expect("a document that fits");
1254        }
1255        (db, orders)
1256    }
1257
1258    #[test]
1259    fn a_document_comes_back_as_the_struct_that_went_in() {
1260        let (_db, orders) = three();
1261        assert_eq!(
1262            orders.get(&1).expect("a read"),
1263            Some(order(1, "open", 12.5))
1264        );
1265        assert_eq!(orders.get(&9).expect("a read"), None);
1266        assert_eq!(orders.len().expect("a count"), 3);
1267        assert!(orders.contains(&2).expect("a read"));
1268    }
1269
1270    #[test]
1271    fn every_field_kind_survives_the_round_trip() {
1272        let db = open(crate::MEMORY).expect("a database in memory");
1273        let orders = db.docs::<Order>("orders").expect("a new collection");
1274        let o = Order {
1275            id: 7,
1276            status: "open".to_owned(),
1277            total: -0.5,
1278            tags: vec!["red".to_owned(), "small".to_owned()],
1279            note: "A red kite".to_owned(),
1280            sent: Some("tuesday".to_owned()),
1281        };
1282        orders.put(&o).expect("a document that fits");
1283        assert_eq!(orders.get(&7).expect("a read"), Some(o));
1284    }
1285
1286    #[test]
1287    fn putting_the_same_id_twice_replaces_it() {
1288        let (_db, orders) = three();
1289        assert!(!orders.put(&order(1, "shut", 1.0)).expect("a write"));
1290        assert_eq!(orders.len().expect("a count"), 3);
1291        assert_eq!(
1292            orders.get(&1).expect("a read").expect("it is there").status,
1293            "shut"
1294        );
1295        // And the old value is out of the index it was under.
1296        assert_eq!(orders.count(Order::STATUS, "open").expect("a count"), 1);
1297    }
1298
1299    #[test]
1300    fn removing_a_document_takes_it_out_of_its_indexes() {
1301        let (_db, orders) = three();
1302        assert!(orders.remove(&1).expect("a write"));
1303        assert!(!orders.remove(&1).expect("a write"));
1304        assert_eq!(orders.len().expect("a count"), 2);
1305        assert_eq!(orders.count(Order::STATUS, "open").expect("a count"), 1);
1306        assert!(orders.find(Order::TOTAL, &12.5).expect("a read").is_empty());
1307    }
1308
1309    #[test]
1310    fn an_equality_index_answers_with_the_documents() {
1311        let (_db, orders) = three();
1312        let mut open = orders.find(Order::STATUS, "open").expect("a read");
1313        open.sort_by_key(|o| o.id);
1314        assert_eq!(open, [order(1, "open", 12.5), order(3, "open", 40.0)]);
1315        assert_eq!(orders.count(Order::STATUS, "gone").expect("a count"), 0);
1316    }
1317
1318    /// This is the query that was wrong before the numeric key encoding was
1319    /// fixed: 12.5 sorted after 99.0, so the range came back empty and the
1320    /// reverse walk came back ascending.
1321    #[test]
1322    fn a_range_over_a_float_field_is_in_numeric_order() {
1323        let (_db, orders) = three();
1324        let cheap = orders.range(Order::TOTAL, 0.0..50.0).expect("a read");
1325        assert_eq!(
1326            cheap.iter().map(|o| o.total).collect::<Vec<_>>(),
1327            [12.5, 40.0]
1328        );
1329
1330        let all = orders.range(Order::TOTAL, ..).expect("a read");
1331        assert_eq!(
1332            all.iter().map(|o| o.total).collect::<Vec<_>>(),
1333            [12.5, 40.0, 99.0]
1334        );
1335
1336        let down = orders.range_rev(Order::TOTAL, ..).expect("a read");
1337        assert_eq!(
1338            down.iter().map(|o| o.total).collect::<Vec<_>>(),
1339            [99.0, 40.0, 12.5]
1340        );
1341
1342        assert_eq!(
1343            orders
1344                .count_range(Order::TOTAL, 12.5..=40.0)
1345                .expect("a count"),
1346            2
1347        );
1348    }
1349
1350    #[test]
1351    fn an_ordered_path_can_still_be_asked_for_equality() {
1352        let (_db, orders) = three();
1353        assert_eq!(orders.find(Order::TOTAL, &40.0).expect("a read").len(), 1);
1354        assert_eq!(orders.count(Order::TOTAL, &99.0).expect("a count"), 1);
1355    }
1356
1357    #[test]
1358    fn a_range_over_a_string_field_takes_a_pair_of_bounds() {
1359        let db = open(crate::MEMORY).expect("a database in memory");
1360        let names = db.docs::<Named>("names").expect("a new collection");
1361        for (id, name) in [(1u64, "banana"), (2, "apple"), (3, "quince")] {
1362            names
1363                .put(&Named {
1364                    id,
1365                    name: name.to_owned(),
1366                })
1367                .expect("a document that fits");
1368        }
1369        let early = names
1370            .range(Named::NAME, (Bound::Included("a"), Bound::Excluded("m")))
1371            .expect("a read");
1372        assert_eq!(
1373            early.iter().map(|n| n.name.as_str()).collect::<Vec<_>>(),
1374            ["apple", "banana"]
1375        );
1376    }
1377
1378    #[derive(Yo, Debug, PartialEq)]
1379    struct Named {
1380        #[yo(id)]
1381        id: u64,
1382        #[yo(ordered)]
1383        name: String,
1384    }
1385
1386    #[test]
1387    fn an_array_index_files_a_document_under_every_element() {
1388        let db = open(crate::MEMORY).expect("a database in memory");
1389        let orders = db.docs::<Order>("orders").expect("a new collection");
1390        let mut o = order(1, "open", 1.0);
1391        o.tags = vec!["red".to_owned(), "small".to_owned()];
1392        orders.put(&o).expect("a document that fits");
1393
1394        assert_eq!(orders.find(Order::TAGS, "red").expect("a read").len(), 1);
1395        assert_eq!(orders.find(Order::TAGS, "small").expect("a read").len(), 1);
1396        assert_eq!(orders.count(Order::TAGS, "large").expect("a count"), 0);
1397    }
1398
1399    #[test]
1400    fn a_text_index_files_a_document_under_every_word() {
1401        let db = open(crate::MEMORY).expect("a database in memory");
1402        let orders = db.docs::<Order>("orders").expect("a new collection");
1403        let mut o = order(1, "open", 1.0);
1404        o.note = "A red kite".to_owned();
1405        orders.put(&o).expect("a document that fits");
1406
1407        // The case is folded on both sides, so the query does not have to match
1408        // how the document happened to be written.
1409        assert_eq!(orders.find(Order::NOTE, "RED").expect("a read").len(), 1);
1410        assert_eq!(orders.find(Order::NOTE, "kite").expect("a read").len(), 1);
1411        assert_eq!(orders.count(Order::NOTE, "blue").expect("a count"), 0);
1412    }
1413
1414    #[test]
1415    fn asking_a_text_index_for_a_phrase_says_so() {
1416        let (_db, orders) = three();
1417        let e = orders
1418            .find(Order::NOTE, "red kite")
1419            .expect_err("not one word");
1420        assert_eq!(e.code(), crate::Code::Invalid);
1421        assert!(e.message().contains("one word"), "{}", e.message());
1422    }
1423
1424    #[test]
1425    fn an_absent_field_reads_back_as_none() {
1426        let (_db, orders) = three();
1427        assert_eq!(
1428            orders.get(&1).expect("a read").expect("it is there").sent,
1429            None
1430        );
1431    }
1432
1433    #[test]
1434    fn a_nested_struct_is_a_field() {
1435        #[derive(Yo, Debug, PartialEq)]
1436        struct Where {
1437            city: String,
1438            postcode: String,
1439        }
1440
1441        #[derive(Yo, Debug, PartialEq)]
1442        struct Person {
1443            #[yo(id)]
1444            id: u64,
1445            home: Where,
1446        }
1447
1448        let db = open(crate::MEMORY).expect("a database in memory");
1449        let people = db.docs::<Person>("people").expect("a new collection");
1450        let p = Person {
1451            id: 1,
1452            home: Where {
1453                city: "Hanoi".to_owned(),
1454                postcode: "100000".to_owned(),
1455            },
1456        };
1457        people.put(&p).expect("a document that fits");
1458        assert_eq!(people.get(&1).expect("a read"), Some(p));
1459    }
1460
1461    #[test]
1462    fn all_walks_every_document() {
1463        let (_db, orders) = three();
1464        let mut ids: Vec<u64> = orders.all().expect("a read").iter().map(|o| o.id).collect();
1465        ids.sort_unstable();
1466        assert_eq!(ids, [1, 2, 3]);
1467    }
1468
1469    #[test]
1470    fn opening_a_collection_as_the_wrong_thing_is_refused() {
1471        let db = open(crate::MEMORY).expect("a database in memory");
1472        let _orders = db.docs::<Order>("orders").expect("a new collection");
1473        let e = db
1474            .map::<String, u64>("orders")
1475            .expect_err("a different shape");
1476        assert_eq!(e.code(), crate::Code::ShapeMismatch);
1477        let e = db.docs::<Named>("orders").expect_err("a different struct");
1478        assert_eq!(e.code(), crate::Code::ShapeMismatch);
1479    }
1480
1481    #[test]
1482    fn reopening_a_collection_hands_back_the_same_documents() {
1483        let (db, orders) = three();
1484        let again = db.docs::<Order>("orders").expect("the same collection");
1485        assert_eq!(again.len().expect("a count"), 3);
1486        assert_eq!(again.count(Order::STATUS, "open").expect("a count"), 2);
1487        drop(orders);
1488    }
1489
1490    #[test]
1491    fn a_u64_past_what_json_can_hold_is_refused() {
1492        let db = open(crate::MEMORY).expect("a database in memory");
1493        let orders = db.docs::<Order>("orders").expect("a new collection");
1494        let e = orders
1495            .put(&order(u64::MAX, "open", 1.0))
1496            .expect_err("too big");
1497        assert_eq!(e.code(), crate::Code::Invalid);
1498    }
1499
1500    // ---- the vector index
1501
1502    #[derive(Yo, Debug, Clone, PartialEq)]
1503    struct Note {
1504        #[yo(id)]
1505        id: u64,
1506        #[yo(index)]
1507        lang: String,
1508        #[yo(vector = 4)]
1509        embedding: Vec<f32>,
1510    }
1511
1512    fn note(id: u64, lang: &str, embedding: [f32; 4]) -> Note {
1513        Note {
1514            id,
1515            lang: lang.to_owned(),
1516            embedding: embedding.to_vec(),
1517        }
1518    }
1519
1520    /// A collection holding four notes, two of them in each language.
1521    fn notes() -> (crate::Db, Docs<Note>) {
1522        let db = open(crate::MEMORY).expect("a database in memory");
1523        let notes = db.docs::<Note>("notes").expect("a new collection");
1524        for n in [
1525            note(1, "en", [1.0, 0.0, 0.0, 0.0]),
1526            note(2, "fr", [0.0, 1.0, 0.0, 0.0]),
1527            note(3, "en", [0.0, 0.0, 1.0, 0.0]),
1528            note(4, "fr", [0.0, 0.0, 0.0, 1.0]),
1529        ] {
1530            notes.put(&n).expect("a document that fits");
1531        }
1532        (db, notes)
1533    }
1534
1535    #[test]
1536    fn a_derived_vector_field_is_indexed_and_comes_back_whole() {
1537        let (_db, notes) = notes();
1538        assert_eq!(
1539            Note::VECTORS,
1540            [("$.embedding", 4usize)],
1541            "the derive declares the path and the width"
1542        );
1543        assert_eq!(Note::EMBEDDING.path(), "$.embedding");
1544        assert_eq!(Note::EMBEDDING.dim(), 4);
1545
1546        // The embedding is a field of the document like any other, so the
1547        // struct that comes out is the struct that went in.
1548        assert_eq!(
1549            notes.get(&2).expect("a read"),
1550            Some(note(2, "fr", [0.0, 1.0, 0.0, 0.0]))
1551        );
1552
1553        let near = notes
1554            .nearest(Note::EMBEDDING, &[0.9, 0.1, 0.0, 0.0], 2)
1555            .expect("a search");
1556        assert_eq!(near.iter().map(|n| n.id).collect::<Vec<_>>(), [1, 2]);
1557    }
1558
1559    #[test]
1560    fn a_filter_narrows_the_search_and_not_the_answers() {
1561        let (_db, notes) = notes();
1562        let french = notes
1563            .near(Note::EMBEDDING, &[0.9, 0.1, 0.0, 0.0])
1564            .filter(Note::LANG, "fr")
1565            .take(2)
1566            .expect("a search");
1567        assert_eq!(french.iter().map(|n| n.id).collect::<Vec<_>>(), [2, 4]);
1568
1569        // The scores come back in the same order, nearer first.
1570        let scored = notes
1571            .near(Note::EMBEDDING, &[0.9, 0.1, 0.0, 0.0])
1572            .filter(Note::LANG, "fr")
1573            .scored(2)
1574            .expect("a search");
1575        assert_eq!(scored[0].0.id, 2);
1576        assert!(scored[0].1 <= scored[1].1);
1577
1578        // A filter on a path with no index says so rather than answering
1579        // nothing, and it says so when the search runs.
1580        let e = notes
1581            .near(Note::EMBEDDING, &[1.0, 0.0, 0.0, 0.0])
1582            .filter(
1583                Path::<Note, String>::new("$.author", IndexKind::Equality),
1584                "me",
1585            )
1586            .take(1)
1587            .expect_err("no index there");
1588        assert_eq!(e.code(), crate::Code::Invalid);
1589    }
1590
1591    #[test]
1592    fn more_like_this_leaves_the_document_itself_out() {
1593        let (_db, notes) = notes();
1594        let like = notes.nearest_to(Note::EMBEDDING, &1, 2).expect("a search");
1595        assert_eq!(like.len(), 2);
1596        assert!(!like.iter().any(|n| n.id == 1));
1597
1598        // An id that is not in the collection has nothing to be like.
1599        assert!(
1600            notes
1601                .nearest_to(Note::EMBEDDING, &99, 2)
1602                .expect("a search")
1603                .is_empty()
1604        );
1605    }
1606
1607    #[test]
1608    fn an_embedding_of_the_wrong_width_is_refused() {
1609        let db = open(crate::MEMORY).expect("a database in memory");
1610        let notes = db.docs::<Note>("notes").expect("a new collection");
1611        let e = notes
1612            .put(&Note {
1613                id: 1,
1614                lang: "en".to_owned(),
1615                embedding: vec![1.0, 0.0],
1616            })
1617            .expect_err("two coordinates where four were declared");
1618        assert_eq!(e.code(), crate::Code::Invalid);
1619        assert_eq!(notes.len().expect("a count"), 0);
1620
1621        // And so is a query vector of the wrong width.
1622        let e = notes
1623            .nearest(Note::EMBEDDING, &[1.0, 0.0], 1)
1624            .expect_err("two coordinates");
1625        assert_eq!(e.code(), crate::Code::Invalid);
1626    }
1627
1628    #[test]
1629    fn reopening_a_collection_keeps_the_vector_index() {
1630        let (db, notes) = notes();
1631        let again = db.docs::<Note>("notes").expect("the same collection");
1632        assert_eq!(
1633            again
1634                .nearest(Note::EMBEDDING, &[0.0, 0.0, 0.9, 0.1], 1)
1635                .expect("a search")
1636                .first()
1637                .map(|n| n.id),
1638            Some(3)
1639        );
1640        drop(notes);
1641    }
1642}