1 /*
2  * Copyright (C) 2014 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.inputmethod.latin.utils;
18 
19 import com.android.inputmethod.latin.SuggestedWords.SuggestedWordInfo;
20 import com.android.inputmethod.latin.define.ProductionFlags;
21 
22 import java.util.ArrayList;
23 import java.util.Collection;
24 import java.util.Comparator;
25 import java.util.TreeSet;
26 
27 /**
28  * A TreeSet of SuggestedWordInfo that is bounded in size and throws everything that's smaller
29  * than its limit
30  */
31 public final class SuggestionResults extends TreeSet<SuggestedWordInfo> {
32     public final ArrayList<SuggestedWordInfo> mRawSuggestions;
33     // TODO: Instead of a boolean , we may want to include the context of this suggestion results,
34     // such as {@link NgramContext}.
35     public final boolean mIsBeginningOfSentence;
36     public final boolean mFirstSuggestionExceedsConfidenceThreshold;
37     private final int mCapacity;
38 
SuggestionResults(final int capacity, final boolean isBeginningOfSentence, final boolean firstSuggestionExceedsConfidenceThreshold)39     public SuggestionResults(final int capacity, final boolean isBeginningOfSentence,
40             final boolean firstSuggestionExceedsConfidenceThreshold) {
41         this(sSuggestedWordInfoComparator, capacity, isBeginningOfSentence,
42                 firstSuggestionExceedsConfidenceThreshold);
43     }
44 
SuggestionResults(final Comparator<SuggestedWordInfo> comparator, final int capacity, final boolean isBeginningOfSentence, final boolean firstSuggestionExceedsConfidenceThreshold)45     private SuggestionResults(final Comparator<SuggestedWordInfo> comparator, final int capacity,
46             final boolean isBeginningOfSentence,
47             final boolean firstSuggestionExceedsConfidenceThreshold) {
48         super(comparator);
49         mCapacity = capacity;
50         if (ProductionFlags.INCLUDE_RAW_SUGGESTIONS) {
51             mRawSuggestions = new ArrayList<>();
52         } else {
53             mRawSuggestions = null;
54         }
55         mIsBeginningOfSentence = isBeginningOfSentence;
56         mFirstSuggestionExceedsConfidenceThreshold = firstSuggestionExceedsConfidenceThreshold;
57     }
58 
59     @Override
add(final SuggestedWordInfo e)60     public boolean add(final SuggestedWordInfo e) {
61         if (size() < mCapacity) return super.add(e);
62         if (comparator().compare(e, last()) > 0) return false;
63         super.add(e);
64         pollLast(); // removes the last element
65         return true;
66     }
67 
68     @Override
addAll(final Collection<? extends SuggestedWordInfo> e)69     public boolean addAll(final Collection<? extends SuggestedWordInfo> e) {
70         if (null == e) return false;
71         return super.addAll(e);
72     }
73 
74     static final class SuggestedWordInfoComparator implements Comparator<SuggestedWordInfo> {
75         // This comparator ranks the word info with the higher frequency first. That's because
76         // that's the order we want our elements in.
77         @Override
compare(final SuggestedWordInfo o1, final SuggestedWordInfo o2)78         public int compare(final SuggestedWordInfo o1, final SuggestedWordInfo o2) {
79             if (o1.mScore > o2.mScore) return -1;
80             if (o1.mScore < o2.mScore) return 1;
81             if (o1.mCodePointCount < o2.mCodePointCount) return -1;
82             if (o1.mCodePointCount > o2.mCodePointCount) return 1;
83             return o1.mWord.compareTo(o2.mWord);
84         }
85     }
86 
87     private static final SuggestedWordInfoComparator sSuggestedWordInfoComparator =
88             new SuggestedWordInfoComparator();
89 }
90