properly lock variables to loops in registerize

2 files changed, 43 insertions, 23 deletions

diff --git a/tools/js-optimizer.js b/tools/js-optimizer.js
index ee067a19..e76f58f2 100644
--- a/tools/js-optimizer.js
+++ b/tools/js-optimizer.js
@@ -1220,17 +1220,21 @@ function registerize(ast) {
     }
     // Go through the function's code, assigning 'registers'.
     // The only tricky bit is to keep variables locked on a register through loops,
-    // since they can potentially be returned to. We use a simple approach of
-    // locking the register to the topmost loop.
+    // since they can potentially be returned to. Optimizable variables lock onto
+    // loops that they enter, unoptimizable variables lock in a conservative way
+    // into the topmost loop.
     var varRegs = {}; // maps variables to the register they will use all their life
     var freeRegs = [];
     var nextReg = 1;
     var fullNames = {};
-    var loopRegs = [];
-    var loops = 0;
+    var loopRegs = { 1: [] }; // for each loop nesting level, the list of bound variables
+    var loops = 0; // 0 is toplevel, 1 is first loop, etc
     var saved = 0;
+    var activeOptimizables = {};
+    var optimizableLoops = {};
     function decUse(name) {
       if (!varUses[name]) return false; // no uses left, or not a relevant variable
+      if (optimizables[name]) activeOptimizables[name] = 1;
       var reg = varRegs[name];
       if (!reg) {
         // acquire register
@@ -1246,11 +1250,20 @@ function registerize(ast) {
       varUses[name]--;
       assert(varUses[name] >= 0);
       if (varUses[name] == 0) {
-        if (loops == 0 || optimizables[name]) {
+        if (optimizables[name]) delete activeOptimizables[name];
+        // If we are not in a loop, or we are optimizable and not bound to a loop
+        // (we might have been in one but left it), we can free the register now.
+        if (loops == 0 || (optimizables[name] && !optimizableLoops[name])) {
           // free register
           freeRegs.push(reg);
         } else {
-          loopRegs.push(reg);
+          // when the relevant loop is exited, we will free the register
+          if (optimizables[name]) {
+            if (!loopRegs[loops]) loopRegs[loops] = [];
+            loopRegs[loops].push(reg);
+          } else {
+            loopRegs[1].push(reg);
+          }
         }
       }
       return true;
@@ -1263,14 +1276,21 @@ function registerize(ast) {
         }
       } else if (type in LOOP) {
         loops++;
+        // Active optimizables lock onto this loop, if not locked onto one that encloses this one
+        for (var name in activeOptimizables) {
+          if (!optimizableLoops[name]) {
+            optimizableLoops[name] = loops;
+          }
+        }
       }
     }, function(node, type) {
       if (type in LOOP) {
-        loops--;
-        if (loops == 0 && loopRegs.length > 0) {
-          freeRegs = freeRegs.concat(loopRegs);
-          loopRegs = [];
+        // Free registers that were locked to this loop
+        if (loopRegs[loops]) {
+          freeRegs = freeRegs.concat(loopRegs[loops]);
+          loopRegs[loops] = [];
         }
+        loops--;
       }
     });
     // Add vars at the beginning
diff --git a/tools/test-js-optimizer-regs-output.js b/tools/test-js-optimizer-regs-output.js
index 3a0db393..be8af33b 100644
--- a/tools/test-js-optimizer-regs-output.js
+++ b/tools/test-js-optimizer-regs-output.js
@@ -25,7 +25,7 @@ function test() {
   f(r1());
 }
 function primes() {
-  var r1, r2, r3, r4, r5, r6;
+  var r1, r2, r3, r4, r5, r6, r7;
   r1 = 2;
   r2 = 0;
   $_$2 : while (1) {
@@ -39,26 +39,26 @@ function primes() {
         r6 = 1;
         break $_$4;
       }
-      r4 = (r1 | 0) % (r3 | 0);
-      r5 = (r4 | 0) == 0;
-      if (r5) {
+      r5 = (r1 | 0) % (r3 | 0);
+      r7 = (r5 | 0) == 0;
+      if (r7) {
         r6 = 0;
         break $_$4;
       }
-      r5 = r3 + 1 | 0;
-      r3 = r5;
+      r7 = r3 + 1 | 0;
+      r3 = r7;
     }
-    r5 = r6 + r2 | 0;
-    r3 = r1 + 1 | 0;
-    r4 = (r5 | 0) < 1e5;
-    if (r4) {
-      r1 = r3;
-      r2 = r5;
+    r3 = r6 + r2 | 0;
+    r4 = r1 + 1 | 0;
+    r7 = (r3 | 0) < 1e5;
+    if (r7) {
+      r1 = r4;
+      r2 = r3;
     } else {
       break $_$2;
     }
   }
-  r6 = _printf(STRING_TABLE.__str | 0, (tempInt = STACKTOP, STACKTOP += 4, HEAP32[tempInt >> 2] = r1, tempInt));
+  r2 = _printf(STRING_TABLE.__str | 0, (tempInt = STACKTOP, STACKTOP += 4, HEAP32[tempInt >> 2] = r1, tempInt));
   return 1;
   return null;
 }