/* Java and Android Analysis Framework
 * Copyright (C) 2017 Jingbo Lu and Yulei Sui
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 59 Temple Place - Suite 330,
 * Boston, MA 02111-1307, USA.
 */

package pta.solver;

import pag.PAG;
import pag.node.GNode;
import pag.node.alloc.Alloc_Node;
import pag.node.alloc.Constant_Node;
import pag.node.alloc.ContextAlloc_Node;
import pag.node.call.VirtualInvokeSite;
import pag.node.var.AllocDotField_Node;
import pag.node.var.FieldRef_Node;
import pag.node.var.GlobalVar_Node;
import pag.node.var.Var_Node;

import java.util.Collection;
import java.util.Iterator;
import java.util.TreeSet;

import driver.FakeMainFactory;
import driver.PTAOptions;
import pta.PTA.CallGraphBuilder;
import pta.ContextSensPTA;
import pta.ContextSensPTA.SCGBuilder;
import pta.PTA;
import pta.SubPTA;
import pta.pts.PTSetInternal;
import pta.pts.PTSetVisitor;
import soot.Context;
import soot.MethodOrMethodContext;
import soot.RefType;
import soot.SootMethod;
import soot.Unit;
import soot.jimple.spark.pag.SparkField;
import soot.jimple.spark.solver.Propagator;
import soot.toolkits.scalar.Pair;
import soot.util.queue.QueueReader;

public class Solver extends Propagator {
	protected PAG pag;
	protected PTA pta;
	protected CallGraphBuilder cgb;
	protected final TreeSet<Var_Node> varNodeWorkList = new TreeSet<>();
	protected final TreeSet<GlobalVar_Node> globalVarNodeWorkList = new TreeSet<>();

	public Solver(PTA _pta) {
		cgb = _pta.getCgb();
		pag = _pta.getPag();
		pta = _pta;
	}

	@Override
	public 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();
		
		pta.build();
		
