1 files changed, 157 insertions, 0 deletions
diff --git a/src/optimizer/restorer.rs b/src/optimizer/restorer.rs
new file mode 100644
index 0000000..34a710e
--- /dev/null
+++ b/src/optimizer/restorer.rs
@@ -0,0 +1,157 @@
+// pest. The Elegant Parser
+// Copyright (c) 2018 Dragoș Tiselice
+//
+// Licensed under the Apache License, Version 2.0
+// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
+// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. All files in the project carrying such notice may not be copied,
+// modified, or distributed except according to those terms.
+use std::collections::HashMap;
+
+use optimizer::*;
+
+pub fn restore_on_err(
+    rule: OptimizedRule,
+    rules: &HashMap<String, OptimizedExpr>,
+) -> OptimizedRule {
+    match rule {
+        OptimizedRule { name, ty, expr } => {
+            let expr = expr.map_bottom_up(|expr| wrap_branching_exprs(expr, rules));
+
+            OptimizedRule { name, ty, expr }
+        }
+    }
+}
+
+fn wrap_branching_exprs(
+    expr: OptimizedExpr,
+    rules: &HashMap<String, OptimizedExpr>,
+) -> OptimizedExpr {
+    match expr {
+        OptimizedExpr::Opt(expr) => {
+            if child_modifies_state(&expr, rules, &mut HashMap::new()) {
+                OptimizedExpr::Opt(Box::new(OptimizedExpr::RestoreOnErr(expr)))
+            } else {
+                OptimizedExpr::Opt(expr)
+            }
+        }
+        OptimizedExpr::Choice(lhs, rhs) => {
+            let wrapped_lhs = if child_modifies_state(&lhs, rules, &mut HashMap::new()) {
+                Box::new(OptimizedExpr::RestoreOnErr(lhs))
+            } else {
+                lhs
+            };
+            let wrapped_rhs = if child_modifies_state(&rhs, rules, &mut HashMap::new()) {
+                Box::new(OptimizedExpr::RestoreOnErr(rhs))
+            } else {
+                rhs
+            };
+            OptimizedExpr::Choice(wrapped_lhs, wrapped_rhs)
+        }
+        OptimizedExpr::Rep(expr) => {
+            if child_modifies_state(&expr, rules, &mut HashMap::new()) {
+                OptimizedExpr::Rep(Box::new(OptimizedExpr::RestoreOnErr(expr)))
+            } else {
+                OptimizedExpr::Rep(expr)
+            }
+        }
+        _ => expr,
+    }
+}
+
+fn child_modifies_state(
+    expr: &OptimizedExpr,
+    rules: &HashMap<String, OptimizedExpr>,
+    cache: &mut HashMap<String, Option<bool>>,
+) -> bool {
+    expr.iter_top_down().any(|expr| match expr {
+        OptimizedExpr::Push(_) => true,
+        OptimizedExpr::Ident(ref name) if name == "DROP" => true,
+        OptimizedExpr::Ident(ref name) if name == "POP" => true,
+        OptimizedExpr::Ident(ref name) => match cache.get(name).cloned() {
+            Some(option) => match option {
+                Some(cached) => cached,
+                None => {
+                    cache.insert(name.to_owned(), Some(false));
+                    false
+                }
+            },
+            None => {
+                cache.insert(name.to_owned(), None);
+
+                let result = match rules.get(name) {
+                    Some(expr) => child_modifies_state(expr, rules, cache),
+                    None => false,
+                };
+
+                cache.insert(name.to_owned(), Some(result));
+
+                result
+            }
+        },
+        _ => false,
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use optimizer::OptimizedExpr::*;
+
+    #[test]
+    fn restore_no_stack_children() {
+        let rules = vec![OptimizedRule {
+            name: "rule".to_owned(),
+            ty: RuleType::Normal,
+            expr: box_tree!(Opt(Str("a".to_string()))),
+        }];
+
+        assert_eq!(
+            restore_on_err(rules[0].clone(), &to_hash_map(&rules)),
+            rules[0].clone()
+        );
+    }
+
+    #[test]
+    fn restore_with_child_stack_ops() {
+        let rules = vec![OptimizedRule {
+            name: "rule".to_owned(),
+            ty: RuleType::Normal,
+            expr: box_tree!(Rep(Push(Str("a".to_string())))),
+        }];
+
+        let restored = OptimizedRule {
+            name: "rule".to_owned(),
+            ty: RuleType::Normal,
+            expr: box_tree!(Rep(RestoreOnErr(Push(Str("a".to_string()))))),
+        };
+
+        assert_eq!(
+            restore_on_err(rules[0].clone(), &to_hash_map(&rules)),
+            restored
+        );
+    }
+
+    #[test]
+    fn restore_choice_branch_with_and_branch_without() {
+        let rules = vec![OptimizedRule {
+            name: "rule".to_owned(),
+            ty: RuleType::Normal,
+            expr: box_tree!(Choice(Push(Str("a".to_string())), Str("a".to_string()))),
+        }];
+
+        let restored = OptimizedRule {
+            name: "rule".to_owned(),
+            ty: RuleType::Normal,
+            expr: box_tree!(Choice(
+                RestoreOnErr(Push(Str("a".to_string()))),
+                Str("a".to_string())
+            )),
+        };
+
+        assert_eq!(
+            restore_on_err(rules[0].clone(), &to_hash_map(&rules)),
+            restored
+        );
+    }
+}