package pta.solver;

import driver.PTAOptions;
import pag.MethodPAG;
import pag.PAG;
import pag.builder.MtdPAGBuilder;
import pag.node.GNode;
import pag.node.alloc.Alloc_Node;
import pag.node.alloc.SummaryNode;
import pag.node.var.AllocDotField_Node;
import pag.node.var.FieldRef_Node;
import pag.node.var.GlobalVar_Node;
import pag.node.var.LocalVar_Node;
import pag.node.var.Var_Node;
import pta.PTA;
import pta.SubPTAInterface;
import pta.pts.PTSetInternal;
import pta.pts.PTSetVisitor;
import soot.PointsToAnalysis;
import soot.RefLikeType;
import soot.SootMethod;
import soot.Type;
import soot.jimple.CastExpr;
import soot.jimple.Constant;
import soot.jimple.spark.pag.SparkField;
import soot.toolkits.scalar.Pair;
import soot.util.queue.QueueReader;

public class SubSolver extends Solver {
	
	public SubSolver(SubPTAInterface _pta) {
		super((PTA) _pta);
	}
	
	@Override
	public final void propagate() {
		final QueueReader<Pair<Alloc_Node, Var_Node>> addedAllocEdges = pag.allocReader();
		final QueueReader<Pair<Var_Node, Var_Node>> addedSimpleEdges = pag.simpleReader();
		final QueueReader<Pair<Var_Node, GlobalVar_Node>> addedStaticStoreEdges = pag.staticStoreReader();
		final QueueReader<Pair<GlobalVar_Node, Var_Node>> addedStaticLoadEdges = pag.staticLoadReader();
		final QueueReader<Pair<FieldRef_Node, Var_Node>> addedStoreInvEdges = pag.storeInvReader();
		final QueueReader<Pair<FieldRef_Node, Var_Node>> addedLoadEdges = pag.loadReader();
		final QueueReader<Var_Node> unprocessedArgReader = ((SubPTAInterface) pta).unprocessedArgQueueReader();
		
		// handle parameters of entry method.
		SootMethod entryMethod = ((SubPTAInterface)pta).getEntryMethod();
		MethodPAG mpag = pta.getMethodPAG(entryMethod);
		MtdPAGBuilder nodeFactory = mpag.nodeFactory();
		int parmnum = entryMethod.getParameterCount();
		for(int i=0;i<parmnum;i++) {
			Type type = entryMethod.getParameterType(i);
			if (!(type instanceof RefLikeType))
				continue;
			LocalVar_Node parm = (LocalVar_Node)nodeFactory.caseParm(i);
			Alloc_Node argv = ((SubPTAInterface)pta).getSummaryNode(parm);
			mpag.addInternalEdge(argv, parm);
		}
		if(!entryMethod.isStatic()) {
			LocalVar_Node thisParm = (LocalVar_Node)nodeFactory.caseThis();
			Alloc_Node argv = ((SubPTAInterface)pta).getSummaryNode(thisParm);
			mpag.addInternalEdge(argv, thisParm);
		}
		
		pta.build();
		if(((SubPTAInterface)pta).isIgnore())
			return;
		
		new TopoSorter(pag, false).sort();
		 
		handleAddedSimpleEdges(addedAllocEdges,addedStaticStoreEdges,addedStaticLoadEdges,addedSimpleEdges);
		
		
		while (!varNodeWorkList.isEmpty()||!globalVarNodeWorkList.isEmpty()){
			
			while (!varNodeWorkList.isEmpty()) {
				final Var_Node src = varNodeWorkList.pollFirst();
				
				propagateFromSrc(src);
				
				propagateArg(src);
				
				handleStoreAndLoadOnBase(src);
				
				updateCallGraph(src);
				
				if(((SubPTAInterface)pta).isIgnore())
					return;
				if(cgb.getReachableMethods().size()>PTAOptions.maxReachables) {
					((SubPTAInterface)pta).setIgnore();
					return;
				}
				
				src.getP2Set().flushNew();
				
				handleAddedComplexEdges(addedStoreInvEdges,addedLoadEdges);
				
				handleAddedArg(unprocessedArgReader);
				
				handleAddedSimpleEdges(addedAllocEdges,addedStaticStoreEdges,addedStaticLoadEdges,addedSimpleEdges);
			}
			
			while (!globalVarNodeWorkList.isEmpty()) {
				final GlobalVar_Node src = globalVarNodeWorkList.pollFirst();
				
				final PTSetInternal pts = src.getP2Set();
				final PTSetInternal newset = pts.getNewSet();
				pag.staticLoadLookup(src).forEach(element->{
					if (addAll(element, newset)) {
						varNodeWorkList.add(element);
					}
				});
				pts.flushNew();
			}
		}
	}
	
