package s3.services.simplejit.optimization;

import ovm.core.OVMBase;
import ovm.core.repository.ExceptionHandler;
import ovm.core.repository.TypeName;
import ovm.core.repository.Binder;
import ovm.core.repository.RepositoryString;
import ovm.core.domain.Type;
import ovm.core.domain.ConstantResolvedStaticMethodref;
import ovm.core.domain.ConstantResolvedInstanceMethodref;
import ovm.core.domain.LinkageException;
import ovm.core.domain.MemberResolver;
import ovm.services.bytecode.JVMConstants;
import s3.core.domain.S3Constants;
import s3.core.domain.S3ConstantsBuilder;
import s3.core.domain.S3Method;
import s3.core.domain.S3ByteCode;
import s3.services.simplejit.ir.*;
import s3.services.simplejit.util.*;
import s3.services.simplejit.ir.SimpleJITInstruction.*;
import s3.services.simplejit.ir.SimpleJITOperand.*;
import s3.services.simplejit.CompilerVMInterface;
import s3.services.simplejit.SimpleJITLogger;
import java.util.Date;

public class SimpleJITPreInliner {

    public static class Driver {

	private final SimpleJITBC2IR.Driver bc2ir_driver;

	public Driver(SimpleJITBC2IR.Driver bc2ir_driver) {
	    this.bc2ir_driver = bc2ir_driver;
	}

	public SimpleJITIRCode run
	    (S3Method method,
	     CompilerVMInterface compilerVMInterface) {

	    SimpleJITIRCode sjMethod = null;
	    sjMethod = bc2ir_driver.run(method,
					compilerVMInterface);
	    if (sjMethod == null)
		return null;

	    runInliner(sjMethod,
		       compilerVMInterface);

	    return sjMethod;
	}	    
	private void runInliner
	    (SimpleJITIRCode sjMethod,
	     CompilerVMInterface compilerVMInterface) {
	    SimpleJITPreInliner inliner = 
		new SimpleJITPreInliner(sjMethod, 
					bc2ir_driver,
					compilerVMInterface);
	    inliner.run();
	    inliner.print();
	}
    }

    /** This option enables logging in this abstract intepreter, if
     * this is turned on, the abstract interpretation gets way
     * slow. */
    public static final boolean LOG_ENABLED = false;

    /** If this option is enabled, full log messages will be logged,
     * whereas if it is disabled, only certain amount of the latest
     * log messages will be logged. This option is meaningful only
     * when LOG_ENABLED == true */
    private static final boolean FULL_LOG = true;

    /** This option specifies the size of the log message buffer when
     * fullLog == false */
    private static final int LOG_SIZE = 8 * 8192;

    private SimpleJITBasicBlock.Set blockSet;

    private SimpleJITLogger logger;

    private SimpleJITIRCode sjMethod;

    private SimpleJITBC2IR.Driver bc2ir_driver;
    private CompilerVMInterface compilerVMInterface;

    private long inliningTime = 0L;
    private Date startTime;
    private Date endTime;

    public SimpleJITPreInliner(SimpleJITIRCode sjMethod,
			       SimpleJITBC2IR.Driver bc2ir_driver,
			       CompilerVMInterface compilerVMInterface) {
	this.sjMethod = sjMethod;
	this.blockSet = sjMethod.getBlockSet();
	this.bc2ir_driver = bc2ir_driver;
	this.compilerVMInterface = compilerVMInterface;
	this.logger = 
	    FULL_LOG 
	    ? new SimpleJITLogger("SJPreInliner", System.err)
	    : new SimpleJITLogger("SJPreInliner", LOG_SIZE);
    }

    private void log(String message) {
	if (LOG_ENABLED)
	    logger.addLog(message);
    }

    private void __assert__(boolean shouldBeTrue) {
	OVMBase.OVMassert(shouldBeTrue);
    }

    private void run() {
	if (LOG_ENABLED) {
	    log("[PreInliner " 
		+ sjMethod.getS3Method() + "]");
	}

	startTime = new Date();
	inlineStaticCalls();
	inlineSpecialCalls();
	endTime = new Date();
	inliningTime += endTime.getTime() - startTime.getTime();

	sjMethod.optimizationTime += inliningTime;
    }

