Merge pull request #1901 from xlsynth:cdleary/2025-01-30-constexpr-match-namerefs

PiperOrigin-RevId: 721824651

Author: copybara-github

xls/dslx/bytecode/bytecode_emitter_test.cc

@@ -508,8 +508,7 @@ fn unops() {
 // Tests array creation.
 TEST(BytecodeEmitterTest, Arrays) {
   constexpr std::string_view kProgram = R"(#[test]
-fn arrays() -> u32[3] {
-  let a = u32:32;
+fn arrays(a: u32) -> u32[3] {
   u32[3]:[u32:0, u32:1, a]
 }
 )";
@@ -519,8 +518,8 @@ fn arrays() -> u32[3] {
                            EmitBytecodes(&import_data, kProgram, "arrays"));
 
   const std::vector<Bytecode>& bytecodes = bf->bytecodes();
-  ASSERT_EQ(bytecodes.size(), 6);
-  const Bytecode* bc = &bytecodes[5];
+  ASSERT_EQ(bytecodes.size(), 4);
+  const Bytecode* bc = &bytecodes[3];
   ASSERT_EQ(bc->op(), Bytecode::Op::kCreateArray);
   ASSERT_TRUE(bc->has_data());
   XLS_ASSERT_OK_AND_ASSIGN(Bytecode::NumElements num_elements,

xls/dslx/frontend/ast.h

@@ -638,7 +638,7 @@ class AnyTypeAnnotation : public TypeAnnotation {
     return {};
   }
 
-  std::string ToString() const { return "Any"; };
+  std::string ToString() const override { return "Any"; };
 };
 
 // Represents an array type annotation; e.g. `u32[5]`.

xls/dslx/type_system/deduce.cc

@@ -1953,6 +1953,11 @@ absl::StatusOr<std::unique_ptr<Type>> DeduceNameRef(const NameRef* node,
   AstNode* name_def = ToAstNode(node->name_def());
   XLS_RET_CHECK(name_def != nullptr);
 
+  if (std::optional<InterpValue> const_expr =
+          ctx->type_info()->GetConstExprOption(name_def)) {
+    ctx->type_info()->NoteConstExpr(node, const_expr.value());
+  }
+
   std::optional<Type*> item = ctx->type_info()->GetItem(name_def);
   if (item.has_value()) {
     auto type = (*item)->CloneToUnique();

xls/dslx/type_system/typecheck_module_test.cc

@@ -4235,6 +4235,56 @@ fn test_f() {
   XLS_EXPECT_OK(Typecheck(kProgram));
 }
 
+TEST(TypecheckTest, MatchPackageLevelConstant) {
+  constexpr std::string_view kProgram = R"(
+const FOO = u8:0xff;
+fn f(x: u8) -> u2 {
+  match x {
+    FOO => u2:0,
+    _ => u2:0,
+  }
+}
+)";
+  XLS_ASSERT_OK_AND_ASSIGN(TypecheckResult result, Typecheck(kProgram));
+  const TypeInfo& type_info = *result.tm.type_info;
+  // Get the pattern match for the first arm of the match expression inside of
+  // function `f`.
+  Function* f = result.tm.module->GetFunction("f").value();
+  Statement* stmt = f->body()->statements()[0];
+  Expr* expr = std::get<Expr*>(stmt->wrapped());
+  Match* m = dynamic_cast<Match*>(expr);
+  ASSERT_NE(m, nullptr);
+  const MatchArm* arm = m->arms()[0];
+  const NameDefTree* pattern = arm->patterns()[0];
+
+  // Check that the pattern is just a leaf NameRef.
+  NameRef* name_ref = std::get<NameRef*>(pattern->leaf());
+  ASSERT_NE(name_ref, nullptr);
+  EXPECT_EQ(name_ref->identifier(), "FOO");
+
+  std::optional<InterpValue> const_expr =
+      type_info.GetConstExprOption(name_ref);
+  ASSERT_TRUE(const_expr.has_value());
+  EXPECT_EQ(const_expr->ToString(), "u8:255");
+}
+
+TEST(TypecheckTest, MatchPackageLevelConstantIntoTypeAlias) {
+  constexpr std::string_view kProgram = R"(
+const C = u32:2;
+fn f(x: u32) -> u2 {
+  match x {
+    C => {
+      const D = C;
+      type T = uN[D];
+      T::MAX
+    },
+    _ => u2:0,
+  }
+}
+)";
+  XLS_EXPECT_OK(Typecheck(kProgram));
+}
+
 // Table-oriented test that lets us validate that *types on parameters* are
 // compatible with *particular values* that should be type-compatible.
 TEST(PassValueToIdentityFnTest, ParameterVsValue) {