package s3.services.bootimage;

import ovm.core.domain.Blueprint;
import ovm.core.domain.Code;
import ovm.core.domain.ConstantResolvedFieldref;
import ovm.core.domain.ConstantResolvedMethodref;
import ovm.core.domain.Domain;
import ovm.core.domain.DomainDirectory;
import ovm.core.domain.Field;
import ovm.core.domain.Method;
import ovm.core.domain.Oop;
import ovm.core.domain.Type;
import ovm.core.domain.UserDomain;
import ovm.core.repository.ConstantClass;
import ovm.core.repository.ConstantFieldref;
import ovm.core.repository.ConstantMethodref;
import ovm.core.repository.JavaNames;
import ovm.core.repository.UnboundSelector;
import ovm.core.services.memory.VM_Address;
import ovm.services.bytecode.Instruction;
import ovm.util.ArrayList;
import ovm.util.BitSet;
import ovm.util.HashMap;
import ovm.util.HashSet;
import ovm.util.Iterator;
import ovm.util.S3Set;
import s3.core.domain.S3Blueprint;
import s3.core.domain.S3Domain;

public class RapidTypeAnalysis extends Analysis {
    // set to trace call graph walk
    String indent = "";

    // set of all classes that are ever instantiated
    final BitSet[]  liveClasses;

    // set of all classes that may be heap allocated, a subset of
    // liveClasses
    final BitSet[] heapClasses;
    
    // the set of all selectors for virtual call sites
    final HashSet   calledVtSelectors;

    // the set of all virtual methods in live classes
    final BitSet[]  liveVtMethods;

    // maps a selector to a list of virtual method in live classes
    final HashMap   selector2liveVtMethods;

    // Remember when <clinit> has already been checked.
    final BitSet[]  initializedClasses;

    static class MethodNode {
	Method method;
	MethodNode next;
	MethodNode(Method method, MethodNode next) {
	    this.method = method;
	    this.next = next;
	}
    }

    RapidTypeAnalysis(Domain d, IRewriter rewriter, NativeCallGenerator ncg) {
	super(d, rewriter, ncg);
	int nctx = DomainDirectory.maxContextID() + 1;
	liveClasses = new BitSet[nctx];
	heapClasses = new BitSet[nctx];
	liveVtMethods = new BitSet[nctx];
	initializedClasses = clinitIsLive ? new BitSet[nctx] : null;

	calledVtSelectors = new HashSet();
	selector2liveVtMethods = new HashMap();

	for (int i = 0; i < nctx; i++) {
	    Type.Context tc = DomainDirectory.getContext(i);
	    if (tc != null && tc.getDomain() == domain) {
		liveClasses[i] = new BitSet(tc.getBlueprintCount());
		heapClasses[i] = new BitSet(tc.getBlueprintCount());
		liveVtMethods[i] = new BitSet(tc.getMethodCount());
		if (clinitIsLive)
		    initializedClasses[i] = new BitSet(tc.getBlueprintCount());
	    }
	}
    }

    public void addConcreteType(Blueprint bp) {
	// Arrays don't create any more work, but we need to remember
	// them for checkcast elimination
	int cid = bp.getCID();
	int bpid = bp.getUID();
	if (bp.getType().isArray())
	    liveClasses[cid].set(bpid);
	if (shouldCompile(bp) &&!liveClasses[cid].get(bpid)) {
	    Type.Compound t = bp.getType().asCompound();
	    if (TRACE)
		System.err.println(indent + t.getUnrefinedName()
				   + " is concrete");
	    liveClasses[cid].set(bpid);
	    String saveIndent = indent; 
	    indent = "    ";
	    Code[] vt = ((S3Blueprint) bp).getVTable();
	    for (int i = 0; i < vt.length; i++) {
		Method m = vt[i].getMethod();
		if (m.isNonVirtual())
		    continue;
		if (shouldAnalyze(m)) {
		    UnboundSelector ubs
			= m.getSelector().getUnboundSelector();
		    if (liveVtMethods[m.getCID()].get(m.getUID())) {
			// method is inherited from a live class,
			// nothing to do
			continue;
		    }
		    else if (calledVtSelectors.contains(ubs)) {
			// If we have already seen a call, add the
			// virtual method to the work stack

			// FIXME: to build a stack trace for call and
			// for allocation TracingWorkList needs to
			// know why isVirtual is tree, and to see
			// every allocation point
			processMethod(m, true);
		    }
		    else  {
			// Otherwise, remember to add it if we do
			// encounter a call
			liveVtMethods[m.getCID()].set(m.getUID());
			MethodNode next
			    = (MethodNode) selector2liveVtMethods.get(ubs);
			selector2liveVtMethods.put(ubs,
						   new MethodNode(m, next));
		    }
		}
	    }
	    indent = saveIndent;
	}

	if (clinitIsLive)
	    addClinit(bp.getSharedState().getBlueprint());
    }

