/** Attempts to colour locals according to Briggs' heuristic.
 *
 * This locals colourer roughly implements 'Coloring Heuristics for
 * Register Allocation' by Briggs, Cooper, Kennedy and Torczon.
 */

package soot.toolkits.scalar;

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

public class BriggsColorer extends LocalColorer {
private static int REGISTERS = 4;

	public static void assignColorsToLocals(Body b, Map localToGroup,
			Map localToColor, Map groupToColorCount){
		System.out.println("Briggs colorer doing " + b.getMethod());
		// preamble, set things up
		int k = REGISTERS;
		UnitGraph g = new ExceptionalUnitGraph(b);
		SimpleLiveLocals liveness = new SimpleLiveLocals(g);
		InterferenceGraph ig = new InterferenceGraph(b,localToGroup,liveness);
		NestDepth nd = new NestDepth(g);

		List locals = ig.getLocals();
		Map localToNeighb = mapLocalsToNeighbours(locals,ig);
		List[] N = orderByInterference(locals,localToNeighb);

		Stack toColor = new Stack();
		Stack toSpill = new Stack();
		// until we know what the packing is
		while(true){
			int i = 0;
			// colour the nodes in order.
			while(i < N.length){
				if(N[i]==null || N[i].isEmpty()){
					i++;
					continue;
				}
				Object l = N[i].get(0);
				if(i < k){
					toColor.push(l);
					N[i].remove(l);
				}else{
					// spill a node, it's got more than k uncoloured neighbours.
					// the best node to spill is the one with the smallest spill cost.
					float bestcost = Float.POSITIVE_INFINITY;
					Object bestfit = null;
					List bestbox = null;
					for(int j = i; j < N.length; j++){
						if(N[j] == null || N[j].isEmpty())
							continue;
						Iterator it = N[j].iterator();
						while(it.hasNext()){
							Local cur = (Local)it.next();
							float cost = (float) nd.weightOfLocal(cur) / j;
							if(cost < bestcost){
								bestcost = cost;
								bestfit = cur;
								bestbox = N[j];
							}
						}
					}
					if(bestfit == null)
						throw new RuntimeException(
										"Eehm, expected to spill but found no node at " + i);
					toSpill.push(bestfit);
					bestbox.remove(bestfit);
					l = (Object) bestfit;
				}
				updateNeighbours(N,(Local)l,localToNeighb,ig);
				if(i != 0)
					i--; // start from cur_pos-1 next time
			}
			// now, we must figure out what to do with spilt nodes.
			if(!toSpill.isEmpty()){
				/* increase register size and rerun the algo on just the spilt nodes */
				// localToNeighb = mapLocalsToNeighbours(toSpill,ig);
				N = orderByInterference(toSpill,localToNeighb);
				toSpill = new Stack();
				// find the lowest possible degree now.
				for(k = 0;k < N.length && N[k] == null;k++);
				k++;
			}else break;
		}
		// assign non-interfering colours to the nodes in stack order.
		while(!toColor.isEmpty()){
			Object l = toColor.pop();
			Object group = localToGroup.get(l);
			int used = ((Integer)groupToColorCount.get(group)).intValue();
			int free[] = new int[used+1];
			Arrays.fill(free,1);
			Local neighb[] = ig.getInterferencesOf((Local)l);
			for(int x = 0; x < neighb.length;x++){
				if(localToColor.containsKey(neighb[x])){
					int color = ((Integer)localToColor.get(neighb[x])).intValue();
					free[color] = 0;
				}
			}
			int color;
			for(color = 0;color<free.length && free[color] == 0;color++);
			localToColor.put(l,new Integer(color));
			if(color >= used)
				groupToColorCount.put(group,new Integer(color));
		}
	}

	private static Map mapLocalsToNeighbours(List locals,InterferenceGraph ig){
		Map localToNeighb = new HashMap(locals.size()*2 + 1);
		Iterator it = locals.iterator();
		while(it.hasNext()){
			Local l = (Local) it.next();
			Local[] intfs = ig.getInterferencesOf(l);
			localToNeighb.put(l,new ArrayList(Arrays.asList(intfs)));
		}
	return localToNeighb;
	}

	/** Order locals by interference size.
	 * if N is an array of lists, N[i] is a list of locals l
	 * with deg(l) == i. */
	private static List[] orderByInterference(List locals, Map localToNeighb){
		// we don't care about groups yet, just throw the whole bunch in

		// find the max degree
		int maxinterf = 0;
		Iterator it = localToNeighb.values().iterator();
		while(it.hasNext()){
			int size = ((List)it.next()).size();
			maxinterf = (maxinterf > size) ? maxinterf : size;
		}

		List N[] = new List[maxinterf+1];
		it = locals.iterator();
		while(it.hasNext()){
			Object l = it.next();
			int i = ((List)localToNeighb.get(l)).size();
			if(N[i] == null)
				N[i] = new ArrayList();
			if(N[i].contains(l))
				throw new RuntimeException("ordering locals: encountered again " + l
									+ " already in " + Arrays.asList(N));
			else
			N[i].add(l);
		}
		return N;
	}

	/* jack all neighbours of l down a notch */
	private static void updateNeighbours(List N[], Local l,Map localToNeighb,
												InterferenceGraph ig){
		// update the position of neighbours in the list
		//Local[] neighb= ig.getInterferencesOf(l);
		List neighb = (List)localToNeighb.get(l);
		for(int j = 0; j < neighb.size(); j++){
			Object n = neighb.get(j);
			List nlist = (List)localToNeighb.get(n);
			int deg = nlist.size();
			nlist.remove(l);
			if(N[deg] != null && N[deg].contains(n)){
				N[deg].remove(n);
				if(N[--deg] == null)
					N[deg] = new ArrayList();
				if(N[deg].contains(n))
					throw new RuntimeException(Arrays.asList(N) + " already has " + n);
				N[deg].add(n);
			}
		}
	}

	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 an interference graph.
	* Should probably be put in its own class at this point */
	public static class InterferenceGraph{
		Map localToRange;// Maps a local to its live units.
		Map localToLocals; // Maps a local to locals simultaneously live
		List locals;

		private InterferenceGraph(){
		}

		public List getLocals(){
			return locals;
		}

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

			// Initialize localToLocals
			localToLocals = new HashMap(body.getLocalCount() * 2 + 1,0.7f);
			Iterator localIt = body.getLocals().iterator();
			while(localIt.hasNext()){
				Object local = localIt.next();
				localToLocals.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);

				// 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;
		}
	}
}