typed-sled table joins

join.rs

use std::{marker::PhantomData, sync::Arc};
use typed_sled::{Tree, KV};

#[derive(Clone)]
pub struct JoinTree<I, T> {
    inner: I,
    joins: T,
}

pub struct InnerJoin<T, F, FP>(T, F, PhantomData<fn(FP)>);

pub trait TreeMeta {
    type Key: KV;
    type Value: KV;
}

impl<I: TreeMeta, T> TreeMeta for JoinTree<I, T> {
    type Key = I::Key;
    type Value = I::Value;
}

pub trait Join<K, V>: TreeMeta {
    type InnerJoin<F>;

    fn join<F>(self, tree: Tree<K, V>, on: F) -> Self::InnerJoin<F>
    where
        F: Fn(&Self::Value) -> &[K];
}

impl<K: KV, V: KV> TreeMeta for Tree<K, V> {
    type Key = K;
    type Value = V;
}

impl<K, V, A, B> Join<K, V> for Tree<A, B>
where
    K: KV,
    V: KV,
    A: KV,
    B: KV,
{
    type InnerJoin<F> = JoinTree<Self, (InnerJoin<Tree<K, V>, F, B>,)>;

    fn join<F>(self, tree: Tree<K, V>, on: F) -> Self::InnerJoin<F>
    where
        F: Fn(&Self::Value) -> &[K],
    {
        JoinTree {
            inner: self,
            joins: (InnerJoin(tree, on, PhantomData),),
        }
    }
}

impl<K, V, I, A> Join<K, V> for JoinTree<I, (A,)>
where
    I: TreeMeta,
{
    type InnerJoin<F> = JoinTree<I, (A, InnerJoin<Tree<K, V>, F, V>)>;

    fn join<F>(self, tree: Tree<K, V>, on: F) -> Self::InnerJoin<F>
    where
        F: Fn(&Self::Value) -> &[K],
    {
        let (a,) = self.joins;
        JoinTree {
            inner: self.inner,
            joins: (a, InnerJoin(tree, on, PhantomData)),
        }
    }
}
/// Implemented by [`InnerJoin`] and friends to encapsulate the view type.
pub trait JoinView {
    type Tree: TreeMeta;
    type FkSource;
    type Output<K, V>;

    fn get(
        &self,
        key: &<Self::Tree as TreeMeta>::Key,
    ) -> sled::Result<Self::Output<<Self::Tree as TreeMeta>::Key, <Self::Tree as TreeMeta>::Value>>;

    fn relate<'src>(&self, source: &'src Self::FkSource) -> &'src [<Self::Tree as TreeMeta>::Key];

    fn produce(&self, source: &Self::FkSource) -> sled::Result<JoinOutput<Self>> {
        Ok(self
            .relate(source)
            .iter()
            .map(|fk| self.get(fk))
            .collect::<sled::Result<Vec<_>>>()?)
    }
}

impl<K, V, F, P> JoinView for InnerJoin<Tree<K, V>, F, P>
where
    K: KV + Clone,
    V: KV,
    F: Fn(&P) -> &[K],
{
    type Tree = Tree<K, V>;
    type FkSource = P;
    type Output<OK, OV> = Option<(OK, OV)>;

    fn get(&self, key: &K) -> sled::Result<Self::Output<K, V>> {
        Ok(self.0.get(key)?.map(|value| (key.clone(), value)))
    }

    fn relate<'src>(&self, source: &'src Self::FkSource) -> &'src [<Self::Tree as TreeMeta>::Key] {
        self.1(source)
    }
}

pub type JoinOutputSingle<T> = <T as JoinView>::Output<
    <<T as JoinView>::Tree as TreeMeta>::Key,
    <<T as JoinView>::Tree as TreeMeta>::Value,
>;

pub type JoinOutput<T> = Vec<JoinOutputSingle<T>>;

impl<K, V, A> JoinTree<Tree<K, V>, (A,)>
where
    K: KV,
    V: KV,
    A: JoinView<FkSource = V>,
{
    pub fn get(&self, key: &K) -> sled::Result<Option<(V, JoinOutput<A>)>> {
        let Some(v) = self.inner.get(key)? else {return Ok(None)};
        let (a,) = &self.joins;
        let a = a.produce(&v)?;
        Ok(Some((v, a)))
    }
    pub fn range<RK>(
        &self,
        range: RK,
    ) -> sled::Result<impl Iterator<Item = (K, Arc<V>, JoinOutputSingle<A>)>>
    where
        RK: std::ops::RangeBounds<K>,
        K: std::fmt::Debug + Clone,
    {
        let (a,) = &self.joins;
        let mut out = Vec::new();
        for value in self.inner.range(range) {
            let (k, v) = value?;
            let v = Arc::new(v);
            let a = a.produce(&v)?;

            out.push(
                std::iter::repeat(k)
                    .zip(a)
                    .map(move |(k, a)| (k, Arc::clone(&v), a)),
            );
        }

        Ok(out.into_iter().flatten())
    }
}