/**
 * Copyright (c) 2003-2004 Craig Setera
 * All Rights Reserved.
 * Licensed under the Eclipse Public License - v 1.0
 * For more information see http://www.eclipse.org/legal/epl-v10.html
 */
package eclipseme.preverifier.internal;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.objectweb.asm.Label;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.attrs.StackMapAttribute;
import org.objectweb.asm.attrs.StackMapFrame;
import org.objectweb.asm.attrs.StackMapType;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.InsnNode;
import org.objectweb.asm.tree.JumpInsnNode;
import org.objectweb.asm.tree.LabelNode;
import org.objectweb.asm.tree.LineNumberNode;
import org.objectweb.asm.tree.LocalVariableNode;
import org.objectweb.asm.tree.LookupSwitchInsnNode;
import org.objectweb.asm.tree.MethodNode;
import org.objectweb.asm.tree.TableSwitchInsnNode;
import org.objectweb.asm.tree.TryCatchBlockNode;
import org.objectweb.asm.tree.VarInsnNode;
import org.objectweb.asm.tree.analysis.Analyzer;
import org.objectweb.asm.tree.analysis.AnalyzerException;
import org.objectweb.asm.tree.analysis.BasicValue;
import org.objectweb.asm.tree.analysis.Frame;
import org.objectweb.asm.tree.analysis.Interpreter;
import org.objectweb.asm.tree.analysis.Value;

/**
 * Handler for a single method in the class.  Capable of 
 * inlining subroutines and creating the Stack map attribute.
 * <p />
 * Copyright (c) 2003-2005 Craig Setera<br>
 * All Rights Reserved.<br>
 * Licensed under the Eclipse Public License - v 1.0<p/>
 * <br>
 * $Revision: 1.7 $
 * <br>
 * $Date: 2005/11/15 00:34:26 $
 * <br>
 * @author Craig Setera
 */
public class MethodRewriter {
	/** A Label instance mapped to the new region and code offset after inlining */
	public class MappedLabel extends Label {
		private Label originalLabel;
		
		/**
		 * Construct a new instance
		 */
		public MappedLabel(Label label) {
			super();
			originalLabel = label;
		}

		public Label getOriginalLabel() {
			return originalLabel;
		}
		
		/**
		 * @see java.lang.Object#toString()
		 */
		public String toString() {
			return originalLabel.toString() + " -> " + super.toString();
		}
	}

	/** A region of instructions to be handled.  */
	class Region {
		protected Region parentRegion;
		protected int startIndex;
		protected int endIndex;
		protected Map labelMap;
		protected Set labels;
		protected List tryCatchBlocks;
		
		/**
		 * Construct a new instance. 
		 */
		Region() {
			this(0, 0);
		}

		/**
		 * Construct a new instance. 
		 * 
		 * @param startIndex
		 * @param endIndex
		 */
		Region(int startIndex, int endIndex) {
			this.startIndex = startIndex;
			this.endIndex = endIndex;
			
			labelMap = new HashMap();
			labels = new HashSet();
			tryCatchBlocks = new ArrayList();
		}
		
		/**
		 * Add a new label to this region's list.
		 * 
		 * @param label
		 */
		void addLabel(Label label) {
			labels.add(label);
		}
		
		/**
		 * Add a new try catch block that is contained within
		 * this region.
		 * 
		 * @param tryCatchBlock
		 */
		void addTryCatchBlock(TryCatchBlockNode tryCatchBlock) {
			tryCatchBlocks.add(tryCatchBlock);
		}
		
		/**
		 * Return a boolean indicating whether the specified 
		 * TryCatchBlockNode is enclosed within this instruction
		 * region.
		 * 
		 * @param tryCatchBlock
		 * @return
		 */
		boolean encloses(TryCatchBlockNode tryCatchBlock) {
			return
				labels.contains(tryCatchBlock.start) &&
				labels.contains(tryCatchBlock.end);
		}
		
		/**
		 * Enter the region.  Do any setup for this region.
		 * @param method 
		 */
		void enter(MethodNode method) {
			labelMap.clear();
			copyTryCatchBlocks(method, this, tryCatchBlocks);
		}
		
		/**
		 * Exit the region.  Do any cleanup for this region.
		 * @param method 
		 */
		void exit(MethodNode method) {
		}

		/**
		 * Return the length (in AbstractInsnNode instances) of
		 * this region.
		 * 
		 * @return
		 */
		int getLength() {
			return endIndex - startIndex;
		}
		
		/**
		 * Get the mapped label for the specified label based on the
		 * current region.
		 * 
		 * @param originalLabel
		 * @return
		 */
		Label getMappedLabel(Label originalLabel) {
			Label mappedLabel = findMappedLabel(originalLabel);
			if (mappedLabel == null) {
				mappedLabel = new MappedLabel(originalLabel);
				
				Map map = findLabelMap(originalLabel);
				map.put(originalLabel, mappedLabel);
			}
			
			return mappedLabel;
		}
		
		/**
		 * Returnn a boolean indicating whether the specified opcode
		 * at the specified index in the method is a subroutine store
		 * opcode.
		 * 
		 * @param methodNode
		 * @param index
		 * @return
		 */
		boolean isSubroutineReturnStore(MethodNode methodNode, int index) {
			return false;
		}
		
		/**
		 * Store the parent region.
		 * 
		 * @param parentRegion
		 */
		void setParentRegion(Region parentRegion) {
			this.parentRegion = parentRegion;
		}
		
		/**
		 * Find the label map recursively as necessary.
		 * 
		 * @param originalLabel
		 * @return
		 */
		protected Map findLabelMap(Label originalLabel) {
			Map map = findLabelMapRecursive(originalLabel);
			return (map != null) ? map : labelMap;
		}
		
