FTL: Generalize SmallMap lookup transformer

Bug: 185536303
Test: ftl_test
Change-Id: Idde9d842a4404095e839fbdbfca1ded416d95263
diff --git a/include/ftl/small_map.h b/include/ftl/small_map.h
index 5217e76..49cde7f 100644
--- a/include/ftl/small_map.h
+++ b/include/ftl/small_map.h
@@ -17,11 +17,11 @@
 #pragma once
 
 #include <ftl/initializer_list.h>
+#include <ftl/optional.h>
 #include <ftl/small_vector.h>
 
 #include <algorithm>
 #include <functional>
-#include <optional>
 #include <type_traits>
 #include <utility>
 
@@ -47,7 +47,7 @@
 //   assert(!map.dynamic());
 //
 //   assert(map.contains(123));
-//   assert(map.get(42, [](const std::string& s) { return s.size(); }) == 3u);
+//   assert(map.get(42).transform([](const std::string& s) { return s.size(); }) == 3u);
 //
 //   const auto opt = map.get(-1);
 //   assert(opt);
@@ -59,7 +59,7 @@
 //   map.emplace_or_replace(0, "vanilla", 2u, 3u);
 //   assert(map.dynamic());
 //
-//   assert(map == SmallMap(ftl::init::map(-1, "xyz")(0, "nil")(42, "???")(123, "abc")));
+//   assert(map == SmallMap(ftl::init::map(-1, "xyz"sv)(0, "nil"sv)(42, "???"sv)(123, "abc"sv)));
 //
 template <typename K, typename V, std::size_t N, typename KeyEqual = std::equal_to<K>>
 class SmallMap final {
@@ -123,9 +123,7 @@
   const_iterator cend() const { return map_.cend(); }
 
   // Returns whether a mapping exists for the given key.
-  bool contains(const key_type& key) const {
-    return get(key, [](const mapped_type&) {});
-  }
+  bool contains(const key_type& key) const { return get(key).has_value(); }
 
   // Returns a reference to the value for the given key, or std::nullopt if the key was not found.
   //
@@ -139,44 +137,22 @@
   //   ref.get() = 'D';
   //   assert(d == 'D');
   //
-  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 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
-  // for the given key, or std::nullopt if the key was not found. If F has a return type of void,
-  // then the Boolean result indicates whether the key was found.
-  //
-  //   ftl::SmallMap map = ftl::init::map('a', 'x')('b', 'y')('c', 'z');
-  //
-  //   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 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) {
+  auto get(const key_type& key) const -> Optional<std::reference_wrapper<const mapped_type>> {
+    for (const auto& [k, v] : *this) {
       if (KeyEqual{}(k, key)) {
-        if constexpr (std::is_void_v<R>) {
-          f(v);
-          return true;
-        } else {
-          return f(v);
-        }
+        return std::cref(v);
       }
     }
-
     return {};
   }
 
-  template <typename F>
-  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)); });
+  auto get(const key_type& key) -> Optional<std::reference_wrapper<mapped_type>> {
+    for (auto& [k, v] : *this) {
+      if (KeyEqual{}(k, key)) {
+        return std::ref(v);
+      }
+    }
+    return {};
   }
 
   // Returns an iterator to an existing mapping for the given key, or the end() iterator otherwise.
