[FIRRTL] Fix TagExtractOp's type inference

TagExtractOp's type inference predates adding user-defined encodings for
enums in FIRRTL. This fixes the issue by using the bitwidth helpers
defined on the enumeration type.

Author: youngar

lib/Dialect/FIRRTL/FIRRTLOps.cpp

@@ -4863,8 +4863,7 @@ FIRRTLType SubaccessOp::inferReturnType(Type inType, Type indexType,
 FIRRTLType TagExtractOp::inferReturnType(FIRRTLType input,
                                          std::optional<Location> loc) {
   auto inType = type_cast<FEnumType>(input);
-  auto i = llvm::Log2_32_Ceil(inType.getNumElements());
-  return UIntType::get(inType.getContext(), i);
+  return UIntType::get(inType.getContext(), inType.getTagWidth());
 }
 
 ParseResult MultibitMuxOp::parse(OpAsmParser &parser, OperationState &result) {

test/Dialect/FIRRTL/parse-basic.fir

@@ -961,10 +961,10 @@ circuit MyModule :     ; CHECK: firrtl.circuit "MyModule"
 
   ; CHECK-LABEL: firrtl.module private @TagExtract
   module TagExtract:
-    input i : {|A, B, C|}
-    output o : UInt<2>
-    ; CHECK: %0 = firrtl.tagextract %i : !firrtl.enum<A, B, C>
-    ; CHECK: firrtl.matchingconnect %o, %0 : !firrtl.uint<2>
+    input i : {|A, B, C = 8|}
+    output o : UInt<4>
+    ; CHECK: %0 = firrtl.tagextract %i : !firrtl.enum<A, B, C = 8>
+    ; CHECK: firrtl.matchingconnect %o, %0 : !firrtl.uint<4>
     connect o, tagExtract(i)
 
   ; CHECK-LABEL: module private @RefsChild(