5 files changed, 54 insertions, 59 deletions

diff --git a/src/test/java/cuchaz/enigma/TestDeobfed.java b/src/test/java/cuchaz/enigma/TestDeobfed.java
index 3ee86cb..c88b0eb 100644
--- a/src/test/java/cuchaz/enigma/TestDeobfed.java
+++ b/src/test/java/cuchaz/enigma/TestDeobfed.java
@@ -13,6 +13,9 @@ package cuchaz.enigma;
 
 import cuchaz.enigma.analysis.ClassCache;
 import cuchaz.enigma.analysis.index.JarIndex;
+import cuchaz.enigma.source.Decompiler;
+import cuchaz.enigma.source.Decompilers;
+import cuchaz.enigma.source.SourceSettings;
 import org.junit.BeforeClass;
 import org.junit.Test;
 
@@ -68,31 +71,29 @@ public class TestDeobfed {
         @Test
         public void decompile() {
                 EnigmaProject project = new EnigmaProject(enigma, classCache, index);
+                Decompiler decompiler = Decompilers.PROCYON.create(project.getClassCache(), new SourceSettings(false, false));
 
-                CompiledSourceTypeLoader typeLoader = new CompiledSourceTypeLoader(project.getClassCache());
-                SourceProvider sourceProvider = new SourceProvider(SourceProvider.createSettings(), typeLoader);
-
-                sourceProvider.getSources("a");
-                sourceProvider.getSources("b");
-                sourceProvider.getSources("c");
-                sourceProvider.getSources("d");
-                sourceProvider.getSources("d$1");
-                sourceProvider.getSources("e");
-                sourceProvider.getSources("f");
-                sourceProvider.getSources("g");
-                sourceProvider.getSources("g$a");
-                sourceProvider.getSources("g$a$a");
-                sourceProvider.getSources("g$b");
-                sourceProvider.getSources("g$b$a");
-                sourceProvider.getSources("h");
-                sourceProvider.getSources("h$a");
-                sourceProvider.getSources("h$a$a");
-                sourceProvider.getSources("h$b");
-                sourceProvider.getSources("h$b$a");
-                sourceProvider.getSources("h$b$a$a");
-                sourceProvider.getSources("h$b$a$b");
-                sourceProvider.getSources("i");
-                sourceProvider.getSources("i$a");
-                sourceProvider.getSources("i$b");
+                decompiler.getSource("a");
+                decompiler.getSource("b");
+                decompiler.getSource("c");
+                decompiler.getSource("d");
+                decompiler.getSource("d$1");
+                decompiler.getSource("e");
+                decompiler.getSource("f");
+                decompiler.getSource("g");
+                decompiler.getSource("g$a");
+                decompiler.getSource("g$a$a");
+                decompiler.getSource("g$b");
+                decompiler.getSource("g$b$a");
+                decompiler.getSource("h");
+                decompiler.getSource("h$a");
+                decompiler.getSource("h$a$a");
+                decompiler.getSource("h$b");
+                decompiler.getSource("h$b$a");
+                decompiler.getSource("h$b$a$a");
+                decompiler.getSource("h$b$a$b");
+                decompiler.getSource("i");
+                decompiler.getSource("i$a");
+                decompiler.getSource("i$b");
         }
 }
diff --git a/src/test/java/cuchaz/enigma/TestDeobfuscator.java b/src/test/java/cuchaz/enigma/TestDeobfuscator.java
index 2a6fec4..6619d26 100644
--- a/src/test/java/cuchaz/enigma/TestDeobfuscator.java
+++ b/src/test/java/cuchaz/enigma/TestDeobfuscator.java
@@ -11,13 +11,15 @@
 
 package cuchaz.enigma;
 
