2 files changed, 24 insertions, 11 deletions

diff --git a/lib/ast.js b/lib/ast.js
index 4e0c08b1..3ede1d13 100644
--- a/lib/ast.js
+++ b/lib/ast.js
@@ -104,16 +104,11 @@ AST_Node.warn = function(txt, props) {
         AST_Node.warn_function(string_template(txt, props));
 };
 
-var AST_Debugger = DEFNODE("Debugger", null, {
-    $documentation: "Represents a debugger statement"
-});
+/* -----[ statements ]----- */
 
-var AST_Directive = DEFNODE("Directive", "value", {
-    $documentation: "Represents a directive, like \"use strict\";"
+var AST_StatementBase = DEFNODE("StatementBase", null, {
 });
 
-/* -----[ loops ]----- */
-
 var AST_Statement = DEFNODE("Statement", "body", {
     $documentation: "Base class of all statements",
     _walk: function(visitor) {
@@ -121,7 +116,15 @@ var AST_Statement = DEFNODE("Statement", "body", {
             this.body._walk(visitor);
         });
     }
-});
+}, AST_StatementBase);
+
+var AST_Debugger = DEFNODE("Debugger", null, {
+    $documentation: "Represents a debugger statement"
+}, AST_StatementBase);
+
+var AST_Directive = DEFNODE("Directive", "value", {
+    $documentation: "Represents a directive, like \"use strict\";"
+}, AST_StatementBase);
 
 var AST_SimpleStatement = DEFNODE("SimpleStatement", null, {
     $documentation: "A statement consisting of an expression, i.e. a = 1 + 2."
@@ -141,7 +144,8 @@ var AST_BlockStatement = DEFNODE("BlockStatement", null, {
 function walk_body(node, visitor) {
     if (node.body instanceof Array) node.body.forEach(function(stat){
         stat._walk(visitor);
-    }); else node.body._walk(visitor);
+    });
+    else if (node.body instanceof AST_Statement) node.body._walk(visitor);
 };
 
 var AST_Block = DEFNODE("Block", null, {
diff --git a/test/run-tests.js b/test/run-tests.js
index 158c9d35..385c43f1 100755
--- a/test/run-tests.js
+++ b/test/run-tests.js
@@ -44,6 +44,13 @@ function test_directory(dir) {
     return path.resolve(tests_dir, dir);
 }
 
+function as_toplevel(input) {
+    if (input instanceof U.AST_BlockStatement) input = input.body;
+    else if (input instanceof U.AST_StatementBase) input = [ input ];
+    else throw new Error("Unsupported input syntax");
+    return new U.AST_Toplevel({ body: input });
+}
+
 function run_compress_tests() {
     var dir = test_directory("compress");
     log_directory("compress");
@@ -53,8 +60,10 @@ function run_compress_tests() {
         function test_case(test) {
             log_test(test.name);
             var cmp = new U.Compressor(test.options || {}, true);
-            var expect = make_code(test.expect, false);
-            var output = make_code(test.input.squeeze(cmp), false);
+            var expect = make_code(as_toplevel(test.expect), false);
+            var input = as_toplevel(test.input);
+            input.figure_out_scope();
+            var output = make_code(input.squeeze(cmp), false);
             if (expect != output) {
                 log("!!! failed\n---INPUT---\n{input}\n---OUTPUT---\n{output}\n---EXPECTED---\n{expected}\n\n", {
                     input: make_code(test.input),