object-rainbow-append-tree 0.0.0-a.11

use std::{fmt::Debug, future::ready, marker::PhantomData};

use object_rainbow::{
    Enum, ExtraFor, Fetch, FullHash, Inline, InlineOutput, ListHashes, Object, Output, Parse,
    ParseAsInline, ParseInline, ParseInput, PointInput, Tagged, ToOutput, Topological, Traversible,
    assert_impl, numeric::Be,
};
use object_rainbow_point::{IntoPoint, Point};
use typenum::{U256, Unsigned};

#[derive(ToOutput, InlineOutput, Tagged, ListHashes, Topological, Parse)]
#[topology(recursive)]
struct Node<T, N, M> {
    _capacity: PhantomData<N>,
    _marker: PhantomData<M>,
    #[tags(skip)]
    #[parse(unchecked)]
    #[topology(unchecked)]
    prev: Option<Point<Self>>,
    items: Vec<T>,
}

impl<T, N, M> Drop for Node<T, N, M> {
    fn drop(&mut self) {
        self.items.drain(..).rev().for_each(|_| {});
        while let Some(prev) = self.prev.take()
            && let Some(node) = prev.try_unwrap()
        {
            *self = node;
        }
    }
}

impl<T: std::fmt::Debug, N, M> std::fmt::Debug for Node<T, N, M> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Node")
            .field("_capacity", &self._capacity)
            .field("_marker", &self._marker)
            .field("prev", &self.prev)
            .field("items", &self.items)
            .finish()
    }
}

assert_impl!(
    impl<E, T, N, M> Object<E> for Node<T, N, M>
    where
        E: 'static + Send + Sync + Clone,
        N: 'static + Send + Sync + Clone,
        M: 'static + Send + Sync + Clone,
        T: Inline<E>,
    {
    }
);

impl<T: PartialEq, N, M> PartialEq for Node<T, N, M> {
    fn eq(&self, other: &Self) -> bool {
        self.prev == other.prev && self.items == other.items
    }
}

impl<T: Eq, N, M> Eq for Node<T, N, M> {}

trait History: Sized + Send + Sync {
    type History: Send + Sync;
    type Block: Send + Sync;
    const CAPACITY: u64;
}

trait ToContiguousOutput: History {
    fn to_contiguous_output(&self, history: &Self::History, output: &mut impl Output);
}

trait ParseWithLen<I: ParseInput>: History {
    fn parse_with_len(input: &mut I, len: u64) -> object_rainbow::Result<(Self, Self::History)>;
}

#[derive(Debug, thiserror::Error)]
pub enum PushError<T> {
    #[error("leaf overflow")]
    LeafOverflow(T),
    #[error("non-leaf overflow")]
    NonLeafOverflow(T),
    #[error("root overflow")]
    RootOverflow(T),
}

impl<T> PushError<T> {
    pub fn into_value(self) -> T {
        match self {
            PushError::LeafOverflow(value) => value,
            PushError::NonLeafOverflow(value) => value,
            PushError::RootOverflow(value) => value,
        }
    }
}

impl<T: 'static + Send + Sync + Debug> From<PushError<T>> for object_rainbow::Error {
    fn from(value: PushError<T>) -> Self {
        Self::operation(value)
    }
}

trait Push: Clone + History {
    type T: Send + Sync;
    fn get(
        &self,
        index: u64,
        history: Option<&Self::History>,
    ) -> impl Send + Future<Output = object_rainbow::Result<Self::T>>;
    fn push(
        &mut self,
        len: u64,
        value: Self::T,
        history: &mut Self::History,
    ) -> Result<(), PushError<Self::T>>;
    fn last<'a>(&'a self, history: &'a Self::History) -> Option<&'a Self::T>;
    fn from_value(prev: Point<Self>, history: &mut Self::History, value: Self::T) -> Self;
    fn to_point(&self, history: &Self::History) -> Point<Self>;
}

impl<T: Clone, N, M> Clone for Node<T, N, M> {
    fn clone(&self) -> Self {
        Self {
            _capacity: PhantomData,
            _marker: PhantomData,
            prev: self.prev.clone(),
            items: self.items.clone(),
        }
    }
}

impl<T, N, M> Default for Node<T, N, M> {
    fn default() -> Self {
        Self::new(None, Vec::new())
    }
}

impl<T, N, M> Node<T, N, M> {
    const fn new(prev: Option<Point<Self>>, items: Vec<T>) -> Self {
        Self {
            _capacity: PhantomData,
            _marker: PhantomData,
            prev,
            items,
        }
    }
}

