Make sparse array minimal interface into interface functions

Author: mtfishman

src/abstractsparsearrayinterface.jl

@@ -1,19 +1,37 @@
+using Derive: Derive, @derive, @interface, AbstractArrayInterface
+
 # This is to bring `ArrayLayouts.zero!` into the namespace
 # since it is considered part of the sparse array interface.
 using ArrayLayouts: zero!
 
+function eachstoredindex end
+function getstoredindex end
+function getunstoredindex end
+function isstored end
+function setstoredindex! end
+function setunstoredindex! end
+function storedlength end
+function storedpairs end
+function storedvalues end
+
 # Minimal interface for `SparseArrayInterface`.
-isstored(a::AbstractArray, I::Int...) = true
-eachstoredindex(a::AbstractArray) = eachindex(a)
-getstoredindex(a::AbstractArray, I::Int...) = getindex(a, I...)
-function setstoredindex!(a::AbstractArray, value, I::Int...)
+# Fallbacks for dense/non-sparse arrays.
+@interface ::AbstractArrayInterface isstored(a::AbstractArray, I::Int...) = true
+@interface ::AbstractArrayInterface eachstoredindex(a::AbstractArray) = eachindex(a)
+@interface ::AbstractArrayInterface getstoredindex(a::AbstractArray, I::Int...) =
+  getindex(a, I...)
+@interface ::AbstractArrayInterface function setstoredindex!(
+  a::AbstractArray, value, I::Int...
+)
   setindex!(a, value, I...)
   return a
 end
 # TODO: Should this error by default if the value at the index
 # is stored? It could be disabled with something analogous
 # to `checkbounds`, like `checkstored`/`checkunstored`.
-function setunstoredindex!(a::AbstractArray, value, I::Int...)
+@interface ::AbstractArrayInterface function setunstoredindex!(
+  a::AbstractArray, value, I::Int...
+)
   setindex!(a, value, I...)
   return a
 end
@@ -36,11 +54,41 @@ end
 # Interface defaults.
 # TODO: Have a fallback that handles element types
 # that don't define `zero(::Type)`.
-getunstoredindex(a::AbstractArray, I::Int...) = zero(eltype(a))
+@interface ::AbstractArrayInterface getunstoredindex(a::AbstractArray, I::Int...) =
+  zero(eltype(a))
 
 # Derived interface.
-storedlength(a::AbstractArray) = length(storedvalues(a))
-storedpairs(a::AbstractArray) = map(I -> I => getstoredindex(a, I), eachstoredindex(a))
+@interface ::AbstractArrayInterface storedlength(a::AbstractArray) = length(storedvalues(a))
+@interface ::AbstractArrayInterface storedpairs(a::AbstractArray) =
+  map(I -> I => getstoredindex(a, I), eachstoredindex(a))
+
+@interface ::AbstractArrayInterface function eachstoredindex(as::AbstractArray...)
+  return eachindex(as...)
+end
+
+@interface ::AbstractArrayInterface storedvalues(a::AbstractArray) = a
+
+# Automatically derive the interface for all `AbstractArray` subtypes.
+# TODO: Define `SparseArrayInterfaceOps` derivable trait and rewrite this
+# as `@derive AbstractArray SparseArrayInterfaceOps`.
+@derive (T=AbstractArray,) begin
+  SparseArraysBase.eachstoredindex(::T)
+  SparseArraysBase.eachstoredindex(::T...)
+  SparseArraysBase.getstoredindex(::T, ::Int...)
+  SparseArraysBase.getunstoredindex(::T, ::Int...)
+  SparseArraysBase.isstored(::T, ::Int...)
+  SparseArraysBase.setstoredindex!(::T, ::Any, ::Int...)
+  SparseArraysBase.setunstoredindex!(::T, ::Any, ::Int...)
+  SparseArraysBase.storedlength(::T)
+  SparseArraysBase.storedpairs(::T)
+  SparseArraysBase.storedvalues(::T)
+end
+
+# TODO: Add `ndims` type parameter, like `Base.Broadcast.AbstractArrayStyle`.
+# TODO: This isn't used to define interface functions right now.
+# Currently, `@interface` expects an instance, probably it should take a
+# type instead so fallback functions can use abstract types.
+abstract type AbstractSparseArrayInterface <: AbstractArrayInterface end
 
 to_vec(x) = vec(collect(x))
 to_vec(x::AbstractArray) = vec(x)
