[NDTensors] Unwrap module for dispatching on unwrapped types

NDTensors/src/NDTensors.jl

@@ -28,7 +28,8 @@ include("SortedSets/src/SortedSets.jl")
 using .SortedSets
 include("TagSets/src/TagSets.jl")
 using .TagSets
-
+include("Unwrap/src/Unwrap.jl")
+using .Unwrap
 using Base: @propagate_inbounds, ReshapedArray, DimOrInd, OneTo
 
 using Base.Cartesian: @nexprs

NDTensors/src/Unwrap/README.md

@@ -0,0 +1,3 @@
+# Unwrap
+
+A module to unwrap complex array types to assist in the generic programming of array-type based functions.

NDTensors/src/Unwrap/TODO.md

@@ -0,0 +1 @@
+Replace all `leaf_parenttype` calls by wrapping the arrays in this `expose` type

NDTensors/src/Unwrap/src/Unwrap.jl

@@ -0,0 +1,16 @@
+module Unwrap
+using SimpleTraits
+using LinearAlgebra
+using Base: ReshapedArray
+using Strided: StridedView
+
+include("iswrappedarray.jl")
+include("expose.jl")
+
+export IsWrappedArray, is_wrapped_array, parenttype, unwrap_type, expose
+
+## TODO Create functions which take the `Expose` type and launch functions
+## using that type
+## TODO write exposed based functions in the NDTensors Extensions when necessary
+
+end

NDTensors/src/Unwrap/src/expose.jl

@@ -0,0 +1,5 @@
+struct Expose{Unwraped,Object}
+  object::Object
+end
+
+expose(object) = Expose{unwrap_type(object),typeof(object)}(object)

NDTensors/src/Unwrap/src/iswrappedarray.jl

@@ -0,0 +1,57 @@
+# Trait indicating if the AbstractArray type is an array wrapper.
+# Assumes that it implements `NDTensors.parenttype`.
+@traitdef IsWrappedArray{ArrayT}
+
+#! format: off
+@traitimpl IsWrappedArray{ArrayT} <- is_wrapped_array(ArrayT)
+#! format: on
+
+is_wrapped_array(arraytype::Type{<:AbstractArray}) = (parenttype(arraytype) ≠ arraytype)
+
+# TODO: This is only defined because the current design
+# of `Diag` using a `Number` as the data type if it
+# is a uniform diagonal type. Delete this when it is
+# replaced by `DiagonalArray`.
+is_wrapped_array(arraytype::Type{<:Number}) = false
+
+# For working with instances, not used by
+# `SimpleTraits.jl` traits dispatch.
+is_wrapped_array(array::AbstractArray) = is_wrapped_array(typeof(array))
+
+# By default, the `parentype` of an array type is itself
+parenttype(arraytype::Type{<:AbstractArray}) = arraytype
+
+# TODO: Use `SetParameters` here.
+parenttype(::Type{<:Base.ReshapedArray{<:Any,<:Any,P}}) where {P} = P
+parenttype(::Type{<:Transpose{<:Any,P}}) where {P} = P
+parenttype(::Type{<:Adjoint{<:Any,P}}) where {P} = P
+parenttype(::Type{<:Symmetric{<:Any,P}}) where {P} = P
+parenttype(::Type{<:Hermitian{<:Any,P}}) where {P} = P
+parenttype(::Type{<:UpperTriangular{<:Any,P}}) where {P} = P
+parenttype(::Type{<:LowerTriangular{<:Any,P}}) where {P} = P
+parenttype(::Type{<:UnitUpperTriangular{<:Any,P}}) where {P} = P
+parenttype(::Type{<:UnitLowerTriangular{<:Any,P}}) where {P} = P
+parenttype(::Type{<:Diagonal{<:Any,P}}) where {P} = P
+parenttype(::Type{<:SubArray{<:Any,<:Any,P}}) where {P} = P
+parenttype(::Type{<:StridedView{<:Any,<:Any,P}}) where {P} = P
+
+# For working with instances, not used by
+# `SimpleTraits.jl` traits dispatch.
+parenttype(array::AbstractArray) = parenttype(typeof(array))
+
+## These functions will be used in place of leaf_parenttype but will be 
+## call indirectly through the expose function.
+@traitfn function unwrap_type(
+  arraytype::Type{ArrayT}
+) where {ArrayT; IsWrappedArray{ArrayT}}
+  return leaf_parenttype(parenttype(arraytype))
+end
+
+@traitfn function unwrap_type(
+  arraytype::Type{ArrayT}
+) where {ArrayT; !IsWrappedArray{ArrayT}}
+  return arraytype
+end
+
+# For working with instances.
+unwrap_type(array::AbstractArray) = unwrap_type(typeof(array))

NDTensors/src/abstractarray/iswrappedarray.jl

@@ -1,44 +1,3 @@
-# Trait indicating if the AbstractArray type is an array wrapper.
-# Assumes that it implements `NDTensors.parenttype`.
-@traitdef IsWrappedArray{ArrayT}
-
-#! format: off
-@traitimpl IsWrappedArray{ArrayT} <- is_wrapped_array(ArrayT)
-#! format: on
-
-is_wrapped_array(arraytype::Type{<:AbstractArray}) = (parenttype(arraytype) ≠ arraytype)
-
-# TODO: This is only defined because the current design
-# of `Diag` using a `Number` as the data type if it
-# is a uniform diagonal type. Delete this when it is
-# replaced by `DiagonalArray`.
-is_wrapped_array(arraytype::Type{<:Number}) = false
-
-# For working with instances, not used by
-# `SimpleTraits.jl` traits dispatch.
-is_wrapped_array(array::AbstractArray) = is_wrapped_array(typeof(array))
-
-# By default, the `parentype` of an array type is itself
-parenttype(arraytype::Type{<:AbstractArray}) = arraytype
-
-# TODO: Use `SetParameters` here.
-parenttype(::Type{<:ReshapedArray{<:Any,<:Any,P}}) where {P} = P
-parenttype(::Type{<:Transpose{<:Any,P}}) where {P} = P
-parenttype(::Type{<:Adjoint{<:Any,P}}) where {P} = P
-parenttype(::Type{<:Symmetric{<:Any,P}}) where {P} = P
-parenttype(::Type{<:Hermitian{<:Any,P}}) where {P} = P
-parenttype(::Type{<:UpperTriangular{<:Any,P}}) where {P} = P
-parenttype(::Type{<:LowerTriangular{<:Any,P}}) where {P} = P
-parenttype(::Type{<:UnitUpperTriangular{<:Any,P}}) where {P} = P
-parenttype(::Type{<:UnitLowerTriangular{<:Any,P}}) where {P} = P
-parenttype(::Type{<:Diagonal{<:Any,P}}) where {P} = P
-parenttype(::Type{<:SubArray{<:Any,<:Any,P}}) where {P} = P
-parenttype(::Type{<:StridedView{<:Any,<:Any,P}}) where {P} = P
-
-# For working with instances, not used by
-# `SimpleTraits.jl` traits dispatch.
-parenttype(array::AbstractArray) = parenttype(typeof(array))
-
 @traitfn function leaf_parenttype(
   arraytype::Type{ArrayT}
 ) where {ArrayT; IsWrappedArray{ArrayT}}