struct Leaf;

impl<T: Send + Sync, N: Send + Sync + Unsigned> History for Node<T, N, Leaf> {
    type History = ();
    type Block = Self;
    const CAPACITY: u64 = N::U64;
}

impl<T: InlineOutput + Send + Sync, N: Send + Sync + Unsigned> ToContiguousOutput
    for Node<T, N, Leaf>
{
    fn to_contiguous_output(&self, (): &Self::History, output: &mut impl Output) {
        self.to_output(output);
    }
}

impl<T: ParseInline<I> + Send + Sync, N: Send + Sync + Unsigned, I: ParseInput> ParseWithLen<I>
    for Node<T, N, Leaf>
where
    Point<Self>: ParseInline<I>,
{
    fn parse_with_len(input: &mut I, len: u64) -> object_rainbow::Result<(Self, Self::History)> {
        if len > N::U64 {
            return Err(object_rainbow::error_parse!("overflow"));
        }
        Ok((
            Self::new(
                input.parse_inline()?,
                input.parse_vec_n(
                    len.try_into()
                        .map_err(|_| object_rainbow::error_parse!("overflow"))?,
                )?,
            ),
            (),
        ))
    }
}

impl<T: Send + Sync + Clone + Traversible + InlineOutput, N: Send + Sync + Unsigned> Push
    for Node<T, N, Leaf>
{
    type T = T;

    fn get(
        &self,
        index: u64,
        _: Option<&Self::History>,
    ) -> impl Send + Future<Output = object_rainbow::Result<Self::T>> {
        ready(
            usize::try_from(index)
                .map_err(|_| object_rainbow::Error::UnsupportedLength)
                .and_then(|index| {
                    self.items.get(index).cloned().ok_or_else(|| {
                        object_rainbow::error_consistency!(
                            "out of bounds L {index}/{}",
                            self.items.len()
                        )
                    })
                }),
        )
    }

    fn push(
        &mut self,
        len: u64,
        value: Self::T,
        (): &mut Self::History,
    ) -> Result<(), PushError<Self::T>> {
        if len != (self.items.len() as u64) {
            panic!("a node has been parsed with invalid length")
        } else if self.items.len() >= N::USIZE {
            Err(PushError::LeafOverflow(value))
        } else {
            self.items.push(value);
            Ok(())
        }
    }

    fn last<'a>(&'a self, (): &'a Self::History) -> Option<&'a Self::T> {
        self.items.last()
    }

    fn from_value(prev: Point<Self>, (): &mut Self::History, value: Self::T) -> Self {
        Self::new(Some(prev), vec![value])
    }

    fn to_point(&self, (): &Self::History) -> Point<Self> {
        self.clone().point()
    }
}

struct NonLeaf;

impl<T: Send + Sync + History, N: Send + Sync + Unsigned> History for Node<Point<T>, N, NonLeaf> {
    type History = (T, T::History);
    type Block = T::Block;
    const CAPACITY: u64 = N::U64.saturating_mul(T::CAPACITY);
}

impl<T: ToContiguousOutput, N: Send + Sync + Unsigned> ToContiguousOutput
    for Node<Point<T>, N, NonLeaf>
{
    fn to_contiguous_output(&self, (child, history): &Self::History, output: &mut impl Output) {
        self.prev.to_output(output);
        self.items.to_output(output);
        child.to_contiguous_output(history, output);
    }
}

impl<
    T: 'static + Send + Sync + ParseWithLen<I> + FullHash,
    N: Send + Sync + Unsigned,
    I: PointInput<Extra: Send + Sync + ExtraFor<T>>,
> ParseWithLen<I> for Node<Point<T>, N, NonLeaf>
where
    Point<T>: ParseInline<I>,
    Point<Self>: ParseInline<I>,
{
    fn parse_with_len(input: &mut I, len: u64) -> object_rainbow::Result<(Self, Self::History)> {
        let own = len / T::CAPACITY
            + if len.is_multiple_of(T::CAPACITY) {
                0
            } else {
                1
            };
        if own > N::U64 {
            return Err(object_rainbow::error_parse!("overflow"));
        }
        let prev = input.parse_inline()?;
        let items = input.parse_vec_n::<Point<T>>(
            own.saturating_sub(1)
                .try_into()
                .map_err(|_| object_rainbow::error_parse!("overflow"))?,
        )?;
        let history = T::parse_with_len(
            input,
            if len.is_multiple_of(T::CAPACITY) {
                T::CAPACITY
            } else {
                len % T::CAPACITY
            },
        )?;
        Ok((Self::new(prev, items), history))
    }
}

