package accrue.analysis.defuse;

import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import polyglot.ast.ArrayAccess;
import polyglot.ast.Call;
import polyglot.ast.Cast;
import polyglot.ast.CodeNode;
import polyglot.ast.ConstructorCall;
import polyglot.ast.ConstructorDecl;
import polyglot.ast.Expr;
import polyglot.ast.Field;
import polyglot.ast.Formal;
import polyglot.ast.New;
import polyglot.ast.ProcedureCall;
import polyglot.ast.ProcedureDecl;
import polyglot.ast.StringLit;
import polyglot.types.ClassType;
import polyglot.types.FieldInstance;
import polyglot.types.MethodInstance;
import polyglot.types.ProcedureInstance;
import polyglot.types.SemanticException;
import polyglot.types.Type;
import polyglot.types.TypeSystem;
import polyglot.util.InternalCompilerError;
import polyglot.visit.FlowGraph.Peer;
import accrue.AccrueExtensionInfo;
import accrue.analysis.ext.AccrueExt;
import accrue.analysis.ext.AccrueExt_c;
import accrue.analysis.ext.ExtCast;
import accrue.analysis.ext.ExtConstructorDecl;
import accrue.analysis.ext.ExtExpr;
import accrue.analysis.ext.ExtFormal;
import accrue.analysis.interprocanalysis.AbstractLocation;
import accrue.analysis.interprocanalysis.AnalysisContext;
import accrue.analysis.interprocanalysis.AnalysisUnit;
import accrue.analysis.interprocanalysis.DataFlowNodeIdentifier;
import accrue.analysis.interprocanalysis.EdgeIdentifier;
import accrue.analysis.interprocanalysis.EdgeIdentifier.HeapUseKey;
import accrue.analysis.interprocanalysis.EdgeIdentifier_c;
import accrue.analysis.interprocanalysis.ExitMap;
import accrue.analysis.interprocanalysis.ExitMap.Key;
import accrue.analysis.interprocanalysis.WorkQueue;
import accrue.analysis.interprocvarcontext.AnalysisUtilVarContextDataFlow;
import accrue.analysis.interprocvarcontext.FinalVarContext;
import accrue.analysis.interprocvarcontext.VarContext;
import accrue.analysis.pointer.HContext;
import accrue.analysis.pointer.ReferenceVariable;

