[Bf-blender-cvs] [45917d435df] functions: utility functions for set and map

[Jacques Lucke]

Commit: 45917d435df3a1c710347438df3dec82ba6cd3bd
Author: Jacques Lucke
Date:   Mon Apr 8 18:34:35 2019 +0200
Branches: functions
https://developer.blender.org/rB45917d435df3a1c710347438df3dec82ba6cd3bd

utility functions for set and map

===================================================================

M	source/blender/blenlib/BLI_small_map.hpp
M	source/blender/blenlib/BLI_small_set.hpp
M	tests/gtests/blenlib/BLI_small_map_test.cc
M	tests/gtests/blenlib/BLI_small_set_vector_test.cc

===================================================================

diff --git a/source/blender/blenlib/BLI_small_map.hpp b/source/blender/blenlib/BLI_small_map.hpp
index 31a59886df0..0e4e0607c20 100644
--- a/source/blender/blenlib/BLI_small_map.hpp
+++ b/source/blender/blenlib/BLI_small_map.hpp
@@ -107,6 +107,14 @@ namespace BLI {
 			}
 		}
 
+		V *lookup_ptr_or_insert(const K &key, V initial_value)
+		{
+			if (!this->contains(key)) {
+				this->add_new(key, initial_value);
+			}
+			return this->lookup_ptr(key);
+		}
+
 		uint size() const
 		{
 			return m_entries.size();
diff --git a/source/blender/blenlib/BLI_small_set.hpp b/source/blender/blenlib/BLI_small_set.hpp
index 3b9a85e64e2..dacd8f82dec 100644
--- a/source/blender/blenlib/BLI_small_set.hpp
+++ b/source/blender/blenlib/BLI_small_set.hpp
@@ -41,13 +41,22 @@ namespace BLI {
 			return m_lookup.contains(m_elements.begin(), value);
 		}
 
-		void add(const T &value)
+		void add_new(const T &value)
 		{
-			if (!this->contains(value)) {
-				uint index = m_elements.size();
+			BLI_assert(!this->contains(value));
+			uint index = m_elements.size();
+			m_elements.append(value);
+			m_lookup.add_new(m_elements.begin(), index);
+		}
+
+		bool add(const T &value)
+		{
+			uint potential_index = m_elements.size();
+			bool newly_inserted = m_lookup.add(m_elements.begin(), value, potential_index);
+			if (newly_inserted) {
 				m_elements.append(value);
-				m_lookup.add_new(m_elements.begin(), index);
 			}
+			return newly_inserted;
 		}
 
 		T pop()
@@ -59,6 +68,24 @@ namespace BLI {
 			return value;
 		}
 
+		void remove(const T &value)
+		{
+			BLI_assert(this->contains(value));
+			uint index = m_lookup.find(m_elements.begin(), value);
+
+			uint last_index = m_elements.size() - 1;
+			if (index == last_index) {
+				m_elements.remove_last();
+				m_lookup.remove(value, index);
+			}
+			else {
+				m_elements.remove_and_reorder(index);
+				m_lookup.remove(value, index);
+				T &moved_value = m_elements[index];
+				m_lookup.update_index(moved_value, last_index, index);
+			}
+		}
+
 		T any() const
 		{
 			BLI_assert(this->size() > 0);
@@ -77,4 +104,4 @@ namespace BLI {
 
 	};
 
-} /* namespace BLI */
\ No newline at end of file
+} /* namespace BLI */
diff --git a/tests/gtests/blenlib/BLI_small_map_test.cc b/tests/gtests/blenlib/BLI_small_map_test.cc
index 6f866e690c6..b349e503ebb 100644
--- a/tests/gtests/blenlib/BLI_small_map_test.cc
+++ b/tests/gtests/blenlib/BLI_small_map_test.cc
@@ -85,3 +85,13 @@ TEST(small_map, PopItemMany)
 		EXPECT_EQ(map.contains(i), i < 25 || i >= 80);
 	}
 }
+
+TEST(small_map, LookupPtrOrInsert)
+{
+	IntFloatMap map;
+	float *value = map.lookup_ptr_or_insert(3, 5.0f);
+	EXPECT_EQ(*value, 5.0f);
+	*value += 1;
+	value = map.lookup_ptr_or_insert(3, 5.0f);
+	EXPECT_EQ(*value, 6.0f);
+}
diff --git a/tests/gtests/blenlib/BLI_small_set_vector_test.cc b/tests/gtests/blenlib/BLI_small_set_vector_test.cc
index a6fb51832c2..0da54f5130a 100644
--- a/tests/gtests/blenlib/BLI_small_set_vector_test.cc
+++ b/tests/gtests/blenlib/BLI_small_set_vector_test.cc
@@ -137,4 +137,24 @@ TEST(small_set, Pop)
 	EXPECT_EQ(set.size(), 1);
 	EXPECT_EQ(set.pop(), 4);
 	EXPECT_EQ(set.size(), 0);
-}
\ No newline at end of file
+}
+
+TEST(small_set, Remove)
+{
+	IntSet set = {3, 4, 5};
+	EXPECT_TRUE(set.contains(3));
+	EXPECT_TRUE(set.contains(4));
+	EXPECT_TRUE(set.contains(5));
+	set.remove(4);
+	EXPECT_TRUE(set.contains(3));
+	EXPECT_FALSE(set.contains(4));
+	EXPECT_TRUE(set.contains(5));
+	set.remove(3);
+	EXPECT_FALSE(set.contains(3));
+	EXPECT_FALSE(set.contains(4));
+	EXPECT_TRUE(set.contains(5));
+	set.remove(5);
+	EXPECT_FALSE(set.contains(3));
+	EXPECT_FALSE(set.contains(4));
+	EXPECT_FALSE(set.contains(5));
+}