2 files changed, 12 insertions, 6 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 6e88bed9e..fab7208de 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,13 +100,13 @@ bool BufferWithExtendableBuffer::writeCodePointsAndAdvancePosition(const int *co
     return true;
 }
 
-bool BufferWithExtendableBuffer::extendBuffer() {
-    const size_t sizeAfterExtending =
-            mAdditionalBuffer.size() + EXTEND_ADDITIONAL_BUFFER_SIZE_STEP;
+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() + extendSize;
     if (sizeAfterExtending > mMaxAdditionalBufferSize) {
         return false;
     }
-    mAdditionalBuffer.resize(mAdditionalBuffer.size() + EXTEND_ADDITIONAL_BUFFER_SIZE_STEP);
+    mAdditionalBuffer.resize(sizeAfterExtending);
     return true;
 }
 
@@ -133,7 +137,7 @@ bool BufferWithExtendableBuffer::checkAndPrepareWriting(const int pos, const int
     }
     const size_t extendSize = totalRequiredSize -
             std::min(mAdditionalBuffer.size() + mOriginalBufferSize, totalRequiredSize);
-    if (extendSize > 0 && !extendBuffer()) {
+    if (extendSize > 0 && !extendBuffer(extendSize)) {
         // Failed to extend the buffer.
         return false;
     }
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.