package accrue.analysis.visit;

import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import polyglot.ast.ClassDecl;
import polyglot.ast.CodeDecl;
import polyglot.ast.Node;
import polyglot.ast.Stmt;
import polyglot.ast.Term;
import polyglot.visit.NodeVisitor;

public class RegisterExpressions extends NodeVisitor {
    protected Node currentStmt;
    protected Map<Node, Node> enclosingStmt;
    protected Map<Node, Set<Node>> enclosedExprs;

    public RegisterExpressions(Map<Node, Node> enclosingStmt,
            Map<Node, Set<Node>> enclosedExprs) {
        this.enclosingStmt = enclosingStmt;
        this.enclosedExprs = enclosedExprs;
    }

    @Override
    public NodeVisitor enter(Node parent, Node n) {
        if (n instanceof Stmt || n instanceof CodeDecl
                || n instanceof ClassDecl) {
            RegisterExpressions re = (RegisterExpressions) copy();
            re.currentStmt = n;
            return re;
        }
        return this;
    }

    @Override
    public Node leave(Node old, Node n, NodeVisitor v) {
        if (n instanceof Term) {
            RegisterExpressions re = (RegisterExpressions) v;
            enclosingStmt.put(n, re.currentStmt);
            Set<Node> encs = enclosedExprs.get(re.currentStmt);
            if (encs == null) {
                encs = new HashSet<Node>();
                enclosedExprs.put(re.currentStmt, encs);
            }
            encs.add(n);
        }
        return n;
    }

}
