feature: support VarHandle#withInvoke(Exact)Behavior methods (#153)

src/main/kotlin/de/sirywell/handlehints/dfa/SsaAnalyzer.kt

@@ -36,6 +36,7 @@ class SsaAnalyzer(private val controlFlow: ControlFlow, val typeData: TypeData)
     private val methodHandlesMerger = MethodHandlesMerger(this)
     private val methodHandlesInitializer = MethodHandlesInitializer(this)
     private val methodHandleTransformer = MethodHandleTransformer(this)
+    private val varHandleHelper = VarHandleHelper(this)
     private val lookupHelper = LookupHelper(this)
     private val methodTypeHelper = MethodTypeHelper(this)
 
@@ -262,6 +263,8 @@ class SsaAnalyzer(private val controlFlow: ControlFlow, val typeData: TypeData)
                     in objectMethods -> unrelatedType()
                     else -> warnUnsupported(expression, "MethodHandle")
                 }
+            } else if (receiverIsVarHandle(expression)) {
+                return varHandle(expression, arguments, block)
             } else if (receiverIsLookup(expression)) {
                 return lookup(expression, arguments, block)
             } else if (receiverIsMemoryLayout(expression)) {
@@ -510,6 +513,25 @@ class SsaAnalyzer(private val controlFlow: ControlFlow, val typeData: TypeData)
         return clazz.qualifiedName.equals("java.lang.foreign.ValueLayout")
     }
 
+    private fun varHandle(
+        expression: PsiMethodCallExpression,
+        arguments: List<PsiExpression>,
+        block: Block
+    ) : TypeLatticeElement<*>? {
+        val qualifier = expression.methodExpression.qualifierExpression
+        return when (expression.methodName) {
+            "withInvokeBehavior" -> {
+                if (arguments.isNotEmpty() || qualifier == null) noMatch()
+                else varHandleHelper.withInvokeBehavior(qualifier, block)
+            }
+            "withInvokeExactBehavior" -> {
+                if (arguments.isNotEmpty() || qualifier == null) noMatch()
+                else varHandleHelper.withInvokeExactBehavior(qualifier, block)
+            }
+            else -> noMatch()
+        }
+    }
+
     private fun lookup(
         expression: PsiMethodCallExpression,
         arguments: List<PsiExpression>,
@@ -825,6 +847,13 @@ class SsaAnalyzer(private val controlFlow: ControlFlow, val typeData: TypeData)
         return type
     }
 
+    @JvmName("varHandleType_extension")
+    private fun PsiExpression.varHandleType(block: Block): VarHandleType? {
+        val type = resolveType(this, block) as? VarHandleType ?: return null
+        typeData[this] = type
+        return type
+    }
+
     @JvmName("memoryLayoutType_extension")
     private fun PsiExpression.memoryLayoutType(block: Block): MemoryLayoutType? {
         val type = resolveType(this, block) as? MemoryLayoutType ?: return null
@@ -850,6 +879,8 @@ class SsaAnalyzer(private val controlFlow: ControlFlow, val typeData: TypeData)
 
     fun methodHandleType(expression: PsiExpression, block: Block) = expression.methodHandleType(block)
 
+    fun varHandleType(expression: PsiExpression, block: Block) = expression.varHandleType(block)
+
     fun memoryLayoutType(expression: PsiExpression, block: Block) = expression.memoryLayoutType(block)
 
     fun pathElementType(expression: PsiExpression, block: Block) = expression.pathElementType(block)

src/main/kotlin/de/sirywell/handlehints/foreign/MemoryLayoutHelper.kt

@@ -314,19 +314,19 @@ class MemoryLayoutHelper(private val ssaAnalyzer: SsaAnalyzer) : ProblemEmitter(
             path: List<IndexedValue<PathElementType>>,
             coords: MutableList<Type>
         ): VarHandleType {
-            return CompleteVarHandleType(TopType, IncompleteTypeList(coords.toIndexedMap()))
+            return CompleteVarHandleType(TopType, IncompleteTypeList(coords.toIndexedMap()), KnownInvocationBehavior.INVOKE)
         }
 
         private fun onComplete(
             layoutType: ValueLayoutType,
             coords: MutableList<Type>
         ): VarHandleType {
             if (layoutType is NormalValueLayoutType) {
-                return CompleteVarHandleType(layoutType.type, CompleteTypeList(coords))
+                return CompleteVarHandleType(layoutType.type, CompleteTypeList(coords), KnownInvocationBehavior.INVOKE)
             } else {
                 val ctx = contextElement(-1) // good enough for us
                 val type = findPsiType("java.lang.foreign.MemorySegment", ctx)
-                return CompleteVarHandleType(ExactType(type), CompleteTypeList(coords))
+                return CompleteVarHandleType(ExactType(type), CompleteTypeList(coords), KnownInvocationBehavior.INVOKE)
             }
         }
     }

