Fix bug where PSC IR conversion drops trait-derived functions too eagerly.

The problem was that when HandleInvocation would encounter an invocation like `c.to_bits()`, it would create a new visitor using the callee TypeInfo for that invocation, and call `new_visitor->HandleFunctionInternal(callee)`. The latter would then check if it was a derived function with no calls. But the call corresponding to the HandleInvocation in this scenario would not be in the callee TypeInfo.

So this changes the logic to convert the function if we encounter it via a call. We also add better logging.

PiperOrigin-RevId: 896099586

Author: richmckeever

xls/dslx/ir_convert/BUILD

@@ -531,6 +531,7 @@ cc_library(
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
     ],
 )

xls/dslx/ir_convert/get_conversion_records.cc

@@ -14,9 +14,11 @@
 
 #include "xls/dslx/ir_convert/get_conversion_records.h"
 
+#include <cstdint>
 #include <ios>
 #include <memory>
 #include <optional>
+#include <string>
 #include <string_view>
 #include <utility>
 #include <variant>
@@ -26,6 +28,7 @@
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "xls/common/status/ret_check.h"
 #include "xls/common/status/status_macros.h"
@@ -160,7 +163,8 @@ class ConversionRecordVisitor : public AstNodeVisitorWithDefault {
       return absl::OkStatus();
     }
 
-    return HandleFunctionInternal(f, ParametricEnv{});
+    return HandleFunctionInternal(f, ParametricEnv{},
+                                  /*handle_for_invocation=*/false);
   }
 
   // Handles all non-proc functions. For parametric functions, we only get here
@@ -169,7 +173,8 @@ class ConversionRecordVisitor : public AstNodeVisitorWithDefault {
   // non-parametric, standalone functions, the `env` should be empty and the
   // `type_info_` on the visitor can be the root one.
   absl::Status HandleFunctionInternal(const Function* f,
-                                      const ParametricEnv& env) {
+                                      const ParametricEnv& env,
+                                      bool handle_for_invocation) {
     XLS_RET_CHECK(module_ == f->owner());
     XLS_RET_CHECK(!f->IsInProc());
     if (f->IsParametric() || f->IsMethodOnParametricStruct()) {
@@ -181,9 +186,14 @@ class ConversionRecordVisitor : public AstNodeVisitorWithDefault {
 
     std::vector<InvocationCalleeData> calls =
         type_info_->GetUniqueInvocationCalleeData(f);
-    if (f->IsCompilerDerived() && calls.empty()) {
-      VLOG(5) << "No calls to derived function " << f->name_def()->ToString()
-              << "; not traversing for dependencies.";
+    if (f->IsCompilerDerived() && calls.empty() && !handle_for_invocation) {
+      std::string f_name = f->identifier();
+      if (f->impl().has_value()) {
+        f_name =
+            absl::StrCat((*f->impl())->struct_ref()->ToString(), "::", f_name);
+      }
+      VLOG(5) << "No calls to derived function " << f_name << " in TI "
+              << type_info_->name() << "; not traversing for dependencies.";
       return absl::OkStatus();
     }
 
@@ -247,8 +257,8 @@ class ConversionRecordVisitor : public AstNodeVisitorWithDefault {
           proc_id_factory_, top_, resolved_proc_alias_, records_,
           processed_invocations_);
 
-      XLS_RETURN_IF_ERROR(
-          visitor.HandleFunctionInternal(call.callee, call.callee_bindings));
+      XLS_RETURN_IF_ERROR(visitor.HandleFunctionInternal(
+          call.callee, call.callee_bindings, /*handle_for_invocation=*/true));
 
       VLOG(5) << "Processing invocation " << invocation->ToString();
       XLS_ASSIGN_OR_RETURN(

xls/dslx/ir_convert/ir_converter_test.cc

@@ -8626,5 +8626,24 @@ proc main {
   ExpectIr(converted);
 }
 
+TEST_F(IrConverterTest, DerivedToBitsNestedStructs) {
+  constexpr std::string_view kDslxProgram = R"(
+#[derive(ToBits)]
+pub struct A { x: u32 }
+
+#[derive(ToBits)]
+pub struct B { a: A }
+
+#[derive(ToBits)]
+pub struct C { b: B }
+
+pub fn main(c: C) -> bits[32] {
+  c.to_bits()
+}
+)";
+
+  XLS_ASSERT_OK(ConvertModuleForTest(kDslxProgram));
+}
+
 }  // namespace
 }  // namespace xls::dslx

xls/dslx/type_system_v2/inference_table_converter_impl.cc

@@ -683,8 +683,7 @@ class InferenceTableConverterImpl : public InferenceTableConverter,
     XLS_ASSIGN_OR_RETURN(
         TypeInfo * invocation_type_info,
         import_data_.type_info_owner().New(
-            function->owner(),
-            absl::StrCat("invocation_of_", function->identifier()),
+            function->owner(), CreateInvocationTypeInfoName(function),
             base_type_info));
 
     std::optional<const StructDefBase*> target_struct = std::nullopt;
@@ -913,6 +912,17 @@ class InferenceTableConverterImpl : public InferenceTableConverter,
     return absl::OkStatus();
   }
 
+  // Generates a name for a `TypeInfo` object to be used for an invocation of
+  // `f`.
+  std::string CreateInvocationTypeInfoName(const Function* f) {
+    if (f->impl().has_value()) {
+      return absl::Substitute("invocation_of_$0::$1",
+                              (*f->impl())->struct_ref()->ToString(),
+                              f->identifier());
+    }
+    return absl::StrCat("invocation_of_", f->identifier());
+  }
+
   // Gets the output `TypeInfo` corresponding to the given
   // `parametric_context`, which may be `nullopt`, in which case it returns
   // the base type info. If the proposed type info doesn't belong to the