/*
 * Copyright (C) 2017 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 lockedregioncodeinjection;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.List;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.commons.TryCatchBlockSorter;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.InsnList;
import org.objectweb.asm.tree.LabelNode;
import org.objectweb.asm.tree.MethodInsnNode;
import org.objectweb.asm.tree.MethodNode;
import org.objectweb.asm.tree.TryCatchBlockNode;
import org.objectweb.asm.tree.analysis.Analyzer;
import org.objectweb.asm.tree.analysis.AnalyzerException;
import org.objectweb.asm.tree.analysis.BasicValue;
import org.objectweb.asm.tree.analysis.Frame;

import static com.google.common.base.Preconditions.checkElementIndex;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;

/**
 * This visitor does two things:
 *
 * 1. Finds all the MONITOR_ENTER / MONITOR_EXIT in the byte code and insert the corresponding pre
 * and post methods calls should it matches one of the given target type in the Configuration.
 *
 * 2. Find all methods that are synchronized and insert pre method calls in the beginning and post
 * method calls just before all return instructions.
 */
class LockFindingClassVisitor extends ClassVisitor {
    private String className = null;
    private final List<LockTarget> targets;

    public LockFindingClassVisitor(List<LockTarget> targets, ClassVisitor chain) {
        super(Utils.ASM_VERSION, chain);
        this.targets = targets;
    }

    @Override
    public MethodVisitor visitMethod(int access, String name, String desc, String signature,
            String[] exceptions) {
        assert this.className != null;
        MethodNode mn = new TryCatchBlockSorter(null, access, name, desc, signature, exceptions);
        MethodVisitor chain = super.visitMethod(access, name, desc, signature, exceptions);
        return new LockFindingMethodVisitor(this.className, mn, chain);
    }

    @Override
    public void visit(int version, int access, String name, String signature, String superName,
            String[] interfaces) {
        this.className = name;
        super.visit(version, access, name, signature, superName, interfaces);
    }

    class LockFindingMethodVisitor extends MethodVisitor {
        private String owner;
        private MethodVisitor chain;

        public LockFindingMethodVisitor(String owner, MethodNode mn, MethodVisitor chain) {
            super(Utils.ASM_VERSION, mn);
            assert owner != null;
            this.owner = owner;
            this.chain = chain;
        }

        @SuppressWarnings("unchecked")
        @Override
        public void visitEnd() {
            MethodNode mn = (MethodNode) mv;

            Analyzer a = new Analyzer(new LockTargetStateAnalysis(targets));

            LockTarget ownerMonitor = null;
            if ((mn.access & Opcodes.ACC_SYNCHRONIZED) != 0) {
                for (LockTarget t : targets) {
                    if (t.getTargetDesc().equals("L" + owner + ";")) {
                        ownerMonitor = t;
                    }
                }
            }

            try {
                a.analyze(owner, mn);
            } catch (AnalyzerException e) {
                throw new RuntimeException("Locked region code injection: " + e.getMessage(), e);
            }
            InsnList instructions = mn.instructions;

            Frame[] frames = a.getFrames();
            List<Frame> frameMap = new LinkedList<>();
            frameMap.addAll(Arrays.asList(frames));

            List<List<TryCatchBlockNode>> handlersMap = new LinkedList<>();

            for (int i = 0; i < instructions.size(); i++) {
                handlersMap.add(a.getHandlers(i));
            }

            if (ownerMonitor != null) {
                AbstractInsnNode s = instructions.getFirst();
                MethodInsnNode call = new MethodInsnNode(Opcodes.INVOKESTATIC,
                        ownerMonitor.getPreOwner(), ownerMonitor.getPreMethod(), "()V", false);
                insertMethodCallBefore(mn, frameMap, handlersMap, s, 0, call);
            }

            for (int i = 0; i < instructions.size(); i++) {
                AbstractInsnNode s = instructions.get(i);

                if (s.getOpcode() == Opcodes.MONITORENTER) {
                    Frame f = frameMap.get(i);
                    BasicValue operand = (BasicValue) f.getStack(f.getStackSize() - 1);
                    if (operand instanceof LockTargetState) {
                        LockTargetState state = (LockTargetState) operand;
                        for (int j = 0; j < state.getTargets().size(); j++) {
                            LockTarget target = state.getTargets().get(j);
                            MethodInsnNode call = new MethodInsnNode(Opcodes.INVOKESTATIC,
                                    target.getPreOwner(), target.getPreMethod(), "()V", false);
                            insertMethodCallAfter(mn, frameMap, handlersMap, s, i, call);
                        }
                    }
                }

                if (s.getOpcode() == Opcodes.MONITOREXIT) {
                    Frame f = frameMap.get(i);
                    BasicValue operand = (BasicValue) f.getStack(f.getStackSize() - 1);
                    if (operand instanceof LockTargetState) {
                        LockTargetState state = (LockTargetState) operand;
                        for (int j = 0; j < state.getTargets().size(); j++) {
                            // The instruction after a monitor_exit should be a label for the end of the implicit
                            // catch block that surrounds the synchronized block to call monitor_exit when an exception
                            // occurs.
                            checkState(instructions.get(i + 1).getType() == AbstractInsnNode.LABEL,
                                "Expected to find label after monitor exit");

                            int labelIndex = i + 1;
                            checkElementIndex(labelIndex, instructions.size());

                            LabelNode label = (LabelNode)instructions.get(labelIndex);

                            checkNotNull(handlersMap.get(i));
                            checkElementIndex(0, handlersMap.get(i).size());
                            checkState(handlersMap.get(i).get(0).end == label,
                                "Expected label to be the end of monitor exit's try block");

                            LockTarget target = state.getTargets().get(j);
                            MethodInsnNode call = new MethodInsnNode(Opcodes.INVOKESTATIC,
                                    target.getPostOwner(), target.getPostMethod(), "()V", false);
                            insertMethodCallAfter(mn, frameMap, handlersMap, label, labelIndex, call);
                        }
                    }
                }

                if (ownerMonitor != null && (s.getOpcode() == Opcodes.RETURN
                        || s.getOpcode() == Opcodes.ARETURN || s.getOpcode() == Opcodes.DRETURN
                        || s.getOpcode() == Opcodes.FRETURN || s.getOpcode() == Opcodes.IRETURN)) {
                    MethodInsnNode call =
                            new MethodInsnNode(Opcodes.INVOKESTATIC, ownerMonitor.getPostOwner(),
                                    ownerMonitor.getPostMethod(), "()V", false);
                    insertMethodCallBefore(mn, frameMap, handlersMap, s, i, call);
                    i++; // Skip ahead. Otherwise, we will revisit this instruction again.
                }
            }
            super.visitEnd();
            mn.accept(chain);
        }
    }

