Skip to content

Commit c3b5f88

Browse files
Remove more code duplication in DataLayouts
1 parent 1db4997 commit c3b5f88

File tree

1 file changed

+47
-185
lines changed

1 file changed

+47
-185
lines changed

src/DataLayouts/DataLayouts.jl

Lines changed: 47 additions & 185 deletions
Original file line numberDiff line numberDiff line change
@@ -361,31 +361,6 @@ function IJFH{S, Nij, Nh}(::Type{ArrayType}) where {S, Nij, Nh, ArrayType}
361361
IJFH{S, Nij, Nh}(ArrayType(undef, Nij, Nij, typesize(T, S), Nh))
362362
end
363363

364-
@propagate_inbounds function Base.getindex(
365-
data::IJFH{S},
366-
I::CartesianIndex{5},
367-
) where {S}
368-
@inbounds get_struct(
369-
parent(data),
370-
S,
371-
Val(field_dim(data)),
372-
to_data_specific(data, I),
373-
)
374-
end
375-
@propagate_inbounds function Base.setindex!(
376-
data::IJFH{S},
377-
val,
378-
I::CartesianIndex{5},
379-
) where {S}
380-
@inbounds set_struct!(
381-
parent(data),
382-
convert(S, val),
383-
Val(field_dim(data)),
384-
to_data_specific(data, I),
385-
)
386-
end
387-
388-
389364
Base.length(data::IJFH) = get_Nh(data)
390365

391366
@inline function slab(data::IJFH{S, Nij}, h::Integer) where {S, Nij}
@@ -477,27 +452,6 @@ end
477452
Base.@propagate_inbounds column(data::IFH{S, Ni}, i, j, h) where {S, Ni} =
478453
column(data, i, h)
479454

480-
@inline function Base.getindex(data::IFH{S}, I::CartesianIndex{5}) where {S}
481-
@inbounds get_struct(
482-
parent(data),
483-
S,
484-
Val(field_dim(data)),
485-
to_data_specific(data, I),
486-
)
487-
end
488-
@inline function Base.setindex!(
489-
data::IFH{S},
490-
val,
491-
I::CartesianIndex{5},
492-
) where {S}
493-
@inbounds set_struct!(
494-
parent(data),
495-
convert(S, val),
496-
Val(field_dim(data)),
497-
to_data_specific(data, I),
498-
)
499-
end
500-
501455
# ======================
502456
# Data0D DataLayout
503457
# ======================
@@ -619,39 +573,6 @@ end
619573

620574
@inline universal_size(::IJF{S, Nij}) where {S, Nij} = (Nij, Nij, 1, 1, 1)
621575

622-
@inline function Base.getindex(
623-
data::IJF{S, Nij},
624-
I::CartesianIndex,
625-
) where {S, Nij}
626-
i = I.I[1]
627-
j = I.I[2]
628-
@boundscheck (1 <= i <= Nij && 1 <= j <= Nij) ||
629-
throw(BoundsError(data, (i, j)))
630-
@inbounds get_struct(
631-
parent(data),
632-
S,
633-
Val(field_dim(data)),
634-
CartesianIndex(i, j, 1),
635-
)
636-
end
637-
638-
@inline function Base.setindex!(
639-
data::IJF{S, Nij},
640-
val,
641-
I::CartesianIndex,
642-
) where {S, Nij}
643-
i = I.I[1]
644-
j = I.I[2]
645-
@boundscheck (1 <= i <= Nij && 1 <= j <= Nij) ||
646-
throw(BoundsError(data, (i, j)))
647-
@inbounds set_struct!(
648-
parent(data),
649-
convert(S, val),
650-
Val(field_dim(data)),
651-
CartesianIndex(i, j, 1),
652-
)
653-
end
654-
655576
@inline function column(data::IJF{S, Nij}, i, j) where {S, Nij}
656577
@boundscheck (1 <= j <= Nij && 1 <= i <= Nij) ||
657578
throw(BoundsError(data, (i, j)))
@@ -705,32 +626,6 @@ function SArray(data::IF{S, Ni, <:MArray}) where {S, Ni}
705626
IF{S, Ni}(SArray(parent(data)))
706627
end
707628

