From 246f79ab8676bfaaab3ec6c0f28792271e3fb5ad Mon Sep 17 00:00:00 2001
From: "Vitaliy.Bibaev" <vitaliy.bibaev@jetbrains.com>
Date: Thu, 16 Nov 2017 10:41:48 +0300
Subject: [PATCH] Implement handler for filter call for kotlin collections

---
 .../impl/handler/collections/CallUtils.kt     | 21 ++++++
 .../handler/collections/FilterCallHandler.kt  | 69 +++++++++++++++++++
 2 files changed, 90 insertions(+)
 create mode 100644 idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/trace/impl/handler/collections/CallUtils.kt

diff --git a/idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/trace/impl/handler/collections/CallUtils.kt b/idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/trace/impl/handler/collections/CallUtils.kt
new file mode 100644
index 00000000000..641f6202838
--- /dev/null
+++ b/idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/trace/impl/handler/collections/CallUtils.kt
@@ -0,0 +1,21 @@
+package com.intellij.debugger.streams.kotlin.trace.impl.handler.collections
+
+import com.intellij.debugger.streams.kotlin.trace.dsl.KotlinTypes
+import com.intellij.debugger.streams.wrapper.*
+import com.intellij.debugger.streams.wrapper.impl.IntermediateStreamCallImpl
+import com.intellij.debugger.streams.wrapper.impl.TerminatorStreamCallImpl
+
+/**
+ * @author Vitaliy.Bibaev
+ */
+fun IntermediateStreamCall.withArgs(args: List<CallArgument>) =
+    IntermediateStreamCallImpl(name, args, typeBefore, typeAfter, textRange)
+
+fun TerminatorStreamCall.withArgs(args: List<CallArgument>) =
+    TerminatorStreamCallImpl(name, args, typeBefore, resultType, textRange)
+
+fun StreamCall.typeBefore() =
+    if (StreamCall@ this is TypeBeforeAware) StreamCall@ this.typeBefore else KotlinTypes.ANY
+
+fun StreamCall.typeAfter() =
+    if (StreamCall@ this is TypeAfterAware) StreamCall@ this.typeAfter else KotlinTypes.ANY
\ No newline at end of file
diff --git a/idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/trace/impl/handler/collections/FilterCallHandler.kt b/idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/trace/impl/handler/collections/FilterCallHandler.kt
index c2aaa218aa8..32d56202140 100644
--- a/idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/trace/impl/handler/collections/FilterCallHandler.kt
+++ b/idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/trace/impl/handler/collections/FilterCallHandler.kt
@@ -1,7 +1,76 @@
 package com.intellij.debugger.streams.kotlin.trace.impl.handler.collections
 
+import com.intellij.debugger.streams.trace.dsl.*
+import com.intellij.debugger.streams.trace.dsl.impl.TextExpression
+import com.intellij.debugger.streams.trace.impl.handler.type.ClassTypeImpl
+import com.intellij.debugger.streams.wrapper.CallArgument
+import com.intellij.debugger.streams.wrapper.IntermediateStreamCall
+import com.intellij.debugger.streams.wrapper.StreamCall
+import com.intellij.debugger.streams.wrapper.TerminatorStreamCall
+import com.intellij.debugger.streams.wrapper.impl.CallArgumentImpl
+
 /**
  * @author Vitaliy.Bibaev
  */
 class FilterCallHandler : BothSemanticsHandler {
+  private companion object {
+    val VALUES_ARRAY_NAME = "objectsInPredicate"
+    val PREDICATE_RESULT_ARRAY_NAME = "filteringResults"
+
+    fun oldPredicateVariableName(order: Int): String = "filterPredicate" + order
+  }
+
+  override fun variablesDeclaration(call: StreamCall, order: Int, dsl: Dsl): List<VariableDeclaration> {
+    val types = dsl.types
+
+    val timeToObjectMap = dsl.linkedMap(types.INT, call.typeBefore(), call.name + "Values" + order)
+    val predicateResultMap = dsl.linkedMap(types.INT, types.BOOLEAN, call.name + "PredicateValue" + order)
+    val predicate = call.arguments[0]
+    // TODO: use generic types in CallArgument
+    val oldFilterPredicate = dsl.variable(ClassTypeImpl(predicate.type), oldPredicateVariableName(order))
+
+    return listOf(timeToObjectMap.defaultDeclaration(), predicateResultMap.defaultDeclaration(),
+        dsl.declaration(oldFilterPredicate, TextExpression(predicate.text), false))
+  }
+
+  override fun prepareResult(dsl: Dsl, variables: List<Variable>): CodeBlock {
+    val values = variables[0] as MapVariable
+    val filterResult = variables[1] as MapVariable
+    return dsl.block {
+      add(values.convertToArray(dsl, VALUES_ARRAY_NAME))
+      add(filterResult.convertToArray(dsl, PREDICATE_RESULT_ARRAY_NAME))
+    }
+  }
+
+  override fun additionalCallsBefore(call: StreamCall, dsl: Dsl): List<IntermediateStreamCall> = emptyList()
+
+  override fun additionalCallsAfter(call: StreamCall, dsl: Dsl): List<IntermediateStreamCall> = emptyList()
+
+  override fun getResultExpression(call: StreamCall, dsl: Dsl, variables: List<Variable>): Expression {
+    return dsl.newArray(dsl.types.ANY, TextExpression(VALUES_ARRAY_NAME), TextExpression(PREDICATE_RESULT_ARRAY_NAME))
+  }
+
+  override fun transformAsIntermediateCall(call: IntermediateStreamCall, variables: List<Variable>, dsl: Dsl): IntermediateStreamCall {
+    return call.withArgs(listOf(createNewPredicate(variables, dsl)))
+  }
+
+  override fun transformAsTerminalCall(call: TerminatorStreamCall, variables: List<Variable>, dsl: Dsl): TerminatorStreamCall {
+    return call.withArgs(listOf(createNewPredicate(variables, dsl)))
+  }
+
+  private fun createNewPredicate(variables: List<Variable>, dsl: Dsl): CallArgument {
+    val valuesMap = variables[0] as MapVariable
+    val filteringMap = variables[1] as MapVariable
+    val oldPredicate = variables[2]
+    val newPredicate = dsl.lambda("value") {
+      +dsl.updateTime()
+      +valuesMap.set(dsl.currentTime(), lambdaArg)
+      val filterResult = dsl.variable(dsl.types.BOOLEAN, "result")
+      declare(filterResult, oldPredicate.call("invoke", lambdaArg), false)
+      +filteringMap.set(dsl.currentTime(), filterResult)
+      doReturn(filterResult)
+    }.toCode()
+
+    return CallArgumentImpl(oldPredicate.type.genericTypeName, newPredicate)
+  }
 }
\ No newline at end of file