package s3.services.simplejit.bytecode;

import java.io.BufferedOutputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.PrintWriter;

import ovm.core.domain.ConstantResolvedMethodref;
import ovm.core.domain.ConstantResolvedStaticMethodref;
import ovm.core.domain.Domain;
import ovm.core.domain.Method;
import ovm.core.repository.ConstantMethodref;
import ovm.core.repository.ConstantFieldref;
import ovm.core.repository.ConstantClass;
import ovm.core.repository.Constants;
import ovm.core.repository.ConstantsEditor;
import ovm.core.repository.TypeCodes;
import ovm.core.repository.RepositoryString;
import ovm.core.repository.TypeName;
import ovm.services.bytecode.JVMConstants;
import ovm.services.bytecode.ByteCodeGen;
import ovm.services.bytecode.CodeExceptionGen;
import ovm.services.bytecode.Instruction;
import ovm.services.bytecode.InstructionHandle;
import ovm.services.bytecode.InstructionList;
import ovm.services.bytecode.InstructionTargeter;
import ovm.services.bytecode.TargetLostException;
import ovm.services.bytecode.Instruction.LocalWrite;
import ovm.util.ArrayList;
import ovm.util.HTint2Object;
import ovm.util.HashSet;
import ovm.util.HashMap;
import ovm.util.Iterator;
import ovm.util.NumberRanges;
import ovm.util.Vector;
import s3.core.domain.S3ByteCode;
import s3.services.bootimage.Analysis.Visitor;

/**
 * A simple bytecode (ovm IR) inliner
 * @author Hiroshi Yamauchi
 */
public class Inliner {
    
    protected static final boolean DEBUG = true;
    protected static PrintWriter inliningLog = null;
    protected static final String inliningLogFile = "inlining.log";
    static {
	try {
	    if (DEBUG) {
		inliningLog = new PrintWriter(new BufferedOutputStream(new FileOutputStream(inliningLogFile)));
	    }
	} catch (IOException e) { throw new RuntimeException(e); }
    }

    protected static final HashSet inlined = new HashSet();

    protected static final int CALLEE_MAX_SIZE = 10;
    protected static final int MAX_DEPTH = 5;

    protected Method method;
    protected final int depth;
    protected HTint2Object work = new HTint2Object();
    
    public Inliner(final Method method) {
        this.method = method;
	this.depth = 1;
    }

    protected Inliner(final Method method, int depth) {
        this.method = method;
	this.depth = depth;
    }

    public Method method() { return method; }
    
    public void inline() {
	if (depth >= MAX_DEPTH)
	    return;
	if (! inlined.contains(method)) {
	    find();
	    if (work.keys().length > 0) {
		perform();
	    }
	    inlined.add(method);
	}
    }
    
    protected void find() {
        final Domain domain = method.getDeclaringType().getDomain();
        final Visitor visitor = new Visitor() {
            public void visit(Instruction.Invocation i) {
                ConstantMethodref mr = i.getConstantMethodref(buf, cp);
                try {
                    Method m = ((ConstantResolvedMethodref) mr).getMethod();
		    if (domain != m.getDeclaringType().getDomain()
			|| method == m) {
			return;
		    }
		    S3ByteCode calleeBC = (S3ByteCode) m.getCode(S3ByteCode.KIND);
                    if (calleeBC == null
			|| calleeBC.getBytes() == null
			|| ! shouldInline(m)) {
                        return;
		    }
                    if (m.isNonVirtual() || m.getMode().isFinal()
			|| m.getDeclaringType().getMode().isFinal()
			|| (i instanceof Instruction.INVOKESPECIAL)) {
                        work.put(getPC(), m);
                    }
                } catch (ClassCastException e) {
                    warnMissing(i.getSelector(buf, cp).toString());
                }
            }
        };
        visitor.run(method);
    }
    
