Merge pull request #69 from bnorm/infix-receiver

Support transforming infix functions

.github/workflows/lint.yml

@@ -10,7 +10,7 @@ jobs:
       - uses: actions/checkout@v2
 
       - name: Run ktlint
-        uses: ScaCap/action-ktlint@1.3
+        uses: ScaCap/action-ktlint@1.4
         with:
           github_token: ${{ secrets.GITHUB_TOKEN }}
           reporter: github-pr-check

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/PowerAssertCallTransformer.kt

@@ -22,15 +22,14 @@ import com.bnorm.power.delegate.SamConversionLambdaFunctionDelegate
 import com.bnorm.power.delegate.SimpleFunctionDelegate
 import com.bnorm.power.diagram.IrTemporaryVariable
 import com.bnorm.power.diagram.Node
+import com.bnorm.power.diagram.SourceFile
 import com.bnorm.power.diagram.buildDiagramNesting
+import com.bnorm.power.diagram.buildDiagramNestingNullable
 import com.bnorm.power.diagram.buildTree
-import com.bnorm.power.diagram.info
 import com.bnorm.power.diagram.irDiagramString
-import com.bnorm.power.diagram.substring
 import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
-import org.jetbrains.kotlin.cli.common.messages.CompilerMessageLocation
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.ir.IrElement
@@ -38,10 +37,8 @@ import org.jetbrains.kotlin.ir.backend.js.utils.asString
 import org.jetbrains.kotlin.ir.builders.irCall
 import org.jetbrains.kotlin.ir.builders.irString
 import org.jetbrains.kotlin.ir.builders.parent
-import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.IrValueParameter
-import org.jetbrains.kotlin.ir.declarations.path
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrConst
 import org.jetbrains.kotlin.ir.expressions.IrExpression
