Doughlang/src/ast/fold.rs

//! A folder (implementer of the [`Fold`] trait) maps ASTs to ASTs
#![warn(clippy::all, clippy::pedantic)]
#![allow(clippy::wildcard_imports, clippy::missing_errors_doc)]

use std::mem::MaybeUninit;

use super::*;

/// Deconstructs an entire AST, and reconstructs it from parts.
///
/// Each function acts as a customization point.
///
/// Aside from [Annotation]s, each node in the AST implements the [`Foldable`] trait,
/// which provides double dispatch.
pub trait Fold<A: Annotation> {
    type Error;

    fn fold_annotation(&mut self, anno: A) -> Result<A, Self::Error> {
        Ok(anno)
    }

    /// Consumes an [`Expr`], possibly transforms it, and produces a replacement [`Expr`]
    fn fold_expr(&mut self, expr: Expr<A>) -> Result<Expr<A>, Self::Error> {
        expr.children(self)
    }

    /// Consumes a Symbol, possibly transforms it, and produces a replacement Symbol
    fn fold_ident(&mut self, name: String) -> Result<String, Self::Error> {
        name.children(self) // TODO: ^ this should be a symbol
    }

    /// Consumes a [`Path`], possibly transforms it, and produces a replacement [`Path`]
    fn fold_path(&mut self, path: Path) -> Result<Path, Self::Error> {
        path.children(self)
    }

    /// Consumes a [`Literal`], possibly transforms it, and produces a replacement [`Literal`]
    fn fold_literal(&mut self, lit: Literal) -> Result<Literal, Self::Error> {
        lit.children(self)
    }

    /// Consumes a [`Use`], possibly transforms it, and produces a replacement [`Use`]
    fn fold_use(&mut self, item: Use) -> Result<Use, Self::Error> {
        item.children(self)
    }

    /// Consumes a [`Pat`], possibly transforms it, and produces a replacement [`Pat`]
    fn fold_pat(&mut self, pat: Pat<A>) -> Result<Pat<A>, Self::Error> {
        pat.children(self)
    }

    /// Consumes a [`Bind`], possibly transforms it, and produces a replacement [`Bind`]
    fn fold_bind(&mut self, bind: Bind<A>) -> Result<Bind<A>, Self::Error> {
        bind.children(self)
    }

    /// Consumes a [`Make`], possibly transforms it, and produces a replacement [`Make`]
    fn fold_make(&mut self, make: Make<A>) -> Result<Make<A>, Self::Error> {
        make.children(self)
    }

    /// Consumes a [`MakeArm`], possibly transforms it, and produces a replacement [`MakeArm`]
    fn fold_makearm(&mut self, arm: MakeArm<A>) -> Result<MakeArm<A>, Self::Error> {
        arm.children(self)
    }
}

pub trait Foldable<A: Annotation>: Sized {
    /// Calls `Self`'s appropriate [Folder] function(s)
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error>;

    /// Destructures `self`, calling [`Foldable::fold_in`] on all foldable members,
    /// and rebuilds a `Self` out of the results.
    #[allow(unused_variables)]
    fn children<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        Ok(self)
    }
}

impl<A: Annotation> Foldable<A> for Expr<A> {
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        folder.fold_expr(self)
    }

    fn children<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        Ok(match self {
            Self::Omitted => Self::Omitted,
            Self::Id(path) => Self::Id(path.fold_in(folder)?),
            Self::MetId(id) => Self::MetId(id.fold_in(folder)?),
            Self::Lit(literal) => Self::Lit(literal.fold_in(folder)?),
            Self::Use(item) => Self::Use(item.fold_in(folder)?),
            Self::Bind(bind) => Self::Bind(bind.fold_in(folder)?),
            Self::Make(make) => Self::Make(make.fold_in(folder)?),
            Self::Op(op, annos) => Self::Op(op, annos.fold_in(folder)?),
        })
    }
}

impl<A: Annotation> Foldable<A> for String {
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        folder.fold_ident(self)
    }
}

impl<A: Annotation> Foldable<A> for Path {
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        folder.fold_path(self)
    }

    fn children<F: Fold<A> + ?Sized>(self, _folder: &mut F) -> Result<Self, F::Error> {
        Ok(self)
    }
}

impl<A: Annotation> Foldable<A> for Literal {
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        folder.fold_literal(self)
    }

    fn children<F: Fold<A> + ?Sized>(self, _folder: &mut F) -> Result<Self, F::Error> {
        Ok(self)
    }
}

