Internal change

Improve status-builder for internal users

PiperOrigin-RevId: 611129402

xls/common/status/BUILD

@@ -93,6 +93,7 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
+        "@com_google_absl//absl/types:span",
         "//xls/common:source_location",
         "//xls/common:symbolized_stacktrace",
         "//xls/common/logging",

xls/common/status/status_builder.cc

@@ -15,6 +15,8 @@
 #include "xls/common/status/status_builder.h"
 
 #include <cstdio>
+#include <memory>
+#include <optional>
 #include <ostream>
 #include <string>
 #include <string_view>
@@ -187,7 +189,26 @@ void StatusBuilder::CopyPayloads(const absl::Status& src, absl::Status* dst) {
 
 absl::Status StatusBuilder::WithMessage(const absl::Status& status,
                                         std::string_view msg) {
+  // Unfortunately since we can't easily strip the source-location off of this
+  // new status the backtrace can end up with a lot of copies of this line at
+  // the beginning. We manually try to trim them out but we can't actually
+  // remove the first one.
   auto ret = absl::Status(status.code(), msg);
+  std::optional<SourceLocation> first =
+      StatusBuilder::GetSourceLocations(ret).empty()
+          ? std::nullopt
+          : std::make_optional<SourceLocation>(
+                StatusBuilder::GetSourceLocations(ret).front());
+  bool first_non_duplicate = false;
+  for (const SourceLocation& sl : StatusBuilder::GetSourceLocations(status)) {
+    if (!first_non_duplicate && first && first->line() == sl.line() &&
+        std::string_view(first->file_name()) ==
+            std::string_view(sl.file_name())) {
+      continue;
+    }
+    first_non_duplicate = true;
+    StatusBuilder::AddSourceLocation(ret, sl);
+  }
   CopyPayloads(status, &ret);
   return ret;
 }
@@ -213,6 +234,7 @@ absl::Status StatusBuilder::CreateStatusAndConditionallyLog() && {
   absl::Status result = JoinMessageToStatus(
       std::move(status_), rep_->stream.str(), rep_->message_join_style);
   ConditionallyLog(result);
+  StatusBuilder::AddSourceLocation(result, loc_);
 
   // We consumed the status above, we set it to some error just to prevent
   // people relying on it become OK or something.

xls/common/status/status_builder.h

@@ -28,6 +28,7 @@
 #include "absl/log/log_sink.h"
 #include "absl/status/status.h"
 #include "absl/time/time.h"
+#include "absl/types/span.h"
 #include "xls/common/source_location.h"
 
 // The xabsl namespace has types that are anticipated to become available in
@@ -418,6 +419,11 @@ class ABSL_MUST_USE_RESULT StatusBuilder {
   //
   // This is only for internal use by the StatusBuilder.
   static void AddSourceLocation(absl::Status& status, SourceLocation loc);
+  // Get the current source-location set for the given status.
+  //
+  // This is only for internal use by the StatusBuilder.
+  static absl::Span<const SourceLocation> GetSourceLocations(
+      const absl::Status& status);
 };
 
 // Implicitly converts `builder` to `Status` and write it to `os`.

xls/common/status/status_builder_oss.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "xls/common/source_location.h"
+#include "absl/types/span.h"
 #include "xls/common/status/status_builder.h"
 
 namespace xabsl {
@@ -21,5 +22,10 @@ namespace xabsl {
                                                  xabsl::SourceLocation loc) {
   // Not supported in OSS absl::Status at the moment.
 }
+/* static */ absl::Span<const SourceLocation> StatusBuilder::GetSourceLocations(
+    const absl::Status& status) {
+  // Not supported in OSS absl::Status at the moment.
+  return {};
+}
 
 }  // namespace xabsl

xls/common/status/status_macros_test.cc

@@ -289,6 +289,33 @@ TEST(AssignOrReturn, KeepsBackTrace) {
           "(.*xls/common/status/status_macros_test.cc:[0-9]+.*\n?){3}"));
 }
 
+namespace assign_or_return_with_message {
+absl::Status CallThree() { return absl::InternalError("Clap"); }
+absl::StatusOr<int> CallTwo() {
+  XLS_RETURN_IF_ERROR(CallThree()) << "Your";
+  return 1;
+}
+absl::Status CallOne() {
+  XLS_ASSIGN_OR_RETURN(auto abc, CallTwo(), _ << "Hands");
+  (void)abc;
+  return absl::OkStatus();
+}
+}  // namespace assign_or_return_with_message
+TEST(AssignOrReturn, KeepsBackTraceWithMessage) {
+#ifndef XLS_USE_ABSL_SOURCE_LOCATION
+  GTEST_SKIP() << "Back trace not recorded";
+#endif
+  auto result = assign_or_return_with_message::CallOne();
+  RecordProperty("result", testing::PrintToString(result));
+  EXPECT_THAT(result.message(), Eq("Clap; Your; Hands"));
+  // Expect 3 lines in the stack trace with status_macros_test.cc in them for
+  // the three deep call stack.
+  EXPECT_THAT(
+      result.ToString(absl::StatusToStringMode::kWithEverything),
+      ContainsRegex(
+          "(.*xls/common/status/status_macros_test.cc:[0-9]+.*\n?){3}"));
+}
+
 TEST(ReturnIfError, Works) {
   auto func = []() -> absl::Status {
     XLS_RETURN_IF_ERROR(ReturnOk());
@@ -380,6 +407,33 @@ TEST(ReturnIfError, KeepsBackTrace) {
           "(.*xls/common/status/status_macros_test.cc:[0-9]+.*\n?){3}"));
 }
 
+namespace return_if_error_with_message {
+absl::Status CallThree() { return absl::InternalError("Clap"); }
+absl::Status CallTwo() {
+  XLS_RETURN_IF_ERROR(CallThree()) << "Your";
+  return absl::OkStatus();
+}
+absl::Status CallOne() {
+  XLS_RETURN_IF_ERROR(CallTwo()) << "Hands";
+  return absl::OkStatus();
+}
+}  // namespace return_if_error_with_message
+
+TEST(ReturnIfError, KeepsBackTraceWithMessage) {
+#ifndef XLS_USE_ABSL_SOURCE_LOCATION
+  GTEST_SKIP() << "Back trace not recorded";
+#endif
+  auto result = return_if_error_with_message::CallOne();
+  RecordProperty("result", testing::PrintToString(result));
+  EXPECT_THAT(result.message(), Eq("Clap; Your; Hands"));
+  // Expect 3 lines in the stack trace with status_macros_test.cc in them for
+  // the three deep call stack.
+  EXPECT_THAT(
+      result.ToString(absl::StatusToStringMode::kWithEverything),
+      ContainsRegex(
+          "(.*xls/common/status/status_macros_test.cc:[0-9]+.*\n?){3}"));
+}
+
 // Basis for XLS_RETURN_IF_ERROR and XLS_ASSIGN_OR_RETURN benchmarks.  Derived
 // classes override LoopAgain() with the macro invocation(s).
 template <class T>