package pta;

import java.util.ArrayList;
import java.util.Collection;
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 driver.PTAOptions;
import pag.MethodPAG;
import pag.builder.MtdPAGBuilder;
import pag.node.GNode;
import pag.node.alloc.Alloc_Node;
import pag.node.alloc.SummaryNode;
import pag.node.call.VirtualInvokeSite;
import pag.node.var.AllocDotField_Node;
import pag.node.var.FieldRef_Node;
import pag.node.var.Var_Node;
import pta.pts.PTSetInternal;
import pta.pts.PTSetVisitor;
import soot.ArrayType;
import soot.FastHierarchy;
import soot.Kind;
import soot.Local;
import soot.MethodOrMethodContext;
import soot.PointsToAnalysis;
import soot.RefLikeType;
import soot.RefType;
import soot.Scene;
import soot.SootClass;
import soot.SootField;
import soot.SootMethod;
import soot.SootMethodRef;
import soot.Type;
import soot.Value;
import soot.jimple.AssignStmt;
import soot.jimple.Constant;
import soot.jimple.InstanceInvokeExpr;
import soot.jimple.IntConstant;
import soot.jimple.InvokeExpr;
import soot.jimple.Jimple;
import soot.jimple.NewArrayExpr;
import soot.jimple.NewExpr;
import soot.jimple.NewMultiArrayExpr;
import soot.jimple.NullConstant;
import soot.jimple.ParameterRef;
import soot.jimple.SpecialInvokeExpr;
import soot.jimple.Stmt;
import soot.jimple.ThisRef;
import soot.jimple.internal.JAssignStmt;
import soot.jimple.internal.JIdentityStmt;
import soot.jimple.internal.JInterfaceInvokeExpr;
import soot.jimple.internal.JInvokeStmt;
import soot.jimple.internal.JReturnStmt;
import soot.jimple.internal.JVirtualInvokeExpr;
import soot.jimple.internal.JimpleLocal;
import soot.jimple.spark.pag.ArrayElement;
import soot.jimple.spark.pag.Parm;
import soot.jimple.spark.pag.SparkField;
import soot.jimple.spark.solver.Propagator;
import soot.jimple.toolkits.callgraph.VirtualCalls;
import soot.toolkits.scalar.Pair;
import soot.util.queue.ChunkedQueue;
import soot.util.queue.QueueReader;
import util.Debug;
import util.PTAEvaluator;

/**sub-PTA. will generate summaries*/
public class SubPTA extends PTA implements SubPTAInterface{
	public boolean ignore;
	@Override
	public boolean isIgnore() {
		return ignore;
	}
	@Override
	public void setIgnore() {
		ignore=true;
	}
	
	SootMethod entryMethod;
	@Override
	public SootMethod getEntryMethod() {
		return entryMethod;
	}
	Collection<Stmt> summary;
	
