1 files changed, 113 insertions, 0 deletions
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/serializers/CanonicalReferenceResolver.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/serializers/CanonicalReferenceResolver.java
new file mode 100644
index 0000000000..daa338984b
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/serializers/CanonicalReferenceResolver.java
@@ -0,0 +1,113 @@
+// 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.serializers;
+
+import com.esotericsoftware.kryo.Kryo;
+import com.esotericsoftware.kryo.ReferenceResolver;
+import com.esotericsoftware.kryo.util.Util;
+import java.util.ArrayList;
+import java.util.HashMap;
+
+/**
+ * {@link ReferenceResolver} implementation that uses object equality.
+ *
+ * <p>Provided that underlying objects implement equals correctly, this produces canonical, stable,
+ * serialized representations.
+ *
+ * <p>References must match also on class for stability of serialized representation. For example,
+ * {@code ArrayList} and {@code LinkedList} objects might evaluate equals, but Kryo has different
+ * serialized representations because each distinct type has a different registration index.
+ *
+ * <p>TODO(shahan): consider changing ClassResolver to allow multiple classes to share a
+ * registration in Kryo so this can be avoided.
+ */
+public class CanonicalReferenceResolver implements ReferenceResolver {
+
+  private final HashMap<Element, Integer> written;
+  private final ArrayList<Object> read;
+
+  CanonicalReferenceResolver() {
+    this.written = new HashMap<>();
+    this.read = new ArrayList<>();
+  }
+
+  @Override
+  public void setKryo(Kryo unusedKryo) {}
+
+  @Override
+  public int getWrittenId(Object object) {
+    return written.getOrDefault(new Element(object), -1);
+  }
+
+  @Override
+  public int addWrittenObject(Object object) {
+    int id = written.size();
+    written.put(new Element(object), id);
+    return id;
+  }
+
+  @Override
+  public int nextReadId(Class unusedType) {
+    int id = read.size();
+    read.add(null);
+    return id;
+  }
+
+  @Override
+  public void setReadObject(int id, Object object) {
+    read.set(id, object);
+  }
+
+  @Override
+  public Object getReadObject(Class unusedType, int id) {
+    return read.get(id);
+  }
+
+  @Override
+  public void reset() {
+    written.clear();
+    read.clear();
+  }
+
+  @Override
+  public boolean useReferences(Class type) {
+    return !Util.isWrapperClass(type);
+  }
+
+  /**
+   * Wrapper to ensure types are equal in addition to underlying objects.
+   *
+   * <p>Despite the overhead introduced by churn, this is empirically more efficient than {@link
+   * com.google.common.collect.Table}.
+   */
+  private static class Element {
+    public final Object contents;
+
+    private Element(Object contents) {
+      this.contents = contents;
+    }
+
+    @Override
+    public int hashCode() {
+      return contents.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      Element that = (Element) obj;
+      return contents.getClass().equals(that.contents.getClass()) && contents.equals(that.contents);
+    }
+  }
+}