    private void inlineStaticCalls() {
	S3Method method = sjMethod.getS3Method();
	S3ByteCode bytecode = (S3ByteCode)method.getCode(S3ByteCode.KIND);
	ExceptionHandler[] handlers = bytecode.getExceptionHandlers();
	if (handlers.length > 0) {
 	    if (LOG_ENABLED)
 		log("Not try inlining in " + method 
 		    + " because it has exception handlers.");
	    return;
	}
	S3Constants cp = sjMethod.getConstantPool();
	SimpleJITIdentitySet workset = new SimpleJITIdentitySet();
	SimpleJITBasicBlock.List blockList = blockSet.getBlockList();
	workset.addAll(blockList.toArray());
	while (! workset.empty()) {
	    SimpleJITBasicBlock b = 
		(SimpleJITBasicBlock)workset.remove();
	    SimpleJITInstruction[] insts = b.getInstructions();
	    for(int i = 0; i < insts.length; i++) {
		SimpleJITInstruction inst = insts[i];
		if (! (inst instanceof Call))
		    continue;
		Call call = (Call)inst;
		if (call.getOperator() != Operators.INVOKESTATIC)
		    continue;
		// Now it is a static call

		// Get the target Method and compile it
		ConstantResolvedStaticMethodref smi = null;
		try {
		    smi = cp.resolveStaticMethod(call.getCPIndex());
		} catch (LinkageException e) {
		    if (LOG_ENABLED)
			log("Linkage error : " + call.getSelector());
		    continue;
		}
		S3Method inlinee_method = (S3Method)smi.getMethod();
		inlinee_method = findTargetStaticMethod(inlinee_method);
		if (inlinee_method == null
		    || method == inlinee_method /* no self-recursive */)
		    continue;

		endTime = new Date();
		inliningTime += endTime.getTime() - startTime.getTime();
		SimpleJITIRCode inlinee_ircode =
		    bc2ir_driver.run(inlinee_method, compilerVMInterface);
		startTime = new Date();

		SimpleJITBasicBlock.Set inlinee_blockSet =
		    inlinee_ircode.getBlockSet();
		SimpleJITBasicBlock.List inlinee_blockList =
		    inlinee_blockSet.getBlockList();

		renumberInlineeTemporaries(bytecode.getMaxStack(),
					   bytecode.getMaxLocals(), 
					   inlinee_blockSet);

		// merge CP
		cp = mergeConstantPool(method.getDeclaringType(),
				       cp, 
				       inlinee_blockSet,
				       inlinee_ircode.getConstantPool());

		// split the block into three
		SimpleJITBasicBlock[] block123 =
		    new SimpleJITBasicBlock[3];
		boolean wasSplit = b.splitAtInstruction(inst, block123);

		// Invariant at this point: block123[1] has only one
		// instruction which is the Call instruction

		// Replace block123[1] with the target method's blocks
		block123[1].replaceWith(inlinee_blockSet);

		// Replace GetArgs and Returns with Moves
		connectArgsAndReturnStatic(call, inlinee_blockSet);

		blockList.addUnlessContained(block123[0]);
		blockList.addUnlessContained(block123[1]);
		blockList.addUnlessContained(block123[2]);

		blockList.addAll(inlinee_blockList);
		blockList.remove(block123[1]);

		if (LOG_ENABLED)
		    log("PreInlined " + call.getSelector()
			+ " in " + method);

		// Record that an inlining happened
		sjMethod.setInliningFlag();
		sjMethod.addInlinedMethod(inlinee_ircode);

		if (wasSplit) {
		    // We don't want to go though block123[1]
		    if (block123[2] != block123[1])
			workset.add(block123[2]);
		    break; // block structure changed
		}
	    }
	}
	sjMethod.setConstantPool(cp);
    }

