conjure_cp_core/ast/

moo.rs

// NOTE
//
// we use a wrapper type over Arc, instead of just using Arc, so that we can implement traits on it
// (e.g. Uniplate, Serialize).
//
// As we are just using Arc for copy on write, not shared ownership, it is safe to break shared
// ownership in Moo's Uniplate implementation, by calling Arc::make_mut and Arc::new on modified
// values. In general, this is not safe for all Rc/Arc types, e.g. those that use Cell / RefCell
// internally.
//
// ~niklasdewally 13/08/25

use funcmap::{FuncMap, TryFuncMap};
use polyquine::Quine;
use proc_macro2::TokenStream;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_with::de::DeserializeAsWrap;
use serde_with::ser::SerializeAsWrap;
use serde_with::{DeserializeAs, SerializeAs};
use std::hash::{Hash, Hasher};
use std::ops::DerefMut;
use std::{collections::VecDeque, fmt::Display, ops::Deref, sync::Arc};
use uniplate::{
    Biplate, Tree, Uniplate,
    impl_helpers::{transmute_if_same_type, try_transmute_if_same_type},
    spez::try_biplate_to,
};

/// A clone-on-write, reference counted pointer to an AST type.
///
/// Cloning values of this type will not clone the underlying value until it is modified, e.g.,
/// with [`Moo::make_mut`].
///
/// Unlike `Rc` and `Arc`, trait implementations on this type do not need to preserve shared
/// ownership - that is, two pointers that used to point to the same value may not do so after
/// calling a trait method on them. In particular, calling Uniplate methods may cause a
/// clone-on-write to occur.
///
/// **Note:** like Box` and `Rc`, methods on `Moo` are all associated functions, which means you
/// have to call them as, e.g. `Moo::make_mut(&value)` instead of `value.make_mut()`. This is so
/// that there are no conflicts with the inner type `T`, which this type dereferences to.
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct Moo<T> {
    inner: Arc<T>,
}

impl<T: Quine> Quine for Moo<T> {
    fn ctor_tokens(&self) -> TokenStream {
        let inner = self.inner.as_ref().ctor_tokens();
        quote::quote! { ::conjure_cp::ast::Moo::new(#inner) }
    }
}

impl<T> Moo<T> {
    /// Constructs a new `Moo<T>`.
    pub fn new(value: T) -> Moo<T> {
        Moo {
            inner: Arc::new(value),
        }
    }
}

impl<T: Clone> Moo<T> {
    /// Makes a mutable reference into the given `Moo`.
    ///
    /// If there are other `Moo` pointers to the same allocation, then `make_mut` will `clone` the
    /// inner value to a new allocation to ensure unique ownership. This is also referred to as
    /// clone-on-write.
    pub fn make_mut(this: &mut Moo<T>) -> &mut T {
        Arc::make_mut(&mut this.inner)
    }

    /// If we have the only reference to T then unwrap it. Otherwise, clone T and return the clone.
    ///
    /// Assuming moo_t is of type `Moo<T>`, this function is functionally equivalent to
    /// `(*moo_t).clone()`, but will avoid cloning the inner value where possible.
    pub fn unwrap_or_clone(this: Moo<T>) -> T {
        Arc::unwrap_or_clone(this.inner)
    }
}

impl<T> AsRef<T> for Moo<T> {
    fn as_ref(&self) -> &T {
        self.inner.as_ref()
    }
}

impl<T> Deref for Moo<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.inner.deref()
    }
}

impl<T: Clone> DerefMut for Moo<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        Moo::make_mut(self)
    }
}

