common/collect/TreeMultiset.java

/*
 * Copyright (C) 2007 The Guava Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.common.collect;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.BstSide.LEFT;
import static com.google.common.collect.BstSide.RIGHT;

import java.io.Serializable;
import java.util.Comparator;
import java.util.ConcurrentModificationException;
import java.util.Iterator;

import javax.annotation.Nullable;

import com.google.common.annotations.GwtCompatible;
import com.google.common.primitives.Ints;

/**
 * A multiset which maintains the ordering of its elements, according to either
 * their natural order or an explicit {@link Comparator}. In all cases, this
 * implementation uses {@link Comparable#compareTo} or {@link
 * Comparator#compare} instead of {@link Object#equals} to determine
 * equivalence of instances.
 *
 * <p><b>Warning:</b> The comparison must be <i>consistent with equals</i> as
 * explained by the {@link Comparable} class specification. Otherwise, the
 * resulting multiset will violate the {@link java.util.Collection} contract,
 * which is specified in terms of {@link Object#equals}.
 *
 * @author Louis Wasserman
 * @author Jared Levy
 * @since 2.0 (imported from Google Collections Library)
 */
@GwtCompatible(emulated = true)
public final class TreeMultiset<E> extends AbstractSortedMultiset<E>
    implements Serializable {

  /**
   * Creates a new, empty multiset, sorted according to the elements' natural
   * order. All elements inserted into the multiset must implement the
   * {@code Comparable} interface. Furthermore, all such elements must be
   * <i>mutually comparable</i>: {@code e1.compareTo(e2)} must not throw a
   * {@code ClassCastException} for any elements {@code e1} and {@code e2} in
   * the multiset. If the user attempts to add an element to the multiset that
   * violates this constraint (for example, the user attempts to add a string
   * element to a set whose elements are integers), the {@code add(Object)}
   * call will throw a {@code ClassCastException}.
   *
   * <p>The type specification is {@code <E extends Comparable>}, instead of the
   * more specific {@code <E extends Comparable<? super E>>}, to support
   * classes defined without generics.
   */
  public static <E extends Comparable> TreeMultiset<E> create() {
    return new TreeMultiset<E>(Ordering.natural());
  }

  /**
   * Creates a new, empty multiset, sorted according to the specified
   * comparator. All elements inserted into the multiset must be <i>mutually
   * comparable</i> by the specified comparator: {@code comparator.compare(e1,
   * e2)} must not throw a {@code ClassCastException} for any elements {@code
   * e1} and {@code e2} in the multiset. If the user attempts to add an element
   * to the multiset that violates this constraint, the {@code add(Object)} call
   * will throw a {@code ClassCastException}.
   *
   * @param comparator the comparator that will be used to sort this multiset. A
   *     null value indicates that the elements' <i>natural ordering</i> should
   *     be used.
   */
  @SuppressWarnings("unchecked")
  public static <E> TreeMultiset<E> create(
      @Nullable Comparator<? super E> comparator) {
    return (comparator == null)
           ? new TreeMultiset<E>((Comparator) Ordering.natural())
           : new TreeMultiset<E>(comparator);
  }

  /**
   * Creates an empty multiset containing the given initial elements, sorted
   * according to the elements' natural order.
   *
   * <p>This implementation is highly efficient when {@code elements} is itself
   * a {@link Multiset}.
   *
   * <p>The type specification is {@code <E extends Comparable>}, instead of the
   * more specific {@code <E extends Comparable<? super E>>}, to support
   * classes defined without generics.
   */
  public static <E extends Comparable> TreeMultiset<E> create(
      Iterable<? extends E> elements) {
    TreeMultiset<E> multiset = create();
    Iterables.addAll(multiset, elements);
    return multiset;
  }

  /**
   * Returns an iterator over the elements contained in this collection.
   */
  @Override
  public Iterator<E> iterator() {
    // Needed to avoid Javadoc bug.
    return super.iterator();
  }

  private TreeMultiset(Comparator<? super E> comparator) {
    super(comparator);
    this.range = GeneralRange.all(comparator);
    this.rootReference = new Reference<Node<E>>();
  }

  private TreeMultiset(GeneralRange<E> range, Reference<Node<E>> root) {
    super(range.comparator());
    this.range = range;
    this.rootReference = root;
  }

  @SuppressWarnings("unchecked")
  E checkElement(Object o) {
    return (E) o;
  }

  private transient final GeneralRange<E> range;

  private transient final Reference<Node<E>> rootReference;

  static final class Reference<T> {
    T value;

    public Reference() {}

    public T get() {
      return value;
    }

    public boolean compareAndSet(T expected, T newValue) {
      if (value == expected) {
        value = newValue;
        return true;
      }
      return false;
    }
  }

  @Override
  int distinctElements() {
    Node<E> root = rootReference.get();
    return Ints.checkedCast(BstRangeOps.totalInRange(distinctAggregate(), range, root));
  }

  @Override
  public int size() {
    Node<E> root = rootReference.get();
    return Ints.saturatedCast(BstRangeOps.totalInRange(sizeAggregate(), range, root));
  }

  @Override
  public int count(@Nullable Object element) {
    try {
      E e = checkElement(element);
      if (range.contains(e)) {
        Node<E> node = BstOperations.seek(comparator(), rootReference.get(), e);
        return countOrZero(node);
      }
      return 0;
    } catch (ClassCastException e) {
      return 0;
    } catch (NullPointerException e) {
      return 0;
    }
  }

  private int mutate(@Nullable E e, MultisetModifier modifier) {
    BstMutationRule<E, Node<E>> mutationRule = BstMutationRule.createRule(
        modifier,
        BstCountBasedBalancePolicies.
          <E, Node<E>>singleRebalancePolicy(distinctAggregate()),
        nodeFactory());
    BstMutationResult<E, Node<E>> mutationResult =
        BstOperations.mutate(comparator(), mutationRule, rootReference.get(), e);
    if (!rootReference.compareAndSet(
        mutationResult.getOriginalRoot(), mutationResult.getChangedRoot())) {
      throw new ConcurrentModificationException();
    }
    Node<E> original = mutationResult.getOriginalTarget();
    return countOrZero(original);
  }

  @Override
  public int add(E element, int occurrences) {
    checkElement(element);
    if (occurrences == 0) {
      return count(element);
    }
    checkArgument(range.contains(element));
    return mutate(element, new AddModifier(occurrences));
  }

  @Override
  public int remove(@Nullable Object element, int occurrences) {
    if (element == null) {
      return 0;
    } else if (occurrences == 0) {
      return count(element);
    }
    try {
      E e = checkElement(element);
      return range.contains(e) ? mutate(e, new RemoveModifier(occurrences)) : 0;
    } catch (ClassCastException e) {
      return 0;
    }
  }

  @Override
  public boolean setCount(E element, int oldCount, int newCount) {
    checkElement(element);
    checkArgument(range.contains(element));
    return mutate(element, new ConditionalSetCountModifier(oldCount, newCount))
        == oldCount;
  }

  @Override
  public int setCount(E element, int count) {
    checkElement(element);
    checkArgument(range.contains(element));
    return mutate(element, new SetCountModifier(count));
  }

  private BstPathFactory<Node<E>, BstInOrderPath<Node<E>>> pathFactory() {
    return BstInOrderPath.inOrderFactory();
  }

  @Override
  Iterator<Entry<E>> entryIterator() {
    Node<E> root = rootReference.get();
    final BstInOrderPath<Node<E>> startingPath =
        BstRangeOps.furthestPath(range, LEFT, pathFactory(), root);
    return iteratorInDirection(startingPath, RIGHT);
  }

  @Override
  Iterator<Entry<E>> descendingEntryIterator() {
    Node<E> root = rootReference.get();
    final BstInOrderPath<Node<E>> startingPath =
        BstRangeOps.furthestPath(range, RIGHT, pathFactory(), root);
    return iteratorInDirection(startingPath, LEFT);
  }

  private Iterator<Entry<E>> iteratorInDirection(
      @Nullable BstInOrderPath<Node<E>> start, final BstSide direction) {
    final Iterator<BstInOrderPath<Node<E>>> pathIterator =
        new AbstractLinkedIterator<BstInOrderPath<Node<E>>>(start) {
          @Override
          protected BstInOrderPath<Node<E>> computeNext(BstInOrderPath<Node<E>> previous) {
            if (!previous.hasNext(direction)) {
              return null;
            }
            BstInOrderPath<Node<E>> next = previous.next(direction);
            // TODO(user): only check against one side
            return range.contains(next.getTip().getKey()) ? next : null;
          }
        };
    return new Iterator<Entry<E>>() {
      E toRemove = null;

      @Override
      public boolean hasNext() {
        return pathIterator.hasNext();
      }

      @Override
      public Entry<E> next() {
        BstInOrderPath<Node<E>> path = pathIterator.next();
        return new LiveEntry(
            toRemove = path.getTip().getKey(), path.getTip().elemCount());
      }

      @Override
      public void remove() {
        checkState(toRemove != null);
        setCount(toRemove, 0);
        toRemove = null;
      }
    };
  }

  class LiveEntry extends Multisets.AbstractEntry<E> {
    private Node<E> expectedRoot;
    private final E element;
    private int count;

    private LiveEntry(E element, int count) {
      this.expectedRoot = rootReference.get();
      this.element = element;
      this.count = count;
    }

    @Override
    public E getElement() {
      return element;
    }

    @Override
    public int getCount() {
      if (rootReference.get() == expectedRoot) {
        return count;
      } else {
        // check for updates
        expectedRoot = rootReference.get();
        return count = TreeMultiset.this.count(element);
      }
    }
  }

  @Override
  public void clear() {
    Node<E> root = rootReference.get();
    Node<E> cleared = BstRangeOps.minusRange(range,
        BstCountBasedBalancePolicies.<E, Node<E>>fullRebalancePolicy(distinctAggregate()),
        nodeFactory(), root);
    if (!rootReference.compareAndSet(root, cleared)) {
      throw new ConcurrentModificationException();
    }
  }

  @Override
  public SortedMultiset<E> headMultiset(E upperBound, BoundType boundType) {
    checkNotNull(upperBound);
    return new TreeMultiset<E>(
        range.intersect(GeneralRange.upTo(comparator, upperBound, boundType)), rootReference);
  }

  @Override
  public SortedMultiset<E> tailMultiset(E lowerBound, BoundType boundType) {
    checkNotNull(lowerBound);
    return new TreeMultiset<E>(
        range.intersect(GeneralRange.downTo(comparator, lowerBound, boundType)), rootReference);
  }

  /**
   * {@inheritDoc}
   *
   * @since 11.0
   */
  @Override
  public Comparator<? super E> comparator() {
    return super.comparator();
  }

  private static final class Node<E> extends BstNode<E, Node<E>> implements Serializable {
    private final long size;
    private final int distinct;

    private Node(E key, int elemCount, @Nullable Node<E> left,
        @Nullable Node<E> right) {
      super(key, left, right);
      checkArgument(elemCount > 0);
      this.size = (long) elemCount + sizeOrZero(left) + sizeOrZero(right);
      this.distinct = 1 + distinctOrZero(left) + distinctOrZero(right);
    }

    int elemCount() {
      long result = size - sizeOrZero(childOrNull(LEFT))
          - sizeOrZero(childOrNull(RIGHT));
      return Ints.checkedCast(result);
    }

    private Node(E key, int elemCount) {
      this(key, elemCount, null, null);
    }

    private static final long serialVersionUID = 0;
  }

  private static long sizeOrZero(@Nullable Node<?> node) {
    return (node == null) ? 0 : node.size;
  }

  private static int distinctOrZero(@Nullable Node<?> node) {
    return (node == null) ? 0 : node.distinct;
  }

  private static int countOrZero(@Nullable Node<?> entry) {
    return (entry == null) ? 0 : entry.elemCount();
  }

  @SuppressWarnings("unchecked")
  private BstAggregate<Node<E>> distinctAggregate() {
    return (BstAggregate) DISTINCT_AGGREGATE;
  }

  private static final BstAggregate<Node<Object>> DISTINCT_AGGREGATE =
      new BstAggregate<Node<Object>>() {
    @Override
    public int entryValue(Node<Object> entry) {
      return 1;
    }

    @Override
    public long treeValue(@Nullable Node<Object> tree) {
      return distinctOrZero(tree);
    }
  };

  @SuppressWarnings("unchecked")
  private BstAggregate<Node<E>> sizeAggregate() {
    return (BstAggregate) SIZE_AGGREGATE;
  }

  private static final BstAggregate<Node<Object>> SIZE_AGGREGATE =
      new BstAggregate<Node<Object>>() {
        @Override
        public int entryValue(Node<Object> entry) {
          return entry.elemCount();
        }

        @Override
        public long treeValue(@Nullable Node<Object> tree) {
          return sizeOrZero(tree);
        }
      };

  @SuppressWarnings("unchecked")
  private BstNodeFactory<Node<E>> nodeFactory() {
    return (BstNodeFactory) NODE_FACTORY;
  }

  private static final BstNodeFactory<Node<Object>> NODE_FACTORY =
      new BstNodeFactory<Node<Object>>() {
        @Override
        public Node<Object> createNode(Node<Object> source, @Nullable Node<Object> left,
            @Nullable Node<Object> right) {
          return new Node<Object>(source.getKey(), source.elemCount(), left, right);
        }
      };

  private abstract class MultisetModifier implements BstModifier<E, Node<E>> {
    abstract int newCount(int oldCount);

    @Nullable
    @Override
    public BstModificationResult<Node<E>> modify(E key, @Nullable Node<E> originalEntry) {
      int oldCount = countOrZero(originalEntry);
      int newCount = newCount(oldCount);
      if (oldCount == newCount) {
        return BstModificationResult.identity(originalEntry);
      } else if (newCount == 0) {
        return BstModificationResult.rebalancingChange(originalEntry, null);
      } else if (oldCount == 0) {
        return BstModificationResult.rebalancingChange(null, new Node<E>(key, newCount));
      } else {
        return BstModificationResult.rebuildingChange(originalEntry,
            new Node<E>(originalEntry.getKey(), newCount));
      }
    }
  }

  private final class AddModifier extends MultisetModifier {
    private final int countToAdd;

    private AddModifier(int countToAdd) {
      checkArgument(countToAdd > 0);
      this.countToAdd = countToAdd;
    }

    @Override
    int newCount(int oldCount) {
      checkArgument(countToAdd <= Integer.MAX_VALUE - oldCount, "Cannot add this many elements");
      return oldCount + countToAdd;
    }
  }

  private final class RemoveModifier extends MultisetModifier {
    private final int countToRemove;

    private RemoveModifier(int countToRemove) {
      checkArgument(countToRemove > 0);
      this.countToRemove = countToRemove;
    }

    @Override
    int newCount(int oldCount) {
      return Math.max(0, oldCount - countToRemove);
    }
  }

  private final class SetCountModifier extends MultisetModifier {
    private final int countToSet;

    private SetCountModifier(int countToSet) {
      checkArgument(countToSet >= 0);
      this.countToSet = countToSet;
    }

    @Override
    int newCount(int oldCount) {
      return countToSet;
    }
  }

  private final class ConditionalSetCountModifier extends MultisetModifier {
    private final int expectedCount;
    private final int setCount;

    private ConditionalSetCountModifier(int expectedCount, int setCount) {
      checkArgument(setCount >= 0 & expectedCount >= 0);
      this.expectedCount = expectedCount;
      this.setCount = setCount;
    }

    @Override
    int newCount(int oldCount) {
      return (oldCount == expectedCount) ? setCount : oldCount;
    }
  }

  /*
   * TODO(jlevy): Decide whether entrySet() should return entries with an
   * equals() method that calls the comparator to compare the two keys. If that
   * change is made, AbstractMultiset.equals() can simply check whether two
   * multisets have equal entry sets.
   */
}