    private void inlineSpecialCalls() {
	S3Method method = sjMethod.getS3Method();
	S3ByteCode bytecode = (S3ByteCode)method.getCode(S3ByteCode.KIND);
	ExceptionHandler[] handlers = bytecode.getExceptionHandlers();
	if (handlers.length > 0) {
 	    if (LOG_ENABLED)
 		log("Not try inlining in " + method 
 		    + " because it has exception handlers.");
	    return;
	}
	S3Constants cp = sjMethod.getConstantPool();
	SimpleJITIdentitySet workset = new SimpleJITIdentitySet();
	SimpleJITBasicBlock.List blockList = blockSet.getBlockList();
	workset.addAll(blockList.toArray());
	while (! workset.empty()) {
	    SimpleJITBasicBlock b = 
		(SimpleJITBasicBlock)workset.remove();
	    SimpleJITInstruction[] insts = b.getInstructions();
	    for(int i = 0; i < insts.length; i++) {
		SimpleJITInstruction inst = insts[i];
		if (! (inst instanceof Call))
		    continue;
		Call call = (Call)inst;
		if (call.getOperator() != Operators.INVOKESPECIAL)
		    continue;

		// Get the target Method and compile it
		ConstantResolvedInstanceMethodref imi = null;
		try {
		    imi = cp.resolveInstanceMethod(call.getCPIndex());
		} catch (LinkageException e) {
		    if (LOG_ENABLED)
			log("Linkage error : " + call.getSelector());
		    continue;
		}
		if (! imi.isNonVirtual) // only deal with 'special' calls
		    continue;
		S3Method inlinee_method = (S3Method)imi.getMethod();
		if (inlinee_method == null
		    || method == inlinee_method /* no self-recursive */)
		    continue;
		if (inlinee_method.getCode(S3ByteCode.KIND) == null)
		    throw new Error("Found special method without bytecode");

		endTime = new Date();
		inliningTime += endTime.getTime() - startTime.getTime();
		SimpleJITIRCode inlinee_ircode =
		    bc2ir_driver.run(inlinee_method, compilerVMInterface);
		startTime = new Date();

		SimpleJITBasicBlock.Set inlinee_blockSet =
		    inlinee_ircode.getBlockSet();
		SimpleJITBasicBlock.List inlinee_blockList =
		    inlinee_blockSet.getBlockList();

		renumberInlineeTemporaries(bytecode.getMaxStack(),
					   bytecode.getMaxLocals(), 
					   inlinee_blockSet);

		// merge CP
		cp = mergeConstantPool(method.getDeclaringType(),
				       cp, 
				       inlinee_blockSet,
				       inlinee_ircode.getConstantPool());

		// split the block into three
		SimpleJITBasicBlock[] block123 =
		    new SimpleJITBasicBlock[3];
		boolean wasSplit = b.splitAtInstruction(inst, block123);

		// Invariant at this point: block123[1] has only one
		// instruction which is the Call instruction

		// Replace block123[1] with the target method's blocks
		block123[1].replaceWith(inlinee_blockSet);

		// Replace GetArgs and Returns with Moves
		connectArgsAndReturnInstance(call, inlinee_blockSet);

		blockList.addUnlessContained(block123[0]);
		blockList.addUnlessContained(block123[1]);
		blockList.addUnlessContained(block123[2]);

		blockList.addAll(inlinee_blockList);
		blockList.remove(block123[1]);

		if (LOG_ENABLED)
		    log("PreInlined " + call.getSelector()
			+ " in " + method);

		// Record that an inlining happened
		sjMethod.setInliningFlag();
		sjMethod.addInlinedMethod(inlinee_ircode);

		if (wasSplit) {
		    // We don't want to go though block123[1]
		    if (block123[2] != block123[1])
			workset.add(block123[2]);
		    break; // block structure changed
		}
	    }
	}
	sjMethod.setConstantPool(cp);
    }

    private S3Method findTargetStaticMethod(S3Method method) {
	try {
	    MemberResolver resolver = 
		method.getDeclaringType().getDomain().getMemberResolver();
	    S3Method target =
		(S3Method)resolver
		.resolveStaticMethod(method.getDeclaringType(),
				     method.getSelector().getUnboundSelector(),
				     null);
	    //	    if (LOG_ENABLED)
	    //		log("findTargetStaticMethod : " + method + " -> " + target);
	    if (target.getMode().isNative())
		return null;
	    else
		return target;
	} catch (Throwable t) {
	    throw new Error
		("The inliner can't find the target of " + method + " : " + t);
	}
    }