impl<T: Push + Traversible, N: Send + Sync + Unsigned> Push for Node<Point<T>, N, NonLeaf> {
    type T = T::T;

    fn get(
        &self,
        index: u64,
        history: Option<&Self::History>,
    ) -> impl Send + Future<Output = object_rainbow::Result<Self::T>> {
        async move {
            let n = usize::try_from(index / T::CAPACITY)
                .map_err(|_| object_rainbow::Error::UnsupportedLength)?;
            let r = index % T::CAPACITY;
            println!("g {n} {}", self.items.len());
            if let Some((node, history)) = history
                && n == self.items.len()
            {
                println!("AAA");
                node.get(r, Some(history)).await
            } else {
                println!("BBB");
                self.items
                    .get(n)
                    .ok_or_else(|| {
                        object_rainbow::error_consistency!(
                            "out of bounds N {n}/{}",
                            self.items.len(),
                        )
                    })?
                    .fetch()
                    .await?
                    .get(r, None)
                    .await
            }
        }
    }

    fn push(
        &mut self,
        len: u64,
        value: Self::T,
        (node, history): &mut Self::History,
    ) -> Result<(), PushError<Self::T>> {
        if len.is_multiple_of(T::CAPACITY) {
            if len / T::CAPACITY != (self.items.len() as u64 + 1) {
                panic!("a node has been parsed with invalid length");
            }
            if (self.items.len() + 1) >= N::USIZE {
                return Err(PushError::NonLeafOverflow(value));
            }
            let node =
                std::mem::replace(node, T::from_value(node.to_point(history), history, value));
            self.items.push(node.point());
            Ok(())
        } else {
            node.push(len % T::CAPACITY, value, history)?;
            Ok(())
        }
    }

    fn last<'a>(&'a self, (node, history): &'a Self::History) -> Option<&'a Self::T> {
        node.last(history)
    }

    fn from_value(prev: Point<Self>, (child, history): &mut Self::History, value: Self::T) -> Self {
        *child = T::from_value(child.clone().point(), history, value);
        Self::new(Some(prev), vec![])
    }

    fn to_point(&self, (child, history): &Self::History) -> Point<Self> {
        let mut node = self.clone();
        node.items.push(child.to_point(history));
        node.point()
    }
}

impl<T: Push<History: Clone> + Traversible, N: Send + Sync + Unsigned> Node<Point<T>, N, NonLeaf> {
    fn from_inner(
        inner: T,
        history: &mut T::History,
        value: T::T,
    ) -> (Self, <Self as History>::History) {
        let inner = inner.point();
        let next = T::from_value(inner.clone(), history, value);
        let parent = Self::new(None, vec![inner]);
        (parent, (next, history.clone()))
    }
}

type N1<T> = Node<T, U256, Leaf>;
type N2<T> = Node<Point<N1<T>>, U256, NonLeaf>;
type N3<T> = Node<Point<N2<T>>, U256, NonLeaf>;
type N4<T> = Node<Point<N3<T>>, U256, NonLeaf>;
type N5<T> = Node<Point<N4<T>>, U256, NonLeaf>;
type N6<T> = Node<Point<N5<T>>, U256, NonLeaf>;
type N7<T> = Node<Point<N6<T>>, U256, NonLeaf>;
type N8<T> = Node<Point<N7<T>>, U256, NonLeaf>;

assert_impl!(
    impl<T> Push for N1<T> where T: Send + Sync + Clone + Traversible + InlineOutput {}
);
assert_impl!(
    impl<T> Push for N2<T> where T: Send + Sync + Clone + Traversible + InlineOutput {}
);
assert_impl!(
    impl<T> Push for N3<T> where T: Send + Sync + Clone + Traversible + InlineOutput {}
);
assert_impl!(
    impl<T> Push for N4<T> where T: Send + Sync + Clone + Traversible + InlineOutput {}
);
assert_impl!(
    impl<T> Push for N5<T> where T: Send + Sync + Clone + Traversible + InlineOutput {}
);
assert_impl!(
    impl<T> Push for N6<T> where T: Send + Sync + Clone + Traversible + InlineOutput {}
);
assert_impl!(
    impl<T> Push for N7<T> where T: Send + Sync + Clone + Traversible + InlineOutput {}
);
assert_impl!(
    impl<T> Push for N8<T> where T: Send + Sync + Clone + Traversible + InlineOutput {}
);

