feature: support MethodHandles#byte(Array|Buffer)ViewVarHandle methods (

#148)
SirYwell · Dec 6, 2024 · 3aa0615 · 3aa0615
1 parent d93a1b5
commit 3aa0615
Show file tree

Hide file tree

Showing 14 changed files with 209 additions and 44 deletions.
diff --git a/src/main/kotlin/de/sirywell/handlehints/dfa/SsaAnalyzer.kt b/src/main/kotlin/de/sirywell/handlehints/dfa/SsaAnalyzer.kt
@@ -775,8 +775,14 @@ class SsaAnalyzer(private val controlFlow: ControlFlow, val typeData: TypeData)
                 methodHandlesInitializer.zero(arguments[0].asType())
             }
 
-            "byteArrayViewVarHandle",
-            "byteBufferViewVarHandle",
+            "byteArrayViewVarHandle" -> {
+                if (arguments.size != 2) return noMatch()
+                methodHandlesInitializer.byteArrayViewVarHandle(arguments[0], arguments[1])
+            }
+            "byteBufferViewVarHandle" -> {
+                if (arguments.size != 2) return noMatch()
+                methodHandlesInitializer.byteBufferViewVarHandle(arguments[0], arguments[1])
+            }
             "classData",
             "classDataAt",
             "lookup",

diff --git a/src/main/kotlin/de/sirywell/handlehints/foreign/FunctionDescriptorHelper.kt b/src/main/kotlin/de/sirywell/handlehints/foreign/FunctionDescriptorHelper.kt
@@ -4,6 +4,7 @@ import com.intellij.psi.PsiExpression
 import com.intellij.psi.PsiType
 import de.sirywell.handlehints.dfa.SsaAnalyzer
 import de.sirywell.handlehints.dfa.SsaConstruction
+import de.sirywell.handlehints.findPsiType
 import de.sirywell.handlehints.getConstantOfType
 import de.sirywell.handlehints.inspection.ProblemEmitter
 import de.sirywell.handlehints.type.*
@@ -70,11 +71,7 @@ class FunctionDescriptorHelper(private val ssaAnalyzer: SsaAnalyzer) : ProblemEm
     fun toMethodType(qualifier: PsiExpression, block: SsaConstruction.Block): MethodHandleType {
         val q = ssaAnalyzer.functionDescriptorType(qualifier, block) ?: TopFunctionDescriptorType
         val memorySegmentType = ExactType(
-            PsiType.getTypeByName(
-                "java.lang.foreign.MemorySegment",
-                qualifier.project,
-                qualifier.resolveScope
-            )
+            findPsiType("java.lang.foreign.MemorySegment", qualifier)
         )
         return q.toMethodType(memorySegmentType)
     }

diff --git a/src/main/kotlin/de/sirywell/handlehints/foreign/LinkerHelper.kt b/src/main/kotlin/de/sirywell/handlehints/foreign/LinkerHelper.kt
@@ -5,6 +5,7 @@ import com.intellij.psi.PsiType
 import de.sirywell.handlehints.TriState
 import de.sirywell.handlehints.dfa.SsaAnalyzer
 import de.sirywell.handlehints.dfa.SsaConstruction
+import de.sirywell.handlehints.findPsiType
 import de.sirywell.handlehints.inspection.ProblemEmitter
 import de.sirywell.handlehints.type.*
 