    private S3Constants mergeConstantPool
	(Type.Compound type,
	 S3Constants inliner_cp,
	 SimpleJITBasicBlock.Set inlinee_blockSet,
	 S3Constants inlinee_cp) {

	S3ConstantsBuilder s3cb = new S3ConstantsBuilder(type, inliner_cp);

	SimpleJITIdentitySet workset = new SimpleJITIdentitySet();
	workset.addAll(inlinee_blockSet.getBlockList().toArray());
	while (! workset.empty()) {
	    SimpleJITBasicBlock b = 
		(SimpleJITBasicBlock)workset.remove();
	    SimpleJITInstruction[] insts = b.getInstructions();
	    for(int i = 0; i < insts.length; i++) {
		SimpleJITInstruction inst = insts[i];
		if (inst.isConstantPoolReferer()) {
		    ConstantPoolReferer cpr = inst.asConstantPoolReferer();
		    int cpindex = cpr.getCPIndex();
		    if (cpindex < 0)
			continue;
		    byte tag = inlinee_cp.getTagAt(cpindex);
		    int newCPIndex;
		    switch(tag) {
		    case JVMConstants.CONSTANT_Class:
		    case JVMConstants.CONSTANT_ResolvedClass:
			newCPIndex = s3cb.addClass(inlinee_cp.getClassAt(cpindex));
			break;
		    case JVMConstants.CONSTANT_Fieldref:
		    case JVMConstants.CONSTANT_ResolvedStaticField:
		    case JVMConstants.CONSTANT_ResolvedInstanceField:
			newCPIndex = 
			    s3cb.addFieldref(inlinee_cp.getFieldrefAt(cpindex));
			break;
		    case JVMConstants.CONSTANT_Methodref:
		    case JVMConstants.CONSTANT_ResolvedStaticMethod:
		    case JVMConstants.CONSTANT_ResolvedInstanceMethod:
			newCPIndex = 
			    s3cb.addMethodref(inlinee_cp.getMethodrefAt(cpindex));
			break;
		    case JVMConstants.CONSTANT_InterfaceMethodref:
		    case JVMConstants.CONSTANT_ResolvedInterfaceMethod:
			newCPIndex = s3cb.addInterfaceMethodref
			    (inlinee_cp.getMethodrefAt(cpindex));
			break;
		    case JVMConstants.CONSTANT_String:
			newCPIndex = s3cb.addUnresolvedString
			    ((RepositoryString)inlinee_cp.getConstantAt(cpindex));
			break;
		    case JVMConstants.CONSTANT_Reference:
			newCPIndex = s3cb.addResolvedConstant
			    (inlinee_cp.getConstantAt(cpindex));
			break;
		    case JVMConstants.CONSTANT_Binder:
			newCPIndex = s3cb.addUnresolvedBinder
			    ((Binder)inlinee_cp.getConstantAt(cpindex));
			break;
		    case JVMConstants.CONSTANT_SharedState:
			newCPIndex = s3cb.addUnresolvedSharedState
			    ((TypeName.Gemeinsam)inlinee_cp.getConstantAt(cpindex));
			break;
		    case 0: // empty
			newCPIndex = 0;
			break;
		    default:
			throw new Error("Unexpected tag " + tag);
		    }
		    cpr.setCPIndex(newCPIndex);
		}
	    }
	}
	return s3cb.build();
    }

    private void renumberInlineeTemporaries(int inliner_maxstack,
					    int inliner_maxlocals,
					    SimpleJITBasicBlock.Set inlinee_blockSet) {
	inliner_maxstack = (inliner_maxstack + 3) * 2;
	TemporaryFactory tempFactory = new TemporaryFactory();
	SimpleJITIdentitySet workset = new SimpleJITIdentitySet();
	workset.addAll(inlinee_blockSet.getBlockList().toArray());
	while (! workset.empty()) {
	    SimpleJITBasicBlock b = 
		(SimpleJITBasicBlock)workset.remove();
	    SimpleJITInstruction[] insts = b.getInstructions();
	    for(int i = 0; i < insts.length; i++) {
		SimpleJITInstruction inst = insts[i];
		if (inst.isUser()) {
		    SimpleJITOperand[] uses = 
			inst.asUser().getUseOperands();
		    for(int j = 0; j < uses.length; j++) {
			if (uses[j].isStackTemporaryOperand()) {
			    SimpleJITStackTemporaryOperand old_stack_op =
				uses[j].asStackTemporaryOperand();
			    SimpleJITStackTemporaryOperand new_stack_op = 
				tempFactory.createStackTemporary
				(old_stack_op.getTypeCode(),
				 old_stack_op.getIndex() + inliner_maxstack);
			    uses[j] = new_stack_op;
			} else if (uses[j].isLocalTemporaryOperand()) {
			    SimpleJITLocalTemporaryOperand old_local_op =
				uses[j].asLocalTemporaryOperand();
			    SimpleJITLocalTemporaryOperand new_local_op = 
				tempFactory.createLocalTemporary
				(old_local_op.getTypeCode(),
				 old_local_op.getIndex() + inliner_maxlocals);
			    uses[j] = new_local_op;
			}
		    }
		    inst.asUser().setUseOperands(uses);
		}
		if (inst.isDefiner()) {
		    SimpleJITOperand[] defs = 
			inst.asDefiner().getDefOperands();
		    for(int j = 0; j < defs.length; j++) {
			if (defs[j].isStackTemporaryOperand()) {
			    SimpleJITStackTemporaryOperand old_stack_op =
				defs[j].asStackTemporaryOperand();
			    SimpleJITStackTemporaryOperand new_stack_op = 
				tempFactory.createStackTemporary
				(old_stack_op.getTypeCode(),
				 old_stack_op.getIndex() + inliner_maxstack);
			    defs[j] = new_stack_op;
			} else if (defs[j].isLocalTemporaryOperand()) {
			    SimpleJITLocalTemporaryOperand old_local_op =
				defs[j].asLocalTemporaryOperand();
			    SimpleJITLocalTemporaryOperand new_local_op = 
				tempFactory.createLocalTemporary
				(old_local_op.getTypeCode(),
				 old_local_op.getIndex() + inliner_maxlocals);
			    defs[j] = new_local_op;
			}
		    }
		    inst.asDefiner().setDefOperands(defs);
		}
	    }
	}
    }

