/*******************************************************************************
 * Copyright (c) 2014 Jeff Martin.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the GNU Public License v3.0
 * which accompanies this distribution, and is available at
 * http://www.gnu.org/licenses/gpl.html
 * 
 * Contributors:
 *     Jeff Martin - initial API and implementation
 ******************************************************************************/
package cuchaz.enigma.bytecode;

import java.util.Iterator;

import javassist.bytecode.BadBytecode;
import javassist.bytecode.Bytecode;
import javassist.bytecode.CodeAttribute;
import javassist.bytecode.CodeIterator;
import javassist.bytecode.Opcode;

public class BytecodeIndexIterator implements Iterator<BytecodeIndexIterator.Index>
{
	public static class Index
	{
		private CodeIterator m_iter;
		private int m_pos;
		private boolean m_isWide;
		
		protected Index( CodeIterator iter, int pos, boolean isWide )
		{
			m_iter = iter;
			m_pos = pos;
			m_isWide = isWide;
		}
		
		public int getIndex( )
		{
			if( m_isWide )
			{
				return m_iter.s16bitAt( m_pos );
			}
			else
			{
				return m_iter.byteAt( m_pos );
			}
		}
		
		public void setIndex( int val )
		throws BadBytecode
		{
			if( m_isWide )
			{
				m_iter.write16bit( val, m_pos );
			}
			else
			{
				if( val < 256 )
				{
					// we can write the byte
					m_iter.writeByte( val, m_pos );
				}
				else
				{
					// we need to upgrade this instruction to LDC_W
					assert( m_iter.byteAt( m_pos - 1 ) == Opcode.LDC );
					m_iter.insertGap( m_pos - 1, 1 );
					m_iter.writeByte( Opcode.LDC_W, m_pos - 1 );
					m_iter.write16bit( val, m_pos );
					m_isWide = true;
					
					// move the iterator to the next opcode
					m_iter.move( m_pos + 2 );
				}
			}
			
			// sanity check
			assert( val == getIndex() );
		}
		
		public boolean isValid( Bytecode bytecode )
		{
			return getIndex() >= 0 && getIndex() < bytecode.getConstPool().getSize();
		}
	}
	
	private Bytecode m_bytecode;
	private CodeAttribute m_attribute;
	private CodeIterator m_iter;
	private Index m_next;
	
	public BytecodeIndexIterator( Bytecode bytecode )
	throws BadBytecode
	{
		m_bytecode = bytecode;
		m_attribute = bytecode.toCodeAttribute();
		m_iter = m_attribute.iterator();
		
		m_next = getNext();
	}
	
	@Override
	public boolean hasNext( )
	{
		return m_next != null;
	}

	@Override
	public Index next( )
	{
		Index out = m_next;
		try
		{
			m_next = getNext();
		}
		catch( BadBytecode ex )
		{
			throw new Error( ex );
		}
		return out;
	}

	@Override
	public void remove( )
	{
		throw new UnsupportedOperationException();
	}
	
	private Index getNext( )
	throws BadBytecode
	{
		while( m_iter.hasNext() )
		{
			int pos = m_iter.next();
			int opcode = m_iter.byteAt( pos );
			switch( opcode )
			{
				// for only these opcodes, the next two bytes are a const pool reference
				case Opcode.ANEWARRAY:
				case Opcode.CHECKCAST:
				case Opcode.INSTANCEOF:
				case Opcode.INVOKEDYNAMIC:
				case Opcode.INVOKEINTERFACE:
				case Opcode.INVOKESPECIAL:
				case Opcode.INVOKESTATIC:
				case Opcode.INVOKEVIRTUAL:
				case Opcode.LDC_W:
				case Opcode.LDC2_W:
				case Opcode.MULTIANEWARRAY:
				case Opcode.NEW:
				case Opcode.PUTFIELD:
				case Opcode.PUTSTATIC:
				case Opcode.GETFIELD:
				case Opcode.GETSTATIC:
					return new Index( m_iter, pos + 1, true );
				
				case Opcode.LDC:
					return new Index( m_iter, pos + 1, false );
			}
		}
		
		return null;
	}
	
	public Iterable<Index> indices( )
	{
		return new Iterable<Index>( )
		{
			@Override
			public Iterator<Index> iterator( )
			{
				return BytecodeIndexIterator.this;
			}
		};
	}
	
	public void saveChangesToBytecode( )
	{
		BytecodeTools.setBytecode( m_bytecode, m_attribute.getCode() );
	}
}