/*

 * Copyright (c) 2017, Oracle and/or its affiliates. All rights reserved.

 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.

 *

 * This code is free software; you can redistribute it and/or modify it

 * under the terms of the GNU General Public License version 2 only, as

 * published by the Free Software Foundation.

 *

 * This code is distributed in the hope that it will be useful, but WITHOUT

 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or

 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License

 * version 2 for more details (a copy is included in the LICENSE file that

 * accompanied this code).

 *

 * You should have received a copy of the GNU General Public License version

 * 2 along with this work; if not, write to the Free Software Foundation,

 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

 *

 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA

 * or visit www.oracle.com if you need additional information or have any

 * questions.

 */

package jdk.experimental.bytecode;

import java.util.ArrayList;

import java.util.HashMap;

import java.util.Iterator;

import java.util.List;

import java.util.Map;

import java.util.Map.Entry;

import java.util.Vector;

import java.util.function.Consumer;

import java.util.function.Supplier;

import java.util.function.ToIntFunction;

public class TypedCodeBuilder<S, T, E, C extends TypedCodeBuilder<S, T, E, C>> extends MacroCodeBuilder<S, T, E, C> {

    State lastStackMapState;

    int lastStackMapPc = -1;

    Map<CharSequence, LocalVarInfo> lvarOffsets = new HashMap<>();

    protected State state;

    int depth = 0;

    int currLocalOffset = 0;

    class StatefulPendingJump extends PendingJump {

        State state;

        StatefulPendingJump(CharSequence label, int pc, State state) {

            super(label, pc);

            this.state = state;

        }

        @Override

        boolean resolve(CharSequence label, int pc) {

            boolean b = super.resolve(label, pc);

            if (b) {

                TypedCodeBuilder.this.state = TypedCodeBuilder.this.state.merge(state);

            }

            return b;

        }

    }

    class LocalVarInfo {

        CharSequence name;

        int offset;

        int depth;

        TypeTag type;

        LocalVarInfo(CharSequence name, int offset, int depth, TypeTag type) {

            this.name = name;

            this.offset = offset;

            this.depth = depth;

            this.type = type;

        }

    }

    public TypedCodeBuilder(MethodBuilder<S, T, E> methodBuilder) {

        super(methodBuilder);

        T t = methodBuilder.desc;

        state = new State();

        if ((methodBuilder.flags & Flag.ACC_STATIC.flag) == 0) {

            T clazz = typeHelper.type(methodBuilder.thisClass);

            state.load(clazz, currLocalOffset++); //TODO: uninit??

        }

        Iterator<T> paramsIt = typeHelper.parameterTypes(t);

        while (paramsIt.hasNext()) {

            T p = paramsIt.next();

            state.load(p, currLocalOffset);

            currLocalOffset += typeHelper.tag(p).width;

        }

        lastStackMapState = state.dup();

        stacksize = state.stack.size();

        localsize = state.locals.size();

    }

    @Override

    protected C emitOp(Opcode opcode, Object optPoolValue) {

        updateState(opcode, optPoolValue);

        return super.emitOp(opcode, optPoolValue);

    }

    @Override

    protected SwitchBuilder makeSwitchBuilder() {

        return new TypedSwitchBuilder();

    }

    class TypedSwitchBuilder extends SwitchBuilder {

        @Override

        public SwitchBuilder withCase(int value, Consumer<? super C> case_, boolean fallthrough) {

            super.withCase(value, c -> {

                withLocalScope(() -> {

                    State prevState = state;

                    state = prevState.dup();

                    emitStackMap(c.offset());

                    case_.accept(c);

                    state = prevState;

                });

            }, fallthrough);

            return this;

        }

        @Override

        public SwitchBuilder withDefault(Consumer<? super C> defaultCase) {

            super.withDefault(c -> {

                withLocalScope(() -> {

                    State prevState = state;

                    state = prevState.dup();

                    emitStackMap(c.offset());

                    defaultCase.accept(c);

                    state = prevState;

                });

            });

            return this;

        }

    }

    @Override

    public StatefulTypedBuilder typed(TypeTag tag) {

        return super.typed(tag, StatefulTypedBuilder::new);

    }

    public class StatefulTypedBuilder extends LabelledTypedBuilder {

        TypeTag tag;

        StatefulTypedBuilder(TypeTag tag) {

            this.tag = tag;

        }

        @Override

        public C astore_0() {

            return storeAndUpdate(super::astore_0);

        }

        @Override

        public C astore_1() {

            return storeAndUpdate(super::astore_1);

        }

        @Override

