FTL: Add SmallMap mutators

Add try_emplace, try_replace, emplace_or_replace, and erase. Rename
getters so that `find` returns an iterator like standard containers.

Bug: 185536303
Test: ftl_test
Change-Id: I095bb800c6bd786b9fceaf632aa06b2d6cdadb71

3 files changed, 340 insertions, 37 deletions

diff --git a/include/ftl/small_map.h b/include/ftl/small_map.h
index 84c15ebdca..bcaba82934 100644
--- a/include/ftl/small_map.h
+++ b/include/ftl/small_map.h
@@ -19,6 +19,7 @@
 #include <ftl/initializer_list.h>
 #include <ftl/small_vector.h>
 
+#include <algorithm>
 #include <functional>
 #include <optional>
 #include <type_traits>
@@ -28,7 +29,10 @@ namespace android::ftl {
 
 // Associative container with unique, unordered keys. Unlike std::unordered_map, key-value pairs are
 // stored in contiguous storage for cache efficiency. The map is allocated statically until its size
-// exceeds N, at which point mappings are relocated to dynamic memory.
+// exceeds N, at which point mappings are relocated to dynamic memory. The try_emplace operation has
+// a non-standard analogue try_replace that destructively emplaces. The API also defines an in-place
+// counterpart to insert_or_assign: emplace_or_replace. Lookup is done not via a subscript operator,
+// but immutable getters that can optionally transform the value.
 //
 // SmallMap<K, V, 0> unconditionally allocates on the heap.
 //
@@ -43,16 +47,19 @@ namespace android::ftl {
 //   assert(!map.dynamic());
 //
 //   assert(map.contains(123));
-//   assert(map.find(42, [](const std::string& s) { return s.size(); }) == 3u);
+//   assert(map.get(42, [](const std::string& s) { return s.size(); }) == 3u);
 //
-//   const auto opt = map.find(-1);
+//   const auto opt = map.get(-1);
 //   assert(opt);
 //
 //   std::string& ref = *opt;
 //   assert(ref.empty());
 //   ref = "xyz";
 //
-//   assert(map == SmallMap(ftl::init::map(-1, "xyz")(42, "???")(123, "abc")));
+//   map.emplace_or_replace(0, "vanilla", 2u, 3u);
+//   assert(map.dynamic());
+//
+//   assert(map == SmallMap(ftl::init::map(-1, "xyz")(0, "nil")(42, "???")(123, "abc")));
 //
 template <typename K, typename V, std::size_t N>
 class SmallMap final {
@@ -80,12 +87,7 @@ class SmallMap final {
   // The syntax for listing pairs is as follows:
   //
   //   ftl::SmallMap map = ftl::init::map<int, std::string>(123, "abc")(-1)(42, 3u, '?');
-  //
   //   static_assert(std::is_same_v<decltype(map), ftl::SmallMap<int, std::string, 3>>);
-  //   assert(map.size() == 3u);
-  //   assert(map.contains(-1) && map.find(-1)->get().empty());
-  //   assert(map.contains(42) && map.find(42)->get() == "???");
-  //   assert(map.contains(123) && map.find(123)->get() == "abc");
   //
   // The types of the key and value are deduced if the first pair contains exactly two arguments:
   //
@@ -95,7 +97,7 @@ class SmallMap final {
   template <typename U, std::size_t... Sizes, typename... Types>
   SmallMap(InitializerList<U, std::index_sequence<Sizes...>, Types...>&& list)
       : map_(std::move(list)) {
-    // TODO: Enforce unique keys.
+    deduplicate();
   }
 
   size_type max_size() const { return map_.max_size(); }
@@ -115,27 +117,27 @@ class SmallMap final {
 
   // Returns whether a mapping exists for the given key.
   bool contains(const key_type& key) const {
-    return find(key, [](const mapped_type&) {});
+    return get(key, [](const mapped_type&) {});
   }
 
   // Returns a reference to the value for the given key, or std::nullopt if the key was not found.
   //
   //   ftl::SmallMap map = ftl::init::map('a', 'A')('b', 'B')('c', 'C');
   //
-  //   const auto opt = map.find('c');
+  //   const auto opt = map.get('c');
   //   assert(opt == 'C');
   //
   //   char d = 'd';
-  //   const auto ref = map.find('d').value_or(std::ref(d));
+  //   const auto ref = map.get('d').value_or(std::ref(d));
   //   ref.get() = 'D';
   //   assert(d == 'D');
   //
-  auto find(const key_type& key) const -> std::optional<std::reference_wrapper<const mapped_type>> {
-    return find(key, [](const mapped_type& v) { return std::cref(v); });
+  auto get(const key_type& key) const -> std::optional<std::reference_wrapper<const mapped_type>> {
+    return get(key, [](const mapped_type& v) { return std::cref(v); });
   }
 
-  auto find(const key_type& key) -> std::optional<std::reference_wrapper<mapped_type>> {
-    return find(key, [](mapped_type& v) { return std::ref(v); });
+  auto get(const key_type& key) -> std::optional<std::reference_wrapper<mapped_type>> {
+    return get(key, [](mapped_type& v) { return std::ref(v); });
   }
 
   // Returns the result R of a unary operation F on (a constant or mutable reference to) the value
@@ -144,11 +146,11 @@ class SmallMap final {
   //
   //   ftl::SmallMap map = ftl::init::map('a', 'x')('b', 'y')('c', 'z');
   //
-  //   assert(map.find('c', [](char c) { return std::toupper(c); }) == 'Z');
-  //   assert(map.find('c', [](char& c) { c = std::toupper(c); }));
+  //   assert(map.get('c', [](char c) { return std::toupper(c); }) == 'Z');
+  //   assert(map.get('c', [](char& c) { c = std::toupper(c); }));
   //
   template <typename F, typename R = std::invoke_result_t<F, const mapped_type&>>
-  auto find(const key_type& key, F f) const
+  auto get(const key_type& key, F f) const
       -> std::conditional_t<std::is_void_v<R>, bool, std::optional<R>> {
     for (auto& [k, v] : *this) {
       if (k == key) {
@@ -165,12 +167,96 @@ class SmallMap final {
   }
 
   template <typename F>
-  auto find(const key_type& key, F f) {
-    return std::as_const(*this).find(
+  auto get(const key_type& key, F f) {
+    return std::as_const(*this).get(
         key, [&f](const mapped_type& v) { return f(const_cast<mapped_type&>(v)); });
   }
 
+  // Returns an iterator to an existing mapping for the given key, or the end() iterator otherwise.
+  const_iterator find(const key_type& key) const { return const_cast<SmallMap&>(*this).find(key); }
+  iterator find(const key_type& key) { return find(key, begin()); }
+
+  // Inserts a mapping unless it exists. Returns an iterator to the inserted or existing mapping,
+  // and whether the mapping was inserted.
+  //
+  // On emplace, if the map reaches its static or dynamic capacity, then all iterators are
+  // invalidated. Otherwise, only the end() iterator is invalidated.
+  //
+  template <typename... Args>
+  std::pair<iterator, bool> try_emplace(const key_type& key, Args&&... args) {
+    if (const auto it = find(key); it != end()) {
+      return {it, false};
+    }
+
+    auto& ref = map_.emplace_back(std::piecewise_construct, std::forward_as_tuple(key),
+                                  std::forward_as_tuple(std::forward<Args>(args)...));
+    return {&ref, true};
+  }
+
+  // Replaces a mapping if it exists, and returns an iterator to it. Returns the end() iterator
+  // otherwise.
+  //
+  // The value is replaced via move constructor, so type V does not need to define copy/move
+  // assignment, e.g. its data members may be const.
+  //
+  // The arguments may directly or indirectly refer to the mapping being replaced.
+  //
+  // Iterators to the replaced mapping point to its replacement, and others remain valid.
+  //
+  template <typename... Args>
+  iterator try_replace(const key_type& key, Args&&... args) {
+    const auto it = find(key);
+    if (it == end()) return it;
+    map_.replace(it, std::piecewise_construct, std::forward_as_tuple(key),
+                 std::forward_as_tuple(std::forward<Args>(args)...));
+    return it;
+  }
+
+  // In-place counterpart of std::unordered_map's insert_or_assign. Returns true on emplace, or
+  // false on replace.
+  //
+  // The value is emplaced and replaced via move constructor, so type V does not need to define
+  // copy/move assignment, e.g. its data members may be const.
+  //
+  // On emplace, if the map reaches its static or dynamic capacity, then all iterators are
+  // invalidated. Otherwise, only the end() iterator is invalidated. On replace, iterators
+  // to the replaced mapping point to its replacement, and others remain valid.
+  //
+  template <typename... Args>
+  std::pair<iterator, bool> emplace_or_replace(const key_type& key, Args&&... args) {
+    const auto [it, ok] = try_emplace(key, std::forward<Args>(args)...);
+    if (ok) return {it, ok};
+    map_.replace(it, std::piecewise_construct, std::forward_as_tuple(key),
+                 std::forward_as_tuple(std::forward<Args>(args)...));
+    return {it, ok};
+  }
+
+  // Removes a mapping if it exists, and returns whether it did.
+  //
+  // The last() and end() iterators, as well as those to the erased mapping, are invalidated.
+  //
+  bool erase(const key_type& key) { return erase(key, begin()); }
+
  private:
+  iterator find(const key_type& key, iterator first) {
+    return std::find_if(first, end(), [&key](const auto& pair) { return pair.first == key; });
+  }
+
+  bool erase(const key_type& key, iterator first) {
+    const auto it = find(key, first);
+    if (it == end()) return false;
+    map_.unstable_erase(it);
+    return true;
+  }
+
+  void deduplicate() {
+    for (auto it = begin(); it != end();) {
+      if (const auto key = it->first; ++it != end()) {
+        while (erase(key, it));
+      }
+    }
+  }
+
   Map map_;
 };
 
@@ -186,7 +272,7 @@ bool operator==(const SmallMap<K, V, N>& lhs, const SmallMap<Q, W, M>& rhs) {
 
   for (const auto& [k, v] : lhs) {
     const auto& lv = v;
-    if (!rhs.find(k, [&lv](const auto& rv) { return lv == rv; }).value_or(false)) {
+    if (!rhs.get(k, [&lv](const auto& rv) { return lv == rv; }).value_or(false)) {
       return false;
     }
   }
diff --git a/include/ftl/small_vector.h b/include/ftl/small_vector.h
index cb0ae359eb..0341435813 100644
--- a/include/ftl/small_vector.h
+++ b/include/ftl/small_vector.h
@@ -348,7 +348,7 @@ class SmallVector<T, 0> final : ArrayTraits<T>,
   using Impl::pop_back;
 
   void unstable_erase(iterator it) {
-    if (it != last()) std::iter_swap(it, last());
+    if (it != last()) replace(it, std::move(back()));
     pop_back();
   }
 
diff --git a/libs/ftl/small_map_test.cpp b/libs/ftl/small_map_test.cpp
index 323b9f91e7..2e81022f38 100644
--- a/libs/ftl/small_map_test.cpp
+++ b/libs/ftl/small_map_test.cpp
@@ -18,6 +18,9 @@
 #include <gtest/gtest.h>
 
 #include <cctype>
+#include <string>
+
+using namespace std::string_literals;
 
 namespace android::test {
 
@@ -35,16 +38,19 @@ TEST(SmallMap, Example) {
 
   EXPECT_TRUE(map.contains(123));
 
-  EXPECT_EQ(map.find(42, [](const std::string& s) { return s.size(); }), 3u);
+  EXPECT_EQ(map.get(42, [](const std::string& s) { return s.size(); }), 3u);
 
-  const auto opt = map.find(-1);
+  const auto opt = map.get(-1);
   ASSERT_TRUE(opt);
 
   std::string& ref = *opt;
   EXPECT_TRUE(ref.empty());
   ref = "xyz";
 
-  EXPECT_EQ(map, SmallMap(ftl::init::map(-1, "xyz")(42, "???")(123, "abc")));
+  map.emplace_or_replace(0, "vanilla", 2u, 3u);
+  EXPECT_TRUE(map.dynamic());
+
+  EXPECT_EQ(map, SmallMap(ftl::init::map(-1, "xyz")(0, "nil")(42, "???")(123, "abc")));
 }
 
 TEST(SmallMap, Construct) {
@@ -90,42 +96,253 @@ TEST(SmallMap, Construct) {
   }
 }
 
+TEST(SmallMap, UniqueKeys) {
+  {
+    // Duplicate mappings are discarded.
+    const SmallMap map = ftl::init::map<int, float>(1)(2)(3)(2)(3)(1)(3)(2)(1);
+
+    EXPECT_EQ(map.size(), 3u);
+    EXPECT_EQ(map.max_size(), 9u);
+
+    using Map = decltype(map);
+    EXPECT_EQ(map, Map(ftl::init::map(1, 0.f)(2, 0.f)(3, 0.f)));
+  }
+  {
+    // Duplicate mappings may be reordered.
+    const SmallMap map = ftl::init::map('a', 'A')(
+        'b', 'B')('b')('b')('c', 'C')('a')('d')('c')('e', 'E')('d', 'D')('a')('f', 'F');
+
+    EXPECT_EQ(map.size(), 6u);
+    EXPECT_EQ(map.max_size(), 12u);
+
+    using Map = decltype(map);
+    EXPECT_EQ(map, Map(ftl::init::map('a', 'A')('b', 'B')('c', 'C')('d', 'D')('e', 'E')('f', 'F')));
+  }
+}
+
 TEST(SmallMap, Find) {
   {
     // Constant reference.
-    const ftl::SmallMap map = ftl::init::map('a', 'A')('b', 'B')('c', 'C');
+    const SmallMap map = ftl::init::map('a', 'A')('b', 'B')('c', 'C');
 
-    const auto opt = map.find('b');
+    const auto opt = map.get('b');
     EXPECT_EQ(opt, 'B');
 
     const char d = 'D';
-    const auto ref = map.find('d').value_or(std::cref(d));
+    const auto ref = map.get('d').value_or(std::cref(d));
     EXPECT_EQ(ref.get(), 'D');
   }
   {
     // Mutable reference.
-    ftl::SmallMap map = ftl::init::map('a', 'A')('b', 'B')('c', 'C');
+    SmallMap map = ftl::init::map('a', 'A')('b', 'B')('c', 'C');
 
-    const auto opt = map.find('c');
+    const auto opt = map.get('c');
     EXPECT_EQ(opt, 'C');
 
     char d = 'd';
-    const auto ref = map.find('d').value_or(std::ref(d));
+    const auto ref = map.get('d').value_or(std::ref(d));
     ref.get() = 'D';
     EXPECT_EQ(d, 'D');
   }
   {
     // Constant unary operation.
-    const ftl::SmallMap map = ftl::init::map('a', 'x')('b', 'y')('c', 'z');
-    EXPECT_EQ(map.find('c', [](char c) { return std::toupper(c); }), 'Z');
+    const SmallMap map = ftl::init::map('a', 'x')('b', 'y')('c', 'z');
+    EXPECT_EQ(map.get('c', [](char c) { return std::toupper(c); }), 'Z');
   }
   {
     // Mutable unary operation.
-    ftl::SmallMap map = ftl::init::map('a', 'x')('b', 'y')('c', 'z');
-    EXPECT_TRUE(map.find('c', [](char& c) { c = std::toupper(c); }));
+    SmallMap map = ftl::init::map('a', 'x')('b', 'y')('c', 'z');
+    EXPECT_TRUE(map.get('c', [](char& c) { c = std::toupper(c); }));
 
     EXPECT_EQ(map, SmallMap(ftl::init::map('c', 'Z')('b', 'y')('a', 'x')));
   }
 }
 
+TEST(SmallMap, TryEmplace) {
+  SmallMap<int, std::string, 3> map;
+  using Pair = decltype(map)::value_type;
+
+  {
+    const auto [it, ok] = map.try_emplace(123, "abc");
+    ASSERT_TRUE(ok);
+    EXPECT_EQ(*it, Pair(123, "abc"s));
+  }
+  {
+    const auto [it, ok] = map.try_emplace(42, 3u, '?');
+    ASSERT_TRUE(ok);
+    EXPECT_EQ(*it, Pair(42, "???"s));
+  }
+  {
+    const auto [it, ok] = map.try_emplace(-1);
+    ASSERT_TRUE(ok);
+    EXPECT_EQ(*it, Pair(-1, std::string()));
+    EXPECT_FALSE(map.dynamic());
+  }
+  {
+    // Insertion fails if mapping exists.
+    const auto [it, ok] = map.try_emplace(42, "!!!");
+    EXPECT_FALSE(ok);
+    EXPECT_EQ(*it, Pair(42, "???"));
+    EXPECT_FALSE(map.dynamic());
+  }
+  {
+    // Insertion at capacity promotes the map.
+    const auto [it, ok] = map.try_emplace(999, "xyz");
+    ASSERT_TRUE(ok);
+    EXPECT_EQ(*it, Pair(999, "xyz"));
+    EXPECT_TRUE(map.dynamic());
+  }
+
+  EXPECT_EQ(map, SmallMap(ftl::init::map(-1, ""s)(42, "???"s)(123, "abc"s)(999, "xyz"s)));
+}
+
+namespace {
+
+// The mapped type does not require a copy/move assignment operator.
+struct String {
+  template <typename... Args>
+  String(Args... args) : str(args...) {}
+  const std::string str;
+
+  bool operator==(const String& other) const { return other.str == str; }
+};
+
+}  // namespace
+
+TEST(SmallMap, TryReplace) {
+  SmallMap<int, String, 3> map = ftl::init::map(1, "a")(2, "B");
+  using Pair = decltype(map)::value_type;
+
+  {
+    // Replacing fails unless mapping exists.
+    const auto it = map.try_replace(3, "c");
+    EXPECT_EQ(it, map.end());
+  }
+  {
+    // Replacement arguments can refer to the replaced mapping.
+    const auto ref = map.get(2, [](const auto& s) { return s.str[0]; });
+    ASSERT_TRUE(ref);
+
+    // Construct std::string from one character.
+    const auto it = map.try_replace(2, 1u, static_cast<char>(std::tolower(*ref)));
+    ASSERT_NE(it, map.end());
+    EXPECT_EQ(*it, Pair(2, "b"));
+  }
+
+  EXPECT_FALSE(map.dynamic());
+  EXPECT_TRUE(map.try_emplace(3, "abc").second);
+  EXPECT_TRUE(map.try_emplace(4, "d").second);
+  EXPECT_TRUE(map.dynamic());
+
+  {
+    // Replacing fails unless mapping exists.
+    const auto it = map.try_replace(5, "e");
+    EXPECT_EQ(it, map.end());
+  }
+  {
+    // Replacement arguments can refer to the replaced mapping.
+    const auto ref = map.get(3);
+    ASSERT_TRUE(ref);
+
+    // Construct std::string from substring.
+    const auto it = map.try_replace(3, ref->get().str, 2u, 1u);
+    ASSERT_NE(it, map.end());
+    EXPECT_EQ(*it, Pair(3, "c"));
+  }
+
+  EXPECT_EQ(map, SmallMap(ftl::init::map(4, "d"s)(3, "c"s)(2, "b"s)(1, "a"s)));
+}
+
+TEST(SmallMap, EmplaceOrReplace) {
+  SmallMap<int, String, 3> map = ftl::init::map(1, "a")(2, "B");
+  using Pair = decltype(map)::value_type;
+
+  {
+    // New mapping is emplaced.
+    const auto [it, emplace] = map.emplace_or_replace(3, "c");
+    EXPECT_TRUE(emplace);
+    EXPECT_EQ(*it, Pair(3, "c"));
+  }
+  {
+    // Replacement arguments can refer to the replaced mapping.
+    const auto ref = map.get(2, [](const auto& s) { return s.str[0]; });
+    ASSERT_TRUE(ref);
+
+    // Construct std::string from one character.
+    const auto [it, emplace] = map.emplace_or_replace(2, 1u, static_cast<char>(std::tolower(*ref)));
+    EXPECT_FALSE(emplace);
+    EXPECT_EQ(*it, Pair(2, "b"));
+  }
+
+  EXPECT_FALSE(map.dynamic());
+  EXPECT_FALSE(map.emplace_or_replace(3, "abc").second);  // Replace.
+  EXPECT_TRUE(map.emplace_or_replace(4, "d").second);     // Emplace.
+  EXPECT_TRUE(map.dynamic());
+
+  {
+    // New mapping is emplaced.
+    const auto [it, emplace] = map.emplace_or_replace(5, "e");
+    EXPECT_TRUE(emplace);
+    EXPECT_EQ(*it, Pair(5, "e"));
+  }
+  {
+    // Replacement arguments can refer to the replaced mapping.
+    const auto ref = map.get(3);
+    ASSERT_TRUE(ref);
+
+    // Construct std::string from substring.
+    const auto [it, emplace] = map.emplace_or_replace(3, ref->get().str, 2u, 1u);
+    EXPECT_FALSE(emplace);
+    EXPECT_EQ(*it, Pair(3, "c"));
+  }
+
+  EXPECT_EQ(map, SmallMap(ftl::init::map(5, "e"s)(4, "d"s)(3, "c"s)(2, "b"s)(1, "a"s)));
+}
+
+TEST(SmallMap, Erase) {
+  {
+    SmallMap map = ftl::init::map(1, '1')(2, '2')(3, '3')(4, '4');
+    EXPECT_FALSE(map.dynamic());
+
+    EXPECT_FALSE(map.erase(0));  // Key not found.
+
+    EXPECT_TRUE(map.erase(2));
+    EXPECT_EQ(map, SmallMap(ftl::init::map(1, '1')(3, '3')(4, '4')));
+
+    EXPECT_TRUE(map.erase(1));
+    EXPECT_EQ(map, SmallMap(ftl::init::map(3, '3')(4, '4')));
+
+    EXPECT_TRUE(map.erase(4));
+    EXPECT_EQ(map, SmallMap(ftl::init::map(3, '3')));
+
+    EXPECT_TRUE(map.erase(3));
+    EXPECT_FALSE(map.erase(3));  // Key not found.
+
+    EXPECT_TRUE(map.empty());
+    EXPECT_FALSE(map.dynamic());
+  }
+  {
+    SmallMap map = ftl::init::map(1, '1')(2, '2')(3, '3');
+    map.try_emplace(4, '4');
+    EXPECT_TRUE(map.dynamic());
+
+    EXPECT_FALSE(map.erase(0));  // Key not found.
+
+    EXPECT_TRUE(map.erase(2));
+    EXPECT_EQ(map, SmallMap(ftl::init::map(1, '1')(3, '3')(4, '4')));
+
+    EXPECT_TRUE(map.erase(1));
+    EXPECT_EQ(map, SmallMap(ftl::init::map(3, '3')(4, '4')));
+
+    EXPECT_TRUE(map.erase(4));
+    EXPECT_EQ(map, SmallMap(ftl::init::map(3, '3')));
+
+    EXPECT_TRUE(map.erase(3));
+    EXPECT_FALSE(map.erase(3));  // Key not found.
+
+    EXPECT_TRUE(map.empty());
+    EXPECT_TRUE(map.dynamic());
+  }
+}
+
 }  // namespace android::test