3% speedup with LazyMap and MetadataStorage

engine/runtime-parser/src/main/java/org/enso/compiler/core/ir/IrLazyMap.java

@@ -0,0 +1,85 @@
+package org.enso.compiler.core.ir;
+
+import java.io.IOException;
+import java.util.AbstractMap;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+import org.enso.persist.Persistance;
+import org.enso.persist.Persistance.Reference;
+
+
+final class IrLazyMap<K, V> extends AbstractMap<K, V> {
+  private final Map<K, Entry<K, V>> delegate;
+
+  @SuppressWarnings("unchecked")
+  IrLazyMap(Persistance.Input in) throws IOException {
+    var map = new LinkedHashMap<K, Entry<K, V>>();
+    var n = in.readInt();
+    for (var i = 0; i < n; i++) {
+      var key = (K) in.readObject();
+      var ref = (Reference<V>) in.readReference(Object.class);
+      var en = new En<>(key, ref);
+      map.put(key, en);
+    }
+    this.delegate = map;
+  }
+
+  @Override
+  public Set<Entry<K, V>> entrySet() {
+    return new LinkedHashSet<>(this.delegate.values());
+  }
+
+  @Override
+  public V get(Object key) {
+    var entry = this.delegate.get(key);
+    return entry == null ? null : entry.getValue();
+  }
+
+  private static final class En<K, V> implements Entry<K, V> {
+    private final K key;
+    private final Reference<V> ref;
+
+    En(K key, Reference<V> ref) {
+      this.key = key;
+      this.ref = ref;
+    }
+
+    @Override
+    public K getKey() {
+      return key;
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public V getValue() {
+      return (V) ref.get(Object.class);
+    }
+
+    @Override
+    public V setValue(V value) {
+      throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public int hashCode() {
+      int hash = 7;
+      hash = 29 * hash + Objects.hashCode(this.key);
+      return hash;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      if (this == obj) {
+        return true;
+      }
+      if (obj instanceof En<?, ?> other) {
+        return Objects.equals(this.key, other.key);
+      }
+      return false;
+    }
+  }
+}

engine/runtime-parser/src/main/java/org/enso/compiler/core/ir/IrPersistance.java

@@ -1,7 +1,6 @@
 package org.enso.compiler.core.ir;
 
 import java.io.IOException;
-import java.util.HashMap;
 import java.util.Map;
 import java.util.UUID;
 import org.enso.compiler.core.ir.expression.Application;
@@ -230,14 +229,9 @@ protected void writeObject(scala.collection.immutable.Map map, Output out) throw
     @SuppressWarnings("unchecked")
     protected scala.collection.immutable.Map readObject(Input in)
         throws IOException, ClassNotFoundException {
-      var size = in.readInt();
-      var map = scala.collection.immutable.Map$.MODULE$.empty();
-      for (var i = 0; i < size; i++) {
-        var key = in.readObject();
-        var value = in.readObject();
-        map = map.$plus(new Tuple2(key, value));
-      }
-      return map;
+      var map = new IrLazyMap(in);
+      var immutableMap = IrLazyImMap$.MODULE$.apply(map);
+      return immutableMap;
     }
   }
 
@@ -309,13 +303,13 @@ protected scala.collection.immutable.Set readObject(Input in)
   }
 
   @ServiceProvider(service = Persistance.class)
