[PowerAssert] Find the earliest starting offset for receiver expression

Calls of infix functions have a start offset which doesn't include the
receiver. Property accessors have a start offset at the start of the
property name and do not include their receiver expression. This
combination can cause problems when the entire expression needs to be
displayed, including the entirety of a complex receiver. Make sure to
navigate expressions to find their earliest starting offset and use that
instead.

^KT-65810 Fixed

plugins/power-assert/power-assert.backend/src/org/jetbrains/kotlin/powerassert/IrUtils.kt

@@ -21,6 +21,7 @@ package org.jetbrains.kotlin.powerassert
 
 import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
@@ -28,6 +29,8 @@ import org.jetbrains.kotlin.ir.expressions.IrFunctionExpression
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionExpressionImpl
 import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
+import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.name.Name
 
 fun IrBuilderWithScope.irString(builderAction: StringBuilder.() -> Unit) =
@@ -57,3 +60,17 @@ fun IrBuilderWithScope.irLambda(
     }
     return IrFunctionExpressionImpl(startOffset, endOffset, lambdaType, lambda, IrStatementOrigin.LAMBDA)
 }
+
+val IrElement.earliestStartOffset: Int
+    get() {
+        var offset = startOffset
+        this.acceptChildrenVoid(
+            object : IrElementVisitorVoid {
+                override fun visitElement(element: IrElement) {
+                    if (element.startOffset < offset) offset = element.startOffset
+                    element.acceptChildrenVoid(this)
+                }
+            },
+        )
+        return offset
+    }

plugins/power-assert/power-assert.backend/src/org/jetbrains/kotlin/powerassert/PowerAssertCallTransformer.kt

