new type instability during restrict on OffsetArray

[johnnychen94]

Observed from c09f576

It seems that JuliaArrays/OffsetArrays.jl#213 breaks this:

using ImageTransformations
using OffsetArrays
using Test

A = reshape(collect(1:60), 5, 4, 3)
Ao = OffsetArray(A, (-2,1,0))
@inferred(restrict(Ao, 1)) # inferred Any

Yet I haven't investigated where the patch should be put.

cc: @jishnub

[jishnub]

The problem appears to be that newinds doesn't have a concrete type.

julia> @code_warntype restrict(Ao, 1)
MethodInstance for ImageTransformations.restrict(::OffsetArray{Int64, 3, Array{Int64, 3}}, ::Int64)
  from restrict(A::AbstractArray{T, N}, dim::Integer) where {T, N} in ImageTransformations at /home/jishnu/.julia/packages/ImageTransformations/xYRLH/src/resizing.jl:78
Static Parameters
  T = Int64
  N = 3
Arguments
  #self#::Core.Const(ImageTransformations.restrict)
  A::OffsetArray{Int64, 3, Array{Int64, 3}}
  dim::Int64
Locals
  #13::ImageTransformations.var"#13#14"{Int64, Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}}
  out::Any
  newinds::Tuple{Union{UnitRange{Int64}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}, Union{UnitRange{Int64}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}, Union{UnitRange{Int64}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}}
  indsA::Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}
Body::Any
1 ─       (indsA = ImageTransformations.axes(A))
│   %2  = ImageTransformations.:(var"#13#14")::Core.Const(ImageTransformations.var"#13#14")
│   %3  = Core.typeof(dim)::Core.Const(Int64)
│   %4  = Core.typeof(indsA)::Core.Const(Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}})
│   %5  = Core.apply_type(%2, %3, %4)::Core.Const(ImageTransformations.var"#13#14"{Int64, Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}})
│         (#13 = %new(%5, dim, indsA))
│   %7  = #13::ImageTransformations.var"#13#14"{Int64, Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}}
│   %8  = ImageTransformations.Val($(Expr(:static_parameter, 2)))::Core.Const(Val{3}())
│         (newinds = ImageTransformations.ntuple(%7, %8))
│   %10 = ImageTransformations.first(A)::Int64
│   %11 = ImageTransformations.restrict_eltype(%10)::Core.Const(Float64)
│   %12 = Core.apply_type(ImageTransformations.Array, %11, $(Expr(:static_parameter, 2)))::Core.Const(Array{Float64, 3})
│         (out = ImageTransformations.similar(%12, newinds))
│         ImageTransformations.restrict!(out, A, dim)
└──       return out

Casting the axes to the same type appears to solve this:

julia> function restrict2(A::AbstractArray{T,N}, dim::Integer) where {T,N}
           indsA = axes(A)
           newinds = map(UnitRange, ntuple(i->i==dim ? ImageTransformations.restrict_indices(indsA[dim]) : indsA[i], Val(N)))
           out = similar(Array{ImageTransformations.restrict_eltype(first(A)),N}, newinds)
           ImageTransformations.restrict!(out, A, dim)
           out
       end
restrict2 (generic function with 1 method)

julia> @code_warntype restrict2(Ao, 1)
MethodInstance for restrict2(::OffsetArray{Int64, 3, Array{Int64, 3}}, ::Int64)
  from restrict2(A::AbstractArray{T, N}, dim::Integer) where {T, N} in Main at REPL[49]:1
Static Parameters
  T = Int64
  N = 3
Arguments
  #self#::Core.Const(restrict2)
  A::OffsetArray{Int64, 3, Array{Int64, 3}}
  dim::Int64
Locals
  #23::var"#23#24"{Int64, Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}}
  out::OffsetArray{Float64, 3, Array{Float64, 3}}
  newinds::Tuple{UnitRange{Int64}, UnitRange{Int64}, UnitRange{Int64}}
  indsA::Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}
