Move segment merging to fit web limits into its own pass (#1980)

It was previously part of writing a binary, but changing the number of
segments at such a late stage would not work in the presence of bulk
memory's datacount section. Also updates the memory packing pass
to respect the web's limits on the number of data segments.

8 files changed, 228 insertions, 138 deletions

diff --git a/src/ir/memory-utils.h b/src/ir/memory-utils.h
index ea19ca799..0bad50c5a 100644
--- a/src/ir/memory-utils.h
+++ b/src/ir/memory-utils.h
@@ -22,6 +22,7 @@
 
 #include "literal.h"
 #include "wasm.h"
+#include "wasm-binary.h"
 
 namespace wasm {
 
@@ -59,7 +60,109 @@ namespace MemoryUtils {
       memory.initial = memory.max = 1;
     }
   }
-};
+
+  // Try to merge segments until they fit into web limitations.
+  // Return true if successful.
+  inline bool ensureLimitedSegments(Module& module) {
+    Memory& memory = module.memory;
+    if (memory.segments.size() <= WebLimitations::MaxDataSegments) {
+      return true;
+    }
+
+    auto isEmpty = [](Memory::Segment& segment) {
+      return segment.data.size() == 0;
+    };
+
+    auto isConstantOffset = [](Memory::Segment& segment) {
+      return segment.offset && segment.offset->is<Const>();
+    };
+
+    Index numConstant = 0,
+           numDynamic = 0;
+    bool hasPassiveSegments = false;
+    for (auto& segment : memory.segments) {
+      if (!isEmpty(segment)) {
+        if (isConstantOffset(segment)) {
+          numConstant++;
+        } else {
+          numDynamic++;
+        }
+      }
+      hasPassiveSegments |= segment.isPassive;
+    }
+
+    if (hasPassiveSegments) {
+      return false;
+    }
+
+    // check if we have too many dynamic data segments, which we can do nothing about
+    auto num = numConstant + numDynamic;
+    if (numDynamic + 1 >= WebLimitations::MaxDataSegments) {
+      return false;
+    }
+
+    // we'll merge constant segments if we must
+    if (numConstant + numDynamic >= WebLimitations::MaxDataSegments) {
+      numConstant = WebLimitations::MaxDataSegments - numDynamic - 1;
+      num = numConstant + numDynamic;
+      assert(num == WebLimitations::MaxDataSegments - 1);
+    }
+
+    std::vector<Memory::Segment> mergedSegments;
+    mergedSegments.reserve(WebLimitations::MaxDataSegments);
+
+    // drop empty segments and pass through dynamic-offset segments
+    for (auto& segment : memory.segments) {
+      if (isEmpty(segment)) continue;
+      if (isConstantOffset(segment)) continue;
+      mergedSegments.push_back(segment);
+    }
+
+    // from here on, we concern ourselves with non-empty constant-offset
+    // segments, the ones which we may need to merge
+    auto isRelevant = [&](Memory::Segment& segment) {
+      return !isEmpty(segment) && isConstantOffset(segment);
+    };
+    for (Index i = 0; i < memory.segments.size(); i++) {
+      auto& segment = memory.segments[i];
+      if (!isRelevant(segment)) continue;
+      if (mergedSegments.size() + 2 < WebLimitations::MaxDataSegments) {
+        mergedSegments.push_back(segment);
+        continue;
+      }
+      // we can emit only one more segment! merge everything into one
+      // start the combined segment at the bottom of them all
+      auto start = segment.offset->cast<Const>()->value.getInteger();
+      for (Index j = i + 1; j < memory.segments.size(); j++) {
+        auto& segment = memory.segments[j];
+        if (!isRelevant(segment)) continue;
+        auto offset = segment.offset->cast<Const>()->value.getInteger();
+        start = std::min(start, offset);
+      }
+      // create the segment and add in all the data
+      auto* c = module.allocator.alloc<Const>();
+      c->value = Literal(int32_t(start));
+      c->type = i32;
+
+      Memory::Segment combined(c);
+      for (Index j = i; j < memory.segments.size(); j++) {
+        auto& segment = memory.segments[j];
+        if (!isRelevant(segment)) continue;
+        auto offset = segment.offset->cast<Const>()->value.getInteger();
+        auto needed = offset + segment.data.size() - start;
+        if (combined.data.size() < needed) {
+          combined.data.resize(needed);
+        }
+        std::copy(segment.data.begin(), segment.data.end(), combined.data.begin() + (offset - start));
+      }
+      mergedSegments.push_back(combined);
+      break;
+    }
+
+    memory.segments.swap(mergedSegments);
+    return true;
+  }
+} // namespace MemoryUtils
 
 } // namespace wasm
 
