maybe: add `JUST_MSG` and `CHECK_JUST_MSG` (#5904)

* maybe: add JUST_MSG and CHECK_JUST_MSG

* maybe: add test

* maybe: support ostream && fix test

* maybe: put to private_details

* maybe: fix name

* error: clang format
Co-authored-by: Noneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

oneflow/core/common/error.h

@@ -112,6 +112,12 @@ Error&& operator<<(Error&& error, const T& x) {
   return std::move(error);
 }
 
+template<>
+inline Error&& operator<<(Error&& error, const std::ostream& x) {
+  error << x.rdbuf();
+  return std::move(error);
+}
+
 template<>
 inline Error&& operator<<(Error&& error, const Error& other) {
   error.Assign(other);

oneflow/core/common/maybe.h

@@ -269,33 +269,72 @@ inline bool MaybeIsOk(Maybe<void>&& maybe) {
 
 #if defined(__GNUC__) || defined(__CUDACC__) || defined(__clang__)
 
+namespace private_details {
+
+inline void MaybeErrorAddStackFrame(const std::shared_ptr<cfg::ErrorProto>& err,
+                                    const std::string& file, int64_t line, const std::string& func,
+                                    const std::string& message) {
+  auto* stack_frame = err->add_stack_frame();
+  stack_frame->set_file(file);
+  stack_frame->set_line(line);
+  stack_frame->set_function(func);
+  stack_frame->set_error_msg(message);
+}
+
+template<typename... T>
+Error&& MaybeErrorAddMessage(Error&& err, T&&... msg) {
+  __attribute__((unused)) int dummy[] = {((void)(std::move(err) << std::forward<T>(msg)), 0)...};
+  return std::move(err);
+}
+
+}  // namespace private_details
+
 #define TRY(...) __MaybeErrorStackCheckWrapper__(__VA_ARGS__)
-#define JUST(...)                                                 \
-  ({                                                              \
-    auto&& maybe = __MaybeErrorStackCheckWrapper__(__VA_ARGS__);  \
-    if (!maybe.IsOk()) {                                          \
-      auto* stack_frame = maybe.error()->add_stack_frame();       \
-      stack_frame->set_file(__FILE__);                            \
-      stack_frame->set_line(__LINE__);                            \
-      stack_frame->set_function(__FUNCTION__);                    \
-      stack_frame->set_error_msg(OF_PP_STRINGIZE((__VA_ARGS__))); \
-      return maybe.error();                                       \
-    }                                                             \
-    std::move(maybe);                                             \
+#define JUST(...)                                                                           \
+  ({                                                                                        \
+    auto&& maybe = __MaybeErrorStackCheckWrapper__(__VA_ARGS__);                            \
+    if (!maybe.IsOk()) {                                                                    \
+      ::oneflow::private_details::MaybeErrorAddStackFrame(                                  \
+          maybe.error(), __FILE__, __LINE__, __FUNCTION__, OF_PP_STRINGIZE((__VA_ARGS__))); \
+      return maybe.error();                                                                 \
+    }                                                                                       \
+    std::move(maybe);                                                                       \
   }).Data_YouAreNotAllowedToCallThisFuncOutsideThisFile()
-#define CHECK_JUST(...)                                           \
-  ([&](const char* func_name) {                                   \
-    auto&& maybe = __MaybeErrorStackCheckWrapper__(__VA_ARGS__);  \
-    if (!maybe.IsOk()) {                                          \
-      auto* stack_frame = maybe.error()->add_stack_frame();       \
-      stack_frame->set_file(__FILE__);                            \
-      stack_frame->set_line(__LINE__);                            \
-      stack_frame->set_function(func_name);                       \
-      stack_frame->set_error_msg(OF_PP_STRINGIZE((__VA_ARGS__))); \
-      LOG(FATAL) << maybe.GetSerializedError();                   \
-    }                                                             \
-    return std::move(maybe);                                      \
-  })(__FUNCTION__)                                                \
+#define CHECK_JUST(...)                                                                  \
+  ([&](const char* func_name) {                                                          \
+    auto&& maybe = __MaybeErrorStackCheckWrapper__(__VA_ARGS__);                         \
+    if (!maybe.IsOk()) {                                                                 \
+      ::oneflow::private_details::MaybeErrorAddStackFrame(                               \
+          maybe.error(), __FILE__, __LINE__, func_name, OF_PP_STRINGIZE((__VA_ARGS__))); \
+      LOG(FATAL) << maybe.GetSerializedError();                                          \
+    }                                                                                    \
+    return std::move(maybe);                                                             \
+  })(__FUNCTION__)                                                                       \
+      .Data_YouAreNotAllowedToCallThisFuncOutsideThisFile()
+
+#define JUST_MSG(value, ...)                                                               \
+  ({                                                                                       \
+    auto&& maybe = (value);                                                                \
+    if (!maybe.IsOk()) {                                                                   \
+      return ::oneflow::private_details::MaybeErrorAddMessage(                             \
+          ::oneflow::Error(maybe.error()).AddStackFrame(__FILE__, __LINE__, __FUNCTION__), \
+          OF_PP_STRINGIZE((value)), ": ", __VA_ARGS__);                                    \
+    }                                                                                      \
+    std::move(maybe);                                                                      \
+  }).Data_YouAreNotAllowedToCallThisFuncOutsideThisFile()
+
+#define CHECK_JUST_MSG(value, ...)                                                             \
+  ([&](const char* func_name) {                                                                \
+    auto&& maybe = (value);                                                                    \
+    if (!maybe.IsOk()) {                                                                       \
+      LOG(FATAL)                                                                               \
+          << ::oneflow::private_details::MaybeErrorAddMessage(                                 \
+                 ::oneflow::Error(maybe.error()).AddStackFrame(__FILE__, __LINE__, func_name), \
+                 OF_PP_STRINGIZE((value)), ": ", __VA_ARGS__)                                  \
+                 ->DebugString();                                                              \
+    }                                                                                          \
+    return std::move(maybe);                                                                   \
+  })(__FUNCTION__)                                                                             \
       .Data_YouAreNotAllowedToCallThisFuncOutsideThisFile()
 
 #define CHECK_OK(...) CHECK(MaybeIsOk(__VA_ARGS__))

oneflow/core/common/maybe_test.cpp

+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "oneflow/core/common/maybe.h"
+#include "oneflow/core/common/util.h"
+
+namespace oneflow {
+namespace test {
+
+TEST(Maybe, JUST_MSG) {
+  auto f = [](int x) -> Maybe<int> {
+    if (x > 10) { return Error::ValueError("") << "input value " << x; }
+
+    return 233;
+  };
+
+  auto g = [](int x) { return x * x - 5 * x + 3; };
+
+  auto h = [&](int x) -> Maybe<int> {
+    auto y = g(x);
+    return JUST_MSG(f(y), "input value g(", x, ")");
+  };
+
+  auto i = [&](float x) -> Maybe<int> {
+    int y = x;
+    return JUST_MSG(h(y), std::stringstream() << "input value int(" << x << ")");
+  };
+
+  auto data = CHECK_JUST(i(1));
+  ASSERT_EQ(data, 233);
+
+  auto err = i(10.123).error();
+  ASSERT_EQ(err->msg(), "input value 53");
+  ASSERT_EQ(err->stack_frame(0).error_msg(), "(f(y)): input value g(10)");
+  ASSERT_EQ(err->stack_frame(1).error_msg(), "(h(y)): input value int(10.123)");
+}
+
+}  // namespace test
+}  // namespace oneflow