From dea1a8e379dca787c4b9fe19a1baf9b0ed45c6be Mon Sep 17 00:00:00 2001
From: Schuyler Eldridge <schuyler.eldridge@sifive.com>
Date: Fri, 12 Jan 2024 18:35:03 -0500
Subject: [PATCH] fixup! [FIRRTL] Parse Layer-Associated Probes

---
 lib/Dialect/FIRRTL/Import/FIRParser.cpp | 8 +++++---
 test/Dialect/FIRRTL/parse-errors.fir    | 8 ++++++++
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/lib/Dialect/FIRRTL/Import/FIRParser.cpp b/lib/Dialect/FIRRTL/Import/FIRParser.cpp
index 7b423d3bc29d..36141a15cb82 100644
--- a/lib/Dialect/FIRRTL/Import/FIRParser.cpp
+++ b/lib/Dialect/FIRRTL/Import/FIRParser.cpp
@@ -942,14 +942,16 @@ ParseResult FIRParser::parseType(FIRRTLType &result, const Twine &message) {
       return failure();
     while (getToken().getKind() == FIRToken::identifier) {
       StringRef layer;
-      if (parseId(layer, "expected layer name"))
-        return failure();
+      loc = getToken().getLoc();
+      (void)parseId(layer, "expected layer name");
       layers.push_back(layer);
       if (getToken().getKind() == FIRToken::period)
         consumeToken();
     }
-    if (parseToken(FIRToken::greater, "expected '>' in reference type"))
+    if (!consumeIf(FIRToken::greater)) {
+      emitError(loc, "expected '>' to end reference type");
       return failure();
+    }
 
     bool forceable = kind == FIRToken::kw_RWProbe;
 
diff --git a/test/Dialect/FIRRTL/parse-errors.fir b/test/Dialect/FIRRTL/parse-errors.fir
index 5e63b774cd48..cd7e83760788 100644
--- a/test/Dialect/FIRRTL/parse-errors.fir
+++ b/test/Dialect/FIRRTL/parse-errors.fir
@@ -1246,3 +1246,11 @@ circuit Foo:
   option Platform:
     FPGA
     FPGA
+
+;// -----
+FIRRTL version 4.0.0
+circuit Foo:
+  layer A bind:
+  module Foo:
+    ; expected-error @below {{expected '>' to end reference type}}
+    output a: Probe<UInt<1>