/*
 * Copyright (C) 2015 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package dodola.anole.lib;


import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.GeneratorAdapter;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.LabelNode;
import org.objectweb.asm.tree.LineNumberNode;
import org.objectweb.asm.tree.LocalVariableNode;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.MethodNode;
import org.objectweb.asm.tree.TryCatchBlockNode;
import org.objectweb.asm.tree.VarInsnNode;
import org.objectweb.asm.tree.analysis.Analyzer;
import org.objectweb.asm.tree.analysis.AnalyzerException;
import org.objectweb.asm.tree.analysis.BasicInterpreter;
import org.objectweb.asm.tree.analysis.BasicValue;
import org.objectweb.asm.tree.analysis.Frame;
import org.objectweb.asm.tree.analysis.Value;

import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * Utilities to detect and manipulate constructor methods.
 * <p>
 * A constructor of a non static inner class usually has the form:
 * <p>
 * ALOAD_0              // push this to the stack
 * ...                  // Code to set up $this
 * ALOAD_0              // push this to the stack
 * ...                  // Code to set up the arguments (aka "args") for the delegation
 * ...                  // via super() or this(). Note that here we can have INVOKESPECIALS
 * ...                  // for all the new calls here.
 * INVOKESPECIAL <init> // super() or this() call
 * ...                  // the "body" of the constructor goes here.
 * <p>
 * This class has the utilities to detect which instruction is the right INVOKESPECIAL call before
 * the "body".
 */
public class ConstructorDelegationDetector {

    /**
     * A specialized value used to track the first local variable (this) on the
     * constructor.
     */
    public static class LocalValue extends BasicValue {
        public LocalValue(Type type) {
            super(type);
        }

        @Override
        public String toString() {
            return "*";
        }
    }

    /**
     * A deconstructed constructor, split up in the parts mentioned above.
     */
    static class Constructor {

        /**
         * The last LOAD_0 instruction of the original code, before the call to the delegated
         * constructor.
         */
        public final VarInsnNode loadThis;

        /**
         * Line number of LOAD_0. Used to set the line number in the generated constructor call
         * so that a break point may be set at this(...) or super(...)
         */
        public final int lineForLoad;

        /**
         * The "args" part of the constructor. Described above.
         */
        public final MethodNode args;

        /**
         * The INVOKESPECIAL instruction of the original code that calls the delegation.
         */
        public final MethodInsnNode delegation;

        /**
         * A copy of the body of the constructor.
         */
        public final MethodNode body;

        Constructor(VarInsnNode loadThis, int lineForLoad, MethodNode args, MethodInsnNode delegation, MethodNode body) {
            this.loadThis = loadThis;
            this.lineForLoad = lineForLoad;
            this.args = args;
            this.delegation = delegation;
            this.body = body;
        }
    }

