Add additional classes and interfaces to infer types of chain items

idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/psi/CallTypeExtractor.kt

@@ -0,0 +1,15 @@
+package com.intellij.debugger.streams.kotlin.psi
+
+import com.intellij.debugger.streams.trace.impl.handler.type.GenericType
+import org.jetbrains.kotlin.psi.KtCallExpression
+
+/**
+ * @author Vitaliy.Bibaev
+ */
+interface CallTypeExtractor {
+  fun extractIntermediateCallTypes(call: KtCallExpression): IntermediateCallTypes
+  fun extractTerminalCallTypes(call: KtCallExpression): TerminatorCallTypes
+
+  data class IntermediateCallTypes(val typeBefore: GenericType, val typeAfter: GenericType)
+  data class TerminatorCallTypes(val typeBefore: GenericType, val resultType: GenericType)
+}

idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/psi/impl/JavaStreamChainTypeExtractor.kt

@@ -0,0 +1,18 @@
+package com.intellij.debugger.streams.kotlin.psi.impl
+
+import com.intellij.debugger.streams.kotlin.psi.CallTypeExtractor
+import com.intellij.debugger.streams.kotlin.psi.CallTypeExtractor.IntermediateCallTypes
+import com.intellij.debugger.streams.kotlin.psi.CallTypeExtractor.TerminatorCallTypes
+import com.intellij.debugger.streams.kotlin.trace.dsl.KotlinTypes
+import org.jetbrains.kotlin.psi.KtCallExpression
+
+/**
+ * @author Vitaliy.Bibaev
+ */
+class JavaStreamChainTypeExtractor : CallTypeExtractor {
+  override fun extractIntermediateCallTypes(call: KtCallExpression): IntermediateCallTypes =
+      IntermediateCallTypes(KotlinTypes.NULLABLE_ANY, KotlinTypes.NULLABLE_ANY)
+
+  override fun extractTerminalCallTypes(call: KtCallExpression): TerminatorCallTypes =
+      TerminatorCallTypes(KotlinTypes.NULLABLE_ANY, KotlinTypes.NULLABLE_ANY)
+}

idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/psi/impl/KotlinChainTransformerImpl.kt

@@ -1,9 +1,9 @@
 // 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.psi.impl
 
+import com.intellij.debugger.streams.kotlin.psi.CallTypeExtractor
 import com.intellij.debugger.streams.kotlin.psi.callName
 import com.intellij.debugger.streams.kotlin.psi.resolveType
-import com.intellij.debugger.streams.kotlin.trace.dsl.KotlinTypes.NULLABLE_ANY
 import com.intellij.debugger.streams.psi.ChainTransformer
 import com.intellij.debugger.streams.wrapper.CallArgument
 import com.intellij.debugger.streams.wrapper.IntermediateStreamCall
@@ -23,26 +23,30 @@ import org.jetbrains.kotlin.types.KotlinType
 /**
  * @author Vitaliy.Bibaev
  */