+import cuchaz.enigma.source.Decompiler;
+import cuchaz.enigma.source.Decompilers;
+import cuchaz.enigma.source.SourceSettings;
 import org.junit.Test;
 
 import java.io.IOException;
 import java.nio.file.Paths;
 
 public class TestDeobfuscator {
-
         private EnigmaProject openProject() throws IOException {
                 Enigma enigma = Enigma.create();
                 return enigma.openJar(Paths.get("build/test-obf/loneClass.jar"), ProgressListener.none());
@@ -32,10 +34,8 @@ public class TestDeobfuscator {
         @Test
         public void decompileClass() throws Exception {
                 EnigmaProject project = openProject();
+                Decompiler decompiler = Decompilers.PROCYON.create(project.getClassCache(), new SourceSettings(false, false));
 
-                CompiledSourceTypeLoader typeLoader = new CompiledSourceTypeLoader(project.getClassCache());
-                SourceProvider sourceProvider = new SourceProvider(SourceProvider.createSettings(), typeLoader);
-
-                sourceProvider.writeSourceToString(sourceProvider.getSources("a"));
+                decompiler.getSource("a").asString();
         }
 }
diff --git a/src/test/java/cuchaz/enigma/TestInnerClasses.java b/src/test/java/cuchaz/enigma/TestInnerClasses.java
index 18e4936..85c72f8 100644
--- a/src/test/java/cuchaz/enigma/TestInnerClasses.java
+++ b/src/test/java/cuchaz/enigma/TestInnerClasses.java
@@ -13,6 +13,9 @@ package cuchaz.enigma;
 
 import cuchaz.enigma.analysis.ClassCache;
 import cuchaz.enigma.analysis.index.JarIndex;
+import cuchaz.enigma.source.Decompiler;
+import cuchaz.enigma.source.Decompilers;
+import cuchaz.enigma.source.SourceSettings;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
 import org.junit.Test;
 
@@ -32,15 +35,13 @@ public class TestInnerClasses {
         private static final ClassEntry ClassTreeLevel1 = newClass("f$a");
         private static final ClassEntry ClassTreeLevel2 = newClass("f$a$a");
         private static final ClassEntry ClassTreeLevel3 = newClass("f$a$a$a");
-        private JarIndex index;
-        private SourceProvider sourceProvider;
+        private final JarIndex index;
+        private final Decompiler decompiler;
 
         public TestInnerClasses() throws Exception {
                 ClassCache classCache = ClassCache.of(Paths.get("build/test-obf/innerClasses.jar"));
                 index = classCache.index(ProgressListener.none());
-
-                CompiledSourceTypeLoader typeLoader = new CompiledSourceTypeLoader(classCache);
-                sourceProvider = new SourceProvider(SourceProvider.createSettings(), typeLoader);
+                decompiler = Decompilers.PROCYON.create(classCache, new SourceSettings(false, false));
         }
 
         @Test
@@ -79,6 +80,6 @@ public class TestInnerClasses {
         }
 
         private void decompile(ClassEntry classEntry) {
-                sourceProvider.getSources(classEntry.getName());
+                decompiler.getSource(classEntry.getName());
         }
 }
diff --git a/src/test/java/cuchaz/enigma/TestSourceIndex.java b/src/test/java/cuchaz/enigma/TestSourceIndex.java
index b1d1182..b201608 100644
--- a/src/test/java/cuchaz/enigma/TestSourceIndex.java
+++ b/src/test/java/cuchaz/enigma/TestSourceIndex.java
@@ -12,9 +12,8 @@
 package cuchaz.enigma;
 
 import com.google.common.collect.Sets;
-import com.strobel.decompiler.languages.java.ast.CompilationUnit;
 import cuchaz.enigma.analysis.ClassCache;
-import cuchaz.enigma.analysis.SourceIndex;
+import cuchaz.enigma.source.*;
 import cuchaz.enigma.analysis.index.JarIndex;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
 import org.junit.Test;
@@ -53,8 +52,7 @@ public class TestSourceIndex {
                 ClassCache classCache = project.getClassCache();
                 JarIndex index = project.getJarIndex();
 
-                CompiledSourceTypeLoader typeLoader = new CompiledSourceTypeLoader(classCache);
-                SourceProvider sourceProvider = new SourceProvider(SourceProvider.createSettings(), typeLoader);
+                Decompiler decompiler = Decompilers.PROCYON.create(classCache, new SourceSettings(false, false));
 
                 // get all classes that aren't inner classes
                 Set<ClassEntry> classEntries = Sets.newHashSet();
@@ -66,10 +64,8 @@ public class TestSourceIndex {
 
                 for (ClassEntry obfClassEntry : classEntries) {
                         try {
-                                CompilationUnit tree = sourceProvider.getSources(obfClassEntry.getName());
-                                String source = sourceProvider.writeSourceToString(tree);
-
-                                SourceIndex.buildIndex(source, tree, true);
+                                Source source = decompiler.getSource(obfClassEntry.getName());
+                                source.index();
                         } catch (Throwable t) {
                                 throw new Error("Unable to index " + obfClassEntry, t);
                         }
diff --git a/src/test/java/cuchaz/enigma/TokenChecker.java b/src/test/java/cuchaz/enigma/TokenChecker.java
index 1dde034..48d0c83 100644
--- a/src/test/java/cuchaz/enigma/TokenChecker.java
+++ b/src/test/java/cuchaz/enigma/TokenChecker.java
@@ -12,10 +12,10 @@
 package cuchaz.enigma;
 
 import com.google.common.collect.Lists;
-import com.strobel.decompiler.languages.java.ast.CompilationUnit;
 import cuchaz.enigma.analysis.ClassCache;
 import cuchaz.enigma.analysis.EntryReference;
-import cuchaz.enigma.analysis.SourceIndex;
+import cuchaz.enigma.source.SourceIndex;
+import cuchaz.enigma.source.*;
 import cuchaz.enigma.analysis.Token;
 import cuchaz.enigma.translation.representation.entry.Entry;
 
@@ -25,43 +25,40 @@ import java.util.Collection;
 import java.util.List;
 
 public class TokenChecker {
-
-        private SourceProvider sourceProvider;
+        private final Decompiler decompiler;
 
         protected TokenChecker(Path path) throws IOException {
                 ClassCache classCache = ClassCache.of(path);
-
-                CompiledSourceTypeLoader typeLoader = new CompiledSourceTypeLoader(classCache);
-                sourceProvider = new SourceProvider(SourceProvider.createSettings(), typeLoader);
+                decompiler = Decompilers.PROCYON.create(classCache, new SourceSettings(false, false));
         }
 
         protected String getDeclarationToken(Entry<?> entry) {
                 // decompile the class
-                CompilationUnit tree = sourceProvider.getSources(entry.getContainingClass().getFullName());
+                Source source = decompiler.getSource(entry.getContainingClass().getFullName());
                 // DEBUG
                 // tree.acceptVisitor( new TreeDumpVisitor( new File( "tree." + entry.getClassName().replace( '/', '.' ) + ".txt" ) ), null );
-                String source = sourceProvider.writeSourceToString(tree);
-                SourceIndex index = SourceIndex.buildIndex(source, tree, true);
+                String string = source.asString();
+                SourceIndex index = source.index();
 
                 // get the token value
                 Token token = index.getDeclarationToken(entry);
                 if (token == null) {
                         return null;
                 }
-                return source.substring(token.start, token.end);
+                return string.substring(token.start, token.end);
         }
 
         @SuppressWarnings("unchecked")
         protected Collection<String> getReferenceTokens(EntryReference<? extends Entry<?>, ? extends Entry<?>> reference) {
                 // decompile the class
-                CompilationUnit tree = sourceProvider.getSources(reference.context.getContainingClass().getFullName());
-                String source = sourceProvider.writeSourceToString(tree);
-                SourceIndex index = SourceIndex.buildIndex(source, tree, true);
+                Source source = decompiler.getSource(reference.context.getContainingClass().getFullName());
+                String string = source.asString();
+                SourceIndex index = source.index();
 
                 // get the token values
                 List<String> values = Lists.newArrayList();
                 for (Token token : index.getReferenceTokens((EntryReference<Entry<?>, Entry<?>>) reference)) {
-                        values.add(source.substring(token.start, token.end));
+                        values.add(string.substring(token.start, token.end));
                 }
                 return values;
         }