Add parent tracking and mutation to AST visitors

Fixes #916

1 files changed, 78 insertions, 35 deletions

diff --git a/lib/handlebars/compiler/visitor.js b/lib/handlebars/compiler/visitor.js
index c2480ad..e98d0bf 100644
--- a/lib/handlebars/compiler/visitor.js
+++ b/lib/handlebars/compiler/visitor.js
@@ -1,66 +1,109 @@
-/*jshint unused: false */
-function Visitor() {}
+import Exception from "../exception";
+import AST from "./AST";
+
+function Visitor() {
+  this.parents = [];
+}
 
 Visitor.prototype = {
   constructor: Visitor,
+  mutating: false,
 
-  accept: function(object) {
-    return object && this[object.type](object);
+  // Visits a given value. If mutating, will replace the value if necessary.
+  acceptKey: function(node, name) {
+    var value = this.accept(node[name]);
+    if (this.mutating) {
+      // Hacky sanity check:
+      if (value && (!value.type || !AST[value.type])) {
+        throw new Exception('Unexpected node type "' + value.type + '" found when accepting ' + name + ' on ' + node.type);
+      }
+      node[name] = value;
+    }
   },
 
-  Program: function(program) {
-    var body = program.body,
-        i, l;
+  // Performs an accept operation with added sanity check to ensure
+  // required keys are not removed.
+  acceptRequired: function(node, name) {
+    this.acceptKey(node, name);
 
-    for(i=0, l=body.length; i<l; i++) {
-      this.accept(body[i]);
+    if (!node[name]) {
+      throw new Exception(node.type + ' requires ' + name);
     }
   },
 
+  // Traverses a given array. If mutating, empty respnses will be removed
+  // for child elements.
+  acceptArray: function(array) {
+    for (var i = 0, l = array.length; i < l; i++) {
+      this.acceptKey(array, i);
+
+      if (!array[i]) {
+        array.splice(i, 1);
+        i--;
+        l--;
+      }
+    }
+  },
+
+  accept: function(object) {
+    if (!object) {
+      return;
+    }
+
+    if (this.current) {
+      this.parents.unshift(this.current);
+    }
+    this.current = object;
+
+    var ret = this[object.type](object);
+
+    this.current = this.parents.shift();
+
+    if (!this.mutating || ret) {
+      return ret;
+    } else if (ret !== false) {
+      return object;
+    }
+  },
+
+  Program: function(program) {
+    this.acceptArray(program.body);
+  },
+
   MustacheStatement: function(mustache) {
-    this.accept(mustache.sexpr);
+    this.acceptRequired(mustache, 'sexpr');
   },
 
   BlockStatement: function(block) {
-    this.accept(block.sexpr);
-    this.accept(block.program);
-    this.accept(block.inverse);
+    this.acceptRequired(block, 'sexpr');
+    this.acceptKey(block, 'program');
+    this.acceptKey(block, 'inverse');
   },
 
   PartialStatement: function(partial) {
-    this.accept(partial.partialName);
-    this.accept(partial.context);
-    this.accept(partial.hash);
+    this.acceptRequired(partial, 'sexpr');
   },
 
-  ContentStatement: function(content) {},
-  CommentStatement: function(comment) {},
+  ContentStatement: function(/* content */) {},
+  CommentStatement: function(/* comment */) {},
 
   SubExpression: function(sexpr) {
-    var params = sexpr.params, paramStrings = [], hash;
-
-    this.accept(sexpr.path);
-    for(var i=0, l=params.length; i<l; i++) {
-      this.accept(params[i]);
-    }
-    this.accept(sexpr.hash);
+    this.acceptRequired(sexpr, 'path');
+    this.acceptArray(sexpr.params);
+    this.acceptKey(sexpr, 'hash');
   },
 
-  PathExpression: function(path) {},
+  PathExpression: function(/* path */) {},
 
-  StringLiteral: function(string) {},
-  NumberLiteral: function(number) {},
-  BooleanLiteral: function(bool) {},
+  StringLiteral: function(/* string */) {},
+  NumberLiteral: function(/* number */) {},
+  BooleanLiteral: function(/* bool */) {},
 
   Hash: function(hash) {
-    var pairs = hash.pairs;
-
-    for(var i=0, l=pairs.length; i<l; i++) {
-      this.accept(pairs[i]);
-    }
+    this.acceptArray(hash.pairs);
   },
   HashPair: function(pair) {
-    this.accept(pair.value);
+    this.acceptRequired(pair, 'value');
   }
 };