    public void addHeapType(Blueprint bp) {
	int cid = bp.getCID();
	int uid = bp.getUID();
	if (!heapClasses[cid].get(uid)) {
	    addConcreteType(bp);
	    if (liveClasses[cid].get(uid))
		heapClasses[cid].set(uid);
	}
    }
    
    private void addClinit(Blueprint shSt) {
	if (!initializedClasses[shSt.getCID()].get(shSt.getUID())) {
	    initializedClasses[shSt.getCID()].set(shSt.getUID());
	    Method clinit = getClinit(shSt);
	    if (clinit != null)
		addNonvirtualCall(clinit);

	    Blueprint ibp = shSt.getInstanceBlueprint();
	    if (ibp != ((S3Domain) domain).ROOT_BLUEPRINT) {
		try {
		    Type.Class sup = ibp.getType().getSuperclass();
		    Oop supShSt = sup.getSharedStateType().getSingleton();
		    addClinit(supShSt.getBlueprint());
		} catch (NullPointerException npe) {
		    System.err.println("Something is wrong with the blueprint " + ibp);
		}
						    
	    }
	}
    }

    public void addVirtualCall(Method method) {
	UnboundSelector ubs = method.getSelector().getUnboundSelector();
	if (!calledVtSelectors.contains(ubs)) {
	    if (TRACE)
		System.err.println(indent + "virtual call to    " + ubs);
	    calledVtSelectors.add(ubs);
	    String saveIndent = indent;
	    indent = indent == "" ? "  " : "    ";
	    for (MethodNode n = (MethodNode) selector2liveVtMethods.get(ubs);
		 n != null; n = n.next)
		processMethod(n.method, true);
	    // this entry no longer needed
	    selector2liveVtMethods.put(ubs, null);
	    indent = saveIndent;
	}
    }

    public void addAllocation(Type t) {
	addHeapType(blueprintFor(t));
    }

    public void addNonvirtualCall(Method method) {
	processMethod(method, false);
	if (clinitIsLive
	    && method.getSelector().getUnboundSelector() != JavaNames.CLINIT
	    && method.getDeclaringType().getDomain() == domain)
	{
	    Blueprint bp = blueprintFor(method.getDeclaringType());
	    addClinit(bp.isSharedState() ? bp
		      : (Blueprint) bp.getSharedState().getBlueprint());
	}
    }

    public void addFieldRead(Field f) {
	Type t = f.getDeclaringType();
	if (clinitIsLive && t.isSharedState()) {
	    addClinit(blueprintFor(t));
	}
    }
    public void addFieldWrite(Field f) {
	Type t = f.getDeclaringType();
	if (clinitIsLive && t.isSharedState()) {
	    addClinit(blueprintFor(t));
	}
    }
    public void addArrayAccess(Type.Array _) { }

    private void processMethod(Method method, boolean isVirtual) {
	int cnum = method.getCID();
	int num = method.getUID();
	if (!calledMethods[cnum].get(num))
	    if (shouldAnalyze(method)) {
		if (TRACE)
		    System.err.println(indent +
				       (isVirtual
					? "push "
					: "nonvirtual call to ")
				       + method.getSelector());
		addMethod(method);
	    }
	    else
		System.err.println("skip bad method " + method);
    }