	/**Implementation of Algorithm 2*/
	@Override
	public Collection<Stmt> makeSummary() {
		if(summary==null) {
			if(ignore) {
				Debug.println("reachables: >"+PTAOptions.maxReachables);
				ignores.add(entryMethod);
				return getMethodStmts(entryMethod);
			}
			Debug.println("reachables: "+cgb.reachables.size());
			Debug.println("make Summary for:"+entryMethod);
			int[] oldSize=new int[1];
			cgb.reachables.listener().forEachRemaining(rm->{
				oldSize[0]+=getMethodStmts(rm.method()).size();
			});
			Debug.println("old size: "+oldSize[0]);
			summary = new HashSet<>();
			MtdPAGBuilder nodeFactory = getMethodPAG(entryMethod).nodeFactory();
			
			//* -> ret:
			if(entryMethod.getReturnType() instanceof RefLikeType) {
				Var_Node ret = (Var_Node) nodeFactory.caseRet();
				ret.getP2Set().forall(new PTSetVisitor() {
					@Override
					public void visit(GNode n) {
						if(!(((Alloc_Node) n).getNewExpr() instanceof Alloc_Node)) // ignore arg node
							addAssign(ret, (Alloc_Node) n);// process n and n... // r = n;
					}
				});
			}
			
			//* -> global:
			pag.getStaticStore().values().stream().flatMap(s->s.stream()).distinct().forEach(g->{
				if(!(g.getVariable() instanceof Constant)&&!(g.getVariable() instanceof Type))
					g.getP2Set().forall(new PTSetVisitor() {
						@Override
						public void visit(GNode n) {
							if(!(((Alloc_Node) n).getNewExpr() instanceof Alloc_Node)) // ignore arg node
								addAssign(g, (Alloc_Node) n);// process n and n... // g = n;
						}
					});
			});
			
			//* -*-> sum:
			modified.forEach(odf->{
				SummaryNode baseSummaryNode = (SummaryNode) odf.getBase();
				SparkField field = odf.getField();
				Value base = getValue(baseSummaryNode);
				odf.getP2Set().forall(new PTSetVisitor() {
					@Override
					public void visit(GNode n) {
						if(!(((Alloc_Node) n).getNewExpr() instanceof Alloc_Node)) // ignore arg node
							addStore(base, field, getValue((Alloc_Node) n));
					}
				});
			});
			
			//full summary: -> receiver:
			getCgb().getUnprocessedCallsites().forEach(site->{
				MtdPAGBuilder srcnf = getMethodPAG(site.container().method()).nodeFactory();
				Stmt s = site.stmt();
				InvokeExpr ie = s.getInvokeExpr();
				SootMethod tgtmtd = ie.getMethod();
				int numArgs = ie.getArgCount();
				for (int i = 0; i < numArgs; i++) {
					Value arg = ie.getArg(i);
					if (!(arg.getType() instanceof RefLikeType)|| arg instanceof NullConstant)
						continue;
					Type tgtType = tgtmtd.getParameterType(i);
					if (!(tgtType instanceof RefLikeType) )
						continue;
					if(arg instanceof Constant) // no other value should store into const
						continue;
					Var_Node argNode = (Var_Node) srcnf.getNode(arg);
					argNode.getP2Set().forall(new PTSetVisitor() {
						@Override
						public void visit(GNode n) {
							if(n instanceof SummaryNode) // all normal nodes have been changed to summary nodes
								addAssign(argNode, (Alloc_Node) n);// process n and n... // arg = n;
						}
					});
				}
				InstanceInvokeExpr iie = (InstanceInvokeExpr) ie;
				Local receiver = (Local) iie.getBase();
				Var_Node recNode = getReceiverVarNode(receiver, site.container());
				Value newRecValue = newRecLocal(recNode.getType());
				recNode.getP2Set().forall(new PTSetVisitor() {
					@Override
					public void visit(GNode n) {
						if(isFullSummaryNode(n))
							addAssign(newRecValue, getValue((Alloc_Node) n));// process n and n... // arg = n;
					}
				});
				addNewInvoke(newRecValue, s);
			});
			
			Debug.println("new size: "+summary.size());
			
			SumPTA.scalableMethods.add(entryMethod);
		}
		return summary;
	}
	
	Map<Alloc_Node, Value> a2v = new HashMap<>();
	
