[xla:ffi] Add ErrorOr container for returning value or error

PiperOrigin-RevId: 652558243
ROCm · Jul 15, 2024 · 83e3852 · 83e3852
1 parent 8d2a60d
commit 83e3852
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 2 deletions.
diff --git a/xla/ffi/api/api.h b/xla/ffi/api/api.h
@@ -975,6 +975,7 @@ class Expected {
   constexpr const T* operator->() const { return &value(); }
 
   constexpr bool has_value() const { return std::holds_alternative<T>(data_); }
+  constexpr bool has_error() const { return std::holds_alternative<E>(data_); }
 
   constexpr T& value() & { return std::get<T>(data_); }
   constexpr const T& value() const& { return std::get<T>(data_); }
@@ -993,7 +994,7 @@ class Expected {
 template <typename E>
 class Unexpected {
  public:
-  explicit constexpr Unexpected(E error) : error_(std::move(error)) {}
+  constexpr Unexpected(E error) : error_(std::move(error)) {}  // NOLINT
 
  private:
   template <typename, typename>

diff --git a/xla/ffi/api/ffi.h b/xla/ffi/api/ffi.h
@@ -144,7 +144,7 @@ class Span {
 };
 
 //===----------------------------------------------------------------------===//
-// Error
+// Error and ErrorOr
 //===----------------------------------------------------------------------===//
 
 enum class ErrorCode : uint8_t {
@@ -190,6 +190,12 @@ class Error {
   std::string message_;
 };
 
+template <typename T>
+class ErrorOr : public Expected<T, Error> {
+ public:
+  using Expected<T, Error>::Expected;
+};
+
 //===----------------------------------------------------------------------===//
 // Arguments
 //===----------------------------------------------------------------------===//

diff --git a/xla/ffi/api/ffi_test.cc b/xla/ffi/api/ffi_test.cc
@@ -187,6 +187,18 @@ TEST(FfiTest, ErrorEnumValue) {
             encoded(ErrorCode::kUnauthenticated));
 }
 
+TEST(FfiTest, Expected) {
+  ErrorOr<int32_t> value(42);
+  EXPECT_TRUE(value.has_value());
+  EXPECT_FALSE(value.has_error());
+  EXPECT_EQ(*value, 42);
+
+  ErrorOr<int32_t> error(Error(ErrorCode::kInternal, "Test error"));
+  EXPECT_FALSE(error.has_value());
+  EXPECT_TRUE(error.has_error());
+  EXPECT_THAT(error.error().message(), HasSubstr("Test error"));
+}
+
 TEST(FfiTest, ReturnError) {
   CallFrameBuilder builder(/*num_args=*/0, /*num_rets=*/0);
   auto call_frame = builder.Build();