    private final Visitor visitor = new Visitor() {
	    public void visit(Instruction.Invocation i) {
		ConstantMethodref mr = i.getConstantMethodref(buf, cp);
		try {
		    Method m = ((ConstantResolvedMethodref) mr).getMethod();
		    if (m.getDeclaringType().getDomain() != domain)
			return;
		    if (m.isNonVirtual() 
			|| m.getMode().isFinal()
			|| m.getDeclaringType().getMode().isFinal()
			|| i instanceof Instruction.INVOKESPECIAL)
			addNonvirtualCall(m);
		    else
			addVirtualCall(m);
		} catch (ClassCastException e) {
		    warnMissing(i.getSelector(buf, cp).toString());
		}
	    }

	    // This catches both reads and writes, which is actually what
	    // we want.
	    public void visit(Instruction.FieldAccess i) {
		ConstantFieldref r = i.getConstantFieldref(buf, cp);
		try {
		    Field f = ((ConstantResolvedFieldref) r).getField();
		    addFieldRead(f); // read/write doesn't matter
		} catch (ClassCastException e) {
		    warnMissing(i.getSelector(buf, cp).toString());
		}
	    }

	    // Force resolution of instance fields, or move that into
	    // IRewriter?
	    public void visit(Instruction.Allocation i) {
		ConstantClass cl = i.getResultType(buf, cp);
		try {
		    addHeapType((Blueprint) cl);
		} catch (ClassCastException e) {
		    warnMissing("in " + i + ": " + cl.toString());
		    //e.printStackTrace(System.err);
		}
	    }
	};

    public void analyzeMethod(Method m) {
	super.analyzeMethod(m);
	if (TRACE)
	    System.err.println("anlayzing " + m.getSelector());
	indent = "  ";
	visitor.run(m);
	indent = "";
    }

    public boolean analyzePartialHeap(S3Set image) {
	image.forAll(new S3Set.Visitor() {
	    public void visit(Object o) {
		    VM_Address addr = VM_Address.fromObject(o);
		    if (addr.isNull())
			return;
		    Oop oop = addr.asOop();
		    try {
			addConcreteType(oop.getBlueprint());
		    } catch (ClassCastException _) {
			//System.err.println("bad object " + o);
		    }
		}
	    });
	return analyze();
    }

    // Code that uses analysis results to answer questions such as
    // "can this call be devirtualized," and "does this subtype test
    // suceed."  At one point, there was a seperate class called
    // RapidTypeAnalysis.Result that answered these questions.

    // A refrence to a possibly null Method object
    static class Devirt {
	final Method method;
	Devirt(Method method) { this.method = method; }
    }
    static final Devirt VIRTUAL = new Devirt(null);

    Blueprint[]   liveClassList;
    BitSet[]      liveVTables;
    BitSet[]      liveIFTables;

    int maxIFTable;
	
    Blueprint[][][] liveSubtypeCache;
    byte[][][][]      safeCastCache;
    Devirt[][][]      cacheDevirtV;
    Devirt[][][]      cacheDevirtI;
	
    // Maintain counters to monitor caching and accuracy
    int subtypeCacheHits;

    int devirtCalls;
    int devirtSuccess;

    int castChecks;
    int castSuccess;
    int castFailure;

    int nclasses;
    int nmethods;
	
    int sumCardinality(BitSet[] bs) {
	int ret = 0;
	for (int i = 0; i < bs.length; i++)
	    if (bs[i] != null)
		ret += bs[i].cardinality();
	return ret;
    }

