|
| 1 | +# zero-preserving Traits |
| 2 | +# ---------------------- |
| 3 | +""" |
| 4 | + abstract type ZeroPreserving <: Function end |
| 5 | +
|
| 6 | +Holy Trait to indicate how a function interacts with abstract zero values: |
| 7 | +
|
| 8 | +- `StrongPreserving` : output is guaranteed to be zero if **any** input is. |
| 9 | +- `WeakPreserving` : output is guaranteed to be zero if **all** inputs are. |
| 10 | +- `NonPreserving` : no guarantees on output. |
| 11 | +
|
| 12 | +To attempt to automatically determine this, either `ZeroPreserving(f, A::AbstractArray...)` or |
| 13 | +`ZeroPreserving(f, T::Type...)` can be used/overloaded. |
| 14 | +
|
| 15 | +!!! warning |
| 16 | + incorrectly registering a function to be zero-preserving will lead to silently wrong results. |
| 17 | +""" |
| 18 | +abstract type ZeroPreserving <: Function end |
| 19 | + |
| 20 | +struct StrongPreserving{F} <: ZeroPreserving |
| 21 | + f::F |
| 22 | +end |
| 23 | +struct WeakPreserving{F} <: ZeroPreserving |
| 24 | + f::F |
| 25 | +end |
| 26 | +struct NonPreserving{F} <: ZeroPreserving |
| 27 | + f::F |
| 28 | +end |
| 29 | + |
| 30 | +# Backport: remove in 1.12 |
| 31 | +@static if !isdefined(Base, :haszero) |
| 32 | + _haszero(T::Type) = false |
| 33 | + _haszero(::Type{<:Number}) = true |
| 34 | +else |
| 35 | + _haszero = Base.haszero |
| 36 | +end |
| 37 | + |
| 38 | +# warning: cannot automatically detect WeakPreserving since this would mean checking all values |
| 39 | +function ZeroPreserving(f, A::AbstractArray, Bs::AbstractArray...) |
| 40 | + return ZeroPreserving(f, eltype(A), eltype.(Bs)...) |
| 41 | +end |
| 42 | +# TODO: the following might not properly specialize on the types |
| 43 | +# TODO: non-concrete element types |
| 44 | +function ZeroPreserving(f, T::Type, Ts::Type...) |
| 45 | + if all(_haszero, (T, Ts...)) |
| 46 | + return iszero(f(zero(T), zero.(Ts)...)) ? WeakPreserving(f) : NonPreserving(f) |
| 47 | + else |
| 48 | + return NonPreserving(f) |
| 49 | + end |
| 50 | +end |
| 51 | + |
| 52 | +const _WEAK_FUNCTIONS = (:+, :-) |
| 53 | +for f in _WEAK_FUNCTIONS |
| 54 | + @eval begin |
| 55 | + ZeroPreserving(::typeof($f), ::Type{<:Number}, ::Type{<:Number}...) = WeakPreserving($f) |
| 56 | + end |
| 57 | +end |
| 58 | + |
| 59 | +const _STRONG_FUNCTIONS = (:*,) |
| 60 | +for f in _STRONG_FUNCTIONS |
| 61 | + @eval begin |
| 62 | + ZeroPreserving(::typeof($f), ::Type{<:Number}, ::Type{<:Number}...) = StrongPreserving( |
| 63 | + $f |
| 64 | + ) |
| 65 | + end |
| 66 | +end |
| 67 | + |
| 68 | +# map(!) |
| 69 | +# ------ |
| 70 | +@interface I::AbstractSparseArrayInterface function Base.map( |
| 71 | + f, A::AbstractArray, Bs::AbstractArray... |
| 72 | +) |
| 73 | + f_pres = ZeroPreserving(f, A, Bs...) |
| 74 | + return @interface I map(f_pres, A, Bs...) |
| 75 | +end |
| 76 | +@interface I::AbstractSparseArrayInterface function Base.map( |
| 77 | + f::ZeroPreserving, A::AbstractArray, Bs::AbstractArray... |
| 78 | +) |
| 79 | + T = Base.Broadcast.combine_eltypes(f.f, (A, Bs...)) |
| 80 | + C = similar(I, T, size(A)) |
| 81 | + return @interface I map!(f, C, A, Bs...) |
| 82 | +end |
| 83 | + |
| 84 | +@interface I::AbstractSparseArrayInterface function Base.map!( |
| 85 | + f, C::AbstractArray, A::AbstractArray, Bs::AbstractArray... |
| 86 | +) |
| 87 | + f_pres = ZeroPreserving(f, A, Bs...) |
| 88 | + return @interface I map!(f_pres, C, A, Bs...) |
| 89 | +end |
| 90 | + |
| 91 | +@interface ::AbstractSparseArrayInterface function Base.map!( |
| 92 | + f::ZeroPreserving, C::AbstractArray, A::AbstractArray, Bs::AbstractArray... |
| 93 | +) |
| 94 | + checkshape(C, A, Bs...) |
| 95 | + unaliased = map(Base.Fix1(Base.unalias, C), (A, Bs...)) |
| 96 | + |
| 97 | + if f isa StrongPreserving |
| 98 | + style = IndexStyle(C, unaliased...) |
| 99 | + inds = intersect(eachstoredindex.(Ref(style), unaliased)...) |
| 100 | + zero!(C) |
| 101 | + elseif f isa WeakPreserving |
| 102 | + style = IndexStyle(C, unaliased...) |
| 103 | + inds = union(eachstoredindex.(Ref(style), unaliased)...) |
| 104 | + zero!(C) |
| 105 | + elseif f isa NonPreserving |
| 106 | + inds = eachindex(C, unaliased...) |
| 107 | + else |
| 108 | + error(lazy"unknown zero-preserving type $(typeof(f))") |
| 109 | + end |
| 110 | + |
| 111 | + @inbounds for I in inds |
| 112 | + C[I] = f.f(ith_all(I, unaliased)...) |
| 113 | + end |
| 114 | + |
| 115 | + return C |
| 116 | +end |
| 117 | + |
| 118 | +# Derived functions |
| 119 | +# ----------------- |
| 120 | +@interface I::AbstractSparseArrayInterface Base.copyto!(C::AbstractArray, A::AbstractArray) = @interface I map!( |
| 121 | + identity, C, A |
| 122 | +) |
| 123 | + |
| 124 | +# Only map the stored values of the inputs. |
| 125 | +function map_stored! end |
| 126 | + |
| 127 | +@interface interface::AbstractArrayInterface function map_stored!( |
| 128 | + f, a_dest::AbstractArray, as::AbstractArray... |
| 129 | +) |
| 130 | + @interface interface map!(WeakPreserving(f), a_dest, as...) |
| 131 | + return a_dest |
| 132 | +end |
| 133 | + |
| 134 | +# Only map all values, not just the stored ones. |
| 135 | +function map_all! end |
| 136 | + |
| 137 | +@interface interface::AbstractArrayInterface function map_all!( |
| 138 | + f, a_dest::AbstractArray, as::AbstractArray... |
| 139 | +) |
| 140 | + @interface interface map!(NonPreserving(f), a_dest, as...) |
| 141 | + return a_dest |
| 142 | +end |
| 143 | + |
| 144 | +# Utility functions |
| 145 | +# ----------------- |
| 146 | +# shape check similar to checkbounds |
| 147 | +checkshape(::Type{Bool}, A::AbstractArray) = true |
| 148 | +checkshape(::Type{Bool}, A::AbstractArray, B::AbstractArray) = size(A) == size(B) |
| 149 | +function checkshape(::Type{Bool}, A::AbstractArray, Bs::AbstractArray...) |
| 150 | + return allequal(size, (A, Bs...)) |
| 151 | +end |
| 152 | + |
| 153 | +function checkshape(A::AbstractArray, Bs::AbstractArray...) |
| 154 | + return checkshape(Bool, A, Bs...) || |
| 155 | + throw(DimensionMismatch("argument shapes must match")) |
| 156 | +end |
| 157 | + |
| 158 | +@inline ith_all(i, ::Tuple{}) = () |
| 159 | +function ith_all(i, as) |
| 160 | + @_propagate_inbounds_meta |
| 161 | + return (as[1][i], ith_all(i, Base.tail(as))...) |
| 162 | +end |
0 commit comments