Propagate instrumentationScope via State

enso-org · Jun 20, 2023 · 32582bb · 32582bb
1 parent 15e7413
commit 32582bb
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 14 deletions.
diff --git a/...untime-instrument-common/src/main/java/org/enso/interpreter/service/ExecutionService.java b/...untime-instrument-common/src/main/java/org/enso/interpreter/service/ExecutionService.java
@@ -113,7 +113,7 @@ public FunctionCallInstrumentationNode.FunctionCall prepareFunctionCall(
       throw new MethodNotFoundException(module.getName().toString(), type, methodName);
     }
     Object[] arguments = MAIN_METHOD.equals(methodName) ? new Object[] {} : new Object[] {type};
-    return new FunctionCallInstrumentationNode.FunctionCall(
+    return new FunctionCallInstrumentationNode.FunctionCall(module.getScope(),
         function, State.create(EnsoContext.get(null)), arguments);
   }
 

diff --git a/...ent-id-execution/src/main/java/org/enso/interpreter/instrument/IdExecutionInstrument.java b/...ent-id-execution/src/main/java/org/enso/interpreter/instrument/IdExecutionInstrument.java
@@ -34,6 +34,7 @@
 
 import java.util.function.Consumer;
 import org.enso.interpreter.node.ClosureRootNode;
+import org.enso.interpreter.node.EnsoRootNode;
 import org.enso.interpreter.runtime.tag.AvoidIdInstrumentationTag;
 
 /** An instrument for getting values from AST-identified expressions. */
