Add utility for determining if a TypeAnnotation represents a bitvecto…

…r type.

Use utility in fuzzer.

PiperOrigin-RevId: 569015286

xls/dslx/frontend/BUILD

@@ -158,11 +158,12 @@ cc_test(
     deps = [
         ":ast",
         ":ast_utils",
-        "@com_google_absl//absl/status:statusor",
         "//xls/common:xls_gunit",
         "//xls/common:xls_gunit_main",
+        "//xls/common/logging",
         "//xls/common/status:matchers",
         "//xls/dslx:import_data",
+        "//xls/dslx:parse_and_typecheck",
     ],
 )
 

xls/dslx/frontend/ast_utils.cc

@@ -467,4 +467,74 @@ int64_t DetermineIndentLevel(const AstNode& n) {
   }
 }
 
+std::optional<BitVectorMetadata> ExtractBitVectorMetadata(
+    const TypeAnnotation* type_annotation) {
+  bool is_enum = false;
+  bool is_alias = false;
+  const TypeAnnotation* type = type_annotation;
+  while (dynamic_cast<const TypeRefTypeAnnotation*>(type) != nullptr) {
+    auto type_ref = dynamic_cast<const TypeRefTypeAnnotation*>(type);
+    if (std::holds_alternative<TypeAlias*>(
+            type_ref->type_ref()->type_definition())) {
+      is_alias = true;
+      type = std::get<TypeAlias*>(type_ref->type_ref()->type_definition())
+                 ->type_annotation();
+    } else if (std::holds_alternative<EnumDef*>(
+                   type_ref->type_ref()->type_definition())) {
+      is_enum = true;
+      type = std::get<EnumDef*>(type_ref->type_ref()->type_definition())
+                 ->type_annotation();
+    } else {
+      break;
+    }
+  }
+
+  BitVectorKind kind;
+  if (is_enum && is_alias) {
+    kind = BitVectorKind::kEnumTypeAlias;
+  } else if (is_enum && !is_alias) {
+    kind = BitVectorKind::kEnumType;
+  } else if (!is_enum && is_alias) {
+    kind = BitVectorKind::kBitTypeAlias;
+  } else {
+    kind = BitVectorKind::kBitType;
+  }
+
+  if (const BuiltinTypeAnnotation* builtin =
+          dynamic_cast<const BuiltinTypeAnnotation*>(type);
+      builtin != nullptr) {
+    switch (builtin->builtin_type()) {
+      case BuiltinType::kToken:
+      case BuiltinType::kChannelIn:
+      case BuiltinType::kChannelOut:
+        return std::nullopt;
+      default:
+        break;
+    }
+    return BitVectorMetadata{.bit_count = builtin->GetBitCount(),
+                             .is_signed = builtin->GetSignedness(),
+                             .kind = kind};
+  }
+  if (const ArrayTypeAnnotation* array_type =
+          dynamic_cast<const ArrayTypeAnnotation*>(type);
+      array_type != nullptr) {
+    // bits[..], uN[..], and sN[..] are bit-vector types but a represented with
+    // ArrayTypeAnnotations.
+    const BuiltinTypeAnnotation* builtin_element_type =
+        dynamic_cast<const BuiltinTypeAnnotation*>(array_type->element_type());
+    if (builtin_element_type == nullptr) {
+      return std::nullopt;
+    }
+    if (builtin_element_type->builtin_type() == BuiltinType::kBits ||
+        builtin_element_type->builtin_type() == BuiltinType::kUN ||
+        builtin_element_type->builtin_type() == BuiltinType::kSN) {
+      return BitVectorMetadata{
+          .bit_count = array_type->dim(),
+          .is_signed = builtin_element_type->builtin_type() == BuiltinType::kSN,
+          .kind = kind};
+    }
+  }
+  return std::nullopt;
+}
+
 }  // namespace xls::dslx

xls/dslx/frontend/ast_utils.h