//
public class DefUseAnalysisUtil extends
        AnalysisUtilVarContextDataFlow<VarContext<DefUseAbsVal>, DefUseAbsVal> {
    public static enum Mode {
        /** interprocedural dependencies are edges between actual statements */
        INTER_EDGES,
        /** interprocedural dependencies are summarized at call sites */
        INTRA_EDGES
    }

    protected TypeSystem ts;
    protected final Mode mode;
    protected final boolean trackStrings;
    public static final String STRING_CONTENTS = "[str contents]";

    public DefUseAnalysisUtil(WorkQueue<VarContext<DefUseAbsVal>> wq,
            AnalysisUnit currentAnalysisUnit, AccrueExtensionInfo extInfo) {
        this(wq, currentAnalysisUnit, extInfo, Mode.INTER_EDGES, false);
    }

    public DefUseAnalysisUtil(WorkQueue<VarContext<DefUseAbsVal>> wq,
            AnalysisUnit currentAnalysisUnit, AccrueExtensionInfo extInfo,
            Mode mode, boolean trackStrings) {
        super(wq, currentAnalysisUnit, extInfo);
        this.ts = extInfo.typeSystem();
        this.mode = mode;
        this.trackStrings = trackStrings;
    }

    public Mode mode() {
        return mode;
    }

    @Override
    public FinalVarContext<DefUseAbsVal> analyze(CodeNode n,
            VarContext<DefUseAbsVal> input) throws SemanticException {
        DefUseDataFlow d =
                new DefUseDataFlow(this,
                                   this.workQueue(),
                                   this.extInfo.scheduler().currentJob(),
                                   this.extInfo.typeSystem(),
                                   this.extInfo.nodeFactory());
        if (mode == Mode.INTRA_EDGES) {
            d.dataflow(n, createInitContext());
            return createFinalVarContext(createInitContext());
        }
        return d.dataflow(n, input);
    }

    @Override
    public VarContext<DefUseAbsVal> createInitContext() {
        return new VarContext<DefUseAbsVal>(analysisName(), DefUseAbsVal.EMPTY);
    }

    protected Set<ProcedureInstance> getCallTargets(Call n) {
        // this is a virtual call.
        // find out all the method bodies this may call, and join them all
        // together.
        MethodInstance mi = n.methodInstance();

        if (!(n.target() instanceof Expr))
            return Collections.<ProcedureInstance> singleton(mi);

        Set<ProcedureInstance> targets = new HashSet<ProcedureInstance>();
        // analyze all method implementations that may potentially be called
        for (ClassType receiverType : getReceiverTypes(n)) {
            // analyze the implementing method
            MethodInstance omi = null;
            if (receiverType != null) {
                omi =
                        extInfo.typeSystem()
                               .findImplementingMethod(receiverType, mi);
            }
            else {
                // System.err.println(" *** no receiver for virtual method invocation to "
                // + mi + " " + n.position());
                throw new InternalCompilerError("no receiver for virtual method invocation to "
                                                        + mi,
                                                n.position());
            }
            if (omi == null) {
                throw new InternalCompilerError("No implementing method for : "
                        + receiverType + "->" + mi);
            }
            targets.add(omi);
        }
        return targets;
    }

    protected Set<ClassType> getReceiverTypes(Call n) {
        // this is a virtual call.
        // find out all the method bodies this may call, and join them all
        // together.
        if (!(n.target() instanceof Expr)) return Collections.emptySet();
        Expr receiver = (Expr) n.target();
        Set<HContext> pointsTo = pointsTo(receiver);
        MethodInstance mi = n.methodInstance();

        if (pointsTo == null) {
            throw new InternalCompilerError("No points to info for the receiver: "
                    + n);
        }

        if (pointsTo.isEmpty()) {
            System.err.println(" *** Receiver doesn't point to anything! May be a bug in the program? ");
            System.err.println("     " + n + " " + mi.container() + "."
                    + mi.name() + "  at " + n.position());
        }

        Set<ClassType> targetTypes = new HashSet<ClassType>();
        // analyze all method implementations that may potentially be called
        for (HContext o : pointsTo) {
            // analyze the implementing method
            Type receiverType = o.type();
            if (receiverType != null) {
                targetTypes.add(receiverType.toClass());
            }
            else {
                throw new InternalCompilerError("no receiver for virtual method invocation to "
                                                        + mi,
                                                n.position());
            }
        }
        return targetTypes;
    }

    /**
     * Returns analysis results for a specific target.
     */
    protected <T> Set<T> getNonVirtualCallResults(String analysisKey,
            ProcedureCall n, ProcedureInstance pi, HContext o) {
        AnalysisContext context =
                workQueue.factory().createAnalysisContext(this, pi, o, n);
        ProcedureDecl pd = (ProcedureDecl) workQueue.registrar().getCode(pi);
        if (pd == null || pd.body() == null)
            return getNativeCallResults(analysisKey, n, pi, o);
        AccrueExt ext = AccrueExt_c.ext(pd);
        Set<T> result = ext.getAnalysisResult(analysisKey, context);
        if (result == null) return Collections.emptySet();
        return result;
    }

    /**
     * Returns analysis results for a native call target. Native analysis
     * results are recorded on the call site, but use the AnalysisContext of the
     * callee.
     */
    protected <T> Set<T> getNativeCallResults(String analysisKey,
            ProcedureCall n, ProcedureInstance pi, HContext o) {
        AnalysisContext context =
                workQueue.factory().createAnalysisContext(this, pi, o, n);
        AccrueExt ext = AccrueExt_c.ext(n);
        Set<T> result = ext.getAnalysisResult(analysisKey, context);
        if (result == null) return Collections.emptySet();
        return result;
    }

    /**
     * Return analysis results for a procedure call. For virtual calls, returns
     * the union of results for each potential call target.
     */
    public <T> Set<T> getCallResults(String analysisKey, ProcedureCall pcall) {
        // this is a virtual call.
        // find out all the method bodies this may call, and join them all
        // together.
        if (pcall instanceof New) {
            New n = (New) pcall;
            Set<HContext> pointsTo = pointsTo(n);
            if (pointsTo == null || pointsTo.size() != 1)
                throw new InternalCompilerError("Invalid pointsTo set for allocation:"
                        + n);
            return getNonVirtualCallResults(analysisKey,
                                            n,
                                            pcall.procedureInstance(),
                                            pointsTo.iterator().next());
        }
        else if (pcall instanceof Call) {
            Call c = (Call) pcall;
            if (!(c.target() instanceof Expr)) {
                return getNonVirtualCallResults(analysisKey,
                                                pcall,
                                                pcall.procedureInstance(),
                                                null);
            }

            Expr receiver = (Expr) c.target();
            Set<HContext> pointsTo = pointsTo(receiver);
            MethodInstance mi = c.methodInstance();

            if (pointsTo == null) {
                throw new InternalCompilerError("No points to info for the receiver: "
                        + c);
            }

            if (pointsTo.isEmpty()) {
                System.err.println(" *** Receiver doesn't point to anything! May be a bug in the program? ");
                System.err.println("     " + c + " " + mi.container() + "."
                        + mi.name() + "  at " + c.position());
            }

            Set<T> results = new HashSet<T>();
            // analyze all method implementations that may potentially be called
            for (HContext o : pointsTo) {
                // analyze the implementing method
                Type receiverType = o.type();
                // analyze the implementing method
                MethodInstance omi = null;
                if (receiverType != null) {
                    omi =
                            extInfo.typeSystem()
                                   .findImplementingMethod(receiverType.toClass(),
                                                           mi);
                }
                else {
                    throw new InternalCompilerError("no receiver for virtual method invocation to "
                                                            + mi,
                                                    c.position());
                }
                if (omi == null) {
                    // No implementing method. Probably, we have a signature for
                    // an interface. Just use the callers methodInstance.
                    omi = mi;
                }
                Set<T> target_locs =
                        getNonVirtualCallResults(analysisKey, pcall, omi, o);
                if (target_locs != null) {
                    results.addAll(target_locs);
                }
            }
            return results;
        }
        else if (pcall instanceof ConstructorCall) {
            // FIXME: This is mostly guesswork. Untested!
            ConstructorDecl currentCtor =
                    (ConstructorDecl) this.currentAnalysisUnit().codeNode();
            ExtConstructorDecl current_ext =
                    (ExtConstructorDecl) AccrueExt_c.ext(currentCtor);
            Set<HContext> pointsTo =
                    pointsTo(current_ext.getThisNode(),
                             currentContext(),
                             extInfo);
            if (pointsTo.isEmpty()) {
                // TODO: Native constructors
                return Collections.emptySet();
            }
            else {
                AnalysisContext context = currentContext();
                ProcedureInstance pi = pcall.procedureInstance();
                ProcedureDecl pd =
                        (ProcedureDecl) workQueue.registrar().getCode(pi);
                AccrueExt ext;
                if (pd == null || pd.body() == null) {
                    ext = AccrueExt_c.ext(pcall);
                }
                else {
                    ext = AccrueExt_c.ext(pd);
                }
                Set<T> result = ext.getAnalysisResult(analysisKey, context);
                if (result == null) return Collections.emptySet();
                return result;
            }
        }
        else {
            throw new InternalCompilerError("Unexpected procedure call:"
                    + pcall);
        }
    }

    private void scheduleAnalysis(AnalysisContext calleeContext,
            ProcedureInstance pi) {
        // get the implementing method body to add the argument names...
        ProcedureDecl pd = (ProcedureDecl) workQueue.registrar().getCode(pi);
        AnalysisUnit callee = null;
        if (pd != null && pd.body() != null) {
            callee = new AnalysisUnit(pd, calleeContext);

            try {
                this.workQueue.getAnalysisResult(callee,
                                                 createInitContext(),
                                                 this);
            }
            catch (SemanticException e) {
                throw new InternalCompilerError("Unexpected semantic exception",
                                                e);
            }
        }
    }

    @Override
    public FinalVarContext<DefUseAbsVal> callNonVirtual(ProcedureInstance pi,
            AnalysisContext calleeContext, HContext receiver,
            DefUseAbsVal receiverAbsVal, VarContext<DefUseAbsVal> input,
            DataFlowNodeIdentifier callSite, boolean canTerminateNormally)
            throws SemanticException {
        if (mode() == Mode.INTER_EDGES) {
            return super.callNonVirtual(pi,
                                        calleeContext,
                                        receiver,
                                        receiverAbsVal,
                                        input,
                                        callSite,
                                        canTerminateNormally);
        }
        else {
            // Schedule call targets
            scheduleAnalysis(calleeContext, pi);

            // pop arguments
            int numArgs = pi.formalTypes().size();
            if (pi instanceof MethodInstance
                    || callSite.getNode() instanceof New) numArgs++;
            input = input.popExprResults(numArgs);

            if (pi instanceof MethodInstance) {
                MethodInstance mi = (MethodInstance) pi;
                if (!mi.returnType().isVoid()) {
                    input =
                            input.pushExprResult(bottomAbstractValue(),
                                                 "returnX2-" + pi);
                }
            }

            AccrueExt ext = AccrueExt_c.ext(callSite.getNode());
            Set<AbstractLocation> obs =
                    ext.getAnalysisResult(DefUseAnalysisFactory.MODIFIED_LOCATIONS,
                                          currentContext());
            if (obs == null || obs.isEmpty()) {
                return createFinalVarContext(input);
            }

            HeapUseKey key = new HeapUseKey(obs, true);
            Set<EdgeIdentifier> nodeID =
                    Collections.<EdgeIdentifier> singleton(new EdgeIdentifier_c(key,
                                                                                callSite));
            DefUseAbsVal uses = new DefUseAbsVal(nodeID);
            input = addLocations(input, obs, uses);
            return createFinalVarContext(input);
        }
    }

    protected boolean mayBeNull(DefUseAbsVal exprDFI, Expr e,
            Peer<VarContext<DefUseAbsVal>> peer) {
        ExtExpr ext = (ExtExpr) AccrueExt_c.ext(e);
        if (exprDFI instanceof NewAllocAbsVal)
            return false;
        else return ext.mayBeNull(currentContext(), peer.peerKey());
    }

    protected boolean mayThrowClassCastException(DefUseAbsVal exprDFI,
            Cast cast, Peer<VarContext<DefUseAbsVal>> peer) {
        ExtCast ext = (ExtCast) AccrueExt_c.ext(cast);
        return ext.mayThrowClassCastException(currentContext());
    }

    protected Set<AbstractLocation> abstractLocations(DefUseAbsVal targetDFI,
            Field f) {
        return abstractLocations(f);
    }

    protected Set<AbstractLocation> abstractLocationsImplicitTarget(
            DefUseAbsVal targetDFI, FieldInstance fi) {
        return abstractLocationsImplicitTarget(fi);
    }

    protected Set<AbstractLocation> abstractLocations(DefUseAbsVal arrayDFI,
            ArrayAccess aa) {
        return abstractLocations(aa);
    }

    protected Set<AbstractLocation> stringLocations(DefUseAbsVal stringDFI,
            Expr stringExpr) {
        if (stringExpr instanceof StringLit) {
            return Collections.emptySet();
        }
        Set<HContext> pt = pointsTo(stringDFI, stringExpr);
        Set<AbstractLocation> locs = new HashSet<AbstractLocation>();
        for (HContext hc : pt) {
            locs.add(AbstractLocation.create(hc, STRING_CONTENTS));
        }
        return locs;
    }

    protected boolean hasReferenceVariable(AnalysisContext context, Expr e) {
        ExtExpr ee = (ExtExpr) AccrueExt_c.ext(e);
        ReferenceVariable ln =
                ee.getReferenceVariable(context.getConstructorContext());
        return ln != null;
    }

    protected boolean hasReferenceVariable(AnalysisContext context, Formal f) {
        ExtFormal ee = (ExtFormal) AccrueExt_c.ext(f);
        ReferenceVariable ln =
                ee.getReferenceVariable(context.getConstructorContext());
        return ln != null;
    }

    protected Set<HContext> pointsTo(DefUseAbsVal exprDFI, Expr e) {
        Set<HContext> pointsTo = new HashSet<HContext>();
        if (exprDFI instanceof NewAllocAbsVal) {
            pointsTo = ((NewAllocAbsVal) exprDFI).objects();
        }
        else {
            return pointsTo(e);
        }
        return pointsTo;
    }

    @Override
    protected FinalVarContext<DefUseAbsVal> getAnalysisResultsFromWorkQueue(
            AnalysisUnit callee, VarContext<DefUseAbsVal> beforeCall,
            DataFlowNodeIdentifier callSite) throws SemanticException {

        return (FinalVarContext<DefUseAbsVal>) workQueue.getAnalysisResult(callee,
                                                                           beforeCall,
                                                                           this);
    }

    @Override
    public DefUseAbsVal bottomAbstractValue() {
        return DefUseAbsVal.EMPTY;
    }

    @Override
    protected FinalVarContext<DefUseAbsVal> guessAnalysisForMissingReceiver(
            ProcedureInstance pi, VarContext<DefUseAbsVal> input,
            DefUseAbsVal receiverAbsVal, List<DefUseAbsVal> argAbsVals,
            Set<Key> expectedExits, DataFlowNodeIdentifier callSiteNode) {
        DefUseAbsVal ret = this.bottomAbstractValue();
        VarContext<DefUseAbsVal> nt = input.returnResult(ret);
        VarContext<DefUseAbsVal> ex =
                input.clearExprResultStack().returnResult(ret);
        return this.makeFinalVarContext(expectedExits, nt, ex);
    }

    @Override
    protected FinalVarContext<DefUseAbsVal> guessAnalysisForMissingCode(
            ProcedureInstance pi, AnalysisContext calleeContext,
            HContext receiver, VarContext<DefUseAbsVal> input,
            DefUseAbsVal receiverAbsVal, List<DefUseAbsVal> argAbsVals,
            Set<ExitMap.Key> expectedExits, DataFlowNodeIdentifier callSite) {
        System.err.println("sig for procedure = " + pi + "apparently not found");
        return DefUseLibrarySignature.NOOP.process(pi,
                                                   calleeContext,
                                                   receiver,
                                                   input,
                                                   receiverAbsVal,
                                                   argAbsVals,
                                                   this,
                                                   expectedExits,
                                                   callSite);
    }

    @Override
    public FinalVarContext<DefUseAbsVal> analyzeInitializersForDataFlow(
            VarContext<DefUseAbsVal> before) throws SemanticException {
        if (mode.equals(Mode.INTER_EDGES))
            return (FinalVarContext<DefUseAbsVal>) this.analyzeInitializers(before);
        else {
            VarContext<DefUseAbsVal> empty = createInitContext();
            analyzeInitializers(empty);
            return createFinalVarContext(empty);
        }
    }

    public boolean trackStrings() {
        return trackStrings;
    }
}
