From 4fe0f87711106a681e62741550309103bcd0b5c7 Mon Sep 17 00:00:00 2001
From: nharmata <nharmata@google.com>
Date: Tue, 23 Feb 2021 13:56:56 -0800
Subject: [PATCH] Optimize `ImmutableSortedSetCodec#deserialize`.

Avoid re-sorting the list of elements that was already sorted on serialization.

PiperOrigin-RevId: 359127225
---
 .../ImmutableSortedSetCodec.java              | 135 +++++++++++++++++-
 1 file changed, 131 insertions(+), 4 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableSortedSetCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableSortedSetCodec.java
index 1517a6158b5eb3..be54fdf95439ce 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableSortedSetCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableSortedSetCodec.java
@@ -18,6 +18,10 @@
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.Iterator;
+import java.util.SortedSet;
 
 /**
  * {@link ObjectCodec} for {@link ImmutableSortedSet}. Comparator must be serializable, ideally a
@@ -44,12 +48,135 @@ public void serialize(
   @Override
   public ImmutableSortedSet<E> deserialize(DeserializationContext context, CodedInputStream codedIn)
       throws SerializationException, IOException {
-    ImmutableSortedSet.Builder<E> builder =
-        ImmutableSortedSet.orderedBy(context.deserialize(codedIn));
+    Comparator<E> comparator = context.deserialize(codedIn);
     int size = codedIn.readInt32();
+    @SuppressWarnings("unchecked")
+    E[] sortedElementsArray = (E[]) new Object[size];
     for (int i = 0; i < size; i++) {
-      builder.add(context.<E>deserialize(codedIn));
+      sortedElementsArray[i] = context.<E>deserialize(codedIn);
+    }
+    SortedSetShimForEfficientDeserialization<E> sortedSetShimForEfficientDeserialization =
+        new SortedSetShimForEfficientDeserialization<>(comparator, sortedElementsArray);
+    return ImmutableSortedSet.copyOfSorted(sortedSetShimForEfficientDeserialization);
+  }
+
+  /**
+   * Implementation of parts of the {@link SortedSet} interface minimally needed for efficient
+   * {@link ImmutableSortedSet} construction that avoids re-sorting the list of elements.
+   */
+  private static class SortedSetShimForEfficientDeserialization<E> implements SortedSet<E> {
+    private final Comparator<E> comparator;
+    private final Object[] sortedElementsArray;
+
+    private SortedSetShimForEfficientDeserialization(
+        Comparator<E> comparator, Object[] sortedElementsArray) {
+      this.comparator = comparator;
+      this.sortedElementsArray = sortedElementsArray;
+    }
+
+    @Override
+    public Comparator<? super E> comparator() {
+      return comparator;
+    }
+
+    @Override
+    public Object[] toArray() {
+      return sortedElementsArray;
+    }
+
+    @Override
+    public <T> T[] toArray(T[] a) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public SortedSet<E> subSet(E fromElement, E toElement) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public SortedSet<E> headSet(E toElement) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public SortedSet<E> tailSet(E fromElement) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public E first() {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public E last() {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public int size() {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean isEmpty() {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean contains(Object o) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public Iterator<E> iterator() {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean add(E e) {
+      return false;
+    }
+
+    @Override
+    public boolean remove(Object o) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean containsAll(Collection<?> c) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean addAll(Collection<? extends E> c) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean retainAll(Collection<?> c) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean removeAll(Collection<?> c) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public void clear() {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public int hashCode() {
+      throw new UnsupportedOperationException();
     }
-    return builder.build();
   }
 }