Index traits

Project.toml

@@ -19,7 +19,10 @@ ArrayInterfaceBandedMatricesExt = "BandedMatrices"
 ArrayInterfaceBlockBandedMatricesExt = "BlockBandedMatrices"
 ArrayInterfaceCUDAExt = "CUDA"
 ArrayInterfaceGPUArraysCoreExt = "GPUArraysCore"
+ArrayInterfaceOffsetArraysExt = "OffsetArrays"
 ArrayInterfaceStaticArraysCoreExt = "StaticArraysCore"
+ArrayInterfaceStaticArraysExt = "StaticArrays"
+ArrayInterfaceStaticExt = "Static"
 ArrayInterfaceTrackerExt = "Tracker"
 
 [extras]
@@ -30,21 +33,26 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["SafeTestsets", "Pkg", "Test", "Aqua", "Random", "SparseArrays", "SuiteSparse", "BandedMatrices", "BlockBandedMatrices", "GPUArraysCore", "StaticArrays", "StaticArraysCore", "Tracker"]
+test = ["SafeTestsets", "Pkg", "Test", "Aqua", "Random", "SparseArrays", "SuiteSparse", "BandedMatrices", "BlockBandedMatrices", "GPUArraysCore", "OffsetArrays", "StaticArrays", "StaticArraysCore", "Static", "Tracker"]
 
 [weakdeps]
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"

docs/src/indexing.md

@@ -14,6 +14,7 @@ ArrayInterface.can_change_size
 ArrayInterface.can_setindex
 ArrayInterface.fast_scalar_indexing
 ArrayInterface.ismutable
+ArrayInterface.is_splat_index
 ArrayInterface.ndims_index
 ArrayInterface.ndims_shape
 ArrayInterface.defines_strides
@@ -22,6 +23,11 @@ ArrayInterface.ensures_sorted
 ArrayInterface.indices_do_not_alias
 ArrayInterface.instances_do_not_alias
 ArrayInterface.device
+ArrayInterface.known_first
+ArrayInterface.known_step
+ArrayInterface.known_last
+ArrayInterface.known_size
+ArrayInterface.known_length
 ```
 
 ## Allowed Indexing Functions
@@ -46,4 +52,4 @@ and index translations.
 ArrayInterface.ArrayIndex
 ArrayInterface.GetIndex
 ArrayInterface.SetIndex!
-```
+```

ext/ArrayInterfaceOffsetArraysExt.jl

@@ -0,0 +1,21 @@
+module ArrayInterfaceOffsetArraysExt
+
+if isdefined(Base, :get_extension) 
+    using ArrayInterface
+    using OffsetArrays
+else 
+    using ..ArrayInterface
+    using ..OffsetArrays
+end
+
+ArrayInterface.parent_type(@nospecialize T::Type{<:OffsetArrays.IdOffsetRange}) = fieldtype(T, :parent)
+ArrayInterface.parent_type(@nospecialize T::Type{<:OffsetArray}) = fieldtype(T, :parent)
+
+function ArrayInterface.known_size(@nospecialize T::Type{<:OffsetArrays.IdOffsetRange})
+    ArrayInterface.known_size(ArrayInterface.parent_type(T))
+end
+function ArrayInterface.known_size(@nospecialize T::Type{<:OffsetArray})
+    ArrayInterface.known_size(ArrayInterface.parent_type(T))
+end
+
+end

ext/ArrayInterfaceStaticArraysCoreExt.jl

@@ -32,4 +32,13 @@ end
 
 ArrayInterface.restructure(x::StaticArraysCore.SArray{S}, y) where {S} = StaticArraysCore.SArray{S}(y)
 
+function ArrayInterface.known_size(::Type{<:StaticArraysCore.StaticArray{S}}) where {S}
+    @isdefined(S) ? tuple(S.parameters...) : ntuple(_-> nothing, ndims(T))
+end
+
+function ArrayInterface.known_length(T::Type{<:StaticArraysCore.StaticArray})
+    sz = ArrayInterface.known_size(T)
+    isa(sz, Tuple{Vararg{Nothing}}) ? nothing : prod(sz)
+end
+
 end

ext/ArrayInterfaceStaticArraysExt.jl

@@ -0,0 +1,26 @@
+module ArrayInterfaceStaticArraysExt
+
+if isdefined(Base, :get_extension)
+    import ArrayInterface
+    import StaticArrays
+else
+    import ..ArrayInterface
+    import ..StaticArrays
+end
+
+ArrayInterface.known_first(@nospecialize T::Type{<:StaticArrays.SOneTo}) = 1
+ArrayInterface.known_last(::Type{StaticArrays.SOneTo{N}}) where {N} = @isdefined(N) ? N::Int : nothing
+
+function ArrayInterface.known_first(::Type{<:StaticArrays.SUnitRange{S}}) where {S}
+    @isdefined(S) ? S::Int : nothing
+end
+function ArrayInterface.known_size(::Type{<:StaticArrays.SUnitRange{<:Any, L}}) where {L}
+    @isdefined(L) ? (L::Int,) : (nothing,)
+end
+function ArrayInterface.known_last(::Type{<:StaticArrays.SUnitRange{S, L}}) where {S, L}
+    start = @isdefined(S) ? S::Int : nothing
+    len = @isdefined(L) ? L::Int : nothing
+    (start === nothing || len === nothing) ? nothing : (start + len - 1)
+end
+
+end

ext/ArrayInterfaceStaticExt.jl

@@ -0,0 +1,19 @@
+module ArrayInterfaceStaticExt
+
+if isdefined(Base, :get_extension)
+    import ArrayInterface
+    import Static
+else
+    import ..ArrayInterface
+    import ..Static
+end
+
+ArrayInterface.known_first(::Type{<:Static.OptionallyStaticUnitRange{Static.StaticInt{F}}}) where {F} = F::Int
+ArrayInterface.known_first(::Type{<:Static.OptionallyStaticStepRange{Static.StaticInt{F}}}) where {F} = F::Int
+
+ArrayInterface.known_step(::Type{<:Static.OptionallyStaticStepRange{<:Any,Static.StaticInt{S}}}) where {S} = S::Int
+
+ArrayInterface.known_last(::Type{<:Static.OptionallyStaticUnitRange{<:Any,Static.StaticInt{L}}}) where {L} = L::Int
+ArrayInterface.known_last(::Type{<:Static.OptionallyStaticStepRange{<:Any,<:Any,Static.StaticInt{L}}}) where {L} = L::Int
+
+end