Body::OffsetArray{Float64, 3, Array{Float64, 3}}
1 ─       (indsA = Main.axes(A))
│   %2  = Main.:(var"#23#24")::Core.Const(var"#23#24")
│   %3  = Core.typeof(dim)::Core.Const(Int64)
│   %4  = Core.typeof(indsA)::Core.Const(Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}})
│   %5  = Core.apply_type(%2, %3, %4)::Core.Const(var"#23#24"{Int64, Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}})
│         (#23 = %new(%5, dim, indsA))
│   %7  = #23::var"#23#24"{Int64, Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}}
│   %8  = Main.Val($(Expr(:static_parameter, 2)))::Core.Const(Val{3}())
│   %9  = Main.ntuple(%7, %8)::Tuple{Union{UnitRange{Int64}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}, Union{UnitRange{Int64}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}, Union{UnitRange{Int64}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}}
│         (newinds = Main.map(Main.UnitRange, %9))
│   %11 = ImageTransformations.restrict_eltype::Core.Const(ImageTransformations.restrict_eltype)
│   %12 = Main.first(A)::Int64
│   %13 = (%11)(%12)::Core.Const(Float64)
│   %14 = Core.apply_type(Main.Array, %13, $(Expr(:static_parameter, 2)))::Core.Const(Array{Float64, 3})
│         (out = Main.similar(%14, newinds))
│         Main.restrict!(out, A, dim)
└──       return out

julia> function restrict2(A::AbstractArray{T,N}, dim::Integer) where {T,N}
                  indsA = axes(A)
                  newinds = ntuple(i->i==dim ? oftype(first(indsA), ImageTransformations.restrict_indices(indsA[dim])) : indsA[i], Val(N))
                  out = similar(Array{ImageTransformations.restrict_eltype(first(A)),N}, newinds)
                  ImageTransformations.restrict!(out, A, dim)
                  out
              end

julia> @code_warntype restrict2(Ao, 1)
Variables
  #self#::Core.Const(restrict2)
  A::OffsetArray{Int64, 3, Array{Int64, 3}}
  dim::Int64
  #387::var"#387#388"{Int64, Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}}
  out::OffsetArray{Float64, 3, Array{Float64, 3}}
  newinds::Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}
  indsA::Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}

Body::OffsetArray{Float64, 3, Array{Float64, 3}}
1 ─       (indsA = Main.axes(A))
│   %2  = Main.:(var"#387#388")::Core.Const(var"#387#388")
│   %3  = Core.typeof(dim)::Core.Const(Int64)
│   %4  = Core.typeof(indsA)::Core.Const(Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}})
│   %5  = Core.apply_type(%2, %3, %4)::Core.Const(var"#387#388"{Int64, Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}})
│         (#387 = %new(%5, dim, indsA))
│   %7  = #387::var"#387#388"{Int64, Tuple{OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}, OffsetArrays.IdOffsetRange{Int64, Base.OneTo{Int64}}}}
│   %8  = Main.Val($(Expr(:static_parameter, 2)))::Core.Const(Val{3}())
│         (newinds = Main.ntuple(%7, %8))
│   %10 = ImageTransformations.restrict_eltype::Core.Const(ImageTransformations.restrict_eltype)
│   %11 = Main.first(A)::Int64
│   %12 = (%10)(%11)::Core.Const(Float64)
│   %13 = Core.apply_type(Main.Array, %12, $(Expr(:static_parameter, 2)))::Core.Const(Array{Float64, 3})
│         (out = Main.similar(%13, newinds))
│   %15 = ImageTransformations.restrict!::Core.Const(ImageTransformations.restrict!)
│   %16 = out::OffsetArray{Float64, 3, Array{Float64, 3}}
│         (%15)(%16, A, dim)
└──       return out

I'm a bit surprised that this isn't inferred though. Shouldn't this just be an OffsetArray{Float64,3,Array{Float64,3}}?