diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableTableCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableTableCodec.java index 5113b500135de0..434fe0c7ac6d0e 100644 --- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableTableCodec.java +++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableTableCodec.java @@ -11,33 +11,35 @@ // 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.skyframe.serialization; +import static sun.misc.Unsafe.ARRAY_OBJECT_BASE_OFFSET; +import static sun.misc.Unsafe.ARRAY_OBJECT_INDEX_SCALE; + +import com.google.common.collect.ImmutableSet; import com.google.common.collect.ImmutableTable; import com.google.common.collect.Table.Cell; import com.google.protobuf.CodedInputStream; import com.google.protobuf.CodedOutputStream; import java.io.IOException; -import java.util.Set; +import java.util.function.Supplier; /** Codec for {@link ImmutableTable}. */ -public class ImmutableTableCodec implements ObjectCodec> { +@SuppressWarnings({"rawtypes", "unchecked"}) +public class ImmutableTableCodec extends DeferredObjectCodec { - @SuppressWarnings("unchecked") @Override - public Class> getEncodedClass() { - // Compiler doesn't like to do a direct cast. - return (Class>) ((Class) ImmutableTable.class); + public Class getEncodedClass() { + return ImmutableTable.class; } @Override public void serialize( - SerializationContext context, ImmutableTable object, CodedOutputStream codedOut) + SerializationContext context, ImmutableTable object, CodedOutputStream codedOut) throws SerializationException, IOException { - Set> cellSet = object.cellSet(); + ImmutableSet cellSet = object.cellSet(); codedOut.writeInt32NoTag(cellSet.size()); - for (Cell cell : cellSet) { + for (Cell cell : cellSet) { context.serialize(cell.getRowKey(), codedOut); context.serialize(cell.getColumnKey(), codedOut); context.serialize(cell.getValue(), codedOut); @@ -45,20 +47,51 @@ public void serialize( } @Override - public ImmutableTable deserialize( - DeserializationContext context, CodedInputStream codedIn) + public Supplier deserializeDeferred( + AsyncDeserializationContext context, CodedInputStream codedIn) throws SerializationException, IOException { - int length = codedIn.readInt32(); - if (length < 0) { - throw new SerializationException("Expected non-negative length: " + length); + int size = codedIn.readInt32(); + if (size < 0) { + throw new SerializationException("Expected non-negative size: " + size); + } + if (size == 0) { + return ImmutableTable::of; + } + + EntryBuffer buffer = new EntryBuffer(size); + long offset = ARRAY_OBJECT_BASE_OFFSET; + for (int i = 0; i < size; i++) { + context.deserializeFully(codedIn, buffer.rowKeys, offset); + context.deserializeFully(codedIn, buffer.columnKeys, offset); + context.deserialize(codedIn, buffer.values, offset); + offset += ARRAY_OBJECT_INDEX_SCALE; + } + return buffer; + } + + private static final class EntryBuffer implements Supplier { + private final Object[] rowKeys; + private final Object[] columnKeys; + private final Object[] values; + + private EntryBuffer(int size) { + this.rowKeys = new Object[size]; + this.columnKeys = new Object[size]; + this.values = new Object[size]; } - ImmutableTable.Builder builder = ImmutableTable.builder(); - for (int i = 0; i < length; i++) { - builder.put( - /*rowKey=*/ context.deserialize(codedIn), - /*columnKey=*/ context.deserialize(codedIn), - /*value=*/ context.deserialize(codedIn)); + + @Override + public ImmutableTable get() { + ImmutableTable.Builder builder = ImmutableTable.builder(); + for (int i = 0; i < size(); i++) { + builder.put( + /* rowKey= */ rowKeys[i], /* columnKey= */ columnKeys[i], /* value= */ values[i]); + } + return builder.buildOrThrow(); + } + + private int size() { + return rowKeys.length; } - return builder.buildOrThrow(); } }