registerize function parameters too

3 files changed, 61 insertions, 31 deletions

diff --git a/tools/js-optimizer.js b/tools/js-optimizer.js
index 7ceade86..193556f9 100644
--- a/tools/js-optimizer.js
+++ b/tools/js-optimizer.js
@@ -1344,6 +1344,19 @@ function denormalizeAsm(func, data) {
 function registerize(ast, asm) {
   traverseGeneratedFunctions(ast, function(fun) {
     if (asm) var asmData = normalizeAsm(fun);
+    // Add parameters as a first (fake) var (with assignment), so they get taken into consideration
+    var params = {};
+    if (fun[2] && fun[2].length) {
+      var assign = ['num', 0];
+      fun[3].unshift(['var', fun[2].map(function(param) {
+        params[param] = 1;
+        return [param, assign];
+      })]);
+    }
+    if (asm) {
+      // copy params into vars
+      for (var p in asmData.params) asmData.vars[p] = asmData.params[p];
+    }
     // Replace all var definitions with assignments; we will add var definitions at the top after we registerize
     // We also mark local variables - i.e., having a var definition
     var localVars = {};
@@ -1496,23 +1509,41 @@ function registerize(ast, asm) {
         loops--;
       }
     });
-    // Add vars at the beginning
+    var paramRegs = {}; // true if the register is used by a parameter (and so needs no def at start of function)
+    if (fun[2] && fun[2].length) {
+      for (var p in params) {
+        paramRegs[varRegs[p]] = 1;
+      }
+      fun[2].length = 0; // clear params, we will fill with registers
+      fun[3].shift(); // remove fake initial var
+    }
     if (!asm) {
       if (nextReg > 1) {
         var vars = [];
         for (var i = 1; i < nextReg; i++) {
-          vars.push([fullNames[i]]);
+          var reg = fullNames[i];
+          if (!paramRegs[i]) {
+            vars.push([reg]);
+          } else {
+            fun[2].push(reg);
+          }
         }
         getStatements(fun).unshift(['var', vars]);
       }
     } else {
       var finalAsmData = {
-        params: asmData.params,
+        params: {},
         vars: {},
       };
       for (var i = 1; i < nextReg; i++) {
         var reg = fullNames[i];
-        finalAsmData.vars[reg] = reg[0] == 'i' ? ASM_INT : ASM_DOUBLE;
+        var type = reg[0] == 'i' ? ASM_INT : ASM_DOUBLE
+        if (!paramRegs[i]) {
+          finalAsmData.vars[reg] = type;
+        } else {
+          finalAsmData.params[reg] = type;
+          fun[2].push(reg);
+        }
       }
       denormalizeAsm(fun, finalAsmData);
     }
@@ -2090,7 +2121,6 @@ if (metadata) setGeneratedFunctions(metadata);
 arguments_.slice(1).forEach(function(arg) {
   passes[arg](ast);
 });
-
 var js = astToSrc(ast, compress), old;
 
 // remove unneeded newlines+spaces, and print
diff --git a/tools/test-js-optimizer-asm-regs-output.js b/tools/test-js-optimizer-asm-regs-output.js
index 9dcb4f70..00f032f7 100644
--- a/tools/test-js-optimizer-asm-regs-output.js
+++ b/tools/test-js-optimizer-asm-regs-output.js
@@ -1,12 +1,12 @@
-function asm(x, y) {
-  x = +x;
-  y = y | 0;
-  var i1 = 0, d2 = +0, i3 = 0, d4 = +0;
-  i1 = x + x | 0;
-  d2 = d(Math_max(10, Math_min(5, f())));
-  i3 = i1 + 2 | 0;
+function asm(d1, i2) {
+  d1 = +d1;
+  i2 = i2 | 0;
+  var i3 = 0, d4 = +0;
+  i2 = d1 + d1 | 0;
+  d1 = d(Math_max(10, Math_min(5, f())));
+  i3 = i2 + 2 | 0;
   print(i3);
-  d4 = d2 * 5;
+  d4 = d1 * 5;
   return d4;
 }
 
diff --git a/tools/test-js-optimizer-regs-output.js b/tools/test-js-optimizer-regs-output.js
index 90b67a47..36006b7c 100644
--- a/tools/test-js-optimizer-regs-output.js
+++ b/tools/test-js-optimizer-regs-output.js
@@ -197,32 +197,32 @@ function ex() {
     }))();
   }
 }
-function switchey(x) {
-  var r1, r2, r3, r4, r5, r6, r7, r8;
-  r1 = 5;
+function switchey(r1) {
+  var r2, r3, r4, r5, r6, r7, r8, r9;
+  r2 = 5;
   while (1) {
-    switch (x = f(x, r1)) {
+    switch (r1 = f(r1, r2)) {
      case 1:
-      g(r1);
-      r2 = x + 1;
-      x--;
+      g(r2);
+      r3 = r1 + 1;
+      r1--;
       break;
      case 2:
-      g(r1 * 2);
-      r3 = x + 22;
-      r4 = r3 + 5;
-      x -= 20;
+      g(r2 * 2);
+      r4 = r1 + 22;
+      r5 = r4 + 5;
+      r1 -= 20;
       break;
      default:
-      r5 = x + 22;
-      r6 = r3 + 5;
-      ch(r5, r6 * r3);
+      r6 = r1 + 22;
+      r7 = r4 + 5;
+      ch(r6, r7 * r4);
       throw 99;
     }
   }
-  r7 = x + 1;
-  p(r1, r7);
-  r8 = x + 2;
-  pp(r8);
+  r8 = r1 + 1;
+  p(r2, r8);
+  r9 = r1 + 2;
+  pp(r9);
 }