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    public final boolean mFirstSuggestionExceedsConfidenceThreshold;
37    private final int mCapacity;
38
39    public SuggestionResults(final int capacity, final boolean isBeginningOfSentence,
40            final boolean firstSuggestionExceedsConfidenceThreshold) {
41        this(sSuggestedWordInfoComparator, capacity, isBeginningOfSentence,
42                firstSuggestionExceedsConfidenceThreshold);
43    }
44
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
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
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
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