Merge pull request #668 from lixitrixi/refactor-treemorph

Refactor treemorph

crates/tree_morph/README.md

@@ -0,0 +1,5 @@
+# tree-morph
+
+A library to perform boilerplate-free generic tree transformations.
+
+<!-- TODO: Finish README -->

crates/tree_morph/src/commands.rs

@@ -1,54 +1,102 @@
 use std::collections::VecDeque;
 use uniplate::Uniplate;
 
-enum Command<T, M>
-where
-    T: Uniplate,
-{
-    Transform(fn(&T) -> T),
+enum Command<T: Uniplate, M> {
+    Transform(fn(T) -> T),
     MutMeta(fn(&mut M)),
 }
 
-/// A queue of commands (side effects) to be applied after every successful rule application.
-pub struct Commands<T, M>
-where
-    T: Uniplate,
-{
+/// A queue of commands (side-effects) to be applied after a successful rule application.
+///
+/// A rule is given a mutable reference to a [`Commands`] and can use it to register side-effects.
+/// These side-effects are applied in order of registration **after** the rule itself updates
+/// a part of the tree.
+///
+/// # Application
+///
+/// A rule may not be applied due to different reasons, for example:
+/// - It does not return a new subtree (i.e. it returns `None`).
+/// - It returns a new subtree but the resulting [`Update`](crate::update::Update) is not chosen
+/// by the user-defined selector function. The function may select a different rule's update or
+/// no update at all.
+/// - It is part of a lower-priority rule group and a higher-priority rule is applied first.
+///
+/// In these cases, any side-effects which are registered by the rule are not applied and are
+/// dropped by the engine.
+///
+/// # Example
+/// ```rust
+/// use tree_morph::prelude::*;
+/// use uniplate::derive::Uniplate;
+///
+/// #[derive(Debug, Clone, PartialEq, Eq, Uniplate)]
+/// #[uniplate()]
+/// enum Expr {
+///     A,
+///     B,
+///     C,
+/// }
+///
+/// fn rule(cmds: &mut Commands<Expr, bool>, subtree: &Expr, meta: &bool) -> Option<Expr> {
+///     cmds.transform(|t| match t { // A pure transformation (no other side-effects)
+///         Expr::B => Expr::C,
+///         _ => t,
+///     });
+///     cmds.mut_meta(|m| *m = true); // Set the metadata to 'true'
+///
+///     match subtree {
+///         Expr::A => Some(Expr::B),
+///         _ => None,
+///     }
+/// }
+///
+/// // Start with the expression 'A' and a metadata value of 'false'
+/// let (result, meta) = morph(vec![rule_fns![rule]], select_first, Expr::A, false);
+///
+/// // After applying the rule itself, the commands are applied in order
+/// assert_eq!(result, Expr::C);
+/// assert_eq!(meta, true);
+/// ```
+
+pub struct Commands<T: Uniplate, M> {
     commands: VecDeque<Command<T, M>>,
 }
 
-impl<T, M> Commands<T, M>
-where
-    T: Uniplate,
-{
+impl<T: Uniplate, M> Commands<T, M> {
     pub(crate) fn new() -> Self {
         Self {
             commands: VecDeque::new(),
         }
     }
 
-    /// Apply the given transformation to the root node.
-    /// Commands are applied in order after the rule is applied.
-    pub fn transform(&mut self, f: fn(&T) -> T) {
+    /// Registers a pure transformation of the whole tree.
+    ///
+    /// In this case, "pure" means that the transformation cannot register additional side-effects.
+    /// The transformation function is given ownership of the tree and should return the updated
+    /// tree.
+    ///
+    /// Side-effects are applied in order of registration after the rule is applied.
+    pub fn transform(&mut self, f: fn(T) -> T) {
         self.commands.push_back(Command::Transform(f));
     }
 
-    /// Update the associated metadata.
-    /// Commands are applied in order after the rule is applied.
+    /// Updates the global metadata in-place via a mutable reference.
+    ///
+    /// Side-effects are applied in order of registration after the rule is applied.
     pub fn mut_meta(&mut self, f: fn(&mut M)) {
         self.commands.push_back(Command::MutMeta(f));
     }
 
-    /// Remove all commands in the queue.
+    /// Removes all side-effects previously registered by the rule.
     pub fn clear(&mut self) {
         self.commands.clear();
     }
 
-    /// Consume and apply the commands currently in the queue.
+    /// Consumes and apply the side-effects currently in the queue.
     pub(crate) fn apply(&mut self, mut tree: T, mut meta: M) -> (T, M) {
         while let Some(cmd) = self.commands.pop_front() {
             match cmd {
-                Command::Transform(f) => tree = f(&tree),
+                Command::Transform(f) => tree = f(tree),
                 Command::MutMeta(f) => f(&mut meta),
             }
         }

crates/tree_morph/src/engine.rs

@@ -0,0 +1,161 @@
+use crate::{helpers::one_or_select, Commands, Rule, Update};
+use uniplate::Uniplate;
+
+/// Exhaustively rewrites a tree using a set of transformation rules.
+///
+/// Rewriting is complete when all rules have been attempted with no change. Rules may be organised
+/// into groups to control the order in which they are attempted.
+///
+/// # Rule Groups
+/// If all rules are treated equally, those which apply higher in the tree will take precedence
+/// because of the left-most outer-most traversal order of the engine.
+///
+/// This can cause problems if a rule which should ideally be applied early (e.g. evaluating
+/// constant expressions) is left until later.
+///
+/// To solve this, rules can be organised into different collections in the `groups` argument.
+/// The engine will apply rules in an earlier group to the entire tree before trying later groups.
+/// That is, no rule is attempted if a rule in an earlier group is applicable to any part of the
+/// tree.
+///
+/// # Selector Functions
+///
+/// If multiple rules in the same group are applicable to an expression, the user-defined
+/// selector function is used to choose one. This function is given an iterator over pairs of
+/// rules and the engine-created [`Update`] values which contain their modifications to the tree.
+///
+/// Some useful selector functions are available in the [`helpers`](crate::helpers) module. One
+/// common function is [`select_first`](crate::helpers::select_first), which simply returns the
+/// first applicable rule.
+///
+/// # Example
+/// ```rust
+/// use tree_morph::{prelude::*, helpers::select_panic};
+/// use uniplate::derive::Uniplate;
+///
+/// // A simple language of multiplied and squared constant expressions
+/// #[derive(Debug, Clone, PartialEq, Eq, Uniplate)]
+/// #[uniplate()]
+/// enum Expr {
+///     Val(i32),
+///     Mul(Box<Expr>, Box<Expr>),
+///     Sqr(Box<Expr>),
+/// }
+///
+/// // a * b ~> (value of) a * b, where 'a' and 'b' are literal values
+/// fn rule_eval_mul(cmds: &mut Commands<Expr, i32>, subtree: &Expr, meta: &i32) -> Option<Expr> {
+///     cmds.mut_meta(|m| *m += 1);
+///
+///     if let Expr::Mul(a, b) = subtree {
+///         if let (Expr::Val(a_v), Expr::Val(b_v)) = (a.as_ref(), b.as_ref()) {
+///             return Some(Expr::Val(a_v * b_v));
+///         }
+///     }
+///     None
+/// }
+///
+/// // e ^ 2 ~> e * e, where e is an expression
+/// // If this rule is applied before the sub-expression is fully evaluated, duplicate work
+/// // will be done on the resulting two identical sub-expressions.
+/// fn rule_expand_sqr(cmds: &mut Commands<Expr, i32>, subtree: &Expr, meta: &i32) -> Option<Expr> {
+///     cmds.mut_meta(|m| *m += 1);
+///
+///     if let Expr::Sqr(expr) = subtree {
+///         return Some(Expr::Mul(
+///             Box::new(*expr.clone()),
+///             Box::new(*expr.clone())
+///         ));
+///     }
+///     None
+/// }
+///
+/// // (1 * 2) ^ 2
+/// let expr = Expr::Sqr(
+///     Box::new(Expr::Mul(
+///         Box::new(Expr::Val(1)),
+///         Box::new(Expr::Val(2))
+///     ))
+/// );
+///
+/// // Try with both rules in the same group, keeping track of the number of rule applications
+/// let (result, num_applications) = morph(
+///     vec![rule_fns![rule_eval_mul, rule_expand_sqr]],
+///     select_panic,
+///     expr.clone(),
+///     0
+/// );
+/// assert_eq!(result, Expr::Val(4));
+/// assert_eq!(num_applications, 4); // The `Sqr` is expanded first, causing duplicate work
+///
+/// // Move the evaluation rule to an earlier group
+/// let (result, num_applications) = morph(
+///     vec![rule_fns![rule_eval_mul], rule_fns![rule_expand_sqr]],
+///     select_panic,
+///     expr.clone(),
+///     0
+/// );
+/// assert_eq!(result, Expr::Val(4));
+/// assert_eq!(num_applications, 3); // Now the sub-expression (1 * 2) is evaluated first
+/// ```
+pub fn morph<T, M, R>(
+    groups: Vec<Vec<R>>,
+    select: impl Fn(&T, &mut dyn Iterator<Item = (&R, Update<T, M>)>) -> Option<Update<T, M>>,
+    tree: T,
+    meta: M,
+) -> (T, M)
+where
+    T: Uniplate,
+    R: Rule<T, M>,
+{
+    let transforms: Vec<_> = groups
+        .iter()
+        .map(|group| {
+            |subtree: &T, meta: &M| {
+                let applicable = &mut group.iter().filter_map(|rule| {
+                    let mut commands = Commands::new();
+                    let new_tree = rule.apply(&mut commands, &subtree, &meta)?;
+                    Some((
+                        rule,
+                        Update {
+                            new_subtree: new_tree,
+                            commands,
+                        },
+                    ))
+                });
+                one_or_select(&select, subtree, applicable)
+            }
+        })
+        .collect();
+    morph_impl(transforms, tree, meta)
+}
+
+/// This implements the core rewriting logic for the engine.
+///
+/// Iterate over rule groups and apply each one to the tree. If any changes are
+/// made, restart with the first rule group.
+fn morph_impl<T: Uniplate, M>(
+    transforms: Vec<impl Fn(&T, &M) -> Option<Update<T, M>>>,
+    mut tree: T,
+    mut meta: M,
+) -> (T, M) {
+    let mut new_tree = tree;
+
+    'main: loop {
+        tree = new_tree;
+        for transform in transforms.iter() {
+            // Try each transform on the entire tree before moving to the next
+            for (node, ctx) in tree.contexts() {
+                if let Some(mut update) = transform(&node, &meta) {
+                    let whole_tree = ctx(update.new_subtree);
+                    (new_tree, meta) = update.commands.apply(whole_tree, meta);
+
+                    // Restart with the first transform every time a change is made
+                    continue 'main;
+                }
+            }
+        }
+        // All transforms were attempted without change
+        break;
+    }
+    (tree, meta)
+}