Add BatchedIterator

src/main/java/lanchon/multidexlib2/BatchedIterator.java

@@ -0,0 +1,61 @@
+/*
+ * DexPatcher - Copyright 2015-2017 Rodrigo Balerdi
+ * (GNU General Public License version 3 or later)
+ *
+ * DexPatcher is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published
+ * by the Free Software Foundation, either version 3 of the License,
+ * or (at your option) any later version.
+ */
+
+package lanchon.multidexlib2;
+
+import java.util.ArrayDeque;
+import java.util.Iterator;
+
+import com.google.common.collect.PeekingIterator;
+import com.google.common.collect.UnmodifiableIterator;
+
+public class BatchedIterator<E> extends UnmodifiableIterator<E> implements PeekingIterator<E> {
+
+	private final Iterator<? extends E> iterator;
+	private final Object iteratorLock;
+	private final int batchSize;
+	private final ArrayDeque<E> batch;
+
+	public BatchedIterator(Iterator<? extends E> iterator, Object iteratorLock, int batchSize) {
+		if (batchSize < 1) throw new IllegalArgumentException("batchSize");
+		this.iterator = iterator;
+		this.iteratorLock = iteratorLock;
+		this.batchSize = batchSize;
+		batch = new ArrayDeque<>(batchSize);
+		preloadBatch();
+	}
+
+	@Override
+	public boolean hasNext() {
+		return !batch.isEmpty();
+	}
+
+	@Override
+	public E peek() {
+		return batch.element();                     // element throws NoSuchElementException if batch is empty
+	}
+
+	@Override
+	public E next() {
+		E item = batch.remove();                    // remove throws NoSuchElementException if batch is empty
+		if (batch.isEmpty()) preloadBatch();
+		return item;
+	}
+
+	public void preloadBatch() {
+		synchronized (iteratorLock) {
+			for (int n = batchSize - batch.size(); n > 0; n--) {
+				if (!iterator.hasNext()) break;
+				batch.add(iterator.next());         // add throws NullPointerException if element is null
+			}
+		}
+	}
+
+}

src/main/java/lanchon/multidexlib2/DexIO.java

@@ -14,12 +14,9 @@ import java.io.File;
 import java.io.IOException;
 import java.io.InterruptedIOException;
 import java.lang.reflect.UndeclaredThrowableException;
-import java.util.ArrayDeque;
 import java.util.ArrayList;
-import java.util.Deque;
 import java.util.Iterator;
 import java.util.List;
-import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
@@ -27,6 +24,8 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 
+import com.google.common.collect.Iterators;
+import com.google.common.collect.PeekingIterator;
 import org.jf.dexlib2.Opcodes;
 import org.jf.dexlib2.iface.ClassDef;
 import org.jf.dexlib2.iface.DexFile;
@@ -68,24 +67,28 @@ public class DexIO {
 		}
 		Object lock = new Object();
 		synchronized (lock) {       // avoid multiple synchronizations in single-threaded mode
-			writeCommon(base, nameIterator, currentName, currentFile, classes.iterator(), minMainDexClassCount,
+			writeCommon(base, nameIterator, currentName, currentFile, Iterators.peekingIterator(classes.iterator()),
-					minimalMainDex, dexFile.getOpcodes(), maxDexPoolSize, logger, lock);
+					minMainDexClassCount, minimalMainDex, dexFile.getOpcodes(), maxDexPoolSize, logger, lock);
 		}
 	}
 
 	// Multi-Threaded Write
 
