Track mapping loading progress for the Enigma directory format (#577)

2 files changed, 77 insertions, 3 deletions

diff --git a/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java b/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java
index b891f438..fb52ac0f 100644
--- a/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java
+++ b/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java
@@ -93,10 +93,11 @@ public enum MappingFormat {
                         loadingMessage = I18n.translate("progress.mappings.loading_directory");
                 }
 
-                progressListener.init(1, loadingMessage);
-
                 VisitableMappingTree mappingTree = new MemoryMappingTree();
-                MappingReader.read(path, mappingIoCounterpart, mappingTree);
+                ProgressTrackingMappingVisitor.trackLoadingProgress(mappingTree, path, this, progressListener, (visitor, totalWork) -> {
+                        progressListener.init(totalWork, loadingMessage);
+                        MappingReader.read(path, mappingIoCounterpart, visitor);
+                });
                 EntryTree<EntryMapping> mappings = MappingIoConverter.fromMappingIo(mappingTree, progressListener, index);
 
                 return this == PROGUARD ? MappingOperations.invert(mappings) : mappings;
diff --git a/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/ProgressTrackingMappingVisitor.java b/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/ProgressTrackingMappingVisitor.java
new file mode 100644
index 00000000..0ee90e48
--- /dev/null
+++ b/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/ProgressTrackingMappingVisitor.java
@@ -0,0 +1,73 @@
+package cuchaz.enigma.translation.mapping.serde;
+
+import java.io.IOException;
+import java.nio.file.FileVisitResult;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.attribute.BasicFileAttributes;
+import java.util.HashSet;
+import java.util.Set;
+
+import net.fabricmc.mappingio.MappedElementKind;
+import net.fabricmc.mappingio.MappingVisitor;
+import net.fabricmc.mappingio.adapter.ForwardingMappingVisitor;
+
+import cuchaz.enigma.ProgressListener;
+
+final class ProgressTrackingMappingVisitor extends ForwardingMappingVisitor {
+        private final Set<String> classNames;
+        private final ProgressListener progressListener;
+        private int progress = 0;
+
+        ProgressTrackingMappingVisitor(MappingVisitor next, Set<String> classNames, ProgressListener progressListener) {
+                super(next);
+                this.classNames = classNames;
+                this.progressListener = progressListener;
+        }
+
+        @Override
+        public void visitDstName(MappedElementKind targetKind, int namespace, String name) throws IOException {
+                if (targetKind == MappedElementKind.CLASS && classNames.contains(name)) {
+                        progressListener.step(++progress, name);
+                }
+
+                super.visitDstName(targetKind, namespace, name);
+        }
+
+        static void trackLoadingProgress(MappingVisitor next, Path path, MappingFormat format, ProgressListener progressListener, VisitorWithProgressConsumer consumer) throws IOException {
+                if (format != MappingFormat.ENIGMA_DIRECTORY) {
+                        consumer.accept(next, 1);
+                }
+
+                Set<String> classNames = collectClassNames(path);
+                consumer.accept(new ProgressTrackingMappingVisitor(next, classNames, progressListener), classNames.size());
+        }
+
+        private static Set<String> collectClassNames(Path dir) throws IOException {
+                Set<String> names = new HashSet<>();
+
+                Files.walkFileTree(dir, new SimpleFileVisitor<>() {
+                        @Override
+                        public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) {
+                                // Matches the logic for finding files in mapping-io's EnigmaDirReader.
+                                String extension = "." + net.fabricmc.mappingio.format.MappingFormat.ENIGMA_FILE.fileExt;
+
+                                if (file.getFileName().toString().endsWith(extension)) {
+                                        String filePath = dir.relativize(file).toString().replace(dir.getFileSystem().getSeparator(), "/");
+                                        String className = filePath.substring(0, filePath.length() - extension.length());
+                                        names.add(className);
+                                }
+
+                                return FileVisitResult.CONTINUE;
+                        }
+                });
+
+                return names;
+        }
+
+        @FunctionalInterface
+        interface VisitorWithProgressConsumer {
+                void accept(MappingVisitor visitor, int totalWork) throws IOException;
+        }
+}