/** Attempts to colour locals according to their priorities.
 *
 * This colouring algorithm is similar to the one described by Chow and
 * Hennessey in the paper "Register Allocationn by Priority-based Coloring".
 *
 * Priorities are assigned according to the number of uses of a variable,
 * weighted according to their nest depths.
 *
 * The whole spilling business doesn't really apply for the JVM, but we're
 * looking for efficient packings here, with lower numbers assigned to
 * 'livelier' locals. We're aiming at a 4-coloring here, to hopefully use
 * the cheaper i{load,store}_x instructions more often.
 */

package soot.toolkits.scalar;

import soot.*;
import soot.util.*;
import soot.options.*;
import java.util.*;
import soot.toolkits.graph.*;

public class PriorityColorer extends LocalColorer {

private static int REGISTERS = 4;

	static Map localToPriority;
	public static void assignColorsToLocals(Body b, Map localToGroup,
			Map localToColor, Map groupToColorCount){
		Iterator lit; // temporary iterator
		int localslim = REGISTERS;
		if(Options.v().verbose()){
			G.v().out.println("Assigning colors to locals using priority method..");
		}
		UnitGraph g = new ExceptionalUnitGraph(b);
		SimpleLiveLocals liveness = new SimpleLiveLocals(g);
		LiveRangeInterference ig =
				new LiveRangeInterference(b,localToGroup,liveness);
		NestDepth nd = new NestDepth(b);
		List locals = ig.getLocals();
		if(Options.v().verbose()){
			G.v().out.println("Locals: " + locals);
		}
		Set unconstrained = new HashSet();
		Set constrained = new HashSet();
		Set coloured = new HashSet();
		localToPriority = new HashMap(locals.size()*2 +1);

		// 1. Set aside locals with < REGISTERS neighbours ("unconstrained" nodes)
		constrain_nodes(
				locals,localslim,localToColor,constrained,unconstrained,coloured,ig);
		if(Options.v().verbose()){
			G.v().out.println("Unconstrained: " + unconstrained +
				" Constrained: " + constrained + " Coloured :" + coloured);
		}
		TreeSet ord = new TreeSet( new Comparator() {
			public int compare(Object o1, Object o2){
				//if(((Local)o1).equivTo((Local)o2) )
				if(o1 == o2)
					return 0;
				Float f1 = (Float) localToPriority.get(o1);
				Float f2 = (Float) localToPriority.get(o2);
				int ret;
				// if the two priorities are equal, order alphabetically
				if( (ret = f1.compareTo(f2)) == 0)
					return o1.toString().compareTo(o2.toString());
				else
					return ret;
			}
		});
		// 2. until all constrained are assigned or no registers left
		Iterator i = constrained.iterator();
		while(i.hasNext()){
			Local l = (Local)i.next();
			ArraySet iferences = new ArraySet(ig.getInterferencesOf(l));
			int neighbours = 0;
			// find the coloured neighbours
			lit = coloured.iterator();
			while(lit.hasNext()){
				Object neigh = lit.next();
				if(iferences.contains(neigh))
					neighbours++;
			}
			Object group = localToGroup.get(l);
			int assigned = ((Integer)groupToColorCount.get(group)).intValue();
			//if(neighbours >= localslim - assigned){ // FIXME: could bug?
			if(neighbours >= localslim){
				// split case in original algo.. we just increase the locals limit
				localslim++;
				// check if some nodes are now unconstrained
				constrain_nodes(
						locals,localslim,localToColor,constrained,unconstrained,coloured,ig);
				if(Options.v().verbose()){
					G.v().out.println("Unconstrained: " + unconstrained +
						" Constrained: " + constrained + " Coloured :" + coloured);
				}
				// and do it again
				i = constrained.iterator();
			}else{
				// Mike suggests we toString the Integer(int), then use
				// Float.parseFloat to convert it to a Float, then get floatValue().
				// I'm tempted.
				float priority = nd.weightOfLocal(l);
				priority /= ig.getRangeOf(l).size();
				localToPriority.put(l,new Float(priority));
				//System.out.println(i + " adding " + l + " with priority " + priority);
				ord.add(l);
				// this concludes computing the ADJSAVE (=priority)
			}
		}
		// Colour constrained nodes according to priority

		colorByPriority(ord,localToColor,localToGroup,groupToColorCount,coloured,
					ig,localslim);
		ord = new TreeSet(new Comparator() {
			public int compare(Object o1, Object o2){
				return o1.toString().compareTo(o2.toString());
			}
		});
		ord.addAll(unconstrained);
		colorByPriority(ord,localToColor,localToGroup,groupToColorCount,coloured,
					ig,localslim);

		if(locals.size() != coloured.size()){
			if(Options.v().verbose()){
				G.v().out.println("Coloured :" + coloured);
			}
			Iterator err = locals.iterator();
			while(err.hasNext()){
				Object errl = err.next();
				if(!coloured.contains(errl))
					System.out.println("Left uncoloured: " + errl);
			}
			throw new
				RuntimeException("We can't be done yet; we have locals yet uncolored!");
		}
	}