	@Override
	protected void propagateFromSrc(Var_Node src) {
		final PTSetInternal newset = src.getP2Set().getNewSet();
		pag.staticStoreLookup(src).forEach(element->{
			if (addAll(element, newset)) {
				globalVarNodeWorkList.add(element);
			}
		});
		pag.simpleLookup(src).forEach(element -> {
			if(element instanceof AllocDotField_Node){
				Alloc_Node baseAlloc = ((AllocDotField_Node) element).getBase();
				if(isConstant(baseAlloc))
					throw new RuntimeException("Modifying ConstNode:" + element + " with " +src);
			}else if(element.getVariable() instanceof CastExpr) { // handle cast edges.
				if (castAddAll(element, ((CastExpr)element.getVariable()).getType(), newset,src)) {
					varNodeWorkList.add(element);
				}
				return;
			}
			if (addAll(element, newset)) {
				varNodeWorkList.add(element);
			}
		});
	}
	
	@Override
	protected void handleStoreAndLoadOnBase(Var_Node src) {
		final PTSetInternal newP2Set = src.getP2Set().getNewSet();
		for (final FieldRef_Node fr : src.getAllFieldRefs()) {
			final SparkField fld = fr.getField();
			/// foreach src.fld = v do add simple from v-->o.fld where o\in
			/// pts(src)
			for (final Var_Node v : pag.storeInvLookup(fr)) {
				newP2Set.forall(new PTSetVisitor() {
					public final void visit(GNode n) {
						Alloc_Node a = (Alloc_Node) n;
						if(isConstant(a))
							return; //	cannot modify a const!
						final AllocDotField_Node oDotF = (AllocDotField_Node) pta.parameterize(fld, a);
						if(a instanceof SummaryNode) {
							if(a.getNewExpr() instanceof Alloc_Node)
								return; // ignore arg summary
							((SubPTAInterface) pta).recordModified(oDotF);
						}
						pag.addSimpleEdge(v, oDotF);	
					}
				});
			}
			/// foreach v = src.fld do add simple from o.fld-->v where o\in
			/// pts(src)
			for (final Var_Node element : pag.loadLookup(fr))
				if(newP2Set.forall(new PTSetVisitor() {
					public final void visit(GNode n) {
						if(n instanceof SummaryNode) {
							returnValue=true;
						}else {
							final Var_Node oDotF = pta.parameterize(fld, (Alloc_Node) n);
							pag.addSimpleEdge(oDotF, element);
						}
					}
				}))
					handleUnprocessedLoadEdges(fr,element);
		}
	}

	@Override
	protected void handleAddedComplexEdges(
			QueueReader<Pair<FieldRef_Node, Var_Node>> addedStoreInvEdges,
			QueueReader<Pair<FieldRef_Node, Var_Node>> addedLoadEdges) {
		
		while(addedStoreInvEdges.hasNext()){
			Pair<FieldRef_Node, Var_Node> storeEdge = addedStoreInvEdges.next();
			final FieldRef_Node tgtfrn = storeEdge.getO1();
			final Var_Node srcv = storeEdge.getO2();
			final SparkField fld = tgtfrn.getField();
			tgtfrn.getBase().getP2Set().getOldSet().forall(new PTSetVisitor() {
				public void visit(GNode n) {
					Alloc_Node a = (Alloc_Node) n;
					if(isConstant(a))
						return; //	cannot modify a const!
					final AllocDotField_Node oDotF = (AllocDotField_Node) pta.parameterize(fld, a);
					if(a instanceof SummaryNode) {
						if(a.getNewExpr() instanceof Alloc_Node)
							return; // ignore arg summary
						((SubPTAInterface) pta).recordModified(oDotF);
					}
					pag.addSimpleEdge(srcv, oDotF);
				}
			});
		}		
		while(addedLoadEdges.hasNext()){
			Pair<FieldRef_Node, Var_Node> loadEdge = addedLoadEdges.next();
			final FieldRef_Node srcfrn = loadEdge.getO1();
			final Var_Node tgtv = loadEdge.getO2();
			final SparkField fld = srcfrn.getField();
			if(srcfrn.getBase().getP2Set().getOldSet().forall(new PTSetVisitor() {
				public void visit(GNode n) {
					if(n instanceof SummaryNode) {
						returnValue=true;
					}else {
						final Var_Node oDotF = pta.parameterize(fld, (Alloc_Node) n);
						pag.addSimpleEdge(oDotF, tgtv);
					}
				}
			}))
				handleUnprocessedLoadEdges(srcfrn, tgtv);
		}
	}
	
