1 files changed, 129 insertions, 76 deletions
diff --git a/guava/src/com/google/common/collect/RegularImmutableSortedMultiset.java b/guava/src/com/google/common/collect/RegularImmutableSortedMultiset.java
index 3db3405..7318d9b 100644
--- a/guava/src/com/google/common/collect/RegularImmutableSortedMultiset.java
+++ b/guava/src/com/google/common/collect/RegularImmutableSortedMultiset.java
@@ -15,11 +15,16 @@
 package com.google.common.collect;
 
 import static com.google.common.base.Preconditions.checkNotNull;
-import static com.google.common.base.Preconditions.checkPositionIndexes;
-import static com.google.common.collect.BoundType.CLOSED;
+import static com.google.common.collect.SortedLists.KeyAbsentBehavior.INVERTED_INSERTION_INDEX;
+import static com.google.common.collect.SortedLists.KeyAbsentBehavior.NEXT_HIGHER;
+import static com.google.common.collect.SortedLists.KeyAbsentBehavior.NEXT_LOWER;
+import static com.google.common.collect.SortedLists.KeyPresentBehavior.ANY_PRESENT;
 
 import com.google.common.primitives.Ints;
 
+import java.util.Comparator;
+import java.util.List;
+
 import javax.annotation.Nullable;
 
 /**
@@ -27,119 +32,167 @@ import javax.annotation.Nullable;
  *
  * @author Louis Wasserman
  */