impl<T> Uniplate for Moo<T>
where
    T: Uniplate,
{
    fn uniplate(
        &self,
    ) -> (
        uniplate::Tree<Self>,
        Box<dyn Fn(uniplate::Tree<Self>) -> Self>,
    ) {
        let this = Moo::clone(self);

        // do not need to preserve shared ownership, so treat this identically to values of the
        // inner type.
        let (tree, ctx) = try_biplate_to!((**self).clone(), Moo<T>);
        (
            Tree::Many(VecDeque::from([tree.clone()])),
            Box::new(move |x| {
                let Tree::Many(trees) = x else { panic!() };
                let new_tree = trees.into_iter().next().unwrap();
                let mut this = Moo::clone(&this);

                // Only update the pointer with the new value if the value has changed. Without
                // this check, writing to the pointer might trigger a clone on write, even
                // though the value inside the pointer remained the same.
                if new_tree != tree {
                    let this = Moo::make_mut(&mut this);
                    *this = ctx(new_tree)
                }

                this
            }),
        )
    }
}

impl<To, U> Biplate<To> for Moo<U>
where
    To: Uniplate,
    U: Uniplate + Biplate<To>,
{
    fn biplate(&self) -> (Tree<To>, Box<dyn Fn(Tree<To>) -> Self>) {
        if let Some(self_as_to) = transmute_if_same_type::<Self, To>(self) {
            // To = Self -> return self
            let tree = Tree::One(self_as_to.clone());
            let ctx = Box::new(move |x| {
                let Tree::One(self_as_to) = x else { panic!() };

                let self_as_self = try_transmute_if_same_type::<To, Self>(&self_as_to);

                Moo::clone(self_as_self)
            });

            (tree, ctx)
        } else {
            // To != Self -> return children of type To

            let this = Moo::clone(self);

            // Do not need to preserve shared ownership, so treat this identically to values of the
            // inner type.
            let (tree, ctx) = try_biplate_to!((**self).clone(), To);
            (
                Tree::Many(VecDeque::from([tree.clone()])),
                Box::new(move |x| {
                    let Tree::Many(trees) = x else { panic!() };
                    let new_tree = trees.into_iter().next().unwrap();
                    let mut this = Moo::clone(&this);

                    // Only update the pointer with the new value if the value has changed. Without
                    // this check, writing to the pointer might trigger a clone on write, even
                    // though the value inside the pointer remained the same.
                    if new_tree != tree {
                        let this = Moo::make_mut(&mut this);
                        *this = ctx(new_tree)
                    }

                    this
                }),
            )
        }
    }

    /// Counts children through the pointer, without cloning the pointee.
    ///
    /// The default implementation materialises `children_bi`, which goes via [`Biplate::biplate`]
    /// and therefore clones the whole inner value. Payload syncs call this once per tree level, so
    /// the default turns a single rewrite under a wide node (e.g. `or` over a large matrix) into
    /// O(n) work, and a full rewrite pass into O(n^2).
    fn children_bi_count(&self) -> usize {
        if std::any::TypeId::of::<Self>() == std::any::TypeId::of::<To>() {
            // Biplate<T> for T treats the value as its only child.
            return 1;
        }
        <U as Biplate<To>>::children_bi_count(&**self)
    }

    /// Replaces a child through the pointer, cloning the pointee only when it is shared.
    ///
    /// Mirrors [`Biplate::children_bi_count`] above: the default implementation clones the inner
    /// value twice (once to list children, once to rebuild) and structurally compares the old and
    /// new trees.
    fn try_replace_child_at_bi(&mut self, index: usize, child: To) -> bool {
        if std::any::TypeId::of::<Self>() == std::any::TypeId::of::<To>() {
            if index != 0 {
                return false;
            }
            // SAFETY: TypeId equality means `Self` and `To` are the same type.
            unsafe {
                let child_as_self = std::mem::transmute_copy::<To, Self>(&child);
                std::mem::forget(child);
                *self = child_as_self;
            }
            return true;
        }
        if std::any::TypeId::of::<U>() == std::any::TypeId::of::<To>() {
            if index != 0 {
                return false;
            }
            // The child is the whole pointee, not a child inside it. Replace the pointer
            // directly: make_mut would clone the old shared value only to discard it.
            // SAFETY: TypeId equality means U and To are the same type.
            let value = unsafe { std::mem::transmute_copy::<To, U>(&child) };
            std::mem::forget(child);
            *self = Moo::new(value);
            return true;
        }
        <U as Biplate<To>>::try_replace_child_at_bi(Moo::make_mut(self), index, child)
    }
}

