add keepat! for in-place logical filtering

NEWS.md

@@ -50,6 +50,8 @@ New library functions
 * New `keepat!(vector, inds)` function which is the inplace equivalent of `vector[inds]`
   for a list `inds` of integers ([#36229]).
 * 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 `keepat!(a::AbstractVector, m::AbstractVector{Bool})` which provides an in-place equivalent to the logical indexing expression `a = a[m]`
 * Two arguments method `lock(f, lck)` now accepts a `Channel` as the second argument. ([#39312])
 * New functor `Returns(value)`, which returns `value` for any arguments ([#39794])
 * New macro `Base.@invoke f(arg1::T1, arg2::T2; kwargs...)` provides an easier syntax to call `invoke(f, Tuple{T1,T2}, arg1, arg2; kwargs...)` ([#38438])

base/array.jl

@@ -2500,6 +2500,43 @@ function filter!(f, a::AbstractVector)
     return a
 end
 
+"""
+     keepat!(a::AbstractVector, m::AbstractVector{Bool})
+
+The inplace version of logical indexing `a = a[m]`. That is, `keepat!(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> keepat!(a, [true, false, true])
+2-element Vector{Symbol}:
+ :a
+ :c
+
+julia> a
+2-element Vector{Symbol}:
+ :a
+ :c
+```
+"""
+function keepat!(a::AbstractVector, m::AbstractVector{Bool})
+    j = firstindex(a)
+    for i in eachindex(a, m)
+        @inbounds begin
+            if m[i]
+                i == j || (a[j] = a[i])
+                j = nextind(a, j)
+            end
+        end
+    end
+    j > lastindex(a) && return a
+    deleteat!(a, j:lastindex(a))
+    return a
+end
+
 # set-like operators for vectors
 # These are moderately efficient, preserve order, and remove dupes.
 

base/exports.jl

@@ -534,6 +534,7 @@ export
     mapfoldl,
     mapfoldr,
     mapreduce,
+    keepat!,
     merge!,
     mergewith!,
     merge,

doc/src/base/collections.md

@@ -135,6 +135,7 @@ Base.collect(::Any)
 Base.collect(::Type, ::Any)
 Base.filter
 Base.filter!
+Base.keepat!
 Base.replace(::Any, ::Pair...)
 Base.replace(::Base.Callable, ::Any)
 Base.replace!

test/arrayops.jl

@@ -1446,6 +1446,32 @@ end
     @test isempty(eoa)
 end
 
+@testset "keepat!" begin
+    # base case w/ Vector
+    a = Vector(1:10)
+    keepat!(a, [falses(5); trues(5)])
+    @test a == 6:10
+
+    # different subtype of AbstractVector
+    ba = rand(10) .> 0.5 #
+    @test isa(ba, BitArray)
+    keepat!(ba, ba)
+    @test all(ba)
+
+    # empty array
+    ea = []
+    keepat!(ea, Bool[])
+    @test isempty(ea)
+
+    # non-1-indexed array
+    # deleteat! is not supported for OffsetArrays
+
+    # empty non-1-indexed array
+    eoa = OffsetArray([], -5)
+    keepat!(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]