1 files changed, 150 insertions, 0 deletions

diff --git a/src/cuchaz/enigma/bytecode/CheckCastIterator.java b/src/cuchaz/enigma/bytecode/CheckCastIterator.java
new file mode 100644
index 0000000..7ed5d7f
--- /dev/null
+++ b/src/cuchaz/enigma/bytecode/CheckCastIterator.java
@@ -0,0 +1,150 @@
+/*******************************************************************************
+ * 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.CodeAttribute;
+import javassist.bytecode.CodeIterator;
+import javassist.bytecode.ConstPool;
+import javassist.bytecode.Descriptor;
+import javassist.bytecode.Opcode;
+import cuchaz.enigma.bytecode.CheckCastIterator.CheckCast;
+import cuchaz.enigma.mapping.ClassEntry;
+import cuchaz.enigma.mapping.MethodEntry;
+
+public class CheckCastIterator implements Iterator<CheckCast>
+{
+        public static class CheckCast
+        {
+                public String className;
+                public MethodEntry prevMethodEntry;
+                
+                public CheckCast( String className, MethodEntry prevMethodEntry )
+                {
+                        this.className = className;
+                        this.prevMethodEntry = prevMethodEntry;
+                }
+        }
+        
+        private ConstPool m_constants;
+        private CodeAttribute m_attribute;
+        private CodeIterator m_iter;
+        private CheckCast m_next;
+        
+        public CheckCastIterator( CodeAttribute codeAttribute )
+        throws BadBytecode
+        {
+                m_constants = codeAttribute.getConstPool();
+                m_attribute = codeAttribute;
+                m_iter = m_attribute.iterator();
+                
+                m_next = getNext();
+        }
+        
+        @Override
+        public boolean hasNext( )
+        {
+                return m_next != null;
+        }
+
+        @Override
+        public CheckCast next( )
+        {
+                CheckCast out = m_next;
+                try
+                {
+                        m_next = getNext();
+                }
+                catch( BadBytecode ex )
+                {
+                        throw new Error( ex );
+                }
+                return out;
+        }
+
+        @Override
+        public void remove( )
+        {
+                throw new UnsupportedOperationException();
+        }
+        
+        private CheckCast getNext( )
+        throws BadBytecode
+        {
+                int prevPos = 0;
+                while( m_iter.hasNext() )
+                {
+                        int pos = m_iter.next();
+                        int opcode = m_iter.byteAt( pos );
+                        switch( opcode )
+                        {
+                                case Opcode.CHECKCAST:
+                                        
+                                        // get the type of this op code (next two bytes are a classinfo index)
+                                        MethodEntry prevMethodEntry = getMethodEntry( prevPos );
+                                        if( prevMethodEntry != null )
+                                        {
+                                                return new CheckCast(
+                                                        m_constants.getClassInfo( m_iter.s16bitAt( pos + 1 ) ),
+                                                        prevMethodEntry
+                                                );
+                                        }
+                                break;
+                        }
+                        prevPos = pos;
+                }
+                return null;
+        }
+        
+        private MethodEntry getMethodEntry( int pos )
+        {
+                switch( m_iter.byteAt( pos ) )
+                {
+                        case Opcode.INVOKEVIRTUAL:
+                        case Opcode.INVOKESTATIC:
+                        case Opcode.INVOKEDYNAMIC:
+                        case Opcode.INVOKESPECIAL:
+                        {
+                                int index = m_iter.s16bitAt( pos + 1 );
+                                return new MethodEntry(
+                                        new ClassEntry( Descriptor.toJvmName( m_constants.getMethodrefClassName( index ) ) ),
+                                        m_constants.getMethodrefName( index ),
+                                        m_constants.getMethodrefType( index )
+                                );
+                        }
+                        
+                        case Opcode.INVOKEINTERFACE:
+                        {
+                                int index = m_iter.s16bitAt( pos + 1 );
+                                return new MethodEntry(
+                                        new ClassEntry( Descriptor.toJvmName( m_constants.getInterfaceMethodrefClassName( index ) ) ),
+                                        m_constants.getInterfaceMethodrefName( index ),
+                                        m_constants.getInterfaceMethodrefType( index )
+                                );
+                        }
+                }
+                return null;
+        }
+
+        public Iterable<CheckCast> casts( )
+        {
+                return new Iterable<CheckCast>( )
+                {
+                        @Override
+                        public Iterator<CheckCast> iterator( )
+                        {
+                                return CheckCastIterator.this;
+                        }
+                };
+        }
+}