type H1 = ();
type H2<T> = (N1<T>, H1);
type H3<T> = (N2<T>, H2<T>);
type H4<T> = (N3<T>, H3<T>);
type H5<T> = (N4<T>, H4<T>);
type H6<T> = (N5<T>, H5<T>);
type H7<T> = (N6<T>, H6<T>);
type H8<T> = (N7<T>, H7<T>);

#[derive(Enum, Tagged, ListHashes, Topological, Clone, PartialEq, Eq, Debug)]
enum TreeKind<T> {
    N1((N1<T>, H1)),
    N2((N2<T>, H2<T>)),
    N3((N3<T>, H3<T>)),
    N4((N4<T>, H4<T>)),
    N5((N5<T>, H5<T>)),
    N6((N6<T>, H6<T>)),
    N7((N7<T>, H7<T>)),
    N8((N8<T>, H8<T>)),
}

#[derive(Tagged, ListHashes, Topological, Clone, ParseAsInline, PartialEq, Eq, Debug)]
pub struct AppendTree<T> {
    len: Be<u64>,
    kind: TreeKind<T>,
}

impl<T: Send + Sync + InlineOutput> ToOutput for AppendTree<T> {
    fn to_output(&self, output: &mut impl Output) {
        self.len.to_output(output);
        match &self.kind {
            TreeKind::N1((node, history)) => node.to_contiguous_output(history, output),
            TreeKind::N2((node, history)) => node.to_contiguous_output(history, output),
            TreeKind::N3((node, history)) => node.to_contiguous_output(history, output),
            TreeKind::N4((node, history)) => node.to_contiguous_output(history, output),
            TreeKind::N5((node, history)) => node.to_contiguous_output(history, output),
            TreeKind::N6((node, history)) => node.to_contiguous_output(history, output),
            TreeKind::N7((node, history)) => node.to_contiguous_output(history, output),
            TreeKind::N8((node, history)) => node.to_contiguous_output(history, output),
        }
    }
}

impl<T: Send + Sync + InlineOutput> InlineOutput for AppendTree<T> {}

const C1: u64 = 256;
const C2: u64 = 256 * C1;
const C3: u64 = 256 * C2;
const C4: u64 = 256 * C3;
const C5: u64 = 256 * C4;
const C6: u64 = 256 * C5;
const C7: u64 = 256 * C6;
const C8: u64 = C7.saturating_mul(256);
const C2_MIN: u64 = C1 + 1;
const C3_MIN: u64 = C2 + 1;
const C4_MIN: u64 = C3 + 1;
const C5_MIN: u64 = C4 + 1;
const C6_MIN: u64 = C5 + 1;
const C7_MIN: u64 = C6 + 1;
const C8_MIN: u64 = C7 + 1;

impl<T, I: ParseInput> ParseInline<I> for AppendTree<T>
where
    N1<T>: ParseWithLen<I, History = H1>,
    N2<T>: ParseWithLen<I, History = H2<T>>,
    N3<T>: ParseWithLen<I, History = H3<T>>,
    N4<T>: ParseWithLen<I, History = H4<T>>,
    N5<T>: ParseWithLen<I, History = H5<T>>,
    N6<T>: ParseWithLen<I, History = H6<T>>,
    N7<T>: ParseWithLen<I, History = H7<T>>,
    N8<T>: ParseWithLen<I, History = H8<T>>,
{
    fn parse_inline(input: &mut I) -> object_rainbow::Result<Self> {
        let len = input.parse_inline::<Be<u64>>()?;
        let kind = match len.0 {
            0..=C1 => TreeKind::N1(N1::<T>::parse_with_len(input, len.0)?),
            C2_MIN..=C2 => TreeKind::N2(N2::<T>::parse_with_len(input, len.0)?),
            C3_MIN..=C3 => TreeKind::N3(N3::<T>::parse_with_len(input, len.0)?),
            C4_MIN..=C4 => TreeKind::N4(N4::<T>::parse_with_len(input, len.0)?),
            C5_MIN..=C5 => TreeKind::N5(N5::<T>::parse_with_len(input, len.0)?),
            C6_MIN..=C6 => TreeKind::N6(N6::<T>::parse_with_len(input, len.0)?),
            C7_MIN..=C7 => TreeKind::N7(N7::<T>::parse_with_len(input, len.0)?),
            C8_MIN..=C8 => TreeKind::N8(N8::<T>::parse_with_len(input, len.0)?),
        };
        Ok(Self { len, kind })
    }
}

assert_impl!(
    impl<T, E> Inline<E> for AppendTree<T>
    where
        E: 'static + Send + Sync + Clone,
        T: Inline<E>,
    {
    }
);

