From ae5333f41a008242c334f555ede3edd5528fa2ba Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Samuel=20K=2E=20Guti=C3=A9rrez?= <samuel@lanl.gov>
Date: Wed, 24 Apr 2024 16:34:34 -0600
Subject: [PATCH] Lift color space size restriction. (#115)

Now we allow for more colors than number of pieces for a split.
qvi_scope_ksplit() is getting closer to the semantics we need.

Signed-off-by: Samuel K. Gutierrez <samuel@lanl.gov>
---
 src/qvi-common.h               |  1 +
 src/qvi-scope.cc               | 26 +++++---------------------
 tests/test-mpi-scope-ksplit.cc |  2 +-
 3 files changed, 7 insertions(+), 22 deletions(-)

diff --git a/src/qvi-common.h b/src/qvi-common.h
index 447a9ae6..bbf98a5e 100644
--- a/src/qvi-common.h
+++ b/src/qvi-common.h
@@ -48,6 +48,7 @@
 #include <chrono>
 #include <map>
 #include <new>
+#include <numeric>
 #include <set>
 #include <stack>
 #include <stdexcept>
diff --git a/src/qvi-scope.cc b/src/qvi-scope.cc
index d6e8415a..d503eae5 100644
--- a/src/qvi-scope.cc
+++ b/src/qvi-scope.cc
@@ -963,30 +963,14 @@ agg_split_affinity_preserving(
 }
 
 /**
- * Takes a vector of colors and clamps their values to [0, max) in place. If
- * this function finds more colors than slots available from [0, max), then it
- * returns an error code.
+ * Takes a vector of colors and clamps their values to [0, values.size()) in
+ * place.
  */
-// TODO(skg) Lift size restriction.
 static int
 clamp_colors(
-    std::vector<int> &values,
-    uint_t max
+    std::vector<int> &values
 ) {
-    // Recall: sets are ordered.
-    std::set<int> valset(values.begin(), values.end());
-    // Too many colors for the given space.
-    if (valset.size() > max) return QV_ERR_INVLD_ARG;
-    // Maps the input vector colors to their clamped values.
-    std::map<int, int> colors2clamped;
-    // color': the clamped color.
-    int colorp = 0;
-    for (auto val : valset) {
-        colors2clamped.insert({val, colorp++});
-    }
-    for (uint_t i = 0; i < values.size(); ++i) {
-        values[i] = colors2clamped[values[i]];
-    }
+    std::iota(values.begin(), values.end(), 0);
     return QV_SUCCESS;
 }
 
@@ -1019,7 +1003,7 @@ agg_split(
 
     // All colors are positive.
     if (tcolors.front() >= 0) {
-        rc = clamp_colors(splitagg.colors, splitagg.split_size);
+        rc = clamp_colors(splitagg.colors);
         if (rc != QV_SUCCESS) return rc;
     }
     // Some values are negative.
diff --git a/tests/test-mpi-scope-ksplit.cc b/tests/test-mpi-scope-ksplit.cc
index 930f0185..f0d2005e 100644
--- a/tests/test-mpi-scope-ksplit.cc
+++ b/tests/test-mpi-scope-ksplit.cc
@@ -82,7 +82,7 @@ main(void)
 
     // Test internal APIs
     const int npieces = ncores / 2;
-    std::vector<int> colors(npieces);
+    std::vector<int> colors(npieces * 2);
     std::iota(colors.begin(), colors.end(), 0);
 
     qv_scope_t **subscopes = nullptr;