[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
2024-02-26 11:59:36 -06:00
parent d53d5dddc5
commit 721d02f4f3
10 changed files with 77 additions and 51 deletions
@@ -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
+    }
@@ -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 ->
-            buildDiagramNestingNullable(extensionRoot, dispatchNewVariables) { extension, extensionNewVariables ->
+        return buildDiagramNestingNullable(sourceFile, dispatchRoot) { dispatch, dispatchNewVariables ->
+            buildDiagramNestingNullable(sourceFile, extensionRoot, dispatchNewVariables) { extension, extensionNewVariables ->
                recursive(0, dispatch, extension, emptyList(), extensionNewVariables)
            }
        }
@@ -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 AndNode -> nest(node, 0, variables, call)
-    is OrNode -> nest(node, 0, variables, call)
+    is ExpressionNode -> add(sourceFile, node, variables, call)
+    is AndNode -> nest(sourceFile, 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
            )
        }
    }
@@ -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 = info.startColumnNumber - callIndent
-    var row = info.startLineNumber - originalInfo.startLineNumber
+    var indent = sourceRangeInfo.startColumnNumber - callIndent
+    var row = sourceRangeInfo.startLineNumber - originalInfo.startLineNumber

-    val source = fileSource.getText(info)
-        .replace("\n" + " ".repeat(callIndent), "\n") // Remove additional indentation
-    val columnOffset = findDisplayOffset(fileSource, original, source)
+    val source = text.replace("\n" + " ".repeat(callIndent), "\n") // Remove additional indentation
+    val columnOffset = findDisplayOffset(original, sourceRangeInfo, 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 - expressionInfo.startOffset + 1
+        var offset = receiver.endOffset - sourceRangeInfo.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
@@ -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)
@@ -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
-                    range = receiver.startOffset..element.endOffset
+                    // Force the range to include the receiver, so it is always present.
+                    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
                    }
                }
            }