more custom source transformers

1 files changed, 197 insertions, 0 deletions

diff --git a/src/main/java/oml/ast/transformers/VaragsFixer.java b/src/main/java/oml/ast/transformers/VaragsFixer.java
new file mode 100644
index 00000000..cd711ae0
--- /dev/null
+++ b/src/main/java/oml/ast/transformers/VaragsFixer.java
@@ -0,0 +1,197 @@
+package oml.ast.transformers;
+
+import com.strobel.assembler.metadata.MemberReference;
+import com.strobel.assembler.metadata.MetadataFilters;
+import com.strobel.assembler.metadata.MetadataHelper;
+import com.strobel.assembler.metadata.MethodBinder;
+import com.strobel.assembler.metadata.MethodDefinition;
+import com.strobel.assembler.metadata.MethodReference;
+import com.strobel.assembler.metadata.TypeReference;
+import com.strobel.core.StringUtilities;
+import com.strobel.core.VerifyArgument;
+import com.strobel.decompiler.DecompilerContext;
+import com.strobel.decompiler.languages.java.ast.ArrayCreationExpression;
+import com.strobel.decompiler.languages.java.ast.ArrayInitializerExpression;
+import com.strobel.decompiler.languages.java.ast.AstNode;
+import com.strobel.decompiler.languages.java.ast.AstNodeCollection;
+import com.strobel.decompiler.languages.java.ast.CastExpression;
+import com.strobel.decompiler.languages.java.ast.ContextTrackingVisitor;
+import com.strobel.decompiler.languages.java.ast.DepthFirstAstVisitor;
+import com.strobel.decompiler.languages.java.ast.Expression;
+import com.strobel.decompiler.languages.java.ast.InvocationExpression;
+import com.strobel.decompiler.languages.java.ast.JavaResolver;
+import com.strobel.decompiler.languages.java.ast.Keys;
+import com.strobel.decompiler.languages.java.ast.MemberReferenceExpression;
+import com.strobel.decompiler.languages.java.ast.ObjectCreationExpression;
+import com.strobel.decompiler.languages.java.ast.transforms.IAstTransform;
+import com.strobel.decompiler.semantics.ResolveResult;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Created by Thiakil on 12/07/2018.
+ */
+public class VaragsFixer implements IAstTransform {
+        private final DecompilerContext _context;
+
+        public VaragsFixer(final DecompilerContext context) {
+                _context = VerifyArgument.notNull(context, "context");
+        }
+
+        @Override
+        public void run(AstNode compilationUnit) {
+                compilationUnit.acceptVisitor(new Visitor(_context), null);
+        }
+
+        class Visitor extends ContextTrackingVisitor<Void> {
+                private final JavaResolver _resolver;
+                protected Visitor(DecompilerContext context) {
+                        super(context);
+                        _resolver = new JavaResolver(context);
+                }
+
+                //remove `new Object[0]` on varagrs as the normal tranformer doesnt do them
+                @Override
+                public Void visitInvocationExpression(InvocationExpression node, Void data) {
+                        super.visitInvocationExpression(node, data);
+                        MemberReference definition = node.getUserData(Keys.MEMBER_REFERENCE);
+                        if (definition instanceof MethodDefinition && ((MethodDefinition) definition).isVarArgs()){
+                                AstNodeCollection<Expression> arguments = node.getArguments();
+                                Expression lastParam = arguments.lastOrNullObject();
+                                if (!lastParam.isNull() && lastParam instanceof ArrayCreationExpression){
+                                        ArrayCreationExpression varargArray = (ArrayCreationExpression)lastParam;
+                                        if (varargArray.getInitializer().isNull() || varargArray.getInitializer().getElements().isEmpty()){
+                                                lastParam.remove();
+                                        } else {
+                                                for (Expression e : varargArray.getInitializer().getElements()){
+                                                        arguments.insertBefore(varargArray, e.clone());
+                                                }
+                                                varargArray.remove();
+                                        }
+                                }
+                        }
+                        return null;
+                }
+
+                //applies the vararg transform to object creation
+                @Override
+                public Void visitObjectCreationExpression(ObjectCreationExpression node, Void data) {
+                        super.visitObjectCreationExpression(node, data);
+                        final AstNodeCollection<Expression> arguments = node.getArguments();
+                        final Expression lastArgument = arguments.lastOrNullObject();
+
+                        Expression arrayArg = lastArgument;
+
+                        if (arrayArg instanceof CastExpression)
+                                arrayArg = ((CastExpression) arrayArg).getExpression();
+
+                        if (arrayArg == null ||
+                                        arrayArg.isNull() ||
+                                        !(arrayArg instanceof ArrayCreationExpression &&
+                                                        node.getTarget() instanceof MemberReferenceExpression)) {
+
+                                return null;
+                        }
+
+                        final ArrayCreationExpression newArray = (ArrayCreationExpression) arrayArg;
+                        final MemberReferenceExpression target = (MemberReferenceExpression) node.getTarget();
+
+                        if (!newArray.getAdditionalArraySpecifiers().hasSingleElement()) {
+                                return null;
+                        }
+
+                        final MethodReference method = (MethodReference) node.getUserData(Keys.MEMBER_REFERENCE);
+
+                        if (method == null) {
+                                return null;
+                        }
+
+                        final MethodDefinition resolved = method.resolve();
+
+                        if (resolved == null || !resolved.isVarArgs()) {
+                                return null;
+                        }
+
+                        final List<MethodReference> candidates;
+                        final Expression invocationTarget = target.getTarget();
+
+                        if (invocationTarget == null || invocationTarget.isNull()) {
+                                candidates = MetadataHelper.findMethods(
+                                                context.getCurrentType(),
+                                                MetadataFilters.matchName(resolved.getName())
+                                );
+                        }
+                        else {
+                                final ResolveResult targetResult = _resolver.apply(invocationTarget);
+
+                                if (targetResult == null || targetResult.getType() == null) {
+                                        return null;
+                                }
+
+                                candidates = MetadataHelper.findMethods(
+                                                targetResult.getType(),
+                                                MetadataFilters.matchName(resolved.getName())
+                                );
+                        }
+
+                        final List<TypeReference> argTypes = new ArrayList<>();
+
+                        for (final Expression argument : arguments) {
+                                final ResolveResult argResult = _resolver.apply(argument);
+
+                                if (argResult == null || argResult.getType() == null) {
+                                        return null;
+                                }
+
+                                argTypes.add(argResult.getType());
+                        }
+
+                        final MethodBinder.BindResult c1 = MethodBinder.selectMethod(candidates, argTypes);
+
+                        if (c1.isFailure() || c1.isAmbiguous()) {
+                                return null;
+                        }
+
+                        argTypes.remove(argTypes.size() - 1);
+
+                        final ArrayInitializerExpression initializer = newArray.getInitializer();
+                        final boolean hasElements = !initializer.isNull() && !initializer.getElements().isEmpty();
+
+                        if (hasElements) {
+                                for (final Expression argument : initializer.getElements()) {
+                                        final ResolveResult argResult = _resolver.apply(argument);
+
+                                        if (argResult == null || argResult.getType() == null) {
+                                                return null;
+                                        }
+
+                                        argTypes.add(argResult.getType());
+                                }
+                        }
+
+                        final MethodBinder.BindResult c2 = MethodBinder.selectMethod(candidates, argTypes);
+
+                        if (c2.isFailure() ||
+                                        c2.isAmbiguous() ||
+                                        !StringUtilities.equals(c2.getMethod().getErasedSignature(), c1.getMethod().getErasedSignature())) {
+
+                                return null;
+                        }
+
+                        lastArgument.remove();
+
+                        if (!hasElements) {
+                                lastArgument.remove();
+                                return null;
+                        }
+
+                        for (final Expression newArg : initializer.getElements()) {
+                                newArg.remove();
+                                arguments.add(newArg);
+                        }
+
+                        return null;
+                }
+        }
+}