Make baksmali thread safe, and add -j option

baksmali/src/main/java/org/jf/baksmali/baksmali.java

@@ -30,6 +30,7 @@ package org.jf.baksmali;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
 import org.jf.baksmali.Adaptors.ClassDefinition;
 import org.jf.dexlib2.analysis.ClassPath;
 import org.jf.dexlib2.iface.ClassDef;
@@ -40,10 +41,11 @@ import org.jf.util.IndentingWriter;
 
 import java.io.*;
 import java.util.*;
+import java.util.concurrent.*;
 
 public class baksmali {
 
-    public static void disassembleDexFile(DexFile dexFile, baksmaliOptions options) {
+    public static void disassembleDexFile(DexFile dexFile, final baksmaliOptions options) {
         if (options.registerInfo != 0 || options.deodex) {
             try {
                 Iterable<String> extraClassPathEntries;
@@ -87,11 +89,34 @@ public class baksmali {
             options.syntheticAccessorResolver = new SyntheticAccessorResolver(classDefs);
         }
 
-        ClassFileNameHandler fileNameHandler = new ClassFileNameHandler(outputDirectoryFile, ".smali");
+        final ClassFileNameHandler fileNameHandler = new ClassFileNameHandler(outputDirectoryFile, ".smali");
 
-        for (ClassDef classDef: classDefs) {
+        ExecutorService executor = Executors.newFixedThreadPool(options.jobs);
+        List<Future<Void>> tasks = Lists.newArrayList();
+
+        for (final ClassDef classDef: classDefs) {
+            tasks.add(executor.submit(new Callable<Void>() {
+                @Override public Void call() throws Exception {
                     disassembleClass(classDef, fileNameHandler, options);
+                    return null;
                 }
+            }));
+        }
+
+        for (Future<Void> task: tasks) {
+            while(true) {
+                try {
+                    task.get();
+                } catch (InterruptedException ex) {
+                    continue;
+                } catch (ExecutionException ex) {
+                    throw new RuntimeException(ex);
+                }
+                break;
+            }
+        }
+
+        executor.shutdown();
     }
 
     private static void disassembleClass(ClassDef classDef, ClassFileNameHandler fileNameHandler,
@@ -124,10 +149,13 @@ public class baksmali {
             File smaliParent = smaliFile.getParentFile();
             if (!smaliParent.exists()) {
                 if (!smaliParent.mkdirs()) {
+                    // check again, it's likely it was created in a different thread
+                    if (!smaliParent.exists()) {
                         System.err.println("Unable to create directory " + smaliParent.toString() + " - skipping class");
                         return;
                     }
                 }
+            }
 
             if (!smaliFile.exists()){
                 if (!smaliFile.createNewFile()) {

baksmali/src/main/java/org/jf/baksmali/baksmaliOptions.java

@@ -68,6 +68,7 @@ public class baksmaliOptions {
     public InlineMethodResolver inlineResolver = null;
     public int registerInfo = 0;
     public ClassPath classPath = null;
+    public int jobs = -1;
 
     public SyntheticAccessorResolver syntheticAccessorResolver = null;
 

baksmali/src/main/java/org/jf/baksmali/main.java

@@ -193,6 +193,9 @@ public class main {
                 case 'a':
                     options.apiLevel = Integer.parseInt(commandLine.getOptionValue("a"));
                     break;
+                case 'j':
+                    options.jobs = Integer.parseInt(commandLine.getOptionValue("j"));
+                    break;
                 case 'N':
                     disassemble = false;
                     break;
@@ -219,6 +222,13 @@ public class main {
             return;
         }
 
+        if (options.jobs <= 0) {
+            options.jobs = Runtime.getRuntime().availableProcessors();
+            if (options.jobs > 6) {
+                options.jobs = 6;
+            }
+        }
+
         String inputDexFileName = remainingArgs[0];
 
         File dexFileFile = new File(inputDexFileName);
@@ -375,6 +385,13 @@ public class main {
                 .withArgName("API_LEVEL")
                 .create("a");
 
+        Option jobsOption = OptionBuilder.withLongOpt("jobs")
+                .withDescription("The number of threads to use. Defaults to the number of cores available, up to a " +
+                        "maximum of 6")
+                .hasArg()
+                .withArgName("NUM_THREADS")
+                .create("j");
+
         Option dumpOption = OptionBuilder.withLongOpt("dump-to")
                 .withDescription("dumps the given dex file into a single annotated dump file named FILE" +
                         " (<dexfile>.dump by default), along with the normal disassembly")
@@ -418,6 +435,7 @@ public class main {
         basicOptions.addOption(codeOffsetOption);
         basicOptions.addOption(noAccessorCommentsOption);
         basicOptions.addOption(apiLevelOption);
+        basicOptions.addOption(jobsOption);
 
         debugOptions.addOption(dumpOption);
         debugOptions.addOption(ignoreErrorsOption);

dexlib2/src/main/java/org/jf/dexlib2/analysis/ClassPath.java

@@ -31,6 +31,9 @@
 
 package org.jf.dexlib2.analysis;
 
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
@@ -53,10 +56,8 @@ import java.util.regex.Pattern;
 
 public class ClassPath {
     @Nonnull private final TypeProto unknownClass;
-    @Nonnull private DexFile[] dexFiles;
-    @Nonnull private HashMap<String, TypeProto> loadedClasses = Maps.newHashMap();
     @Nonnull private HashMap<String, ClassDef> availableClasses = Maps.newHashMap();
-    @Nonnull private int api;
+    private int api;
 
     /**
      * Creates a new ClassPath instance that can load classes from the given dex files
@@ -78,6 +79,7 @@ public class ClassPath {
     }
 
     private ClassPath(@Nonnull DexFile[] classPath, boolean copyArray, int api) {
+        DexFile[] dexFiles;
         if (copyArray) {
             dexFiles = new DexFile[classPath.length+1];
             System.arraycopy(classPath, 0, dexFiles, 0, classPath.length);
@@ -128,22 +130,20 @@ public class ClassPath {
 
     @Nonnull
     public TypeProto getClass(CharSequence type) {
-        String typeString = type.toString();
+        return loadedClasses.getUnchecked(type.toString());
-        TypeProto typeProto = loadedClasses.get(typeString);
-        if (typeProto != null) {
-            return typeProto;
     }
 
+    private final CacheLoader<String, TypeProto> classLoader = new CacheLoader<String, TypeProto>() {
+        @Override public TypeProto load(String type) throws Exception {
             if (type.charAt(0) == '[') {
-            typeProto = new ArrayProto(this, typeString);
+                return new ArrayProto(ClassPath.this, type);
             } else {
-            typeProto = new ClassProto(this, typeString);
+                return new ClassProto(ClassPath.this, type);
             }
-        // All primitive types are preloaded into loadedClasses, so we don't need to check for that here
+        }
+    };
 
-        loadedClasses.put(typeString, typeProto);
+    @Nonnull private LoadingCache<String, TypeProto> loadedClasses = CacheBuilder.newBuilder().build(classLoader);
-        return typeProto;
-    }
 
     @Nonnull
     public ClassDef getClassDef(String type) {

dexlib2/src/main/java/org/jf/dexlib2/analysis/ClassProto.java

@@ -32,7 +32,12 @@
 package org.jf.dexlib2.analysis;
 
 import com.google.common.base.Predicates;
-import com.google.common.collect.*;
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
+import com.google.common.collect.FluentIterable;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
 import org.jf.dexlib2.AccessFlags;
 import org.jf.dexlib2.analysis.util.TypeProtoUtils;
 import org.jf.dexlib2.iface.ClassDef;
@@ -57,10 +62,7 @@ import java.util.List;
 public class ClassProto implements TypeProto {
     @Nonnull protected final ClassPath classPath;
     @Nonnull protected final String type;
-    @Nullable protected ClassDef classDef;
+
-    @Nullable protected LinkedHashMap<String, ClassDef> interfaces;
-    @Nullable protected List<Method> vtable;
-    @Nullable protected SparseArray<FieldReference> instanceFields;
     protected boolean vtableFullyResolved = true;
     protected boolean interfacesFullyResolved = true;
 
@@ -78,27 +80,15 @@ public class ClassProto implements TypeProto {
 
     @Nonnull
     public ClassDef getClassDef() {
-        if (classDef == null) {
+        return classDefSupplier.get();
-            classDef = classPath.getClassDef(type);
-        }
-        return classDef;
     }
 
-    @Nonnull
-    List<Method> getVtable() {
-        if (vtable == null) {
-            loadVtable();
-        }
-        return vtable;
-    }
 
-    @Nonnull
+    @Nonnull private final Supplier<ClassDef> classDefSupplier = Suppliers.memoize(new Supplier<ClassDef>() {
-    SparseArray<FieldReference> getInstanceFields() {
+        @Override public ClassDef get() {
-        if (instanceFields == null) {
+            return classPath.getClassDef(type);
-            instanceFields = loadFields();
-        }
-        return instanceFields;
         }
+    });
 
     /**
      * Returns true if this class is an interface.
@@ -128,11 +118,14 @@ public class ClassProto implements TypeProto {
      */
     @Nonnull
     protected LinkedHashMap<String, ClassDef> getInterfaces() {
-        if (interfaces != null) {
+        return interfacesSupplier.get();
-            return interfaces;
     }
 
-        interfaces = Maps.newLinkedHashMap();
+    @Nonnull
+    private final Supplier<LinkedHashMap<String, ClassDef>> interfacesSupplier =
+            Suppliers.memoize(new Supplier<LinkedHashMap<String, ClassDef>>() {
+                @Override public LinkedHashMap<String, ClassDef> get() {
+                    LinkedHashMap<String, ClassDef> interfaces = Maps.newLinkedHashMap();
 
                     try {
                         for (String interfaceType: getClassDef().getInterfaces()) {
@@ -187,6 +180,7 @@ public class ClassProto implements TypeProto {
 
                     return interfaces;
                 }
+            });
 
     /**
      * Gets the interfaces directly implemented by this class, or the interfaces they transitively implement.
@@ -194,15 +188,18 @@ public class ClassProto implements TypeProto {
      * This does not include any interfaces that are only implemented by a superclass
      *
      * @return An iterables of ClassDefs representing the directly or transitively implemented interfaces
-     * @throws UnresolvedClassException if any interface could not be resolved
+     * @throws UnresolvedClassException if interfaces could not be fully resolved
      */
     @Nonnull
     protected Iterable<ClassDef> getDirectInterfaces() {
+        Iterable<ClassDef> directInterfaces =
+                FluentIterable.from(getInterfaces().values()).filter(Predicates.notNull());
+
         if (!interfacesFullyResolved) {
             throw new UnresolvedClassException("Interfaces for class %s not fully resolved", getType());
         }
 
-        return FluentIterable.from(getInterfaces().values()).filter(Predicates.notNull());
+        return directInterfaces;
     }
 
     /**
@@ -213,6 +210,8 @@ public class ClassProto implements TypeProto {
      *
      * @param iface The interface to check for
      * @return true if this class implements the given interface, otherwise false
+     * @throws UnresolvedClassException if the interfaces for this class could not be fully resolved, and the interface
+     * is not one of the interfaces that were successfully resolved
      */
     @Override
     public boolean implementsInterface(@Nonnull String iface) {
@@ -354,8 +353,13 @@ public class ClassProto implements TypeProto {
         return vtable.get(vtableIndex);
     }
 
-    @Nonnull
+    @Nonnull SparseArray<FieldReference> getInstanceFields() {
-    private SparseArray<FieldReference> loadFields() {
+        return instanceFieldsSupplier.get();
+    }
+
+    @Nonnull private final Supplier<SparseArray<FieldReference>> instanceFieldsSupplier =
+            Suppliers.memoize(new Supplier<SparseArray<FieldReference>>() {
+                @Override public SparseArray<FieldReference> get() {
                     //This is a bit of an "involved" operation. We need to follow the same algorithm that dalvik uses to
                     //arrange fields, so that we end up with the same field offsets (which is needed for deodexing).
                     //See mydroid/dalvik/vm/oo/Class.c - computeFieldOffsets()
@@ -448,10 +452,13 @@ public class ClassProto implements TypeProto {
                         }
                     }
 
-        int superFieldCount = 0;
+                    SparseArray<FieldReference> superFields;
                     if (superclass != null) {
-            superFieldCount = superclass.instanceFields.size();
+                        superFields = superclass.getInstanceFields();
+                    } else {
+                        superFields = new SparseArray<FieldReference>();
                     }
+                    int superFieldCount = superFields.size();
 
                     //now the fields are in the correct order. Add them to the SparseArray and lookup, and calculate the offsets
                     int totalFieldCount = superFieldCount + fieldCount;
@@ -461,12 +468,12 @@ public class ClassProto implements TypeProto {
 
                     if (superclass != null && superFieldCount > 0) {
                         for (int i=0; i<superFieldCount; i++) {
-                instanceFields.append(superclass.instanceFields.keyAt(i), superclass.instanceFields.valueAt(i));
+                            instanceFields.append(superFields.keyAt(i), superFields.valueAt(i));
                         }
 
                         fieldOffset = instanceFields.keyAt(superFieldCount-1);
 
-            FieldReference lastSuperField = superclass.instanceFields.valueAt(superFieldCount-1);
+                        FieldReference lastSuperField = superFields.valueAt(superFieldCount-1);
                         char fieldType = lastSuperField.getType().charAt(0);
                         if (fieldType == 'J' || fieldType == 'D') {
                             fieldOffset += 8;
@@ -505,7 +512,7 @@ public class ClassProto implements TypeProto {
                 }
 
                 @Nonnull
-    private static ArrayList<Field> getSortedInstanceFields(@Nonnull ClassDef classDef) {
+                private ArrayList<Field> getSortedInstanceFields(@Nonnull ClassDef classDef) {
                     ArrayList<Field> fields = Lists.newArrayList(classDef.getInstanceFields());
                     Collections.sort(fields);
                     return fields;
@@ -532,6 +539,7 @@ public class ClassProto implements TypeProto {
                     Field tempField = fields.set(position1, fields.get(position2));
                     fields.set(position2, tempField);
                 }
+            });
 
     private int getNextFieldOffset() {
         SparseArray<FieldReference> instanceFields = getInstanceFields();
@@ -552,9 +560,14 @@ public class ClassProto implements TypeProto {
         }
     }
 
+    @Nonnull List<Method> getVtable() {
+        return vtableSupplier.get();
+    }
+
     //TODO: check the case when we have a package private method that overrides an interface method
-    private void loadVtable() {
+    @Nonnull private final Supplier<List<Method>> vtableSupplier = Suppliers.memoize(new Supplier<List<Method>>() {
-        vtable = Lists.newArrayList();
+        @Override public List<Method> get() {
+            List<Method> vtable = Lists.newArrayList();
 
             //copy the virtual methods from the superclass
             String superclassType;
@@ -563,7 +576,7 @@ public class ClassProto implements TypeProto {
             } catch (UnresolvedClassException ex) {
                 vtable.addAll(((ClassProto)classPath.getClass("Ljava/lang/Object;")).getVtable());
                 vtableFullyResolved = false;
-            return;
+                return vtable;
             }
 
             if (superclassType != null) {
@@ -572,9 +585,9 @@ public class ClassProto implements TypeProto {
 
                 // if the superclass's vtable wasn't fully resolved, then we can't know where the new methods added by this
                 // class should start, so we just propagate what we can from the parent and hope for the best.
-            if (!superclass.interfacesFullyResolved) {
+                if (!superclass.vtableFullyResolved) {
                     vtableFullyResolved = false;
-                return;
+                    return vtable;
                 }
             }
 
@@ -587,6 +600,7 @@ public class ClassProto implements TypeProto {
                     addToVtable(interfaceDef.getVirtualMethods(), vtable, false);
                 }
             }
+            return vtable;
         }
 
         private void addToVtable(@Nonnull Iterable<? extends Method> localMethods,
@@ -611,7 +625,7 @@ public class ClassProto implements TypeProto {
             }
         }
 
-    private static boolean methodSignaturesMatch(@Nonnull Method a, @Nonnull Method b) {
+        private boolean methodSignaturesMatch(@Nonnull Method a, @Nonnull Method b) {
             return (a.getName().equals(b.getName()) &&
                     a.getReturnType().equals(b.getReturnType()) &&
                     a.getParameters().equals(b.getParameters()));
@@ -628,7 +642,7 @@ public class ClassProto implements TypeProto {
         }
 
         @Nonnull
-    private static String getPackage(@Nonnull String classType) {
+        private String getPackage(@Nonnull String classType) {
             int lastSlash = classType.lastIndexOf('/');
             if (lastSlash < 0) {
                 return "";
@@ -636,9 +650,10 @@ public class ClassProto implements TypeProto {
             return classType.substring(1, lastSlash);
         }
 
-    private static boolean methodIsPackagePrivate(int accessFlags) {
+        private boolean methodIsPackagePrivate(int accessFlags) {
             return (accessFlags & (AccessFlags.PRIVATE.getValue() |
                     AccessFlags.PROTECTED.getValue() |
                     AccessFlags.PUBLIC.getValue())) == 0;
         }
+    });
 }

util/src/main/java/org/jf/util/ClassFileNameHandler.java

@@ -159,7 +159,7 @@ public class ClassFileNameHandler {
         }
 
         @Override
-        public File addUniqueChild(String[] pathElements, int pathElementsIndex) {
+        public synchronized File addUniqueChild(String[] pathElements, int pathElementsIndex) {
             String elementName;
             String elementNameLower;