SimplifyGlobals: Apply constant globals to segment offsets (#6226)

We already applied such globals to other globals, but can do the same to offsets
of data and element segments.

Suggested in #6220

2 files changed, 117 insertions, 8 deletions

diff --git a/src/passes/SimplifyGlobals.cpp b/src/passes/SimplifyGlobals.cpp
index 0e66b5926..3baec4409 100644
--- a/src/passes/SimplifyGlobals.cpp
+++ b/src/passes/SimplifyGlobals.cpp
@@ -646,23 +646,47 @@ struct SimplifyGlobals : public Pass {
   // since we do know the value during startup, it can't be modified until
   // code runs.
   void propagateConstantsToGlobals() {
-    // Go over the list of globals in order, which is the order of
-    // initialization as well, tracking their constant values.
+    Builder builder(*module);
+
+    // We will note constant globals here as we compute them.
     std::map<Name, Literals> constantGlobals;
+
+    // Given an init expression (something like the init of a global or a
+    // segment), see if it is a simple global.get of a constant that we can
+    // apply.
+    auto applyGlobals = [&](Expression*& init) {
+      if (!init) {
+        // This is the init of a passive segment, which is null.
+        return;
+      }
+      if (auto* get = init->dynCast<GlobalGet>()) {
+        auto iter = constantGlobals.find(get->name);
+        if (iter != constantGlobals.end()) {
+          init = builder.makeConstantExpression(iter->second);
+        }
+      }
+    };
+
+    // Go over the list of globals first, and note their constant values as we
+    // go, as well as applying them where possible.
     for (auto& global : module->globals) {
       if (!global->imported()) {
         if (Properties::isConstantExpression(global->init)) {
           constantGlobals[global->name] =
             getLiteralsFromConstExpression(global->init);
-        } else if (auto* get = global->init->dynCast<GlobalGet>()) {
-          auto iter = constantGlobals.find(get->name);
-          if (iter != constantGlobals.end()) {
-            Builder builder(*module);
-            global->init = builder.makeConstantExpression(iter->second);
-          }
+        } else {
+          applyGlobals(global->init);
         }
       }
     }
+
+    // Go over other things with inits and apply globals there.
+    for (auto& elementSegment : module->elementSegments) {
+      applyGlobals(elementSegment->offset);
+    }
+    for (auto& dataSegment : module->dataSegments) {
+      applyGlobals(dataSegment->offset);
+    }
   }
 
   // Constant propagation part 2: apply the values of immutable globals
diff --git a/test/lit/passes/simplify-globals-offsets.wast b/test/lit/passes/simplify-globals-offsets.wast
new file mode 100644
index 000000000..077c51e73
--- /dev/null
+++ b/test/lit/passes/simplify-globals-offsets.wast
@@ -0,0 +1,85 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; NOTE: This test was ported using port_passes_tests_to_lit.py and could be cleaned up.
+
+;; RUN: foreach %s %t wasm-opt --simplify-globals -all -S -o - | filecheck %s
+
+;; Apply constant globals to segment offsets. The non-imported global.gets will
+;; be applied in the segments named $use-defined.
+
+(module
+  ;; CHECK:      (type $0 (func))
+
+  ;; CHECK:      (import "env" "memory" (memory $memory 256))
+  (import "env" "memory" (memory $memory 256))
+
+  ;; CHECK:      (import "env" "table" (table $table 0 funcref))
+  (import "env" "table" (table $table 0 funcref))
+
+  ;; CHECK:      (import "env" "imported" (global $imported i32))
+  (import "env" "imported" (global $imported i32))
+
+  ;; CHECK:      (global $defined i32 (i32.const 42))
+  (global $defined i32 (i32.const 42))
+
+  ;; CHECK:      (global $use-defined i32 (i32.const 42))
+  (global $use-defined i32 (global.get $defined))
+
+  ;; CHECK:      (data $use-imported (global.get $imported) "hello, world!")
+  (data $use-imported (global.get $imported) "hello, world!")
+
+  ;; CHECK:      (data $use-defined (i32.const 42) "hello, world!")
+  (data $use-defined (global.get $defined) "hello, world!")
+
+  ;; A passive segment has no offset to test, which we should not error on.
+  ;; CHECK:      (data $dropped "hello, world!")
+  (data $dropped "hello, world!")
+
+  ;; CHECK:      (elem $use-imported (global.get $imported) $func)
+  (elem $use-imported (global.get $imported) $func)
+
+  ;; CHECK:      (elem $use-defined (i32.const 42) $func $func)
+  (elem $use-defined (global.get $defined) $func $func)
+
+  ;; CHECK:      (export "func" (func $func))
+  (export "func" (func $func))
+
+  ;; CHECK:      (func $func (type $0)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.load
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (table.get $table
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (global.get $imported)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.const 42)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (data.drop $dropped)
+  ;; CHECK-NEXT: )
+  (func $func
+    ;; Use things to avoid DCE.
+    (drop
+      (i32.load
+        (i32.const 0)
+      )
+    )
+    (drop
+      (table.get $table
+        (i32.const 0)
+      )
+    )
+    (drop
+      (global.get $imported)
+    )
+    (drop
+      (global.get $use-defined)
+    )
+    (data.drop $dropped)
+  )
+)