From 53bec88d9b4153618b6854fabdfc320dd233e0f1 Mon Sep 17 00:00:00 2001
From: Brian Norman <brian@bnorm.com>
Date: Thu, 28 Jan 2021 23:05:28 -0600
Subject: [PATCH] Fix extension infix function diagrams

---
 .../kotlin/com/bnorm/power/IrStackVariable.kt | 20 ++++++++---
 .../kotlin/com/bnorm/power/RegexMatchTest.kt  | 36 +++++++++++++++++++
 2 files changed, 51 insertions(+), 5 deletions(-)
 create mode 100644 kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/RegexMatchTest.kt

diff --git a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/IrStackVariable.kt b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/IrStackVariable.kt
index aa0a5b4a259..cfe7b25ca63 100644
--- a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/IrStackVariable.kt
+++ b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/IrStackVariable.kt
@@ -25,6 +25,8 @@ 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.descriptors.IrBuiltIns
+import org.jetbrains.kotlin.ir.expressions.IrConst
+import org.jetbrains.kotlin.ir.expressions.IrConstKind
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
 import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
@@ -175,15 +177,23 @@ private fun findDisplayOffset(
 
   if (owner.isInfix || owner.isOperator || owner.origin == IrBuiltIns.BUILTIN_OPERATOR) {
     // Ignore single value operators
-    if (expression.dispatchReceiver != null && expression.valueArgumentsCount == 0) return 0
+    val singleReceiver = (expression.dispatchReceiver != null) xor (expression.extensionReceiver != null)
+    if (singleReceiver && expression.valueArgumentsCount == 0) return 0
 
     // Start after the dispatcher or first argument
-    val dispatcher = expression.dispatchReceiver ?: expression.getValueArgument(0) ?: return 0
-    var offset = dispatcher.endOffset - expression.startOffset
-    if (offset < 0) return 0 // infix function called using non-infix syntax
+    val receiver = expression.dispatchReceiver
+      ?: expression.extensionReceiver
+      ?: expression.getValueArgument(0).takeIf { owner.origin == IrBuiltIns.BUILTIN_OPERATOR }
+      ?: return 0
+    var offset = receiver.endOffset - expression.startOffset
+    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()) offset++
+    while (source[offset].isWhitespace()) {
+      offset++
+      if (offset >= source.length) return 0
+    }
     return offset
   }
 
diff --git a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/RegexMatchTest.kt b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/RegexMatchTest.kt
new file mode 100644
index 00000000000..ac40b07389e
--- /dev/null
+++ b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/RegexMatchTest.kt
@@ -0,0 +1,36 @@
+package com.bnorm.power
+
+import org.junit.Test
+import kotlin.test.assertEquals
+
+class RegexMatchTest {
+  @Test
+  fun `regex matches`() {
+    val actual = executeMainAssertion("""assert("Hello, World".matches("[A-Za-z]+".toRegex()))""")
+    assertEquals(
+      """
+      Assertion failed
+      assert("Hello, World".matches("[A-Za-z]+".toRegex()))
+                            |                   |
+                            |                   [A-Za-z]+
+                            false
+      """.trimIndent(),
+      actual
+    )
+  }
+
+  @Test
+  fun `infix regex matches`() {
+    val actual = executeMainAssertion("""assert("Hello, World" matches "[A-Za-z]+".toRegex())""")
+    assertEquals(
+      """
+      Assertion failed
+      assert("Hello, World" matches "[A-Za-z]+".toRegex())
+                            |                   |
+                            |                   [A-Za-z]+
+                            false
+      """.trimIndent(),
+      actual
+    )
+  }
+}