From 684adafea2a3142bde9730f2414ff99fc30abed3 Mon Sep 17 00:00:00 2001
From: Elias Carvalho <73039601+eliascarv@users.noreply.github.com>
Date: Tue, 26 Dec 2023 18:08:47 -0300
Subject: [PATCH] Add support for `Colon` (`:`) in grid indexing (#673)

* Add support for 'Colon' in grid indexing

* Rename argument
---
 src/mesh.jl  | 11 +++++++++--
 test/mesh.jl | 21 +++++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/mesh.jl b/src/mesh.jl
index 3bc849f66..4d611c74b 100644
--- a/src/mesh.jl
+++ b/src/mesh.jl
@@ -195,8 +195,15 @@ Base.eltype(g::Grid) = typeof(first(g))
 
 Base.getindex(g::Grid{Dim}, ijk::Vararg{Int,Dim}) where {Dim} = element(g, LinearIndices(size(g))[ijk...])
 
-Base.getindex(g::Grid{Dim}, ijk::Vararg{Union{UnitRange{Int},Int},Dim}) where {Dim} =
-  getindex(g, CartesianIndex(first.(ijk)):CartesianIndex(last.(ijk)))
+function Base.getindex(g::Grid{Dim}, ijk::Vararg{Union{UnitRange{Int},Colon,Int},Dim}) where {Dim}
+  dims = size(g)
+  ranges = ntuple(i -> _asrange(dims[i], ijk[i]), Dim)
+  getindex(g, CartesianIndices(ranges))
+end
+
+_asrange(::Int, r::UnitRange{Int}) = r
+_asrange(d::Int, ::Colon) = 1:d
+_asrange(::Int, i::Int) = i:i
 
 # ----------------
 # IMPLEMENTATIONS
diff --git a/test/mesh.jl b/test/mesh.jl
index e757a7574..ec31d7c9f 100644
--- a/test/mesh.jl
+++ b/test/mesh.jl
@@ -119,6 +119,11 @@
     @test spacing(sub) == spacing(grid)
     @test minimum(sub) == P2(2, 3)
     @test maximum(sub) == P2(3, 8)
+    sub = grid[:, 3:7]
+    @test size(sub) == (10, 5)
+    @test spacing(sub) == spacing(grid)
+    @test minimum(sub) == P2(1, 3)
+    @test maximum(sub) == P2(11, 8)
 
     # subgrid with comparable vertices of grid
     grid = CartesianGrid((10, 10), P2(0.0, 0.0), T.((1.2, 1.2)))
@@ -246,6 +251,10 @@
     @test size(sub) == (1, 3)
     @test minimum(sub) == P2(0.2, 0.3)
     @test maximum(sub) == P2(0.4, 1.0)
+    sub = grid[:, 3:5]
+    @test size(sub) == (5, 3)
+    @test minimum(sub) == P2(0.0, 0.3)
+    @test maximum(sub) == P2(1.0, 1.0)
     @test Meshes.xyz(grid) == (x, y)
     @test Meshes.XYZ(grid) == (repeat(x, 1, 6), repeat(y', 6, 1))
 
@@ -293,6 +302,10 @@
     @test minimum(sub) == P2(0.2, 0.3)
     @test maximum(sub) == P2(0.4, 1.0)
     @test Meshes.XYZ(grid) == (X, Y)
+    sub = grid[:, 3:5]
+    @test size(sub) == (5, 3)
+    @test minimum(sub) == P2(0.0, 0.3)
+    @test maximum(sub) == P2(1.0, 1.0)
 
     # conversion
     cg = CartesianGrid{T}(10, 10)
@@ -440,6 +453,10 @@
     @test size(sub) == (1, 5)
     @test minimum(sub) == P2(1, 2)
     @test maximum(sub) == P2(2, 7)
+    sub = mesh[:, 3:7]
+    @test size(sub) == (10, 5)
+    @test minimum(sub) == P2(0, 2)
+    @test maximum(sub) == P2(10, 7)
 
     # test for https://github.com/JuliaGeometry/Meshes.jl/issues/261
     points = rand(P2, 5)
@@ -529,5 +546,9 @@
     @test size(sub) == (1, 5)
     @test minimum(sub) == P2(1, 2)
     @test maximum(sub) == P2(2, 7)
+    sub = tgrid[:, 3:7]
+    @test size(sub) == (10, 5)
+    @test minimum(sub) == P2(0, 2)
+    @test maximum(sub) == P2(10, 7)
   end
 end