        public C astore_2() {

            return storeAndUpdate(super::astore_2);

        }

        @Override

        public C astore_3() {

            return storeAndUpdate(super::astore_3);

        }

        @Override

        public C astore(int n) {

            return storeAndUpdate(() -> super.astore(n));

        }

        @Override

        public C aastore() {

            return storeAndUpdate(super::aastore);

        }

        @Override

        public C areturn() {

            state.pop(tag);

            state.push(typeHelper.nullType());

            return super.areturn();

        }

        @Override

        public C anewarray(S s) {

            super.anewarray(s);

            state.pop();

            state.push(typeHelper.arrayOf(typeHelper.type(s)));

            return thisBuilder();

        }

        @Override

        public C aconst_null() {

            super.aconst_null();

            state.pop();

            state.push(tag);

            return thisBuilder();

        }

        public C if_acmpeq(CharSequence label) {

            return jumpAndUpdate(() -> super.if_acmpeq(label));

        }

        public C if_acmpne(CharSequence label) {

            return jumpAndUpdate(() -> super.if_acmpne(label));

        }

        private C storeAndUpdate(Supplier<C> op) {

            state.pop(tag);

            state.push(typeHelper.nullType());

            return op.get();

        }

        private C jumpAndUpdate(Supplier<C> op) {

            state.pop(tag);

            state.pop(tag);

            state.push(typeHelper.nullType());

            state.push(typeHelper.nullType());

            return op.get();

        }

    }

    public class State {

        public final ArrayList<T> stack;

        public final Vector<T> locals;

        boolean alive;

        State(ArrayList<T> stack, Vector<T> locals) {

            this.stack = stack;

            this.locals = locals;

        }

        State() {

            this(new ArrayList<>(), new Vector<>());

        }

        void push(TypeTag tag) {

            switch (tag) {

                case A:

                case V:

                    throw new IllegalStateException("Bad type tag");

                default:

                    push(typeHelper.fromTag(tag));

            }

        }

        void push(T t) {

            stack.add(t);

            if (width(t) == 2) {

                stack.add(null);

            }

            if (stack.size() > stacksize) {

                stacksize = stack.size();

            }

        }

        T peek() {

            return stack.get(stack.size() - 1);

        }

        T tosType() {

            T tos = peek();

            if (tos == null) {

                //double slot

                tos = stack.get(stack.size() - 2);

            }

            return tos;

        }

        T popInternal() {

            return stack.remove(stack.size() - 1);

        }

        @SuppressWarnings("unchecked")

        T pop() {

            if (stack.size() == 0 || peek() == null) throw new IllegalStateException();

            return popInternal();

        }

        T pop2() {

            T o = stack.get(stack.size() - 2);

            TypeTag t = typeHelper.tag(o);

            if (t.width != 2) throw new IllegalStateException();

            popInternal();

            popInternal();

            return o;

        }

        T pop(TypeTag t) {

            return (t.width() == 2) ?

                pop2() : pop();

        }

        void load(TypeTag tag, int index) {

            if (tag == TypeTag.A) throw new IllegalStateException("Bad type tag");

            load(typeHelper.fromTag(tag), index);

        }

        void load(T t, int index) {

            ensureDefined(index);

            locals.set(index, t);

            if (width(t) == 2) {

                locals.add(null);

            }

            if (locals.size() > localsize) {

                localsize = locals.size();

            }

        }

        void ensureDefined(int index) {

            if (index >= locals.size()) {

                locals.setSize(index + 1);

            }

        }

        State dup() {

            State newState = new State(new ArrayList<>(stack), new Vector<>(locals));

            return newState;

        }

        State merge(State that) {

            if (!alive) { return that; }

            if (that.stack.size() != stack.size()) {

                throw new IllegalStateException("Bad stack size at merge point");

            }

            for (int i = 0; i < stack.size(); i++) {

                T t1 = stack.get(i);

                T t2 = that.stack.get(i);

                stack.set(i, merge(t1, t2, "Bad stack type at merge point"));

            }

            int nlocals = locals.size() > that.locals.size() ? that.locals.size() : locals.size();

            for (int i = 0; i < nlocals; i++) {

                T t1 = locals.get(i);

                T t2 = that.locals.get(i);

                locals.set(i, merge(t1, t2, "Bad local type at merge point"));

            }

            if (locals.size() > nlocals) {

                for (int i = nlocals; i < locals.size(); i++) {

                    locals.remove(i);

                }

            }

            return this;

        }

