Move a workaround for collections into CollectionChainTransformer

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

@@ -0,0 +1,44 @@
+// 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.collections
+
+import com.intellij.debugger.streams.kotlin.psi.impl.KotlinChainTransformerImpl
+import com.intellij.debugger.streams.kotlin.psi.resolveType
+import com.intellij.debugger.streams.psi.ChainTransformer
+import com.intellij.debugger.streams.wrapper.QualifierExpression
+import com.intellij.debugger.streams.wrapper.StreamChain
+import com.intellij.debugger.streams.wrapper.impl.StreamChainImpl
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.types.KotlinType
+
+/**
+ * @author Vitaliy.Bibaev
+ */
+class CollectionChainTransformer : ChainTransformer<KtCallExpression> {
+  override fun transform(chainCalls: List<KtCallExpression>, context: PsiElement): StreamChain {
+    val transformer = KotlinChainTransformerImpl(KotlinCollectionsTypeExtractor())
+    val chain = transformer.transform(chainCalls, context)
+
+    if (chainCalls.first().resolveType().isArray) {
+      val qualifier = WrappedQualifier(chain.qualifierExpression)
+      return StreamChainImpl(qualifier, chain.intermediateCalls, chain.terminationCall, chain.context)
+    }
+
+    return chain
+  }
+
+  /**
+   * Kotlin arrays have not {@code onEach} extension. But current implementation uses onEach to increment a time counter.
+   * We use asIterable to avoid further issues with the transformed expression evaluation
+   * TODO: Avoid showing "asIterable()" in the tab name in trace window
+   */
+  private class WrappedQualifier(private val qualifierExpression: QualifierExpression)
+    : QualifierExpression by qualifierExpression {
+    override val text: String
+      get() = qualifierExpression.text + ".asIterable()"
+  }
+
+  private val KotlinType.isArray: Boolean
+    get() = KotlinBuiltIns.isArray(this) || KotlinBuiltIns.isPrimitiveArray(this)
+}

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

@@ -3,7 +3,6 @@ package com.intellij.debugger.streams.kotlin.psi.collections
 
 import com.intellij.debugger.streams.kotlin.psi.KotlinPsiUtil
 import com.intellij.debugger.streams.kotlin.psi.impl.KotlinChainBuilderBase
-import com.intellij.debugger.streams.kotlin.psi.impl.KotlinChainTransformerImpl
 import com.intellij.debugger.streams.kotlin.psi.previousCall
 import com.intellij.debugger.streams.kotlin.psi.receiverType
 import org.jetbrains.kotlin.psi.KtCallExpression
@@ -14,7 +13,7 @@ import org.jetbrains.kotlin.types.typeUtil.supertypes
  * @author Vitaliy.Bibaev
  */
 class KotlinCollectionChainBuilder
-  : KotlinChainBuilderBase(KotlinChainTransformerImpl(KotlinCollectionsTypeExtractor())) {
+  : KotlinChainBuilderBase(CollectionChainTransformer()) {
   private companion object {
     // TODO: Avoid enumeration of all available types
     val SUPPORTED_RECEIVERS = setOf("kotlin.collections.Iterable", "kotlin.CharSequence", "kotlin.Array",

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

@@ -14,14 +14,12 @@ import com.intellij.debugger.streams.wrapper.StreamChain
 import com.intellij.debugger.streams.wrapper.impl.*
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.idea.caches.resolve.analyzeFully
 import org.jetbrains.kotlin.psi.KtCallExpression
 import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
 import org.jetbrains.kotlin.psi.KtValueArgument
 import org.jetbrains.kotlin.resolve.calls.callUtil.getParameterForArgument
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
-import org.jetbrains.kotlin.types.KotlinType
 
 /**
  * @author Vitaliy.Bibaev
@@ -65,25 +63,7 @@ class KotlinChainTransformerImpl(private val typeExtractor: CallTypeExtractor) :
     val parent = expression.parent as? KtDotQualifiedExpression
         ?: return QualifierExpressionImpl("", TextRange.EMPTY_RANGE, typeAfter)
     val receiver = parent.receiverExpression
-    val qualifier = QualifierExpressionImpl(receiver.text, receiver.textRange, typeAfter)
-    if (receiver.resolveType().isArray) {
-      return WrappedQualifier(qualifier)
-    }
 
-    return qualifier
+    return QualifierExpressionImpl(receiver.text, receiver.textRange, typeAfter)
   }
-
-  /**
-   * Kotlin arrays has not {@code onEach} extension. But current implementation uses onEach to increment a time counter.
-   * We use asIterable to avoid further issues with the transformed expression evaluation
-   * TODO: Avoid showing "asIterable()" in the tab name in trace window
-   */
-  private class WrappedQualifier(private val qualifierExpression: QualifierExpression)
-    : QualifierExpression by qualifierExpression {
-    override val text: String
-      get() = qualifierExpression.text + ".asIterable()"
-  }
-
-  private val KotlinType.isArray: Boolean
-    get() = KotlinBuiltIns.isArray(this) || KotlinBuiltIns.isPrimitiveArray(this)
 }