package accrue.analysis.pointer;

import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicReference;

import polyglot.main.Report;
import polyglot.types.ClassType;
import polyglot.types.CodeInstance;
import polyglot.types.MemberInstance;
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.util.Position;
import accrue.analysis.AnalysisTopics;
import accrue.analysis.pointer.PointsToEngine.StmtAndContext;
import accrue.analysis.pointer.StmtProcedureCall.ExceptionHandlerPointsToGraphNode;
import accrue.analysis.pointer.analyses.HeapAbstractionFactory;

public class PointsToLibrarySigs {
    /**
     * Returns a singleton.
     */
    public static class Singleton extends LibraryPointsToSignature {
        protected static Map<String, AllocSiteNode> singletons =
                new HashMap<String, AllocSiteNode>();
        /**
         * The identifier for the singleton. There is one abstract object (i.e., AllocSiteNode) for
         * a given identifier.
         */
        protected final String id;
        /**
         * The type of the abstract object. 
         */
        protected String typeName = null;
        protected Type type = null;

        protected boolean isPreciseType;

        public Singleton(String id, Type type, boolean isPreciseType) {
            this.id = id;
            this.type = type;
            this.isPreciseType = isPreciseType;
        }

        public Singleton(String id, String typeName, boolean isPreciseType) {
            this.id = id;
            this.typeName = typeName;
            this.isPreciseType = isPreciseType;
        }

        @Override
        public boolean process(CodeInstance caller, CContext callerContext,
                CallSiteNode callSiteNode, ProcedureInstance mi,
                List<ReferenceVariableReplica> args, HContext receiver,
                ReferenceVariableReplica result,
                List<ExceptionHandlerPointsToGraphNode> handlers,
                Map<Type, ReferenceVariableReplica> procThrows,
                StmtRegistrar stmtRegistrar, PointsToGraph g, Position origin,
                HeapAbstractionFactory af, StmtAndContext sac)
                throws SemanticException {

            AllocSiteNode s = singletons.get(id);
            if (s == null) {
                TypeSystem ts = mi.typeSystem();
                if (type == null) type = ts.typeForName(typeName);
                s =
                        new AllocSiteNode(type,
                                          ((MemberInstance) mi).container()
                                                               .toClass(),
                                          isPreciseType,
                                          "singleton-for-" + type,
                                          Position.COMPILER_GENERATED);
                singletons.put(id, s);
            }
            if (result == null) return false;
            return g.addEdge(sac, result, af.record(callerContext, s), origin);

        }
    }

    public static class AssignsArgToField extends LibraryPointsToSignature {
        public AssignsArgToField(String fieldname, String fieldTypeName,
                int argIndex) {
            this.fieldname = fieldname;
            this.fieldTypeName = fieldTypeName;
            this.argIndex = argIndex;
        }

        final String fieldname;
        final String fieldTypeName;
        final int argIndex;

        @Override
        public boolean process(CodeInstance caller, CContext callerContext,
                CallSiteNode callSiteNode, ProcedureInstance mi,
                List<ReferenceVariableReplica> args, HContext receiver,
                ReferenceVariableReplica result,
                List<ExceptionHandlerPointsToGraphNode> handlers,
                Map<Type, ReferenceVariableReplica> procThrows,
                StmtRegistrar stmtRegistrar, PointsToGraph g, Position origin,
                HeapAbstractionFactory af, StmtAndContext sac)
                throws SemanticException {

            ObjectField node;
            Type receiverType = receiver.type();
            Type fieldType = mi.typeSystem().typeForName(fieldTypeName);
            if (receiverType.isReference()
                    && receiverType.toReference().fieldNamed(fieldname) != null) {
                node =
                        ObjectField.create(receiver,
                                           receiverType.toReference()
                                                       .fieldNamed(fieldname));
            }
            else {
                node = ObjectField.create(receiver, fieldname, fieldType);
            }

            // read the node to record a dependency
            g.pointsTo(sac, node);

            boolean changed = false;
            for (HContext p : g.pointsTo(sac, args.get(argIndex))) {
                changed = g.addEdge(sac, node, p, origin) || changed;
            }
            return changed;

        }

    }

