Re-land #2235 with fixes (#2245)

#2242 had exposed the bug that the `Trapper` pass was defining `walkFunction` when it should have been defining `doWalkFunction`.

1 files changed, 42 insertions, 4 deletions

diff --git a/src/passes/MemoryPacking.cpp b/src/passes/MemoryPacking.cpp
index 0b24379e5..a831c0efc 100644
--- a/src/passes/MemoryPacking.cpp
+++ b/src/passes/MemoryPacking.cpp
@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+#include "ir/manipulation.h"
+#include "ir/utils.h"
 #include "pass.h"
 #include "wasm-binary.h"
 #include "wasm-builder.h"
@@ -32,11 +34,13 @@ struct MemoryPacking : public Pass {
       return;
     }
 
-    // Conservatively refuse to change segments if any are passive to avoid
-    // invalidating segment indices or segment contents referenced from
-    // memory.init instructions.
-    // TODO: optimize in the presence of memory.init instructions
     if (module->features.hasBulkMemory()) {
+      // Remove any references to active segments that might be invalidated.
+      optimizeTrappingBulkMemoryOps(runner, module);
+      // Conservatively refuse to change segments if any are passive to avoid
+      // invalidating segment indices or segment contents referenced from
+      // memory.init and data.drop instructions.
+      // TODO: optimize in the presence of memory.init and  data.drop
       for (auto segment : module->memory.segments) {
         if (segment.isPassive) {
           return;
@@ -120,6 +124,40 @@ struct MemoryPacking : public Pass {
     }
     module->memory.segments.swap(packed);
   }
+
+  void optimizeTrappingBulkMemoryOps(PassRunner* runner, Module* module) {
+    struct Trapper : WalkerPass<PostWalker<Trapper>> {
+      bool isFunctionParallel() override { return true; }
+      bool changed;
+
+      Pass* create() override { return new Trapper; }
+
+      void visitMemoryInit(MemoryInit* curr) {
+        if (!getModule()->memory.segments[curr->segment].isPassive) {
+          Builder builder(*getModule());
+          replaceCurrent(builder.blockify(builder.makeDrop(curr->dest),
+                                          builder.makeDrop(curr->offset),
+                                          builder.makeDrop(curr->size),
+                                          builder.makeUnreachable()));
+          changed = true;
+        }
+      }
+      void visitDataDrop(DataDrop* curr) {
+        if (!getModule()->memory.segments[curr->segment].isPassive) {
+          ExpressionManipulator::unreachable(curr);
+          changed = true;
+        }
+      }
+      void doWalkFunction(Function* func) {
+        changed = false;
+        super::doWalkFunction(func);
+        if (changed) {
+          ReFinalize().walkFunctionInModule(func, getModule());
+        }
+      }
+    } trapper;
+    trapper.run(runner, module);
+  }
 };
 
 Pass* createMemoryPackingPass() { return new MemoryPacking(); }