2 files changed, 40 insertions, 23 deletions

diff --git a/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.cpp b/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.cpp
index 259dae4c6..825b72c6a 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.cpp
@@ -49,6 +49,10 @@ void BufferWithExtendableBuffer::readCodePointsAndAdvancePosition(const int maxC
     }
 }
 
+bool BufferWithExtendableBuffer::extend(const int size) {
+    return checkAndPrepareWriting(getTailPosition(), size);
+}
+
 bool BufferWithExtendableBuffer::writeUint(const uint32_t data, const int size, const int pos) {
     int writingPos = pos;
     return writeUintAndAdvancePosition(data, size, &writingPos);
@@ -96,38 +100,49 @@ bool BufferWithExtendableBuffer::writeCodePointsAndAdvancePosition(const int *co
     return true;
 }
 
-bool BufferWithExtendableBuffer::extendBuffer() {
+bool BufferWithExtendableBuffer::extendBuffer(const size_t size) {
+    const size_t extendSize = std::max(EXTEND_ADDITIONAL_BUFFER_SIZE_STEP, size);
     const size_t sizeAfterExtending =
-            mAdditionalBuffer.size() + EXTEND_ADDITIONAL_BUFFER_SIZE_STEP;
-    if (sizeAfterExtending > mMaxAdditionalBufferSize) {
+            std::min(mAdditionalBuffer.size() + extendSize, mMaxAdditionalBufferSize);
+    if (sizeAfterExtending < mAdditionalBuffer.size() + size) {
         return false;
     }
-    mAdditionalBuffer.resize(mAdditionalBuffer.size() + EXTEND_ADDITIONAL_BUFFER_SIZE_STEP);
+    mAdditionalBuffer.resize(sizeAfterExtending);
     return true;
 }
 
 bool BufferWithExtendableBuffer::checkAndPrepareWriting(const int pos, const int size) {
-    if (isInAdditionalBuffer(pos)) {
-        const int tailPosition = getTailPosition();
-        if (pos == tailPosition) {
-            // Append data to the tail.
-            if (pos + size > static_cast<int>(mAdditionalBuffer.size()) + mOriginalBufferSize) {
-                // Need to extend buffer.
-                if (!extendBuffer()) {
-                    return false;
-                }
-            }
-            mUsedAdditionalBufferSize += size;
-        } else if (pos + size > tailPosition) {
-            // The access will beyond the tail of used region.
-            return false;
-        }
-    } else {
-        if (pos < 0 || mOriginalBufferSize < pos + size) {
-            // Invalid position or violate the boundary.
+    if (pos < 0 || size < 0) {
+        // Invalid position or size.
+        return false;
+    }
+    const size_t totalRequiredSize = static_cast<size_t>(pos + size);
+    if (!isInAdditionalBuffer(pos)) {
+        // Here don't need to care about the additional buffer.
+        if (static_cast<size_t>(mOriginalBufferSize) < totalRequiredSize) {
+            // Violate the boundary.
             return false;
         }
+        // The buffer has sufficient capacity.
+        return true;
+    }
+    // Hereafter, pos is in the additional buffer.
+    const size_t tailPosition = static_cast<size_t>(getTailPosition());
+    if (totalRequiredSize <= tailPosition) {
+        // The buffer has sufficient capacity.
+        return true;
+    }
+    if (static_cast<size_t>(pos) != tailPosition) {
+        // The additional buffer must be extended from the tail position.
+        return false;
+    }
+    const size_t extendSize = totalRequiredSize -
+            std::min(mAdditionalBuffer.size() + mOriginalBufferSize, totalRequiredSize);
+    if (extendSize > 0 && !extendBuffer(extendSize)) {
+        // Failed to extend the buffer.
+        return false;
     }
+    mUsedAdditionalBufferSize += size;
     return true;
 }
 
diff --git a/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.h b/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.h
index a2e88a46c..5e1362eee 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/utils/buffer_with_extendable_buffer.h
@@ -87,6 +87,8 @@ class BufferWithExtendableBuffer {
                 * NEAR_BUFFER_LIMIT_THRESHOLD_PERCENTILE) / 100);
     }
 
+    bool extend(const int size);
+
     /**
      * For writing.
      *
@@ -115,7 +117,7 @@ class BufferWithExtendableBuffer {
     const size_t mMaxAdditionalBufferSize;
 
     // Return if the buffer is successfully extended or not.
-    bool extendBuffer();
+    bool extendBuffer(const size_t size);
 
     // Returns if it is possible to write size-bytes from pos. When pos is at the tail position of
     // the additional buffer, try extending the buffer.