    public static class ReturnsField extends LibraryPointsToSignature {
        public ReturnsField(String fieldname, String fieldTypeName) {
            this.fieldname = fieldname;
            this.fieldTypeName = fieldTypeName;
        }

        final String fieldname;
        final String fieldTypeName;

        @Override
        public boolean process(CodeInstance caller, CContext callerContext,
                CallSiteNode callSiteNode, ProcedureInstance mi,
                List<ReferenceVariableReplica> args, HContext receiver,
                ReferenceVariableReplica result,
                List<ExceptionHandlerPointsToGraphNode> handlers,
                Map<Type, ReferenceVariableReplica> procThrows,
                StmtRegistrar stmtRegistrar, PointsToGraph g, Position origin,
                HeapAbstractionFactory af, StmtAndContext sac)
                throws SemanticException {

            if (result == null) return false;

            ObjectField node;
            Type receiverType = receiver.type();
            Type fieldType = mi.typeSystem().typeForName(fieldTypeName);

            if (receiverType.isReference()
                    && receiverType.toReference().fieldNamed(fieldname) != null) {
                node =
                        ObjectField.create(receiver,
                                           receiverType.toReference()
                                                       .fieldNamed(fieldname));
            }
            else {
                node = ObjectField.create(receiver, fieldname, fieldType);
            }

            // read the node to record a dependency
            Set<HContext> fieldPointsTo = g.pointsTo(sac, node);
            boolean changed = false;
            for (HContext hc : fieldPointsTo) {
                changed |= g.addEdge(sac, result, hc, origin);
            }
            return changed;
        }
    }

    /**
     * Returns a singleton array, containing an element of the specified type.
     */
    private static class SingletonArray extends LibraryPointsToSignature {
        static Map<String, AllocSiteNode> arraySingletons =
                new HashMap<String, AllocSiteNode>();
        static Map<String, AllocSiteNode> arrayContentsSingletons =
                new HashMap<String, AllocSiteNode>();
        private final String id;
        private final String type;

        public SingletonArray(String id, String type) {
            this.id = id;
            this.type = type;
        }

        @Override
        public boolean process(CodeInstance caller, CContext callerContext,
                CallSiteNode callSiteNode, ProcedureInstance mi,
                List<ReferenceVariableReplica> args, HContext receiver,
                ReferenceVariableReplica result,
                List<ExceptionHandlerPointsToGraphNode> handlers,
                Map<Type, ReferenceVariableReplica> procThrows,
                StmtRegistrar stmtRegistrar, PointsToGraph g, Position origin,
                HeapAbstractionFactory af, StmtAndContext sac)
                throws SemanticException {

            if (result == null) return false;

            AllocSiteNode s = arraySingletons.get(id);
            TypeSystem ts = mi.typeSystem();
            Type baseType = ts.typeForName(type);
            if (s == null) {
                s =
                        new AllocSiteNode(ts.arrayOf(baseType),
                                          ((MemberInstance) mi).container()
                                                               .toClass(),
                                          true,
                                          "array-singleton-for-" + type,
                                          Position.COMPILER_GENERATED);
                arraySingletons.put(id, s);
            }
            AllocSiteNode t = arrayContentsSingletons.get(id);
            if (t == null) {
                t =
                        new AllocSiteNode(baseType,
                                          ((MemberInstance) mi).container()
                                                               .toClass(),
                                          true,
                                          "array-singleton-contents-for-"
                                                  + type,
                                          Position.COMPILER_GENERATED);
                arrayContentsSingletons.put(id, t);
            }
            ObjectField node =
                    ObjectField.create(af.record(callerContext, s),
                                       StmtRegistrar.ARRAY_CONTENTS,
                                       baseType);
            boolean changed =
                    g.addEdge(sac, node, af.record(callerContext, t), origin);
            // read the node to record a dependency
            g.pointsTo(sac, node);

            return g.addEdge(sac, result, af.record(callerContext, s), origin)
                    || changed;

        }

    }

