Turing allTypesWith method into EnsoMultiType.AllTypesWith node

enso-org · Dec 28, 2024 · ebe0553 · ebe0553
1 parent 4dacf53
commit ebe0553
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 21 deletions.
diff --git a/engine/runtime/src/main/java/org/enso/interpreter/runtime/data/EnsoMultiType.java b/engine/runtime/src/main/java/org/enso/interpreter/runtime/data/EnsoMultiType.java
@@ -38,18 +38,6 @@ final int typesLength() {
     return types.length;
   }
 
-  final Type[] allTypesWith(EnsoMultiType nextOrNull) {
-    if (nextOrNull == null || nextOrNull.types.length == 0) {
-      return this.types.clone();
-    } else {
-      var next = nextOrNull;
-      var arr = new Type[this.types.length + next.types.length];
-      System.arraycopy(this.types, 0, arr, 0, types.length);
-      System.arraycopy(next.types, 0, arr, types.length, next.types.length);
-      return arr;
-    }
-  }
-
   final Type firstType() {
     return types[0];
   }
@@ -130,4 +118,49 @@ int findsAnIndexOfAType(Type type, EnsoMultiType mt) {
       return index;
     }
   }
+
+  @GenerateUncached
+  abstract static class AllTypesWith extends Node {
+    private static final String INLINE_CACHE_LIMIT = "5";
+
+    @NeverDefault
+    static AllTypesWith getUncached() {
+      return EnsoMultiTypeFactory.AllTypesWithNodeGen.getUncached();
+    }
+
+    @NeverDefault
+    static AllTypesWith create() {
+      return EnsoMultiTypeFactory.AllTypesWithNodeGen.create();
+    }
+
+    abstract Type[] executeAllTypes(EnsoMultiType first, EnsoMultiType second);
+
+    @Specialization(
+        limit = INLINE_CACHE_LIMIT,
+        guards = {
+          "self == cachedSelf",
+          "nextOrNull == cachedNextOrNull",
+        })
+    Type[] optimizeForTypes(
+        EnsoMultiType self,
+        EnsoMultiType nextOrNull,
+        @Cached("self") EnsoMultiType cachedSelf,
+        @Cached("nextOrNull") EnsoMultiType cachedNextOrNull,
+        @Cached("slowlyComputeTypes(self, nextOrNull)") Type[] result) {
+      return result;
+    }
+
+    @Specialization(replaces = "optimizeForTypes")
+    Type[] slowlyComputeTypes(EnsoMultiType self, EnsoMultiType nextOrNull) {
+      if (nextOrNull == null || nextOrNull.types.length == 0) {
+        return self.types.clone();
+      } else {
+        var next = nextOrNull;
+        var arr = new Type[self.types.length + next.types.length];
+        System.arraycopy(self.types, 0, arr, 0, self.types.length);
+        System.arraycopy(next.types, 0, arr, self.types.length, next.types.length);
+        return arr;
+      }
+    }
+  }
 }
diff --git a/engine/runtime/src/main/java/org/enso/interpreter/runtime/data/EnsoMultiValue.java b/engine/runtime/src/main/java/org/enso/interpreter/runtime/data/EnsoMultiValue.java
@@ -32,6 +32,7 @@
 import org.enso.interpreter.runtime.EnsoContext;
 import org.enso.interpreter.runtime.callable.UnresolvedSymbol;
 import org.enso.interpreter.runtime.callable.function.Function;
+import org.enso.interpreter.runtime.data.EnsoMultiType.AllTypesWith;
 import org.enso.interpreter.runtime.data.vector.ArrayLikeHelpers;
 import org.enso.interpreter.runtime.library.dispatch.TypesLibrary;
 import org.graalvm.collections.Pair;
@@ -158,11 +159,12 @@ final Type getType() {
   }
 
   @ExportMessage
-  final Type[] allTypes(boolean includeExtraTypes) {
+  final Type[] allTypes(
+      boolean includeExtraTypes, @Cached EnsoMultiType.AllTypesWith allTypesWith) {
     if (!includeExtraTypes) {
-      return dispatch.allTypesWith(null);
+      return allTypesWith.executeAllTypes(dispatch, null);
     } else {
-      return dispatch.allTypesWith(extra);
+      return allTypesWith.executeAllTypes(dispatch, extra);
     }
   }
 
@@ -494,25 +496,31 @@ Object invokeMember(
   @TruffleBoundary
   @Override
   public String toString() {
-    var both = dispatch.allTypesWith(extra);
+    var both = EnsoMultiType.AllTypesWith.getUncached().executeAllTypes(dispatch, extra);
     return Stream.of(both).map(t -> t.getName()).collect(Collectors.joining(" & "));
   }
 
   /** Casts {@link EnsoMultiValue} to requested type effectively. */
   public static final class CastToNode extends Node {
     private static final CastToNode UNCACHED =
-        new CastToNode(EnsoMultiType.FindIndexNode.getUncached(), NewNode.getUncached());
+        new CastToNode(
+            EnsoMultiType.FindIndexNode.getUncached(),
+            NewNode.getUncached(),
+            AllTypesWith.getUncached());
     @Child private EnsoMultiType.FindIndexNode findNode;
     @Child private NewNode newNode;
+    @Child private AllTypesWith allTypesWith;
 
-    private CastToNode(EnsoMultiType.FindIndexNode f, NewNode n) {
+    private CastToNode(EnsoMultiType.FindIndexNode f, NewNode n, AllTypesWith a) {
       this.findNode = f;
       this.newNode = n;
+      this.allTypesWith = a;
     }
 
     @NeverDefault
     public static CastToNode create() {
-      return new CastToNode(EnsoMultiType.FindIndexNode.create(), NewNode.create());
+      return new CastToNode(
+          EnsoMultiType.FindIndexNode.create(), NewNode.create(), AllTypesWith.create());
     }
 
     @NeverDefault
@@ -540,7 +548,7 @@ public final Object findTypeOrNull(
       }
       if (i != -1) {
         if (reorderOnly) {
-          var copyTypes = mv.dispatch.allTypesWith(mv.extra);
+          var copyTypes = allTypesWith.executeAllTypes(mv.dispatch, mv.extra);
           var copyValues = mv.values.clone();
           copyTypes[0] = copyTypes[i];
           copyValues[0] = copyValues[i];
@@ -567,7 +575,7 @@ public final Pair<Function, Type> resolveSymbol(
       MethodResolverNode node, UnresolvedSymbol symbol) {
     var ctx = EnsoContext.get(node);
     Pair<Function, Type> foundAnyMethod = null;
-    for (var t : dispatch.allTypesWith(null)) {
+    for (var t : EnsoMultiType.AllTypesWith.getUncached().executeAllTypes(dispatch, null)) {
       var fnAndType = node.execute(t, symbol);
       if (fnAndType != null) {
         if (dispatch.typesLength() == 1 || fnAndType.getRight() != ctx.getBuiltins().any()) {