Index more class references (#531)

* Index more class references

* fix checkstyle

* index checkcast as well

5 files changed, 45 insertions, 1 deletions

diff --git a/enigma/src/main/java/cuchaz/enigma/analysis/index/IndexReferenceVisitor.java b/enigma/src/main/java/cuchaz/enigma/analysis/index/IndexReferenceVisitor.java
index 97fec47b..cb71275a 100644
--- a/enigma/src/main/java/cuchaz/enigma/analysis/index/IndexReferenceVisitor.java
+++ b/enigma/src/main/java/cuchaz/enigma/analysis/index/IndexReferenceVisitor.java
@@ -10,7 +10,9 @@ import org.objectweb.asm.Type;
 import org.objectweb.asm.tree.AbstractInsnNode;
 import org.objectweb.asm.tree.FieldInsnNode;
 import org.objectweb.asm.tree.InvokeDynamicInsnNode;
+import org.objectweb.asm.tree.LdcInsnNode;
 import org.objectweb.asm.tree.MethodInsnNode;
+import org.objectweb.asm.tree.TypeInsnNode;
 import org.objectweb.asm.tree.analysis.Analyzer;
 import org.objectweb.asm.tree.analysis.AnalyzerException;
 import org.objectweb.asm.tree.analysis.BasicValue;
@@ -65,7 +67,7 @@ public class IndexReferenceVisitor extends ClassVisitor {
 
         private static class MethodInterpreter extends InterpreterPair<BasicValue, SourceValue> {
                 private final MethodDefEntry callerEntry;
-                private JarIndexer indexer;
+                private final JarIndexer indexer;
 
                 MethodInterpreter(MethodDefEntry callerEntry, JarIndexer indexer, EntryIndex entryIndex, InheritanceIndex inheritanceIndex) {
                         super(new IndexSimpleVerifier(entryIndex, inheritanceIndex), new SourceInterpreter());
@@ -80,6 +82,15 @@ public class IndexReferenceVisitor extends ClassVisitor {
                                 indexer.indexFieldReference(callerEntry, FieldEntry.parse(field.owner, field.name, field.desc), ReferenceTargetType.none());
                         }
 
+                        if (insn.getOpcode() == Opcodes.LDC) {
+                                LdcInsnNode ldc = (LdcInsnNode) insn;
+
+                                if (ldc.getType() == Type.ARRAY && ldc.cst instanceof Type type) {
+                                        String className = type.getClassName().replace(".", "/");
+                                        indexer.indexClassReference(callerEntry, ClassEntry.parse(className), ReferenceTargetType.none());
+                                }
+                        }
+
                         return super.newOperation(insn);
                 }
 
@@ -95,6 +106,17 @@ public class IndexReferenceVisitor extends ClassVisitor {
                                 indexer.indexFieldReference(callerEntry, FieldEntry.parse(field.owner, field.name, field.desc), getReferenceTargetType(value, insn));
                         }
 
+                        if (insn.getOpcode() == Opcodes.INSTANCEOF) {
+                                TypeInsnNode type = (TypeInsnNode) insn;
+                                // Note: type.desc is actually the name
+                                indexer.indexClassReference(callerEntry, ClassEntry.parse(type.desc), ReferenceTargetType.none());
+                        }
+
+                        if (insn.getOpcode() == Opcodes.CHECKCAST) {
+                                TypeInsnNode type = (TypeInsnNode) insn;
+                                indexer.indexClassReference(callerEntry, ClassEntry.parse(type.desc), ReferenceTargetType.none());
+                        }
+
                         return super.unaryOperation(insn, value);
                 }
 
diff --git a/enigma/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java b/enigma/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
index 336ad58c..967983a8 100644
--- a/enigma/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
+++ b/enigma/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
@@ -153,6 +153,15 @@ public class JarIndex implements JarIndexer {
         }
 
         @Override
+        public void indexClassReference(MethodDefEntry callerEntry, ClassEntry referencedEntry, ReferenceTargetType targetType) {
+                if (callerEntry.getParent().isJre()) {
+                        return;
+                }
+
+                indexers.forEach(indexer -> indexer.indexClassReference(callerEntry, referencedEntry, targetType));
+        }
+
+        @Override
         public void indexMethodReference(MethodDefEntry callerEntry, MethodEntry referencedEntry, ReferenceTargetType targetType) {
                 if (callerEntry.getParent().isJre()) {
                         return;
diff --git a/enigma/src/main/java/cuchaz/enigma/analysis/index/JarIndexer.java b/enigma/src/main/java/cuchaz/enigma/analysis/index/JarIndexer.java
index 8726fb56..066000d0 100644
--- a/enigma/src/main/java/cuchaz/enigma/analysis/index/JarIndexer.java
+++ b/enigma/src/main/java/cuchaz/enigma/analysis/index/JarIndexer.java
@@ -3,6 +3,7 @@ package cuchaz.enigma.analysis.index;
 import cuchaz.enigma.analysis.ReferenceTargetType;
 import cuchaz.enigma.translation.representation.Lambda;
 import cuchaz.enigma.translation.representation.entry.ClassDefEntry;
+import cuchaz.enigma.translation.representation.entry.ClassEntry;
 import cuchaz.enigma.translation.representation.entry.FieldDefEntry;
 import cuchaz.enigma.translation.representation.entry.FieldEntry;
 import cuchaz.enigma.translation.representation.entry.MethodDefEntry;
@@ -18,6 +19,9 @@ public interface JarIndexer {
         default void indexMethod(MethodDefEntry methodEntry) {
         }
 
+        default void indexClassReference(MethodDefEntry callerEntry, ClassEntry referencedEntry, ReferenceTargetType targetType) {
+        }
+
         default void indexMethodReference(MethodDefEntry callerEntry, MethodEntry referencedEntry, ReferenceTargetType targetType) {
         }
 
diff --git a/enigma/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java b/enigma/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java
index 332a9674..c55d36ad 100644
--- a/enigma/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java
+++ b/enigma/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java
@@ -65,6 +65,11 @@ public class ReferenceIndex implements JarIndexer {
         }
 
         @Override
+        public void indexClassReference(MethodDefEntry callerEntry, ClassEntry referencedEntry, ReferenceTargetType targetType) {
+                referencesToClasses.put(referencedEntry, new EntryReference<>(referencedEntry, referencedEntry.getName(), callerEntry, targetType));
+        }
+
+        @Override
         public void indexMethodReference(MethodDefEntry callerEntry, MethodEntry referencedEntry, ReferenceTargetType targetType) {
                 referencesToMethods.put(referencedEntry, new EntryReference<>(referencedEntry, referencedEntry.getName(), callerEntry, targetType));
                 methodReferences.put(callerEntry, referencedEntry);
diff --git a/enigma/src/main/java/cuchaz/enigma/translation/representation/entry/ClassEntry.java b/enigma/src/main/java/cuchaz/enigma/translation/representation/entry/ClassEntry.java
index 8659b402..0e179d7f 100644
--- a/enigma/src/main/java/cuchaz/enigma/translation/representation/entry/ClassEntry.java
+++ b/enigma/src/main/java/cuchaz/enigma/translation/representation/entry/ClassEntry.java
@@ -50,6 +50,10 @@ public class ClassEntry extends ParentedEntry<ClassEntry> implements Comparable<
                 }
         }
 
+        public static ClassEntry parse(String name) {
+                return new ClassEntry(name);
+        }
+
         @Override
         public Class<ClassEntry> getParentType() {
                 return ClassEntry.class;