use core::cell::RefCell;
use std::rc::Rc;
use yo_common::{Code, Error, Result};
use yo_index::RawMap;
use yo_shape::{Desc, Shape, Tag};
use crate::counter::Counter;
use crate::doc::{Docs, Document, Documents};
use crate::graph::Graph;
use crate::keys::Keys;
use crate::keyspace::Strings;
use crate::map::Map;
use crate::sets::{Set, Sets};
use crate::store::Decode;
use crate::vector::Vectors;
pub const MEMORY: &str = ":memory:";
pub fn open(path: &str) -> Result<Db> {
if path != MEMORY {
return Err(Error::fmt(
Code::Unsupported,
format_args!(
"this build holds a database in memory only, so the path has to be \"{MEMORY}\", not \"{path}\". A file backed database arrives with the .yo format in M5"
),
));
}
Ok(Db {
db: Handle {
inner: Rc::new(RefCell::new(Inner {
collections: Vec::new(),
strings: yo_kv::Keyspace::new(),
deadlines: false,
})),
},
})
}
#[derive(Clone)]
pub struct Db {
db: Handle,
}
impl core::fmt::Debug for Db {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
let mut d = f.debug_struct("Db");
match self.collections() {
Ok(names) => d.field("collections", &names).finish(),
Err(_) => d.finish_non_exhaustive(),
}
}
}
pub(crate) struct Inner {
pub(crate) collections: Vec<Collection>,
pub(crate) strings: yo_kv::Keyspace,
pub(crate) deadlines: bool,
}
pub(crate) struct Collection {
pub(crate) name: String,
pub(crate) desc: Desc,
pub(crate) data: Data,
}
#[allow(clippy::large_enum_variant)]
pub(crate) enum Data {
Map(RawMap),
Docs(Box<Documents>),
Graph(Box<crate::graph::Store>),
Vectors(Box<yo_vector::Collection>),
}
impl Data {
pub(crate) fn memory_bytes(&self) -> usize {
match self {
Data::Map(m) => m.memory_bytes(),
Data::Docs(d) => d.docs.memory_bytes(),
Data::Graph(g) => g.memory_bytes(),
Data::Vectors(v) => v.memory_bytes(),
}
}
#[track_caller]
pub(crate) fn map(&self) -> &RawMap {
match self {
Data::Map(m) => m,
_ => wrong_kind(),
}
}
#[track_caller]
pub(crate) fn map_mut(&mut self) -> &mut RawMap {
match self {
Data::Map(m) => m,
_ => wrong_kind(),
}
}
#[track_caller]
pub(crate) fn docs(&self) -> &yo_doc::Docs {
match self {
Data::Docs(d) => &d.docs,
_ => wrong_kind(),
}
}
#[track_caller]
pub(crate) fn docs_mut(&mut self) -> &mut Documents {
match self {
Data::Docs(d) => d,
_ => wrong_kind(),
}
}
#[track_caller]
pub(crate) fn graph(&self) -> &crate::graph::Store {
match self {
Data::Graph(g) => g,
_ => wrong_kind(),
}
}
#[track_caller]
pub(crate) fn graph_mut(&mut self) -> &mut crate::graph::Store {
match self {
Data::Graph(g) => g,
_ => wrong_kind(),
}
}
#[track_caller]
pub(crate) fn vectors(&self) -> &yo_vector::Collection {
match self {
Data::Vectors(v) => v,
_ => wrong_kind(),
}
}
#[track_caller]
pub(crate) fn vectors_mut(&mut self) -> &mut yo_vector::Collection {
match self {
Data::Vectors(v) => v,
_ => wrong_kind(),
}
}
}
#[track_caller]
fn wrong_kind() -> ! {
panic!(
"this handle and the collection it names hold different things, which the shape check is there to make impossible. Please report this as a bug"
)
}
#[derive(Clone)]
pub(crate) struct Handle {
inner: Rc<RefCell<Inner>>,
}
impl Handle {
pub(crate) fn run<R>(&self, f: impl FnOnce(&mut Inner) -> Result<R>) -> Result<R> {
let mut inner = self.inner.try_borrow_mut().map_err(|_| reentrant())?;
if inner.deadlines {
inner.strings.clock_mut().refresh();
}
f(&mut inner)
}
pub(crate) fn deadlines<R>(&self, f: impl FnOnce(&mut Inner) -> Result<R>) -> Result<R> {
let mut inner = self.inner.try_borrow_mut().map_err(|_| reentrant())?;
inner.deadlines = true;
inner.strings.clock_mut().refresh();
f(&mut inner)
}
pub(crate) fn read<R>(&self, f: impl FnOnce(&Inner) -> Result<R>) -> Result<R> {
let inner = self.inner.try_borrow().map_err(|_| reentrant())?;
f(&inner)
}
pub(crate) fn write<R>(&self, f: impl FnOnce(&mut Inner) -> Result<R>) -> Result<R> {
let mut inner = self.inner.try_borrow_mut().map_err(|_| reentrant())?;
f(&mut inner)
}
pub(crate) fn is(&self, other: &Handle) -> bool {
Rc::ptr_eq(&self.inner, &other.inner)
}
}
fn declare<T: Document>(data: &mut Documents) -> Result<()> {
for (path, kind) in T::INDEXES {
data.docs.create_index_bytes(path.as_bytes(), *kind)?;
}
for (path, dim) in T::VECTORS {
data.docs.create_vector_index(path, *dim)?;
}
Ok(())
}
struct AGraph;
impl Shape for AGraph {
fn describe(d: &mut Desc) {
d.strukt("graph", &[]);
}
}
pub(crate) fn reentrant() -> Error {
Error::new(
Code::Invalid,
"this database is already in use by the call above this one. A closure passed to with() or update() cannot call back into the same database, so read what you need first and write after the closure returns",
)
}
impl Db {
pub fn map<K: Decode, V: Decode>(&self, name: &str) -> Result<Map<K, V>> {
let mut desc = Desc::new();
desc.map(K::describe, V::describe);
let tag = desc.tag();
let at =
self.db.write(
|inner| match inner.collections.iter().position(|c| c.name == name) {
Some(at) => {
yo_shape::check(name, &inner.collections[at].desc, &desc, None)?;
Ok(at)
}
None => {
inner.collections.push(Collection {
name: name.to_owned(),
desc,
data: Data::Map(RawMap::new()),
});
Ok(inner.collections.len() - 1)
}
},
)?;
Ok(Map::new(self.db.clone(), at, tag))
}
pub fn docs<T: Document>(&self, name: &str) -> Result<Docs<T>> {
let mut desc = Desc::new();
T::describe(&mut desc);
let tag = desc.tag();
let at =
self.db.write(
|inner| match inner.collections.iter().position(|c| c.name == name) {
Some(at) => {
yo_shape::check(name, &inner.collections[at].desc, &desc, None)?;
declare::<T>(inner.collections[at].data.docs_mut())?;
Ok(at)
}
None => {
let mut data = Documents::new();
declare::<T>(&mut data)?;
inner.collections.push(Collection {
name: name.to_owned(),
desc,
data: Data::Docs(Box::new(data)),
});
Ok(inner.collections.len() - 1)
}
},
)?;
Ok(Docs::new(self.db.clone(), at, tag))
}
pub fn graph(&self, name: &str) -> Result<Graph> {
let mut desc = Desc::new();
AGraph::describe(&mut desc);
let at =
self.db.write(
|inner| match inner.collections.iter().position(|c| c.name == name) {
Some(at) => {
yo_shape::check(name, &inner.collections[at].desc, &desc, None)?;
Ok(at)
}
None => {
inner.collections.push(Collection {
name: name.to_owned(),
desc,
data: Data::Graph(Box::new(crate::graph::Store::new())),
});
Ok(inner.collections.len() - 1)
}
},
)?;
Ok(Graph::new(self.db.clone(), at))
}
pub fn vectors(&self, name: &str, dim: usize) -> Result<Vectors> {
self.vectors_with(name, dim, yo_shape::Metric::L2)
}
pub fn vectors_with(
&self,
name: &str,
dim: usize,
metric: yo_shape::Metric,
) -> Result<Vectors> {
let width = yo_vector::collection::width(dim)?;
yo_vector::collection::check_metric(metric)?;
let mut desc = Desc::new();
desc.vector(width, metric);
let at =
self.db.write(
|inner| match inner.collections.iter().position(|c| c.name == name) {
Some(at) => {
yo_shape::check(name, &inner.collections[at].desc, &desc, None)?;
Ok(at)
}
None => {
inner.collections.push(Collection {
name: name.to_owned(),
desc,
data: Data::Vectors(Box::new(yo_vector::Collection::new(dim, metric)?)),
});
Ok(inner.collections.len() - 1)
}
},
)?;
Ok(Vectors {
db: self.db.clone(),
at,
dim,
metric,
})
}
#[must_use]
pub fn strings(&self) -> Strings {
Strings {
db: self.db.clone(),
}
}
#[must_use]
pub fn counter(&self, key: impl Into<Vec<u8>>) -> Counter {
Counter {
db: self.db.clone(),
key: key.into(),
}
}
#[must_use]
pub fn sets(&self) -> Sets {
Sets {
db: self.db.clone(),
}
}
#[must_use]
pub fn set(&self, key: impl Into<Vec<u8>>) -> Set {
Set {
sets: self.sets(),
key: key.into(),
}
}
#[must_use]
pub fn keys(&self) -> Keys {
Keys {
db: self.db.clone(),
}
}
pub fn collections(&self) -> Result<Vec<String>> {
self.db
.read(|inner| Ok(inner.collections.iter().map(|c| c.name.clone()).collect()))
}
pub fn shape(&self, name: &str) -> Result<Option<Tag>> {
self.db.read(|inner| {
Ok(inner
.collections
.iter()
.find(|c| c.name == name)
.map(|c| c.desc.tag()))
})
}
pub fn memory_bytes(&self) -> Result<usize> {
self.db.read(|inner| {
Ok(inner.strings.memory_bytes()
+ inner
.collections
.iter()
.map(|c| c.data.memory_bytes())
.sum::<usize>())
})
}
#[must_use]
pub fn reads_the_clock(&self) -> bool {
self.db.read(|inner| Ok(inner.deadlines)).unwrap_or(false)
}
#[must_use]
pub fn is(&self, other: &Db) -> bool {
self.db.is(&other.db)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_path_on_disk_says_which_build_would_take_it() {
let e = open("app.yo").expect_err("no file format yet");
assert_eq!(e.code(), Code::Unsupported);
assert!(e.message().contains("M5"), "{e}");
}
#[test]
fn opening_the_same_name_twice_gives_the_same_collection() {
let db = open(MEMORY).unwrap();
let a = db.map::<String, u64>("hits").unwrap();
let b = db.map::<String, u64>("hits").unwrap();
a.set("home", &1).unwrap();
assert_eq!(b.get("home").unwrap(), Some(1));
assert_eq!(db.collections().unwrap(), vec!["hits".to_owned()]);
}
#[test]
fn opening_the_same_name_with_another_type_is_a_shape_mismatch() {
let db = open(MEMORY).unwrap();
let _first = db.map::<String, u64>("hits").unwrap();
let e = db
.map::<String, String>("hits")
.expect_err("that is a different shape");
assert_eq!(e.code(), Code::ShapeMismatch);
assert!(
e.message().contains("the type changed from u64 to str"),
"{e}"
);
assert_eq!(e.detail(), Some("change=breaking"));
}
#[test]
fn two_collections_are_two_keyspaces() {
let db = open(MEMORY).unwrap();
let a = db.map::<String, u64>("a").unwrap();
let b = db.map::<String, u64>("b").unwrap();
a.set("k", &1).unwrap();
b.set("k", &2).unwrap();
assert_eq!(a.get("k").unwrap(), Some(1));
assert_eq!(b.get("k").unwrap(), Some(2));
assert_eq!(db.collections().unwrap().len(), 2);
}
#[test]
fn a_typed_collection_and_the_keyspace_are_not_the_same_store() {
let db = open(MEMORY).unwrap();
let map = db.map::<String, u64>("hits").unwrap();
map.set("home", &1).unwrap();
db.strings().set("home", "elsewhere").unwrap();
assert_eq!(map.get("home").unwrap(), Some(1));
assert_eq!(
db.strings().get("home").unwrap().as_deref(),
Some(&b"elsewhere"[..])
);
}
#[test]
fn a_shape_can_be_read_back_and_an_unopened_name_has_none() {
let db = open(MEMORY).unwrap();
let map = db.map::<String, u64>("hits").unwrap();
assert_eq!(db.shape("hits").unwrap(), Some(map.tag()));
assert_eq!(db.shape("misses").unwrap(), None);
}
#[test]
fn a_clone_is_the_same_database() {
let db = open(MEMORY).unwrap();
let map = db.map::<String, u64>("hits").unwrap();
map.set("home", &3).unwrap();
let same = db.clone();
assert_eq!(
same.map::<String, u64>("hits")
.unwrap()
.get("home")
.unwrap(),
Some(3)
);
assert!(db.is(&same));
assert!(!db.is(&open(MEMORY).unwrap()));
assert!(db.memory_bytes().unwrap() > 0);
assert!(format!("{db:?}").contains("hits"));
}
}