Move StaticArrayInterface support into a package extension

Project.toml

@@ -1,26 +1,32 @@
 name = "EllipsisNotation"
 uuid = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
-version = "1.9.0"
+version = "1.10.0"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
 
 [deps]
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+
+[weakdeps]
 StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
 
+[extensions]
+EllipsisNotationStaticArrayInterfaceExt = "StaticArrayInterface"
+
 [compat]
 AllocCheck = "0.2"
 PrecompileTools = "1.0.1"
 SafeTestsets = "0.1"
 SciMLTesting = "1"
-StaticArrayInterface = "1.2.1"
+StaticArrayInterface = "1.4.1"
 Test = "1"
 julia = "1.10"
 
 [extras]
 AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SciMLTesting = "09d9d899-5365-40a9-917a-5f67fddea283"
+StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AllocCheck", "Test", "SafeTestsets", "SciMLTesting"]
+test = ["AllocCheck", "Test", "SafeTestsets", "SciMLTesting", "StaticArrayInterface"]

ext/EllipsisNotationStaticArrayInterfaceExt.jl

@@ -0,0 +1,15 @@
+module EllipsisNotationStaticArrayInterfaceExt
+
+using EllipsisNotation: Ellipsis
+using StaticArrayInterface: StaticArrayInterface
+
+# Integrate `..` with StaticArrayInterface's indexing machinery so that `..`
+# works through `StaticArrayInterface.getindex` / `setindex!`. Loaded only when
+# StaticArrayInterface is available.
+StaticArrayInterface.is_splat_index(::Type{Ellipsis}) = StaticArrayInterface.static(true)
+StaticArrayInterface.ndims_index(::Type{Ellipsis}) = StaticArrayInterface.static(1)
+function StaticArrayInterface.to_index(x, ::Ellipsis)
+    return ntuple(i -> StaticArrayInterface.indices(x, i), Val(ndims(x)))
+end
+
+end

src/EllipsisNotation.jl

@@ -1,7 +1,6 @@
 module EllipsisNotation
 
 using PrecompileTools: @compile_workload, @setup_workload
-using StaticArrayInterface: StaticArrayInterface
 
 import Base: to_indices, tail
 
@@ -54,21 +53,22 @@ const .. = Ellipsis()
         inds::NTuple{M, Any},
         I::Tuple{Ellipsis, Vararg{Any, N}}
     ) where {M, N}
-    # Align the remaining indices to the tail of the `inds`
-    colons = ntuple(n -> Colon(), M - _ndims_index(I) + 1)
+    # Align the remaining indices to the tail of the `inds`. `Val` keeps the
+    # number of inserted colons in the type domain so the result is type stable.
+    colons = ntuple(Returns(Colon()), Val(M - _ndims_index(I) + 1))
     return to_indices(A, inds, (colons..., tail(I)...))
 end
 
-@inline _ndims_index(inds::Tuple{}) = StaticArrayInterface.static(0)
-@inline function _ndims_index(inds::Tuple)
-    return StaticArrayInterface.ndims_index(inds[1]) + _ndims_index(tail(inds))
-end
+# Number of array dimensions consumed by a tuple of indices. Most index types
+# consume a single dimension; the exceptions below consume several at once.
+@inline _ndims_index(inds::Tuple{}) = 0
+@inline _ndims_index(inds::Tuple) = _ndims_index1(inds[1]) + _ndims_index(tail(inds))
 
-StaticArrayInterface.is_splat_index(::Type{Ellipsis}) = StaticArrayInterface.static(true)
-StaticArrayInterface.ndims_index(::Type{Ellipsis}) = StaticArrayInterface.static(1)
-function StaticArrayInterface.to_index(x, ::Ellipsis)
-    return ntuple(i -> StaticArrayInterface.indices(x, i), Val(ndims(x)))
-end
+@inline _ndims_index1(::Any) = 1
+@inline _ndims_index1(::CartesianIndex{N}) where {N} = N
+@inline _ndims_index1(::AbstractArray{Bool, N}) where {N} = N
+@inline _ndims_index1(::AbstractArray{<:CartesianIndex{N}}) where {N} = N
+@inline _ndims_index1(::Ellipsis) = 1
 
 export ..
 

