geam 0.1.1

Experimental Rust-embedded execution runtime for typed Gleam programs
use super::schema::StringTreeSchema;
use crate::gleam_stdlib::GleamStdlibHostProfile;
use crate::{
    HostExternalEquality, HostExternalHashing, HostExternalInspection, HostExternalStorage,
    HostExternalStore,
};
use ecow::EcoString;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::mem;
use std::rc::Rc;

#[derive(Default)]
pub(in crate::gleam_stdlib) struct Stores {
    values: HostExternalStore<StringTreePayload>,
}

pub(crate) struct StringTreePayload {
    pub(crate) tree: StringTree,
}

#[derive(Clone)]
pub(crate) struct StringTree {
    root: Rc<StringTreeNode>,
}

struct StringTreeNode {
    byte_len: usize,
    kind: StringTreeNodeKind,
}

enum StringTreeNodeKind {
    Text(EcoString),
    Sequence(Box<[Rc<StringTreeNode>]>),
}

impl StringTree {
    pub(crate) fn text(text: EcoString) -> Self {
        Self {
            root: Rc::new(StringTreeNode {
                byte_len: text.len(),
                kind: StringTreeNodeKind::Text(text),
            }),
        }
    }

    pub(crate) fn sequence(trees: impl IntoIterator<Item = Self>) -> Self {
        let children = trees
            .into_iter()
            .map(|tree| tree.root)
            .collect::<Box<[_]>>();
        let byte_len = children.iter().map(|child| child.byte_len).sum();
        Self {
            root: Rc::new(StringTreeNode {
                byte_len,
                kind: StringTreeNodeKind::Sequence(children),
            }),
        }
    }

    pub(super) fn append(&self, suffix: &Self) -> Self {
        Self::sequence([self.clone(), suffix.clone()])
    }

    pub(super) fn byte_len(&self) -> usize {
        self.root.byte_len
    }

    pub(crate) fn flatten(&self) -> EcoString {
        let mut output = String::with_capacity(self.byte_len());
        let mut pending = vec![self.root.as_ref()];
        while let Some(node) = pending.pop() {
            match &node.kind {
                StringTreeNodeKind::Text(text) => output.push_str(text),
                StringTreeNodeKind::Sequence(children) => {
                    pending.extend(children.iter().rev().map(Rc::as_ref));
                }
            }
        }
        output.into()
    }

    pub(crate) fn structurally_equal(&self, other: &Self) -> bool {
        if self.byte_len() != other.byte_len() {
            return false;
        }

        let mut pending = vec![(self.root.as_ref(), other.root.as_ref())];
        while let Some((left, right)) = pending.pop() {
            if std::ptr::eq(left, right) {
                continue;
            }
            match (&left.kind, &right.kind) {
                (StringTreeNodeKind::Text(left), StringTreeNodeKind::Text(right)) => {
                    if left != right {
                        return false;
                    }
                }
                (StringTreeNodeKind::Sequence(left), StringTreeNodeKind::Sequence(right))
                    if left.len() == right.len() =>
                {
                    pending.extend(
                        left.iter()
                            .zip(right.iter())
                            .rev()
                            .map(|(left, right)| (left.as_ref(), right.as_ref())),
                    );
                }
                _ => return false,
            }
        }
        true
    }

    pub(crate) fn structural_hash(&self) -> u64 {
        let mut hasher = DefaultHasher::new();
        let mut pending = vec![self.root.as_ref()];
        while let Some(node) = pending.pop() {
            match &node.kind {
                StringTreeNodeKind::Text(text) => {
                    0_u8.hash(&mut hasher);
                    text.hash(&mut hasher);
                }
                StringTreeNodeKind::Sequence(children) => {
                    1_u8.hash(&mut hasher);
                    children.len().hash(&mut hasher);
                    pending.extend(children.iter().rev().map(Rc::as_ref));
                }
            }
        }
        hasher.finish()
    }

    pub(super) fn inspect(&self) -> EcoString {
        format!("string_tree.from_string({:?})", self.flatten()).into()
    }
}

impl Drop for StringTreeNode {
    fn drop(&mut self) {
        let mut pending = take_children(&mut self.kind);
        while let Some(child) = pending.pop() {
            let Ok(mut child) = Rc::try_unwrap(child) else {
                continue;
            };
            pending.extend(take_children(&mut child.kind));
        }
    }
}

fn take_children(kind: &mut StringTreeNodeKind) -> Vec<Rc<StringTreeNode>> {
    match mem::replace(kind, StringTreeNodeKind::Text(EcoString::new())) {
        StringTreeNodeKind::Text(_) => Vec::new(),
        StringTreeNodeKind::Sequence(children) => children.into_vec(),
    }
}

pub(crate) struct StringTreeExternalStorage;

impl<Profile> HostExternalStorage<Profile, StringTreeSchema> for StringTreeExternalStorage
where
    Profile: GleamStdlibHostProfile,
{
    type Payload = StringTreePayload;

    fn store(stores: &Profile::ExternalStores) -> &HostExternalStore<Self::Payload> {
        &Profile::gleam_stdlib_stores(stores).string_tree.values
    }

    fn source_equal(
        _context: &HostExternalEquality<'_>,
        left: &Self::Payload,
        right: &Self::Payload,
    ) -> bool {
        left.tree.structurally_equal(&right.tree)
    }

    fn source_hash(_context: &HostExternalHashing<'_>, value: &Self::Payload) -> u64 {
        value.tree.structural_hash()
    }

    fn inspect(_context: &HostExternalInspection<'_>, value: &Self::Payload) -> EcoString {
        value.tree.inspect()
    }
}

