Merge pull request #1775 from xlsynth:cdleary-2024-12-07-use-construct

PiperOrigin-RevId: 704896302

xls/dslx/fmt/ast_fmt.cc

@@ -2596,6 +2596,12 @@ DocRef Formatter::Format(const Import& n) {
   return ConcatNGroup(arena_, pieces);
 }
 
+DocRef Formatter::Format(const Use& n) {
+  // TODO(cdleary): 2024-12-07 This is just a stopgap, we should add reflow
+  // capability.
+  return arena_.MakeText(n.ToString());
+}
+
 DocRef Formatter::Format(const Let& n, bool trailing_semi) {
   std::vector<DocRef> leader_pieces = {
       arena_.Make(n.is_const() ? Keyword::kConst : Keyword::kLet),
@@ -2703,6 +2709,7 @@ DocRef Formatter::Format(const ModuleMember& n) {
                          [&](const ConstantDef* n) { return Format(*n); },
                          [&](const EnumDef* n) { return Format(*n); },
                          [&](const Import* n) { return Format(*n); },
+                         [&](const Use* n) { return Format(*n); },
                          [&](const ConstAssert* n) {
                            return arena_.MakeConcat(Format(*n), arena_.semi());
                          },

xls/dslx/fmt/ast_fmt.h

@@ -59,6 +59,7 @@ class Formatter {
   DocRef Format(const Impl& n);
   DocRef Format(const ImplMember& n);
   DocRef Format(const Import& n);
+  DocRef Format(const Use& n);
   DocRef Format(const ModuleMember& n);
   DocRef Format(const ParametricBinding& n);
   DocRef Format(const ParametricBinding* n);

xls/dslx/frontend/ast.cc

@@ -313,6 +313,8 @@ std::string_view AstNodeKindToString(AstNodeKind kind) {
       return "tuple index";
     case AstNodeKind::kUnrollFor:
       return "unroll-for";
+    case AstNodeKind::kUse:
+      return "use";
     case AstNodeKind::kVerbatimNode:
       return "verbatim-node";
   }
@@ -783,6 +785,102 @@ std::string Import::ToString() const {
   return absl::StrFormat("import %s;", absl::StrJoin(subject_, "."));
 }
 
+// -- class Use
+
+UseInteriorEntry::UseInteriorEntry(
+    std::string identifier, std::vector<std::unique_ptr<UseTreeEntry>> subtrees)
+    : identifier_(std::move(identifier)), subtrees_(std::move(subtrees)) {
+  CHECK(!subtrees_.empty());
+  for (const auto& subtree : subtrees_) {
+    DCHECK(subtree != nullptr);
+  }
+}
+
+std::vector<NameDef*> UseInteriorEntry::GetLeafNameDefs() const {
+  std::vector<NameDef*> result;
+  for (const auto& subtree : subtrees_) {
+    std::vector<NameDef*> subtree_leaf_name_defs = subtree->GetLeafNameDefs();
+    result.insert(result.end(), subtree_leaf_name_defs.begin(),
+                  subtree_leaf_name_defs.end());
+  }
+  return result;
+}
+
+std::string UseInteriorEntry::ToString() const {
+  std::string subtrees_str;
+  if (subtrees_.size() == 1) {
+    subtrees_str = subtrees_.front()->ToString();
+  } else {
+    subtrees_str =
+        absl::StrCat("{",
+                     absl::StrJoin(subtrees_, ", ",
+                                   [](std::string* out, const auto& subtree) {
+                                     absl::StrAppend(out, subtree->ToString());
+                                   }),
+                     "}");
+  }
+  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 UseInteriorEntry& interior) { return interior.ToString(); }},
+      entry_);
+}
+
+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_);
+}
+
+std::vector<std::string> UseTreeEntry::GetLeafIdentifiers() const {
+  std::vector<NameDef*> leaf_name_defs = GetLeafNameDefs();
+  std::vector<std::string> result;
+  result.reserve(leaf_name_defs.size());
+  for (const NameDef* name_def : leaf_name_defs) {
+    result.push_back(name_def->identifier());
+  }
+  return result;
+}
+
+Use::Use(Module* owner, Span span, std::unique_ptr<UseTreeEntry> root)
+    : AstNode(owner), span_(std::move(span)), root_(std::move(root)) {}
+
+Use::~Use() = default;
+
+std::string Use::ToString() const {
+  return absl::StrCat("use ", root_->ToString(), ";");
+}
+
+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;
+}
+
 // -- class ColonRef
 
 ColonRef::ColonRef(Module* owner, Span span, Subject subject, std::string attr,

xls/dslx/frontend/ast.h

@@ -107,6 +107,7 @@
   X(TestProc)                     \
   X(TypeAlias)                    \
   X(TypeRef)                      \
+  X(Use)                          \
   X(WidthSlice)                   \
   X(WildcardPattern)              \
   /* type annotations */          \
@@ -1313,6 +1314,115 @@ class Import : public AstNode {
   std::optional<std::string> alias_;
 };
 
+// -- Use
+
+class UseTreeEntry;  // forward decl
+
+// Interior node in the `use` construct tree.
+//
+//              v--v identifier
+// e.g. in `use foo::{bar, baz}`
+//                   ^--------^ subtrees
+//
+// 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);
+
+  // Move-only.
+  UseInteriorEntry(UseInteriorEntry&&) = default;
+  UseInteriorEntry& operator=(UseInteriorEntry&&) = default;
+
+  std::vector<NameDef*> GetLeafNameDefs() const;
+
+  absl::Span<const std::unique_ptr<UseTreeEntry>> subtrees() const {
+    return subtrees_;
+  }
+
+  std::string ToString() const;
+
+ private:
+  std::string identifier_;
+  std::vector<std::unique_ptr<UseTreeEntry>> subtrees_;
+};
+
+// Leaf node in the `use` construct tree.
+struct UseLeafEntry {
+  NameDef* name_def;
+};
+
+// Arbitrary entry (interior or leaf) in the `use` construct tree.
+class UseTreeEntry {
+ 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)) {}
+
+  std::string ToString() const;
+
+  std::vector<std::string> GetLeafIdentifiers() const;
+  std::vector<NameDef*> GetLeafNameDefs() const;
+
+ private:
+  std::variant<UseInteriorEntry, UseLeafEntry> entry_;
+  Span span_;
+};
+
+// Represents a use statement; e.g.
+//  use foo::bar::{baz::{bat, qux}, ipsum};
+//
+// In that case there are 3 leafs that create name definitions for the module:
+// `bat`, `qux`, and `ipsum`.
+//
+// You can also use a module directly instead of an item within it; e.g.
+//  use foo;
+//
+// And then subsequently refer to `foo::STUFF`.
+//
+// Note we DO NOT support multiple use at the "root" level; e.g.
+// ```
+//  use {bar, baz};
+// ```
+// is invalid.
+//
+// Attributes:
+//  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() override;
+
+  AstNodeKind kind() const override { return AstNodeKind::kUse; }
+
+  absl::Status Accept(AstNodeVisitor* v) const override {
+    return v->HandleUse(this);
+  }
+  std::string_view GetNodeTypeName() const override { return "Use"; }
+  std::string ToString() const override;
+  std::optional<Span> GetSpan() const override { return span_; }
+
+  std::vector<AstNode*> GetChildren(bool want_types) const override;
+
+  std::vector<std::string> GetLeafIdentifiers() const {
+    return root_->GetLeafIdentifiers();
+  }
+  std::vector<NameDef*> GetLeafNameDefs() const {
+    return root_->GetLeafNameDefs();
+  }
+
+  const Span& span() const { return span_; }
+
+ private:
+  Span span_;
+  std::unique_ptr<UseTreeEntry> root_;
+};
+
 // Represents a module-value or enum-value style reference when the LHS
 // expression is unknown; e.g. when accessing a member in a module:
 //

xls/dslx/frontend/ast_cloner.cc

@@ -399,6 +399,10 @@ class AstCloner : public AstNodeVisitor {
     return absl::OkStatus();
   }
 
+  absl::Status HandleUse(const Use* n) override {
+    return absl::UnimplementedError("Not implemented: clone use");
+  }
+
   absl::Status HandleIndex(const Index* n) override {
     XLS_RETURN_IF_ERROR(VisitChildren(n));
 

xls/dslx/frontend/ast_node.h

@@ -94,6 +94,7 @@ enum class AstNodeKind : uint8_t {
   kTypeRef,
   kUnop,
   kUnrollFor,
+  kUse,
   kVerbatimNode,
   kWidthSlice,
   kWildcardPattern,