From fdb4eb471649aa8cac7e29a3cd9a6280481b5b29 Mon Sep 17 00:00:00 2001 From: JamesWrigley Date: Mon, 22 Jun 2026 12:41:31 +0200 Subject: [PATCH] Move StaticArrayInterface support into a package extension This makes the package much more lightweight, at the cost of a few extra lines of code. --- Project.toml | 12 +++-- ...EllipsisNotationStaticArrayInterfaceExt.jl | 15 +++++++ src/EllipsisNotation.jl | 24 +++++----- test/basic.jl | 41 ----------------- test/static_array_interface.jl | 45 +++++++++++++++++++ 5 files changed, 81 insertions(+), 56 deletions(-) create mode 100644 ext/EllipsisNotationStaticArrayInterfaceExt.jl create mode 100644 test/static_array_interface.jl diff --git a/Project.toml b/Project.toml index 909db39..3fb858a 100644 --- a/Project.toml +++ b/Project.toml @@ -1,18 +1,23 @@ name = "EllipsisNotation" uuid = "da5c29d0-fa7d-589e-88eb-ea29b0a81949" -version = "1.9.0" +version = "1.10.0" authors = ["Chris Rackauckas "] [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" @@ -20,7 +25,8 @@ julia = "1.10" 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"] diff --git a/ext/EllipsisNotationStaticArrayInterfaceExt.jl b/ext/EllipsisNotationStaticArrayInterfaceExt.jl new file mode 100644 index 0000000..061650c --- /dev/null +++ b/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 diff --git a/src/EllipsisNotation.jl b/src/EllipsisNotation.jl index 02698e8..e440366 100644 --- a/src/EllipsisNotation.jl +++ b/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 .. diff --git a/test/basic.jl b/test/basic.jl index 79b52ff..5ef4a11 100644 --- a/test/basic.jl +++ b/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 diff --git a/test/static_array_interface.jl b/test/static_array_interface.jl new file mode 100644 index 0000000..6b2c049 --- /dev/null +++ b/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