SuggestionResults.java revision 5f00fe09e9a611b647592188316e5999465df4d3 
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
17package com.android.inputmethod.latin.utils;
18
19import com.android.inputmethod.latin.SuggestedWords.SuggestedWordInfo;
20import com.android.inputmethod.latin.define.ProductionFlags;
21
22import java.util.ArrayList;
23import java.util.Collection;
24import java.util.Comparator;
25import 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 */
31public 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    private final int mCapacity;
37
38    public SuggestionResults(final int capacity, final boolean isBeginningOfSentence) {
39        this(sSuggestedWordInfoComparator, capacity, isBeginningOfSentence);
40    }
41
42    private SuggestionResults(final Comparator<SuggestedWordInfo> comparator,
43            final int capacity, final boolean isBeginningOfSentence) {
44        super(comparator);
45        mCapacity = capacity;
46        if (ProductionFlags.INCLUDE_RAW_SUGGESTIONS) {
47            mRawSuggestions = new ArrayList<>();
48        } else {
49            mRawSuggestions = null;
50        }
51        mIsBeginningOfSentence = isBeginningOfSentence;
52    }
53
54    @Override
55    public boolean add(final SuggestedWordInfo e) {
56        if (size() < mCapacity) return super.add(e);
57        if (comparator().compare(e, last()) > 0) return false;
58        super.add(e);
59        pollLast(); // removes the last element
60        return true;
61    }
62
63    @Override
64    public boolean addAll(final Collection<? extends SuggestedWordInfo> e) {
65        if (null == e) return false;
66        return super.addAll(e);
67    }
68
69    static final class SuggestedWordInfoComparator implements Comparator<SuggestedWordInfo> {
70        // This comparator ranks the word info with the higher frequency first. That's because
71        // that's the order we want our elements in.
72        @Override
73        public int compare(final SuggestedWordInfo o1, final SuggestedWordInfo o2) {
74            if (o1.mScore > o2.mScore) return -1;
75            if (o1.mScore < o2.mScore) return 1;
76            if (o1.mCodePointCount < o2.mCodePointCount) return -1;
77            if (o1.mCodePointCount > o2.mCodePointCount) return 1;
78            return o1.mWord.compareTo(o2.mWord);
79        }
80    }
81
82    private static final SuggestedWordInfoComparator sSuggestedWordInfoComparator =
83            new SuggestedWordInfoComparator();
84}
85