state_tree/

tree.rs

pub const DELAY_ADDITIONAL_OFFSET: usize = 2;

/// State Tree structure.

//on attributes, see https://github.com/rkyv/rkyv/blob/main/rkyv/examples/json_like_schema.rs
#[derive(Clone, PartialEq, Eq)]
pub enum StateTree {
    Delay {
        readidx: u64,
        writeidx: u64,
        data: Vec<u64>, //assume we are using only mono f64 data
    },
    Mem {
        data: Vec<u64>, //assume we are using only mono f64 data
    },
    Feed {
        data: Vec<u64>, //assume we are using generic data, might be tuple of float
    },
    FnCall(Vec<StateTree>),
}

impl std::fmt::Display for StateTree {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            StateTree::Delay {
                readidx,
                writeidx,
                data,
            } => write!(
                f,
                "Delay(readidx: {}, writeidx: {}, data: {:?} ...)",
                readidx,
                writeidx,
                data.iter().take(10).collect::<Vec<&u64>>()
            ),
            StateTree::Mem { data } => write!(f, "Mem(data: {data:?})"),
            StateTree::Feed { data } => write!(f, "Feed(data: {data:?})"),
            StateTree::FnCall(children) => {
                let children_str: Vec<String> = children.iter().map(|c| format!("{c}")).collect();
                write!(f, "FnCall([{}])", children_str.join(", "))
            }
        }
    }
}
impl std::fmt::Debug for StateTree {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            StateTree::Delay {
                readidx,
                writeidx,
                data,
            } => write!(
                f,
                "Delay(readidx: {}, writeidx: {}, data: {:?} ...)",
                readidx,
                writeidx,
                data.iter().take(10).collect::<Vec<&u64>>()
            ),
            StateTree::Mem { data } => write!(f, "Mem(data: {data:?})"),
            StateTree::Feed { data } => write!(f, "Feed(data: {data:?})"),
            StateTree::FnCall(children) => {
                let children_str: Vec<String> = children.iter().map(|c| format!("{c:?}")).collect();
                write!(f, "FnCall([{}])", children_str.join(", "))
            }
        }
    }
}
impl<T: SizedType> From<StateTreeSkeleton<T>> for StateTree {
    //create empty StateTree from StateTreeSkeleton
    fn from(skeleton: StateTreeSkeleton<T>) -> Self {
        match skeleton {
            StateTreeSkeleton::Delay { len } => StateTree::Delay {
                readidx: 0,
                writeidx: 0,
                data: vec![0; len as usize],
            },
            StateTreeSkeleton::Mem(t) => StateTree::Mem {
                data: vec![0; t.word_size() as usize],
            },
            StateTreeSkeleton::Feed(t) => StateTree::Feed {
                data: vec![0; t.word_size() as usize],
            },
            StateTreeSkeleton::FnCall(children_layout) => StateTree::FnCall(
                children_layout
                    .into_iter()
                    .map(|child_layout| StateTree::from(*child_layout))
                    .collect(),
            ),
        }
    }
}

impl StateTree {
    /// パスを指定して、イミュータブルなノードへの参照を取得する
    pub fn get_node(&self, path: &[usize]) -> Option<&StateTree> {
        let mut current = self;
        for &index in path {
            if let StateTree::FnCall(children) = current {
                current = children.get(index)?;
            } else {
                // パスが深すぎるか、FnCallではないノードを指している
                return None;
            }
        }
        Some(current)
    }

    /// パスを指定して、ミュータブルなノードへの参照を取得する
    pub fn get_node_mut(&mut self, path: &[usize]) -> Option<&mut StateTree> {
        let mut current = self;
        for &index in path {
            if let StateTree::FnCall(children) = current {
                current = children.get_mut(index)?;
            } else {
                // パスが深すぎるか、FnCallではないノードを指している
                return None;
            }
        }
        Some(current)
    }
}

