feat: Add TraversableFold and Folder

README.md

@@ -18,7 +18,7 @@
 
 ## Overview
 
-This crate provides traits and proc macros to implement the visitor pattern for arbitrary data structures. This pattern is particularly useful when dealing with complex nested data structures, abstract trees and hierarchies of all kinds.
+This crate provides traits and proc macros to implement visitor and folder patterns for arbitrary data structures. These patterns are particularly useful when dealing with complex nested data structures, abstract trees and hierarchies of all kinds.
 
 ## Quick Start
 
@@ -87,12 +87,49 @@ fn main() {
 }
 ```
 
+Use `TraversableFold` for owned bottom-up rewrites:
+
+```rust
+use std::ops::ControlFlow;
+
+use traversable::TraversableFold;
+use traversable::function::folder_leave;
+
+#[derive(TraversableFold)]
+enum Expr {
+    Add(Box<Expr>, Box<Expr>),
+    Literal(i32),
+}
+
+fn simplify(expr: Expr) -> Expr {
+    match expr {
+        Expr::Add(left, right) => match (*left, *right) {
+            (Expr::Literal(0), expr) | (expr, Expr::Literal(0)) => expr,
+            (left, right) => Expr::Add(Box::new(left), Box::new(right)),
+        },
+        expr => expr,
+    }
+}
+
+fn main() {
+    let expr = Expr::Add(Box::new(Expr::Literal(0)), Box::new(Expr::Literal(1)));
+    let mut folder = folder_leave::<Expr, (), _>(|expr| ControlFlow::Continue(simplify(expr)));
+
+    let expr = match expr.traverse_fold(&mut folder) {
+        ControlFlow::Continue(expr) => expr,
+        ControlFlow::Break(()) => unreachable!(),
+    };
+
+    assert!(matches!(expr, Expr::Literal(1)));
+}
+```
+
 ## Attributes
 
 The derive macro supports the following attributes on structs and enums:
 
-*   `#[traverse(skip_self)]`: Skips calling the visitor for the annotated type while still traversing its children.
-*   `#[traverse(skip_children)]`: Calls the visitor for the annotated type without traversing its children.
+*   `#[traverse(skip_self)]`: Skips calling the visitor or folder for the annotated type while still traversing its children.
+*   `#[traverse(skip_children)]`: Calls the visitor or folder for the annotated type without traversing its children.
 
 The derive macro supports the following attributes on fields and variants:
 

traversable-derive/Cargo.toml

@@ -16,7 +16,7 @@
 name = "traversable-derive"
 version = "0.2.0"
 
-description = "Procedural macro to derive Traversable and TraversableMut"
+description = "Procedural macro to derive Traversable, TraversableMut, and TraversableFold"
 documentation = "https://docs.rs/traversable-derive"
 keywords = ["visitor", "traverse", "traversable"]
 readme = "README.md"

traversable-derive/README.md