impl<'de, T: Deserialize<'de>> Deserialize<'de> for Moo<T> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        Ok(Moo::new(T::deserialize(deserializer)?))
    }
}

impl<T: Serialize> Serialize for Moo<T> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        T::serialize(&**self, serializer)
    }
}

impl<T: Display> Display for Moo<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        (**self).fmt(f)
    }
}

impl<T: Hash> Hash for Moo<T> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        (**self).hash(state);
    }
}

impl<T> From<T> for Moo<T> {
    fn from(value: T) -> Self {
        Moo::new(value)
    }
}

impl<A, B> FuncMap<A, B> for Moo<A>
where
    A: Clone,
{
    type Output = Moo<B>;

    fn func_map<F>(self, mut f: F) -> Self::Output
    where
        F: FnMut(A) -> B,
    {
        let inner = Moo::unwrap_or_clone(self);
        Moo::new(f(inner))
    }
}

impl<A, B> TryFuncMap<A, B> for Moo<A>
where
    A: Clone,
{
    type Output = Moo<B>;

    fn try_func_map<E, F>(self, mut f: F) -> Result<Self::Output, E>
    where
        F: FnMut(A) -> Result<B, E>,
    {
        let inner = Moo::unwrap_or_clone(self);
        let res = f(inner)?;
        Ok(Moo::new(res))
    }
}

impl<T, As> SerializeAs<Moo<T>> for Moo<As>
where
    As: SerializeAs<T>,
{
    fn serialize_as<S>(source: &Moo<T>, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let wrap = SerializeAsWrap::<T, As>::new(&**source);
        wrap.serialize(serializer)
    }
}

impl<'de, T, As> DeserializeAs<'de, Moo<T>> for Moo<As>
where
    As: DeserializeAs<'de, T>,
{
    fn deserialize_as<D>(deserializer: D) -> Result<Moo<T>, D::Error>
    where
        D: Deserializer<'de>,
    {
        let wrap = DeserializeAsWrap::<T, As>::deserialize(deserializer)?;
        Ok(Moo::new(wrap.into_inner()))
    }
}

#[cfg(test)]
mod replacement_tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};

    static CLONES: AtomicUsize = AtomicUsize::new(0);

    #[derive(Debug, PartialEq, Eq)]
    struct CloneProbe(i32);

    impl Clone for CloneProbe {
        fn clone(&self) -> Self {
            CLONES.fetch_add(1, Ordering::Relaxed);
            Self(self.0)
        }
    }

    impl Uniplate for CloneProbe {
        fn uniplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
            let value = self.0;
            (Tree::Zero, Box::new(move |_| Self(value)))
        }
    }

    impl Biplate<CloneProbe> for CloneProbe {
        fn biplate(&self) -> (Tree<Self>, Box<dyn Fn(Tree<Self>) -> Self>) {
            (
                Tree::One(self.clone()),
                Box::new(|tree| {
                    let Tree::One(value) = tree else {
                        panic!("expected one child")
                    };
                    value
                }),
            )
        }
    }

    #[test]
    fn replacing_shared_pointee_does_not_clone_the_discarded_value() {
        let mut value = Moo::new(CloneProbe(1));
        let shared = value.clone();
        CLONES.store(0, Ordering::Relaxed);
        assert!(!value.try_replace_child_at_bi(1, CloneProbe(2)));
        assert_eq!(value.0, 1);
        assert!(value.try_replace_child_at_bi(0, CloneProbe(2)));
        assert_eq!(value.0, 2);
        assert_eq!(shared.0, 1);
        assert_eq!(CLONES.load(Ordering::Relaxed), 0);
    }
}