    protected void perform() {
        if (DEBUG) {
            inliningLog.println("Inlining " + method.toString());
            S3ByteCode bc = (S3ByteCode) method.getCode(S3ByteCode.KIND);
            bc.dumpAscii("", inliningLog);
            inliningLog.flush();
        }
        
        S3ByteCode bc = (S3ByteCode) method.getCode(S3ByteCode.KIND);
        ByteCodeGen bcgen = new ByteCodeGen(bc);
	bcgen.removeLineNumbers();
	bcgen.removeLocalVariables();
	runVisitor(bcgen);
        method.addCode(bcgen.getByteCode());
        
        if (DEBUG) {
            inliningLog.println("Inlined " + method.toString());
            S3ByteCode _bc = (S3ByteCode) method.getCode(S3ByteCode.KIND);
            _bc.dumpAscii("", inliningLog);
            inliningLog.flush();
        }
    }

    protected void runVisitor(ByteCodeGen bcgen) {
        new InlineVisitor(bcgen).run();
    }
    
    protected boolean shouldInline(Method callee) {
        S3ByteCode bc = (S3ByteCode) callee.getCode(S3ByteCode.KIND);
        if (bc.getBytes().length < CALLEE_MAX_SIZE && bc.getExceptionHandlers().length == 0
                && ! callee.getMode().isSynchronized()) {
            return true;
        }
        return false;
    }

    protected void inlineRecursively(Method method) {
	if (! inlined.contains(method)) {
	    new Inliner(method, depth + 1).inline();
	    inlined.add(method);
	}
    }

    protected class InlineVisitor extends Visitor {
        ByteCodeGen bcgen;
        InstructionList il;
        InstructionHandle ih;
        HashMap p2nCache;
        final int maxLocals;
        final int maxStack;
	int newMaxLocals;
	int newMaxStack;
        
        public InlineVisitor(ByteCodeGen bcgen) {
            this.bcgen = bcgen;
            this.il = bcgen.getInstructionList();
            this.maxLocals = bcgen.getMaxLocals();
            this.maxStack = bcgen.getMaxStack();
	    newMaxLocals = maxLocals;
	    newMaxStack = maxStack;
        }

        public void run() {
            ih = il.getStart();
            HashMap ih2Callee = new HashMap();
            while (ih != null) {
                InstructionHandle next = ih.getNext();
                Object callee = work.get(ih.getPC());
                if (callee != null) {
		    if (DEBUG) {
			inliningLog.println("@PC" + ih.getPC() + " " + callee.toString());
			S3ByteCode _bc = (S3ByteCode) ((Method)callee).getCode(S3ByteCode.KIND);
			_bc.dumpAscii("", inliningLog);
			inliningLog.flush();
		    }
                    ih2Callee.put(ih, callee);
                }
                ih = next;
            }
            for(Iterator it = ih2Callee.keySet().iterator();
                it.hasNext(); ) {
                ih = (InstructionHandle) it.next();
                ih.accept(this);
            }
            bcgen.setMaxLocals(newMaxLocals);
            bcgen.setMaxStack(newMaxStack);
        }

