fix hexagonal domain repetition; must check 4 neighbors

1 files changed, 97 insertions, 42 deletions

diff --git a/ray_marching.cginc b/ray_marching.cginc
index bf01eb4..9cb8dd0 100644
--- a/ray_marching.cginc
+++ b/ray_marching.cginc
@@ -46,65 +46,120 @@ float map(float3 p) {
 }
 
 #if defined(_RAY_MARCHING_HEX_GRID)
-void hex_nearest_neighbors(float3 cubeId, float3 neighborDir,
-    out float3 id0, out float3 id1, out float3 id2) {
-  // Standard cube-coordinate neighbour offsets (sum equals zero).
-  const float3 axis0 = float3( 1.0f,  0.0f, -1.0f);
-  const float3 axis1 = float3( 0.0f,  1.0f, -1.0f);
-  const float3 axis2 = float3( 1.0f, -1.0f,  0.0f);
-
-  const float3 cube0 = cubeId + axis0 * neighborDir.xxx;
-  const float3 cube1 = cubeId + axis1 * neighborDir.yyy;
-  const float3 cube2 = cubeId + axis2 * neighborDir.zzz;
-
-  // Convert back to our "x = y + z" coordinate representation.
-  id0 = float3(cube0.x + cube0.y, cube0.x, cube0.y);
-  id1 = float3(cube1.x + cube1.y, cube1.x, cube1.y);
-  id2 = float3(cube2.x + cube2.y, cube2.x, cube2.y);
-}
-
 float domain_repeat_hex_grid(inout float3 p) {
-  float d;
-
-  const float3 count = _Ray_Marching_Hex_Grid_Count * 0.5f;
+  const float gridCount = max(_Ray_Marching_Hex_Grid_Count, 1.0f);
+  const float halfCount = gridCount * 0.5f;
+  const float3 count = float3(halfCount, halfCount, halfCount);
   const float3 period = 1.0f / count;
   const float3 hex = cart_to_hex(p.xy);
   // Cell ID.
   const float3 id = round_hex(hex / period);
   // Coordinates within the current cell.
   const float3 local = hex - period * id;
+  float3 candidate = p;
+  candidate.xy = hex_to_cart(local);
+  float d = map(candidate);
+  float3 bestPos = candidate;
 
+#if defined(_RAY_MARCHING_CORRECT_REPETITION)
   const float3 cubeId = float3(id.y, id.z, -id.y - id.z);
-  const float3 cubeLocal = float3(local.y, local.z, -local.y - local.z);
-  float3 neighborDir = sign(cubeLocal);
-  neighborDir += (1.0f.xxx - abs(neighborDir));
+  const float limit = max((gridCount - 1.0f) * 0.5f, 0.0f);
+  const float3 cubeLimit = float3(limit, limit, limit);
+  if (any(abs(cubeId) > cubeLimit)) {
+    p = float3(1e9, 1e9, 1e9);
+    return 1e9;
+  }
+  const float3 cubeOffsets[6] = {
+      float3( 1.0f, -1.0f,  0.0f),
+      float3( 1.0f,  0.0f, -1.0f),
+      float3( 0.0f,  1.0f, -1.0f),
+      float3(-1.0f,  1.0f,  0.0f),
+      float3(-1.0f,  0.0f,  1.0f),
+      float3( 0.0f, -1.0f,  1.0f),
+  };
 
-#if defined(_RAY_MARCHING_CORRECT_REPETITION)
-  float3 ids[4];
-  ids[0] = id;
-  hex_nearest_neighbors(cubeId, neighborDir, ids[1], ids[2], ids[3]);
-  const uint kNumCells = 4u;
-#else
-  float3 ids[1];
-  ids[0] = id;
-  const uint kNumCells = 1u;
-#endif
+  float neighborDistSq[4] = { 1e9, 1e9, 1e9, 1e9 };
+  float3 neighborIds[4] = {
+      float3(0.0f, 0.0f, 0.0f),
+      float3(0.0f, 0.0f, 0.0f),
+      float3(0.0f, 0.0f, 0.0f),
+      float3(0.0f, 0.0f, 0.0f)};
+  float2 neighborLoc[4] = {
+      float2(0.0f, 0.0f),
+      float2(0.0f, 0.0f),
+      float2(0.0f, 0.0f),
+      float2(0.0f, 0.0f)};
 
-  d = 1e9;
-  float3 bestPos = p;
-  for (uint ii = 0; ii < kNumCells; ++ii) {
-    const float3 rid = ids[ii];
-    if (any(abs(rid) > count * 0.5f)) continue;
+  [unroll]
+  for (int idx = 0; idx < 6; ++idx) {
+    const float3 cube = cubeId + cubeOffsets[idx];
+    const float3 rid = float3(cube.x + cube.y, cube.x, cube.y);
+
+    if (any(abs(cube) > cubeLimit)) {
+      continue;
+    }
 
-    float3 candidate = p;
-    candidate.xy = hex_to_cart(hex - period * rid);
+    const float2 localXY = hex_to_cart(hex - period * rid);
+    const float distSq = dot(localXY, localXY);
+
+    if (distSq >= neighborDistSq[3]) {
+      continue;
+    }
+
+    neighborDistSq[3] = distSq;
+    neighborIds[3] = rid;
+    neighborLoc[3] = localXY;
+
+    if (neighborDistSq[3] < neighborDistSq[2]) {
+      const float tmpDist = neighborDistSq[2];
+      const float3 tmpId = neighborIds[2];
+      const float2 tmpLoc = neighborLoc[2];
+      neighborDistSq[2] = neighborDistSq[3];
+      neighborIds[2] = neighborIds[3];
+      neighborLoc[2] = neighborLoc[3];
+      neighborDistSq[3] = tmpDist;
+      neighborIds[3] = tmpId;
+      neighborLoc[3] = tmpLoc;
+    }
+    if (neighborDistSq[2] < neighborDistSq[1]) {
+      const float tmpDist = neighborDistSq[1];
+      const float3 tmpId = neighborIds[1];
+      const float2 tmpLoc = neighborLoc[1];
+      neighborDistSq[1] = neighborDistSq[2];
+      neighborIds[1] = neighborIds[2];
+      neighborLoc[1] = neighborLoc[2];
+      neighborDistSq[2] = tmpDist;
+      neighborIds[2] = tmpId;
+      neighborLoc[2] = tmpLoc;
+    }
+    if (neighborDistSq[1] < neighborDistSq[0]) {
+      const float tmpDist = neighborDistSq[0];
+      const float3 tmpId = neighborIds[0];
+      const float2 tmpLoc = neighborLoc[0];
+      neighborDistSq[0] = neighborDistSq[1];
+      neighborIds[0] = neighborIds[1];
+      neighborLoc[0] = neighborLoc[1];
+      neighborDistSq[1] = tmpDist;
+      neighborIds[1] = tmpId;
+      neighborLoc[1] = tmpLoc;
+    }
+  }
+
+  [unroll]
+  for (int idx = 0; idx < 4; ++idx) {
+    if (neighborDistSq[idx] >= 1e9) {
+      continue;
+    }
 
-    const float d_cur = map(candidate);
+    float3 neighbour = p;
+    neighbour.xy = neighborLoc[idx];
+    const float d_cur = map(neighbour);
     if (d_cur < d) {
       d = d_cur;
-      bestPos = candidate;
+      bestPos = neighbour;
     }
   }
+#endif  // _RAY_MARCHING_CORRECT_REPETITION
 
   p = bestPos;
   return d;