-  public static final class PersistMap extends Persistance<HashMap> {
+  public static final class PersistMap extends Persistance<Map> {
     public PersistMap() {
-      super(HashMap.class, true, 4440);
+      super(Map.class, true, 4440);
     }
 
     @Override
-    protected void writeObject(HashMap m, Output out) throws IOException {
+    protected void writeObject(Map m, Output out) throws IOException {
       var size = m.size();
       out.writeInt(size);
       var it = m.entrySet().iterator();
@@ -328,15 +322,8 @@ protected void writeObject(HashMap m, Output out) throws IOException {
 
     @Override
     @SuppressWarnings("unchecked")
-    protected HashMap readObject(Input in) throws IOException, ClassNotFoundException {
-      var size = in.readInt();
-      var map = new HashMap<Object, Object>();
-      for (var i = 0; i < size; i++) {
-        var key = in.readObject();
-        var value = in.readObject();
-        map.put(key, value);
-      }
-      return map;
+    protected Map readObject(Input in) throws IOException, ClassNotFoundException {
+      return new IrLazyMap(in);
     }
   }
 

engine/runtime-parser/src/main/scala/org/enso/compiler/core/ir/IrLazyImMap.scala

@@ -0,0 +1,43 @@
+package org.enso.compiler.core.ir
+
+import scala.collection.immutable.StrictOptimizedMapOps
+import scala.collection.immutable.Map
+import scala.collection.Iterator
+
+final private class IrLazyImMap[K, V](
+  val underlying: java.util.Map[K, V]
+) extends Map[K, V]
+    with StrictOptimizedMapOps[K, V, Map, Map[K, V]] {
+  def removed(key: K) =
+    new scala.jdk.CollectionConverters.MapHasAsScala(underlying).asScala.toMap
+      .removed(key)
+
+  def updated[V1 >: V](key: K, value: V1) =
+    new scala.jdk.CollectionConverters.MapHasAsScala(underlying).asScala.toMap
+      .updated(key, value)
+
+  def get(key: K) = Option(underlying.get(key))
+
+  def iterator = new Iterator[(K, V)]() {
+    val keys             = underlying.entrySet().iterator()
+    def hasNext: Boolean = keys.hasNext()
+    def next(): (K, V) = {
+      val entry = keys.next()
+      (entry.getKey, entry.getValue)
+    }
+  }
+
+  override def size = underlying.size
+
+  // Has to override the `contains` method because the underlying implementation
+  // calls `get(key)`
+  override def contains(key: K): Boolean = {
+    underlying.containsKey(key)
+  }
+
+}
+
+object IrLazyImMap {
+  def apply[K, V](map: java.util.Map[K, V]): Map[K, V] =
+    new IrLazyImMap[K, V](map)
+}

engine/runtime-parser/src/test/java/org/enso/compiler/core/IrPersistanceTest.java

@@ -16,6 +16,7 @@
 import org.enso.compiler.core.ir.Module;
 import org.enso.persist.Persistable;
 import org.enso.persist.Persistance;
+import org.junit.Before;
 import org.junit.Test;
 import org.openide.util.lookup.ServiceProvider;
 
@@ -25,6 +26,11 @@
 import scala.collection.immutable.Seq;
 
 public class IrPersistanceTest {
+  @Before
+  public void resetDebris() {
+    LazySeq.forbidden = false;
+  }
+
   @Test
   public void locationTest() throws Exception {
     var l = new Location(12, 33);
@@ -61,6 +67,30 @@ public void scalaMap() throws Exception {
     assertEquals(in, out);
   }
 
+  @Test
+  @SuppressWarnings("unchecked")
+  public void scalaImmutableMapIsLazy() throws Exception {
+    var s1 = new LazySeq("Hello");
+    var s2 = new LazySeq("World");
+    var in = (scala.collection.immutable.Map) scala.collection.immutable.Map$.MODULE$.empty()
+      .$plus(new Tuple2("Hello", s1))
+      .$plus(new Tuple2("World", s2));
+
+    LazySeq.forbidden = true;
+    var out = (scala.collection.immutable.Map)
+      serde(scala.collection.immutable.Map.class, in, 64);
+
+    assertEquals("Two pairs element", 2, out.size());
+    assertEquals("Two keys", 2, out.keySet().size());
+    assertTrue(out.keySet().contains("Hello"));
+    assertTrue(out.keySet().contains("World"));
+    LazySeq.forbidden = false;
+
+    assertEquals(s1, out.get("Hello").get());
+    assertEquals(s2, out.get("World").get());
+    assertEquals(in, out);
+  }
+
   @Test
   @SuppressWarnings("unchecked")
   public void scalaHashMap() throws Exception {
@@ -184,11 +214,34 @@ public void hashMap() throws Exception {
     map.put("two", "duo");
     map.put("ten", "tre");
 
-    var out = serde(HashMap.class, map, -1);
+    var out = serde(java.util.Map.class, map, -1);
 
     assertEquals("Same", map, out);
   }
 
+  @Test
+  @SuppressWarnings("unchecked")
+  public void hashMapIsLazy() throws Exception {
+    var s1 = new LazySeq("Hello");
+    var s2 = new LazySeq("World");
+    var in = new HashMap<String, LazySeq>();
+    in.put("Hello", s1);
+    in.put("World", s2);
+
+    LazySeq.forbidden = true;
+    var out = serde(java.util.Map.class, in, 64);
+
+    assertEquals("Two pairs element", 2, out.size());
+    assertEquals("Two keys", 2, out.keySet().size());
+    assertTrue(out.keySet().contains("Hello"));
+    assertTrue(out.keySet().contains("World"));
+    LazySeq.forbidden = false;
+
+    assertEquals(s1, out.get("Hello"));
+    assertEquals(s2, out.get("World"));
+    assertEquals(in, out);
+  }
+
   @Test
   public void readResolve() throws Exception {
     var in = new Service(5);