tint: Add utils::TransformN()

A fork of utils::Transform() that transforms at most 'n' elements of the
input vector.

Change-Id: I0ec546f99ee74817ebfd20abf02db5faf5f8ec5c
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/91302
Reviewed-by: Antonio Maiorano <amaiorano@google.com>
Commit-Queue: Ben Clayton <bclayton@chromium.org>
Kokoro: Kokoro <noreply+kokoro@google.com>

src/tint/utils/transform.h

@@ -27,8 +27,7 @@ namespace tint::utils {
 /// Transform performs an element-wise transformation of a vector.
 /// @param in the input vector.
 /// @param transform the transformation function with signature: `OUT(IN)`
-/// @returns a new vector with each element of the source vector transformed by
+/// @returns a new vector with each element of the source vector transformed by `transform`.
-/// `transform`.
 template <typename IN, typename TRANSFORMER>
 auto Transform(const std::vector<IN>& in, TRANSFORMER&& transform)
     -> std::vector<decltype(transform(in[0]))> {
@@ -41,10 +40,8 @@ auto Transform(const std::vector<IN>& in, TRANSFORMER&& transform)
 
 /// Transform performs an element-wise transformation of a vector.
 /// @param in the input vector.
-/// @param transform the transformation function with signature:
+/// @param transform the transformation function with signature: `OUT(IN, size_t)`
-/// `OUT(IN, size_t)`
+/// @returns a new vector with each element of the source vector transformed by `transform`.
-/// @returns a new vector with each element of the source vector transformed by
-/// `transform`.
 template <typename IN, typename TRANSFORMER>
 auto Transform(const std::vector<IN>& in, TRANSFORMER&& transform)
     -> std::vector<decltype(transform(in[0], 1u))> {
@@ -55,6 +52,40 @@ auto Transform(const std::vector<IN>& in, TRANSFORMER&& transform)
     return result;
 }
 
+/// TransformN performs an element-wise transformation of a vector, transforming and returning at
+/// most `n` elements.
+/// @param in the input vector.
+/// @param n the maximum number of elements to transform.
+/// @param transform the transformation function with signature: `OUT(IN)`
+/// @returns a new vector with at most n-elements of the source vector transformed by `transform`.
+template <typename IN, typename TRANSFORMER>
+auto TransformN(const std::vector<IN>& in, size_t n, TRANSFORMER&& transform)
+    -> std::vector<decltype(transform(in[0]))> {
+    const auto count = std::min(n, in.size());
+    std::vector<decltype(transform(in[0]))> result(count);
+    for (size_t i = 0; i < count; ++i) {
+        result[i] = transform(in[i]);
+    }
+    return result;
+}
+
+/// TransformN performs an element-wise transformation of a vector, transforming and returning at
+/// most `n` elements.
+/// @param in the input vector.
+/// @param n the maximum number of elements to transform.
+/// @param transform the transformation function with signature: `OUT(IN, size_t)`
+/// @returns a new vector with at most n-elements of the source vector transformed by `transform`.
+template <typename IN, typename TRANSFORMER>
+auto TransformN(const std::vector<IN>& in, size_t n, TRANSFORMER&& transform)
+    -> std::vector<decltype(transform(in[0], 1u))> {
+    const auto count = std::min(n, in.size());
+    std::vector<decltype(transform(in[0], 1u))> result(count);
+    for (size_t i = 0; i < count; ++i) {
+        result[i] = transform(in[i], i);
+    }
+    return result;
+}
+
 }  // namespace tint::utils
 
 #endif  // SRC_TINT_UTILS_TRANSFORM_H_

src/tint/utils/transform_test.cc

@@ -30,7 +30,7 @@ TEST(TransformTest, Empty) {
     const std::vector<int> empty{};
     {
         auto transformed = Transform(empty, [](int) -> int {
-            [] { FAIL() << "Transform should not be called for empty vector"; }();
+            [] { FAIL() << "Callback should not be called for empty vector"; }();
             return 0;
         });
         CHECK_ELEMENT_TYPE(transformed, int);
@@ -38,7 +38,7 @@ TEST(TransformTest, Empty) {
     }
     {
         auto transformed = Transform(empty, [](int, size_t) -> int {
-            [] { FAIL() << "Transform should not be called for empty vector"; }();
+            [] { FAIL() << "Callback should not be called for empty vector"; }();
             return 0;
         });
         CHECK_ELEMENT_TYPE(transformed, int);
@@ -48,16 +48,16 @@ TEST(TransformTest, Empty) {
 
 TEST(TransformTest, Identity) {
     const std::vector<int> input{1, 2, 3, 4};
-    {
     auto transformed = Transform(input, [](int i) { return i; });
     CHECK_ELEMENT_TYPE(transformed, int);
     EXPECT_THAT(transformed, testing::ElementsAre(1, 2, 3, 4));
-    }
+}
-    {
+
+TEST(TransformTest, IdentityWithIndex) {
+    const std::vector<int> input{1, 2, 3, 4};
     auto transformed = Transform(input, [](int i, size_t) { return i; });
     CHECK_ELEMENT_TYPE(transformed, int);
     EXPECT_THAT(transformed, testing::ElementsAre(1, 2, 3, 4));
-    }
 }
 
 TEST(TransformTest, Index) {
@@ -87,5 +87,135 @@ TEST(TransformTest, TransformDifferentType) {
     }
 }
 
+TEST(TransformNTest, Empty) {
+    const std::vector<int> empty{};
+    {
+        auto transformed = TransformN(empty, 4u, [](int) -> int {
+            [] { FAIL() << "Callback should not be called for empty vector"; }();
+            return 0;
+        });
+        CHECK_ELEMENT_TYPE(transformed, int);
+        EXPECT_EQ(transformed.size(), 0u);
+    }
+    {
+        auto transformed = TransformN(empty, 4u, [](int, size_t) -> int {
+            [] { FAIL() << "Callback should not be called for empty vector"; }();
+            return 0;
+        });
+        CHECK_ELEMENT_TYPE(transformed, int);
+        EXPECT_EQ(transformed.size(), 0u);
+    }
+}
+
+TEST(TransformNTest, Identity) {
+    const std::vector<int> input{1, 2, 3, 4};
+    {
+        auto transformed = TransformN(input, 0u, [](int) {
+            [] { FAIL() << "Callback should not call the transform when n == 0"; }();
+            return 0;
+        });
+        CHECK_ELEMENT_TYPE(transformed, int);
+        EXPECT_TRUE(transformed.empty());
+    }
+    {
+        auto transformed = TransformN(input, 2u, [](int i) { return i; });
+        CHECK_ELEMENT_TYPE(transformed, int);
+        EXPECT_THAT(transformed, testing::ElementsAre(1, 2));
+    }
+    {
+        auto transformed = TransformN(input, 6u, [](int i) { return i; });
+        CHECK_ELEMENT_TYPE(transformed, int);
+        EXPECT_THAT(transformed, testing::ElementsAre(1, 2, 3, 4));
+    }
+}
+
+TEST(TransformNTest, IdentityWithIndex) {
+    const std::vector<int> input{1, 2, 3, 4};
+    {
+        auto transformed = TransformN(input, 0u, [](int, size_t) {
+            [] { FAIL() << "Callback should not call the transform when n == 0"; }();
+            return 0;
+        });
+        CHECK_ELEMENT_TYPE(transformed, int);
+        EXPECT_TRUE(transformed.empty());
+    }
+    {
+        auto transformed = TransformN(input, 3u, [](int i, size_t) { return i; });
+        CHECK_ELEMENT_TYPE(transformed, int);
+        EXPECT_THAT(transformed, testing::ElementsAre(1, 2, 3));
+    }
+    {
+        auto transformed = TransformN(input, 9u, [](int i, size_t) { return i; });
+        CHECK_ELEMENT_TYPE(transformed, int);
+        EXPECT_THAT(transformed, testing::ElementsAre(1, 2, 3, 4));
+    }
+}
+
+TEST(TransformNTest, Index) {
+    const std::vector<int> input{10, 20, 30, 40};
+    {
+        auto transformed = TransformN(input, 0u, [](int, size_t) {
+            [] { FAIL() << "Callback should not call the transform when n == 0"; }();
+            return static_cast<size_t>(0);
+        });
+        CHECK_ELEMENT_TYPE(transformed, size_t);
+        EXPECT_TRUE(transformed.empty());
+    }
+    {
+        auto transformed = TransformN(input, 2u, [](int, size_t idx) { return idx; });
+        CHECK_ELEMENT_TYPE(transformed, size_t);
+        EXPECT_THAT(transformed, testing::ElementsAre(0u, 1u));
+    }
+    {
+        auto transformed = TransformN(input, 9u, [](int, size_t idx) { return idx; });
+        CHECK_ELEMENT_TYPE(transformed, size_t);
+        EXPECT_THAT(transformed, testing::ElementsAre(0u, 1u, 2u, 3u));
+    }
+}
+
+TEST(TransformNTest, TransformSameType) {
+    const std::vector<int> input{1, 2, 3, 4};
+    {
+        auto transformed = TransformN(input, 0u, [](int, size_t) {
+            [] { FAIL() << "Callback should not call the transform when n == 0"; }();
+            return 0;
+        });
+        CHECK_ELEMENT_TYPE(transformed, int);
+        EXPECT_TRUE(transformed.empty());
+    }
+    {
+        auto transformed = TransformN(input, 2u, [](int i) { return i * 10; });
+        CHECK_ELEMENT_TYPE(transformed, int);
+        EXPECT_THAT(transformed, testing::ElementsAre(10, 20));
+    }
+    {
+        auto transformed = TransformN(input, 9u, [](int i) { return i * 10; });
+        CHECK_ELEMENT_TYPE(transformed, int);
+        EXPECT_THAT(transformed, testing::ElementsAre(10, 20, 30, 40));
+    }
+}
+
+TEST(TransformNTest, TransformDifferentType) {
+    const std::vector<int> input{1, 2, 3, 4};
+    {
+        auto transformed = TransformN(input, 0u, [](int) {
+            [] { FAIL() << "Callback should not call the transform when n == 0"; }();
+            return std::string();
+        });
+        CHECK_ELEMENT_TYPE(transformed, std::string);
+        EXPECT_TRUE(transformed.empty());
+    }
+    {
+        auto transformed = TransformN(input, 2u, [](int i) { return std::to_string(i); });
+        CHECK_ELEMENT_TYPE(transformed, std::string);
+        EXPECT_THAT(transformed, testing::ElementsAre("1", "2"));
+    }
+    {
+        auto transformed = TransformN(input, 9u, [](int i) { return std::to_string(i); });
+        CHECK_ELEMENT_TYPE(transformed, std::string);
+        EXPECT_THAT(transformed, testing::ElementsAre("1", "2", "3", "4"));
+    }
+}
+
 }  // namespace
 }  // namespace tint::utils