Add support for var-args method calls with @PolymorphicSignature causing unmatched method signatures

sootup.callgraph/pom.xml

@@ -56,6 +56,10 @@
 			<artifactId>sootup.jimple.frontend</artifactId>
 			<scope>test</scope>
 		</dependency>
-	</dependencies>
+        <dependency>
+            <groupId>org.soot-oss</groupId>
+            <artifactId>sootup.java.core</artifactId>
+        </dependency>
+    </dependencies>
 
 </project>

sootup.callgraph/src/main/java/sootup/callgraph/AbstractCallGraphAlgorithm.java

@@ -24,6 +24,8 @@
 
 import static sootup.core.jimple.basic.StmtPositionInfo.getNoStmtPositionInfo;
 
+import com.google.common.collect.HashBasedTable;
+import com.google.common.collect.Table;
 import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -49,6 +51,9 @@
 import sootup.core.types.ClassType;
 import sootup.core.types.VoidType;
 import sootup.core.views.View;
+import sootup.java.core.AnnotationUsage;
+import sootup.java.core.JavaSootClass;
+import sootup.java.core.JavaSootMethod;
 
 /**
  * The AbstractCallGraphAlgorithm class is the super class of all call graph algorithm. It provides
@@ -61,10 +66,59 @@ public abstract class AbstractCallGraphAlgorithm implements CallGraphAlgorithm {
 
   @NonNull protected final View view;
   @NonNull protected final TypeHierarchy typeHierarchy;
+  @NonNull protected final Table<String, String, SootMethod> preanalysis;
 
   protected AbstractCallGraphAlgorithm(@NonNull View view) {
     this.view = view;
     this.typeHierarchy = view.getTypeHierarchy();
+    this.preanalysis = getMethodsWithPolymorphicAnnotation();
+  }
+
+  /**
+   * Collects all method signatures from the current view that are annotated with {@code
+   * java.lang.invoke.MethodHandle$PolymorphicSignature}.
+   *
+   * @return a list of method signatures annotated with {@code @PolymorphicSignature}
+   */
+  protected Table<String, String, SootMethod> getMethodsWithPolymorphicAnnotation() {
+    Table<String, String, SootMethod> polymorphicMethods = HashBasedTable.create();
+    ClassType polymorphicAnnotationType =
+        view.getIdentifierFactory()
+            .getClassType("java.lang.invoke.MethodHandle$PolymorphicSignature");
+    Map<ClassType, List<JavaSootMethod>> polymorphicTypeAndMethods = new HashMap<>();
+    view.getClasses()
+        .filter(sootClass -> sootClass instanceof JavaSootClass)
+        .map(sootClass -> (JavaSootClass) sootClass)
+        .flatMap(javaSootClass -> javaSootClass.getMethods().stream())
+        .forEach(
+            javaSootMethod -> {
+              Iterable<AnnotationUsage> annotationUsages = javaSootMethod.getAnnotations();
+              for (AnnotationUsage annotationUsage : annotationUsages) {
+                if (annotationUsage.getAnnotation().equals(polymorphicAnnotationType)) {
+                  ClassType classType = javaSootMethod.getDeclaringClassType();
+                  polymorphicTypeAndMethods
+                      .computeIfAbsent(classType, k -> new ArrayList<>())
+                      .add(javaSootMethod);
+                  polymorphicMethods.put(
+                      javaSootMethod.getName(), classType.getFullyQualifiedName(), javaSootMethod);
+                }
+              }
+            });
+    for (Map.Entry<ClassType, List<JavaSootMethod>> methodAndType :
+        polymorphicTypeAndMethods.entrySet()) {
+      ClassType classType = methodAndType.getKey();
+      List<JavaSootMethod> javaSootMethods = methodAndType.getValue();
+      typeHierarchy
+          .subtypesOf(classType)
+          .forEach(
+              subtype -> {
+                for (JavaSootMethod javaSootMethod : javaSootMethods) {
+                  polymorphicMethods.put(
+                      javaSootMethod.getName(), classType.getFullyQualifiedName(), javaSootMethod);
+                }
+              });
+    }
+    return polymorphicMethods;
   }
 
   /**
@@ -592,18 +646,7 @@ protected static Optional<? extends SootMethod> findConcreteMethod(
     }
 
     // search method in interfaces
-    Optional<? extends SootMethod> defaultMethod = findDefaultMethod(view, startClass, methodSig);
-    if (defaultMethod.isPresent()) {
-      return defaultMethod;
-    }
-
-    logger.warn(
-        "Could not find \""
-            + sig.getSubSignature()
-            + "\" in "
-            + sig.getDeclClassType().getClassName()
-            + " and in its superclasses and interfaces");
-    return Optional.empty();
+    return findDefaultMethod(view, startClass, methodSig);
   }
 
   protected static Optional<SootMethod> findMethodInHierarchy(
@@ -667,4 +710,30 @@ protected static Optional<? extends SootMethod> findDefaultMethod(
   protected boolean isInterface(ClassType classType) {
     return typeHierarchy.isInterface(classType);
   }
+
+  /**
+   * Attempts to resolve a previously unresolved method signature by comparing it with all known
+   * method signatures annotated with {@code @PolymorphicSignature}. If a match is found (ignoring
+   * the parameters), the corresponding SootMethod is returned.
+   *
+   * @param targetMethodSignature the signature of the searched method
+   * @return the found method object, or null if the method was not found.
+   */
+  protected Optional<? extends SootMethod> findMatchingVarArgsMethod(
+      @NonNull MethodSignature targetMethodSignature) {
+    Optional<SootMethod> resolvedVarArgsMethod =
+        Optional.ofNullable(
+            preanalysis.get(
+                targetMethodSignature.getName(),
+                targetMethodSignature.getDeclClassType().getFullyQualifiedName()));
+    if (resolvedVarArgsMethod.isEmpty()) {
+      logger.warn(
+          "Could not find \""
+              + targetMethodSignature.getSubSignature()
+              + "\" in "
+              + targetMethodSignature.getDeclClassType().getClassName()
+              + " and in its superclasses and interfaces");
+    }
+    return resolvedVarArgsMethod;
+  }
 }