    /**
     * Deconstruct a constructor into its components and adds the necessary code to link the components
     * later. The generated bytecode does not correspond exactly to this code, but in essence, for
     * a constructor of this form:
     * <p/>
     * <code>
     * <init>(int x) {
     * super(x = 1, expr2() ? 3 : 7)
     * doSomething(x)
     * }
     * </code>
     * <p/>
     * it creates the two parts:
     * <code>
     * Object[] init$args(Object[] locals, int x) {
     * Object[] args = new Object[2];
     * args[0] = (x = 1)
     * args[1] = expr2() ? 3 : 7;
     * locals[0] = x;
     * return new Object[] {"myclass.<init>(I;I;)V", args};
     * }
     * <p>
     * void init$body(int x) {
     * doSomething(x);
     * }
     * </code>
     *
     * @param owner  the owning class.
     * @param method the constructor method.
     */
    public static Constructor deconstruct(String owner, MethodNode method) {
        // Basic interpreter uses BasicValue.REFERENCE_VALUE for all object types. However
        // we need to distinguish one in particular. The value of the local variable 0, ie. the
        // uninitialized this. By doing it this way we ensure that whenever there is a ALOAD_0
        // a LocalValue instance will be on the stack.
        BasicInterpreter interpreter = new BasicInterpreter() {
            boolean done = false;

            @Override
            // newValue is called first to initialize the frame values of all the local variables
            // we intercept the first one to create our own special value.
            public BasicValue newValue(Type type) {
                if (type == null) {
                    return BasicValue.UNINITIALIZED_VALUE;
                } else if (type.getSort() == Type.VOID) {
                    return null;
                } else {
                    // If this is the first value created (i.e. the first local variable)
                    // we use a special marker.
                    BasicValue ret = done ? super.newValue(type) : new LocalValue(type);
                    done = true;
                    return ret;
                }
            }
        };

        Analyzer analyzer = new Analyzer(interpreter);
        AbstractInsnNode[] instructions = method.instructions.toArray();
        try {
            Frame[] frames = analyzer.analyze(owner, method);
            if (frames.length != instructions.length) {
                // Should never happen.
                throw new IllegalStateException(
                        "The number of frames is not equals to the number of instructions");
            }
            VarInsnNode lastThis = null;
            int stackAtThis = -1;
            boolean poppedThis = false;
            // Records the most recent line number encountered. For javac, there should always be
            // a line number node before the call of interest to this(...) or super(...). For robustness,
            // -1 is recorded as a sentinel to indicate this assumption didn't hold. Upstream consumers
            // should check for -1 and recover in a reasonable way (for example, don't set the line
            // number in generated code).
            int recentLine = -1;
            for (int i = 0; i < instructions.length; i++) {
                AbstractInsnNode insn = instructions[i];
                Frame frame = frames[i];
                if (frame.getStackSize() < stackAtThis) {
                    poppedThis = true;
                }
                if (insn instanceof MethodInsnNode) {
                    // TODO: Do we need to check that the stack is empty after this super call?
                    MethodInsnNode methodhInsn = (MethodInsnNode) insn;
                    Type[] types = Type.getArgumentTypes(methodhInsn.desc);
                    Value value = frame.getStack(frame.getStackSize() - types.length - 1);
                    if (value instanceof LocalValue && methodhInsn.name.equals("<init>")) {
                        if (poppedThis) {
                            throw new IllegalStateException("Unexpected constructor structure.");
                        }
                        return split(owner, method, lastThis, methodhInsn, recentLine);
                    }
                } else if (insn instanceof VarInsnNode) {
                    VarInsnNode var = (VarInsnNode) insn;
                    if (var.var == 0) {
                        lastThis = var;
                        stackAtThis = frame.getStackSize();
                        poppedThis = false;
                    }
                } else if (insn instanceof LineNumberNode) {
                    // Record the most recent line number encountered so that call to this(...)
                    // or super(...) has line number information. Ultimately used to emit a line
                    // number in the generated code.
                    LineNumberNode lineNumberNode = (LineNumberNode) insn;
                    recentLine = lineNumberNode.line;
                }
            }
            throw new IllegalStateException("Unexpected constructor structure.");
        } catch (AnalyzerException e) {
            throw new IllegalStateException(e);
        }
    }

