2 files changed, 55 insertions, 18 deletions

diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
index 2dc1d77a..b62e9cf9 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
@@ -189,9 +189,14 @@ public class GuiController implements ClientPacketHandler {
 
                 return ProgressDialog.runOffThread(this.gui.getFrame(), progress -> {
                         EnigmaProject.JarExport jar = project.exportRemappedJar(progress);
-                        EnigmaProject.SourceExport source = jar.decompile(progress, chp.getDecompilerService());
-
-                        source.write(path, progress);
+                        jar.decompileStream(progress, chp.getDecompilerService(), EnigmaProject.DecompileErrorStrategy.TRACE_AS_SOURCE)
+                                        .forEach(source -> {
+                                                try {
+                                                        source.writeTo(source.resolvePath(path));
+                                                } catch (IOException e) {
+                                                        e.printStackTrace();
+                                                }
+                                        });
                 });
         }
 
diff --git a/enigma/src/main/java/cuchaz/enigma/EnigmaProject.java b/enigma/src/main/java/cuchaz/enigma/EnigmaProject.java
index 3999572f..c3751581 100644
--- a/enigma/src/main/java/cuchaz/enigma/EnigmaProject.java
+++ b/enigma/src/main/java/cuchaz/enigma/EnigmaProject.java
@@ -31,12 +31,14 @@ import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Collection;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.jar.JarEntry;
 import java.util.jar.JarOutputStream;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 public class EnigmaProject {
         private final Enigma enigma;
@@ -173,15 +175,13 @@ public class EnigmaProject {
                                 .filter(Objects::nonNull)
                                 .collect(Collectors.toMap(n -> n.name, Functions.identity()));
 
-                return new JarExport(jarIndex, compiled);
+                return new JarExport(compiled);
         }
 
         public static final class JarExport {
-                private final JarIndex jarIndex;
                 private final Map<String, ClassNode> compiled;
 
-                JarExport(JarIndex jarIndex, Map<String, ClassNode> compiled) {
-                        this.jarIndex = jarIndex;
+                JarExport(Map<String, ClassNode> compiled) {
                         this.compiled = compiled;
                 }
 
@@ -207,6 +207,15 @@ public class EnigmaProject {
                 }
 
                 public SourceExport decompile(ProgressListener progress, DecompilerService decompilerService) {
+                        return this.decompile(progress, decompilerService, DecompileErrorStrategy.PROPAGATE);
+                }
+
+                public SourceExport decompile(ProgressListener progress, DecompilerService decompilerService, DecompileErrorStrategy errorStrategy) {
+                        List<ClassSource> decompiled = this.decompileStream(progress, decompilerService, errorStrategy).collect(Collectors.toList());
+                        return new SourceExport(decompiled);
+                }
+
+                public Stream<ClassSource> decompileStream(ProgressListener progress, DecompilerService decompilerService, DecompileErrorStrategy errorStrategy) {
                         Collection<ClassNode> classes = this.compiled.values().stream()
                                         .filter(classNode -> classNode.name.indexOf('$') == -1)
                                         .collect(Collectors.toList());
@@ -218,16 +227,33 @@ public class EnigmaProject {
 
                         AtomicInteger count = new AtomicInteger();
 
-                        Collection<ClassSource> decompiled = classes.parallelStream()
+                        return classes.parallelStream()
                                         .map(translatedNode -> {
                                                 progress.step(count.getAndIncrement(), translatedNode.name);
 
-                                                String source = decompileClass(translatedNode, decompiler);
+                                                String source = null;
+                                                try {
+                                                        source = decompileClass(translatedNode, decompiler);
+                                                } catch (Throwable throwable) {
+                                                        switch (errorStrategy) {
+                                                                case PROPAGATE: throw throwable;
+                                                                case IGNORE: break;
+                                                                case TRACE_AS_SOURCE: {
+                                                                        StringWriter writer = new StringWriter();
+                                                                        throwable.printStackTrace(new PrintWriter(writer));
+                                                                        source = writer.toString();
+                                                                        break;
+                                                                }
+                                                        }
+                                                }
+
+                                                if (source == null) {
+                                                        return null;
+                                                }
+
                                                 return new ClassSource(translatedNode.name, source);
                                         })
-                                        .collect(Collectors.toList());
-
-                        return new SourceExport(decompiled);
+                                        .filter(Objects::nonNull);
                 }
 
                 private String decompileClass(ClassNode translatedNode, Decompiler decompiler) {
@@ -236,7 +262,7 @@ public class EnigmaProject {
         }
 
         public static final class SourceExport {
-                private final Collection<ClassSource> decompiled;
+                public final Collection<ClassSource> decompiled;
 
                 SourceExport(Collection<ClassSource> decompiled) {
                         this.decompiled = decompiled;
@@ -255,24 +281,30 @@ public class EnigmaProject {
                 }
         }
 
-        private static class ClassSource {
-                private final String name;
-                private final String source;
+        public static class ClassSource {
+                public final String name;
+                public final String source;
 
                 ClassSource(String name, String source) {
                         this.name = name;
                         this.source = source;
                 }
 
-                void writeTo(Path path) throws IOException {
+                public void writeTo(Path path) throws IOException {
                         Files.createDirectories(path.getParent());
                         try (BufferedWriter writer = Files.newBufferedWriter(path)) {
                                 writer.write(source);
                         }
                 }
 
-                Path resolvePath(Path root) {
+                public Path resolvePath(Path root) {
                         return root.resolve(name.replace('.', '/') + ".java");
                 }
         }
+
+        public enum DecompileErrorStrategy {
+                PROPAGATE,
+                TRACE_AS_SOURCE,
+                IGNORE
+        }
 }