import soot.*;
import soot.tagkit.*;
import soot.jimple.*;
import soot.toolkits.scalar.*;
import soot.toolkits.graph.BriefUnitGraph;
import java.util.*;

public class AssignStmtDriver{

	public static void main(String[] args){
		Pack jtp = soot.PackManager.v().getPack("jtp");
		jtp.add(new Transform("jtp.ad", new ExampleTransformer()));

		soot.Main.main(args);
	}
}

class ExampleTransformer extends BodyTransformer{

	protected void internalTransform(Body b, String phase, Map opts){
		// This is where you would create your new flow analysis instance
		//FlowAnalysisClass analysis = new FlowAnalysisClass(new BriefUnitGraph(b));

		Iterator units = b.getUnits().iterator();
		Iterator keysIt = b.getMethod().getDeclaringClass().getTags().iterator();
		boolean keysAdded = false;
		/* check if color key hasn't been added before */
		while (keysIt.hasNext()){
			Object next = keysIt.next();
			if (next instanceof KeyTag){
				if (((KeyTag)next).analysisType().equals("MyUnitTags")){
					keysAdded = true;
					break;
				}
			}
		}
		if (!keysAdded){
			b.getMethod().getDeclaringClass().addTag(new KeyTag(ColorTag.BLUE,
				"Assign Stmt", "MyUnitTags"));
		}

		while(units.hasNext()){
			Unit u = (Unit) units.next();

			if (u instanceof AssignStmt)
				u.addTag(new ColorTag(ColorTag.BLUE,
					"Assign Stmt"));
		}
	}
}