maybe: add JUST_MSG and CHECK_JUST_MSG

oneflow/core/common/maybe.h

@@ -269,33 +269,67 @@ inline bool MaybeIsOk(Maybe<void>&& maybe) {
 
 #if defined(__GNUC__) || defined(__CUDACC__) || defined(__clang__)
 
+inline void maybe_error_add_stack_frame(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&& maybe_error_add_message(Error&& err, T&&... msg) {
+  __attribute__((unused)) int dummy[] = {((void)(std::move(err) << std::forward<T>(msg)), 0)...};
+  return std::move(err);
+}
+
 #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()) {                                                           \
+      maybe_error_add_stack_frame(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()) {                                                        \
+      maybe_error_add_stack_frame(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 maybe_error_add_message(                                           \
+          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) << maybe_error_add_message(                                               \
+                        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

@@ -0,0 +1,46 @@
+/*
+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 data = std::mem_fn(&Maybe<int>::Data_YouAreNotAllowedToCallThisFuncOutsideThisFile);
+
+  EXPECT_EQ(data(h(1)), 233);
+
+  auto err = h(10).error();
+  EXPECT_EQ(err->msg(), "input value 53");
+  EXPECT_EQ(err->stack_frame(0).error_msg(), "(f(y)): input value g(10)");
+}
+
+}  // namespace test
+}  // namespace oneflow