A demo of external compiler passes in Julia

expandci.jl

using Base.Meta: isexpr

# Grabbing the inferred code

function disable_pass!(src)
    for (line, stmt) in enumerate(src.code)
        if isexpr(stmt, :meta, 1)
            a = stmt.args[1]
            if isexpr(a, :external_pass)
                mode = a.args[1]
                mode === :start && return disable_pass!(src, line)
            end
        end
    end
    error("not resolved")
end

function disable_pass!(src, line)
    stmt = src.code[line]
    @assert isexpr(stmt, :meta, 1)
    @assert isexpr(stmt.args[1], :external_pass, 2)
    @assert stmt.args[1].args[1] === :start
    nnest = 0
    for i = line+1:length(src.code)
        stmt = src.code[i]
        if isexpr(stmt, :meta, 1)
            a = stmt.args[1]
            if isexpr(a, :external_pass)
                mode = a.args[1]
                if mode === :start
                    nnest += 1
                elseif mode === :stop
                    if nnest == 0
                        src.code[i] = nothing
                        src.code[line] = nothing
                        return src
                    else
                        nnest -= 1
                    end
                end
            end
        end
    end
    error("finished and didn't see the end")
end

# These let you inspect the code before an after; it's a huge help for debugging!
const _src0 = Ref{Any}(nothing)
const _src1 = Ref{Any}(nothing)
function expandci_pass!(opt, line)
    _src0[] = deepcopy(opt.src)
    disable_pass!(opt.src, line)
    expand_cartesian_loops!(opt.src)
    _src1[] = deepcopy(opt.src)
    return nothing
end

macro expandci(ex)
    esc(Expr(:block, Expr(:meta, Expr(:external_pass, :start, expandci_pass!)), ex, Expr(:meta, Expr(:external_pass, :stop))))
end

function sumcart_iter(A)
    s = zero(eltype(A))
    @expandci for I in CartesianIndices(A)
        @inbounds s += A[I]
    end
    s
end

function sumcart_manual(A::AbstractMatrix)
    s = zero(eltype(A))
    ax1, ax2 = axes(A)
    for j in ax2, i in ax1
        @inbounds s += A[i, j]
    end
    s
end

A = rand(4, 4)
using Test
@test sumcart_iter(A) == sumcart_manual(A)
using BenchmarkTools
@btime sumcart_iter($A)
@btime sumcart_manual($A)