Allow customizing the limit in LimitSegments (#7285)

The custom limit can be passed by e.g.

--limit-segments --pass-arg=limit-segments@1024

Fixes #7229

Author: lorsanta

src/ir/memory-utils.h

@@ -46,12 +46,14 @@ inline void ensureExists(Module* wasm) {
 // Try to merge segments until they fit into web limitations.
 // Return true if successful.
 // Does not yet support multimemory
-inline bool ensureLimitedSegments(Module& module) {
+inline bool
+ensureLimitedSegments(Module& module,
+                      Index maxDataSegments = WebLimitations::MaxDataSegments) {
   if (module.memories.size() > 1) {
     return false;
   }
   auto& dataSegments = module.dataSegments;
-  if (dataSegments.size() <= WebLimitations::MaxDataSegments) {
+  if (dataSegments.size() <= maxDataSegments) {
     return true;
   }
 
@@ -86,19 +88,19 @@ inline bool ensureLimitedSegments(Module& module) {
 
   // check if we have too many dynamic data segments, which we can do nothing
   // about
-  if (numDynamic + 1 >= WebLimitations::MaxDataSegments) {
+  if (numDynamic + 1 >= maxDataSegments) {
     return false;
   }
 
   // we'll merge constant segments if we must
-  if (numConstant + numDynamic >= WebLimitations::MaxDataSegments) {
-    numConstant = WebLimitations::MaxDataSegments - numDynamic - 1;
+  if (numConstant + numDynamic >= maxDataSegments) {
+    numConstant = maxDataSegments - numDynamic - 1;
     [[maybe_unused]] auto num = numConstant + numDynamic;
-    assert(num == WebLimitations::MaxDataSegments - 1);
+    assert(num == maxDataSegments - 1);
   }
 
   std::vector<std::unique_ptr<wasm::DataSegment>> mergedSegments;
-  mergedSegments.reserve(WebLimitations::MaxDataSegments);
+  mergedSegments.reserve(maxDataSegments);
 
   // drop empty segments and pass through dynamic-offset segments
   for (auto& segment : dataSegments) {
@@ -121,7 +123,7 @@ inline bool ensureLimitedSegments(Module& module) {
     if (!isRelevant(*segment)) {
       continue;
     }
-    if (mergedSegments.size() + 2 < WebLimitations::MaxDataSegments) {
+    if (mergedSegments.size() + 2 < maxDataSegments) {
       mergedSegments.push_back(std::move(segment));
       continue;
     }

src/passes/LimitSegments.cpp

@@ -18,11 +18,28 @@
 #include "pass.h"
 #include "wasm.h"
 
+//
+// Attempt to merge segments to fit within a specified limit.
+//
+// By default this limit is equal to the one commonly used by wasm VMs
+// (see wasm-limits.h), but it can be changed with the option below:
+//
+//   --pass-arg=limit-segments@max-data-segments
+//
+//       Specify a custom maximum number of data segments.
+//
+
 namespace wasm {
 
 struct LimitSegments : public Pass {
   void run(Module* module) override {
-    if (!MemoryUtils::ensureLimitedSegments(*module)) {
+    Index maxDataSegments;
+    if (hasArgument("limit-segments")) {
+      maxDataSegments = std::stoul(getArgument("limit-segments", ""));
+    } else {
+      maxDataSegments = WebLimitations::MaxDataSegments;
+    }
+    if (!MemoryUtils::ensureLimitedSegments(*module, maxDataSegments)) {
       std::cerr << "Unable to merge segments. "
                 << "wasm VMs may not accept this binary" << std::endl;
     }

test/lit/passes/custom-max-data-segments.wast

@@ -0,0 +1,13 @@
+;; RUN: wasm-opt %s --limit-segments --pass-arg=limit-segments@3 -S -o - | filecheck %s
+
+;; Test that the data segments custom limit is respected.
+(module
+ (memory 256 256)
+ ;; CHECK:      (data $0 (i32.const 0) "A")
+ (data (i32.const 0) "A")
+ ;; CHECK:      (data $"" (i32.const 1) "AAAA")
+ (data (i32.const 1) "A")
+ (data (i32.const 2) "A")
+ (data (i32.const 3) "A")
+ (data (i32.const 4) "A")
+)