From a3d4a8c5b61905a6511b3099e3132138bb2f06c3 Mon Sep 17 00:00:00 2001
From: Nelson Liang <nelsonliang@google.com>
Date: Thu, 16 Oct 2025 07:51:42 -0700
Subject: [PATCH] Set error payload for type inference errors with single spans

PiperOrigin-RevId: 820232578
---
 xls/dslx/errors.cc      | 10 +++++++---
 xls/dslx/errors_test.cc | 10 ++++++----
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/xls/dslx/errors.cc b/xls/dslx/errors.cc
index 565cf9df93..fdd483b52a 100644
--- a/xls/dslx/errors.cc
+++ b/xls/dslx/errors.cc
@@ -45,9 +45,13 @@ absl::Status TypeInferenceErrorStatusInternal(
   if (type_or_annotation != nullptr) {
     type_str = type_or_annotation->ToString() + " ";
   }
-  return absl::InvalidArgumentError(
-      absl::StrFormat("TypeInferenceError: %s %s%s", span.ToString(file_table),
-                      type_str, message));
+
+  absl::Status status = absl::InvalidArgumentError(
+      absl::StrFormat("TypeInferenceError: %s%s", type_str, message));
+  StatusPayloadProto payload;
+  *payload.add_spans() = ToProto(span, file_table);
+  SetStatusPayload(status, payload);
+  return status;
 }
 
 }  // namespace
diff --git a/xls/dslx/errors_test.cc b/xls/dslx/errors_test.cc
index 68c019a093..c475665f1c 100644
--- a/xls/dslx/errors_test.cc
+++ b/xls/dslx/errors_test.cc
@@ -42,10 +42,12 @@ TEST(ErrorsTest, TypeInferenceErrorMessage) {
   std::unique_ptr<Type> type = BitsType::MakeU32();
   absl::Status status = TypeInferenceErrorStatus(
       span, type.get(), "this is the message!", file_table);
-  EXPECT_EQ(
-      status.ToString(),
-      "INVALID_ARGUMENT: TypeInferenceError: <no-file>:1:1-2:2 uN[32] this "
-      "is the message!");
+  std::optional<StatusPayloadProto> payload = GetStatusPayload(status);
+  ASSERT_TRUE(payload.has_value());
+  EXPECT_THAT(payload->spans(),
+              UnorderedElementsAre(EqualsProto(ToProto(span, file_table))));
+  EXPECT_EQ(status.message(),
+            "TypeInferenceError: uN[32] this is the message!");
 }
 
 TEST(ErrorsTest, SignednessMismatchErrorAnnotationPayload) {