From a5873624aa09ba20c09fe3431b1489d76d33e7ef Mon Sep 17 00:00:00 2001
From: Stephen Seo <seo.disparate@gmail.com>
Date: Wed, 6 Nov 2019 13:32:39 +0900
Subject: [PATCH] Fix size count in TSLQueue when using iterator

---
 cpp_impl/src/TSLQueue.hpp          | 16 ++++++++++++----
 cpp_impl/src/test/TestTSLQueue.cpp |  6 ++++++
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/cpp_impl/src/TSLQueue.hpp b/cpp_impl/src/TSLQueue.hpp
index 5b7410f..d2f09c9 100644
--- a/cpp_impl/src/TSLQueue.hpp
+++ b/cpp_impl/src/TSLQueue.hpp
@@ -62,7 +62,8 @@ class TSLQueue {
     class TSLQIter {
     public:
         TSLQIter(std::mutex &mutex,
-                 std::weak_ptr<TSLQNode> currentNode);
+                 std::weak_ptr<TSLQNode> currentNode,
+                 unsigned long long *msize);
         ~TSLQIter();
 
         std::optional<T> current();
@@ -73,6 +74,8 @@ class TSLQueue {
     private:
         std::lock_guard<std::mutex> lock;
         std::weak_ptr<TSLQNode> currentNode;
+        unsigned long long *msize;
+
     };
 
   public:
@@ -273,9 +276,11 @@ bool TSLQueue<T>::TSLQNode::isNormal() const {
 
 template <typename T>
 TSLQueue<T>::TSLQIter::TSLQIter(std::mutex &mutex,
-                                std::weak_ptr<TSLQNode> currentNode) :
+                                std::weak_ptr<TSLQNode> currentNode,
+                                unsigned long long *msize) :
 lock(mutex),
-currentNode(currentNode)
+currentNode(currentNode),
+msize(msize)
 {
 }
 
@@ -335,12 +340,15 @@ bool TSLQueue<T>::TSLQIter::remove() {
     currentNode->next->prev = parent;
     parent->next = currentNode->next;
 
+    assert(*msize > 0);
+    --(*msize);
+
     return parent->next->isNormal();
 }
 
 template <typename T>
 typename TSLQueue<T>::TSLQIter TSLQueue<T>::begin() {
-    return TSLQIter(mutex, head->next);
+    return TSLQIter(mutex, head->next, &msize);
 }
 
 #endif
diff --git a/cpp_impl/src/test/TestTSLQueue.cpp b/cpp_impl/src/test/TestTSLQueue.cpp
index adffa50..2ac41a5 100644
--- a/cpp_impl/src/test/TestTSLQueue.cpp
+++ b/cpp_impl/src/test/TestTSLQueue.cpp
@@ -123,6 +123,7 @@ TEST(TSLQueue, Iterator) {
     for(int i = 0; i < 10; ++i) {
         q.push(i);
     }
+    EXPECT_EQ(q.size(), 10);
 
     {
         // iteration
@@ -175,6 +176,7 @@ TEST(TSLQueue, Iterator) {
         EXPECT_TRUE(op.has_value());
         EXPECT_EQ(op.value(), 2);
     }
+    EXPECT_EQ(q.size(), 9);
 
     // check that "3" was removed from queue
     int i = 0;
@@ -194,10 +196,12 @@ TEST(TSLQueue, Iterator) {
     q.push(1);
     q.push(2);
     q.push(3);
+    EXPECT_EQ(q.size(), 4);
     {
         auto iter = q.begin();
         EXPECT_TRUE(iter.remove());
     }
+    EXPECT_EQ(q.size(), 3);
     i = 1;
     while(!q.empty()) {
         op = q.top();
@@ -211,6 +215,7 @@ TEST(TSLQueue, Iterator) {
     q.push(1);
     q.push(2);
     q.push(3);
+    EXPECT_EQ(q.size(), 4);
     {
         auto iter = q.begin();
         while(true) {
@@ -223,6 +228,7 @@ TEST(TSLQueue, Iterator) {
             }
         }
     }
+    EXPECT_EQ(q.size(), 3);
     i = 0;
     while(!q.empty()) {
         op = q.top();