[Parser] Parse `local.get` (#5137)

This requires parsing local indices and fixing a bug in `Function::setLocalName`
where it only set up the mapping from index to name and not the mapping from
name to index.

3 files changed, 146 insertions, 48 deletions

diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp
index 27690f43e..6347d83e5 100644
--- a/src/wasm/wasm.cpp
+++ b/src/wasm/wasm.cpp
@@ -1273,6 +1273,7 @@ Name Function::getLocalName(Index index) { return localNames.at(index); }
 void Function::setLocalName(Index index, Name name) {
   assert(index < getNumLocals());
   localNames[index] = name;
+  localIndices[name] = index;
 }
 
 Name Function::getLocalNameOrDefault(Index index) {
diff --git a/src/wasm/wat-parser.cpp b/src/wasm/wat-parser.cpp
index 0928de094..bba751c71 100644
--- a/src/wasm/wat-parser.cpp
+++ b/src/wasm/wat-parser.cpp
@@ -554,12 +554,17 @@ struct NullInstrParserCtx {
   using InstrsT = Ok;
   using ExprT = Ok;
 
+  using LocalT = Ok;
+
   InstrsT makeInstrs() { return Ok{}; }
   void appendInstr(InstrsT&, InstrT) {}
   InstrsT finishInstrs(InstrsT&) { return Ok{}; }
 
   ExprT makeExpr(InstrsT) { return Ok{}; }
 
+  LocalT getLocalFromIdx(uint32_t idx) { return Ok{}; }
+  LocalT getLocalFromName(Name name) { return Ok{}; }
+
   InstrT makeUnreachable(Index pos) { return Ok{}; }
   InstrT makeNop(Index pos) { return Ok{}; }
   InstrT makeBinary(Index pos, BinaryOp op) { return Ok{}; }
@@ -574,6 +579,8 @@ struct NullInstrParserCtx {
   InstrT makeF32Const(Index, float) { return Ok{}; }
   InstrT makeF64Const(Index, double) { return Ok{}; }
 
+  InstrT makeLocalGet(Index pos, LocalT local) { return Ok{}; }
+
   template<typename HeapTypeT> InstrT makeRefNull(Index, HeapTypeT) {
     return {};
   }
@@ -586,6 +593,8 @@ template<typename Ctx> struct InstrParserCtx : TypeParserCtx<Ctx> {
   using InstrsT = std::vector<Expression*>;
   using ExprT = Expression*;
 
+  using LocalT = Index;
+
   Builder builder;
 
   // The stack of parsed expressions, used as the children of newly parsed
@@ -756,7 +765,7 @@ template<typename Ctx> struct InstrParserCtx : TypeParserCtx<Ctx> {
   Result<> makeF64Const(Index pos, double c) {
     return push(pos, builder.makeConst(Literal(c)));
   }
-  Result<> makeRefNull(Index pos, typename TypeParserCtx<Ctx>::HeapTypeT type) {
+  Result<> makeRefNull(Index pos, HeapType type) {
     return push(pos, builder.makeRefNull(type));
   }
 };
@@ -1097,7 +1106,6 @@ struct ParseModuleTypesCtx : TypeParserCtx<ParseModuleTypesCtx>,
         Builder::addVar(f.get(), l.name, l.type);
       }
     }
-    // TODO: local types and names.
     return Ok{};
   }
 
@@ -1150,6 +1158,17 @@ struct ParseDefsCtx : InstrParserCtx<ParseDefsCtx> {
     return types[idx];
   }
 
+  Index getLocalFromIdx(uint32_t idx) { return idx; }
+  Result<Index> getLocalFromName(Name name) {
+    if (!func) {
+      return in.err("cannot access locals outside of a function");
+    }
+    if (!func->hasLocalIndex(name)) {
+      return in.err("local $" + name.toString() + " does not exist");
+    }
+    return func->getLocalIndex(name);
+  }
+
   Result<TypeUseT> makeTypeUse(Index pos,
                                std::optional<HeapTypeT> type,
                                ParamsT* params,
@@ -1228,6 +1247,14 @@ struct ParseDefsCtx : InstrParserCtx<ParseDefsCtx> {
     Name name = Names::getValidLocalName(*func, "scratch");
     return Builder::addVar(func, name, type);
   }
+
+  Result<> makeLocalGet(Index pos, Index local) {
+    if (!func) {
+      return in.err(pos, "local.get must be inside a function");
+    }
+    assert(local < func->getNumLocals());
+    return push(pos, builder.makeLocalGet(local, func->getLocalType(local)));
+  }
 };
 
 // ================
@@ -1406,6 +1433,7 @@ Result<typename Ctx::InstrT> makeStringSliceIter(Ctx&, Index);
 // Modules
 template<typename Ctx> MaybeResult<Index> maybeTypeidx(Ctx& ctx);
 template<typename Ctx> Result<typename Ctx::HeapTypeT> typeidx(Ctx&);
+template<typename Ctx> Result<typename Ctx::LocalT> localidx(Ctx&);
 template<typename Ctx> Result<typename Ctx::TypeUseT> typeuse(Ctx&);
 MaybeResult<ImportNames> inlineImport(ParseInput&);
 Result<std::vector<Name>> inlineExports(ParseInput&);
@@ -1843,7 +1871,9 @@ Result<typename Ctx::InstrT> makeMemoryGrow(Ctx& ctx, Index pos) {
 
 template<typename Ctx>
 Result<typename Ctx::InstrT> makeLocalGet(Ctx& ctx, Index pos) {
-  return ctx.in.err("unimplemented instruction");
+  auto local = localidx(ctx);
+  CHECK_ERR(local);
+  return ctx.makeLocalGet(pos, *local);
 }
 
 template<typename Ctx>
@@ -2338,6 +2368,18 @@ template<typename Ctx> Result<typename Ctx::HeapTypeT> typeidx(Ctx& ctx) {
   return ctx.in.err("expected type index or identifier");
 }
 
+// localidx ::= x:u32 => x
+//            | v:id  => x (if types[x] = v)
+template<typename Ctx> Result<typename Ctx::LocalT> localidx(Ctx& ctx) {
+  if (auto x = ctx.in.takeU32()) {
+    return ctx.getLocalFromIdx(*x);
+  }
+  if (auto id = ctx.in.takeID()) {
+    return ctx.getLocalFromName(*id);
+  }
+  return ctx.in.err("expected local index or identifier");
+}
+
 // typeuse ::= '(' 'type' x:typeidx ')'                                => x, []
 //                 (if typedefs[x] = [t1*] -> [t2*]
 //           | '(' 'type' x:typeidx ')' ((t1,IDs):param)* (t2:result)* => x, IDs
@@ -2364,7 +2406,6 @@ template<typename Ctx> Result<typename Ctx::TypeUseT> typeuse(Ctx& ctx) {
   auto resultTypes = results(ctx);
   CHECK_ERR(resultTypes);
 
-  // TODO: Use `pos` for error reporting rather than `in`.
   return ctx.makeTypeUse(pos, type, namedParams.getPtr(), resultTypes.getPtr());
 }
 
diff --git a/test/lit/wat-kitchen-sink.wast b/test/lit/wat-kitchen-sink.wast
index 62448a02e..1482e87f6 100644
--- a/test/lit/wat-kitchen-sink.wast
+++ b/test/lit/wat-kitchen-sink.wast
@@ -17,6 +17,14 @@
 
   ;; CHECK:      (type $many (func_subtype (param i32 i64 f32 f64) (result anyref (ref func)) func))
 
+  ;; CHECK:      (type $i32_i32_=>_none (func_subtype (param i32 i32) func))
+
+  ;; CHECK:      (type $i32_i32_f64_f64_=>_none (func_subtype (param i32 i32 f64 f64) func))
+
+  ;; CHECK:      (type $i64_=>_none (func_subtype (param i64) func))
+
+  ;; CHECK:      (type $i32_i32_i32_=>_none (func_subtype (param i32 i32 i32) func))
+
   ;; CHECK:      (rec
   ;; CHECK-NEXT:  (type $s0 (struct_subtype  data))
   (type $s0 (sub (struct)))
@@ -551,85 +559,133 @@
   drop
  )
 
- ;; CHECK:      (func $use-insts (type $void)
- ;; CHECK-NEXT:  (nop)
+ ;; CHECK:      (func $locals (type $i32_i32_=>_none) (param $0 i32) (param $x i32)
+ ;; CHECK-NEXT:  (local $2 i32)
+ ;; CHECK-NEXT:  (local $y i32)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $0)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $x)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $2)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $y)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $x)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $y)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $locals (param i32) (param $x i32)
+  (local i32)
+  (local $y i32)
+  local.get 0
+  drop
+  local.get 1
+  drop
+  local.get 2
+  drop
+  local.get 3
+  drop
+  local.get $x
+  drop
+  local.get $y
+  drop
+ )
+
+
+ ;; CHECK:      (func $binary (type $i32_i32_f64_f64_=>_none) (param $0 i32) (param $1 i32) (param $2 f64) (param $3 f64)
  ;; CHECK-NEXT:  (drop
  ;; CHECK-NEXT:   (i32.add
- ;; CHECK-NEXT:    (i32.const 0)
- ;; CHECK-NEXT:    (i32.const 1)
+ ;; CHECK-NEXT:    (local.get $0)
+ ;; CHECK-NEXT:    (local.get $1)
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
  ;; CHECK-NEXT:  (drop
  ;; CHECK-NEXT:   (f64.mul
- ;; CHECK-NEXT:    (f64.const 0)
- ;; CHECK-NEXT:    (f64.const 1)
+ ;; CHECK-NEXT:    (local.get $2)
+ ;; CHECK-NEXT:    (local.get $3)
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $binary (param i32 i32 f64 f64)
+  local.get 0
+  local.get 1
+  i32.add
+  drop
+  local.get 2
+  local.get 3
+  f64.mul
+  drop
+ )
+
+ ;; CHECK:      (func $unary (type $i64_=>_none) (param $0 i64)
  ;; CHECK-NEXT:  (drop
  ;; CHECK-NEXT:   (i64.eqz
- ;; CHECK-NEXT:    (i64.const 0)
+ ;; CHECK-NEXT:    (local.get $0)
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $unary (param i64)
+  local.get 0
+  i64.eqz
+  drop
+ )
+
+ ;; CHECK:      (func $select (type $i32_i32_i32_=>_none) (param $0 i32) (param $1 i32) (param $2 i32)
  ;; CHECK-NEXT:  (drop
  ;; CHECK-NEXT:   (select
- ;; CHECK-NEXT:    (i32.const 0)
- ;; CHECK-NEXT:    (i32.const 1)
- ;; CHECK-NEXT:    (i32.const 2)
+ ;; CHECK-NEXT:    (local.get $0)
+ ;; CHECK-NEXT:    (local.get $1)
+ ;; CHECK-NEXT:    (local.get $2)
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
  ;; CHECK-NEXT:  (drop
  ;; CHECK-NEXT:   (select
- ;; CHECK-NEXT:    (i32.const 0)
- ;; CHECK-NEXT:    (i32.const 1)
- ;; CHECK-NEXT:    (i32.const 2)
+ ;; CHECK-NEXT:    (local.get $0)
+ ;; CHECK-NEXT:    (local.get $1)
+ ;; CHECK-NEXT:    (local.get $2)
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
  ;; CHECK-NEXT:  (drop
  ;; CHECK-NEXT:   (select
- ;; CHECK-NEXT:    (i32.const 0)
- ;; CHECK-NEXT:    (i32.const 1)
- ;; CHECK-NEXT:    (i32.const 2)
+ ;; CHECK-NEXT:    (local.get $0)
+ ;; CHECK-NEXT:    (local.get $1)
+ ;; CHECK-NEXT:    (local.get $2)
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
  ;; CHECK-NEXT:  (drop
  ;; CHECK-NEXT:   (select
- ;; CHECK-NEXT:    (i32.const 0)
- ;; CHECK-NEXT:    (i32.const 1)
- ;; CHECK-NEXT:    (i32.const 2)
+ ;; CHECK-NEXT:    (local.get $0)
+ ;; CHECK-NEXT:    (local.get $1)
+ ;; CHECK-NEXT:    (local.get $2)
  ;; CHECK-NEXT:   )
  ;; CHECK-NEXT:  )
  ;; CHECK-NEXT: )
- (func $use-insts
-  nop
-  i32.const 0
-  i32.const 1
-  i32.add
-  drop
-  f64.const 0
-  f64.const 1
-  f64.mul
-  drop
-  i64.const 0
-  i64.eqz
-  drop
-  i32.const 0
-  i32.const 1
-  i32.const 2
+ (func $select (param i32 i32 i32)
+  local.get 0
+  local.get 1
+  local.get 2
   select
   drop
-  i32.const 0
-  i32.const 1
-  i32.const 2
+  local.get 0
+  local.get 1
+  local.get 2
   select (result)
   drop
-  i32.const 0
-  i32.const 1
-  i32.const 2
+  local.get 0
+  local.get 1
+  local.get 2
   select (result i32)
   drop
-  i32.const 0
-  i32.const 1
-  i32.const 2
+  local.get 0
+  local.get 1
+  local.get 2
   select (result) (result i64) (result)
   drop
  )