Skip to content

Commit

Permalink
feat: initial support for 'invoke-custom' instruction (#384)
Browse files Browse the repository at this point in the history
  • Loading branch information
skylot committed Jan 14, 2021
1 parent 778106c commit 3dfaec5
Show file tree
Hide file tree
Showing 31 changed files with 779 additions and 34 deletions.
2 changes: 1 addition & 1 deletion jadx-core/src/main/java/jadx/core/codegen/ClassGen.java
Original file line number Diff line number Diff line change
Expand Up @@ -254,7 +254,7 @@ public void addClassBody(CodeWriter clsCode, boolean printClassName) throws Code
private void addInnerClsAndMethods(CodeWriter clsCode) {
Stream.of(cls.getInnerClasses(), cls.getMethods())
.flatMap(Collection::stream)
.filter(node -> !node.contains(AFlag.DONT_GENERATE))
.filter(node -> !node.contains(AFlag.DONT_GENERATE) || fallback)
.sorted(Comparator.comparingInt(LineAttrNode::getSourceLine))
.forEach(node -> {
if (node instanceof ClassNode) {
Expand Down
92 changes: 91 additions & 1 deletion jadx-core/src/main/java/jadx/core/codegen/InsnGen.java
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import jadx.core.dex.instructions.IfNode;
import jadx.core.dex.instructions.IndexInsnNode;
import jadx.core.dex.instructions.InsnType;
import jadx.core.dex.instructions.InvokeCustomNode;
import jadx.core.dex.instructions.InvokeNode;
import jadx.core.dex.instructions.InvokeType;
import jadx.core.dex.instructions.NewArrayNode;
Expand Down Expand Up @@ -700,11 +701,15 @@ private void inlineAnonymousConstructor(CodeWriter code, ClassNode cls, Construc
}

private void makeInvoke(InvokeNode insn, CodeWriter code) throws CodegenException {
InvokeType type = insn.getInvokeType();
if (type == InvokeType.CUSTOM) {
makeInvokeLambda(code, (InvokeCustomNode) insn);
return;
}
MethodInfo callMth = insn.getCallMth();
MethodNode callMthNode = mth.root().deepResolveMethod(callMth);

int k = 0;
InvokeType type = insn.getInvokeType();
switch (type) {
case DIRECT:
case VIRTUAL:
Expand Down Expand Up @@ -746,6 +751,91 @@ private void makeInvoke(InvokeNode insn, CodeWriter code) throws CodegenExceptio
generateMethodArguments(code, insn, k, callMthNode);
}

private void makeInvokeLambda(CodeWriter code, InvokeCustomNode customNode) throws CodegenException {
if (fallback || !customNode.isInlineInsn()) {
makeSimpleLambda(code, customNode);
return;
}
MethodNode callMth = (MethodNode) customNode.getCallInsn().get(AType.METHOD_DETAILS);
makeInlinedLambdaMethod(code, customNode, callMth);
}

private void makeSimpleLambda(CodeWriter code, InvokeCustomNode customNode) {
try {
InsnNode callInsn = customNode.getCallInsn();
MethodInfo implMthInfo = customNode.getImplMthInfo();
int implArgsCount = implMthInfo.getArgsCount();
if (implArgsCount == 0) {
code.add("()");
} else {
code.add('(');
// rename lambda args
int callArgsCount = callInsn.getArgsCount();
int startArg = callArgsCount - implArgsCount;
if (startArg < 0) {
System.out.println();
}
for (int i = startArg; i < callArgsCount; i++) {
if (i != startArg) {
code.add(", ");
}
addArg(code, callInsn.getArg(i));
}
code.add(')');
}
code.add(" -> {");
if (fallback) {
code.add(" // ").add(implMthInfo.toString());
}
code.incIndent();
code.startLine();
if (!implMthInfo.getReturnType().isVoid()) {
code.add("return ");
}
makeInsn(callInsn, code, Flags.INLINE);
code.add(";");

code.decIndent();
code.startLine('}');
} catch (Exception e) {
throw new JadxRuntimeException("Failed to generate 'invoke-custom' instruction: " + e.getMessage(), e);
}
}

private void makeInlinedLambdaMethod(CodeWriter code, InvokeCustomNode customNode, MethodNode callMth) throws CodegenException {
MethodGen callMthGen = new MethodGen(mgen.getClassGen(), callMth);
NameGen nameGen = callMthGen.getNameGen();
nameGen.inheritUsedNames(this.mgen.getNameGen());

List<ArgType> implArgs = customNode.getImplMthInfo().getArgumentsTypes();
List<RegisterArg> callArgs = callMth.getArgRegs();
if (implArgs.isEmpty()) {
code.add("()");
} else {
int callArgsCount = callArgs.size();
int startArg = callArgsCount - implArgs.size();
for (int i = startArg; i < callArgsCount; i++) {
if (i != startArg) {
code.add(", ");
}
CodeVar argCodeVar = callArgs.get(i).getSVar().getCodeVar();
code.add(nameGen.assignArg(argCodeVar));
}
}
// force set external arg names into call method args
int extArgsCount = customNode.getArgsCount();
for (int i = 0; i < extArgsCount; i++) {
RegisterArg extArg = (RegisterArg) customNode.getArg(i);
callArgs.get(i).setName(extArg.getName());
}
code.add(" -> {");
code.incIndent();
callMthGen.addInstructions(code);

code.decIndent();
code.startLine('}');
}

@Nullable
private ClassInfo getClassForSuperCall(CodeWriter code, MethodInfo callMth) {
ClassNode useCls = mth.getParentClass();
Expand Down
4 changes: 4 additions & 0 deletions jadx-core/src/main/java/jadx/core/codegen/NameGen.java
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,10 @@ public NameGen(MethodNode mth, boolean fallback) {
addNamesUsedInClass();
}

public void inheritUsedNames(NameGen otherNameGen) {
varNames.addAll(otherNameGen.varNames);
}

private void addNamesUsedInClass() {
ClassNode parentClass = mth.getParentClass();
for (FieldNode field : parentClass.getFields()) {
Expand Down
7 changes: 7 additions & 0 deletions jadx-core/src/main/java/jadx/core/dex/info/MethodInfo.java
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

import org.jetbrains.annotations.Nullable;

import jadx.api.plugins.input.data.IMethodProto;
import jadx.api.plugins.input.data.IMethodRef;
import jadx.core.codegen.TypeGen;
import jadx.core.dex.instructions.args.ArgType;
Expand Down Expand Up @@ -52,6 +53,12 @@ public static MethodInfo fromDetails(RootNode root, ClassInfo declClass, String
return root.getInfoStorage().putMethod(newMth);
}

public static MethodInfo fromMethodProto(RootNode root, ClassInfo declClass, String name, IMethodProto proto) {
List<ArgType> args = Utils.collectionMap(proto.getArgTypes(), ArgType::parse);
ArgType returnType = ArgType.parse(proto.getReturnType());
return fromDetails(root, declClass, name, args, returnType);
}

public String makeSignature(boolean includeRetType) {
return makeSignature(false, includeRetType);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -418,6 +418,8 @@ private InsnNode decode(InsnData insn) throws DecodeException {
return invoke(insn, InvokeType.SUPER, false);
case INVOKE_VIRTUAL:
return invoke(insn, InvokeType.VIRTUAL, false);
case INVOKE_CUSTOM:
return invoke(insn, InvokeType.CUSTOM, false);

case INVOKE_DIRECT_RANGE:
return invoke(insn, InvokeType.DIRECT, true);
Expand All @@ -427,6 +429,8 @@ private InsnNode decode(InsnData insn) throws DecodeException {
return invoke(insn, InvokeType.SUPER, true);
case INVOKE_VIRTUAL_RANGE:
return invoke(insn, InvokeType.VIRTUAL, true);
case INVOKE_CUSTOM_RANGE:
return invoke(insn, InvokeType.CUSTOM, true);

case NEW_INSTANCE:
ArgType clsType = ArgType.parse(insn.getIndexAsType());
Expand Down Expand Up @@ -524,6 +528,9 @@ private InsnNode cast(InsnData insn, ArgType from, ArgType to) {
}

private InsnNode invoke(InsnData insn, InvokeType type, boolean isRange) {
if (type == InvokeType.CUSTOM) {
return InvokeCustomBuilder.build(method, insn, isRange);
}
MethodInfo mthInfo = MethodInfo.fromRef(root, insn.getIndexAsMethod());
return new InvokeNode(mthInfo, insn, type, isRange);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,128 @@
package jadx.core.dex.instructions;

import java.util.List;

import jadx.api.plugins.input.data.ICallSite;
import jadx.api.plugins.input.data.IMethodHandle;
import jadx.api.plugins.input.data.IMethodProto;
import jadx.api.plugins.input.data.IMethodRef;
import jadx.api.plugins.input.data.MethodHandleType;
import jadx.api.plugins.input.data.annotations.EncodedValue;
import jadx.api.plugins.input.insns.InsnData;
import jadx.core.dex.attributes.AFlag;
import jadx.core.dex.info.ClassInfo;
import jadx.core.dex.info.MethodInfo;
import jadx.core.dex.instructions.args.ArgType;
import jadx.core.dex.instructions.args.InsnArg;
import jadx.core.dex.instructions.args.NamedArg;
import jadx.core.dex.nodes.InsnNode;
import jadx.core.dex.nodes.MethodNode;
import jadx.core.dex.nodes.RootNode;
import jadx.core.utils.exceptions.JadxRuntimeException;

public class InvokeCustomBuilder {

public static InsnNode build(MethodNode mth, InsnData insn, boolean isRange) {
try {
ICallSite callSite = insn.getIndexAsCallSite();
callSite.load();
List<EncodedValue> values = callSite.getValues();
if (!checkLinkerMethod(values)) {
throw new JadxRuntimeException("Failed to process invoke-custom instruction: " + callSite);
}
IMethodHandle callMthHandle = (IMethodHandle) values.get(4).getValue();
MethodHandleType methodHandleType = callMthHandle.getType();
if (methodHandleType.isField()) {
throw new JadxRuntimeException("Not yet supported");
}
RootNode root = mth.root();
IMethodProto lambdaProto = (IMethodProto) values.get(2).getValue();
MethodInfo lambdaInfo = MethodInfo.fromMethodProto(root, mth.getParentClass().getClassInfo(), "", lambdaProto);

InvokeCustomNode invokeCustomNode = new InvokeCustomNode(lambdaInfo, insn, false, isRange);
invokeCustomNode.setHandleType(methodHandleType);

ClassInfo implCls = ClassInfo.fromType(root, lambdaInfo.getReturnType());
String implName = (String) values.get(1).getValue();
IMethodProto implProto = (IMethodProto) values.get(3).getValue();
invokeCustomNode.setImplMthInfo(MethodInfo.fromMethodProto(root, implCls, implName, implProto));

MethodInfo callMthInfo = MethodInfo.fromRef(root, callMthHandle.getMethodRef());

InvokeType invokeType = convertInvokeType(methodHandleType);
int callArgsCount = callMthInfo.getArgsCount();
InvokeNode callInsn = new InvokeNode(callMthInfo, invokeType, callArgsCount);
invokeCustomNode.setCallInsn(callInsn);

// copy insn args
int argsCount = invokeCustomNode.getArgsCount();
for (int i = 0; i < argsCount; i++) {
InsnArg arg = invokeCustomNode.getArg(i);
callInsn.addArg(arg.duplicate());
}
if (callArgsCount > argsCount) {
// fill remaining args with NamedArg
for (int i = argsCount; i < callArgsCount; i++) {
ArgType argType = callMthInfo.getArgumentsTypes().get(i);
callInsn.addArg(new NamedArg("v" + i, argType));
}
}

MethodNode callMth = root.resolveMethod(callMthInfo);
if (callMth != null) {
callInsn.addAttr(callMth);
if (callMth.getAccessFlags().isSynthetic()
&& callMth.getUseIn().size() <= 1
&& callMth.getParentClass().equals(mth.getParentClass())) {
// inline only synthetic methods from same class
callMth.add(AFlag.DONT_GENERATE);
invokeCustomNode.setInlineInsn(true);
}
}
// prevent args inlining into not generated invoke custom node
for (InsnArg arg : invokeCustomNode.getArguments()) {
arg.add(AFlag.DONT_INLINE);
}
return invokeCustomNode;
} catch (Exception e) {
throw new JadxRuntimeException("'invoke-custom' instruction processing error: " + e.getMessage(), e);
}
}

/**
* Expect LambdaMetafactory.metafactory method
*/
private static boolean checkLinkerMethod(List<EncodedValue> values) {
if (values.size() < 6) {
return false;
}
IMethodHandle methodHandle = (IMethodHandle) values.get(0).getValue();
if (methodHandle.getType() != MethodHandleType.INVOKE_STATIC) {
return false;
}
IMethodRef methodRef = methodHandle.getMethodRef();
if (!methodRef.getName().equals("metafactory")) {
return false;
}
if (!methodRef.getParentClassType().equals("Ljava/lang/invoke/LambdaMetafactory;")) {
return false;
}
return true;
}

private static InvokeType convertInvokeType(MethodHandleType type) {
switch (type) {
case INVOKE_STATIC:
return InvokeType.STATIC;
case INVOKE_INSTANCE:
return InvokeType.VIRTUAL;
case INVOKE_DIRECT:
return InvokeType.DIRECT;
case INVOKE_INTERFACE:
return InvokeType.INTERFACE;

default:
throw new JadxRuntimeException("Unsupported method handle type: " + type);
}
}
}
Loading

0 comments on commit 3dfaec5

Please sign in to comment.