Exposing just a single Persistance.Reference abstraction (#10117)

* Exposing just a single Persistance.Reference abstraction * Reference.of(obj, deferredWrite) * Encapsulating Reference serde in PerReferencePeristance
enso-org · Jun 1, 2024 · 3f8e144 · 3f8e144
1 parent f0e9ffa
commit 3f8e144
Show file tree

Hide file tree

Showing 14 changed files with 387 additions and 110 deletions.
diff --git a/engine/runtime-parser/src/main/scala/org/enso/compiler/core/ir/Function.scala b/engine/runtime-parser/src/main/scala/org/enso/compiler/core/ir/Function.scala
@@ -46,7 +46,7 @@ object Function {
     */
   sealed case class Lambda(
     override val arguments: List[DefinitionArgument],
-    bodyReference: Persistance.InlineReference[Expression],
+    bodyReference: Persistance.Reference[Expression],
     location: Option[IdentifiedLocation],
     override val canBeTCO: Boolean,
     passData: MetadataStorage,
@@ -65,15 +65,15 @@ object Function {
     ) = {
       this(
         arguments,
-        Persistance.InlineReference.of(body),
+        Persistance.Reference.of(body, true),
         location,
         canBeTCO,
         passData,
         diagnostics
       )
     }
 
-    override lazy val body: Expression = bodyReference.get()
+    override lazy val body: Expression = bodyReference.get(classOf[Expression])
 
     /** Creates a copy of `this`.
       *
@@ -98,7 +98,7 @@ object Function {
       val res =
         Lambda(
           arguments,
-          Persistance.InlineReference.of(body),
+          Persistance.Reference.of(body, false),
           location,
           canBeTCO,
           passData,

diff --git a/...time-parser/src/main/scala/org/enso/compiler/core/ir/module/scope/definition/Method.scala b/...time-parser/src/main/scala/org/enso/compiler/core/ir/module/scope/definition/Method.scala
@@ -49,7 +49,7 @@ object Method {
     */
   sealed case class Explicit(
     override val methodReference: Name.MethodReference,
-    val bodyReference: Persistance.InlineReference[Expression],
+    val bodyReference: Persistance.Reference[Expression],
     val isStatic: Boolean,
     val isStaticWrapperForInstanceMethod: Boolean,
     override val location: Option[IdentifiedLocation],
@@ -67,7 +67,7 @@ object Method {
     ) = {
       this(
         methodReference,
-        Persistance.InlineReference.of(body),
+        Persistance.Reference.of(body, false),
         Explicit.computeIsStatic(body),
         Explicit.computeIsStaticWrapperForInstanceMethod(body),
         location,
@@ -76,7 +76,7 @@ object Method {
       );
     }
 
-    lazy val body: Expression = bodyReference.get()
+    lazy val body: Expression = bodyReference.get(classOf[Expression])
 
     /** Creates a copy of `this`.
       *
@@ -101,7 +101,7 @@ object Method {
     ): Explicit = {
       val res = Explicit(
         methodReference,
-        Persistance.InlineReference.of(body),
+        Persistance.Reference.of(body, false),
         isStatic,
         isStaticWrapperForInstanceMethod,
         location,

diff --git a/engine/runtime-parser/src/test/java/org/enso/compiler/core/IrPersistanceTest.java b/engine/runtime-parser/src/test/java/org/enso/compiler/core/IrPersistanceTest.java
@@ -332,14 +332,14 @@ public void hashMapIsLazy() throws Exception {
   @Test
   public void inlineReferenceIsLazy() throws Exception {
     var s1 = new LazyString("Hello");
-    var in = new InlineReferenceHolder(Persistance.InlineReference.of(s1));
+    var in = new InlineReferenceHolder(Persistance.Reference.of(s1, false));
 
     LazyString.forbidden = true;
     InlineReferenceHolder out = serde(InlineReferenceHolder.class, in, -1);
-    Persistance.InlineReference<CharSequence> ref = out.ref();
+    Persistance.Reference<CharSequence> ref = out.ref();
     LazyString.forbidden = false;
 
-    assertEquals(s1, ref.get());
+    assertEquals(s1, ref.get(Object.class));
   }
 
   @Test
@@ -560,5 +560,5 @@ public record RefHolder(Persistance.Reference<UUID> id) {
   }
 
   @Persistable(clazz = InlineReferenceHolder.class, id = 432437)
-  public record InlineReferenceHolder(Persistance.InlineReference<CharSequence> ref) {}
+  public record InlineReferenceHolder(Persistance.Reference<CharSequence> ref) {}
 }
diff --git a/lib/java/persistance-dsl/src/main/java/org/enso/persist/impl/PersistableProcessor.java b/lib/java/persistance-dsl/src/main/java/org/enso/persist/impl/PersistableProcessor.java
@@ -279,7 +279,7 @@ private int countInlineRef(List<? extends VariableElement> parameters) {
     var cnt = 0;
     for (var p : parameters) {
       var type = tu.asElement(tu.erasure(p.asType()));
-      if (type != null && type.getSimpleName().toString().equals("InlineReference")) {
+      if (type != null && type.getSimpleName().toString().equals("Reference")) {
         cnt++;
       }
     }

diff --git a/lib/java/persistance/src/main/java/org/enso/persist/PerBufferReference.java b/lib/java/persistance/src/main/java/org/enso/persist/PerBufferReference.java
@@ -52,6 +52,11 @@ final <T> T readObject(Class<T> clazz) throws IOException {
     return obj;
   }
 
+  @Override
+  boolean isDeferredWrite() {
+    return true;
+  }
+
   static <V> Reference<V> from(Persistance<V> p, InputCache buffer, int offset) {
     return new PerBufferReference<>(p, buffer, offset, false);
   }

diff --git a/lib/java/persistance/src/main/java/org/enso/persist/PerGenerator.java b/lib/java/persistance/src/main/java/org/enso/persist/PerGenerator.java
@@ -194,8 +194,12 @@ private int writeObjectAndReferences(Object obj) throws IOException {
     return tableAt;
   }
 
-  private static final class ReferenceOutput extends DataOutputStream
-      implements Persistance.Output {
+  static int registerReference(Persistance.Output out, Persistance.Reference<?> ref) {
+    var g = ((ReferenceOutput) out).generator;
+    return g.registerReference(ref);
+  }
+
+  static final class ReferenceOutput extends DataOutputStream implements Persistance.Output {
     private final PerGenerator generator;
 
     ReferenceOutput(PerGenerator g, ByteArrayOutputStream out) {
@@ -205,12 +209,6 @@ private static final class ReferenceOutput extends DataOutputStream
 
     @Override
     public <T> void writeInline(Class<T> clazz, T t) throws IOException {
-      if (Persistance.Reference.class == clazz) {
-        Persistance.Reference<?> ref = (Persistance.Reference<?>) t;
-        var id = this.generator.registerReference(ref);
-        writeInt(id);
-        return;
-      }
       var obj = generator.writeReplace.apply(t);
       var p = generator.map.forType(clazz);
       p.writeInline(obj, this);
@@ -255,5 +253,6 @@ private void register(Class<?> type, int length) {
     }
   }
 
+  static final int INLINED_REFERENCE_ID = -2;
   static final int NULL_REFERENCE_ID = -1;
 }
diff --git a/lib/java/persistance/src/main/java/org/enso/persist/PerInputImpl.java b/lib/java/persistance/src/main/java/org/enso/persist/PerInputImpl.java
@@ -1,5 +1,6 @@
 package org.enso.persist;
 
+import static org.enso.persist.PerGenerator.INLINED_REFERENCE_ID;
 import static org.enso.persist.PerGenerator.NULL_REFERENCE_ID;
 import static org.enso.persist.PerUtils.raise;
 
@@ -43,23 +44,25 @@ static Reference<?> readObject(ByteBuffer buf, Function<Object, Object> readReso
 
     var tableAt = buf.getInt(8);
     buf.position(tableAt);
-    InputCache.RawReferenceMap rawRefMap = InputCache.RawReferenceMap.readFromBuffer(buf);
+    var rawRefMap = InputCache.RawReferenceMap.readFromBuffer(buf);
     var cache = new InputCache(buf, readResolve, map, rawRefMap);
     return cache.getRef(0);
   }
 
+  static Persistance.Reference<?> findReference(Persistance.Input input, int refId) {
+    if (refId == NULL_REFERENCE_ID) {
+      return Persistance.Reference.none();
+    }
+    if (refId != INLINED_REFERENCE_ID) {
+      var impl = (PerInputImpl) input;
+      var ref = impl.cache.getRef(refId);
+      return ref;
+    }
+    return null;
+  }
+
   @Override
   public <T> T readInline(Class<T> clazz) throws IOException {
-    if (clazz == Persistance.Reference.class) {
-      var refId = readInt();
-      if (refId == NULL_REFERENCE_ID) {
-        var nullReference = Persistance.Reference.none();
-        return clazz.cast(nullReference);
-      }
-
-      var ref = cache.getRef(refId);
-      return clazz.cast(ref);
-    }
     Persistance<T> p = cache.map().forType(clazz);
     T res = p.readWith(this);
     var resolve = cache.resolveObject(res);
@@ -244,7 +247,7 @@ static <T> Reference<T> readIndirectAsReference(
       InputCache buffer, PerMap map, Input in, Class<T> clazz) throws IOException {
     var at = in.readInt();
     if (at < 0) {
-      return null;
+      return Reference.none();
     }
     var id = in.readInt();
     var p = map.forId(id);

diff --git a/lib/java/persistance/src/main/java/org/enso/persist/PerMap.java b/lib/java/persistance/src/main/java/org/enso/persist/PerMap.java
@@ -16,6 +16,7 @@ final class PerMap {
     var loader = PerMap.class.getClassLoader();
     var lookup = Lookups.metaInfServices(loader);
     var all = new ArrayList<Persistance>();
+    all.add(PerReferencePeristance.INSTANCE);
     all.addAll(lookup.lookupAll(Persistance.class));
     ALL = all;
   }

diff --git a/lib/java/persistance/src/main/java/org/enso/persist/PerMemoryReference.java b/lib/java/persistance/src/main/java/org/enso/persist/PerMemoryReference.java
@@ -1,6 +1,7 @@
 package org.enso.persist;
 
-final class PerMemoryReference<T> extends Persistance.Reference<T> {
+sealed class PerMemoryReference<T> extends Persistance.Reference<T>
+    permits PerMemoryReference.Deferred {
   static final Persistance.Reference<?> NULL = new PerMemoryReference<>(null);
   private final T value;
 
@@ -11,4 +12,15 @@ final class PerMemoryReference<T> extends Persistance.Reference<T> {
   final T value() {
     return value;
   }
+
+  @Override
+  boolean isDeferredWrite() {
+    return Deferred.class == getClass();
+  }
+
+  static final class Deferred<T> extends PerMemoryReference<T> {
+    Deferred(T obj) {
+      super(obj);
+    }
+  }
 }
diff --git a/lib/java/persistance/src/main/java/org/enso/persist/PerReferencePersistance.java b/lib/java/persistance/src/main/java/org/enso/persist/PerReferencePersistance.java
@@ -0,0 +1,39 @@
+package org.enso.persist;
+
+import static org.enso.persist.PerGenerator.INLINED_REFERENCE_ID;
+
+import java.io.IOException;
+import org.enso.persist.Persistance.Reference;
+
+final class PerReferencePeristance extends Persistance<Reference> {
+  static final Persistance<Reference> INSTANCE = new PerReferencePeristance();
+
+  private PerReferencePeristance() {
+    super(Reference.class, true, 4320);
+  }
+
+  @SuppressWarnings("unchecked")
+  @Override
+  protected void writeObject(Reference ref, Output out) throws IOException {
+    if (ref.isDeferredWrite()) {
+      var refId = PerGenerator.registerReference(out, ref);
+      out.writeInt(refId);
+    } else {
+      out.writeInt(INLINED_REFERENCE_ID);
+      var obj = ref.get(Object.class);
+      out.writeObject(obj);
+    }
+  }
+
+  @SuppressWarnings("unchecked")
+  @Override
+  protected Reference readObject(Input in) throws IOException, ClassNotFoundException {
+    var refId = in.readInt();
+    if (refId != INLINED_REFERENCE_ID) {
+      return PerInputImpl.findReference(in, refId);
+    } else {
+      var ref = in.readReference(Object.class);
+      return ref;
+    }
+  }
+}