708-
@inline function Base.getindex(data::IF{S, Ni}, I::CartesianIndex) where {S, Ni}
709-
i = I.I[1]
710-
@boundscheck (1 <= i <= Ni) || throw(BoundsError(data, (i,)))
711-
@inbounds get_struct(
712-
parent(data),
713-
S,
714-
Val(field_dim(data)),
715-
CartesianIndex(i, 1),
716-
)
717-
end
718-
719-
@inline function Base.setindex!(
720-
data::IF{S, Ni},
721-
val,
722-
I::CartesianIndex,
723-
) where {S, Ni}
724-
i = I.I[1]
725-
@boundscheck (1 <= i <= Ni) || throw(BoundsError(data, (i,)))
726-
@inbounds set_struct!(
727-
parent(data),
728-
convert(S, val),
729-
Val(field_dim(data)),
730-
CartesianIndex(i, 1),
731-
)
732-
end
733-
734629
@inline function column(data::IF{S, Ni}, i) where {S, Ni}
735630
@boundscheck (1 <= i <= Ni) || throw(BoundsError(data, (i,)))
736631
dataview = @inbounds view(parent(data), i, :)
@@ -783,28 +678,6 @@ nlevels(::VF{S, Nv}) where {S, Nv} = Nv
783678
Base.@propagate_inbounds Base.getproperty(data::VF, i::Integer) =
784679
_property_view(data, Val(i))
785680

786-
@inline function Base.getindex(data::VF{S, Nv}, I::CartesianIndex) where {S, Nv}
787-
v = I.I[4]
788-
@boundscheck 1 <= v <= nlevels(data) || throw(BoundsError(data, (v,)))
789-
@inbounds get_struct(
790-
parent(data),
791-
S,
792-
Val(field_dim(data)),
793-
CartesianIndex(v, 1),
794-
)
795-
end
796-
797-
@inline function Base.setindex!(data::VF{S}, val, I::CartesianIndex) where {S}
798-
v = I.I[4]
799-
@boundscheck (1 <= v <= nlevels(data)) || throw(BoundsError(data, (v,)))
800-
@inbounds set_struct!(
801-
parent(data),
802-
convert(S, val),
803-
Val(field_dim(data)),
804-
CartesianIndex(v, 1),
805-
)
806-
end
807-
808681
@inline function column(data::VF, i, h)
809682
@boundscheck (i >= 1 && h >= 1) || throw(BoundsError(data, (i, h)))
810683
data
@@ -899,31 +772,6 @@ end
899772
IJFH{S, Nij, Nh}(dataview)
900773
end
901774

