/*

 * Copyright (c) 2014, 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.  Oracle designates this

 * particular file as subject to the "Classpath" exception as provided

 * by Oracle in the LICENSE file that accompanied this code.

 *

 * 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.nashorn.internal.codegen;

import static jdk.nashorn.internal.ir.Node.NO_FINISH;

import static jdk.nashorn.internal.ir.Node.NO_LINE_NUMBER;

import static jdk.nashorn.internal.ir.Node.NO_TOKEN;

import java.util.ArrayDeque;

import java.util.ArrayList;

import java.util.Arrays;

import java.util.Collections;

import java.util.Deque;

import java.util.List;

import java.util.Objects;

import jdk.nashorn.internal.ir.AccessNode;

import jdk.nashorn.internal.ir.BinaryNode;

import jdk.nashorn.internal.ir.Block;

import jdk.nashorn.internal.ir.BlockLexicalContext;

import jdk.nashorn.internal.ir.BreakNode;

import jdk.nashorn.internal.ir.CallNode;

import jdk.nashorn.internal.ir.CaseNode;

import jdk.nashorn.internal.ir.ContinueNode;

import jdk.nashorn.internal.ir.Expression;

import jdk.nashorn.internal.ir.ExpressionStatement;

import jdk.nashorn.internal.ir.FunctionNode;

import jdk.nashorn.internal.ir.FunctionNode.CompilationState;

import jdk.nashorn.internal.ir.GetSplitState;

import jdk.nashorn.internal.ir.IdentNode;

import jdk.nashorn.internal.ir.IfNode;

import jdk.nashorn.internal.ir.JumpStatement;

import jdk.nashorn.internal.ir.LiteralNode;

import jdk.nashorn.internal.ir.Node;

import jdk.nashorn.internal.ir.ReturnNode;

import jdk.nashorn.internal.ir.SetSplitState;

import jdk.nashorn.internal.ir.SplitNode;

import jdk.nashorn.internal.ir.SplitReturn;

import jdk.nashorn.internal.ir.Statement;

import jdk.nashorn.internal.ir.SwitchNode;

import jdk.nashorn.internal.ir.VarNode;

import jdk.nashorn.internal.ir.visitor.NodeVisitor;

import jdk.nashorn.internal.parser.Token;

import jdk.nashorn.internal.parser.TokenType;

/**

 * A node visitor that replaces {@link SplitNode}s with anonymous function invocations and some additional constructs

 * to support control flow across splits. By using this transformation, split functions are translated into ordinary

 * JavaScript functions with nested anonymous functions. The transformations however introduce several AST nodes that

 * have no JavaScript source representations ({@link GetSplitState}, {@link SetSplitState}, and {@link SplitReturn}),

 * and therefore such function is no longer reparseable from its source. For that reason, split functions and their

 * fragments are serialized in-memory and deserialized when they need to be recompiled either for deoptimization or

 * for type specialization.

 * NOTE: all {@code leave*()} methods for statements are returning their input nodes. That way, they will not mutate

 * the original statement list in the block containing the statement, which is fine, as it'll be replaced by the

 * lexical context when the block is left. If we returned something else (e.g. null), we'd cause a mutation in the

 * enclosing block's statement list that is otherwise overwritten later anyway.

 */

final class SplitIntoFunctions extends NodeVisitor<BlockLexicalContext> {

    private static final int FALLTHROUGH_STATE = -1;

    private static final int RETURN_STATE = 0;

    private static final int BREAK_STATE = 1;

    private static final int FIRST_JUMP_STATE = 2;

    private static final String THIS_NAME = CompilerConstants.THIS.symbolName();

    private static final String RETURN_NAME = CompilerConstants.RETURN.symbolName();

    // Used as the name of the formal parameter for passing the current value of :return symbol into a split fragment.

    private static final String RETURN_PARAM_NAME = RETURN_NAME + "-in";

    private final Deque<FunctionState> functionStates = new ArrayDeque<>();

    private final Deque<SplitState> splitStates = new ArrayDeque<>();

    private final Namespace namespace;

    private boolean artificialBlock = false;

    // -1 is program; we need to use negative ones

    private int nextFunctionId = -2;