@@ -168,6 +168,39 @@ inline T WidenVariantTo(const std::variant<FromTypes...>& v) {
   return TryWidenVariant<sizeof...(FromTypes) - 1, T>(v);
 }
 
+// Enumeration of the different kinds of DSLX types whose underlying data
+// representation is a bitvector.
+enum BitVectorKind : uint8_t {
+  // A native bit type. Examples: u32, s16, bits[3], uN[42], sN[2], bool.
+  kBitType,
+
+  // An alias of a bit type., or an alias of an alias of a bit type, etc.
+  // Example: `type Foo = u32;.
+  kBitTypeAlias,
+
+  // An enum type.
+  kEnumType,
+
+  // An alias of an enum type, or an alias of an alias of a enum type, etc.
+  // Example: `type MyEnumAlias = MyEnum;`
+  kEnumTypeAlias,
+};
+
+// Metadata about a DSLX type whose underlying data representation is a
+// bitvector.
+struct BitVectorMetadata {
+  std::variant<int64_t, Expr*> bit_count;
+  bool is_signed;
+  BitVectorKind kind;
+};
+
+// Returns metadata about the bit-vector type if `type_annotation` refers to a
+// type whose underlying representation is a bit-vector. Examples include u32,
+// s10, uN[42], bits[11], enums, etc, and aliases of these types. Returns
+// std::nullopt otherwise.
+std::optional<BitVectorMetadata> ExtractBitVectorMetadata(
+    const TypeAnnotation* type_annotation);
+
 }  // namespace xls::dslx
 
 #endif  // XLS_DSLX_FRONTEND_AST_UTILS_H_

xls/dslx/frontend/ast_utils_test.cc

@@ -13,16 +13,19 @@
 // limitations under the License.
 #include "xls/dslx/frontend/ast_utils.h"
 
+#include <cstdint>
 #include <memory>
 #include <optional>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "gtest/gtest.h"
-#include "absl/status/statusor.h"
+#include "xls/common/logging/logging.h"
 #include "xls/common/status/matchers.h"
 #include "xls/dslx/frontend/ast.h"
 #include "xls/dslx/import_data.h"
