Keep lambda invocation kind tag in the lambda expr (#269)

UnitTestBot · Sep 18, 2024 · 82a70aa · 82a70aa
1 parent 94784e0
commit 82a70aa
Show file tree

Hide file tree

Showing 5 changed files with 76 additions and 21 deletions.
diff --git a/jacodb-api-jvm/src/main/kotlin/org/jacodb/api/jvm/cfg/JcInst.kt b/jacodb-api-jvm/src/main/kotlin/org/jacodb/api/jvm/cfg/JcInst.kt
@@ -652,8 +652,10 @@ data class JcLambdaExpr(
     val callSiteArgTypes: List<JcType>,
     val callSiteReturnType: JcType,
     val callSiteArgs: List<JcValue>,
-    val isNewInvokeSpecial: Boolean,
+    val lambdaInvokeKind: BsmHandleTag.MethodHandle,
 ) : JcCallExpr {
+    val isNewInvokeSpecial: Boolean
+        get() = lambdaInvokeKind == BsmHandleTag.MethodHandle.NEW_INVOKE_SPECIAL
 
     override val method get() = bsmRef.method
     override val args get() = callSiteArgs

diff --git a/jacodb-api-jvm/src/main/kotlin/org/jacodb/api/jvm/cfg/JcRawInst.kt b/jacodb-api-jvm/src/main/kotlin/org/jacodb/api/jvm/cfg/JcRawInst.kt
@@ -692,8 +692,25 @@ data class BsmMethodTypeArg(val argumentTypes: List<TypeName>, val returnType: T
     override fun toString(): String = "(${argumentTypes.joinToString { it.typeName }}:${returnType.typeName})"
 }
 
+sealed interface BsmHandleTag {
+    enum class FieldHandle : BsmHandleTag {
+        GET_FIELD,
+        GET_STATIC,
+        PUT_FIELD,
+        PUT_STATIC,
+    }
+
+    enum class MethodHandle : BsmHandleTag {
+        INVOKE_VIRTUAL,
+        INVOKE_STATIC,
+        INVOKE_SPECIAL,
+        NEW_INVOKE_SPECIAL,
+        INVOKE_INTERFACE,
+    }
+}
+
 data class BsmHandle(
-    val tag: Int,
+    val tag: BsmHandleTag,
     val declaringClass: TypeName,
     val name: String,
     val argTypes: List<TypeName>,

diff --git a/jacodb-core/src/main/kotlin/org/jacodb/impl/cfg/JcInstListBuilder.kt b/jacodb-core/src/main/kotlin/org/jacodb/impl/cfg/JcInstListBuilder.kt
@@ -287,21 +287,26 @@ class JcInstListBuilder(val method: JcMethod,val instList: JcInstList<JcRawInst>
 
         val argTypes: List<TypeName>
         val tag = implementation.tag
-        var isNewInvokeSpecial = false
-        if (tag == 6) {
-            // Invoke static case
-            argTypes = implementation.argTypes
-        } else if (tag == 8) {
-            isNewInvokeSpecial = true
-            argTypes = implementation.argTypes
-        } else {
-            // Invoke non-static case
-            check(tag == 5 || tag == 7 || tag == 9) {
-                "Unexpected tag for invoke dynamic $tag"
+
+        check(tag is BsmHandleTag.MethodHandle) {
+            "Unexpected tag for invoke dynamic $tag"
+        }
+
+        when (tag) {
+            BsmHandleTag.MethodHandle.INVOKE_STATIC,
+            BsmHandleTag.MethodHandle.NEW_INVOKE_SPECIAL -> {
+                // Invoke static or invoke constructor case
+                argTypes = implementation.argTypes
+            }
+
+            BsmHandleTag.MethodHandle.INVOKE_VIRTUAL,
+            BsmHandleTag.MethodHandle.INVOKE_SPECIAL,
+            BsmHandleTag.MethodHandle.INVOKE_INTERFACE -> {
+                // Invoke non-static case
+                argTypes = implementation.argTypes.toMutableList()
+                // Adding 'this' type as first argument type
+                argTypes.add(0, implementation.declaringClass)
             }
-            argTypes = implementation.argTypes.toMutableList()
-            // Adding 'this' type as first argument type
-            argTypes.add(0, implementation.declaringClass)
         }
 
         // Check implementation signature match (starts with) call site arguments
@@ -323,7 +328,7 @@ class JcInstListBuilder(val method: JcMethod,val instList: JcInstList<JcRawInst>
             expr.callSiteArgTypes.map { it.asType() },
             expr.callSiteReturnType.asType(),
             expr.callSiteArgs.map { it.accept(this) as JcValue },
-            isNewInvokeSpecial
+            tag
         )
     }
 

diff --git a/jacodb-core/src/main/kotlin/org/jacodb/impl/cfg/MethodNodeBuilder.kt b/jacodb-core/src/main/kotlin/org/jacodb/impl/cfg/MethodNodeBuilder.kt
@@ -24,7 +24,6 @@ import org.jacodb.api.jvm.cfg.*
 import org.jacodb.impl.cfg.util.*
 import org.objectweb.asm.Handle
 import org.objectweb.asm.Opcodes
-import org.objectweb.asm.Opcodes.H_GETSTATIC
 import org.objectweb.asm.Type
 import org.objectweb.asm.tree.*
 
@@ -64,6 +63,19 @@ private val TypeName.typeInt
         else -> error("$typeName is not primitive type")
     }
 
+private val BsmHandleTag.tagInt: Int
+    get() = when (this) {
+        BsmHandleTag.FieldHandle.GET_FIELD -> Opcodes.H_GETFIELD
+        BsmHandleTag.FieldHandle.GET_STATIC -> Opcodes.H_GETSTATIC
+        BsmHandleTag.FieldHandle.PUT_FIELD -> Opcodes.H_PUTFIELD
+        BsmHandleTag.FieldHandle.PUT_STATIC -> Opcodes.H_PUTSTATIC
+        BsmHandleTag.MethodHandle.INVOKE_VIRTUAL -> Opcodes.H_INVOKEVIRTUAL
+        BsmHandleTag.MethodHandle.INVOKE_STATIC -> Opcodes.H_INVOKESTATIC
+        BsmHandleTag.MethodHandle.INVOKE_SPECIAL -> Opcodes.H_INVOKESPECIAL
+        BsmHandleTag.MethodHandle.NEW_INVOKE_SPECIAL -> Opcodes.H_NEWINVOKESPECIAL
+        BsmHandleTag.MethodHandle.INVOKE_INTERFACE -> Opcodes.H_INVOKEINTERFACE
+    }
+
 class MethodNodeBuilder(
     val method: JcMethod,
     val instList: JcInstList<JcRawInst>
@@ -597,12 +609,15 @@ class MethodNodeBuilder(
     }
 
 
+    private val BsmHandleTag.isGetFieldOrStaticTag
+        get() = this == BsmHandleTag.FieldHandle.GET_FIELD || this == BsmHandleTag.FieldHandle.GET_STATIC
+
     private val BsmHandle.asAsmHandle: Handle
         get() = Handle(
-            tag,
+            tag.tagInt,
             declaringClass.jvmClassName,
             name,
-            if (argTypes.isEmpty() && tag <= H_GETSTATIC) {
+            if (argTypes.isEmpty() && tag.isGetFieldOrStaticTag) {
                 returnType.jvmTypeName
             } else {
                 "(${argTypes.joinToString("") { it.jvmTypeName }})${returnType.jvmTypeName}"

diff --git a/jacodb-core/src/main/kotlin/org/jacodb/impl/cfg/RawInstListBuilder.kt b/jacodb-core/src/main/kotlin/org/jacodb/impl/cfg/RawInstListBuilder.kt
@@ -27,6 +27,7 @@ import org.jacodb.api.jvm.cfg.BsmArg
 import org.jacodb.api.jvm.cfg.BsmDoubleArg
 import org.jacodb.api.jvm.cfg.BsmFloatArg
 import org.jacodb.api.jvm.cfg.BsmHandle
+import org.jacodb.api.jvm.cfg.BsmHandleTag
 import org.jacodb.api.jvm.cfg.BsmIntArg
 import org.jacodb.api.jvm.cfg.BsmLongArg
 import org.jacodb.api.jvm.cfg.BsmMethodTypeArg
@@ -159,6 +160,21 @@ private fun parsePrimitiveType(opcode: Int) = when (opcode) {
     else -> error("Unknown opcode in primitive type parsing: $opcode")
 }
 
+private fun parseBsmHandleTag(tag: Int): BsmHandleTag = when (tag) {
+    Opcodes.H_GETFIELD -> BsmHandleTag.FieldHandle.GET_FIELD
+    Opcodes.H_GETSTATIC -> BsmHandleTag.FieldHandle.GET_STATIC
+    Opcodes.H_PUTFIELD -> BsmHandleTag.FieldHandle.PUT_FIELD
+    Opcodes.H_PUTSTATIC -> BsmHandleTag.FieldHandle.PUT_STATIC
+
+    Opcodes.H_INVOKEVIRTUAL -> BsmHandleTag.MethodHandle.INVOKE_VIRTUAL
+    Opcodes.H_INVOKESTATIC -> BsmHandleTag.MethodHandle.INVOKE_STATIC
+    Opcodes.H_INVOKESPECIAL -> BsmHandleTag.MethodHandle.INVOKE_SPECIAL
+    Opcodes.H_NEWINVOKESPECIAL -> BsmHandleTag.MethodHandle.NEW_INVOKE_SPECIAL
+    Opcodes.H_INVOKEINTERFACE -> BsmHandleTag.MethodHandle.INVOKE_INTERFACE
+
+    else -> error("Unknown tag in BSM handle: $tag")
+}
+
 private fun parseType(any: Any): TypeName = when (any) {
     is String -> any.typeName()
     is Int -> parsePrimitiveType(any)
@@ -1357,7 +1373,7 @@ class RawInstListBuilder(
 
     private val Handle.bsmHandleArg
         get() = BsmHandle(
-            tag,
+            parseBsmHandleTag(tag),
             owner.typeName(),
             name,
             if (desc.contains("(")) {