	/**Implementation of Algorithm 3*/
	Value getValue(Alloc_Node n) {
		Value ret = a2v.get(n);
		if(ret==null) {
			if(n instanceof SummaryNode) {
				Object newExpr = n.getNewExpr();
				if(newExpr instanceof Alloc_Node) { // arg summary node
					a2v.put(n, ret = getValue((Alloc_Node) newExpr)); // share the local for real alloc
				}else if(newExpr instanceof Var_Node) {
					Var_Node parmNode = (Var_Node) newExpr;
					Object parmVar = parmNode.getVariable();
					if(parmVar instanceof Local) { // call assign target
						ret = (Value) parmVar;
					}else if(parmVar instanceof SootField){ // global
						Value rValue = Jimple.v().newStaticFieldRef(((SootField) parmVar).makeRef());
						ret = new JimpleLocal("g-local:"+parmVar, parmNode.getType());
						addAssign(ret, rValue);
					}else if(parmVar instanceof Parm) { // Parm
						Parm parmParm = (Parm) parmVar;
						int index = parmParm.getIndex();
						Value parmRef = new ParameterRef(parmNode.getType(), index);
						ret = new JimpleLocal("parameter"+index, parmNode.getType());
						addIdentity(ret, parmRef);
					}else if(parmVar instanceof Pair) { // This
						Pair<?, ?> parmPair = (Pair<?, ?>) parmVar;
						if(parmPair.getO2()==PointsToAnalysis.THIS_NODE) {
							if(parmPair.getO1()!=entryMethod)
								throw new RuntimeException("unexpected THIS from: "+parmPair.getO1());
							Value parmRef = Jimple.v().newThisRef((RefType) parmNode.getType());
							ret = Jimple.v().newLocal("this", parmNode.getType());
							addIdentity(ret, parmRef);
						}else
							throw new RuntimeException("unknown pair: "+parmPair);
					}else
						throw new RuntimeException("unknown var:"+parmVar);	
					a2v.put(n, ret);
				}else if(newExpr instanceof FieldRef_Node){ // load
					FieldRef_Node fr = (FieldRef_Node) newExpr;
					Var_Node base = fr.getBase();
					Value baseValue = newTempLocal(base.getType());
					SparkField field = fr.getField();
					ret = newTempLocal(field.getType());
					Value rValue = field instanceof ArrayElement?Jimple.v().newArrayRef(baseValue, IntConstant.v(0)):
						Jimple.v().newInstanceFieldRef(baseValue, ((SootField)field).makeRef());
					summary.add(new JAssignStmt(ret, rValue));
					a2v.put(n, ret);
					base.getP2Set().forall(new PTSetVisitor() {
						@Override
						public void visit(GNode n) {
							if(n instanceof SummaryNode) {
								Value rValue = getValue((Alloc_Node) n);
								addAssign(baseValue, rValue);
							}
						}
					});
				}else if(newExpr instanceof Pair){ // cast / exception
					Pair<Type, ?> pair = (Pair<Type, ?>) newExpr;
					Type castType = pair.getO1();
					Object srcObj = pair.getO2();
					if(srcObj instanceof Var_Node) {  // cast
						Var_Node src = (Var_Node) srcObj;
						Value srcValue = newTempLocal(src.getType());
						ret = newTempLocal(castType);
						Value rValue = Jimple.v().newCastExpr(srcValue, castType);
						summary.add(new JAssignStmt(ret, rValue));
						a2v.put(n, ret);
						src.getP2Set().forall(new PTSetVisitor() {
							@Override
							public void visit(GNode n) {
								if(isFullSummaryNode(n)) {
									Value rValue = getValue((Alloc_Node) n);
									addAssign(srcValue, rValue);
								}
							}
						});
					}else if(srcObj==PointsToAnalysis.EXCEPTION_NODE) {  // exception
						Value rValue = Jimple.v().newCaughtExceptionRef();
						ret = new JimpleLocal("("+castType+")exception", castType);
						addIdentity(ret, rValue);
						a2v.put(n, ret);
					}else
						throw new RuntimeException("unknown pair: "+pair);
				}else
					throw new RuntimeException("unknown expr:"+newExpr);
			}else { // allocNode
				Object newExpr = n.getNewExpr();
				if(newExpr instanceof Constant) {
					a2v.put(n, ret = (Constant)newExpr);
				}else {
					a2v.put(n, ret = newTempLocal(n.getType()));
					Value rValue;
					if(newExpr instanceof Pair) { // multiArray
						Pair<NewMultiArrayExpr, Integer> pair = (Pair<NewMultiArrayExpr, Integer>) newExpr;
						NewMultiArrayExpr newMultiArrayExpr = pair.getO1();
						ArrayType aType = (ArrayType) newMultiArrayExpr.getType();
						int dimensions = aType.numDimensions;
						int currentDimension = pair.getO2();
						if(currentDimension==dimensions) {
							List<IntConstant> list= new ArrayList<>();
							for(int i=0;i<dimensions;i++)
								list.add(IntConstant.v(1));
							rValue = Jimple.v().newNewMultiArrayExpr(aType, list);
							if (PTAOptions.originalAlloc)
								PTAEvaluator.exprMap.put(rValue, newMultiArrayExpr);
						}else {
							Value supValue = getValue(pag.valToAllocNode.get(new Pair<>(newMultiArrayExpr,currentDimension+1)));
							rValue = Jimple.v().newArrayRef(supValue, IntConstant.v(0));
						}
						if(currentDimension!=1) { // if a multiarray>2, its 1-field should also be record.
							Object expr1 = new Pair<>(newMultiArrayExpr,1);
							Alloc_Node node1 = pag.valToAllocNode.get(expr1);
							AllocDotField_Node adf = node1.dot(ArrayElement.v());
							if(adf!=null&&!adf.getP2Set().isEmpty())
								getValue(node1);
							addAssign(ret, rValue);
							return ret;
						}
					}else if(newExpr instanceof Type) // merged
						rValue = Jimple.v().newNewExpr((RefType) newExpr);
					else{	
						if(newExpr instanceof NewArrayExpr)
							rValue = Jimple.v().newNewArrayExpr(((NewArrayExpr)newExpr).getBaseType(), IntConstant.v(0));
						else
							rValue = Jimple.v().newNewExpr(((NewExpr)newExpr).getBaseType());
						if (PTAOptions.originalAlloc)
							PTAEvaluator.exprMap.put(rValue, newExpr);
					}
					addAssign(ret, rValue);
					// further add all dependencies
					final Value retF = ret;
					n.getAllFieldRefs().forEach(adf->{
						SparkField field = adf.getField();
						adf.getP2Set().forall(new PTSetVisitor() {
							@Override
							public void visit(GNode n) {
								if(!(((Alloc_Node) n).getNewExpr() instanceof Alloc_Node)) // ignore arg summary
									addStore(retF, field, getValue((Alloc_Node) n));
							}
						});
					});
				}
			}
		}
		return ret;
	}
	
