MWEMissings.jl

module MWEMissings

using Missings, InteractiveUtils

export TwoVectors, mapreduce_impl_bad, mapreduce_impl_good

struct TwoVectors{V}
    x::V
    others::V
end

Base.eltype(::Type{<:TwoVectors{V}}) where {V} = nonmissingtype(eltype(V))
Base.IndexStyle(itr::TwoVectors) = Base.IndexStyle(itr.x)

@inline function _anymissingindex(others, i)
	others[i] === missing 
end

@noinline function mapreduce_impl_bad(f, op, itr::TwoVectors,
                                       ifirst::Integer, ilast::Integer, blksize::Int)
    A = itr.x
    if ifirst == ilast
        @inbounds a1 = A[ifirst]
        if a1 !== missing && !_anymissingindex(itr.others, ifirst) 
        	return Some(Base.mapreduce_first(f, op, a1))
        end
        return nothing
    elseif ifirst + blksize > ilast
        # sequential portion
        local ai
        i = ifirst
        while i <= ilast
            @inbounds ai = A[i]
            ai === missing || @inbounds _anymissingindex(itr.others, i) || break
            i += 1
        end
        i > ilast && return nothing
        a1 = ai::eltype(itr)
        i += 1
        while i <= ilast
            @inbounds ai = A[i]
            ai === missing || @inbounds _anymissingindex(itr.others, i) || break
            i += 1
        end
        i > ilast && return Some(Base.mapreduce_first(f, op, a1))
        a2 = ai::eltype(itr)
        i += 1
        v = op(f(a1), f(a2))
        @simd for i = i:ilast
            @inbounds ai = A[i]
            if ai !== missing 
                if @inbounds !_anymissingindex(itr.others, i)
                   v = op(v, f(ai))
                end
            end
        end
        return Some(v)
    else
        # pairwise portion
        imid = (ifirst + ilast) >> 1
        v1 = mapreduce_impl_bad(f, op, itr, ifirst, imid, blksize)
        v2 = mapreduce_impl_bad(f, op, itr, imid+1, ilast, blksize)
        if v1 === nothing && v2 === nothing
            return nothing
        elseif v1 === nothing
            return v2
        elseif v2 === nothing
            return v1
        else
            return Some(op(something(v1), something(v2)))
        end
    end
end

@noinline function mapreduce_impl_good(f, op, itr::TwoVectors,
                                       ifirst::Integer, ilast::Integer, blksize::Int)
    A = itr.x
    if ifirst == ilast
        @inbounds a1 = A[ifirst]
        if a1 === missing 
        	return nothing
        elseif _anymissingindex(itr.others, ifirst)
        	return nothing
    	end 
    	return Some(Base.mapreduce_first(f, op, a1))
    elseif ifirst + blksize > ilast
        # sequential portion
        local ai
        i = ifirst
        while i <= ilast
            @inbounds ai = A[i]
            ai === missing || @inbounds _anymissingindex(itr.others, i) || break
            i += 1
        end
        i > ilast && return nothing
        a1 = ai::eltype(itr)
        i += 1
        while i <= ilast
            @inbounds ai = A[i]
            ai === missing || @inbounds _anymissingindex(itr.others, i) || break
            i += 1
        end
        i > ilast && return Some(Base.mapreduce_first(f, op, a1))
        a2 = ai::eltype(itr)
        i += 1
        v = op(f(a1), f(a2))
        @simd for i = i:ilast
            @inbounds ai = A[i]
            if ai !== missing 
                if @inbounds !_anymissingindex(itr.others, i)
                   v = op(v, f(ai))
                end
            end
        end
        return Some(v)
    else
        # pairwise portion
        imid = (ifirst + ilast) >> 1
        v1 = mapreduce_impl_good(f, op, itr, ifirst, imid, blksize)
        v2 = mapreduce_impl_good(f, op, itr, imid+1, ilast, blksize)
        if v1 === nothing && v2 === nothing
            return nothing
        elseif v1 === nothing
            return v2
        elseif v2 === nothing
            return v1
        else
            return Some(op(something(v1), something(v2)))
        end
    end
end

@noinline function mapreduce_impl_good2(f, op, itr::TwoVectors,
                                       ifirst::Integer, ilast::Integer, blksize::Int)
    A = itr.x
    if ifirst == ilast
        @inbounds a1 = A[ifirst]
        a1 === missing || _anymissingindex(itr.others, ifirst) || return Some(Base.mapreduce_first(f, op, a1))
        return nothing
    elseif ifirst + blksize > ilast
        # sequential portion
        local ai
        i = ifirst
        while i <= ilast
            @inbounds ai = A[i]
            ai === missing || @inbounds _anymissingindex(itr.others, i) || break
            i += 1
        end
        i > ilast && return nothing
        a1 = ai::eltype(itr)
        i += 1
        while i <= ilast
            @inbounds ai = A[i]
            ai === missing || @inbounds _anymissingindex(itr.others, i) || break
            i += 1
        end
        i > ilast && return Some(Base.mapreduce_first(f, op, a1))
        a2 = ai::eltype(itr)
        i += 1
        v = op(f(a1), f(a2))
        @simd for i = i:ilast
            @inbounds ai = A[i]
            if ai !== missing 
                if @inbounds !_anymissingindex(itr.others, i)
                   v = op(v, f(ai))
                end
            end
        end
        return Some(v)
    else
        # pairwise portion
        imid = (ifirst + ilast) >> 1
        v1 = mapreduce_impl_good2(f, op, itr, ifirst, imid, blksize)
        v2 = mapreduce_impl_good2(f, op, itr, imid+1, ilast, blksize)
        if v1 === nothing && v2 === nothing
            return nothing
        elseif v1 === nothing
            return v2
        elseif v2 === nothing
            return v1
        else
            return Some(op(something(v1), something(v2)))
        end
    end
end

function easysum(itr::TwoVectors)
	s = zero(eltype(itr.x))
	for i in 1:length(itr.x)
		xi = itr.x[i]
		if xi !== missing && itr.others[i] !== missing 
			s+= xi
		end
	end
	return Some(s)
end

function main()
	x = map(1:100) do _ 
		rand() < .2 ? missing : rand()
	end 

	y = map(1:100) do _ 
		rand() < .2 ? missing : rand()
	end 

	tv = TwoVectors(x, y)

	@assert easysum(tv) == mapreduce_impl_bad(identity, +, tv, 1, 100, 1024) 
	@assert easysum(tv) == mapreduce_impl_good(identity, +, tv, 1, 100, 1024)
	@assert easysum(tv) == mapreduce_impl_good2(identity, +, tv, 1, 100, 1024)

	InteractiveUtils.@code_warntype MWEMissings.mapreduce_impl_bad(identity, +, tv, 1, 100, 1024)
	InteractiveUtils.@code_warntype MWEMissings.mapreduce_impl_good(identity, +, tv, 1, 100, 1024)
	InteractiveUtils.@code_warntype MWEMissings.mapreduce_impl_good2(identity, +, tv, 1, 100, 1024)

	return tv
end

end