Merge "Update children's parent position when the node get moved."

3 files changed, 27 insertions, 2 deletions

diff --git a/native/jni/src/suggest/policyimpl/dictionary/dynamic_patricia_trie_node_reader.cpp b/native/jni/src/suggest/policyimpl/dictionary/dynamic_patricia_trie_node_reader.cpp
index 5674cb48e..5db4dd48e 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/dynamic_patricia_trie_node_reader.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/dynamic_patricia_trie_node_reader.cpp
@@ -28,6 +28,7 @@ void DynamicPatriciaTrieNodeReader::fetchNodeInfoFromBufferAndProcessMovedNode(c
     const bool usesAdditionalBuffer = mBuffer->isInAdditionalBuffer(nodePos);
     const uint8_t *const dictBuf = mBuffer->getBuffer(usesAdditionalBuffer);
     int pos = nodePos;
+    mHeadPos = nodePos;
     if (usesAdditionalBuffer) {
         pos -= mBuffer->getOriginalBufferSize();
     }
diff --git a/native/jni/src/suggest/policyimpl/dictionary/dynamic_patricia_trie_node_reader.h b/native/jni/src/suggest/policyimpl/dictionary/dynamic_patricia_trie_node_reader.h
index 2ee7c2495..30d251f3e 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/dynamic_patricia_trie_node_reader.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/dynamic_patricia_trie_node_reader.h
@@ -39,8 +39,8 @@ class DynamicPatriciaTrieNodeReader {
             const DictionaryBigramsStructurePolicy *const bigramsPolicy,
             const DictionaryShortcutsStructurePolicy *const shortcutsPolicy)
             : mBuffer(buffer), mBigramsPolicy(bigramsPolicy),
-              mShortcutsPolicy(shortcutsPolicy), mNodePos(NOT_A_VALID_WORD_POS), mFlags(0),
-              mParentPos(NOT_A_DICT_POS),  mCodePointCount(0),
+              mShortcutsPolicy(shortcutsPolicy), mNodePos(NOT_A_VALID_WORD_POS),
+              mHeadPos(NOT_A_DICT_POS), mFlags(0), mParentPos(NOT_A_DICT_POS), mCodePointCount(0),
               mProbabilityFieldPos(NOT_A_DICT_POS), mProbability(NOT_A_PROBABILITY),
               mChildrenPosFieldPos(NOT_A_DICT_POS), mChildrenPos(NOT_A_DICT_POS),
               mShortcutPos(NOT_A_DICT_POS), mBigramPos(NOT_A_DICT_POS),
@@ -65,6 +65,11 @@ class DynamicPatriciaTrieNodeReader {
         return mNodePos;
     }
 
+    // HeadPos is different from NodePos when the current PtNode is a moved PtNode.
+    AK_FORCE_INLINE int getHeadPos() const {
+        return mHeadPos;
+    }
+
     // Flags
     AK_FORCE_INLINE bool isDeleted() const {
         return DynamicPatriciaTrieReadingUtils::isDeleted(mFlags);
@@ -136,6 +141,7 @@ class DynamicPatriciaTrieNodeReader {
     const DictionaryBigramsStructurePolicy *const mBigramsPolicy;
     const DictionaryShortcutsStructurePolicy *const mShortcutsPolicy;
     int mNodePos;
+    int mHeadPos;
     DynamicPatriciaTrieReadingUtils::NodeFlags mFlags;
     int mParentPos;
     uint8_t mCodePointCount;
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 7c0b6286c..7dfa9ec5a 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
@@ -125,6 +125,24 @@ bool DynamicPatriciaTrieWritingHelper::markNodeAsMovedAndSetPosition(
             mBuffer, movedPosOffset, &writingPos)) {
         return false;
     }
+    if (originalNode->hasChildren()) {
+        // Update children's parent position.
+        DynamicPatriciaTrieReadingHelper readingHelper(mBuffer, mBigramPolicy, mShortcutPolicy);
+        const DynamicPatriciaTrieNodeReader *const nodeReader = readingHelper.getNodeReader();
+        readingHelper.initWithNodeArrayPos(originalNode->getChildrenPos());
+        while (!readingHelper.isEnd()) {
+            const int childPtNodeWrittenPos = nodeReader->getHeadPos();
+            const int parentOffset = movedPos - childPtNodeWrittenPos;
+            int parentOffsetFieldPos = childPtNodeWrittenPos + 1 /* Flags */;
+            if (!DynamicPatriciaTrieWritingUtils::writeParentOffsetAndAdvancePosition(
+                    mBuffer, parentOffset, &parentOffsetFieldPos)) {
+                // Parent offset cannot be written because of a bug or a broken dictionary; thus,
+                // we give up to update dictionary.
+                return false;
+            }
+            readingHelper.readNextSiblingNode();
+        }
+    }
     return true;
 }