Add memory: init, copy, fill support to Multi-Memory Lowering Pass (#5346)

This PR adds support for memory.init, memory.copy, and memory.fill instructions in the multi-memory lowering pass. Also includes optional bounds checks per the wasm spec guidelines.

2 files changed, 291 insertions, 4 deletions

diff --git a/src/passes/MultiMemoryLowering.cpp b/src/passes/MultiMemoryLowering.cpp
index 1183a7d74..38a5ded93 100644
--- a/src/passes/MultiMemoryLowering.cpp
+++ b/src/passes/MultiMemoryLowering.cpp
@@ -170,6 +170,17 @@ struct MultiMemoryLowering : public Pass {
       return boundsCheck;
     }
 
+    Expression* makeDataSegmentBoundsCheck(MemoryInit* curr,
+                                           Index sizeIdx,
+                                           Index offsetIdx) {
+      auto& segment = parent.wasm->dataSegments[curr->segment];
+      Expression* addGtuTrap = makeAddGtuTrap(
+        builder.makeLocalGet(offsetIdx, parent.pointerType),
+        builder.makeLocalGet(sizeIdx, parent.pointerType),
+        builder.makeConstPtr(segment->data.size(), parent.pointerType));
+      return addGtuTrap;
+    }
+
     template<typename T> Expression* getPtr(T* curr, Index bytes) {
       Expression* ptrValue = addOffsetGlobal(curr->ptr, curr->memory);
       if (parent.checkBounds) {
@@ -183,6 +194,107 @@ struct MultiMemoryLowering : public Pass {
       return ptrValue;
     }
 
+    template<typename T>
+    Expression* getDest(T* curr,
+                        Name memory,
+                        Index sizeIdx = Index(-1),
+                        Expression* localSet = nullptr,
+                        Expression* additionalCheck = nullptr) {
+      Expression* destValue = addOffsetGlobal(curr->dest, memory);
+
+      if (parent.checkBounds) {
+        Expression* sizeSet = builder.makeLocalSet(sizeIdx, curr->size);
+        Index destIdx = Builder::addVar(getFunction(), parent.pointerType);
+        Expression* destSet = builder.makeLocalSet(destIdx, destValue);
+        Expression* boundsCheck = makeAddGtuMemoryTrap(
+          builder.makeLocalGet(destIdx, parent.pointerType),
+          builder.makeLocalGet(sizeIdx, parent.pointerType),
+          memory);
+        std::vector<Expression*> exprs = {
+          destSet, localSet, sizeSet, boundsCheck};
+        if (additionalCheck) {
+          exprs.push_back(additionalCheck);
+        }
+        Expression* destGet = builder.makeLocalGet(destIdx, parent.pointerType);
+        exprs.push_back(destGet);
+        return builder.makeBlock(exprs);
+      }
+
+      return destValue;
+    }
+
+    Expression* getSource(MemoryCopy* curr,
+                          Index sizeIdx = Index(-1),
+                          Index sourceIdx = Index(-1)) {
+      Expression* sourceValue =
+        addOffsetGlobal(curr->source, curr->sourceMemory);
+
+      if (parent.checkBounds) {
+        Expression* boundsCheck = makeAddGtuMemoryTrap(
+          builder.makeLocalGet(sourceIdx, parent.pointerType),
+          builder.makeLocalGet(sizeIdx, parent.pointerType),
+          curr->sourceMemory);
+        Expression* sourceGet =
+          builder.makeLocalGet(sourceIdx, parent.pointerType);
+        std::vector<Expression*> exprs = {boundsCheck, sourceGet};
+        return builder.makeBlock(exprs);
+      }
+
+      return sourceValue;
+    }
+
+    void visitMemoryInit(MemoryInit* curr) {
+      if (parent.checkBounds) {
+        Index offsetIdx = Builder::addVar(getFunction(), parent.pointerType);
+        Index sizeIdx = Builder::addVar(getFunction(), parent.pointerType);
+        curr->dest =
+          getDest(curr,
+                  curr->memory,
+                  sizeIdx,
+                  builder.makeLocalSet(offsetIdx, curr->offset),
+                  makeDataSegmentBoundsCheck(curr, sizeIdx, offsetIdx));
+        curr->offset = builder.makeLocalGet(offsetIdx, parent.pointerType);
+        curr->size = builder.makeLocalGet(sizeIdx, parent.pointerType);
+      } else {
+        curr->dest = getDest(curr, curr->memory);
+      }
+      setMemory(curr);
+    }
+
+    void visitMemoryCopy(MemoryCopy* curr) {
+      if (parent.checkBounds) {
+        Index sourceIdx = Builder::addVar(getFunction(), parent.pointerType);
+        Index sizeIdx = Builder::addVar(getFunction(), parent.pointerType);
+        curr->dest = getDest(curr,
+                             curr->destMemory,
+                             sizeIdx,
+                             builder.makeLocalSet(sourceIdx, curr->source));
+        curr->source = getSource(curr, sizeIdx, sourceIdx);
+        curr->size = builder.makeLocalGet(sizeIdx, parent.pointerType);
+      } else {
+        curr->dest = getDest(curr, curr->destMemory);
+        curr->source = getSource(curr);
+      }
+      curr->destMemory = parent.combinedMemory;
+      curr->sourceMemory = parent.combinedMemory;
+    }
+
+    void visitMemoryFill(MemoryFill* curr) {
+      if (parent.checkBounds) {
+        Index valueIdx = Builder::addVar(getFunction(), parent.pointerType);
+        Index sizeIdx = Builder::addVar(getFunction(), parent.pointerType);
+        curr->dest = getDest(curr,
+                             curr->memory,
+                             sizeIdx,
+                             builder.makeLocalSet(valueIdx, curr->value));
+        curr->value = builder.makeLocalGet(valueIdx, parent.pointerType);
+        curr->size = builder.makeLocalGet(sizeIdx, parent.pointerType);
+      } else {
+        curr->dest = getDest(curr, curr->memory);
+      }
+      setMemory(curr);
+    }
+
     template<typename T> void setMemory(T* curr) {
       curr->memory = parent.combinedMemory;
     }
diff --git a/test/lit/passes/multi-memory-lowering.wast b/test/lit/passes/multi-memory-lowering.wast
index 6cac24ee0..53d50fa0f 100644
--- a/test/lit/passes/multi-memory-lowering.wast
+++ b/test/lit/passes/multi-memory-lowering.wast
@@ -8,14 +8,14 @@
   (memory $memory3 3)
   (data (memory $memory1) (i32.const 0) "a")
   (data (memory $memory3) (i32.const 1) "123")
+  ;; CHECK:      (type $none_=>_none (func))
+
   ;; CHECK:      (type $i32_=>_v128 (func (param i32) (result v128)))
 
   ;; CHECK:      (type $none_=>_i32 (func (result i32)))
 
   ;; CHECK:      (type $i32_=>_i32 (func (param i32) (result i32)))
 
-  ;; CHECK:      (type $none_=>_none (func))
-
   ;; CHECK:      (type $i32_v128_=>_v128 (func (param i32 v128) (result v128)))
 
   ;; CHECK:      (type $i32_i64_=>_none (func (param i32 i64)))
@@ -56,14 +56,14 @@
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
+  ;; BOUNDS:      (type $none_=>_none (func))
+
   ;; BOUNDS:      (type $i32_=>_v128 (func (param i32) (result v128)))
 
   ;; BOUNDS:      (type $none_=>_i32 (func (result i32)))
 
   ;; BOUNDS:      (type $i32_=>_i32 (func (param i32) (result i32)))
 
-  ;; BOUNDS:      (type $none_=>_none (func))
-
   ;; BOUNDS:      (type $i32_v128_=>_v128 (func (param i32 v128) (result v128)))
 
   ;; BOUNDS:      (type $i32_i64_=>_none (func (param i32 i64)))
@@ -628,6 +628,181 @@
    )
   )
   )
