Support distinctBy intermediate call

idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/lib/sequence/KotlinSequencesSupport.kt

@@ -5,6 +5,7 @@ import com.intellij.debugger.streams.kotlin.resolve.ChunkedResolver
 import com.intellij.debugger.streams.kotlin.resolve.FilteredMapResolver
 import com.intellij.debugger.streams.kotlin.resolve.WindowedResolver
 import com.intellij.debugger.streams.kotlin.trace.impl.handler.sequence.FilterIsInstanceHandler
+import com.intellij.debugger.streams.kotlin.trace.impl.handler.sequence.KotlinDistinctByHandler
 import com.intellij.debugger.streams.lib.IntermediateOperation
 import com.intellij.debugger.streams.lib.impl.*
 import com.intellij.debugger.streams.resolve.AppendResolver
@@ -30,7 +31,7 @@ class KotlinSequencesSupport : LibrarySupportBase() {
     addIntermediateOperationsSupport(*sortedOperations("sorted", "sortedBy", "sortedDescending", "sortedWith"))
 
     addIntermediateOperationsSupport(DistinctOperation("distinct", ::DistinctTraceHandler))
-//    addIntermediateOperationsSupport(DistinctOperation("distinctBy", ::DistinctByKeyHandler)) waiting for a fix in the platform
+    addIntermediateOperationsSupport(DistinctOperation("distinctBy", ::KotlinDistinctByHandler))
 
     addIntermediateOperationsSupport(ConcatOperation("plus", AppendResolver()))
     addIntermediateOperationsSupport(ConcatOperation("plusElement", AppendResolver()))

idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/trace/impl/handler/sequence/KotlinDistinctByHandler.kt

@@ -0,0 +1,125 @@
+// Copyright 2000-2017 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.kotlin.trace.impl.handler.sequence
+
+import com.intellij.debugger.streams.kotlin.trace.dsl.KotlinTypes
+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.trace.impl.handler.unified.HandlerBase
+import com.intellij.debugger.streams.trace.impl.handler.unified.PeekTraceHandler
+import com.intellij.debugger.streams.wrapper.CallArgument
+import com.intellij.debugger.streams.wrapper.IntermediateStreamCall
+import com.intellij.debugger.streams.wrapper.impl.CallArgumentImpl
+import com.intellij.debugger.streams.wrapper.impl.IntermediateStreamCallImpl
+import one.util.streamex.StreamEx
+
+/**
+ * Based on com.intellij.debugger.streams.trace.impl.handler.unified.DistinctByKeyHandler
+ * @author Vitaliy.Bibaev
+ */
+class KotlinDistinctByHandler(callNumber: Int, private val call: IntermediateStreamCall, dsl: Dsl)
+  : HandlerBase.Intermediate(dsl) {
+  private companion object {
+    val KEY_EXTRACTOR_VARIABLE_PREFIX = "keyExtractor"
+    val TRANSITIONS_ARRAY_NAME = "transitionsArray"
+  }
+
+  private val peekHandler = PeekTraceHandler(callNumber, "distinctBy", call.typeBefore, call.typeAfter, dsl)
+  private val keyExtractor: CallArgument
+  private val extractorVariable: Variable
+  private val beforeTimes = dsl.list(dsl.types.INT, call.name + callNumber + "BeforeTimes")
+  private val beforeValues = dsl.list(KotlinTypes.NULLABLE_ANY, call.name + callNumber + "BeforeValues")
+  private val keys = dsl.list(KotlinTypes.NULLABLE_ANY, call.name + callNumber + "Keys")
+  private val time2ValueAfter = dsl.linkedMap(dsl.types.INT, KotlinTypes.NULLABLE_ANY, call.name + callNumber + "after")
+
+  init {
+    val arguments = call.arguments
+    assert(arguments.isNotEmpty(), { "Key extractor is not specified" })
+    keyExtractor = arguments.first()
+    extractorVariable = dsl.variable(ClassTypeImpl(keyExtractor.type), KEY_EXTRACTOR_VARIABLE_PREFIX + callNumber)
+  }
+
+  override fun additionalVariablesDeclaration(): List<VariableDeclaration> {
+    val extractor = dsl.declaration(extractorVariable, TextExpression(keyExtractor.text), false)
+    val variables =
+        mutableListOf(extractor, beforeTimes.defaultDeclaration(), beforeValues.defaultDeclaration(),
+            time2ValueAfter.defaultDeclaration(), keys.defaultDeclaration())
+    variables.addAll(peekHandler.additionalVariablesDeclaration())
+
+    return variables
+  }
+
+  override fun transformCall(call: IntermediateStreamCall): IntermediateStreamCall {
+    val newKeyExtractor = dsl.lambda("x") {
+      val key = dsl.variable(KotlinTypes.NULLABLE_ANY, "key")
+      declare(key, extractorVariable.call("invoke", lambdaArg), false)
+      statement { beforeTimes.add(dsl.currentTime()) }
+      statement { beforeValues.add(lambdaArg) }
+      statement { keys.add(key) }
+      doReturn(key)
+    }.toCode()
+    return call.updateArguments(listOf(CallArgumentImpl(keyExtractor.type, newKeyExtractor)))
+  }
+
+  override fun prepareResult(): CodeBlock {
+    val keys2TimesBefore = dsl.map(KotlinTypes.NULLABLE_ANY, dsl.types.list(dsl.types.INT), "keys2Times")
+    val transitions = dsl.map(dsl.types.INT, dsl.types.INT, "transitionsMap")
+    StreamEx.of(1).distinct().toList()
+    return dsl.block {
+      add(peekHandler.prepareResult())
+      declare(keys2TimesBefore.defaultDeclaration())
+      declare(transitions.defaultDeclaration())
+
+      integerIteration(keys.size(), block@ this) {
+        val key = declare(variable(KotlinTypes.NULLABLE_ANY, "key"), keys.get(loopVariable), false)
+        val lst = list(dsl.types.INT, "lst")
+        declare(lst, keys2TimesBefore.computeIfAbsent(key, lambda("k") {
+          doReturn(newList(types.INT))
+        }), false)
+        statement { lst.add(beforeTimes.get(loopVariable)) }
+      }
+
+      forEachLoop(variable(types.INT, "afterTime"), time2ValueAfter.keys()) {
+        val afterTime = loopVariable
+        val valueAfter = declare(variable(KotlinTypes.NULLABLE_ANY, "valueAfter"), time2ValueAfter.get(loopVariable), false)
+        val key = declare(variable(KotlinTypes.NULLABLE_ANY, "key"), nullExpression, true)
+        integerIteration(beforeTimes.size(), forEachLoop@ this) {
+          ifBranch((valueAfter same beforeValues.get(loopVariable)) and !transitions.contains(beforeTimes.get(loopVariable))) {
+            key assign keys.get(loopVariable)
+            statement { breakIteration() }
+          }
+        }
+
+        forEachLoop(variable(types.INT, "beforeTime"), keys2TimesBefore.get(key)) {
+          statement { transitions.set(loopVariable, afterTime) }
+        }
+      }
+
+      add(transitions.convertToArray(this, "transitionsArray"))
+    }
+  }
+
+  override fun getResultExpression(): Expression =
+      dsl.newArray(dsl.types.ANY, peekHandler.resultExpression, TextExpression(TRANSITIONS_ARRAY_NAME))
+
+  override fun additionalCallsBefore(): List<IntermediateStreamCall> = peekHandler.additionalCallsBefore()
+
+  override fun additionalCallsAfter(): List<IntermediateStreamCall> {
+    val callsAfter = ArrayList(peekHandler.additionalCallsAfter())
+    val lambda = dsl.lambda("x") {
+      doReturn(time2ValueAfter.set(dsl.currentTime(), lambdaArg))
+    }
+
+    callsAfter.add(dsl.createPeekCall(call.typeAfter, lambda.toCode()))
+    return callsAfter
+  }
+
+  private fun CodeContext.integerIteration(border: Expression, block: CodeBlock, init: ForLoopBody.() -> Unit) {
+    block.forLoop(declaration(variable(types.INT, "i"), TextExpression("0"), true),
+        TextExpression("i < ${border.toCode()}"),
+        TextExpression("i = i + 1"), init)
+  }
+
+  private fun IntermediateStreamCall.updateArguments(args: List<CallArgument>): IntermediateStreamCall =
+      IntermediateStreamCallImpl("distinctBy", args, typeBefore, typeAfter, textRange)
+}

idea/testData/debugger/sequence/exec/sequence/outs/distinctBy.out

@@ -13,22 +13,22 @@ toList
     after: nothing
 mappings for distinctBy
   direct:
-    1 -> nothing
+    1 -> 2
-    3 -> nothing
+    3 -> 4
-    5 -> nothing
+    5 -> 4
-    6 -> nothing
+    6 -> 4
-    7 -> nothing
+    7 -> 4
-    8 -> nothing
+    8 -> 2
-    9 -> nothing
+    9 -> 4
-    10 -> nothing
+    10 -> 2
-    11 -> nothing
+    11 -> 2
-    12 -> nothing
+    12 -> 13
-    14 -> nothing
+    14 -> 15
   reverse:
-    nothing <- 2
+    1,8,10,11 <- 2
-    nothing <- 4
+    3,5,6,7,9 <- 4
-    nothing <- 13
+    12 <- 13
-    nothing <- 15
+    14 <- 15
 mappings for toList
   direct:
     2 -> nothing
@@ -37,7 +37,6 @@ mappings for toList
     15 -> nothing
   reverse:
     empty
-WRONG!
 Disconnected from the target VM
 
 Process finished with exit code 0

idea/testData/debugger/sequence/exec/sequence/outs/distinctByNullableElement.out

@@ -0,0 +1,32 @@
+LineBreakpoint created at DistinctByNullableElement.kt:5
+Run Java
+Connected to the target VM
+DistinctByNullableElement.kt:5
+sequenceOf(null, 1, 2, 3, null)
+.distinctBy({ it == null })
+.count()
+distinctBy
+    before: 1,3,5,6,7
+    after: 2,4
+count
+    before: 2,4
+    after: nothing
+mappings for distinctBy
+  direct:
+    1 -> 2
+    3 -> 4
+    5 -> 4
+    6 -> 4
+    7 -> 2
+  reverse:
+    1,7 <- 2
+    3,5,6 <- 4
+mappings for count
+  direct:
+    2 -> nothing
+    4 -> nothing
+  reverse:
+    empty
+Disconnected from the target VM
+
+Process finished with exit code 0

idea/testData/debugger/sequence/exec/sequence/outs/distinctByNullableKey.out

@@ -0,0 +1,36 @@
+LineBreakpoint created at DistinctByNullableKey.kt:5
+Run Java
+Connected to the target VM
+DistinctByNullableKey.kt:5
+sequenceOf(1, 2, 3, 4, 5)
+.distinctBy({ if (it % 2 == 0) null else it })
+.count()
+distinctBy
+    before: 1,3,5,7,8
+    after: 2,4,6,9
+count
+    before: 2,4,6,9
+    after: nothing
+mappings for distinctBy
+  direct:
+    1 -> 2
+    3 -> 4
+    5 -> 6
+    7 -> 4
+    8 -> 9
+  reverse:
+    1 <- 2
+    3,7 <- 4
+    5 <- 6
+    8 <- 9
+mappings for count
+  direct:
+    2 -> nothing
+    4 -> nothing
+    6 -> nothing
+    9 -> nothing
+  reverse:
+    empty
+Disconnected from the target VM
+
+Process finished with exit code 0

idea/testData/debugger/sequence/exec/sequence/outs/distinctByNullableKeyAndElement.out

@@ -0,0 +1,35 @@
+LineBreakpoint created at DistinctByNullableKeyAndElement.kt:5
+Run Java
+Connected to the target VM
+DistinctByNullableKeyAndElement.kt:5
+sequenceOf(1, 2, null, null, 3, 1)
+.distinctBy({ if (it == null) 2 else if (it == 3) null else it })
+.count()
+distinctBy
+    before: 1,3,5,6,7,9
+    after: 2,4,8
+count
+    before: 2,4,8
+    after: nothing
+mappings for distinctBy
+  direct:
+    1 -> 2
+    3 -> 4
+    5 -> 4
+    6 -> 4
+    7 -> 8
+    9 -> 2
+  reverse:
+    1,9 <- 2
+    3,5,6 <- 4
+    7 <- 8
+mappings for count
+  direct:
+    2 -> nothing
+    4 -> nothing
+    8 -> nothing
+  reverse:
+    empty
+Disconnected from the target VM
+
+Process finished with exit code 0

idea/testData/debugger/sequence/exec/sequence/src/distinct/DistinctByNullableElement.kt

@@ -0,0 +1,6 @@
+package distinct
+
+fun main(args: Array<String>) {
+  // Breakpoint!
+  sequenceOf(null, 1, 2, 3, null).distinctBy { it == null }.count()
+}

idea/testData/debugger/sequence/exec/sequence/src/distinct/DistinctByNullableKey.kt

@@ -0,0 +1,6 @@
+package distinct
+
+fun main(args: Array<String>) {
+  // Breakpoint!
+  sequenceOf(1, 2, 3, 4, 5).distinctBy { if (it % 2 == 0) null else it }.count()
+}

idea/testData/debugger/sequence/exec/sequence/src/distinct/DistinctByNullableKeyAndElement.kt

@@ -0,0 +1,6 @@
+package distinct
+
+fun main(args: Array<String>) {
+  // Breakpoint!
+  sequenceOf(1, 2, null, null, 3, 1).distinctBy { if (it == null) 2 else if (it == 3) null else it }.count()
+}

idea/tests/com/intellij/debugger/streams/kotlin/exec/sequence/DistinctOperationsTest.kt

@@ -7,5 +7,9 @@ package com.intellij.debugger.streams.kotlin.exec.sequence
 class DistinctOperationsTest : OperationsTestCase("distinct") {
   fun testDistinct() = doTestWithResult()
   fun testDistinctObjects() = doTestWithResult()
+
   fun testDistinctBy() = doTestWithResult()
+  fun testDistinctByNullableElement() = doTestWithResult()
+  fun testDistinctByNullableKey() = doTestWithResult()
+  fun testDistinctByNullableKeyAndElement() = doTestWithResult()
 }