Use to_indices() and @BoundsCheck

JuliaLang · Feb 6, 2017 · fe91e5d · fe91e5d
1 parent 043a2f0
commit fe91e5d
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/base/array.jl b/base/array.jl
@@ -807,7 +807,10 @@ Stacktrace:
  [1] deleteat!(::Array{Int64,1}, ::Tuple{Int64,Int64}) at ./array.jl:808
 ```
 """
-function deleteat!(a::Vector, inds)
+deleteat!(a::Vector, inds) = _deleteat!(a, inds)
+deleteat!(a::Vector, inds::AbstractVector) = _deleteat!(a, to_indices(a, (inds,))[1])
+
+function _deleteat!(a::Vector, inds)
     n = length(a)
     s = start(inds)
     done(inds, s) && return a
@@ -836,9 +839,10 @@ function deleteat!(a::Vector, inds)
     return a
 end
 
+# Simpler and more efficient version for logical indexing
 function deleteat!(a::Vector, inds::AbstractVector{Bool})
     n = length(a)
-    length(inds) == n || throw(BoundsError(a, inds))
+    @boundscheck length(inds) == n || throw(BoundsError(a, inds))
     p = 1
     for (q, i) in enumerate(inds)
         @inbounds a[p] = a[q]

diff --git a/test/arrayops.jl b/test/arrayops.jl
@@ -1067,10 +1067,14 @@ 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]
-        # integer indexing
+        # integer indexing with AbstractArray
         a = [1:10;]; acopy = copy(a)
         @test deleteat!(a, idx) == [acopy[1:(first(idx)-1)]; acopy[(last(idx)+1):end]]
 
+        # integer indexing with non-AbstractArray iterable
+        a = [1:10;]; acopy = copy(a)
+        @test deleteat!(a, (i for i in idx)) == [acopy[1:(first(idx)-1)]; acopy[(last(idx)+1):end]]
+
         # logical indexing
         a = [1:10;]; acopy = copy(a)
         @test deleteat!(a, map(i -> i in idx, 1:length(a))) == [acopy[1:(first(idx)-1)]; acopy[(last(idx)+1):end]]