1 files changed, 394 insertions, 0 deletions
diff --git a/src/main/java/com/google/devtools/build/lib/collect/ImmutableSortedKeyListMultimap.java b/src/main/java/com/google/devtools/build/lib/collect/ImmutableSortedKeyListMultimap.java
new file mode 100644
index 0000000000..216337813c
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/collect/ImmutableSortedKeyListMultimap.java
@@ -0,0 +1,394 @@
+// Copyright 2014 Google Inc. All rights reserved.
+//
+// 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.devtools.build.lib.collect;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.AbstractIterator;
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMultiset;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ListMultimap;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Multiset;
+
+import java.util.AbstractCollection;
+import java.util.AbstractMap;
+import java.util.AbstractMap.SimpleImmutableEntry;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
+
+import javax.annotation.Nullable;
+
+/**
+ * A immutable multimap implementation for multimaps with comparable keys. It uses a sorted array
+ * and binary search to return the correct values. It's only purpose is to save memory - it consumes
+ * only about half the memory of the equivalent ImmutableListMultimap. Only a few methods are
+ * efficiently implemented: {@link #isEmpty} is O(1), {@link #get} and {@link #containsKey} are
+ * O(log(n)), and {@link #asMap} and {@link #values} refer to the parent instance. All other methods
+ * can take O(n) or even make a copy of the contents.
+ *
+ * <p>This implementation supports neither {@code null} keys nor {@code null} values.
+ */
+public final class ImmutableSortedKeyListMultimap<K extends Comparable<K>, V>
+    implements ListMultimap<K, V> {
+
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  private static final ImmutableSortedKeyListMultimap EMPTY_MULTIMAP =
+      new ImmutableSortedKeyListMultimap(new Comparable<?>[0], new List<?>[0]);
+
+  /** Returns the empty multimap. */
+  @SuppressWarnings("unchecked")
+  public static <K extends Comparable<K>, V> ImmutableSortedKeyListMultimap<K, V> of() {
+    // Safe because the multimap will never hold any elements.
+    return EMPTY_MULTIMAP;
+  }
+
+  @SuppressWarnings("unchecked")
+  public static <K extends Comparable<K>, V> ImmutableSortedKeyListMultimap<K, V> copyOf(
+      Multimap<K, V> data) {
+    if (data.isEmpty()) {
+      return EMPTY_MULTIMAP;
+    }
+    if (data instanceof ImmutableSortedKeyListMultimap) {
+      return (ImmutableSortedKeyListMultimap<K, V>) data;
+    }
+    Set<K> keySet = data.keySet();
+    int size = keySet.size();
+    K[] sortedKeys = (K[]) new Comparable<?>[size];
+    int index = 0;
+    for (K key : keySet) {
+      sortedKeys[index++] = Preconditions.checkNotNull(key);
+    }
+    Arrays.sort(sortedKeys);
+    List<V>[] values = (List<V>[]) new List<?>[size];
+    for (int i = 0; i < size; i++) {
+      values[i] = ImmutableList.copyOf(data.get(sortedKeys[i]));
+    }
+    return new ImmutableSortedKeyListMultimap<>(sortedKeys, values);
+  }
+
+  public static <K extends Comparable<K>, V> Builder<K, V> builder() {
+    return new Builder<>();
+  }
+
+  /**
+   * A builder class for ImmutableSortedKeyListMultimap<K, V> instances.
+   */
+  public static final class Builder<K extends Comparable<K>, V> {
+    private final Multimap<K, V> builderMultimap = ArrayListMultimap.create();
+
+    Builder() {
+      // Not public so you must call builder() instead.
+    }
+
+    public ImmutableSortedKeyListMultimap<K, V> build() {
+      return ImmutableSortedKeyListMultimap.copyOf(builderMultimap);
+    }
+
+    public Builder<K, V> put(K key, V value) {
+      builderMultimap.put(Preconditions.checkNotNull(key), Preconditions.checkNotNull(value));
+      return this;
+    }
+
+    public Builder<K, V> putAll(K key, Collection<? extends V> values) {
+      Collection<V> valueList = builderMultimap.get(Preconditions.checkNotNull(key));
+      for (V value : values) {
+        valueList.add(Preconditions.checkNotNull(value));
+      }
+      return this;
+    }
+
+    @SuppressWarnings("unchecked")
+    public Builder<K, V> putAll(K key, V... values) {
+      return putAll(Preconditions.checkNotNull(key), Arrays.asList(values));
+    }
+
+    public Builder<K, V> putAll(Multimap<? extends K, ? extends V> multimap) {
+      for (Map.Entry<? extends K, ? extends Collection<? extends V>> entry
+          : multimap.asMap().entrySet()) {
+        putAll(entry.getKey(), entry.getValue());
+      }
+      return this;
+    }
+  }
+
+  /**
+   * An implementation for the Multimap.asMap method. Note that AbstractMap already provides
+   * implementations for all methods except {@link #entrySet}, but we override a few here because we
+   * can do it much faster than the existing entrySet-based implementations. Also note that it
+   * inherits the type parameters K and V from the parent class.
+   */
+  private class AsMap extends AbstractMap<K, Collection<V>> {
+
+    AsMap() {
+    }
+
+    @Override
+    public int size() {
+      return sortedKeys.length;
+    }
+
+    @Override
+    public boolean containsKey(Object key) {
+      return ImmutableSortedKeyListMultimap.this.containsKey(key);
+    }
+
+    @Override
+    public Collection<V> get(Object key) {
+      int index = Arrays.binarySearch(sortedKeys, key);
+      // Note the different semantic between Map and Multimap.
+      return index >= 0 ? values[index] : null;
+    }
+
+    @Override
+    public Collection<V> remove(Object key) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void clear() {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public Set<Entry<K, Collection<V>>> entrySet() {
+      ImmutableSet.Builder<Entry<K, Collection<V>>> builder = ImmutableSet.builder();
+      for (int i = 0; i < sortedKeys.length; i++) {
+        builder.add(new SimpleImmutableEntry<K, Collection<V>>(sortedKeys[i], values[i]));
+      }
+      return builder.build();
+    }
+  }
+
+  private class ValuesCollection extends AbstractCollection<V> {
+
+    ValuesCollection() {
+    }
+
+    @Override
+    public int size() {
+      return ImmutableSortedKeyListMultimap.this.size();
+    }
+
+    @Override
+    public boolean isEmpty() {
+      return sortedKeys.length == 0;
+    }
+
+    @Override
+    public boolean contains(Object o) {
+      return ImmutableSortedKeyListMultimap.this.containsValue(o);
+    }
+
+    @Override
+    public Iterator<V> iterator() {
+      if (isEmpty()) {
+        return Collections.emptyIterator();
+      }
+      return new AbstractIterator<V>() {
+        private int currentList = 0;
+        private int currentIndex = 0;
+
+        @Override
+        protected V computeNext() {
+          if (currentList >= values.length) {
+            return endOfData();
+          }
+          V result = values[currentList].get(currentIndex);
+          // Find the next list/index pair.
+          currentIndex++;
+          if (currentIndex >= values[currentList].size()) {
+            currentIndex = 0;
+            currentList++;
+          }
+          return result;
+        }
+      };
+    }
+
+    @Override
+    public boolean remove(Object o) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean removeAll(Collection<?> c) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean retainAll(Collection<?> c) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void clear() {
+      throw new UnsupportedOperationException();
+    }
+  }
+
+  private final K[] sortedKeys;
+  private final List<V>[] values;
+
+  private ImmutableSortedKeyListMultimap(K[] sortedKeys, List<V>[] values) {
+    this.sortedKeys = sortedKeys;
+    this.values = values;
+  }
+
+  @Override
+  public int size() {
+    int result = 0;
+    for (List<V> list : values) {
+      result += list.size();
+    }
+    return result;
+  }
+
+  @Override
+  public boolean isEmpty() {
+    return sortedKeys.length == 0;
+  }
+
+  @Override
+  public boolean containsKey(Object key) {
+    int index = Arrays.binarySearch(sortedKeys, key);
+    return index >= 0;
+  }
+
+  @Override
+  public boolean containsValue(Object value) {
+    for (List<V> list : values) {
+      if (list.contains(value)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  @Override
+  public boolean containsEntry(Object key, Object value) {
+    int index = Arrays.binarySearch(sortedKeys, key);
+    if (index >= 0) {
+      return values[index].contains(value);
+    }
+    return false;
+  }
+
+  @Override
+  public boolean put(K key, V value) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public boolean remove(Object key, Object value) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public boolean putAll(K key, Iterable<? extends V> values) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public boolean putAll(Multimap<? extends K, ? extends V> multimap) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public List<V> replaceValues(K key, Iterable<? extends V> values) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public List<V> removeAll(Object key) {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public void clear() {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public List<V> get(K key) {
+    int index = Arrays.binarySearch(sortedKeys, key);
+    return index >= 0 ? values[index] : ImmutableList.<V>of();
+  }
+
+  @Override
+  public Set<K> keySet() {
+    return ImmutableSet.copyOf(sortedKeys);
+  }
+
+  @Override
+  public Multiset<K> keys() {
+    return ImmutableMultiset.copyOf(sortedKeys);
+  }
+
+  @Override
+  public Collection<V> values() {
+    return new ValuesCollection();
+  }
+
+  @Override
+  public Collection<Entry<K, V>> entries() {
+    ImmutableList.Builder<Entry<K, V>> builder = ImmutableList.builder();
+    for (int i = 0; i < sortedKeys.length; i++) {
+      for (V value : values[i]) {
+        builder.add(new SimpleImmutableEntry<K, V>(sortedKeys[i], value));
+      }
+    }
+    return builder.build();
+  }
+
+  /**
+   * {@inheritDoc}
+   *
+   * <p>Note that only {@code get} and {@code containsKey} are implemented efficiently on the
+   * returned map.
+   */
+  @Override
+  public Map<K, Collection<V>> asMap() {
+    return new AsMap();
+  }
+
+  @Override
+  public String toString() {
+    return asMap().toString();
+  }
+
+  @Override
+  public int hashCode() {
+    return asMap().hashCode();
+  }
+
+  @Override
+  public boolean equals(@Nullable Object object) {
+    if (this == object) {
+      return true;
+    }
+    if (object instanceof Multimap) {
+      Multimap<?, ?> that = (Multimap<?, ?>) object;
+      return asMap().equals(that.asMap());
+    }
+    return false;
+  }
+}