+
+  ;; CHECK:      (func $memory.fill
+  ;; CHECK-NEXT:  (memory.fill
+  ;; CHECK-NEXT:   (i32.const 0)
+  ;; CHECK-NEXT:   (i32.const 1)
+  ;; CHECK-NEXT:   (i32.const 2)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; BOUNDS:      (func $memory.fill
+  ;; BOUNDS-NEXT:  (local $0 i32)
+  ;; BOUNDS-NEXT:  (local $1 i32)
+  ;; BOUNDS-NEXT:  (local $2 i32)
+  ;; BOUNDS-NEXT:  (memory.fill
+  ;; BOUNDS-NEXT:   (block (result i32)
+  ;; BOUNDS-NEXT:    (local.set $2
+  ;; BOUNDS-NEXT:     (i32.const 0)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.set $0
+  ;; BOUNDS-NEXT:     (i32.const 1)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.set $1
+  ;; BOUNDS-NEXT:     (i32.const 2)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (if
+  ;; BOUNDS-NEXT:     (i32.gt_u
+  ;; BOUNDS-NEXT:      (i32.add
+  ;; BOUNDS-NEXT:       (local.get $2)
+  ;; BOUNDS-NEXT:       (local.get $1)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:      (call $memory1_size)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (unreachable)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.get $2)
+  ;; BOUNDS-NEXT:   )
+  ;; BOUNDS-NEXT:   (local.get $0)
+  ;; BOUNDS-NEXT:   (local.get $1)
+  ;; BOUNDS-NEXT:  )
+  ;; BOUNDS-NEXT: )
+  (func $memory.fill
+    (memory.fill $memory1
+      (i32.const 0)
+      (i32.const 1)
+      (i32.const 2)
+    )
+  )
+
+  ;; CHECK:      (func $memory.copy
+  ;; CHECK-NEXT:  (memory.copy
+  ;; CHECK-NEXT:   (i32.add
+  ;; CHECK-NEXT:    (global.get $memory2_byte_offset)
+  ;; CHECK-NEXT:    (i32.const 512)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (i32.add
+  ;; CHECK-NEXT:    (global.get $memory3_byte_offset)
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (i32.const 12)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; BOUNDS:      (func $memory.copy
+  ;; BOUNDS-NEXT:  (local $0 i32)
+  ;; BOUNDS-NEXT:  (local $1 i32)
+  ;; BOUNDS-NEXT:  (local $2 i32)
+  ;; BOUNDS-NEXT:  (memory.copy
+  ;; BOUNDS-NEXT:   (block (result i32)
+  ;; BOUNDS-NEXT:    (local.set $2
+  ;; BOUNDS-NEXT:     (i32.add
+  ;; BOUNDS-NEXT:      (global.get $memory2_byte_offset)
+  ;; BOUNDS-NEXT:      (i32.const 512)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.set $0
+  ;; BOUNDS-NEXT:     (i32.const 0)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.set $1
+  ;; BOUNDS-NEXT:     (i32.const 12)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (if
+  ;; BOUNDS-NEXT:     (i32.gt_u
+  ;; BOUNDS-NEXT:      (i32.add
+  ;; BOUNDS-NEXT:       (local.get $2)
+  ;; BOUNDS-NEXT:       (local.get $1)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:      (call $memory2_size)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (unreachable)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.get $2)
+  ;; BOUNDS-NEXT:   )
+  ;; BOUNDS-NEXT:   (block (result i32)
+  ;; BOUNDS-NEXT:    (if
+  ;; BOUNDS-NEXT:     (i32.gt_u
+  ;; BOUNDS-NEXT:      (i32.add
+  ;; BOUNDS-NEXT:       (local.get $0)
+  ;; BOUNDS-NEXT:       (local.get $1)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:      (call $memory3_size)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (unreachable)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.get $0)
+  ;; BOUNDS-NEXT:   )
+  ;; BOUNDS-NEXT:   (local.get $1)
+  ;; BOUNDS-NEXT:  )
+  ;; BOUNDS-NEXT: )
+  (func $memory.copy
+    (memory.copy $memory2 $memory3
+      (i32.const 512)
+      (i32.const 0)
+      (i32.const 12)
+    )
+  )
+
+  ;; CHECK:      (func $memory.init
+  ;; CHECK-NEXT:  (memory.init 0
+  ;; CHECK-NEXT:   (i32.add
+  ;; CHECK-NEXT:    (global.get $memory2_byte_offset)
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (i32.const 1)
+  ;; CHECK-NEXT:   (i32.const 45)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  ;; BOUNDS:      (func $memory.init
+  ;; BOUNDS-NEXT:  (local $0 i32)
+  ;; BOUNDS-NEXT:  (local $1 i32)
+  ;; BOUNDS-NEXT:  (local $2 i32)
+  ;; BOUNDS-NEXT:  (memory.init 0
+  ;; BOUNDS-NEXT:   (block (result i32)
+  ;; BOUNDS-NEXT:    (local.set $2
+  ;; BOUNDS-NEXT:     (i32.add
+  ;; BOUNDS-NEXT:      (global.get $memory2_byte_offset)
+  ;; BOUNDS-NEXT:      (i32.const 0)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.set $0
+  ;; BOUNDS-NEXT:     (i32.const 1)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.set $1
+  ;; BOUNDS-NEXT:     (i32.const 45)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (if
+  ;; BOUNDS-NEXT:     (i32.gt_u
+  ;; BOUNDS-NEXT:      (i32.add
+  ;; BOUNDS-NEXT:       (local.get $2)
+  ;; BOUNDS-NEXT:       (local.get $1)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:      (call $memory2_size)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (unreachable)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (if
+  ;; BOUNDS-NEXT:     (i32.gt_u
+  ;; BOUNDS-NEXT:      (i32.add
+  ;; BOUNDS-NEXT:       (local.get $0)
+  ;; BOUNDS-NEXT:       (local.get $1)
+  ;; BOUNDS-NEXT:      )
+  ;; BOUNDS-NEXT:      (i32.const 1)
+  ;; BOUNDS-NEXT:     )
+  ;; BOUNDS-NEXT:     (unreachable)
+  ;; BOUNDS-NEXT:    )
+  ;; BOUNDS-NEXT:    (local.get $2)
+  ;; BOUNDS-NEXT:   )
+  ;; BOUNDS-NEXT:   (local.get $0)
+  ;; BOUNDS-NEXT:   (local.get $1)
+  ;; BOUNDS-NEXT:  )
+  ;; BOUNDS-NEXT: )
+  (func $memory.init
+    (memory.init $memory2 0
+      (i32.const 0)
+      (i32.const 1)
+      (i32.const 45)
+    )
+  )
 )
 
 ;; CHECK:      (func $memory1_size (result i32)