	int tempLocals=0;
	private JimpleLocal newTempLocal(Type type){
		return new JimpleLocal("t"+tempLocals++, type);
	}
	int recLocals=0;
	private JimpleLocal newRecLocal(Type type){
		return new JimpleLocal("rec"+recLocals++, type);
	}
	
	void addAssign(Var_Node lNode, Alloc_Node rNode) {
		Value rValue = getValue(rNode);
		Object lVar = lNode.getVariable();
		if(lVar instanceof Parm) { // return
			Parm parm = (Parm) lVar;
			if(parm.getIndex()==PointsToAnalysis.RETURN_NODE)
				summary.add(new JReturnStmt(rValue));
			else
				throw new RuntimeException("parm assigned");
		}else{
			Value lValue = null;
			if(lVar instanceof SootField) // global
				lValue = Jimple.v().newStaticFieldRef(((SootField) lVar).makeRef());
			else if(lVar==PointsToAnalysis.EXCEPTION_NODE) { // throw
				summary.add(Jimple.v().newThrowStmt(rValue));
				return;
			}else if(lVar instanceof Pair) // const
				throw new RuntimeException("lvar is pair");
			else
				lValue = (Value) lVar;
			addAssign(lValue, rValue);
		}
	}

	void addAssign(Value lValue, Value rValue) {
		summary.add(new JAssignStmt(lValue, rValue));
	}
	
	void addIdentity(Value lValue, Value rValue) {
		summary.add(new JIdentityStmt(lValue, rValue));
	}
	
	void addStore(Value lBaseValue, SparkField field, Value rValue) {
		Value lValue = field instanceof ArrayElement?Jimple.v().newArrayRef(lBaseValue, IntConstant.v(0)):
			Jimple.v().newInstanceFieldRef(lBaseValue, ((SootField)field).makeRef());
		addAssign(lValue, rValue);
	}
	
