1 files changed, 86 insertions, 22 deletions

diff --git a/extern/mantaflow/preprocessed/particle.h b/extern/mantaflow/preprocessed/particle.h
index 0be141ed26f..da6733b6845 100644
--- a/extern/mantaflow/preprocessed/particle.h
+++ b/extern/mantaflow/preprocessed/particle.h
@@ -100,6 +100,17 @@ class ParticleBase : public PbClass {
   //! threads)
   inline void addBuffered(const Vec3 &pos, int flag = 0);
 
+  virtual void resize(IndexInt size)
+  {
+    assertMsg(false, "Dont use, override...");
+    return;
+  }
+  virtual void resizeAll(IndexInt size)
+  {
+    assertMsg(false, "Dont use, override...");
+    return;
+  }
+
   //! particle data functions
 
   //! create a particle data object
@@ -152,6 +163,20 @@ class ParticleBase : public PbClass {
     return mPartData[i];
   }
 
+  //! expose maximum number of particles to python
+  int mMaxParticles;
+  static PyObject *_GET_mMaxParticles(PyObject *self, void *cl)
+  {
+    ParticleBase *pbo = dynamic_cast<ParticleBase *>(Pb::objFromPy(self));
+    return toPy(pbo->mMaxParticles);
+  }
+  static int _SET_mMaxParticles(PyObject *self, PyObject *val, void *cl)
+  {
+    ParticleBase *pbo = dynamic_cast<ParticleBase *>(Pb::objFromPy(self));
+    pbo->mMaxParticles = fromPy<int>(val);
+    return 0;
+  }
+
  protected:
   //! new particle candidates
   std::vector<Vec3> mNewBufferPos;
@@ -431,8 +456,14 @@ template<class S> class ParticleSystem : public ParticleBase {
   }
   //! insert buffered positions as new particles, update additional particle data
   void insertBufferedParticles();
+  //! resize only the data vector, only use if you know what you're doing, otherwise use
+  //! resizeAll()
+  virtual void resize(IndexInt size)
+  {
+    mData.resize(size);
+  }
   //! resize data vector, and all pdata fields
-  void resizeAll(IndexInt newsize);
+  virtual void resizeAll(IndexInt size);
 
   //! adding and deleting
   inline void kill(IndexInt idx);
@@ -877,11 +908,6 @@ class ParticleIndexSystem : public ParticleSystem<ParticleIndexData> {
       return -1;
     }
   };
-  //! we only need a resize function...
-  void resize(IndexInt size)
-  {
-    mData.resize(size);
-  }
  public:
   PbArgs _args;
 }
@@ -2479,28 +2505,66 @@ template<class S> void ParticleSystem<S>::insertBufferedParticles()
   for (IndexInt i = 0; i < (IndexInt)mData.size(); ++i)
     mData[i].flag &= ~PNEW;
 
-  if (mNewBufferPos.size() == 0)
+  if (mNewBufferPos.empty())
     return;
-  IndexInt newCnt = mData.size();
-  resizeAll(newCnt + mNewBufferPos.size());
-
-  for (IndexInt i = 0; i < (IndexInt)mNewBufferPos.size(); ++i) {
-    int flag = (mNewBufferFlag.size() > 0) ? mNewBufferFlag[i] : 0;
-    // note, other fields are not initialized here...
-    mData[newCnt].pos = mNewBufferPos[i];
-    mData[newCnt].flag = PNEW | flag;
+  IndexInt bufferSize = mNewBufferPos.size();
+  IndexInt partsSize = mData.size();
+
+  if (mMaxParticles > 0)
+    assertMsg(mMaxParticles >= partsSize,
+              "Particle system cannot contain more particles that the maximum allowed number");
+
+  // max number of new particles that can be inserted, adjusted buffer size when using maxParticles
+  // field
+  IndexInt numNewParts = (mMaxParticles > 0) ? mMaxParticles - mData.size() : bufferSize;
+  if (numNewParts > bufferSize)
+    numNewParts = bufferSize;  // upper clamp
+
+  assertMsg(numNewParts >= 0, "Must not have negative number of new particles");
+
+  // new size of particle system
+  IndexInt newSize = mData.size() + numNewParts;
+  if (mMaxParticles > 0)
+    assertMsg(newSize <= mMaxParticles,
+              "Particle system cannot contain more particles that the maximum allowed number");
+  resizeAll(newSize);
+
+  int insertFlag;
+  Vec3 insertPos;
+  static RandomStream mRand(9832);
+  for (IndexInt i = 0; i < numNewParts; ++i) {
+
+    // get random index in newBuffer vector
+    // we are inserting particles randomly so that they are sampled uniformly in the fluid region
+    // otherwise, regions of fluid can remain completely empty once mData.size() == maxParticles is
+    // reached.
+    int randIndex = floor(mRand.getReal() * mNewBufferPos.size());
+
+    // get elements from new buffers with random index
+    std::swap(mNewBufferPos[randIndex], mNewBufferPos.back());
+    insertPos = mNewBufferPos.back();
+    mNewBufferPos.pop_back();
+
+    insertFlag = 0;
+    if (!mNewBufferFlag.empty()) {
+      std::swap(mNewBufferFlag[randIndex], mNewBufferFlag.back());
+      insertFlag = mNewBufferFlag.back();
+      mNewBufferFlag.pop_back();
+    }
+
+    mData[partsSize].pos = insertPos;
+    mData[partsSize].flag = PNEW | insertFlag;
+
     // now init pdata fields from associated grids...
     for (IndexInt pd = 0; pd < (IndexInt)mPdataReal.size(); ++pd)
-      mPdataReal[pd]->initNewValue(newCnt, mNewBufferPos[i]);
+      mPdataReal[pd]->initNewValue(partsSize, insertPos);
     for (IndexInt pd = 0; pd < (IndexInt)mPdataVec3.size(); ++pd)
-      mPdataVec3[pd]->initNewValue(newCnt, mNewBufferPos[i]);
+      mPdataVec3[pd]->initNewValue(partsSize, insertPos);
     for (IndexInt pd = 0; pd < (IndexInt)mPdataInt.size(); ++pd)
-      mPdataInt[pd]->initNewValue(newCnt, mNewBufferPos[i]);
-    newCnt++;
+      mPdataInt[pd]->initNewValue(partsSize, insertPos);
+    partsSize++;
   }
-  if (mNewBufferPos.size() > 0)
-    debMsg("Added & initialized " << (IndexInt)mNewBufferPos.size() << " particles",
-           2);  // debug info
+  debMsg("Added & initialized " << numNewParts << " particles", 2);  // debug info
   mNewBufferPos.clear();
   mNewBufferFlag.clear();
 }