1 files changed, 77 insertions, 13 deletions

diff --git a/src/com/android/tv/dvr/ui/SortedArrayAdapter.java b/src/com/android/tv/dvr/ui/SortedArrayAdapter.java
index 393a5ff3..8c0af9ed 100644
--- a/src/com/android/tv/dvr/ui/SortedArrayAdapter.java
+++ b/src/com/android/tv/dvr/ui/SortedArrayAdapter.java
@@ -20,11 +20,15 @@ import android.support.annotation.VisibleForTesting;
 import android.support.v17.leanback.widget.ArrayObjectAdapter;
 import android.support.v17.leanback.widget.PresenterSelector;
 
+import com.android.tv.common.SoftPreconditions;
+
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 /**
  * Keeps a set of items sorted
@@ -35,16 +39,18 @@ public abstract class SortedArrayAdapter<T> extends ArrayObjectAdapter {
     private final Comparator<T> mComparator;
     private final int mMaxItemCount;
     private int mExtraItemCount;
+    private final Set<Long> mIds = new HashSet<>();
 
-    SortedArrayAdapter(PresenterSelector presenterSelector, Comparator<T> comparator) {
+    public SortedArrayAdapter(PresenterSelector presenterSelector, Comparator<T> comparator) {
         this(presenterSelector, comparator, Integer.MAX_VALUE);
     }
 
-    SortedArrayAdapter(PresenterSelector presenterSelector, Comparator<T> comparator,
+    public SortedArrayAdapter(PresenterSelector presenterSelector, Comparator<T> comparator,
             int maxItemCount) {
         super(presenterSelector);
         mComparator = comparator;
         mMaxItemCount = maxItemCount;
+        setHasStableIds(true);
     }
 
     /**
@@ -56,7 +62,12 @@ public abstract class SortedArrayAdapter<T> extends ArrayObjectAdapter {
     final void setInitialItems(List<T> items) {
         List<T> itemsCopy = new ArrayList<>(items);
         Collections.sort(itemsCopy, mComparator);
-        addAll(0, itemsCopy.subList(0, Math.min(mMaxItemCount, itemsCopy.size())));
+        for (T item : itemsCopy) {
+            add(item, true);
+            if (size() == mMaxItemCount) {
+                break;
+            }
+        }
     }
 
     /**
@@ -82,6 +93,9 @@ public abstract class SortedArrayAdapter<T> extends ArrayObjectAdapter {
      *                    the end to save search time.
      */
     public final void add(T item, boolean insertToEnd) {
+        long newItemId = getId(item);
+        SoftPreconditions.checkState(!mIds.contains(newItemId));
+        mIds.add(newItemId);
         int i;
         if (insertToEnd) {
             i = findInsertPosition(item);
@@ -89,8 +103,9 @@ public abstract class SortedArrayAdapter<T> extends ArrayObjectAdapter {
             i = findInsertPositionBinary(item);
         }
         super.add(i, item);
-        if (size() > mMaxItemCount + mExtraItemCount) {
-            removeItems(mMaxItemCount, size() - mMaxItemCount - mExtraItemCount);
+        if (mMaxItemCount < Integer.MAX_VALUE && size() > mMaxItemCount + mExtraItemCount) {
+            Object removedItem = get(mMaxItemCount);
+            remove(removedItem);
         }
     }
 
@@ -100,48 +115,97 @@ public abstract class SortedArrayAdapter<T> extends ArrayObjectAdapter {
      * They will be presented in their insertion order.
      */
     public int addExtraItem(T item) {
+        long newItemId = getId(item);
+        SoftPreconditions.checkState(!mIds.contains(newItemId));
+        mIds.add(newItemId);
         super.add(item);
         return ++mExtraItemCount;
     }
 
+    @Override
+    public boolean remove(Object item) {
+        return removeWithId((T) item);
+    }
+
     /**
      * Removes an item which has the same ID as {@code item}.
      */
     public boolean removeWithId(T item) {
-        int index = indexWithTypeAndId(item);
-        return index >= 0 && index < size() && remove(get(index));
+        int index = indexWithId(item);
+        return index >= 0 && index < size() && removeItems(index, 1) == 1;
+    }
+
+    @Override
+    public int removeItems(int position, int count) {
+        int upperBound = Math.min(position + count, size());
+        for (int i = position; i < upperBound; i++) {
+            mIds.remove(getId((T) get(i)));
+        }
+        if (upperBound > size() - mExtraItemCount) {
+            mExtraItemCount -= upperBound - Math.max(size() - mExtraItemCount, position);
+        }
+        return super.removeItems(position, count);
+    }
+
+    @Override
+    public void replace(int position, Object item) {
+        boolean wasExtra = position >= size() - mExtraItemCount;
+        removeItems(position, 1);
+        if (!wasExtra) {
+            add(item);
+        } else {
+            addExtraItem((T) item);
+        }
+    }
+
+    @Override
+    public void clear() {
+        mIds.clear();
+        super.clear();
     }
 
     /**
-     * Change an item in the list.
+     * Changes an item in the list.
      * @param item The item to change.
      */
     public final void change(T item) {
-        int oldIndex = indexWithTypeAndId(item);
+        int oldIndex = indexWithId(item);
         if (oldIndex != -1) {
             T old = (T) get(oldIndex);
             if (mComparator.compare(old, item) == 0) {
                 replace(oldIndex, item);
                 return;
             }
-            removeItems(oldIndex, 1);
+            remove(old);
         }
         add(item);
     }
 
     /**
+     * Checks whether the item is in the list.
+     */
+    public final boolean contains(T item) {
+        return indexWithId(item) != -1;
+    }
+
+    @Override
+    public long getId(int position) {
+        return getId((T) get(position));
+    }
+
+    /**
      * Returns the id of the the given {@code item}, which will be used in {@link #change} to
      * decide if the given item is already existed in the adapter.
      *
      * The id must be stable.
      */
-    abstract long getId(T item);
+    protected abstract long getId(T item);
 
-    private int indexWithTypeAndId(T item) {
+    private int indexWithId(T item) {
         long id = getId(item);
         for (int i = 0; i < size() - mExtraItemCount; i++) {
             T r = (T) get(i);
-            if (r.getClass() == item.getClass() && getId(r) == id) {
+            if (getId(r) == id) {
                 return i;
             }
         }