some tweaks

Author: lkdvos

src/abstractsparsearray.jl

@@ -18,21 +18,19 @@ using LinearAlgebra: LinearAlgebra
 const MIMEtextplain = MIME"text/plain"
 
 @derive (T=AnyAbstractSparseArray,) begin
-  Base.getindex(::T, ::Any...)
   Base.getindex(::T, ::Int...)
-  Base.setindex!(::T, ::Any, ::Any...)
   Base.setindex!(::T, ::Any, ::Int...)
   Base.similar(::T, ::Type, ::Tuple{Vararg{Int}})
   Base.similar(::T, ::Type, ::Tuple{Base.OneTo,Vararg{Base.OneTo}})
-  Base.copy(::T)
-  Base.copy!(::AbstractArray, ::T)
-  Base.copyto!(::AbstractArray, ::T)
+  # Base.copy(::T)
+  # Base.copy!(::AbstractArray, ::T)
+  # Base.copyto!(::AbstractArray, ::T)
   Base.map(::Any, ::T...)
   Base.map!(::Any, ::AbstractArray, ::T...)
-  Base.mapreduce(::Any, ::Any, ::T...; kwargs...)
-  Base.reduce(::Any, ::T...; kwargs...)
-  Base.all(::Function, ::T)
-  Base.all(::T)
+  # Base.mapreduce(::Any, ::Any, ::T...; kwargs...)
+  # Base.reduce(::Any, ::T...; kwargs...)
+  # Base.all(::Function, ::T)
+  # Base.all(::T)
   Base.iszero(::T)
   Base.real(::T)
   Base.fill!(::T, ::Any)
@@ -44,7 +42,7 @@ const MIMEtextplain = MIME"text/plain"
   Base.cat(::T...; kwargs...)
   ArrayLayouts.MemoryLayout(::Type{<:T})
   LinearAlgebra.mul!(::AbstractMatrix, ::T, ::T, ::Number, ::Number)
-  Base.show(::IO, ::MIMEtextplain, ::T)
+  # Base.show(::IO, ::MIMEtextplain, ::T)
 end
 
 function Base.replace_in_print_matrix(

src/abstractsparsearrayinterface.jl

@@ -104,11 +104,10 @@ to_vec(x::AbstractArray) = vec(x)
 
 # TODO: This may need to be defined in `sparsearraydok.jl`, after `SparseArrayDOK`
 # is defined. And/or define `default_type(::SparseArrayStyle, T::Type) = SparseArrayDOK{T}`.
-@interface ::AbstractSparseArrayInterface function Base.similar(
-  a::AbstractArray, T::Type, size::Tuple{Vararg{Int}}
-)
-  # TODO: Define `default_similartype` or something like that?
-  return SparseArrayDOK{T}(size...)
+@interface I::AbstractSparseArrayInterface function Base.similar(
+  ::AbstractArray, ::Type{T}, ax
+) where {T}
+  return similar(I, T, ax)
 end
 
 using ArrayLayouts: ArrayLayouts, zero!
@@ -117,13 +116,11 @@ using ArrayLayouts: ArrayLayouts, zero!
 # and is useful for sparse array logic, since it can be used to empty
 # the sparse array storage.
 # We use a single function definition to minimize method ambiguities.
-@interface interface::AbstractSparseArrayInterface function ArrayLayouts.zero!(
-  a::AbstractArray
+@interface interface::AbstractSparseArrayInterface function DerivableInterfaces.zero!(
+  A::AbstractArray
 )
-  # More generally, this codepath could be taking if `zero(eltype(a))`
-  # is defined and the elements are immutable.
-  f = eltype(a) <: Number ? Returns(zero(eltype(a))) : zero!
-  return @interface interface map_stored!(f, a, a)
+  storedvalues(A) .= zero!(storedvalues(A))
+  return A
 end
 
 # `f::typeof(norm)`, `op::typeof(max)` used by `norm`.
@@ -150,23 +147,6 @@ end
   return output
 end
 
-abstract type AbstractSparseArrayStyle{N} <: Broadcast.AbstractArrayStyle{N} end
-
-@derive (T=AbstractSparseArrayStyle,) begin
-  Base.similar(::Broadcast.Broadcasted{<:T}, ::Type, ::Tuple)
-  Base.copyto!(::AbstractArray, ::Broadcast.Broadcasted{<:T})
-end
-
-struct SparseArrayStyle{N} <: AbstractSparseArrayStyle{N} end
-
-SparseArrayStyle{M}(::Val{N}) where {M,N} = SparseArrayStyle{N}()
-
-DerivableInterfaces.interface(::Type{<:AbstractSparseArrayStyle}) = SparseArrayInterface()
-
-@interface ::AbstractSparseArrayInterface function Broadcast.BroadcastStyle(type::Type)
-  return SparseArrayStyle{ndims(type)}()
-end
-
 using ArrayLayouts: ArrayLayouts, MatMulMatAdd
 
 abstract type AbstractSparseLayout <: ArrayLayouts.MemoryLayout end
@@ -190,19 +170,20 @@ using LinearAlgebra: LinearAlgebra, mul!
 @interface ::AbstractSparseArrayInterface function LinearAlgebra.mul!(
   C::AbstractArray, A::AbstractArray, B::AbstractArray, α::Number, β::Number
 )
-  a_dest .*= β
+  C .*= β
   β′ = one(Bool)
-  for I1 in eachstoredindex(a1)
-    for I2 in eachstoredindex(a2)
-      I_dest = mul_indices(I1, I2)
-      if !isnothing(I_dest)
-        a_dest[I_dest] = mul!(a_dest[I_dest], a1[I1], a2[I2], α, β′)
-      end
+  for iA in eachstoredindex(A), iB in eachstoredindex(B)
+    iC = mul_indices(iA, iB)
+    if !isnothing(iC)
+      C[iC] = mul!!(C[iC], A[iA], B[iB], α, β′)
     end
   end
-  return a_dest
+  return C
 end
 
+mul!!(C, A, B, α, β) = mul!(C, A, B, α, β)
+mul!!(C::Number, A::Number, B::Number, α::Number, β::Number) = β * C + α * A * B
+
 function ArrayLayouts.materialize!(
   m::MatMulMatAdd{<:AbstractSparseLayout,<:AbstractSparseLayout,<:AbstractSparseLayout}
 )

src/sparsearraydok.jl

@@ -35,6 +35,10 @@ function SparseArrayDOK{T}(size::Int...) where {T}
   return SparseArrayDOK{T,length(size)}(size...)
 end
 
+function SparseArrayDOK{T}(::UndefInitializer, axes::Tuple) where {T}
+  return SparseArrayDOK{T}(undef, Base.to_shape(axes))
+end
+
 using DerivableInterfaces: @array_aliases
 # Define `SparseMatrixDOK`, `AnySparseArrayDOK`, etc.
 @array_aliases SparseArrayDOK
@@ -46,7 +50,7 @@ storedvalues(a::SparseArrayDOK) = values(storage(a))
 function isstored(a::SparseArrayDOK, I::Int...)
   return CartesianIndex(I) in keys(storage(a))
 end
-function eachstoredindex(a::SparseArrayDOK)
+function eachstoredindex(::IndexCartesian, a::SparseArrayDOK)
   return keys(storage(a))
 end
 function getstoredindex(a::SparseArrayDOK, I::Int...)