/** Attempts to colour locals according to the Linear Scan algorithm.
 *
 * This locals colourer roughly implements 'Linear Scan Register Allocation'
 * by Poletto and Sarkar.
 */

package soot.toolkits.scalar;

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

public class LinearScanColorer extends LocalColorer {

private static int REGISTERS = 4;

	public static void assignColorsToLocals(Body b, Map localToGroup,
			Map localToColor, Map groupToColorCount){
		System.out.println("Linear Scan doing " + b.getMethod());
		// preamble, set things up
		UnitGraph g = new ExceptionalUnitGraph(b);
		SimpleLiveLocals liveness = new SimpleLiveLocals(g);

		Chain locals = b.getLocals();
		// bludgeon liveness into (coarser) ranges
		List order = new PseudoTopologicalOrderer().newList(g);
		final Map localToStart = new HashMap(locals.size()*2+1);
		final Map localToEnd = new HashMap(locals.size()*2+1);

		List colored = new ArrayList();
		// put some "sane" initial values for the ranges.
		Iterator it = locals.iterator();
		while(it.hasNext()){
			Object next = it.next();
			localToStart.put(next,new Integer(order.size()));
			localToEnd.put(next,new Integer(0));
		}
		for(int i = 0; i < order.size(); i++){
			Unit unit = (Unit) order.get(i);
			List livehere = liveness.getLiveLocalsBefore(unit);
			// for each local live here, update the intervals to include this unit
			it = livehere.iterator();
			while(it.hasNext()){
				Object l = it.next();
				int start,end;
				if(localToStart.containsKey(l)){
					start = ((Integer)localToStart.get(l)).intValue();
					if(i < start)
						localToStart.put(l,new Integer(i));
				}else
					localToStart.put(l,new Integer(i));
				if(localToEnd.containsKey(l)){
					end = ((Integer)localToEnd.get(l)).intValue();
					if(i > end)
						localToEnd.put(l,new Integer(i));
				}else
					localToEnd.put(l,new Integer(i));
			}
		}
		Iterator tit = locals.iterator();
		while(tit.hasNext()){
			Object l = tit.next();
			if(localToEnd.get(l) == null){
				System.out.println("No end in sight " + l);
				localToEnd.put(l,new Integer(order.size()));
			}
			if(localToStart.get(l) == null){
				System.out.println("Never beginning story" + l);
				localToStart.put(l,new Integer(order.size()));
			}
		}

		// for each group
		Set lcp = new HashSet(locals);
		while(!lcp.isEmpty()){
			int k = 0;
			// order locals by interval start
			TreeSet ord = new TreeSet( new Comparator(){
				public int compare(Object o1, Object o2){
					if(o1==o2) return 0;
					int ret;
					if( (ret = ((Integer)localToStart.get(o1)).compareTo(
								(Integer)localToStart.get(o2))) == 0)
						return o1.toString().compareTo(o2.toString());
					return ret;
				}
			});

			// add this group's locals to ord
			it = locals.iterator();
			Object curgroup = null;
			while(it.hasNext()){
				Object l = it.next();
				if(localToColor.containsKey(l)){
					colored.add(l);
					lcp.remove(l);
					continue; // already colored, a param or something..
				}
				if(!lcp.contains(l))
					continue; // we've already done this local/group
				Object group = localToGroup.get(l);
				if(curgroup != null && !curgroup.equals(group))
					continue; //local doesn't belong
				if(curgroup == null){
					curgroup = group; // just starting out
					k = ((Integer)groupToColorCount.get(group)).intValue();
				}
				ord.add(l);
				lcp.remove(l);
			}
			System.out.println("Ord: " + ord + "\nLeft: " + lcp + "\nLocals: " + locals);
			TreeSet active = new TreeSet( new Comparator(){
				public int compare(Object o1, Object o2){
					if(o1 == o2) return 0;
					int ret = ((Integer)localToEnd.get(o1)).compareTo(
								(Integer)localToEnd.get(o2));
					if(ret == 0)
						return o1.toString().compareTo(o2.toString());
					else
						return ret;
				}
			});

			List free = new ArrayList();
			for(int i = 0; i < k; i++)
				free.add(new Integer(i));

			// omgf!! 1337 L1n3ar 5can!1!!!!1!
			while(!ord.isEmpty()){
				Object l = ord.first();
				if(l == null)
					throw new RuntimeException("");
				System.out.println(ord);
				ord.remove(l);
				/*if(ord.remove(l))
					throw new RuntimeException("Removed element " + l +
							" was not in list " + ord);
				 */
				System.out.println("K is " + k + " and active is " + active.size() + " free is " + free.size());
				Expire_Old_Intervals(l,active,free,localToStart,localToEnd,localToColor);
				if(active.size() == k){
					k++; // Spill_At_Interval
					System.out.println("Spilt!");
					localToColor.put(l,new Integer(k));
					colored.add(l);
					active.add(l);
					groupToColorCount.put(curgroup,new Integer(k));
				}else{
					active.add(l);
					localToColor.put(l,free.remove(0));
					colored.add(l);
					System.out.println("Free!");
				}
			}
		}
		System.out.println("Colored : " + colored + "\nlocal : " + locals);
	}

	private static void Expire_Old_Intervals(
				Object local, Set active, List free,
				Map localToStart, Map localToEnd, Map localToColor){
		int start = ((Integer)localToStart.get(local)).intValue();
		List toRemove = new ArrayList();
		Iterator it = active.iterator();
		while(it.hasNext()){
			Object l = it.next();
			int end = ((Integer)localToEnd.get(l)).intValue();
			if(end >= start) return;
			toRemove.add(l);
		}
		it = toRemove.iterator();
		while(it.hasNext()){
			Object l = it.next();
			active.remove(l);
			System.out.println("Freeing!");
			free.add(localToColor.get(l));
		}
	}

	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!");
	}
}