Proof-of-concept of method exit value sniffing

plugins/stream-debugger/src/com/intellij/debugger/streams/action/TraceStreamAction.java

@@ -8,6 +8,9 @@
 import com.intellij.debugger.streams.psi.DebuggerPositionResolver;
 import com.intellij.debugger.streams.psi.impl.DebuggerPositionResolverImpl;
 import com.intellij.debugger.streams.trace.*;
+import com.intellij.debugger.streams.trace.breakpoint.BreakpointConfigurator;
+import com.intellij.debugger.streams.trace.breakpoint.BulkBreakpointConfigurator;
+import com.intellij.debugger.streams.trace.breakpoint.MethodBreakpointTracer;
 import com.intellij.debugger.streams.trace.impl.TraceResultInterpreterImpl;
 import com.intellij.debugger.streams.ui.ChooserOption;
 import com.intellij.debugger.streams.ui.impl.ElementChooserImpl;
@@ -110,8 +113,9 @@ private static void runTrace(@NotNull StreamChain chain, @NotNull LibrarySupport
     ApplicationManager.getApplication().invokeLater(window::show);
     final Project project = session.getProject();
     final TraceExpressionBuilder expressionBuilder = provider.getExpressionBuilder(project);
-    final TraceResultInterpreterImpl resultInterpreter = new TraceResultInterpreterImpl(provider.getLibrarySupport().getInterpreterFactory());
-    final StreamTracer tracer = new EvaluateExpressionTracer(session, expressionBuilder, resultInterpreter);
+    final TraceResultInterpreter resultInterpreter = new TraceResultInterpreterImpl(provider.getLibrarySupport().getInterpreterFactory());
+    final BreakpointConfigurator breakpointConfigurator = new BulkBreakpointConfigurator();
+    final StreamTracer tracer = new MethodBreakpointTracer(session, breakpointConfigurator, resultInterpreter);
     tracer.trace(chain, new TracingCallback() {
       @Override
       public void evaluated(@NotNull TracingResult result, @NotNull EvaluationContextImpl context) {

plugins/stream-debugger/src/com/intellij/debugger/streams/trace/breakpoint/BreakpointConfigurator.kt

@@ -0,0 +1,24 @@
+// Copyright 2000-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.intellij.debugger.streams.trace.breakpoint
+
+import com.intellij.debugger.engine.JavaDebugProcess
+import com.intellij.debugger.streams.wrapper.StreamChain
+
+/**
+ * @author Shumaf Lovpache
+ */
+interface BreakpointConfigurator {
+  /**
+   * Хотим предусмотреть возможность подмены алгоритма расстановки брейкпоинтов
+   * На текущий момент идей две:
+   * - ставим последовательно после каждого метода
+   * - ставим все в начале исполнения цепочки, снимаем в конце выполнения цепочки
+   *
+   * TODO: продумать как пробрасывать свой модификатор значений
+   */
+
+  /**
+   * Вызываем перед заходом в цепочку, чтобы настроить брейкпоинты
+   */
+  fun setBreakpoints(process: JavaDebugProcess, chain: StreamChain, chainEvaluatedCallback: ChainEvaluatedCallback)
+}

plugins/stream-debugger/src/com/intellij/debugger/streams/trace/breakpoint/BulkBreakpointConfigurator.kt

@@ -0,0 +1,106 @@
+// Copyright 2000-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.intellij.debugger.streams.trace.breakpoint
+
+import com.intellij.debugger.engine.JavaDebugProcess
+import com.intellij.debugger.engine.SuspendContextImpl
+import com.intellij.debugger.streams.trace.breakpoint.DebuggerUtils.runInDebuggerThread
+import com.intellij.debugger.streams.trace.breakpoint.value.transform.MethodReturnValueTransformer
+import com.intellij.debugger.streams.trace.breakpoint.value.transform.PrintToStdoutMethodReturnValueTransformer
+import com.intellij.debugger.streams.wrapper.StreamCall
+import com.intellij.debugger.streams.wrapper.StreamChain
+import com.intellij.debugger.ui.breakpoints.FilteredRequestor
+import com.intellij.openapi.diagnostic.logger
+import com.intellij.psi.PsiManager
+import com.intellij.psi.PsiMethod
+import com.intellij.psi.PsiMethodCallExpression
+import com.intellij.xdebugger.XDebugSession
+import com.sun.jdi.Value
+import com.sun.jdi.event.MethodExitEvent
+
+private val LOG = logger<BulkBreakpointConfigurator>()
+
+/**
+ * @author Shumaf Lovpache
+ */
+class BulkBreakpointConfigurator : BreakpointConfigurator {
+  override fun setBreakpoints(process: JavaDebugProcess, chain: StreamChain, chainEvaluatedCallback: ChainEvaluatedCallback) {
+    val intermediateMethods = chain.intermediateCalls
+      .map { findStreamCallMethod(process.session, it) }
+
+    if (null in intermediateMethods) {
+      LOG.info("Cannot find declarations for some methods in stream chain")
+      return
+    }
+
+    val terminationMethod = findStreamCallMethod(process.session, chain.terminationCall)
+
+    if (terminationMethod == null) {
+      LOG.info("Cannot find declarations for termination method in stream chain")
+      return
+    }
+
+    runInDebuggerThread(process.debuggerSession.process) {
+      val identityValueModifier: (Value?) -> Value? = { value ->
+        LOG.info("Modifying value of type ${value?.type()?.name()}")
+        value
+      }
+
+      val valueTransformer: MethodReturnValueTransformer = PrintToStdoutMethodReturnValueTransformer()
+
+      // TODO: создаем какую-то машинерию и все такое тут или выносим отдельно? Лучше вынести отдельно и пробрасывать сюда
+
+      val intermediateStepsRequestors = chain.intermediateCalls.zip(intermediateMethods).map {
+        DebuggerUtils.createMethodExitBreakpoint(process, it.second!!, applyReturnValueTranformer(it.first, valueTransformer))
+      }
+
+      DebuggerUtils.createMethodExitBreakpoint(process, terminationMethod) { requestor, suspendContext, ev ->
+        // TODO: у терминального оператора по идее надо что-то другое сделать
+        //  а вообще было бы неплохо высунуть наружу это, т. е. делегировать возню с терминальный/нетерминальный
+        //  операцией и какой именно метод вызывается (может туда peek не подойдет или что-то такое) в MethodReturnValueTransformer
+        applyReturnValueTranformer(chain.terminationCall, valueTransformer)(requestor, suspendContext, ev)
+
+        chainEvaluatedCallback.onChainEvaluated()
+
+        val debugProcess = suspendContext.debugProcess
+        intermediateStepsRequestors.forEach {
+          debugProcess.requestsManager.deleteRequest(it)
+        }
+        debugProcess.requestsManager.deleteRequest(requestor)
+      }
+    }
+  }
+}
+
+// TODO: надо куда-то вкорячить это
+fun applyReturnValueTranformer(chainStep: StreamCall, valueTransformer: MethodReturnValueTransformer) = transformerCallback@{
+  requestor: FilteredRequestor, suspendContext: SuspendContextImpl, event: MethodExitEvent ->
+
+  // TODO: сюда хочется вснуть вызов какого-то красивого DSL для манипулирования значениями
+  //  У DSL хочется, чтобы можно было создавать и в какое-то место складывать переменные разных типов
+  val threadProxy = suspendContext.thread ?: return@transformerCallback
+
+  val originalReturnValue = try {
+    event.returnValue()
+  }
+  catch (e: UnsupportedOperationException) {
+    val vm = event.virtualMachine()
+    LOG.info("Return value interception is not supported in ${vm.name()} ${vm.version()}", e)
+    return@transformerCallback
+  }
+
+  val replacedReturnValue = valueTransformer
+                              .transform(chainStep, threadProxy.threadReference, originalReturnValue) ?: return@transformerCallback
+
+  // TODO: ClassNotLoadedException, IncompatibleThreadStateException, InvalidTypeException
+  threadProxy.forceEarlyReturn(replacedReturnValue)
+}
+
+fun findStreamCallMethod(session: XDebugSession, step: StreamCall): PsiMethod? {
+  val currentFile = session.currentPosition?.file ?: return null
+
+  val psiManager = PsiManager.getInstance(session.project)
+  val psiFile = psiManager.findFile(currentFile) ?: return null
+
+  val methodCallExpression = psiFile.findElementAt(step.textRange.endOffset)?.prevSibling as? PsiMethodCallExpression ?: return null
+  return methodCallExpression.methodExpression.reference?.resolve() as? PsiMethod
+}

plugins/stream-debugger/src/com/intellij/debugger/streams/trace/breakpoint/ChainEvaluatedCallback.kt

@@ -0,0 +1,9 @@
+// Copyright 2000-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.intellij.debugger.streams.trace.breakpoint
+
+/**
+ * @author Shumaf Lovpache
+ */
+fun interface ChainEvaluatedCallback {
+  fun onChainEvaluated()
+}

plugins/stream-debugger/src/com/intellij/debugger/streams/trace/breakpoint/DebuggerUtils.kt

@@ -0,0 +1,106 @@
+// Copyright 2000-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.intellij.debugger.streams.trace.breakpoint
+
+import com.intellij.debugger.engine.DebugProcessImpl
+import com.intellij.debugger.engine.DebuggerManagerThreadImpl
+import com.intellij.debugger.engine.JavaDebugProcess
+import com.intellij.debugger.engine.SuspendContextImpl
+import com.intellij.debugger.engine.events.DebuggerCommandImpl
+import com.intellij.debugger.engine.requests.RequestManagerImpl
+import com.intellij.debugger.impl.PrioritizedTask
+import com.intellij.debugger.ui.breakpoints.FilteredRequestor
+import com.intellij.openapi.diagnostic.logger
+import com.intellij.psi.PsiMethod
+import com.intellij.psi.util.TypeConversionUtil
+import com.intellij.util.containers.JBIterable
+import com.sun.jdi.ClassNotPreparedException
+import com.sun.jdi.Method
+import com.sun.jdi.VirtualMachine
+import com.sun.jdi.event.MethodExitEvent
+import com.sun.jdi.request.MethodExitRequest
+import kotlin.streams.asSequence
+import kotlin.streams.asStream
+
+private val LOG = logger<DebuggerUtils>()
+
+fun interface MethodExitCallback {
+  fun beforeMethodExit(requestor: FilteredRequestor, suspendContext: SuspendContextImpl, event: MethodExitEvent)
+}
+
+/**
+ * @author Shumaf Lovpache
+ */
+object DebuggerUtils {
+  public fun runInDebuggerThread(debugProcess: DebugProcessImpl, action: () -> Unit) {
+    val command = object : DebuggerCommandImpl(PrioritizedTask.Priority.NORMAL) {
+      override fun action() {
+        action()
+      }
+    }
+
+    val managerThread = debugProcess.managerThread
+    if (DebuggerManagerThreadImpl.isManagerThread()) {
+      managerThread.invoke(command)
+    }
+    else {
+      managerThread.schedule(command)
+    }
+  }
+
+  fun createMethodExitBreakpoint(process: JavaDebugProcess, psiMethod: PsiMethod, callback: MethodExitCallback): FilteredRequestor? {
+    val vmMethod = findVmMethod(process, psiMethod) ?: return null
+
+    val javaDebuggerSession = process.debuggerSession
+    val debugProcess = javaDebuggerSession.process
+
+    val requestor = MethodExitRequestor(debugProcess.project, vmMethod, callback)
+    enableMethodExitRequest(debugProcess, vmMethod, requestor)
+    return requestor
+  }
+
+  private fun enableMethodExitRequest(debugProcess: DebugProcessImpl, vmMethod: Method, requestor: FilteredRequestor) {
+    val requestManager: RequestManagerImpl = debugProcess.requestsManager ?: return
+    val methodExitRequest: MethodExitRequest = requestManager.createMethodExitRequest(requestor)
+    methodExitRequest.addClassFilter(vmMethod.declaringType())
+    methodExitRequest.enable()
+  }
+
+  private fun findVmMethod(process: JavaDebugProcess, psiMethod: PsiMethod): Method? {
+    val javaDebuggerSession = process.debuggerSession
+    val debugProcess = javaDebuggerSession.process
+    val vm: VirtualMachine = debugProcess.virtualMachineProxy.virtualMachine
+
+    val fqClassName = psiMethod.containingClass?.qualifiedName
+    val vmClass = vm.classesByName(fqClassName).firstOrNull()
+    if (vmClass == null) {
+      LOG.info("Class $fqClassName not found by jvm")
+      return null
+    }
+
+    try {
+      val vmMethod = vmClass.methods().findByPsiMethodSignature(psiMethod) // TODO: methodsByName(name, jni like signature)
+      if (vmMethod == null) {
+        LOG.info("Can not find method with signature ${psiMethod.signatureText()} in $fqClassName")
+        return null
+      }
+
+      return vmMethod
+    }
+    catch (e: ClassNotPreparedException) {
+      LOG.warn("Failed to retreive $fqClassName method because class not yet been prepared.", e)
+    }
+
+    return null
+  }
+
+  internal fun Method?.equalBySignature(other: Method): Boolean = this != null && this.name() == other.name()
+                                                                  && this.returnTypeName() == other.returnTypeName()
+                                                                  && this.argumentTypeNames() == other.argumentTypeNames()
+
+  private fun List<Method>.findByPsiMethodSignature(psiMethod: PsiMethod) = this.find {
+    it.name() == psiMethod.name
+    && it.returnTypeName() == TypeConversionUtil.erasure(psiMethod.returnType)?.canonicalText
+    && it.argumentTypeNames() == psiMethod.parameterList.parameters
+      .map { param -> TypeConversionUtil.erasure(param.type)?.canonicalText }
+  }
+}

plugins/stream-debugger/src/com/intellij/debugger/streams/trace/breakpoint/MethodBreakpointTracer.kt

@@ -0,0 +1,33 @@
+// Copyright 2000-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.intellij.debugger.streams.trace.breakpoint
+
+import com.intellij.debugger.engine.JavaDebugProcess
+import com.intellij.debugger.streams.trace.StreamTracer
+import com.intellij.debugger.streams.trace.TraceResultInterpreter
+import com.intellij.debugger.streams.trace.TracingCallback
+import com.intellij.debugger.streams.trace.breakpoint.value.transform.MethodReturnValueTransformer
+import com.intellij.debugger.streams.trace.breakpoint.value.transform.PrintToStdoutMethodReturnValueTransformer
+import com.intellij.debugger.streams.wrapper.StreamChain
+import com.intellij.openapi.diagnostic.logger
+import com.intellij.xdebugger.XDebugSession
+
+private val LOG = logger<MethodBreakpointTracer>()
+
+/**
+ * @author Shumaf Lovpache
+ */
+class MethodBreakpointTracer(val mySession: XDebugSession,
+                             val breakpointConfigurator: BreakpointConfigurator,
+                             val myResultInterpreter: TraceResultInterpreter) : StreamTracer {
+  override fun trace(chain: StreamChain, callback: TracingCallback) {
+    val xDebugProcess = mySession.debugProcess as? JavaDebugProcess ?: return
+
+    // TODO: create objects for tracer
+
+    breakpointConfigurator.setBreakpoints(xDebugProcess, chain) {
+      LOG.info("stream chain evaluated")
+    }
+    mySession.resume()
+  }
+}
+

plugins/stream-debugger/src/com/intellij/debugger/streams/trace/breakpoint/MethodExitCallbackImpl.kt

@@ -0,0 +1,22 @@
+// Copyright 2000-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.intellij.debugger.streams.trace.breakpoint
+
+import com.intellij.debugger.engine.SuspendContextImpl
+import com.intellij.debugger.ui.breakpoints.FilteredRequestor
+import com.intellij.openapi.diagnostic.logger
+import com.sun.jdi.event.LocatableEvent
+import com.sun.jdi.event.MethodExitEvent
+
+private val LOG = logger<MethodExitCallbackImpl>()
+
+class MethodExitCallbackImpl : MethodExitCallback {
+  override fun beforeMethodExit(requestor: FilteredRequestor, suspendContext: SuspendContextImpl, event: MethodExitEvent) {
+    TODO("not implemented")
+    val vm = event.virtualMachine()
+
+    // This should be checked inside callback
+    if (!vm.canGetMethodReturnValues()) {
+      LOG.info("Can't modify method return value because vm version (${vm.version()}) does not supports this feature")
+    }
+  }
+}

plugins/stream-debugger/src/com/intellij/debugger/streams/trace/breakpoint/MethodExitRequestor.kt

@@ -0,0 +1,53 @@
+// Copyright 2000-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package com.intellij.debugger.streams.trace.breakpoint
+
+import com.intellij.debugger.engine.events.SuspendContextCommandImpl
+import com.intellij.debugger.settings.DebuggerSettings
+import com.intellij.debugger.streams.trace.breakpoint.DebuggerUtils.equalBySignature
+import com.intellij.debugger.ui.breakpoints.FilteredRequestorImpl
+import com.intellij.openapi.diagnostic.logger
+import com.intellij.openapi.project.Project
+import com.sun.jdi.Method
+import com.sun.jdi.event.LocatableEvent
+import com.sun.jdi.event.MethodExitEvent
+import com.sun.jdi.request.InvalidRequestStateException
+
+private val LOG = logger<MethodExitRequestor>()
+
+/**
+ * @author Shumaf Lovpache
+ */
+class MethodExitRequestor(
+  project: Project,
+  val method: Method,
+  val callback: MethodExitCallback
+) : FilteredRequestorImpl(project) {
+  override fun processLocatableEvent(action: SuspendContextCommandImpl, event: LocatableEvent?): Boolean {
+    if (event == null) return false
+    val context = action.suspendContext ?: return false
+
+    val currentExecutingMethod = event.location().method()
+    if (event !is MethodExitEvent) return false
+
+    if (context.thread?.isSuspended == true && currentExecutingMethod.equalBySignature(method)) {
+      try {
+        callback.beforeMethodExit(this, context, event)
+      }
+      catch (e: Throwable) {
+        LOG.info(e)
+      }
+      finally {
+        try {
+          event.request().disable()
+        }
+        catch (e: InvalidRequestStateException) {
+          LOG.warn(e)
+        }
+      }
+    }
+
+    return false
+  }
+
+  override fun getSuspendPolicy(): String = DebuggerSettings.SUSPEND_ALL
+}