Store a pointer to the impl for functions within impl scope.

xls/dslx/frontend/ast.cc

@@ -1728,16 +1728,25 @@ std::vector<AstNode*> Impl::GetChildren(bool want_types) const {
   return results;
 }
 
-std::vector<ConstantDef*> Impl::GetConstants() const {
-  std::vector<ConstantDef*> results;
+template <typename T>
+std::vector<T> Impl::GetMembersOfType() const {
+  std::vector<T> results;
   for (const auto& member : members_) {
-    if (std::holds_alternative<ConstantDef*>(member)) {
-      results.push_back(std::get<ConstantDef*>(member));
+    if (std::holds_alternative<T>(member)) {
+      results.push_back(std::get<T>(member));
     }
   }
   return results;
 }
 
+std::vector<ConstantDef*> Impl::GetConstants() const {
+  return GetMembersOfType<ConstantDef*>();
+}
+
+std::vector<Function*> Impl::GetFunctions() const {
+  return GetMembersOfType<Function*>();
+}
+
 static std::string ImplMemberIdentifier(const ImplMember member) {
   return absl::visit(
       Visitor{

xls/dslx/frontend/ast.h

@@ -2286,6 +2286,9 @@ class Function : public AstNode {
   void set_proc(Proc* proc) { proc_ = proc; }
   bool IsInProc() const { return proc_.has_value(); }
 
+  std::optional<Impl*> impl() const { return impl_; }
+  void set_impl(Impl* impl) { impl_ = impl; }
+
   std::optional<Span> GetParametricBindingsSpan() const {
     if (parametric_bindings_.empty()) {
       return std::nullopt;
@@ -2304,6 +2307,7 @@ class Function : public AstNode {
   StatementBlock* body_;
   const FunctionTag tag_;
   std::optional<Proc*> proc_;
+  std::optional<Impl*> impl_;
 
   const bool is_public_;
   std::optional<ForeignFunctionData> extern_verilog_module_;
@@ -3146,6 +3150,7 @@ class Impl : public AstNode {
   void set_members(std::vector<ImplMember>& members) { members_ = members; }
 
   std::vector<ConstantDef*> GetConstants() const;
+  std::vector<Function*> GetFunctions() const;
 
   // Returns the member with the given name, if present.
   std::optional<ImplMember> GetMember(std::string_view name) const;
@@ -3164,6 +3169,9 @@ class Impl : public AstNode {
 
   template <typename T>
   std::optional<T> GetMemberOfType(std::string_view name) const;
+
+  template <typename T>
+  std::vector<T> GetMembersOfType() const;
 };
 
 // A virtual base class for nodes that directly instantiate a struct.

xls/dslx/frontend/ast_cloner.cc

@@ -403,6 +403,10 @@ class AstCloner : public AstNodeVisitor {
     new_function->set_disable_format(n->disable_format());
     old_to_new_[n] = new_function;
     new_name_def->set_definer(old_to_new_.at(n));
+    if (n->impl().has_value()) {
+      Impl* new_impl = down_cast<Impl*>(old_to_new_.at(*(n->impl())));
+      new_function->set_impl(new_impl);
+    }
     return absl::OkStatus();
   }
 

xls/dslx/frontend/ast_cloner_test.cc

@@ -591,7 +591,8 @@ impl MyStruct {
 }
 
 fn main() -> u32 {
-    MyStruct::MY_CONST
+    let a = MyStruct::another();
+    a.some_func()
 }
 )";
 
@@ -610,7 +611,8 @@ fn main() -> u32 {
 })";
 
   constexpr std::string_view kExpectedFunction = R"(fn main() -> u32 {
-    MyStruct::MY_CONST
+    let a = MyStruct::another();
+    a.some_func()
 })";
 
   FileTable file_table;
@@ -632,6 +634,9 @@ fn main() -> u32 {
   XLS_ASSERT_OK_AND_ASSIGN(AstNode * clone, CloneAst(f));
   EXPECT_EQ(kExpectedFunction, clone->ToString());
   XLS_ASSERT_OK(VerifyClone(f, clone, *module->file_table()));
+
+  Function* cloned_impl_fn = *cloned_impl->GetFunction("some_func");
+  EXPECT_EQ(cloned_impl_fn->impl(), cloned_impl);
 }
 
 TEST(AstClonerTest, ColonRefToImportedStruct) {

xls/dslx/frontend/parser.cc

@@ -3697,6 +3697,9 @@ absl::StatusOr<Impl*> Parser::ParseImpl(const Pos& start_pos, bool is_public,
   Span span(start_pos, GetPos());
   auto* impl = module_->Make<Impl>(span, type, std::move(members), is_public);
   (*struct_def)->set_impl(impl);
+  for (Function* f : impl->GetFunctions()) {
+    f->set_impl(impl);
+  }
   return impl;
 }