From afc720d6a35616c894b6aa342c49554afe0f35cd Mon Sep 17 00:00:00 2001
From: Rich McKeever <mckeever@google.com>
Date: Wed, 4 Sep 2024 18:04:02 -0700
Subject: [PATCH] Fix ConcretizeStructAnnotation to deal with deps of default
 parametric exprs.

This partially fixes https://github.com/google/xls/issues/1523, and the IR converter test added here is a derivative of the code in that issue, which is completely fixed by this change.

PiperOrigin-RevId: 671166194
---
 xls/dslx/ir_convert/ir_converter_test.cc      | 22 ++++++++++++
 ...onverter_test_ParametricDefaultInStruct.ir | 12 +++++++
 xls/dslx/type_system/deduce.cc                | 35 +++++++++++++++----
 3 files changed, 63 insertions(+), 6 deletions(-)
 create mode 100644 xls/dslx/ir_convert/testdata/ir_converter_test_ParametricDefaultInStruct.ir

diff --git a/xls/dslx/ir_convert/ir_converter_test.cc b/xls/dslx/ir_convert/ir_converter_test.cc
index 5169bb7154..d78e3a8238 100644
--- a/xls/dslx/ir_convert/ir_converter_test.cc
+++ b/xls/dslx/ir_convert/ir_converter_test.cc
@@ -877,6 +877,28 @@ fn f(outer_thing_1: u32, outer_thing_2: u32) -> u32 {
   ExpectIr(converted, TestName());
 }
 
+TEST(IrConverterTest, ParametricDefaultInStruct) {
+  const char* kProgram = R"(
+struct Foo <X: u32, Y: u32 = {X + u32:1}> {
+    a: uN[X],
+    b: uN[Y],
+}
+
+fn make_zero_foo<X: u32>() -> Foo<X> {
+  zero!<Foo<X>>()
+}
+
+fn test() -> Foo<u32:5> {
+ make_zero_foo<u32:5>()
+}
+)";
+
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::string converted,
+      ConvertModuleForTest(kProgram, ConvertOptions{.emit_positions = false}));
+  ExpectIr(converted, TestName());
+}
+
 TEST(IrConverterTest, UnrollForSimple) {
   const char* kProgram = R"(
 fn test() -> u32 {
diff --git a/xls/dslx/ir_convert/testdata/ir_converter_test_ParametricDefaultInStruct.ir b/xls/dslx/ir_convert/testdata/ir_converter_test_ParametricDefaultInStruct.ir
new file mode 100644
index 0000000000..c7d2181508
--- /dev/null
+++ b/xls/dslx/ir_convert/testdata/ir_converter_test_ParametricDefaultInStruct.ir
@@ -0,0 +1,12 @@
+package test_module
+
+file_number 0 "test_module.x"
+
+fn __test_module__make_zero_foo__5() -> (bits[5], bits[6]) {
+  X: bits[32] = literal(value=5, id=1)
+  ret literal.2: (bits[5], bits[6]) = literal(value=(0, 0), id=2)
+}
+
+fn __test_module__test() -> (bits[5], bits[6]) {
+  ret invoke.3: (bits[5], bits[6]) = invoke(to_apply=__test_module__make_zero_foo__5, id=3)
+}
diff --git a/xls/dslx/type_system/deduce.cc b/xls/dslx/type_system/deduce.cc
index 6c9dc78103..4b551824ed 100644
--- a/xls/dslx/type_system/deduce.cc
+++ b/xls/dslx/type_system/deduce.cc
@@ -86,6 +86,29 @@ absl::StatusOr<InterpValue> EvaluateConstexprValue(DeduceCtx* ctx,
       ctx->GetCurrentParametricEnv(), node, type);
 }
 
+// Evaluates the given `dim` to the extent possible with the given `env`.
+// Examples:
+//   - If `dim` is already a constant, the result is `dim`.
+//   - If `dim` is `X + 1`, and `env` says `X` == `5`, then the result is `6`.
+//   - If `dim` is `X + 1`, and `env` says `X` == `Y` and does not have `Y`,
+//     then the result is `Y + 1`.
+//   - If `dim` is `X + 1`, and `env` says `X` == `Y` and `Y` == `5`, then the
+//     result is `6`.
+TypeDim EvaluateTypeDim(TypeDim dim, const ParametricExpression::Env& env) {
+  absl::flat_hash_set<std::string> prev_free_variables;
+  while (std::holds_alternative<TypeDim::OwnedParametric>(dim.value())) {
+    auto& parametric = std::get<TypeDim::OwnedParametric>(dim.value());
+    absl::flat_hash_set<std::string> next_free_variables =
+        parametric->GetFreeVariables();
+    if (prev_free_variables == next_free_variables) {
+      break;
+    }
+    prev_free_variables = std::move(next_free_variables);
+    dim = TypeDim(parametric->Evaluate(env));
+  }
+  return dim;
+}
+
 // Attempts to convert an expression from the full DSL AST into the
 // ParametricExpression sub-AST (a limited form that we can embed into a
 // TypeDim for later instantiation).
@@ -1683,6 +1706,10 @@ static absl::StatusOr<std::unique_ptr<Type>> ConcretizeStructAnnotation(
                           defined_parametric->identifier(),
                           struct_def->identifier()));
     }
+    XLS_ASSIGN_OR_RETURN(std::unique_ptr<ParametricExpression> parametric_expr,
+                         ExprToParametric(defined_parametric->expr(), ctx));
+    parametric_env.emplace(defined_parametric->identifier(),
+                           TypeDim(std::move(parametric_expr)));
   }
 
   ParametricExpression::Env env;
@@ -1698,16 +1725,12 @@ static absl::StatusOr<std::unique_ptr<Type>> ConcretizeStructAnnotation(
   XLS_ASSIGN_OR_RETURN(
       std::unique_ptr<Type> mapped_type,
       base_type.MapSize([&](const TypeDim& dim) -> absl::StatusOr<TypeDim> {
-        if (std::holds_alternative<TypeDim::OwnedParametric>(dim.value())) {
-          auto& parametric = std::get<TypeDim::OwnedParametric>(dim.value());
-          return TypeDim(parametric->Evaluate(env));
-        }
-        return dim;
+        return EvaluateTypeDim(dim, env);
       }));
 
   // Attach the nominal parametrics to the type, so that we will remember the
   // fact that we have instantiated e.g. Foo<M:u32, N:u32> as Foo<5, 6>.
-  return mapped_type->AddNominalTypeDims(std::move(parametric_env));
+  return mapped_type->AddNominalTypeDims(parametric_env);
 }
 
 absl::StatusOr<std::unique_ptr<Type>> DeduceTypeRefTypeAnnotation(