Allow ObjectCodecRegistry to handle MemoizingCodecs. Initially this is just Skylark memoization, but we may extend in future to handle more than just Skylark this way. In fact, we probably want most of our ObjectCodecs to be MemoizingCodecs that can efficiently fall back to ObjectCodec if not using memoization (maybe coming in a follow-up).

At a high level, this CL merges the functionality of MemoizingCodecMap into ObjectCodecRegistry, adds auto-registration of MemoizingCodec, and uses that to get rid of a lot of codecs that were just delegating. The big one to get rid of there is SkylarkValueCodec: all of its delegation duties are implicitly now in ObjectCodecRegistry and friends.

One danger with this CL is that many of the features of Skylark serialization are only being tested in unit tests, which had to be reworked as part of this change. I don't think we've lost any coverage, but I could be wrong. SkylarkValueCodec had a bunch of methods that were effectively test-only, which made it easier to remove.

The plan is to provide a Memoizer.Serializer inside the SerializationContext. At the top level, it will be a DUMMY_SERIALIZER that does no memoization, but a MemoizingCodec can do
context = context.ensureMemoizing()
which will recreate the context with a true memoizing serializer. Then all references to the Serializer in codec code can be cleaned up, and the MemoizingCodec and ObjectCodec signatures will be the same. At that point, we can make all ObjectCodecs compatible with memoization by default (with strategy MEMOIZE_AFTER), and add a "memoize" boolean to @AutoCodec. That should allow us to have full interoperability between all codecs.

