Optimize FindDegenerates so it doesn't explode

code/PostProcessing/FindDegenerates.cpp

@@ -52,12 +52,12 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include "FindDegenerates.h"
 #include <assimp/Exceptional.h>
 
+#include <unordered_map>
+
 using namespace Assimp;
 
-//remove mesh at position 'index' from the scene
-static void removeMesh(aiScene* pScene, unsigned const index);
 //correct node indices to meshes and remove references to deleted mesh
-static void updateSceneGraph(aiNode* pNode, unsigned const index);
+static void updateSceneGraph(aiNode* pNode, const std::unordered_map<unsigned int, unsigned int>& meshMap);
 
 // ------------------------------------------------------------------------------------------------
 // Constructor to be privately used by Importer
@@ -91,50 +91,51 @@ void FindDegeneratesProcess::SetupProperties(const Importer* pImp) {
 // Executes the post processing step on the given imported data.
 void FindDegeneratesProcess::Execute( aiScene* pScene) {
     ASSIMP_LOG_DEBUG("FindDegeneratesProcess begin");
+
+    std::unordered_map<unsigned int, unsigned int> meshMap;
+    meshMap.reserve(pScene->mNumMeshes);
+
+    const unsigned int originalNumMeshes = pScene->mNumMeshes;
+    unsigned int targetIndex = 0;
     for (unsigned int i = 0; i < pScene->mNumMeshes;++i)
     {
-        //Do not process point cloud, ExecuteOnMesh works only with faces data
+        // Do not process point cloud, ExecuteOnMesh works only with faces data
         if ((pScene->mMeshes[i]->mPrimitiveTypes != aiPrimitiveType::aiPrimitiveType_POINT) && ExecuteOnMesh(pScene->mMeshes[i])) {
-            removeMesh(pScene, i);
-            --i; //the current i is removed, do not skip the next one
+            delete pScene->mMeshes[i];
+            // Not strictly required, but clean:
+            pScene->mMeshes[i] = nullptr;
+        }
+        else {
+            meshMap[i] = targetIndex;
+            pScene->mMeshes[targetIndex] = pScene->mMeshes[i];
+            ++targetIndex;
         }
     }
-    ASSIMP_LOG_DEBUG("FindDegeneratesProcess finished");
-}
+    pScene->mNumMeshes = targetIndex;
 
-static void removeMesh(aiScene* pScene, unsigned const index) {
-    //we start at index and copy the pointers one position forward
-    //save the mesh pointer to delete it later
-    auto delete_me = pScene->mMeshes[index];
-    for (unsigned i = index; i < pScene->mNumMeshes - 1; ++i) {
-        pScene->mMeshes[i] = pScene->mMeshes[i+1];
+    if (meshMap.size() < originalNumMeshes)
+    {
+        updateSceneGraph(pScene->mRootNode, meshMap);
     }
-    pScene->mMeshes[pScene->mNumMeshes - 1] = nullptr;
-    --(pScene->mNumMeshes);
-    delete delete_me;
 
-    //removing a mesh also requires updating all references to it in the scene graph
-    updateSceneGraph(pScene->mRootNode, index);
+    ASSIMP_LOG_DEBUG("FindDegeneratesProcess finished");
 }
 
-static void updateSceneGraph(aiNode* pNode, unsigned const index) {
+static void updateSceneGraph(aiNode* pNode, const std::unordered_map<unsigned int, unsigned int>& meshMap) {
+    unsigned int targetIndex = 0;
     for (unsigned i = 0; i < pNode->mNumMeshes; ++i) {
-        if (pNode->mMeshes[i] > index) {
-            --(pNode->mMeshes[i]);
-            continue;
-        }
-        if (pNode->mMeshes[i] == index) {
-            for (unsigned j = i; j < pNode->mNumMeshes -1; ++j) {
-                pNode->mMeshes[j] = pNode->mMeshes[j+1];
-            }
-            --(pNode->mNumMeshes);
-            --i;
-            continue;
+        const unsigned int sourceMeshIndex = pNode->mMeshes[i];
+        auto it = meshMap.find(sourceMeshIndex);
+        if (it != meshMap.end())
+        {
+            pNode->mMeshes[targetIndex] = it->second;
+            ++targetIndex;
         }
     }
+    pNode->mNumMeshes = targetIndex;
     //recurse to all children
     for (unsigned i = 0; i < pNode->mNumChildren; ++i) {
-        updateSceneGraph(pNode->mChildren[i], index);
+        updateSceneGraph(pNode->mChildren[i], meshMap);
     }
 }
 

test/unit/utFindDegenerates.cpp

@@ -40,6 +40,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
 #include "UnitTestPCH.h"
 
+#include "../../include/assimp/scene.h"
 #include "PostProcessing/FindDegenerates.h"
 
 using namespace std;
@@ -147,3 +148,61 @@ TEST_F(FindDegeneratesProcessTest, testDegeneratesRemovalWithAreaCheck) {
 
     EXPECT_EQ(mMesh->mNumUVComponents[1] - 100, mMesh->mNumFaces);
 }
+
+namespace
+{
+    std::unique_ptr<aiMesh> getDegenerateMesh()
+    {
+        std::unique_ptr<aiMesh> mesh = std::make_unique<aiMesh>();
+        mesh->mNumVertices = 2;
+        mesh->mVertices = new aiVector3D[2];
+        mesh->mVertices[0] = aiVector3D{ 0.0f, 0.0f, 0.0f };
+        mesh->mVertices[1] = aiVector3D{ 1.0f, 0.0f, 0.0f };
+        mesh->mNumFaces = 1;
+        mesh->mFaces = new aiFace[1];
+        mesh->mFaces[0].mNumIndices = 3;
+        mesh->mFaces[0].mIndices = new unsigned int[3];
+        mesh->mFaces[0].mIndices[0] = 0;
+        mesh->mFaces[0].mIndices[1] = 1;
+        mesh->mFaces[0].mIndices[2] = 0;
+        return mesh;
+    }
+}
+
+TEST_F(FindDegeneratesProcessTest, meshRemoval) {
+    mProcess->EnableAreaCheck(true);
+    mProcess->EnableInstantRemoval(true);
+    mProcess->ExecuteOnMesh(mMesh);
+
+    std::unique_ptr<aiScene> scene = std::make_unique<aiScene>();
+    scene->mNumMeshes = 5;
+    scene->mMeshes = new aiMesh*[5];
+
+    /// Use the mesh which doesn't get completely stripped of faces from the main test.
+    aiMesh* meshWhichSurvives = mMesh;
+    mMesh = nullptr;
+
+    scene->mMeshes[0] = getDegenerateMesh().release();
+    scene->mMeshes[1] = getDegenerateMesh().release();
+    scene->mMeshes[2] = meshWhichSurvives;
+    scene->mMeshes[3] = getDegenerateMesh().release();
+    scene->mMeshes[4] = getDegenerateMesh().release();
+
+    scene->mRootNode = new aiNode;
+    scene->mRootNode->mNumMeshes = 5;
+    scene->mRootNode->mMeshes = new unsigned int[5];
+    scene->mRootNode->mMeshes[0] = 0;
+    scene->mRootNode->mMeshes[1] = 1;
+    scene->mRootNode->mMeshes[2] = 2;
+    scene->mRootNode->mMeshes[3] = 3;
+    scene->mRootNode->mMeshes[4] = 4;
+
+    mProcess->Execute(scene.get());    
+
+    EXPECT_EQ(scene->mNumMeshes, 1);
+    EXPECT_EQ(scene->mMeshes[0], meshWhichSurvives);
+    EXPECT_EQ(scene->mRootNode->mNumMeshes, 1);
+    EXPECT_EQ(scene->mRootNode->mMeshes[0], 0);    
+
+    EXPECT_EQ(mMesh->mNumUVComponents[1] - 100, mMesh->mNumFaces);
+}