DeNaN improvements (#2888)

Instead of instrumenting every local.get, instrument parameters
on arrival at a function once on entry. After that, every local will
always contain a de-naned value (since we would denan on a
local.set). This is more efficient and also less confusing I think.

Also avoid doing anything to values that fall through as they
have already been fixed up.

4 files changed, 212 insertions, 7 deletions

diff --git a/src/ir/properties.h b/src/ir/properties.h
index 094d90bd6..0c6824e4a 100644
--- a/src/ir/properties.h
+++ b/src/ir/properties.h
@@ -200,7 +200,9 @@ inline Index getZeroExtBits(Expression* curr) {
 }
 
 // Returns a falling-through value, that is, it looks through a local.tee
-// and other operations that receive a value and let it flow through them.
+// and other operations that receive a value and let it flow through them. If
+// there is no value falling through, returns the node itself (as that is the
+// value that trivially falls through, with 0 steps in the middle).
 inline Expression* getFallthrough(Expression* curr,
                                   const PassOptions& passOptions,
                                   FeatureSet features) {
@@ -241,6 +243,25 @@ inline Expression* getFallthrough(Expression* curr,
   return curr;
 }
 
+// Returns whether the resulting value here must fall through without being
+// modified. For example, a tee always does so. That is, this returns false if
+// and only if the return value may have some computation performed on it to
+// change it from the inputs the instruction receives.
+// This differs from getFallthrough() which returns a single value that falls
+// through - here if more than one value can fall through, like in if-else,
+// we can return true. That is, there we care about a value falling through and
+// for us to get that actual value to look at; here we just care whether the
+// value falls through without being changed, even if it might be one of
+// several options.
+inline bool isResultFallthrough(Expression* curr) {
+  // Note that we don't check if there is a return value here; the node may be
+  // unreachable, for example, but then there is no meaningful answer to give
+  // anyhow.
+  return curr->is<LocalSet>() || curr->is<Block>() || curr->is<If>() ||
+         curr->is<Loop>() || curr->is<Try>() || curr->is<Select>() ||
+         curr->is<Break>();
+}
+
 } // namespace Properties
 
 } // namespace wasm
diff --git a/src/passes/DeNaN.cpp b/src/passes/DeNaN.cpp
index 044c13c86..c894cb3f1 100644
--- a/src/passes/DeNaN.cpp
+++ b/src/passes/DeNaN.cpp
@@ -22,6 +22,7 @@
 // differ on wasm's nondeterminism around NaNs.
 //
 
+#include "ir/properties.h"
 #include "pass.h"
 #include "wasm-builder.h"
 #include "wasm.h"
@@ -33,7 +34,18 @@ struct DeNaN : public WalkerPass<
   void visitExpression(Expression* expr) {
     // If the expression returns a floating-point value, ensure it is not a
     // NaN. If we can do this at compile time, do it now, which is useful for
-    // initializations of global (which we can't do a function call in).
+    // initializations of global (which we can't do a function call in). Note
+    // that we don't instrument local.gets, which would cause problems if we
+    // ran this pass more than once (the added functions use gets, and we don't
+    // want to instrument them).
+    if (expr->is<LocalGet>()) {
+      return;
+    }
+    // If the result just falls through without being modified, then we've
+    // already fixed it up earlier.
+    if (Properties::isResultFallthrough(expr)) {
+      return;
+    }
     Builder builder(*getModule());
     Expression* replacement = nullptr;
     auto* c = expr->dynCast<Const>();
@@ -61,6 +73,38 @@ struct DeNaN : public WalkerPass<
     }
   }
 
