unroll hilbert curve correctly

src/tree.jl

@@ -300,9 +300,10 @@ Base.@kwdef struct Tree{N_Dim} <: AbstractTree
   periodic::NTuple{N_Dim,Bool}
 
   function Tree{N_Dim}(dims::Vararg{Int64,N_Dim};
-                       periodic::NTuple{N_Dim,Bool} = Tuple(false for _ in 1:N_Dim)) where N_Dim
+                       periodic::NTuple{N_Dim,Bool} = Tuple(false for _ in 1:N_Dim),
+                       layout=:hilbert) where N_Dim
     @assert all(dims .> 0)
-    coords = coordinates_snake_curve(dims...)
+    coords = coordinates_curve(layout, dims...)
     ncells = prod(dims)
     cells = [ Cell{N_Dim,2*N_Dim}(index=i) for i in 1:ncells ]
     resize!(cells, prod(dims))
@@ -453,6 +454,17 @@ function _connect_periodic!(t::Tree3d, ax::Cart3d.Axis)
 end
 
 
+function coordinates_curve(layout::Symbol, Ns...)
+  if layout === :snake
+    return coordinates_snake_curve(Ns...)
+  elseif layout === :hilbert
+    return coordinates_hilbert_curve(Ns...)
+  else
+    error("Unknown curve layout '$layout' requested!")
+  end
+end
+
+
 """
     coordinates_snake_curve(Nx)
     coordinates_snake_curve(Nx,Ny)
@@ -534,26 +546,54 @@ end
 Generate coordinates of a hilbert curve that fills a 1d/2d/3d space with dimensions
 `(Nx)/(Nx,Ny)/(Nx,Ny,Nz)`.
 """
+function coordinates_hilbert_curve(Nx)
+  return coordinates_snake_curve(Nx)
+end
 function coordinates_hilbert_curve(Ns...)
-  @assert (1 <= length(Ns) <= 3) """
+  @assert (2 <= length(Ns) <= 3) """
   Only 1d/2d/3d dimensional spaces supported.
   """
   @assert all(Ns .> 0) && all(ispow2.(Ns)) """
   Only powers of 2 supported for dimensions.
   """
-  coords = NTuple{length(Ns),Int64}[]
+  # generate hilbert curve in a box with dimensions Ns
+  coords_box = zeros(Int64, Ns)
   alg = Compact(Int64, [ Ni for Ni in Ns ])
   CI = CartesianIndices(Ns)
   for ci in CI
     idx = [ i for i in Tuple(ci) ]
     lin_coord = encode_hilbert(alg, idx)
-    coord = decode_hilbert!(alg, idx, lin_coord)
-    push!(coords, tuple(coord...))
+    coords_box[ci] = lin_coord
   end
+  order = unroll_curve_order(coords_box)
+  coords = [ Tuple(CI[o]) for o in order ]
   return coords
 end
 
 
+"""
+    unroll_curve(curve::AbstractArray)
+
+Return the order in which a linear index `curve` must be unrolled such that its index
+is increasing.
+
+A linear index curve of length `N` is a sequence of integers ranging from `1` to `N`.
+"""
+function unroll_curve_order(curve::AbstractArray)
+  LI = LinearIndices(curve)
+  N = length(curve)
+  order = zeros(Int64, N)
+  for i = 1:N
+    next_i = findfirst(LI) do li
+      curve[li] == i
+    end
+    isnothing(next_i) && error("Failed to find curve index `$i`.")
+    order[i] = LI[next_i]
+  end
+  return order
+end
+
+
 function error_failed_placement(cell, prev_cell)
   error("Failed to place box at '$cell' because it is not a neighbor of '$prev_cell'")
 end