added subsetting for Vec<bool>

Konrad1991 · Jan 10, 2025 · 507aee0 · 507aee0
1 parent e99efa6
commit 507aee0
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 20 deletions.
diff --git a/.development/test_subsetting.cpp b/.development/test_subsetting.cpp
@@ -215,4 +215,13 @@ int main() {
   v = coca(1, 2, 3, 4);
   auto lv_test = subset_test(coca(1, 2, 3, 4), coca(1, 2, 1));
   print(lv_test);
+
+  Vec<bool> indices;
+  indices = coca(true, false, true, false);
+  auto bool_test = subset_test(v, indices);
+  print(v);
+  bool_test = 300;
+  print(v);
+  subset_test(v, coca(true, true, true, false)) = 400;
+  print(v);
 }
diff --git a/inst/include/etr_bits/Subsetting/LazySubsetting.hpp b/inst/include/etr_bits/Subsetting/LazySubsetting.hpp
@@ -18,16 +18,17 @@
 
 namespace etr {
 
-template <typename V, typename I>
-void precalcVecBool(V &vec, Indices &ind, I *idx) {
+template <typename V> void precalcVecInt(const V &vec, Indices &ind) {
   std::size_t sizeTrue = 0;
-  for (std::size_t i = 0; i < vec.size(); i++)
-    if ((*idx)[i % idx->size()])
+  for (std::size_t i = 0; i < vec.size(); i++) {
+    if (vec[i]) {
       sizeTrue++;
+    }
+  }
   ind.resize(sizeTrue);
   std::size_t counter = 0;
   for (std::size_t i = 0; i < vec.size(); i++) {
-    if ((*idx)[i % idx->size()]) {
+    if (vec[i]) {
       ind[counter] = i + 1;
       counter++;
     }
@@ -331,22 +332,38 @@ template <typename L, typename R> inline auto subset_test(L &&l, R &&r) {
   static_assert(!IsArithV<Decayed<L>>,
                 "\n\nYou cannot subset a scalar value\n\n");
   using RetType = typename ExtractDataType<Decayed<L>>::RetType;
-  using SubsetType =
-      SubsetClass<decltype(l.d), decltype(r.d), SubsetClassTrait>;
-  if constexpr (!IsRvalue<L &&> && !IsRvalue<R &&>) {
-    return Vec<RetType,
-               SubsetClass<decltype(l.d), decltype(r.d), SubsetClassTrait>>(
-        SubsetClass<decltype(l.d), decltype(r.d), SubsetClassTrait>(l.d, r.d));
-  } else if constexpr (!IsRvalue<L &&> && IsRvalue<R &&>) {
-    return Vec<RetType, SubsetClass<decltype(l.d), const decltype(r.d),
-                                    SubsetClassTrait>>(
-        SubsetClass<decltype(l.d), const decltype(r.d), SubsetClassTrait>(l.d,
-                                                                          r.d));
+  using RetTypeR = typename ExtractDataType<Decayed<R>>::RetType;
+  if constexpr (IS<RetTypeR, bool>) {
+    Indices ind;
+    precalcVecInt(r, ind);
+    if constexpr (!IsRvalue<L>) {
+      return Vec<RetType, SubsetClass<decltype(l.d), const decltype(ind),
+                                      SubsetClassTrait>>(
+          SubsetClass<decltype(l.d), const decltype(ind), SubsetClassTrait>(
+              l.d, ind));
+    } else {
+      return Vec<RetType, SubsetClass<const decltype(l.d), const decltype(ind),
+                                      SubsetClassTrait>>(
+          SubsetClass<const decltype(l.d), const decltype(ind),
+                      SubsetClassTrait>(l.d, ind));
+    }
   } else {
-    return Vec<RetType, SubsetClass<const decltype(l.d), const decltype(r.d),
-                                    SubsetClassTrait>>(
-        SubsetClass<const decltype(l.d), const decltype(r.d), SubsetClassTrait>(
-            l.d, r.d));
+    if constexpr (!IsRvalue<L &&> && !IsRvalue<R &&>) {
+      return Vec<RetType,
+                 SubsetClass<decltype(l.d), decltype(r.d), SubsetClassTrait>>(
+          SubsetClass<decltype(l.d), decltype(r.d), SubsetClassTrait>(l.d,
+                                                                      r.d));
+    } else if constexpr (!IsRvalue<L &&> && IsRvalue<R &&>) {
+      return Vec<RetType, SubsetClass<decltype(l.d), const decltype(r.d),
+                                      SubsetClassTrait>>(
+          SubsetClass<decltype(l.d), const decltype(r.d), SubsetClassTrait>(
+              l.d, r.d));
+    } else {
+      return Vec<RetType, SubsetClass<const decltype(l.d), const decltype(r.d),
+                                      SubsetClassTrait>>(
+          SubsetClass<const decltype(l.d), const decltype(r.d),
+                      SubsetClassTrait>(l.d, r.d));
+    }
   }
 }