From c225a1899f6bb3d36044ea8261ee0c990230f4f3 Mon Sep 17 00:00:00 2001
From: Hideto Ueno <uenoku.tokotoko@gmail.com>
Date: Fri, 16 Jun 2023 07:09:40 -0700
Subject: [PATCH 1/8] [FIRRTL] Add BaseTypeAliasType

This commit adds minimal changes BaseTypeAliasType for base type alias.
 Implement storage and methods for FIRRTLBaseType and FieldIDTypeInterface.

This doesn't include recursive property chagnes and
`FIRRTLBaseType::getAnonnymousType` whcih requires more broader changes.
---
 .../circt/Dialect/FIRRTL/FIRRTLTypesImpl.td   |  45 ++++++
 lib/Dialect/FIRRTL/FIRRTLTypes.cpp            | 141 ++++++++++++++++--
 test/Dialect/FIRRTL/test.mlir                 |  10 ++
 3 files changed, 187 insertions(+), 9 deletions(-)

diff --git a/include/circt/Dialect/FIRRTL/FIRRTLTypesImpl.td b/include/circt/Dialect/FIRRTL/FIRRTLTypesImpl.td
index 8070483eac5e..5bfd3f160e40 100644
--- a/include/circt/Dialect/FIRRTL/FIRRTLTypesImpl.td
+++ b/include/circt/Dialect/FIRRTL/FIRRTLTypesImpl.td
@@ -481,6 +481,51 @@ def RefImpl : FIRRTLImplType<"Ref",
   }];
 }
 
+def BaseTypeAliasImpl : FIRRTLImplType<"BaseTypeAlias", [FieldIDTypeInterface],
+                                        "::circt::firrtl::FIRRTLBaseType"> {
+  let summary = "type alias for firrtl base types";
+  let parameters =
+    (ins  "StringAttr":$name, TypeParameter<"::circt::firrtl::FIRRTLBaseType",
+                                            "An inner type">:$innerType);
+  let storageClass = "BaseTypeAliasStorage";
+  let genAccessors = true;
+  let skipDefaultBuilders = true;
+  let extraClassDeclaration = [{
+      // FIRRTLBaseType utils.
+      FIRRTLBaseType getAnonymousType();
+      FIRRTLBaseType getPassiveType();
+      FIRRTLBaseType getConstType(bool isConst);
+      FIRRTLBaseType getAllConstDroppedType();
+
+      /// Return the recursive properties of the type.
+      RecursiveTypeProperties getRecursiveTypeProperties() const;
+
+      // If a given `newInnerType` is identical to innerType, return `*this`
+      // because we can reuse the type alias. Otherwise return `newInnerType`.
+      FIRRTLBaseType getModifiedType(FIRRTLBaseType newInnerType);
+
+      /// Implement FieldIDTypeInterface.
+
+      /// Strip off a single layer of this type and return the sub-type and a
+      /// field ID targeting the same field, but rebased on the sub-type.
+      std::pair<circt::hw::FieldIDTypeInterface, uint64_t>
+      getSubTypeByFieldID(uint64_t fieldID);
+
+      /// Get the maximum field ID.
+      uint64_t getMaxFieldID();
+
+      /// Returns the effective field id when treating the index field as the root
+      /// of the type.  Essentially maps a fieldID to a fieldID after a subfield
+      /// op. Returns the new id and whether the id is in the given child.
+      std::pair<uint64_t, bool> rootChildFieldID(uint64_t fieldID, uint64_t index);
+  }];
+
+  let builders = [
+    TypeBuilderWithInferredContext<(ins "::mlir::StringAttr":$name,
+           "::circt::firrtl::FIRRTLBaseType":$innerType)>
+  ];
+}
+
 //===----------------------------------------------------------------------===//
 // Non-Hardware Type Definitions
 //===----------------------------------------------------------------------===//
diff --git a/lib/Dialect/FIRRTL/FIRRTLTypes.cpp b/lib/Dialect/FIRRTL/FIRRTLTypes.cpp
index 216448c7a2f4..c7ecacd793c1 100644
--- a/lib/Dialect/FIRRTL/FIRRTLTypes.cpp
+++ b/lib/Dialect/FIRRTL/FIRRTLTypes.cpp
@@ -106,6 +106,11 @@ static LogicalResult customTypePrinter(Type type, AsmPrinter &os) {
       })
       .Case<StringType>([&](auto stringType) { os << "string"; })
       .Case<BigIntType>([&](auto bigIntType) { os << "bigint"; })