	/** handle arguments passed in to an unknown call.*/
	private void propagateArg(Var_Node src) {
		if(((SubPTAInterface) pta).isUnprocessedArg(src)) {
			final PTSetInternal newset = src.getP2Set().getNewSet();
			newset.forall(new PTSetVisitor() {
				public final void visit(GNode n) {
					if (n instanceof SummaryNode) return;
					Alloc_Node a = (Alloc_Node) n;
					if(!PTAOptions.sumconst && isConstant(a)) return;
					SummaryNode partialSummaryNode = ((SubPTAInterface) pta).getSummaryNode(a);
					PAG.lookup(pag.getAlloc(), a).forEach(v->{
						pag.addAllocEdge(partialSummaryNode, v);
					});
				}
			});
		}
	}
	/** handle arguments passed in to an unknown call.*/
	private void handleAddedArg(QueueReader<Var_Node> unprocessedArgReader) {
		while (unprocessedArgReader.hasNext()) {
			Var_Node arg = unprocessedArgReader.next();
			final PTSetInternal oldset = arg.getP2Set().getOldSet();
			oldset.forall(new PTSetVisitor() {
				public final void visit(GNode n) {
					if (n instanceof SummaryNode) return;
					Alloc_Node a = (Alloc_Node) n;
					if(!PTAOptions.sumconst && isConstant(a)) return;
					SummaryNode partialSummaryNode = ((SubPTAInterface) pta).getSummaryNode(a);
					PAG.lookup(pag.getAlloc(), a).forEach(v->{
						pag.addAllocEdge(partialSummaryNode, v);
					});
				}
			});
		}
	}
	
	/** handle load edges from summary nodes.*/
	private void handleUnprocessedLoadEdges(FieldRef_Node src, Var_Node tgt) {
		SummaryNode loadSummaryNode = ((SubPTAInterface) pta).getSummaryNode(src, tgt.getType());
		pag.addEdge(loadSummaryNode, tgt);
	}
	
	@Override
	protected void handleAddedSimpleEdges(QueueReader<Pair<Alloc_Node, Var_Node>> addedAllocEdges,
			QueueReader<Pair<Var_Node, GlobalVar_Node>> addedStaticStoreEdges, 
			QueueReader<Pair<GlobalVar_Node, Var_Node>> addedStaticLoadEdges, 
			QueueReader<Pair<Var_Node, Var_Node>> addedSimpleEdges){
		while(addedAllocEdges.hasNext()){
			Pair<Alloc_Node, Var_Node> allocEdge = addedAllocEdges.next();
			Alloc_Node src = allocEdge.getO1();
			Var_Node tgt = allocEdge.getO2();
			if (tgt.makeP2Set().add(src))
				varNodeWorkList.add(tgt);
		}
		while(addedStaticStoreEdges.hasNext()){
			Pair<Var_Node, GlobalVar_Node> staticStoreEdge = addedStaticStoreEdges.next();
			final Var_Node srcv = staticStoreEdge.getO1();
			final GlobalVar_Node tgtv = staticStoreEdge.getO2();
			if (addAll(tgtv, srcv.getP2Set().getOldSet())) {
				globalVarNodeWorkList.add(tgtv);
			}
		}
		while(addedStaticLoadEdges.hasNext()){
			Pair<GlobalVar_Node, Var_Node> staticLoadEdge = addedStaticLoadEdges.next();
			final GlobalVar_Node srcv = staticLoadEdge.getO1();
			final Var_Node tgtv = staticLoadEdge.getO2();
			
			// handle global load edges from global variables.
			Object variable = srcv.getVariable();
			if(!(variable instanceof Constant)&&!(variable instanceof Type)) { // constant and merged node no need to sum
				SummaryNode sn = srcv.getVariable()==PointsToAnalysis.EXCEPTION_NODE?
						((SubPTAInterface) pta).getSummaryNode(new Pair<>(tgtv.getType(),PointsToAnalysis.EXCEPTION_NODE), tgtv.getType()):
						((SubPTAInterface) pta).getSummaryNode(srcv);
				if (tgtv.makeP2Set().add(sn))
					varNodeWorkList.add(tgtv);
			}
			
			if (addAll(tgtv, srcv.getP2Set().getOldSet())) {
				varNodeWorkList.add(tgtv);
			}
		}
		
		while(addedSimpleEdges.hasNext()){
			Pair<Var_Node, Var_Node> simpleEdge = addedSimpleEdges.next();
			final Var_Node srcv = simpleEdge.getO1();
			final Var_Node tgtv = simpleEdge.getO2();
			if (addAll(tgtv, srcv.getP2Set().getOldSet()))
				varNodeWorkList.add(tgtv);
		}
	}
	
	protected boolean castAddAll(final Var_Node pointer, Type castType, PTSetInternal other, Var_Node src) {
		final PTSetInternal addTo = pointer.makeP2Set();
		// deprecated because not equal to forall->add
		// return addTo.addAll(other, null);
		boolean[] flag = new boolean[1];
		boolean ret = other.forall(new PTSetVisitor() {
			public final void visit(GNode n) {
				if (addTo.add(n))
					returnValue = true;
				else if(((SubPTAInterface)pta).isFullSummaryNode(n))
					flag[0]=true;
			}
		});
		if(flag[0]) {
			Pair<Type, Var_Node> pair = new Pair<>(castType,src);
			SummaryNode castSummaryNode = ((SubPTAInterface) pta).getSummaryNode(pair, castType);
			pag.addEdge(castSummaryNode, pointer);
		}
		return ret;
	}
}