        public void visit(Instruction.Invocation i) {
	    boolean isInvokeStatic = i instanceof Instruction.INVOKESTATIC;
            Method _method = (Method) work.get(ih.getPC());
	    inlineRecursively(_method);
            ByteCodeGen _bcgen = new ByteCodeGen((S3ByteCode) _method.getCode(S3ByteCode.KIND));
	    ConstantsEditor callerCP_W = bcgen.getConstantPoolEditor();
	    Constants callerCP_R = bcgen.getConstantPool();
            transformCallee(_bcgen, ih.getNext(), callerCP_W, maxLocals);
            InstructionList _il = _bcgen.getInstructionList();
            int _maxLocals = _bcgen.getMaxLocals();
            int _maxStack = _bcgen.getMaxStack();
            ArrayList _argTypeList = new ArrayList();
            int _argSize;
            char[] _argTypes = _bcgen.getArgumentTypes();
            _argTypeList.add(new Character(TypeCodes.REFERENCE));
            for (int j = 0; j < _argTypes.length; j++) {
                char t = toBasicType(_argTypes[j]);
                if (typeCode2Size(t) == 1) {
                    _argTypeList.add(new Character(t));
                } else if (typeCode2Size(t) == 2) {
                    _argTypeList.add(new Character(t));
                    _argTypeList.add(new Character(TypeCodes.VOID));
                } else
                    throw new Error();
            }
            _argSize = _argTypeList.size();
            
            // Insert localWrites to pop the arguments
            int localIndex = maxLocals + _argSize - 1;
            InstructionHandle firstInserted = null;
            for (int j = _argTypes.length - 1; j >= 0; j--) {
                char t = toBasicType(_argTypes[j]);
                if (t == TypeCodes.VOID)
                    continue;
		if (typeCode2Size(t) == 2)
		    localIndex--;
                if (firstInserted == null)
                    firstInserted = il.insert(ih, LocalWrite.make(t, localIndex));
                else
                    il.insert(ih, LocalWrite.make(t, localIndex));
                localIndex--;
            }
	    if (isInvokeStatic) {
		Instruction.INVOKESTATIC invokeStatic = (Instruction.INVOKESTATIC) i;
		int invokeStaticCPIndex = invokeStatic.getCPIndex();		
		Object entry = callerCP_R.getConstantAt(invokeStaticCPIndex);
		byte tag = callerCP_R.getTagAt(invokeStaticCPIndex);
		int shstCPIndex = -1;
		switch(tag) {
		case JVMConstants.CONSTANT_ResolvedStaticMethod: {
		    ConstantResolvedStaticMethodref mr =
			(ConstantResolvedStaticMethodref) entry;
		    shstCPIndex = callerCP_W.addResolvedConstant(mr.getSharedState());
		    break;
		}
		case JVMConstants.CONSTANT_Methodref:
		default:
		    throw new Error();
		}
		if (firstInserted == null) {
		    firstInserted = il.insert(ih, new Instruction.LDC_REF_QUICK(shstCPIndex));
		} else {
		    il.insert(ih, new Instruction.LDC_REF_QUICK(shstCPIndex));
		}
		il.insert(ih, LocalWrite.make(TypeCodes.REFERENCE, localIndex));
	    } else {
		if (firstInserted == null)
		    firstInserted = il.insert(ih, LocalWrite.make(TypeCodes.REFERENCE, localIndex));
		else
		    il.insert(ih, LocalWrite.make(TypeCodes.REFERENCE, localIndex));
		// insert REF_GETFIELD_QUICK to enforce a null check
		il.insert(ih, Instruction.LocalRead.make(TypeCodes.REFERENCE, localIndex));
		il.insert(ih, new Instruction.REF_GETFIELD_QUICK(0));
		il.insert(ih, Instruction.POP.make());
	    }
            
            // Insert the callee instructions into caller
	    if (DEBUG) {
		//inliningLog.println("Caller code : \n" + il.toString());
		//inliningLog.println("Callee code : \n" + _il.toString());
		//inliningLog.flush();
	    }

	    il.insert(ih, _il);

            InstructionHandle last = ih.getPrev();
            // Remove the INVOKE* instruction
            try {
                il.delete(ih);
            } catch (TargetLostException e) {
                InstructionHandle[] targets = e.getTargets();
                for (int k = 0; k < targets.length; k++) {
                    InstructionTargeter[] targeters = targets[k]
                            .getTargeters();
                    for (int j = 0; j < targeters.length; j++) {
			if (targeters[j] instanceof CodeExceptionGen) {
			    CodeExceptionGen ceg = (CodeExceptionGen) targeters[j];
			    if (ceg.getStartPC() == ih
				&& ceg.getEndPC() == ih) {
				ceg.setStartPC(firstInserted);
				ceg.setEndPC(last);
			    } else if (ceg.getStartPC() == ih) {
				targeters[j].updateTarget(targets[k],
							  firstInserted);
			    } else if (ceg.getEndPC() == ih) {
				targeters[j].updateTarget(targets[k],
							  last);
			    } else if (ceg.getHandlerPC() == ih) {
				targeters[j].updateTarget(targets[k],
							  firstInserted);
			    }
			} else {
			    targeters[j].updateTarget(targets[k],
						      firstInserted);
			}
                    }
                }
            }
            
	    if (newMaxLocals < maxLocals + _argSize + _maxLocals) {
		newMaxLocals = maxLocals + _argSize + _maxLocals;
	    }
	    if (newMaxStack < maxStack + _maxStack) {
		newMaxStack = maxStack + _maxStack;
	    }
        }
    }

