From 007e8f6d0fd33e87aa80b285d8d6378c32fd697a Mon Sep 17 00:00:00 2001 From: Sirui Lu Date: Sat, 30 Aug 2025 00:11:32 -0700 Subject: [PATCH 1/5] Fix duplicated cloned type entities bug --- xls/dslx/frontend/ast_cloner.cc | 87 +++++++- xls/dslx/frontend/ast_cloner_test.cc | 308 ++++++++++++++++++++++++++- 2 files changed, 386 insertions(+), 9 deletions(-) diff --git a/xls/dslx/frontend/ast_cloner.cc b/xls/dslx/frontend/ast_cloner.cc index db137831f6..9cc17c9994 100644 --- a/xls/dslx/frontend/ast_cloner.cc +++ b/xls/dslx/frontend/ast_cloner.cc @@ -274,6 +274,9 @@ class AstCloner : public AstNodeVisitor { } absl::Status HandleEnumDef(const EnumDef* n) override { + if (old_to_new_.contains(n)) { + return absl::OkStatus(); + } XLS_RETURN_IF_ERROR(VisitChildren(n)); NameDef* new_name_def = down_cast(old_to_new_.at(n->name_def())); @@ -845,6 +848,9 @@ class AstCloner : public AstNodeVisitor { } absl::Status HandleStructDef(const StructDef* n) override { + if (old_to_new_.contains(n)) { + return absl::OkStatus(); + } absl::Status status = HandleStructDefBaseInternal(n); if (status.ok()) { AstNode* new_node = old_to_new_.at(n); @@ -859,6 +865,9 @@ class AstCloner : public AstNodeVisitor { } absl::Status HandleProcDef(const ProcDef* n) override { + if (old_to_new_.contains(n)) { + return absl::OkStatus(); + } return HandleStructDefBaseInternal(n); } @@ -985,6 +994,13 @@ class AstCloner : public AstNodeVisitor { } absl::Status HandleTypeAlias(const TypeAlias* n) override { + if (old_to_new_.contains(n)) { + // Ensure the NameDef definer points at the already-cloned alias. + NameDef* existing_name_def = + down_cast(old_to_new_.at(&n->name_def())); + existing_name_def->set_definer(old_to_new_.at(n)); + return absl::OkStatus(); + } XLS_RETURN_IF_ERROR(VisitChildren(n)); NameDef* new_name_def = down_cast(old_to_new_.at(&n->name_def())); @@ -1003,14 +1019,71 @@ class AstCloner : public AstNodeVisitor { absl::Status HandleTypeRef(const TypeRef* n) override { TypeDefinition new_type_definition = n->type_definition(); - // A TypeRef doesn't own its referenced type definition, so we have to - // explicitly visit it. + // A TypeRef doesn't own its referenced type definition. Reuse an existing + // type definition in the target module (by identifier) to avoid + // duplicating type defs when cloning subtrees into an already-cloned + // module. If none exists yet, clone the referenced definition now. XLS_RETURN_IF_ERROR(absl::visit( - Visitor{[&](auto* ref) -> absl::Status { - XLS_RETURN_IF_ERROR(ReplaceOrVisit(ref)); - new_type_definition = down_cast(old_to_new_.at(ref)); - return absl::OkStatus(); - }}, + Visitor{ + // Reuse an existing same-named type definition already present + // in the target module for module-scoped type defs. + [&](TypeAlias* ref) -> absl::Status { + absl::StatusOr existing = + module(n)->GetTypeDefinition(ref->identifier()); + if (existing.ok()) { + new_type_definition = *existing; + return absl::OkStatus(); + } + XLS_RETURN_IF_ERROR(ReplaceOrVisit(ref)); + new_type_definition = + down_cast(old_to_new_.at(ref)); + return absl::OkStatus(); + }, + [&](StructDef* ref) -> absl::Status { + absl::StatusOr existing = + module(n)->GetTypeDefinition(ref->identifier()); + if (existing.ok()) { + new_type_definition = *existing; + return absl::OkStatus(); + } + XLS_RETURN_IF_ERROR(ReplaceOrVisit(ref)); + new_type_definition = + down_cast(old_to_new_.at(ref)); + return absl::OkStatus(); + }, + [&](ProcDef* ref) -> absl::Status { + absl::StatusOr existing = + module(n)->GetTypeDefinition(ref->identifier()); + if (existing.ok()) { + new_type_definition = *existing; + return absl::OkStatus(); + } + XLS_RETURN_IF_ERROR(ReplaceOrVisit(ref)); + new_type_definition = + down_cast(old_to_new_.at(ref)); + return absl::OkStatus(); + }, + [&](EnumDef* ref) -> absl::Status { + absl::StatusOr existing = + module(n)->GetTypeDefinition(ref->identifier()); + if (existing.ok()) { + new_type_definition = *existing; + return absl::OkStatus(); + } + XLS_RETURN_IF_ERROR(ReplaceOrVisit(ref)); + new_type_definition = + down_cast(old_to_new_.at(ref)); + return absl::OkStatus(); + }, + // For ColonRef* and UseTreeEntry* we cannot look up by identifier + // in the target module; just clone/visit as before. + [&](auto* ref) -> absl::Status { + XLS_RETURN_IF_ERROR(ReplaceOrVisit(ref)); + new_type_definition = + down_cast(old_to_new_.at(ref)); + return absl::OkStatus(); + }, + }, n->type_definition())); old_to_new_[n] = module(n)->Make(n->span(), new_type_definition); diff --git a/xls/dslx/frontend/ast_cloner_test.cc b/xls/dslx/frontend/ast_cloner_test.cc index 5bf6826be2..b0ffafa380 100644 --- a/xls/dslx/frontend/ast_cloner_test.cc +++ b/xls/dslx/frontend/ast_cloner_test.cc @@ -21,12 +21,12 @@ #include #include -#include "gmock/gmock.h" -#include "gtest/gtest.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "gmock/gmock.h" +#include "gtest/gtest.h" #include "xls/common/casts.h" #include "xls/common/status/matchers.h" #include "xls/common/status/status_macros.h" @@ -1634,6 +1634,310 @@ fn bar() -> u32{ EXPECT_EQ(orig_ref->name_def(), new_ref->name_def()); } +TEST(AstClonerTest, TypeRefDoesntCloneModuleTypeAlias) { + constexpr std::string_view kProgram = R"( +type MyU32 = u32; + +fn foo(x: MyU32) -> MyU32 { + x +} +)"; + + FileTable file_table; + XLS_ASSERT_OK_AND_ASSIGN(auto module, ParseModule(kProgram, "fake_path.x", + "the_module", file_table)); + XLS_ASSERT_OK_AND_ASSIGN(Function * foo, + module->GetMemberOrError("foo")); + XLS_ASSERT_OK_AND_ASSIGN(TypeAlias * alias, + module->GetMemberOrError("MyU32")); + + // Clone just the function into the same target module. + XLS_ASSERT_OK_AND_ASSIGN(AstNode * clone, CloneAst(foo)); + Function* cloned_foo = down_cast(clone); + + // Param type should reference the existing module-level alias, not a cloned + // one. + auto* param_trta = down_cast( + cloned_foo->params()[0]->type_annotation()); + TypeRef* param_tr = param_trta->type_ref(); + ASSERT_TRUE(std::holds_alternative(param_tr->type_definition())); + EXPECT_EQ(std::get(param_tr->type_definition()), alias); + + // Return type should also reference the existing alias. + auto* ret_trta = down_cast(cloned_foo->return_type()); + TypeRef* ret_tr = ret_trta->type_ref(); + ASSERT_TRUE(std::holds_alternative(ret_tr->type_definition())); + EXPECT_EQ(std::get(ret_tr->type_definition()), alias); +} + +TEST(AstClonerTest, EnumTypeRefReusesTopLevelEnumInCloneModule) { + constexpr std::string_view kProgram = R"( +enum MyEnum : u32 { + A = 0, + B = 1, +} + +fn id(x: MyEnum) -> MyEnum { x } +)"; + + FileTable file_table; + XLS_ASSERT_OK_AND_ASSIGN( + std::unique_ptr module, + ParseModule(kProgram, "fake_path.x", "the_module", file_table)); + + // Clone the whole module. + XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr clone, + CloneModule(*module.get())); + + // Get the cloned top-level enum and function. + XLS_ASSERT_OK_AND_ASSIGN(EnumDef * enum_def, + clone->GetMemberOrError("MyEnum")); + XLS_ASSERT_OK_AND_ASSIGN(Function * id, + clone->GetMemberOrError("id")); + + // Param type should reference the cloned top-level enum, not a detached copy. + auto* param_trta = + down_cast(id->params()[0]->type_annotation()); + TypeRef* param_tr = param_trta->type_ref(); + ASSERT_TRUE(std::holds_alternative(param_tr->type_definition())); + EXPECT_EQ(std::get(param_tr->type_definition()), enum_def); + + // Return type should also reference the same cloned top-level enum. + auto* ret_trta = down_cast(id->return_type()); + TypeRef* ret_tr = ret_trta->type_ref(); + ASSERT_TRUE(std::holds_alternative(ret_tr->type_definition())); + EXPECT_EQ(std::get(ret_tr->type_definition()), enum_def); +} + +TEST(AstClonerTest, TypeRefReusesTargetModuleAliasWhenCloningIntoOtherModule) { + constexpr std::string_view kSrc = R"( +type MyU32 = u32; +fn foo(x: MyU32) -> MyU32 { x } +)"; + constexpr std::string_view kDst = R"( +type MyU32 = u32; +)"; + + FileTable ft_src; + XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr src, + ParseModule(kSrc, "src.x", "src_mod", ft_src)); + XLS_ASSERT_OK_AND_ASSIGN(Function * foo, + src->GetMemberOrError("foo")); + + FileTable ft_dst; + XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr dst, + ParseModule(kDst, "dst.x", "dst_mod", ft_dst)); + XLS_ASSERT_OK_AND_ASSIGN(TypeAlias * dst_alias, + dst->GetMemberOrError("MyU32")); + + absl::flat_hash_map clones; + XLS_ASSERT_OK_AND_ASSIGN( + clones, CloneAstAndGetAllPairs(foo, dst.get(), &NoopCloneReplacer)); + + auto* cloned_foo = down_cast(clones.at(foo)); + auto* param_trta = down_cast( + cloned_foo->params()[0]->type_annotation()); + TypeRef* param_tr = param_trta->type_ref(); + ASSERT_TRUE(std::holds_alternative(param_tr->type_definition())); + EXPECT_EQ(std::get(param_tr->type_definition()), dst_alias); + + auto* ret_trta = down_cast(cloned_foo->return_type()); + TypeRef* ret_tr = ret_trta->type_ref(); + ASSERT_TRUE(std::holds_alternative(ret_tr->type_definition())); + EXPECT_EQ(std::get(ret_tr->type_definition()), dst_alias); +} + +TEST(AstClonerTest, TypeRefClonesAliasWhenTargetHasNone) { + constexpr std::string_view kSrc = R"( +type MyU32 = u32; +fn foo(x: MyU32) -> MyU32 { x } +)"; + + FileTable ft_src; + XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr src, + ParseModule(kSrc, "src.x", "src_mod", ft_src)); + XLS_ASSERT_OK_AND_ASSIGN(Function * foo, + src->GetMemberOrError("foo")); + XLS_ASSERT_OK_AND_ASSIGN(TypeAlias * src_alias, + src->GetMemberOrError("MyU32")); + + // Empty destination module. + FileTable ft_dst; + std::unique_ptr dst = + std::make_unique("dst_mod", std::nullopt, ft_dst); + + absl::flat_hash_map clones; + XLS_ASSERT_OK_AND_ASSIGN( + clones, CloneAstAndGetAllPairs(foo, dst.get(), &NoopCloneReplacer)); + + auto* cloned_foo = down_cast(clones.at(foo)); + auto* param_trta = down_cast( + cloned_foo->params()[0]->type_annotation()); + TypeRef* param_tr = param_trta->type_ref(); + ASSERT_TRUE(std::holds_alternative(param_tr->type_definition())); + + // Ensure it's using the cloned alias in the destination module, not the + // source. + TypeAlias* cloned_alias = down_cast(clones.at(src_alias)); + EXPECT_NE(cloned_alias, src_alias); + EXPECT_EQ(std::get(param_tr->type_definition()), cloned_alias); +} + +TEST(AstClonerTest, TypeRefFromUseTreeEntry) { + constexpr std::string_view kProgram = R"( +#![feature(use_syntax)] + +use foo::ImportedStruct; + +fn main(x: ImportedStruct) { + () +})"; + + constexpr std::string_view kExpectedFunction = + R"(fn main(x: ImportedStruct) { + () +})"; + + FileTable file_table; + XLS_ASSERT_OK_AND_ASSIGN(auto module, ParseModule(kProgram, "fake_path.x", + "the_module", file_table)); + XLS_ASSERT_OK_AND_ASSIGN(Function * f, + module->GetMemberOrError("main")); + + // Ensure the parsed TypeRef is backed by a UseTreeEntry. + std::optional type_ref = FindFirstTypeRef(f); + ASSERT_TRUE(type_ref.has_value()); + ASSERT_TRUE( + std::holds_alternative((*type_ref)->type_definition())); + + // Clone and ensure we preserve the UseTreeEntry variant and string form. + XLS_ASSERT_OK_AND_ASSIGN(AstNode * clone, CloneAst(f)); + EXPECT_EQ(kExpectedFunction, clone->ToString()); + XLS_ASSERT_OK(VerifyClone(f, clone, *module->file_table())); + + std::optional cloned_type_ref = FindFirstTypeRef(clone); + ASSERT_TRUE(cloned_type_ref.has_value()); + ASSERT_TRUE(std::holds_alternative( + (*cloned_type_ref)->type_definition())); +} + +TEST(AstClonerTest, TypeRefFromColonRef) { + constexpr std::string_view kProgram = R"( +import my.module as foo; + +fn main(x: foo::ImportedStruct) { + () +})"; + + constexpr std::string_view kExpectedFunction = + R"(fn main(x: foo::ImportedStruct) { + () +})"; + + FileTable file_table; + XLS_ASSERT_OK_AND_ASSIGN(auto module, ParseModule(kProgram, "fake_path.x", + "the_module", file_table)); + XLS_ASSERT_OK_AND_ASSIGN(Function * f, + module->GetMemberOrError("main")); + + // Ensure the parsed TypeRef is backed by a ColonRef. + std::optional type_ref = FindFirstTypeRef(f); + ASSERT_TRUE(type_ref.has_value()); + ASSERT_TRUE(std::holds_alternative((*type_ref)->type_definition())); + + // Clone and ensure we preserve the ColonRef variant and string form. + XLS_ASSERT_OK_AND_ASSIGN(AstNode * clone, CloneAst(f)); + EXPECT_EQ(kExpectedFunction, clone->ToString()); + XLS_ASSERT_OK(VerifyClone(f, clone, *module->file_table())); + + std::optional cloned_type_ref = FindFirstTypeRef(clone); + ASSERT_TRUE(cloned_type_ref.has_value()); + ASSERT_TRUE(std::holds_alternative( + (*cloned_type_ref)->type_definition())); +} + +TEST(AstClonerTest, TypeRefFromStructDef) { + constexpr std::string_view kProgram = R"( +struct MyStruct { a: u32 } + +fn id(x: MyStruct) -> MyStruct { x } +)"; + + constexpr std::string_view kExpectedFunction = + R"(fn id(x: MyStruct) -> MyStruct { + x +})"; + + FileTable file_table; + XLS_ASSERT_OK_AND_ASSIGN(auto module, ParseModule(kProgram, "fake_path.x", + "the_module", file_table)); + + XLS_ASSERT_OK_AND_ASSIGN(StructDef * struct_def, + module->GetMemberOrError("MyStruct")); + XLS_ASSERT_OK_AND_ASSIGN(Function * id, + module->GetMemberOrError("id")); + + // Param and return types should be backed by the StructDef. + auto* param_trta = + down_cast(id->params()[0]->type_annotation()); + TypeRef* param_tr = param_trta->type_ref(); + ASSERT_TRUE(std::holds_alternative(param_tr->type_definition())); + EXPECT_EQ(std::get(param_tr->type_definition()), struct_def); + + auto* ret_trta = down_cast(id->return_type()); + TypeRef* ret_tr = ret_trta->type_ref(); + ASSERT_TRUE(std::holds_alternative(ret_tr->type_definition())); + EXPECT_EQ(std::get(ret_tr->type_definition()), struct_def); + + // Clone and ensure ToString and variant preservation. + XLS_ASSERT_OK_AND_ASSIGN(AstNode * clone, CloneAst(id)); + EXPECT_EQ(kExpectedFunction, clone->ToString()); + XLS_ASSERT_OK(VerifyClone(id, clone, *module->file_table())); + + auto* cloned_id = down_cast(clone); + auto* cloned_param_trta = down_cast( + cloned_id->params()[0]->type_annotation()); + ASSERT_TRUE(std::holds_alternative( + cloned_param_trta->type_ref()->type_definition())); +} + +TEST(AstClonerTest, TypeRefFromProcDefInImpl) { + constexpr std::string_view kProgram = R"( +proc MyProc { + a: u32, +} + +impl MyProc {} +)"; + + constexpr std::string_view kExpectedImpl = R"(impl MyProc { +})"; + + FileTable file_table; + XLS_ASSERT_OK_AND_ASSIGN(auto module, ParseModule(kProgram, "fake_path.x", + "the_module", file_table)); + + XLS_ASSERT_OK_AND_ASSIGN(ProcDef * proc_def, + module->GetMemberOrError("MyProc")); + Impl* impl = module->GetImpls().at(0); + + // The impl's struct_ref should be a TypeRefTypeAnnotation backed by ProcDef. + auto* trta = down_cast(impl->struct_ref()); + TypeRef* tr = trta->type_ref(); + ASSERT_TRUE(std::holds_alternative(tr->type_definition())); + EXPECT_EQ(std::get(tr->type_definition()), proc_def); + + // Clone and ensure ToString and variant preservation. + XLS_ASSERT_OK_AND_ASSIGN(AstNode * clone, CloneAst(impl)); + EXPECT_EQ(kExpectedImpl, clone->ToString()); + XLS_ASSERT_OK(VerifyClone(impl, clone, *module->file_table())); + + auto* cloned_impl = down_cast(clone); + auto* cloned_trta = down_cast(cloned_impl->struct_ref()); + ASSERT_TRUE(std::holds_alternative( + cloned_trta->type_ref()->type_definition())); +} + TEST(AstClonerTest, CloneAstClonesVerbatimNode) { constexpr std::string_view kProgram = "const FOO = u32:42;"; FileTable file_table; From 22f4a81afb4a6475b33b247cb408e9ae60d4da3b Mon Sep 17 00:00:00 2001 From: Sirui Lu Date: Mon, 1 Sep 2025 12:38:07 -0700 Subject: [PATCH 2/5] Remove the buggy test cases --- xls/dslx/frontend/ast_cloner.cc | 71 +++------------------------- xls/dslx/frontend/ast_cloner_test.cc | 60 +++++------------------ 2 files changed, 19 insertions(+), 112 deletions(-) diff --git a/xls/dslx/frontend/ast_cloner.cc b/xls/dslx/frontend/ast_cloner.cc index 9cc17c9994..faf423d9b1 100644 --- a/xls/dslx/frontend/ast_cloner.cc +++ b/xls/dslx/frontend/ast_cloner.cc @@ -1019,71 +1019,14 @@ class AstCloner : public AstNodeVisitor { absl::Status HandleTypeRef(const TypeRef* n) override { TypeDefinition new_type_definition = n->type_definition(); - // A TypeRef doesn't own its referenced type definition. Reuse an existing - // type definition in the target module (by identifier) to avoid - // duplicating type defs when cloning subtrees into an already-cloned - // module. If none exists yet, clone the referenced definition now. + // A TypeRef doesn't own its referenced type definition, so we have to + // explicitly visit it. XLS_RETURN_IF_ERROR(absl::visit( - Visitor{ - // Reuse an existing same-named type definition already present - // in the target module for module-scoped type defs. - [&](TypeAlias* ref) -> absl::Status { - absl::StatusOr existing = - module(n)->GetTypeDefinition(ref->identifier()); - if (existing.ok()) { - new_type_definition = *existing; - return absl::OkStatus(); - } - XLS_RETURN_IF_ERROR(ReplaceOrVisit(ref)); - new_type_definition = - down_cast(old_to_new_.at(ref)); - return absl::OkStatus(); - }, - [&](StructDef* ref) -> absl::Status { - absl::StatusOr existing = - module(n)->GetTypeDefinition(ref->identifier()); - if (existing.ok()) { - new_type_definition = *existing; - return absl::OkStatus(); - } - XLS_RETURN_IF_ERROR(ReplaceOrVisit(ref)); - new_type_definition = - down_cast(old_to_new_.at(ref)); - return absl::OkStatus(); - }, - [&](ProcDef* ref) -> absl::Status { - absl::StatusOr existing = - module(n)->GetTypeDefinition(ref->identifier()); - if (existing.ok()) { - new_type_definition = *existing; - return absl::OkStatus(); - } - XLS_RETURN_IF_ERROR(ReplaceOrVisit(ref)); - new_type_definition = - down_cast(old_to_new_.at(ref)); - return absl::OkStatus(); - }, - [&](EnumDef* ref) -> absl::Status { - absl::StatusOr existing = - module(n)->GetTypeDefinition(ref->identifier()); - if (existing.ok()) { - new_type_definition = *existing; - return absl::OkStatus(); - } - XLS_RETURN_IF_ERROR(ReplaceOrVisit(ref)); - new_type_definition = - down_cast(old_to_new_.at(ref)); - return absl::OkStatus(); - }, - // For ColonRef* and UseTreeEntry* we cannot look up by identifier - // in the target module; just clone/visit as before. - [&](auto* ref) -> absl::Status { - XLS_RETURN_IF_ERROR(ReplaceOrVisit(ref)); - new_type_definition = - down_cast(old_to_new_.at(ref)); - return absl::OkStatus(); - }, - }, + Visitor{[&](auto* ref) -> absl::Status { + XLS_RETURN_IF_ERROR(ReplaceOrVisit(ref)); + new_type_definition = down_cast(old_to_new_.at(ref)); + return absl::OkStatus(); + }}, n->type_definition())); old_to_new_[n] = module(n)->Make(n->span(), new_type_definition); diff --git a/xls/dslx/frontend/ast_cloner_test.cc b/xls/dslx/frontend/ast_cloner_test.cc index b0ffafa380..70d8b6ee4b 100644 --- a/xls/dslx/frontend/ast_cloner_test.cc +++ b/xls/dslx/frontend/ast_cloner_test.cc @@ -1655,19 +1655,19 @@ fn foo(x: MyU32) -> MyU32 { XLS_ASSERT_OK_AND_ASSIGN(AstNode * clone, CloneAst(foo)); Function* cloned_foo = down_cast(clone); - // Param type should reference the existing module-level alias, not a cloned - // one. + // Param type and the return type should reference the same alias. auto* param_trta = down_cast( cloned_foo->params()[0]->type_annotation()); TypeRef* param_tr = param_trta->type_ref(); ASSERT_TRUE(std::holds_alternative(param_tr->type_definition())); - EXPECT_EQ(std::get(param_tr->type_definition()), alias); - - // Return type should also reference the existing alias. + EXPECT_EQ(param_tr->owner(), clone->owner()); + TypeAlias* param_alias = std::get(param_tr->type_definition()); + EXPECT_EQ(param_alias->owner(), clone->owner()); auto* ret_trta = down_cast(cloned_foo->return_type()); TypeRef* ret_tr = ret_trta->type_ref(); ASSERT_TRUE(std::holds_alternative(ret_tr->type_definition())); - EXPECT_EQ(std::get(ret_tr->type_definition()), alias); + EXPECT_EQ(std::get(param_tr->type_definition()), + std::get(ret_tr->type_definition())); } TEST(AstClonerTest, EnumTypeRefReusesTopLevelEnumInCloneModule) { @@ -1709,44 +1709,6 @@ fn id(x: MyEnum) -> MyEnum { x } EXPECT_EQ(std::get(ret_tr->type_definition()), enum_def); } -TEST(AstClonerTest, TypeRefReusesTargetModuleAliasWhenCloningIntoOtherModule) { - constexpr std::string_view kSrc = R"( -type MyU32 = u32; -fn foo(x: MyU32) -> MyU32 { x } -)"; - constexpr std::string_view kDst = R"( -type MyU32 = u32; -)"; - - FileTable ft_src; - XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr src, - ParseModule(kSrc, "src.x", "src_mod", ft_src)); - XLS_ASSERT_OK_AND_ASSIGN(Function * foo, - src->GetMemberOrError("foo")); - - FileTable ft_dst; - XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr dst, - ParseModule(kDst, "dst.x", "dst_mod", ft_dst)); - XLS_ASSERT_OK_AND_ASSIGN(TypeAlias * dst_alias, - dst->GetMemberOrError("MyU32")); - - absl::flat_hash_map clones; - XLS_ASSERT_OK_AND_ASSIGN( - clones, CloneAstAndGetAllPairs(foo, dst.get(), &NoopCloneReplacer)); - - auto* cloned_foo = down_cast(clones.at(foo)); - auto* param_trta = down_cast( - cloned_foo->params()[0]->type_annotation()); - TypeRef* param_tr = param_trta->type_ref(); - ASSERT_TRUE(std::holds_alternative(param_tr->type_definition())); - EXPECT_EQ(std::get(param_tr->type_definition()), dst_alias); - - auto* ret_trta = down_cast(cloned_foo->return_type()); - TypeRef* ret_tr = ret_trta->type_ref(); - ASSERT_TRUE(std::holds_alternative(ret_tr->type_definition())); - EXPECT_EQ(std::get(ret_tr->type_definition()), dst_alias); -} - TEST(AstClonerTest, TypeRefClonesAliasWhenTargetHasNone) { constexpr std::string_view kSrc = R"( type MyU32 = u32; @@ -1843,7 +1805,8 @@ fn main(x: foo::ImportedStruct) { // Ensure the parsed TypeRef is backed by a ColonRef. std::optional type_ref = FindFirstTypeRef(f); ASSERT_TRUE(type_ref.has_value()); - ASSERT_TRUE(std::holds_alternative((*type_ref)->type_definition())); + ASSERT_TRUE( + std::holds_alternative((*type_ref)->type_definition())); // Clone and ensure we preserve the ColonRef variant and string form. XLS_ASSERT_OK_AND_ASSIGN(AstNode * clone, CloneAst(f)); @@ -1852,8 +1815,8 @@ fn main(x: foo::ImportedStruct) { std::optional cloned_type_ref = FindFirstTypeRef(clone); ASSERT_TRUE(cloned_type_ref.has_value()); - ASSERT_TRUE(std::holds_alternative( - (*cloned_type_ref)->type_definition())); + ASSERT_TRUE( + std::holds_alternative((*cloned_type_ref)->type_definition())); } TEST(AstClonerTest, TypeRefFromStructDef) { @@ -1933,7 +1896,8 @@ impl MyProc {} XLS_ASSERT_OK(VerifyClone(impl, clone, *module->file_table())); auto* cloned_impl = down_cast(clone); - auto* cloned_trta = down_cast(cloned_impl->struct_ref()); + auto* cloned_trta = + down_cast(cloned_impl->struct_ref()); ASSERT_TRUE(std::holds_alternative( cloned_trta->type_ref()->type_definition())); } From 0a6f7656d35d8de71cb95c28b02d3fb173c17017 Mon Sep 17 00:00:00 2001 From: Sirui Lu Date: Mon, 1 Sep 2025 12:55:42 -0700 Subject: [PATCH 3/5] Move old_to_new_ check to ReplaceOrVisit, clean up test cases --- xls/dslx/frontend/ast_cloner.cc | 23 +--- xls/dslx/frontend/ast_cloner_test.cc | 155 +-------------------------- 2 files changed, 8 insertions(+), 170 deletions(-) diff --git a/xls/dslx/frontend/ast_cloner.cc b/xls/dslx/frontend/ast_cloner.cc index faf423d9b1..a2dd97ec07 100644 --- a/xls/dslx/frontend/ast_cloner.cc +++ b/xls/dslx/frontend/ast_cloner.cc @@ -274,9 +274,6 @@ class AstCloner : public AstNodeVisitor { } absl::Status HandleEnumDef(const EnumDef* n) override { - if (old_to_new_.contains(n)) { - return absl::OkStatus(); - } XLS_RETURN_IF_ERROR(VisitChildren(n)); NameDef* new_name_def = down_cast(old_to_new_.at(n->name_def())); @@ -848,9 +845,6 @@ class AstCloner : public AstNodeVisitor { } absl::Status HandleStructDef(const StructDef* n) override { - if (old_to_new_.contains(n)) { - return absl::OkStatus(); - } absl::Status status = HandleStructDefBaseInternal(n); if (status.ok()) { AstNode* new_node = old_to_new_.at(n); @@ -865,9 +859,6 @@ class AstCloner : public AstNodeVisitor { } absl::Status HandleProcDef(const ProcDef* n) override { - if (old_to_new_.contains(n)) { - return absl::OkStatus(); - } return HandleStructDefBaseInternal(n); } @@ -994,13 +985,6 @@ class AstCloner : public AstNodeVisitor { } absl::Status HandleTypeAlias(const TypeAlias* n) override { - if (old_to_new_.contains(n)) { - // Ensure the NameDef definer points at the already-cloned alias. - NameDef* existing_name_def = - down_cast(old_to_new_.at(&n->name_def())); - existing_name_def->set_definer(old_to_new_.at(n)); - return absl::OkStatus(); - } XLS_RETURN_IF_ERROR(VisitChildren(n)); NameDef* new_name_def = down_cast(old_to_new_.at(&n->name_def())); @@ -1206,9 +1190,7 @@ class AstCloner : public AstNodeVisitor { // already been processed. absl::Status VisitChildren(const AstNode* node) { for (const auto& child : node->GetChildren(/*want_types=*/true)) { - if (!old_to_new_.contains(child)) { - XLS_RETURN_IF_ERROR(ReplaceOrVisit(child)); - } + XLS_RETURN_IF_ERROR(ReplaceOrVisit(child)); } return absl::OkStatus(); } @@ -1217,6 +1199,9 @@ class AstCloner : public AstNodeVisitor { if (node == nullptr) { return absl::OkStatus(); } + if (old_to_new_.contains(node)) { + return absl::OkStatus(); + } XLS_ASSIGN_OR_RETURN(std::optional replacement, replacer_(node)); if (replacement.has_value()) { old_to_new_[node] = *replacement; diff --git a/xls/dslx/frontend/ast_cloner_test.cc b/xls/dslx/frontend/ast_cloner_test.cc index 70d8b6ee4b..4bfa63da55 100644 --- a/xls/dslx/frontend/ast_cloner_test.cc +++ b/xls/dslx/frontend/ast_cloner_test.cc @@ -1634,7 +1634,7 @@ fn bar() -> u32{ EXPECT_EQ(orig_ref->name_def(), new_ref->name_def()); } -TEST(AstClonerTest, TypeRefDoesntCloneModuleTypeAlias) { +TEST(AstClonerTest, DoesntCloneTypeAliasTwice) { constexpr std::string_view kProgram = R"( type MyU32 = u32; @@ -1670,7 +1670,7 @@ fn foo(x: MyU32) -> MyU32 { std::get(ret_tr->type_definition())); } -TEST(AstClonerTest, EnumTypeRefReusesTopLevelEnumInCloneModule) { +TEST(AstClonerTest, DoesntCloneEnumTwice) { constexpr std::string_view kProgram = R"( enum MyEnum : u32 { A = 0, @@ -1688,6 +1688,7 @@ fn id(x: MyEnum) -> MyEnum { x } // Clone the whole module. XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr clone, CloneModule(*module.get())); + XLS_ASSERT_OK(VerifyClone(module.get(), clone.get(), file_table)); // Get the cloned top-level enum and function. XLS_ASSERT_OK_AND_ASSIGN(EnumDef * enum_def, @@ -1709,117 +1710,7 @@ fn id(x: MyEnum) -> MyEnum { x } EXPECT_EQ(std::get(ret_tr->type_definition()), enum_def); } -TEST(AstClonerTest, TypeRefClonesAliasWhenTargetHasNone) { - constexpr std::string_view kSrc = R"( -type MyU32 = u32; -fn foo(x: MyU32) -> MyU32 { x } -)"; - - FileTable ft_src; - XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr src, - ParseModule(kSrc, "src.x", "src_mod", ft_src)); - XLS_ASSERT_OK_AND_ASSIGN(Function * foo, - src->GetMemberOrError("foo")); - XLS_ASSERT_OK_AND_ASSIGN(TypeAlias * src_alias, - src->GetMemberOrError("MyU32")); - - // Empty destination module. - FileTable ft_dst; - std::unique_ptr dst = - std::make_unique("dst_mod", std::nullopt, ft_dst); - - absl::flat_hash_map clones; - XLS_ASSERT_OK_AND_ASSIGN( - clones, CloneAstAndGetAllPairs(foo, dst.get(), &NoopCloneReplacer)); - - auto* cloned_foo = down_cast(clones.at(foo)); - auto* param_trta = down_cast( - cloned_foo->params()[0]->type_annotation()); - TypeRef* param_tr = param_trta->type_ref(); - ASSERT_TRUE(std::holds_alternative(param_tr->type_definition())); - - // Ensure it's using the cloned alias in the destination module, not the - // source. - TypeAlias* cloned_alias = down_cast(clones.at(src_alias)); - EXPECT_NE(cloned_alias, src_alias); - EXPECT_EQ(std::get(param_tr->type_definition()), cloned_alias); -} - -TEST(AstClonerTest, TypeRefFromUseTreeEntry) { - constexpr std::string_view kProgram = R"( -#![feature(use_syntax)] - -use foo::ImportedStruct; - -fn main(x: ImportedStruct) { - () -})"; - - constexpr std::string_view kExpectedFunction = - R"(fn main(x: ImportedStruct) { - () -})"; - - FileTable file_table; - XLS_ASSERT_OK_AND_ASSIGN(auto module, ParseModule(kProgram, "fake_path.x", - "the_module", file_table)); - XLS_ASSERT_OK_AND_ASSIGN(Function * f, - module->GetMemberOrError("main")); - - // Ensure the parsed TypeRef is backed by a UseTreeEntry. - std::optional type_ref = FindFirstTypeRef(f); - ASSERT_TRUE(type_ref.has_value()); - ASSERT_TRUE( - std::holds_alternative((*type_ref)->type_definition())); - - // Clone and ensure we preserve the UseTreeEntry variant and string form. - XLS_ASSERT_OK_AND_ASSIGN(AstNode * clone, CloneAst(f)); - EXPECT_EQ(kExpectedFunction, clone->ToString()); - XLS_ASSERT_OK(VerifyClone(f, clone, *module->file_table())); - - std::optional cloned_type_ref = FindFirstTypeRef(clone); - ASSERT_TRUE(cloned_type_ref.has_value()); - ASSERT_TRUE(std::holds_alternative( - (*cloned_type_ref)->type_definition())); -} - -TEST(AstClonerTest, TypeRefFromColonRef) { - constexpr std::string_view kProgram = R"( -import my.module as foo; - -fn main(x: foo::ImportedStruct) { - () -})"; - - constexpr std::string_view kExpectedFunction = - R"(fn main(x: foo::ImportedStruct) { - () -})"; - - FileTable file_table; - XLS_ASSERT_OK_AND_ASSIGN(auto module, ParseModule(kProgram, "fake_path.x", - "the_module", file_table)); - XLS_ASSERT_OK_AND_ASSIGN(Function * f, - module->GetMemberOrError("main")); - - // Ensure the parsed TypeRef is backed by a ColonRef. - std::optional type_ref = FindFirstTypeRef(f); - ASSERT_TRUE(type_ref.has_value()); - ASSERT_TRUE( - std::holds_alternative((*type_ref)->type_definition())); - - // Clone and ensure we preserve the ColonRef variant and string form. - XLS_ASSERT_OK_AND_ASSIGN(AstNode * clone, CloneAst(f)); - EXPECT_EQ(kExpectedFunction, clone->ToString()); - XLS_ASSERT_OK(VerifyClone(f, clone, *module->file_table())); - - std::optional cloned_type_ref = FindFirstTypeRef(clone); - ASSERT_TRUE(cloned_type_ref.has_value()); - ASSERT_TRUE( - std::holds_alternative((*cloned_type_ref)->type_definition())); -} - -TEST(AstClonerTest, TypeRefFromStructDef) { +TEST(AstClonerTest, DoesntCloneStructDefTwice) { constexpr std::string_view kProgram = R"( struct MyStruct { a: u32 } @@ -1864,44 +1755,6 @@ fn id(x: MyStruct) -> MyStruct { x } cloned_param_trta->type_ref()->type_definition())); } -TEST(AstClonerTest, TypeRefFromProcDefInImpl) { - constexpr std::string_view kProgram = R"( -proc MyProc { - a: u32, -} - -impl MyProc {} -)"; - - constexpr std::string_view kExpectedImpl = R"(impl MyProc { -})"; - - FileTable file_table; - XLS_ASSERT_OK_AND_ASSIGN(auto module, ParseModule(kProgram, "fake_path.x", - "the_module", file_table)); - - XLS_ASSERT_OK_AND_ASSIGN(ProcDef * proc_def, - module->GetMemberOrError("MyProc")); - Impl* impl = module->GetImpls().at(0); - - // The impl's struct_ref should be a TypeRefTypeAnnotation backed by ProcDef. - auto* trta = down_cast(impl->struct_ref()); - TypeRef* tr = trta->type_ref(); - ASSERT_TRUE(std::holds_alternative(tr->type_definition())); - EXPECT_EQ(std::get(tr->type_definition()), proc_def); - - // Clone and ensure ToString and variant preservation. - XLS_ASSERT_OK_AND_ASSIGN(AstNode * clone, CloneAst(impl)); - EXPECT_EQ(kExpectedImpl, clone->ToString()); - XLS_ASSERT_OK(VerifyClone(impl, clone, *module->file_table())); - - auto* cloned_impl = down_cast(clone); - auto* cloned_trta = - down_cast(cloned_impl->struct_ref()); - ASSERT_TRUE(std::holds_alternative( - cloned_trta->type_ref()->type_definition())); -} - TEST(AstClonerTest, CloneAstClonesVerbatimNode) { constexpr std::string_view kProgram = "const FOO = u32:42;"; FileTable file_table; From a08c44073d9c7a175cade2b30e7ef0866c682452 Mon Sep 17 00:00:00 2001 From: Sirui Lu Date: Mon, 1 Sep 2025 13:22:11 -0700 Subject: [PATCH 4/5] Use CastIfNotVerbatim for TestFunction and TestProc to fix fmt test failure --- xls/dslx/frontend/ast_cloner.cc | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/xls/dslx/frontend/ast_cloner.cc b/xls/dslx/frontend/ast_cloner.cc index a2dd97ec07..ccccebf7d2 100644 --- a/xls/dslx/frontend/ast_cloner.cc +++ b/xls/dslx/frontend/ast_cloner.cc @@ -949,16 +949,19 @@ class AstCloner : public AstNodeVisitor { XLS_RETURN_IF_ERROR(VisitChildren(n)); XLS_RETURN_IF_ERROR(ReplaceOrVisit(&n->fn())); - old_to_new_[n] = module(n)->Make( - n->span(), *down_cast(old_to_new_.at(&n->fn()))); + XLS_ASSIGN_OR_RETURN(Function * new_fn, CastIfNotVerbatim( + old_to_new_.at(&n->fn()))); + old_to_new_[n] = module(n)->Make(n->span(), *new_fn); return absl::OkStatus(); } absl::Status HandleTestProc(const TestProc* n) override { XLS_RETURN_IF_ERROR(VisitChildren(n)); - old_to_new_[n] = module(n)->Make( - down_cast(old_to_new_.at(n->proc())), n->expected_fail_label()); + XLS_ASSIGN_OR_RETURN(Proc * new_proc, + CastIfNotVerbatim(old_to_new_.at(n->proc()))); + old_to_new_[n] = + module(n)->Make(new_proc, n->expected_fail_label()); return absl::OkStatus(); } From 1211546fbdfb8c0699a70856154d5b93eb41b664 Mon Sep 17 00:00:00 2001 From: Sirui Lu Date: Mon, 8 Sep 2025 17:21:02 -0700 Subject: [PATCH 5/5] Remove unused variable --- xls/dslx/frontend/ast_cloner_test.cc | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/xls/dslx/frontend/ast_cloner_test.cc b/xls/dslx/frontend/ast_cloner_test.cc index 4bfa63da55..f4c8d87c15 100644 --- a/xls/dslx/frontend/ast_cloner_test.cc +++ b/xls/dslx/frontend/ast_cloner_test.cc @@ -1648,8 +1648,7 @@ fn foo(x: MyU32) -> MyU32 { "the_module", file_table)); XLS_ASSERT_OK_AND_ASSIGN(Function * foo, module->GetMemberOrError("foo")); - XLS_ASSERT_OK_AND_ASSIGN(TypeAlias * alias, - module->GetMemberOrError("MyU32")); + XLS_ASSERT_OK(module->GetMemberOrError("MyU32")); // Clone just the function into the same target module. XLS_ASSERT_OK_AND_ASSIGN(AstNode * clone, CloneAst(foo));