[multibyte] Support offset and align for array load and store

src/parser/contexts.h

@@ -578,14 +578,19 @@ struct NullInstrParserCtx {
                      MemoryOrder) {
     return Ok{};
   }
-  template<typename HeapTypeT>
-  Result<> makeArrayLoad(
-    Index, const std::vector<Annotation>&, Type, int, bool, HeapTypeT) {
+  template<typename MemargT, typename HeapTypeT>
+  Result<> makeArrayLoad(Index,
+                         const std::vector<Annotation>&,
+                         Type,
+                         int,
+                         bool,
+                         MemargT,
+                         HeapTypeT) {
     return Ok{};
   }
-  template<typename HeapTypeT>
-  Result<>
-  makeArrayStore(Index, const std::vector<Annotation>&, Type, int, HeapTypeT) {
+  template<typename MemargT, typename HeapTypeT>
+  Result<> makeArrayStore(
+    Index, const std::vector<Annotation>&, Type, int, MemargT, HeapTypeT) {
     return Ok{};
   }
   Result<> makeAtomicRMW(Index,
@@ -2388,17 +2393,23 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {
                          Type type,
                          int bytes,
                          bool signed_,
+                         Memarg memarg,
                          HeapTypeT arrayType) {
-    return withLoc(pos,
-                   irBuilder.makeArrayLoad(arrayType, bytes, signed_, type));
+    return withLoc(
+      pos,
+      irBuilder.makeArrayLoad(
+        arrayType, bytes, signed_, memarg.offset, memarg.align, type));
   }
 
   Result<> makeArrayStore(Index pos,
                           const std::vector<Annotation>& annotations,
                           Type type,
                           int bytes,
+                          Memarg memarg,
                           HeapTypeT arrayType) {
-    return withLoc(pos, irBuilder.makeArrayStore(arrayType, bytes, type));
+    return withLoc(pos,
+                   irBuilder.makeArrayStore(
+                     arrayType, bytes, memarg.offset, memarg.align, type));
   }
 
   Result<> makeAtomicRMW(Index pos,

src/parser/parsers.h

@@ -1807,22 +1807,27 @@ Result<> makeLoad(Ctx& ctx,
                   bool signed_,
                   int bytes,
                   bool isAtomic) {
+  auto mem = maybeMemidx(ctx);
+  CHECK_ERR(mem);
 
   if (ctx.in.takeSExprStart("type"sv)) {
+    if (mem) {
+      return ctx.in.err("memory index is not allowed for array load");
+    }
     auto arrayType = typeidx(ctx);
     CHECK_ERR(arrayType);
 
     if (!ctx.in.takeRParen()) {
       return ctx.in.err("expected end of type use");
     }
 
+    auto arg = memarg(ctx, bytes);
+    CHECK_ERR(arg);
+
     return ctx.makeArrayLoad(
-      pos, annotations, type, bytes, signed_, *arrayType);
+      pos, annotations, type, bytes, signed_, *arg, *arrayType);
   }
 
-  auto mem = maybeMemidx(ctx);
-  CHECK_ERR(mem);
-
   // We could only parse this when `isAtomic`, but this way gives a clearer
   // error when a memorder is given for non-atomic operations
   // since the next token can never be mistaken for a `memOrder`.
@@ -1855,18 +1860,25 @@ Result<> makeStore(Ctx& ctx,
                    Type type,
                    int bytes,
                    bool isAtomic) {
+  auto mem = maybeMemidx(ctx);
+  CHECK_ERR(mem);
+
   if (ctx.in.takeSExprStart("type"sv)) {
+    if (mem) {
+      return ctx.in.err("memory index is not allowed for array store");
+    }
     auto arrayType = typeidx(ctx);
     CHECK_ERR(arrayType);
 
     if (!ctx.in.takeRParen()) {
       return ctx.in.err("expected end of type use");
     }
 
-    return ctx.makeArrayStore(pos, annotations, type, bytes, *arrayType);
+    auto arg = memarg(ctx, bytes);
+    CHECK_ERR(arg);
+
+    return ctx.makeArrayStore(pos, annotations, type, bytes, *arg, *arrayType);
   }
-  auto mem = maybeMemidx(ctx);
-  CHECK_ERR(mem);
 
   auto maybeOrder = maybeMemOrder(ctx);
   CHECK_ERR(maybeOrder);

src/passes/DeAlign.cpp

@@ -36,6 +36,10 @@ struct DeAlign : public WalkerPass<PostWalker<DeAlign>> {
   void visitStore(Store* curr) { curr->align = 1; }
 
   void visitSIMDLoad(SIMDLoad* curr) { curr->align = 1; }
+
+  void visitArrayLoad(ArrayLoad* curr) { curr->align = 1; }
+
+  void visitArrayStore(ArrayStore* curr) { curr->align = 1; }
 };
 
 Pass* createDeAlignPass() { return new DeAlign(); }

src/passes/Print.cpp

@@ -2479,6 +2479,12 @@ struct PrintExpressionContents
     printMinor(o, "type ");
     printHeapTypeName(curr->ref->type.getHeapType());
     o << ')';
+    if (curr->offset) {
+      o << " offset=" << curr->offset;
+    }
+    if (curr->align != curr->bytes) {
+      o << " align=" << curr->align;
+    }
   }
 
   void visitArrayStore(ArrayStore* curr) {
@@ -2492,6 +2498,12 @@ struct PrintExpressionContents
     printMinor(o, "type ");
     printHeapTypeName(curr->ref->type.getHeapType());
     o << ')';
+    if (curr->offset) {
+      o << " offset=" << curr->offset;
+    }
+    if (curr->align != curr->bytes) {
+      o << " align=" << curr->align;
+    }
   }
   void visitArrayLen(ArrayLen* curr) { printMedium(o, "array.len"); }
   void visitArrayCopy(ArrayCopy* curr) {

src/wasm-builder.h

@@ -1162,12 +1162,16 @@ class Builder {
   }
   ArrayLoad* makeArrayLoad(unsigned bytes,
                            bool signed_,
+                           Address offset,
+                           Address align,
                            Expression* ref,
                            Expression* index,
                            Type type) {
     auto* ret = wasm.allocator.alloc<ArrayLoad>();
     ret->bytes = bytes;
     ret->signed_ = signed_;
+    ret->offset = offset;
+    ret->align = align ? align : Address(bytes);
     ret->ref = ref;
     ret->index = index;
     ret->type = type;
@@ -1176,11 +1180,15 @@ class Builder {
   }
 
   ArrayStore* makeArrayStore(unsigned bytes,
+                             Address offset,
+                             Address align,
                              Expression* ref,
                              Expression* index,
                              Expression* value) {
     auto* ret = wasm.allocator.alloc<ArrayStore>();
     ret->bytes = bytes;
+    ret->offset = offset;
+    ret->align = align ? align : Address(bytes);
     ret->ref = ref;
     ret->index = index;
     ret->value = value;

src/wasm-delegations-fields.def

@@ -762,6 +762,8 @@ DELEGATE_FIELD_CASE_END(ArraySet)
 DELEGATE_FIELD_CASE_START(ArrayLoad)
 DELEGATE_FIELD_CHILD(ArrayLoad, index)
 DELEGATE_FIELD_IMMEDIATE_TYPED_CHILD(ArrayLoad, ref)
+DELEGATE_FIELD_ADDRESS(ArrayLoad, offset)
+DELEGATE_FIELD_ADDRESS(ArrayLoad, align)
 DELEGATE_FIELD_INT(ArrayLoad, bytes)
 DELEGATE_FIELD_INT(ArrayLoad, signed_)
 DELEGATE_FIELD_TYPE(ArrayLoad, type)
@@ -771,6 +773,8 @@ DELEGATE_FIELD_CASE_START(ArrayStore)
 DELEGATE_FIELD_CHILD(ArrayStore, value)
 DELEGATE_FIELD_CHILD(ArrayStore, index)
 DELEGATE_FIELD_IMMEDIATE_TYPED_CHILD(ArrayStore, ref)
+DELEGATE_FIELD_ADDRESS(ArrayStore, offset)
+DELEGATE_FIELD_ADDRESS(ArrayStore, align)
 DELEGATE_FIELD_INT(ArrayStore, bytes)
 DELEGATE_FIELD_CASE_END(ArrayStore)
 

src/wasm-ir-builder.h

@@ -253,9 +253,17 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   Result<> makeArrayNewFixed(HeapType type, uint32_t arity);
   Result<> makeArrayGet(HeapType type, bool signed_, MemoryOrder order);
   Result<> makeArraySet(HeapType type, MemoryOrder order);
-  Result<>
-  makeArrayLoad(HeapType arrayType, unsigned bytes, bool signed_, Type type);
-  Result<> makeArrayStore(HeapType arrayType, unsigned bytes, Type type);
+  Result<> makeArrayLoad(HeapType arrayType,
+                         unsigned bytes,
+                         bool signed_,
+                         Address offset,
+                         Address align,
+                         Type type);
+  Result<> makeArrayStore(HeapType arrayType,
+                          unsigned bytes,
+                          Address offset,
+                          Address align,
+                          Type type);
   Result<> makeArrayLen();
   Result<> makeArrayCopy(HeapType destType, HeapType srcType);
   Result<> makeArrayFill(HeapType type);

src/wasm.h

@@ -1918,6 +1918,8 @@ class ArrayLoad : public SpecificExpression<Expression::ArrayLoadId> {
 
   uint8_t bytes;
   bool signed_ = false;
+  Address offset = 0;
+  Address align = 0;
   Expression* ref;
   Expression* index;
 
@@ -1930,6 +1932,8 @@ class ArrayStore : public SpecificExpression<Expression::ArrayStoreId> {
   ArrayStore(MixedArena& allocator) {}
 
   uint8_t bytes;
+  Address offset = 0;
+  Address align = 0;
   Expression* ref;
   Expression* index;
   Expression* value;

src/wasm/wasm-binary.cpp

@@ -3394,7 +3394,8 @@ Result<> WasmBinaryReader::readLoad(unsigned bytes, bool signed_, Type type) {
   auto [mem, align, offset, backing] = getMemarg();
   if (backing == BackingType::Array) {
     HeapType arrayType = getIndexedHeapType();
-    return builder.makeArrayLoad(arrayType, bytes, signed_, type);
+    return builder.makeArrayLoad(
+      arrayType, bytes, signed_, offset, align, type);
   }
   return builder.makeLoad(bytes, signed_, offset, align, type, mem);
 }
@@ -3403,7 +3404,7 @@ Result<> WasmBinaryReader::readStore(unsigned bytes, Type type) {
   auto [mem, align, offset, backing] = getMemarg();
   if (backing == BackingType::Array) {
     HeapType arrayType = getIndexedHeapType();
-    return builder.makeArrayStore(arrayType, bytes, type);
+    return builder.makeArrayStore(arrayType, bytes, offset, align, type);
   }
   return builder.makeStore(bytes, offset, align, type, mem);
 }
@@ -4111,12 +4112,10 @@ Result<> WasmBinaryReader::readInst() {
           return builder.makeElemDrop(elem);
         }
         case BinaryConsts::F32_F16LoadMem: {
-          auto [mem, align, offset, backing] = getMemarg();
-          return builder.makeLoad(2, false, offset, align, Type::f32, mem);
+          return readLoad(2, false, Type::f32);
         }
         case BinaryConsts::F32_F16StoreMem: {
-          auto [mem, align, offset, backing] = getMemarg();
-          return builder.makeStore(2, offset, align, Type::f32, mem);
+          return readStore(2, Type::f32);
         }
       }
       return Err{"unknown misc operation: " + std::to_string(op)};
@@ -4666,12 +4665,10 @@ Result<> WasmBinaryReader::readInst() {
         case BinaryConsts::V128Const:
           return builder.makeConst(getVec128Literal());
         case BinaryConsts::V128Store: {
-          auto [mem, align, offset, backing] = getMemarg();
-          return builder.makeStore(16, offset, align, Type::v128, mem);
+          return readStore(16, Type::v128);
         }
         case BinaryConsts::V128Load: {
-          auto [mem, align, offset, backing] = getMemarg();
-          return builder.makeLoad(16, false, offset, align, Type::v128, mem);
+          return readLoad(16, false, Type::v128);
         }
         case BinaryConsts::V128Load8Splat: {
           auto [mem, align, offset, backing] = getMemarg();
@@ -5759,6 +5756,7 @@ WasmBinaryReader::readMemoryAccess(bool isAtomic, bool isRMW) {
       throwError(
         "Memory index and memory order are not allowed for array backing.");
     }
+    offset = getU32LEB();
   } else {
     WASM_UNREACHABLE("Invalid backing type");
   }

src/wasm/wasm-ir-builder.cpp

@@ -2463,8 +2463,11 @@ Result<> IRBuilder::makeArraySet(HeapType type, MemoryOrder order) {
   return Ok{};
 }
 
-Result<>
-IRBuilder::makeArrayStore(HeapType arrayType, unsigned bytes, Type type) {
+Result<> IRBuilder::makeArrayStore(HeapType arrayType,
+                                   unsigned bytes,
+                                   Address offset,
+                                   Address align,
+                                   Type type) {
   if (!arrayType.isArray()) {
     return Err{"expected array type annotation on array store"};
   }
@@ -2473,13 +2476,16 @@ IRBuilder::makeArrayStore(HeapType arrayType, unsigned bytes, Type type) {
   CHECK_ERR(ChildPopper{*this}.visitArrayStore(&curr, arrayType, type));
 
   CHECK_ERR(validateTypeAnnotation(arrayType, curr.ref));
-  push(builder.makeArrayStore(bytes, curr.ref, curr.index, curr.value));
+  push(builder.makeArrayStore(
+    bytes, offset, align, curr.ref, curr.index, curr.value));
   return Ok{};
 }
 
 Result<> IRBuilder::makeArrayLoad(HeapType arrayType,
                                   unsigned bytes,
                                   bool signed_,
+                                  Address offset,
+                                  Address align,
                                   Type type) {
   if (!arrayType.isArray()) {
     return Err{"expected array type annotation on array load"};
@@ -2489,7 +2495,8 @@ Result<> IRBuilder::makeArrayLoad(HeapType arrayType,
   CHECK_ERR(ChildPopper{*this}.visitArrayLoad(&curr, arrayType));
 
   CHECK_ERR(validateTypeAnnotation(arrayType, curr.ref));
-  push(builder.makeArrayLoad(bytes, signed_, curr.ref, curr.index, type));
+  push(builder.makeArrayLoad(
+    bytes, signed_, offset, align, curr.ref, curr.index, type));
   return Ok{};
 }
 

src/wasm/wasm-stack.cpp

@@ -2828,24 +2828,36 @@ void BinaryInstWriter::visitArraySet(ArraySet* curr) {
 }
 
 void BinaryInstWriter::visitArrayLoad(ArrayLoad* curr) {
+  if (curr->type == Type::unreachable) {
+    return;
+  }
   if (curr->ref->type.isNull()) {
     emitUnreachable();
     return;
   }
   emitLoadOpcode(curr->bytes, curr->signed_, curr->type);
-  uint32_t alignmentBits = BinaryConsts::HasBackingArrayMask;
+  uint32_t alignmentBits =
+    Bits::log2(curr->align ? curr->align : Address(curr->bytes)) |
+    BinaryConsts::HasBackingArrayMask;
   o << U32LEB(alignmentBits);
+  o << U32LEB(curr->offset);
   parent.writeIndexedHeapType(curr->ref->type.getHeapType());
 }
 
 void BinaryInstWriter::visitArrayStore(ArrayStore* curr) {
+  if (curr->type == Type::unreachable) {
+    return;
+  }
   if (curr->ref->type.isNull()) {
     emitUnreachable();
     return;
   }
   emitStoreOpcode(curr->bytes, curr->value->type);
-  uint32_t alignmentBits = BinaryConsts::HasBackingArrayMask;
+  uint32_t alignmentBits =
+    Bits::log2(curr->align ? curr->align : Address(curr->bytes)) |
+    BinaryConsts::HasBackingArrayMask;
   o << U32LEB(alignmentBits);
+  o << U32LEB(curr->offset);
   parent.writeIndexedHeapType(curr->ref->type.getHeapType());
 }