#[cfg(test)]
mod tests {
    use super::{
        StringTree, StringTreeExternalStorage, StringTreeNode, StringTreeNodeKind,
        StringTreePayload,
    };
    use crate::gleam_stdlib::{GleamStdlibProfile, GleamStdlibStores};
    use crate::{HostExternalEquality, HostExternalHashing, HostExternalStorage};
    use ecow::EcoString;
    use std::rc::Rc;

    fn children(tree: &StringTree) -> Option<&[Rc<StringTreeNode>]> {
        match &tree.root.kind {
            StringTreeNodeKind::Text(_) => None,
            StringTreeNodeKind::Sequence(children) => Some(children),
        }
    }

    #[test]
    fn preserves_structural_identity_separately_from_textual_content() {
        let flat = StringTree::text("ab".into());
        let segmented =
            StringTree::sequence([StringTree::text("a".into()), StringTree::text("b".into())]);
        let same =
            StringTree::sequence([StringTree::text("a".into()), StringTree::text("b".into())]);

        assert_eq!(flat.flatten(), segmented.flatten());
        assert!(!flat.structurally_equal(&segmented));
        assert!(segmented.structurally_equal(&same));
        assert!(segmented.structurally_equal(&segmented.clone()));
        assert!(!StringTree::text("a".into()).structurally_equal(&StringTree::text("bb".into())));
        assert!(!StringTree::text("a".into()).structurally_equal(&StringTree::text("b".into())));
        assert_eq!(segmented.structural_hash(), same.structural_hash());
        assert_eq!(segmented.byte_len(), 2);
        assert_eq!(segmented.inspect(), r#"string_tree.from_string("ab")"#);
    }

    #[test]
    fn append_shares_existing_subtrees_and_deep_chains_drop_iteratively() {
        let prefix = StringTree::text("a".into());
        let suffix = StringTree::text("b".into());
        let appended = prefix.append(&suffix);
        assert!(children(&prefix).is_none());
        let children = children(&appended).expect("append should create a sequence node");
        assert!(Rc::ptr_eq(&children[0], &prefix.root));
        assert!(Rc::ptr_eq(&children[1], &suffix.root));

        let mut deep = StringTree::text(EcoString::new());
        for _ in 0..50_000 {
            deep = deep.append(&StringTree::text("x".into()));
        }
        assert_eq!(deep.byte_len(), 50_000);
        assert_eq!(deep.flatten().len(), 50_000);
        drop(deep);
    }

    #[test]
    fn empty_text_and_empty_sequence_are_textually_but_not_structurally_equal() {
        let text = StringTree::text(EcoString::new());
        let sequence = StringTree::sequence([]);

        assert_eq!(text.flatten(), "");
        assert_eq!(sequence.flatten(), "");
        assert!(!text.structurally_equal(&sequence));
        assert_ne!(text.structural_hash(), sequence.structural_hash());
    }

    #[test]
    fn external_storage_delegates_structural_source_semantics() {
        type Profile = GleamStdlibProfile;

        let left = StringTreePayload {
            tree: StringTree::sequence([
                StringTree::text("a".into()),
                StringTree::text("b".into()),
            ]),
        };
        let right = StringTreePayload {
            tree: StringTree::sequence([
                StringTree::text("a".into()),
                StringTree::text("b".into()),
            ]),
        };
        let stored_equal =
            |_: &crate::runtime::StoredRuntimeValue, _: &crate::runtime::StoredRuntimeValue| true;
        let stored_hash = |_: &crate::runtime::StoredRuntimeValue| 0;
        let equality = HostExternalEquality::new(&stored_equal);
        let hashing = HostExternalHashing::new(&stored_hash);

        assert!(<StringTreeExternalStorage as HostExternalStorage<
            Profile,
            super::StringTreeSchema,
        >>::source_equal(&equality, &left, &right,));
        assert_eq!(
            <StringTreeExternalStorage as HostExternalStorage<
                Profile,
                super::StringTreeSchema,
            >>::source_hash(&hashing, &left,),
            <StringTreeExternalStorage as HostExternalStorage<
                Profile,
                super::StringTreeSchema,
            >>::source_hash(
                &hashing, &right,
            ),
        );
        let stores = GleamStdlibStores::default();
        let store = <StringTreeExternalStorage as HostExternalStorage<
            Profile,
            super::StringTreeSchema,
        >>::store(&stores);
        let lease =
            store.insert(
                left,
                <StringTreeExternalStorage as HostExternalStorage<
                    Profile,
                    super::StringTreeSchema,
                >>::source_equal,
                <StringTreeExternalStorage as HostExternalStorage<
                    Profile,
                    super::StringTreeSchema,
                >>::source_hash,
                <StringTreeExternalStorage as HostExternalStorage<
                    Profile,
                    super::StringTreeSchema,
                >>::inspect,
            );
        assert_eq!(store.view(&lease).tree.flatten(), "ab");
    }
}