@@ -286,7 +262,7 @@
 
   for (const auto& [k, v] : lhs) {
     const auto& lv = v;
-    if (!rhs.get(k, [&lv](const auto& rv) { return lv == rv; }).value_or(false)) {
+    if (!rhs.get(k).transform([&lv](const W& rv) { return lv == rv; }).value_or(false)) {
       return false;
     }
   }
diff --git a/include/ftl/unit.h b/include/ftl/unit.h
new file mode 100644
index 0000000..e38230b
--- /dev/null
+++ b/include/ftl/unit.h
@@ -0,0 +1,61 @@
+/*
+ * Copyright 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <type_traits>
+#include <utility>
+
+namespace android::ftl {
+
+// The unit type, and its only value.
+constexpr struct Unit {
+} unit;
+
+constexpr bool operator==(Unit, Unit) {
+  return true;
+}
+
+constexpr bool operator!=(Unit, Unit) {
+  return false;
+}
+
+// Adapts a function object F to return Unit. The return value of F is ignored.
+//
+// As a practical use, the function passed to ftl::Optional<T>::transform is not allowed to return
+// void (cf. https://wg21.link/P0798R8#mapping-functions-returning-void), but may return Unit if
+// only its side effects are meaningful:
+//
+//   ftl::Optional opt = "food"s;
+//   opt.transform(ftl::unit_fn([](std::string& str) { str.pop_back(); }));
+//   assert(opt == "foo"s);
+//
+template <typename F>
+struct UnitFn {
+  F f;
+
+  template <typename... Args>
+  Unit operator()(Args&&... args) {
+    return f(std::forward<Args>(args)...), unit;
+  }
+};
+
+template <typename F>
+constexpr auto unit_fn(F&& f) -> UnitFn<std::decay_t<F>> {
+  return {std::forward<F>(f)};
+}
+
+}  // namespace android::ftl
diff --git a/libs/ftl/optional_test.cpp b/libs/ftl/optional_test.cpp
index 6a8c8f9..ede159a 100644
--- a/libs/ftl/optional_test.cpp
+++ b/libs/ftl/optional_test.cpp
@@ -17,6 +17,7 @@
 #include <ftl/optional.h>
 #include <ftl/static_vector.h>
 #include <ftl/string.h>
+#include <ftl/unit.h>
 #include <gtest/gtest.h>
 
 #include <functional>
@@ -62,6 +63,13 @@
     EXPECT_EQ(out, "abc"s);
   }
 
+  // No return value.
+  {
+    Optional opt = "food"s;
+    EXPECT_EQ(ftl::unit, opt.transform(ftl::unit_fn([](std::string& str) { str.pop_back(); })));
+    EXPECT_EQ(opt, "foo"s);
+  }
+
   // Chaining.
   EXPECT_EQ(14u, Optional(StaticVector{"upside"s, "down"s})
                      .transform([](StaticVector<std::string, 3>&& v) {
diff --git a/libs/ftl/small_map_test.cpp b/libs/ftl/small_map_test.cpp
index 1740a2b..634877f 100644
--- a/libs/ftl/small_map_test.cpp
+++ b/libs/ftl/small_map_test.cpp
@@ -15,12 +15,15 @@
  */
 
 #include <ftl/small_map.h>
+#include <ftl/unit.h>
 #include <gtest/gtest.h>
 
 #include <cctype>
 #include <string>
+#include <string_view>
 
 using namespace std::string_literals;
+using namespace std::string_view_literals;
 
 namespace android::test {
 
@@ -38,7 +41,7 @@
 
   EXPECT_TRUE(map.contains(123));
 
-  EXPECT_EQ(map.get(42, [](const std::string& s) { return s.size(); }), 3u);
+  EXPECT_EQ(map.get(42).transform([](const std::string& s) { return s.size(); }), 3u);
 
   const auto opt = map.get(-1);
   ASSERT_TRUE(opt);
@@ -50,7 +53,7 @@
   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")));
+  EXPECT_EQ(map, SmallMap(ftl::init::map(-1, "xyz"sv)(0, "nil"sv)(42, "???"sv)(123, "abc"sv)));
 }
 
 TEST(SmallMap, Construct) {
@@ -70,7 +73,7 @@
     EXPECT_EQ(map.max_size(), 5u);
     EXPECT_FALSE(map.dynamic());
 
-    EXPECT_EQ(map, SmallMap(ftl::init::map(123, "abc")(456, "def")(789, "ghi")));
+    EXPECT_EQ(map, SmallMap(ftl::init::map(123, "abc"sv)(456, "def"sv)(789, "ghi"sv)));
   }
   {
     // In-place constructor with different types.
@@ -81,7 +84,7 @@
     EXPECT_EQ(map.max_size(), 5u);
     EXPECT_FALSE(map.dynamic());
 
-    EXPECT_EQ(map, SmallMap(ftl::init::map(42, "???")(123, "abc")(-1, "\0\0\0")));
+    EXPECT_EQ(map, SmallMap(ftl::init::map(42, "???"sv)(123, "abc"sv)(-1, ""sv)));
   }
   {
     // In-place constructor with implicit size.
@@ -92,7 +95,7 @@
     EXPECT_EQ(map.max_size(), 3u);
     EXPECT_FALSE(map.dynamic());
 
-    EXPECT_EQ(map, SmallMap(ftl::init::map(-1, "\0\0\0")(42, "???")(123, "abc")));
+    EXPECT_EQ(map, SmallMap(ftl::init::map(-1, ""sv)(42, "???"sv)(123, "abc"sv)));
   }
 }
 
