Merge "Re-call getBuffer after writing in bigram/shortcut policy."

4 files changed, 24 insertions, 12 deletions

diff --git a/native/jni/src/suggest/policyimpl/dictionary/bigram/dynamic_bigram_list_policy.cpp b/native/jni/src/suggest/policyimpl/dictionary/bigram/dynamic_bigram_list_policy.cpp
index e31a91069..936dc9c5d 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/bigram/dynamic_bigram_list_policy.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/bigram/dynamic_bigram_list_policy.cpp
@@ -20,12 +20,13 @@ namespace latinime {
 
 bool DynamicBigramListPolicy::copyAllBigrams(int *const fromPos, int *const toPos) {
     const bool usesAdditionalBuffer = mBuffer->isInAdditionalBuffer(*fromPos);
-    const uint8_t *const buffer = mBuffer->getBuffer(usesAdditionalBuffer);
     if (usesAdditionalBuffer) {
         *fromPos -= mBuffer->getOriginalBufferSize();
     }
     BigramListReadWriteUtils::BigramFlags flags;
     do {
+        // The buffer address can be changed after calling buffer writing methods.
+        const uint8_t *const buffer = mBuffer->getBuffer(usesAdditionalBuffer);
         flags = BigramListReadWriteUtils::getFlagsAndForwardPointer(buffer, fromPos);
         int bigramPos = BigramListReadWriteUtils::getBigramAddressAndForwardPointer(
                 buffer, flags, fromPos);
@@ -63,7 +64,6 @@ bool DynamicBigramListPolicy::copyAllBigrams(int *const fromPos, int *const toPo
 bool DynamicBigramListPolicy::addBigramEntry(const int bigramPos, const int probability,
         int *const pos) {
     const bool usesAdditionalBuffer = mBuffer->isInAdditionalBuffer(*pos);
-    const uint8_t *const buffer = mBuffer->getBuffer(usesAdditionalBuffer);
     if (usesAdditionalBuffer) {
         *pos -= mBuffer->getOriginalBufferSize();
     }
@@ -73,6 +73,8 @@ bool DynamicBigramListPolicy::addBigramEntry(const int bigramPos, const int prob
         if (usesAdditionalBuffer) {
             entryPos += mBuffer->getOriginalBufferSize();
         }
+        // The buffer address can be changed after calling buffer writing methods.
+        const uint8_t *const buffer = mBuffer->getBuffer(usesAdditionalBuffer);
         flags = BigramListReadWriteUtils::getFlagsAndForwardPointer(buffer, pos);
         BigramListReadWriteUtils::getBigramAddressAndForwardPointer(buffer, flags, pos);
         if (BigramListReadWriteUtils::hasNext(flags)) {
@@ -118,13 +120,14 @@ bool DynamicBigramListPolicy::addBigramEntry(const int bigramPos, const int prob
 
 bool DynamicBigramListPolicy::removeBigram(const int bigramListPos, const int targetBigramPos) {
     const bool usesAdditionalBuffer = mBuffer->isInAdditionalBuffer(bigramListPos);
-    const uint8_t *const buffer = mBuffer->getBuffer(usesAdditionalBuffer);
     int pos = bigramListPos;
     if (usesAdditionalBuffer) {
         pos -= mBuffer->getOriginalBufferSize();
     }
     BigramListReadWriteUtils::BigramFlags flags;
     do {
+        // The buffer address can be changed after calling buffer writing methods.
+        const uint8_t *const buffer = mBuffer->getBuffer(usesAdditionalBuffer);
         flags = BigramListReadWriteUtils::getFlagsAndForwardPointer(buffer, &pos);
         int bigramOffsetFieldPos = pos;
         if (usesAdditionalBuffer) {
@@ -139,8 +142,7 @@ bool DynamicBigramListPolicy::removeBigram(const int bigramListPos, const int ta
             continue;
         }
         // Target entry is found. Write 0 into the bigram pos field to mark the bigram invalid.
-        const int bigramOffsetFieldSize =
-                BigramListReadWriteUtils::attributeAddressSize(flags);
+        const int bigramOffsetFieldSize = BigramListReadWriteUtils::attributeAddressSize(flags);
         if (!mBuffer->writeUintAndAdvancePosition(0 /* data */, bigramOffsetFieldSize,
                 &bigramOffsetFieldPos)) {
             return false;
diff --git a/native/jni/src/suggest/policyimpl/dictionary/dynamic_patricia_trie_writing_helper.cpp b/native/jni/src/suggest/policyimpl/dictionary/dynamic_patricia_trie_writing_helper.cpp
index e24421219..6d1bd02c2 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/dynamic_patricia_trie_writing_helper.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/dynamic_patricia_trie_writing_helper.cpp
@@ -186,7 +186,9 @@ bool DynamicPatriciaTrieWritingHelper::writeNodeToBuffer(const bool isBlackliste
     // Copy shortcut list when the originalShortcutListPos is valid dictionary position.
     if (originalShortcutListPos != NOT_A_DICT_POS) {
         int fromPos = originalShortcutListPos;
-        mShortcutPolicy->copyAllShortcuts(&fromPos, writingPos);
+        if (!mShortcutPolicy->copyAllShortcutsAndReturnIfSucceededOrNot(&fromPos, writingPos)) {
+            return false;
+        }
     }
     // Copy bigram list when the originalBigramListPos is valid dictionary position.
     if (originalBigramListPos != NOT_A_DICT_POS) {
diff --git a/native/jni/src/suggest/policyimpl/dictionary/shortcut/dynamic_shortcut_list_policy.h b/native/jni/src/suggest/policyimpl/dictionary/shortcut/dynamic_shortcut_list_policy.h
index 5e9c52950..1803c09cb 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/shortcut/dynamic_shortcut_list_policy.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/shortcut/dynamic_shortcut_list_policy.h
@@ -83,8 +83,8 @@ class DynamicShortcutListPolicy : public DictionaryShortcutsStructurePolicy {
     }
 
     // Copy shortcuts from the shortcut list that starts at fromPos to toPos and advance these
-    // positions after the shortcut lists.
-    void copyAllShortcuts(int *const fromPos, int *const toPos) {
+    // positions after the shortcut lists. This returns whether the copy was succeeded or not.
+    bool copyAllShortcutsAndReturnIfSucceededOrNot(int *const fromPos, int *const toPos) {
         const bool usesAdditionalBuffer = mBuffer->isInAdditionalBuffer(*fromPos);
         const uint8_t *const buffer = mBuffer->getBuffer(usesAdditionalBuffer);
         if (usesAdditionalBuffer) {
@@ -93,16 +93,23 @@ class DynamicShortcutListPolicy : public DictionaryShortcutsStructurePolicy {
         const int shortcutListSize = ShortcutListReadingUtils
                 ::getShortcutListSizeAndForwardPointer(buffer, fromPos);
         // Copy shortcut list size.
-        mBuffer->writeUintAndAdvancePosition(
+        if (!mBuffer->writeUintAndAdvancePosition(
                 shortcutListSize + ShortcutListReadingUtils::getShortcutListSizeFieldSize(),
-                ShortcutListReadingUtils::getShortcutListSizeFieldSize(), toPos);
+                ShortcutListReadingUtils::getShortcutListSizeFieldSize(), toPos)) {
+            return false;
+        }
+        // Copy shortcut list.
         for (int i = 0; i < shortcutListSize; ++i) {
-            const uint8_t data = ByteArrayUtils::readUint8AndAdvancePosition(buffer, fromPos);
-            mBuffer->writeUintAndAdvancePosition(data, 1 /* size */, toPos);
+            const uint8_t data = ByteArrayUtils::readUint8AndAdvancePosition(
+                    mBuffer->getBuffer(usesAdditionalBuffer), fromPos);
+            if (!mBuffer->writeUintAndAdvancePosition(data, 1 /* size */, toPos)) {
+                return false;
+            }
         }
         if (usesAdditionalBuffer) {
             *fromPos += mBuffer->getOriginalBufferSize();
         }
+        return true;
     }
 
  private:
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 b35b47d7a..c6a484131 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
@@ -47,6 +47,7 @@ class BufferWithExtendableBuffer {
         return position >= mOriginalBufferSize;
     }
 
+    // TODO: Resolve the issue that the address can be changed when the vector is resized.
     // CAVEAT!: Be careful about array out of bound access with buffers
     AK_FORCE_INLINE const uint8_t *getBuffer(const bool usesAdditionalBuffer) const {
         if (usesAdditionalBuffer) {