package accrue.analysis.visit;

import java.util.LinkedList;
import java.util.Stack;

import polyglot.ast.Branch;
import polyglot.ast.CodeDecl;
import polyglot.ast.Labeled;
import polyglot.ast.Loop;
import polyglot.ast.Node;
import polyglot.ast.NodeFactory;
import polyglot.ast.Stmt;
import polyglot.ast.Switch;
import polyglot.types.TypeSystem;
import polyglot.visit.NodeVisitor;

/**
 * Makes targets of break and continue explicit.
 *
 */
public class MakeTargetsExplicit extends NodeVisitor {
    NodeFactory nf;
    TypeSystem ts;

    Stack<String> nearestContinueTarget;
    Stack<String> nearestBreakTarget;
    LinkedList<Integer> varCount;

    public MakeTargetsExplicit(NodeFactory nf, TypeSystem ts) {
        super();
        this.nf = nf;
        this.ts = ts;
        nearestContinueTarget = new Stack<String>();
        nearestBreakTarget = new Stack<String>();
        varCount = new LinkedList<Integer>();
    }

    private String freshName(String desc) {
        int count = varCount.removeLast();
        varCount.addLast(count + 1);
        if (count == 0) {
            return "label$" + desc;
        }
        return "label$" + desc + "$" + count;
    }

    @Override
    public NodeVisitor enter(Node parent, Node n) {
        if (n instanceof Loop || n instanceof Switch) {
            if (parent instanceof Labeled) {
                if (n instanceof Loop) {
                    nearestContinueTarget.push(((Labeled) parent).labelNode()
                                                                 .id());
                }
                nearestBreakTarget.push(((Labeled) parent).labelNode().id());
            }
            else {
                // Need a fresh label!
                String newName = freshName("target");
                if (n instanceof Loop) {
                    nearestContinueTarget.push(newName);
                }
                nearestBreakTarget.push(newName);
            }
        }
        if (n instanceof CodeDecl) {
            varCount.addLast(0);
        }
        return this;
    }

    @Override
    public Node leave(Node parent, Node old, Node n, NodeVisitor v) {
        if (n instanceof CodeDecl) {
            varCount.removeLast();
        }
        if (n instanceof Branch) {
            Branch b = (Branch) n;
            // Make the branch explicit if it isn't already
            if (b.labelNode() == null) {
                // Insert an explicit target for the nearest appropriate structure
                if (b.kind().equals(Branch.CONTINUE)) {
                    return b.labelNode(nf.Id(b.position(),
                                             nearestContinueTarget.peek()));
                }
                else {
                    return b.labelNode(nf.Id(b.position(),
                                             nearestBreakTarget.peek()));
                }
            }
            return b;
        }
        else if (n instanceof Loop || n instanceof Switch) {
            if (!(parent instanceof Labeled)) {
                // Need to insert a new label and pop it off the stack
                String label;
                if (n instanceof Loop) {
                    label = nearestContinueTarget.pop();
                    nearestBreakTarget.pop();
                }
                else {
                    label = nearestBreakTarget.pop();
                }
                return nf.Labeled(n.position(),
                                  nf.Id(n.position(), label),
                                  (Stmt) n);
            }
            else {
                if (n instanceof Loop) {
                    nearestContinueTarget.pop();
                }
                nearestBreakTarget.pop();
                return n;
            }
        }
        return super.leave(parent, old, n, v);
    }

}
