package accrue.analysis.readwritesets;

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

import polyglot.ast.ArrayAccess;
import polyglot.ast.ArrayAccessAssign;
import polyglot.ast.Call;
import polyglot.ast.CodeDecl;
import polyglot.ast.CodeNode;
import polyglot.ast.ConstructorCall;
import polyglot.ast.Field;
import polyglot.ast.FieldAssign;
import polyglot.ast.New;
import polyglot.ast.Node;
import polyglot.types.SemanticException;
import polyglot.util.InternalCompilerError;
import polyglot.visit.NodeVisitor;
import accrue.analysis.ext.AccrueExt_c;
import accrue.analysis.ext.ExtCall;
import accrue.analysis.ext.ExtCodeDecl;
import accrue.analysis.ext.ExtConstructorCall;
import accrue.analysis.ext.ExtNew;
import accrue.analysis.interprocanalysis.AbstractLocation;

public class ReadWriteSetVisitor extends NodeVisitor {

    private final ReadWriteSetsAnalysisUtil autil;
    /**
     * Set of read locations
     */
    private Set<AbstractLocation> readLocs = new HashSet<AbstractLocation>();
    /**
     * Set of written locations
     */
    private Set<AbstractLocation> writtenLocs = new HashSet<AbstractLocation>();

    /**
     * The final result for this code decl.
     */
    private ReadWriteSet finalResult = null;

    public ReadWriteSetVisitor(ReadWriteSetsAnalysisUtil autil) {
        this.autil = autil;
    }

    public ReadWriteSet performVisit(CodeNode cd) {
        this.finalResult = null;
        cd.visit(this);
//        System.err.println("Visted " + cd.codeInstance() + "\n   "
//                + finalResult);
        return this.finalResult;
    }

    @Override
    public Node leave(Node old, Node n, NodeVisitor v) {
        if (n instanceof ArrayAccess) {
            ArrayAccess aa = (ArrayAccess) n;
            Set<AbstractLocation> locs = autil.abstractLocations(aa);
            addReads(locs);
        }
        else if (n instanceof ArrayAccessAssign) {
            ArrayAccessAssign aaa = (ArrayAccessAssign) n;
            ArrayAccess aa = (ArrayAccess) aaa.left();
            Set<AbstractLocation> locs = autil.abstractLocations(aa);
            addReads(locs);
            addWrites(locs);
        }
        else if (n instanceof Field) {
            Field f = (Field) n;
            Set<AbstractLocation> locs = autil.abstractLocations(f);
            addReads(locs);
        }
        else if (n instanceof FieldAssign) {
            FieldAssign fa = (FieldAssign) n;
            Field f = (Field) fa.left();
            Set<AbstractLocation> locs = autil.abstractLocations(f);
            addReads(locs);
            addWrites(locs);
        }
        else if (n instanceof Call) {
            Call c = (Call) n;
            ReadWriteSet s = autil.call(c.methodInstance(), c);

            ExtCall ext = (ExtCall) AccrueExt_c.ext(n);
            ext.recordAnalysisResult(ReadWriteSetsAnalysisFactory.ANALYSIS_NAME,
                                     autil.currentContext(),
                                     s);

            addReads(s.readSet());
            addWrites(s.writeSet());
        }
        else if (n instanceof ConstructorCall) {
            /* The read and write sets of this constructor call */
            Set<AbstractLocation> ccr = new HashSet<AbstractLocation>();
            Set<AbstractLocation> ccw = new HashSet<AbstractLocation>();

            ConstructorCall cc = (ConstructorCall) n;
            ReadWriteSet s =
                    autil.constructorCall(cc.constructorInstance(), cc);
            addReads(s.readSet());
            addWrites(s.writeSet());

            ccr.addAll(s.readSet());
            ccw.addAll(s.writeSet());

            if (cc.kind().equals(ConstructorCall.SUPER)) {
                // Analyze initializers and add to the read write set
                try {
                    for (ReadWriteSet r : autil.analyzeInitializers(ReadWriteSet.EMPTY)
                                               .allItems()) {
                        addReads(r.readSet());
                        addWrites(r.writeSet());

                        ccr.addAll(r.readSet());
                        ccw.addAll(r.writeSet());
                    }
                }
                catch (SemanticException e) {
                    throw new InternalCompilerError(e);
                }
            }
            ExtConstructorCall ext = (ExtConstructorCall) AccrueExt_c.ext(cc);
            ext.recordAnalysisResult(ReadWriteSetsAnalysisFactory.ANALYSIS_NAME,
                                     autil.currentContext(),
                                     new ReadWriteSet(ccr, ccw));
        }
        else if (n instanceof New) {
            New nn = (New) n;
            ReadWriteSet s = autil.call(nn.constructorInstance(), nn);
            addReads(s.readSet());
            addWrites(s.writeSet());

            ExtNew ext = (ExtNew) AccrueExt_c.ext(nn);
            ext.recordAnalysisResult(ReadWriteSetsAnalysisFactory.ANALYSIS_NAME,
                                     autil.currentContext(),
                                     s);
        }
        else if (n instanceof CodeDecl) {
            CodeDecl cd = (CodeDecl) n;
            ExtCodeDecl ext = (ExtCodeDecl) AccrueExt_c.ext(cd);
            this.finalResult = new ReadWriteSet(readLocs, writtenLocs);
            // now that we have set the final result, make the readLocs and writtenLocs
            // field null, so we don't accidentally add more locations in.
            readLocs = null;
            writtenLocs = null;

            ext.recordAnalysisResult(ReadWriteSetsAnalysisFactory.ANALYSIS_NAME,
                                     autil.currentContext(),
                                     finalResult);
        }
        return n;
    }

    private void addReads(Set<AbstractLocation> locs) {
        readLocs.addAll(locs);
    }

    private void addWrites(Set<AbstractLocation> locs) {
        writtenLocs.addAll(locs);
    }
}
