Node added inside collapsed node is never executed (#11341)

close #11251

Changelog:
- add: `CachePreferences` configuration of the expressions that are marked for caching
- update: Invalidate the `self` keywords in the parent frames on function edit
- update: persistance config

# Important Notes
The demo shows that:

1. The new node created in the collapsed function is highlighted (GUI receives the expression update)
2. When the collapsed function is edited, the nodes in the `main` method are updated correctly

https://github.com/user-attachments/assets/41c406ed-ba76-41c8-9e3f-89ac9ff63c3f

engine/common/src/main/java/org/enso/common/CachePreferences.java

@@ -0,0 +1,98 @@
+package org.enso.common;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+
+/**
+ * Contains settings telling which nodes are marked for caching and their kinds.
+ *
+ * @param preferences the config with cached preferences
+ */
+public record CachePreferences(Map<UUID, Kind> preferences) {
+
+  /**
+   * @return empty cache preferences config
+   */
+  public static CachePreferences empty() {
+    return new CachePreferences(new HashMap<>());
+  }
+
+  /** A kind of cached value. */
+  public enum Kind {
+    BINDING_EXPRESSION,
+    SELF_ARGUMENT
+  }
+
+  /**
+   * Get the cache preference for the provided node id.
+   *
+   * @param id the node id
+   * @return the kind of cached value if available
+   */
+  public Kind get(UUID id) {
+    return preferences.get(id);
+  }
+
+  /**
+   * Get all the node ids marked for caching for a particular node kind.
+   *
+   * @param kind the node kind
+   * @return the set of node ids marked for caching
+   */
+  public Set<UUID> get(Kind kind) {
+    var result = new HashSet<UUID>();
+    for (var entry : preferences.entrySet()) {
+      if (entry.getValue().equals(kind)) {
+        result.add(entry.getKey());
+      }
+    }
+
+    return result;
+  }
+
+  /**
+   * Add a cache preference.
+   *
+   * @param id the node id
+   * @param kind the node kind
+   */
+  public void set(UUID id, Kind kind) {
+    preferences.put(id, kind);
+  }
+
+  /**
+   * Remove the cache preference.
+   *
+   * @param id the node id to remove
+   */
+  public void remove(UUID id) {
+    preferences.remove(id);
+  }
+
+  /**
+   * Check if the provided node id is marked for caching.
+   *
+   * @param id the node id
+   * @return return {@code true} if the provided node id is marked for caching
+   */
+  public boolean contains(UUID id) {
+    return preferences.containsKey(id);
+  }
+
+  /** Clear all the cache preferences. */
+  public void clear() {
+    preferences.clear();
+  }
+
+  /**
+   * Make a copy of this cache preferences.
+   *
+   * @return a copy of this cache preferences
+   */
+  public CachePreferences copy() {
+    return new CachePreferences(new HashMap<>(preferences));
+  }
+}

engine/runtime-compiler/src/main/java/org/enso/compiler/pass/analyse/PassPersistance.java

@@ -1,6 +1,7 @@
 package org.enso.compiler.pass.analyse;
 
 import java.io.IOException;
+import org.enso.common.CachePreferences;
 import org.enso.compiler.pass.analyse.alias.AliasMetadata;
 import org.enso.compiler.pass.analyse.alias.graph.Graph;
 import org.enso.compiler.pass.analyse.alias.graph.GraphOccurrence;
@@ -67,6 +68,7 @@ import scala.Tuple2$;
 @Persistable(clazz = TypeInference.class, id = 1280)
 @Persistable(clazz = FramePointerAnalysis$.class, id = 1281)
 @Persistable(clazz = TailCall.TailPosition.class, id = 1282)
+@Persistable(clazz = CachePreferences.class, id = 1284)
 public final class PassPersistance {
   private PassPersistance() {}
 

engine/runtime-compiler/src/main/scala/org/enso/compiler/pass/analyse/CachePreferenceAnalysis.scala

@@ -1,5 +1,6 @@
 package org.enso.compiler.pass.analyse
 
+import org.enso.common.CachePreferences
 import org.enso.compiler.context.{InlineContext, ModuleContext}
 import org.enso.compiler.core.Implicits.AsMetadata
 import org.enso.compiler.core.ir.CallArgument.Specified
@@ -20,19 +21,12 @@ import org.enso.compiler.pass.IRPass
 import org.enso.compiler.pass.IRProcessingPass
 import org.enso.compiler.pass.desugar._
 
-import java.util
 import java.util.UUID
-import scala.collection.mutable
-import scala.jdk.CollectionConverters._
 
 /** This pass implements the preference analysis for caching.
   *
-  * The pass assigns weights to the expressions. The greater the weight, the
-  * more preferable the expression for caching.
-  *
-  * Weights:
-  *
-  *   - `1` - Right hand side expressions
+  * The pass assigns preferences to the expressions. To mark the expression for
+  * caching, it should be added to the [[CachePreferences]] configuration.
   */
 case object CachePreferenceAnalysis extends IRPass {
 
@@ -143,9 +137,8 @@ case object CachePreferenceAnalysis extends IRPass {
   ): Expression = {
     expression.transformExpressions {
       case binding: Expression.Binding =>
-        binding.getExternalId.foreach(weights.update(_, Weight.Never))
         binding.expression.getExternalId
-          .foreach(weights.update(_, Weight.Always))
+          .foreach(weights.update(_, CachePreferences.Kind.BINDING_EXPRESSION))
         binding
           .copy(
             name       = binding.name.updateMetadata(new MetadataPair(this, weights)),
@@ -163,10 +156,6 @@ case object CachePreferenceAnalysis extends IRPass {
             app
         }
       case expr =>
-        expr.getExternalId.foreach {
-          case id if !weights.contains(id) => weights.update(id, Weight.Never)
-          case _                           =>
-        }
         expr
           .mapExpressions(analyseExpression(_, weights))
           .updateMetadata(new MetadataPair(this, weights))
@@ -177,13 +166,13 @@ case object CachePreferenceAnalysis extends IRPass {
     callArgument: CallArgument,
     weights: WeightInfo
   ): CallArgument = {
-    callArgument.value.getExternalId.foreach(weights.update(_, Weight.Always))
+    callArgument.value.getExternalId
+      .foreach(weights.update(_, CachePreferences.Kind.SELF_ARGUMENT))
     callArgument match {
       case arg: Specified =>
         arg.copy(value =
-          analyseExpression(arg.value, weights).updateMetadata(
-            new MetadataPair(this, weights)
-          )
+          analyseExpression(arg.value, weights)
+            .updateMetadata(new MetadataPair(this, weights))
         )
     }
   }
@@ -213,7 +202,7 @@ case object CachePreferenceAnalysis extends IRPass {
     * @param weights the storage for weights of the program components
     */
   sealed case class WeightInfo(
-    weights: mutable.HashMap[UUID @ExternalID, Double] = mutable.HashMap()
+    preferences: CachePreferences = CachePreferences.empty()
   ) extends IRPass.IRMetadata {
 
     /** The name of the metadata as a string. */
@@ -230,35 +219,13 @@ case object CachePreferenceAnalysis extends IRPass {
     /** Assign the weight to an id.
       *
       * @param id the external id
-      * @param weight the assigned weight
+      * @param kind the assigned expression kind
       */
-    def update(id: UUID @ExternalID, weight: Double): Unit =
-      weights.put(id, weight)
+    def update(id: UUID @ExternalID, kind: CachePreferences.Kind): Unit =
+      preferences.set(id, kind)
 
-    /** Get the weight associated with given id */
-    def get(id: UUID @ExternalID): Double =
-      weights.getOrElse(id, Weight.Never)
-
-    /** Check if the weight is assigned to this id. */
-    def contains(id: UUID @ExternalID): Boolean =
-      weights.contains(id)
-
-    /** @return weights as the Java collection */
-    def asJavaWeights: util.Map[UUID @ExternalID, java.lang.Double] =
-      weights.asJava
-        .asInstanceOf[util.Map[UUID @ExternalID, java.lang.Double]]
-
-    override def duplicate(): Option[IRPass.IRMetadata] =
-      Some(copy(weights = this.weights))
+    override def duplicate(): Option[IRPass.IRMetadata] = {
+      Some(copy(preferences = this.preferences.copy()))
     }
-
-  /** Weight constants */
-  object Weight {
-
-    /** Maximum weight meaning that the program component is always cached. */
-    val Always: Double = 1.0
-
-    /** Minimum weight meaning that the program component is never cached. */
-    val Never: Double = 0.0
   }
 }

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

@@ -10,6 +10,7 @@ import java.util.Set;
 import java.util.UUID;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
+import org.enso.common.CachePreferences;
 import org.enso.interpreter.service.ExecutionService;
 
 /** A storage for computed values. */
@@ -18,7 +19,7 @@ public final class RuntimeCache implements java.util.function.Function<String, O
   private final Map<UUID, Reference<Object>> expressions = new HashMap<>();
   private final Map<UUID, String> types = new HashMap<>();
   private final Map<UUID, ExecutionService.FunctionCallInfo> calls = new HashMap<>();
-  private Map<UUID, Double> weights = new HashMap<>();
+  private CachePreferences preferences = CachePreferences.empty();
   private Consumer<UUID> observer;
 
   /**
@@ -30,8 +31,7 @@ public final class RuntimeCache implements java.util.function.Function<String, O
    */
   @CompilerDirectives.TruffleBoundary
   public boolean offer(UUID key, Object value) {
-    var weight = weights.get(key);
-    if (weight != null && weight > 0) {
+    if (preferences.contains(key)) {
       var ref = new SoftReference<>(value);
       cache.put(key, ref);
       expressions.put(key, new WeakReference<>(value));
@@ -87,6 +87,20 @@ public final class RuntimeCache implements java.util.function.Function<String, O
     cache.clear();
   }
 
+  /**
+   * Clear cached values of the provided kind.
+   *
+   * @param kind the kind of cached value to clear
+   * @return the set of cleared keys
+   */
+  public Set<UUID> clear(CachePreferences.Kind kind) {
+    var keys = preferences.get(kind);
+    for (var key : keys) {
+      cache.remove(key);
+    }
+    return keys;
+  }
+
   /**
    * Cache the type of expression.
    *
@@ -161,25 +175,33 @@ public final class RuntimeCache implements java.util.function.Function<String, O
   }
 
   /**
-   * @return the weights of this cache.
+   * @return the preferences of this cache.
    */
-  public Map<UUID, Double> getWeights() {
-    return weights;
+  public CachePreferences getPreferences() {
+    return preferences;
   }
 
-  /** Set the new weights. */
-  public void setWeights(Map<UUID, Double> weights) {
-    this.weights = weights;
+  /**
+   * Set the new cache preferences.
+   *
+   * @param preferences the new cache preferences
+   */
+  public void setPreferences(CachePreferences preferences) {
+    this.preferences = preferences;
   }
 
-  /** Remove the weight associated with the provided key. */
-  public void removeWeight(UUID key) {
-    weights.remove(key);
+  /**
+   * Remove the cache preference associated with the provided key.
+   *
+   * @param key the preference to remove
+   */
+  public void removePreference(UUID key) {
+    preferences.remove(key);
   }
 
-  /** Clear the weights. */
-  public void clearWeights() {
-    weights.clear();
+  /** Clear the cache preferences. */
+  public void clearPreferences() {
+    preferences.clear();
   }
 
   /**

engine/runtime-instrument-common/src/main/scala/org/enso/interpreter/instrument/CacheInvalidation.scala

@@ -1,7 +1,8 @@
 package org.enso.interpreter.instrument
 
-import java.util.UUID
+import org.enso.common.CachePreferences
 
+import java.util.UUID
 import org.enso.compiler.pass.analyse.CachePreferenceAnalysis
 import org.enso.polyglot.runtime.Runtime.Api
 
@@ -57,6 +58,13 @@ object CacheInvalidation {
       */
     case class InvalidateKeys(keys: Iterable[UUID]) extends Command
 
+    /** A command to invalidate cache entries by preference kinds.
+      *
+      * @param kinds the kinds of entries that should be invalidated
+      */
+    case class InvalidateByKind(kinds: Iterable[CachePreferences.Kind])
+        extends Command
+
     /** A command to invalidate stale entries from the cache.
       *
       * @param scope all ids of the source
@@ -93,6 +101,9 @@ object CacheInvalidation {
 
     /** Select top stack element. */
     case object Top extends StackSelector
+
+    /** Select all except the top stack element. */
+    case object Tail extends StackSelector
   }
 
   /** Create an invalidation instruction using a stack selector and an
@@ -154,7 +165,7 @@ object CacheInvalidation {
     indexes: Set[IndexSelector] = Set()
   ): Unit =
     visualizations.foreach { visualization =>
-      run(visualization.cache, command, indexes)
+      run(visualization.cache, None, command, indexes)
     }
 
   /** Run a cache invalidation instruction on an execution stack.
@@ -169,6 +180,7 @@ object CacheInvalidation {
     val frames = instruction.elements match {
       case StackSelector.All  => stack
       case StackSelector.Top  => stack.headOption.toSeq
+      case StackSelector.Tail => stack.tail
     }
     run(frames, instruction.command, instruction.indexes)
   }
@@ -184,37 +196,9 @@ object CacheInvalidation {
     command: Command,
     indexes: Set[IndexSelector]
   ): Unit = {
-    frames.foreach(frame => run(frame.cache, frame.syncState, command, indexes))
-  }
-
-  /** Run cache invalidation of a single instrument frame.
-    *
-    * @param cache the cache to invalidate
-    * @param command the invalidation instruction
-    * @param indexes the list of indexes to invalidate
-    */
-  private def run(
-    cache: RuntimeCache,
-    command: Command,
-    indexes: Set[IndexSelector]
-  ): Unit =
-    command match {
-      case Command.InvalidateAll =>
-        cache.clear()
-        indexes.foreach(clearIndex(_, cache))
-      case Command.InvalidateKeys(keys) =>
-        keys.foreach { key =>
-          cache.remove(key)
-          indexes.foreach(clearIndexKey(key, _, cache))
-        }
-      case Command.InvalidateStale(scope) =>
-        val staleKeys = cache.getKeys.asScala.diff(scope.toSet)
-        staleKeys.foreach { key =>
-          cache.remove(key)
-          indexes.foreach(clearIndexKey(key, _, cache))
-        }
-      case Command.SetMetadata(metadata) =>
-        cache.setWeights(metadata.asJavaWeights)
+    frames.foreach(frame =>
+      run(frame.cache, Some(frame.syncState), command, indexes)
+    )
   }
 
   /** Run cache invalidation of a single instrument frame.
@@ -226,7 +210,7 @@ object CacheInvalidation {
     */
   private def run(
     cache: RuntimeCache,
-    syncState: UpdatesSynchronizationState,
+    syncState: Option[UpdatesSynchronizationState],
     command: Command,
     indexes: Set[IndexSelector]
   ): Unit =
@@ -239,15 +223,22 @@ object CacheInvalidation {
           cache.remove(key)
           indexes.foreach(clearIndexKey(key, _, cache))
         }
+      case Command.InvalidateByKind(kinds) =>
+        kinds.foreach { kind =>
+          val keys = cache.clear(kind)
+          keys.forEach { key =>
+            indexes.foreach(clearIndexKey(key, _, cache))
+          }
+        }
       case Command.InvalidateStale(scope) =>
         val staleKeys = cache.getKeys.asScala.diff(scope.toSet)
         staleKeys.foreach { key =>
           cache.remove(key)
           indexes.foreach(clearIndexKey(key, _, cache))
-          syncState.invalidate(key)
+          syncState.foreach(_.invalidate(key))
         }
       case Command.SetMetadata(metadata) =>
-        cache.setWeights(metadata.asJavaWeights)
+        cache.setPreferences(metadata.preferences)
     }
 
   /** Clear the selected index.
@@ -259,10 +250,10 @@ object CacheInvalidation {
     selector match {
       case IndexSelector.All =>
         cache.clearTypes()
-        cache.clearWeights()
+        cache.clearPreferences()
         cache.clearCalls()
       case IndexSelector.Weights =>
-        cache.clearWeights()
+        cache.clearPreferences()
       case IndexSelector.Types =>
         cache.clearTypes()
       case IndexSelector.Calls =>
@@ -283,10 +274,10 @@ object CacheInvalidation {
     selector match {
       case IndexSelector.All =>
         cache.removeType(key)
-        cache.removeWeight(key)
+        cache.removePreference(key)
         cache.removeCall(key)
       case IndexSelector.Weights =>
-        cache.removeWeight(key)
+        cache.removePreference(key)
       case IndexSelector.Types =>
         cache.removeType(key)
       case IndexSelector.Calls =>

engine/runtime-instrument-common/src/main/scala/org/enso/interpreter/instrument/command/RecomputeContextCmd.scala

@@ -186,7 +186,9 @@ object RecomputeContextCmd {
       case CacheInvalidation.Command.InvalidateAll =>
         stack.headOption
           .map { frame =>
-            frame.cache.getWeights.keySet().forEach(builder.addOne)
+            frame.cache.getPreferences.preferences
+              .keySet()
+              .forEach(builder.addOne)
           }
       case CacheInvalidation.Command.InvalidateKeys(expressionIds) =>
         builder ++= expressionIds

engine/runtime-instrument-common/src/main/scala/org/enso/interpreter/instrument/job/EnsureCompiledJob.scala

@@ -1,7 +1,7 @@
 package org.enso.interpreter.instrument.job
 
 import com.oracle.truffle.api.TruffleLogger
-import org.enso.common.CompilationStage
+import org.enso.common.{CachePreferences, CompilationStage}
 import org.enso.compiler.{data, CompilerResult}
 import org.enso.compiler.context._
 import org.enso.compiler.core.Implicits.AsMetadata
@@ -403,6 +403,12 @@ class EnsureCompiledJob(
         CacheInvalidation.StackSelector.All,
         invalidateStaleCommand,
         Set(CacheInvalidation.IndexSelector.All)
+      ),
+      CacheInvalidation(
+        CacheInvalidation.StackSelector.Tail,
+        CacheInvalidation.Command.InvalidateByKind(
+          Seq(CachePreferences.Kind.SELF_ARGUMENT)
+        )
       )
     )
   }

engine/runtime-instrument-common/src/test/java/org/enso/interpreter/instrument/RuntimeCacheTest.java

@@ -9,11 +9,12 @@ import static org.junit.Assert.assertTrue;
 import java.lang.ref.Reference;
 import java.lang.ref.WeakReference;
 import java.util.HashSet;
-import java.util.Map;
 import java.util.UUID;
+import org.enso.common.CachePreferences;
 import org.junit.Test;
 
 public class RuntimeCacheTest {
+
   @Test
   public void cacheItems() {
     var cache = new RuntimeCache();
@@ -23,7 +24,7 @@ public class RuntimeCacheTest {
     assertFalse(cache.offer(key, obj));
     assertNull(cache.get(key));
 
-    cache.setWeights(Map.of(key, 1.0));
+    cache.setPreferences(of(key, CachePreferences.Kind.BINDING_EXPRESSION));
     assertTrue(cache.offer(key, obj));
     assertEquals(obj, cache.get(key));
   }
@@ -34,7 +35,7 @@ public class RuntimeCacheTest {
     var key = UUID.randomUUID();
     var obj = new Object();
 
-    cache.setWeights(Map.of(key, 1.0));
+    cache.setPreferences(of(key, CachePreferences.Kind.BINDING_EXPRESSION));
     assertTrue(cache.offer(key, obj));
     assertEquals(obj, cache.remove(key));
     assertNull(cache.get(key));
@@ -60,7 +61,7 @@ public class RuntimeCacheTest {
     var exprKey = UUID.randomUUID();
     var obj = new Object();
 
-    cache.setWeights(Map.of(key, 1.0));
+    cache.setPreferences(of(key, CachePreferences.Kind.BINDING_EXPRESSION));
 
     assertFalse("Not inserted, as the value isn't in the map yet", cache.offer(exprKey, obj));
     assertNull("No UUID for exprKey in cache", cache.get(exprKey));
@@ -82,7 +83,7 @@ public class RuntimeCacheTest {
     var exprKey = UUID.randomUUID();
     var obj = new Object();
 
-    cache.setWeights(Map.of(key, 1.0));
+    cache.setPreferences(of(key, CachePreferences.Kind.BINDING_EXPRESSION));
 
     assertFalse("Not inserted, as the value isn't in the map yet", cache.offer(exprKey, obj));
     assertNull("No UUID for exprKey in cache", cache.get(exprKey));
@@ -112,7 +113,7 @@ public class RuntimeCacheTest {
     var exprKey = UUID.randomUUID();
     var obj = new Object();
 
-    cache.setWeights(Map.of(key, 1.0));
+    cache.setPreferences(of(key, CachePreferences.Kind.BINDING_EXPRESSION));
 
     assertFalse("Not inserted, as the value isn't in the map yet", cache.offer(exprKey, obj));
     assertNull("No UUID for exprKey in cache", cache.get(exprKey));
@@ -163,4 +164,10 @@ public class RuntimeCacheTest {
       assertNotNull(msg + " ref has been cleaned", obj);
     }
   }
+
+  private static CachePreferences of(UUID key, CachePreferences.Kind value) {
+    var preferences = CachePreferences.empty();
+    preferences.set(key, value);
+    return preferences;
+  }
 }

engine/runtime-integration-tests/src/test/scala/org/enso/interpreter/test/instrument/RuntimeServerTest.scala

@@ -7206,6 +7206,272 @@ class RuntimeServerTest
     context.consumeOut shouldEqual List("Hello World!")
   }
 
+  it should "edit local call with cached self" in {
+    val contextId  = UUID.randomUUID()
+    val requestId  = UUID.randomUUID()
+    val moduleName = "Enso_Test.Test.Main"
+
+    val metadata   = new Metadata
+    val idX        = metadata.addItem(46, 11, "aa")
+    val idSelfMain = metadata.addItem(46, 4, "ab")
+    val idIncZ     = new UUID(0, 1)
+
+    val code =
+      """from Standard.Base import all
+        |
+        |main =
+        |    x = Main.inc 41
+        |    IO.println x
+        |
+        |inc a =
+        |    y = 1
+        |    r = a + y
+        |    r
+        |""".stripMargin.linesIterator.mkString("\n")
+    val contents = metadata.appendToCode(code)
+    val mainFile = context.writeMain(contents)
+
+    // create context
+    context.send(Api.Request(requestId, Api.CreateContextRequest(contextId)))
+    context.receive shouldEqual Some(
+      Api.Response(requestId, Api.CreateContextResponse(contextId))
+    )
+
+    // open file
+    context.send(
+      Api.Request(requestId, Api.OpenFileRequest(mainFile, contents))
+    )
+    context.receive shouldEqual Some(
+      Api.Response(Some(requestId), Api.OpenFileResponse)
+    )
+
+    // push main
+    context.send(
+      Api.Request(
+        requestId,
+        Api.PushContextRequest(
+          contextId,
+          Api.StackItem.ExplicitCall(
+            Api.MethodPointer(moduleName, moduleName, "main"),
+            None,
+            Vector()
+          )
+        )
+      )
+    )
+    context.receiveNIgnoreStdLib(4) should contain theSameElementsAs Seq(
+      Api.Response(requestId, Api.PushContextResponse(contextId)),
+      TestMessages.update(
+        contextId,
+        idX,
+        ConstantsGen.INTEGER,
+        methodCall =
+          Some(Api.MethodCall(Api.MethodPointer(moduleName, moduleName, "inc")))
+      ),
+      TestMessages.update(contextId, idSelfMain, moduleName),
+      context.executionComplete(contextId)
+    )
+    context.consumeOut shouldEqual List("42")
+
+    // push inc call
+    context.send(
+      Api.Request(
+        requestId,
+        Api.PushContextRequest(contextId, Api.StackItem.LocalCall(idX))
+      )
+    )
+    context.receiveNIgnoreStdLib(2) should contain theSameElementsAs Seq(
+      Api.Response(requestId, Api.PushContextResponse(contextId)),
+      context.executionComplete(contextId)
+    )
+    context.consumeOut shouldEqual List("42")
+
+    // Modify the file
+    context.send(
+      Api.Request(
+        Api.EditFileNotification(
+          mainFile,
+          Seq(
+            model.TextEdit(
+              model.Range(model.Position(8, 4), model.Position(8, 9)),
+              "z = 2\n    r = a + z"
+            )
+          ),
+          execute = true,
+          idMap   = Some(model.IdMap(Vector(model.Span(102, 103) -> idIncZ)))
+        )
+      )
+    )
+    context.receiveNIgnorePendingExpressionUpdates(2, 10) shouldEqual Seq(
+      TestMessages.update(
+        contextId,
+        idIncZ,
+        ConstantsGen.INTEGER
+      ),
+      context.executionComplete(contextId)
+    )
+    context.consumeOut shouldEqual List("44")
+  }
+
+  it should "edit two local calls with cached self" in {
+    val contextId  = UUID.randomUUID()
+    val requestId  = UUID.randomUUID()
+    val moduleName = "Enso_Test.Test.Main"
+
+    val metadata    = new Metadata
+    val idX         = metadata.addItem(46, 10, "aa")
+    val idY         = metadata.addItem(65, 10, "ab")
+    val idXSelfMain = metadata.addItem(46, 4, "ac")
+    val idYSelfMain = metadata.addItem(65, 4, "ad")
+    val idIncZ      = new UUID(0, 1)
+
+    val code =
+      """from Standard.Base import all
+        |
+        |main =
+        |    x = Main.inc 3
+        |    y = Main.inc 7
+        |    IO.println x+y
+        |
+        |inc a =
+        |    y = 1
+        |    r = a + y
+        |    r
+        |""".stripMargin.linesIterator.mkString("\n")
+    val contents = metadata.appendToCode(code)
+    val mainFile = context.writeMain(contents)
+
+    // create context
+    context.send(Api.Request(requestId, Api.CreateContextRequest(contextId)))
+    context.receive shouldEqual Some(
+      Api.Response(requestId, Api.CreateContextResponse(contextId))
+    )
+
+    // open file
+    context.send(
+      Api.Request(requestId, Api.OpenFileRequest(mainFile, contents))
+    )
+    context.receive shouldEqual Some(
+      Api.Response(Some(requestId), Api.OpenFileResponse)
+    )
+
+    // push main
+    context.send(
+      Api.Request(
+        requestId,
+        Api.PushContextRequest(
+          contextId,
+          Api.StackItem.ExplicitCall(
+            Api.MethodPointer(moduleName, moduleName, "main"),
+            None,
+            Vector()
+          )
+        )
+      )
+    )
+    context.receiveNIgnoreStdLib(6) should contain theSameElementsAs Seq(
+      Api.Response(requestId, Api.PushContextResponse(contextId)),
+      TestMessages.update(
+        contextId,
+        idX,
+        ConstantsGen.INTEGER,
+        methodCall =
+          Some(Api.MethodCall(Api.MethodPointer(moduleName, moduleName, "inc")))
+      ),
+      TestMessages.update(
+        contextId,
+        idY,
+        ConstantsGen.INTEGER,
+        methodCall =
+          Some(Api.MethodCall(Api.MethodPointer(moduleName, moduleName, "inc")))
+      ),
+      TestMessages.update(contextId, idXSelfMain, moduleName),
+      TestMessages.update(contextId, idYSelfMain, moduleName),
+      context.executionComplete(contextId)
+    )
+    context.consumeOut shouldEqual List("12")
+
+    // push inc call
+    context.send(
+      Api.Request(
+        requestId,
+        Api.PushContextRequest(contextId, Api.StackItem.LocalCall(idX))
+      )
+    )
+    context.receiveNIgnoreStdLib(2) should contain theSameElementsAs Seq(
+      Api.Response(requestId, Api.PushContextResponse(contextId)),
+      context.executionComplete(contextId)
+    )
+    context.consumeOut shouldEqual List("12")
+
+    // Modify the file
+    context.send(
+      Api.Request(
+        Api.EditFileNotification(
+          mainFile,
+          Seq(
+            model.TextEdit(
+              model.Range(model.Position(9, 4), model.Position(9, 9)),
+              "z = 2\n    r = a + z"
+            )
+          ),
+          execute = true,
+          idMap   = Some(model.IdMap(Vector(model.Span(122, 123) -> idIncZ)))
+        )
+      )
+    )
+    context.receiveNIgnorePendingExpressionUpdates(2, 10) shouldEqual Seq(
+      TestMessages.update(
+        contextId,
+        idIncZ,
+        ConstantsGen.INTEGER
+      ),
+      context.executionComplete(contextId)
+    )
+    context.consumeOut shouldEqual List("16")
+
+    // pop the inc call
+    context.send(Api.Request(requestId, Api.PopContextRequest(contextId)))
+
+    context.receiveNIgnoreStdLib(6) should contain theSameElementsAs Seq(
+      Api.Response(requestId, Api.PopContextResponse(contextId)),
+      TestMessages.update(
+        contextId,
+        idX,
+        ConstantsGen.INTEGER,
+        fromCache   = false,
+        typeChanged = false,
+        methodCall =
+          Some(Api.MethodCall(Api.MethodPointer(moduleName, moduleName, "inc")))
+      ),
+      TestMessages.update(
+        contextId,
+        idY,
+        ConstantsGen.INTEGER,
+        fromCache   = false,
+        typeChanged = false,
+        methodCall =
+          Some(Api.MethodCall(Api.MethodPointer(moduleName, moduleName, "inc")))
+      ),
+      TestMessages.update(
+        contextId,
+        idXSelfMain,
+        moduleName,
+        fromCache   = true,
+        typeChanged = false
+      ),
+      TestMessages.update(
+        contextId,
+        idYSelfMain,
+        moduleName,
+        fromCache   = true,
+        typeChanged = false
+      ),
+      context.executionComplete(contextId)
+    )
+    context.consumeOut shouldEqual List("16")
+  }
+
 }
 object RuntimeServerTest {