src/main/kotlin/de/sirywell/handlehints/mhtype/LookupHelper.kt

@@ -135,12 +135,12 @@ class LookupHelper(private val ssaAnalyzer: SsaAnalyzer) : ProblemEmitter(ssaAna
     fun findStaticVarHandle(declExpr: PsiExpression, typeExpr: PsiExpression): VarHandleType {
         declExpr.asReferenceType() // problem reporting only, not needed otherwise
         val type = typeExpr.asNonVoidType()
-        return CompleteVarHandleType(type, CompleteTypeList(listOf()))
+        return CompleteVarHandleType(type, CompleteTypeList(listOf()), KnownInvocationBehavior.INVOKE)
     }
 
     fun findVarHandle(recvExpr: PsiExpression, typeExpr: PsiExpression): VarHandleType {
         val recv = recvExpr.asReferenceType()
         val type = typeExpr.asNonVoidType()
-        return CompleteVarHandleType(type, CompleteTypeList(listOf(recv)))
+        return CompleteVarHandleType(type, CompleteTypeList(listOf(recv)), KnownInvocationBehavior.INVOKE)
     }
 }

src/main/kotlin/de/sirywell/handlehints/mhtype/MethodHandlesInitializer.kt

@@ -51,7 +51,11 @@ class MethodHandlesInitializer(private val ssaAnalyzer: SsaAnalyzer) : ProblemEm
     fun arrayElementVarHandle(arrayClass: PsiExpression): VarHandleType {
         val arrayType = arrayClass.asArrayType()
         val componentType = getComponentType(arrayType)
-        return CompleteVarHandleType(componentType, CompleteTypeList(listOf(arrayType, ExactType.intType)))
+        return CompleteVarHandleType(
+            componentType,
+            CompleteTypeList(listOf(arrayType, ExactType.intType)),
+            KnownInvocationBehavior.INVOKE
+        )
     }
 
     private fun getComponentType(arrayType: Type) = if (arrayType !is ExactType) {
@@ -79,7 +83,11 @@ class MethodHandlesInitializer(private val ssaAnalyzer: SsaAnalyzer) : ProblemEm
         if (componentType is ExactType && !isSupportedViewHandleComponentType(componentType.psiType)) {
             componentType = emitMustBeViewHandleSupportedComponentType(arrayClass, componentType)
         }
-        return CompleteVarHandleType(componentType, CompleteTypeList(listOf(collectionType, ExactType.intType)))
+        return CompleteVarHandleType(
+            componentType,
+            CompleteTypeList(listOf(collectionType, ExactType.intType)),
+            KnownInvocationBehavior.INVOKE
+        )
     }
 
     private fun isSupportedViewHandleComponentType(type: PsiType): Boolean {

src/main/kotlin/de/sirywell/handlehints/mhtype/VarHandleHelper.kt

@@ -0,0 +1,42 @@
+package de.sirywell.handlehints.mhtype
+
+import com.intellij.psi.PsiExpression
+import com.intellij.psi.PsiMethodCallExpression
+import com.intellij.psi.util.parentOfType
+import de.sirywell.handlehints.MethodHandleBundle.message
+import de.sirywell.handlehints.dfa.SsaAnalyzer
+import de.sirywell.handlehints.dfa.SsaConstruction
+import de.sirywell.handlehints.inspection.ProblemEmitter
+import de.sirywell.handlehints.inspection.RedundantInvocationFix
+import de.sirywell.handlehints.type.BotVarHandleType
+import de.sirywell.handlehints.type.KnownInvocationBehavior
+import de.sirywell.handlehints.type.KnownInvocationBehavior.INVOKE
+import de.sirywell.handlehints.type.KnownInvocationBehavior.INVOKE_EXACT
+import de.sirywell.handlehints.type.VarHandleType
+
+class VarHandleHelper(private val ssaAnalyzer: SsaAnalyzer) : ProblemEmitter(ssaAnalyzer.typeData) {
+    fun withInvokeBehavior(targetExpr: PsiExpression, block: SsaConstruction.Block) : VarHandleType {
+        return withBehavior(targetExpr, block, INVOKE)
+    }
+
+    fun withInvokeExactBehavior(targetExpr: PsiExpression, block: SsaConstruction.Block) : VarHandleType {
+        return withBehavior(targetExpr, block, INVOKE_EXACT)
+    }
+
+    private fun withBehavior(
+        targetExpr: PsiExpression,
+        block: SsaConstruction.Block,
+        behavior: KnownInvocationBehavior
+    ): VarHandleType {
+        val varHandleType = ssaAnalyzer.varHandleType(targetExpr, block) ?: BotVarHandleType
+        val invocationBehavior = varHandleType.invocationBehavior
+        if (invocationBehavior == behavior) {
+            emitRedundant(
+                targetExpr.parentOfType<PsiMethodCallExpression>()!!,
+                message("problem.transforming.varHandleInvokeBehavior.redundant", behavior.readableName),
+                RedundantInvocationFix()
+            )
+        }
+        return varHandleType.withInvokeBehavior(behavior)
+    }
+}

src/main/kotlin/de/sirywell/handlehints/presentation/TypePrinter.kt

@@ -287,4 +287,16 @@ class TypePrinter : TypeVisitor<TypePrinter.PrintContext, Unit> {
         context.append("⊥")
     }
 
+    override fun visit(type: KnownInvocationBehavior, context: PrintContext) {
+
+    }
+
+    override fun visit(type: TopInvocationBehavior, context: PrintContext) {
+
+    }
+
+    override fun visit(type: BotInvocationBehavior, context: PrintContext) {
+
+    }
+
 }

