Use a TreeSet to sort suggestions instead of doing it by hand (A1)

Change-Id: I16ba39321107e87ad48a99c2410a15995a66f23c

2 files changed, 57 insertions, 28 deletions

diff --git a/java/src/com/android/inputmethod/latin/BoundedTreeSet.java b/java/src/com/android/inputmethod/latin/BoundedTreeSet.java
new file mode 100644
index 000000000..cf977617d
--- /dev/null
+++ b/java/src/com/android/inputmethod/latin/BoundedTreeSet.java
@@ -0,0 +1,49 @@
+/*
+ * Copyright (C) 2012 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package com.android.inputmethod.latin;
+
+import com.android.inputmethod.latin.SuggestedWords.SuggestedWordInfo;
+
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.TreeSet;
+
+/**
+ * A TreeSet that is bounded in size and throws everything that's smaller than its limit
+ */
+public class BoundedTreeSet extends TreeSet<SuggestedWordInfo> {
+    private final int mCapacity;
+    public BoundedTreeSet(final Comparator<SuggestedWordInfo> comparator, final int capacity) {
+        super(comparator);
+        mCapacity = capacity;
+    }
+
+    @Override
+    public boolean add(final SuggestedWordInfo e) {
+        if (size() < mCapacity) return super.add(e);
+        if (comparator().compare(e, last()) > 0) return false;
+        super.add(e);
+        pollLast(); // removes the last element
+        return true;
+    }
+
+    @Override
+    public boolean addAll(final Collection<? extends SuggestedWordInfo> e) {
+        if (null == e) return false;
+        return super.addAll(e);
+    }
+}
diff --git a/java/src/com/android/inputmethod/latin/Suggest.java b/java/src/com/android/inputmethod/latin/Suggest.java
index dbdb24ada..08a7b869d 100644
--- a/java/src/com/android/inputmethod/latin/Suggest.java
+++ b/java/src/com/android/inputmethod/latin/Suggest.java
@@ -26,7 +26,6 @@ import com.android.inputmethod.latin.SuggestedWords.SuggestedWordInfo;
 
 import java.io.File;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.Comparator;
 import java.util.HashSet;
 import java.util.Locale;
@@ -190,8 +189,8 @@ public class Suggest {
                 !isPrediction && wordComposer.isFirstCharCapitalized();
         final boolean isAllUpperCase = !isPrediction && wordComposer.isAllUpperCase();
         final int trailingSingleQuotesCount = wordComposer.trailingSingleQuotesCount();
-        final ArrayList<SuggestedWordInfo> suggestionsContainer =
-                new ArrayList<SuggestedWordInfo>(MAX_SUGGESTIONS);
+        final BoundedTreeSet suggestionsSet = new BoundedTreeSet(sSuggestedWordInfoComparator,
+                MAX_SUGGESTIONS);
 
         final String typedWord = wordComposer.getTypedWord();
         final String consideredWord = trailingSingleQuotesCount > 0
@@ -211,13 +210,9 @@ public class Suggest {
                 }
                 for (final String key : mDictionaries.keySet()) {
                     final Dictionary dictionary = mDictionaries.get(key);
-                    final ArrayList<SuggestedWordInfo> localSuggestions =
-                            dictionary.getBigrams(wordComposer, prevWordForBigram);
+                    suggestionsSet.addAll(dictionary.getBigrams(wordComposer, prevWordForBigram));
                     if (null != lowerPrevWord) {
-                        localSuggestions.addAll(dictionary.getBigrams(wordComposer, lowerPrevWord));
-                    }
-                    for (final SuggestedWordInfo localSuggestion : localSuggestions) {
-                        addWord(localSuggestion, suggestionsContainer);
+                        suggestionsSet.addAll(dictionary.getBigrams(wordComposer, lowerPrevWord));
                     }
                 }
             }
@@ -238,14 +233,13 @@ public class Suggest {
                         || key.equals(Dictionary.TYPE_WHITELIST))
                     continue;
                 final Dictionary dictionary = mDictionaries.get(key);
-                final ArrayList<SuggestedWordInfo> localSuggestions = dictionary.getWords(
-                        wordComposerForLookup, prevWordForBigram, proximityInfo);
-                for (final SuggestedWordInfo suggestion : localSuggestions) {
-                    addWord(suggestion, suggestionsContainer);
-                }
+                suggestionsSet.addAll(dictionary.getWords(
+                        wordComposerForLookup, prevWordForBigram, proximityInfo));
             }
         }
 
+        final ArrayList<SuggestedWordInfo> suggestionsContainer =
+                new ArrayList<SuggestedWordInfo>(suggestionsSet);
         for (int i = 0; i < suggestionsContainer.size(); ++i) {
             final SuggestedWordInfo wordInfo = suggestionsContainer.get(i);
             final SuggestedWordInfo transformedWordInfo = getTransformedSuggestedWordInfo(wordInfo,
@@ -375,20 +369,6 @@ public class Suggest {
     private static final SuggestedWordInfoComparator sSuggestedWordInfoComparator =
             new SuggestedWordInfoComparator();
 
-    private static void addWord(final SuggestedWordInfo wordInfo,
-            final ArrayList<SuggestedWordInfo> suggestions) {
-        final int index =
-                Collections.binarySearch(suggestions, wordInfo, sSuggestedWordInfoComparator);
-        // binarySearch returns the index of an equal word info if found. If not found
-        // it returns -insertionPoint - 1. We want the insertion point, so:
-        final int pos = index >= 0 ? index : -index - 1;
-        if (pos >= MAX_SUGGESTIONS) return;
-        suggestions.add(pos, wordInfo);
-        if (suggestions.size() > MAX_SUGGESTIONS) {
-            suggestions.remove(MAX_SUGGESTIONS);
-        }
-    }
-
     private static SuggestedWordInfo getTransformedSuggestedWordInfo(
             final SuggestedWordInfo wordInfo, final Locale locale, final boolean isAllUpperCase,
             final boolean isFirstCharCapitalized, final int trailingSingleQuotesCount) {