		if(PTAOptions.staticcontext == PTAOptions.THIS && pta instanceof ContextSensPTA && PTAOptions.ptaPattern.isObjSens()){
			ContextSensPTA csPTA = (ContextSensPTA)pta;
			Context emptyContext = csPTA.emptyContext();
			final Alloc_Node rootNode = new ContextAlloc_Node(pag, new Alloc_Node(pag, "ROOT", RefType.v("java.lang.Object"), null), emptyContext);
			SootMethod fakeMain = FakeMainFactory.getFakeMain();
			GNode thisRef = pta.getMethodPAG(fakeMain).nodeFactory().caseThis();
			thisRef = csPTA.parameterize(thisRef, emptyContext);
			pag.addEdge(rootNode, thisRef);
		}
		
		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);
				
				handleStoreAndLoadOnBase(src);
				
				updateCallGraph(src);
				
				src.getP2Set().flushNew();
				
				handleAddedComplexEdges(addedStoreInvEdges,addedLoadEdges);
				
				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();
			}
		}
	}

	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);
			}
			if (addAll(element, newset)) {
				varNodeWorkList.add(element);
			}
		});
	}

	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 Var_Node oDotF = pta.parameterize(fld, a);
						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)) {
				newP2Set.forall(new PTSetVisitor() {
					public final void visit(GNode n) {
						final Var_Node oDotF = pta.parameterize(fld, (Alloc_Node) n);
						pag.addSimpleEdge(oDotF, element);
					}
				});
			}
		}
	}

	protected void updateCallGraph(Var_Node src) {
		final QueueReader<VirtualInvokeSite> addedSites = cgb.edgeReader();
		final PTSetInternal newP2Set = src.getP2Set().getNewSet();
		Collection<VirtualInvokeSite> sites = cgb.callSitesLookUp(src);
		for (VirtualInvokeSite site : sites)
			cgb.updateCallGraph(newP2Set, site);

		if(PTAOptions.staticcontext == PTAOptions.THIS && pta instanceof ContextSensPTA && PTAOptions.ptaPattern.isObjSens()){
			SCGBuilder scgb = (SCGBuilder) cgb;
			final QueueReader<Pair<MethodOrMethodContext, Unit>> addedStaticSites = ((SCGBuilder) cgb).staticEdgeReader();
			Collection<Pair<MethodOrMethodContext, Unit>> staticSites = ((SCGBuilder) cgb).staticCallSitesLookUp(src);
			for (Pair<MethodOrMethodContext, Unit> site : staticSites)
				scgb.updateCallGraph(newP2Set, site);
			pta.build();
			
			while(addedSites.hasNext()||addedStaticSites.hasNext()){
				while(addedSites.hasNext()){
					final VirtualInvokeSite site = addedSites.next();
					final Var_Node receiver = site.recNode();
					cgb.updateCallGraph(receiver.getP2Set().getOldSet(), site);
					if(receiver==src)
						cgb.updateCallGraph(receiver.getP2Set().getNewSet(), site);
				}
				while(addedStaticSites.hasNext()){
					final Pair<MethodOrMethodContext, Unit> site = addedStaticSites.next();
					MethodOrMethodContext caller = site.getO1();
					GNode thisRef = pta.getMethodPAG(caller.method()).nodeFactory().caseThis();
					final Var_Node receiver = (Var_Node) ((ContextSensPTA)pta).parameterize(thisRef, caller.context());
					scgb.updateCallGraph(receiver.getP2Set().getOldSet(), site);
					if(receiver==src)
						scgb.updateCallGraph(receiver.getP2Set().getNewSet(), site);
				}
				pta.build();
			}
		}else{
			pta.build();
			
			while(addedSites.hasNext()){
				final VirtualInvokeSite site = addedSites.next();
				final Var_Node receiver = site.recNode();
				cgb.updateCallGraph(receiver.getP2Set().getOldSet(), site);
				if(receiver==src)
					cgb.updateCallGraph(receiver.getP2Set().getNewSet(), site);
				
				pta.build();
			}
		}
	}

	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 Var_Node oDotF = pta.parameterize(fld, a);
					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();
			srcfrn.getBase().getP2Set().getOldSet().forall(new PTSetVisitor() {
				public void visit(GNode n) {
					final Var_Node oDotF = pta.parameterize(fld, (Alloc_Node) n);
					pag.addSimpleEdge(oDotF, tgtv);
				}
			});
		}
	}
	
	protected boolean isConstant(Alloc_Node a) {
		if(a instanceof ContextAlloc_Node)
			a = ((ContextAlloc_Node) a).base();
		return a instanceof Constant_Node;
	}

	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();
			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 addAll(final Var_Node pointer, PTSetInternal other) {
		final PTSetInternal addTo = pointer.makeP2Set();
		// deprecated because not equal to forall->add
		// return addTo.addAll(other, null);
		return other.forall(new PTSetVisitor() {
			public final void visit(GNode n) {
				if (addTo.add(n))
					returnValue = true;
			}
		});
	}
	
	public static boolean is(Object o, String what) {
		String toString = o.toString();
		return switch (what) {
		case "InstantiationException": {
			yield toString.contains("java.lang.InstantiationException");
		}
		case "12 truncate": {
			yield multipleContains(toString, "$r12 <org.python.core.PyFile$RFileWrapper: void truncate(long)>");
		}
		case "getTargetException()": {
			yield multipleContains(toString, "Parm -2 to <java.lang.reflect.InvocationTargetException: java.lang.Throwable getTargetException()> <java.lang.reflect.InvocationTargetException: java.lang.Throwable getTargetException()>");
		}
		case "t1 getTargetException": {
			yield multipleContains(toString, " <java.lang.reflect.InvocationTargetException: java.lang.Throwable getTargetException()>");
		}
		case "?.target": {
			yield multipleContains(toString, "java.lang.reflect.InvocationTargetException in method null.<java.lang.reflect.InvocationTargetException: java.lang.Throwable target>");
		}
		case "(InstantiationException)exception": {
			yield multipleContains(toString, "(java.lang.InstantiationException)"," <sun.reflect.BootstrapConstructorAccessorImpl: java.lang.Object newInstance(java.lang.Object[])>");
		}
		case "EXCEPTION_NODE": {
			yield multipleContains(toString, "EXCEPTION_NODE");
		}
		case "newInstance": {
			yield multipleContains(toString, "<sun.reflect.BootstrapConstructorAccessorImpl: java.lang.Object newInstance(java.lang.Object[])>");
		}
		case "Constructor.newInstance": {
			yield multipleContains(toString, "<java.lang.reflect.Constructor: java.lang.Object newInstance(java.lang.Object[])>");
		}
		case "newInstanceGeneric": {
			yield multipleContains(toString, "<java.security.Provider$Service: java.lang.Object newInstanceGeneric(java.lang.Object)>");
		}
		case "Constructor newInstance call": {
			yield multipleContains(toString, "r6 = virtualinvoke r4.<java.lang.reflect.Constructor: java.lang.Object newInstance(java.lang.Object[])>($r17)");
		}
		case "receiverobj": {
			yield multipleContains(toString, "new java.lang.reflect.Constructor");
		}
		case "newarray[]": {
			yield multipleContains(toString, "newarray (java.lang.reflect.Constructor)[0] in method <java.lang.Class: java.lang.reflect.Constructor[] copyConstructors(java.lang.reflect.Constructor[])>.soot.jimple.spark.pag.ArrayElement");
		}
		case "Class.copyConstructors()": {
			yield multipleContains(toString, "<java.lang.Class: java.lang.reflect.Constructor[] copyConstructors(java.lang.reflect.Constructor[])>");
		}
		case "ReflectionFactory.copyConstructor()": {
			yield multipleContains(toString, "<sun.reflect.ReflectionFactory: java.lang.reflect.Constructor copyConstructor(java.lang.reflect.Constructor)>");
		}
		case "ReflectAccess.copyConstructor()": {
			yield multipleContains(toString, "<java.lang.reflect.ReflectAccess: java.lang.reflect.Constructor copyConstructor(java.lang.reflect.Constructor)>");
		}
		case "copy()": {
			yield multipleContains(toString, "<java.lang.reflect.Constructor: java.lang.reflect.Constructor copy()>");
		}
		case "copy call": {
			yield multipleContains(toString, "$r2 = virtualinvoke r1.<java.lang.reflect.Constructor: java.lang.reflect.Constructor copy()>()");
		}
		case "receiverobj2": {
			yield multipleContains(toString, "new java.lang.reflect.Constructor");
		}
		case "Entry[]": {
			yield multipleContains(toString,
				"newarray (java.util.Hashtable$Entry)[i2] in method <java.util.Hashtable: void <init>(int,float)>",
				"new java.util.Hashtable in method <org.eclipse.core.runtime.internal.adaptor.PluginConverterImpl: void init()>]).soot.jimple.spark.pag.ArrayElement@");
		}
		case "t7[]": {
			yield multipleContains(toString,
				" <java.util.Hashtable$Enumerator: java.lang.Object nextElement()>",
				"new java.util.Hashtable$Enumerator",
				".soot.jimple.spark.pag.ArrayElement@");
		}
		case "container_next": {
			yield multipleContains(toString,
					"Method <java.util.Hashtable$Enumerator: java.lang.Object next()> in context [AllocNode ", 
					"new java.util.Hashtable$Enumerator");
		}
		case "container_getBundleProperties": {
			yield multipleContains(toString,
					"Method <sun.jkernel.DownloadManager: java.util.Map getBundleProperties()>", 
					"new sun.jkernel.DownloadManager$4 in method <sun.jkernel.DownloadManager: java.lang.String getBootClassPathEntryForResource(java.lang.String)>", 
					"new sun.misc.Launcher$AppClassLoader in method <sun.misc.Launcher$AppClassLoader$1: java.lang.Object run()>]");
		}
		case "nextrec": {
			yield multipleContains(toString,
					"new java.util.Hashtable$Enumerator in method <java.util.Hashtable$EntrySet: java.util.Iterator iterator()>, ", 
					"new java.util.Hashtable$EntrySet in method <java.util.Hashtable: java.util.Set entrySet()>]");
		}
		case "nextrec2": {
			yield multipleContains(toString,
					"new java.util.Hashtable$Enumerator");
		}
		case "r1": {
			yield multipleContains(toString,
					"r1 <sun.jkernel.DownloadManager: java.util.Map getBundleProperties()>", 
					"new sun.jkernel.DownloadManager$4 in method <sun.jkernel.DownloadManager: java.lang.String getBootClassPathEntryForResource(java.lang.String)>", 
					"new sun.misc.Launcher$AppClassLoader in method <sun.misc.Launcher$AppClassLoader$1: java.lang.Object run()>]");
		}
		case "Collection.iterator()": {
			yield multipleContains(toString,
					"Parm -2 to <java.util.Collections$SynchronizedCollection: java.util.Iterator iterator()> <java.util.Collections$SynchronizedCollection: java.util.Iterator iterator()>",
					"new java.util.Collections$SynchronizedSet in method <java.util.Hashtable: java.util.Set entrySet()>, ",
					"new java.util.Properties in method <sun.jkernel.DownloadManager: java.util.Map getBundleProperties()>");
		}
		case "t0_iterator()": {
			yield multipleContains(toString,
					"<java.util.Collections$SynchronizedCollection: java.util.Iterator iterator()>",
					"new java.util.Collections$SynchronizedSet in method <java.util.Hashtable: java.util.Set entrySet()>",
					"new java.util.Properties in method <sun.jkernel.DownloadManager: java.util.Map getBundleProperties()>");
		}
		case "EntrySet.iterator()": {
			yield multipleContains(toString,
					"Parm -2 to <java.util.Hashtable$EntrySet: java.util.Iterator iterator()> <java.util.Hashtable$EntrySet: java.util.Iterator iterator()>",
					" new java.util.Hashtable$EntrySet in method <java.util.Hashtable: java.util.Set entrySet()>,",
					" new java.util.Properties in method <sun.jkernel.DownloadManager: java.util.Map getBundleProperties()>])");
		}
		case "t0_EntrySet.iterator()": {
			yield multipleContains(toString,
					"<java.util.Hashtable$EntrySet: java.util.Iterator iterator()>",
					"new java.util.Hashtable$EntrySet in method <java.util.Hashtable: java.util.Set entrySet()>",
					"new java.util.Properties in method <sun.jkernel.DownloadManager: java.util.Map getBundleProperties()>");
		}
		default:
			throw new IllegalArgumentException("Unexpected value: " + what);
		};
	}
	private static boolean multipleContains(String string, CharSequence... css) {
		for (CharSequence cs:css) {
			if(!string.contains(cs))
				return false;
		}
		return true;
	}
}
