Merge pull request #992 from simeonschaub/offset-idx

properly support offset indices

Author: ChrisRackauckas

src/array-lib.jl

@@ -22,12 +22,12 @@ function Base.getindex(x::SymArray, idx...)
     meta = metadata(unwrap(x))
     if shape(x) !== Unknown() && all(i -> i isa Integer, idx)
         II = CartesianIndices(axes(x))
+        ii = CartesianIndex(idx)
         @boundscheck begin
-            if !checkbounds(Bool, II, idx...)
+            if !in(ii, II)
                 throw(BoundsError(x, idx))
             end
         end
-        ii = II[idx...]
         res = Term{eltype(symtype(x))}(getindex, [x, Tuple(ii)...]; metadata = meta)
     elseif all(i -> symtype(i) <: Integer, idx)
         shape(x) !== Unknown() && @boundscheck begin

src/arrays.jl

@@ -81,6 +81,13 @@ function Base.show(io::IO, aop::ArrayOp)
     end
 end
 
+Base.summary(io::IO, aop::ArrayOp) = Base.array_summary(io, aop, shape(aop))
+function Base.showarg(io::IO, aop::ArrayOp, toplevel)
+    show(io, aop)
+    toplevel && print(io, "::", typeof(aop))
+    return nothing
+end
+
 symtype(a::ArrayOp{T}) where {T} = T
 istree(a::ArrayOp) = true
 function operation(a::ArrayOp)
@@ -208,7 +215,9 @@ function make_shape(output_idx, expr, ranges=Dict())
             end
             mi = matches[i]
             @assert !isempty(mi)
-            return get_extents(mi)
+            ext = get_extents(mi)
+            ext isa Unknown && return Unknown()
+            return Base.OneTo(length(ext))
         elseif i isa Integer
             return Base.OneTo(1)
         end
@@ -526,7 +535,7 @@ end
 function axes(A::Union{Arr, SymArray})
     s = shape(unwrap(A))
     s === Unknown() && error("axes of $A not known")
-    return map(x->1:length(x), s)
+    return s
 end
 
 

test/arrays.jl

@@ -393,3 +393,19 @@ end
     @test !isequal(a, b) && !isequal(b, c) && !isequal(a, c)
     @test hash(a) != hash(b) && hash(b) != hash(c) && hash(a) != hash(c)
 end
+
+@testset "Offset Indices" begin
+    @variables k[0:3]
+
+    @testset "i = $i" for i in 0:3
+        sym = unwrap(k[i])
+        @test operation(sym) === getindex
+        args = arguments(sym)
+        @test length(args) == 2
+        @test args[1] === unwrap(k)
+        @test args[2] === i
+    end
+
+    @test_throws BoundsError k[-1]
+    @test_throws BoundsError k[4]
+end

test/build_function_tests/stencil-broadcast-inplace.jl

@@ -5,8 +5,8 @@
               ˍ₋out_1 = (view)(ˍ₋out, 1:6, 1:6)
               ˍ₋out_1 .= 0
               ˍ₋out_2 = (view)(ˍ₋out, 2:5, 2:5)
-              for (j, j′) = zip(1:4, reset_to_one(1:4))
-                  for (i, i′) = zip(1:4, reset_to_one(1:4))
+              for (j, j′) = zip(Base.OneTo(4), reset_to_one(Base.OneTo(4)))
+                  for (i, i′) = zip(Base.OneTo(4), reset_to_one(Base.OneTo(4)))
                       ˍ₋out_2[i′, j′] = (+)(ˍ₋out_2[i′, j′], (+)(1, (getindex)(ˍ₋out_2_input_1, i, j)))
                   end
               end

test/build_function_tests/stencil-broadcast-outplace.jl

@@ -4,8 +4,8 @@
           ˍ₋out_1 = (view)(ˍ₋out, 1:6, 1:6)
           ˍ₋out_1 .= 0
           ˍ₋out_2 = (view)(ˍ₋out, 2:5, 2:5)
-          for (j, j′) = zip(1:4, reset_to_one(1:4))
-              for (i, i′) = zip(1:4, reset_to_one(1:4))
+          for (j, j′) = zip(Base.OneTo(4), reset_to_one(Base.OneTo(4)))
+              for (i, i′) = zip(Base.OneTo(4), reset_to_one(Base.OneTo(4)))
                   ˍ₋out_2[i′, j′] = (+)(ˍ₋out_2[i′, j′], (+)(1, (getindex)(ˍ₋out_2_input_1, i, j)))
               end
           end

test/build_function_tests/transpose-inplace.jl

@@ -1,5 +1,5 @@
 :(function (x,)
-      let ˍ₋out = zeros(Float64, map(length, (1:4, 1:4)))
+      let ˍ₋out = zeros(Float64, map(length, (Base.OneTo(4), Base.OneTo(4))))
           begin
               for (j, j′) = zip(1:4, reset_to_one(1:4))
                   for (i, i′) = zip(1:4, reset_to_one(1:4))