Add Call Graph [WIP]

build.gradle

@@ -92,6 +92,7 @@ subprojects {
         javassist = 'org.javassist:javassist:3.28.0-GA'
         javax_annos = 'javax.annotation:javax.annotation-api:1.3.2'
         jelf = 'net.fornwall:jelf:0.7.0'
+        jlinker = 'com.github.xxDark:jlinker:1.0.3'
         jphantom = 'com.github.Col-E:jphantom:1.4.3'
         junit_api = "org.junit.jupiter:junit-jupiter-api:$junitVersion"
         junit_engine = "org.junit.jupiter:junit-jupiter-engine:$junitVersion"

recaf-core/build.gradle

@@ -36,6 +36,7 @@ dependencies {
     api ssvm
     api procyon_core
     api procyon_compiler_tools
+    api jlinker
 
     // Android
     api smali

recaf-core/src/main/java/me/coley/recaf/Services.java

@@ -2,14 +2,17 @@
 
 import me.coley.recaf.compile.CompilerManager;
 import me.coley.recaf.decompile.DecompileManager;
-import me.coley.recaf.graph.InheritanceGraph;
+import me.coley.recaf.graph.call.CallGraphRegistry;
+import me.coley.recaf.graph.inheritance.InheritanceGraph;
 import me.coley.recaf.mapping.MappingsManager;
 import me.coley.recaf.parse.JavaParserHelper;
 import me.coley.recaf.parse.WorkspaceSymbolSolver;
 import me.coley.recaf.ssvm.SsvmIntegration;
 import me.coley.recaf.util.WorkspaceTreeService;
 import me.coley.recaf.workspace.Workspace;
 
+import javax.annotation.Nullable;
+
 /**
  * Wrapper of multiple services that are provided by a controller.
  * Placing them in here keeps the actual {@link Controller} class minimal.
@@ -25,6 +28,7 @@ public class Services {
 	private InheritanceGraph inheritanceGraph;
 	private WorkspaceSymbolSolver symbolSolver;
 	private JavaParserHelper javaParserHelper;
+	private CallGraphRegistry callGraphRegistry;
 
 	/**
 	 * Initialize services.
@@ -95,6 +99,14 @@ public SsvmIntegration getSsvmIntegration() {
 	 */
 	public WorkspaceTreeService getTreeService() {
 		return treeService;
+  }
+
+  /**
+	 * @return Registry for all calls made mby all methods of the {@link Controller#getWorkspace() current workspace}.
+	 * If no workspace is set, then this will be {@code null}.
+	 */
+	public @Nullable CallGraphRegistry getCallGraphRegistry() {
+		return callGraphRegistry;
 	}
 
 	/**
@@ -114,12 +126,14 @@ void updateWorkspace(Workspace workspace) {
 			javaParserHelper = null;
 			ssvmIntegration = null;
 			treeService = null;
+			callGraphRegistry = null;
 		} else {
 			inheritanceGraph = new InheritanceGraph(workspace);
 			symbolSolver = WorkspaceSymbolSolver.create(workspace);
 			javaParserHelper = JavaParserHelper.create(symbolSolver);
 			ssvmIntegration = new SsvmIntegration(workspace);
 			treeService = new WorkspaceTreeService(workspace);
+			callGraphRegistry = CallGraphRegistry.createAndLoad(workspace);
 		}
 	}
 }

recaf-core/src/main/java/me/coley/recaf/graph/call/CallGraphRegistry.java

@@ -0,0 +1,251 @@
+package me.coley.recaf.graph.call;
+
+import dev.xdark.jlinker.LinkResolver;
+import dev.xdark.jlinker.Result;
+import me.coley.recaf.code.ClassInfo;
+import me.coley.recaf.code.FieldInfo;
+import me.coley.recaf.code.MethodInfo;
+import me.coley.recaf.util.MemoizedFunction;
+import me.coley.recaf.workspace.Workspace;
+import me.coley.recaf.workspace.resource.Resources;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+public final class CallGraphRegistry {
+	private Map<ClassInfo, Map<Descriptor, MethodInfo>> methodMap = new HashMap<>();
+	private final Map<MethodInfo, MutableCallGraphVertex> vertexMap = new HashMap<>();
+	private final Workspace workspace;
+
+	public CallGraphRegistry(Workspace workspace) {
+		this.workspace = workspace;
+	}
+
+	public static CallGraphRegistry createAndLoad(Workspace workspace) {
+		CallGraphRegistry registry = new CallGraphRegistry(workspace);
+		registry.load();
+		return registry;
+	}
+
+	public @Nullable CallGraphVertex getVertex(MethodInfo info) {
+		return vertexMap.get(info);
+	}
+
+	public void load() {
+		Resources resources = workspace.getResources();
+		Function<String, ClassInfo> classInfoFromPathResolver = MemoizedFunction.memoize(path -> workspace.getResources().getClass(path));
+		Function<ClassInfo, BiFunction<String, String, MethodInfo>> methodMapGetter
+				= MemoizedFunction.memoize(clazz -> (name, descriptor) -> getMethodMap(clazz).get(new Descriptor(name, descriptor)));
+		// seems like a hack tho, needs feedback!
+		resources.getClasses().forEach(info -> visitClass(info, classInfoFromPathResolver, methodMapGetter));
+		methodMap = null;
+	}
+
+	private Map<Descriptor, MethodInfo> getMethodMap(ClassInfo info) {
+		return methodMap.computeIfAbsent(info, k ->
+				k.getMethods()
+						.stream()
+						.collect(Collectors.toMap(
+								Descriptor::new,
+								Function.identity()
+						))
+		);
+	}
+
+	private void visitClass(
+			ClassInfo info,
+			Function<String, ClassInfo> classInfoFromPathResolver,
+			Function<ClassInfo, BiFunction<String, String, MethodInfo>> otherMethodInfoResolver
+	) {
+		visitClass(info, classInfoFromPathResolver, otherMethodInfoResolver.apply(info), otherMethodInfoResolver);
+	}
+
+	private void visitClass(
+			ClassInfo info,
+			Function<String, ClassInfo> classInfoFromPathResolver,
+			BiFunction<String, String, MethodInfo> thisClassMethodInfoResolver,
+			Function<ClassInfo, BiFunction<String, String, MethodInfo>> otherMethodInfoResolver
+	) {
+		LinkResolver<ClassInfo, MethodInfo, FieldInfo> resolver = LinkResolver.jvm();
+		info.getClassReader().accept(new ClassVisitor(Opcodes.ASM9) {
+			@Override
+			public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+				MethodInfo info = thisClassMethodInfoResolver.apply(name, descriptor);
+				if (info == null)
+					return null;
+				Map<MethodInfo, MutableCallGraphVertex> vertexMap = CallGraphRegistry.this.vertexMap;
+				MutableCallGraphVertex vertex = vertexMap.computeIfAbsent(info, MutableCallGraphVertex::new);
+				if (!vertex.visited) {
+					vertex.visited = true;
+					return new MethodVisitor(Opcodes.ASM9, super.visitMethod(access, name, descriptor, signature, exceptions)) {
+						@Override
+						public void visitMethodInsn(int opcode, String owner, String name, String descriptor, boolean isInterface) {
+							ClassInfo callClassInfo = classInfoFromPathResolver.apply(owner);
+							if (callClassInfo == null)
+								return;
+							MethodInfo call = otherMethodInfoResolver.apply(callClassInfo).apply(name, descriptor);
+							if (call == null) {
+								Result<dev.xdark.jlinker.Resolution<ClassInfo, MethodInfo>> result;
+								switch (opcode) {
+									case Opcodes.INVOKESPECIAL:
+									case Opcodes.INVOKEVIRTUAL:
+										result = resolver.resolveVirtualMethod(classInfo(callClassInfo, classInfoFromPathResolver), name, descriptor);
+										break;
+									case Opcodes.INVOKESTATIC:
+										result = resolver.resolveStaticMethod(classInfo(callClassInfo, classInfoFromPathResolver), name, descriptor);
+										break;
+									case Opcodes.INVOKEINTERFACE:
+										result = resolver.resolveInterfaceMethod(classInfo(callClassInfo, classInfoFromPathResolver), name, descriptor);
+										break;
+									default:
+										throw new IllegalArgumentException("Opcode in visitMethodInsn must be INVOKEVIRTUAL, INVOKESPECIAL, INVOKESTATIC or INVOKEINTERFACE.");
+								}
+								if (result.isSuccess()) {
+									call = result.value().member().innerValue();
+								}
+								// should it log on else here? or would it be spam?
+							}
+							if (call == null)
+								return;
+							MutableCallGraphVertex nestedVertex = vertexMap.computeIfAbsent(call, MutableCallGraphVertex::new);
+							vertex.getCalls().add(nestedVertex);
+							nestedVertex.getCallers().add(vertex);
+						}
+					};
+				}
+				return null;
+			}
+		}, ClassReader.SKIP_FRAMES | ClassReader.SKIP_DEBUG);
+	}
+
+	private static final class Descriptor {
+		private final String name, desc;
+
+		Descriptor(String name, String desc) {
+			this.name = name;
+			this.desc = desc;
+		}
+
+		Descriptor(MethodInfo info) {
+			this.name = info.getName();
+			this.desc = info.getDescriptor();
+		}
+
+		@Override
+		public boolean equals(Object o) {
+			if (this == o) return true;
+			if (o == null || getClass() != o.getClass()) return false;
+
+			Descriptor that = (Descriptor) o;
+
+			if (!name.equals(that.name)) return false;
+			return desc.equals(that.desc);
+		}
+
+		@Override
+		public int hashCode() {
+			int result = name.hashCode();
+			result = 31 * result + desc.hashCode();
+			return result;
+		}
+	}
+
+	private static dev.xdark.jlinker.ClassInfo<ClassInfo> classInfo(@Nonnull ClassInfo node, Function<String, ClassInfo> fn) {
+		return new dev.xdark.jlinker.ClassInfo<>() {
+			@Override
+			public ClassInfo innerValue() {
+				return node;
+			}
+
+			@Override
+			public int accessFlags() {
+				return node.getAccess();
+			}
+
+			@Override
+			public dev.xdark.jlinker.ClassInfo<ClassInfo> superClass() {
+				String superName = node.getSuperName();
+				if (superName == null) return null;
+				final ClassInfo classInfo = fn.apply(superName);
+				return classInfo == null ? null : classInfo(classInfo, fn);
+			}
+
+			@Override
+			public List<dev.xdark.jlinker.ClassInfo<ClassInfo>> interfaces() {
+				return node.getInterfaces().stream().map(x -> {
+					final ClassInfo classInfo = fn.apply(x);
+					return classInfo == null ? null : classInfo(classInfo, fn);
+				}).filter(Objects::nonNull).collect(Collectors.toList());
+			}
+
+			@Override
+			public dev.xdark.jlinker.MemberInfo<?> getMethod(String name, String descriptor) {
+				for (MethodInfo method : node.getMethods()) {
+					if (name.equals(method.getName()) && descriptor.equals(method.getDescriptor())) {
+						return methodInfo(method);
+					}
+				}
+				return null;
+			}
+
+			@Override
+			public dev.xdark.jlinker.MemberInfo<?> getField(String name, String descriptor) {
+				for (FieldInfo field : node.getFields()) {
+					if (name.equals(field.getName()) && descriptor.equals(field.getDescriptor())) {
+						return fieldInfo(field);
+					}
+				}
+				return null;
+			}
+		};
+	}
+
+	private static dev.xdark.jlinker.MemberInfo<MethodInfo> methodInfo(MethodInfo node) {
+		return new dev.xdark.jlinker.MemberInfo<>() {
+			@Override
+			public MethodInfo innerValue() {
+				return node;
+			}
+
+			@Override
+			public int accessFlags() {
+				return node.getAccess();
+			}
+
+			@Override
+			public boolean isPolymorphic() {
+				return false;
+			}
+		};
+	}
+
+	private static dev.xdark.jlinker.MemberInfo<FieldInfo> fieldInfo(FieldInfo node) {
+		return new dev.xdark.jlinker.MemberInfo<>() {
+			@Override
+			public FieldInfo innerValue() {
+				return node;
+			}
+
+			@Override
+			public int accessFlags() {
+				return node.getAccess();
+			}
+
+			@Override
+			public boolean isPolymorphic() {
+				return false;
+			}
+		};
+	}
+}

recaf-core/src/main/java/me/coley/recaf/graph/call/CallGraphVertex.java

@@ -0,0 +1,14 @@
+package me.coley.recaf.graph.call;
+
+import me.coley.recaf.code.MethodInfo;
+
+import java.util.Collection;
+
+public interface CallGraphVertex {
+
+	MethodInfo getMethodInfo();
+
+	Collection<CallGraphVertex> getCallers();
+
+	Collection<CallGraphVertex> getCalls();
+}

recaf-core/src/main/java/me/coley/recaf/graph/call/MutableCallGraphVertex.java

@@ -0,0 +1,31 @@
+package me.coley.recaf.graph.call;
+
+import me.coley.recaf.code.MethodInfo;
+
+import java.util.*;
+
+public final class MutableCallGraphVertex implements CallGraphVertex {
+	private final Set<CallGraphVertex> callers = Collections.newSetFromMap(new LinkedHashMap<>());
+	private final Set<CallGraphVertex> calls = Collections.newSetFromMap(new LinkedHashMap<>());
+	private final MethodInfo methodInfo;
+	boolean visited;
+
+	public MutableCallGraphVertex(MethodInfo methodInfo) {
+		this.methodInfo = methodInfo;
+	}
+
+	@Override
+	public MethodInfo getMethodInfo() {
+		return methodInfo;
+	}
+
+	@Override
+	public Collection<CallGraphVertex> getCallers() {
+		return callers;
+	}
+
+	@Override
+	public Collection<CallGraphVertex> getCalls() {
+		return calls;
+	}
+}