From 8f9bb75e606296015b412aea99989857d3c60213 Mon Sep 17 00:00:00 2001
From: Brian Norman <brian.norman@jetbrains.com>
Date: Tue, 16 Jan 2024 10:32:21 -0600
Subject: [PATCH] [FIR] Allow identity equality to be used in contract
 implications

^KT-63256 Fixed
---
 .../fir/resolve/transformers/contracts/ConeEffectExtractor.kt | 2 ++
 .../tests/smartCasts/varnotnull/equalityAndIdentity.fir.kt    | 4 ++--
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/ConeEffectExtractor.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/ConeEffectExtractor.kt
index e60dd61960a..6d2ea05c4b2 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/ConeEffectExtractor.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/ConeEffectExtractor.kt
@@ -124,6 +124,8 @@ class ConeEffectExtractor(
         val isNegated = when (val operation = equalityOperatorCall.operation) {
             FirOperation.EQ -> false
             FirOperation.NOT_EQ -> true
+            FirOperation.IDENTITY -> false
+            FirOperation.NOT_IDENTITY -> true
             else -> return ConeContractDescriptionError.IllegalEqualityOperator(operation).asElement()
         }
 
diff --git a/compiler/testData/diagnostics/tests/smartCasts/varnotnull/equalityAndIdentity.fir.kt b/compiler/testData/diagnostics/tests/smartCasts/varnotnull/equalityAndIdentity.fir.kt
index e166d1926e5..fe93b85d949 100644
--- a/compiler/testData/diagnostics/tests/smartCasts/varnotnull/equalityAndIdentity.fir.kt
+++ b/compiler/testData/diagnostics/tests/smartCasts/varnotnull/equalityAndIdentity.fir.kt
@@ -5,7 +5,7 @@
 import kotlin.contracts.*
 
 fun <E> checkEquality(arg: E?) { contract { returns() implies (arg != null) } }
-fun <E> checkIdentity(arg: E?) { contract { <!ERROR_IN_CONTRACT_DESCRIPTION("NOT_IDENTITY operator call is illegal in contract description")!>returns() implies (arg !== null)<!> } }
+fun <E> checkIdentity(arg: E?) { contract { returns() implies (arg !== null) } }
 fun checkTrue(statement: Boolean) { contract { returns() implies statement } }
 fun checkFalse(statement: Boolean) { contract { returns() implies !statement } }
 
@@ -18,7 +18,7 @@ fun test(a: String?) {
     }
     a.let {
         checkIdentity(it)
-        consume(<!ARGUMENT_TYPE_MISMATCH!>it<!>)
+        consume(it)
     }
     a.let {
         checkTrue(it != null)