From cd2cca1fd246316af843baef240ada60447fd141 Mon Sep 17 00:00:00 2001
From: Marco Barbone <mbarbone@flatironinstitute.org>
Date: Tue, 28 Jul 2026 14:32:09 -0400
Subject: [PATCH] Add runtime batch_bool mask overloads for complex batch
 load/store

The scalar batch<T, A> already accepted both batch_bool_constant (compile-time)
and batch_bool (runtime) masks for load/store, but batch<std::complex<T>, A>
was missing the runtime overloads.  This made it impossible to use a dynamic
mask with complex batches.

Changes:
- Declare and define load(value_type const*, batch_bool<T,A>, Mode)
  and store(value_type*, batch_bool<T,A>, Mode) for batch<std::complex<T>>.
- Implement both with a scalar fallback: load zeroes inactive lanes,
  store only writes active lanes.
- Add runtime-mask load/store tests in test_batch_complex.cpp.

The existing free-function API in xsimd_api.hpp already documented that
batch_bool overloads exist and are slower than batch_bool_constant; no
documentation or comment churn in the headers was needed.
---
 include/xsimd/types/xsimd_batch.hpp | 25 +++++++++++++++++++++
 test/test_batch_complex.cpp         | 34 +++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/include/xsimd/types/xsimd_batch.hpp b/include/xsimd/types/xsimd_batch.hpp
index e97dd24e0..8d4721fa9 100644
--- a/include/xsimd/types/xsimd_batch.hpp
+++ b/include/xsimd/types/xsimd_batch.hpp
@@ -553,6 +553,8 @@ namespace xsimd
         // Compile-time mask overloads
         template <class U, bool... Values, class Mode = aligned_mode>
         [[nodiscard]] static XSIMD_INLINE batch load(U const* mem, batch_bool_constant<value_type, A, Values...> mask, Mode = {}) noexcept;
+        template <class Mode = aligned_mode>
+        [[nodiscard]] static XSIMD_INLINE batch load(value_type const* mem, batch_bool<T, A> mask, Mode = {}) noexcept;
         template <class U>
         [[nodiscard]] static XSIMD_INLINE batch load(U const* mem, stream_mode) noexcept;
         template <class U>
@@ -562,6 +564,8 @@ namespace xsimd
         // Compile-time mask overloads
         template <class U, bool... Values, class Mode = aligned_mode>
         XSIMD_INLINE void store(U* mem, batch_bool_constant<value_type, A, Values...> mask, Mode = {}) const noexcept;
+        template <class Mode = aligned_mode>
+        XSIMD_INLINE void store(value_type* mem, batch_bool<T, A> mask, Mode = {}) const noexcept;
         template <class U>
         XSIMD_INLINE void store(U* mem, stream_mode) const noexcept;
 
@@ -1506,6 +1510,17 @@ namespace xsimd
         kernel::store_masked<A>(mem, *this, mask, mode, A {});
     }
 
+    template <class T, class A>
+    template <class Mode>
+    XSIMD_INLINE void batch<std::complex<T>, A>::store(value_type* mem, batch_bool<T, A> mask, Mode) const noexcept
+    {
+        alignas(A::alignment()) std::array<value_type, size> buffer;
+        store_aligned(buffer.data());
+        for (std::size_t i = 0; i < size; ++i)
+            if (mask.get(i))
+                mem[i] = buffer[i];
+    }
+
     template <class T, class A>
     XSIMD_INLINE void batch<std::complex<T>, A>::store_aligned(T* real_dst, T* imag_dst) const noexcept
     {
@@ -1544,6 +1559,16 @@ namespace xsimd
         return kernel::load_masked<A>(mem, mask, kernel::convert<value_type> {}, mode, A {});
     }
 
+    template <class T, class A>
+    template <class Mode>
+    XSIMD_INLINE batch<std::complex<T>, A> batch<std::complex<T>, A>::load(value_type const* mem, batch_bool<T, A> mask, Mode) noexcept
+    {
+        alignas(A::alignment()) std::array<value_type, size> buffer {};
+        for (std::size_t i = 0; i < size; ++i)
+            buffer[i] = mask.get(i) ? mem[i] : value_type(0);
+        return load_aligned(buffer.data());
+    }
+
     template <class T, class A>
     template <class U>
     XSIMD_INLINE batch<std::complex<T>, A> batch<std::complex<T>, A>::load(U const* mem, stream_mode) noexcept
diff --git a/test/test_batch_complex.cpp b/test/test_batch_complex.cpp
index 9db85cef6..cccb4bee1 100644
--- a/test/test_batch_complex.cpp
+++ b/test/test_batch_complex.cpp
@@ -114,6 +114,40 @@ struct batch_complex_test
             CHECK_EQ(ares_real, areal);
             CHECK_EQ(ares_imag, aimag);
         }
+        {
+            // Runtime-masked load / store
+            alignas(arch_type::alignment()) array_type amem;
+            for (size_t i = 0; i < size; ++i)
+                amem[i] = lhs[i];
+
+            uint64_t bits = 0;
+            for (size_t i = 0; i < size; ++i)
+                if (i % 2 == 0)
+                    bits |= uint64_t(1) << i;
+            using mask_type = typename batch_type::batch_bool_type;
+            auto mask = mask_type::from_mask(bits);
+
+            batch_type loaded = batch_type::load(amem.data(), mask, xsimd::aligned_mode());
+            for (size_t i = 0; i < size; ++i)
+            {
+                if (i % 2 == 0)
+                    CHECK_EQ(loaded.get(i), amem[i]);
+                else
+                    CHECK_EQ(loaded.get(i), value_type(0, 0));
+            }
+
+            batch_type to_store = batch_type::load_unaligned(rhs.data());
+            alignas(arch_type::alignment()) array_type adest;
+            std::fill(adest.begin(), adest.end(), value_type(0, 0));
+            to_store.store(adest.data(), mask, xsimd::aligned_mode());
+            for (size_t i = 0; i < size; ++i)
+            {
+                if (i % 2 == 0)
+                    CHECK_EQ(adest[i], rhs[i]);
+                else
+                    CHECK_EQ(adest[i], value_type(0, 0));
+            }
+        }
     }
 #ifdef XSIMD_ENABLE_XTL_COMPLEX
     void test_load_store_xtl() const