    public SplitIntoFunctions(final Compiler compiler) {

        super(new BlockLexicalContext() {

            @Override

            protected Block afterSetStatements(Block block) {

                for(Statement stmt: block.getStatements()) {

                    assert !(stmt instanceof SplitNode);

                }

                return block;

            }

        });

        namespace = new Namespace(compiler.getScriptEnvironment().getNamespace());

    }

    @Override

    public boolean enterFunctionNode(final FunctionNode functionNode) {

        functionStates.push(new FunctionState(functionNode));

        return true;

    }

    @Override

    public Node leaveFunctionNode(final FunctionNode functionNode) {

        functionStates.pop();

        return functionNode;

    }

    @Override

    protected Node leaveDefault(final Node node) {

        if (node instanceof Statement) {

            appendStatement((Statement)node);

        }

        return node;

    }

    @Override

    public boolean enterSplitNode(final SplitNode splitNode) {

        getCurrentFunctionState().splitDepth++;

        splitStates.push(new SplitState(splitNode));

        return true;

    }

    @Override

    public Node leaveSplitNode(final SplitNode splitNode) {

        // Replace the split node with an anonymous function expression call.

        final FunctionState fnState = getCurrentFunctionState();

        final String name = splitNode.getName();

        Block body = splitNode.getBody();

        final int firstLineNumber = body.getFirstStatementLineNumber();

        final long token = body.getToken();

        final int finish = body.getFinish();

        final FunctionNode originalFn = fnState.fn;

        assert originalFn == lc.getCurrentFunction();

        final boolean isProgram = originalFn.isProgram();

        // Change SplitNode({...}) into "function () { ... }", or "function (:return-in) () { ... }" (for program)

        final long newFnToken = Token.toDesc(TokenType.FUNCTION, nextFunctionId--, 0);

        final FunctionNode fn = new FunctionNode(

                originalFn.getSource(),

                body.getFirstStatementLineNumber(),

                newFnToken,

                finish,

                newFnToken,

                NO_TOKEN,

                namespace,

                createIdent(name),

                originalFn.getName() + "$" + name,

                isProgram ? Collections.singletonList(createReturnParamIdent()) : Collections.<IdentNode>emptyList(),

                FunctionNode.Kind.NORMAL,

                // We only need IS_SPLIT conservatively, in case it contains any array units so that we force

                // the :callee's existence, to force :scope to never be in a slot lower than 2. This is actually

                // quite a horrible hack to do with CodeGenerator.fixScopeSlot not trampling other parameters

                // and should go away once we no longer have array unit handling in codegen. Note however that

                // we still use IS_SPLIT as the criteria in CompilationPhase.SERIALIZE_SPLIT_PHASE.

                FunctionNode.IS_ANONYMOUS | FunctionNode.USES_ANCESTOR_SCOPE | FunctionNode.IS_SPLIT,

                body,

                CompilationState.INITIALIZED,

                null

        )

        .setCompileUnit(lc, splitNode.getCompileUnit())

        .copyCompilationState(lc, originalFn);

        // Call the function:

        //     either "(function () { ... }).call(this)"

        //     or     "(function (:return-in) { ... }).call(this, :return)"

        // NOTE: Function.call() has optimized linking that basically does a pass-through to the function being invoked.

        // NOTE: CompilationPhase.PROGRAM_POINT_PHASE happens after this, so these calls are subject to optimistic

        // assumptions on their return value (when they return a value), as they should be.

        final IdentNode thisIdent = createIdent(THIS_NAME);

        final CallNode callNode = new CallNode(firstLineNumber, token, finish, new AccessNode(NO_TOKEN, NO_FINISH, fn, "call"),

                isProgram ? Arrays.<Expression>asList(thisIdent, createReturnIdent())

                          : Collections.<Expression>singletonList(thisIdent),

                false);

        final SplitState splitState = splitStates.pop();

        fnState.splitDepth--;

        final Expression callWithReturn;

        final boolean hasReturn = splitState.hasReturn;

        if (hasReturn && fnState.splitDepth > 0) {

            final SplitState parentSplit = splitStates.peek();

            if (parentSplit != null) {

                // Propagate hasReturn to parent split

                parentSplit.hasReturn = true;

            }

        }

        if (hasReturn || isProgram) {

            // capture return value: ":return = (function () { ... })();"

            callWithReturn = new BinaryNode(Token.recast(token, TokenType.ASSIGN), createReturnIdent(), callNode);

        } else {

            // no return value, just call : "(function () { ... })();"

            callWithReturn = callNode;

        }

        appendStatement(new ExpressionStatement(firstLineNumber, token, finish, callWithReturn));

        Statement splitStateHandler;

        final List<JumpStatement> jumpStatements = splitState.jumpStatements;

        final int jumpCount = jumpStatements.size();

        // There are jumps (breaks or continues) that need to be propagated outside the split node. We need to

        // set up a switch statement for them:

        // switch(:scope.getScopeState()) { ... }

        if (jumpCount > 0) {

            final List<CaseNode> cases = new ArrayList<>(jumpCount + (hasReturn ? 1 : 0));

            if (hasReturn) {

                // If the split node also contained a return, we'll slip it as a case in the switch statement

                addCase(cases, RETURN_STATE, createReturnFromSplit());

            }

            int i = FIRST_JUMP_STATE;

            for (final JumpStatement jump: jumpStatements) {

                addCase(cases, i++, enblockAndVisit(jump));

            }

            splitStateHandler = new SwitchNode(NO_LINE_NUMBER, token, finish, GetSplitState.INSTANCE, cases, null);

        } else {

            splitStateHandler = null;

        }

        // As the switch statement itself is breakable, an unlabelled break can't be in the switch statement,

        // so we need to test for it separately.

        if (splitState.hasBreak) {

            // if(:scope.getScopeState() == Scope.BREAK) { break; }

            splitStateHandler = makeIfStateEquals(firstLineNumber, token, finish, BREAK_STATE,

                    enblockAndVisit(new BreakNode(NO_LINE_NUMBER, token, finish, null)), splitStateHandler);

        }

        // Finally, if the split node had a return statement, but there were no external jumps, we didn't have

        // the switch statement to handle the return, so we need a separate if for it.

        if (hasReturn && jumpCount == 0) {

            // if (:scope.getScopeState() == Scope.RETURN) { return :return; }

            splitStateHandler = makeIfStateEquals(NO_LINE_NUMBER, token, finish, RETURN_STATE,

                    createReturnFromSplit(), splitStateHandler);

        }

        if (splitStateHandler != null) {

            appendStatement(splitStateHandler);

        }

        return splitNode;

    }

