simplify shift optimizer and make it safer by not optimizing out >> << combos

5 files changed, 56 insertions, 93 deletions

diff --git a/tools/js-optimizer.js b/tools/js-optimizer.js
index 48139e82..7b489b33 100644
--- a/tools/js-optimizer.js
+++ b/tools/js-optimizer.js
@@ -372,8 +372,7 @@ function optimizeShiftsInternal(ast, conservative) {
             uses: 0,
             timesShifted: [0, 0, 0, 0], // zero shifts of size 0, 1, 2, 3
             benefit: 0,
-            primaryShift: -1,
-            replaceOriginal: false
+            primaryShift: -1
           };
         }
       }
@@ -447,16 +446,8 @@ function optimizeShiftsInternal(ast, conservative) {
         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 do it
-        data.replaceOriginal = false; // Let later passes remove unneeded new variables. The only case where that might not
-                                      // work perfectly is with parameters, but inlining can fix even that
-        if (data.replaceOriginal) {
-          assert(0);
-          //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
-          if (data.defs == 1) {
-            data.benefit = totalTimesShifted - 2*(data.defs + (data.param ? 1 : 0));
-          }
+        if (data.defs == 1) {
+          data.benefit = totalTimesShifted - 2*(data.defs + (data.param ? 1 : 0));
         }
         if (conservative) data.benefit = 0;
         if (data.benefit > 0) {
@@ -485,11 +476,7 @@ function optimizeShiftsInternal(ast, conservative) {
         var data = vars[name];
         if (needsShift(name)) {
           if (data.param) {
-            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]]]]]);
-            }
+            fun[3].unshift(['var', [[name + '$s' + data.primaryShift, ['binary', '>>', ['name', name], ['num', data.primaryShift]]]]]);
           } else {
             fun[3].unshift(['var', [[name + '$s' + data.primaryShift]]]);
           }
@@ -500,22 +487,17 @@ function optimizeShiftsInternal(ast, conservative) {
         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;
+          var parent = stack[stack.length-3];
+          var parentIndex;
+          if (parent[0] == 'defun') {
+            parentIndex = 3;
+          } else if (parent[0] == 'block') {
+            parentIndex = 1;
           } 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]]]]);
+            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') {
           var args = node[1];
           for (var i = 0; i < args.length; i++) {
@@ -523,74 +505,43 @@ function optimizeShiftsInternal(ast, conservative) {
             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]]]);
-              }
+              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
+      traverse(fun, function(node, type, stack) { // replace shifted name with new variable
         stack.push(node);
-        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];
+        if (node[0] == 'binary' && node[1] == '>>' && node[2][0] == 'name' && needsShift(node[2][1]) && node[3][0] == 'num') {
+          var name = node[2][1];
           var data = vars[name];
-          node[2] = true;
-          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.
+          var parent = stack[stack.length-2];
+          // Don't modify in |x$sN = x >> 2|, in normal assigns and in var assigns
+          if (parent[0] == 'assign' && parent[2][0] == 'name' && parent[2][1] == name + '$s' + data.primaryShift) return;
+          if (parent[0] == name + '$s' + data.primaryShift) return;
+          if (node[3][1] == data.primaryShift) {
+            return ['name', name + '$s' + data.primaryShift];
           }
-          return ['binary', '<<', node, ['num', data.primaryShift]];
         }
       }, null, []);
       cleanNotes();
       var SIMPLE_SHIFTS = set('<<', '>>');
       var more = true;