+  void visitFunction(Function* func) {
+    if (func->imported()) {
+      return;
+    }
+    // Instrument all locals as they enter the function.
+    Builder builder(*getModule());
+    std::vector<Expression*> fixes;
+    auto num = func->getNumParams();
+    for (Index i = 0; i < num; i++) {
+      if (func->getLocalType(i) == Type::f32) {
+        fixes.push_back(builder.makeLocalSet(
+          i,
+          builder.makeCall(
+            "deNan32", {builder.makeLocalGet(i, Type::f32)}, Type::f32)));
+      } else if (func->getLocalType(i) == Type::f64) {
+        fixes.push_back(builder.makeLocalSet(
+          i,
+          builder.makeCall(
+            "deNan64", {builder.makeLocalGet(i, Type::f64)}, Type::f64)));
+      }
+    }
+    if (!fixes.empty()) {
+      fixes.push_back(func->body);
+      func->body = builder.makeBlock(fixes);
+      // Merge blocks so we don't add an unnecessary one.
+      PassRunner runner(getModule(), getPassOptions());
+      runner.setIsNested(true);
+      runner.add("merge-blocks");
+      runner.run();
+    }
+  }
+
   void visitModule(Module* module) {
     // Add helper functions.
     Builder builder(*module);
diff --git a/test/passes/denan.txt b/test/passes/denan.txt
index 67f4989e4..10794b03b 100644
--- a/test/passes/denan.txt
+++ b/test/passes/denan.txt
@@ -1,26 +1,139 @@
 (module
  (type $f32_=>_f32 (func (param f32) (result f32)))
  (type $f64_=>_f64 (func (param f64) (result f64)))
+ (type $i32_f32_i64_f64_=>_none (func (param i32 f32 i64 f64)))
+ (type $f32_f64_=>_none (func (param f32 f64)))
  (global $global$1 (mut f32) (f32.const 0))
  (global $global$2 (mut f32) (f32.const 12.34000015258789))
  (func $foo32 (param $x f32) (result f32)
+  (local.set $x
+   (call $deNan32
+    (local.get $x)
+   )
+  )
   (call $deNan32
    (call $foo32
-    (call $deNan32
-     (local.get $x)
-    )
+    (local.get $x)
    )
   )
  )
  (func $foo64 (param $x f64) (result f64)
+  (local.set $x
+   (call $deNan64
+    (local.get $x)
+   )
+  )
   (call $deNan64
    (call $foo64
-    (call $deNan64
-     (local.get $x)
+    (local.get $x)
+   )
+  )
+ )
+ (func $various (param $x i32) (param $y f32) (param $z i64) (param $w f64)
+  (local.set $y
+   (call $deNan32
+    (local.get $y)
+   )
+  )
+  (local.set $w
+   (call $deNan64
+    (local.get $w)
+   )
+  )
+  (nop)
+ )
+ (func $ignore-local.get (param $f f32) (param $d f64)
+  (local.set $f
+   (call $deNan32
+    (local.get $f)
+   )
+  )
+  (local.set $d
+   (call $deNan64
+    (local.get $d)
+   )
+  )
+  (drop
+   (local.get $f)
+  )
+  (drop
+   (local.get $d)
+  )
+  (local.set $f
+   (local.get $f)
+  )
+  (local.set $d
+   (local.get $d)
+  )
+  (drop
+   (local.get $f)
+  )
+  (drop
+   (local.get $d)
+  )
+  (drop
+   (call $deNan32
+    (f32.abs
+     (local.get $f)
+    )
+   )
+  )
+  (drop
+   (call $deNan64
+    (f64.abs
+     (local.get $d)
+    )
+   )
+  )
+  (local.set $f
+   (call $deNan32
+    (f32.abs
+     (local.get $f)
+    )
+   )
+  )
+  (local.set $d
+   (call $deNan64
+    (f64.abs
+     (local.get $d)
+    )
+   )
+  )
+  (drop
+   (local.get $f)
+  )
+  (drop
+   (local.get $d)
+  )
+ )
+ (func $tees (param $x f32) (result f32)
+  (local.set $x
+   (call $deNan32
+    (local.get $x)
+   )
+  )
+  (local.tee $x
+   (local.tee $x
+    (local.tee $x
+     (local.tee $x
+      (local.get $x)
+     )
     )
    )
   )
  )
+ (func $select (param $x f32) (result f32)
+  (local.set $x
+   (call $deNan32
+    (local.get $x)
+   )
+  )
+  (select
+   (local.get $x)
+   (local.get $x)
+   (i32.const 1)
+  )
+ )
  (func $deNan32 (param $0 f32) (result f32)
   (if (result f32)
    (f32.eq
diff --git a/test/passes/denan.wast b/test/passes/denan.wast
index 73c3c8f1c..7ab0964b2 100644
--- a/test/passes/denan.wast
+++ b/test/passes/denan.wast
@@ -7,4 +7,31 @@
   (func $foo64 (param $x f64) (result f64)
     (call $foo64 (local.get $x))
   )
+  (func $various (param $x i32) (param $y f32) (param $z i64) (param $w f64)
+  )
+  (func $ignore-local.get (param $f f32) (param $d f64)
+    (drop (local.get $f))
+    (drop (local.get $d))
+    (local.set $f (local.get $f))
+    (local.set $d (local.get $d))
+    (drop (local.get $f))
+    (drop (local.get $d))
+    (drop (f32.abs (local.get $f)))
+    (drop (f64.abs (local.get $d)))
+    (local.set $f (f32.abs (local.get $f)))
+    (local.set $d (f64.abs (local.get $d)))
+    (drop (local.get $f))
+    (drop (local.get $d))
+  )
+  (func $tees (param $x f32) (result f32)
+    (local.tee $x
+      (local.tee $x
+        (local.tee $x
+          (local.tee $x
+            (local.get $x))))))
+  (func $select (param $x f32) (result f32)
+    (select
+      (local.get $x)
+      (local.get $x)
+      (i32.const 1)))
 )