    public void analysisComplete() {

	// We don't really need to check live classes by number,
	// but we do need to iterate over them.  Copy all live
	// blueprints into an array

	liveClassList = new Blueprint[sumCardinality(liveClasses)];
	int idx = 0;
	for (Iterator it = domain.getBlueprintIterator();
	     it.hasNext(); ) {
	    Blueprint bp = (Blueprint) it.next();
	    if (liveClasses[bp.getCID()].get(bp.getUID())) {
		liveClassList[idx++] = bp;
		Code[] tab = ((S3Blueprint) bp).getIFTable();
		if (tab != null && tab.length > maxIFTable)
		    maxIFTable = tab.length;
	    }
	}


	int nctx = DomainDirectory.maxContextID() + 1;
	liveVTables = new BitSet[nctx];
	liveIFTables = new BitSet[nctx];
	// Hmm.  Would it make sense to create a linked list of
	// each blueprint's direct subclasses and classes that
	// directly or indirectly implement an interface?
	// (Somehow I find the thought of listing an interface's
	// subclasses a little scary.  It seems like there may be
	// some way to generate cycles.)
	liveSubtypeCache = new Blueprint[nctx][][];

	// Try to speed up evalCast and devirtualize by memoizing
	// results.  I have a feeling that these caches won't
	// suffice.
	safeCastCache = new byte[nctx][][][];
	cacheDevirtV = new Devirt[nctx][][];
	cacheDevirtI = new Devirt[nctx][][];

	for (int i = 0; i < nctx; i++) {
	    Type.Context tc = DomainDirectory.getContext(i);
	    if (tc != null && tc.getDomain() == domain) {
		int nbp = tc.getBlueprintCount();
		nclasses += nbp;
		nmethods += tc.getMethodCount();
		liveVTables[i] = new BitSet(nbp);
		liveIFTables[i] = new BitSet(nbp);
		liveSubtypeCache[i] = new Blueprint[nbp][];
		safeCastCache[i] = new byte[nbp][][];
		cacheDevirtV[i] = new Devirt[nbp][];
		cacheDevirtI[i] = new Devirt[nbp][];
	    }
	}
    }


    public void printStats() {
	System.err.println("\nRapid Type Analysis for " + domain + ":");
	System.err.println("\t" + liveClassList.length + " of " + nclasses + " blueprints are concrete");
	System.err.println("\t" + sumCardinality(calledMethods) + " of " + nmethods + " methods are called");
	System.err.println("\t" + devirtSuccess + " of " + devirtCalls + " calls devirtualized");
	System.err.println("\t" + castSuccess + " of " + castChecks + " casts known to succeed");
	System.err.println("\t" + castFailure + " of " + castChecks + " casts known to fail");
	System.err.println("\t" + sumCardinality(liveVTables) + " virtual tables used");
	System.err.println("\t" + sumCardinality(liveIFTables) + " interface tables used");
	// ... still need to print cache stats
    }

    ArrayList subtypeList = new ArrayList();

    public Blueprint[] concreteSubtypesOf(Blueprint bp) {
	int cid = bp.getCID();
	int uid = bp.getUID();
	Blueprint[] ret = liveSubtypeCache[cid][uid];
	if (ret == null) {
	    if (bp == ((S3Domain) domain).ROOT_BLUEPRINT)
		ret = liveSubtypeCache[cid][uid] = liveClassList;
	    else {
		for (int i = 0; i < liveClassList.length; i++)
		    if (liveClassList[i].isSubtypeOf(bp))
			subtypeList.add(liveClassList[i]);
		ret = new Blueprint[subtypeList.size()];
		subtypeList.toArray(ret);
		liveSubtypeCache[cid][uid] = ret;
		subtypeList.clear();
	    }
	}
	subtypeCacheHits++;
	return ret;
    }

    public boolean isBlueprintConcrete(Blueprint bp) {
	return liveClasses[bp.getCID()].get(bp.getUID());
    }

    public boolean isHeapAllocated(Blueprint bp) {
	Blueprint[] sub = concreteSubtypesOf(bp);
	for (int i = 0; i < sub.length; i++)
	    if (heapClasses[bp.getCID()].get(sub[i].getUID()))
		return true;
	return false;
    }