+	private static final int PER_THREAD_BATCH_SIZE = 100;
+
 	static void writeMultiDexDirectoryMultiThread(int threadCount, final File directory,
 			final DexFileNameIterator nameIterator, final DexFile dexFile, final int maxDexPoolSize,
 			final DexIO.Logger logger) throws IOException {
-		final Iterator<? extends ClassDef> classIterator = dexFile.getClasses().iterator();
+		Iterator<? extends ClassDef> classIterator = dexFile.getClasses().iterator();
 		final Object lock = new Object();
 		List<Callable<Void>> callables = new ArrayList<>(threadCount);
 		for (int i = 0; i < threadCount; i++) {
+			final BatchedIterator<ClassDef> batchedIterator =
+					new BatchedIterator<>(classIterator, lock, PER_THREAD_BATCH_SIZE);
 			callables.add(new Callable<Void>() {
 				@Override
 				public Void call() throws IOException {
-					writeCommon(directory, nameIterator, null, null, classIterator, 0, false, dexFile.getOpcodes(),
+					writeCommon(directory, nameIterator, null, null, batchedIterator, 0, false, dexFile.getOpcodes(),
 							maxDexPoolSize, logger, lock);
 					return null;
 				}
@@ -116,24 +119,21 @@ public class DexIO {
 
 	// Common Code
 
-	private static final int PER_THREAD_BATCH_SIZE = 100;
-
 	private static void writeCommon(File base, DexFileNameIterator nameIterator, String currentName, File currentFile,
-			Iterator<? extends ClassDef> classIterator, int minMainDexClassCount, boolean minimalMainDex,
+			PeekingIterator<? extends ClassDef> classIterator, int minMainDexClassCount, boolean minimalMainDex,
 			Opcodes opcodes, int maxDexPoolSize, DexIO.Logger logger, Object lock) throws IOException {
-		Deque<ClassDef> queue = new ArrayDeque<>(PER_THREAD_BATCH_SIZE);
-		ClassDef currentClass = getQueueItem(queue, classIterator, lock);
 		do {
 			DexPool dexPool = new DexPool(opcodes);
 			int fileClassCount = 0;
-			while (currentClass != null) {
+			while (classIterator.hasNext()) {
 				if (minimalMainDex && fileClassCount >= minMainDexClassCount) break;
-				if (!internClass(dexPool, currentClass, maxDexPoolSize)) {
+				ClassDef classDef = classIterator.peek();
-					checkDexPoolOverflow(currentClass, fileClassCount, minMainDexClassCount);
+				if (!internClass(dexPool, classDef, maxDexPoolSize)) {
+					checkDexPoolOverflow(classDef, fileClassCount, minMainDexClassCount);
 					break;
 				}
+				classIterator.next();
 				fileClassCount++;
-				currentClass = getQueueItem(queue, classIterator, lock);
 			}
 			synchronized (lock) {
 				if (currentFile == null) {
@@ -141,13 +141,13 @@ public class DexIO {
 					currentFile = new File(base, currentName);
 				}
 				if (logger != null) logger.log(base, currentName, fileClassCount);
-				fillQueue(queue, classIterator, PER_THREAD_BATCH_SIZE - 1);
+				if (classIterator instanceof BatchedIterator) ((BatchedIterator) classIterator).preloadBatch();
 			}
 			dexPool.writeTo(new FileDataStore(currentFile));
 			currentFile = null;
 			minMainDexClassCount = 0;
 			minimalMainDex = false;
-		} while (currentClass != null);
+		} while (classIterator.hasNext());
 	}
 
 	private static boolean internClass(DexPool dexPool, ClassDef classDef, int maxDexPoolSize) {
@@ -174,22 +174,4 @@ public class DexIO {
 				"Type too big for dex pool: " + classDef.getType());
 	}
 
-	private static <T> T getQueueItem(Queue<T> queue, Iterator<? extends T> iterator, Object lock) {
-		T item = queue.poll();
-		if (item == null) {
-			synchronized (lock) {
-				fillQueue(queue, iterator, PER_THREAD_BATCH_SIZE);
-			}
-			item = queue.poll();
-		}
-		return item;
-	}
-
-	private static <T> void fillQueue(Queue<T> queue, Iterator<? extends T> iterator, int targetSize) {
-		for (int i = queue.size(); i < targetSize; i++) {
-			if (!iterator.hasNext()) break;
-			queue.add(iterator.next());
-		}
-	}
-
 }