src/main/kotlin/de/sirywell/handlehints/psiSupport.kt

@@ -27,6 +27,9 @@ fun receiverIsInvokeClass(element: PsiMethodCallExpression, className: String) =
 fun receiverIsMethodHandle(element: PsiMethodCallExpression) =
     receiverIsInvokeClass(element, "MethodHandle")
 
+fun receiverIsVarHandle(element: PsiMethodCallExpression) =
+    receiverIsInvokeClass(element, "VarHandle")
+
 fun receiverIsMethodHandles(element: PsiMethodCallExpression) =
     receiverIsInvokeClass(element, "MethodHandles")
 

src/main/kotlin/de/sirywell/handlehints/type/TypeVisitor.kt

@@ -42,4 +42,7 @@ interface TypeVisitor<C, R> {
     fun visit(type: CompleteFunctionDescriptorType, context: C): R
     fun visit(type: TopFunctionDescriptorType, context: C): R
     fun visit(type: BotFunctionDescriptorType, context: C): R
+    fun visit(type: KnownInvocationBehavior, context: C): R
+    fun visit(type: TopInvocationBehavior, context: C): R
+    fun visit(type: BotInvocationBehavior, context: C): R
 }

src/main/kotlin/de/sirywell/handlehints/type/VarHandleType.kt

