Add stats generation (#177)

* Add stats generation

* Parameters and inner classes too

* Fixes

6 files changed, 305 insertions, 4 deletions

diff --git a/src/main/java/cuchaz/enigma/gui/GuiController.java b/src/main/java/cuchaz/enigma/gui/GuiController.java
index 54b5c92a..69f12e2e 100644
--- a/src/main/java/cuchaz/enigma/gui/GuiController.java
+++ b/src/main/java/cuchaz/enigma/gui/GuiController.java
@@ -20,6 +20,8 @@ import cuchaz.enigma.api.service.ObfuscationTestService;
 import cuchaz.enigma.bytecode.translators.SourceFixVisitor;
 import cuchaz.enigma.config.Config;
 import cuchaz.enigma.gui.dialog.ProgressDialog;
+import cuchaz.enigma.gui.stats.StatsGenerator;
+import cuchaz.enigma.gui.stats.StatsMember;
 import cuchaz.enigma.gui.util.History;
 import cuchaz.enigma.throwables.MappingParseException;
 import cuchaz.enigma.translation.Translator;
@@ -31,17 +33,19 @@ import cuchaz.enigma.translation.representation.entry.Entry;
 import cuchaz.enigma.translation.representation.entry.FieldEntry;
 import cuchaz.enigma.translation.representation.entry.MethodEntry;
 import cuchaz.enigma.utils.ReadableToken;
+import cuchaz.enigma.utils.Utils;
 import org.objectweb.asm.Opcodes;
 
 import javax.annotation.Nullable;
 import javax.swing.*;
+import java.awt.*;
 import java.awt.event.ItemEvent;
-import java.io.PrintWriter;
-import java.io.StringWriter;
+import java.io.*;
 import java.nio.file.Path;
 import java.util.Collection;
 import java.util.List;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -527,4 +531,25 @@ public class GuiController {
 
                 refreshCurrentClass(reference);
         }
+
+        public void openStats(Set<StatsMember> includedMembers) {
+                ProgressDialog.runOffThread(gui.getFrame(), progress -> {
+                        String data = new StatsGenerator(project).generate(progress, includedMembers);
+
+                        try {
+                                File statsFile = File.createTempFile("stats", ".html");
+
+                                try (FileWriter w = new FileWriter(statsFile)) {
+                                        w.write(
+                                                        Utils.readResourceToString("/stats.html")
+                                                                 .replace("/*data*/", data)
+                                        );
+                                }
+
+                                Desktop.getDesktop().open(statsFile);
+                        } catch (IOException e) {
+                                throw new Error(e);
+                        }
+                });
+        }
 }