-      while (more) { // collapse shifts and unshifts, completing the optimization
+      while (more) { // combine shifts in the same direction as an optimization
         more = false;
         traverse(fun, function(node, type) {
-          if (node[0] == 'binary' && node[1] in SIMPLE_SHIFTS && node[2][0] == 'binary' && node[2][1] in SIMPLE_SHIFTS &&
-              node[3][0] == 'num' && node[2][3][0] == 'num' && // do not turn a << b << c into a << b + c; while logically identical, it is slower
-              Math.abs(node[3][1]) <= MAX_SHIFTS && Math.abs(node[2][3][1]) <= MAX_SHIFTS) { // do not modify things like x << 24 >> 24 (which is reSign)
+          if (node[0] == 'binary' && node[1] in SIMPLE_SHIFTS && node[2][0] == 'binary' && node[2][1] == node[1] &&
+              node[3][0] == 'num' && node[2][3][0] == 'num') { // do not turn a << b << c into a << b + c; while logically identical, it is slower
             more = true;
-            var combinedShift = '>>';
-            var sign1 = node[1] == '>>' ? 1 : -1;
-            var sign2 = node[2][1] == '>>' ? 1 : -1;
-            var total = 0;
-            total = ['num', sign1*node[3][1] + sign2*node[2][3][1]];
-            if (total[1] < 0) {
-              combinedShift = '<<';
-              total[1] *= -1;
-            }
-            if (total[1] == 0) {
-              // XXX this may be wrong, if we removed a |0 that assumed we had this >> which |0's anyhow!
-              return node[2][2];
-            }
-            return ['binary', combinedShift, node[2][2], total];
+            return ['binary', node[1], node[2][2], ['num', node[3][1] + node[2][3][1]]];
           }
         });
       }
-      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) {
+        // Apply constant shifts onto constants
         if (type == 'binary' && node[1] == '>>' && node[2][0] == 'num' && node[3][0] == 'num' && node[3][1] <= MAX_SHIFTS) {
           var subNode = node[2];
           var shifts = node[3][1];
diff --git a/tools/test-js-optimizer-t2-output.js b/tools/test-js-optimizer-t2-output.js
index 11806014..9b22905d 100644
--- a/tools/test-js-optimizer-t2-output.js
+++ b/tools/test-js-optimizer-t2-output.js
@@ -37,17 +37,17 @@ function shifty($id2) {
   pause(1);
   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$s2 + 1]);
-  q(HEAP32[($i - 1 << 1) + $a_addr$s2 + 1]);
-  q(HEAP32[($i << 1) + $a_addr$s2]);
-  q(HEAP32[$a_addr$s2 + $i]);
+  q(HEAP32[(($sn_addr - 1 << 1) + 1 << 2 >> 2) + $a_addr$s2]);
+  q(HEAP32[(($i - 1 << 1) + 1 << 2 >> 2) + $a_addr$s2]);
+  q(HEAP32[($i << 3 >> 2) + $a_addr$s2]);
+  q(HEAP32[($i << 2 >> 2) + $a_addr$s2]);
   q($a_addr$s2, z($a_addr$s2));
   pause(2);
   var $level = HEAP[get(322) >> 2];
   var _dwt_norms_real = get("a"), _dwt_norms_real$s2 = _dwt_norms_real >> 2, $orient = get("cheez");
-  q(HEAP32[($level << 1) + _dwt_norms_real$s2 + ($orient * 20 | 0)]);
-  q(HEAP32[($level << 1) + _dwt_norms_real$s2 + ($orient * 20 | 0) + 1]);
-  q(HEAP32[($level << 1) + _dwt_norms_real$s2 + ($orient * 20 | 0) + 2]);
+  q(HEAP32[($level << 3 >> 2) + _dwt_norms_real$s2 + ($orient * 20 | 0)]);
+  q(HEAP32[(($level << 3) + 4 >> 2) + _dwt_norms_real$s2 + ($orient * 20 | 0)]);
+  q(HEAP32[(($level << 3) + 8 >> 2) + _dwt_norms_real$s2 + ($orient * 20 | 0)]);
   pause(3);
   var $wavelet38 = get(38), $wavelet38$s2 = $wavelet38 >> 2;
   $k = $a_addr;
@@ -56,12 +56,12 @@ function shifty($id2) {
   q(HEAPF32[HEAP32[$wavelet38$s2] + ($k << 4) + 400 >> 2]);
   pause(4);
   var $p = $k, $parameters_addr = get("burger"), $parameters_addr$s2 = $parameters_addr >> 2;
-  q(HEAP32[$parameters_addr$s2 + $p + 1406]);
-  q(HEAP32[$parameters_addr$s2 + $p + 1411]);
-  q(HEAP32[$parameters_addr$s2 + $p + 1416]);
+  q(HEAP32[(($p << 2) + 5624 >> 2) + $parameters_addr$s2]);
+  q(HEAP32[(($p << 2) + 5644 >> 2) + $parameters_addr$s2]);
+  q(HEAP32[(($p << 2) + 5664 >> 2) + $parameters_addr$s2]);
   pause(5);
   var $res_spec242 = get($real), $cp = get("b"), $tileno = arguments[2];
-  q(HEAP32[$parameters_addr$s2 + ($res_spec242 - 1) + 1406]);
+  q(HEAP32[(($res_spec242 - 1 << 2) + 5624 >> 2) + $parameters_addr$s2]);
   q(HEAP32[(HEAP32[$cp + 108 >> 2] + 420 >> 2) + ($tileno * 1397 | 0)]);
   pause(6);
   q($idx << 3);
@@ -78,9 +78,14 @@ function shifty($id2) {
   HEAP32[HEAP[$tp$s2]] = 5;
   HEAP32[HEAP[$tp$s2] >> 2] = 5;
   pause(7);
-  q(go());
+  q(go() >> 1 << 1);
+  q(go() << 1 >> 1);
+  q(go() >> 2);
+  q(go() << 2);
   q(go() >> 8 << 8);
+  q(go() << 8 >> 8);
+  q(go() >> 16);
+  q(go() << 16);
   q(go() + 2 >> 2);
-  pause(7.5);
 }
 // EMSCRIPTEN_GENERATED_FUNCTIONS: ["shifty"]
diff --git a/tools/test-js-optimizer-t2.js b/tools/test-js-optimizer-t2.js
index 3bc7cea5..1582dce7 100644
--- a/tools/test-js-optimizer-t2.js
+++ b/tools/test-js-optimizer-t2.js
@@ -80,8 +80,13 @@ function shifty($id2) {
   HEAP32[HEAP[$tp >> 2] >> 2] = 5;
   pause(7);
   q(go() >> 1 << 1);
+  q(go() << 1 >> 1);
+  q(go() >> 1 >> 1);
+  q(go() << 1 << 1);
   q(go() >> 8 << 8);
+  q(go() << 8 >> 8);
+  q(go() >> 8 >> 8);
+  q(go() << 8 << 8);
   q((go() + 2) >> 2); // the 2 >> 2 can't be simplified
-  pause(7.5);
 }
 // EMSCRIPTEN_GENERATED_FUNCTIONS: ["shifty"]
diff --git a/tools/test-js-optimizer-t2c-output.js b/tools/test-js-optimizer-t2c-output.js
index ee81b62f..e0f95fcc 100644
--- a/tools/test-js-optimizer-t2c-output.js
+++ b/tools/test-js-optimizer-t2c-output.js
@@ -1,6 +1,6 @@
 function shifty() {
-  $pPage = HEAP32[($26 << 16 >> 16) + ($pCur_addr + 116 >> 2)];
-  var $ead_192394 = HEAP32[($26 << 16 >> 16) + ($pCur_addr + 116 >> 2)];
+  $pPage = HEAP32[$pCur_addr + ($26 << 16 >> 16 << 2) + 116 >> 2];
+  var $ead_192394 = HEAP32[$pCur_addr + ($26 << 16 >> 16 << 2) + 116 >> 2];
   $pPage2 = HEAP32[($26 << 16 >> 16 << 2) + $pCur_addr + 116];
   var $ead_192394b = HEAP32[($26 << 16 >> 16 << 2) + $pCur_addr + 116];
   $pPage2 = HEAP32[($26 << 16 >> 16) + $pCur_addr + 116];
@@ -12,5 +12,6 @@ function shifty() {
   q($13 + $15 + 12 >> 2);
   q(HEAPF32[$output + ($j37 << 4) + 4 >> 2]);
   q($13 + 13 << 2);
+  q(h() >> 2 << 2);
 }
 // EMSCRIPTEN_GENERATED_FUNCTIONS: ["shifty"]
diff --git a/tools/test-js-optimizer-t2c.js b/tools/test-js-optimizer-t2c.js
index e820548c..52e9b4f9 100644
--- a/tools/test-js-optimizer-t2c.js
+++ b/tools/test-js-optimizer-t2c.js
@@ -13,5 +13,6 @@ function shifty() {
   q($13 + $15 + 12 >> 2);
   q(HEAPF32[$output + ($j37 << 4) + 4 >> 2]);
   q(5 + $13 + 8 << 2);
+  q(((h() | 0) >> 2) << 2); // removing the shifts is dangerous
 }
 // EMSCRIPTEN_GENERATED_FUNCTIONS: ["shifty"]