@@ -20,4 +20,5 @@ This is an implementation crate for the [`traversable`](https://crates.io/crates
 
 Please refer to the main [`traversable`](https://crates.io/crates/traversable) crate for documentation and usage examples.
 
-This crate contains procedural macros that derive `Traversable` and `TraversableMut` implementations.
+This crate contains procedural macros that derive `Traversable`, `TraversableMut`, and
+`TraversableFold` implementations.

traversable-derive/src/lib.rs

@@ -55,6 +55,11 @@ pub fn derive_traversable_mut(input: proc_macro::TokenStream) -> proc_macro::Tok
     expand_with(input, |stream| impl_traversable(stream, true))
 }
 
+#[proc_macro_derive(TraversableFold, attributes(traverse))]
+pub fn derive_traversable_fold(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    expand_with(input, impl_traversable_fold)
+}
+
 fn expand_with(
     input: proc_macro::TokenStream,
     handler: impl Fn(DeriveInput) -> Result<TokenStream>,
@@ -203,20 +208,16 @@ fn resolve_crate_name() -> Path {
     parse_quote!(::traversable)
 }
 
+fn take_unit_param(params: &mut Params, name: &str) -> Result<bool> {
+    Ok(params.param(name)?.map(Param::unit).transpose()?.is_some())
+}
+
 fn impl_traversable(input: DeriveInput, mutable: bool) -> Result<TokenStream> {
     let mut params = Params::from_attrs(input.attrs, "traverse")?;
     params.validate(&["skip_self", "skip_children"])?;
 
-    let skip_visit_self = params
-        .param("skip_self")?
-        .map(Param::unit)
-        .transpose()?
-        .is_some();
-    let skip_children = params
-        .param("skip_children")?
-        .map(Param::unit)
-        .transpose()?
-        .is_some();
+    let skip_visit_self = take_unit_param(&mut params, "skip_self")?;
+    let skip_children = take_unit_param(&mut params, "skip_children")?;
 
     let name = input.ident;
     let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
@@ -348,7 +349,7 @@ fn traverse_enum(e: DataEnum, mutable: bool) -> Result<TokenStream> {
 fn traverse_variant(v: Variant, mutable: bool) -> Result<TokenStream> {
     let mut params = Params::from_attrs(v.attrs, "traverse")?;
     params.validate(&["skip"])?;
-    if params.param("skip")?.map(Param::unit).is_some() {
+    if take_unit_param(&mut params, "skip")? {
         return Ok(TokenStream::new());
     }
     let name = v.ident;
@@ -385,7 +386,7 @@ fn destructure_fields(fields: Fields) -> Result<TokenStream> {
                 .map(|field| {
                     let mut params = Params::from_attrs(field.attrs, "traverse")?;
                     let field_name = field.ident.unwrap();
-                    Ok(if params.param("skip")?.map(Param::unit).is_some() {
+                    Ok(if take_unit_param(&mut params, "skip")? {
                         quote! { #field_name: _ }
                     } else {
                         field_name.into_token_stream()
@@ -403,7 +404,7 @@ fn destructure_fields(fields: Fields) -> Result<TokenStream> {
                 .enumerate()
                 .map(|(index, field)| {
                     let mut params = Params::from_attrs(field.attrs, "traverse")?;
-                    Ok(if params.param("skip")?.map(Param::unit).is_some() {
+                    Ok(if take_unit_param(&mut params, "skip")? {
                         quote! { _ }
                     } else {
                         Ident::new(&format!("i{index}",), Span::call_site()).into_token_stream()
@@ -422,7 +423,7 @@ fn traverse_field(value: &TokenStream, field: Field, mutable: bool) -> Result<To
     let mut params = Params::from_attrs(field.attrs, "traverse")?;
     params.validate(&["skip", "with"])?;
 
-    if params.param("skip")?.map(Param::unit).is_some() {
+    if take_unit_param(&mut params, "skip")? {
         return Ok(TokenStream::new());
     }
 
@@ -442,3 +443,203 @@ fn traverse_field(value: &TokenStream, field: Field, mutable: bool) -> Result<To
         }
     }
 }
+
+fn impl_traversable_fold(input: DeriveInput) -> Result<TokenStream> {
+    let mut params = Params::from_attrs(input.attrs, "traverse")?;
+    params.validate(&["skip_self", "skip_children"])?;
+
+    let skip_visit_self = take_unit_param(&mut params, "skip_self")?;
+    let skip_children = take_unit_param(&mut params, "skip_children")?;
+
+    let name = input.ident;
+    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
+    let crate_name = resolve_crate_name();
+
+    let enter_self = if skip_visit_self {
+        quote! {
+            let this = self;
+        }
+    } else {
+        quote! {
+            let this = #crate_name::Folder::enter(folder, self)?;
+        }
+    };
+
+    let fold_children = match input.data {
+        Data::Struct(struct_) => {
+            if skip_children {
+                Ok(TokenStream::new())
+            } else {
+                fold_struct(struct_)
+            }
+        }
+        Data::Enum(enum_) => {
+            if skip_children {
+                Ok(TokenStream::new())
+            } else {
+                fold_enum(enum_)
+            }
+        }
+        Data::Union(union_) => {
+            return Err(Error::new_spanned(
+                union_.union_token,
+                "unions are not supported",
+            ));
+        }
+    }?;
+
+    let leave_self = if skip_visit_self {
+        TokenStream::new()
+    } else {
+        quote! {
+            let this = #crate_name::Folder::leave(folder, this)?;
+        }
+    };
+
+    Ok(quote! {
+        impl #impl_generics #crate_name::TraversableFold for #name #ty_generics #where_clause {
+            fn traverse_fold<V: #crate_name::Folder>(
+                self,
+                folder: &mut V
+            ) -> ::core::ops::ControlFlow<V::Break, Self> {
+                #enter_self
+                #fold_children
+                #leave_self
+                ::core::ops::ControlFlow::Continue(this)
+            }
+        }
+    })
+}
+
+fn fold_struct(s: DataStruct) -> Result<TokenStream> {
+    Ok(match s.fields {
+        Fields::Named(fields) => {
+            let mut field_names = Vec::new();
+            let mut fold_fields = Vec::new();
+
+            for field in fields.named {
+                let field_name = field.ident.clone().unwrap();
+                fold_fields.push(fold_field(&field_name.to_token_stream(), field)?);
+                field_names.push(field_name);
+            }
+
+            quote! {
+                let this = match this {
+                    Self { #( #field_names ),* } => {
+                        #( #fold_fields )*
+                        Self { #( #field_names ),* }
+                    }
+                };
+            }
+        }
+        Fields::Unnamed(fields) => {
+            let mut field_names = Vec::new();
+            let mut fold_fields = Vec::new();
+
+            for (index, field) in fields.unnamed.into_iter().enumerate() {
+                let field_name = Ident::new(&format!("i{index}"), Span::call_site());
+                fold_fields.push(fold_field(&field_name.to_token_stream(), field)?);
+                field_names.push(field_name);
+            }
+
+            quote! {
+                let this = match this {
+                    Self( #( #field_names ),* ) => {
+                        #( #fold_fields )*
+                        Self( #( #field_names ),* )
+                    }
+                };
+            }
+        }
+        Fields::Unit => TokenStream::new(),
+    })
+}
+
+fn fold_enum(e: DataEnum) -> Result<TokenStream> {
+    let variants = e
+        .variants
+        .into_iter()
+        .map(fold_variant)
+        .collect::<Result<TokenStream>>()?;
+    Ok(quote! {
+        let this = match this {
+            #variants
+        };
+    })
+}
+
+fn fold_variant(v: Variant) -> Result<TokenStream> {
+    let mut params = Params::from_attrs(v.attrs, "traverse")?;
+    params.validate(&["skip"])?;
+    let skip = take_unit_param(&mut params, "skip")?;
+
+    let name = v.ident;
+    Ok(match v.fields {
+        Fields::Named(fields) => {
+            let mut field_names = Vec::new();
+            let mut fold_fields = Vec::new();
+
+            for field in fields.named {
+                let field_name = field.ident.clone().unwrap();
+                if !skip {
+                    fold_fields.push(fold_field(&field_name.to_token_stream(), field)?);
+                }
+                field_names.push(field_name);
+            }
+
+            quote! {
+                Self::#name { #( #field_names ),* } => {
+                    #( #fold_fields )*
+                    Self::#name { #( #field_names ),* }
+                }
+            }
+        }
+        Fields::Unnamed(fields) => {
+            let mut field_names = Vec::new();
+            let mut fold_fields = Vec::new();
+
+            for (index, field) in fields.unnamed.into_iter().enumerate() {
+                let field_name = Ident::new(&format!("i{index}"), Span::call_site());
+                if !skip {
+                    fold_fields.push(fold_field(&field_name.to_token_stream(), field)?);
+                }
+                field_names.push(field_name);
+            }
+
+            quote! {
+                Self::#name( #( #field_names ),* ) => {
+                    #( #fold_fields )*
+                    Self::#name( #( #field_names ),* )
+                }
+            }
+        }
+        Fields::Unit => {
+            quote! {
+                Self::#name => Self::#name
+            }
+        }
+    })
+}
+
+fn fold_field(value: &TokenStream, field: Field) -> Result<TokenStream> {
+    let mut params = Params::from_attrs(field.attrs, "traverse")?;
+    params.validate(&["skip", "with"])?;
+
+    if take_unit_param(&mut params, "skip")? {
+        return Ok(TokenStream::new());
+    }
+
+    let crate_name = resolve_crate_name();
+
+    match params.param("with")? {
+        None => Ok(quote! {
+            let #value = #crate_name::TraversableFold::traverse_fold(#value, folder)?;
+        }),
+        Some(traverse_fn) => {
+            let traverse_fn = traverse_fn.string_literal()?.parse::<Path>()?;
+            Ok(quote! {
+                let #value = #traverse_fn(#value, folder)?;
+            })
+        }
+    }
+}