Speedup DataflowError.withDefaultTrace (#11153)

Improves the speed of `ExecutionEnvironment.hasContextEnabled`.

# Important Notes
Local speedup of `Map_Error_Benchmark_Vector_ignore.Map_Id_All_Errors` benchmark is roughly ???.

Author: Akirathan

engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/error/ThrowErrorNode.java

@@ -4,6 +4,7 @@
 import com.oracle.truffle.api.nodes.Node;
 import org.enso.interpreter.dsl.BuiltinMethod;
 import org.enso.interpreter.runtime.error.DataflowError;
+import org.enso.interpreter.runtime.state.HasContextEnabledNode;
 import org.enso.interpreter.runtime.state.State;
 
 @BuiltinMethod(
@@ -12,7 +13,9 @@
     description = "Returns a new value error with given payload.",
     inlineable = true)
 public class ThrowErrorNode extends Node {
+  private @Child HasContextEnabledNode hasContextEnabledNode = HasContextEnabledNode.create();
+
   public Object execute(VirtualFrame giveMeAStackFrame, State state, Object payload) {
-    return DataflowError.withDefaultTrace(state, payload, this);
+    return DataflowError.withDefaultTrace(state, payload, this, hasContextEnabledNode);
   }
 }

engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/ContextIsEnabledNode.java

@@ -7,13 +7,15 @@
 import org.enso.interpreter.runtime.data.atom.Atom;
 import org.enso.interpreter.runtime.error.PanicException;
 import org.enso.interpreter.runtime.state.ExecutionEnvironment;
+import org.enso.interpreter.runtime.state.HasContextEnabledNode;
 
 @BuiltinMethod(
     type = "Context",
     name = "is_enabled_builtin",
     description = "Check if the context is enabled in the provided execution environment.")
 public class ContextIsEnabledNode extends Node {
   private @Child ExpectStringNode expectStringNode = ExpectStringNode.build();
+  private @Child HasContextEnabledNode hasContextEnabledNode = HasContextEnabledNode.create();
 
   Object execute(Atom self, Object environmentName) {
     String envName = expectStringNode.execute(environmentName);
@@ -26,6 +28,6 @@ Object execute(Atom self, Object environmentName) {
               .makeUnimplemented("execution environment mismatch");
       throw new PanicException(error, this);
     }
-    return currentEnv.hasContextEnabled(self.getConstructor().getName());
+    return hasContextEnabledNode.executeHasContextEnabled(currentEnv, self.getConstructor());
   }
 }

engine/runtime/src/main/java/org/enso/interpreter/runtime/EnsoContext.java

@@ -59,6 +59,7 @@
 import org.enso.interpreter.runtime.scope.TopLevelScope;
 import org.enso.interpreter.runtime.state.ExecutionEnvironment;
 import org.enso.interpreter.runtime.state.State;
+import org.enso.interpreter.runtime.state.WithContextNode;
 import org.enso.interpreter.runtime.util.TruffleFileSystem;
 import org.enso.librarymanager.ProjectLoadingFailure;
 import org.enso.librarymanager.resolved.LibraryRoot;
@@ -894,7 +895,9 @@ public void setExecutionEnvironment(ExecutionEnvironment executionEnvironment) {
   public ExecutionEnvironment enableExecutionEnvironment(Atom context, String environmentName) {
     ExecutionEnvironment original = globalExecutionEnvironment;
     if (original.getName().equals(environmentName)) {
-      setExecutionEnvironment(original.withContextEnabled(context));
+      var newExecEnv =
+          WithContextNode.getUncached().executeEnvironmentUpdate(original, context, true);
+      setExecutionEnvironment(newExecEnv);
     }
     return original;
   }
@@ -909,7 +912,9 @@ public ExecutionEnvironment enableExecutionEnvironment(Atom context, String envi
   public ExecutionEnvironment disableExecutionEnvironment(Atom context, String environmentName) {
     ExecutionEnvironment original = globalExecutionEnvironment;
     if (original.getName().equals(environmentName)) {
-      setExecutionEnvironment(original.withContextDisabled(context));
+      var newExecEnv =
+          WithContextNode.getUncached().executeEnvironmentUpdate(original, context, false);
+      setExecutionEnvironment(newExecEnv);
     }
     return original;
   }

engine/runtime/src/main/java/org/enso/interpreter/runtime/error/DataflowError.java

@@ -24,6 +24,7 @@
 import org.enso.interpreter.runtime.data.Type;
 import org.enso.interpreter.runtime.data.vector.ArrayLikeHelpers;
 import org.enso.interpreter.runtime.library.dispatch.TypesLibrary;
+import org.enso.interpreter.runtime.state.HasContextEnabledNode;
 import org.enso.interpreter.runtime.state.State;
 
 /**
@@ -51,15 +52,18 @@ public final class DataflowError extends AbstractTruffleException implements Ens
    * @param location the node in which the error was created
    * @return a new dataflow error
    */
-  public static DataflowError withDefaultTrace(State state, Object payload, Node location) {
+  public static DataflowError withDefaultTrace(
+      State state, Object payload, Node location, HasContextEnabledNode hasContextEnabledNode) {
     assert payload != null;
+    var ensoCtx = EnsoContext.get(location);
+    var dataflowStacktraceCtx = ensoCtx.getBuiltins().context().getDataflowStackTrace();
     boolean attachFullStackTrace =
         state == null
-            || EnsoContext.get(location)
-                .getExecutionEnvironment()
-                .hasContextEnabled("Dataflow_Stack_Trace");
+            || hasContextEnabledNode.executeHasContextEnabled(
+                ensoCtx.getExecutionEnvironment(), dataflowStacktraceCtx);
     if (attachFullStackTrace) {
-      var result = new DataflowError(payload, UNLIMITED_STACK_TRACE, location);
+      var result =
+          new DataflowError(payload, AbstractTruffleException.UNLIMITED_STACK_TRACE, location);
       TruffleStackTrace.fillIn(result);
       return result;
     } else {
@@ -68,8 +72,9 @@ public static DataflowError withDefaultTrace(State state, Object payload, Node l
     }
   }
 
+  /** Slow version of {@link #withDefaultTrace(State, Object, Node, HasContextEnabledNode)}. */
   public static DataflowError withDefaultTrace(Object payload, Node location) {
-    return withDefaultTrace(null, payload, location);
+    return withDefaultTrace(null, payload, location, HasContextEnabledNode.getUncached());
   }
 
   /**

engine/runtime/src/main/java/org/enso/interpreter/runtime/state/ContextPermissions.java

@@ -0,0 +1,4 @@
+package org.enso.interpreter.runtime.state;
+
+/** Fields correspond to the constructors of {@code Standard.Base.Runtime.Context} builtin type. */
+record ContextPermissions(boolean input, boolean output, boolean dataflowStacktrace) {}

engine/runtime/src/main/java/org/enso/interpreter/runtime/state/ExecutionEnvironment.java

@@ -1,18 +1,10 @@
 package org.enso.interpreter.runtime.state;
 
-import org.enso.interpreter.node.expression.builtin.runtime.Context;
-import org.enso.interpreter.runtime.EnsoContext;
-import org.enso.interpreter.runtime.data.atom.Atom;
-
 public class ExecutionEnvironment {
 
   private final String name;
 
-  // Ideally we would "just" use a map here. But that leads
-  // to native image build problems. This in turn leads to
-  // TruffleBoundary annotations which in turn leads to slow path.
-  private final String[] keys;
-  private final Boolean[] permissions;
+  final ContextPermissions permissions;
 
   public static final String LIVE_ENVIRONMENT_NAME = "live";
 
@@ -22,82 +14,25 @@ public class ExecutionEnvironment {
   public static final ExecutionEnvironment DESIGN =
       new ExecutionEnvironment(DESIGN_ENVIRONMENT_NAME);
 
-  private static final ExecutionEnvironment initLive(String name) {
-    String[] keys = new String[] {Context.INPUT_NAME, Context.OUTPUT_NAME};
-    Boolean[] permissions = new Boolean[] {true, true};
-    return new ExecutionEnvironment(name, keys, permissions);
+  private static ExecutionEnvironment initLive(String name) {
+    var permissions = new ContextPermissions(true, true, false);
+    return new ExecutionEnvironment(name, permissions);
   }
 
   public ExecutionEnvironment(String name) {
     this.name = name;
-    this.keys = new String[0];
-    this.permissions = new Boolean[0];
+    this.permissions = new ContextPermissions(false, false, false);
   }
 
-  private ExecutionEnvironment(String name, String[] keys, Boolean[] permissions) {
+  ExecutionEnvironment(String name, ContextPermissions permissions) {
     this.name = name;
-    this.keys = keys;
     this.permissions = permissions;
   }
 
   public String getName() {
     return this.name;
   }
 
-  public ExecutionEnvironment withContextEnabled(Atom context) {
-    return update(context, true);
-  }
-
-  public ExecutionEnvironment withContextDisabled(Atom context) {
-    return update(context, false);
-  }
-
-  private ExecutionEnvironment update(Atom context, boolean value) {
-    assert context.getConstructor().getType()
-        == EnsoContext.get(null).getBuiltins().context().getType();
-    int keyFound = -1;
-    for (int i = 0; i < keys.length; i++) {
-      if (keys[i].equals(context.getConstructor().getName())) {
-        keyFound = i;
-      }
-    }
-    String[] keys1;
-    Boolean[] permissions1;
-    if (keyFound != -1) {
-      keys1 = cloneArray(keys, new String[keys.length]);
-      permissions1 = cloneArray(permissions, new Boolean[keys.length]);
-      permissions1[keyFound] = value;
-    } else {
-      keys1 = cloneArray(keys, new String[keys.length + 1]);
-      permissions1 = cloneArray(permissions, new Boolean[keys.length + 1]);
-      keyFound = keys.length;
-      keys1[keyFound] = context.getConstructor().getName();
-      permissions1[keyFound] = value;
-    }
-    return new ExecutionEnvironment(name, keys1, permissions1);
-  }
-
-  private <T> T[] cloneArray(T[] fromArray, T[] toArray) {
-    for (int i = 0; i < fromArray.length; i++) {
-      toArray[i] = fromArray[i];
-    }
-    return toArray;
-  }
-
-  public Boolean hasContextEnabled(String context) {
-    int keyFound = -1;
-    for (int i = 0; i < keys.length; i++) {
-      if (keys[i].equals(context)) {
-        keyFound = i;
-      }
-    }
-    if (keyFound != -1) {
-      return permissions[keyFound];
-    } else {
-      return false;
-    }
-  }
-
   public static ExecutionEnvironment forName(String name) {
     switch (name) {
       case LIVE_ENVIRONMENT_NAME:

engine/runtime/src/main/java/org/enso/interpreter/runtime/state/HasContextEnabledNode.java

@@ -0,0 +1,62 @@
+package org.enso.interpreter.runtime.state;
+
+import com.oracle.truffle.api.CompilerDirectives;
+import com.oracle.truffle.api.dsl.Cached;
+import com.oracle.truffle.api.dsl.GenerateUncached;
+import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.nodes.Node;
+import org.enso.interpreter.runtime.EnsoContext;
+import org.enso.interpreter.runtime.data.atom.AtomConstructor;
+
+/**
+ * A node representing the functionality done by {@code Standard.Base.Runtime.Context.is_enabled}.
+ */
+@GenerateUncached
+public abstract class HasContextEnabledNode extends Node {
+
+  public static HasContextEnabledNode getUncached() {
+    return HasContextEnabledNodeGen.getUncached();
+  }
+
+  public static HasContextEnabledNode create() {
+    return HasContextEnabledNodeGen.create();
+  }
+
+  /**
+   * Returns true if the context represented by the given {@code runtimeCtxCtor} is enabled in the
+   * given {@code executionEnvironment}.
+   *
+   * @param runtimeCtxCtor Constructor of {@code Standard.Base.Runtime.Context}.
+   */
+  public abstract boolean executeHasContextEnabled(
+      ExecutionEnvironment executionEnvironment, AtomConstructor runtimeCtxCtor);
+
+  @Specialization(guards = "executionEnvironment == cachedEnv", limit = "3")
+  boolean cachedHasContextEnabled(
+      ExecutionEnvironment executionEnvironment,
+      AtomConstructor runtimeCtxCtor,
+      @Cached("executionEnvironment") ExecutionEnvironment cachedEnv) {
+    return doIt(cachedEnv, runtimeCtxCtor);
+  }
+
+  @Specialization(replaces = "cachedHasContextEnabled")
+  boolean uncachedHasContextEnabled(
+      ExecutionEnvironment executionEnvironment, AtomConstructor runtimeCtxCtor) {
+    return doIt(executionEnvironment, runtimeCtxCtor);
+  }
+
+  private boolean doIt(ExecutionEnvironment executionEnvironment, AtomConstructor runtimeCtxCtor) {
+    var ensoCtx = EnsoContext.get(this);
+    var contextBuiltin = ensoCtx.getBuiltins().context();
+    if (runtimeCtxCtor == contextBuiltin.getInput()) {
+      return executionEnvironment.permissions.input();
+    } else if (runtimeCtxCtor == contextBuiltin.getOutput()) {
+      return executionEnvironment.permissions.output();
+    } else if (runtimeCtxCtor == contextBuiltin.getDataflowStackTrace()) {
+      return executionEnvironment.permissions.dataflowStacktrace();
+    } else {
+      CompilerDirectives.transferToInterpreter();
+      throw ensoCtx.raiseAssertionPanic(this, "Unknown context: " + runtimeCtxCtor, null);
+    }
+  }
+}

engine/runtime/src/main/java/org/enso/interpreter/runtime/state/WithContextNode.java

@@ -0,0 +1,60 @@
+package org.enso.interpreter.runtime.state;
+
+import com.oracle.truffle.api.dsl.GenerateUncached;
+import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.nodes.Node;
+import org.enso.interpreter.runtime.EnsoContext;
+import org.enso.interpreter.runtime.data.atom.Atom;
+
+/**
+ * A node representing functionality done by {@code Standard.Base.Runtime.Context.with_enabled} and
+ * {@code Standard.Base.Runtime.Context.with_disabled}. That is, it enables or disables the given
+ * context in the current {@link ExecutionEnvironment execution environment}.
+ */
+@GenerateUncached
+public abstract class WithContextNode extends Node {
+  public static WithContextNode getUncached() {
+    return WithContextNodeGen.getUncached();
+  }
+
+  public static WithContextNode create() {
+    return WithContextNodeGen.create();
+  }
+
+  /**
+   * Returns a new {@link ExecutionEnvironment} with the given context enabled or disabled.
+   *
+   * @param current Current execution environment.
+   * @param context Atom of type {@code Standard.Base.Runtime.Context}.
+   * @param enabled Whether to enable or disable the context.
+   */
+  public abstract ExecutionEnvironment executeEnvironmentUpdate(
+      ExecutionEnvironment current, Atom context, boolean enabled);
+
+  @Specialization
+  ExecutionEnvironment doIt(ExecutionEnvironment current, Atom context, boolean enabled) {
+    var ensoCtx = EnsoContext.get(this);
+    var contextBuiltin = ensoCtx.getBuiltins().context();
+    if (context.getConstructor().getType() != contextBuiltin.getType()) {
+      throw ensoCtx.raiseAssertionPanic(this, "Invalid context type", null);
+    }
+    var ctor = context.getConstructor();
+    ContextPermissions newPermissions;
+    if (ctor == contextBuiltin.getInput()) {
+      newPermissions =
+          new ContextPermissions(
+              enabled, current.permissions.output(), current.permissions.dataflowStacktrace());
+    } else if (ctor == contextBuiltin.getOutput()) {
+      newPermissions =
+          new ContextPermissions(
+              current.permissions.input(), enabled, current.permissions.dataflowStacktrace());
+    } else if (ctor == contextBuiltin.getDataflowStackTrace()) {
+      newPermissions =
+          new ContextPermissions(
+              current.permissions.input(), current.permissions.output(), enabled);
+    } else {
+      throw ensoCtx.raiseAssertionPanic(this, "Unknown context: " + ctor, null);
+    }
+    return new ExecutionEnvironment(current.getName(), newPermissions);
+  }
+}