    /**
     * Splits the constructor in two methods, the "set up" and the "body" parts (see above).
     */
    private static Constructor split(String owner, MethodNode method, VarInsnNode loadThis, MethodInsnNode delegation, int loadThisLine) {
        String[] exceptions = ((List<String>) method.exceptions).toArray(new String[method.exceptions.size()]);
        String newDesc = method.desc.replaceAll("\\((.*)\\)V", "([Ljava/lang/Object;$1)Ljava/lang/Object;");

        MethodNode initArgs = new MethodNode(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC, "init$args", newDesc, null, exceptions);
        AbstractInsnNode insn = loadThis.getNext();
        while (insn != delegation) {
            insn.accept(initArgs);
            insn = insn.getNext();
        }
        LabelNode labelBefore = new LabelNode();
        labelBefore.accept(initArgs);

        GeneratorAdapter mv = new GeneratorAdapter(initArgs, initArgs.access, initArgs.name, initArgs.desc);
        // Copy the arguments back to the argument array
        // The init_args part cannot access the "this" object and can have side effects on the
        // local variables. Because of this we use the first argument (which we want to keep
        // so all the other arguments remain unchanged) as a reference to the array where to
        // return the values of the modified local variables.
        Type[] types = Type.getArgumentTypes(initArgs.desc);
        int stack = 1; // Skip the first one which is a reference to the local array.
        for (int i = 1; i < types.length; i++) {
            Type type = types[i];
            // This is not this, but the array of local arguments final values.
            mv.visitVarInsn(Opcodes.ALOAD, 0);
            mv.push(i);
            mv.visitVarInsn(type.getOpcode(Opcodes.ILOAD), stack);
            mv.box(type);
            mv.arrayStore(Type.getType(Object.class));
            stack += type.getSize();
        }
        // Create the args array with the values to send to the delegated constructor
        Type[] returnTypes = Type.getArgumentTypes(delegation.desc);
        // The extra element for the qualified name of the constructor.
        mv.push(returnTypes.length + 1);
        mv.newArray(Type.getType(Object.class));
        int args = mv.newLocal(Type.getType("[Ljava/lang/Object;"));
        mv.storeLocal(args);
        for (int i = returnTypes.length - 1; i >= 0; i--) {
            Type type = returnTypes[i];
            mv.loadLocal(args);
            mv.swap(type, Type.getType(Object.class));
            mv.push(i + 1);
            mv.swap(type, Type.INT_TYPE);
            mv.box(type);
            mv.arrayStore(Type.getType(Object.class));
        }

        // Store the qualified name of the constructor in the first element of the array.
        mv.loadLocal(args);
        mv.push(0);
        mv.push(delegation.owner + "." + delegation.desc); // Name of the constructor to be called.
        mv.arrayStore(Type.getType(Object.class));

        mv.loadLocal(args);
        mv.returnValue();

        newDesc = method.desc.replace("(", "(L" + owner + ";");
        MethodNode body = new MethodNode(Opcodes.ACC_PUBLIC | Opcodes.ACC_STATIC,
                "init$body", newDesc, null, exceptions);
        LabelNode labelAfter = new LabelNode();
        labelAfter.accept(body);
        Set<LabelNode> bodyLabels = new HashSet<LabelNode>();

        insn = delegation.getNext();
        while (insn != null) {
            if (insn instanceof LabelNode) {
                bodyLabels.add((LabelNode) insn);
            }
            insn.accept(body);
            insn = insn.getNext();
        }

        // manually transfer the exception table from the existing constructor to the new
        // "init$body" method. The labels were transferred just above so we can reuse them.

        //noinspection unchecked
        for (TryCatchBlockNode tryCatch : (List<TryCatchBlockNode>) method.tryCatchBlocks) {
            tryCatch.accept(body);
        }

        //noinspection unchecked
        for (LocalVariableNode variable : (List<LocalVariableNode>) method.localVariables) {
            boolean startsInBody = bodyLabels.contains(variable.start);
            boolean endsInBody = bodyLabels.contains(variable.end);
            if (!startsInBody && !endsInBody) {
                if (variable.index != 0) { // '#0' on init$args is not 'this'
                    variable.accept(initArgs);
                }
            } else if (startsInBody && endsInBody) {
                variable.accept(body);
            } else if (!startsInBody && endsInBody) {
                // The variable spans from the args to the end of the method, create two:
                if (variable.index != 0) { // '#0' on init$args is not 'this'
                    LocalVariableNode var0 = new LocalVariableNode(variable.name,
                            variable.desc, variable.signature,
                            variable.start, labelBefore, variable.index);
                    var0.accept(initArgs);
                }
                LocalVariableNode var1 = new LocalVariableNode(variable.name,
                        variable.desc, variable.signature,
                        labelAfter, variable.end, variable.index);
                var1.accept(body);
            } else {
                throw new IllegalStateException("Local variable starts after it ends.");
            }
        }

        return new Constructor(loadThis, loadThisLine, initArgs, delegation, body);
    }
}