diff --git a/src/main/java/cuchaz/enigma/gui/dialog/ProgressDialog.java b/src/main/java/cuchaz/enigma/gui/dialog/ProgressDialog.java
index cf5c2c86..ae30667c 100644
--- a/src/main/java/cuchaz/enigma/gui/dialog/ProgressDialog.java
+++ b/src/main/java/cuchaz/enigma/gui/dialog/ProgressDialog.java
@@ -89,7 +89,12 @@ public class ProgressDialog implements ProgressListener, AutoCloseable {
         @Override
         public void step(int numDone, String message) {
                 this.labelText.setText(message);
-                this.progress.setValue(numDone);
+                if (numDone != -1) {
+                        this.progress.setValue(numDone);
+                        this.progress.setIndeterminate(false);
+                } else {
+                        this.progress.setIndeterminate(true);
+                }
 
                 // update the frame
                 this.frame.validate();
diff --git a/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java b/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
index 2e10bfb5..cdcad05c 100644
--- a/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
+++ b/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
@@ -5,6 +5,7 @@ import cuchaz.enigma.config.Themes;
 import cuchaz.enigma.gui.Gui;
 import cuchaz.enigma.gui.dialog.AboutDialog;
 import cuchaz.enigma.gui.dialog.SearchDialog;
+import cuchaz.enigma.gui.stats.StatsMember;
 import cuchaz.enigma.translation.mapping.serde.MappingFormat;
 import cuchaz.enigma.utils.Utils;
 
@@ -19,8 +20,9 @@ import java.net.URL;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.ArrayList;
 import java.util.List;
+import java.util.*;
+import java.util.stream.Collectors;
 
 public class MenuBar extends JMenuBar {
 
@@ -154,6 +156,45 @@ public class MenuBar extends JMenuBar {
                         }
                         menu.addSeparator();
                         {
+                                JMenuItem stats = new JMenuItem("Mapping Stats...");
+
+                                stats.addActionListener(event -> {
+                    JFrame frame = new JFrame("Choose Included Members");
+                                        Container pane = frame.getContentPane();
+                                        pane.setLayout(new FlowLayout());
+
+                                        Map<StatsMember, JCheckBox> checkboxes = Arrays
+                                                        .stream(StatsMember.values())
+                                                        .collect(Collectors.toMap(m -> m, m -> {
+                                                                JCheckBox checkbox = new JCheckBox(Utils.caplisiseCamelCase(m.name()));
+                                                                pane.add(checkbox);
+                                                                return checkbox;
+                                                        }));
+
+                                        JButton button = new JButton("Generate Stats");
+
+                                        button.addActionListener(e -> {
+                                                Set<StatsMember> includedMembers = checkboxes
+                                                                .entrySet()
+                                                                .stream()
+                                                                .filter(entry -> entry.getValue().isSelected())
+                                                                .map(Map.Entry::getKey)
+                                                                .collect(Collectors.toSet());
+
+                                                frame.setVisible(false);
+                                                frame.dispose();
+                                                gui.getController().openStats(includedMembers);
+                                        });
+
+                                        pane.add(button);
+                    frame.pack();
+                    frame.setVisible(true);
+                });
+
+                                menu.add(stats);
+                        }
+                        menu.addSeparator();
+                        {
                                 JMenuItem item = new JMenuItem("Exit");
                                 menu.add(item);
                                 item.addActionListener(event -> this.gui.close());
diff --git a/src/main/java/cuchaz/enigma/gui/stats/StatsGenerator.java b/src/main/java/cuchaz/enigma/gui/stats/StatsGenerator.java
new file mode 100644
index 00000000..485daf55
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/gui/stats/StatsGenerator.java
@@ -0,0 +1,188 @@
+package cuchaz.enigma.gui.stats;
+
+import com.google.gson.GsonBuilder;
+import cuchaz.enigma.EnigmaProject;
+import cuchaz.enigma.ProgressListener;
+import cuchaz.enigma.analysis.index.EntryIndex;
+import cuchaz.enigma.api.service.NameProposalService;
+import cuchaz.enigma.api.service.ObfuscationTestService;
+import cuchaz.enigma.translation.mapping.EntryRemapper;
+import cuchaz.enigma.translation.mapping.EntryResolver;
+import cuchaz.enigma.translation.mapping.ResolutionStrategy;
+import cuchaz.enigma.translation.representation.TypeDescriptor;
+import cuchaz.enigma.translation.representation.entry.*;
+
+import java.util.*;
+
+public class StatsGenerator {
+    private final EntryIndex entryIndex;
+    private final EntryRemapper mapper;
+    private final EntryResolver entryResolver;
+    private final ObfuscationTestService obfuscationTestService;
+    private final NameProposalService nameProposalService;
+
+    public StatsGenerator(EnigmaProject project) {
+        entryIndex = project.getJarIndex().getEntryIndex();
+        mapper = project.getMapper();
+        entryResolver = project.getJarIndex().getEntryResolver();
+        obfuscationTestService = project.getEnigma().getServices().get(ObfuscationTestService.TYPE).orElse(null);
+        nameProposalService = project.getEnigma().getServices().get(NameProposalService.TYPE).orElse(null);
+    }
+
+    public String generate(ProgressListener progress, Set<StatsMember> includedMembers) {
+        includedMembers = EnumSet.copyOf(includedMembers);
+        int totalWork = 0;
+
+        if (includedMembers.contains(StatsMember.METHODS) || includedMembers.contains(StatsMember.PARAMETERS)) {
+            totalWork += entryIndex.getMethods().size();
+        }
+
+        if (includedMembers.contains(StatsMember.FIELDS)) {
+            totalWork += entryIndex.getFields().size();
+        }
+
+        if (includedMembers.contains(StatsMember.CLASSES)) {
+            totalWork += entryIndex.getClasses().size();
+        }
+
+        progress.init(totalWork, "Generating stats");
+
+        Map<String, Integer> counts = new HashMap<>();
+
+        int numDone = 0;
+        if (includedMembers.contains(StatsMember.METHODS) || includedMembers.contains(StatsMember.PARAMETERS)) {
+            for (MethodEntry method : entryIndex.getMethods()) {
+                progress.step(numDone++, "Methods");
+                MethodEntry root = entryResolver
+                        .resolveEntry(method, ResolutionStrategy.RESOLVE_ROOT)
+                        .stream()
+                        .findFirst()
+                        .orElseThrow(AssertionError::new);
+
+                if (root == method && !((MethodDefEntry) method).getAccess().isSynthetic()) {
+                    if (includedMembers.contains(StatsMember.METHODS)) {
+                        update(counts, method);
+                    }
+
+                    if (includedMembers.contains(StatsMember.PARAMETERS)) {
+                        int index = ((MethodDefEntry) method).getAccess().isStatic() ? 0 : 1;
+                        for (TypeDescriptor argument : method.getDesc().getArgumentDescs()) {
+                            update(counts, new LocalVariableEntry(method, index, "", true));
+                            index += argument.getSize();
+                        }
+                    }
+                }
+            }
+        }
+
+        if (includedMembers.contains(StatsMember.FIELDS)) {
+            for (FieldEntry field : entryIndex.getFields()) {
+                progress.step(numDone++, "Fields");
+                update(counts, field);
+            }
+        }
+
+        if (includedMembers.contains(StatsMember.CLASSES)) {
+            for (ClassEntry clazz : entryIndex.getClasses()) {
+                progress.step(numDone++, "Classes");
+                update(counts, clazz);
+            }
+        }
+
+        progress.step(-1, "Generating data");
+
+        Tree<Integer> tree = new Tree<>();
+
+        for (Map.Entry<String, Integer> entry : counts.entrySet()) {
+            if (entry.getKey().startsWith("com.mojang")) continue; // just a few unmapped names, no point in having a subsection
+            tree.getNode(entry.getKey()).value = entry.getValue();
+        }
+
+        tree.collapse(tree.root);
+        return new GsonBuilder().setPrettyPrinting().create().toJson(tree.root);
+    }
+
+    private void update(Map<String, Integer> counts, Entry<?> entry) {
+        if (isObfuscated(entry)) {
+            String parent = mapper.deobfuscate(entry.getAncestry().get(0)).getName().replace('/', '.');
+            counts.put(parent, counts.getOrDefault(parent, 0) + 1);
+        }
+    }
+
+    private boolean isObfuscated(Entry<?> entry) {
+        String name = entry.getName();
+
+        if (obfuscationTestService != null && obfuscationTestService.testDeobfuscated(entry)) {
+            return false;
+        }
+
+        if (nameProposalService != null && nameProposalService.proposeName(entry, mapper).isPresent()) {
+            return false;
+        }
+
+        String mappedName = mapper.deobfuscate(entry).getName();
+        if (mappedName != null && !mappedName.isEmpty() && !mappedName.equals(name)) {
+            return false;
+        }
+
+        return true;
+    }
+
+    private static class Tree<T> {
+        public final Node<T> root;
+        private final Map<String, Node<T>> nodes = new HashMap<>();
+
+        public static class Node<T> {
+            public String name;
+            public T value;
+            public List<Node<T>> children = new ArrayList<>();
+            private final transient Map<String, Node<T>> namedChildren = new HashMap<>();
+
+            public Node(String name, T value) {
+                this.name = name;
+                this.value = value;
+            }
+        }
+
+        public Tree() {
+            root = new Node<>("", null);
+        }
+
+        public Node<T> getNode(String name) {
+            Node<T> node = nodes.get(name);
+
+            if (node == null) {
+                node = root;
+
+                for (String part : name.split("\\.")) {
+                    Node<T> child = node.namedChildren.get(part);
+
+                    if (child == null) {
+                        child = new Node<>(part, null);
+                        node.namedChildren.put(part, child);
+                        node.children.add(child);
+                    }
+
+                    node = child;
+                }
+
+                nodes.put(name, node);
+            }
+
+            return node;
+        }
+
+        public void collapse(Node<T> node) {
+            while (node.children.size() == 1) {
+                Node<T> child = node.children.get(0);
+                node.name = node.name.isEmpty() ? child.name : node.name + "." + child.name;
+                node.children = child.children;
+                node.value = child.value;
+            }
+
+            for (Node<T> child : node.children) {
+                collapse(child);
+            }
+        }
+    }
+}
diff --git a/src/main/java/cuchaz/enigma/gui/stats/StatsMember.java b/src/main/java/cuchaz/enigma/gui/stats/StatsMember.java
new file mode 100644
index 00000000..70b4f40d
--- /dev/null
+++ b/src/main/java/cuchaz/enigma/gui/stats/StatsMember.java
@@ -0,0 +1,8 @@
+package cuchaz.enigma.gui.stats;
+
+public enum StatsMember {
+    METHODS,
+    FIELDS,
+    PARAMETERS,
+    CLASSES
+}
diff --git a/src/main/resources/stats.html b/src/main/resources/stats.html
new file mode 100644
index 00000000..fcff7c0f
--- /dev/null
+++ b/src/main/resources/stats.html
@@ -0,0 +1,34 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta content="width=device-width, initial-scale=1" name="viewport">
+    <title>Stats</title>
+    <link href="https://cdn.anychart.com/releases/v8/css/anychart-ui.min.css" rel="stylesheet" type="text/css">
+    <style>
+        html, body, #container {
+            width: 100%;
+            height: 100%;
+            margin: 0;
+            padding: 0;
+        }
+    </style>
+</head>
+
+<body>
+<div id="container"></div>
+<script src="https://cdn.anychart.com/releases/v8/js/anychart-base.min.js"></script>
+<script src="https://cdn.anychart.com/releases/v8/js/anychart-sunburst.min.js"></script>
+<script src="https://cdn.anychart.com/releases/v8/js/anychart-exports.min.js"></script>
+<script src="https://cdn.anychart.com/releases/v8/js/anychart-ui.min.js"></script>
+<script>
+    anychart.onDocumentReady(function () {
+        var chart = anychart.sunburst([/*data*/], "as-tree");
+        chart.sort("desc");
+        chart.calculationMode("parent-independent");
+        chart.container("container");
+        chart.draw();
+    });
+</script>
+</body>
+</html>