pub fn serialize_tree_untagged(tree: StateTree) -> Vec<u64> {
    match tree {
        StateTree::Delay {
            readidx,
            writeidx,
            data,
        } => itertools::concat([vec![readidx, writeidx], data]),
        StateTree::Mem { data } | StateTree::Feed { data } => data,
        StateTree::FnCall(state_trees) => {
            itertools::concat(state_trees.into_iter().map(serialize_tree_untagged))
        }
    }
}

pub trait SizedType: std::fmt::Debug {
    fn word_size(&self) -> u64;
}

/// This data represents just a memory layout on a flat array, do not own actual data.
#[derive(Debug, Clone)]
pub enum StateTreeSkeleton<T: SizedType> {
    Delay {
        len: u64, //assume we are using only mono f64 data
    },
    Mem(T),
    Feed(T),
    FnCall(Vec<Box<StateTreeSkeleton<T>>>),
}
impl<T: SizedType> StateTreeSkeleton<T> {
    pub fn total_size(&self) -> u64 {
        match self {
            StateTreeSkeleton::Delay { len } => DELAY_ADDITIONAL_OFFSET as u64 + *len,
            StateTreeSkeleton::Mem(t) | StateTreeSkeleton::Feed(t) => t.word_size(),
            StateTreeSkeleton::FnCall(children_layout) => children_layout
                .iter()
                .map(|child_layout| child_layout.total_size())
                .sum(),
        }
    }
}
impl<T: SizedType> PartialEq for StateTreeSkeleton<T> {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::Delay { len: l_len }, Self::Delay { len: r_len }) => l_len == r_len,
            (Self::Mem(l0), Self::Mem(r0)) => l0.word_size() == r0.word_size(),
            (Self::Feed(l0), Self::Feed(r0)) => l0.word_size() == r0.word_size(),
            (Self::FnCall(l0), Self::FnCall(r0)) => l0 == r0,
            _ => false,
        }
    }
}

fn deserialize_tree_untagged_rec<T: SizedType>(
    data: &[u64],
    data_layout: &StateTreeSkeleton<T>,
) -> Option<(StateTree, usize)> {
    match data_layout {
        StateTreeSkeleton::Delay { len } => {
            let readidx = data.first().copied()?;
            let writeidx = data.get(1).copied()?;
            let d = data
                .get(DELAY_ADDITIONAL_OFFSET..DELAY_ADDITIONAL_OFFSET + (*len as usize))?
                .to_vec();
            Some((
                StateTree::Delay {
                    readidx,
                    writeidx,
                    data: d,
                },
                DELAY_ADDITIONAL_OFFSET + (*len as usize),
            ))
        }
        StateTreeSkeleton::Mem(t) => {
            let size = t.word_size() as usize;
            let data = data.get(0..size)?.to_vec();
            Some((StateTree::Mem { data }, size))
        }
        StateTreeSkeleton::Feed(t) => {
            let size = t.word_size() as usize;
            let data = data.get(0..size)?.to_vec();
            Some((StateTree::Feed { data }, size))
        }
        StateTreeSkeleton::FnCall(children_layout) => {
            let (children, used) =
                children_layout
                    .iter()
                    .try_fold((vec![], 0), |(v, last_used), child_layout| {
                        let (child, used) =
                            deserialize_tree_untagged_rec(&data[last_used..], child_layout)?;

                        Some(([v, vec![child]].concat(), last_used + used))
                    })?;

            Some((StateTree::FnCall(children), used))
        }
    }
}

pub fn deserialize_tree_untagged<T: SizedType>(
    data: &[u64],
    data_layout: &StateTreeSkeleton<T>,
) -> Option<StateTree> {
    log::trace!("Deserializing  with layout: {data_layout:?}");
    if let Some((tree, used)) = deserialize_tree_untagged_rec(data, data_layout) {
        if used == data.len() { Some(tree) } else { None }
    } else {
        None
    }
}