-@SuppressWarnings("serial") // uses writeReplace, not default serialization
 final class RegularImmutableSortedMultiset<E> extends ImmutableSortedMultiset<E> {
-  private final transient RegularImmutableSortedSet<E> elementSet;
-  private final transient int[] counts;
-  private final transient long[] cumulativeCounts;
-  private final transient int offset;
-  private final transient int length;
+  private static final class CumulativeCountEntry<E> extends Multisets.AbstractEntry<E> {
+    final E element;
+    final int count;
+    final long cumulativeCount;
+
+    CumulativeCountEntry(E element, int count, @Nullable CumulativeCountEntry<E> previous) {
+      this.element = element;
+      this.count = count;
+      this.cumulativeCount = count + ((previous == null) ? 0 : previous.cumulativeCount);
+    }
+
+    @Override
+    public E getElement() {
+      return element;
+    }
+
+    @Override
+    public int getCount() {
+      return count;
+    }
+  }
+
+  static <E> RegularImmutableSortedMultiset<E> createFromSorted(Comparator<? super E> comparator,
+      List<? extends Multiset.Entry<E>> entries) {
+    List<CumulativeCountEntry<E>> newEntries = Lists.newArrayListWithCapacity(entries.size());
+    CumulativeCountEntry<E> previous = null;
+    for (Multiset.Entry<E> entry : entries) {
+      newEntries.add(
+          previous = new CumulativeCountEntry<E>(entry.getElement(), entry.getCount(), previous));
+    }
+    return new RegularImmutableSortedMultiset<E>(comparator, ImmutableList.copyOf(newEntries));
+  }
+
+  final transient ImmutableList<CumulativeCountEntry<E>> entries;
 
   RegularImmutableSortedMultiset(
-      RegularImmutableSortedSet<E> elementSet,
-      int[] counts,
-      long[] cumulativeCounts,
-      int offset,
-      int length) {
-    this.elementSet = elementSet;
-    this.counts = counts;
-    this.cumulativeCounts = cumulativeCounts;
-    this.offset = offset;
-    this.length = length;
+      Comparator<? super E> comparator, ImmutableList<CumulativeCountEntry<E>> entries) {
+    super(comparator);
+    this.entries = entries;
+    assert !entries.isEmpty();
   }
 
-  private Entry<E> getEntry(int index) {
-    return Multisets.immutableEntry(
-        elementSet.asList().get(index),
-        counts[offset + index]);
+  ImmutableList<E> elementList() {
+    return new TransformedImmutableList<CumulativeCountEntry<E>, E>(entries) {
+      @Override
+      E transform(CumulativeCountEntry<E> entry) {
+        return entry.getElement();
+      }
+    };
   }
 
   @Override
-  public Entry<E> firstEntry() {
-    return getEntry(0);
+  ImmutableSortedSet<E> createElementSet() {
+    return new RegularImmutableSortedSet<E>(elementList(), comparator());
   }
 
   @Override
-  public Entry<E> lastEntry() {
-    return getEntry(length - 1);
+  ImmutableSortedSet<E> createDescendingElementSet() {
+    return new RegularImmutableSortedSet<E>(elementList().reverse(), reverseComparator());
   }
 
+  @SuppressWarnings("unchecked")
   @Override
-  public int count(@Nullable Object element) {
-    int index = elementSet.indexOf(element);
-    return (index == -1) ? 0 : counts[index + offset];
+  UnmodifiableIterator<Multiset.Entry<E>> entryIterator() {
+    return (UnmodifiableIterator) entries.iterator();
   }
 
+  @SuppressWarnings("unchecked")
   @Override
-  public int size() {
-    long size = cumulativeCounts[offset + length] - cumulativeCounts[offset];
-    return Ints.saturatedCast(size);
+  UnmodifiableIterator<Multiset.Entry<E>> descendingEntryIterator() {
+    return (UnmodifiableIterator) entries.reverse().iterator();
   }
 
   @Override
-  public ImmutableSortedSet<E> elementSet() {
-    return elementSet;
+  public CumulativeCountEntry<E> firstEntry() {
+    return entries.get(0);
   }
 
   @Override
-  public ImmutableSortedMultiset<E> headMultiset(E upperBound, BoundType boundType) {
-    return getSubMultiset(0, elementSet.headIndex(upperBound, checkNotNull(boundType) == CLOSED));
+  public CumulativeCountEntry<E> lastEntry() {
+    return entries.get(entries.size() - 1);
   }
 
   @Override
-  public ImmutableSortedMultiset<E> tailMultiset(E lowerBound, BoundType boundType) {
-    return getSubMultiset(elementSet.tailIndex(lowerBound, checkNotNull(boundType) == CLOSED),
-        length);
+  public int size() {
+    CumulativeCountEntry<E> firstEntry = firstEntry();
+    CumulativeCountEntry<E> lastEntry = lastEntry();
+    return Ints.saturatedCast(
+        lastEntry.cumulativeCount - firstEntry.cumulativeCount + firstEntry.count);
   }
 
-  ImmutableSortedMultiset<E> getSubMultiset(int from, int to) {
-    checkPositionIndexes(from, to, length);
-    if (from == to) {
-      return emptyMultiset(comparator());
-    } else if (from == 0 && to == length) {
-      return this;
-    } else {
-      RegularImmutableSortedSet<E> subElementSet =
-          (RegularImmutableSortedSet<E>) elementSet.getSubSet(from, to);
-      return new RegularImmutableSortedMultiset<E>(
-          subElementSet, counts, cumulativeCounts, offset + from, to - from);
-    }
+  @Override
+  int distinctElements() {
+    return entries.size();
   }
 
   @Override
-  ImmutableSet<Entry<E>> createEntrySet() {
-    return new EntrySet();
+  boolean isPartialView() {
+    return entries.isPartialView();
   }
 
-  private final class EntrySet extends ImmutableMultiset<E>.EntrySet {
-    @Override
-    public int size() {
-      return length;
+  @SuppressWarnings("unchecked")
+  @Override
+  public int count(@Nullable Object element) {
+    if (element == null) {
+      return 0;
     }
-
-    @Override
-    public UnmodifiableIterator<Entry<E>> iterator() {
-      return asList().iterator();
+    try {
+      int index = SortedLists.binarySearch(
+          elementList(), (E) element, comparator(), ANY_PRESENT, INVERTED_INSERTION_INDEX);
+      return (index >= 0) ? entries.get(index).getCount() : 0;
+    } catch (ClassCastException e) {
+      return 0;
     }
+  }
 
-    @Override
-    ImmutableList<Entry<E>> createAsList() {
-      return new ImmutableAsList<Entry<E>>() {
-        @Override
-        public Entry<E> get(int index) {
-          return getEntry(index);
-        }
-
-        @Override
-        ImmutableCollection<Entry<E>> delegateCollection() {
-          return EntrySet.this;
-        }
-      };
+  @Override
+  public ImmutableSortedMultiset<E> headMultiset(E upperBound, BoundType boundType) {
+    int index;
+    switch (boundType) {
+      case OPEN:
+        index = SortedLists.binarySearch(
+            elementList(), checkNotNull(upperBound), comparator(), ANY_PRESENT, NEXT_HIGHER);
+        break;
+      case CLOSED:
+        index = SortedLists.binarySearch(
+            elementList(), checkNotNull(upperBound), comparator(), ANY_PRESENT, NEXT_LOWER) + 1;
+        break;
+      default:
+        throw new AssertionError();
     }
+    return createSubMultiset(0, index);
   }
 
   @Override
-  boolean isPartialView() {
-    return offset > 0 || length < counts.length;
+  public ImmutableSortedMultiset<E> tailMultiset(E lowerBound, BoundType boundType) {
+    int index;
+    switch (boundType) {
+      case OPEN:
+        index = SortedLists.binarySearch(
+            elementList(), checkNotNull(lowerBound), comparator(), ANY_PRESENT, NEXT_LOWER) + 1;
+        break;
+      case CLOSED:
+        index = SortedLists.binarySearch(
+            elementList(), checkNotNull(lowerBound), comparator(), ANY_PRESENT, NEXT_HIGHER);
+        break;
+      default:
+        throw new AssertionError();
+    }
+    return createSubMultiset(index, distinctElements());
+  }
+
+  private ImmutableSortedMultiset<E> createSubMultiset(int newFromIndex, int newToIndex) {
+    if (newFromIndex == 0 && newToIndex == entries.size()) {
+      return this;
+    } else if (newFromIndex >= newToIndex) {
+      return emptyMultiset(comparator());
+    } else {
+      return new RegularImmutableSortedMultiset<E>(
+          comparator(), entries.subList(newFromIndex, newToIndex));
+    }
   }
 }