Merge pull request probmods#297 from null-a/freevars-fix

freevars bug fix

src/transforms/freevars.js

@@ -16,28 +16,26 @@ var makeGensym = require('../util').makeGensym;
 var genid = null;
 var boundVarsStack = null;
 var freeVarsStack = null;
-var nodeStack = null;
 
 var literalIdentifiers = {
   undefined: true,
   NaN: true,
   Infinity: true
 };
 
-function identifierIsVar(node) {
+function identifierIsVar(node, parent) {
   // esprima represents some special literals as Identifer nodes; skip those
   if (literalIdentifiers[node.name]) return false;
   // exprima also represents non-computed object member access with an
   //    Identifer node, so skip those as well.
-  var ntop = nodeStack[nodeStack.length - 1];
-  if (ntop.type === Syntax.MemberExpression && !ntop.computed &&
-      node === ntop.property) return false;
+  if (parent.type === Syntax.MemberExpression && !parent.computed &&
+      node === parent.property) return false;
   // Property keys in object literal expressions are also Identifer nodes
-  if (ntop.type === Syntax.Property && node === ntop.key) return false;
+  if (parent.type === Syntax.Property && node === parent.key) return false;
   return true;
 }
 
-function enter(node) {
+function enter(node, parent) {
   switch (node.type) {
     case Syntax.FunctionExpression:
       // Bind the formal parameters of the function
@@ -54,15 +52,14 @@ function enter(node) {
         boundVarsStack[boundVarsStack.length - 1][node.id.name] = true;
       break;
     case Syntax.Identifier:
-      if (boundVarsStack.length > 0 && identifierIsVar(node)) {
+      if (boundVarsStack.length > 0 && identifierIsVar(node, parent)) {
         // If the Identifier isn't already bound, then it's a free var
         if (!boundVarsStack[boundVarsStack.length - 1][node.name])
           freeVarsStack[freeVarsStack.length - 1][node.name] = true;
       }
       break;
     default:
   }
-  nodeStack.push(node);
 }
 
 function exit(node) {
@@ -96,14 +93,12 @@ function exit(node) {
       return wrappedFn;
     default:
   }
-  nodeStack.pop();
 }
 
 function freevarsMain(node) {
   genid = makeGensym();
   boundVarsStack = [];
   freeVarsStack = [];
-  nodeStack = [];
 
   return replace(node, { enter: enter, leave: exit });
 }

tests/test-transforms.js

@@ -12,6 +12,7 @@ var store = require('../src/transforms/store').store;
 var optimize = require('../src/transforms/optimize').optimize;
 var trampoline = require('../src/transforms/trampoline').trampoline;
 var varargs = require('../src/transforms/varargs').varargs;
+var freevars = require('../src/transforms/freevars').freevars;
 
 var fooObj = {
   bar: 1,
@@ -92,6 +93,35 @@ var runVarargs = runOptimize;
 var transformAstTrampoline = compose(trampoline, transformAstVarargs);
 var runTrampoline = runVarargs;
 
+var transformAstFreevars = compose(freevars, function(node) {
+  // By thunkifying we ensure that freevars is exercised (by
+  // identifying the free variables of the thunk) even when the test
+  // code doesn't contain any functions.
+  return thunkify(node, fail('transform', node));
+});
+function runFreevars(test, code, newCode, expected) {
+  check(test, code, newCode, expected, eval(newCode)());
+}
+
+var selectFreevarsPrimitives = function() {
+  // Set global definitions
+  plus = function(x, y) {
+    return (x + y);
+  };
+  minus = function(x, y) {
+    return (x - y);
+  };
+  times = function(x, y) {
+    return (x * y);
+  };
+  and = function(x, y) {
+    return (x && y);
+  };
+  plusTwo = function(x, y) {
+    return (x + 2);
+  };
+};
+
 var selectNamingPrimitives = function() {
   // Set global definitions
   plus = function(a, x, y) {
@@ -183,6 +213,11 @@ function runTrampolineTest(test, code, expected) {
   return runTest(test, code, expected, transformAstTrampoline, runTrampoline);
 }
 
+function runFreevarsTest(test, code, expected) {
+  selectFreevarsPrimitives();
+  return runTest(test, code, expected, transformAstFreevars, runFreevars);
+}
+
 function generateTestFunctions(allTests, testRunner) {
   var exports = {};
   for (var testClassName in allTests) {
@@ -514,7 +549,11 @@ var tests = {
 
     { name: 'testMember3',
       code: 'var a = [1,2]; a[1]',
-      expected: 2 }
+      expected: 2 },
+
+    { name: 'testMember4',
+      code: '(function() { return fooObj; })().bar',
+      expected: 1 }
 
   ],
 
@@ -635,3 +674,4 @@ exports.testOptimize = generateTestFunctions(tests, runOptimizeTest);
 exports.testTrampoline = generateTestFunctions(tests, runTrampolineTest);
 exports.testVarargs = generateTestFunctions(tests, runVarargsTest);
 exports.testTrampoline = generateTestFunctions(tests, runTrampolineTest);
+exports.testFreevars = generateTestFunctions(tests, runFreevarsTest);