Run node in a different execution environment (#11173)

close #10719

Changelog:
- add: optional `expressionConfigs` parameter to the `executionContext/recompute` request
- update: `IdExecutionInstrument` allowing to run a single node in a specified execution environment
- refactor: move tests related to the recompute request to a separate test suite. Otherwise the `RuntimeServerTest` is becoming too bloated

# Important Notes
The updated `executionContext/recompute` request.

```typescript
interface ExecutionContextRecomputeParameters {
/** The execution context identifier. */
contextId: ContextId;

/** The expressions that will be invalidated before the execution.
*
*  Only the provided expression ids are invalidated excluding the dependencies.
*/
invalidatedExpressions?: "all" | ExpressionId[];

/** The execution environment that will be used in the execution. */
executionEnvironment?: ExecutionEnvironment;

/** The execution configurations for particular expressions.
*
*  The provided expressions will be invalidated from the cache with the
*  dependencies. The result of the execution will stay in the cache until the
*  cache is invalidated by editing the node or other means.
*/
expressionConfigs?: ExpressionConfig[];
}

interface ExpressionConfig {
/** The expression identifier. */
expressionId: ExpressionId;
/** The execution environment that should be used to run this expression. */
executionEnvironment?: ExecutionEnvironment;
}
```

#### Use cases

- to re-run a single node without re-running the dependent nodes (subtree), put the node id in the `invalidatedExpressions` parameter.
- to re-run a node with dependent nodes (subtree), put the node id in the `expressionConfigs` parameter with empty `executionEnvironment`
- to re-run a node in a different execution environment, put the node  id in the `expressionConfigs` and specify the `executionEnvieronment`

docs/language-server/protocol-language-server.md

@@ -25,6 +25,8 @@ transport formats, please look [here](./protocol-architecture).
   - [`MethodCall`](#methodcall)
   - [`MethodPointer`](#methodpointer)
   - [`ProfilingInfo`](#profilinginfo)
+  - [`ExecutionEnvironment`](#executionenvironment)
+  - [`ExpressionConfig`](#expressionConfig)
   - [`ExpressionUpdate`](#expressionupdate)
   - [`ExpressionUpdatePayload`](#expressionupdatepayload)
   - [`VisualizationConfiguration`](#visualizationconfiguration)
@@ -343,6 +345,19 @@ The execution environment of Enso runtime.
 type ExecutionEnvironment = Design | Live;
 ```
 
+### `ExpressionConfig`
+
+The expression configuration used in the recompute request.
+
+```typescript
+interface ExpressionConfig {
+  /** The expression identifier. */
+  expressionId: ExpressionId;
+  /** The execution environment that should be used to run this expression. */
+  executionEnvironment?: ExecutionEnvironment;
+}
+```
+
 ### `ExpressionUpdate`
 
 An update about the computed expression.
@@ -3633,10 +3648,23 @@ May include a list of expressions for which caches should be invalidated.
 interface ExecutionContextRecomputeParameters {
   /** The execution context identifier. */
   contextId: ContextId;
-  /** The expressions that will be invalidated before the execution. */
+
+  /** The expressions that will be invalidated before the execution.
+   *
+   *  Only the provided expression ids are invalidated excluding the dependencies.
+   */
   invalidatedExpressions?: "all" | ExpressionId[];
+
   /** The execution environment that will be used in the execution. */
   executionEnvironment?: ExecutionEnvironment;
+
+  /** The execution configurations for particular expressions.
+   *
+   *  The provided expressions will be invalidated from the cache with the
+   *  dependencies. The result of the execution will stay in the cache until the
+   *  cache is invalidated by editing the node or other means.
+   */
+  expressionConfigs?: ExpressionConfig[];
 }
 ```
 

engine/language-server/src/main/scala/org/enso/languageserver/requesthandler/executioncontext/RecomputeHandler.scala

@@ -48,7 +48,8 @@ class RecomputeHandler(
       session,
       msg.params.contextId,
       msg.params.invalidatedExpressions,
-      msg.params.executionEnvironment
+      msg.params.executionEnvironment,
+      msg.params.expressionConfigs.getOrElse(Seq())
     )
 
   override protected def positiveResponse(

engine/language-server/src/main/scala/org/enso/languageserver/runtime/ContextRegistry.scala

@@ -193,7 +193,8 @@ final class ContextRegistry(
             client,
             contextId,
             expressions,
-            environment
+            environment,
+            expressionConfigs
           ) =>
         if (store.hasContext(client.clientId, contextId)) {
           val handler =
@@ -210,7 +211,8 @@ final class ContextRegistry(
             Api.RecomputeContextRequest(
               contextId,
               invalidatedExpressions,
-              environment.map(ExecutionEnvironments.toApi)
+              environment.map(ExecutionEnvironments.toApi),
+              expressionConfigs.map(_.toApi)
             )
           )
         } else {

engine/language-server/src/main/scala/org/enso/languageserver/runtime/ContextRegistryProtocol.scala

@@ -98,12 +98,14 @@ object ContextRegistryProtocol {
     * @param contextId execution context identifier
     * @param invalidatedExpressions the expressions that should be invalidated
     * @param executionEnvironment the environment that should be used for execution
+    * @param expressionConfigs the execution configurations for particular expressions
     */
   case class RecomputeContextRequest(
     rpcSession: JsonSession,
     contextId: ContextId,
     invalidatedExpressions: Option[InvalidatedExpressions],
-    executionEnvironment: Option[ExecutionEnvironments.ExecutionEnvironment]
+    executionEnvironment: Option[ExecutionEnvironments.ExecutionEnvironment],
+    expressionConfigs: Seq[ExpressionConfig]
   )
 
   /** A response about recomputing the context.

engine/language-server/src/main/scala/org/enso/languageserver/runtime/ExecutionApi.scala

@@ -108,7 +108,8 @@ object ExecutionApi {
     case class Params(
       contextId: ContextId,
       invalidatedExpressions: Option[InvalidatedExpressions],
-      executionEnvironment: Option[ExecutionEnvironments.ExecutionEnvironment]
+      executionEnvironment: Option[ExecutionEnvironments.ExecutionEnvironment],
+      expressionConfigs: Option[Seq[ExpressionConfig]]
     )
 
     implicit

engine/language-server/src/main/scala/org/enso/languageserver/runtime/ExpressionConfig.scala

@@ -0,0 +1,23 @@
+package org.enso.languageserver.runtime
+
+import org.enso.polyglot.runtime.Runtime.Api
+
+import java.util.UUID
+
+/** The expression configuration used in the recompute request.
+  *
+  * @param expressionId the expression identifier
+  * @param executionEnvironment the execution environment used to run this expression
+  */
+case class ExpressionConfig(
+  expressionId: UUID,
+  executionEnvironment: Option[ExecutionEnvironments.ExecutionEnvironment]
+) {
+
+  /** Convert this expression config to the runtime API. */
+  def toApi: Api.ExpressionConfig =
+    Api.ExpressionConfig(
+      expressionId,
+      executionEnvironment.map(ExecutionEnvironments.toApi)
+    )
+}

engine/language-server/src/main/scala/org/enso/languageserver/runtime/InvalidatedExpressions.scala

@@ -3,8 +3,7 @@ package org.enso.languageserver.runtime
 import io.circe.syntax._
 import io.circe.{Decoder, DecodingFailure, Encoder}
 
-/** A request to invalidate expressions.
-  */
+/** A request to invalidate expressions. */
 sealed trait InvalidatedExpressions
 
 object InvalidatedExpressions {

engine/language-server/src/test/scala/org/enso/languageserver/websocket/json/ContextRegistryTest.scala

@@ -449,7 +449,8 @@ class ContextRegistryTest extends BaseServerTest with ReportLogsOnFailure {
                 Api.RecomputeContextRequest(
                   `contextId`,
                   Some(Api.InvalidatedExpressions.All()),
-                  None
+                  None,
+                  Seq()
                 )
               ) =>
             requestId
@@ -527,7 +528,92 @@ class ContextRegistryTest extends BaseServerTest with ReportLogsOnFailure {
                       Vector(`expressionId`)
                     )
                   ),
-                  None
+                  None,
+                  Seq()
+                )
+              ) =>
+            requestId
+          case msg =>
+            fail(s"Unexpected message: $msg")
+        }
+      runtimeConnectorProbe.lastSender ! Api.Response(
+        requestId3,
+        Api.RecomputeContextResponse(contextId)
+      )
+      client.expectJson(json.ok(3))
+    }
+
+    "recompute with expression configs" in {
+      val client = getInitialisedWsClient()
+
+      // create context
+      client.send(json.executionContextCreateRequest(1))
+      val (requestId, contextId) =
+        runtimeConnectorProbe.receiveN(1).head match {
+          case Api.Request(requestId, Api.CreateContextRequest(contextId)) =>
+            (requestId, contextId)
+          case msg =>
+            fail(s"Unexpected message: $msg")
+        }
+      runtimeConnectorProbe.lastSender ! Api.Response(
+        requestId,
+        Api.CreateContextResponse(contextId)
+      )
+      client.expectJson(json.executionContextCreateResponse(1, contextId))
+
+      // push stack item
+      val expressionId = UUID.randomUUID()
+      client.send(json.executionContextPushRequest(2, contextId, expressionId))
+      val requestId2 =
+        runtimeConnectorProbe.receiveN(1).head match {
+          case Api.Request(
+                requestId,
+                Api.PushContextRequest(
+                  `contextId`,
+                  Api.StackItem.LocalCall(`expressionId`)
+                )
+              ) =>
+            requestId
+          case msg =>
+            fail(s"Unexpected message: $msg")
+        }
+      runtimeConnectorProbe.lastSender ! Api.Response(
+        requestId2,
+        Api.PushContextResponse(contextId)
+      )
+      client.expectJson(json.ok(2))
+
+      // recompute
+      client.send(
+        json"""
+          { "jsonrpc": "2.0",
+            "method": "executionContext/recompute",
+            "id": 3,
+            "params": {
+              "contextId": $contextId,
+              "expressionConfigs": [
+                 { "expressionId":  $expressionId,
+                   "executionEnvironment": "Live"
+                 }
+               ]
+            }
+          }
+          """
+      )
+      val requestId3 =
+        runtimeConnectorProbe.receiveN(1).head match {
+          case Api.Request(
+                requestId,
+                Api.RecomputeContextRequest(
+                  `contextId`,
+                  None,
+                  None,
+                  Seq(
+                    Api.ExpressionConfig(
+                      `expressionId`,
+                      Some(Api.ExecutionEnvironment.Live())
+                    )
+                  )
                 )
               ) =>
             requestId

engine/polyglot-api/src/main/java/org/enso/polyglot/debugger/IdExecutionService.java

@@ -69,6 +69,8 @@ public interface Callbacks {
      *     the execution and return the value as a result.
      */
     Object onFunctionReturn(Info info);
+
+    Object getExecutionEnvironment(Info info);
   }
 
   /**

engine/polyglot-api/src/main/scala/org/enso/polyglot/runtime/Runtime.scala

@@ -691,6 +691,16 @@ object Runtime {
       }
     }
 
+    /** The configuration of how to execute the expression.
+      *
+      * @param expressionId the expression identifier
+      * @param executionEnvironment the execution environment for the expression
+      */
+    sealed case class ExpressionConfig(
+      expressionId: ExpressionId,
+      executionEnvironment: Option[ExecutionEnvironment]
+    )
+
     /** The notification about the execution status.
       *
       * @param contextId the context's id
@@ -914,12 +924,14 @@ object Runtime {
       * @param expressions the selector specifying which expressions should be
       * recomputed.
       * @param executionEnvironment the environment used for execution
+      * @param expressionConfigs execution configurations for selected expressions
       */
     @named("recomputeContextRequest")
     final case class RecomputeContextRequest(
       contextId: ContextId,
       expressions: Option[InvalidatedExpressions],
-      executionEnvironment: Option[ExecutionEnvironment]
+      executionEnvironment: Option[ExecutionEnvironment],
+      expressionConfigs: Seq[ExpressionConfig]
     ) extends ApiRequest
 
     /** A response sent from the server upon handling the

engine/runtime-instrument-common/src/main/java/org/enso/interpreter/instrument/ExecutionConfig.java

@@ -0,0 +1,55 @@
+package org.enso.interpreter.instrument;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.UUID;
+import org.enso.interpreter.runtime.state.ExecutionEnvironment;
+import org.enso.polyglot.runtime.Runtime$Api$ExecutionEnvironment;
+import org.enso.polyglot.runtime.Runtime$Api$ExpressionConfig;
+import scala.Option;
+import scala.collection.immutable.Seq;
+
+/**
+ * The program execution config.
+ *
+ * @param executionEnvironment the global execution environment of the program
+ * @param expressionConfigs execution configs for each expression
+ */
+public record ExecutionConfig(
+    ExecutionEnvironment executionEnvironment, Map<UUID, ExecutionEnvironment> expressionConfigs) {
+
+  public static ExecutionConfig empty() {
+    return new ExecutionConfig(null, Collections.emptyMap());
+  }
+
+  @SuppressWarnings("unchecked")
+  public static ExecutionConfig create(
+      Object executionEnvironmentOption1, Object expressionConfigs1) {
+    Map<UUID, ExecutionEnvironment> expressionConfigsBuilder = new HashMap<>();
+    Option<Runtime$Api$ExecutionEnvironment> executionEnvironmentOption =
+        (Option<Runtime$Api$ExecutionEnvironment>) executionEnvironmentOption1;
+    Seq<Runtime$Api$ExpressionConfig> expressionConfigs =
+        (Seq<Runtime$Api$ExpressionConfig>) expressionConfigs1;
+    expressionConfigs.foreach(
+        expressionConfig -> {
+          expressionConfig
+              .executionEnvironment()
+              .foreach(
+                  executionEnvironment -> {
+                    expressionConfigsBuilder.put(
+                        expressionConfig.expressionId(),
+                        ExecutionEnvironment.forName(executionEnvironment.name()));
+                    return null;
+                  });
+          return null;
+        });
+
+    ExecutionEnvironment executionEnvironment =
+        executionEnvironmentOption
+            .map(env -> ExecutionEnvironment.forName(env.name()))
+            .getOrElse(() -> null);
+
+    return new ExecutionConfig(executionEnvironment, expressionConfigsBuilder);
+  }
+}

engine/runtime-instrument-common/src/main/java/org/enso/interpreter/instrument/ExpressionExecutionState.java

@@ -0,0 +1,31 @@
+package org.enso.interpreter.instrument;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.UUID;
+import org.enso.interpreter.runtime.state.ExecutionEnvironment;
+
+public final class ExpressionExecutionState {
+
+  private final Map<UUID, ExecutionEnvironment> expressionConfigs;
+
+  public ExpressionExecutionState() {
+    this.expressionConfigs = new HashMap<>();
+  }
+
+  public ExpressionExecutionState(Map<UUID, ExecutionEnvironment> expressionConfigs) {
+    this.expressionConfigs = expressionConfigs;
+  }
+
+  public void setExpressionConfigs(Map<UUID, ExecutionEnvironment> expressionConfigs) {
+    this.expressionConfigs.putAll(expressionConfigs);
+  }
+
+  public void setExpressionExecuted(UUID expressionId) {
+    expressionConfigs.remove(expressionId);
+  }
+
+  public ExecutionEnvironment getExecutionEnvironment(UUID expressionId) {
+    return expressionConfigs.get(expressionId);
+  }
+}