@@ -108,7 +111,7 @@
   {
     // Convertible types; same capacity.
     SmallMap map1 = ftl::init::map<char, std::string>('M', "mega")('G', "giga");
-    const SmallMap map2 = ftl::init::map('T', "tera")('P', "peta");
+    const SmallMap map2 = ftl::init::map('T', "tera"sv)('P', "peta"sv);
 
     map1 = map2;
     EXPECT_EQ(map1, map2);
@@ -147,7 +150,7 @@
   }
 }
 
-TEST(SmallMap, Find) {
+TEST(SmallMap, Get) {
   {
     // Constant reference.
     const SmallMap map = ftl::init::map('a', 'A')('b', 'B')('c', 'C');
@@ -172,14 +175,15 @@
     EXPECT_EQ(d, 'D');
   }
   {
-    // Constant unary operation.
+    // Immutable transform operation.
     const SmallMap map = ftl::init::map('a', 'x')('b', 'y')('c', 'z');
-    EXPECT_EQ(map.get('c', [](char c) { return std::toupper(c); }), 'Z');
+    EXPECT_EQ(map.get('c').transform([](char c) { return std::toupper(c); }), 'Z');
   }
   {
-    // Mutable unary operation.
+    // Mutable transform operation.
     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.get('c').transform(ftl::unit_fn([](char& c) { c = std::toupper(c); })),
+              ftl::unit);
 
     EXPECT_EQ(map, SmallMap(ftl::init::map('c', 'Z')('b', 'y')('a', 'x')));
   }
@@ -247,7 +251,7 @@
   }
   {
     // Replacement arguments can refer to the replaced mapping.
-    const auto ref = map.get(2, [](const auto& s) { return s.str[0]; });
+    const auto ref = map.get(2).transform([](const String& s) { return s.str[0]; });
     ASSERT_TRUE(ref);
 
     // Construct std::string from one character.
@@ -292,7 +296,7 @@
   }
   {
     // Replacement arguments can refer to the replaced mapping.
-    const auto ref = map.get(2, [](const auto& s) { return s.str[0]; });
+    const auto ref = map.get(2).transform([](const String& s) { return s.str[0]; });
     ASSERT_TRUE(ref);
 
     // Construct std::string from one character.
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index 96598cd..2d1220c 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -3247,10 +3247,11 @@
     mDrawingState.traverseInReverseZOrder([&](Layer* layer) {
         if (!layer->needsInputInfo()) return;
 
-        const auto opt = displayInputInfos.get(layer->getLayerStack(),
-                                               [](const auto& info) -> Layer::InputDisplayArgs {
-                                                   return {&info.transform, info.isSecure};
-                                               });
+        const auto opt = displayInputInfos.get(layer->getLayerStack())
+                                 .transform([](const DisplayDevice::InputInfo& info) {
+                                     return Layer::InputDisplayArgs{&info.transform, info.isSecure};
+                                 });
+
         outWindowInfos.push_back(layer->fillInputInfo(opt.value_or(Layer::InputDisplayArgs{})));
     });
 
diff --git a/services/surfaceflinger/fuzzer/surfaceflinger_scheduler_fuzzer.cpp b/services/surfaceflinger/fuzzer/surfaceflinger_scheduler_fuzzer.cpp
index f507ef0..9584492 100644
--- a/services/surfaceflinger/fuzzer/surfaceflinger_scheduler_fuzzer.cpp
+++ b/services/surfaceflinger/fuzzer/surfaceflinger_scheduler_fuzzer.cpp
@@ -379,7 +379,8 @@
     RefreshRateStats refreshRateStats(timeStats, Fps::fromValue(mFdp.ConsumeFloatingPoint<float>()),
                                       PowerMode::OFF);
 
-    const auto fpsOpt = displayModes.get(modeId, [](const auto& mode) { return mode->getFps(); });
+    const auto fpsOpt = displayModes.get(modeId).transform(
+            [](const DisplayModePtr& mode) { return mode->getFps(); });
     refreshRateStats.setRefreshRate(*fpsOpt);
 
     refreshRateStats.setPowerMode(mFdp.PickValueInArray(kPowerModes));