add keepat! for in-place logical filtering

NEWS.md

@@ -34,6 +34,7 @@ New library functions
 * `isunordered(x)` returns true if `x` is value that is normally unordered, such as `NaN` or `missing`.
 * New macro `Base.@invokelatest f(args...; kwargs...)` provides a convenient way to call `Base.invokelatest(f, args...; kwargs...)` ([#37971])
 * New macro `Base.@invoke f(arg1::T1, arg2::T2; kwargs...)` provides an easier syntax to call `invoke(f, Tuple{T1,T2}; kwargs...)` ([#38438])
+* New function `mask!(a::AbstractVector, m::AbstractVector{Bool})` which provides an in-place equivalent to the logical indexing expression `a = a[m]`
 
 New library features
 --------------------

base/array.jl

@@ -2436,6 +2436,48 @@ function filter!(f, a::AbstractVector)
     return a
 end
 
+"""
+     mask!(a::AbstractVector, m::AbstractVector{Bool})
+
+The inplace version of logical indexing `a = a[m]`. That is, `mask!(a, m)` on
+vectors of equal length `a` and `m` will remove all elements from `a` for which
+`m` at the corresponding index is `false`.
+
+# Examples
+```jldoctest
+julia> a = [:a, :b, :c];
+
+julia> mask!(a, [true, false, true])
+2-element Vector{Symbol}:
+ :a
+ :c
+
+julia> a
+2-element Vector{Symbol}:
+ :a
+ :c
+```
+"""
+function mask!(a::AbstractVector, m::AbstractVector{Bool})
+    j = firstindex(a)
+    for i in eachindex(a, m)
+        @inbounds begin
+            ai = a[i]
+            mi = m[i]
+            a[j] = ai
+        end
+        j = ifelse(mi, nextind(a, j), j)
+    end
+    j > lastindex(a) && return a
+    if a isa Vector
+        resize!(a, j-1)
+        sizehint!(a, j-1)
+    else
+        deleteat!(a, j:lastindex(a))
+    end
+    return a
+end
+
 # set-like operators for vectors
 # These are moderately efficient, preserve order, and remove dupes.
 

base/exports.jl

@@ -528,6 +528,7 @@ export
     mapfoldl,
     mapfoldr,
     mapreduce,
+    mask!,
     merge!,
     mergewith!,
     merge,

test/arrayops.jl

@@ -1442,6 +1442,32 @@ end
     @test isempty(eoa)
 end
 
+@testset "mask!" begin
+    # base case w/ Vector
+    a = Vector(1:10)
+    mask!(a, [falses(5); trues(5)])
+    @test a == 6:10
+
+    # different subtype of AbstractVector
+    ba = rand(10) .> 0.5 #
+    @test isa(ba, BitArray)
+    mask!(ba, ba)
+    @test all(ba)
+
+    # empty array
+    ea = []
+    mask!(ea, Bool[])
+    @test isempty(ea)
+
+    # non-1-indexed array
+    # deleteat! is not supported for OffsetArrays
+
+    # empty non-1-indexed array
+    eoa = OffsetArray([], -5)
+    mask!(eoa, Bool[])
+    @test isempty(eoa)
+end
+
 @testset "deleteat!" begin
     for idx in Any[1, 2, 5, 9, 10, 1:0, 2:1, 1:1, 2:2, 1:2, 2:4, 9:8, 10:9, 9:9, 10:10,
                    8:9, 9:10, 6:9, 7:10]