Refactor aliasing logic

Author: mtfishman

src/abstractsparsearrayinterface.jl

@@ -207,7 +207,7 @@ end
 function map_indices! end
 
 @interface interface::AbstractArrayInterface function map_indices!(
-  f, a_dest::AbstractArray, indices, as::AbstractArray...
+  indices, f, a_dest::AbstractArray, as::AbstractArray...
 )
   for I in indices
     a_dest[I] = f(map(a -> a[I], as)...)
@@ -221,7 +221,7 @@ function map_stored! end
 @interface interface::AbstractArrayInterface function map_stored!(
   f, a_dest::AbstractArray, as::AbstractArray...
 )
-  @interface interface map_indices!(f, a_dest, eachstoredindex(as...), as...)
+  @interface interface map_indices!(eachstoredindex(as...), f, a_dest, as...)
   return a_dest
 end
 
@@ -231,7 +231,7 @@ function map_all! end
 @interface interface::AbstractArrayInterface function map_all!(
   f, a_dest::AbstractArray, as::AbstractArray...
 )
-  @interface interface map_indices!(f, a_dest, eachindex(as...), as...)
+  @interface interface map_indices!(eachindex(as...), f, a_dest, as...)
   return a_dest
 end
 
@@ -250,37 +250,32 @@ using ArrayLayouts: ArrayLayouts, zero!
   return @interface interface map_stored!(f, a, a)
 end
 
+# Determines if a function preserves the stored values
+# of the destination sparse array.
+# The current code may be inefficient since it actually
+# accesses an unstored element, which in the case of a
+# sparse array of arrays can allocate an array.
+# Sparse arrays could be expected to define a cheap
+# unstored element allocator, for example
+# `get_prototypical_unstored(a::AbstractArray)`.
+function preserves_unstored(f, a_dest::AbstractArray, as::AbstractArray...)
+  I = first(eachindex(as...))
+  return iszero(f(map(a -> getunstoredindex(a, I), as)...))
+end
+
 @interface interface::AbstractSparseArrayInterface function Base.map!(
   f, a_dest::AbstractArray, as::AbstractArray...
 )
-  # TODO: Define a function `preserves_unstored(a_dest, f, as...)`
-  # to determine if a function preserves the stored values
-  # of the destination sparse array.
-  # The current code may be inefficient since it actually
-  # accesses an unstored element, which in the case of a
-  # sparse array of arrays can allocate an array.
-  # Sparse arrays could be expected to define a cheap
-  # unstored element allocator, for example
-  # `get_prototypical_unstored(a::AbstractArray)`.
-  I = first(eachindex(as...))
-  preserves_unstored = iszero(f(map(a -> getunstoredindex(a, I), as)...))
-  if !preserves_unstored
-    # Doesn't preserve unstored values, loop over all elements.
-    @interface interface map_all!(f, a_dest, as...)
-    return a_dest
-  end
-  # Unalias the inputs from the destination
-  # to make sure the inputs aren't overwritten incorrectly.
-  # See: https://github.com/JuliaLang/julia/blob/v1.11.2/base/broadcast.jl#L935-L948
-  as = map(a -> a_dest === a ? a : Base.unalias(a_dest, a), as)
-  indices_stored = eachstoredindex(as...)
-  if eachstoredindex(a_dest) ⊈ indices_stored
-    # If not all indices being mapped over are stored in the destination,
-    # zero out the destination. An extreme example of this is when
-    # the sources are sparse but the destination is dense.
+  indices = if !preserves_unstored(f, a_dest, as...)
+    eachindex(a_dest)
+  elseif any(a -> a_dest !== a, as)
+    as = map(a -> Base.unalias(a_dest, a), as)
     @interface interface zero!(a_dest)
+    eachstoredindex(as...)
+  else
+    eachstoredindex(a_dest)
   end
-  @interface interface map_indices!(f, a_dest, indices_stored, as...)
+  @interface interface map_indices!(indices, f, a_dest, as...)
   return a_dest
 end