+      .Case<BaseTypeAliasType>([&](BaseTypeAliasType alias) {
+        os << "alias<" << alias.getName().getValue() << ", ";
+        printNestedType(alias.getInnerType(), os);
+        os << '>';
+      })
       .Default([&](auto) { anyFailed = true; });
   return failure(anyFailed);
 }
@@ -370,6 +375,18 @@ static OptionalParseResult customTypeParser(AsmParser &parser, StringRef name,
     result = BigIntType::get(parser.getContext());
     return success();
   }
+  if (name.equals("alias")) {
+    FIRRTLBaseType type;
+    StringRef name;
+    if (parser.parseLess() || parser.parseKeyword(&name) ||
+        parser.parseComma() || parseNestedBaseType(type, parser) ||
+        parser.parseGreater())
+      return failure();
+
+    return result =
+               BaseTypeAliasType::get(StringAttr::get(context, name), type),
+           success();
+  }
 
   return {};
 }
@@ -501,6 +518,9 @@ bool FIRRTLType::isGround() {
             AnalogType>([](Type) { return true; })
       .Case<BundleType, FVectorType, FEnumType, OpenBundleType, OpenVectorType>(
           [](Type) { return false; })
+      .Case<BaseTypeAliasType>([](BaseTypeAliasType alias) {
+        return alias.getAnonymousType().isGround();
+      })
       // Not ground per spec, but leaf of aggregate.
       .Case<PropertyType, RefType>([](Type) { return false; })
       .Default([](Type) {
@@ -534,7 +554,7 @@ RecursiveTypeProperties FIRRTLType::getRecursiveTypeProperties() const {
             true, false, true, type.isConst(), !type.hasWidth(), false};
       })
       .Case<BundleType, FVectorType, FEnumType, OpenBundleType, OpenVectorType,
-            RefType>(
+            RefType, BaseTypeAliasType>(
           [](auto type) { return type.getRecursiveTypeProperties(); })
       .Case<StringType, BigIntType>([](auto type) {
         return RecursiveTypeProperties{true, false, false, false, false, false};
@@ -550,7 +570,7 @@ FIRRTLBaseType FIRRTLBaseType::getPassiveType() {
   return TypeSwitch<FIRRTLBaseType, FIRRTLBaseType>(*this)
       .Case<ClockType, ResetType, AsyncResetType, SIntType, UIntType,
             AnalogType, FEnumType>([&](Type) { return *this; })
-      .Case<BundleType, FVectorType, FEnumType>(
+      .Case<BundleType, FVectorType, FEnumType, BaseTypeAliasType>(
           [](auto type) { return type.getPassiveType(); })
       .Default([](Type) {
         llvm_unreachable("unknown FIRRTL type");
@@ -562,7 +582,7 @@ FIRRTLBaseType FIRRTLBaseType::getPassiveType() {
 FIRRTLBaseType FIRRTLBaseType::getConstType(bool isConst) {
   return TypeSwitch<FIRRTLBaseType, FIRRTLBaseType>(*this)
       .Case<ClockType, ResetType, AsyncResetType, AnalogType, SIntType,
-            UIntType, BundleType, FVectorType, FEnumType>(
+            UIntType, BundleType, FVectorType, FEnumType, BaseTypeAliasType>(
           [&](auto type) { return type.getConstType(isConst); })
       .Default([](Type) {
         llvm_unreachable("unknown FIRRTL type");
@@ -575,7 +595,7 @@ FIRRTLBaseType FIRRTLBaseType::getAllConstDroppedType() {
   return TypeSwitch<FIRRTLBaseType, FIRRTLBaseType>(*this)
       .Case<ClockType, ResetType, AsyncResetType, AnalogType, SIntType,
             UIntType>([&](auto type) { return type.getConstType(false); })
-      .Case<BundleType, FVectorType, FEnumType>(
+      .Case<BundleType, FVectorType, FEnumType, BaseTypeAliasType>(
           [&](auto type) { return type.getAllConstDroppedType(); })
       .Default([](Type) {
         llvm_unreachable("unknown FIRRTL type");
@@ -605,6 +625,9 @@ FIRRTLBaseType FIRRTLBaseType::getMaskType() {
                                 vectorType.getNumElements(),
                                 vectorType.isConst());
       })
+      .Case<BaseTypeAliasType>([](BaseTypeAliasType base) {
+        return base.getModifiedType(base.getInnerType().getMaskType());
+      })
       .Default([](Type) {
         llvm_unreachable("unknown FIRRTL type");
         return FIRRTLBaseType();
@@ -637,6 +660,9 @@ FIRRTLBaseType FIRRTLBaseType::getWidthlessType() {
           newElements.push_back({elt.name, elt.type.getWidthlessType()});
         return FEnumType::get(this->getContext(), newElements, a.isConst());
       })
+      .Case<BaseTypeAliasType>([](BaseTypeAliasType type) {
+        return type.getModifiedType(type.getInnerType().getWidthlessType());
+      })
       .Default([](auto) {
         llvm_unreachable("unknown FIRRTL type");
         return FIRRTLBaseType();
@@ -655,6 +681,10 @@ int32_t FIRRTLBaseType::getBitWidthOrSentinel() {
       .Case<AnalogType>(
           [](AnalogType analogType) { return analogType.getWidthOrSentinel(); })
       .Case<BundleType, FVectorType, FEnumType>([](Type) { return -2; })
+      .Case<BaseTypeAliasType>([](BaseTypeAliasType type) {
+        // It's faster to use its anonymous type.
+        return type.getAnonymousType().getBitWidthOrSentinel();
+      })
       .Default([](Type) {
         llvm_unreachable("unknown FIRRTL type");
         return -2;
@@ -669,6 +699,8 @@ bool FIRRTLBaseType::isResetType() {
       .Case<ResetType, AsyncResetType>([](Type) { return true; })
       .Case<UIntType>(
           [](UIntType a) { return !a.hasWidth() || a.getWidth() == 1; })
+      .Case<BaseTypeAliasType>(
+          [](auto type) { return type.getInnerType().isResetType(); })
       .Default([](Type) { return false; });
 }
 
@@ -676,7 +708,7 @@ uint64_t FIRRTLBaseType::getMaxFieldID() {
   return TypeSwitch<FIRRTLBaseType, uint64_t>(*this)
       .Case<AnalogType, ClockType, ResetType, AsyncResetType, SIntType,
             UIntType>([](Type) { return 0; })
-      .Case<BundleType, FVectorType, FEnumType>(
+      .Case<BundleType, FVectorType, FEnumType, BaseTypeAliasType>(
           [](auto type) { return type.getMaxFieldID(); })
       .Default([](Type) {
         llvm_unreachable("unknown FIRRTL type");
@@ -693,7 +725,7 @@ FIRRTLBaseType::getSubTypeByFieldID(uint64_t fieldID) {
         assert(!fieldID && "non-aggregate types must have a field id of 0");
         return std::pair(t, 0);
       })
-      .Case<BundleType, FVectorType, FEnumType>(
+      .Case<BundleType, FVectorType, FEnumType, BaseTypeAliasType>(
           [&](auto type) { return type.getSubTypeByFieldID(fieldID); })
       .Default([](Type) {
         llvm_unreachable("unknown FIRRTL type");
@@ -714,7 +746,7 @@ std::pair<uint64_t, bool> FIRRTLBaseType::rootChildFieldID(uint64_t fieldID,
   return TypeSwitch<FIRRTLBaseType, std::pair<uint64_t, bool>>(*this)
       .Case<AnalogType, ClockType, ResetType, AsyncResetType, SIntType,
             UIntType>([&](Type) { return std::make_pair(0, fieldID == 0); })
-      .Case<BundleType, FVectorType, FEnumType>(
+      .Case<BundleType, FVectorType, FEnumType, BaseTypeAliasType>(
           [&](auto type) { return type.rootChildFieldID(fieldID, index); })
       .Default([](Type) {
         llvm_unreachable("unknown FIRRTL type");
@@ -2112,6 +2144,97 @@ auto FEnumType::verify(function_ref<InFlightDiagnostic()> emitErrorFn,
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// BaseTypeAliasType
+//===----------------------------------------------------------------------===//
+
+struct circt::firrtl::detail::BaseTypeAliasStorage
+    : circt::firrtl::detail::FIRRTLBaseTypeStorage {
+  using KeyTy = std::tuple<StringAttr, FIRRTLBaseType>;
+
+  BaseTypeAliasStorage(StringAttr name, FIRRTLBaseType innerType)
+      : detail::FIRRTLBaseTypeStorage(innerType.isConst()), name(name),
+        innerType(innerType) {}
+
+  bool operator==(const KeyTy &key) const {
+    return key == KeyTy(name, innerType);
+  }
+
+  static llvm::hash_code hashKey(const KeyTy &key) {
+    return llvm::hash_combine(key);
+  }
+
+  static BaseTypeAliasStorage *construct(TypeStorageAllocator &allocator,
+                                         KeyTy key) {
+    return new (allocator.allocate<BaseTypeAliasStorage>())
+        BaseTypeAliasStorage(std::get<0>(key), std::get<1>(key));
+  }
+  StringAttr name;
+  FIRRTLBaseType innerType;
+  FIRRTLBaseType anonymousType;
+};
+
+auto BaseTypeAliasType::get(StringAttr name, FIRRTLBaseType innerType)
+    -> BaseTypeAliasType {
+  return Base::get(name.getContext(), name, innerType);
+}
+
+auto BaseTypeAliasType::getName() const -> StringAttr {
+  return getImpl()->name;
+}
+
+auto BaseTypeAliasType::getInnerType() const -> FIRRTLBaseType {
+  return getImpl()->innerType;
+}
+
+FIRRTLBaseType BaseTypeAliasType::getAnonymousType() {
+  // FIXME: Update to use FIRRTLBaseType::getAnonymousType()
+  return getInnerType();
+}
+
+FIRRTLBaseType BaseTypeAliasType::getPassiveType() {
+  return getModifiedType(getInnerType().getPassiveType());
+}
+
+RecursiveTypeProperties BaseTypeAliasType::getRecursiveTypeProperties() const {
+  auto rtp = getInnerType().getRecursiveTypeProperties();
+  // TODO: Enable containsTypeAlias
+  return rtp;
+}
+
+// If a given `newInnerType` is identical to innerType, return `*this`
+// because we can reuse the type alias. Otherwise return `newInnerType`.
+FIRRTLBaseType BaseTypeAliasType::getModifiedType(FIRRTLBaseType newInnerType) {
+  if (newInnerType == getInnerType())
+    return *this;
+  return newInnerType;
+}
+
+// FieldIDTypeInterface implementation.
+FIRRTLBaseType BaseTypeAliasType::getAllConstDroppedType() {
+  return getModifiedType(getInnerType().getAllConstDroppedType());
+}
+
+FIRRTLBaseType BaseTypeAliasType::getConstType(bool isConst) {
+  return getModifiedType(getInnerType().getConstType(isConst));
+}
+
+std::pair<circt::hw::FieldIDTypeInterface, uint64_t>
+BaseTypeAliasType::getSubTypeByFieldID(uint64_t fieldID) {
+  return getInnerType().getSubTypeByFieldID(fieldID);
+}
+
+uint64_t BaseTypeAliasType::getMaxFieldID() {
+  // We can use anonymous type.
+  return getAnonymousType().getMaxFieldID();
+}
+
+std::pair<uint64_t, bool> BaseTypeAliasType::rootChildFieldID(uint64_t fieldID,
+                                                              uint64_t index) {
+  // We can use anonymous type.
+  return getAnonymousType().rootChildFieldID(fieldID, index);
+}
+
 //===----------------------------------------------------------------------===//
 // RefType
 //===----------------------------------------------------------------------===//
@@ -2225,8 +2348,8 @@ AsyncResetType AsyncResetType::getConstType(bool isConst) {
 void FIRRTLDialect::registerTypes() {
   addTypes<SIntType, UIntType, ClockType, ResetType, AsyncResetType, AnalogType,
            // Derived Types
-           BundleType, FVectorType, FEnumType, RefType, OpenBundleType,
-           OpenVectorType, StringType, BigIntType>();
+           BundleType, FVectorType, FEnumType, BaseTypeAliasType, RefType,
+           OpenBundleType, OpenVectorType, StringType, BigIntType>();
 }
 
 // Get the bit width for this type, return None  if unknown. Unlike
diff --git a/test/Dialect/FIRRTL/test.mlir b/test/Dialect/FIRRTL/test.mlir
index 9132a4e02bc5..ad77d9ed89ed 100644
--- a/test/Dialect/FIRRTL/test.mlir
+++ b/test/Dialect/FIRRTL/test.mlir
@@ -261,4 +261,14 @@ firrtl.module @BigIntTest(in %in: !firrtl.bigint, out %out: !firrtl.bigint) {
   // CHECK: %1 = firrtl.bigint -4
   %1 = firrtl.bigint -4
 }
+
+// CHECK-LABEL: TypeAlias
+// CHECK-SAME: %in: !firrtl.alias<bar, uint<1>>
+// CHECK-SAME: %const: !firrtl.const.alias<baz, const.uint<1>>
+// CHECK-SAME: %r: !firrtl.openbundle<a: alias<baz, uint<1>>>
+firrtl.module @TypeAlias(in %in: !firrtl.alias<bar, uint<1>>,
+                         in %const: !firrtl.const.alias<baz, const.uint<1>>,
+                         out %r : !firrtl.openbundle<a: alias<baz, uint<1>>>) {
+}
+
 }

From 9987245505c44aa01d8d6d7b7228af82b70dbae8 Mon Sep 17 00:00:00 2001
From: Hideto Ueno <uenoku.tokotoko@gmail.com>
Date: Fri, 16 Jun 2023 08:59:13 -0700
Subject: [PATCH 2/8] [FIRRTL] Implement alias-aware type casts

This implements `type_isa/type_cast/type_dyn_cast/type_dyn_cast_or_null`.
When a given type is an alias, these casts will visit its inner type recursively.
These type casts can be used anywhere and if the requested type is
non-base type or FIRRTLBaseType itself, then it will optimized in the
compile time,
---
 include/circt/Dialect/FIRRTL/FIRRTLTypes.h | 69 ++++++++++++++++++++++
 unittests/Dialect/FIRRTL/TypesTest.cpp     | 23 ++++++++
 2 files changed, 92 insertions(+)

diff --git a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
index 47b336045649..b1742bb7de86 100644
--- a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
+++ b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
@@ -46,6 +46,7 @@ class RefType;
 class PropertyType;
 class StringType;
 class BigIntType;
+class BaseTypeAliasType;
 
 /// A collection of bits indicating the recursive properties of a type.
 struct RecursiveTypeProperties {
@@ -367,4 +368,72 @@ struct DenseMapInfo<circt::firrtl::FIRRTLType> {
 
 } // namespace llvm
 
+namespace circt {
+namespace firrtl {
+//===--------------------------------------------------------------------===//
+// Utility for type aliases
+//===--------------------------------------------------------------------===//
+
+/// A struct to check if there is a type derived from FIRRTLBaseType.
+/// `ContainBaseSubTypes<BaseTy>::value` returns true if `BaseTy` is a derived
+/// from `FIRRTLBaseType` and not `FIRRTLBaseType` itself.
+template <typename head, typename... tail>
+struct ContainBaseSubTypes {
+  static constexpr bool value =
+      ContainBaseSubTypes<head>::value || ContainBaseSubTypes<tail...>::value;
+};
+
+template <typename BaseTy>
+struct ContainBaseSubTypes<BaseTy> {
+  static constexpr bool value =
+      std::is_base_of<FIRRTLBaseType, BaseTy>::value &&
+      !std::is_same_v<FIRRTLBaseType, BaseTy>;
+};
+
+template <typename... BaseTy>
+bool type_isa(Type type) { // NOLINT(readability-identifier-naming)
+  // First check if the type is the requested type.
+  if (isa<BaseTy...>(type))
+    return true;
+
+  // If the requested type is a subtype of FIRRTLBaseType, then check if it is a
+  // type alias wrapping the requested type.
+  if constexpr (ContainBaseSubTypes<BaseTy...>::value) {
+    if (auto alias = dyn_cast<BaseTypeAliasType>(type))
+      return type_isa<BaseTy...>(alias.getInnerType());
+  }
+
+  return false;
+}
+
+// type_isa for a nullable argument.
+template <typename... BaseTy>
+bool type_isa_and_nonnull(Type type) { // NOLINT(readability-identifier-naming)
+  if (!type)
+    return false;
+  return type_isa<BaseTy...>(type);
+}
+
+template <typename BaseTy>
+BaseTy type_cast(Type type) { // NOLINT(readability-identifier-naming)
+  assert(type_isa<BaseTy>(type) && "type must convert to requested type");
+
+  // If the type is the requested type, return it.
+  if (isa<BaseTy>(type))
+    return cast<BaseTy>(type);
+
+  // Otherwise, it must be a type alias wrapping the requested type.
+  if constexpr (ContainBaseSubTypes<BaseTy>::value) {
+    if (auto alias = dyn_cast<BaseTypeAliasType>(type))
+      return type_cast<BaseTy>(alias.getInnerType());
+  }
+
+  // Otherwise, it should fail. `cast` should cause a better assertion failure,
+  // so just use it.
+  return cast<BaseTy>(type);
+}
+
+} // namespace firrtl
+} // namespace circt
+
 #endif // CIRCT_DIALECT_FIRRTL_TYPES_H
diff --git a/unittests/Dialect/FIRRTL/TypesTest.cpp b/unittests/Dialect/FIRRTL/TypesTest.cpp
index 1cde290ae4e1..f7b9219946e8 100644
--- a/unittests/Dialect/FIRRTL/TypesTest.cpp
+++ b/unittests/Dialect/FIRRTL/TypesTest.cpp
@@ -22,4 +22,27 @@ TEST(TypesTest, AnalogContainsAnalog) {
   ASSERT_TRUE(AnalogType::get(&context).containsAnalog());
 }
 
+TEST(TypesTest, TypeAliasCast) {
+  MLIRContext context;
+  context.loadDialect<FIRRTLDialect>();
+  // Check containBaseSubTypes.
+  static_assert(!ContainBaseSubTypes<FIRRTLType>::value);
+  // Return false for FIRRTLBaseType.
+  static_assert(!ContainBaseSubTypes<FIRRTLBaseType>::value);
+  static_assert(!ContainBaseSubTypes<StringType>::value);
+  static_assert(ContainBaseSubTypes<FVectorType>::value);
+  static_assert(ContainBaseSubTypes<UIntType, StringType>::value);
+  AnalogType analog = AnalogType::get(&context);
+  BaseTypeAliasType alias1 =
+      BaseTypeAliasType::get(StringAttr::get(&context, "foo"), analog);
+  BaseTypeAliasType alias2 =
+      BaseTypeAliasType::get(StringAttr::get(&context, "bar"), alias1);
+  ASSERT_TRUE(!type_isa<FVectorType>(analog));
+  ASSERT_TRUE(type_isa<AnalogType>(analog));
+  ASSERT_TRUE(type_isa<AnalogType>(alias1));
+  ASSERT_TRUE(type_isa<AnalogType>(alias2));
+  ASSERT_TRUE(!type_isa<FVectorType>(alias2));
+  ASSERT_TRUE((type_isa<AnalogType, StringType>(alias2)));
+}
+
 } // namespace

From bbb4acd83e127d9378673d95bea661745dd10b11 Mon Sep 17 00:00:00 2001
From: Hideto Ueno <uenoku.tokotoko@gmail.com>
Date: Thu, 22 Jun 2023 03:44:42 -0700
Subject: [PATCH 3/8] Add missing casts

---
 include/circt/Dialect/FIRRTL/FIRRTLTypes.h | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
index b1742bb7de86..c3e5bb4f0ea5 100644
--- a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
+++ b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
@@ -433,6 +433,21 @@ BaseTy type_cast(Type type) { // NOLINT(readability-identifier-naming)
   return cast<BaseTy>(type);
 }
 
+template <typename BaseTy>
+BaseTy type_dyn_cast(Type type) { // NOLINT(readability-identifier-naming)
+  if (type_isa<BaseTy>(type))
+    return type_cast<BaseTy>(type);
+  return {};
+}
+
+template <typename BaseTy>
+BaseTy
+type_dyn_cast_or_null(Type type) { // NOLINT(readability-identifier-naming)
+  if (type_isa_and_nonnull<BaseTy>(type))
+    return type_cast<BaseTy>(type);
+  return {};
+}
+
 } // namespace firrtl
 } // namespace circt
 

From d63d2903b68e4a830607219a5911535511f8ef59 Mon Sep 17 00:00:00 2001
From: Hideto Ueno <uenoku.tokotoko@gmail.com>
Date: Thu, 22 Jun 2023 06:11:41 -0700
Subject: [PATCH 4/8] [FIRRTL] Add FIRRTLTypeSwitch

---
 include/circt/Dialect/FIRRTL/FIRRTLTypes.h | 92 ++++++++++++++++++++++
 1 file changed, 92 insertions(+)

diff --git a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
index c3e5bb4f0ea5..5563edf16827 100644
--- a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
+++ b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
@@ -16,6 +16,7 @@
 #include "circt/Dialect/FIRRTL/FIRRTLDialect.h"
 #include "circt/Dialect/HW/HWTypeInterfaces.h"
 #include "circt/Support/LLVM.h"
+#include "llvm/ADT/TypeSwitch.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/Types.h"
 
@@ -448,6 +449,97 @@ type_dyn_cast_or_null(Type type) { // NOLINT(readability-identifier-naming)
   return {};
 }
 
+//===--------------------------------------------------------------------===//
+// Type alias aware TypeSwitch.
+//===--------------------------------------------------------------------===//
+
+/// This class implements the same functionality as TypeSwitch except that
+/// it uses firrtl::type_dyn_cast for dynamic cast. llvm::TypeSwitch is not
+/// customizable so this class currently duplicates the code.
+template <typename T, typename ResultT = void>
+class FIRRTLTypeSwitch
+    : public llvm::detail::TypeSwitchBase<FIRRTLTypeSwitch<T, ResultT>, T> {
+public:
+  using BaseT = llvm::detail::TypeSwitchBase<FIRRTLTypeSwitch<T, ResultT>, T>;
+  using BaseT::BaseT;
+  using BaseT::Case;
+  FIRRTLTypeSwitch(FIRRTLTypeSwitch &&other) = default;
+
+  /// Add a case on the given type.
+  template <typename CaseT, typename CallableT>
+  FIRRTLTypeSwitch<T, ResultT> &Case(CallableT &&caseFn) {
+    if (result)
+      return *this;
+
+    // Check to see if CaseT applies to 'value'. Use `type_dyn_cast` here.
+    if (auto caseValue = circt::firrtl::type_dyn_cast<CaseT>(this->value))
+      result.emplace(caseFn(caseValue));
+    return *this;
+  }
+
+  /// As a default, invoke the given callable within the root value.
+  template <typename CallableT>
+  [[nodiscard]] ResultT Default(CallableT &&defaultFn) {
+    if (result)
+      return std::move(*result);
+    return defaultFn(this->value);
+  }
+
+  /// As a default, return the given value.
+  [[nodiscard]] ResultT Default(ResultT defaultResult) {
+    if (result)
+      return std::move(*result);
+    return defaultResult;
+  }
+
+  [[nodiscard]] operator ResultT() {
+    assert(result && "Fell off the end of a type-switch");
+    return std::move(*result);
+  }
+
+private:
+  /// The pointer to the result of this switch statement, once known,
+  /// null before that.
+  std::optional<ResultT> result;
+};
+
+/// Specialization of FIRRTLTypeSwitch for void returning callables.
+template <typename T>
+class FIRRTLTypeSwitch<T, void>
+    : public llvm::detail::TypeSwitchBase<FIRRTLTypeSwitch<T, void>, T> {
+public:
+  using BaseT = llvm::detail::TypeSwitchBase<FIRRTLTypeSwitch<T, void>, T>;
+  using BaseT::BaseT;
+  using BaseT::Case;
+  FIRRTLTypeSwitch(FIRRTLTypeSwitch &&other) = default;
+
+  /// Add a case on the given type.
+  template <typename CaseT, typename CallableT>
+  FIRRTLTypeSwitch<T, void> &Case(CallableT &&caseFn) {
+    if (foundMatch)
+      return *this;
+
+    // Check to see if any of the types apply to 'value'.
+    if (auto caseValue = circt::firrtl::type_dyn_cast<CaseT>(this->value)) {
+      caseFn(caseValue);
+      foundMatch = true;
+    }
+    return *this;
+  }
+
+  /// As a default, invoke the given callable within the root value.
+  template <typename CallableT>
+  void Default(CallableT &&defaultFn) {
+    if (!foundMatch)
+      defaultFn(this->value);
+  }
+
+private:
+  /// A flag detailing if we have already found a match.
+  bool foundMatch = false;
+};
+
+
 } // namespace firrtl
 } // namespace circt
 

From de354b1478d5d558b98d1dc33687f2887bf7ce26 Mon Sep 17 00:00:00 2001
From: Hideto Ueno <uenoku.tokotoko@gmail.com>
Date: Thu, 22 Jun 2023 23:31:40 +0900
Subject: [PATCH 5/8] clang-format

---
 include/circt/Dialect/FIRRTL/FIRRTLTypes.h | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
index 5563edf16827..0a33a3371347 100644
--- a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
+++ b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
@@ -16,9 +16,9 @@
 #include "circt/Dialect/FIRRTL/FIRRTLDialect.h"
 #include "circt/Dialect/HW/HWTypeInterfaces.h"
 #include "circt/Support/LLVM.h"
-#include "llvm/ADT/TypeSwitch.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/Types.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 namespace circt {
 namespace firrtl {
@@ -539,7 +539,6 @@ class FIRRTLTypeSwitch<T, void>
   bool foundMatch = false;
 };
 
-
 } // namespace firrtl
 } // namespace circt
 

From 76e8a12a13348395cf05cf42e198198a9d8c028d Mon Sep 17 00:00:00 2001
From: Hideto Ueno <uenoku.tokotoko@gmail.com>
Date: Fri, 23 Jun 2023 02:58:32 +0900
Subject: [PATCH 6/8] Remove conflict marker

---
 include/circt/Dialect/FIRRTL/FIRRTLTypes.h | 1 -
 1 file changed, 1 deletion(-)

diff --git a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
index 379b507443af..4c41e301519b 100644
--- a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
+++ b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
@@ -449,7 +449,6 @@ type_dyn_cast_or_null(Type type) { // NOLINT(readability-identifier-naming)
   return {};
 }
 
-<<<<<<< HEAD
 //===--------------------------------------------------------------------===//
 // Type alias aware TypeSwitch.
 //===--------------------------------------------------------------------===//

From 2fb088ce56949527e18c317928f1fe44d1a3f4d3 Mon Sep 17 00:00:00 2001
From: Hideto Ueno <uenoku.tokotoko@gmail.com>
Date: Mon, 26 Jun 2023 21:53:09 +0900
Subject: [PATCH 7/8] Update FIRRTLTypes.h

---
 include/circt/Dialect/FIRRTL/FIRRTLTypes.h | 2 --
 1 file changed, 2 deletions(-)

diff --git a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
index 4c41e301519b..363f93021379 100644
--- a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
+++ b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
@@ -539,8 +539,6 @@ class FIRRTLTypeSwitch<T, void>
   bool foundMatch = false;
 };
 
-=======
->>>>>>> a4aabbca0d1961054dc4353c4f3c3fe8cb57d6ea
 } // namespace firrtl
 } // namespace circt
 

From eb373da8b03525b1af1404f676d432dfe761b51d Mon Sep 17 00:00:00 2001
From: Hideto Ueno <uenoku.tokotoko@gmail.com>
Date: Thu, 29 Jun 2023 06:56:26 +0900
Subject: [PATCH 8/8] Apply suggestions from code review

Co-authored-by: Prithayan Barua <prithayan@gmail.com>
---
 include/circt/Dialect/FIRRTL/FIRRTLTypes.h | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
index 363f93021379..f16174b634cf 100644
--- a/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
+++ b/include/circt/Dialect/FIRRTL/FIRRTLTypes.h
@@ -467,7 +467,8 @@ class FIRRTLTypeSwitch
 
   /// Add a case on the given type.
   template <typename CaseT, typename CallableT>
-  FIRRTLTypeSwitch<T, ResultT> &Case(CallableT &&caseFn) {
+  FIRRTLTypeSwitch<T, ResultT> &
+  Case(CallableT &&caseFn) { // NOLINT(readability-identifier-naming)
     if (result)
       return *this;
 
@@ -479,14 +480,16 @@ class FIRRTLTypeSwitch
 
   /// As a default, invoke the given callable within the root value.
   template <typename CallableT>
-  [[nodiscard]] ResultT Default(CallableT &&defaultFn) {
+  [[nodiscard]] ResultT
+  Default(CallableT &&defaultFn) { // NOLINT(readability-identifier-naming)
     if (result)
       return std::move(*result);
     return defaultFn(this->value);
   }
 
   /// As a default, return the given value.
-  [[nodiscard]] ResultT Default(ResultT defaultResult) {
+  [[nodiscard]] ResultT
+  Default(ResultT defaultResult) { // NOLINT(readability-identifier-naming)
     if (result)
       return std::move(*result);
     return defaultResult;
@@ -515,7 +518,8 @@ class FIRRTLTypeSwitch<T, void>
 
   /// Add a case on the given type.
   template <typename CaseT, typename CallableT>
-  FIRRTLTypeSwitch<T, void> &Case(CallableT &&caseFn) {
+  FIRRTLTypeSwitch<T, void> &
+  Case(CallableT &&caseFn) { // NOLINT(readability-identifier-naming)
     if (foundMatch)
       return *this;
 
@@ -529,7 +533,7 @@ class FIRRTLTypeSwitch<T, void>
 
   /// As a default, invoke the given callable within the root value.
   template <typename CallableT>
-  void Default(CallableT &&defaultFn) {
+  void Default(CallableT &&defaultFn) { // NOLINT(readability-identifier-naming)
     if (!foundMatch)
       defaultFn(this->value);
   }