	private static void colorByPriority(
				SortedSet ord,Map localToColor,Map localToGroup, Map groupToColorCount,
				Set coloured,LiveRangeInterference ig,int localslim){
		if(Options.v().verbose()){
			G.v().out.println("Colouring: " + ord);
		}
		while(!ord.isEmpty()){
			Local l = (Local)ord.last();
			Object group = localToGroup.get(l);
			int usedcolours = ((Integer)groupToColorCount.get(group)).intValue();
			int freecolours[] = new int[usedcolours+1];
			Arrays.fill(freecolours,1);
			if(!ord.remove(l))
				throw new RuntimeException("Error removing elements from ordered list of locals.");
			Local interference[] = ig.getInterferencesOf(l);
			if(Options.v().verbose()){
				G.v().out.println("Interferences of " + l + ": " + Arrays.asList(interference));
			}
			for(int x = 0; x< interference.length; x++){
				if(localToColor.containsKey(interference[x])){
					int colour = ((Integer)localToColor.get(interference[x])).intValue();
					if(Options.v().verbose()){
						G.v().out.println("Local " + interference[x] + " has colour " + colour);
					}
					freecolours[colour] = 0;
				}
			}
			// TODO: could do better and assign unconstrained by priority
			// but it doesn't matter, as we've either got >4 registers or <4 locals.
			int colour = 0;
			while(colour < usedcolours+1){
				if(freecolours[colour] == 1)
					break;
				colour++;
			}
			localToColor.put(l,new Integer(colour));

			if(colour > usedcolours){
				groupToColorCount.put(group,new Integer(colour));
			}
			coloured.add(l);
		}
	}

	private static void constrain_nodes(List locals,int localslim,Map localToColor,Set constrained,Set unconstrained,Set coloured, LiveRangeInterference ig){
		Iterator lit = locals.iterator();
		while(lit.hasNext()){
			Local l = (Local)lit.next();
			if(localToColor.containsKey(l)){
				int colour = ((Integer)localToColor.get(l)).intValue();
				coloured.add(l);
				continue;
			}
			ArraySet iferences = new ArraySet(ig.getInterferencesOf(l));
			if(iferences.size() < localslim)
				unconstrained.add(l);
			else
				constrained.add(l);
			// also init priorities.. unless we don't
			// localToPriority.put(l,new Float(0.0));
		}
	}

	public static void unsplitAssignColorsToLocals(Body b,
												Map localToGroup,
												Map localToColor,
												Map groupToColorCount){
		/* the Itchy and Scratchy show! */
		//assignColorsToLocals(b,localToGroup,localToColor,groupToColorCount);
		throw new
			RuntimeException("Unsplit priority-based coloring is not implemented!");
	}

	/** Implementation of local live ranges. */
	public static class LiveRangeInterference{
		Map localToRange;// Maps a local to its live units.
		Map localToLocals; // Maps a local to locals simultaneously live
		List locals;

		private LiveRangeInterference(){
		}

		public List getLocals(){
			return locals;
		}

		public LiveRangeInterference(Body body, Map localToGroup,
									 LiveLocals liveLocals){
			locals = new ArrayList();
			locals.addAll(body.getLocals());

			// Initialize localToLocals & localToRange
			localToLocals = new HashMap(body.getLocalCount() * 2 + 1,0.7f);
			localToRange = new HashMap(body.getLocalCount() * 2 + 1,0.7f);
			Iterator localIt = body.getLocals().iterator();
			while(localIt.hasNext()){
				Local local = (Local) localIt.next();
				localToLocals.put(local, new ArraySet());
				localToRange.put(local, new ArraySet());
			}
			// Go through code, noting interferences
			Iterator codeIt = body.getUnits().iterator();
			while(codeIt.hasNext()){
				Unit unit = (Unit) codeIt.next();
				List liveLocalsAtUnit =
					liveLocals.getLiveLocalsAfter(unit);
				Iterator liveit = liveLocalsAtUnit.iterator();
				// add this unit to the local's live range
				while(liveit.hasNext()){
					ArraySet range = (ArraySet)localToRange.get((Local)liveit.next());
					range.add(unit);
				}
				// Note interferences if this stmt is a definition
				List defBoxes = unit.getDefBoxes();
				if(!defBoxes.isEmpty()) {
					if(!(defBoxes.size()==1))
						throw new RuntimeException
							("invalid number of def boxes");
					if(((ValueBox)defBoxes.get(0)).getValue()
						 instanceof Local){
						Local defLocal = (Local) ((ValueBox)defBoxes.get(0)).getValue();
						localIt = liveLocalsAtUnit.iterator();
						while(localIt.hasNext()){
							Local otherLocal = (Local) localIt.next();
							if(localToGroup.get(otherLocal)
									.equals(localToGroup.get(defLocal)))
								setInterference(defLocal, otherLocal);
						}
					}
				}
			}
		}

		public boolean localsInterfere(Local l1, Local l2){
			return ((Set) localToLocals.get(l1)).contains(l2);
		}

		public void setInterference(Local l1, Local l2){
			if(l1==l2) return;
			((Set) localToLocals.get(l1)).add(l2);
			((Set) localToLocals.get(l2)).add(l1);
		}

		public Set getRangeOf(Local l){
			return (Set)localToRange.get(l);
		}

		public boolean isEmpty(){
			return localToLocals.isEmpty();
		}

		Local[] getInterferencesOf(Local l){
			Object[] objects = ((Set) localToLocals.get(l)).toArray();
			Local[] locals = new Local[objects.length];

			for(int i = 0; i < objects.length; i++)
				locals[i] = (Local) objects[i];
			return locals;
		}
	}
}