-class KotlinChainTransformerImpl : ChainTransformer<KtCallExpression> {
+class KotlinChainTransformerImpl(private val typeExtractor: CallTypeExtractor) : ChainTransformer<KtCallExpression> {
   override fun transform(callChain: List<KtCallExpression>, context: PsiElement): StreamChain {
-    val firstCall = callChain.first()
-    val parent = firstCall.parent
-    val qualifier = if (parent is KtDotQualifiedExpression)
-      QualifierExpressionImpl(parent.receiverExpression.text, parent.receiverExpression.textRange, NULLABLE_ANY)
-    else
-      QualifierExpressionImpl("", TextRange.EMPTY_RANGE, NULLABLE_ANY) // This is possible only when this inherits Stream
-
     val intermediateCalls = mutableListOf<IntermediateStreamCall>()
     for (call in callChain.subList(0, callChain.size - 1)) {
+      val (typeBefore, typeAfter) = typeExtractor.extractIntermediateCallTypes(call)
       intermediateCalls += IntermediateStreamCallImpl(call.callName(),
-          call.valueArguments.map { createCallArgument(call, it) }, NULLABLE_ANY, NULLABLE_ANY, call.textRange)
+          call.valueArguments.map { createCallArgument(call, it) }, typeBefore, typeAfter, call.textRange)
     }
 
     val terminationsPsiCall = callChain.last()
-    // TODO: infer true types
+    val (typeBeforeTerminator, resultType) = typeExtractor.extractTerminalCallTypes(terminationsPsiCall)
     val terminationCall = TerminatorStreamCallImpl(terminationsPsiCall.callName(),
         terminationsPsiCall.valueArguments.map { createCallArgument(terminationsPsiCall, it) },
-        NULLABLE_ANY, NULLABLE_ANY, terminationsPsiCall.textRange)
+        typeBeforeTerminator, resultType, terminationsPsiCall.textRange)
+
+    val qualifierTypeAfter =
+        if (intermediateCalls.isEmpty()) typeBeforeTerminator else intermediateCalls.first().typeBefore
+
+    val firstCall = callChain.first()
+    val parent = firstCall.parent
+    val qualifier = if (parent is KtDotQualifiedExpression)
+      QualifierExpressionImpl(parent.receiverExpression.text, parent.receiverExpression.textRange, qualifierTypeAfter)
+    else // This is possible only when {@code this} inherits Stream/Iterable/Sequence/etc
+      QualifierExpressionImpl("", TextRange.EMPTY_RANGE, qualifierTypeAfter)
 
     return StreamChainImpl(qualifier, intermediateCalls, terminationCall, context)
   }

idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/psi/impl/KotlinCollectionChainBuilder.kt

@@ -9,7 +9,8 @@ import org.jetbrains.kotlin.psi.KtCallExpression
 /**
  * @author Vitaliy.Bibaev
  */
-class KotlinCollectionChainBuilder : KotlinChainBuilderBase(KotlinChainTransformerImpl()) {
+class KotlinCollectionChainBuilder
+  : KotlinChainBuilderBase(KotlinChainTransformerImpl(KotlinCollectionsTypeExtractor())) {
   private companion object {
     // TODO: Avoid enumeration of all available types
     val SUPPORTED_RECEIVERS = setOf("kotlin.collections.List", "kotlin.collections.Set")

idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/psi/impl/KotlinCollectionsTypeExtractor.kt

@@ -0,0 +1,16 @@
+package com.intellij.debugger.streams.kotlin.psi.impl
+
+import com.intellij.debugger.streams.kotlin.psi.CallTypeExtractor
+import com.intellij.debugger.streams.kotlin.trace.dsl.KotlinTypes
+import org.jetbrains.kotlin.psi.KtCallExpression
+
+/**
+ * @author Vitaliy.Bibaev
+ */
+class KotlinCollectionsTypeExtractor : CallTypeExtractor {
+  override fun extractIntermediateCallTypes(call: KtCallExpression): CallTypeExtractor.IntermediateCallTypes =
+      CallTypeExtractor.IntermediateCallTypes(KotlinTypes.NULLABLE_ANY, KotlinTypes.NULLABLE_ANY)
+
+  override fun extractTerminalCallTypes(call: KtCallExpression): CallTypeExtractor.TerminatorCallTypes =
+      CallTypeExtractor.TerminatorCallTypes(KotlinTypes.NULLABLE_ANY, KotlinTypes.NULLABLE_ANY)
+}

idea/idea-jvm/src/com/intellij/debugger/streams/kotlin/psi/impl/KotlinJavaStreamChainBuilder.kt

@@ -4,13 +4,13 @@ package com.intellij.debugger.streams.kotlin.psi.impl
 import com.intellij.debugger.streams.kotlin.psi.StreamCallChecker
 import com.intellij.debugger.streams.kotlin.psi.previousCall
 import org.jetbrains.kotlin.psi.KtCallExpression
-import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
 import java.util.*
 
 /**
  * @author Vitaliy.Bibaev
  */
-class KotlinJavaStreamChainBuilder(private val callChecker: StreamCallChecker) : KotlinChainBuilderBase(KotlinChainTransformerImpl()) {
+class KotlinJavaStreamChainBuilder(private val callChecker: StreamCallChecker)
+  : KotlinChainBuilderBase(KotlinChainTransformerImpl(JavaStreamChainTypeExtractor())) {
   override val existenceChecker: ExistenceChecker = MyExistenceChecker()
 
   override fun createChainsBuilder(): ChainBuilder = MyBuilderVisitor()