Broadcasting performance improvements

Author: jonniedie

.gitignore

@@ -10,6 +10,6 @@ research/
 /examples/wip
 /examples/ignore
 /examples/.ipynb_checkpoints
-.vscode/settings.json
+.vscode
 /issue_debug
 Manifest.toml

src/array_interface.jl

@@ -59,14 +59,14 @@ Base.@propagate_inbounds function Base.getindex(x::ComponentArray, idx::FlatOrCo
     return ComponentArray(getdata(x)[idx...], axs...)
 end
 Base.@propagate_inbounds Base.getindex(x::ComponentArray, ::Colon) = getdata(x)[:]
-@inline Base.getindex(x::ComponentArray, ::Colon...) = x
-Base.@propagate_inbounds Base.getindex(x::ComponentArray, idx...) = getindex(x, toval.(idx)...)
+Base.@propagate_inbounds Base.getindex(x::ComponentArray, ::Colon...) = x
+@inline Base.getindex(x::ComponentArray, idx...) = getindex(x, toval.(idx)...)
 @inline Base.getindex(x::ComponentArray, idx::Val...) = _getindex(x, idx...)
 
 # Set ComponentArray index
-@inline Base.setindex!(x::ComponentArray, v, idx::FlatIdx...) = setindex!(getdata(x), v, idx...)
+Base.@propagate_inbounds Base.setindex!(x::ComponentArray, v, idx::FlatOrColonIdx...) = setindex!(getdata(x), v, idx...)
 Base.@propagate_inbounds Base.setindex!(x::ComponentArray, v, ::Colon) = setindex!(getdata(x), v, :)
-Base.@propagate_inbounds Base.setindex!(x::ComponentArray, v, idx...) = setindex!(x, v, toval.(idx)...)
+@inline Base.setindex!(x::ComponentArray, v, idx...) = setindex!(x, v, toval.(idx)...)
 @inline Base.setindex!(x::ComponentArray, v, idx::Val...) = _setindex!(x, v, idx...)
 
 # Explicitly view

src/broadcasting.jl

@@ -69,7 +69,7 @@ function Base.similar(bc::BC.Broadcasted{<:CAStyle{<:BC.Unknown, Axes, N}}, T::T
 end
 
 
-BC.broadcasted(f, x::ComponentArray) = ComponentArray(map(f, getdata(x)), getaxes(x))
+# BC.broadcasted(f, x::ComponentArray) = ComponentArray(map(f, getdata(x)), getaxes(x))
 
 # Need a special case here because `map` doesn't follow same rules as normal broadcasting. To be safe and avoid ambiguities,
 # we'll just handle the case where everything is a ComponentArray. Else it falls back to a plain Array output.

src/similar_convert_copy.jl

@@ -2,14 +2,21 @@
 # Similar
 Base.similar(x::ComponentArray) = ComponentArray(similar(getdata(x)), getaxes(x)...)
 Base.similar(x::ComponentArray, ::Type{T}) where T = ComponentArray(similar(getdata(x), T), getaxes(x)...)
-Base.similar(x::ComponentArray, ::Type{T}, ax::Tuple{Vararg{Int64,N}}) where {T,N} = similar(x, T, ax...)
-function Base.similar(x::ComponentArray, ::Type{T}, ax::Union{Integer, Base.OneTo}...) where T
-    A = similar(getdata(x), T, ax...)
-    if size(getdata(x)) == size(A)
-        return ComponentArray(A, getaxes(x))
-    else
-        return A
-    end
+# Base.similar(x::ComponentArray, ::Type{T}, ax::Tuple{Vararg{Int64,N}}) where {T,N} = similar(x, T, ax...)
+# function Base.similar(x::ComponentArray, ::Type{T}, ax::Union{Integer, Base.OneTo}...) where T
+#     A = similar(getdata(x), T, ax...)
+#     if size(getdata(x)) == size(A)
+#         return ComponentArray(A, getaxes(x))
+#     else
+#         return A
+#     end
+# end
+function Base.similar(x::ComponentArray{T1,N,A,Ax}, ::Type{T}, dims::NTuple{N,Int}) where {T,T1,N,A,Ax}
+    arr = similar(getdata(x), T, dims)
+    return ComponentArray(arr, getaxes(x))
+end
+function Base.similar(x::ComponentArray{T1,N1,A,Ax}, ::Type{T}, dims::NTuple{N2,Int}) where {T,T1,N1,N2,A,Ax}
+    return similar(getdata(x), T, dims)
 end
 
 ## TODO: write length method for AbstractAxis so we can do this?