    private static void addCase(final List<CaseNode> cases, final int i, final Block body) {

        cases.add(new CaseNode(NO_TOKEN, NO_FINISH, intLiteral(i), body));

    }

    private static LiteralNode<Number> intLiteral(final int i) {

        return LiteralNode.newInstance(NO_TOKEN, NO_FINISH, i);

    }

    private static Block createReturnFromSplit() {

        return new Block(NO_TOKEN, NO_FINISH, createReturnReturn());

    }

    private static ReturnNode createReturnReturn() {

        return new ReturnNode(NO_LINE_NUMBER, NO_TOKEN, NO_FINISH, createReturnIdent());

    }

    private static IdentNode createReturnIdent() {

        return createIdent(RETURN_NAME);

    }

    private static IdentNode createReturnParamIdent() {

        return createIdent(RETURN_PARAM_NAME);

    }

    private static IdentNode createIdent(final String name) {

        return new IdentNode(NO_TOKEN, NO_FINISH, name);

    }

    private Block enblockAndVisit(final JumpStatement jump) {

        artificialBlock = true;

        final Block block = (Block)new Block(NO_TOKEN, NO_FINISH, jump).accept(this);

        artificialBlock = false;

        return block;

    }

    private static IfNode makeIfStateEquals(final int lineNumber, final long token, final int finish,

            final int value, final Block pass, final Statement fail) {

        return new IfNode(lineNumber, token, finish,

                new BinaryNode(Token.recast(token, TokenType.EQ_STRICT),

                        GetSplitState.INSTANCE, intLiteral(value)),

                pass,

                fail == null ? null : new Block(NO_TOKEN, NO_FINISH, fail));

    }

    @Override

    public boolean enterVarNode(VarNode varNode) {

        if (!inSplitNode()) {

            return super.enterVarNode(varNode);

        }

        assert !varNode.isBlockScoped(); //TODO: we must handle these too, but we currently don't

        final Expression init = varNode.getInit();

        if (varNode.isAnonymousFunctionDeclaration()) {

            // We ain't moving anonymous function declarations.

            return super.enterVarNode(varNode);

        }

        // Move a declaration-only var statement to the top of the outermost function.

        getCurrentFunctionState().varStatements.add(varNode.setInit(null));

        // If it had an initializer, replace it with an assignment expression statement. Note that "var" is a

        // statement, so it doesn't contribute to :return of the programs, therefore we are _not_ adding a

        // ":return = ..." assignment around the original assignment.

        if (init != null) {

            final long token = Token.recast(varNode.getToken(), TokenType.ASSIGN);

            new ExpressionStatement(varNode.getLineNumber(), token, varNode.getFinish(),

                    new BinaryNode(token, varNode.getName(), varNode.getInit())).accept(this);

        }

        return false;

    }

