Change cache capacity depending on the dictionary size.

Bug: 10699291

Change-Id: I7042d4c1307da2d991d4dd10d637f18026acb996

8 files changed, 48 insertions, 21 deletions

diff --git a/java/src/com/android/inputmethod/latin/BinaryDictionary.java b/java/src/com/android/inputmethod/latin/BinaryDictionary.java
index e8b06570f..834d3ed53 100644
--- a/java/src/com/android/inputmethod/latin/BinaryDictionary.java
+++ b/java/src/com/android/inputmethod/latin/BinaryDictionary.java
@@ -43,6 +43,7 @@ public final class BinaryDictionary extends Dictionary {
 
     private long mNativeDict;
     private final Locale mLocale;
+    private final long mDictSize;
     private final int[] mInputCodePoints = new int[MAX_WORD_LENGTH];
     private final int[] mOutputCodePoints = new int[MAX_WORD_LENGTH * MAX_RESULTS];
     private final int[] mSpaceIndices = new int[MAX_RESULTS];
@@ -62,7 +63,7 @@ public final class BinaryDictionary extends Dictionary {
             if (traverseSession == null) {
                 traverseSession = mDicTraverseSessions.get(traverseSessionId);
                 if (traverseSession == null) {
-                    traverseSession = new DicTraverseSession(mLocale, mNativeDict);
+                    traverseSession = new DicTraverseSession(mLocale, mNativeDict, mDictSize);
                     mDicTraverseSessions.put(traverseSessionId, traverseSession);
                 }
             }
@@ -85,6 +86,7 @@ public final class BinaryDictionary extends Dictionary {
             final boolean isUpdatable) {
         super(dictType);
         mLocale = locale;
+        mDictSize = length;
         mNativeSuggestOptions.setUseFullEditDistance(useFullEditDistance);
         loadDictionary(filename, offset, length, isUpdatable);
     }
diff --git a/java/src/com/android/inputmethod/latin/DicTraverseSession.java b/java/src/com/android/inputmethod/latin/DicTraverseSession.java
index 45b281318..8d295adee 100644
--- a/java/src/com/android/inputmethod/latin/DicTraverseSession.java
+++ b/java/src/com/android/inputmethod/latin/DicTraverseSession.java
@@ -25,16 +25,16 @@ public final class DicTraverseSession {
         JniUtils.loadNativeLibrary();
     }
 
-    private static native long setDicTraverseSessionNative(String locale);
+    private static native long setDicTraverseSessionNative(String locale, long dictSize);
     private static native void initDicTraverseSessionNative(long nativeDicTraverseSession,
             long dictionary, int[] previousWord, int previousWordLength);
     private static native void releaseDicTraverseSessionNative(long nativeDicTraverseSession);
 
     private long mNativeDicTraverseSession;
 
-    public DicTraverseSession(Locale locale, long dictionary) {
+    public DicTraverseSession(Locale locale, long dictionary, long dictSize) {
         mNativeDicTraverseSession = createNativeDicTraverseSession(
-                locale != null ? locale.toString() : "");
+                locale != null ? locale.toString() : "", dictSize);
         initSession(dictionary);
     }
 
@@ -51,8 +51,8 @@ public final class DicTraverseSession {
                 mNativeDicTraverseSession, dictionary, previousWord, previousWordLength);
     }
 
-    private final long createNativeDicTraverseSession(String locale) {
-        return setDicTraverseSessionNative(locale);
+    private final long createNativeDicTraverseSession(String locale, long dictSize) {
+        return setDicTraverseSessionNative(locale, dictSize);
     }
 
     private void closeInternal() {
diff --git a/native/jni/com_android_inputmethod_latin_DicTraverseSession.cpp b/native/jni/com_android_inputmethod_latin_DicTraverseSession.cpp
index 72e625836..386643332 100644
--- a/native/jni/com_android_inputmethod_latin_DicTraverseSession.cpp
+++ b/native/jni/com_android_inputmethod_latin_DicTraverseSession.cpp
@@ -25,8 +25,9 @@
 
 namespace latinime {
 class Dictionary;
-static jlong latinime_setDicTraverseSession(JNIEnv *env, jclass clazz, jstring localeJStr) {
-    void *traverseSession = DicTraverseSession::getSessionInstance(env, localeJStr);
+static jlong latinime_setDicTraverseSession(JNIEnv *env, jclass clazz, jstring localeJStr,
+        jlong dictSize) {
+    void *traverseSession = DicTraverseSession::getSessionInstance(env, localeJStr, dictSize);
     return reinterpret_cast<jlong>(traverseSession);
 }
 
@@ -53,7 +54,7 @@ static void latinime_releaseDicTraverseSession(JNIEnv *env, jclass clazz, jlong
 static const JNINativeMethod sMethods[] = {
     {
         const_cast<char *>("setDicTraverseSessionNative"),
-        const_cast<char *>("(Ljava/lang/String;)J"),
+        const_cast<char *>("(Ljava/lang/String;J)J"),
         reinterpret_cast<void *>(latinime_setDicTraverseSession)
     },
     {
diff --git a/native/jni/src/defines.h b/native/jni/src/defines.h
index 07f1e52c6..4605890c7 100644
--- a/native/jni/src/defines.h
+++ b/native/jni/src/defines.h
@@ -32,8 +32,6 @@
 #define MAX_WORD_LENGTH 48
 // Must be equal to BinaryDictionary.MAX_RESULTS in Java
 #define MAX_RESULTS 18
-// The biggest value among MAX_CACHE_DIC_NODE_SIZE, MAX_CACHE_DIC_NODE_SIZE_FOR_SINGLE_POINT, ...
-#define MAX_DIC_NODE_PRIORITY_QUEUE_CAPACITY 310
 // Must be equal to ProximityInfo.MAX_PROXIMITY_CHARS_SIZE in Java
 #define MAX_PROXIMITY_CHARS_SIZE 16
 #define ADDITIONAL_PROXIMITY_CHAR_DELIMITER_CODE 2
diff --git a/native/jni/src/suggest/core/dicnode/dic_nodes_cache.cpp b/native/jni/src/suggest/core/dicnode/dic_nodes_cache.cpp
index c3d2a2e74..b6be47e90 100644
--- a/native/jni/src/suggest/core/dicnode/dic_nodes_cache.cpp
+++ b/native/jni/src/suggest/core/dicnode/dic_nodes_cache.cpp
@@ -23,6 +23,11 @@
 
 namespace latinime {
 
+// The biggest value among MAX_CACHE_DIC_NODE_SIZE, MAX_CACHE_DIC_NODE_SIZE_FOR_SINGLE_POINT, ...
+const int DicNodesCache::LARGE_PRIORITY_QUEUE_CAPACITY = 310;
+// Capacity for reducing memory footprint.
+const int DicNodesCache::SMALL_PRIORITY_QUEUE_CAPACITY = 100;
+
 /**
  * Truncates all of the dicNodes so that they start at the given commit point.
  * Only called for multi-word typing input.
diff --git a/native/jni/src/suggest/core/dicnode/dic_nodes_cache.h b/native/jni/src/suggest/core/dicnode/dic_nodes_cache.h
index f085848aa..8493b6a8b 100644
--- a/native/jni/src/suggest/core/dicnode/dic_nodes_cache.h
+++ b/native/jni/src/suggest/core/dicnode/dic_nodes_cache.h
@@ -31,10 +31,11 @@ class DicNode;
  */
 class DicNodesCache {
  public:
-    AK_FORCE_INLINE DicNodesCache()
-            : mDicNodePriorityQueue0(MAX_DIC_NODE_PRIORITY_QUEUE_CAPACITY),
-              mDicNodePriorityQueue1(MAX_DIC_NODE_PRIORITY_QUEUE_CAPACITY),
-              mDicNodePriorityQueue2(MAX_DIC_NODE_PRIORITY_QUEUE_CAPACITY),
+    AK_FORCE_INLINE explicit DicNodesCache(const bool usesLargeCapacityCache)
+            : mUsesLargeCapacityCache(usesLargeCapacityCache),
+              mDicNodePriorityQueue0(getCacheCapacity()),
+              mDicNodePriorityQueue1(getCacheCapacity()),
+              mDicNodePriorityQueue2(getCacheCapacity()),
               mDicNodePriorityQueueForTerminal(MAX_RESULTS),
               mActiveDicNodes(&mDicNodePriorityQueue0),
               mNextActiveDicNodes(&mDicNodePriorityQueue1),
@@ -50,7 +51,8 @@ class DicNodesCache {
         // We want to use the max capacity for the current active dic node queue.
         mActiveDicNodes->clearAndResizeToCapacity();
         // nextActiveSize is used to limit the next iteration's active dic node size.
-        mNextActiveDicNodes->clearAndResize(nextActiveSize);
+        const int nextActiveSizeFittingToTheCapacity = min(nextActiveSize, getCacheCapacity());
+        mNextActiveDicNodes->clearAndResize(nextActiveSizeFittingToTheCapacity);
         mTerminalDicNodes->clearAndResize(terminalSize);
         // We want to use the max capacity for the cached dic nodes that will be used for the
         // continuous suggestion.
@@ -162,12 +164,21 @@ class DicNodesCache {
         return tmp;
     }
 
+    AK_FORCE_INLINE int getCacheCapacity() const {
+        return mUsesLargeCapacityCache ?
+                LARGE_PRIORITY_QUEUE_CAPACITY : SMALL_PRIORITY_QUEUE_CAPACITY;
+    }
+
     AK_FORCE_INLINE void resetTemporaryCaches() {
         mActiveDicNodes->clear();
         mNextActiveDicNodes->clear();
         mTerminalDicNodes->clear();
     }
 
+    static const int LARGE_PRIORITY_QUEUE_CAPACITY;
+    static const int SMALL_PRIORITY_QUEUE_CAPACITY;
+
+    const bool mUsesLargeCapacityCache;
     // Instances
     DicNodePriorityQueue mDicNodePriorityQueue0;
     DicNodePriorityQueue mDicNodePriorityQueue1;
diff --git a/native/jni/src/suggest/core/session/dic_traverse_session.cpp b/native/jni/src/suggest/core/session/dic_traverse_session.cpp
index e7b386b2d..2c2259214 100644
--- a/native/jni/src/suggest/core/session/dic_traverse_session.cpp
+++ b/native/jni/src/suggest/core/session/dic_traverse_session.cpp
@@ -23,6 +23,11 @@
 
 namespace latinime {
 
+// 256K bytes threshold is heuristically used to distinguish dictionaries containing many unigrams
+// (e.g. main dictionary) from small dictionaries (e.g. contacts...)
+const int DicTraverseSession::DICTIONARY_SIZE_THRESHOLD_TO_USE_LARGE_CACHE_FOR_SUGGESTION =
+        256 * 1024;
+
 void DicTraverseSession::init(const Dictionary *const dictionary, const int *prevWord,
         int prevWordLength, const SuggestOptions *const suggestOptions) {
     mDictionary = dictionary;
diff --git a/native/jni/src/suggest/core/session/dic_traverse_session.h b/native/jni/src/suggest/core/session/dic_traverse_session.h
index b25580b96..fe8893590 100644
--- a/native/jni/src/suggest/core/session/dic_traverse_session.h
+++ b/native/jni/src/suggest/core/session/dic_traverse_session.h
@@ -37,8 +37,12 @@ class DicTraverseSession {
  public:
 
     // A factory method for DicTraverseSession
-    static AK_FORCE_INLINE void *getSessionInstance(JNIEnv *env, jstring localeStr) {
-        return new DicTraverseSession(env, localeStr);
+    static AK_FORCE_INLINE void *getSessionInstance(JNIEnv *env, jstring localeStr,
+            jlong dictSize) {
+        // To deal with the trade-off between accuracy and memory space, large cache is used for
+        // dictionaries larger that the threshold
+        return new DicTraverseSession(env, localeStr,
+                dictSize >= DICTIONARY_SIZE_THRESHOLD_TO_USE_LARGE_CACHE_FOR_SUGGESTION);
     }
 
     static AK_FORCE_INLINE void initSessionInstance(DicTraverseSession *traverseSession,
@@ -54,10 +58,10 @@ class DicTraverseSession {
         delete traverseSession;
     }
 
-    AK_FORCE_INLINE DicTraverseSession(JNIEnv *env, jstring localeStr)
+    AK_FORCE_INLINE DicTraverseSession(JNIEnv *env, jstring localeStr, bool usesLargeCache)
             : mPrevWordPos(NOT_A_VALID_WORD_POS), mProximityInfo(0),
-              mDictionary(0), mSuggestOptions(0), mDicNodesCache(), mMultiBigramMap(),
-              mInputSize(0), mPartiallyCommited(false), mMaxPointerCount(1),
+              mDictionary(0), mSuggestOptions(0), mDicNodesCache(usesLargeCache),
+              mMultiBigramMap(), mInputSize(0), mPartiallyCommited(false), mMaxPointerCount(1),
               mMultiWordCostMultiplier(1.0f) {
         // NOTE: mProximityInfoStates is an array of instances.
         // No need to initialize it explicitly here.
@@ -181,6 +185,7 @@ class DicTraverseSession {
     DISALLOW_IMPLICIT_CONSTRUCTORS(DicTraverseSession);
     // threshold to start caching
     static const int CACHE_START_INPUT_LENGTH_THRESHOLD;
+    static const int DICTIONARY_SIZE_THRESHOLD_TO_USE_LARGE_CACHE_FOR_SUGGESTION;
     void initializeProximityInfoStates(const int *const inputCodePoints, const int *const inputXs,
             const int *const inputYs, const int *const times, const int *const pointerIds,
             const int inputSize, const float maxSpatialDistance, const int maxPointerCount);