        T merge(T t1, T t2, String msg) {

            if (t1 == null && t2 == null) {

                return t1;

            }

            T res;

            TypeTag tag1 = typeHelper.tag(t1);

            TypeTag tag2 = typeHelper.tag(t2);

            if (tag1 != TypeTag.A && tag2 != TypeTag.A &&

                    tag1 != TypeTag.Q && tag2 != TypeTag.Q) {

                res = typeHelper.fromTag(TypeTag.commonSupertype(tag1, tag2));

            } else if (t1 == typeHelper.nullType()) {

                res = t2;

            } else if (t2 == typeHelper.nullType()) {

                res = t1;

            } else {

                res = typeHelper.commonSupertype(t1, t2);

            }

            if (res == null) {

                throw new IllegalStateException(msg);

            }

            return res;

        }

        @Override

        public String toString() {

            return String.format("[locals = %s, stack = %s]", locals, stack);

        }

    }

    int width(T o) {

        return o == typeHelper.nullType() ?

                TypeTag.A.width() :

                typeHelper.tag(o).width;

    }

    @SuppressWarnings("unchecked")

    public void updateState(Opcode op, Object optValue) {

        switch (op) {

            case VALOAD:

            case AALOAD:

                state.pop();

                state.push(typeHelper.elemtype(state.pop()));

                break;

            case GOTO_:

                state.alive = false;

                break;

            case NOP:

            case INEG:

            case LNEG:

            case FNEG:

            case DNEG:

                break;

            case ACONST_NULL:

                state.push(typeHelper.nullType());

                break;

            case ICONST_M1:

            case ICONST_0:

            case ICONST_1:

            case ICONST_2:

            case ICONST_3:

            case ICONST_4:

            case ICONST_5:

                state.push(TypeTag.I);

                break;

            case LCONST_0:

            case LCONST_1:

                state.push(TypeTag.J);

                break;

            case FCONST_0:

            case FCONST_1:

            case FCONST_2:

                state.push(TypeTag.F);

                break;

            case DCONST_0:

            case DCONST_1:

                state.push(TypeTag.D);

                break;

            case ILOAD_0:

            case FLOAD_0:

            case ALOAD_0:

            case LLOAD_0:

            case DLOAD_0:

                state.push(state.locals.get(0));

                break;

            case ILOAD_1:

            case FLOAD_1:

            case ALOAD_1:

            case LLOAD_1:

            case DLOAD_1:

                state.push(state.locals.get(1));

                break;

            case ILOAD_2:

            case FLOAD_2:

            case ALOAD_2:

            case LLOAD_2:

            case DLOAD_2:

                state.push(state.locals.get(2));

                break;

            case ILOAD_3:

            case FLOAD_3:

            case ALOAD_3:

            case LLOAD_3:

            case DLOAD_3:

                state.push(state.locals.get(3));

                break;

            case ILOAD:

            case FLOAD:

            case ALOAD:

            case LLOAD:

            case DLOAD:

            case VLOAD:

                state.push(state.locals.get((Integer) optValue));

                break;

            case IALOAD:

            case BALOAD:

            case CALOAD:

            case SALOAD:

                state.pop();

                state.pop();

                state.push(TypeTag.I);

                break;

            case LALOAD:

                state.pop();

                state.pop();

                state.push(TypeTag.J);

                break;

            case FALOAD:

                state.pop();

                state.pop();

                state.push(TypeTag.F);

                break;

            case DALOAD:

                state.pop();

                state.pop();

                state.push(TypeTag.D);

                break;

            case ISTORE_0:

            case FSTORE_0:

            case ASTORE_0:

                state.load(state.pop(), 0);

                break;

            case ISTORE_1:

            case FSTORE_1:

            case ASTORE_1:

                state.load(state.pop(), 1);

                break;

            case ISTORE_2:

            case FSTORE_2:

            case ASTORE_2:

                state.load(state.pop(), 2);

                break;

            case ISTORE_3:

            case FSTORE_3:

            case ASTORE_3:

                state.load(state.pop(), 3);

                break;

            case ISTORE:

            case FSTORE:

            case ASTORE:

            case VSTORE:

                state.load(state.pop(), (int) optValue);

                break;

            case LSTORE_0:

            case DSTORE_0:

                state.load(state.pop2(), 0);

                break;

            case LSTORE_1:

            case DSTORE_1:

                state.load(state.pop2(), 1);

                break;

            case LSTORE_2:

            case DSTORE_2:

                state.load(state.pop2(), 2);

                break;

            case LSTORE_3:

            case DSTORE_3:

                state.load(state.pop2(), 3);

                break;

            case LSTORE: