SuggestionResults.java revision 5a5ee95faead8a2ae749067716481e86faf5f113 
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.Locale;
26import java.util.TreeSet;
27
28/**
29 * A TreeSet of SuggestedWordInfo that is bounded in size and throws everything that's smaller
30 * than its limit
31 */
32public final class SuggestionResults extends TreeSet<SuggestedWordInfo> {
33    public final Locale mLocale;
34    public final ArrayList<SuggestedWordInfo> mRawSuggestions;
35    private final int mCapacity;
36
37    public SuggestionResults(final Locale locale, final int capacity) {
38        this(locale, sSuggestedWordInfoComparator, capacity);
39    }
40
41    public SuggestionResults(final Locale locale, final Comparator<SuggestedWordInfo> comparator,
42            final int capacity) {
43        super(comparator);
44        mLocale = locale;
45        mCapacity = capacity;
46        if (ProductionFlags.INCLUDE_RAW_SUGGESTIONS) {
47            mRawSuggestions = new ArrayList<>();
48        } else {
49            mRawSuggestions = null;
50        }
51    }
52
53    @Override
54    public boolean add(final SuggestedWordInfo e) {
55        if (size() < mCapacity) return super.add(e);
56        if (comparator().compare(e, last()) > 0) return false;
57        super.add(e);
58        pollLast(); // removes the last element
59        return true;
60    }
61
62    @Override
63    public boolean addAll(final Collection<? extends SuggestedWordInfo> e) {
64        if (null == e) return false;
65        return super.addAll(e);
66    }
67
68    private static final class SuggestedWordInfoComparator
69            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