From 79b3a972399644a3652b7dc65f166d7d7f999ee3 Mon Sep 17 00:00:00 2001
From: Thomas Lively <tlively@google.com>
Date: Wed, 28 Feb 2024 17:34:49 -0800
Subject: [PATCH] [Parser] Support inline data in 64-bit memory declarations

This new form of the abbreviated memory declaration with inline data is
introduced in the memory64 proposal.
---
 src/parser/parsers.h           | 31 ++++++++++++++++++++++++-------
 test/lit/wat-kitchen-sink.wast | 12 +++++++++++-
 2 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/src/parser/parsers.h b/src/parser/parsers.h
index 4c210624fe5..4a2b99d8459 100644
--- a/src/parser/parsers.h
+++ b/src/parser/parsers.h
@@ -41,6 +41,8 @@ template<typename Ctx> MaybeResult<typename Ctx::ArrayT> arraytype(Ctx&);
 template<typename Ctx> Result<typename Ctx::LimitsT> limits32(Ctx&);
 template<typename Ctx> Result<typename Ctx::LimitsT> limits64(Ctx&);
 template<typename Ctx> Result<typename Ctx::MemTypeT> memtype(Ctx&);
+template<typename Ctx>
+Result<typename Ctx::MemTypeT> memtypeContinued(Ctx&, Type indexType);
 template<typename Ctx> Result<typename Ctx::TableTypeT> tabletype(Ctx&);
 template<typename Ctx> Result<typename Ctx::GlobalTypeT> globaltype(Ctx&);
 template<typename Ctx> Result<uint32_t> tupleArity(Ctx&);
@@ -770,20 +772,28 @@ template<typename Ctx> Result<typename Ctx::LimitsT> limits64(Ctx& ctx) {
 }
 
 // memtype ::= (limits32 | 'i32' limits32 | 'i64' limit64) shared?
+//  note: the index type 'i32' or 'i64' is already parsed to simplify parsing of
+//  memory abbreviations.
 template<typename Ctx> Result<typename Ctx::MemTypeT> memtype(Ctx& ctx) {
-  auto type = Type::i32;
+  Type indexType = Type::i32;
   if (ctx.in.takeKeyword("i64"sv)) {
-    type = Type::i64;
+    indexType = Type::i64;
   } else {
     ctx.in.takeKeyword("i32"sv);
   }
-  auto limits = type == Type::i32 ? limits32(ctx) : limits64(ctx);
+  return memtypeContinued(ctx, indexType);
+}
+
+template<typename Ctx>
+Result<typename Ctx::MemTypeT> memtypeContinued(Ctx& ctx, Type indexType) {
+  assert(indexType == Type::i32 || indexType == Type::i64);
+  auto limits = indexType == Type::i32 ? limits32(ctx) : limits64(ctx);
   CHECK_ERR(limits);
   bool shared = false;
   if (ctx.in.takeKeyword("shared"sv)) {
     shared = true;
   }
-  return ctx.makeMemType(type, *limits, shared);
+  return ctx.makeMemType(indexType, *limits, shared);
 }
 
 // tabletype ::= limits32 reftype
@@ -3058,7 +3068,7 @@ template<typename Ctx> MaybeResult<> table(Ctx& ctx) {
   return Ok{};
 }
 
-// mem ::= '(' 'memory' id? ('(' 'export' name ')')*
+// mem ::= '(' 'memory' id? ('(' 'export' name ')')* index_type?
 //             ('(' 'data' b:datastring ')' | memtype) ')'
 //       | '(' 'memory' id? ('(' 'export' name ')')*
 //             '(' 'import' mod:name nm:name ')' memtype ')'
@@ -3079,6 +3089,13 @@ template<typename Ctx> MaybeResult<> memory(Ctx& ctx) {
   auto import = inlineImport(ctx.in);
   CHECK_ERR(import);
 
+  auto indexType = Type::i32;
+  if (ctx.in.takeKeyword("i64"sv)) {
+    indexType = Type::i64;
+  } else {
+    ctx.in.takeKeyword("i32"sv);
+  }
+
   std::optional<typename Ctx::MemTypeT> mtype;
   std::optional<typename Ctx::DataStringT> data;
   if (ctx.in.takeSExprStart("data"sv)) {
@@ -3090,10 +3107,10 @@ template<typename Ctx> MaybeResult<> memory(Ctx& ctx) {
     if (!ctx.in.takeRParen()) {
       return ctx.in.err("expected end of inline data");
     }
-    mtype = ctx.makeMemType(Type::i32, ctx.getLimitsFromData(*datastr), false);
+    mtype = ctx.makeMemType(indexType, ctx.getLimitsFromData(*datastr), false);
     data = *datastr;
   } else {
-    auto type = memtype(ctx);
+    auto type = memtypeContinued(ctx, indexType);
     CHECK_ERR(type);
     mtype = *type;
   }
diff --git a/test/lit/wat-kitchen-sink.wast b/test/lit/wat-kitchen-sink.wast
index ab81fa4c764..3476195fdc6 100644
--- a/test/lit/wat-kitchen-sink.wast
+++ b/test/lit/wat-kitchen-sink.wast
@@ -339,10 +339,20 @@
  ;; CHECK:      (memory $mem-init 1 1)
  (memory $mem-init (data "hello inline data"))
 
+ ;; CHECK:      (memory $mem-init-32 1 1)
+ (memory $mem-init-32 i32 (data "hello i32 inline data"))
+
+ ;; CHECK:      (memory $mem-init-64 i64 1 1)
+ (memory $mem-init-64 i64 (data "hello i64 inline data"))
+
  ;; data segments
  (data "hello world")
  ;; CHECK:      (data $implicit-data (memory $mem-init) (i32.const 0) "hello inline data")
 
+ ;; CHECK:      (data $implicit-data_1 (memory $mem-init-32) (i32.const 0) "hello i32 inline data")
+
+ ;; CHECK:      (data $implicit-data_2 (memory $mem-init-64) (i64.const 0) "hello i64 inline data")
+
  ;; CHECK:      (data $0 "hello world")
 
  ;; CHECK:      (data $passive "hello again")
@@ -3500,7 +3510,7 @@
   i64.const 0
   local.get 1
   local.get 2
-  memory.init 5 1
+  memory.init 5 3
   local.get 0
   local.get 1
   local.get 2