	/**change receiver of arg2 to arg1 and add the new invoke*/
	void addNewInvoke(Value receiver, Stmt s){
		InstanceInvokeExpr expr = (InstanceInvokeExpr) s.getInvokeExpr();
		SootMethodRef methodRef = expr.getMethodRef();
		List<Value> argsL = expr.getArgs();
		Value invoke = methodRef.declaringClass().isInterface()?new JInterfaceInvokeExpr(receiver,methodRef,argsL):new JVirtualInvokeExpr(receiver,methodRef,argsL);
		Stmt stmt = null;
		if (s instanceof AssignStmt as) {
			Value dest = as.getLeftOp();
			if (dest.getType() instanceof RefLikeType) {
				stmt = new JAssignStmt(dest, invoke);
			}else
				stmt = new JInvokeStmt(invoke);
		}else
			stmt = new JInvokeStmt(invoke);
		summary.add(stmt);
		if(PTAOptions.originalCall)
			PTAEvaluator.exprMap.put(stmt, s);
	}
	
	Set<AllocDotField_Node> modified = new HashSet<>();
	/** record modified summary nodes which is to be handled in Algorithm 2(Lines 9-14).*/
	@Override
	public void recordModified(AllocDotField_Node oDotF) {
		modified.add(oDotF);
	}
	
	static Collection<SootMethod> visit = new HashSet<>(), ignores = new HashSet<>();
	
	public SubPTA(SootMethod entryMethod) {
		this.entryMethod = entryMethod;
		visit.add(entryMethod);
	}
	
	@Override
	public MethodPAG getMethodPAG(SootMethod m){
		if(!PTAOptions.topo)
			return super.getMethodPAG(m);

		MethodPAG ret = methodToPag.get(m);
		if (ret == null) {
			Collection<Stmt> stmts;
			if(!SumPTA.isLib(m)) {
				stmts = PTA.getMethodStmts(m);
			}else { // lib method
				if(!visit.contains(m)) {
					Debug.incOffset();
					stmts = SumPTA.getMethodStmts(m);
					Debug.decOffset();
					if(ignores.contains(m))
						ignore=true;
				}else { // visiting or visited
					stmts = SumPTA.methodToSumStmts.get(m);
					if(stmts==null) { // visiting
						stmts = PTA.getMethodStmts(m);
					}else { // visited
						if(ignores.contains(m))
							ignore=true;
					}
				}
			}
			methodToPag.put(m, ret = new MethodPAG(pag, m, stmts));
		}
		return ret;
	}
	
	@Override
	public SubCallGraphBuilder getCgb() {
		return (SubCallGraphBuilder) cgb;
	}
	
	@Override
	protected CallGraphBuilder createCallGraphBuilder() {
		return new SubCallGraphBuilder();
	}
	
	@Override
	protected Propagator getPropagator(){
		return new pta.solver.SubSolver(this);
	}
	
	Map<Object, SummaryNode> node2summary = new HashMap<>();
	
	/**Implementation of get*Node*/
	@Override
	public SummaryNode getSummaryNode(Object n, Type type) {
		SummaryNode ret = node2summary.get(n);
		if(ret==null)
			node2summary.put(n, ret=new SummaryNode(pag, n, type, null));
		return ret;
	}
	@Override
	public boolean isFullSummaryNode(GNode n) {
        return n instanceof SummaryNode sn && !sn.isPartial();
    }
	
	@Override
	public void build() {
		while (true) {
			if(!reachablesReader.hasNext()){
				cgb.reachables.update();
				if(!reachablesReader.hasNext())
					break;
			}
			MethodOrMethodContext momc = reachablesReader.next();
			SootMethod method = momc.method();
			if (method.isPhantom())
				continue;
			if (ignores.contains(method)) {
				setIgnore();
				return;
			}

			MethodPAG mpag = getMethodPAG(method);
			if(ignore)
				return;
			addToPAG(mpag, momc.context());
			
			if(PTAOptions.sumclinits)
				updateClinits(mpag.clinits);
			
			cgb.handleInvoke(momc, mpag.invokeStmts);
		}
		callEdges.forEachRemaining(e -> processCallEdge(e));
	}
	Set<SootMethod> clinits=new HashSet<>();
	private void updateClinits(Set<SootMethod> clinits) {
		clinits.stream().filter(this.clinits::add).forEach(cgb.reachables::addMethod);
	}
	