@@ -324,6 +325,7 @@ private void onTailCallReturn(Throwable exception, State state) {
             TailCallException tailCallException = (TailCallException) exception;
             FunctionCallInstrumentationNode.FunctionCall functionCall =
                     new FunctionCallInstrumentationNode.FunctionCall(
+                        this,
                         tailCallException.getFunction(),
                         state,
                         tailCallException.getArguments());

diff --git a/...ime/src/main/java/org/enso/interpreter/node/callable/FunctionCallInstrumentationNode.java b/...ime/src/main/java/org/enso/interpreter/node/callable/FunctionCallInstrumentationNode.java
@@ -11,17 +11,21 @@
 import com.oracle.truffle.api.instrumentation.Tag;
 import com.oracle.truffle.api.interop.InteropLibrary;
 import com.oracle.truffle.api.interop.TruffleObject;
+import com.oracle.truffle.api.library.CachedLibrary;
 import com.oracle.truffle.api.library.ExportLibrary;
 import com.oracle.truffle.api.library.ExportMessage;
 import com.oracle.truffle.api.nodes.Node;
 import com.oracle.truffle.api.nodes.NodeInfo;
+import com.oracle.truffle.api.object.DynamicObjectLibrary;
 import com.oracle.truffle.api.source.SourceSection;
 import org.enso.interpreter.runtime.callable.function.Function;
 import org.enso.interpreter.runtime.tag.IdentifiedTag;
 
 import java.util.Arrays;
 import java.util.UUID;
 import org.enso.interpreter.node.ClosureRootNode;
+import org.enso.interpreter.node.EnsoRootNode;
+import org.enso.interpreter.runtime.scope.ModuleScope;
 import org.enso.interpreter.runtime.state.State;
 import org.enso.interpreter.runtime.tag.AvoidIdInstrumentationTag;
 
@@ -58,23 +62,46 @@ public boolean isInstrumentable() {
   /** A simple value class for function call information. */
   @ExportLibrary(InteropLibrary.class)
   public static final class FunctionCall implements TruffleObject {
+    private final ModuleScope ctx;
     private final Function function;
     private final State state;
     private final @CompilerDirectives.CompilationFinal(dimensions = 1) Object[] arguments;
 
     /**
      * Creates an instance of this class.
      *
+     * @param ctx the scope where the instrumentation happens
      * @param function the function being called.
      * @param state the monadic state to pass to the function.
      * @param arguments the arguments passed to the function.
      */
-    public FunctionCall(Function function, State state, Object[] arguments) {
+    public FunctionCall(ModuleScope ctx, Function function, State state, Object[] arguments) {
+      this.ctx = ctx;
       this.function = function;
       this.state = state;
       this.arguments = arguments;
     }
 
+    /**
+     * Creates an instance of this class.
+     *
+     * @param ctx the node where the instrumentation happens
+     * @param function the function being called.
+     * @param state the monadic state to pass to the function.
+     * @param arguments the arguments passed to the function.
+     */
+    public FunctionCall(Node ctx, Function function, State state, Object[] arguments) {
+      this(findModuleScope(ctx), function, state, arguments);
+    }
+
+    private static ModuleScope findModuleScope(Node ctx) {
+      return switch (ctx.getRootNode()) {
+        case EnsoRootNode root -> root.getModuleScope();
+        case null -> null;
+        default -> null;
+      };
+    }
+
     /**
      * Marks this object as executable for the purposes of the polyglot API.
      *
@@ -93,16 +120,24 @@ boolean isExecutable() {
     static class Execute {
       @Specialization
       static Object callCached(
-          FunctionCall functionCall,
-          Object[] arguments,
-          @Cached InteropApplicationNode interopApplicationNode) {
-        Object[] callArguments =
-            Arrays.copyOf(
-                functionCall.getArguments(), functionCall.getArguments().length + arguments.length);
-        System.arraycopy(
-            arguments, 0, callArguments, functionCall.getArguments().length, arguments.length);
-        return interopApplicationNode.execute(
-            functionCall.function, functionCall.state, callArguments);
+        FunctionCall functionCall,
+        Object[] arguments,
+        @Cached InteropApplicationNode interopApplicationNode,
+        @CachedLibrary(limit = "10") DynamicObjectLibrary objects
+      ) {
+        var args = Arrays.copyOf(
+          functionCall.getArguments(), functionCall.getArguments().length + arguments.length
+        );
+        System.arraycopy(arguments, 0, args, functionCall.getArguments().length, arguments.length);
+
+        var state = functionCall.state;
+        var old = objects.getOrDefault(state.getContainer(), 43L, null);
+        try {
+          objects.put(state.getContainer(), 43L, functionCall.ctx);
+          return interopApplicationNode.execute(functionCall.function, state, args);
+        } finally {
+          objects.put(state.getContainer(), 43L, old);
+        }
       }
     }
 
@@ -132,7 +167,7 @@ public Object[] getArguments() {
    * @return an instance of {@link FunctionCall} containing the function, state and arguments.
    */
   public Object execute(VirtualFrame frame, Function function, State state, Object[] arguments) {
-    return new FunctionCall(function, state, arguments);
+    return new FunctionCall(this, function, state, arguments);
   }
 
   /**

diff --git a/...me/src/main/java/org/enso/interpreter/node/expression/builtin/meta/GetAnnotationNode.java b/...me/src/main/java/org/enso/interpreter/node/expression/builtin/meta/GetAnnotationNode.java
@@ -18,6 +18,7 @@
 import com.oracle.truffle.api.dsl.Specialization;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.library.CachedLibrary;
+import com.oracle.truffle.api.object.DynamicObjectLibrary;
 
 @BuiltinMethod(
     type = "Meta",
@@ -38,7 +39,9 @@ Object doExecute(
       Object parameter,
       @Cached ThunkExecutorNode thunkExecutorNode,
       @Cached ExpectStringNode expectStringNode,
-      @Cached TypeOfNode typeOfNode) {
+      @Cached TypeOfNode typeOfNode,
+      @CachedLibrary(limit = "10") DynamicObjectLibrary objects
+    ) {
     String methodName = expectStringNode.execute(method);
 
     Object targetTypeResult = typeOfNode.execute(target);
@@ -49,6 +52,11 @@ Object doExecute(
     if (targetTypeResult instanceof Type targetType) {
       ModuleScope scope = targetType.getDefinitionScope();
       Function methodFunction = scope.lookupMethodDefinition(targetType, methodName);
+      if (methodFunction == null) {
+        if (objects.getOrDefault(state.getContainer(), 43L, null) instanceof ModuleScope instrumentationScope) {
+          methodFunction = instrumentationScope.lookupMethodDefinition(targetType, methodName);
+        }
+      }
       if (methodFunction != null) {
         String parameterName = expectStringNode.execute(parameter);
         Annotation annotation = methodFunction.getSchema().getAnnotation(parameterName);