impl<T: Send + Sync + Clone + Traversible + InlineOutput> AppendTree<T> {
    pub const fn new() -> Self {
        Self {
            len: Be::<u64>::new(0u64),
            kind: TreeKind::N1((Node::new(None, Vec::new()), ())),
        }
    }

    pub async fn get(&self, index: u64) -> object_rainbow::Result<Option<T>> {
        if index < self.len.0 {
            match &self.kind {
                TreeKind::N1((node, history)) => node.get(index, Some(history)).await,
                TreeKind::N2((node, history)) => node.get(index, Some(history)).await,
                TreeKind::N3((node, history)) => node.get(index, Some(history)).await,
                TreeKind::N4((node, history)) => node.get(index, Some(history)).await,
                TreeKind::N5((node, history)) => node.get(index, Some(history)).await,
                TreeKind::N6((node, history)) => node.get(index, Some(history)).await,
                TreeKind::N7((node, history)) => node.get(index, Some(history)).await,
                TreeKind::N8((node, history)) => node.get(index, Some(history)).await,
            }
            .map(Some)
        } else {
            Ok(None)
        }
    }

    pub fn push(&mut self, value: T) -> Result<(), PushError<T>> {
        let len = self.len.0;
        macro_rules! upgrade {
            ($history:ident, $node:ident, $child:ident, $parent:ident) => {
                if len == $child::<T>::CAPACITY {
                    self.kind =
                        TreeKind::$parent(Node::from_inner(std::mem::take($node), $history, value));
                } else {
                    $node.push(len, value, $history)?;
                }
            };
        }
        match &mut self.kind {
            TreeKind::N1((node, history)) => upgrade!(history, node, N1, N2),
            TreeKind::N2((node, history)) => upgrade!(history, node, N2, N3),
            TreeKind::N3((node, history)) => upgrade!(history, node, N3, N4),
            TreeKind::N4((node, history)) => upgrade!(history, node, N4, N5),
            TreeKind::N5((node, history)) => upgrade!(history, node, N5, N6),
            TreeKind::N6((node, history)) => upgrade!(history, node, N6, N7),
            TreeKind::N7((node, history)) => upgrade!(history, node, N7, N8),
            TreeKind::N8((node, history)) => {
                if len == N8::<T>::CAPACITY {
                    return Err(PushError::RootOverflow(value));
                } else {
                    node.push(len, value, history)?;
                }
            }
        }
        self.len.0 += 1;
        Ok(())
    }

    pub fn is_empty(&self) -> bool {
        self.len.0 == 0
    }

    pub fn len(&self) -> u64 {
        self.len.0
    }

    pub fn last(&self) -> Option<&T> {
        match &self.kind {
            TreeKind::N1((node, history)) => Push::last(node, history),
            TreeKind::N2((node, history)) => Push::last(node, history),
            TreeKind::N3((node, history)) => Push::last(node, history),
            TreeKind::N4((node, history)) => Push::last(node, history),
            TreeKind::N5((node, history)) => Push::last(node, history),
            TreeKind::N6((node, history)) => Push::last(node, history),
            TreeKind::N7((node, history)) => Push::last(node, history),
            TreeKind::N8((node, history)) => Push::last(node, history),
        }
    }
}

impl<T: Send + Sync + Clone + Traversible + InlineOutput> Default for AppendTree<T> {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod test {
    use macro_rules_attribute::apply;
    use object_rainbow::{ParseSlice, numeric::Le};
    use smol_macros::test;

    use crate::AppendTree;

    #[apply(test!)]
    async fn reparse() -> object_rainbow::Result<()> {
        let mut tree = AppendTree::<Le<u64>>::new();
        for i in 0..100_000u64 {
            tree.push(Le(i))?;
            let new = tree.reparse()?;
            assert_eq!(new, tree);
            tree = new;
            assert_eq!(tree.get(i).await?.unwrap().0, i);
        }
        Ok(())
    }

    #[apply(test!)]
    async fn get() -> object_rainbow::Result<()> {
        let mut tree = AppendTree::<Le<u64>>::new();
        for i in 0..100_000u64 {
            tree.push(Le(i))?;
            assert_eq!(tree.get(i).await?.unwrap().0, i);
            assert_eq!(tree.get(i / 2).await?.unwrap().0, i / 2);
            assert_eq!(tree.get(i / 3).await?.unwrap().0, i / 3);
            assert_eq!(tree.get(i / 17).await?.unwrap().0, i / 17);
            assert_eq!(tree.get(i / 101).await?.unwrap().0, i / 101);
        }
        Ok(())
    }
}