    /**
     * Returns a new node, and maybe also the receiver.
     * The canonical example of this is java.lang.String.substring(), which may return this, or it may allocate a new String.
     */
    public static class AllocatorOrReceiver extends LibraryPointsToSignature {
        protected final String allocatedType;
        protected final boolean mayReturnReceiver;
        protected final boolean isPreciseType;
        protected AtomicReference<AllocSiteNode> allocSite;

        public AllocatorOrReceiver(String allocatedType,
                boolean mayReturnReceiver) {
            this(allocatedType, mayReturnReceiver, true);
        }

        public AllocatorOrReceiver(String allocatedType,
                boolean mayReturnReceiver, boolean isPreciseType) {
            this.allocatedType = allocatedType;
            this.mayReturnReceiver = mayReturnReceiver;
            this.isPreciseType = isPreciseType;
            this.allocSite = new AtomicReference<AllocSiteNode>(null);
        }

        protected static Type typeForName(TypeSystem ts, String s)
                throws SemanticException {
            if (s.endsWith("[]")) {
                Type t = typeForName(ts, s.substring(0, s.length() - 2));
                return ts.arrayOf(t);
            }
            if (s.contains(".")) {
                return ts.typeForName(s);
            }
            else {
                return ts.primitiveForName(s);
            }
        }

        protected AllocSiteNode allocSiteNode(TypeSystem ts,
                ClassType container, Position pos) {
            AllocSiteNode asn = this.allocSite.get();
            if (asn == null) {
                // create the alloc site
                try {
                    asn =
                            new AllocSiteNode(typeForName(ts, allocatedType),
                                              container,
                                              isPreciseType,
                                              "allcatorOrReceiver",
                                              pos);
                    if (!this.allocSite.compareAndSet(null, asn)) {
                        // this.allocSite got updated! Use the new value
                        asn = this.allocSite.get();
                    }
                }
                catch (SemanticException e) {
                    throw new InternalCompilerError("Unexpected", e);
                }

            }
            return asn;
        }

        @Override
        public boolean process(CodeInstance caller, CContext callerContext,
                CallSiteNode callSiteNode, ProcedureInstance mi,
                List<ReferenceVariableReplica> args, HContext receiver,
                ReferenceVariableReplica result,
                List<ExceptionHandlerPointsToGraphNode> handlers,
                Map<Type, ReferenceVariableReplica> procThrows,
                StmtRegistrar stmtRegistrar, PointsToGraph g, Position origin,
                HeapAbstractionFactory af, StmtAndContext sac)
                throws SemanticException {
            if (result == null) return false;
            AllocSiteNode asn =
                    allocSiteNode(mi.typeSystem(),
                                  ((MemberInstance) mi).container().toClass(),
                                  origin);
            HContext newObjectName;
            newObjectName = af.record(callerContext, asn);

            boolean changed = g.addEdge(sac, result, newObjectName, origin);
            if (receiver != null && mayReturnReceiver) {
                changed = g.addEdge(sac, result, receiver, origin) || changed;
            }

            return changed;

        }

    }

    /**
     * Returns a new node, and also the receiver.
     */
    private static class ReturnsReceiver extends LibraryPointsToSignature {

        public ReturnsReceiver() {
        }

        @Override
        public boolean process(CodeInstance caller, CContext callerContext,
                CallSiteNode callSiteNode, ProcedureInstance mi,
                List<ReferenceVariableReplica> args, HContext receiver,
                ReferenceVariableReplica result,
                List<ExceptionHandlerPointsToGraphNode> handlers,
                Map<Type, ReferenceVariableReplica> procThrows,
                StmtRegistrar stmtRegistrar, PointsToGraph g, Position origin,
                HeapAbstractionFactory af, StmtAndContext sac)
                throws SemanticException {

            if (result == null) return false;
            if (receiver != null) {
                return g.addEdge(sac, result, receiver, origin);
            }

            return false;

        }

    }

