From 6746a712d5157a4d69a318348b0c8f388273b02d Mon Sep 17 00:00:00 2001
From: Kim Hoang <kimhoang1229@gmail.com>
Date: Wed, 8 Jul 2026 12:23:26 -0700
Subject: [PATCH] test cases for concat vectors, extend vectors, generate
 vectors

---
 .../src/utils/containers/concat_vectors.cc    | 38 +++++++++++++++++++
 .../src/utils/containers/extend_vector.cc     | 30 +++++++++++++++
 .../src/utils/containers/generate_vector.cc   | 31 +++++++++++++++
 3 files changed, 99 insertions(+)
 create mode 100644 lib/utils/test/src/utils/containers/concat_vectors.cc
 create mode 100644 lib/utils/test/src/utils/containers/extend_vector.cc
 create mode 100644 lib/utils/test/src/utils/containers/generate_vector.cc

diff --git a/lib/utils/test/src/utils/containers/concat_vectors.cc b/lib/utils/test/src/utils/containers/concat_vectors.cc
new file mode 100644
index 0000000000..e22b8e8c6c
--- /dev/null
+++ b/lib/utils/test/src/utils/containers/concat_vectors.cc
@@ -0,0 +1,38 @@
+#include "utils/containers/concat_vectors.h"
+#include "test/utils/doctest/fmt/vector.h"
+#include <doctest/doctest.h>
+#include <vector>
+
+using namespace ::FlexFlow;
+
+TEST_SUITE(FF_TEST_SUITE) {
+  TEST_CASE("concat_vectors") {
+    SUBCASE("two vectors") {
+      std::vector<int> prefix = {1, 2};
+      std::vector<int> postfix = {3, 4};
+
+      std::vector<int> result = concat_vectors(prefix, postfix);
+      std::vector<int> correct = {1, 2, 3, 4};
+
+      CHECK(result == correct);
+    }
+
+    SUBCASE("vector of vectors") {
+      std::vector<std::vector<int>> vecs = {{1, 2}, {3}, {4, 5}};
+
+      std::vector<int> result = concat_vectors(vecs);
+      std::vector<int> correct = {1, 2, 3, 4, 5};
+
+      CHECK(result == correct);
+    }
+
+    SUBCASE("empty vector of vectors") {
+      std::vector<std::vector<int>> vecs = {};
+
+      std::vector<int> result = concat_vectors(vecs);
+      std::vector<int> correct = {};
+
+      CHECK(result == correct);
+    }
+  }
+}
diff --git a/lib/utils/test/src/utils/containers/extend_vector.cc b/lib/utils/test/src/utils/containers/extend_vector.cc
new file mode 100644
index 0000000000..a394cdf22f
--- /dev/null
+++ b/lib/utils/test/src/utils/containers/extend_vector.cc
@@ -0,0 +1,30 @@
+#include "utils/containers/extend_vector.h"
+#include "test/utils/doctest/fmt/vector.h"
+#include <doctest/doctest.h>
+#include <vector>
+
+using namespace ::FlexFlow;
+// checks rhs gets added to the end of lhs
+TEST_SUITE(FF_TEST_SUITE) {
+  TEST_CASE("extend_vector") {
+    SUBCASE("non-empty vectors") {
+      std::vector<int> lhs = {1, 2};
+      std::vector<int> rhs = {3, 4};
+
+      extend_vector(lhs, rhs);
+      std::vector<int> correct = {1, 2, 3, 4};
+
+      CHECK(lhs == correct);
+    }
+    // checks that lhs stays the same when the rhs is empty
+    SUBCASE("empty rhs") {
+      std::vector<int> lhs = {1, 2};
+      std::vector<int> rhs = {};
+
+      extend_vector(lhs, rhs);
+      std::vector<int> correct = {1, 2};
+
+      CHECK(lhs == correct);
+    }
+  }
+}
diff --git a/lib/utils/test/src/utils/containers/generate_vector.cc b/lib/utils/test/src/utils/containers/generate_vector.cc
new file mode 100644
index 0000000000..e560ba80bd
--- /dev/null
+++ b/lib/utils/test/src/utils/containers/generate_vector.cc
@@ -0,0 +1,31 @@
+#include "utils/containers/generate_vector.h"
+#include "test/utils/doctest/fmt/vector.h"
+#include "utils/exception.h"
+#include "utils/nonnegative_int/nonnegative_int.h"
+#include <doctest/doctest.h>
+#include <vector>
+
+using namespace ::FlexFlow;
+
+TEST_SUITE(FF_TEST_SUITE) {
+  TEST_CASE("generate_vector") {
+    SUBCASE("empty vector") {
+      std::vector<int> result = generate_vector(0_n, [](nonnegative_int) -> int {
+        PANIC("lambda should not be called");
+      });
+      std::vector<int> correct = {};
+
+      CHECK(result == correct);
+    }
+    SUBCASE("non-empty vector") {
+      std::vector<int> result = generate_vector(5_n, [](nonnegative_int idx) {
+        int i = idx.unwrap_nonnegative();
+        return i * i;
+
+      });
+      std::vector<int> correct = {0, 1, 4, 9, 16};
+
+      CHECK(result == correct);
+    }
+  }
+}