From 29dcf92f870ee51cce8d68f8fcfe228942e8dfe1 Mon Sep 17 00:00:00 2001
From: JeffWegher <jeff_wegher@yahoo.com>
Date: Fri, 14 Feb 2020 12:38:58 -0800
Subject: [PATCH] Fix Node::size() after removing node from sequence (#820)

---
 include/yaml-cpp/node/detail/impl.h | 17 +++++++++----
 test/node/node_test.cpp             | 39 +++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/include/yaml-cpp/node/detail/impl.h b/include/yaml-cpp/node/detail/impl.h
index bfc1d6f..a5cbc0a 100644
--- a/include/yaml-cpp/node/detail/impl.h
+++ b/include/yaml-cpp/node/detail/impl.h
@@ -58,7 +58,9 @@ struct get_idx<Key, typename std::enable_if<std::is_signed<Key>::value>::type> {
 
 template <typename Key, typename Enable = void>
 struct remove_idx {
-  static bool remove(std::vector<node*>&, const Key&) { return false; }
+  static bool remove(std::vector<node*>&, const Key&, std::size_t&) {
+    return false;
+  }
 };
 
 template <typename Key>
@@ -66,11 +68,15 @@ struct remove_idx<
     Key, typename std::enable_if<std::is_unsigned<Key>::value &&
                                  !std::is_same<Key, bool>::value>::type> {
 
-  static bool remove(std::vector<node*>& sequence, const Key& key) {
+  static bool remove(std::vector<node*>& sequence, const Key& key,
+                     std::size_t& seqSize) {
     if (key >= sequence.size()) {
       return false;
     } else {
       sequence.erase(sequence.begin() + key);
+      if (seqSize > key) {
+          --seqSize;
+      }
       return true;
     }
   }
@@ -80,9 +86,10 @@ template <typename Key>
 struct remove_idx<Key,
                   typename std::enable_if<std::is_signed<Key>::value>::type> {
 
-  static bool remove(std::vector<node*>& sequence, const Key& key) {
+  static bool remove(std::vector<node*>& sequence, const Key& key,
+                     std::size_t& seqSize) {
     return key >= 0 ? remove_idx<std::size_t>::remove(
-                          sequence, static_cast<std::size_t>(key))
+                          sequence, static_cast<std::size_t>(key), seqSize)
                     : false;
   }
 };
@@ -161,7 +168,7 @@ inline node& node_data::get(const Key& key, shared_memory_holder pMemory) {
 template <typename Key>
 inline bool node_data::remove(const Key& key, shared_memory_holder pMemory) {
   if (m_type == NodeType::Sequence) {
-    return remove_idx<Key>::remove(m_sequence, key);
+    return remove_idx<Key>::remove(m_sequence, key, m_seqSize);
   } else if (m_type == NodeType::Map) {
     kv_pairs::iterator it = m_undefinedPairs.begin();
     while (it != m_undefinedPairs.end()) {
diff --git a/test/node/node_test.cpp b/test/node/node_test.cpp
index 0677ad5..d13a4bb 100644
--- a/test/node/node_test.cpp
+++ b/test/node/node_test.cpp
@@ -59,6 +59,19 @@ TEST(NodeTest, SequenceElementRemoval) {
   EXPECT_EQ("c", node[1].as<std::string>());
 }
 
+TEST(NodeTest, SequenceElementRemovalSizeCheck) {
+  Node node;
+  node[0] = "a";
+  node[1] = "b";
+  node[2] = "c";
+  EXPECT_EQ(3, node.size());
+  node.remove(1);
+  EXPECT_TRUE(node.IsSequence());
+  EXPECT_EQ(2, node.size());
+  EXPECT_EQ("a", node[0].as<std::string>());
+  EXPECT_EQ("c", node[1].as<std::string>());
+}
+
 TEST(NodeTest, SequenceFirstElementRemoval) {
   Node node;
   node[0] = "a";
@@ -71,6 +84,19 @@ TEST(NodeTest, SequenceFirstElementRemoval) {
   EXPECT_EQ("c", node[1].as<std::string>());
 }
 
+TEST(NodeTest, SequenceFirstElementRemovalSizeCheck) {
+  Node node;
+  node[0] = "a";
+  node[1] = "b";
+  node[2] = "c";
+  EXPECT_EQ(3, node.size());
+  node.remove(0);
+  EXPECT_TRUE(node.IsSequence());
+  EXPECT_EQ(2, node.size());
+  EXPECT_EQ("b", node[0].as<std::string>());
+  EXPECT_EQ("c", node[1].as<std::string>());
+}
+
 TEST(NodeTest, SequenceLastElementRemoval) {
   Node node;
   node[0] = "a";
@@ -83,6 +109,19 @@ TEST(NodeTest, SequenceLastElementRemoval) {
   EXPECT_EQ("b", node[1].as<std::string>());
 }
 
+TEST(NodeTest, SequenceLastElementRemovalSizeCheck) {
+  Node node;
+  node[0] = "a";
+  node[1] = "b";
+  node[2] = "c";
+  EXPECT_EQ(3, node.size());
+  node.remove(2);
+  EXPECT_TRUE(node.IsSequence());
+  EXPECT_EQ(2, node.size());
+  EXPECT_EQ("a", node[0].as<std::string>());
+  EXPECT_EQ("b", node[1].as<std::string>());
+}
+
 TEST(NodeTest, NodeAssignment) {
   Node node1;
   Node node2;