    /**
     * Library signature that does not add any points to statements.
     */
    public static LibraryPointsToSignature NOOP =
            new LibraryPointsToSignature() {
                @Override
                public boolean process(CodeInstance caller,
                        CContext callerContext, CallSiteNode callSiteNode,
                        ProcedureInstance mi,
                        List<ReferenceVariableReplica> args, HContext receiver,
                        ReferenceVariableReplica result,
                        List<ExceptionHandlerPointsToGraphNode> handlers,
                        Map<Type, ReferenceVariableReplica> procThrows,
                        StmtRegistrar stmtRegistrar, PointsToGraph g,
                        Position origin, HeapAbstractionFactory af,
                        StmtAndContext sac) throws SemanticException {

                    return false;
                }

            };

    private static class InvokeMethodOfArg extends LibraryPointsToSignature {
        final MethodInstance invokeMethod;
        final int argNum;

        InvokeMethodOfArg(MethodInstance invokeMethod, int argNum) {
            this.invokeMethod = invokeMethod;
            this.argNum = argNum;
        }

        @Override
        public boolean process(CodeInstance caller, CContext callerContext,
                CallSiteNode callSiteNode, ProcedureInstance mi,
                List<ReferenceVariableReplica> args, HContext receiver,
                ReferenceVariableReplica result,
                List<ExceptionHandlerPointsToGraphNode> handlers,
                Map<Type, ReferenceVariableReplica> procThrows,
                StmtRegistrar stmtRegistrar, PointsToGraph g, Position origin,
                HeapAbstractionFactory af, StmtAndContext sac)
                throws SemanticException {

            return StmtVirtualMethodCall.process(result,
                                                 args.get(0),
                                                 callSiteNode,
                                                 invokeMethod,
                                                 Collections.<ReferenceVariableReplica> emptyList(),
                                                 handlers,
                                                 procThrows,
                                                 caller,
                                                 callerContext,
                                                 g,
                                                 stmtRegistrar,
                                                 origin,
                                                 af,
                                                 sac);
        }

    }

    /**
     * Implements java.lang.System.arrayCopy(Object src,  int  srcPos,
                                        Object dest, int destPos,
                                        int length)
     */
    private static class ArrayCopy extends LibraryPointsToSignature {
        @Override
        public boolean process(CodeInstance caller, CContext callerContext,
                CallSiteNode callSiteNode, ProcedureInstance mi,
                List<ReferenceVariableReplica> args, HContext receiver,
                ReferenceVariableReplica result,
                List<ExceptionHandlerPointsToGraphNode> handlers,
                Map<Type, ReferenceVariableReplica> procThrows,
                StmtRegistrar stmtRegistrar, PointsToGraph g, Position origin,
                HeapAbstractionFactory af, StmtAndContext sac)
                throws SemanticException {
            ReferenceVariableReplica src = args.get(0);
            ReferenceVariableReplica dest = args.get(1); // primitives aren't added to the list of args.

            // fake a dest.contents = src.contents
            // should be like a t = src.contents,  dest.contents = t
            //

            if (Report.should_report(AnalysisTopics.POINTSTO, 3))
                Report.report(3, "Processing array copy");

            Set<HContext> trgs = new HashSet<HContext>();
            for (HContext j : g.pointsTo(sac, src)) {
                ObjectField jn =
                        ObjectField.create(j,
                                           StmtRegistrar.ARRAY_CONTENTS,
                                           src.expectedType().toArray().base());
                for (HContext i : g.pointsTo(sac, jn)) {
                    trgs.add(i);
                }
            }

            boolean changed = false;
            for (HContext i : g.pointsTo(sac, dest)) {
                ObjectField node =
                        ObjectField.create(i,
                                           StmtRegistrar.ARRAY_CONTENTS,
                                           dest.expectedType().toArray().base());
                changed = g.addEdges(sac, node, trgs, origin) || changed;
            }

            return changed;
        }

    }

}
