StringGathering pass

src/ir/type-updating.h

@@ -443,7 +443,6 @@ class TypeMapper : public GlobalTypeRewriter {
 
   std::unordered_map<HeapType, Signature> newSignatures;
 
-public:
   TypeMapper(Module& wasm, const TypeUpdates& mapping)
     : GlobalTypeRewriter(wasm), mapping(mapping) {}
 

src/passes/CMakeLists.txt

@@ -90,6 +90,7 @@ set(passes_SOURCES
   SignaturePruning.cpp
   SignatureRefining.cpp
   SignExtLowering.cpp
+  StringLowering.cpp
   Strip.cpp
   StripTargetFeatures.cpp
   RedundantSetElimination.cpp

src/passes/StringLowering.cpp

@@ -0,0 +1,180 @@
+/*
+ * Copyright 2024 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//
+// Utilities for lowering strings into simpler things.
+//
+// StringGathering collects all string.const operations and stores them in
+// globals, avoiding them appearing in code that can run more than once (which
+// can have overhead in VMs).
+//
+// Building on that, an extended version of StringGathering will also replace
+// those new globals with imported globals of type externref, for use with the
+// string imports proposal. String operations will likewise need to be lowered.
+// TODO
+//
+
+#include <algorithm>
+
+#include "ir/module-utils.h"
+#include "ir/names.h"
+#include "pass.h"
+#include "wasm-builder.h"
+#include "wasm.h"
+
+namespace wasm {
+
+struct StringGathering : public Pass {
+  // All the strings we found in the module.
+  std::vector<Name> strings;
+
+  // Pointers to all StringConsts, so that we can replace them.
+  using StringPtrs = std::vector<Expression**>;
+  StringPtrs stringPtrs;
+
+  // Main entry point.
+  void run(Module* module) override {
+    processModule(module);
+    addGlobals(module);
+    replaceStrings(module);
+  }
+
+  // Scan the entire wasm to find the relevant strings to populate our global
+  // data structures.
+  void processModule(Module* module) {
+    struct StringWalker : public PostWalker<StringWalker> {
+      StringPtrs& stringPtrs;
+
+      StringWalker(StringPtrs& stringPtrs) : stringPtrs(stringPtrs) {}
+
+      void visitStringConst(StringConst* curr) {
+        stringPtrs.push_back(getCurrentPointer());
+      }
+    };
+
+    ModuleUtils::ParallelFunctionAnalysis<StringPtrs> analysis(
+      *module, [&](Function* func, StringPtrs& stringPtrs) {
+        if (!func->imported()) {
+          StringWalker(stringPtrs).walk(func->body);
+        }
+      });
+
+    // Also walk the global module code (for simplicity, also add it to the
+    // function map, using a "function" key of nullptr).
+    auto& globalStrings = analysis.map[nullptr];
+    StringWalker(globalStrings).walkModuleCode(module);
+
+    // Combine all the strings.
+    std::unordered_set<Name> stringSet;
+    for (auto& [_, currStringPtrs] : analysis.map) {
+      for (auto** stringPtr : currStringPtrs) {
+        stringSet.insert((*stringPtr)->cast<StringConst>()->string);
+        stringPtrs.push_back(stringPtr);
+      }
+    }
+
+    // Sort the strings for determinism (alphabetically).
+    strings = std::vector<Name>(stringSet.begin(), stringSet.end());
+    std::sort(strings.begin(), strings.end());
+  }
+
+  // For each string, the name of the global that replaces it.
+  std::unordered_map<Name, Name> stringToGlobalName;
+
+  Type nnstringref = Type(HeapType::string, NonNullable);
+
+  // Existing globals already in the form we emit can be reused. That is, if
+  // we see
+  //
+  //  (global $foo (ref string) (string.const ..))
+  //
+  // then we can just use that as the global for that string. This avoids
+  // repeated executions of the pass adding more and more globals.
+  //
+  // Note that we don't note these in newNames: They are already in the right
+  // sorted position, before any uses, as we use the first of them for each
+  // string. Only actually new names need sorting.
+  //
+  // Any time we reuse a global, we must not modify its body (or else we'd
+  // replace the global that all others read from); we note them here and
+  // avoid them in replaceStrings later to avoid such trampling.
+  std::unordered_set<Expression**> stringPtrsToPreserve;
+
+  void addGlobals(Module* module) {
+    // Note all the new names we create for the sorting later.
+    std::unordered_set<Name> newNames;
+
+    // Find globals to reuse (see comment on stringPtrsToPreserve for context).
+    for (auto& global : module->globals) {
+      if (global->type == nnstringref && !global->imported()) {
+        if (auto* stringConst = global->init->dynCast<StringConst>()) {
+          auto& globalName = stringToGlobalName[stringConst->string];
+          if (!globalName.is()) {
+            // This is the first global for this string, use it.
+            globalName = global->name;
+            stringPtrsToPreserve.insert(&global->init);
+          }
+        }
+      }
+    }
+
+    Builder builder(*module);
+    for (Index i = 0; i < strings.size(); i++) {
+      auto& globalName = stringToGlobalName[strings[i]];
+      if (globalName.is()) {
+        // We are reusing a global for this one.
+        continue;
+      }
+
+      auto& string = strings[i];
+      auto name = Names::getValidGlobalName(
+        *module, std::string("string.const_") + std::string(string.str));
+      globalName = name;
+      newNames.insert(name);
+      auto* stringConst = builder.makeStringConst(string);
+      auto global =
+        builder.makeGlobal(name, nnstringref, stringConst, Builder::Immutable);
+      module->addGlobal(std::move(global));
+    }
+
+    // Sort our new globals to the start, as other global initializers may use
+    // them (and it would be invalid for us to appear after a use). This sort is
+    // a simple way to ensure that we validate, but it may be unoptimal (we
+    // leave that for reorder-globals).
+    std::stable_sort(
+      module->globals.begin(),
+      module->globals.end(),
+      [&](const std::unique_ptr<Global>& a, const std::unique_ptr<Global>& b) {
+        return newNames.count(a->name) && !newNames.count(b->name);
+      });
+  }
+
+  void replaceStrings(Module* module) {
+    Builder builder(*module);
+    for (auto** stringPtr : stringPtrs) {
+      if (stringPtrsToPreserve.count(stringPtr)) {
+        continue;
+      }
+      auto* stringConst = (*stringPtr)->cast<StringConst>();
+      auto globalName = stringToGlobalName[stringConst->string];
+      *stringPtr = builder.makeGlobalGet(globalName, nnstringref);
+    }
+  }
+};
+
+Pass* createStringGatheringPass() { return new StringGathering(); }
+
+} // namespace wasm

src/passes/pass.cpp

@@ -475,6 +475,9 @@ void PassRegistry::registerPasses() {
     "ssa-nomerge",
     "ssa-ify variables so that they have a single assignment, ignoring merges",
     createSSAifyNoMergePass);
+  registerPass("string-gathering",
+               "gathers wasm strings to globals",
+               createStringGatheringPass);
   registerPass(
     "strip", "deprecated; same as strip-debug", createStripDebugPass);
   registerPass("stack-check",
@@ -710,6 +713,11 @@ void PassRunner::addDefaultGlobalOptimizationPostPasses() {
     addIfNoDWARFIssues("simplify-globals");
   }
   addIfNoDWARFIssues("remove-unused-module-elements");
+  if (options.optimizeLevel >= 2 && wasm->features.hasStrings()) {
+    // Gather strings to globals right before reorder-globals, which will then
+    // sort them properly.
+    addIfNoDWARFIssues("string-gathering");
+  }
   if (options.optimizeLevel >= 2 || options.shrinkLevel >= 1) {
     addIfNoDWARFIssues("reorder-globals");
   }

src/passes/passes.h

@@ -153,6 +153,7 @@ Pass* createSimplifyLocalsNoTeePass();
 Pass* createSimplifyLocalsNoStructurePass();
 Pass* createSimplifyLocalsNoTeeNoStructurePass();
 Pass* createStackCheckPass();
+Pass* createStringGatheringPass();
 Pass* createStripDebugPass();
 Pass* createStripDWARFPass();
 Pass* createStripProducersPass();

test/lit/help/wasm-opt.test

@@ -467,6 +467,8 @@
 ;; CHECK-NEXT:   --stack-check                                 enforce limits on llvm's
 ;; CHECK-NEXT:                                                 __stack_pointer global
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --string-gathering                            gathers wasm strings to globals
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --strip                                       deprecated; same as strip-debug
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --strip-debug                                 strip debug info (including the

test/lit/help/wasm2js.test

@@ -426,6 +426,8 @@
 ;; CHECK-NEXT:   --stack-check                                 enforce limits on llvm's
 ;; CHECK-NEXT:                                                 __stack_pointer global
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --string-gathering                            gathers wasm strings to globals
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --strip                                       deprecated; same as strip-debug
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --strip-debug                                 strip debug info (including the