checkmappings command (#137)

* Use expected map sizes for remapped multimaps

* Index method and field types

* Add package visibility index

* Add checkmappings command and use System.err for error messages

* Use exit codes for errors

* Remove outer class check for package visible only refs

* Throw exception on mapping error instead of exiting

5 files changed, 247 insertions, 15 deletions

diff --git a/src/main/java/cuchaz/enigma/CommandMain.java b/src/main/java/cuchaz/enigma/CommandMain.java
index c9f83828..db4fd125 100644
--- a/src/main/java/cuchaz/enigma/CommandMain.java
+++ b/src/main/java/cuchaz/enigma/CommandMain.java
@@ -11,35 +11,47 @@
 
 package cuchaz.enigma;
 
+import cuchaz.enigma.analysis.index.JarIndex;
 import cuchaz.enigma.translation.mapping.EntryMapping;
 import cuchaz.enigma.translation.mapping.serde.MappingFormat;
 import cuchaz.enigma.translation.mapping.tree.EntryTree;
+import cuchaz.enigma.translation.representation.entry.ClassEntry;
 
 import java.io.File;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.Locale;
+import java.util.Set;
 import java.util.jar.JarFile;
+import java.util.stream.Collectors;
 
 public class CommandMain {
 
         public static void main(String[] args) throws Exception {
                 try {
                         // process the command
-                        String command = getArg(args, 0, "command", true);
-                        if (command.equalsIgnoreCase("deobfuscate")) {
-                                deobfuscate(args);
-                        } else if (command.equalsIgnoreCase("decompile")) {
-                                decompile(args);
-                        } else if (command.equalsIgnoreCase("convertmappings")) {
-                                convertMappings(args);
-                        } else {
-                                throw new IllegalArgumentException("Command not recognized: " + command);
+                        String command = getArg(args, 0, "command", true).toLowerCase(Locale.ROOT);
+                        switch (command) {
+                                case "deobfuscate":
+                                        deobfuscate(args);
+                                        break;
+                                case "decompile":
+                                        decompile(args);
+                                        break;
+                                case "convertmappings":
+                                        convertMappings(args);
+                                        break;
+                                case "checkmappings":
+                                        checkMappings(args);
+                                        break;
+                                default:
+                                        throw new IllegalArgumentException("Command not recognized: " + command);
                         }
                 } catch (IllegalArgumentException ex) {
-                        System.out.println(ex.getMessage());
+                        System.err.println(ex.getMessage());
                         printHelp();
+                        System.exit(1);
                 }
         }
 
@@ -51,6 +63,7 @@ public class CommandMain {
                 System.out.println("\t\tdeobfuscate <in jar> <out jar> [<mappings file>]");
                 System.out.println("\t\tdecompile <in jar> <out folder> [<mappings file>]");
                 System.out.println("\t\tconvertmappings <enigma mappings> <converted mappings> <ENIGMA_FILE|ENIGMA_DIRECTORY|SRG_FILE>");
+                System.out.println("\t\tcheckmappings <in jar> <mappings file>");
         }
 
         private static void decompile(String[] args) throws Exception {
@@ -100,6 +113,35 @@ public class CommandMain {
                 saveFormat.write(mappings, result.toPath(), new ConsoleProgressListener());
         }
 
+        private static void checkMappings(String[] args) throws Exception {
+                File fileJarIn = getReadableFile(getArg(args, 1, "in jar", true));
+                Path fileMappings = getReadablePath(getArg(args, 2, "enigma mapping", true));
+
+                System.out.println("Reading JAR...");
+                Deobfuscator deobfuscator = new Deobfuscator(new JarFile(fileJarIn));
+                System.out.println("Reading mappings...");
+
+                MappingFormat format = chooseEnigmaFormat(fileMappings);
+                EntryTree<EntryMapping> mappings = format.read(fileMappings, ProgressListener.VOID);
+                deobfuscator.setMappings(mappings);
+
+                JarIndex idx = deobfuscator.getJarIndex();
+
+                boolean error = false;
+
+                for (Set<ClassEntry> partition : idx.getPackageVisibilityIndex().getPartitions()) {
+                        long packages = partition.stream().map(deobfuscator.getMapper()::deobfuscate).map(ClassEntry::getPackageName).distinct().count();
+                        if (packages > 1) {
+                                error = true;
+                                System.err.println("ERROR: Must be in one package:\n" + partition.stream().map(deobfuscator.getMapper()::deobfuscate).map(ClassEntry::toString).sorted().collect(Collectors.joining("\n")));
+                        }
+                }
+
+                if (error) {
+                        throw new Exception("Access violations detected");
+                }
+        }
+
         private static MappingFormat chooseEnigmaFormat(Path path) {
                 if (Files.isDirectory(path)) {
                         return MappingFormat.ENIGMA_DIRECTORY;
diff --git a/src/main/java/cuchaz/enigma/analysis/index/EntryIndex.java b/src/main/java/cuchaz/enigma/analysis/index/EntryIndex.java
index 773eaf18..31c6f54f 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/EntryIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/EntryIndex.java
@@ -65,6 +65,11 @@ public class EntryIndex implements JarIndexer {
         }
 
         @Nullable
+        public AccessFlags getClassAccess(ClassEntry entry) {
+                return classes.get(entry);
+        }
+
+        @Nullable
         public AccessFlags getEntryAccess(Entry<?> entry) {
                 if (entry instanceof MethodEntry) {
                         return getMethodAccess((MethodEntry) entry);
diff --git a/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java b/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
index a429ff6e..ac907af0 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/JarIndex.java
@@ -29,18 +29,20 @@ public class JarIndex implements JarIndexer {
         private final InheritanceIndex inheritanceIndex;
         private final ReferenceIndex referenceIndex;
         private final BridgeMethodIndex bridgeMethodIndex;
+        private final PackageVisibilityIndex packageVisibilityIndex;
         private final EntryResolver entryResolver;
 
         private final Collection<JarIndexer> indexers;
 
         private final Multimap<String, MethodDefEntry> methodImplementations = HashMultimap.create();
 
-        public JarIndex(EntryIndex entryIndex, InheritanceIndex inheritanceIndex, ReferenceIndex referenceIndex, BridgeMethodIndex bridgeMethodIndex) {
+        public JarIndex(EntryIndex entryIndex, InheritanceIndex inheritanceIndex, ReferenceIndex referenceIndex, BridgeMethodIndex bridgeMethodIndex, PackageVisibilityIndex packageVisibilityIndex) {
                 this.entryIndex = entryIndex;
                 this.inheritanceIndex = inheritanceIndex;
                 this.referenceIndex = referenceIndex;
                 this.bridgeMethodIndex = bridgeMethodIndex;
-                this.indexers = Arrays.asList(entryIndex, inheritanceIndex, referenceIndex, bridgeMethodIndex);
+                this.packageVisibilityIndex = packageVisibilityIndex;
+                this.indexers = Arrays.asList(entryIndex, inheritanceIndex, referenceIndex, bridgeMethodIndex, packageVisibilityIndex);
                 this.entryResolver = new IndexEntryResolver(this);
         }
 
@@ -49,7 +51,8 @@ public class JarIndex implements JarIndexer {
                 InheritanceIndex inheritanceIndex = new InheritanceIndex(entryIndex);
                 ReferenceIndex referenceIndex = new ReferenceIndex();
                 BridgeMethodIndex bridgeMethodIndex = new BridgeMethodIndex(entryIndex, inheritanceIndex, referenceIndex);
-                return new JarIndex(entryIndex, inheritanceIndex, referenceIndex, bridgeMethodIndex);
+                PackageVisibilityIndex packageVisibilityIndex = new PackageVisibilityIndex();
+                return new JarIndex(entryIndex, inheritanceIndex, referenceIndex, bridgeMethodIndex, packageVisibilityIndex);
         }
 
         public void indexJar(ParsedJar jar, Consumer<String> progress) {
@@ -142,6 +145,10 @@ public class JarIndex implements JarIndexer {
                 return bridgeMethodIndex;
         }
 
+        public PackageVisibilityIndex getPackageVisibilityIndex() {
+                return packageVisibilityIndex;
+        }
+
         public EntryResolver getEntryResolver() {
                 return entryResolver;
         }
diff --git a/src/main/java/cuchaz/enigma/analysis/index/PackageVisibilityIndex.java b/src/main/java/cuchaz/enigma/analysis/index/PackageVisibilityIndex.java
new file mode 100644
index 00000000..9e9115fe
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/analysis/index/PackageVisibilityIndex.java
@@ -0,0 +1,127 @@
+package cuchaz.enigma.analysis.index;

+

+import com.google.common.collect.HashMultimap;

+import com.google.common.collect.Lists;

+import com.google.common.collect.Maps;

+import com.google.common.collect.Sets;

+import cuchaz.enigma.analysis.EntryReference;

+import cuchaz.enigma.translation.representation.AccessFlags;

+import cuchaz.enigma.translation.representation.entry.*;

+

+import java.util.*;

+

+public class PackageVisibilityIndex implements JarIndexer {

+        private static boolean isPackageVisibleOnlyRef(AccessFlags entryAcc, EntryReference ref, InheritanceIndex inheritanceIndex) {

+                if (entryAcc.isPublic()) return false;

+                if (entryAcc.isProtected()) {

+                        Set<ClassEntry> callerAncestors = inheritanceIndex.getAncestors(ref.context.getContainingClass());

+                        return !callerAncestors.contains(ref.entry.getContainingClass());

+                }

+                return !entryAcc.isPrivate(); // if isPrivate is false, it must be package-private

+        }

+

+        private final HashMultimap<ClassEntry, ClassEntry> connections = HashMultimap.create();

+        private final List<Set<ClassEntry>> partitions = Lists.newArrayList();

+        private final Map<ClassEntry, Set<ClassEntry>> classPartitions = Maps.newHashMap();

+

+        private void addConnection(ClassEntry classA, ClassEntry classB) {

+                connections.put(classA, classB);

+                connections.put(classB, classA);

+        }

+

+        private void buildPartition(Set<ClassEntry> unassignedClasses, Set<ClassEntry> partition, ClassEntry member) {

+                for (ClassEntry connected : connections.get(member)) {

+                        if (unassignedClasses.remove(connected)) {

+                                partition.add(connected);

+                                buildPartition(unassignedClasses, partition, connected);

+                        }

+                }

+        }

+

+        private void addConnections(EntryIndex entryIndex, ReferenceIndex referenceIndex, InheritanceIndex inheritanceIndex) {

+                for (FieldEntry entry : entryIndex.getFields()) {

+                        AccessFlags entryAcc = entryIndex.getFieldAccess(entry);

+                        if (!entryAcc.isPublic() && !entryAcc.isPrivate()) {

+                                for (EntryReference<FieldEntry, MethodDefEntry> ref : referenceIndex.getReferencesToField(entry)) {

+                                        if (isPackageVisibleOnlyRef(entryAcc, ref, inheritanceIndex)) {

+                                                addConnection(ref.entry.getContainingClass(), ref.context.getContainingClass());

+                                        }

+                                }

+                        }

+                }

+

+                for (MethodEntry entry : entryIndex.getMethods()) {

+                        AccessFlags entryAcc = entryIndex.getMethodAccess(entry);

+                        if (!entryAcc.isPublic() && !entryAcc.isPrivate()) {

+                                for (EntryReference<MethodEntry, MethodDefEntry> ref : referenceIndex.getReferencesToMethod(entry)) {

+                                        if (isPackageVisibleOnlyRef(entryAcc, ref, inheritanceIndex)) {

+                                                addConnection(ref.entry.getContainingClass(), ref.context.getContainingClass());

+                                        }

+                                }

+                        }

+                }

+

+                for (ClassEntry entry : entryIndex.getClasses()) {

+                        AccessFlags entryAcc = entryIndex.getClassAccess(entry);

+                        if (!entryAcc.isPublic() && !entryAcc.isPrivate()) {

+                                for (EntryReference<ClassEntry, FieldDefEntry> ref : referenceIndex.getFieldTypeReferencesToClass(entry)) {

+                                        if (isPackageVisibleOnlyRef(entryAcc, ref, inheritanceIndex)) {

+                                                addConnection(ref.entry.getContainingClass(), ref.context.getContainingClass());

+                                        }

+                                }

+

+                                for (EntryReference<ClassEntry, MethodDefEntry> ref : referenceIndex.getMethodTypeReferencesToClass(entry)) {

+                                        if (isPackageVisibleOnlyRef(entryAcc, ref, inheritanceIndex)) {

+                                                addConnection(ref.entry.getContainingClass(), ref.context.getContainingClass());

+                                        }

+                                }

+                        }

+

+                        for (ClassEntry parent : inheritanceIndex.getParents(entry)) {

+                                AccessFlags parentAcc = entryIndex.getClassAccess(parent);

+                                if (parentAcc != null && !parentAcc.isPublic() && !parentAcc.isPrivate()) {

+                                        addConnection(entry, parent);

+                                }

+                        }

+

+                        ClassEntry outerClass = entry.getOuterClass();

+                        if (outerClass != null) {

+                                addConnection(entry, outerClass);

+                        }

+                }

+        }

+

+        private void addPartitions(EntryIndex entryIndex) {

+                Set<ClassEntry> unassignedClasses = Sets.newHashSet(entryIndex.getClasses());

+                while (!unassignedClasses.isEmpty()) {

+                        Iterator<ClassEntry> iterator = unassignedClasses.iterator();

+                        ClassEntry initialEntry = iterator.next();

+                        iterator.remove();

+

+                        HashSet<ClassEntry> partition = Sets.newHashSet();

+                        partition.add(initialEntry);

+                        buildPartition(unassignedClasses, partition, initialEntry);

+                        partitions.add(partition);

+                        for (ClassEntry entry : partition) {

+                                classPartitions.put(entry, partition);

+                        }

+                }

+        }

+

+        public Collection<Set<ClassEntry>> getPartitions() {

+                return partitions;

+        }

+

+        public Set<ClassEntry> getPartition(ClassEntry classEntry) {

+                return classPartitions.get(classEntry);

+        }

+

+        @Override

+        public void processIndex(JarIndex index) {

+                EntryIndex entryIndex = index.getEntryIndex();

+                ReferenceIndex referenceIndex = index.getReferenceIndex();

+                InheritanceIndex inheritanceIndex = index.getInheritanceIndex();

+                addConnections(entryIndex, referenceIndex, inheritanceIndex);

+                addPartitions(entryIndex);

+        }

+}

diff --git a/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java b/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java
index 2b63c5d4..6764ac0c 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/ReferenceIndex.java
@@ -4,6 +4,8 @@ import com.google.common.collect.HashMultimap;
 import com.google.common.collect.Multimap;
 import cuchaz.enigma.analysis.EntryReference;
 import cuchaz.enigma.translation.mapping.ResolutionStrategy;
+import cuchaz.enigma.translation.representation.MethodDescriptor;
+import cuchaz.enigma.translation.representation.TypeDescriptor;
 import cuchaz.enigma.translation.representation.entry.*;
 
 import java.util.Collection;
@@ -15,6 +17,43 @@ public class ReferenceIndex implements JarIndexer {
         private Multimap<MethodEntry, EntryReference<MethodEntry, MethodDefEntry>> referencesToMethods = HashMultimap.create();
         private Multimap<ClassEntry, EntryReference<ClassEntry, MethodDefEntry>> referencesToClasses = HashMultimap.create();
         private Multimap<FieldEntry, EntryReference<FieldEntry, MethodDefEntry>> referencesToFields = HashMultimap.create();
+        private Multimap<ClassEntry, EntryReference<ClassEntry, FieldDefEntry>> fieldTypeReferences = HashMultimap.create();
+        private Multimap<ClassEntry, EntryReference<ClassEntry, MethodDefEntry>> methodTypeReferences = HashMultimap.create();
+
+        @Override
+        public void indexMethod(MethodDefEntry methodEntry) {
+            indexMethodDescriptor(methodEntry, methodEntry.getDesc());
+        }
+
+        private void indexMethodDescriptor(MethodDefEntry entry, MethodDescriptor descriptor) {
+                for (TypeDescriptor typeDescriptor : descriptor.getArgumentDescs()) {
+                        indexMethodTypeDescriptor(entry, typeDescriptor);
+                }
+                indexMethodTypeDescriptor(entry, descriptor.getReturnDesc());
+        }
+
+        private void indexMethodTypeDescriptor(MethodDefEntry method, TypeDescriptor typeDescriptor) {
+                if (typeDescriptor.isType()) {
+                        ClassEntry referencedClass = typeDescriptor.getTypeEntry();
+                        methodTypeReferences.put(referencedClass, new EntryReference<>(referencedClass, referencedClass.getName(), method));
+                } else if (typeDescriptor.isArray()) {
+                        indexMethodTypeDescriptor(method, typeDescriptor.getArrayType());
+                }
+        }
+
+        @Override
+        public void indexField(FieldDefEntry fieldEntry) {
+            indexFieldTypeDescriptor(fieldEntry, fieldEntry.getDesc());
+        }
+
+        private void indexFieldTypeDescriptor(FieldDefEntry field, TypeDescriptor typeDescriptor) {
+                if (typeDescriptor.isType()) {
+                        ClassEntry referencedClass = typeDescriptor.getTypeEntry();
+                        fieldTypeReferences.put(referencedClass, new EntryReference<>(referencedClass, referencedClass.getName(), field));
+                } else if (typeDescriptor.isArray()) {
+                    indexFieldTypeDescriptor(field, typeDescriptor.getArrayType());
+                }
+        }
 
         @Override
         public void indexMethodReference(MethodDefEntry callerEntry, MethodEntry referencedEntry) {
@@ -25,6 +64,8 @@ public class ReferenceIndex implements JarIndexer {
                         ClassEntry referencedClass = referencedEntry.getParent();
                         referencesToClasses.put(referencedClass, new EntryReference<>(referencedClass, referencedEntry.getName(), callerEntry));
                 }
+
+                indexMethodDescriptor(callerEntry, referencedEntry.getDesc());
         }
 
         @Override
@@ -38,10 +79,12 @@ public class ReferenceIndex implements JarIndexer {
                 referencesToMethods = remapReferencesTo(index, referencesToMethods);
                 referencesToClasses = remapReferencesTo(index, referencesToClasses);
                 referencesToFields = remapReferencesTo(index, referencesToFields);
+                fieldTypeReferences = remapReferencesTo(index, fieldTypeReferences);
+                methodTypeReferences = remapReferencesTo(index, methodTypeReferences);
         }
 
         private <K extends Entry<?>, V extends Entry<?>> Multimap<K, V> remapReferences(JarIndex index, Multimap<K, V> multimap) {
-                Multimap<K, V> resolved = HashMultimap.create();
+                Multimap<K, V> resolved = HashMultimap.create(multimap.keySet().size(), multimap.size() / multimap.keySet().size());
                 for (Map.Entry<K, V> entry : multimap.entries()) {
                         resolved.put(remap(index, entry.getKey()), remap(index, entry.getValue()));
                 }
@@ -49,7 +92,7 @@ public class ReferenceIndex implements JarIndexer {
         }
 
         private <E extends Entry<?>, C extends Entry<?>> Multimap<E, EntryReference<E, C>> remapReferencesTo(JarIndex index, Multimap<E, EntryReference<E, C>> multimap) {
-                Multimap<E, EntryReference<E, C>> resolved = HashMultimap.create();
+                Multimap<E, EntryReference<E, C>> resolved = HashMultimap.create(multimap.keySet().size(), multimap.size() / multimap.keySet().size());
                 for (Map.Entry<E, EntryReference<E, C>> entry : multimap.entries()) {
                         resolved.put(remap(index, entry.getKey()), remap(index, entry.getValue()));
                 }
@@ -79,4 +122,12 @@ public class ReferenceIndex implements JarIndexer {
         public Collection<EntryReference<MethodEntry, MethodDefEntry>> getReferencesToMethod(MethodEntry entry) {
                 return referencesToMethods.get(entry);
         }
+
+        public Collection<EntryReference<ClassEntry, FieldDefEntry>> getFieldTypeReferencesToClass(ClassEntry entry) {
+                return fieldTypeReferences.get(entry);
+        }
+
+        public Collection<EntryReference<ClassEntry, MethodDefEntry>> getMethodTypeReferencesToClass(ClassEntry entry) {
+                return methodTypeReferences.get(entry);
+        }
 }