[DSLX:FE] Switch UseTreeEntry to an AstNode, make a cloner implementa…

…tion.
google · Dec 11, 2024 · aae520f · aae520f
1 parent 98f2b8d
commit aae520f
Show file tree

Hide file tree

Showing 10 changed files with 129 additions and 76 deletions.
diff --git a/xls/dslx/frontend/ast.cc b/xls/dslx/frontend/ast.cc
@@ -315,6 +315,8 @@ std::string_view AstNodeKindToString(AstNodeKind kind) {
       return "unroll-for";
     case AstNodeKind::kUse:
       return "use";
+    case AstNodeKind::kUseTreeEntry:
+      return "use-tree-entry";
     case AstNodeKind::kVerbatimNode:
       return "verbatim-node";
   }
@@ -787,8 +789,8 @@ std::string Import::ToString() const {
 
 // -- class Use
 
-UseInteriorEntry::UseInteriorEntry(
-    std::string identifier, std::vector<std::unique_ptr<UseTreeEntry>> subtrees)
+UseInteriorEntry::UseInteriorEntry(std::string identifier,
+                                   std::vector<UseTreeEntry*> subtrees)
     : identifier_(std::move(identifier)), subtrees_(std::move(subtrees)) {
   CHECK(!subtrees_.empty());
   for (const auto& subtree : subtrees_) {
@@ -822,34 +824,35 @@ std::string UseInteriorEntry::ToString() const {
   return absl::StrCat(identifier_, "::", subtrees_str);
 }
 
-/* static */ std::unique_ptr<UseTreeEntry> UseTreeEntry::MakeLeaf(
-    NameDef* name_def, Span span) {
-  return std::make_unique<UseTreeEntry>(UseLeafEntry{name_def}, span);
-}
-
-/* static */ std::unique_ptr<UseTreeEntry> UseTreeEntry::MakeInterior(
-    std::string identifier, std::vector<std::unique_ptr<UseTreeEntry>> subtrees,
-    Span span) {
-  return std::make_unique<UseTreeEntry>(
-      UseInteriorEntry{std::move(identifier), std::move(subtrees)}, span);
-}
-
 std::string UseTreeEntry::ToString() const {
   return absl::visit(
       Visitor{
-          [](const UseLeafEntry& leaf) { return leaf.name_def->identifier(); },
+          [](const NameDef* leaf) { return leaf->identifier(); },
           [](const UseInteriorEntry& interior) { return interior.ToString(); }},
-      entry_);
+      payload_);
+}
+
+std::vector<AstNode*> UseTreeEntry::GetChildren(bool want_types) const {
+  return absl::visit(
+      Visitor{[](NameDef* leaf) { return std::vector<AstNode*>{leaf}; },
+              [](const UseInteriorEntry& interior) {
+                std::vector<AstNode*> results;
+                results.reserve(interior.subtrees().size());
+                for (UseTreeEntry* subtree : interior.subtrees()) {
+                  results.push_back(subtree);
+                }
+                return results;
+              }},
+      payload_);
 }
 
 std::vector<NameDef*> UseTreeEntry::GetLeafNameDefs() const {
-  return absl::visit(Visitor{[](const UseLeafEntry& leaf) {
-                               return std::vector<NameDef*>{leaf.name_def};
-                             },
-                             [](const UseInteriorEntry& interior) {
-                               return interior.GetLeafNameDefs();
-                             }},
-                     entry_);
+  return absl::visit(
+      Visitor{[](NameDef* leaf) { return std::vector<NameDef*>{leaf}; },
+              [](const UseInteriorEntry& interior) {
+                return interior.GetLeafNameDefs();
+              }},
+      payload_);
 }
 
 std::vector<std::string> UseTreeEntry::GetLeafIdentifiers() const {
@@ -862,8 +865,12 @@ std::vector<std::string> UseTreeEntry::GetLeafIdentifiers() const {
   return result;
 }
 
-Use::Use(Module* owner, Span span, std::unique_ptr<UseTreeEntry> root)
-    : AstNode(owner), span_(std::move(span)), root_(std::move(root)) {}
+absl::Status UseTreeEntry::Accept(AstNodeVisitor* v) const {
+  return v->HandleUseTreeEntry(this);
+}
+
+Use::Use(Module* owner, Span span, UseTreeEntry& root)
+    : AstNode(owner), span_(std::move(span)), root_(&root) {}
 
 Use::~Use() = default;
 
@@ -872,13 +879,7 @@ std::string Use::ToString() const {
 }
 
 std::vector<AstNode*> Use::GetChildren(bool want_types) const {
-  std::vector<NameDef*> name_defs = root_->GetLeafNameDefs();
-  std::vector<AstNode*> results;
-  results.reserve(name_defs.size());
-  for (NameDef* name_def : name_defs) {
-    results.push_back(name_def);
-  }
-  return results;
+  return {root_};
 }
 
 // -- class ColonRef

diff --git a/xls/dslx/frontend/ast.h b/xls/dslx/frontend/ast.h
@@ -108,6 +108,7 @@
   X(TypeAlias)                    \
   X(TypeRef)                      \
   X(Use)                          \
+  X(UseTreeEntry)                 \
   X(WidthSlice)                   \
   X(WildcardPattern)              \
   /* type annotations */          \
@@ -1318,7 +1319,8 @@ class Import : public AstNode {
 
 class UseTreeEntry;  // forward decl
 
-// Interior node in the `use` construct tree.
+// Data structure that holds the payload for an interior node in the `use`
+// construct tree.
 //
 //              v--v identifier
 // e.g. in `use foo::{bar, baz}`
@@ -1327,49 +1329,53 @@ class UseTreeEntry;  // forward decl
 // For a node to be considered "interior" there must be subtrees.
 class UseInteriorEntry {
  public:
-  UseInteriorEntry(std::string identifier,
-                   std::vector<std::unique_ptr<UseTreeEntry>> subtrees);
+  UseInteriorEntry(std::string identifier, std::vector<UseTreeEntry*> subtrees);
 
   // Move-only.
   UseInteriorEntry(UseInteriorEntry&&) = default;
   UseInteriorEntry& operator=(UseInteriorEntry&&) = default;
 
+  // Transitively retrieves a vector of all the name defs at the leaf positions
+  // underneath this interior node.
   std::vector<NameDef*> GetLeafNameDefs() const;
 
-  absl::Span<const std::unique_ptr<UseTreeEntry>> subtrees() const {
-    return subtrees_;
-  }
+  absl::Span<UseTreeEntry* const> subtrees() const { return subtrees_; }
 
   std::string ToString() const;
+  std::string_view identifier() const { return identifier_; }
 
  private:
   std::string identifier_;
-  std::vector<std::unique_ptr<UseTreeEntry>> subtrees_;
-};
-
-// Leaf node in the `use` construct tree.
-struct UseLeafEntry {
-  NameDef* name_def;
+  std::vector<UseTreeEntry*> subtrees_;
 };
 
 // Arbitrary entry (interior or leaf) in the `use` construct tree.
-class UseTreeEntry {
+class UseTreeEntry : public AstNode {
  public:
-  static std::unique_ptr<UseTreeEntry> MakeInterior(
-      std::string identifier,
-      std::vector<std::unique_ptr<UseTreeEntry>> subtrees, Span span);
-  static std::unique_ptr<UseTreeEntry> MakeLeaf(NameDef* name_def, Span span);
-
-  UseTreeEntry(std::variant<UseInteriorEntry, UseLeafEntry> entry, Span span)
-      : entry_(std::move(entry)), span_(std::move(span)) {}
+  UseTreeEntry(Module* owner, std::variant<UseInteriorEntry, NameDef*> payload,
+               Span span)
+      : AstNode(owner), payload_(std::move(payload)), span_(std::move(span)) {}
 
-  std::string ToString() const;
+  std::string ToString() const override;
 
   std::vector<std::string> GetLeafIdentifiers() const;
   std::vector<NameDef*> GetLeafNameDefs() const;
 
+  AstNodeKind kind() const override { return AstNodeKind::kUseTreeEntry; }
+  std::string_view GetNodeTypeName() const override { return "UseTreeEntry"; }
+  std::optional<Span> GetSpan() const override { return span_; }
+  std::vector<AstNode*> GetChildren(bool want_types) const override;
+  absl::Status Accept(AstNodeVisitor* v) const override;
+
+  // The payload of this tree node -- it is either an interior node or a leaf
+  // node.
+  const std::variant<UseInteriorEntry, NameDef*>& payload() const {
+    return payload_;
+  }
+  const Span& span() const { return span_; }
+
  private:
-  std::variant<UseInteriorEntry, UseLeafEntry> entry_;
+  std::variant<UseInteriorEntry, NameDef*> payload_;
   Span span_;
 };
 
@@ -1394,7 +1400,7 @@ class UseTreeEntry {
 //  span: Span of the overall `use` statement in the text.
 class Use : public AstNode {
  public:
-  Use(Module* owner, Span span, std::unique_ptr<UseTreeEntry> root);
+  Use(Module* owner, Span span, UseTreeEntry& root);
 
   ~Use() override;
 
@@ -1416,11 +1422,13 @@ class Use : public AstNode {
     return root_->GetLeafNameDefs();
   }
 
+  UseTreeEntry& root() { return *root_; }
+  const UseTreeEntry& root() const { return *root_; }
   const Span& span() const { return span_; }
 
  private:
   Span span_;
-  std::unique_ptr<UseTreeEntry> root_;
+  UseTreeEntry* root_;
 };
 
 // Represents a module-value or enum-value style reference when the LHS

diff --git a/xls/dslx/frontend/ast_cloner.cc b/xls/dslx/frontend/ast_cloner.cc
@@ -400,7 +400,38 @@ class AstCloner : public AstNodeVisitor {
   }
 
   absl::Status HandleUse(const Use* n) override {
-    return absl::UnimplementedError("Not implemented: clone use");
+    XLS_RETURN_IF_ERROR(VisitChildren(n));
+
+    const UseTreeEntry* old_root = &n->root();
+    UseTreeEntry& new_root =
+        *down_cast<UseTreeEntry*>(old_to_new_.at(old_root));
+    old_to_new_[n] = module_->Make<Use>(n->span(), new_root);
+    return absl::OkStatus();
+  }
+
+  absl::Status HandleUseTreeEntry(const UseTreeEntry* n) override {
+    XLS_RETURN_IF_ERROR(VisitChildren(n));
+
+    using PayloadT = std::variant<UseInteriorEntry, NameDef*>;
+    PayloadT new_payload = absl::visit(
+        Visitor{[&](const NameDef* name_def) -> PayloadT {
+                  return down_cast<NameDef*>(old_to_new_.at(name_def));
+                },
+                [&](const UseInteriorEntry& interior) -> PayloadT {
+                  std::vector<UseTreeEntry*> new_subtrees;
+                  new_subtrees.reserve(interior.subtrees().size());
+                  for (UseTreeEntry* subtree : interior.subtrees()) {
+                    new_subtrees.push_back(
+                        down_cast<UseTreeEntry*>(old_to_new_.at(subtree)));
+                  }
+                  return UseInteriorEntry{std::string{interior.identifier()},
+                                          std::move(new_subtrees)};
+                }},
+        n->payload());
+
+    old_to_new_[n] =
+        module_->Make<UseTreeEntry>(std::move(new_payload), n->span());
+    return absl::OkStatus();
   }
 
   absl::Status HandleIndex(const Index* n) override {

diff --git a/xls/dslx/frontend/ast_cloner_test.cc b/xls/dslx/frontend/ast_cloner_test.cc
@@ -1714,5 +1714,17 @@ fn divmod() {})*";
   EXPECT_EQ(kProgram, clone->ToString());
 }
 
+TEST(AstClonerTest, Use) {
+  constexpr std::string_view kProgram =
+      R"(use foo::bar::{baz::{bat, qux}, ipsum};)";
+
+  FileTable file_table;
+  XLS_ASSERT_OK_AND_ASSIGN(auto module, ParseModule(kProgram, "fake_path.x",
+                                                    "the_module", file_table));
+  XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr<Module> clone,
+                           CloneModule(*module.get()));
+  EXPECT_EQ(kProgram, clone->ToString());
+}
+
 }  // namespace
 }  // namespace xls::dslx
diff --git a/xls/dslx/frontend/ast_node.h b/xls/dslx/frontend/ast_node.h
@@ -95,6 +95,7 @@ enum class AstNodeKind : uint8_t {
   kUnop,
   kUnrollFor,
   kUse,
+  kUseTreeEntry,
   kVerbatimNode,
   kWidthSlice,
   kWildcardPattern,

diff --git a/xls/dslx/frontend/parser.cc b/xls/dslx/frontend/parser.cc
@@ -1574,8 +1574,7 @@ absl::StatusOr<Match*> Parser::ParseMatch(Bindings& bindings) {
   return module_->Make<Match>(span, matched, std::move(arms));
 }
 
-absl::StatusOr<std::unique_ptr<UseTreeEntry>> Parser::ParseUseTreeEntry(
-    Bindings& bindings) {
+absl::StatusOr<UseTreeEntry*> Parser::ParseUseTreeEntry(Bindings& bindings) {
   // Get the identifier for this level of the tree.
   XLS_ASSIGN_OR_RETURN(Token tok, PopTokenOrError(TokenKind::kIdentifier));
   std::string identifier = *tok.GetValue();
@@ -1589,7 +1588,7 @@ absl::StatusOr<std::unique_ptr<UseTreeEntry>> Parser::ParseUseTreeEntry(
     // subsequent level.
     XLS_ASSIGN_OR_RETURN(NameDef * name_def, TokenToNameDef(tok));
     bindings.Add(name_def->identifier(), name_def);
-    return UseTreeEntry::MakeLeaf(name_def, tok.span());
+    return module_->Make<UseTreeEntry>(name_def, tok.span());
   }
 
   // If we've gotten here we know there's a next level, we're just looking to
@@ -1598,11 +1597,10 @@ absl::StatusOr<std::unique_ptr<UseTreeEntry>> Parser::ParseUseTreeEntry(
   if (saw_obrace) {
     // Multiple peer subtrees -- present as children (subtrees) to make this
     // level.
-    std::vector<std::unique_ptr<UseTreeEntry>> children;
+    std::vector<UseTreeEntry*> children;
     while (true) {
-      XLS_ASSIGN_OR_RETURN(std::unique_ptr<UseTreeEntry> child,
-                           ParseUseTreeEntry(bindings));
-      children.push_back(std::move(child));
+      XLS_ASSIGN_OR_RETURN(UseTreeEntry * child, ParseUseTreeEntry(bindings));
+      children.push_back(child);
       XLS_ASSIGN_OR_RETURN(bool saw_cbrace, TryDropToken(TokenKind::kCBrace));
       if (saw_cbrace) {
         break;
@@ -1611,19 +1609,18 @@ absl::StatusOr<std::unique_ptr<UseTreeEntry>> Parser::ParseUseTreeEntry(
           TokenKind::kComma, /*start=*/nullptr,
           "Expect a ',' to separate multiple entries in a `use` statement"));
     }
-    return UseTreeEntry::MakeInterior(identifier, std::move(children),
-                                      tok.span());
+    return module_->Make<UseTreeEntry>(
+        UseInteriorEntry{identifier, std::move(children)}, tok.span());
   }
 
   // Must be a single child in the subsequent level -- we recurse here to look
   // for additional levels after it.
-  XLS_ASSIGN_OR_RETURN(std::unique_ptr<UseTreeEntry> child,
-                       ParseUseTreeEntry(bindings));
-  std::vector<std::unique_ptr<UseTreeEntry>> children;
+  XLS_ASSIGN_OR_RETURN(UseTreeEntry * child, ParseUseTreeEntry(bindings));
+  std::vector<UseTreeEntry*> children;
   children.reserve(1);
   children.push_back(std::move(child));
-  return UseTreeEntry::MakeInterior(identifier, std::move(children),
-                                    tok.span());
+  return module_->Make<UseTreeEntry>(
+      UseInteriorEntry{identifier, std::move(children)}, tok.span());
 }
 
 absl::StatusOr<Use*> Parser::ParseUse(Bindings& bindings) {
@@ -1638,13 +1635,12 @@ absl::StatusOr<Use*> Parser::ParseUse(Bindings& bindings) {
         "is not allowed -- please break into multiple statements");
   }
 
-  XLS_ASSIGN_OR_RETURN(std::unique_ptr<UseTreeEntry> root,
-                       ParseUseTreeEntry(bindings));
+  XLS_ASSIGN_OR_RETURN(UseTreeEntry * root, ParseUseTreeEntry(bindings));
   XLS_RETURN_IF_ERROR(
       DropTokenOrError(TokenKind::kSemi, /*start=*/&kw,
                        "Expect a ';' at end of `use` statement"));
   Span span(kw.span().start(), GetPos());
-  return module_->Make<Use>(span, std::move(root));
+  return module_->Make<Use>(span, *root);
 }
 
 absl::StatusOr<Import*> Parser::ParseImport(Bindings& bindings) {

diff --git a/xls/dslx/frontend/parser.h b/xls/dslx/frontend/parser.h
@@ -561,8 +561,7 @@ class Parser : public TokenParser {
 
   // Parses a single entry in a `use` tree -- this can be a leaf or an interior
   // entry.
-  absl::StatusOr<std::unique_ptr<UseTreeEntry>> ParseUseTreeEntry(
-      Bindings& bindings);
+  absl::StatusOr<UseTreeEntry*> ParseUseTreeEntry(Bindings& bindings);
 
   // Parses a use statement into a `Use` AST node.
   absl::StatusOr<Use*> ParseUse(Bindings& bindings);

diff --git a/xls/dslx/ir_convert/function_converter.cc b/xls/dslx/ir_convert/function_converter.cc
@@ -373,6 +373,7 @@ class FunctionConverterVisitor : public AstNodeVisitor {
   INVALID(Impl)
   INVALID(Import)
   INVALID(Use)
+  INVALID(UseTreeEntry)
   INVALID(Module)
   INVALID(Proc)
   INVALID(ProcMember)

diff --git a/xls/dslx/type_system/deduce.cc b/xls/dslx/type_system/deduce.cc
@@ -2110,6 +2110,9 @@ class DeduceVisitor : public AstNodeVisitor {
   absl::Status HandleSlice(const Slice* n) override { return Fatal(n); }
   absl::Status HandleImport(const Import* n) override { return Fatal(n); }
   absl::Status HandleUse(const Use* n) override { return Fatal(n); }
+  absl::Status HandleUseTreeEntry(const UseTreeEntry* n) override {
+    return Fatal(n);
+  }
   absl::Status HandleFunction(const Function* n) override { return Fatal(n); }
   absl::Status HandleQuickCheck(const QuickCheck* n) override {
     return Fatal(n);