@@ -6,37 +6,81 @@ import de.sirywell.handlehints.TriState
 interface VarHandleType : TypeLatticeElement<VarHandleType> {
     val variableType: Type
     val coordinateTypes: TypeList
+    val invocationBehavior: InvocationBehavior
+    fun withInvokeBehavior(behavior: InvocationBehavior): VarHandleType
+}
+
+sealed interface InvocationBehavior : TypeLatticeElement<InvocationBehavior>
+data object BotInvocationBehavior : InvocationBehavior, BotTypeLatticeElement<InvocationBehavior> {
+    override fun <C, R> accept(visitor: TypeVisitor<C, R>, context: C): R {
+        return visitor.visit(this, context)
+    }
+}
+
+data object TopInvocationBehavior : InvocationBehavior, TopTypeLatticeElement<InvocationBehavior> {
+    override fun self() = this
+
+    override fun <C, R> accept(visitor: TypeVisitor<C, R>, context: C): R {
+        return visitor.visit(this, context)
+    }
+}
+
+enum class KnownInvocationBehavior(val readableName: String) : InvocationBehavior {
+    INVOKE("invoke"),
+    INVOKE_EXACT("invokeExact");
+
+    override fun joinIdentical(other: InvocationBehavior): Pair<InvocationBehavior, TriState> {
+        return when (other) {
+            is BotInvocationBehavior -> this to TriState.UNKNOWN
+            is TopInvocationBehavior -> TopInvocationBehavior to TriState.UNKNOWN
+            else -> if (this == other) this to TriState.YES else TopInvocationBehavior to TriState.NO
+        }
+    }
+
+    override fun <C, R> accept(visitor: TypeVisitor<C, R>, context: C): R {
+        return visitor.visit(this, context)
+    }
 }
 
 data object BotVarHandleType : VarHandleType, BotTypeLatticeElement<VarHandleType> {
+    override fun withInvokeBehavior(behavior: InvocationBehavior) =
+        CompleteVarHandleType(this.variableType, this.coordinateTypes, behavior)
+    override fun <C, R> accept(visitor: TypeVisitor<C, R>, context: C) = visitor.visit(this, context)
+
     override val variableType = BotType
     override val coordinateTypes = BotTypeList
-    override fun <C, R> accept(visitor: TypeVisitor<C, R>, context: C) = visitor.visit(this, context)
+    override val invocationBehavior = BotInvocationBehavior
 }
 
 data object TopVarHandleType : VarHandleType, TopTypeLatticeElement<VarHandleType> {
     override fun self() = this
     override fun <C, R> accept(visitor: TypeVisitor<C, R>, context: C) = visitor.visit(this, context)
-
+    override fun withInvokeBehavior(behavior: InvocationBehavior) = this
     override val variableType = TopType
     override val coordinateTypes = TopTypeList
+    override val invocationBehavior = TopInvocationBehavior
 }
 
 data class CompleteVarHandleType(
     override val variableType: Type,
-    override val coordinateTypes: TypeList
+    override val coordinateTypes: TypeList,
+    override val invocationBehavior: InvocationBehavior
 ) : VarHandleType {
+    override fun withInvokeBehavior(behavior: InvocationBehavior) =
+        CompleteVarHandleType(this.variableType, this.coordinateTypes, behavior)
+
     override fun joinIdentical(other: VarHandleType): Pair<VarHandleType, TriState> {
         val (vt, identicalVt) = variableType.joinIdentical(other.variableType)
         val (ct, identicalCt) = coordinateTypes.joinIdentical(other.coordinateTypes)
-        val identical = identicalVt.sharpenTowardsNo(identicalCt)
+        val (ib, identicalIb) = invocationBehavior.joinIdentical(other.invocationBehavior)
+        val identical = identicalVt.sharpenTowardsNo(identicalCt).sharpenTowardsNo(identicalIb)
         if (vt == TopType && ct == TopTypeList) {
             return TopVarHandleType to identical
         }
         if (vt == BotType && ct == BotTypeList) {
             return BotVarHandleType to identical
         }
-        return CompleteVarHandleType(vt, ct) to identical
+        return CompleteVarHandleType(vt, ct, ib) to identical
     }
 
     override fun <C, R> accept(visitor: TypeVisitor<C, R>, context: C) = visitor.visit(this, context)

src/main/resources/messages/MethodHandleMessages.properties

@@ -65,3 +65,4 @@ problem.general.invocation.append=Append method call ''{0}()''
 problem.general.invocation.wrap=Wrap with call to ''{0}''
 problem.general.array.dimension.add=Add one dimension to the array type
 problem.general.array.unsupportedViewHandleComponentType=Expected one of ''short'', ''char'', ''int'', ''long'', ''float'', ''double'', got ''{0}''
+problem.transforming.varHandleInvokeBehavior.redundant=VarHandle already has ''{0}'' behavior

src/test/kotlin/de/sirywell/handlehints/mhtype/VarHandleInspectionsTest.kt

@@ -0,0 +1,7 @@
+package de.sirywell.handlehints.mhtype
+
+class VarHandleInspectionsTest : TypeAnalysisTestBase() {
+
+    fun testVarHandleWithInvokeBehavior() = doInspectionAndTypeCheckingTest()
+
+}

src/test/testData/VarHandleWithInvokeBehavior.java

@@ -0,0 +1,17 @@
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
+
+class VarHandleWithInvokeBehavior {
+
+  void invokeBehavior() {
+    <info descr="(double[],int)(double)">VarHandle vh0 = <info descr="(double[],int)(double)">MethodHandles.arrayElementVarHandle(double[].class)</info>;</info>
+    // redundant - vh0 has invoke behavior
+    <info descr="(double[],int)(double)">VarHandle vh1 = <warning descr="VarHandle already has 'invoke' behavior"><info descr="(double[],int)(double)">vh0.withInvokeBehavior()</info></warning>;</info>
+    // not redundant
+    <info descr="(double[],int)(double)">VarHandle vh2 = <info descr="(double[],int)(double)">vh0.withInvokeExactBehavior()</info>;</info>
+    // redundant - vh2 has invokeExact behavior
+    <info descr="(double[],int)(double)">VarHandle vh3 = <warning descr="VarHandle already has 'invokeExact' behavior"><info descr="(double[],int)(double)">vh2.withInvokeExactBehavior()</info></warning>;</info>
+    // not redundant
+    <info descr="(double[],int)(double)">VarHandle vh4 = <info descr="(double[],int)(double)">vh2.withInvokeBehavior()</info>;</info>
+  }
+}