Start to support adding shortcuts.

Bug: 11073222
Bug: 11956652

Change-Id: Iea81603a140697594cfea4f4939e82cd1d3963ca

10 files changed, 71 insertions, 10 deletions

diff --git a/native/jni/src/suggest/policyimpl/dictionary/shortcut/ver4_shortcut_list_policy.h b/native/jni/src/suggest/policyimpl/dictionary/shortcut/ver4_shortcut_list_policy.h
index b7bd08531..ae863af57 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/shortcut/ver4_shortcut_list_policy.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/shortcut/ver4_shortcut_list_policy.h
@@ -88,11 +88,13 @@ class Ver4ShortcutListPolicy : public DictionaryShortcutsStructurePolicy {
             return mShortcutDictContent->copyShortcutList(shortcutListPos, writingPos);
         }
         // Overwrite existing entry.
-        int writingPos = entryPos;
-        if (!mShortcutDictContent->writeShortcutEntryAndAdvancePosition(codePoints,
-                codePointCount, probability, true /* hasNext */, &writingPos)) {
+        bool hasNext = false;
+        mShortcutDictContent->getShortcutEntry(MAX_WORD_LENGTH, 0 /* outCodePoint */,
+                0 /* outCodePointCount */ , 0 /* probability */, &hasNext, entryPos);
+        if (!mShortcutDictContent->writeShortcutEntry(codePoints,
+                codePointCount, probability, hasNext, entryPos)) {
             AKLOGE("Cannot overwrite shortcut entry. terminal id: %d, pos: %d", terminalId,
-                    writingPos);
+                    entryPos);
             return false;
         }
         return true;
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/pt_common/pt_node_writer.h b/native/jni/src/suggest/policyimpl/dictionary/structure/pt_common/pt_node_writer.h
index 9957387bf..c7a36e796 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/pt_common/pt_node_writer.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/pt_common/pt_node_writer.h
@@ -74,6 +74,10 @@ class PtNodeWriter {
             const DictPositionRelocationMap *const dictPositionRelocationMap,
             int *const outBigramEntryCount) = 0;
 
+    virtual bool addShortcutTarget(const PtNodeParams *const ptNodeParams,
+            const int *const targetCodePoints, const int targetCodePointCount,
+            const int shortcutProbability) = 0;
+
  protected:
     PtNodeWriter() {};
 
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v3/dynamic_patricia_trie_updating_helper.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/v3/dynamic_patricia_trie_updating_helper.cpp
index e88938fcc..f12a8995e 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v3/dynamic_patricia_trie_updating_helper.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v3/dynamic_patricia_trie_updating_helper.cpp
@@ -105,6 +105,14 @@ bool DynamicPatriciaTrieUpdatingHelper::removeBigramWords(const int word0Pos, co
     return mPtNodeWriter->removeBigramEntry(&sourcePtNodeParams, &targetPtNodeParams);
 }
 
