From c644f9d71979143674686d88534bf738fd27bfd6 Mon Sep 17 00:00:00 2001
From: Malcolm Tyrrell <matyrr@microsoft.com>
Date: Wed, 20 Oct 2021 15:08:08 +0100
Subject: [PATCH] Subtract the centroid

---
 code/Common/SpatialSort.cpp        | 12 +++++++----
 include/assimp/SpatialSort.h       |  7 ++++--
 test/unit/Common/utSpatialSort.cpp | 34 ++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/code/Common/SpatialSort.cpp b/code/Common/SpatialSort.cpp
index a8afbfaf6..08a7c55df 100644
--- a/code/Common/SpatialSort.cpp
+++ b/code/Common/SpatialSort.cpp
@@ -89,8 +89,12 @@ void SpatialSort::Fill(const aiVector3D *pPositions, unsigned int pNumPositions,
 
 // ------------------------------------------------------------------------------------------------
 void SpatialSort::Finalize() {
+    const ai_real scale = 1.0f / mPositions.size();
     for (unsigned int i = 0; i < mPositions.size(); i++) {
-        mPositions[i].mDistance = mPositions[i].mPosition * mPlaneNormal;
+        mCentroid += scale * mPositions[i].mPosition; 
+    }
+    for (unsigned int i = 0; i < mPositions.size(); i++) {
+        mPositions[i].mDistance = (mPositions[i].mPosition - mCentroid) * mPlaneNormal;
     }
     std::sort(mPositions.begin(), mPositions.end());
     mFinalized = true;
@@ -121,7 +125,7 @@ void SpatialSort::Append(const aiVector3D *pPositions, unsigned int pNumPosition
 void SpatialSort::FindPositions(const aiVector3D &pPosition,
         ai_real pRadius, std::vector<unsigned int> &poResults) const {
     ai_assert(mFinalized && "The SpatialSort object must be finalized before FindPositions can be called.");
-    const ai_real dist = pPosition * mPlaneNormal;
+    const ai_real dist = (pPosition - mCentroid) * mPlaneNormal;
     const ai_real minDist = dist - pRadius, maxDist = dist + pRadius;
 
     // clear the array
@@ -262,7 +266,7 @@ void SpatialSort::FindIdenticalPositions(const aiVector3D &pPosition, std::vecto
 
     // Convert the plane distance to its signed integer representation so the ULPs tolerance can be
     //  applied. For some reason, VC won't optimize two calls of the bit pattern conversion.
-    const BinFloat minDistBinary = ToBinary(pPosition * mPlaneNormal) - distanceToleranceInULPs;
+    const BinFloat minDistBinary = ToBinary((pPosition - mCentroid) * mPlaneNormal) - distanceToleranceInULPs;
     const BinFloat maxDistBinary = minDistBinary + 2 * distanceToleranceInULPs;
 
     // clear the array in this strange fashion because a simple clear() would also deallocate
@@ -312,7 +316,7 @@ unsigned int SpatialSort::GenerateMappingTable(std::vector<unsigned int> &fill,
     unsigned int t = 0;
     const ai_real pSquared = pRadius * pRadius;
     for (size_t i = 0; i < mPositions.size();) {
-        dist = mPositions[i].mPosition * mPlaneNormal;
+        dist = (mPositions[i].mPosition - mCentroid) * mPlaneNormal;
         maxDist = dist + pRadius;
 
         fill[mPositions[i].mIndex] = t;
diff --git a/include/assimp/SpatialSort.h b/include/assimp/SpatialSort.h
index 475a11d7c..39f863b15 100644
--- a/include/assimp/SpatialSort.h
+++ b/include/assimp/SpatialSort.h
@@ -143,7 +143,7 @@ public:
             ai_real pRadius) const;
 
 protected:
-    /** Normal of the sorting plane, normalized. The center is always at (0, 0, 0) */
+    /** Normal of the sorting plane, normalized. */
     aiVector3D mPlaneNormal;
 
     /** An entry in a spatially sorted position array. Consists of a vertex index,
@@ -171,8 +171,11 @@ protected:
     // all positions, sorted by distance to the sorting plane
     std::vector<Entry> mPositions;
 
-    // False until the Finalize method is called.
+    /// false until the Finalize method is called.
     bool mFinalized;
+
+    /// The centroid of the positions. Calculated in Finalize.
+    aiVector3D mCentroid;
 };
 
 } // end of namespace Assimp
diff --git a/test/unit/Common/utSpatialSort.cpp b/test/unit/Common/utSpatialSort.cpp
index 94dfd2dc4..ba411d99f 100644
--- a/test/unit/Common/utSpatialSort.cpp
+++ b/test/unit/Common/utSpatialSort.cpp
@@ -82,3 +82,37 @@ TEST_F(utSpatialSort, findPositionsTest) {
     sSort.FindPositions(vecs[0], 0.01f, indices);
     EXPECT_EQ(1u, indices.size());
 }
+
+TEST_F(utSpatialSort, highlyDisplacedPositionsTest) {
+    constexpr unsigned int verticesPerAxis = 10;
+    constexpr ai_real step = 0.001f;
+    constexpr ai_real offset = 5000.0f - (0.5f * verticesPerAxis * step);
+    constexpr unsigned int totalNumPositions = verticesPerAxis * verticesPerAxis * verticesPerAxis;
+    aiVector3D* positions = new aiVector3D[totalNumPositions];
+    for (unsigned int x = 0; x < verticesPerAxis; ++x) {
+        for (unsigned int y = 0; y < verticesPerAxis; ++y) {
+            for (unsigned int z = 0; z < verticesPerAxis; ++z) {
+                const unsigned int index = (x * verticesPerAxis * verticesPerAxis) + (y * verticesPerAxis) + z;
+                positions[index] = aiVector3D(offset + (x * step), offset + (y * step), offset + (z * step));
+            }
+        }
+    }
+
+    SpatialSort sSort;
+    sSort.Fill(positions, totalNumPositions, sizeof(aiVector3D));
+
+    // Enough to pick up the neighboring vertices, but not the neighbors' neighbors.
+    const ai_real epsilon = 1.1f * step;
+    std::vector<unsigned int> indices;
+    // Iterate through the interior points of the cube.
+    for (unsigned int x = 1; x < verticesPerAxis - 1; ++x) {
+        for (unsigned int y = 1; y < verticesPerAxis - 1; ++y) {
+            for (unsigned int z = 1; z < verticesPerAxis - 1; ++z) {
+                const unsigned int index = (x * verticesPerAxis * verticesPerAxis) + (y * verticesPerAxis) + z;
+                sSort.FindPositions(positions[index], epsilon, indices);
+                ASSERT_EQ(7u, indices.size());
+            }
+        }
+    }
+    delete[] positions;
+}