simplify nested ifs with identical elses

4 files changed, 111 insertions, 16 deletions

diff --git a/tests/test_other.py b/tests/test_other.py
index 263bc376..4a62384c 100644
--- a/tests/test_other.py
+++ b/tests/test_other.py
@@ -2578,23 +2578,44 @@ int main()
     assert 'incorrect target triple' in err, err
 
   def test_simplify_ifs(self):
-    open('src.c', 'w').write(r'''
+    def test(src, nums):
+      open('src.c', 'w').write(src)
+      for opts, ifs in [
+        [['-g2'], nums[0]],
+        [['-profiling'], nums[1]],
+        [['-profiling', '-g2'], nums[2]]
+      ]:
+        print opts, ifs
+        Popen([PYTHON, EMCC, 'src.c', '-O2'] + opts, stdout=PIPE, stderr=PIPE).communicate()
+        src = open('a.out.js').read()
+        main = src[src.find('function _main'):src.find('\n}', src.find('function _main'))]
+        actual_ifs = main.count('if (')
+        assert ifs == actual_ifs, main + ' : ' + str([ifs, actual_ifs])
+
+    test(r'''
       #include <stdio.h>
       #include <string.h>
       int main(int argc, char **argv) {
         if (argc > 5 && strlen(argv[0]) > 1 && strlen(argv[1]) > 2) printf("halp");
         return 0;
       }
-    ''')
-    for opts, ifs in [
-      [['-g2'], 3],
-      [['-profiling'], 1],
-      [['-profiling', '-g2'], 1]
-    ]:
-      print opts, ifs
-      Popen([PYTHON, EMCC, 'src.c', '-O2'] + opts, stdout=PIPE, stderr=PIPE).communicate()
-      src = open('a.out.js').read()
-      main = src[src.find('function _main'):src.find('\n}', src.find('function _main'))]
-      actual_ifs = main.count('if (')
-      assert ifs == actual_ifs, main
+    ''', [3, 1, 1])
+
+    test(r'''
+      #include <stdio.h>
+      #include <string.h>
+      int main(int argc, char **argv) {
+        while (argc % 3 == 0) {
+          if (argc > 5 && strlen(argv[0]) > 1 && strlen(argv[1]) > 2) {
+            printf("halp");
+            argc++;
+          } else {
+            while (argc > 0) {
+              printf("%d\n", argc--);
+            }
+          }
+        }
+        return 0;
+      }
+    ''', [8, 6, 6])
 
diff --git a/tools/js-optimizer.js b/tools/js-optimizer.js
index bb6acf73..4b63ab9b 100644
--- a/tools/js-optimizer.js
+++ b/tools/js-optimizer.js
@@ -303,6 +303,16 @@ function removeAllEmptySubNodes(ast) {
   traverse(ast, removeEmptySubNodes);
 }
 
+function astCompare(x, y) {
+  if (!Array.isArray(x)) return x === y;
+  if (!Array.isArray(y)) return false;
+  if (x.length !== y.length) return false;
+  for (var i = 0; i < x.length; i++) {
+    if (!astCompare(x[i], y[i])) return false;
+  }
+  return true;
+}
+
 // Passes
 
 // Dump the AST. Useful for debugging. For example,
@@ -830,14 +840,14 @@ function simplifyExpressions(ast) {
 function simplifyIfs(ast) {
   traverse(ast, function(node, type) {
     // simplify   if (x) { if (y) { .. } }   to   if (x ? y : 0) { .. }
-    if (type === 'if' && !node[3]) {
+    if (type === 'if') {
       var body = node[2];
       // recurse to handle chains
       while (body[0] === 'block') {
         var stats = body[1];
         if (stats.length > 1) {
           var last = stats[stats.length-1];
-          if (last[0] === 'if' && !last[3]) {
+          if (last[0] === 'if') {
             // try to commaify - turn everything between the ifs into a comma operator inside the second if
             var ok = true;
             for (var i = 0; i < stats.length-1; i++) {
@@ -857,7 +867,9 @@ function simplifyIfs(ast) {
         }
         if (stats.length !== 1) break;
         var singleton = stats[0];
-        if (singleton[0] === 'if' && !singleton[3]) {
+        // we can handle elses, but must be fully identical
+        if (!astCompare(node[3], singleton[3])) break;
+        if (singleton[0] === 'if') {
           node[1] = ['conditional', node[1], singleton[1], ['num', 0]];
           body = node[2] = singleton[2];
         } else {
diff --git a/tools/test-js-optimizer-si-output.js b/tools/test-js-optimizer-si-output.js
index bd36ef10..8b93803d 100644
--- a/tools/test-js-optimizer-si-output.js
+++ b/tools/test-js-optimizer-si-output.js
@@ -66,5 +66,34 @@ function a() {
    g();
   }
  }
+ andNowForElses();
+ if (x ? y : 0) {
+  f();
+ } else {
+  label = 5;
+ }
+ if (x) {
+  if (y) {
+   f();
+  } else {
+   label = 5;
+  }
+ } else {
+  label = 6;
+ }
+ if (x) {
+  if (y) {
+   f();
+  } else {
+   label = 5;
+  }
+ }
+ if (x) {
+  if (y) {
+   f();
+  }
+ } else {
+  label = 5;
+ }
 }
 
diff --git a/tools/test-js-optimizer-si.js b/tools/test-js-optimizer-si.js
index 5f6b334f..6952b183 100644
--- a/tools/test-js-optimizer-si.js
+++ b/tools/test-js-optimizer-si.js
@@ -84,5 +84,38 @@ function a() {
       g();
     }
   }
+  andNowForElses();
+  if (x) {
+    if (y) {
+      f();
+    } else {
+      label = 5;
+    }
+  } else {
+    label = 5;
+  }
+  if (x) {
+    if (y) {
+      f();
+    } else {
+      label = 5;
+    }
+  } else {
+    label = 6;
+  }
+  if (x) {
+    if (y) {
+      f();
+    } else {
+      label = 5;
+    }
+  }
+  if (x) {
+    if (y) {
+      f();
+    }
+  } else {
+    label = 5;
+  }
 }