    public static void insertMethodCallBefore(MethodNode mn, List<Frame> frameMap,
            List<List<TryCatchBlockNode>> handlersMap, AbstractInsnNode node, int index,
            MethodInsnNode call) {
        List<TryCatchBlockNode> handlers = handlersMap.get(index);
        InsnList instructions = mn.instructions;
        LabelNode end = new LabelNode();
        instructions.insert(node, end);
        frameMap.add(index, null);
        handlersMap.add(index, null);
        instructions.insertBefore(node, call);
        frameMap.add(index, null);
        handlersMap.add(index, null);

        LabelNode start = new LabelNode();
        instructions.insert(node, start);
        frameMap.add(index, null);
        handlersMap.add(index, null);
        updateCatchHandler(mn, handlers, start, end, handlersMap);
    }

    public static void insertMethodCallAfter(MethodNode mn, List<Frame> frameMap,
            List<List<TryCatchBlockNode>> handlersMap, AbstractInsnNode node, int index,
            MethodInsnNode call) {
        List<TryCatchBlockNode> handlers = handlersMap.get(index + 1);
        InsnList instructions = mn.instructions;

        LabelNode end = new LabelNode();
        instructions.insert(node, end);
        frameMap.add(index + 1, null);
        handlersMap.add(index + 1, null);

        instructions.insert(node, call);
        frameMap.add(index + 1, null);
        handlersMap.add(index + 1, null);

        LabelNode start = new LabelNode();
        instructions.insert(node, start);
        frameMap.add(index + 1, null);
        handlersMap.add(index + 1, null);

        updateCatchHandler(mn, handlers, start, end, handlersMap);
    }

    @SuppressWarnings("unchecked")
    public static void updateCatchHandler(MethodNode mn, List<TryCatchBlockNode> handlers,
            LabelNode start, LabelNode end, List<List<TryCatchBlockNode>> handlersMap) {
        if (handlers == null || handlers.size() == 0) {
            return;
        }

        InsnList instructions = mn.instructions;
        List<TryCatchBlockNode> newNodes = new ArrayList<>(handlers.size());
        for (TryCatchBlockNode handler : handlers) {
            if (!(instructions.indexOf(handler.start) <= instructions.indexOf(start)
                    && instructions.indexOf(end) <= instructions.indexOf(handler.end))) {
                TryCatchBlockNode newNode =
                        new TryCatchBlockNode(start, end, handler.handler, handler.type);
                newNodes.add(newNode);
                for (int i = instructions.indexOf(start); i <= instructions.indexOf(end); i++) {
                    if (handlersMap.get(i) == null) {
                        handlersMap.set(i, new ArrayList<>());
                    }
                    handlersMap.get(i).add(newNode);
                }
            } else {
                for (int i = instructions.indexOf(start); i <= instructions.indexOf(end); i++) {
                    if (handlersMap.get(i) == null) {
                        handlersMap.set(i, new ArrayList<>());
                    }
                    handlersMap.get(i).add(handler);
                }
            }
        }
        mn.tryCatchBlocks.addAll(0, newNodes);
    }
}