Fixes for importing struct constants and related TypeAliases.

#1277 PiperOrigin-RevId: 691892390
google · Oct 31, 2024 · 45204c9 · 45204c9
1 parent 9b67419
commit 45204c9
Show file tree

Hide file tree

Showing 4 changed files with 115 additions and 3 deletions.
diff --git a/xls/dslx/bytecode/bytecode_emitter_test.cc b/xls/dslx/bytecode/bytecode_emitter_test.cc
@@ -880,6 +880,92 @@ fn struct_const_ref() -> u4 {
               IsOkAndHolds(InterpValue::MakeSBits(4, 7)));
 }
 
+TEST(BytecodeEmitterTest, ImportedStructImplConstant) {
+  constexpr std::string_view kImportedProgram = R"(
+pub struct Empty<N: u32> {}
+
+impl Empty<N> {
+  const MY_CONST = uN[N]:7;
+}
+
+  )";
+  constexpr std::string_view kBaseProgram = R"(
+import import_0;
+
+#[test]
+fn struct_const_ref() -> u4 {
+  type MyEmpty = import_0::Empty<u32:4>;
+  MyEmpty::MY_CONST
+}
+)";
+
+  auto import_data = CreateImportDataForTest();
+
+  {
+    XLS_ASSERT_OK_AND_ASSIGN(TypecheckedModule tm,
+                             ParseAndTypecheck(kImportedProgram, "import_0.x",
+                                               "import_0", &import_data));
+  }
+
+  XLS_ASSERT_OK_AND_ASSIGN(
+      TypecheckedModule tm,
+      ParseAndTypecheck(kBaseProgram, "test.x", "test", &import_data));
+
+  XLS_ASSERT_OK_AND_ASSIGN(TestFunction * tf,
+                           tm.module->GetTest("struct_const_ref"));
+  XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr<BytecodeFunction> bf,
+                           BytecodeEmitter::Emit(&import_data, tm.type_info,
+                                                 tf->fn(), ParametricEnv()));
+}
+
+TEST(BytecodeEmitterTest, ImportedStructAliasWithConstant) {
+  constexpr std::string_view kImportedProgram = R"(
+pub struct Empty<N: u32> {}
+
+impl Empty<N> {
+  const MY_CONST = uN[N]:7;
+}
+
+  )";
+  constexpr std::string_view kImportedAlias = R"(
+import import_0;
+
+pub type MyEmpty = import_0::Empty<u32:4>;
+
+  )";
+  constexpr std::string_view kBaseProgram = R"(
+import import_1;
+
+#[test]
+fn struct_const_ref() -> u4 {
+  import_1::MyEmpty::MY_CONST
+}
+)";
+
+  auto import_data = CreateImportDataForTest();
+
+  {
+    XLS_ASSERT_OK_AND_ASSIGN(TypecheckedModule tm,
+                             ParseAndTypecheck(kImportedProgram, "import_0.x",
+                                               "import_0", &import_data));
+  }
+  {
+    XLS_ASSERT_OK_AND_ASSIGN(TypecheckedModule tm,
+                             ParseAndTypecheck(kImportedAlias, "import_1.x",
+                                               "import_1", &import_data));
+  }
+
+  XLS_ASSERT_OK_AND_ASSIGN(
+      TypecheckedModule tm,
+      ParseAndTypecheck(kBaseProgram, "test.x", "test", &import_data));
+
+  XLS_ASSERT_OK_AND_ASSIGN(TestFunction * tf,
+                           tm.module->GetTest("struct_const_ref"));
+  XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr<BytecodeFunction> bf,
+                           BytecodeEmitter::Emit(&import_data, tm.type_info,
+                                                 tf->fn(), ParametricEnv()));
+}
+
 TEST(BytecodeEmitterTest, ImportedConstant) {
   constexpr std::string_view kImportedProgram = R"(pub const MY_CONST = u3:2;)";
   constexpr std::string_view kBaseProgram = R"(

diff --git a/xls/dslx/type_system/deduce_colon_ref.cc b/xls/dslx/type_system/deduce_colon_ref.cc
@@ -342,9 +342,10 @@ absl::StatusOr<std::unique_ptr<Type>> DeduceColonRef(const ColonRef* node,
                 XLS_ASSIGN_OR_RETURN(
                     StructDef * struct_def,
                     DerefToStruct(node->span(), struct_ref->ToString(),
-                                  *struct_ref, ctx->type_info()));
+                                  *struct_ref, subject_type_info));
+
                 return DeduceColonRefToStructType(
-                    struct_def, ctx->type_info()->GetItem(struct_ref), node,
+                    struct_def, subject_type_info->GetItem(struct_ref), node,
                     ctx);
               },
               [&](ColonRef* colon_ref) -> ReturnT {

diff --git a/xls/dslx/type_system/deduce_utils.cc b/xls/dslx/type_system/deduce_utils.cc
@@ -444,9 +444,10 @@ ResolveColonRefSubjectAfterTypeChecking(ImportData* import_data,
             return impl.value();
           },
           [&](TypeRefTypeAnnotation* x) -> ReturnT {
+            const TypeInfo* ti = *(import_data->GetRootTypeInfo(x->owner()));
             XLS_ASSIGN_OR_RETURN(
                 StructDef * struct_def,
-                DerefToStruct(colon_ref->span(), x->ToString(), *x, type_info));
+                DerefToStruct(colon_ref->span(), x->ToString(), *x, ti));
             XLS_RET_CHECK(struct_def->impl().has_value());
             return struct_def->impl().value();
           },

diff --git a/xls/dslx/type_system/impl_typecheck_test.cc b/xls/dslx/type_system/impl_typecheck_test.cc
@@ -361,5 +361,29 @@ fn main() -> uN[10] {
       ParseAndTypecheck(kProgram, "fake_main_path.x", "main", &import_data));
 }
 
+TEST(TypecheckTest, ImportedTypeAlias) {
+  constexpr std::string_view kImported = R"(
+pub struct Empty<X: u32> { }
+
+impl Empty<X> {
+   const IMPORTED = u32:2 * X;
+}
+
+pub type MyEmpty = Empty<u32:5>;
+)";
+  constexpr std::string_view kProgram = R"(
+import imported;
+
+fn main() -> uN[10] {
+    uN[imported::MyEmpty::IMPORTED]:0
+})";
+  auto import_data = CreateImportDataForTest();
+  XLS_ASSERT_OK_AND_ASSIGN(
+      TypecheckedModule module,
+      ParseAndTypecheck(kImported, "imported.x", "imported", &import_data));
+  XLS_EXPECT_OK(
+      ParseAndTypecheck(kProgram, "fake_main_path.x", "main", &import_data));
+}
+
 }  // namespace
 }  // namespace xls::dslx