From 7ef97177aad0097a99373c2b2fce01b93a2491f4 Mon Sep 17 00:00:00 2001
From: Brian Norman <brian@bnorm.com>
Date: Mon, 10 Apr 2023 15:40:51 -0500
Subject: [PATCH] Correctly align contains operators (#96)

---
 .../com/bnorm/power/diagram/ExpressionTree.kt |  5 ++
 .../com/bnorm/power/diagram/IrDiagram.kt      | 10 ++-
 .../kotlin/com/bnorm/power/OperatorTest.kt    | 90 +++++++++++++++++++
 3 files changed, 103 insertions(+), 2 deletions(-)
 create mode 100644 kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/OperatorTest.kt

diff --git a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/ExpressionTree.kt b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/ExpressionTree.kt
index b5e735fd634..f5e95357130 100644
--- a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/ExpressionTree.kt
+++ b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/ExpressionTree.kt
@@ -114,6 +114,11 @@ fun buildTree(expression: IrExpression): Node? {
         if (expression.symbol.owner.name.asString() == "EQEQ" && expression.origin == IrStatementOrigin.EXCLEQ) {
           // Skip the EQEQ part of a EXCLEQ call
           expression.acceptChildren(this, data)
+        } else if (expression.origin == IrStatementOrigin.NOT_IN) {
+          // Exclude the wrapped "contains" call for `!in` operator expressions and only display the final result
+          val node = data as? ExpressionNode ?: ExpressionNode().also { data.addChild(it) }
+          node.add(expression)
+          expression.dispatchReceiver!!.acceptChildren(this, node)
         } else {
           super.visitCall(expression, data)
         }
diff --git a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/IrDiagram.kt b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/IrDiagram.kt
index 0fd1dfed5f8..96c80ab28c2 100644
--- a/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/IrDiagram.kt
+++ b/kotlin-power-assert-plugin/src/main/kotlin/com/bnorm/power/diagram/IrDiagram.kt
@@ -172,9 +172,14 @@ private fun memberAccessOffset(
   expression: IrMemberAccessExpression<*>,
   source: String,
 ): Int {
-  if (expression.origin == IrStatementOrigin.EXCLEQ || expression.origin == IrStatementOrigin.EXCLEQEQ) {
+  when (expression.origin) {
     // special case to handle `value != null`
-    return source.indexOf("!=")
+    IrStatementOrigin.EXCLEQ, IrStatementOrigin.EXCLEQEQ -> return source.indexOf("!=")
+    // special case to handle `in` operator
+    IrStatementOrigin.IN -> return source.indexOf(" in ") + 1
+    // special case to handle `in` operator
+    IrStatementOrigin.NOT_IN -> return source.indexOf(" !in ") + 1
+    else -> Unit
   }
 
   val owner = expression.symbol.owner
@@ -193,6 +198,7 @@ private fun memberAccessOffset(
     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() || source[offset] == '.') {
diff --git a/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/OperatorTest.kt b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/OperatorTest.kt
new file mode 100644
index 00000000000..dbfb11d4056
--- /dev/null
+++ b/kotlin-power-assert-plugin/src/test/kotlin/com/bnorm/power/OperatorTest.kt
@@ -0,0 +1,90 @@
+/*
+ * Copyright (C) 2023 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 kotlin.test.Test
+
+class OperatorTest {
+  @Test
+  fun `contains operator is correctly aligned`() {
+    assertMessage(
+      """
+      fun main() {
+        assert("Name" in listOf("Hello", "World"))
+      }""",
+      """
+      Assertion failed
+      assert("Name" in listOf("Hello", "World"))
+                    |  |
+                    |  [Hello, World]
+                    false
+      """.trimIndent(),
+    )
+  }
+
+  @Test
+  fun `contains function is correctly aligned`() {
+    assertMessage(
+      """
+      fun main() {
+        assert(listOf("Hello", "World").contains("Name"))
+      }""",
+      """
+      Assertion failed
+      assert(listOf("Hello", "World").contains("Name"))
+             |                        |
+             |                        false
+             [Hello, World]
+      """.trimIndent(),
+    )
+  }
+
+  @Test
+  fun `negative contains operator is correctly aligned`() {
+    assertMessage(
+      """
+      fun main() {
+        assert("Hello" !in listOf("Hello", "World"))
+      }""",
+      """
+      Assertion failed
+      assert("Hello" !in listOf("Hello", "World"))
+                     |   |
+                     |   [Hello, World]
+                     false
+      """.trimIndent(),
+    )
+  }
+
+  @Test
+  fun `negative contains function is correctly aligned`() {
+    assertMessage(
+      """
+      fun main() {
+        assert(!listOf("Hello", "World").contains("Hello"))
+      }""",
+      """
+      Assertion failed
+      assert(!listOf("Hello", "World").contains("Hello"))
+             ||                        |
+             ||                        true
+             |[Hello, World]
+             false
+      """.trimIndent(),
+    )
+  }
+}