Remove setjmp/longjmp

5 files changed, 37 insertions, 34 deletions

diff --git a/src/common.h b/src/common.h
index 7e75bda..d7e7554 100644
--- a/src/common.h
+++ b/src/common.h
@@ -95,25 +95,12 @@ std::string WideToUtf8(const std::wstring &input);
 #endif
 
 namespace error {
-extern jmp_buf gTestJmp;
-extern bool gTestMode;
-
-inline void Abort() {
-  if (error::gTestMode) {
-    longjmp(error::gTestJmp, 0);
-  } else {
-    std::cerr << "Program terminated with an unrecoverable error." << std::endl;
-    exit(-1);
-  }
-}
 
-inline void Exit(int code) {
-  if (error::gTestMode) {
-    longjmp(error::gTestJmp, 0);
-  } else {
-    exit(code);
-  }
-}
+void Abort();
+void Exit(int code);
+void SetTestCounter(int c);
+void ResetTestMode();
+bool GetTestCounter();
 
 class Die {
  public:
diff --git a/src/error.cc b/src/error.cc
index 93baced..af1147d 100644
--- a/src/error.cc
+++ b/src/error.cc
@@ -18,8 +18,27 @@
 
 namespace sentencepiece {
 namespace error {
-jmp_buf gTestJmp;
-bool gTestMode = false;
+int gTestCounter = 0;
+
+void Abort() {
+  if (GetTestCounter() == 1) {
+    SetTestCounter(2);
+  } else {
+    std::cerr << "Program terminated with an unrecoverable error." << std::endl;
+    exit(-1);
+  }
+}
+
+void Exit(int code) {
+  if (GetTestCounter() == 1) {
+    SetTestCounter(2);
+  } else {
+    exit(code);
+  }
+}
+
+void SetTestCounter(int c) { gTestCounter = c; }
+bool GetTestCounter() { return gTestCounter; }
 }  // namespace error
 
 namespace util {
diff --git a/src/testharness.h b/src/testharness.h
index 5e12573..63ff510 100644
--- a/src/testharness.h
+++ b/src/testharness.h
@@ -135,17 +135,11 @@ class Tester {
 #define EXPECT_OK(c) EXPECT_EQ(c, ::sentencepiece::util::OkStatus())
 #define EXPECT_NOT_OK(c) EXPECT_NE(c, ::sentencepiece::util::OkStatus())
 
-#define EXPECT_DEATH(statement)         \
-  {                                     \
-    error::gTestMode = true;            \
-    if (setjmp(error::gTestJmp) == 0) { \
-      do {                              \
-        statement;                      \
-      } while (false);                  \
-      EXPECT_TRUE(false);               \
-    } else {                            \
-      error::gTestMode = false;         \
-    }                                   \
+#define EXPECT_DEATH(statement) \
+  {                             \
+    error::SetTestCounter(1);   \
+    statement;                  \
+    error::SetTestCounter(0);   \
   };
 
 #define TCONCAT(a, b, c) TCONCAT1(a, b, c)
diff --git a/src/unigram_model.cc b/src/unigram_model.cc
index 34f0595..5e85dc4 100644
--- a/src/unigram_model.cc
+++ b/src/unigram_model.cc
@@ -155,7 +155,10 @@ std::vector<Lattice::Node *> Lattice::Viterbi() {
           best_score = score;
         }
       }
-      CHECK(best_node);
+      if (best_node == nullptr) {
+        LOG(ERROR) << "Failed to find the best path in Viterbi.";
+        return {};
+      }
       rnode->prev = best_node;
       rnode->backtrace_score = best_score;
     }
diff --git a/src/unigram_model_test.cc b/src/unigram_model_test.cc
index 7f99550..66c9c75 100644
--- a/src/unigram_model_test.cc
+++ b/src/unigram_model_test.cc
@@ -161,11 +161,11 @@ TEST(LatticeTest, InsertTest) {
 TEST(LatticeTest, ViterbiFromIncompleteLatticeTest) {
   Lattice lattice;
   lattice.SetSentence("ABC");
-  EXPECT_DEATH(lattice.Viterbi());
+  EXPECT_TRUE(lattice.Viterbi().empty());
 
   // Still incomplete
   lattice.Insert(0, 1);
-  EXPECT_DEATH(lattice.Viterbi());
+  EXPECT_TRUE(lattice.Viterbi().empty());
 
   lattice.Insert(1, 1);
   lattice.Insert(2, 1);