		/**
		 * Find the label map recursively.
		 * 
		 * @param originalLabel
		 * @return
		 */
		protected Map findLabelMapRecursive(Label originalLabel) {
			Map map = null;
			
			if (labels.contains(originalLabel)) {
				map = labelMap;
			} else if (parentRegion != null) {
				map = parentRegion.findLabelMapRecursive(originalLabel);
			}
			
			return map;
		}
		
		/**
		 * Find and return the specified label mapped appropriately for the
		 * region or <code>null</code> if not found.
		 * 
		 * @param originalLabel
		 * @return
		 */
		protected Label findMappedLabel(Label originalLabel) {
			Label mappedLabel = (Label) labelMap.get(originalLabel);
			if ((mappedLabel == null) && (parentRegion != null)) {
				mappedLabel = parentRegion.findMappedLabel(originalLabel);
			}
			
			return mappedLabel;
		}
		
		/**
		 * Return the try catch blocks in this region.
		 * 
		 * @return
		 */
		List getTryCatchBlocks() {
			return tryCatchBlocks;
		}
	}
	
	/** Holder for information about a subroutine in a method */
	class Subroutine extends Region {
		Label label;
		int returnVariable;
		
		/**
		 * Construct a new subroutine instance.
		 * 
		 * @param label
		 */
		Subroutine(Label label) {
			super();
			this.label = label;
		}
		
		/**
		 * @return Returns the returnVariable.
		 */
		public int getReturnVariable() {
			return returnVariable;
		}
		
		/**
		 * @see eclipseme.preverifier.internal.MethodRewriter.Region#isSubroutineReturnStore(org.objectweb.asm.tree.MethodNode, int)
		 */
		boolean isSubroutineReturnStore(MethodNode methodNode, int index) {
			boolean isReturnStore = false;
			
			AbstractInsnNode insnNode = getInstruction(methodNode, index);
			if (insnNode.getOpcode() == Opcodes.ASTORE) {
				VarInsnNode varInsnNode = (VarInsnNode) insnNode;
				isReturnStore = (varInsnNode.var == returnVariable);
			}
			
			return isReturnStore;
		}

		/**
		 * @param returnVariable The returnVariable to set.
		 */
		public void setReturnVariable(int returnVariable) {
			this.returnVariable = returnVariable;
		}
	}

	private PreverificationClassNode classNode;
	private PreverifierMethodNode srcMethod;
	private MethodNode updatedMethod;
	
	private Map subroutineMap;
	private Map lineNumberMap;
	private Map localVariableByStartLabelMap;
	private Map localVariableByEndLabelMap;
	
	/**
	 * Construct a new rewriter instance.
	 * 
	 * @param classNode
	 * @param srcMethod
	 */
	public MethodRewriter(PreverificationClassNode classNode, PreverifierMethodNode srcMethod)
	{
		super();
		
		this.classNode = classNode;
		this.srcMethod = srcMethod;
		localVariableByEndLabelMap = new HashMap();
	}

	/**
	 * Return the method with subroutines inlined and
	 * an associated StackMapAttribute.
	 * 
	 * @return
	 * @throws AnalyzerException
	 */
	public MethodNode getUpdatedMethod() 
		throws AnalyzerException 
	{
		boolean inliningRequired = (srcMethod.getJsrInstructionIndices().size() > 0); 
		if (inliningRequired) {
			inlineSubroutines();
		} else {
			updatedMethod = srcMethod;
		}
		
		createStackMapAttribute();
		return updatedMethod;
	}

	/**
	 * Create a new StackMapAttribute for the method.
	 * This method also removes any dead code that would
	 * cause the stack map attribute to be incorrect.
	 * 
	 * @throws AnalyzerException 
	 */
	private void createStackMapAttribute() 
		throws AnalyzerException 
	{
		StackMapAttribute stackMapAttribute = new StackMapAttribute();
		
		Set targetLabels = findTargetLabels();
		
		// We need to have the verifier operating with the classpath
		// of the class being rewritten
		Interpreter interpreter = 
			new SimpleVerifierPlusClassloader(classNode.getClassLoader());
		Analyzer analyzer = new Analyzer(interpreter);
		Frame[] frames = analyzer.analyze(classNode.name, updatedMethod);
		
		int deadCodeFrameCount = 0;
		int[] deadCodeIndices = new int[frames.length];
		
		for (int i = 0; i < updatedMethod.instructions.size(); i++) {
			Frame frame = frames[i];

			// We only need to add stack map attributes for labels
			// that are the target of certain instructions
			AbstractInsnNode insnNode = getInstruction(updatedMethod, i); 
			if (insnNode.getType() == AbstractInsnNode.LABEL) {
				LabelNode labelNode = (LabelNode) insnNode;
			
				if (targetLabels.contains(labelNode.label) && (frame != null)) {
					stackMapAttribute.frames.add(newStackMapFrame(labelNode.label, frame));
				}
			} else if (frame == null) {
				// Track the indices of dead code for removal after
				// building the stack map
				deadCodeIndices[deadCodeFrameCount++] = i;
			}
			
		}
		
		// Write out the newly created stack map attribute
		// and remove any unnecessary dead code that would
		// have an adverse effect on the stack map attribute
		if (stackMapAttribute.frames.size() != 0) {
			updatedMethod.visitAttribute(stackMapAttribute);
			
			// Remove the dead code in reverse order so that
			// the indices remain correct after each removal
			for (int i = deadCodeFrameCount - 1; i >= 0; i--) {
				updatedMethod.instructions.remove(deadCodeIndices[i]);
			}
		}
	}

	/**
	 * Inline all subroutines.
	 * 
	 * @throws AnalyzerException