From bb7bd19f8288a3c054b9d7efef273f332356d6cd Mon Sep 17 00:00:00 2001
From: Schuyler Eldridge <schuyler.eldridge@sifive.com>
Date: Thu, 4 Jan 2024 18:01:40 -0500
Subject: [PATCH] [FIRRTL] Parse Layer-Associated Probes

Add support to the FIRRTL parse to handle probes associated with layers.

Signed-off-by: Schuyler Eldridge <schuyler.eldridge@sifive.com>
---
 lib/Dialect/FIRRTL/Import/FIRParser.cpp | 26 ++++++++++++++++++++++---
 test/Dialect/FIRRTL/parse-basic.fir     |  6 +++++-
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/lib/Dialect/FIRRTL/Import/FIRParser.cpp b/lib/Dialect/FIRRTL/Import/FIRParser.cpp
index c47bc0d9448e..7b423d3bc29d 100644
--- a/lib/Dialect/FIRRTL/Import/FIRParser.cpp
+++ b/lib/Dialect/FIRRTL/Import/FIRParser.cpp
@@ -935,10 +935,20 @@ ParseResult FIRParser::parseType(FIRRTLType &result, const Twine &message) {
     auto loc = getToken().getLoc();
     consumeToken();
     FIRRTLType type;
+    SmallVector<StringRef> layers;
 
     if (parseToken(FIRToken::less, "expected '<' in reference type") ||
-        parseType(type, "expected probe data type") ||
-        parseToken(FIRToken::greater, "expected '>' in reference type"))
+        parseType(type, "expected probe data type"))
+      return failure();
+    while (getToken().getKind() == FIRToken::identifier) {
+      StringRef layer;
+      if (parseId(layer, "expected layer name"))
+        return failure();
+      layers.push_back(layer);
+      if (getToken().getKind() == FIRToken::period)
+        consumeToken();
+    }
+    if (parseToken(FIRToken::greater, "expected '>' in reference type"))
       return failure();
 
     bool forceable = kind == FIRToken::kw_RWProbe;
@@ -953,7 +963,17 @@ ParseResult FIRParser::parseType(FIRRTLType &result, const Twine &message) {
     if (forceable && innerType.containsConst())
       return emitError(loc, "rwprobe cannot contain const");
 
-    result = RefType::get(innerType, forceable);
+    SymbolRefAttr layer;
+    if (!layers.empty()) {
+      auto nestedLayers =
+          llvm::map_range(ArrayRef(layers).drop_front(), [&](StringRef a) {
+            return FlatSymbolRefAttr::get(getContext(), a);
+          });
+      layer = SymbolRefAttr::get(getContext(), layers.front(),
+                                 llvm::to_vector(nestedLayers));
+    }
+
+    result = RefType::get(innerType, forceable, layer);
     break;
   }
 
diff --git a/test/Dialect/FIRRTL/parse-basic.fir b/test/Dialect/FIRRTL/parse-basic.fir
index fd2aaa44fa98..284b721bd222 100644
--- a/test/Dialect/FIRRTL/parse-basic.fir
+++ b/test/Dialect/FIRRTL/parse-basic.fir
@@ -1671,9 +1671,13 @@ circuit Layers:
     ; CHECK-NEXT:   }
     ; CHECK-NEXT: }
 
-  ; CHECK: firrtl.module @Layers
+  ; CHECK:      firrtl.module @Layers
+  ; CHECK-SAME:   out %b: !firrtl.probe<uint<1>, @A>
+  ; CHECK-SAME:   out %c: !firrtl.rwprobe<uint<1>, @A::@B>
   module Layers:
     input a: UInt<1>
+    output b: Probe<UInt<1>, A>
+    output c: RWProbe<UInt<1>, A.B>
 
     layerblock A:
       node A_a = a