    private static void transformCallee(ByteCodeGen bcgen, 
            InstructionHandle exit, ConstantsEditor callerCP, int callerMaxLocals) {
        TransformCalleeVisitor v = new TransformCalleeVisitor(bcgen, exit, callerCP, callerMaxLocals);
        v.run();
    }
    
    /**
     * Shift locals, replace returns with gotos, copy the CP entries to the caller's
     */
    protected static class TransformCalleeVisitor extends Instruction.Visitor {
        ByteCodeGen bcgen;
        Constants cp;
        InstructionList il;
        InstructionHandle ih;
        final int callerMaxLocals;
        Vector argTypeList;
        int argSize;
        InstructionHandle exit;
        ConstantsEditor callerCP;

        public TransformCalleeVisitor(ByteCodeGen bcgen, 
                InstructionHandle exit,
                ConstantsEditor callerCP,
		int callerMaxLocals) {
            this.bcgen = bcgen;
            this.il = bcgen.getInstructionList();
            this.callerMaxLocals = callerMaxLocals;
            cp = bcgen.getConstantPool();
            argTypeList = new Vector();
            this.exit = exit;
            this.callerCP = callerCP;
        }

        public void run() {
	    bcgen.removeLineNumbers();
	    bcgen.removeLocalVariables();

            char[] argTypes = bcgen.getArgumentTypes();
            argTypeList.add(new Character(TypeCodes.REFERENCE));
            for (int i = 0; i < argTypes.length; i++) {
                char t = toBasicType(argTypes[i]);
                if (typeCode2Size(t) == 1) {
                    argTypeList.add(new Character(t));
                } else if (typeCode2Size(t) == 2) {
                    argTypeList.add(new Character(t));
                    argTypeList.add(new Character(TypeCodes.VOID));
                } else
                    throw new Error();
            }
            argSize = argTypeList.size();

            ih = il.getStart();
            while (ih != null) {
                InstructionHandle next = ih.getNext();
                ih.accept(this);
                ih = next;
            }

            int maxSize = 0;
            maxSize = 1;
            for (int i = 0; i < argTypes.length; i++) {
                char t = toBasicType(argTypes[i]);
                int size = typeCode2Size(t);
                if (maxSize < size)
                    maxSize = size;
            }
        }

        public void visit(Instruction.LocalWrite o) {
            if (argSize > 0) {
                ih.setInstruction(Instruction.LocalWrite.make(o.getTypeCode(), o
                        .getLocalVariableOffset()
                        + callerMaxLocals));
            }
        }

        public void visit(Instruction.LocalRead o) {
            if (argSize > 0) {
                ih.setInstruction(Instruction.LocalRead.make(o.getTypeCode(), o
                        .getLocalVariableOffset()
                        + callerMaxLocals));
            }
        }
        
        public void visit(Instruction.IINC o) {
            ih.setInstruction(Instruction.IINC.make(o.getLocalVariableOffset() + callerMaxLocals,
                    o.getValue()));
        }
        
        public void visit(Instruction.RET o) {
            ih.setInstruction(Instruction.RET.make(NumberRanges.checkUnsignedShort(o.getLocalVariableOffset() + callerMaxLocals)));
        }
        public void visit(Instruction.WIDE_RET o) {
            ih.setInstruction(Instruction.WIDE_RET.make(NumberRanges.checkUnsignedShort(o.getLocalVariableOffset() + callerMaxLocals)));
        }
        
        public void visit(Instruction.ReturnValue o) {
	    if (ih != il.getEnd()) {
		ih.setInstruction(Instruction.GOTO.make(exit));
	    } else {
		try {
		    il.delete(ih);
		} catch (TargetLostException e) {
		    InstructionHandle[] targets = e.getTargets();
		    for (int k = 0; k < targets.length; k++) {
			InstructionTargeter[] targeters = targets[k]
			    .getTargeters();
			for (int j = 0; j < targeters.length; j++) {
			    targeters[j].updateTarget(targets[k],
						      exit);
			}
		    }
		}
	    }
        }
        
