Testing an iterator for skipping a particular index in a for loop

test_skipat.jl

using Base.Iterators
using Base: @propagate_inbounds

#-------------------------------------------------------------------------------
# Iterator
#-------------------------------------------------------------------------------

struct SkipAt{I}
    iter::I
    skip::Int
end

skipat(iter, skip::Integer) = SkipAt(iter, Int(skip))

@propagate_inbounds function Base.iterate(s::SkipAt, i::Int=1)
    n = length(s.iter)
    if i != s.skip
        if i <= n
            return (s.iter[i], i+1)
        else
            return nothing
        end
    else
        if i < n
            return (s.iter[i+1], i+2)
        else
            return nothing
        end
    end
end

#-------------------------------------------------------------------------------
# Test functions
#-------------------------------------------------------------------------------

@inline cheap_operation(x) = x+one(x)
@inline expensive_operation(x) =  sqrt(exp(sin(x)*log(x) + x*x*x))

function test_if(a, j::Int, op::Function)
    c = zero(eltype(a))
    @inbounds for i in eachindex(a)
        if i != j
            c += op(a[i])
        end
    end
    return c
end

function test_two_loops(a, j::Int, op::Function)
    c = zero(eltype(a))
    @inbounds for i = 1:j-1
        c += op(a[i])
    end
    @inbounds for i = j+1:length(a)
        c += op(a[i])
    end
    return c
end

function test_take_and_drop(a, j::Int, op::Function)
   c = zero(eltype(a))
   @inbounds for b in take(a, j-1)
       c += op(b)
   end
   @inbounds for b in drop(a, j)
       c += op(b)
   end
   return c
end

function test_flatten_take_drop(a, j::Int, op::Function)
   c = zero(eltype(a))
   @inbounds for b in flatten((take(a, j-1), drop(a, j)))
       c += op(b)
   end
   return c
end

function test_skipat(a, j::Int, op::Function)
   c = zero(eltype(a))
   @inbounds for b in skipat(a, j)
       c += op(b)
   end
   return c
end

#-------------------------------------------------------------------------------
# Run tests
#-------------------------------------------------------------------------------

function main()
   a = rand(10000000)
   op = cheap_operation

   j = round(Int, length(a)/2)

   @time c1 = test_if(a, j, op)
   @time c2 = test_two_loops(a, j, op)
   @time c3 = test_take_and_drop(a, j, op)
   @time c4 = test_flatten_take_drop(a, j, op)
   @time c5 = test_skipat(a, j, op)

   println([c1 c2 c3 c4 c5], "\n")
end

println("Warming up:")
main();

println("Actual timing:")
main();