How to "annotate" a recursive AST?

[jmackie]

Hi there 👋

This is really awesome, thanks so much for sharing. I really enjoyed reading the blog posts!

I hope you don't mind me reaching out for help via the issues, but maybe this is useful for someone else too.

I'm trying to get my head around how to use this. But I'm stuck on how I might reconstruct a recursive structure from the "top down", modifying each layer of the structure with some contextual information.

For example, say I adapted your Expr example to include some annotation type T:

pub enum Expr<T, A> {
    Add(T, A, A),
    Sub(T, A, A),
    Mul(T, A, A),
    LiteralInt(T, i64),
}

How would I then implement:

type Depth = usize;
fn tag_depth<T, A>(e: Expr<T, A>) -> Expr<Depth, A> {
    // ¯\_(ツ)_/¯
}

Any help would be much appreciated! 🙏

[inanna-malick]

Thank you! I'm really glad you liked it.

I would use the new recursion-schemes crate (uses a more recent rustc GAT feature to implement cleaner Functors) to do something like this. The general idea is that each layer is made up of the composition of an annotation and a layer of the expression

mod example {
    use recursion_schemes::functor::*;
    use recursion_schemes::recursive::*;

    #[derive(Debug, Clone, Copy)]
    pub enum Expr<A> {
        Add(A, A),
        Sub(A, A),
        Mul(A, A),
        LiteralInt(i64),
    }

    impl Functor for Expr<PartiallyApplied> {
        type Layer<X> = Expr<X>;

        #[inline(always)]
        fn fmap<F, A, B>(input: Self::Layer<A>, mut f: F) -> Self::Layer<B>
        where
            F: FnMut(A) -> B,
        {
            match input {
                Expr::Add(a, b) => Expr::Add(f(a), f(b)),
                Expr::Sub(a, b) => Expr::Sub(f(a), f(b)),
                Expr::Mul(a, b) => Expr::Mul(f(a), f(b)),
                Expr::LiteralInt(x) => Expr::LiteralInt(x),
            }
        }
    }

    type AnnotationFunctor<T> = (T, PartiallyApplied);
    type ExprFunctor = Expr<PartiallyApplied>;

    type Depth = usize;

    type AnnotatedExprFunctor = Compose<AnnotationFunctor<Depth>, ExprFunctor>;

    fn tag_depth(e: Fix<ExprFunctor>) -> Fix<AnnotatedExprFunctor> {
        e.fold_recursive(|layer| {
            let mut max_depth = 0;
            let layer = ExprFunctor::fmap(
                layer,
                |Fix(annotated_subnode): Fix<AnnotatedExprFunctor>| {
                    let subnode_depth = (*annotated_subnode).0;
                    max_depth = max_depth.max(subnode_depth);
                    Fix::<AnnotatedExprFunctor>(annotated_subnode)
                },
            );
            Fix(Box::new((max_depth + 1, layer)))
        })
    }

    #[test]
    fn test_example() {
        let x = Fix::new(Expr::Add(
            Fix::new(Expr::LiteralInt(1)),
            Fix::new(Expr::Mul(
                Fix::new(Expr::LiteralInt(2)),
                Fix::new(Expr::LiteralInt(2)),
            )),
        ));

        let annotated = tag_depth(x);

        assert_eq!(annotated.as_ref().0, 3);
    }
}

[jmackie]

Ahh neat, thanks very much for the detailed answer 🙏

Possible n00b question/observation, but would I be right in thinking using Fix loses some of the performance wins of the old recursion crate? Based purely on the use of Box, which the internet assures me is bad...