@@ -28,11 +29,7 @@ class LinkerHelper(
 
     private fun buildLeadingParams(methodType: MethodHandleType, context: PsiExpression): TypeList {
         val memorySegmentType = ExactType(
-            PsiType.getTypeByName(
-                "java.lang.foreign.MemorySegment",
-                context.project,
-                context.resolveScope
-            )
+            findPsiType("java.lang.foreign.MemorySegment", context)
         )
         val (_, identical) = methodType.returnType.joinIdentical(memorySegmentType)
         return when (identical) {
@@ -50,11 +47,7 @@ class LinkerHelper(
                     listOf(
                         memorySegmentType,
                         ExactType(
-                            PsiType.getTypeByName(
-                                "java.lang.foreign.SegmentAllocator",
-                                context.project,
-                                context.resolveScope
-                            )
+                            findPsiType("java.lang.foreign.SegmentAllocator", context)
                         )
                     )
                 )

diff --git a/src/main/kotlin/de/sirywell/handlehints/foreign/MemoryLayoutHelper.kt b/src/main/kotlin/de/sirywell/handlehints/foreign/MemoryLayoutHelper.kt
@@ -3,13 +3,10 @@ package de.sirywell.handlehints.foreign
 import com.intellij.codeInspection.LocalQuickFix.from
 import com.intellij.psi.PsiExpression
 import com.intellij.psi.PsiType
+import de.sirywell.handlehints.*
 import de.sirywell.handlehints.MethodHandleBundle.message
-import de.sirywell.handlehints.TriState
-import de.sirywell.handlehints.TypeData
 import de.sirywell.handlehints.dfa.SsaAnalyzer
 import de.sirywell.handlehints.dfa.SsaConstruction
-import de.sirywell.handlehints.getConstantLong
-import de.sirywell.handlehints.getConstantOfType
 import de.sirywell.handlehints.inspection.AdjustAlignmentFix
 import de.sirywell.handlehints.inspection.AdjustPaddingFix
 import de.sirywell.handlehints.inspection.ProblemEmitter
@@ -202,7 +199,7 @@ class MemoryLayoutHelper(private val ssaAnalyzer: SsaAnalyzer) : ProblemEmitter(
         val layoutType = ssaAnalyzer.memoryLayoutType(qualifier, block) ?: TopMemoryLayoutType
         val path = toPath(arguments, block)
         val memorySegmentType =
-            PsiType.getTypeByName("java.lang.foreign.MemorySegment", qualifier.project, qualifier.resolveScope)
+            findPsiType("java.lang.foreign.MemorySegment", qualifier)
         return VarHandlePathTraverser(typeData) {
             if (it == -1) methodExpr
             else arguments[it]
@@ -328,7 +325,7 @@ class MemoryLayoutHelper(private val ssaAnalyzer: SsaAnalyzer) : ProblemEmitter(
                 return CompleteVarHandleType(layoutType.type, CompleteTypeList(coords))
             } else {
                 val ctx = contextElement(-1) // good enough for us
-                val type = PsiType.getTypeByName("java.lang.foreign.MemorySegment", ctx.project, ctx.resolveScope)
+                val type = findPsiType("java.lang.foreign.MemorySegment", ctx)
                 return CompleteVarHandleType(ExactType(type), CompleteTypeList(coords))
             }
         }
@@ -374,7 +371,7 @@ class MemoryLayoutHelper(private val ssaAnalyzer: SsaAnalyzer) : ProblemEmitter(
     ): VarHandleType {
         val layoutType = ssaAnalyzer.memoryLayoutType(qualifier, block) ?: TopMemoryLayoutType
         val memorySegmentType =
-            PsiType.getTypeByName("java.lang.foreign.MemorySegment", qualifier.project, qualifier.resolveScope)
+            findPsiType("java.lang.foreign.MemorySegment", qualifier)
         val coords = mutableListOf(ExactType(memorySegmentType), *SCALE_HANDLE_PARAMETERS.typeList.toTypedArray())
         val path = toPath(arguments, block)
         return VarHandlePathTraverser(typeData) {

diff --git a/src/main/kotlin/de/sirywell/handlehints/inspection/AddArrayDimensionFix.kt b/src/main/kotlin/de/sirywell/handlehints/inspection/AddArrayDimensionFix.kt
@@ -0,0 +1,28 @@
+package de.sirywell.handlehints.inspection
+
+import com.intellij.codeInspection.LocalQuickFix
+import com.intellij.codeInspection.ProblemDescriptor
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiClassObjectAccessExpression
+import com.intellij.psi.PsiExpression
+import com.siyeh.ig.PsiReplacementUtil
+import com.siyeh.ig.psiutils.CommentTracker
+import de.sirywell.handlehints.MethodHandleBundle
+
+class AddArrayDimensionFix : LocalQuickFix {
+    override fun getFamilyName(): String {
+        return MethodHandleBundle.message("problem.general.array.dimension.add")
+    }
+
+    override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
+        val commentTracker = CommentTracker()
+        val expression = descriptor.psiElement as? PsiClassObjectAccessExpression ?: return
+        val operand = expression.operand
+        val arrayType = operand.type.createArrayType()
+        PsiReplacementUtil.replaceExpression(
+            descriptor.psiElement as PsiExpression,
+            "${arrayType.presentableText}.class",
+            commentTracker
+        )
+    }
+}
diff --git a/src/main/kotlin/de/sirywell/handlehints/inspection/ProblemEmitter.kt b/src/main/kotlin/de/sirywell/handlehints/inspection/ProblemEmitter.kt
@@ -2,6 +2,7 @@ package de.sirywell.handlehints.inspection
 
 import com.intellij.codeInspection.LocalQuickFix
 import com.intellij.codeInspection.ProblemHighlightType
+import com.intellij.psi.PsiClassObjectAccessExpression
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiExpression
 import com.intellij.psi.PsiTypes
@@ -51,14 +52,21 @@ abstract class ProblemEmitter(protected val typeData: TypeData) {
 
     protected inline fun <reified T : TypeLatticeElement<T>> emitMustBeArrayType(
         refc: PsiExpression,
-        referenceClass: Type
+        referenceClass: Type,
+        suggestFix: Boolean
     ): T {
-        return emitProblem(
-            refc, message(
-                "problem.merging.general.arrayTypeExpected",
-                referenceClass,
-            )
-        )
+        val message = message("problem.general.arrayTypeExpected", referenceClass)
+        return if (suggestFix) {
+            val fix = if (refc is PsiClassObjectAccessExpression) {
+                AddArrayDimensionFix()
+            } else {
+                WrapWithInvocationFix("arrayType", false)
+            }
+            emitProblem(refc, message, fix)
+
+        } else {
+            emitProblem(refc, message)
+        }
     }
 
     protected fun emitIncompatibleReturnTypes(

diff --git a/src/main/kotlin/de/sirywell/handlehints/inspection/WrapWithInvocationFix.kt b/src/main/kotlin/de/sirywell/handlehints/inspection/WrapWithInvocationFix.kt
@@ -0,0 +1,35 @@
+package de.sirywell.handlehints.inspection
+
+import com.intellij.codeInspection.LocalQuickFix
+import com.intellij.codeInspection.ProblemDescriptor
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiExpression
+import com.siyeh.ig.PsiReplacementUtil
+import com.siyeh.ig.psiutils.CommentTracker
+import de.sirywell.handlehints.MethodHandleBundle
+
+class WrapWithInvocationFix(private val methodName: String, private val static: Boolean) : LocalQuickFix {
+    override fun getFamilyName(): String {
+        return if (static) {
+            MethodHandleBundle.message("problem.general.invocation.wrap", methodName)
+        } else {
+            MethodHandleBundle.message("problem.general.invocation.append", methodName)
+        }
+    }
+
+    override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
+        val commentTracker = CommentTracker()
+        val text = descriptor.psiElement.text
+        val newText = if (static) {
+            "$methodName($text)"
+        } else {
+            "$text.$methodName()"
+        }
+        PsiReplacementUtil.replaceExpression(
+            descriptor.psiElement as PsiExpression,
+            newText,
+            commentTracker
+        )
+
+    }
+}
diff --git a/src/main/kotlin/de/sirywell/handlehints/mhtype/MethodHandlesInitializer.kt b/src/main/kotlin/de/sirywell/handlehints/mhtype/MethodHandlesInitializer.kt
@@ -65,7 +65,35 @@ class MethodHandlesInitializer(private val ssaAnalyzer: SsaAnalyzer) : ProblemEm
         return complete(ExactType.intType, listOf(arrayType))
     }
 
-    // byteArray/BufferViewVarHandle() no VarHandle support
+    fun byteArrayViewVarHandle(arrayClass: PsiExpression, byteOrder: PsiExpression): VarHandleType {
+        return varHandleForCollectionType(arrayClass, ExactType(PsiTypes.byteType().createArrayType()))
+    }
+
+    fun byteBufferViewVarHandle(arrayClass: PsiExpression, byteOrder: PsiExpression): VarHandleType {
+        return varHandleForCollectionType(arrayClass, ExactType(findPsiType("java.nio.ByteBuffer", arrayClass)))
+    }
+
+    private fun varHandleForCollectionType(arrayClass: PsiExpression, collectionType: Type): CompleteVarHandleType {
+        val arrayType = arrayClass.asArrayType { isSupportedViewHandleComponentType(it) }
+        var componentType = arrayType.componentType()
+        if (componentType is ExactType && !isSupportedViewHandleComponentType(componentType.psiType)) {
+            componentType = emitMustBeViewHandleSupportedComponentType(arrayClass, componentType)
+        }
+        return CompleteVarHandleType(componentType, CompleteTypeList(listOf(collectionType, ExactType.intType)))
+    }
+
+    private fun isSupportedViewHandleComponentType(type: PsiType): Boolean {
+        return type == PsiTypes.shortType()
+                || type == PsiTypes.charType()
+                || type == PsiTypes.intType()
+                || type == PsiTypes.longType()
+                || type == PsiTypes.floatType()
+                || type == PsiTypes.doubleType()
+    }
+
+    private fun emitMustBeViewHandleSupportedComponentType(expr: PsiExpression, type: Type): Type {
+        return emitProblem<Type>(expr, message("problem.general.array.unsupportedViewHandleComponentType", type))
+    }
 
     fun constant(typeExpr: PsiExpression, valueExpr: PsiExpression): MethodHandleType {
         val type = typeExpr.asType()
@@ -123,10 +151,10 @@ class MethodHandlesInitializer(private val ssaAnalyzer: SsaAnalyzer) : ProblemEm
         return complete(type, listOf())
     }
 
-    private fun PsiExpression.asArrayType(): Type {
+    private fun PsiExpression.asArrayType(validComponentTypeCheck: (PsiType) -> Boolean = { false }): Type {
         val referenceClass = this.asType()
         if (referenceClass is ExactType && referenceClass.psiType !is PsiArrayType) {
-            return emitMustBeArrayType(this, referenceClass)
+            return emitMustBeArrayType(this, referenceClass, validComponentTypeCheck(referenceClass.psiType))
         }
         return referenceClass
     }
@@ -156,7 +184,7 @@ class MethodHandlesInitializer(private val ssaAnalyzer: SsaAnalyzer) : ProblemEm
         if (accessType != null) {
             type = checkAccessModeType(accessType, type, exact, methodTypeExpr)
         }
-        val varHandleType = PsiType.getTypeByName(VAR_HANDLE_FQN, methodTypeExpr.project, methodTypeExpr.resolveScope)
+        val varHandleType = findPsiType(VAR_HANDLE_FQN, methodTypeExpr)
         return type.withParameterTypes(
             type.parameterTypes.addAllAt(0, CompleteTypeList(listOf(ExactType(varHandleType))))
         )

diff --git a/src/main/kotlin/de/sirywell/handlehints/psiSupport.kt b/src/main/kotlin/de/sirywell/handlehints/psiSupport.kt
@@ -15,6 +15,9 @@ import java.lang.invoke.MethodType.methodType
 val PsiMethodCallExpression.methodName
     get() = this.methodExpression.referenceName
 
+fun findPsiType(qName: String, context: PsiElement) =
+    PsiType.getTypeByName(qName, context.project, context.resolveScope)
+
 fun isJavaLangInvoke(element: PsiMethodCallExpression) =
     (element.resolveMethod()?.containingClass?.containingFile as? PsiJavaFile)?.packageName == "java.lang.invoke"
 
@@ -32,7 +35,7 @@ fun receiverIsMethodType(element: PsiMethodCallExpression) = receiverIsInvokeCla
 fun receiverIsLookup(element: PsiMethodCallExpression) = receiverIsInvokeClass(element, "MethodHandles.Lookup")
 
 fun receiverIsMemoryLayout(element: PsiMethodCallExpression): Boolean {
-    val superType = PsiType.getTypeByName("java.lang.foreign.MemoryLayout", element.project, element.resolveScope)
+    val superType = findPsiType("java.lang.foreign.MemoryLayout", element)
     val actual = PsiTypesUtil.getClassType(element.resolveMethod()?.containingClass ?: return false)
     return superType.isAssignableFrom(actual)
 }
@@ -50,23 +53,23 @@ fun receiverIsLinker(element: PsiMethodCallExpression): Boolean {
 }
 
 fun methodHandleType(element: PsiElement): PsiClassType {
-    return PsiType.getTypeByName("java.lang.invoke.MethodHandle", element.project, element.resolveScope)
+    return findPsiType("java.lang.invoke.MethodHandle", element)
 }
 
 fun varHandleType(element: PsiElement): PsiClassType {
-    return PsiType.getTypeByName("java.lang.invoke.VarHandle", element.project, element.resolveScope)
+    return findPsiType("java.lang.invoke.VarHandle", element)
 }
 
 fun methodTypeType(element: PsiElement): PsiClassType {
-    return PsiType.getTypeByName("java.lang.invoke.MethodType", element.project, element.resolveScope)
+    return findPsiType("java.lang.invoke.MethodType", element)
 }
 
 fun pathElementType(element: PsiElement): PsiClassType {
-    return PsiType.getTypeByName("java.lang.foreign.MemoryLayout.PathElement", element.project, element.resolveScope)
+    return findPsiType("java.lang.foreign.MemoryLayout.PathElement", element)
 }
 
 fun functionDescriptorType(element: PsiElement): PsiClassType {
-    return PsiType.getTypeByName("java.lang.foreign.FunctionDescriptor", element.project, element.resolveScope)
+    return findPsiType("java.lang.foreign.FunctionDescriptor", element)
 }
 
 fun objectType(element: PsiElement): PsiType {
@@ -95,7 +98,7 @@ fun memoryLayoutTypes(context: PsiElement): Set<PsiType> {
             "java.lang.foreign.ValueLayout.OfDouble",
             "java.lang.foreign.AddressLayout",
         )
-            .map { PsiType.getTypeByName(it, context.project, context.resolveScope) }
+            .map { findPsiType(it, context) }
             .toSet()
     }.memoryLayoutTypes
 }
@@ -143,7 +146,7 @@ fun PsiExpression.getConstantLong(): Long? {
 }
 
 
-fun Collection<PsiExpression>.mapToTypes(check: (PsiExpression, Type) -> Type = { _, t -> t}): List<Type> {
+fun Collection<PsiExpression>.mapToTypes(check: (PsiExpression, Type) -> Type = { _, t -> t }): List<Type> {
     return this.map { element -> check(element, element.asType()) }
 }
 

diff --git a/src/main/kotlin/de/sirywell/handlehints/type/Type.kt b/src/main/kotlin/de/sirywell/handlehints/type/Type.kt
@@ -1,5 +1,6 @@
 package de.sirywell.handlehints.type
 
+import com.intellij.psi.PsiArrayType
 import com.intellij.psi.PsiManager
 import com.intellij.psi.PsiPrimitiveType
 import com.intellij.psi.PsiType
@@ -20,6 +21,8 @@ sealed interface Type : TypeLatticeElement<Type> {
     fun match(psiType: PsiType): TriState
 
     fun isPrimitive(): TriState
+
+    fun componentType(): Type
 }
 
 data object BotType : Type, BotTypeLatticeElement<Type> {
@@ -28,6 +31,7 @@ data object BotType : Type, BotTypeLatticeElement<Type> {
     override fun erase(manager: PsiManager, scope: GlobalSearchScope) = this
     override fun match(psiType: PsiType) = TriState.UNKNOWN
     override fun isPrimitive() = TriState.UNKNOWN
+    override fun componentType() = this
 }
 
 data object TopType : Type, TopTypeLatticeElement<Type> {
@@ -38,6 +42,7 @@ data object TopType : Type, TopTypeLatticeElement<Type> {
     override fun erase(manager: PsiManager, scope: GlobalSearchScope) = this
     override fun match(psiType: PsiType) = TriState.UNKNOWN
     override fun isPrimitive() = TriState.UNKNOWN
+    override fun componentType() = this
 }
 
 @JvmRecord
@@ -90,6 +95,13 @@ data class ExactType(val psiType: PsiType) : Type {
         return (psiType is PsiPrimitiveType).toTriState()
     }
 
+    override fun componentType(): Type {
+        if (psiType is PsiArrayType) {
+            return ExactType(psiType.componentType)
+        }
+        return TopType
+    }
+
     override fun toString(): String {
         return psiType.presentableText
     }