This CL also makes CodecScanner deterministic in the order of classes that it processes (there was a lurking bug here where constants must be deterministically ordered but that wasn't enforced at all).

PiperOrigin-RevId: 188559983

3 files changed, 365 insertions, 0 deletions

diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/DeserializationContextTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/DeserializationContextTest.java
new file mode 100644
index 0000000000..764284a2ea
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/DeserializationContextTest.java
@@ -0,0 +1,167 @@
+// Copyright 2018 The Bazel Authors. 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.skyframe.serialization;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.when;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.protobuf.CodedInputStream;
+import java.io.IOException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+/** Tests for {@link DeserializationContext}. */
+@RunWith(JUnit4.class)
+public class DeserializationContextTest {
+  @Test
+  public void nullDeserialize() throws Exception {
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    CodedInputStream codedInputStream = Mockito.mock(CodedInputStream.class);
+    when(codedInputStream.readSInt32()).thenReturn(0);
+    DeserializationContext deserializationContext =
+        new DeserializationContext(registry, ImmutableMap.of());
+    assertThat((Object) deserializationContext.deserialize(codedInputStream)).isNull();
+    Mockito.verify(codedInputStream).readSInt32();
+    Mockito.verifyZeroInteractions(registry);
+  }
+
+  @Test
+  public void constantDeserialize() throws Exception {
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    Object constant = new Object();
+    when(registry.maybeGetConstantByTag(1)).thenReturn(constant);
+    CodedInputStream codedInputStream = Mockito.mock(CodedInputStream.class);
+    when(codedInputStream.readSInt32()).thenReturn(1);
+    DeserializationContext deserializationContext =
+        new DeserializationContext(registry, ImmutableMap.of());
+    assertThat((Object) deserializationContext.deserialize(codedInputStream)).isSameAs(constant);
+    Mockito.verify(codedInputStream).readSInt32();
+    Mockito.verify(registry).maybeGetConstantByTag(1);
+  }
+
+  @Test
+  public void descriptorDeserialize() throws Exception {
+    ObjectCodecRegistry.CodecDescriptor codecDescriptor =
+        Mockito.mock(ObjectCodecRegistry.CodecDescriptor.class);
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    when(registry.getCodecDescriptorByTag(1)).thenReturn(codecDescriptor);
+    CodedInputStream codedInputStream = Mockito.mock(CodedInputStream.class);
+    when(codedInputStream.readSInt32()).thenReturn(1);
+    DeserializationContext deserializationContext =
+        new DeserializationContext(registry, ImmutableMap.of());
+    Object returnValue = new Object();
+    when(codecDescriptor.deserialize(deserializationContext, codedInputStream))
+        .thenReturn(returnValue);
+    assertThat((Object) deserializationContext.deserialize(codedInputStream)).isSameAs(returnValue);
+    Mockito.verify(codedInputStream).readSInt32();
+    Mockito.verify(registry).getCodecDescriptorByTag(1);
+    Mockito.verify(codecDescriptor).deserialize(deserializationContext, codedInputStream);
+  }
+
+  @Test
+  public void memoizingDeserialize_null() throws SerializationException, IOException {
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    CodedInputStream codedInputStream = Mockito.mock(CodedInputStream.class);
+    DeserializationContext deserializationContext =
+        new DeserializationContext(registry, ImmutableMap.of());
+    when(codedInputStream.readSInt32()).thenReturn(0);
+    assertThat(
+            deserializationContext
+                .getMemoizingCodecRecordedInInput(codedInputStream)
+                .deserializePayload(
+                    deserializationContext,
+                    /*initial=*/ null,
+                    codedInputStream,
+                    new Memoizer.Deserializer()))
+        .isEqualTo(null);
+    Mockito.verify(codedInputStream).readSInt32();
+    Mockito.verifyZeroInteractions(registry);
+  }
+
+  @Test
+  public void memoizingDeserialize_constant() throws SerializationException, IOException {
+    Object constant = new Object();
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    when(registry.maybeGetConstantByTag(1)).thenReturn(constant);
+    CodedInputStream codedInputStream = Mockito.mock(CodedInputStream.class);
+    DeserializationContext deserializationContext =
+        new DeserializationContext(registry, ImmutableMap.of());
+    when(codedInputStream.readSInt32()).thenReturn(1);
+    assertThat(
+            deserializationContext
+                .getMemoizingCodecRecordedInInput(codedInputStream)
+                .deserializePayload(
+                    deserializationContext,
+                    /*initial=*/ null,
+                    codedInputStream,
+                    new Memoizer.Deserializer()))
+        .isEqualTo(constant);
+    Mockito.verify(codedInputStream).readSInt32();
+    Mockito.verify(registry).maybeGetConstantByTag(1);
+  }
+
+  @Test
+  public void memoizingDeserialize_memoize() throws SerializationException, IOException {
+    @SuppressWarnings("unchecked")
+    Memoizer.MemoizingCodec<?> memoizingCodec = Mockito.mock(Memoizer.MemoizingCodec.class);
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    doReturn(memoizingCodec).when(registry).maybeGetMemoizingCodecByTag(1);
+    CodedInputStream codedInputStream = Mockito.mock(CodedInputStream.class);
+    DeserializationContext deserializationContext =
+        new DeserializationContext(registry, ImmutableMap.of());
+    when(codedInputStream.readSInt32()).thenReturn(1);
+    assertThat(deserializationContext.getMemoizingCodecRecordedInInput(codedInputStream))
+        .isSameAs(memoizingCodec);
+    Mockito.verify(codedInputStream).readSInt32();
+    Mockito.verify(registry).maybeGetConstantByTag(1);
+    Mockito.verify(registry).maybeGetMemoizingCodecByTag(1);
+  }
+
+  @Test
+  public void memoizingDeserialize_codec() throws SerializationException, IOException {
+    Object returned = new Object();
+    @SuppressWarnings("unchecked")
+    ObjectCodec<Object> codec = Mockito.mock(ObjectCodec.class);
+    ObjectCodecRegistry.CodecDescriptor codecDescriptor =
+        Mockito.mock(ObjectCodecRegistry.CodecDescriptor.class);
+    doReturn(codec).when(codecDescriptor).getCodec();
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    when(registry.getCodecDescriptorByTag(1)).thenReturn(codecDescriptor);
+    CodedInputStream codedInputStream = Mockito.mock(CodedInputStream.class);
+    DeserializationContext deserializationContext =
+        new DeserializationContext(registry, ImmutableMap.of());
+    when(codec.deserialize(deserializationContext, codedInputStream)).thenReturn(returned);
+    when(codedInputStream.readSInt32()).thenReturn(1);
+    assertThat(
+            deserializationContext
+                .getMemoizingCodecRecordedInInput(codedInputStream)
+                .deserializePayload(
+                    deserializationContext,
+                    /*initial=*/ null,
+                    codedInputStream,
+                    new Memoizer.Deserializer()))
+        .isEqualTo(returned);
+    Mockito.verify(codedInputStream).readSInt32();
+    Mockito.verify(registry).maybeGetConstantByTag(1);
+    Mockito.verify(registry).maybeGetMemoizingCodecByTag(1);
+    Mockito.verify(registry).getCodecDescriptorByTag(1);
+    Mockito.verify(codecDescriptor).getCodec();
+    Mockito.verify(codec).deserialize(deserializationContext, codedInputStream);
+  }
+}
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistryTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistryTest.java
index 5724ccacea..6c6cc99604 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistryTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistryTest.java
@@ -111,20 +111,71 @@ public class ObjectCodecRegistryTest {
   }
 
   @Test
+  public void constantsOrderedByLastOccurrenceInIteration() {
+    Object constant1 = new Object();
+    Object constant2 = new Object();
+    ObjectCodecRegistry underTest1 =
+        ObjectCodecRegistry.newBuilder()
+            .setAllowDefaultCodec(false)
+            .addConstant(constant1)
+            .addConstant(constant2)
+            .addConstant(constant1)
+            .build();
+    ObjectCodecRegistry underTest2 =
+        ObjectCodecRegistry.newBuilder()
+            .setAllowDefaultCodec(false)
+            .addConstant(constant1)
+            .addConstant(constant2)
+            .build();
+    assertThat(underTest1.maybeGetTagForConstant(constant1)).isEqualTo(3);
+    assertThat(underTest1.maybeGetTagForConstant(constant2)).isEqualTo(2);
+    assertThat(underTest2.maybeGetTagForConstant(constant1)).isEqualTo(1);
+    assertThat(underTest2.maybeGetTagForConstant(constant2)).isEqualTo(2);
+  }
+
+  @Test
+  public void memoizingOrderedByClassName() {
+    Memoizer.MemoizingCodec<String> memoizingCodec1 =
+        new Memoizer.ObjectCodecAdaptor<>(SingletonCodec.of("value1", "mnemonic1"));
+    Memoizer.MemoizingCodec<Object> memoizingCodec2 =
+        new Memoizer.ObjectCodecAdaptor<>(SingletonCodec.of(new Object(), "mnemonic2"));
+    ObjectCodecRegistry underTest =
+        ObjectCodecRegistry.newBuilder()
+            .setAllowDefaultCodec(false)
+            .addMemoizing(memoizingCodec1)
+            .addMemoizing(memoizingCodec2)
+            .build();
+
+    assertThat(underTest.getMemoizingCodecDescriptor(Object.class).getMemoizingCodec())
+        .isEqualTo(memoizingCodec2);
+    assertThat(underTest.getMemoizingCodecDescriptor(String.class).getMemoizingCodec())
+        .isEqualTo(memoizingCodec1);
+    assertThat(underTest.maybeGetMemoizingCodecByTag(1)).isEqualTo(memoizingCodec2);
+    assertThat(underTest.maybeGetMemoizingCodecByTag(2)).isEqualTo(memoizingCodec1);
+  }
+
+  @Test
   public void testGetBuilder() throws NoCodecException {
     SingletonCodec<String> codec1 = SingletonCodec.of("value1", "mnemonic1");
     SingletonCodec<String> codec2 = SingletonCodec.of("value2", "mnemonic2");
+    Object constant = new Object();
+    Memoizer.MemoizingCodec<String> memoizingCodec = new Memoizer.ObjectCodecAdaptor<>(codec1);
 
     ObjectCodecRegistry underTest =
         ObjectCodecRegistry.newBuilder()
             .setAllowDefaultCodec(false)
             .add("foo", codec1)
             .add("bar", codec2)
+            .addConstant(constant)
+            .addMemoizing(memoizingCodec)
             .build();
 
     ObjectCodecRegistry copy = underTest.getBuilder().build();
     assertThat(copy.getCodecDescriptor("bar").getTag()).isEqualTo(1);
     assertThat(copy.getCodecDescriptor("foo").getTag()).isEqualTo(2);
+    assertThat(copy.getMemoizingCodecDescriptor(String.class).getMemoizingCodec())
+        .isEqualTo(memoizingCodec);
+    assertThat(copy.maybeGetTagForConstant(constant)).isNotNull();
     assertThrows(NoCodecException.class, () -> copy.getCodecDescriptor("baz"));
   }
 }
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContextTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContextTest.java
new file mode 100644
index 0000000000..bd05444e21
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContextTest.java
@@ -0,0 +1,147 @@
+// Copyright 2018 The Bazel Authors. 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.skyframe.serialization;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.when;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.protobuf.CodedOutputStream;
+import java.io.IOException;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
+
+/** Tests for {@link SerializationContext}. */
+@RunWith(JUnit4.class)
+public class SerializationContextTest {
+  @Test
+  public void nullSerialize() throws IOException, SerializationException {
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    CodedOutputStream codedOutputStream = Mockito.mock(CodedOutputStream.class);
+    SerializationContext serializationContext =
+        new SerializationContext(registry, ImmutableMap.of());
+    serializationContext.serialize(null, codedOutputStream);
+    Mockito.verify(codedOutputStream).writeSInt32NoTag(0);
+    Mockito.verifyZeroInteractions(registry);
+  }
+
+  @Test
+  public void constantSerialize() throws IOException, SerializationException {
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    when(registry.maybeGetTagForConstant(Mockito.anyObject())).thenReturn(1);
+    CodedOutputStream codedOutputStream = Mockito.mock(CodedOutputStream.class);
+    SerializationContext serializationContext =
+        new SerializationContext(registry, ImmutableMap.of());
+    Object constant = new Object();
+    serializationContext.serialize(constant, codedOutputStream);
+    Mockito.verify(codedOutputStream).writeSInt32NoTag(1);
+    Mockito.verify(registry).maybeGetTagForConstant(constant);
+  }
+
+  @Test
+  public void descriptorSerialize() throws SerializationException, IOException {
+    ObjectCodecRegistry.CodecDescriptor codecDescriptor =
+        Mockito.mock(ObjectCodecRegistry.CodecDescriptor.class);
+    when(codecDescriptor.getTag()).thenReturn(1);
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    when(registry.maybeGetTagForConstant(Mockito.anyObject())).thenReturn(null);
+    when(registry.getCodecDescriptor(String.class)).thenReturn(codecDescriptor);
+    CodedOutputStream codedOutputStream = Mockito.mock(CodedOutputStream.class);
+    SerializationContext underTest = new SerializationContext(registry, ImmutableMap.of());
+    underTest.serialize("string", codedOutputStream);
+    Mockito.verify(codedOutputStream).writeSInt32NoTag(1);
+    Mockito.verify(registry).maybeGetTagForConstant("string");
+    Mockito.verify(registry).getCodecDescriptor(String.class);
+    Mockito.verify(codecDescriptor).getTag();
+    Mockito.verify(codecDescriptor).serialize(underTest, "string", codedOutputStream);
+  }
+
+  @Test
+  public void memoizingSerialize_null() throws IOException, SerializationException {
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    CodedOutputStream codedOutputStream = Mockito.mock(CodedOutputStream.class);
+    SerializationContext serializationContext =
+        new SerializationContext(registry, ImmutableMap.of());
+    assertThat(serializationContext.recordAndMaybeGetMemoizingCodec(null, codedOutputStream))
+        .isNull();
+    Mockito.verify(codedOutputStream).writeSInt32NoTag(0);
+    Mockito.verifyZeroInteractions(registry);
+  }
+
+  @Test
+  public void memoizingSerialize_constant() throws IOException, SerializationException {
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    when(registry.maybeGetTagForConstant(Mockito.anyObject())).thenReturn(1);
+    CodedOutputStream codedOutputStream = Mockito.mock(CodedOutputStream.class);
+    SerializationContext serializationContext =
+        new SerializationContext(registry, ImmutableMap.of());
+    Object constant = new Object();
+    assertThat(serializationContext.recordAndMaybeGetMemoizingCodec(constant, codedOutputStream))
+        .isNull();
+    Mockito.verify(codedOutputStream).writeSInt32NoTag(1);
+    Mockito.verify(registry).maybeGetTagForConstant(constant);
+  }
+
+  @Test
+  public void memoizingSerialize_memoize() throws SerializationException, IOException {
+    @SuppressWarnings("unchecked")
+    Memoizer.MemoizingCodec<Object> memoizingCodec = Mockito.mock(Memoizer.MemoizingCodec.class);
+    @SuppressWarnings("unchecked")
+    ObjectCodecRegistry.MemoizingCodecDescriptor<Object> memoizingCodecDescriptor =
+        Mockito.mock(ObjectCodecRegistry.MemoizingCodecDescriptor.class);
+    when(memoizingCodecDescriptor.getTag()).thenReturn(1);
+    doReturn(memoizingCodec).when(memoizingCodecDescriptor).getMemoizingCodec();
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    when(registry.maybeGetTagForConstant(Mockito.anyObject())).thenReturn(null);
+    doReturn(memoizingCodecDescriptor).when(registry).getMemoizingCodecDescriptor(String.class);
+    CodedOutputStream codedOutputStream = Mockito.mock(CodedOutputStream.class);
+    SerializationContext underTest = new SerializationContext(registry, ImmutableMap.of());
+    assertThat(underTest.recordAndMaybeGetMemoizingCodec("string", codedOutputStream))
+        .isSameAs(memoizingCodec);
+    Mockito.verify(codedOutputStream).writeSInt32NoTag(1);
+    Mockito.verify(registry).maybeGetTagForConstant("string");
+    Mockito.verify(registry).getMemoizingCodecDescriptor(String.class);
+    Mockito.verifyNoMoreInteractions(registry);
+  }
+
+  @Test
+  public void memoizingSerialize_descriptor() throws SerializationException, IOException {
+    @SuppressWarnings("unchecked")
+    ObjectCodec<Object> codec = Mockito.mock(ObjectCodec.class);
+    ObjectCodecRegistry.CodecDescriptor codecDescriptor =
+        Mockito.mock(ObjectCodecRegistry.CodecDescriptor.class);
+    when(codecDescriptor.getTag()).thenReturn(1);
+    doReturn(codec).when(codecDescriptor).getCodec();
+    ObjectCodecRegistry registry = Mockito.mock(ObjectCodecRegistry.class);
+    when(registry.maybeGetTagForConstant(Mockito.anyObject())).thenReturn(null);
+    when(registry.getMemoizingCodecDescriptor(Mockito.anyObject())).thenReturn(null);
+    when(registry.getCodecDescriptor(String.class)).thenReturn(codecDescriptor);
+    CodedOutputStream codedOutputStream = Mockito.mock(CodedOutputStream.class);
+    SerializationContext underTest = new SerializationContext(registry, ImmutableMap.of());
+    underTest
+        .recordAndMaybeGetMemoizingCodec("string", codedOutputStream)
+        .serializePayload(underTest, "string", codedOutputStream, null);
+    Mockito.verify(codedOutputStream).writeSInt32NoTag(1);
+    Mockito.verify(registry).maybeGetTagForConstant("string");
+    Mockito.verify(registry).getMemoizingCodecDescriptor(String.class);
+    Mockito.verify(registry).getCodecDescriptor(String.class);
+    Mockito.verify(codecDescriptor).getTag();
+    Mockito.verify(codecDescriptor).getCodec();
+    Mockito.verify(codec).serialize(underTest, "string", codedOutputStream);
+  }
+}