sootup.callgraph/src/main/java/sootup/callgraph/ClassHierarchyAnalysisAlgorithm.java

@@ -93,9 +93,14 @@ protected Stream<MethodSignature> resolveCall(SootMethod method, InvokableStmt i
     if (actualTargetMethod == null) {
       // method not implemented, search for implementation in super classes or ínterfaces
       actualTargetMethod = findConcreteMethod(view, targetMethodSignature).orElse(null);
-      // method implementation isn't contained in the view. return the called method as target
       if (actualTargetMethod == null) {
-        return Stream.of(targetMethodSignature);
+        // method implementation isn't contained in the view.
+        // check if method got the PolymorphicSignature annotation
+        actualTargetMethod = findMatchingVarArgsMethod(targetMethodSignature).orElse(null);
+        if (actualTargetMethod == null) {
+          // method couldn't be resolved, return the called method as target
+          return Stream.of(targetMethodSignature);
+        }
       }
     }
 
@@ -126,7 +131,7 @@ protected Stream<MethodSignature> resolveCall(SootMethod method, InvokableStmt i
               targets,
               resolveAllDefaultTargets(
                   subclasses,
-                  actualTargetMethod.getDeclClassType(),
+                  actualTargetMethod.getDeclaringClassType(),
                   targetMethodSignature.getSubSignature()));
       // all subtypes that do not have an implementation of the method
       // can have an implementation in the supertype
@@ -166,7 +171,7 @@ private Stream<MethodSignature> resolveAllDefaultTargets(
                         view.getIdentifierFactory()
                             .getMethodSignature(classType, targetSubSignature))
                     .stream())
-        .filter(sootMethod -> interfaces.contains(sootMethod.getDeclClassType()))
+        .filter(sootMethod -> interfaces.contains(sootMethod.getDeclaringClassType()))
         .map(SootMethod::getSignature);
   }
 

sootup.callgraph/src/main/java/sootup/callgraph/RapidTypeAnalysisAlgorithm.java

@@ -98,16 +98,16 @@ protected List<ClassType> collectInstantiatedClassesInMethod(SootMethod method)
     }
 
     Set<ClassType> instantiated =