    public Method devirtualize(Blueprint receiver, Method m) {
	int rc = receiver.getCID();
	int rNum = receiver.getUID();
	Devirt[] dv;
	BitSet[] liveTables;
	int index;
	String table;
	
	if (m.isVirtual()) {
	    dv = cacheDevirtV[rc][rNum];
	    liveTables = liveVTables;
	    if (dv == null) {
		dv = new Devirt[((S3Blueprint) receiver).getVTable().length];
		cacheDevirtV[rc][rNum] = dv;
	    }
	    index = receiver.getVirtualMethodOffset(m);
	    table = "vTable";
	}
	else if (m.isInterface()) {
	    dv = cacheDevirtI[rc][rNum];
	    liveTables = liveIFTables;
	    if (dv == null) {
		dv = new Devirt[receiver.getType().isInterface()
				? maxIFTable
				: ((S3Blueprint) receiver).getIFTable().length];
		cacheDevirtI[rc][rNum] = dv;
	    }
	    index = receiver.getInterfaceMethodOffset(m);
	    table = "ifTable";
	}
	else
	    throw new RuntimeException("bad table name");

	if (index >= dv.length)
	    /* This indicates dead code.  An attempt is being made
	     * to call through an interface that is never
	     * instantiated.
	     */
	    return null;

	Devirt result = dv[index];
	if (result == null) {
	    Blueprint[] bp = concreteSubtypesOf(receiver);
	    Method candidate = null;
	    boolean valid = true;
	    for (int i = 0; i < bp.length; i++) {
		Method c2 = ((S3Blueprint) bp[i]).lookupMethod(index, table);
		if (candidate == null)
		    candidate = c2;
		else if (candidate != c2) {
		    valid = false;
		    break;
		}
	    }
	    // Hmm, another symptom of unreachable code
	    result = (valid && candidate != null
		      ? new Devirt(candidate)
		      : VIRTUAL);

	    if (result == VIRTUAL)
		for (int i = 0; i < bp.length; i++)
		    if (bp[i].getType().isScalar())
			liveTables[bp[i].getCID()].set(bp[i].getUID());
		
	    dv[index] = result;
	}
	devirtCalls++;
	if (result != VIRTUAL)
	    devirtSuccess++;
	return result.method;
    }

    public int evalCast(Blueprint from, Blueprint to) {
	int fc = from.getCID();
	int fNumber = from.getUID();
	int tc = to.getCID();
	int tNumber = to.getUID();
	    
	if (safeCastCache[fc][fNumber] == null) {
	    byte[][] tcs = safeCastCache[fc][fNumber]
		= new byte[DomainDirectory.maxContextID()+1][];
	    for (int i = 0; i < tcs.length; i++) {
		Type.Context tctx = DomainDirectory.getContext(i);
		if (tctx != null && tctx.getDomain() == domain)
		    tcs[i] = new byte[tctx.getBlueprintCount()];
	    }
	}
	int isSafe = safeCastCache[fc][fNumber][tc][tNumber];
	if (isSafe == 0) {
	    // casting from Object is common.  Make it fast
	    if (from == ((S3Domain) domain).ROOT_BLUEPRINT) {
		int len = concreteSubtypesOf(to).length;
		isSafe = (len == liveClassList.length ? GOOD_CAST
			  : len == 0 ? BAD_CAST
			  : UNKNOWN);
	    } else {
		boolean canSucceed = false;
		boolean canFail = false;
		Blueprint[] bp = concreteSubtypesOf(from);
		for (int i = 0; i < bp.length; i++) {
		    if (bp[i].isSubtypeOf(to))
			canSucceed = true;
		    else
			canFail = true;

		    if (canSucceed && canFail)
			break;
		}
		isSafe = (canSucceed
			  ? (canFail ? UNKNOWN : GOOD_CAST)
			  : BAD_CAST);
	    }
	    safeCastCache[fc][fNumber][tc][tNumber] = (byte) isSafe;
	}
	castChecks++;
	if (isSafe == GOOD_CAST)
	    castSuccess++;
	else if (isSafe == BAD_CAST)
	    castFailure++;
	return isSafe;
    }

    public boolean isVTableUsed(Blueprint bp) {
	return liveVTables[bp.getCID()].get(bp.getUID());
    }

    public boolean isIFTableUsed(Blueprint bp) {
	return liveIFTables[bp.getCID()].get(bp.getUID());
    }

    public static class Factory extends Analysis.Factory {
	public Analysis make(UserDomain d, IRewriter rewriter,
			     NativeCallGenerator ncg) {
	    return new RapidTypeAnalysis(d, rewriter, ncg);
	}
    }
}