@@ -143,7 +143,7 @@ class PowerAssertCallTransformer(
                     val newArguments = arguments + call.getValueArgument(index)
                     return recursive(index + 1, dispatch, extension, newArguments, variables)
                 } else {
-                    return buildDiagramNesting(root, variables) { argument, newVariables ->
+                    return buildDiagramNesting(sourceFile, root, variables) { argument, newVariables ->
                         val newArguments = arguments + argument
                         recursive(index + 1, dispatch, extension, newArguments, newVariables)
                     }
@@ -151,8 +151,8 @@ class PowerAssertCallTransformer(
             }
         }
 
-        return buildDiagramNestingNullable(dispatchRoot) { dispatch, dispatchNewVariables ->
+        return buildDiagramNestingNullable(sourceFile, dispatchRoot) { dispatch, dispatchNewVariables ->
-            buildDiagramNestingNullable(extensionRoot, dispatchNewVariables) { extension, extensionNewVariables ->
+            buildDiagramNestingNullable(sourceFile, extensionRoot, dispatchNewVariables) { extension, extensionNewVariables ->
                 recursive(0, dispatch, extension, emptyList(), extensionNewVariables)
             }
         }

plugins/power-assert/power-assert.backend/src/org/jetbrains/kotlin/powerassert/diagram/DiagramBuilder.kt

@@ -24,31 +24,34 @@ import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 
 fun IrBuilderWithScope.buildDiagramNesting(
+    sourceFile: SourceFile,
     root: Node,
     variables: List<IrTemporaryVariable> = emptyList(),
     call: IrBuilderWithScope.(IrExpression, List<IrTemporaryVariable>) -> IrExpression,
 ): IrExpression {
-    return buildExpression(root, variables) { argument, subStack ->
+    return buildExpression(sourceFile, root, variables) { argument, subStack ->
         call(argument, subStack)
     }
 }
 
 fun IrBuilderWithScope.buildDiagramNestingNullable(
+    sourceFile: SourceFile,
     root: Node?,
     variables: List<IrTemporaryVariable> = emptyList(),
     call: IrBuilderWithScope.(IrExpression?, List<IrTemporaryVariable>) -> IrExpression,
 ): IrExpression {
-    return if (root != null) buildDiagramNesting(root, variables, call) else call(null, variables)
+    return if (root != null) buildDiagramNesting(sourceFile, root, variables, call) else call(null, variables)
 }
 
 private fun IrBuilderWithScope.buildExpression(
+    sourceFile: SourceFile,
     node: Node,
     variables: List<IrTemporaryVariable>,
     call: IrBuilderWithScope.(IrExpression, List<IrTemporaryVariable>) -> IrExpression,
 ): IrExpression = when (node) {
-    is ExpressionNode -> add(node, variables, call)
+    is ExpressionNode -> add(sourceFile, node, variables, call)
-    is AndNode -> nest(node, 0, variables, call)
+    is AndNode -> nest(sourceFile, node, 0, variables, call)
-    is OrNode -> nest(node, 0, variables, call)
+    is OrNode -> nest(sourceFile, node, 0, variables, call)
     else -> TODO("Unknown node type=$node")
 }
 
@@ -66,6 +69,7 @@ private fun IrBuilderWithScope.buildExpression(
  * ```
  */
 private fun IrBuilderWithScope.add(
+    sourceFile: SourceFile,
     node: ExpressionNode,
     variables: List<IrTemporaryVariable>,
     call: IrBuilderWithScope.(IrExpression, List<IrTemporaryVariable>) -> IrExpression,
@@ -73,7 +77,7 @@ private fun IrBuilderWithScope.add(
     return irBlock {
         val head = node.expressions.first().deepCopyWithSymbols(scope.getLocalDeclarationParent())
         val expressions = (buildTree(head) as ExpressionNode).expressions
-        val transformer = IrTemporaryExtractionTransformer(this@irBlock, expressions.toSet())
+        val transformer = IrTemporaryExtractionTransformer(this@irBlock, expressions.toSet(), sourceFile)
         val transformed = expressions.first().transform(transformer, null)
         +call(transformed, variables + transformer.variables)
     }
@@ -96,6 +100,7 @@ private fun IrBuilderWithScope.add(
  * ```
  */
 private fun IrBuilderWithScope.nest(
+    sourceFile: SourceFile,
     node: AndNode,
     index: Int,
     variables: List<IrTemporaryVariable>,
@@ -103,14 +108,14 @@ private fun IrBuilderWithScope.nest(
 ): IrExpression {
     val children = node.children
     val child = children[index]
-    return buildExpression(child, variables) { argument, newVariables ->
+    return buildExpression(sourceFile, child, variables) { argument, newVariables ->
         if (index + 1 == children.size) {
             call(argument, newVariables) // last expression, result is false
         } else {
             irIfThenElse(
                 context.irBuiltIns.anyType,
                 argument,
-                nest(node, index + 1, newVariables, call), // more expressions, continue nesting
+                nest(sourceFile, node, index + 1, newVariables, call), // more expressions, continue nesting
                 call(irFalse(), newVariables), // short-circuit result to false
             )
         }
@@ -134,6 +139,7 @@ private fun IrBuilderWithScope.nest(
  * ```
  */
 private fun IrBuilderWithScope.nest(
+    sourceFile: SourceFile,
     node: OrNode,
     index: Int,
     variables: List<IrTemporaryVariable>,
@@ -141,7 +147,7 @@ private fun IrBuilderWithScope.nest(
 ): IrExpression {
     val children = node.children
     val child = children[index]
-    return buildExpression(child, variables) { argument, newVariables ->
+    return buildExpression(sourceFile, child, variables) { argument, newVariables ->
         if (index + 1 == children.size) {
             call(argument, newVariables) // last expression, result is false
         } else {
@@ -149,7 +155,7 @@ private fun IrBuilderWithScope.nest(
                 context.irBuiltIns.anyType,
                 argument,
                 call(irTrue(), newVariables), // short-circuit result to true
-                nest(node, index + 1, newVariables, call), // more expressions, continue nesting
+                nest(sourceFile, node, index + 1, newVariables, call), // more expressions, continue nesting
             )
         }
     }

plugins/power-assert/power-assert.backend/src/org/jetbrains/kotlin/powerassert/diagram/IrDiagram.kt

@@ -38,7 +38,7 @@ fun IrBuilderWithScope.irDiagramString(
     val callInfo = sourceFile.getSourceRangeInfo(call)
     val callIndent = callInfo.startColumnNumber
 
-    val stackValues = variables.map { it.toValueDisplay(sourceFile, callIndent, callInfo) }
+    val stackValues = variables.map { it.toValueDisplay(callIndent, callInfo) }
 
     val valuesByRow = stackValues.groupBy { it.row }
     val rows = sourceFile.getText(callInfo)
@@ -94,17 +94,14 @@ private data class ValueDisplay(
 )
 
 private fun IrTemporaryVariable.toValueDisplay(
-    fileSource: SourceFile,
     callIndent: Int,
     originalInfo: SourceRangeInfo,
 ): ValueDisplay {
-    val info = fileSource.getSourceRangeInfo(original)
+    var indent = sourceRangeInfo.startColumnNumber - callIndent
-    var indent = info.startColumnNumber - callIndent
+    var row = sourceRangeInfo.startLineNumber - originalInfo.startLineNumber
-    var row = info.startLineNumber - originalInfo.startLineNumber
 
-    val source = fileSource.getText(info)
+    val source = text.replace("\n" + " ".repeat(callIndent), "\n") // Remove additional indentation
-        .replace("\n" + " ".repeat(callIndent), "\n") // Remove additional indentation
+    val columnOffset = findDisplayOffset(original, sourceRangeInfo, source)
-    val columnOffset = findDisplayOffset(fileSource, original, source)
 
     val prefix = source.substring(0, columnOffset)
     val rowShift = prefix.count { it == '\n' }
@@ -151,20 +148,20 @@ private fun IrTemporaryVariable.toValueDisplay(
  * ```
  */
 private fun findDisplayOffset(
-    sourceFile: SourceFile,
     expression: IrExpression,
+    sourceRangeInfo: SourceRangeInfo,
     source: String,
 ): Int {
     return when (expression) {
-        is IrMemberAccessExpression<*> -> memberAccessOffset(sourceFile, expression, source)
+        is IrMemberAccessExpression<*> -> memberAccessOffset(expression, sourceRangeInfo, source)
         is IrTypeOperatorCall -> typeOperatorOffset(expression, source)
         else -> 0
     }
 }
 
 private fun memberAccessOffset(
-    sourceFile: SourceFile,
     expression: IrMemberAccessExpression<*>,
+    sourceRangeInfo: SourceRangeInfo,
     source: String,
 ): Int {
     when (expression.origin) {
@@ -190,8 +187,7 @@ private fun memberAccessOffset(
             ?: expression.extensionReceiver
             ?: expression.getValueArgument(0).takeIf { owner.origin == IrBuiltIns.BUILTIN_OPERATOR }
             ?: return 0
-        val expressionInfo = sourceFile.getSourceRangeInfo(expression)
+        var offset = receiver.endOffset - sourceRangeInfo.startOffset + 1
-        var offset = receiver.endOffset - expressionInfo.startOffset + 1
         if (offset < 0 || offset >= source.length) return 0 // infix function called using non-infix syntax
 
         // Continue until there is a non-whitespace character

plugins/power-assert/power-assert.backend/src/org/jetbrains/kotlin/powerassert/diagram/IrTemporaryVariable.kt

@@ -19,6 +19,7 @@
 
 package org.jetbrains.kotlin.powerassert.diagram
 
+import org.jetbrains.kotlin.ir.SourceRangeInfo
 import org.jetbrains.kotlin.ir.builders.IrStatementsBuilder
 import org.jetbrains.kotlin.ir.builders.irGet
 import org.jetbrains.kotlin.ir.builders.irTemporary
@@ -30,11 +31,14 @@ import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 data class IrTemporaryVariable(
     val temporary: IrVariable,
     val original: IrExpression,
+    val sourceRangeInfo: SourceRangeInfo,
+    val text: String,
 )
 
 class IrTemporaryExtractionTransformer(
     private val builder: IrStatementsBuilder<*>,
     private val transform: Set<IrExpression>,
+    private val sourceFile: SourceFile,
 ) : IrElementTransformerVoid() {
     private val _variables = mutableListOf<IrTemporaryVariable>()
     val variables: List<IrTemporaryVariable> = _variables
@@ -43,7 +47,9 @@ class IrTemporaryExtractionTransformer(
         return if (expression in transform) {
             val copy = expression.deepCopyWithSymbols(builder.scope.getLocalDeclarationParent())
             val variable = builder.irTemporary(super.visitExpression(expression))
-            _variables.add(IrTemporaryVariable(variable, copy))
+            val sourceRangeInfo = sourceFile.getSourceRangeInfo(copy)
+            val text = sourceFile.getText(sourceRangeInfo)
+            _variables.add(IrTemporaryVariable(variable, copy, sourceRangeInfo, text))
             builder.irGet(variable)
         } else {
             super.visitExpression(expression)

plugins/power-assert/power-assert.backend/src/org/jetbrains/kotlin/powerassert/diagram/SourceFile.kt

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.ir.SourceRangeInfo
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.path
 import org.jetbrains.kotlin.ir.expressions.IrCall
+import org.jetbrains.kotlin.powerassert.earliestStartOffset
 import java.io.File
 
 data class SourceFile(
@@ -59,14 +60,14 @@ data class SourceFile(
                 val receiver = element.extensionReceiver ?: element.dispatchReceiver
                 if (element.symbol.owner.isInfix && receiver != null) {
                     // When an infix function is called *not* with infix notation, the startOffset will not include the receiver.
-                    // Force the range to include the receiver, so it is always present
+                    // Force the range to include the receiver, so it is always present.
-                    range = receiver.startOffset..element.endOffset
+                    range = element.earliestStartOffset..element.endOffset
 
                     // The offsets of the receiver will *not* include surrounding parentheses so these need to be checked for
                     // manually.
-                    val substring = safeSubstring(receiver.startOffset - 1, receiver.endOffset + 1)
+                    val substring = safeSubstring(range.first - 1, receiver.endOffset + 1)
                     if (substring.startsWith('(') && substring.endsWith(')')) {
-                        range = receiver.startOffset - 1..element.endOffset
+                        range = (range.first - 1)..range.last
                     }
                 }
             }

plugins/power-assert/testData/codegen/infix/DispatchInfixComplexReceiver.box.txt

@@ -1,6 +1,6 @@
-wrapper mustEqual "world".length
+Complex.holder.wrapper mustEqual "world".length
-       |      |                         |
+|       |      |                         |
-       |      |                         5
+|       |      |                         5
-       |      Wrapper
+|       |      Wrapper
-       Holder(wrapper=Wrapper)
+|       Holder(wrapper=Wrapper)
 Complex

plugins/power-assert/testData/codegen/infix/DispatchNonInfixComplexReceiver.box.txt

@@ -1,6 +1,6 @@
-wrapper.mustEqual("world".length)
+Complex.holder.wrapper.mustEqual("world".length)
-       |      |                         |
+|       |      |                         |
-       |      |                         5
+|       |      |                         5
-       |      Wrapper
+|       |      Wrapper
-       Holder(wrapper=Wrapper)
+|       Holder(wrapper=Wrapper)
 Complex

plugins/power-assert/testData/codegen/infix/ExtensionInfixComplexReceiver.box.txt

@@ -1,5 +1,5 @@
-length mustEqual "world".length
+"hello".substring(1, 4).length mustEqual "world".length
-               |                        |
+        |               |                        |
-               |                        5
+        |               |                        5
-               3
+        |               3
-ell
+        ell

plugins/power-assert/testData/codegen/infix/ExtensionNonInfixComplexReceiver.box.txt

@@ -1,5 +1,5 @@
-length.mustEqual("world".length)
+"hello".substring(1, 4).length.mustEqual("world".length)
-               |                        |
+        |               |                        |
-               |                        5
+        |               |                        5
-               3
+        |               3
-ell
+        ell