Fix ptr reference invalidation if last reference gets freed before new reference assigned.

2 files changed, 57 insertions, 2 deletions

diff --git a/include/klee/util/Ref.h b/include/klee/util/Ref.h
index c77149aa..44f3f49e 100644
--- a/include/klee/util/Ref.h
+++ b/include/klee/util/Ref.h
@@ -74,16 +74,41 @@ public:
    * despite a redundant template. */
   ref<T> &operator= (const ref<T> &r) {
     r.inc();
+    // Create a copy of the pointer as the
+    // referenced object might get destroyed by the following dec(),
+    // like in the following example:
+    // ````````````````````````
+    //    Expr {
+    //        ref<Expr> next;
+    //    }
+    //
+    //    ref<Expr> root;
+    //    root = root->next;
+    // ````````````````````````
+    T *saved_ptr = r.ptr;
     dec();
-    ptr = r.ptr;
+    ptr = saved_ptr;
 
     return *this;
   }
 
   template<class U> ref<T> &operator= (const ref<U> &r) {
     r.inc();
+    // Create a copy of the pointer as the currently
+    // referenced object might get destroyed by the following dec(),
+    // like in the following example:
+    // ````````````````````````
+    //    Expr {
+    //        ref<Expr> next;
+    //    }
+    //
+    //    ref<Expr> root;
+    //    root = root->next;
+    // ````````````````````````
+
+    U *saved_ptr = r.ptr;
     dec();
-    ptr = r.ptr;
+    ptr = saved_ptr;
 
     return *this;
   }
diff --git a/unittests/Ref/RefTest.cpp b/unittests/Ref/RefTest.cpp
index 48a15885..7d587f69 100644
--- a/unittests/Ref/RefTest.cpp
+++ b/unittests/Ref/RefTest.cpp
@@ -16,6 +16,7 @@
 using klee::ref;
 
 int finished = 0;
+int finished_counter = 0;
 
 struct Expr
 {
@@ -38,3 +39,32 @@ TEST(RefTest, SelfAssign)
   EXPECT_EQ(r_e->refCount, 1);
   finished = 1;
 }
+struct SelfRefExpr {
+  /// @brief Required by klee::ref-managed objects
+  struct klee::ReferenceCounter __refCount;
+  ref<SelfRefExpr> next_;
+
+  explicit SelfRefExpr(ref<SelfRefExpr> next) : next_(next) {}
+  SelfRefExpr(const SelfRefExpr &) = delete;
+  SelfRefExpr &operator=(const SelfRefExpr &) = delete;
+
+  ~SelfRefExpr() { finished_counter++; }
+};
+TEST(RefTest, SelfRef) {
+  struct SelfRefExpr *e_1 = new SelfRefExpr(nullptr);
+  ref<SelfRefExpr> r_e_1(e_1);
+  EXPECT_EQ(1u, r_e_1->__refCount.refCount);
+  ref<SelfRefExpr> r_root = r_e_1;
+  EXPECT_EQ(2u, r_e_1->__refCount.refCount);
+
+  {
+    ref<SelfRefExpr> r2(new SelfRefExpr(r_e_1));
+    EXPECT_EQ(3u, r_e_1->__refCount.refCount);
+
+    r_root = r2;
+    EXPECT_EQ(2u, r_e_1->__refCount.refCount);
+  }
+
+  r_root = r_root->next_;
+  EXPECT_EQ(2u, r_e_1->__refCount.refCount);
+}
\ No newline at end of file