S2Cell::GetDistance: Optimize same-face case

src/s2/s2cell.cc

@@ -497,31 +497,79 @@ S1ChordAngle S2Cell::GetMaxDistance(const S2Point& a, const S2Point& b) const {
   return S1ChordAngle::Straight() - GetDistance(-a, -b);
 }
 
+namespace {
+// Finds the index `ai` of the edge of A that is furthest away from the
+// opposite edge of B in (u,v)-space such that B is in that direction from A.
+// For example, if A is to the left of B, but not above or below it, then the
+// right edge of A is the furthest edge from the opposite edge of B.  If A is
+// both above and to the left of B, then furthest edge will either be the right
+// or bottom edge of A, depending on which has greater uv distance.  Returns -1
+// if A and B intersect (including edge and vertex intersections).  Used by
+// `GetDistance(const S2Cell&)` and `IsDistanceLess(const S2Cell&)` below.
+ABSL_ATTRIBUTE_ALWAYS_INLINE int FindFurthestEdge(const S2Cell& a,
+                                                  const S2Cell& b) {
+  int ai = -1;
+  double max_dist = 0;
+  auto check_edge = [&](double dist, int a_edge) ABSL_ATTRIBUTE_ALWAYS_INLINE {
+    if (dist > max_dist) {
+      max_dist = dist;
+      ai = a_edge;
+    }
+  };
+  R2Rect a_uv = a.GetBoundUV();
+  R2Rect b_uv = b.GetBoundUV();
+  check_edge(a_uv[0][0] - b_uv[0][1], S2Cell::kLeftEdge);
+  check_edge(b_uv[0][0] - a_uv[0][1], S2Cell::kRightEdge);
+  check_edge(a_uv[1][0] - b_uv[1][1], S2Cell::kBottomEdge);
+  check_edge(b_uv[1][0] - a_uv[1][1], S2Cell::kTopEdge);
+  return ai;
+}
+}  // namespace
+
 S1ChordAngle S2Cell::GetDistance(const S2Cell& target) const {
-  // If the cells intersect, the distance is zero.  We use the (u,v) ranges
-  // rather S2CellId::intersects() so that cells that share a partial edge or
-  // corner are considered to intersect.
-  if (face_ == target.face_ && uv_.Intersects(target.uv_)) {
-    return S1ChordAngle::Zero();
+  S1ChordAngle min_dist = S1ChordAngle::Infinity();
+
+  // If cells are on the same face, we can speed up computation by pruning
+  // edge pairs whose cells are not adjacent in UV space.  Instead of 32
+  // vertex-edge pairs, we only need to check four.  Further optimizations
+  // could include handling other faces and identifying when a corner-corner
+  // distance is the minimum distance.
+  // https://github.com/google/s2geometry/issues/467#issuecomment-3579220154
+  if (face_ == target.face_) {
+    int ai = FindFurthestEdge(*this, target);
+    if (ai < 0) {
+      // A and B intersect (including edge and vertex intersections).
+      return S1ChordAngle::Zero();
+    }
+    // Otherwise the minimum distance always occurs between an endpoint of the
+    // edge `ai` and the opposite edge (`ai ^ 2`) of B, or symmetrically, an
+    // endpoint of the opposite edge of B and the edge `ai`.
+    int bi = ai ^ 2;
+    const S2Point va1 = GetVertex(ai);
+    const S2Point va2 = GetVertex(ai + 1);
+    const S2Point vb1 = target.GetVertex(bi);
+    const S2Point vb2 = target.GetVertex(bi + 1);
+    S2::UpdateMinDistance(va1, vb1, vb2, &min_dist);
+    S2::UpdateMinDistance(va2, vb1, vb2, &min_dist);
+    // The endpoints of `[va1, va2]` have been checked by `UpdateMinDistance`,
+    // so we can use `UpdateMinInteriorDistance`, which is faster.
+    S2::UpdateMinInteriorDistance(vb1, va1, va2, &min_dist);
+    S2::UpdateMinInteriorDistance(vb2, va1, va2, &min_dist);
+    return min_dist;
   }
 
   // Otherwise, the minimum distance always occurs between a vertex of one
   // cell and an edge of the other cell (including the edge endpoints).  This
   // represents a total of 32 possible (vertex, edge) pairs.
-  //
-  // TODO(ericv): This could be optimized to be at least 5x faster by pruning
-  // the set of possible closest vertex/edge pairs using the faces and (u,v)
-  // ranges of both cells.
   S2Point va[4], vb[4];
   for (int i = 0; i < 4; ++i) {
     va[i] = GetVertex(i);
     vb[i] = target.GetVertex(i);
   }
-  S1ChordAngle min_dist = S1ChordAngle::Infinity();
   for (int i = 0; i < 4; ++i) {
     for (int j = 0; j < 4; ++j) {
       S2::UpdateMinDistance(va[i], vb[j], vb[(j + 1) & 3], &min_dist);
-      S2::UpdateMinDistance(vb[i], va[j], va[(j + 1) & 3], &min_dist);
+      S2::UpdateMinInteriorDistance(vb[i], va[j], va[(j + 1) & 3], &min_dist);
     }
   }
   return min_dist;
@@ -549,8 +597,8 @@ S1ChordAngle S2Cell::GetMaxDistance(const S2Cell& target) const {
   // cell and an edge of the other cell (including the edge endpoints).  This
   // represents a total of 32 possible (vertex, edge) pairs.
   //
-  // TODO(user): When the maximum distance is at most Pi/2, the maximum is
-  // always attained between a pair of vertices, and this could be made much
+  // TODO: b/284470618 - When the maximum distance is at most Pi/2, the maximum
+  // is always attained between a pair of vertices, and this could be made much
   // faster by testing each vertex pair once rather than the current 4 times.
   S2Point va[4], vb[4];
   for (int i = 0; i < 4; ++i) {
@@ -568,8 +616,61 @@ S1ChordAngle S2Cell::GetMaxDistance(const S2Cell& target) const {
 }
 
 bool S2Cell::IsDistanceLess(const S2Cell& target, S1ChordAngle limit) const {
-  // TODO: b/284470618 - Use more efficient implementation
-  return GetDistance(target) < limit;
+  // Implementation of this function is similar to `GetDistance(target)`, but
+  // instead of returning the distance, we return true if the distance is less
+  // than the limit, and we can exit early if the limit or distance is zero.
+
+  // It's silly to call this with a zero limit, and in practice, no one does.
+  // One comparison and a well-predicted branch avoids some complexity below.
+  if (ABSL_PREDICT_FALSE(limit <= S1ChordAngle::Zero())) {
+    return false;
+  }
+
+  S1ChordAngle min_dist = S1ChordAngle::Infinity();
+
+  if (face_ == target.face_) {
+    // See `GetDistance(const S2Cell&)` for comments on the implementation.
+    int ai = FindFurthestEdge(*this, target);
+    if (ai < 0) {
+      return true;
+    }
+    int bi = ai ^ 2;
+    const S2Point va1 = GetVertex(ai);
+    const S2Point va2 = GetVertex(ai + 1);
+    const S2Point vb1 = target.GetVertex(bi);
+    const S2Point vb2 = target.GetVertex(bi + 1);
+    return (S2::UpdateMinDistance(va1, vb1, vb2, &min_dist) &&
+            min_dist < limit) ||
+           (S2::UpdateMinDistance(va2, vb1, vb2, &min_dist) &&
+            min_dist < limit) ||
+           (S2::UpdateMinInteriorDistance(vb1, va1, va2, &min_dist) &&
+            min_dist < limit) ||
+           (S2::UpdateMinInteriorDistance(vb2, va1, va2, &min_dist) &&
+            min_dist < limit);
+  }
+
+  // Otherwise, the minimum distance always occurs between a vertex of one
+  // cell and an edge of the other cell (including the edge endpoints).  This
+  // represents a total of 32 possible (vertex, edge) pairs.
+  S2Point va[4], vb[4];
+  for (int i = 0; i < 4; ++i) {
+    va[i] = GetVertex(i);
+    vb[i] = target.GetVertex(i);
+  }
+  for (int i = 0; i < 4; ++i) {
+    for (int j = 0; j < 4; ++j) {
+      if (S2::UpdateMinDistance(va[i], vb[j], vb[(j + 1) & 3], &min_dist) &&
+          min_dist < limit) {
+        return true;
+      }
+      if (S2::UpdateMinInteriorDistance(vb[i], va[j], va[(j + 1) & 3],
+                                        &min_dist) &&
+          min_dist < limit) {
+        return true;
+      }
+    }
+  }
+  return false;
 }
 
 bool S2Cell::IsDistanceLessOrEqual(const S2Cell& target,

src/s2/s2cell_test.cc

@@ -31,6 +31,7 @@
 #include "absl/container/flat_hash_map.h"
 #include "absl/flags/flag.h"
 #include "absl/log/absl_check.h"
+#include "absl/log/absl_log.h"
 #include "absl/log/log_streamer.h"
 #include "absl/random/bit_gen_ref.h"
 #include "absl/random/random.h"
@@ -539,6 +540,101 @@ static S1ChordAngle GetMaxDistanceToPointBruteForce(const S2Cell& cell,
   return max_distance;
 }
 
+// Previous implementation of `S2Cell::GetDistance(const S2Cell&)`.
+static S1ChordAngle GetDistanceToCellBruteForce(const S2Cell& a,
+                                                const S2Cell& b) {
+  // If the cells intersect, the distance is zero.  We use the (u,v) ranges
+  // rather S2CellId::intersects() so that cells that share a partial edge or
+  // corner are considered to intersect.
+  if (a.face() == b.face() && a.GetBoundUV().Intersects(b.GetBoundUV())) {
+    return S1ChordAngle::Zero();
+  }
+
+  // Otherwise, the minimum distance always occurs between a vertex of one
+  // cell and an edge of the other cell (including the edge endpoints).  This
+  // represents a total of 32 possible (vertex, edge) pairs.
+  S2Point va[4], vb[4];
+  for (int i = 0; i < 4; ++i) {
+    va[i] = a.GetVertex(i);
+    vb[i] = b.GetVertex(i);
+  }
+  S1ChordAngle min_dist = S1ChordAngle::Infinity();
+  for (int i = 0; i < 4; ++i) {
+    for (int j = 0; j < 4; ++j) {
+      if (S2::UpdateMinDistance(va[i], vb[j], vb[(j + 1) & 3], &min_dist)) {
+        ABSL_VLOG(3) << "from a vertex " << i << " to edge b " << j << " "
+                     << ((j + 1) & 3) << " min_dist: " << min_dist.degrees();
+      }
+      if (S2::UpdateMinDistance(vb[i], va[j], va[(j + 1) & 3], &min_dist)) {
+        ABSL_VLOG(3) << "from b vertex " << i << " to edge a " << j << " "
+                     << ((j + 1) & 3) << " min_dist: " << min_dist.degrees();
+      }
+    }
+  }
+  return min_dist;
+}
+
+TEST(S2Cell, GetDistanceToCell) {
+  absl::BitGen bitgen(S2Testing::MakeTaggedSeedSeq(
+      "GET_DISTANCE_TO_CELL", absl::LogInfoStreamer(__FILE__, __LINE__).stream()));
+  for (int iter = 0; iter < 10'000; ++iter) {
+    SCOPED_TRACE(StrCat("Iteration ", iter));
+    S2Cell cell1(s2random::CellId(bitgen));
+    S2Cell cell2(s2random::CellId(bitgen));
+    S1ChordAngle expected = GetDistanceToCellBruteForce(cell1, cell2);
+    S1ChordAngle actual = cell1.GetDistance(cell2);
+    EXPECT_NEAR(expected.radians(), actual.radians(), 1e-15)
+        << "cell1: " << cell1.id() << " cell2: " << cell2.id()
+        << " cell1.uv: " << cell1.GetBoundUV()
+        << " cell2.uv: " << cell2.GetBoundUV();
+  }
+}
+
+TEST(S2Cell, GetDistanceToCellHighDifferenceExample) {
+  // This is a test case extracted from `GetDistanceToCell`; it achieved
+  // the maximum difference to the brute-force implementation over 100M
+  // iterations, so is useful as a shortcut for testing.
+  S2Cell cell1(S2CellId::FromDebugString("4/0112122"));
+  S2Cell cell2(S2CellId::FromDebugString("4/2110333"));
+  S1ChordAngle expected = GetDistanceToCellBruteForce(cell1, cell2);
+  S1ChordAngle actual = cell1.GetDistance(cell2);
+  EXPECT_NEAR(expected.radians(), actual.radians(), 1e-15);
+}
+
+TEST(S2Cell, GetDistanceToCellProjectionExample1) {
+  // In uv space, `cell1` is to the slightly to the right right of `cell2` and
+  // significantly below it.  The minimum distance is achieved with the lower
+  // *left* vertex of `cell2`, even though `cell1` is to the right of `cell2`.
+  // This demonstrates we need to use the direction with the max uv distance.
+  // After projection onto the sphere, `cell1`'s shape is:
+  //    ╱│
+  //   ╱ │
+  //  │ ╱
+  //  │╱
+  S2Cell cell1(S2CellId::FromDebugString("1/00100000113012032112132121101"));
+  S2Cell cell2(S2CellId::FromDebugString("1/333"));
+  S1ChordAngle expected = GetDistanceToCellBruteForce(cell1, cell2);
+  S1ChordAngle actual = cell1.GetDistance(cell2);
+  EXPECT_NEAR(expected.radians(), actual.radians(), 1e-15);
+}
+
+TEST(S2Cell, GetDistanceToCellProjectionExample2) {
+  // In uv space, `cell1` is far to the left of `cell2` and slightly below it.
+  // The minimum distance is achieved with the *upper* left vertex of `cell2`,
+  // even though `cell1` is below `cell2`.  This demonstrates we need to take
+  // the direction with the max uv distance.  After projection onto the sphere,
+  // both cells are shaped like:
+  //    ╱╲
+  //   ╱  ╲
+  //   ╲  ╱
+  //    ╲╱
+  S2Cell cell1(S2CellId::FromDebugString("2/11033230030133"));
+  S2Cell cell2(S2CellId::FromDebugString("2/222"));
+  S1ChordAngle expected = GetDistanceToCellBruteForce(cell1, cell2);
+  S1ChordAngle actual = cell1.GetDistance(cell2);
+  EXPECT_NEAR(expected.radians(), actual.radians(), 1e-15);
+}
+
 TEST(S2Cell, GetDistanceToPoint) {
   absl::BitGen bitgen(S2Testing::MakeTaggedSeedSeq(
       "GET_DISTANCE_TO_POINT", absl::LogInfoStreamer(__FILE__, __LINE__).stream()));