@@ -68,8 +65,7 @@ import org.jetbrains.kotlin.ir.util.parentClassOrNull
 import org.jetbrains.kotlin.name.FqName
 
 class PowerAssertCallTransformer(
-  private val file: IrFile,
+  private val sourceFile: SourceFile,
-  private val fileSource: String,
   private val context: IrPluginContext,
   private val messageCollector: MessageCollector,
   private val functions: Set<FqName>
@@ -104,6 +100,10 @@ class PowerAssertCallTransformer(
       return super.visitCall(expression)
     }
 
+    val dispatchRoot =
+      if (expression.symbol.owner.isInfix) expression.dispatchReceiver?.let { buildTree(it) } else null
+    val extensionRoot =
+      if (expression.symbol.owner.isInfix) expression.extensionReceiver?.let { buildTree(it) } else null
     val messageArgument: IrExpression?
     val roots: List<Node?>
     if (delegate.function.valueParameters.size == function.valueParameters.size) {
@@ -119,42 +119,62 @@ class PowerAssertCallTransformer(
     }
 
     // If all roots are null, there are no transformable parameters
-    if (roots.all { it == null }) {
+    if (dispatchRoot == null && extensionRoot == null && roots.all { it == null }) {
       messageCollector.info(expression, "Expression is constant and will not be power-assert transformed")
       return super.visitCall(expression)
     }
 
     val symbol = currentScope!!.scope.scopeOwnerSymbol
     val builder = DeclarationIrBuilder(context, symbol, expression.startOffset, expression.endOffset)
-    return builder.diagram(expression, delegate, messageArgument, roots)
+    return builder.diagram(
+      call = expression,
+      delegate = delegate,
+      messageArgument = messageArgument,
+      roots = roots,
+      dispatchRoot = dispatchRoot,
+      extensionRoot = extensionRoot
+    )
   }
 
   private fun DeclarationIrBuilder.diagram(
-    original: IrCall,
+    call: IrCall,
     delegate: FunctionDelegate,
     messageArgument: IrExpression?,
     roots: List<Node?>,
-    index: Int = 0,
+    dispatchRoot: Node? = null,
-    arguments: List<IrExpression?> = listOf(),
+    extensionRoot: Node? = null
-    variables: List<IrTemporaryVariable> = listOf()
   ): IrExpression {
-    if (index >= roots.size) {
+    fun recursive(
-      val prefix = buildMessagePrefix(messageArgument, delegate.messageParameter, roots, original)
+      index: Int,
-        ?.deepCopyWithSymbols(parent)
+      dispatch: IrExpression?,
-      val diagram = irDiagramString(file, fileSource, prefix, original, variables)
+      extension: IrExpression?,
-      return delegate.buildCall(this, original, arguments, diagram)
+      arguments: List<IrExpression?>,
-    } else {
+      variables: List<IrTemporaryVariable>
-      val root = roots[index]
+    ): IrExpression {
-      if (root == null) {
+      if (index >= roots.size) {
-        val newArguments = arguments + original.getValueArgument(index)
+        val prefix = buildMessagePrefix(messageArgument, delegate.messageParameter, roots, call)
-        return diagram(original, delegate, messageArgument, roots, index + 1, newArguments, variables)
+          ?.deepCopyWithSymbols(parent)
+        val diagram = irDiagramString(sourceFile, prefix, call, variables)
+        return delegate.buildCall(this, call, dispatch, extension, arguments, diagram)
       } else {
-        return buildDiagramNesting(root) { argument, newVariables ->
+        val root = roots[index]
-          val newArguments = arguments + argument
+        if (root == null) {
-          diagram(original, delegate, messageArgument, roots, index + 1, newArguments, variables + newVariables)
+          val newArguments = arguments + call.getValueArgument(index)
+          return recursive(index + 1, dispatch, extension, newArguments, variables)
+        } else {
+          return buildDiagramNesting(root, variables) { argument, newVariables ->
+            val newArguments = arguments + argument
+            recursive(index + 1, dispatch, extension, newArguments, newVariables)
+          }
         }
       }
     }
+
+    return buildDiagramNestingNullable(dispatchRoot) { dispatch, newVariables ->
+      buildDiagramNestingNullable(extensionRoot, newVariables) { extension, newVariables ->
+        recursive(0, dispatch, extension, emptyList(), newVariables)
+      }
+    }
   }
 
   private fun DeclarationIrBuilder.buildMessagePrefix(
@@ -204,23 +224,32 @@ class PowerAssertCallTransformer(
     val possible = (context.referenceFunctions(function.kotlinFqName) + parentClassFunctions)
       .distinct()
 
-    return possible
+    return possible.mapNotNull { overload ->
-      .mapNotNull { overload ->
+      // Dispatch receivers must always match exactly
-        val parameters = overload.owner.valueParameters
+      if (function.dispatchReceiverParameter?.type != overload.owner.dispatchReceiverParameter?.type) {
-        if (parameters.size !in values.size..values.size + 1) return@mapNotNull null
+        return@mapNotNull null
-        if (!parameters.zip(values).all { (param, value) -> value.type.isAssignableTo(param.type) }) {
-          return@mapNotNull null
-        }
-
-        val messageParameter = parameters.last()
-        return@mapNotNull when {
-          isStringSupertype(messageParameter.type) -> SimpleFunctionDelegate(overload, messageParameter)
-          isStringFunction(messageParameter.type) -> LambdaFunctionDelegate(overload, messageParameter)
-          isStringJavaSupplierFunction(messageParameter.type) ->
-            SamConversionLambdaFunctionDelegate(overload, messageParameter)
-          else -> null
-        }
       }
+
+      // Extension receiver may only be assignable
+      if (!function.extensionReceiverParameter?.type.isAssignableTo(overload.owner.extensionReceiverParameter?.type)) {
+        return@mapNotNull null
+      }
+
+      val parameters = overload.owner.valueParameters
+      if (parameters.size !in values.size..values.size + 1) return@mapNotNull null
+      if (!parameters.zip(values).all { (param, value) -> value.type.isAssignableTo(param.type) }) {
+        return@mapNotNull null
+      }
+
+      val messageParameter = parameters.last()
+      return@mapNotNull when {
+        isStringSupertype(messageParameter.type) -> SimpleFunctionDelegate(overload, messageParameter)
+        isStringFunction(messageParameter.type) -> LambdaFunctionDelegate(overload, messageParameter)
+        isStringJavaSupplierFunction(messageParameter.type) ->
+          SamConversionLambdaFunctionDelegate(overload, messageParameter)
+        else -> null
+      }
+    }
   }
 
   private fun isStringFunction(type: IrType): Boolean =
@@ -238,10 +267,14 @@ class PowerAssertCallTransformer(
   private fun isStringSupertype(type: IrType): Boolean =
     context.irBuiltIns.stringType.isSubtypeOf(type, irTypeSystemContext)
 
-  private fun IrType.isAssignableTo(type: IrType): Boolean {
+  private fun IrType?.isAssignableTo(type: IrType?): Boolean {
-    if (isSubtypeOf(type, irTypeSystemContext)) return true
+    if (this != null && type != null) {
-    val superTypes = (type.classifierOrNull as? IrTypeParameterSymbol)?.owner?.superTypes
+      if (isSubtypeOf(type, irTypeSystemContext)) return true
-    return superTypes != null && superTypes.all { isSubtypeOf(it, irTypeSystemContext) }
+      val superTypes = (type.classifierOrNull as? IrTypeParameterSymbol)?.owner?.superTypes
+      return superTypes != null && superTypes.all { isSubtypeOf(it, irTypeSystemContext) }
+    } else {
+      return this == null && type == null
+    }
   }
 
   private fun MessageCollector.info(expression: IrElement, message: String) {
@@ -253,12 +286,6 @@ class PowerAssertCallTransformer(
   }
 
   private fun MessageCollector.report(expression: IrElement, severity: CompilerMessageSeverity, message: String) {
-    report(severity, message, expression.toCompilerMessageLocation())
+    report(severity, message, sourceFile.getCompilerMessageLocation(expression))
-  }
-
-  private fun IrElement.toCompilerMessageLocation(): CompilerMessageLocation {
-    val info = file.info(this)
-    val lineContent = fileSource.substring(this)
-    return CompilerMessageLocation.create(file.path, info.startLineNumber, info.startColumnNumber, lineContent)!!
   }
 }

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/PowerAssertIrGenerationExtension.kt

@@ -16,13 +16,12 @@
 
 package com.bnorm.power
 
+import com.bnorm.power.diagram.SourceFile
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
-import org.jetbrains.kotlin.ir.declarations.path
 import org.jetbrains.kotlin.name.FqName
-import java.io.File
 
 class PowerAssertIrGenerationExtension(
   private val messageCollector: MessageCollector,
@@ -30,10 +29,7 @@ class PowerAssertIrGenerationExtension(
 ) : IrGenerationExtension {
   override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
     for (file in moduleFragment.files) {
-      val fileSource = File(file.path).readText()
+      PowerAssertCallTransformer(SourceFile(file), pluginContext, messageCollector, functions)
-        .replace("\r\n", "\n") // https://youtrack.jetbrains.com/issue/KT-41888
-
-      PowerAssertCallTransformer(file, fileSource, pluginContext, messageCollector, functions)
         .visitFile(file)
     }
   }

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/FunctionDelegate.kt

@@ -33,26 +33,30 @@ interface FunctionDelegate {
   fun buildCall(
     builder: IrBuilderWithScope,
     original: IrCall,
-    arguments: List<IrExpression?>,
+    dispatchReceiver: IrExpression?,
-    message: IrExpression
+    extensionReceiver: IrExpression?,
+    valueArguments: List<IrExpression?>,
+    messageArgument: IrExpression
   ): IrExpression
 
   fun IrBuilderWithScope.irCallCopy(
     overload: IrSimpleFunctionSymbol,
     original: IrCall,
-    arguments: List<IrExpression?>,
+    dispatchReceiver: IrExpression?,
-    expression: IrExpression
+    extensionReceiver: IrExpression?,
+    valueArguments: List<IrExpression?>,
+    messageArgument: IrExpression
   ): IrExpression {
     return irCall(overload, type = original.type).apply {
-      dispatchReceiver = original.dispatchReceiver?.deepCopyWithSymbols(parent)
+      this.dispatchReceiver = original.dispatchReceiver?.deepCopyWithSymbols(parent)
-      extensionReceiver = original.extensionReceiver?.deepCopyWithSymbols(parent)
+      this.extensionReceiver = (extensionReceiver ?: original.extensionReceiver)?.deepCopyWithSymbols(parent)
       for (i in 0 until original.typeArgumentsCount) {
         putTypeArgument(i, original.getTypeArgument(i))
       }
-      for ((i, argument) in arguments.withIndex()) {
+      for ((i, argument) in valueArguments.withIndex()) {
         putValueArgument(i, argument?.deepCopyWithSymbols(parent))
       }
-      putValueArgument(arguments.size, expression.deepCopyWithSymbols(parent))
+      putValueArgument(valueArguments.size, messageArgument.deepCopyWithSymbols(parent))
     }
   }
 }

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/LambdaFunctionDelegate.kt

@@ -33,12 +33,21 @@ class LambdaFunctionDelegate(
   override fun buildCall(
     builder: IrBuilderWithScope,
     original: IrCall,
-    arguments: List<IrExpression?>,
+    dispatchReceiver: IrExpression?,
-    message: IrExpression
+    extensionReceiver: IrExpression?,
+    valueArguments: List<IrExpression?>,
+    messageArgument: IrExpression
   ): IrExpression = with(builder) {
     val expression = irLambda(context.irBuiltIns.stringType, messageParameter.type) {
-      +irReturn(message)
+      +irReturn(messageArgument)
     }
-    irCallCopy(overload, original, arguments, expression)
+    irCallCopy(
+      overload = overload,
+      original = original,
+      dispatchReceiver = dispatchReceiver,
+      extensionReceiver = extensionReceiver,
+      valueArguments = valueArguments,
+      messageArgument = expression
+    )
   }
 }

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/SamConversionLambdaFunctionDelegate.kt

@@ -19,11 +19,10 @@ package com.bnorm.power.delegate
 import com.bnorm.power.irLambda
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.irReturn
-import org.jetbrains.kotlin.ir.builders.typeOperator
+import org.jetbrains.kotlin.ir.builders.irSamConversion
 import org.jetbrains.kotlin.ir.declarations.IrValueParameter
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 
 class SamConversionLambdaFunctionDelegate(
@@ -35,13 +34,22 @@ class SamConversionLambdaFunctionDelegate(
   override fun buildCall(
     builder: IrBuilderWithScope,
     original: IrCall,
-    arguments: List<IrExpression?>,
+    dispatchReceiver: IrExpression?,
-    message: IrExpression
+    extensionReceiver: IrExpression?,
+    valueArguments: List<IrExpression?>,
+    messageArgument: IrExpression
   ): IrExpression = with(builder) {
     val lambda = irLambda(context.irBuiltIns.stringType, messageParameter.type) {
-      +irReturn(message)
+      +irReturn(messageArgument)
     }
-    val expression = typeOperator(messageParameter.type, lambda, IrTypeOperator.SAM_CONVERSION, messageParameter.type)
+    val expression = irSamConversion(lambda, messageParameter.type)
-    irCallCopy(overload, original, arguments, expression)
+    irCallCopy(
+      overload = overload,
+      original = original,
+      dispatchReceiver = dispatchReceiver,
+      extensionReceiver = extensionReceiver,
+      valueArguments = valueArguments,
+      messageArgument = expression
+    )
   }
 }

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/delegate/SimpleFunctionDelegate.kt

@@ -31,7 +31,16 @@ class SimpleFunctionDelegate(
   override fun buildCall(
     builder: IrBuilderWithScope,
     original: IrCall,
-    arguments: List<IrExpression?>,
+    dispatchReceiver: IrExpression?,
-    message: IrExpression
+    extensionReceiver: IrExpression?,
-  ): IrExpression = builder.irCallCopy(overload, original, arguments, message)
+    valueArguments: List<IrExpression?>,
+    messageArgument: IrExpression
+  ): IrExpression = builder.irCallCopy(
+    overload = overload,
+    original = original,
+    dispatchReceiver = dispatchReceiver,
+    extensionReceiver = extensionReceiver,
+    valueArguments = valueArguments,
+    messageArgument = messageArgument
+  )
 }

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/DiagramBuilder.kt

@@ -26,13 +26,22 @@ import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 
 fun IrBuilderWithScope.buildDiagramNesting(
   root: Node,
+  variables: List<IrTemporaryVariable> = emptyList(),
   call: IrBuilderWithScope.(IrExpression, List<IrTemporaryVariable>) -> IrExpression
 ): IrExpression {
-  return buildExpression(root, listOf()) { argument, subStack ->
+  return buildExpression(root, variables) { argument, subStack ->
     call(argument, subStack)
   }
 }
 
+fun IrBuilderWithScope.buildDiagramNestingNullable(
+  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)
+}
+
 private fun IrBuilderWithScope.buildExpression(
   node: Node,
   variables: List<IrTemporaryVariable>,

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/IrDiagram.kt

@@ -18,14 +18,13 @@ package com.bnorm.power.diagram
 
 import com.bnorm.power.irString
 import org.jetbrains.kotlin.ir.IrBuiltIns
-import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.SourceRangeInfo
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.irConcat
 import org.jetbrains.kotlin.ir.builders.irGet
-import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrConst
 import org.jetbrains.kotlin.ir.expressions.IrConstKind
 import org.jetbrains.kotlin.ir.expressions.IrExpression
@@ -36,20 +35,18 @@ import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
 import org.jetbrains.kotlin.ir.expressions.addArgument
 
 fun IrBuilderWithScope.irDiagramString(
-  file: IrFile,
+  sourceFile: SourceFile,
-  fileSource: String,
   prefix: IrExpression? = null,
-  original: IrExpression,
+  call: IrCall,
   variables: List<IrTemporaryVariable>
 ): IrExpression {
+  val callInfo = sourceFile.getSourceRangeInfo(call)
+  val callIndent = callInfo.startColumnNumber
 
-  val originalInfo = file.info(original)
+  val stackValues = variables.map { it.toValueDisplay(sourceFile, callIndent, callInfo) }
-  val callIndent = originalInfo.startColumnNumber
-
-  val stackValues = variables.map { it.toValueDisplay(fileSource, callIndent, file, originalInfo) }
 
   val valuesByRow = stackValues.groupBy { it.row }
-  val rows = fileSource.substring(original)
+  val rows = sourceFile.getText(callInfo)
     .replace("\n" + " ".repeat(callIndent), "\n") // Remove additional indentation
     .split("\n")
 
@@ -102,19 +99,17 @@ private data class ValueDisplay(
 )
 
 private fun IrTemporaryVariable.toValueDisplay(
-  fileSource: String,
+  fileSource: SourceFile,
   callIndent: Int,
-  file: IrFile,
   originalInfo: SourceRangeInfo
 ): ValueDisplay {
-  val source = fileSource.substring(original)
+  val info = fileSource.getSourceRangeInfo(original)
-    .replace("\n" + " ".repeat(callIndent), "\n") // Remove additional indentation
-
-  val info = file.info(original)
   var indent = info.startColumnNumber - callIndent
   var row = info.startLineNumber - originalInfo.startLineNumber
 
-  val columnOffset = findDisplayOffset(original, source)
+  val source = fileSource.getText(info)
+    .replace("\n" + " ".repeat(callIndent), "\n") // Remove additional indentation
+  val columnOffset = findDisplayOffset(fileSource, original, source)
 
   val prefix = source.substring(0, columnOffset)
   val rowShift = prefix.count { it == '\n' }
@@ -161,17 +156,19 @@ private fun IrTemporaryVariable.toValueDisplay(
  * ```
  */
 private fun findDisplayOffset(
+  sourceFile: SourceFile,
   expression: IrExpression,
   source: String
 ): Int {
   return when (expression) {
-    is IrMemberAccessExpression<*> -> memberAccessOffset(expression, source)
+    is IrMemberAccessExpression<*> -> memberAccessOffset(sourceFile, expression, source)
     is IrTypeOperatorCall -> typeOperatorOffset(expression, source)
     else -> 0
   }
 }
 
 private fun memberAccessOffset(
+  sourceFile: SourceFile,
   expression: IrMemberAccessExpression<*>,
   source: String
 ): Int {
@@ -193,12 +190,12 @@ private fun memberAccessOffset(
       ?: expression.extensionReceiver
       ?: expression.getValueArgument(0).takeIf { owner.origin == IrBuiltIns.BUILTIN_OPERATOR }
       ?: return 0
-    var offset = receiver.endOffset - expression.startOffset
+    val expressionInfo = sourceFile.getSourceRangeInfo(expression)
+    var offset = receiver.endOffset - expressionInfo.startOffset + 1
     if (receiver is IrConst<*> && receiver.kind == IrConstKind.String) offset++ // String constants don't include the quote
-    if (offset < 0 || offset >= source.length) return 0 // infix function called using non-infix syntax
 
     // Continue until there is a non-whitespace character
-    while (source[offset].isWhitespace()) {
+    while (source[offset].isWhitespace() || source[offset] == '.') {
       offset++
       if (offset >= source.length) return 0
     }
@@ -219,9 +216,6 @@ private fun typeOperatorOffset(
   }
 }
 
-fun String.substring(expression: IrElement) = substring(expression.startOffset, expression.endOffset)
-fun IrFile.info(expression: IrElement) = fileEntry.getSourceRangeInfo(expression.startOffset, expression.endOffset)
-
 fun StringBuilder.indent(indentation: Int): StringBuilder {
   repeat(indentation) { append(" ") }
   return this

kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/SourceFile.kt

@@ -0,0 +1,51 @@
+package com.bnorm.power.diagram
+
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageLocation
+import org.jetbrains.kotlin.ir.IrElement
+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 java.io.File
+
+data class SourceFile(
+  private val irFile: IrFile
+) {
+  private val source: String = File(irFile.path).readText()
+    .replace("\r\n", "\n") // https://youtrack.jetbrains.com/issue/KT-41888
+
+  fun getSourceRangeInfo(element: IrElement): SourceRangeInfo {
+    var range = element.startOffset..element.endOffset
+    when (element) {
+      is IrCall -> {
+        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
+
+          // 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)
+          if (substring.startsWith('(') && substring.endsWith(')')) {
+            range = receiver.startOffset - 1..element.endOffset
+          }
+        }
+      }
+    }
+    return irFile.fileEntry.getSourceRangeInfo(range.first, range.last)
+  }
+
+  fun getText(info: SourceRangeInfo): String {
+    return safeSubstring(info.startOffset, info.endOffset)
+  }
+
+  private fun safeSubstring(start: Int, end: Int): String =
+    source.substring(maxOf(start, 0), minOf(end, source.length))
+
+  fun getCompilerMessageLocation(element: IrElement): CompilerMessageLocation {
+    val info = getSourceRangeInfo(element)
+    val lineContent = getText(info)
+    return CompilerMessageLocation.create(irFile.path, info.startLineNumber, info.startColumnNumber, lineContent)!!
+  }
+}

kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/InfixFunctionTest.kt

@@ -0,0 +1,369 @@
+/*
+ * Copyright (C) 2022 Brian Norman
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.bnorm.power
+
+import com.tschuchort.compiletesting.KotlinCompilation
+import com.tschuchort.compiletesting.SourceFile
+import org.jetbrains.kotlin.name.FqName
+import java.lang.reflect.InvocationTargetException
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.fail
+
+class InfixFunctionTest {
+  @Test
+  fun `extension infix function call includes receiver`() {
+    val actual = runExtensionInfix(
+      """
+      (1 + 1) mustEqual (2 + 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      (1 + 1) mustEqual (2 + 4)
+         |                 |
+         |                 6
+         2
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `extension infix function call with constant receiver`() {
+    val actual = runExtensionInfix(
+      """
+      1 mustEqual (2 + 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      1 mustEqual (2 + 4)
+                     |
+                     6
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `extension infix function call with constant parameter`() {
+    val actual = runExtensionInfix(
+      """
+      (1 + 1) mustEqual 6
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      (1 + 1) mustEqual 6
+         |
+         2
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `extension infix function call with only constants`() {
+    val actual = runExtensionInfix(
+      """
+      2 mustEqual 6
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Assertion failed
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `extension non-infix function call includes receiver`() {
+    val actual = runExtensionInfix(
+      """
+      (1 + 1).mustEqual(2 + 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      (1 + 1).mustEqual(2 + 4)
+         |                |
+         |                6
+         2
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `extension non-infix function call with constant receiver`() {
+    val actual = runExtensionInfix(
+      """
+      1.mustEqual(2 + 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      1.mustEqual(2 + 4)
+                    |
+                    6
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `extension non-infix function call with constant parameter`() {
+    val actual = runExtensionInfix(
+      """
+      (1 + 1).mustEqual(6)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      (1 + 1).mustEqual(6)
+         |
+         2
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `extension non-infix function call with only constants`() {
+    val actual = runExtensionInfix(
+      """
+      2.mustEqual(6)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Assertion failed
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch infix function call includes receiver`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(1 + 1) mustEqual (2 + 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(1 + 1) mustEqual (2 + 4)
+      |         |                 |
+      |         |                 6
+      |         2
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch infix function call with constant receiver`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(1) mustEqual (2 + 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(1) mustEqual (2 + 4)
+      |                       |
+      |                       6
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch infix function call with constant parameter`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(1 + 1) mustEqual 6
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(1 + 1) mustEqual 6
+      |         |
+      |         2
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch infix function call with only constants`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(2) mustEqual 6
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(2) mustEqual 6
+      |
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch non-infix function call includes receiver`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(1 + 1).mustEqual(2 + 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(1 + 1).mustEqual(2 + 4)
+      |         |                |
+      |         |                6
+      |         2
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch non-infix function call with constant receiver`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(1).mustEqual(2 + 4)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(1).mustEqual(2 + 4)
+      |                      |
+      |                      6
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch non-infix function call with constant parameter`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(1 + 1).mustEqual(6)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(1 + 1).mustEqual(6)
+      |         |
+      |         2
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  @Test
+  fun `dispatch non-infix function call with only constants`() {
+    val actual = runDispatchInfix(
+      """
+      Wrapper(2).mustEqual(6)
+      """.trimIndent()
+    )
+    assertEquals(
+      """
+      Wrapper(2).mustEqual(6)
+      |
+      Wrapper
+      """.trimIndent(),
+      actual.trim()
+    )
+  }
+
+  private fun runExtensionInfix(mainBody: String): String {
+    return run(
+      SourceFile.kotlin(
+        name = "main.kt",
+        contents = """
+        infix fun <V> V.mustEqual(expected: V): Unit = assert(this == expected)
+        
+        fun <V> V.mustEqual(expected: V, message: () -> String): Unit =
+          assert(this == expected, message)
+        
+        fun main() {
+          $mainBody
+        }
+        """.trimIndent(),
+        trimIndent = false
+      ),
+      setOf(FqName("mustEqual"))
+    )
+  }
+
+  private fun runDispatchInfix(mainBody: String): String {
+    return run(
+      SourceFile.kotlin(
+        name = "main.kt",
+        contents = """
+        class Wrapper<V>(
+          private val value: V
+        ) {
+          infix fun mustEqual(expected: V): Unit = assert(value == expected)
+        
+          fun mustEqual(expected: V, message: () -> String): Unit =
+            assert(value == expected, message)
+        
+          override fun toString() = "Wrapper"
+        }
+        
+        fun main() {
+          $mainBody
+        }
+        """.trimIndent(),
+        trimIndent = false
+      ),
+      setOf(FqName("Wrapper.mustEqual"))
+    )
+  }
+
+  private fun run(file: SourceFile, fqNames: Set<FqName>): String {
+    val result = compile(listOf(file), PowerAssertComponentRegistrar(fqNames))
+    assertEquals(KotlinCompilation.ExitCode.OK, result.exitCode, "Failed with messages: " + result.messages)
+
+    val kClazz = result.classLoader.loadClass("MainKt")
+    val main = kClazz.declaredMethods.single { it.name == "main" && it.parameterCount == 0 }
+    try {
+      try {
+        main.invoke(null)
+      } catch (t: InvocationTargetException) {
+        throw t.cause!!
+      }
+      fail("should have thrown assertion")
+    } catch (t: Throwable) {
+      return t.message ?: ""
+    }
+  }
+}