+#include "xls/dslx/parse_and_typecheck.h"
 
 namespace xls::dslx {
 namespace {
@@ -151,5 +154,102 @@ TEST(ProcConfigIrConverterTest, ResolveProcColonRef) {
   EXPECT_EQ(p, original_proc);
 }
 
+TEST(ProcConfigIrConverterTest, BitVectorTests) {
+  constexpr std::string_view kDslxText = R"(type MyType = u37;
+type MyTuple = (u23, u1);
+type MyArray = bits[22][33];
+enum MyEnum : u7 { kA = 0, }
+type MyEnumAlias = MyEnum;
+struct MyStruct { x: u17 }
+
+// A single struct which has many members of different types. This is an easy
+// way of gathering together many TypeAnnotations and referring to them.
+struct TheStruct {
+  a: u32,
+  b: s12,
+  c: uN[111],
+  d: sN[32],
+  e: bits[312],
+  f: bool,
+  g: MyType,
+  h: MyEnum,
+  i: MyEnumAlias,
+  j: (u23, u1),
+  k: bits[22][33],
+  l: MyStruct,
+  m: MyTuple,
+  n: MyArray,
+}
+ )";
+  XLS_ASSERT_OK_AND_ASSIGN(auto module,
+                           ParseModule(kDslxText, "fake_path.x", "the_module"));
+
+  XLS_ASSERT_OK_AND_ASSIGN(std::vector<AstNode*> nodes,
+                           CollectUnder(module.get(), /*want_types=*/true));
+  StructDef* the_struct_def = nullptr;
+  for (AstNode* node : nodes) {
+    if (auto* struct_def = dynamic_cast<StructDef*>(node);
+        struct_def != nullptr && struct_def->identifier() == "TheStruct") {
+      the_struct_def = struct_def;
+      break;
+    }
+  }
+  ASSERT_NE(the_struct_def, nullptr);
+  auto get_type_metadata = [&](std::string_view name) {
+    for (std::pair<NameDef*, TypeAnnotation*> member :
+         the_struct_def->members()) {
+      if (member.first->identifier() == name) {
+        return ExtractBitVectorMetadata(member.second);
+      }
+    }
+    XLS_LOG(FATAL) << "Unknown field: " << name;
+  };
+
+  EXPECT_EQ(std::get<int64_t>(get_type_metadata("a")->bit_count), 32);
+  EXPECT_FALSE(get_type_metadata("a")->is_signed);
+  EXPECT_EQ(get_type_metadata("a")->kind, BitVectorKind::kBitType);
+
+  EXPECT_EQ(std::get<int64_t>(get_type_metadata("b")->bit_count), 12);
+  EXPECT_TRUE(get_type_metadata("b")->is_signed);
+  EXPECT_EQ(get_type_metadata("b")->kind, BitVectorKind::kBitType);
+
+  EXPECT_EQ(std::get<Expr*>(get_type_metadata("c")->bit_count)->ToString(),
+            "111");
+  EXPECT_FALSE(get_type_metadata("c")->is_signed);
+  EXPECT_EQ(get_type_metadata("c")->kind, BitVectorKind::kBitType);
+
+  EXPECT_EQ(std::get<Expr*>(get_type_metadata("d")->bit_count)->ToString(),
+            "32");
+  EXPECT_TRUE(get_type_metadata("d")->is_signed);
+  EXPECT_EQ(get_type_metadata("d")->kind, BitVectorKind::kBitType);
+
+  EXPECT_EQ(std::get<Expr*>(get_type_metadata("e")->bit_count)->ToString(),
+            "312");
+  EXPECT_FALSE(get_type_metadata("e")->is_signed);
+  EXPECT_EQ(get_type_metadata("e")->kind, BitVectorKind::kBitType);
+
+  EXPECT_EQ(std::get<int64_t>(get_type_metadata("f")->bit_count), 1);
+  EXPECT_FALSE(get_type_metadata("f")->is_signed);
+  EXPECT_EQ(get_type_metadata("f")->kind, BitVectorKind::kBitType);
+
+  EXPECT_EQ(std::get<int64_t>(get_type_metadata("g")->bit_count), 37);
+  EXPECT_FALSE(get_type_metadata("g")->is_signed);
+  EXPECT_EQ(get_type_metadata("g")->kind, BitVectorKind::kBitTypeAlias);
+
+  EXPECT_EQ(std::get<int64_t>(get_type_metadata("h")->bit_count), 7);
+  EXPECT_FALSE(get_type_metadata("h")->is_signed);
+  EXPECT_EQ(get_type_metadata("h")->kind, BitVectorKind::kEnumType);
+
+  EXPECT_EQ(std::get<int64_t>(get_type_metadata("i")->bit_count), 7);
+  EXPECT_FALSE(get_type_metadata("i")->is_signed);
+  EXPECT_EQ(get_type_metadata("i")->kind, BitVectorKind::kEnumTypeAlias);
+
+  EXPECT_FALSE(get_type_metadata("j").has_value());
+  EXPECT_FALSE(get_type_metadata("k").has_value());
+  EXPECT_FALSE(get_type_metadata("l").has_value());
+  EXPECT_FALSE(get_type_metadata("m").has_value());
+  EXPECT_FALSE(get_type_metadata("n").has_value());
+}
+
 }  // namespace
 }  // namespace xls::dslx

xls/fuzzer/BUILD

@@ -632,6 +632,7 @@ cc_library(
         "//xls/dslx:interp_value",
         "//xls/dslx/frontend:ast",
         "//xls/dslx/frontend:ast_cloner",
+        "//xls/dslx/frontend:ast_utils",
         "//xls/dslx/frontend:scanner",
         "//xls/ir:bits",
         "//xls/ir:bits_ops",