diff --git a/tensorflow/stream_executor/lib/statusor.h b/tensorflow/stream_executor/lib/statusor.h index 738abf9589300905cd493df686d5ddb2889bd4b5..2243fb1b34aa7d43b9330594e07426762d695e1b 100644 --- a/tensorflow/stream_executor/lib/statusor.h +++ b/tensorflow/stream_executor/lib/statusor.h @@ -195,6 +195,27 @@ class StatusOr : private internal_statusor::StatusOrData, const T&& ValueOrDie() const &&; T&& ValueOrDie() &&; + // Returns a reference to the current value. + // + // REQUIRES: this->ok() == true, otherwise the behavior is undefined. + // + // Use this->ok() or `operator bool()` to verify that there is a current + // value. Alternatively, see ValueOrDie() for a similar API that guarantees + // CHECK-failing if there is no current value. + const T& operator*() const&; + T& operator*() &; + const T&& operator*() const&&; + T&& operator*() &&; + + // Returns a pointer to the current value. + // + // REQUIRES: this->ok() == true, otherwise the behavior is undefined. + // + // Use this->ok() or `operator bool()` to verify that there is a current + // value. + const T* operator->() const; + T* operator->(); + T ConsumeValueOrDie() { return std::move(ValueOrDie()); } // Ignores any errors. This method does nothing except potentially suppress @@ -303,6 +324,42 @@ T&& StatusOr::ValueOrDie() && { return std::move(this->data_); } +template +const T* StatusOr::operator->() const { + this->EnsureOk(); + return &this->data_; +} + +template +T* StatusOr::operator->() { + this->EnsureOk(); + return &this->data_; +} + +template +const T& StatusOr::operator*() const& { + this->EnsureOk(); + return this->data_; +} + +template +T& StatusOr::operator*() & { + this->EnsureOk(); + return this->data_; +} + +template +const T&& StatusOr::operator*() const&& { + this->EnsureOk(); + return std::move(this->data_); +} + +template +T&& StatusOr::operator*() && { + this->EnsureOk(); + return std::move(this->data_); +} + template void StatusOr::IgnoreError() const { // no-op diff --git a/tensorflow/stream_executor/lib/statusor_test.cc b/tensorflow/stream_executor/lib/statusor_test.cc index 16480b3078937b683a6120125fe68034f48f112b..46bdb9d208fbebda645ed04005580c8096185bcd 100644 --- a/tensorflow/stream_executor/lib/statusor_test.cc +++ b/tensorflow/stream_executor/lib/statusor_test.cc @@ -413,6 +413,26 @@ TEST(StatusOr, TestPointerValueConst) { EXPECT_EQ(&kI, thing.ValueOrDie()); } +TEST(StatusOr, TestArrowOperator) { + StatusOr> uptr = ReturnUniquePtr(); + EXPECT_EQ(*uptr->get(), 0); +} + +TEST(StatusOr, TestArrowOperatorNotOk) { + StatusOr error(Status(tensorflow::error::CANCELLED, "cancelled")); + EXPECT_DEATH(error->pad_++, "cancelled"); +} + +TEST(StatusOr, TestStarOperator) { + StatusOr> uptr = ReturnUniquePtr(); + EXPECT_EQ(**uptr, 0); +} + +TEST(StatusOr, TestStarOperatorDeath) { + StatusOr error(Status(tensorflow::error::CANCELLED, "cancelled")); + EXPECT_DEATH(*error, "cancelled"); +} + // NOTE(tucker): StatusOr does not support this kind // of resize op. // TEST(StatusOr, StatusOrVectorOfUniquePointerCanResize) {