Implement decorator helper method

2 files changed, 16 insertions, 15 deletions

diff --git a/lib/handlebars/runtime.js b/lib/handlebars/runtime.js
index 6ee5c84..6b31a7b 100644
--- a/lib/handlebars/runtime.js
+++ b/lib/handlebars/runtime.js
@@ -29,6 +29,8 @@ export function template(templateSpec, env) {
     throw new Exception('Unknown template object: ' + typeof templateSpec);
   }
 
+  templateSpec.main.decorator = templateSpec.main_d;
+
   // Note: Using env.VM references rather than local var references throughout this section to allow
   // for external users to override these as psuedo-supported APIs.
   env.VM.checkRevision(templateSpec.compiler);
@@ -147,13 +149,7 @@ export function template(templateSpec, env) {
     function main(context/*, options*/) {
       return '' + templateSpec.main(container, context, container.helpers, container.partials, data, blockParams, depths);
     }
-
-    if (templateSpec.main_d) {
-      // Note that we are ignoring the props value here as we apply things slightly differently
-      // when applying decorators to the root function.
-      main = templateSpec.main_d(main, {}, container, undefined, data, blockParams, depths);
-    }
-
+    main = executeDecorators(templateSpec.main, main, container, options.depths || [], data, blockParams);
     return main(context, options);
   }
   ret.isTop = true;
@@ -203,11 +199,7 @@ export function wrapProgram(container, i, fn, data, declaredBlockParams, blockPa
         currentDepths);
   }
 
-  if (fn.decorator) {
-    let props = {};
-    prog = fn.decorator(prog, props, container, depths && depths[0], data, blockParams, depths);
-    Utils.extend(prog, props);
-  }
+  prog = executeDecorators(fn, prog, container, depths, data, blockParams);
 
   prog.program = i;
   prog.depth = depths ? depths.length : 0;
@@ -265,3 +257,12 @@ function initData(context, data) {
   }
   return data;
 }
+
+function executeDecorators(fn, prog, container, depths, data, blockParams) {
+  if (fn.decorator) {
+    let props = {};
+    prog = fn.decorator(prog, props, container, depths && depths[0], data, blockParams, depths);
+    Utils.extend(prog, props);
+  }
+  return prog;
+}
diff --git a/spec/runtime.js b/spec/runtime.js
index 502a843..a4830ad 100644
--- a/spec/runtime.js
+++ b/spec/runtime.js
@@ -14,19 +14,19 @@ describe('runtime', function() {
     it('should throw on version mismatch', function() {
       shouldThrow(function() {
         Handlebars.template({
-          main: true,
+          main: {},
           compiler: [Handlebars.COMPILER_REVISION + 1]
         });
       }, Error, /Template was precompiled with a newer version of Handlebars than the current runtime/);
       shouldThrow(function() {
         Handlebars.template({
-          main: true,
+          main: {},
           compiler: [Handlebars.COMPILER_REVISION - 1]
         });
       }, Error, /Template was precompiled with an older version of Handlebars than the current runtime/);
       shouldThrow(function() {
         Handlebars.template({
-          main: true
+          main: {}
         });
       }, Error, /Template was precompiled with an older version of Handlebars than the current runtime/);
     });