From 945ef97ee0b23f491d0790331541fae67d6adfb0 Mon Sep 17 00:00:00 2001
From: Ryan Bennink <benninkrs@ornl.gov>
Date: Fri, 3 Jan 2020 12:00:12 -0500
Subject: [PATCH] Preliminary support for matrix multiplication. - New methods
 for * for combinations of OffsetArray and AbstractArray   Unfortunately this
 introduces method ambiguities with Base. - As an alternative, tried defining
 a new promote_rule, but somehow it doesn't get invoked.

New convenience constructors
- Wrap any (non-offset) array as an OffsetArray with standard indices
- Remove layer of indirection when constructing with explicit offsets

Cosmetic edits to constructor section for improved readability
---
 src/OffsetArrays.jl | 38 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 36 insertions(+), 2 deletions(-)

diff --git a/src/OffsetArrays.jl b/src/OffsetArrays.jl
index 503b860c..ecaf7dfe 100644
--- a/src/OffsetArrays.jl
+++ b/src/OffsetArrays.jl
@@ -3,32 +3,43 @@ VERSION < v"0.7.0-beta2.199" && __precompile__()
 module OffsetArrays
 
 using Base: Indices, tail, @propagate_inbounds
+import Base: (*), convert, promote_rule
+
 @static if !isdefined(Base, :IdentityUnitRange)
     const IdentityUnitRange = Base.Slice
 else
     using Base: IdentityUnitRange
 end
 
-export OffsetArray, OffsetVector
+export OffsetArray, OffsetVector, OffsetMatrix
 
 struct OffsetArray{T,N,AA<:AbstractArray} <: AbstractArray{T,N}
     parent::AA
     offsets::NTuple{N,Int}
 end
 OffsetVector{T,AA<:AbstractArray} = OffsetArray{T,1,AA}
+OffsetMatrix{T,AA<:AbstractArray} = OffsetArray{T,2,AA}
 
+# Construct from existing array
 OffsetArray(A::AbstractArray{T,N}, offsets::NTuple{N,Int}) where {T,N} =
     OffsetArray{T,N,typeof(A)}(A, offsets)
 OffsetArray(A::AbstractArray{T,N}, offsets::Vararg{Int,N}) where {T,N} =
     OffsetArray(A, offsets)
+OffsetArray(A::AbstractArray{T,0}) where {T} = OffsetArray{T,0,typeof(A)}(A, ())
+OffsetArray(A::AbstractArray) = convert(OffsetArray, A)
+OffsetArray(A::OffsetArray) = A
+
+# Convert from existing array
+convert(::Type{OffsetArray}, A::AbstractArray{T,N}) where {T,N} = OffsetArray(A, ntuple(x->0, N))
 
+# Create an uninitialized OffsetArray with given element type
 const ArrayInitializer = Union{UndefInitializer, Missing, Nothing}
 OffsetArray{T,N}(init::ArrayInitializer, inds::Indices{N}) where {T,N} =
     OffsetArray{T,N,Array{T,N}}(Array{T,N}(init, map(indexlength, inds)), map(indexoffset, inds))
 OffsetArray{T}(init::ArrayInitializer, inds::Indices{N}) where {T,N} = OffsetArray{T,N}(init, inds)
+# Same thing, but taking multiple args for offsets/indices
 OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = OffsetArray{T,N}(init, inds)
 OffsetArray{T}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = OffsetArray{T,N}(init, inds)
-OffsetArray(A::AbstractArray{T,0}) where {T} = OffsetArray{T,0,typeof(A)}(A, ())
 
 # OffsetVector constructors
 OffsetVector(A::AbstractVector, offset) = OffsetArray(A, offset)
@@ -59,6 +70,7 @@ OffsetArray(A::AbstractArray{T,N}, inds::Vararg{AbstractUnitRange,N}) where {T,N
     OffsetArray(A, inds)
 
 # avoid a level of indirection when nesting OffsetArrays
+OffsetArray(A::OffsetArray, offsets::NTuple{N,Int}) where {N} = OffsetArray(parent(A), offsets)
 function OffsetArray(A::OffsetArray, inds::NTuple{N,AbstractUnitRange}) where {N}
     OffsetArray(parent(A), inds)
 end
@@ -283,6 +295,28 @@ end
 no_offset_view(A::OffsetArray) = no_offset_view(parent(A))
 
 
+# Quick hack for matrix multiplication.
+# Ideally, one would instead improve LinearAlgebra's support of custom indexing.
+function (*)(A::OffsetMatrix, B::OffsetMatrix)
+    matmult_check_axes(A, B)
+    C = OffsetArray(parent(A) * parent(B), (axes(A,1), axes(B,2)))
+end
+
+function (*)(A::OffsetMatrix, B::OffsetVector)
+    matmult_check_axes(A, B)
+    C = OffsetArray(parent(A) * parent(B), axes(A,1))
+end
+matmult_check_axes(A, B) = axes(A, 2) == axes(B, 1) || error("axes(A,2) must equal axes(B,1)")
+
+(*)(A::OffsetMatrix, B::AbstractMatrix) = A * OffsetArray(B)
+(*)(A::OffsetMatrix, B::AbstractVector) = A * OffsetArray(B)
+(*)(A::AbstractMatrix, B::OffsetArray) = OffsetArray(A) * B
+(*)(A::AbstractVector, B::OffsetArray) = OffsetArray(A) * B
+
+# An alternative to the above four methods would be to use promote_rule, but it doesn't get invoked
+# promote_rule(::Type{A1}, ::Type{A2}) where A1<:AbstractArray{<:Any,N}  where A2<:OffsetArray{<:Any,N,A3} where {N,A3} = OffsetArray{eltype(promote_type(A1, A3)), N, promote_type(A1, A3)}
+
+
 ####
 # work around for segfault in searchsorted*
 #  https://github.com/JuliaLang/julia/issues/33977