-        method.getBody().getStmts().stream()
-            .filter(stmt -> stmt instanceof JAssignStmt)
-            .map(stmt -> ((JAssignStmt) stmt).getRightOp())
-            .filter(value -> value instanceof JNewExpr)
-            .map(value -> ((JNewExpr) value).getType())
-            .collect(Collectors.toSet());
+            method.getBody().getStmts().stream()
+                    .filter(stmt -> stmt instanceof JAssignStmt)
+                    .map(stmt -> ((JAssignStmt) stmt).getRightOp())
+                    .filter(value -> value instanceof JNewExpr)
+                    .map(value -> ((JNewExpr) value).getType())
+                    .collect(Collectors.toSet());
     List<ClassType> newInstantiatedClassTypes =
-        instantiated.stream()
-            .filter(classType -> !instantiatedClasses.contains(classType))
-            .collect(Collectors.toList());
+            instantiated.stream()
+                    .filter(classType -> !instantiatedClasses.contains(classType))
+                    .collect(Collectors.toList());
     instantiatedClasses.addAll(instantiated);
     return newInstantiatedClassTypes;
   }
@@ -141,9 +141,14 @@ protected Stream<MethodSignature> resolveCall(
     if (actualTargetMethod == null) {
       // method not implemented, search for implementation in super classes or ínterfaces
       actualTargetMethod = findConcreteMethod(view, targetMethodSignature).orElse(null);
-      // method implementation isn't contained in the view. return the called method as target
       if (actualTargetMethod == null) {
-        return Stream.of(targetMethodSignature);
+        // method implementation isn't contained in the view.
+        // check if method got the PolymorphicSignature annotation
+        actualTargetMethod = findMatchingVarArgsMethod(targetMethodSignature).orElse(null);
+        if (actualTargetMethod == null) {
+          // method couldn't be resolved, return the called method as target
+          return Stream.of(targetMethodSignature);
+        }
       }
     }
     // special invokes and static invokes can only have one specific target
@@ -159,7 +164,6 @@ protected Stream<MethodSignature> resolveCall(
             .subtypesOf(targetMethodSignature.getDeclClassType())
             .flatMap(classType -> view.getClass(classType).stream())
             .toList();
-
     // get all targets of these subtypes
     Stream<MethodSignature> targets =
         resolveAllCallTargets(
@@ -317,7 +321,6 @@ protected void preProcessingMethod(
     if (method == null) {
       return;
     }
-
     List<ClassType> newInstantiatedClasses = collectInstantiatedClassesInMethod(method);
     newInstantiatedClasses.forEach(
         classType -> includeIgnoredCallsToClass(classType, cg, workList));

sootup.callgraph/src/test/java/sootup/callgraph/CallGraphTestBase.java

@@ -2,9 +2,10 @@
 
 import static org.junit.jupiter.api.Assertions.*;
 
-import java.util.*;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 import org.junit.jupiter.api.Test;
 import sootup.core.inputlocation.AnalysisInputLocation;
 import sootup.core.jimple.common.Value;
@@ -1265,4 +1266,56 @@ public void testImplicitExample6() {
             identifierFactory.getClassType("t6.NoThread"), "run", "void", Collections.emptyList());
     assertFalse(cg.containsMethod(runMethodSig));
   }
+
+  @Test
+  public void testMethodHandleDiffParamExample1() {
+    CallGraph cg = loadCallGraph("Polymorphic", "e1.MethodHandleDiffParamExample1");
+    MethodSignature invokeMethodSig =
+        identifierFactory.getMethodSignature(
+            identifierFactory.getClassType("java.lang.invoke.MethodHandle"),
+            "invoke",
+            "java.lang.Object",
+            Collections.singletonList("java.lang.Object[]"));
+    Set<MethodSignature> callSourcesMethodSigs = cg.callSourcesTo(invokeMethodSig);
+    assertTrue(callSourcesMethodSigs.contains(mainMethodSignature));
+  }
+
+  @Test
+  public void testMethodHandleInvokeExample2() {
+    CallGraph cg = loadCallGraph("Polymorphic", "e2.MethodHandleInvokeExample2");
+    MethodSignature invokeMethodSig =
+        identifierFactory.getMethodSignature(
+            identifierFactory.getClassType("java.lang.invoke.MethodHandle"),
+            "invoke",
+            "java.lang.Object",
+            Collections.singletonList("java.lang.Object[]"));
+    Set<MethodSignature> callSourcesMethodSigs = cg.callSourcesTo(invokeMethodSig);
+    assertTrue(callSourcesMethodSigs.contains(mainMethodSignature));
+  }
+
+  @Test
+  public void testVarHandleGetExample3() {
+    CallGraph cg = loadCallGraph("Polymorphic", "e3.VarHandleGetExample3");
+    MethodSignature getMethodSig =
+        identifierFactory.getMethodSignature(
+            identifierFactory.getClassType("java.lang.invoke.VarHandle"),
+            "get",
+            "java.lang.Object",
+            Collections.singletonList("java.lang.Object[]"));
+    Set<MethodSignature> callSourcesMethodSigs = cg.callSourcesTo(getMethodSig);
+    assertTrue(callSourcesMethodSigs.contains(mainMethodSignature));
+  }
+
+  @Test
+  public void testMethodHandleSuperClassExample4() {
+    CallGraph cg = loadCallGraph("Polymorphic", "e4.MethodHandleSuperClassExample4");
+    MethodSignature invokeMethodSig =
+        identifierFactory.getMethodSignature(
+            identifierFactory.getClassType("java.lang.invoke.MethodHandle"),
+            "invoke",
+            "java.lang.Object",
+            Collections.singletonList("java.lang.Object[]"));
+    Set<MethodSignature> callSourcesMethodSigs = cg.callSourcesTo(invokeMethodSig);
+    assertTrue(callSourcesMethodSigs.contains(mainMethodSignature));
+  }
 }

sootup.callgraph/src/test/resources/callgraph/Polymorphic/source/MethodHandleDiffParamExample1.java

@@ -0,0 +1,18 @@
+package e1;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+
+public class MethodHandleDiffParamExample1 {
+    public static void main(String[] args) throws Throwable {
+        MethodHandle handle = MethodHandles.lookup()
+                .findVirtual(
+                        String.class,
+                        "indexOf",
+                        MethodType.methodType(int.class, String.class, int.class)
+                );
+
+        handle.invoke("HelloWorld", "lo", 3);
+    }
+}

sootup.callgraph/src/test/resources/callgraph/Polymorphic/source/MethodHandleInvokeExample2.java

@@ -0,0 +1,14 @@
+package e2;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+
+public class MethodHandleInvokeExample2 {
+    public static void main(String[] args) throws Throwable {
+        MethodHandle handle = MethodHandles.lookup()
+                .findVirtual(String.class, "substring", MethodType.methodType(String.class, int.class, int.class));
+
+        handle.invoke("HelloWorld", 0, 5);
+    }
+}

sootup.callgraph/src/test/resources/callgraph/Polymorphic/source/MethodHandleSuperClassExample4.java

@@ -0,0 +1,22 @@
+package e4;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+
+public class MethodHandleSuperClassExample4 {
+    static class Base {
+        protected void method1(String dummy) {
+            System.out.println("Method1: " + dummy);
+        }
+    }
+
+    static class Derived extends Base {}
+
+    public static void main(String[] args) throws Throwable {
+        MethodHandle handle = MethodHandles.lookup()
+                .findVirtual(Base.class, "method1", MethodType.methodType(void.class, String.class));
+
+        handle.invoke(new Derived(), "SubClass!");
+    }
+}

sootup.callgraph/src/test/resources/callgraph/Polymorphic/source/VarHandleGetExample3.java

@@ -0,0 +1,15 @@
+package e3;
+
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
+
+public class VarHandleGetExample3 {
+    static int value = 42;
+
+    public static void main(String[] args) throws Throwable {
+        VarHandle handle = MethodHandles.lookup()
+                .findStaticVarHandle(VarHandleGetExample3.class, "value", int.class);
+
+        handle.get();
+    }
+}

sootup.core/src/main/java/sootup/core/model/SootMethod.java

@@ -81,8 +81,6 @@ public interface SootMethod extends Method {
 
   @NonNull MethodSubSignature getSubSignature();
 
-  @NonNull ClassType getDeclClassType();
-
   @NonNull String getName();
 
   @NonNull List<ClassType> getExceptionSignatures();