	@Override
	public QueueReader<Var_Node> unprocessedArgQueueReader(){
		return getCgb().unprocessedArgQueueReader();
	}
	@Override
	public boolean isUnprocessedArg(Var_Node arg) {
		return getCgb().isUnprocessedArg(arg);
	}
	
	@Override
	public String toString() {
		return "subpta for: "+entryMethod;
	}
	
	public class SubCallGraphBuilder extends CallGraphBuilder{
		protected Set<VirtualInvokeSite> unprocessedCallsites = new HashSet<>();
		final Set<Var_Node> unprocessedArgs = new HashSet<>();
		final ChunkedQueue<Var_Node> unprocessedArgQueue = new ChunkedQueue<>();
		
		public Set<VirtualInvokeSite> getUnprocessedCallsites(){
			return unprocessedCallsites;
		}
		public QueueReader<Var_Node> unprocessedArgQueueReader(){
			return unprocessedArgQueue.reader();
		}
		public boolean isUnprocessedArg(Var_Node arg) {
			return unprocessedArgs.contains(arg);
		}
		
		@Override
		protected List<? extends MethodOrMethodContext> getEntryPoints() {
			return Collections.singletonList(entryMethod);
		}
		@Override
		public void updateCallGraph(PTSetInternal p2set, VirtualInvokeSite site) {
			p2set.forall(new PTSetVisitor() {
				public final void visit(GNode n) {
					Alloc_Node receiverNode = (Alloc_Node) n;
					dispatch(receiverNode, site);
					while (targets.hasNext())
						addVirtualEdge(site.container(), site.stmt(), targets.next(), site.kind(), receiverNode);
				}
			});
			reflectionModel.updateNode(site.recNode(), p2set);
		}
		/**replace dispatch(type, site) as alloc*/
		protected void dispatch(Alloc_Node alloc, VirtualInvokeSite site) {
			Type type = alloc.getType(); 
			FastHierarchy fh = Scene.v().getOrMakeFastHierarchy();
			
			MethodOrMethodContext container = site.container();
			if (site.kind() == Kind.THREAD && !fh.canStoreType(type, clRunnable)) {
			} else if (site.iie() instanceof SpecialInvokeExpr && site.kind() != Kind.THREAD) {
				SootMethod target = VirtualCalls.v().resolveSpecial((SpecialInvokeExpr) site.iie(), site.subSig(),container.method());
				if (target != null)
					targetsQueue.add(target);
			} else if(SubPTAInterface.isPoly(site) && isFullSummaryNode(alloc)) {
				addUnprocessedCallsite(site);
			} else
				VirtualCalls.v().resolve(type, site.recNode().getType(), site.subSig(), container.method(), targetsQueue);
		}
		
		/**handle unknown calls*/
		protected void addUnprocessedCallsite(VirtualInvokeSite site) {
			if(unprocessedCallsites.add(site)) {
				MtdPAGBuilder srcnf = getMethodPAG(site.container().method()).nodeFactory();
				Stmt s = site.stmt();
				InvokeExpr ie = s.getInvokeExpr();
				SootMethod tgtmtd = ie.getMethod();
				int numArgs = ie.getArgCount();
				for (int i = 0; i < numArgs; i++) {
					Value arg = ie.getArg(i);
					if (!(arg.getType() instanceof RefLikeType)|| arg instanceof NullConstant)
						continue;
					Type tgtType = tgtmtd.getParameterType(i);
					if (!(tgtType instanceof RefLikeType) )
						continue;
					Var_Node argNode = (Var_Node) srcnf.getNode(arg);
					
					if(unprocessedArgs.add(argNode))
						unprocessedArgQueue.add(argNode);
				}
				if (s instanceof AssignStmt) {
					Value dest = ((AssignStmt) s).getLeftOp();
					if (dest.getType() instanceof RefLikeType && tgtmtd.getReturnType()instanceof RefLikeType) {
						Var_Node destNode = (Var_Node) srcnf.getNode(dest);
						SummaryNode retSummary = getSummaryNode(destNode);
						pag.addAllocEdge(retSummary, destNode);
					}
				}
			}
		}
	}
}
