let shiftOptimizer either replace the original variable, or keep it and add a new shifted variable

3 files changed, 118 insertions, 52 deletions

diff --git a/tools/js-optimizer.js b/tools/js-optimizer.js
index 9eaf1916..1764f4b7 100644
--- a/tools/js-optimizer.js
+++ b/tools/js-optimizer.js
@@ -358,6 +358,7 @@ function simplifyExpressionsPre(ast) {
 function optimizeShifts(ast) {
   traverseGeneratedFunctions(ast, function(fun) {
     var funMore = true;
+    var funFinished = {};
     while (funMore) {
       funMore = false;
       // Recognize variables and parameters
@@ -370,7 +371,8 @@ function optimizeShifts(ast) {
             uses: 0,
             timesShifted: [0, 0, 0, 0], // zero shifts of size 0, 1, 2, 3
             benefit: 0,
-            primaryShift: -1
+            primaryShift: -1,
+            replaceOriginal: false
           };
         }
       }
@@ -441,12 +443,16 @@ function optimizeShifts(ast) {
           // TODO: Handle multiple different shifts
           continue;
         }
+        if (funFinished[name]) continue;
         // We have one shift size (and possible unshifted uses). Consider replacing this variable with a shifted clone. If
-        // the estimated benefit is >0, we will replace
-        // Each time we are shifted, we can save 1 shift. Each addition nonshifted use is a cost of an additional shift.
-        // Each definition is an additional shift. If we are a parameter, we have an additional cost since we add an
-        // entire statement to shift ourselves
-        data.benefit = totalTimesShifted - (data.uses-totalTimesShifted) - data.defs - (data.param ? 1 : 0);
+        // the estimated benefit is >0, we will do it
+        data.replaceOriginal = data.uses == totalTimesShifted; // if always shifted, replace the original
+        if (!data.replaceOriginal) {
+          data.benefit = totalTimesShifted - (data.uses-totalTimesShifted) - data.defs - (data.param ? 1 : 0);
+        } else {
+          // We keep the original around, so there is some additional cost for each def, but none for unshifted uses
+          data.benefit = totalTimesShifted - 2*(data.defs + (data.param ? 1 : 0));
+        }
         if (data.benefit > 0) {
           funMore = true; // We will reprocess this function
           for (var i = 0; i < 4; i++) {
@@ -466,34 +472,74 @@ function optimizeShifts(ast) {
       }
       cleanNotes();
       // Apply changes
+      function needsShift(name) {
+        return vars[name] && vars[name].primaryShift >= 0;
+      }
       for (var name in vars) { // add shifts for params
         var data = vars[name];
-        if (data.primaryShift >= 0 && data.param) {
-          fun[3].unshift(['stat', ['assign', '>>', ['name', name], ['num', data.primaryShift]]]);
+        if (data.param && needsShift(name)) {
+          if (data.replaceOriginal) {
+            fun[3].unshift(['stat', ['assign', '>>', ['name', name], ['num', data.primaryShift]]]);
+          } else {
+            fun[3].unshift(['var', [[name + '$s' + data.primaryShift, ['binary', '>>', ['name', name], ['num', data.primaryShift]]]]]);
+          }
         }
       }
-      function needsShift(name) {
-        return vars[name] && vars[name].primaryShift >= 0;
-      }
-      traverse(fun, function(node, type) { // add shift to assignments
-        if (node[0] == 'assign' && node[1] === true && node[2][0] == 'name' && needsShift(node[2][1])) {
-          node[3] = ['binary', '>>', node[3], ['num', vars[node[2][1]].primaryShift]];
-          return node;
+      traverse(fun, function(node, type, stack) { // add shift to assignments
+        stack.push(node);
+        if (node[0] == 'assign' && node[1] === true && node[2][0] == 'name' && needsShift(node[2][1]) && !node[2][2]) {
+          var name = node[2][1];
+          var data = vars[name];
+          if (data.replaceOriginal) {
+            node[3] = ['binary', '>>', node[3], ['num', vars[name].primaryShift]];
+            return node;
+          } else {
+            var parent = stack[stack.length-3];
+            var parentIndex;
+            if (parent[0] == 'defun') {
+              parentIndex = 3;
+            } else if (parent[0] == 'block') {
+              parentIndex = 1;
+            } else {
+              throw 'Invalid parent for assign-shift: ' + dump(parent);
+            }
+            var i = parent[parentIndex].indexOf(stack[stack.length-2]);
+            parent[parentIndex].splice(i+1, 0, ['stat', ['assign', true, ['name', name + '$s' + data.primaryShift], ['binary', '>>', ['name', name, true], ['num', data.primaryShift]]]]);
+          }
         } else if (node[0] == 'var') {
-          node[1].forEach(function(arg) {
-            if (arg[1] && needsShift(arg[0])) {
-              arg[1] = ['binary', '>>', arg[1], ['num', vars[arg[0]].primaryShift]];
+          var args = node[1];
+          for (var i = 0; i < args.length; i++) {
+            var arg = args[i];
+            var name = arg[0];
+            var data = vars[name];
+            if (arg[1] && needsShift(name)) {
+              if (data.replaceOriginal) {
+                arg[1] = ['binary', '>>', arg[1], ['num', data.primaryShift]];
+              } else {
+                args.splice(i+1, 0, [name + '$s' + data.primaryShift, ['binary', '>>', ['name', name, true], ['num', data.primaryShift]]]);
+              }
             }
-          });
+          }
           return node;
         }
-      });
+      }, null, []);
+      cleanNotes();
       traverse(fun, function(node, type, stack) { // replace name with unshifted name, making everything valid again by balancing the assign changes
         stack.push(node);
-        if (node[0] == 'name' && !node[2] && vars[node[1]] && vars[node[1]].primaryShift >= 0 &&
-            (stack[stack.length-2][0] != 'assign' || stack[stack.length-2][2] != node)) { // ignore changing VAR in |VAR = something|
+        if (node[0] == 'name' && !node[2] && needsShift(node[1]) &&
+            !((stack[stack.length-2][0] == 'assign' && stack[stack.length-2][2] == node) || // ignore changing VAR in |VAR = something|
+              (stack[stack.length-2][0] == 'binary' && stack[stack.length-3][0] == 'assign' && // ignore changing name$sX = name >> X;
+               stack[stack.length-3][2][0] == 'name' && stack[stack.length-3][2][1] == node[1] + '$s' + vars[node[1]].primaryShift) ||
+              (stack[stack.length-2][0] == 'binary' && // ignore changing name$sX = name >> X;
+               stack[stack.length-3][0] == node[1] + '$s' + vars[node[1]].primaryShift) )) {
+          var name = node[1];
+          var data = vars[name];
           node[2] = true;
-          return ['binary', '<<', node, ['num', vars[node[1]].primaryShift]];
+          if (!data.replaceOriginal) {
+            node[1] += '$s' + data.primaryShift; // When this node is shifted, the collapse will get to the optimal case. But when it is
+                                                 // unshifted, we just added an (un)shift. We remove that later.
+          }
+          return ['binary', '<<', node, ['num', data.primaryShift]];
         }
       }, null, []);
       cleanNotes();
@@ -523,6 +569,16 @@ function optimizeShifts(ast) {
           }
         });
       }
+      traverse(fun, function(node, type) { // collapse name$sX << X into name
+        if (type == 'binary' && node[1] == '<<' && node[2][0] == 'name' && node[3][0] == 'num') {
+          var shifts = node[3][1];
+          var full = node[2][1]
+          var raw = full.substr(0, full.length-3);
+          if (raw + '$s' + shifts == full) {
+            return ['name', raw];
+          }
+        }
+      });
       // Before recombining, do some additional optimizations
       traverse(fun, function(node, type) {
         if (type == 'binary' && node[1] == '>>' && node[2][0] == 'num' && node[3][0] == 'num') {
@@ -606,6 +662,10 @@ function optimizeShifts(ast) {
           return ret;
         }
       }, null, []);
+      // Note finished variables
+      for (var name in vars) {
+        funFinished[name] = true;
+      }
     }
   });
 }
diff --git a/tools/test-js-optimizer-t2-output.js b/tools/test-js-optimizer-t2-output.js
index 04f1d0f4..ea384e07 100644
--- a/tools/test-js-optimizer-t2-output.js
+++ b/tools/test-js-optimizer-t2-output.js
@@ -1,38 +1,42 @@
 function shifty($id) {
-  $id >>= 2;
-  q(HEAP32[$id]);
-  q(HEAP32[$id + 10]);
-  q(HEAP32[$id + 20]);
-  q(HEAP32[(unknown1 + unknown2 >> 2) + $id]);
-  q(HEAP32[(unknown1 + unknown2 >> 2) + $id]);
+  var $id$s2 = $id >> 2;
+  q(HEAP32[$id$s2]);
+  q(HEAP32[$id$s2 + 10]);
+  q(HEAP32[$id$s2 + 20]);
+  q(HEAP32[(unknown1 + unknown2 >> 2) + $id$s2]);
+  q(HEAP32[(unknown1 + unknown2 >> 2) + $id$s2]);
   var localUnchanged1 = get(1), localUnchanged2 = get(1);
-  q(HEAP32[(localUnchanged1 + localUnchanged2 >> 2) + $id]);
-  q($id << 2 >> _something_);
-  q($id << 2 << _somethingElse_);
+  q(HEAP32[(localUnchanged1 + localUnchanged2 >> 2) + $id$s2]);
+  q($id >> _something_);
+  $id = q("..");
+  $id$s2 = $id >> 2;
+  q($id << _somethingElse_);
   pause(-1);
   var $id2;
   $id2 = get(54) >> 1;
   q(HEAP32[$id2]);
   q(HEAP32[$id2 + 20]);
   q(HEAP32[$id2 + 40]);
-  var $id3 = get(74) >> 3;
-  q(HEAP32[$id3]);
-  q(HEAP32[$id3 + 5]);
-  q(HEAP32[$id3 + 10]);
+  var $id3 = get(74), $id3$s3 = $id3 >> 3;
+  q(HEAP32[$id3$s3]);
+  q(HEAP32[$id3$s3 + 5]);
+  q(HEAP32[$id3$s3 + 10]);
+  q($id3);
   pause(0);
   var _idents = get("abc") >> 2;
   q(HEAP32[(HEAP32[_idents] >> 2) + 2]);
   q(HEAP32[(HEAP32[_idents] >> 2) + 2]);
   q(HEAP32[(HEAP32[_idents] >> 2) + 2]);
   pause(1);
-  var $sn_addr = get(12), $a_addr = get(999) >> 2;
+  var $sn_addr = get(12), $a_addr = get(999), $a_addr$s2 = $a_addr >> 2;
   var $i = get(112233);
-  q(HEAP32[($sn_addr - 1 << 1) + $a_addr + 1]);
-  q(HEAP32[($i - 1 << 1) + $a_addr + 1]);
-  $a_addr = $a_addr + 1;
-  q(HEAP32[($i << 1) + $a_addr]);
-  q(HEAP32[$a_addr + $i]);
-  q($a_addr, z($a_addr));
+  q(HEAP32[($sn_addr - 1 << 1) + $a_addr$s2 + 1]);
+  q(HEAP32[($i - 1 << 1) + $a_addr$s2 + 1]);
+  $a_addr = $a_addr + 4;
+  $a_addr$s2 = $a_addr >> 2;
+  q(HEAP32[($i << 1) + $a_addr$s2]);
+  q(HEAP32[$a_addr$s2 + $i]);
+  q($a_addr$s2, z($a_addr$s2));
   pause(2);
   var $level = HEAP[get(322) >> 2];
   var _dwt_norms_real = get("a") >> 2, $orient = get("cheez");
@@ -41,7 +45,7 @@ function shifty($id) {
   q(HEAP32[($level << 1) + _dwt_norms_real + $orient * 20 + 2]);
   pause(3);
   var $wavelet38 = get(38) >> 2;
-  $k = $a_addr << 2;
+  $k = $a_addr;
   q(HEAPF32[(HEAP32[$wavelet38] >> 2) + ($k << 2) + 2]);
   q(HEAPF32[(HEAP32[$wavelet38] >> 2) + ($k << 2) + 3]);
   q(HEAPF32[(HEAP32[$wavelet38] >> 2) + ($k << 2) + 100]);
@@ -59,15 +63,15 @@ function shifty($id) {
   q(1 << $idx << 1);
   print(INDENT + "Entering: _main" + "hi");
   pause(7);
-  var $tp = get("tp") >> 2;
-  q($tp);
-  q($tp);
-  q($tp);
-  HEAP32[$H400] = $tp << 2;
-  HEAP32[$tp << 2] = 5;
+  var $tp = get("tp"), $tp$s2 = $tp >> 2;
+  q($tp$s2);
+  q($tp$s2);
+  q($tp$s2);
+  HEAP32[$H400] = $tp;
   HEAP32[$tp] = 5;
-  HEAP32[HEAP[$tp]] = 5;
-  HEAP32[HEAP[$tp] >> 2] = 5;
+  HEAP32[$tp$s2] = 5;
+  HEAP32[HEAP[$tp$s2]] = 5;
+  HEAP32[HEAP[$tp$s2] >> 2] = 5;
   pause(7);
   q(go());
   q(go() >> 8 << 8);
diff --git a/tools/test-js-optimizer-t2.js b/tools/test-js-optimizer-t2.js
index ffda2703..f7bda9ea 100644
--- a/tools/test-js-optimizer-t2.js
+++ b/tools/test-js-optimizer-t2.js
@@ -10,6 +10,7 @@ function shifty($id) {
   var localUnchanged1 = get(1), localUnchanged2 = get(1);
   q(HEAP32[(localUnchanged1 + $id + localUnchanged2) >> 2]); // unknowns should be shifted together
   q($id >> _something_); // non-fixed shift
+  $id = q('..');
   q($id << _somethingElse_); // non-fixed shift
   pause(-1);
   var $id2;
@@ -21,6 +22,7 @@ function shifty($id) {
   q(HEAP32[$id3 >> 3]);
   q(HEAP32[($id3 + 40) >> 3]);
   q(HEAP32[($id3 + 80 | 0) >> 3]);
+  q($id3);
   pause(0);
   // similar, but inside another HEAP
   var _idents = get('abc');