5 files changed, 40 insertions, 33 deletions

diff --git a/src/test/java/cuchaz/enigma/TestDeobfed.java b/src/test/java/cuchaz/enigma/TestDeobfed.java
index 25cb60c..14b1418 100644
--- a/src/test/java/cuchaz/enigma/TestDeobfed.java
+++ b/src/test/java/cuchaz/enigma/TestDeobfed.java
@@ -68,27 +68,28 @@ public class TestDeobfed {
         public void decompile()
                 throws Exception {
                 Deobfuscator deobfuscator = new Deobfuscator(jar);
-                deobfuscator.getSourceTree("a");
-                deobfuscator.getSourceTree("b");
-                deobfuscator.getSourceTree("c");
-                deobfuscator.getSourceTree("d");
-                deobfuscator.getSourceTree("d$1");
-                deobfuscator.getSourceTree("e");
-                deobfuscator.getSourceTree("f");
-                deobfuscator.getSourceTree("g");
-                deobfuscator.getSourceTree("g$a");
-                deobfuscator.getSourceTree("g$a$a");
-                deobfuscator.getSourceTree("g$b");
-                deobfuscator.getSourceTree("g$b$a");
-                deobfuscator.getSourceTree("h");
-                deobfuscator.getSourceTree("h$a");
-                deobfuscator.getSourceTree("h$a$a");
-                deobfuscator.getSourceTree("h$b");
-                deobfuscator.getSourceTree("h$b$a");
-                deobfuscator.getSourceTree("h$b$a$a");
-                deobfuscator.getSourceTree("h$b$a$b");
-                deobfuscator.getSourceTree("i");
-                deobfuscator.getSourceTree("i$a");
-                deobfuscator.getSourceTree("i$b");
+                SourceProvider sourceProvider = deobfuscator.getObfSourceProvider();
+                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");
         }
 }
diff --git a/src/test/java/cuchaz/enigma/TestDeobfuscator.java b/src/test/java/cuchaz/enigma/TestDeobfuscator.java
index 5f11721..e070b66 100644
--- a/src/test/java/cuchaz/enigma/TestDeobfuscator.java
+++ b/src/test/java/cuchaz/enigma/TestDeobfuscator.java
@@ -51,6 +51,7 @@ public class TestDeobfuscator {
         public void decompileClass()
                 throws Exception {
                 Deobfuscator deobfuscator = getDeobfuscator();
-                deobfuscator.getSource(deobfuscator.getSourceTree("a"));
+                SourceProvider sourceProvider = deobfuscator.getObfSourceProvider();
+                sourceProvider.writeSourceToString(sourceProvider.getSources("a"));
         }
 }
diff --git a/src/test/java/cuchaz/enigma/TestInnerClasses.java b/src/test/java/cuchaz/enigma/TestInnerClasses.java
index 0319cf6..8738fd7 100644
--- a/src/test/java/cuchaz/enigma/TestInnerClasses.java
+++ b/src/test/java/cuchaz/enigma/TestInnerClasses.java
@@ -79,6 +79,6 @@ public class TestInnerClasses {
         }
 
         private void decompile(ClassEntry classEntry) {
-                deobfuscator.getSourceTree(classEntry.getName());
+                deobfuscator.getObfSourceProvider().getSources(classEntry.getName());
         }
 }
diff --git a/src/test/java/cuchaz/enigma/TestSourceIndex.java b/src/test/java/cuchaz/enigma/TestSourceIndex.java
index a5f5f71..ce5d631 100644
--- a/src/test/java/cuchaz/enigma/TestSourceIndex.java
+++ b/src/test/java/cuchaz/enigma/TestSourceIndex.java
@@ -13,6 +13,7 @@ package cuchaz.enigma;
 
 import com.google.common.collect.Sets;
 import com.strobel.decompiler.languages.java.ast.CompilationUnit;
+import cuchaz.enigma.analysis.SourceIndex;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
 import org.junit.Test;
 
@@ -50,11 +51,13 @@ public class TestSourceIndex {
                         }
                 }
 
+                SourceProvider sourceProvider = deobfuscator.getObfSourceProvider();
                 for (ClassEntry obfClassEntry : classEntries) {
                         try {
-                                CompilationUnit tree = deobfuscator.getSourceTree(obfClassEntry.getName());
-                                String source = deobfuscator.getSource(tree);
-                                deobfuscator.getSourceIndex(tree, source);
+                                CompilationUnit tree = sourceProvider.getSources(obfClassEntry.getName());
+                                String source = sourceProvider.writeSourceToString(tree);
+
+                                SourceIndex.buildIndex(source, tree, true);
                         } 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 aac2866..c4670a2 100644
--- a/src/test/java/cuchaz/enigma/TokenChecker.java
+++ b/src/test/java/cuchaz/enigma/TokenChecker.java
@@ -34,11 +34,12 @@ public class TokenChecker {
 
         protected String getDeclarationToken(Entry<?> entry) {
                 // decompile the class
-                CompilationUnit tree = deobfuscator.getSourceTree(entry.getContainingClass().getFullName());
+                SourceProvider sourceProvider = deobfuscator.getObfSourceProvider();
+                CompilationUnit tree = sourceProvider.getSources(entry.getContainingClass().getFullName());
                 // DEBUG
                 // tree.acceptVisitor( new TreeDumpVisitor( new File( "tree." + entry.getClassName().replace( '/', '.' ) + ".txt" ) ), null );
-                String source = deobfuscator.getSource(tree);
-                SourceIndex index = deobfuscator.getSourceIndex(tree, source);
+                String source = sourceProvider.writeSourceToString(tree);
+                SourceIndex index = SourceIndex.buildIndex(source, tree, true);
 
                 // get the token value
                 Token token = index.getDeclarationToken(entry);
@@ -51,9 +52,10 @@ public class TokenChecker {
         @SuppressWarnings("unchecked")
         protected Collection<String> getReferenceTokens(EntryReference<? extends Entry<?>, ? extends Entry<?>> reference) {
                 // decompile the class
-                CompilationUnit tree = deobfuscator.getSourceTree(reference.context.getContainingClass().getFullName());
-                String source = deobfuscator.getSource(tree);
-                SourceIndex index = deobfuscator.getSourceIndex(tree, source);
+                SourceProvider sourceProvider = deobfuscator.getObfSourceProvider();
+                CompilationUnit tree = sourceProvider.getSources(reference.context.getContainingClass().getFullName());
+                String source = sourceProvider.writeSourceToString(tree);
+                SourceIndex index = SourceIndex.buildIndex(source, tree, true);
 
                 // get the token values
                 List<String> values = Lists.newArrayList();