@@ -66,22 +114,14 @@ function Base.setindex!(a::StoredValues, value, I::Int)
   return setstoredindex!(a.array, value, a.storedindices[I])
 end
 
-storedvalues(a::AbstractArray) = StoredValues(a)
+@interface ::AbstractSparseArrayInterface storedvalues(a::AbstractArray) = StoredValues(a)
 
-function eachstoredindex(a1, a2, a_rest...)
+@interface ::AbstractSparseArrayInterface function eachstoredindex(as::AbstractArray...)
   # TODO: Make this more customizable, say with a function
   # `combine/promote_storedindices(a1, a2)`.
-  return union(eachstoredindex.((a1, a2, a_rest...))...)
+  return union(eachstoredindex.(as)...)
 end
 
-using Derive: Derive, @derive, @interface, AbstractArrayInterface
-
-# TODO: Add `ndims` type parameter.
-# TODO: This isn't used to define interface functions right now.
-# Currently, `@interface` expects an instance, probably it should take a
-# type instead so fallback functions can use abstract types.
-abstract type AbstractSparseArrayInterface <: AbstractArrayInterface end
-
 # We restrict to `I::Vararg{Int,N}` to allow more general functions to handle trailing
 # indices and linear indices.
 @interface ::AbstractSparseArrayInterface function Base.getindex(

src/wrappers.jl

@@ -18,21 +18,21 @@ function eachstoredparentindex(a::SubArray)
     return all(d -> I[d] ∈ parentindices(a)[d], 1:ndims(parent(a)))
   end
 end
-function storedvalues(a::SubArray)
+@interface ::AbstractSparseArrayInterface function storedvalues(a::SubArray)
   # We use `StoredValues` rather than `@view`/`SubArray` so that
   # it gets interpreted as a dense array.
   return StoredValues(parent(a), collect(eachstoredparentindex(a)))
 end
-function isstored(a::SubArray, I::Int...)
+@interface ::AbstractSparseArrayInterface function isstored(a::SubArray, I::Int...)
   return isstored(parent(a), Base.reindex(parentindices(a), I)...)
 end
-function getstoredindex(a::SubArray, I::Int...)
+@interface ::AbstractSparseArrayInterface function getstoredindex(a::SubArray, I::Int...)
   return getstoredindex(parent(a), Base.reindex(parentindices(a), I)...)
 end
-function getunstoredindex(a::SubArray, I::Int...)
+@interface ::AbstractSparseArrayInterface function getunstoredindex(a::SubArray, I::Int...)
   return getunstoredindex(parent(a), Base.reindex(parentindices(a), I)...)
 end
-function eachstoredindex(a::SubArray)
+@interface ::AbstractSparseArrayInterface function eachstoredindex(a::SubArray)
   nonscalardims = filter(tuple_oneto(ndims(parent(a)))) do d
     return !(parentindices(a)[d] isa Real)
   end
@@ -48,27 +48,36 @@ end
 perm(::PermutedDimsArray{<:Any,<:Any,p}) where {p} = p
 iperm(::PermutedDimsArray{<:Any,<:Any,<:Any,ip}) where {ip} = ip
 
-storedvalues(a::PermutedDimsArray) = storedvalues(parent(a))
-function isstored(a::PermutedDimsArray, I::Int...)
+@interface ::AbstractSparseArrayInterface storedvalues(a::PermutedDimsArray) =
+  storedvalues(parent(a))
+@interface ::AbstractSparseArrayInterface function isstored(a::PermutedDimsArray, I::Int...)
   return isstored(parent(a), genperm(I, iperm(a))...)
 end
-function getstoredindex(a::PermutedDimsArray, I::Int...)
+@interface ::AbstractSparseArrayInterface function getstoredindex(
+  a::PermutedDimsArray, I::Int...
+)
   return getstoredindex(parent(a), genperm(I, iperm(a))...)
 end
-function getunstoredindex(a::PermutedDimsArray, I::Int...)
+@interface ::AbstractSparseArrayInterface function getunstoredindex(
+  a::PermutedDimsArray, I::Int...
+)
   return getunstoredindex(parent(a), genperm(I, iperm(a))...)
 end
-function setstoredindex!(a::PermutedDimsArray, value, I::Int...)
+@interface ::AbstractSparseArrayInterface function setstoredindex!(
+  a::PermutedDimsArray, value, I::Int...
+)
   # TODO: Should this be `iperm(a)`?
   setstoredindex!(parent(a), value, genperm(I, perm(a))...)
   return a
 end
-function setunstoredindex!(a::PermutedDimsArray, value, I::Int...)
+@interface ::AbstractSparseArrayInterface function setunstoredindex!(
+  a::PermutedDimsArray, value, I::Int...
+)
   # TODO: Should this be `iperm(a)`?
   setunstoredindex!(parent(a), value, genperm(I, perm(a))...)
   return a
 end
-function eachstoredindex(a::PermutedDimsArray)
+@interface ::AbstractSparseArrayInterface function eachstoredindex(a::PermutedDimsArray)
   # TODO: Make lazy with `Iterators.map`.
   return map(collect(eachstoredindex(parent(a)))) do I
     return CartesianIndex(genperm(I, perm(a)))
@@ -78,25 +87,34 @@ end
 for (type, func) in ((:Adjoint, :adjoint), (:Transpose, :transpose))
   @eval begin
     using LinearAlgebra: $type
-    storedvalues(a::$type) = storedvalues(parent(a))
-    function isstored(a::$type, i::Int, j::Int)
+    @interface ::AbstractSparseArrayInterface storedvalues(a::$type) =
+      storedvalues(parent(a))
+    @interface ::AbstractSparseArrayInterface function isstored(a::$type, i::Int, j::Int)
       return isstored(parent(a), j, i)
     end
-    function eachstoredindex(a::$type)
+    @interface ::AbstractSparseArrayInterface function eachstoredindex(a::$type)
       # TODO: Make lazy with `Iterators.map`.
       return map(cartesianindex_reverse, collect(eachstoredindex(parent(a))))
     end
-    function getstoredindex(a::$type, i::Int, j::Int)
+    @interface ::AbstractSparseArrayInterface function getstoredindex(
+      a::$type, i::Int, j::Int
+    )
       return $func(getstoredindex(parent(a), j, i))
     end
-    function getunstoredindex(a::$type, i::Int, j::Int)
+    @interface ::AbstractSparseArrayInterface function getunstoredindex(
+      a::$type, i::Int, j::Int
+    )
       return $func(getunstoredindex(parent(a), j, i))
     end
-    function setstoredindex!(a::$type, value, i::Int, j::Int)
+    @interface ::AbstractSparseArrayInterface function setstoredindex!(
+      a::$type, value, i::Int, j::Int
+    )
       setstoredindex!(parent(a), $func(value), j, i)
       return a
     end
-    function setunstoredindex!(a::$type, value, i::Int, j::Int)
+    @interface ::AbstractSparseArrayInterface function setunstoredindex!(
+      a::$type, value, i::Int, j::Int
+    )
       setunstoredindex!(parent(a), $func(value), j, i)
       return a
     end

test/basics/test_basics.jl

@@ -32,7 +32,7 @@ arrayts = (Array, JLArray)
   # TODO: We should be specializing these for dense/strided arrays,
   # probably we can have a trait for that. It could be based
   # on the `ArrayLayouts.MemoryLayout`.
-  @allowscalar @test storedvalues(a) == vec(a)
+  @allowscalar @test storedvalues(a) == a
   @allowscalar @test storedpairs(a) == collect(pairs(vec(a)))
   @allowscalar for I in eachindex(a)
     @test getstoredindex(a, I) == a[I]