        public void visit(Instruction.RETURN o) {
	    if (ih != il.getEnd()) {
		ih.setInstruction(Instruction.GOTO.make(exit));
	    } else {
		try {
		    il.delete(ih);
		} catch (TargetLostException e) {
		    InstructionHandle[] targets = e.getTargets();
		    for (int k = 0; k < targets.length; k++) {
			InstructionTargeter[] targeters = targets[k]
			    .getTargeters();
			for (int j = 0; j < targeters.length; j++) {
			    targeters[j].updateTarget(targets[k],
						      exit);
			}
		    }
		}
	    }
        }
        
        public void visit(Instruction.ConstantPoolRead o) {
            int oldIndex = o.getCPIndex();
            Object entry = cp.getConstantAt(oldIndex);
	    byte tag = cp.getTagAt(oldIndex);
	    int newIndex = -1;
	    switch(tag) {
	    case JVMConstants.CONSTANT_ResolvedStaticMethod:
	    case JVMConstants.CONSTANT_ResolvedInstanceMethod:
	    case JVMConstants.CONSTANT_Methodref:
		newIndex = callerCP.addMethodref((ConstantMethodref)entry);
		break;
	    case JVMConstants.CONSTANT_ResolvedInterfaceMethod:
	    case JVMConstants.CONSTANT_InterfaceMethodref:
		newIndex = callerCP.addInterfaceMethodref((ConstantMethodref)entry);
		break;
	    case JVMConstants.CONSTANT_ResolvedStaticField:
	    case JVMConstants.CONSTANT_ResolvedInstanceField:
	    case JVMConstants.CONSTANT_Fieldref:
		newIndex = callerCP.addFieldref((ConstantFieldref)entry);
		break;
	    case JVMConstants.CONSTANT_Class:
	    case JVMConstants.CONSTANT_ResolvedClass:
		newIndex = callerCP.addClass((ConstantClass)entry);
		break;
	    case JVMConstants.CONSTANT_SharedState:
		newIndex = callerCP.addUnresolvedSharedState((TypeName.Gemeinsam)entry);
		break;
	    case JVMConstants.CONSTANT_String:
		newIndex = callerCP.addUnresolvedString((RepositoryString)entry);
		break;
	    default:
		newIndex = callerCP.addResolvedConstant(entry);
		break;
	    }
	    if (newIndex < 0) {
		throw new Error(o.getClass().toString());
	    }
	    byte newTag = ((Constants)callerCP).getTagAt(newIndex);
	    Object newEntry = ((Constants)callerCP).getConstantAt(newIndex);
	    if ( !(newTag == tag && entry == newEntry)
		 //		 || !(tag == JVMConstants.CONSTANT_String && entry.equals(newEntry))) {
		 && !(newTag == tag && entry.equals(newEntry))) {
		throw new Error("tag=" + tag + ", newtag=" + newTag +
				", entry=" + entry.getClass() + ":" + entry + 
				", newentry=" + newEntry.getClass() + ":" + newEntry);
	    }
            o.setCPIndex(newIndex);
        }
    }

    protected static char toBasicType(char type) {
        if (type == TypeCodes.REFERENCE || type == TypeCodes.OBJECT || type == TypeCodes.ARRAY)
            return TypeCodes.REFERENCE;
        else if (type == TypeCodes.BOOLEAN || type == TypeCodes.BYTE || type == TypeCodes.CHAR
                || type == TypeCodes.SHORT)
            return TypeCodes.INT;
        else {
            return type;
        }
    }
    protected static int typeCode2Size(char t) {
        if (t == TypeCodes.LONG || t == TypeCodes.DOUBLE)
            return 2;
        else if (t == TypeCodes.INT || t == TypeCodes.FLOAT || t == TypeCodes.REFERENCE
                || t == TypeCodes.SHORT || t == TypeCodes.CHAR || t == TypeCodes.BYTE 
                || t == TypeCodes.BOOLEAN 
                || t == TypeCodes.OBJECT 
                || t == TypeCodes.ARRAY) 
            return 1;
        else
            throw new Error();
    }
}