902-
@propagate_inbounds function Base.getindex(
903-
data::VIJFH{S},
904-
I::CartesianIndex{5},
905-
) where {S}
906-
@inbounds get_struct(
907-
parent(data),
908-
S,
909-
Val(field_dim(data)),
910-
to_data_specific(data, I),
911-
)
912-
end
913-
914-
@propagate_inbounds function Base.setindex!(
915-
data::VIJFH{S},
916-
val,
917-
I::CartesianIndex{5},
918-
) where {S}
919-
@inbounds set_struct!(
920-
parent(data),
921-
convert(S, val),
922-
Val(field_dim(data)),
923-
to_data_specific(data, I),
924-
)
925-
end
926-
927775
function gather(
928776
ctx::ClimaComms.AbstractCommsContext,
929777
data::VIJFH{S, Nv, Nij},
@@ -1021,33 +869,6 @@ end
1021869
IFH{S, Nij, Nh}(dataview)
1022870
end
1023871

1024-
@propagate_inbounds function Base.getindex(
1025-
data::VIFH{S},
1026-
I::CartesianIndex{5},
1027-
) where {S}
1028-
i, _, _, v, h = I.I
1029-
@inbounds get_struct(
1030-
parent(data),
1031-
S,
1032-
Val(field_dim(data)),
1033-
CartesianIndex(v, i, 1, h),
1034-
)
1035-
end
1036-
1037-
@inline function Base.setindex!(
1038-
data::VIFH{S},
1039-
val,
1040-
I::CartesianIndex{5},
1041-
) where {S}
1042-
i, _, _, v, h = I.I
1043-
@inbounds set_struct!(
1044-
parent(data),
1045-
convert(S, val),
1046-
Val(field_dim(data)),
1047-
CartesianIndex(v, i, 1, h),
1048-
)
1049-
end
1050-
1051872
# =========================================
1052873
# Special DataLayouts for regular gridding
1053874
# =========================================
@@ -1190,12 +1011,26 @@ type parameters.
11901011
@inline field_dim(::Type{<:VIJFH}) = 4
11911012
@inline field_dim(::Type{<:VIFH}) = 3
11921013

1193-
@inline to_data_specific(::IJF, I::CartesianIndex) = CartesianIndex(I.I[1], I.I[2], 1, 1)
1194-
@inline to_data_specific(::IJFH, I::CartesianIndex) = CartesianIndex(I.I[1], I.I[2], 1, I.I[5])
1195-
@inline to_data_specific(::IFH, I::CartesianIndex) = CartesianIndex(I.I[1], 1, I.I[5])
1196-
@inline to_data_specific(::VIJFH, I::CartesianIndex) = CartesianIndex(I.I[4], I.I[1], I.I[2], 1, I.I[5])
1197-
@inline to_data_specific(::VIFH, I::CartesianIndex) = CartesianIndex(I.I[4], I.I[1], 1, I.I[5])
1198-
@inline to_data_specific(::DataSlab1D, I::CartesianIndex) = CartesianIndex(I.I[1], I.I[1], 1, I.I[5])
1014+
@inline to_data_specific(data::AbstractData, I::CartesianIndex) =
1015+
CartesianIndex(_to_data_specific(data, I.I))
1016+
@inline _to_data_specific(::VF, I::Tuple) = (I[4], 1)
1017+
@inline _to_data_specific(::IF, I::Tuple) = (I[1], 1)
1018+
@inline _to_data_specific(::IJF, I::Tuple) = (I[1], I[2], 1)
1019+
@inline _to_data_specific(::IJFH, I::Tuple) = (I[1], I[2], 1, I[5])
1020+
@inline _to_data_specific(::IFH, I::Tuple) = (I[1], 1, I[5])
1021+
@inline _to_data_specific(::VIJFH, I::Tuple) = (I[4], I[1], I[2], 1, I[5])
1022+
@inline _to_data_specific(::VIFH, I::Tuple) = (I[4], I[1], 1, I[5])
1023+
1024+
"""
1025+
bounds_condition(data::AbstractData, I::Tuple)
1026+
1027+
Returns the condition used for `@boundscheck`
1028+
inside `getindex` with `CartesianIndex`s.
1029+
"""
1030+
@inline bounds_condition(data::AbstractData, I::CartesianIndex) = true # TODO: add more support
1031+
@inline bounds_condition(data::IJF, I::CartesianIndex) = (1 <= I.I[1] <= get_Nij(data) && 1 <= I.I[2] <= get_Nij(data))
1032+
@inline bounds_condition(data::VF, I::CartesianIndex) = 1 <= I.I[4] <= nlevels(data)
1033+
@inline bounds_condition(data::IF, I::CartesianIndex) = 1 <= I.I[1] <= get_Nij(data)
11991034

12001035
"""
12011036
type_params(data::AbstractData)
@@ -1331,6 +1166,33 @@ Base.ndims(data::AbstractData) = Base.ndims(typeof(data))
13311166
Base.ndims(::Type{T}) where {T <: AbstractData} =
13321167
Base.ndims(parent_array_type(T))
13331168

1169+
@inline function Base.getindex(
1170+
data::Union{IJF, IJFH, IFH, VIJFH, VIFH, VF, IF},
1171+
I::CartesianIndex,
1172+
)
1173+
@boundscheck bounds_condition(data, I) || throw(BoundsError(data, I))
1174+
@inbounds get_struct(
1175+
parent(data),
1176+
eltype(data),
1177+
Val(field_dim(data)),
1178+
to_data_specific(data, I),
1179+
)
1180+
end
1181+
1182+
@inline function Base.setindex!(
1183+
data::Union{IJF, IJFH, IFH, VIJFH, VIFH, VF, IF},
1184+
val,
1185+
I::CartesianIndex,
1186+
)
1187+
@boundscheck bounds_condition(data, I) || throw(BoundsError(data, I))
1188+
@inbounds set_struct!(
1189+
parent(data),
1190+
convert(eltype(data), val),
1191+
Val(field_dim(data)),
1192+
to_data_specific(data, I),
1193+
)
1194+
end
1195+
13341196
"""
13351197
data2array(::AbstractData)
13361198

0 commit comments

Comments
 (0)