1 files changed, 126 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..b6efbd4
--- /dev/null
+++ b/src/cuchaz/enigma/bytecode/CheckCastIterator.java
@@ -0,0 +1,126 @@
+/*******************************************************************************
+ * 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;
+                        }
+                };
+        }
+}