use core::marker::PhantomData;
use std::collections::HashMap;
use yo_common::{Code, Error, Result};
use yo_doc::Builder;
use yo_graph::Dir;
use yo_shape::{Desc, Shape, Tag};
use crate::db::Handle;
use crate::doc::{Asked, Document, Field, IndexKind, Indexed, Path, key_of};
#[diagnostic::on_unimplemented(
message = "`{Self}` is not a node type",
label = "this type has no label",
note = "give it a label with an `impl Node` block naming a `const LABEL`, and make sure it derives Yo with one field marked `#[yo(id)]`"
)]
pub trait Node: Document {
const LABEL: &'static str;
}
#[diagnostic::on_unimplemented(
message = "`{Self}` is not an edge type",
label = "this type does not say where it goes",
note = "say where it goes with an `impl Edge` block naming a `From`, a `To` and a `const LABEL`"
)]
pub trait Edge: Field + Indexed {
type From: Node;
type To: Node;
const LABEL: &'static str;
}
pub struct Id<N> {
raw: u64,
marker: PhantomData<fn() -> N>,
}
impl<N> Clone for Id<N> {
fn clone(&self) -> Id<N> {
*self
}
}
impl<N> Copy for Id<N> {}
impl<N> PartialEq for Id<N> {
fn eq(&self, other: &Id<N>) -> bool {
self.raw == other.raw
}
}
impl<N> Eq for Id<N> {}
impl<N> core::hash::Hash for Id<N> {
fn hash<H: core::hash::Hasher>(&self, h: &mut H) {
self.raw.hash(h);
}
}
impl<N: Node> core::fmt::Debug for Id<N> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "{}#{}", N::LABEL, self.raw)
}
}
impl<N> Id<N> {
fn new(raw: u64) -> Id<N> {
Id {
raw,
marker: PhantomData,
}
}
}
pub struct EdgeId<E> {
slot: u32,
marker: PhantomData<fn() -> E>,
}
impl<E> Clone for EdgeId<E> {
fn clone(&self) -> EdgeId<E> {
*self
}
}
impl<E> Copy for EdgeId<E> {}
impl<E> PartialEq for EdgeId<E> {
fn eq(&self, other: &EdgeId<E>) -> bool {
self.slot == other.slot
}
}
impl<E> Eq for EdgeId<E> {}
impl<E: Edge> core::fmt::Debug for EdgeId<E> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "{}#{}", E::LABEL, self.slot)
}
}
pub struct Hop<E: Edge> {
pub to: Id<E::To>,
pub edge: EdgeId<E>,
}
impl<E: Edge> Clone for Hop<E> {
fn clone(&self) -> Hop<E> {
*self
}
}
impl<E: Edge> Copy for Hop<E> {}
impl<E: Edge> core::fmt::Debug for Hop<E> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "{:?} by {:?}", self.to, self.edge)
}
}
pub(crate) struct Store {
g: yo_graph::Graph,
nodes: Vec<Kind>,
edges: Vec<Kind>,
ids: HashMap<Box<[u8]>, u64>,
of: Vec<u32>,
next: u64,
node_paths: HashMap<&'static str, IndexKind>,
edge_paths: HashMap<&'static str, IndexKind>,
scratch: Builder,
}
const GONE: u32 = u32::MAX;
struct Kind {
label: &'static str,
shape: Tag,
live: usize,
}
impl Store {
pub(crate) fn new() -> Store {
Store {
g: yo_graph::Graph::new(),
nodes: Vec::new(),
edges: Vec::new(),
ids: HashMap::new(),
of: Vec::new(),
next: 0,
node_paths: HashMap::new(),
edge_paths: HashMap::new(),
scratch: Builder::new(),
}
}
pub(crate) fn memory_bytes(&self) -> usize {
self.g.memory_bytes()
+ self.of.capacity() * size_of::<u32>()
+ self.ids.capacity() * (size_of::<Box<[u8]>>() + size_of::<u64>() + 16)
}
fn node_kind<N: Node>(&mut self) -> Result<u32> {
let want = tag_of::<N>();
let at = register(&mut self.nodes, N::LABEL, want)?;
declare(&mut self.node_paths, N::INDEXES, N::LABEL, |path, kind| {
self.g.index_nodes(path, kind)
})?;
Ok(at)
}
fn edge_kind<E: Edge>(&mut self) -> Result<u32> {
let want = tag_of::<E>();
let at = register(&mut self.edges, E::LABEL, want)?;
declare(&mut self.edge_paths, E::INDEXES, E::LABEL, |path, kind| {
self.g.index_edges(path, kind)
})?;
Ok(at)
}
fn edge_seen<E: Edge>(&self) -> Option<u32> {
seen(&self.edges, E::LABEL)
}
fn node_seen<N: Node>(&self) -> Option<u32> {
seen(&self.nodes, N::LABEL)
}
fn is(&self, raw: u64, kind: u32) -> bool {
usize::try_from(raw).is_ok_and(|i| self.of.get(i).copied() == Some(kind))
}
}
fn seen(kinds: &[Kind], label: &'static str) -> Option<u32> {
kinds
.iter()
.position(|k| k.label == label)
.map(|at| at as u32)
}
fn register(kinds: &mut Vec<Kind>, label: &'static str, want: Tag) -> Result<u32> {
if let Some(at) = kinds.iter().position(|k| k.label == label) {
if kinds[at].shape != want {
return Err(Error::fmt(
Code::ShapeMismatch,
format_args!(
"this graph already holds {label} under another shape, so the two types cannot share the label"
),
));
}
return Ok(at as u32);
}
if kinds.len() >= u32::MAX as usize {
return Err(Error::new(Code::Full, "this graph has no labels left"));
}
kinds.push(Kind {
label,
shape: want,
live: 0,
});
Ok((kinds.len() - 1) as u32)
}
fn declare(
have: &mut HashMap<&'static str, IndexKind>,
want: &'static [(&'static str, IndexKind)],
label: &'static str,
mut create: impl FnMut(&str, IndexKind) -> Result<()>,
) -> Result<()> {
for (path, kind) in want {
match have.get(path) {
Some(already) if already == kind => {}
Some(already) => {
return Err(Error::fmt(
Code::Invalid,
format_args!(
"{label} asks for {path} to be indexed for {kind:?}, and another type in this graph already indexes it for {already:?}. One path is one index, so the two types have to agree"
),
));
}
None => {
create(path, *kind)?;
have.insert(path, *kind);
}
}
}
Ok(())
}
fn tag_of<T: Shape>() -> Tag {
let mut d = Desc::new();
T::describe(&mut d);
d.tag()
}
#[derive(Clone)]
pub struct Graph {
db: Handle,
at: usize,
}
impl core::fmt::Debug for Graph {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
let name = self
.db
.read(|inner| Ok(inner.collections[self.at].name.clone()))
.unwrap_or_else(|_| "?".to_owned());
f.debug_struct("Graph").field("name", &name).finish()
}
}
impl Graph {
pub(crate) fn new(db: Handle, at: usize) -> Graph {
Graph { db, at }
}
pub fn name(&self) -> Result<String> {
self.db
.read(|inner| Ok(inner.collections[self.at].name.clone()))
}
fn write<R>(&self, f: impl FnOnce(&mut Store) -> Result<R>) -> Result<R> {
self.db
.write(|inner| f(inner.collections[self.at].data.graph_mut()))
}
fn read<R>(&self, f: impl FnOnce(&Store) -> Result<R>) -> Result<R> {
self.db
.read(|inner| f(inner.collections[self.at].data.graph()))
}
pub fn add<N: Node>(&self, node: &N) -> Result<Id<N>> {
let key = key_of(node.id(), IndexKind::Equality, "the id")?;
self.write(|s| {
let kind = s.node_kind::<N>()?;
let tagged = tagged(kind, key.as_bytes());
let raw = match s.ids.get(&tagged[..]) {
Some(raw) => *raw,
None => {
let raw = s.next;
s.next += 1;
s.ids.insert(tagged.into_boxed_slice(), raw);
s.of.push(kind);
s.nodes[kind as usize].live += 1;
raw
}
};
s.scratch.clear();
Field::write(node, &mut s.scratch)?;
let bytes = s.scratch.finish()?;
s.g.put_node(raw, bytes)?;
Ok(Id::new(raw))
})
}
pub fn id_of<N: Node>(&self, id: &<N::Id as Asked>::Ask) -> Result<Option<Id<N>>> {
let key = key_of(id, IndexKind::Equality, "the id")?;
self.read(|s| {
let Some(kind) = s.node_seen::<N>() else {
return Ok(None);
};
Ok(s.ids
.get(&tagged(kind, key.as_bytes())[..])
.map(|raw| Id::new(*raw)))
})
}
pub fn get<N: Node>(&self, id: Id<N>) -> Result<Option<N>> {
self.read(|s| match s.g.node(id.raw) {
Some(doc) => N::read(doc).map(Some),
None => Ok(None),
})
}
pub fn has<N: Node>(&self, id: Id<N>) -> Result<bool> {
self.read(|s| Ok(s.g.has_node(id.raw)))
}
pub fn remove<N: Node>(&self, id: Id<N>) -> Result<bool> {
self.write(|s| {
let Some(doc) = s.g.node(id.raw) else {
return Ok(false);
};
let node = N::read(doc)?;
let key = key_of(node.id(), IndexKind::Equality, "the id")?;
let Some(kind) = s.node_seen::<N>() else {
return Ok(false);
};
s.ids.remove(&tagged(kind, key.as_bytes())[..]);
if let Ok(i) = usize::try_from(id.raw)
&& let Some(slot) = s.of.get_mut(i)
{
*slot = GONE;
}
s.nodes[kind as usize].live -= 1;
s.g.remove_node(id.raw)
})
}
pub fn count<N: Node>(&self) -> Result<usize> {
self.read(|s| {
Ok(s.node_seen::<N>()
.map_or(0, |kind| s.nodes[kind as usize].live))
})
}
pub fn nodes(&self) -> Result<usize> {
self.read(|s| Ok(s.g.nodes()))
}
pub fn edges(&self) -> Result<usize> {
self.read(|s| Ok(s.g.edges()))
}
pub fn link<E: Edge>(&self, from: Id<E::From>, to: Id<E::To>, edge: &E) -> Result<EdgeId<E>> {
self.write(|s| {
let label = s.edge_kind::<E>()?;
let from_kind = s.node_kind::<E::From>()?;
let to_kind = s.node_kind::<E::To>()?;
if !s.is(from.raw, from_kind) {
return Err(gone::<E::From>(from.raw));
}
if !s.is(to.raw, to_kind) {
return Err(gone::<E::To>(to.raw));
}
s.scratch.clear();
Field::write(edge, &mut s.scratch)?;
let bytes = s.scratch.finish()?;
let slot = s.g.link(from.raw, to.raw, label, bytes)?;
Ok(EdgeId {
slot,
marker: PhantomData,
})
})
}
pub fn unlink<E: Edge>(&self, from: Id<E::From>, to: Id<E::To>) -> Result<bool> {
self.write(|s| {
let Some(label) = s.edge_seen::<E>() else {
return Ok(false);
};
Ok(s.g.unlink(from.raw, to.raw, label).is_some())
})
}
pub fn edge<E: Edge>(&self, id: EdgeId<E>) -> Result<Option<E>> {
self.read(|s| match s.g.edge(id.slot) {
Some(doc) => E::read(doc).map(Some),
None => Ok(None),
})
}
pub fn out<E: Edge>(&self, from: Id<E::From>) -> Result<Vec<Id<E::To>>> {
self.step::<E>(from.raw, Dir::Out).map(ids)
}
pub fn incoming<E: Edge>(&self, to: Id<E::To>) -> Result<Vec<Id<E::From>>> {
self.step::<E>(to.raw, Dir::In).map(ids)
}
pub fn out_edges<E: Edge>(&self, from: Id<E::From>) -> Result<Vec<Hop<E>>> {
self.read(|s| {
let Some(label) = s.edge_seen::<E>() else {
return Ok(Vec::new());
};
Ok(s.g
.hop(from.raw, label, Dir::Out)
.map(|(node, slot)| Hop {
to: Id::new(node),
edge: EdgeId {
slot,
marker: PhantomData,
},
})
.collect())
})
}
pub fn degree<E: Edge>(&self, from: Id<E::From>) -> Result<usize> {
self.read(|s| {
Ok(s.edge_seen::<E>()
.map_or(0, |label| s.g.degree(from.raw, label, Dir::Out)))
})
}
#[must_use]
pub fn walk<N: Node>(&self, from: Id<N>) -> Walk<'_, N> {
Walk {
g: self,
at: vec![from.raw],
marker: PhantomData,
}
}
pub fn walk_from<N: Node, V: Asked>(
&self,
path: impl Into<Path<N, V>>,
value: &V::Ask,
) -> Result<Walk<'_, N>> {
let at = self.matching::<N, V>(path.into(), value)?;
Ok(Walk {
g: self,
at,
marker: PhantomData,
})
}
pub fn find<N: Node, V: Asked>(
&self,
path: impl Into<Path<N, V>>,
value: &V::Ask,
) -> Result<Vec<N>> {
let path = path.into();
let key = key_of(value, path.kind(), "this value")?;
self.read(|s| {
let Some(kind) = s.node_seen::<N>() else {
return Ok(Vec::new());
};
let mut out = Vec::new();
let mut bad = None;
s.g.find_nodes(path.path(), &key, |raw, doc| {
if !s.is(raw, kind) {
return;
}
match N::read(doc) {
Ok(node) => out.push(node),
Err(e) => bad = bad.take().or(Some(e)),
}
})?;
match bad {
Some(e) => Err(e),
None => Ok(out),
}
})
}
pub fn count_at<N: Node, V: Asked>(
&self,
path: impl Into<Path<N, V>>,
value: &V::Ask,
) -> Result<usize> {
Ok(self.matching::<N, V>(path.into(), value)?.len())
}
pub fn memory_bytes(&self) -> Result<usize> {
self.read(|s| Ok(s.memory_bytes()))
}
fn matching<N: Node, V: Asked>(&self, path: Path<N, V>, value: &V::Ask) -> Result<Vec<u64>> {
let key = key_of(value, path.kind(), "this value")?;
self.read(|s| {
let Some(kind) = s.node_seen::<N>() else {
return Ok(Vec::new());
};
let mut out = Vec::new();
s.g.find_nodes(path.path(), &key, |raw, _| {
if s.is(raw, kind) {
out.push(raw);
}
})?;
Ok(out)
})
}
fn step<E: Edge>(&self, from: u64, dir: Dir) -> Result<Vec<u64>> {
self.read(|s| {
Ok(s.edge_seen::<E>()
.map_or_else(Vec::new, |label| s.g.neighbours(from, label, dir).to_vec()))
})
}
fn frontier<E: Edge>(&self, at: &[u64], dir: Dir) -> Result<Vec<u64>> {
self.read(|s| {
let Some(label) = s.edge_seen::<E>() else {
return Ok(Vec::new());
};
for node in at {
s.g.prefetch(*node, label, dir);
}
let mut out = Vec::new();
for node in at {
out.extend_from_slice(s.g.neighbours(*node, label, dir));
}
out.sort_unstable();
out.dedup();
Ok(out)
})
}
}
pub struct Walk<'a, N> {
g: &'a Graph,
at: Vec<u64>,
marker: PhantomData<fn() -> N>,
}
impl<N: Node> core::fmt::Debug for Walk<'_, N> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Walk")
.field("on", &N::LABEL)
.field("len", &self.at.len())
.finish()
}
}
impl<'a, N: Node> Walk<'a, N> {
pub fn out<E: Edge<From = N>>(self) -> Result<Walk<'a, E::To>> {
let at = self.g.frontier::<E>(&self.at, Dir::Out)?;
Ok(Walk {
g: self.g,
at,
marker: PhantomData,
})
}
pub fn incoming<E: Edge<To = N>>(self) -> Result<Walk<'a, E::From>> {
let at = self.g.frontier::<E>(&self.at, Dir::In)?;
Ok(Walk {
g: self.g,
at,
marker: PhantomData,
})
}
pub fn filter(self, mut keep: impl FnMut(&N) -> bool) -> Result<Walk<'a, N>> {
let mut at = Vec::with_capacity(self.at.len());
self.g.read(|s| {
for raw in &self.at {
if let Some(doc) = s.g.node(*raw)
&& keep(&N::read(doc)?)
{
at.push(*raw);
}
}
Ok(())
})?;
Ok(Walk {
g: self.g,
at,
marker: PhantomData,
})
}
#[must_use]
pub fn len(&self) -> usize {
self.at.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.at.is_empty()
}
#[must_use]
pub fn ids(self) -> Vec<Id<N>> {
ids(self.at)
}
pub fn nodes(self) -> Result<Vec<N>> {
self.g.read(|s| {
let mut out = Vec::with_capacity(self.at.len());
for raw in &self.at {
if let Some(doc) = s.g.node(*raw) {
out.push(N::read(doc)?);
}
}
Ok(out)
})
}
}
fn ids<N>(raw: Vec<u64>) -> Vec<Id<N>> {
raw.into_iter().map(Id::new).collect()
}
fn tagged(kind: u32, key: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(4 + key.len());
out.extend_from_slice(&kind.to_le_bytes());
out.extend_from_slice(key);
out
}
fn gone<N: Node>(raw: u64) -> Error {
Error::fmt(
Code::NotFound,
format_args!(
"{}#{raw} is not in this graph, so there is nothing to put an edge on. Add the node first, or use the id that add() answered with",
N::LABEL
),
)
}