package s3.services.simplejit.bytecode;

import ovm.services.bytecode.Instruction;
import ovm.services.bytecode.InstructionBuffer;
import ovm.services.bytecode.JVMConstants;
import ovm.services.bytecode.SpecificationIR;
import ovm.services.bytecode.analysis.AbstractValue;
import ovm.services.bytecode.analysis.Frame;
import ovm.services.bytecode.analysis.State;
import ovm.services.bytecode.editor.CodeFragmentEditor;
import ovm.services.bytecode.editor.InstructionEditVisitor;
import s3.core.domain.S3ByteCode;
import s3.core.domain.S3Constants;
import s3.core.domain.S3ConstantsBuilder;
import s3.services.bytecode.analysis.S3Frame;
import s3.services.simplejit.bytecode.SimpleJITValue.SimpleJITDouble;
import s3.services.simplejit.bytecode.SimpleJITValue.SimpleJITLong;
import s3.services.simplejit.bytecode.SimpleJITValue.SimpleJITReference;

public class SimpleJITJSRProblemSolver 
    implements JVMConstants {

    public static class ReferenceMap {
	boolean[] local_map;
	boolean[] stack_map;
	private ReferenceMap(boolean[] local_map,
			     boolean[] stack_map) {
	    this.local_map = local_map;
	    this.stack_map = stack_map;
	}
	public String toString() {
	    String ret = "locals: ";
	    for(int i = 0; i < local_map.length; i++)
		if (local_map[i])
		    ret += "[" + i + "]1";
		else
		    ret += "[" + i + "]0";
	    ret += " stack: ";
	    for(int i = 0; i < stack_map.length; i++)
		if (stack_map[i])
		    ret += "[" + i + "]1";
		else
		    ret += "[" + i + "]0";
	    return "REFMAP " + ret; 
	}
    }

    public static class BytecodeLocalSplitter 
	extends InstructionEditVisitor {
	int[] mapping_;
	private BytecodeLocalSplitter(int[] mapping) {
	    mapping_ = mapping;
	}
	public void visit(Instruction.LocalAccess i) {
	    if (mapping_ == null)
		return;
	    int old_index = i.getLocalVariableOffset(buf);
	    int new_index = mapping_[old_index];
	    if (new_index == -1)
		return;
	    if (i instanceof Instruction.LocalRead) { // xLOAD_y
		SpecificationIR.Value v = i.stackOuts[0];
		if (v instanceof SpecificationIR.RefValue)
		    cfe.replaceInstruction()
			.addALoad((char)new_index);
	    } else if (i instanceof Instruction.LocalWrite) { // xSTORE_y
		SpecificationIR.Value v = i.stackIns[0];
		if (v instanceof SpecificationIR.RefValue)
		    cfe.replaceInstruction()
			.addAStore((char)new_index);
	    }
	}
    }


    private S3ByteCode originalByteCode;
    private State state;
    private S3ByteCode rewrittenByteCode;
    private ReferenceMap[] refMaps;
    private int[] local_split_map;
    private int codeLength;
    private int newMaxLocal;
    private SimpleJITEphemeralTypeMap ephemeralTypeMap;

    public SimpleJITJSRProblemSolver
	(S3ByteCode bytecode,
	 State state,
	 SimpleJITEphemeralTypeMap ephemeralTypeMap) {
	this.originalByteCode = bytecode;
	this.state = state;
	this.codeLength = bytecode.getBytes().length;
	local_split_map = new int[bytecode.getMaxLocals()];
	for(int i = 0; i < local_split_map.length; i++)
	    local_split_map[i] = -1;
	newMaxLocal = bytecode.getMaxLocals();
	this.ephemeralTypeMap = ephemeralTypeMap;
    }

    public S3ByteCode getRewrittenByteCode() {
	return rewrittenByteCode;
    }

    public void run() {
	computeSplitReferenceMaps();
	if (local_split_map != null) {
	    localSplitInState();
	    localSplitInByteCode();
	}
    }

    private Frame copyAndIncreaseMaxLocal(Frame f) {
	S3Frame s3f = (S3Frame)f;
	int stackHeight = s3f.getStackHeight();
	int maxStack = s3f.getMaxStack();
	int oldMaxLocal = s3f.getMaxLocals();
	AbstractValue[] opStack = 
	    new AbstractValue[maxStack];
	AbstractValue[] locals =
	    new AbstractValue[newMaxLocal];
	for(int i = 0; i < stackHeight; i++) {
	    opStack[i] = s3f.peek(stackHeight - 1 - i);
	}
	for(int i = 0; i < oldMaxLocal; i++) {
	    locals[i] = s3f.load(i);
	}
	return s3f.make(opStack, 
			locals, 
			stackHeight, 
			s3f.getHeap());
    }

    private void localSplitInState() {
	for(int pc = 0; pc < codeLength; pc++) {
	    State.Iterator it = state.iterator(pc);
	    for (Frame f = it.currentFrame(); 
		 f != null; 
		 f = it.nextFrame()) {
		f = copyAndIncreaseMaxLocal(f);
		localSplitInFrame(f);
		it.replaceFrame(f);
	    }
	}
    }

    private void localSplitInFrame(Frame f) {
	for(int i = 0; i < local_split_map.length; i++) {
	    int map = local_split_map[i];
	    if (map == -1)
		continue;
	    AbstractValue av = f.load(i);
	    if (av == null)
		continue;
	    SimpleJITValue sjv = 
		ephemeralTypeMap.map((SimpleJITValue)av);
	    if (sjv instanceof SimpleJITReference) {
		f.store(i, null);
		f.store(map, av);
	    }
	}
    }

    public String getResultsAsString() {
	String r = "LOCAL_SPLIT_MAP : ";
	for(int i = 0; i < local_split_map.length; i++)
	    r += "[" + i + "] " + local_split_map[i] + " ";
	r += "\n";
	r += "REFMAP : \n";
	for(int i = 0; i < refMaps.length; i++)
	    if (refMaps[i] != null)
		r += "PC " + i + " "
		    + refMaps[i].toString() + "\n";
	return r;
    }

    public ReferenceMap[] getSplitReferenceMaps() {
	return refMaps;
    }

    public int[] getLocalSplitMapper() {
	return local_split_map;
    }

    private void localSplitInByteCode() {
	if (refMaps == null)
	    throw new Error();
	S3ConstantsBuilder rcpb =
	    (S3ConstantsBuilder)
	    ((S3Constants)originalByteCode.getConstantPool())
	    .getBuilder();
	CodeFragmentEditor cfe = 
	    new CodeFragmentEditor(originalByteCode, rcpb);
	InstructionBuffer codebuf = cfe.getOriginalCode();
	codebuf.rewind();
	BytecodeLocalSplitter splitter = 
	    new BytecodeLocalSplitter(local_split_map);
	splitter.beginEditing(codebuf, cfe);
	try {
	    while(codebuf.hasRemaining()) {
		splitter.visitAppropriate(codebuf.get());
	    }
	} finally {
	    splitter.endEditing();
	}
	S3ByteCode.Builder builder = 
	    new S3ByteCode.Builder(originalByteCode);
	cfe.commit(builder, 
		   builder.getMaxStack(), 
		   (char)newMaxLocal);
	builder.setConstantPool(rcpb.build());
	rewrittenByteCode = builder.build();
    }

    private void computeSplitReferenceMaps() {
	ReferenceMap[] rmaps = new ReferenceMap[codeLength];
	for(int pc = 0; pc < codeLength; pc++) {
	    State.Iterator it = state.iterator(pc);
	    int frameCount = state.getFrameCountAt(pc);
	    if (frameCount == 0)
		continue;
	    Frame[] frames = new Frame[frameCount];
	    int i = 0;
	    for (Frame f = it.currentFrame(); 
		 f != null; 
		 f = it.nextFrame()) {
		frames[i++] = f;
	    }
	    rmaps[pc] = makeSplitReferenceMap(frames, pc);
	}
	// Adjust the size of local_map arrays
	for(int pc = 0; pc < codeLength; pc++) {
	    if (rmaps[pc] == null)
		continue;
	    boolean[] local_map = rmaps[pc].local_map;
	    if (local_map.length > newMaxLocal) {
		boolean[] new_local_map =
		    new boolean[newMaxLocal];
		for(int i = 0; i < newMaxLocal; i++) {
		    new_local_map[i] = local_map[i];
		}
		rmaps[pc].local_map = new_local_map;
	    } else if (local_map.length < newMaxLocal) {
		boolean[] new_local_map =
		    new boolean[newMaxLocal];
		for(int i = 0; i < local_map.length; i++) {
		    new_local_map[i] = local_map[i];
		}
		rmaps[pc].local_map = new_local_map;
	    }
	}
	refMaps = rmaps;
    }

    private ReferenceMap makeSplitReferenceMap(Frame[] frames, 
					       int pc) {
	int frameCount = frames.length;
	if (frameCount < 1)
	    throw new Error();
	int stackHeight = frames[0].getStackHeight();
	boolean[] stack_map = new boolean[stackHeight];
	int maxlocal = originalByteCode.getMaxLocals();
	boolean[] local_map;
	if (frameCount == 1) {
	    local_map = new boolean[maxlocal];
	} else {
	    local_map = new boolean[maxlocal * 2];
	}
	// Initialize local_map with frame 0
	for(int j = 0; j < maxlocal; j++) {
	    AbstractValue av = frames[0].load(j);
	    if (av == null)
		continue;
	    SimpleJITValue sjv =
		ephemeralTypeMap.map((SimpleJITValue)av);
	    if (sjv instanceof SimpleJITReference)
		local_map[j] = true;
	}
	// Initialize stack_map with frame 0
	for(int j = 0; j < stackHeight; j++) {
	    AbstractValue av = frames[0].peek(stackHeight - 1 - j);
	    if (av == null)
		throw new Error("Stack should not have null");
	    SimpleJITValue sjv =
		ephemeralTypeMap.map((SimpleJITValue)av);
	    if (sjv instanceof SimpleJITReference) {
		stack_map[j] = true;
	    } else if (sjv instanceof SimpleJITLong
		       || sjv instanceof SimpleJITDouble)
		j++; // skip null
	}
	// Check the consistency of stack
	for(int i = 1; i < frameCount; i++) {
	    if (frames[i].getStackHeight() != stackHeight)
		throw new Error("Wrong stack height");
	    for(int j = 0; j < stackHeight; j++) {
		AbstractValue av = frames[i].peek(stackHeight - 1 - j);
		if (av == null)
		    throw new Error
			("Stack should not have null");
		SimpleJITValue sjv =
		    ephemeralTypeMap.map((SimpleJITValue)av);
		if (sjv instanceof SimpleJITReference) {
		    if (stack_map[j] == false)
			throw new Error
			    ("Stack should not have " +
			     "a mix of ref/non-ref values");
		} else if (sjv instanceof SimpleJITLong
			   || sjv instanceof SimpleJITDouble)
		    j++; // skip null
	    }
	}
	// update local_map and stack_map by other frames if more than
	// one frame
	for(int i = 1; i < frameCount; i++) {
	    for(int j = 0; j < maxlocal; j++) {
		Frame f = frames[i];
		AbstractValue av = f.load(j);
		if (local_split_map[j] != -1)
		    continue;
		if (av == null) { // null case
		    if (local_map[j]) {
			local_map[j] = false;
			local_map[newMaxLocal] = true;
			local_split_map[j] = newMaxLocal++;
		    }
		} else {
		    SimpleJITValue sjv =
			ephemeralTypeMap.map((SimpleJITValue)av);
		    if (sjv instanceof SimpleJITReference) {
			if (! local_map[j]) {
			    local_map[newMaxLocal] = true;
			    local_split_map[j] = newMaxLocal++;
			}
		    } else { // non-ref case
			if (local_map[j]) {
			    local_map[j] = false;
			    local_map[newMaxLocal] = true;
			    local_split_map[j] = newMaxLocal++;
			}
		    }
		}
	    }
	}
	return new ReferenceMap(local_map, stack_map);
    }

}