diff --git a/src/passes/CMakeLists.txt b/src/passes/CMakeLists.txt
index 2b1b64d84..d2e713e2b 100644
--- a/src/passes/CMakeLists.txt
+++ b/src/passes/CMakeLists.txt
@@ -22,6 +22,7 @@ SET(passes_SOURCES
   InstrumentLocals.cpp
   InstrumentMemory.cpp
   LegalizeJSInterface.cpp
+  LimitSegments.cpp
   LocalCSE.cpp
   LogExecution.cpp
   LoopInvariantCodeMotion.cpp
diff --git a/src/passes/LimitSegments.cpp b/src/passes/LimitSegments.cpp
new file mode 100644
index 000000000..521969a28
--- /dev/null
+++ b/src/passes/LimitSegments.cpp
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2019 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.
+ */
+
+#include "ir/memory-utils.h"
+#include "pass.h"
+#include "wasm.h"
+
+using namespace std;
+
+namespace wasm {
+
+struct LimitSegments : public Pass {
+  void run(PassRunner* runner, Module* module) override {
+    if (!MemoryUtils::ensureLimitedSegments(*module)) {
+      std::cerr << "Unable to merge segments. "
+                << "wasm VMs may not accept this binary"
+                << std::endl;
+    }
+  }
+};
+
+Pass *createLimitSegmentsPass() {
+  return new LimitSegments();
+}
+
+} // namespace wasm
diff --git a/src/passes/MemoryPacking.cpp b/src/passes/MemoryPacking.cpp
index 2c0dac395..ef150e2de 100644
--- a/src/passes/MemoryPacking.cpp
+++ b/src/passes/MemoryPacking.cpp
@@ -14,9 +14,10 @@
  * limitations under the License.
  */
 
-#include <wasm.h>
-#include <pass.h>
-#include <wasm-builder.h>
+#include "pass.h"
+#include "wasm.h"
+#include "wasm-binary.h"
+#include "wasm-builder.h"
 
 namespace wasm {
 
@@ -27,49 +28,82 @@ struct MemoryPacking : public Pass {
   bool modifiesBinaryenIR() override { return false; }
 
   void run(PassRunner* runner, Module* module) override {
-    if (!module->memory.exists) return;
+    // Conservatively refuse to change segments if bulk memory is enabled to
+    // avoid invalidating segment indices or segment contents referenced from
+    // memory.init instructions.
+    // TODO: optimize in the presence of memory.init instructions
+    if (!module->memory.exists || runner->options.features.hasBulkMemory()) {
+      return;
+    }
+
     std::vector<Memory::Segment> packed;
+
+    // we can only handle a constant offset for splitting
+    auto isSplittable = [&](const Memory::Segment& segment) {
+      return segment.offset->is<Const>();
+    };
+
+    for (auto& segment : module->memory.segments) {
+      if (!isSplittable(segment)) {
+        packed.push_back(segment);
+      }
+    }
+
+    size_t numRemaining = module->memory.segments.size() - packed.size();
+
+    // Split only if we have room for more segments
+    auto shouldSplit = [&]() {
+      return WebLimitations::MaxDataSegments > packed.size() + numRemaining;
+    };
+
     for (auto& segment : module->memory.segments) {
+      if (!isSplittable(segment)) continue;
+
       // skip final zeros
       while (segment.data.size() > 0 && segment.data.back() == 0) {
         segment.data.pop_back();
       }
-      // we can only handle a constant offset for splitting
-      Const* offset;
-      if (!segment.isPassive && (offset = segment.offset->dynCast<Const>())) {
-        // Find runs of zeros, and split
-        auto& data = segment.data;
-        auto base = offset->value.geti32();
-        Index start = 0;
-        // create new segments
-        while (start < data.size()) {
-          // skip initial zeros
-          while (start < data.size() && data[start] == 0) {
-            start++;
-          }
-          Index end = start; // end of data-containing part
-          Index next = end; // after zeros we can skip. preserves next >= end
-          while (next < data.size() && (next - end < OVERHEAD)) {
-            if (data[end] != 0) {
-              end++;
-              next = end; // we can try to skip zeros from here
+
+      if (!shouldSplit()) {
+        packed.push_back(segment);
+        continue;
+      }
+
+      auto* offset = segment.offset->cast<Const>();
+      // Find runs of zeros, and split
+      auto& data = segment.data;
+      auto base = offset->value.geti32();
+      Index start = 0;
+      // create new segments
+      while (start < data.size()) {
+        // skip initial zeros
+        while (start < data.size() && data[start] == 0) {
+          start++;
+        }
+        Index end = start; // end of data-containing part
+        Index next = end; // after zeros we can skip. preserves next >= end
+        if (!shouldSplit()) {
+          next = end = data.size();
+        }
+        while (next < data.size() && (next - end < OVERHEAD)) {
+          if (data[end] != 0) {
+            end++;
+            next = end; // we can try to skip zeros from here
+          } else {
+            // end is on a zero, we are looking to skip
+            if (data[next] != 0) {
+              end = next; // we must extend the segment, including some zeros
             } else {
-              // end is on a zero, we are looking to skip
-              if (data[next] != 0) {
-                end = next; // we must extend the segment, including some zeros
-              } else {
-                next++;
-              }
+              next++;
             }
           }
-          if (end != start) {
-            packed.emplace_back(Builder(*module).makeConst(Literal(int32_t(base + start))), &data[start], end - start);
-          }
-          start = next;
         }
-      } else {
-        packed.push_back(segment);
+        if (end != start) {
+          packed.emplace_back(Builder(*module).makeConst(Literal(int32_t(base + start))), &data[start], end - start);
+        }
+        start = next;
       }
+      numRemaining--;
     }
     module->memory.segments.swap(packed);
   }
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index fecb644b5..786a39474 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -93,6 +93,7 @@ void PassRegistry::registerPasses() {
   registerPass("instrument-locals", "instrument the build with code to intercept all loads and stores", createInstrumentLocalsPass);
   registerPass("instrument-memory", "instrument the build with code to intercept all loads and stores", createInstrumentMemoryPass);
   registerPass("licm", "loop invariant code motion", createLoopInvariantCodeMotionPass);
+  registerPass("limit-segments", "attempt to merge segments to fit within web limits", createLimitSegmentsPass);
   registerPass("memory-packing", "packs memory into separate segments, skipping zeros", createMemoryPackingPass);
   registerPass("merge-blocks", "merges blocks to their parents", createMergeBlocksPass);
   registerPass("merge-locals", "merges locals when beneficial", createMergeLocalsPass);
diff --git a/src/passes/passes.h b/src/passes/passes.h
index fea06e388..e4111c43b 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -44,6 +44,7 @@ Pass* createInliningPass();
 Pass* createInliningOptimizingPass();
 Pass* createLegalizeJSInterfacePass();
 Pass* createLegalizeJSInterfaceMinimallyPass();
+Pass* createLimitSegmentsPass();
 Pass* createLocalCSEPass();
 Pass* createLogExecutionPass();
 Pass* createInstrumentLocalsPass();
diff --git a/src/tools/wasm-ctor-eval.cpp b/src/tools/wasm-ctor-eval.cpp
index 8b8eaa4ad..be1de7684 100644
--- a/src/tools/wasm-ctor-eval.cpp
+++ b/src/tools/wasm-ctor-eval.cpp
@@ -25,9 +25,9 @@
 #include <memory>
 
 #include "pass.h"
-#include "support/command-line.h"
 #include "support/file.h"
 #include "support/colors.h"
+#include "tool-options.h"
 #include "wasm-io.h"
 #include "wasm-interpreter.h"
 #include "wasm-builder.h"
@@ -383,7 +383,7 @@ int main(int argc, const char* argv[]) {
   bool debugInfo = false;
   std::string ctorsString;
 
-  Options options("wasm-ctor-eval", "Execute C++ global constructors ahead of time");
+  ToolOptions options("wasm-ctor-eval", "Execute C++ global constructors ahead of time");
   options
       .add("--output", "-o", "Output file (stdout if not specified)",
            Options::Arguments::One,
@@ -425,6 +425,8 @@ int main(int argc, const char* argv[]) {
     }
   }
 
+  options.calculateFeatures(wasm);
+
   if (!WasmValidator().validate(wasm)) {
     WasmPrinter::printModule(&wasm);
     Fatal() << "error in validating input";
@@ -442,6 +444,7 @@ int main(int argc, const char* argv[]) {
   // Do some useful optimizations after the evalling
   {
     PassRunner passRunner(&wasm);
+    passRunner.setFeatures(options.passOptions.features);
     passRunner.add("memory-packing"); // we flattened it, so re-optimize
     passRunner.add("remove-unused-names");
     passRunner.add("dce");
diff --git a/src/wasm/wasm-binary.cpp b/src/wasm/wasm-binary.cpp
index 63adbb08f..b1fc03017 100644
--- a/src/wasm/wasm-binary.cpp
+++ b/src/wasm/wasm-binary.cpp
@@ -311,19 +311,16 @@ void WasmBinaryWriter::writeExports() {
   finishSection(start);
 }
 
-static bool isEmpty(Memory::Segment& segment) {
-  return segment.data.size() == 0;
-}
-
-static bool isConstantOffset(Memory::Segment& segment) {
-  return segment.offset && segment.offset->is<Const>();
-}
-
 void WasmBinaryWriter::writeDataSegments() {
   if (wasm->memory.segments.size() == 0) return;
-
-  Index emitted = 0;
-  auto emit = [&](Memory::Segment& segment) {
+  if (wasm->memory.segments.size() > WebLimitations::MaxDataSegments) {
+    std::cerr << "Some VMs may not accept this binary because it has a large "
+              << "number of data segments. Run the limit-segments pass to "
+              << "merge segments." << std::endl;
+  }
+  auto start = startSection(BinaryConsts::Section::Data);
+  o << U32LEB(wasm->memory.segments.size());
+  for (auto& segment : wasm->memory.segments) {
     uint32_t flags = 0;
     if (segment.isPassive) {
       flags |= BinaryConsts::IsPassive;
@@ -334,95 +331,6 @@ void WasmBinaryWriter::writeDataSegments() {
       o << int8_t(BinaryConsts::End);
     }
     writeInlineBuffer(&segment.data[0], segment.data.size());
-    emitted++;
-  };
-
-  Index numConstant = 0,
-        numDynamic = 0;
-  bool hasPassiveSegments = false;
-  for (auto& segment : wasm->memory.segments) {
-    if (!isEmpty(segment)) {
-      if (isConstantOffset(segment)) {
-        numConstant++;
-      } else {
-        numDynamic++;
-      }
-    }
-    hasPassiveSegments |= segment.isPassive;
-  }
-
-  if (hasPassiveSegments) {
-    if (wasm->memory.segments.size() > WebLimitations::MaxDataSegments) {
-      std::cerr << "too many data segments, wasm VMs may not accept this binary" << std::endl;
-    }
-    // TODO: merge segments in a pass that can fix up memory.init instructions
-    auto section = startSection(BinaryConsts::Section::Data);
-    o << U32LEB(wasm->memory.segments.size());
-    for (auto& segment : wasm->memory.segments) {
-      emit(segment);
-    }
-    finishSection(section);
-    return;
-  }
-  // check if we have too many dynamic data segments, which we can do nothing about
-  auto num = numConstant + numDynamic;
-  if (numDynamic + 1 >= WebLimitations::MaxDataSegments) {
-    std::cerr << "too many non-constant-offset data segments, wasm VMs may not accept this binary" << std::endl;
-  }
-  // we'll merge constant segments if we must
-  if (numConstant + numDynamic >= WebLimitations::MaxDataSegments) {
-    numConstant = WebLimitations::MaxDataSegments - numDynamic - 1;
-    num = numConstant + numDynamic;
-    assert(num == WebLimitations::MaxDataSegments - 1);
-  }
-  auto start = startSection(BinaryConsts::Section::Data);
-  o << U32LEB(num);
-  // first, emit all non-constant-offset segments; then emit the constants,
-  // which we may merge if forced to
-  auto& segments = wasm->memory.segments;
-  for (auto& segment : segments) {
-    if (isEmpty(segment)) continue;
-    if (isConstantOffset(segment)) continue;
-    emit(segment);
-  }
-  // from here on, we concern ourselves with non-empty constant-offset
-  // segments, the ones which we may need to merge
-  auto isRelevant = [](Memory::Segment& segment) {
-    return !isEmpty(segment) && isConstantOffset(segment);
-  };
-  for (Index i = 0; i < segments.size(); i++) {
-    auto& segment = segments[i];
-    if (!isRelevant(segment)) continue;
-    if (emitted + 2 < WebLimitations::MaxDataSegments) {
-      emit(segment);
-    } else {
-      // we can emit only one more segment! merge everything into one
-      // start the combined segment at the bottom of them all
-      auto start = segment.offset->cast<Const>()->value.getInteger();
-      for (Index j = i + 1; j < segments.size(); j++) {
-        auto& segment = segments[j];
-        if (!isRelevant(segment)) continue;
-        auto offset = segment.offset->cast<Const>()->value.getInteger();
-        start = std::min(start, offset);
-      }
-      // create the segment and add in all the data
-      Const c;
-      c.value = Literal(int32_t(start));
-      c.type = i32;
-      Memory::Segment combined(&c);
-      for (Index j = i; j < segments.size(); j++) {
-        auto& segment = segments[j];
-        if (!isRelevant(segment)) continue;
-        auto offset = segment.offset->cast<Const>()->value.getInteger();
-        auto needed = offset + segment.data.size() - start;
-        if (combined.data.size() < needed) {
-          combined.data.resize(needed);
-        }
-        std::copy(segment.data.begin(), segment.data.end(), combined.data.begin() + (offset - start));
-      }
-      emit(combined);
-      break;
-    }
   }
   finishSection(start);
 }