    /**
     * For a static call inlining, replace GetArgs and Return with
     * moves. Discard GetArg0 (shst).
     * FIXME: synchronized static call
     */
    private void connectArgsAndReturnStatic
	(Call call,
	 SimpleJITBasicBlock.Set inlinee_blockSet) {
	boolean isVoidCall = (call instanceof CallVoid);
	SimpleJITOperand result = 
	    isVoidCall 
	    ? null
	    : ((CallNonVoid)call).getResult();
	SimpleJITOperand[] actualArgs = call.getArgs();
	SimpleJITIdentitySet workset = new SimpleJITIdentitySet();
	workset.addAll(inlinee_blockSet.getBlockList().toArray());
	while (! workset.empty()) {
	    SimpleJITBasicBlock b = 
		(SimpleJITBasicBlock)workset.remove();
	    SimpleJITInstruction[] insts = b.getInstructions();
	    for(int i = 0; i < insts.length; i++) {
		SimpleJITInstruction inst = insts[i];
		if (inst instanceof GetArg) {
		    GetArg ga = (GetArg)inst;
		    int index = ga.getIndex();
		    if (index == 0) {
			b.removeInstruction(inst);
		    } else {
			b.replaceInstruction
			    (ga, 
			     Move.create(ga.getType(),
					 ga.getDestination(),
					 actualArgs[index - 1],
					 null));
		    }
		} else if (inst instanceof ReturnValue) {
		    if (isVoidCall)
			throw new Error();
		    ReturnValue rv = (ReturnValue)inst;
		    b.replaceInstruction
			(rv,
			 Move.create(rv.getType(),
				     result,
				     rv.getReturnValueOperand(),
				     null));
		    b.setFlowEndBlock(false);
		} else if (inst instanceof Return) {
		    if (! isVoidCall)
			throw new Error();
		    b.removeInstruction(inst);
		    b.setFlowEndBlock(false);
		}
	    }
	}
    }

    /**
     * For an instance call inlining, replace GetArgs and Return with
     * moves.
     * FIXME: synchronized call
     */
    private void connectArgsAndReturnInstance
	(Call call,
	 SimpleJITBasicBlock.Set inlinee_blockSet) {
	boolean isVoidCall = (call instanceof CallVoid);
	SimpleJITOperand result = 
	    isVoidCall 
	    ? null
	    : ((CallNonVoid)call).getResult();
	SimpleJITOperand[] actualArgs = call.getArgs();
	SimpleJITIdentitySet workset = new SimpleJITIdentitySet();
	workset.addAll(inlinee_blockSet.getBlockList().toArray());
	while (! workset.empty()) {
	    SimpleJITBasicBlock b = 
		(SimpleJITBasicBlock)workset.remove();
	    SimpleJITInstruction[] insts = b.getInstructions();
	    for(int i = 0; i < insts.length; i++) {
		SimpleJITInstruction inst = insts[i];
		if (inst instanceof GetArg) {
		    GetArg ga = (GetArg)inst;
		    int index = ga.getIndex();
		    b.replaceInstruction
			(ga, 
			 Move.create(ga.getType(),
				     ga.getDestination(),
				     actualArgs[index],
				     null));
		} else if (inst instanceof ReturnValue) {
		    if (isVoidCall)
			throw new Error();
		    ReturnValue rv = (ReturnValue)inst;
		    b.replaceInstruction
			(rv,
			 Move.create(rv.getType(),
				     result,
				     rv.getReturnValueOperand(),
				     null));
		    b.setFlowEndBlock(false);
		} else if (inst instanceof Return) {
		    if (! isVoidCall)
			throw new Error();
		    b.removeInstruction(inst);
		    b.setFlowEndBlock(false);
		}
	    }
	}
    }

    private void print() {
	if (LOG_ENABLED) {
	    printBlockSet();
	}
    }

    private void printBlockSet() {
	String b = blockSet.toString();
	log(b);
    }
}
