Introduce WTF::flatMap

https://bugs.webkit.org/show_bug.cgi?id=261780 rdar://problem/115751217 Reviewed by Richard Robinson. Introduced WTF::flatMap. Similar to map and compactMap, flatMap takes a collection that may contain nested collections and allows the caller to specify a map function that produces a flat Vector. Added test cases to TestWTF. A follow-on change will adopt this in CDMInstanceFairPlayStreamingAVFObjC. * Source/WTF/wtf/Vector.h: (WTF::FlatMapper::flatMap): (WTF::flatMap): * Tools/TestWebKitAPI/Tests/WTF/Vector.cpp: (TestWebKitAPI::mapVector): (TestWebKitAPI::TEST): (TestWebKitAPI::mapInnerStruct): Canonical link: https://commits.webkit.org/268171@main
WebKit · Sep 20, 2023 · f3d2011 · f3d2011
1 parent 95289b1
commit f3d2011
Show file tree

Hide file tree

Showing 2 changed files with 133 additions and 0 deletions.
diff --git a/Source/WTF/wtf/Vector.h b/Source/WTF/wtf/Vector.h
@@ -1898,6 +1898,42 @@ Vector<typename CompactMapper<MapFunction, SourceType>::DestinationItemType> com
     return CompactMapper<MapFunction, SourceType>::compactMap(std::forward<SourceType>(source), std::forward<MapFunction>(mapFunction));
 }
 
+template<typename MapFunction, typename SourceType>
+struct FlatMapper {
+    using SourceItemType = typename CollectionInspector<SourceType>::SourceItemType;
+    using DestinationItemType = typename CollectionInspector<typename std::invoke_result<MapFunction, SourceItemType&>::type>::SourceItemType;
+
+    static Vector<DestinationItemType> flatMap(const SourceType& source, const MapFunction& mapFunction)
+    {
+        Vector<DestinationItemType> result;
+        for (auto&& item : source)
+            result.appendVector(mapFunction(item));
+        result.shrinkToFit();
+        return result;
+    }
+};
+
+template<typename MapFunction, typename SourceType> requires std::is_rvalue_reference<SourceType&&>::value
+struct FlatMapper<MapFunction, SourceType> {
+    using SourceItemType = typename CollectionInspector<SourceType>::SourceItemType;
+    using DestinationItemType = typename CollectionInspector<typename std::invoke_result<MapFunction, SourceItemType&&>::type>::SourceItemType;
+
+    static Vector<DestinationItemType> flatMap(SourceType&& source, const MapFunction& mapFunction)
+    {
+        Vector<DestinationItemType> result;
+        for (auto&& item : source)
+            result.appendVector(mapFunction(WTFMove(item)));
+        result.shrinkToFit();
+        return result;
+    }
+};
+
+template<typename MapFunction, typename SourceType>
+Vector<typename FlatMapper<MapFunction, SourceType>::DestinationItemType> flatMap(SourceType&& source, MapFunction&& mapFunction)
+{
+    return FlatMapper<MapFunction, SourceType>::flatMap(std::forward<SourceType>(source), std::forward<MapFunction>(mapFunction));
+}
+
 template<typename DestinationVector, typename Collection>
 inline auto copyToVectorSpecialization(const Collection& collection) -> DestinationVector
 {

diff --git a/Tools/TestWebKitAPI/Tests/WTF/Vector.cpp b/Tools/TestWebKitAPI/Tests/WTF/Vector.cpp
@@ -1910,4 +1910,101 @@ TEST(WTF_Vector, MoveAssignmentOperator)
     }
 }
 
+static Vector<int> mapVector(Vector<int> vector)
+{
+    return vector;
+}
+
+TEST(WTF_Vector, FlatMapCopy)
+{
+    Vector<Vector<int>> vector {
+        { 1, 2 },
+        { 3, 4 },
+    };
+
+    static_assert(std::is_same<decltype(WTF::flatMap(vector, mapVector)), typename WTF::Vector<int>>::value,
+        "WTF::flatMap returns Vector<int>");
+    auto mapped = WTF::flatMap(vector, mapVector);
+
+    EXPECT_EQ(2U, vector.size());
+    EXPECT_EQ(2U, vector[0].size());
+    EXPECT_EQ(2U, vector[1].size());
+
+    EXPECT_EQ(4U, mapped.size());
+    EXPECT_EQ(1, mapped[0]);
+    EXPECT_EQ(2, mapped[1]);
+    EXPECT_EQ(3, mapped[2]);
+    EXPECT_EQ(4, mapped[3]);
+}
+
+TEST(WTF_Vector, FlatMapMove)
+{
+    Vector<Vector<int>> vector {
+        { 1, 2 },
+        { 3, 4 },
+    };
+
+    static_assert(std::is_same<decltype(WTF::flatMap(WTFMove(vector), mapVector)), typename WTF::Vector<int>>::value,
+        "WTF::flatMap returns Vector<int>");
+    auto mapped = WTF::flatMap(WTFMove(vector), mapVector);
+
+    EXPECT_EQ(2U, vector.size());
+    EXPECT_EQ(0U, vector[0].size());
+    EXPECT_EQ(0U, vector[1].size());
+
+    EXPECT_EQ(4U, mapped.size());
+    EXPECT_EQ(1, mapped[0]);
+    EXPECT_EQ(2, mapped[1]);
+    EXPECT_EQ(3, mapped[2]);
+    EXPECT_EQ(4, mapped[3]);
+}
+
+TEST(WTF_Vector, FlatMapEmptyInnerVector)
+{
+    Vector<Vector<int>> vector {
+        { },
+        { 1, 2 },
+        { },
+        { 3, 4 },
+        { },
+    };
+
+    static_assert(std::is_same<decltype(WTF::flatMap(vector, mapVector)), typename WTF::Vector<int>>::value,
+        "WTF::flatMap returns Vector<int>");
+    auto mapped = WTF::flatMap(vector, mapVector);
+
+    EXPECT_EQ(4U, mapped.size());
+    EXPECT_EQ(1, mapped[0]);
+    EXPECT_EQ(2, mapped[1]);
+    EXPECT_EQ(3, mapped[2]);
+    EXPECT_EQ(4, mapped[3]);
+}
+
+struct InnerStruct {
+    Vector<int> vector;
+};
+
+static Vector<int> mapInnerStruct(InnerStruct innerStruct)
+{
+    return innerStruct.vector;
+}
+
+TEST(WTF_Vector, FlatMapInnerStruct)
+{
+    Vector<InnerStruct> vector {
+        { { 1, 2 } },
+        { { 3, 4 } },
+    };
+
+    static_assert(std::is_same<decltype(WTF::flatMap(vector, mapInnerStruct)), typename WTF::Vector<int>>::value,
+        "WTF::flatMap returns Vector<int>");
+    auto mapped = WTF::flatMap(vector, mapInnerStruct);
+
+    EXPECT_EQ(4U, mapped.size());
+    EXPECT_EQ(1, mapped[0]);
+    EXPECT_EQ(2, mapped[1]);
+    EXPECT_EQ(3, mapped[2]);
+    EXPECT_EQ(4, mapped[3]);
+}
+
 } // namespace TestWebKitAPI