impl<A: Annotation> Foldable<A> for Use {
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        folder.fold_use(self)
    }

    fn children<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        Ok(match self {
            Self::Glob => Self::Glob,
            Self::Name(name) => Self::Name(name),
            Self::Alias(name, alias) => Self::Alias(name, alias),
            Self::Path(name, rest) => Self::Path(name, rest.fold_in(folder)?),
            Self::Tree(items) => Self::Tree(items.fold_in(folder)?),
        })
    }
}

impl<A: Annotation> Foldable<A> for Pat<A> {
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        folder.fold_pat(self)
    }

    fn children<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        Ok(match self {
            Self::Ignore => Self::Ignore,
            Self::Never => Self::Never,
            Self::MetId(name) => Self::MetId(name.fold_in(folder)?),
            Self::Name(name) => Self::Name(name.fold_in(folder)?),
            Self::Value(expr) => Self::Value(expr.fold_in(folder)?),
            Self::Op(op, pats) => Self::Op(op, pats.fold_in(folder)?),
        })
    }
}

impl<A: Annotation> Foldable<A> for Bind<A> {
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        folder.fold_bind(self)
    }

    fn children<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        let Self(op, gens, pat, exprs) = self;
        Ok(Self(
            op,
            gens.fold_in(folder)?,
            pat.fold_in(folder)?,
            exprs.fold_in(folder)?,
        ))
    }
}

impl<A: Annotation> Foldable<A> for Make<A> {
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        folder.fold_make(self)
    }

    fn children<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        let Self(expr, arms) = self;
        Ok(Self(expr.fold_in(folder)?, arms.fold_in(folder)?))
    }
}

impl<A: Annotation> Foldable<A> for MakeArm<A> {
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        folder.fold_makearm(self)
    }

    fn children<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        let Self(name, expr) = self;
        Ok(Self(name, expr.fold_in(folder)?))
    }
}

impl<T: Annotation + Foldable<A>, A: Annotation> Foldable<A> for Anno<T, A> {
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        let Self(expr, anno) = self;
        let anno = folder.fold_annotation(anno)?;
        Ok(Self(expr.fold_in(folder)?, anno))
    }

    fn children<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        let Self(expr, anno) = self;
        Ok(Self(expr.children(folder)?, anno))
    }
}

//////////////////////////////////////////////
//  GENERIC IMPLEMENTATIONS ON COLLECTIONS  //
//////////////////////////////////////////////

/// Maps the value in the box across `f()` without deallocating
fn box_try_map<T, E>(boxed: Box<T>, f: impl FnOnce(T) -> Result<T, E>) -> Result<Box<T>, E> {
    // TODO: replace with Box::take when it stabilizes.
    let rawbox = Box::into_raw(boxed);

    // Safety: `rawbox` came from a Box, so it is aligned and initialized.
    //     To prevent further reuse and deallocate on failure, rawbox is
    //     shadowed by a Box<MaybeUninit<T>>.
    // Safety: MaybeUninit<T> has the same size and alignment as T.
    let (value, rawbox) = (unsafe { rawbox.read() }, unsafe {
        Box::from_raw(rawbox.cast::<MaybeUninit<T>>())
    });

    // rawbox is reinitialized with f(value)
    Ok(Box::write(rawbox, f(value)?))
}

impl<T: Foldable<A>, A: Annotation> Foldable<A> for Box<T> {
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        box_try_map(self, |t| t.fold_in(folder))
    }

    fn children<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        box_try_map(self, |t| t.children(folder))
    }
}

impl<T: Foldable<A>, A: Annotation> Foldable<A> for Vec<T> {
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        self.into_iter().map(|e| e.fold_in(folder)).collect()
    }

    fn children<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        // TODO: is this correct for generic data structures?
        self.into_iter().map(|e| e.fold_in(folder)).collect()
    }
}

impl<T: Foldable<A>, A: Annotation> Foldable<A> for Option<T> {
    fn fold_in<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        Ok(match self {
            Self::Some(value) => Some(value.fold_in(folder)?),
            Self::None => None,
        })
    }

    fn children<F: Fold<A> + ?Sized>(self, folder: &mut F) -> Result<Self, F::Error> {
        Ok(match self {
            Self::Some(value) => Some(value.children(folder)?),
            Self::None => None,
        })
    }
}