test/basic.jl

@@ -1,7 +1,6 @@
 using Test
 using Test: detect_ambiguities
 using EllipsisNotation
-using StaticArrayInterface: StaticArrayInterface
 @test isempty(detect_ambiguities(EllipsisNotation))
 
 A = Array{Int}(undef, 2, 4, 2)
@@ -85,43 +84,3 @@ C[1, 1] += 1
 
 @test B[:] == B[..]
 @test B[:] !== B[..]
-
-@testset "StaticArrayInterface" begin
-    A = Array{Int}(undef, 2, 4, 2)
-    StaticArrayInterface.setindex!(A, [2 1 4 5; 2 2 3 6], .., 1)
-    StaticArrayInterface.setindex!(A, [3 2 6 5; 3 2 6 6], .., 2)
-
-    @test StaticArrayInterface.getindex(A, :, :, 1) == [2 1 4 5; 2 2 3 6]
-    @test StaticArrayInterface.getindex(A, :, :, 2) == [3 2 6 5; 3 2 6 6]
-
-    @test StaticArrayInterface.getindex(A, :, .., 1) == [2 1 4 5; 2 2 3 6]
-    @test StaticArrayInterface.getindex(A, :, .., 2) == [3 2 6 5; 3 2 6 6]
-
-    StaticArrayInterface.setindex!(A, reshape([3 4; 5 6; 4 5; 6 7], 4, 2), 1, ..)
-
-    B = [
-        3 4
-        5 6
-        4 5
-        6 7
-    ]
-
-    @test B ==
-        reshape(StaticArrayInterface.getindex(A, 1, ..), 4, 2) ==
-        reshape(view(A, 1, ..), 4, 2)
-
-    @test A[:, 1, 2] == StaticArrayInterface.getindex(A, .., 1, 2)
-
-    # [..]
-    C = zero(B)
-
-    C[:] = StaticArrayInterface.getindex(B, ..)
-    @test B == C
-    C[1, 1] += 1
-    @test B != C
-
-    StaticArrayInterface.setindex!(C, StaticArrayInterface.getindex(B, ..), ..)
-    @test B == C
-    C[1, 1] += 1
-    @test B != C
-end

test/static_array_interface.jl

@@ -0,0 +1,45 @@
+using Test
+using EllipsisNotation
+using StaticArrayInterface: StaticArrayInterface
+
+# Exercises the EllipsisNotationStaticArrayInterfaceExt extension: `..` working
+# through StaticArrayInterface's own getindex/setindex!.
+@testset "StaticArrayInterface" begin
+    A = Array{Int}(undef, 2, 4, 2)
+    StaticArrayInterface.setindex!(A, [2 1 4 5; 2 2 3 6], .., 1)
+    StaticArrayInterface.setindex!(A, [3 2 6 5; 3 2 6 6], .., 2)
+
+    @test StaticArrayInterface.getindex(A, :, :, 1) == [2 1 4 5; 2 2 3 6]
+    @test StaticArrayInterface.getindex(A, :, :, 2) == [3 2 6 5; 3 2 6 6]
+
+    @test StaticArrayInterface.getindex(A, :, .., 1) == [2 1 4 5; 2 2 3 6]
+    @test StaticArrayInterface.getindex(A, :, .., 2) == [3 2 6 5; 3 2 6 6]
+
+    StaticArrayInterface.setindex!(A, reshape([3 4; 5 6; 4 5; 6 7], 4, 2), 1, ..)
+
+    B = [
+        3 4
+        5 6
+        4 5
+        6 7
+    ]
+
+    @test B ==
+        reshape(StaticArrayInterface.getindex(A, 1, ..), 4, 2) ==
+        reshape(view(A, 1, ..), 4, 2)
+
+    @test A[:, 1, 2] == StaticArrayInterface.getindex(A, .., 1, 2)
+
+    # [..]
+    C = zero(B)
+
+    C[:] = StaticArrayInterface.getindex(B, ..)
+    @test B == C
+    C[1, 1] += 1
+    @test B != C
+
+    StaticArrayInterface.setindex!(C, StaticArrayInterface.getindex(B, ..), ..)
+    @test B == C
+    C[1, 1] += 1
+    @test B != C
+end