    @Override

    public Node leaveBlock(final Block block) {

        if (!artificialBlock) {

            if (lc.isFunctionBody()) {

                // Prepend declaration-only var statements to the top of the statement list.

                lc.prependStatements(getCurrentFunctionState().varStatements);

            } else if (lc.isSplitBody()) {

                appendSplitReturn(FALLTHROUGH_STATE, NO_LINE_NUMBER);

                if (getCurrentFunctionState().fn.isProgram()) {

                    // If we're splitting the program, make sure every shard ends with "return :return" and

                    // begins with ":return = :return-in;".

                    lc.prependStatement(new ExpressionStatement(NO_LINE_NUMBER, NO_TOKEN, NO_FINISH,

                            new BinaryNode(Token.toDesc(TokenType.ASSIGN, 0, 0), createReturnIdent(), createReturnParamIdent())));

                }

            }

        }

        return block;

    }

    @Override

    public Node leaveBreakNode(final BreakNode breakNode) {

        return leaveJumpNode(breakNode);

    }

    @Override

    public Node leaveContinueNode(final ContinueNode continueNode) {

        return leaveJumpNode(continueNode);

    }

    private JumpStatement leaveJumpNode(final JumpStatement jump) {

        if (inSplitNode()) {

            final SplitState splitState = getCurrentSplitState();

            final SplitNode splitNode = splitState.splitNode;

            if (lc.isExternalTarget(splitNode, jump.getTarget(lc))) {

                appendSplitReturn(splitState.getSplitStateIndex(jump), jump.getLineNumber());

                return jump;

            }

        }

        appendStatement(jump);

        return jump;

    }

    private void appendSplitReturn(final int splitState, final int lineNumber) {

        appendStatement(new SetSplitState(splitState, lineNumber));

        if (getCurrentFunctionState().fn.isProgram()) {

            // If we're splitting the program, make sure every fragment passes back :return

            appendStatement(createReturnReturn());

        } else {

            appendStatement(SplitReturn.INSTANCE);

        }

    }

    @Override

    public Node leaveReturnNode(final ReturnNode returnNode) {

        if(inSplitNode()) {

            appendStatement(new SetSplitState(RETURN_STATE, returnNode.getLineNumber()));

            getCurrentSplitState().hasReturn = true;

        }

        appendStatement(returnNode);

        return returnNode;

    }

    private void appendStatement(final Statement statement) {

        lc.appendStatement(statement);

    }

    private boolean inSplitNode() {

        return getCurrentFunctionState().splitDepth > 0;

    }

    private FunctionState getCurrentFunctionState() {

        return functionStates.peek();

    }

    private SplitState getCurrentSplitState() {

        return splitStates.peek();

    }

    private static class FunctionState {

        final FunctionNode fn;

        final List<Statement> varStatements = new ArrayList<>();

        int splitDepth;

        FunctionState(final FunctionNode fn) {

            this.fn = fn;

        }

    }

    private static class SplitState {

        final SplitNode splitNode;

        boolean hasReturn;

        boolean hasBreak;

        final List<JumpStatement> jumpStatements = new ArrayList<>();

        int getSplitStateIndex(final JumpStatement jump) {

            if (jump instanceof BreakNode && jump.getLabelName() == null) {

                // Unlabelled break is a special case

                hasBreak = true;

                return BREAK_STATE;

            }

            int i = 0;

            for(final JumpStatement exJump: jumpStatements) {

                if (jump.getClass() == exJump.getClass() && Objects.equals(jump.getLabelName(), exJump.getLabelName())) {

                    return i + FIRST_JUMP_STATE;

                }

                ++i;

            }

            jumpStatements.add(jump);

            return i + FIRST_JUMP_STATE;

        }

        SplitState(final SplitNode splitNode) {

            this.splitNode = splitNode;

        }

    }

}