MemoryPacking: Handle non-empty trapping segments

src/passes/MemoryPacking.cpp

@@ -34,6 +34,7 @@
 #include "ir/utils.h"
 #include "pass.h"
 #include "support/space.h"
+#include "support/stdckdint.h"
 #include "wasm-binary.h"
 #include "wasm-builder.h"
 #include "wasm.h"
@@ -46,6 +47,7 @@ namespace {
 // will be used instead of memory.init for this range.
 struct Range {
  bool isZero;
+ // The range [start, end) - that is, start is included while end is not.
  size_t start;
  size_t end;
 };
@@ -107,7 +109,8 @@ struct MemoryPacking : public Pass {
  void dropUnusedSegments(Module* module,
  std::vector<std::unique_ptr<DataSegment>>& segments,
  ReferrersMap& referrers);
- bool canSplit(const std::unique_ptr<DataSegment>& segment,
+ bool canSplit(Module* module,
+ const std::unique_ptr<DataSegment>& segment,
  const Referrers& referrers);
  void calculateRanges(const std::unique_ptr<DataSegment>& segment,
  const Referrers& referrers,
@@ -161,7 +164,7 @@ void MemoryPacking::run(Module* module) {
 
  std::vector<Range> ranges;
 
- if (canSplit(segment, currReferrers)) {
+ if (canSplit(module, segment, currReferrers)) {
  calculateRanges(segment, currReferrers, ranges);
  } else {
  // A single range covers the entire segment. Set isZero to false so the
@@ -260,7 +263,8 @@ bool MemoryPacking::canOptimize(
  return true;
 }
 
-bool MemoryPacking::canSplit(const std::unique_ptr<DataSegment>& segment,
+bool MemoryPacking::canSplit(Module* module,
+ const std::unique_ptr<DataSegment>& segment,
  const Referrers& referrers) {
  // Don't mess with segments related to llvm coverage tools such as
  // __llvm_covfun. There segments are expected/parsed by external downstream
@@ -277,6 +281,24 @@ bool MemoryPacking::canSplit(const std::unique_ptr<DataSegment>& segment,
  return false;
  }
 
+ // Do not split an active segment that might trap.
+ if (!getPassOptions().trapsNeverHappen && segment->offset) {
+ auto* c = segment->offset->dynCast<Const>();
+ if (!c) {
+ // We can't tell if this will trap or not.
+ return false;
+ }
+ auto* memory = module->getMemory(segment->memory);
+ auto memorySize = memory->initial * Memory::kPageSize;
+ Index start = c->value.getUnsigned();
+ Index size = segment->data.size();
+ Index end;
+ if (std::ckd_add(&end, start, size) || end > memorySize) {
+ // This traps.
+ return false;
+ }
+ }
+
  for (auto* referrer : referrers) {
  if (auto* curr = referrer->dynCast<MemoryInit>()) {
  if (segment->isPassive) {

test/lit/passes/memory-packing_traps.wast

@@ -0,0 +1,96 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: foreach %s %t wasm-opt --memory-packing -all --zero-filled-memory -S -o - | filecheck %s
+;; RUN: foreach %s %t wasm-opt --memory-packing -all --zero-filled-memory -tnh -S -o - | filecheck %s --check-prefix=_TNH_
+
+(module
+ ;; We should not optimize out a segment that will trap, as that is an effect
+ ;; we need to preserve (unless TrapsNeverHappen).
+ (memory 1 2 shared)
+ (data (i32.const -1) "\00")
+)
+
+;; CHECK: (memory $0 1 2 shared)
+
+;; CHECK: (data $0 (i32.const -1) "\00")
+
+;; _TNH_: (memory $0 1 2 shared)
+(module
+ ;; We should handle the possible overflow in adding the offset and size, and
+ ;; see this might trap.
+ (memory 1 2 shared)
+ (data (i32.const -2) "\00\00\00")
+)
+
+;; CHECK: (memory $0 1 2 shared)
+
+;; CHECK: (data $0 (i32.const -2) "\00\00\00")
+
+;; _TNH_: (memory $0 1 2 shared)
+(module
+ ;; This segment will almost trap, but not.
+ (memory 1 2 shared)
+ (data (i32.const 65535) "\00")
+)
+
+;; CHECK: (memory $0 1 2 shared)
+
+;; _TNH_: (memory $0 1 2 shared)
+(module
+ ;; This one is slightly larger, and will trap.
+ (memory 1 2 shared)
+ (data (i32.const 65535) "\00\00")
+)
+
+;; CHECK: (memory $0 1 2 shared)
+
+;; CHECK: (data $0 (i32.const 65535) "\00\00")
+
+;; _TNH_: (memory $0 1 2 shared)
+(module
+ ;; This one is slightly larger, but the offset is lower so it will not trap.
+ (memory 1 2 shared)
+ (data (i32.const 65534) "\00\00")
+)
+
+;; CHECK: (memory $0 1 2 shared)
+
+;; _TNH_: (memory $0 1 2 shared)
+(module
+ ;; This one's offset is just large enough to trap.
+ (memory 1 2 shared)
+ (data (i32.const 65536) "\00")
+)
+
+;; CHECK: (memory $0 1 2 shared)
+
+;; CHECK: (data $0 (i32.const 65536) "\00")
+
+;; _TNH_: (memory $0 1 2 shared)
+(module
+ ;; This offset is unknown, so assume the worst.
+ ;; TODO: We could remove it in TNH mode
+
+ ;; CHECK: (import "a" "b" (global $g i32))
+ ;; _TNH_: (import "a" "b" (global $g i32))
+ (import "a" "b" (global $g i32))
+ (memory 1 2 shared)
+ (data (global.get $g) "\00")
+)
+
+;; CHECK: (memory $0 1 2 shared)
+
+;; CHECK: (data $0 (global.get $g) "\00")
+
+;; _TNH_: (memory $0 1 2 shared)
+
+;; _TNH_: (data $0 (global.get $g) "\00")
+(module
+ ;; Passive segments cannot trap during startup and are removable if they have
+ ;; no uses, like here.
+ (memory 1 2 shared)
+ (data $data "\00\00\00")
+)
+;; CHECK: (memory $0 1 2 shared)
+
+;; _TNH_: (memory $0 1 2 shared)