+bool DynamicPatriciaTrieUpdatingHelper::addShortcutTarget(const int wordPos,
+        const int *const targetCodePoints, const int targetCodePointCount,
+        const int shortcutProbability) {
+    const PtNodeParams ptNodeParams(mPtNodeReader->fetchNodeInfoInBufferFromPtNodePos(wordPos));
+    return mPtNodeWriter->addShortcutTarget(&ptNodeParams, targetCodePoints, targetCodePointCount,
+            shortcutProbability);
+}
+
 bool DynamicPatriciaTrieUpdatingHelper::createAndInsertNodeIntoPtNodeArray(const int parentPos,
         const int *const nodeCodePoints, const int nodeCodePointCount,
         const bool isNotAWord, const bool isBlacklisted, const int probability,
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v3/dynamic_patricia_trie_updating_helper.h b/native/jni/src/suggest/policyimpl/dictionary/structure/v3/dynamic_patricia_trie_updating_helper.h
index 799f07b6a..f02635fe2 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v3/dynamic_patricia_trie_updating_helper.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v3/dynamic_patricia_trie_updating_helper.h
@@ -52,6 +52,10 @@ class DynamicPatriciaTrieUpdatingHelper {
     // Remove a bigram relation from word0Pos to word1Pos.
     bool removeBigramWords(const int word0Pos, const int word1Pos);
 
+    // Add a shortcut target.
+    bool addShortcutTarget(const int wordPos, const int *const targetCodePoints,
+            const int targetCodePointCount, const int shortcutProbability);
+
  private:
     DISALLOW_IMPLICIT_CONSTRUCTORS(DynamicPatriciaTrieUpdatingHelper);
 
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/shortcut_dict_content.h b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/shortcut_dict_content.h
index a4f817e28..670e6eab6 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/shortcut_dict_content.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/content/shortcut_dict_content.h
@@ -38,6 +38,14 @@ class ShortcutDictContent : public SparseTableDictContent {
             : SparseTableDictContent(Ver4DictConstants::SHORTCUT_ADDRESS_TABLE_BLOCK_SIZE,
                       Ver4DictConstants::SHORTCUT_ADDRESS_TABLE_DATA_SIZE) {}
 
+    void getShortcutEntry(const int maxCodePointCount, int *const outCodePoint,
+            int *const outCodePointCount, int *const outProbability, bool *const outhasNext,
+            const int shortcutEntryPos) {
+        int readingPos = shortcutEntryPos;
+        return getShortcutEntryAndAdvancePosition(maxCodePointCount, outCodePoint,
+                outCodePointCount, outProbability, outhasNext, &readingPos);
+    }
+
     void getShortcutEntryAndAdvancePosition(const int maxCodePointCount,
             int *const outCodePoint, int *const outCodePointCount, int *const outProbability,
             bool *const outhasNext, int *const shortcutEntryPos) const;
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.cpp
index f6ea3b731..92bec30f5 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.cpp
@@ -217,6 +217,13 @@ bool Ver4PatriciaTrieNodeWriter::updateAllPositionFields(
     return true;
 }
 
+bool Ver4PatriciaTrieNodeWriter::addShortcutTarget(const PtNodeParams *const ptNodeParams,
+        const int *const targetCodePoints, const int targetCodePointCount,
+        const int shortcutProbability) {
+    return mShortcutPolicy->addNewShortcut(ptNodeParams->getTerminalId(),
+            targetCodePoints, targetCodePointCount, shortcutProbability);
+}
+
 bool Ver4PatriciaTrieNodeWriter::writePtNodeAndGetTerminalIdAndAdvancePosition(
         const PtNodeParams *const ptNodeParams, int *const outTerminalId,
         int *const ptNodeWritingPos) {
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.h b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.h
index 31b47c148..4a2a79259 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.h
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_node_writer.h
@@ -83,6 +83,10 @@ class Ver4PatriciaTrieNodeWriter : public PtNodeWriter {
             const DictPositionRelocationMap *const dictPositionRelocationMap,
             int *const outBigramEntryCount);
 
+    virtual bool addShortcutTarget(const PtNodeParams *const ptNodeParams,
+            const int *const targetCodePoints, const int targetCodePointCount,
+            const int shortcutProbability);
+
  private:
     DISALLOW_COPY_AND_ASSIGN(Ver4PatriciaTrieNodeWriter);
 
diff --git a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp
index d1ba1877c..9c8db3bac 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/structure/v4/ver4_patricia_trie_policy.cpp
@@ -155,12 +155,26 @@ bool Ver4PatriciaTriePolicy::addUnigramWord(const int *const word, const int len
     DynamicPatriciaTrieReadingHelper readingHelper(mDictBuffer, &mNodeReader);
     readingHelper.initWithPtNodeArrayPos(getRootPosition());
     bool addedNewUnigram = false;
-    // TODO: Add shortcut.
     if (mUpdatingHelper.addUnigramWord(&readingHelper, word, length, probability, isNotAWord,
             isBlacklisted, timestamp,  &addedNewUnigram)) {
         if (addedNewUnigram) {
             mUnigramCount++;
         }
+        if (shortcutLength > 0) {
+            // Add shortcut target.
+            const int wordPos = getTerminalPtNodePositionOfWord(word, length,
+                    false /* forceLowerCaseSearch */);
+            if (wordPos == NOT_A_DICT_POS) {
+                AKLOGE("Cannot find terminal PtNode position to add shortcut target.");
+                return false;
+            }
+            if (!mUpdatingHelper.addShortcutTarget(wordPos, shortcutTargetCodePoints,
+                    shortcutLength, shortcutProbability)) {
+                AKLOGE("Cannot add new shortcut target. PtNodePos: %d, length: %d, probability: %d",
+                        wordPos, shortcutLength, shortcutProbability);
+                return false;
+            }
+        }
         return true;
     } else {
         return false;
@@ -308,12 +322,12 @@ const UnigramProperty Ver4PatriciaTriePolicy::getUnigramProperty(const int *cons
     // Fetch shortcut information.
     std::vector<std::vector<int> > shortcutTargets;
     std::vector<int> shortcutProbabilities;
-    if (ptNodeParams.hasShortcutTargets()) {
+    int shortcutPos = getShortcutPositionOfPtNode(ptNodePos);
+    if (shortcutPos != NOT_A_DICT_POS) {
         int shortcutTarget[MAX_WORD_LENGTH];
         const ShortcutDictContent *const shortcutDictContent =
                 mBuffers.get()->getShortcutDictContent();
         bool hasNext = true;
-        int shortcutPos = getShortcutPositionOfPtNode(ptNodePos);
         while (hasNext) {
             int shortcutTargetLength = 0;
             int shortcutProbability = NOT_A_PROBABILITY;
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 4b537da8a..259dae4c6 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
@@ -43,7 +43,7 @@ void BufferWithExtendableBuffer::readCodePointsAndAdvancePosition(const int maxC
         *pos -= mOriginalBufferSize;
     }
     *outCodePointCount = ByteArrayUtils::readStringAndAdvancePosition(
-            getBuffer(readingPosIsInAdditionalBuffer), maxCodePointCount, outCodePointCount, pos);
+            getBuffer(readingPosIsInAdditionalBuffer), maxCodePointCount, outCodePoints, pos);
     if (readingPosIsInAdditionalBuffer) {
         *pos += mOriginalBufferSize;
     }
diff --git a/native/jni/src/suggest/policyimpl/dictionary/utils/sparse_table.cpp b/native/jni/src/suggest/policyimpl/dictionary/utils/sparse_table.cpp
index 9be35620c..4ad82f9f7 100644
--- a/native/jni/src/suggest/policyimpl/dictionary/utils/sparse_table.cpp
+++ b/native/jni/src/suggest/policyimpl/dictionary/utils/sparse_table.cpp
@@ -44,6 +44,7 @@ bool SparseTable::set(const int id, const uint32_t value) {
         int tailPos = mIndexTableBuffer->getTailPosition();
         while(tailPos < posInIndexTable) {
             if (!mIndexTableBuffer->writeUintAndAdvancePosition(NOT_EXIST, INDEX_SIZE, &tailPos)) {
+                AKLOGE("cannot extend index table. tailPos: %d to: %d", tailPos, posInIndexTable);
                 return false;
             }
         }
@@ -51,12 +52,19 @@ bool SparseTable::set(const int id, const uint32_t value) {
     if (contains(id)) {
         // The entry is already in the content table.
         const int index = mIndexTableBuffer->readUint(INDEX_SIZE, posInIndexTable);
-        return mContentTableBuffer->writeUint(value, mDataSize, getPosInContentTable(id, index));
+        if (!mContentTableBuffer->writeUint(value, mDataSize, getPosInContentTable(id, index))) {
+            AKLOGE("cannot update value %d. pos: %d, tailPos: %d, mDataSize: %d", value,
+                    getPosInContentTable(id, index), mContentTableBuffer->getTailPosition(),
+                    mDataSize);
+            return false;
+        }
+        return true;
     }
     // The entry is not in the content table.
     // Create new entry in the content table.
     const int index = getIndexFromContentTablePos(mContentTableBuffer->getTailPosition());
     if (!mIndexTableBuffer->writeUint(index, INDEX_SIZE, posInIndexTable)) {
+        AKLOGE("cannot write index %d. pos %d", index, posInIndexTable);
         return false;
     }
     // Write a new block that containing the entry to be set.
@@ -64,6 +72,8 @@ bool SparseTable::set(const int id, const uint32_t value) {
     for (int i = 0; i < mBlockSize; ++i) {
         if (!mContentTableBuffer->writeUintAndAdvancePosition(NOT_A_DICT_POS, mDataSize,
                 &writingPos)) {
+            AKLOGE("cannot write content table to extend. writingPos: %d, tailPos: %d, "
+                    "mDataSize: %d", writingPos, mContentTableBuffer->getTailPosition(), mDataSize);
             return false;
         }
     }
@@ -80,7 +90,7 @@ int SparseTable::getPosInIndexTable(const int id) const {
 
 int SparseTable::getPosInContentTable(const int id, const int index) const {
     const int offset = id % mBlockSize;
-    return (index * mDataSize + offset) * mBlockSize;
+    return (index * mBlockSize + offset) * mDataSize;
 }
 
 } // namespace latinime