Rewriter bug fixes (#72)

* Changing PoolClassDef to sort virtualMethods on constructor. * Separating ChainedIterator and TransformedIterator into its iterable part and the iterator. This should fix writer/rewriter issues. * Fixing indentation in TransformedIterable
google · Sep 12, 2024 · c44a1de · c44a1de
1 parent a890d24
commit c44a1de
Show file tree

Hide file tree

Showing 12 changed files with 107 additions and 96 deletions.
diff --git a/...src/main/java/com/android/tools/smali/dexlib2/analysis/reflection/ReflectionClassDef.java b/...src/main/java/com/android/tools/smali/dexlib2/analysis/reflection/ReflectionClassDef.java
@@ -35,9 +35,9 @@
 import com.android.tools.smali.dexlib2.iface.ClassDef;
 import com.android.tools.smali.dexlib2.iface.Field;
 import com.android.tools.smali.dexlib2.iface.Method;
-import com.android.tools.smali.util.ChainedIterator;
+import com.android.tools.smali.util.ChainedIterable.ChainedIterator;
 import com.android.tools.smali.util.IteratorUtils;
-import com.android.tools.smali.util.TransformedIterator;
+import com.android.tools.smali.util.TransformedIterable.TransformedIterator;
 
 import com.android.tools.smali.dexlib2.base.reference.BaseTypeReference;
 

diff --git a/dexlib2/src/main/java/com/android/tools/smali/dexlib2/dexbacked/DexBackedClassDef.java b/dexlib2/src/main/java/com/android/tools/smali/dexlib2/dexbacked/DexBackedClassDef.java
@@ -43,7 +43,8 @@
 import com.android.tools.smali.dexlib2.immutable.reference.ImmutableFieldReference;
 import com.android.tools.smali.dexlib2.immutable.reference.ImmutableMethodReference;
 import com.android.tools.smali.dexlib2.writer.DexWriter;
-import com.android.tools.smali.util.ChainedIterator;
+import com.android.tools.smali.util.ChainedIterable;
+import com.android.tools.smali.util.ChainedIterable.ChainedIterator;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -73,8 +74,8 @@ public class DexBackedClassDef extends BaseTypeReference implements ClassDef {
     @Nullable private AnnotationsDirectory annotationsDirectory;
 
     public DexBackedClassDef(@Nonnull DexBackedDexFile dexFile,
-                             int classDefOffset,
-                             int hiddenApiRestrictionsOffset) {
+                            int classDefOffset,
+                            int hiddenApiRestrictionsOffset) {
         this.dexFile = dexFile;
         this.classDefOffset = classDefOffset;
 
@@ -301,7 +302,7 @@ protected DexBackedField readNextItem(@Nonnull DexReader<? extends DexBuffer> re
     @Nonnull
     @Override
     public Iterable<? extends DexBackedField> getFields() {
-        return new ChainedIterator(getStaticFields(), getInstanceFields());
+        return new ChainedIterable(getStaticFields(), getInstanceFields());
     }
 
     @Nonnull
@@ -448,7 +449,7 @@ public Iterable<? extends DexBackedMethod> getVirtualMethods() {
     @Nonnull
     @Override
     public Iterable<? extends DexBackedMethod> getMethods() {
-        return new ChainedIterator(getDirectMethods(), getVirtualMethods());
+        return new ChainedIterable(getDirectMethods(), getVirtualMethods());
     }
 
     private AnnotationsDirectory getAnnotationsDirectory() {

diff --git a/dexlib2/src/main/java/com/android/tools/smali/dexlib2/dexbacked/OatFile.java b/dexlib2/src/main/java/com/android/tools/smali/dexlib2/dexbacked/OatFile.java
@@ -37,7 +37,7 @@
 import com.android.tools.smali.dexlib2.util.DexUtil;
 import com.android.tools.smali.util.AbstractForwardSequentialList;
 import com.android.tools.smali.util.InputStreamUtil;
-import com.android.tools.smali.util.TransformedIterator;
+import com.android.tools.smali.util.TransformedIterable.TransformedIterator;
 
 import java.util.function.Function;
 import com.android.tools.smali.dexlib2.dexbacked.OatFile.SymbolTable.Symbol;

diff --git a/dexlib2/src/main/java/com/android/tools/smali/dexlib2/immutable/ImmutableClassDef.java b/dexlib2/src/main/java/com/android/tools/smali/dexlib2/immutable/ImmutableClassDef.java
@@ -39,7 +39,7 @@
 import com.android.tools.smali.dexlib2.iface.Method;
 import com.android.tools.smali.dexlib2.util.FieldUtil;
 import com.android.tools.smali.dexlib2.util.MethodUtil;
-import com.android.tools.smali.util.ChainedIterator;
+import com.android.tools.smali.util.ChainedIterable;
 import com.android.tools.smali.util.ImmutableConverter;
 import com.android.tools.smali.util.ImmutableUtils;
 import com.android.tools.smali.util.IteratorUtils;
@@ -168,13 +168,13 @@ public static ImmutableClassDef of(ClassDef classDef) {
     @Nonnull
     @Override
     public Iterable<? extends ImmutableField> getFields() {
-        return new ChainedIterator(staticFields, instanceFields);
+        return new ChainedIterable(staticFields, instanceFields);
     }
 
     @Nonnull
     @Override
     public Iterable<? extends ImmutableMethod> getMethods() {
-        return new ChainedIterator(directMethods, virtualMethods);
+        return new ChainedIterable(directMethods, virtualMethods);
     }
 
     @Nonnull

diff --git a/dexlib2/src/main/java/com/android/tools/smali/dexlib2/rewriter/ClassDefRewriter.java b/dexlib2/src/main/java/com/android/tools/smali/dexlib2/rewriter/ClassDefRewriter.java
@@ -35,11 +35,10 @@
 import com.android.tools.smali.dexlib2.iface.ClassDef;
 import com.android.tools.smali.dexlib2.iface.Field;
 import com.android.tools.smali.dexlib2.iface.Method;
-import com.android.tools.smali.util.ChainedIterator;
+import com.android.tools.smali.util.ChainedIterable;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Set;
 
@@ -96,14 +95,7 @@ public RewrittenClassDef(@Nonnull ClassDef classdef) {
         @Nonnull
         @Override
         public Iterable<? extends Field> getFields() {
-            return new Iterable<Field>() {
-                @Nonnull
-                @Override
-                public Iterator<Field> iterator() {
-                    return new ChainedIterator(
-                        getStaticFields().iterator(), getInstanceFields().iterator());
-                }
-            };
+            return new ChainedIterable(getStaticFields(), getInstanceFields());
         }
 
         @Override @Nonnull public Iterable<? extends Method> getDirectMethods() {
@@ -117,14 +109,7 @@ public Iterator<Field> iterator() {
         @Nonnull
         @Override
         public Iterable<? extends Method> getMethods() {
-            return new Iterable<Method>() {
-                @Nonnull
-                @Override
-                public Iterator<Method> iterator() {
-                    return new ChainedIterator(
-                        getDirectMethods().iterator(), getVirtualMethods().iterator());
-                }
-            };
+            return new ChainedIterable(getDirectMethods(), getVirtualMethods());
         }
     }
 }
diff --git a/dexlib2/src/main/java/com/android/tools/smali/dexlib2/writer/DexWriter.java b/dexlib2/src/main/java/com/android/tools/smali/dexlib2/writer/DexWriter.java
@@ -111,7 +111,7 @@
 import com.android.tools.smali.dexlib2.writer.io.DexDataStore;
 import com.android.tools.smali.dexlib2.writer.io.MemoryDeferredOutputStream;
 import com.android.tools.smali.dexlib2.writer.util.TryListBuilder;
-import com.android.tools.smali.util.ChainedIterator;
+import com.android.tools.smali.util.ChainedIterable;
 import com.android.tools.smali.util.CollectionUtils;
 import com.android.tools.smali.util.ExceptionWithContext;
 import com.android.tools.smali.util.IteratorUtils;
@@ -1059,7 +1059,7 @@ private void writeDebugAndCodeItems(@Nonnull DexDataWriter offsetWriter,
             Collection<? extends MethodKey> directMethods = classSection.getSortedDirectMethods(classKey);
             Collection<? extends MethodKey> virtualMethods = classSection.getSortedVirtualMethods(classKey);
 
-            Iterable<MethodKey> methods = new ChainedIterator<MethodKey>(
+            Iterable<MethodKey> methods = new ChainedIterable<MethodKey>(
                 (Collection<MethodKey>)directMethods, (Collection<MethodKey>)virtualMethods);
 
             for (MethodKey methodKey: methods) {

diff --git a/dexlib2/src/main/java/com/android/tools/smali/dexlib2/writer/builder/DexBuilder.java b/dexlib2/src/main/java/com/android/tools/smali/dexlib2/writer/builder/DexBuilder.java
@@ -88,7 +88,6 @@
 import com.android.tools.smali.util.CollectionUtils;
 import com.android.tools.smali.util.ExceptionWithContext;
 import com.android.tools.smali.util.IteratorUtils;
-import com.android.tools.smali.util.TransformedIterator;
 import com.android.tools.smali.dexlib2.writer.util.StaticInitializerUtil;
 
 import javax.annotation.Nonnull;
@@ -100,7 +99,6 @@
 import java.util.List;
 import java.util.Set;
 import java.util.SortedSet;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 
 public class DexBuilder extends DexWriter<BuilderStringReference, BuilderStringReference, BuilderTypeReference,

diff --git a/dexlib2/src/main/java/com/android/tools/smali/dexlib2/writer/pool/ClassPool.java b/dexlib2/src/main/java/com/android/tools/smali/dexlib2/writer/pool/ClassPool.java
@@ -61,7 +61,8 @@
 import com.android.tools.smali.util.AbstractForwardSequentialList;
 import com.android.tools.smali.util.CollectionUtils;
 import com.android.tools.smali.util.ExceptionWithContext;
-import com.android.tools.smali.util.TransformedIterator;
+import com.android.tools.smali.util.TransformedIterable;
+import com.android.tools.smali.util.TransformedIterable.TransformedIterator;
 import com.android.tools.smali.dexlib2.builder.MutableMethodImplementation;
 import com.android.tools.smali.dexlib2.formatter.DexFormatter;
 import com.android.tools.smali.dexlib2.iface.instruction.DualReferenceInstruction;
@@ -402,7 +403,7 @@ public Set<? extends Annotation> apply(MethodParameter input) {
     }
 
     @Nullable @Override public Iterable<CharSequence> getParameterNames(@Nonnull PoolMethod method) {
-        return new TransformedIterator(method.getParameters(), new Function<MethodParameter, CharSequence>() {
+        return new TransformedIterable(method.getParameters(), new Function<MethodParameter, CharSequence>() {
             @Nullable @Override public CharSequence apply(MethodParameter input) {
                 return input.getName();
             }

diff --git a/dexlib2/src/main/java/com/android/tools/smali/dexlib2/writer/pool/PoolClassDef.java b/dexlib2/src/main/java/com/android/tools/smali/dexlib2/writer/pool/PoolClassDef.java
@@ -72,7 +72,7 @@ class PoolClassDef extends BaseTypeReference implements ClassDef {
         directMethods = ArraySortedSet.copyOf(CollectionUtils.naturalOrdering(),
                 IteratorUtils.toList(classDef.getDirectMethods()).stream().map(PoolMethod.TRANSFORM)
                         .collect(Collectors.toList()));
-        virtualMethods = ArraySortedSet.of(CollectionUtils.naturalOrdering(),
+        virtualMethods = ArraySortedSet.copyOf(CollectionUtils.naturalOrdering(),
                 IteratorUtils.toList(classDef.getVirtualMethods()).stream().map(PoolMethod.TRANSFORM)
                 .collect(Collectors.toList()));
     }

diff --git a/...oid/tools/smali/util/ChainedIterator.java → ...oid/tools/smali/util/ChainedIterable.java b/...oid/tools/smali/util/ChainedIterator.java → ...oid/tools/smali/util/ChainedIterable.java
@@ -34,46 +34,54 @@
 import java.util.Iterator;
 
 /**
- * Combines two iterators into a single iterator. The returned iterator iterates across the elements
- * in {@code a}, followed by the elements in {@code b}. The source iterators are not polled until
- * necessary.
- * <p>
- * The returned iterator does not support {@code remove()}.
+ * An iterable wrapping a {@code ChainedIterator}.
  */
-public class ChainedIterator<T extends Object> implements Iterator<T>, Iterable<T> {
-    Iterator<T> iteratorA;
-    Iterator<T> iteratorB;
+public class ChainedIterable<T extends Object> implements Iterable<T> {
+    Iterable<T> iterableA;
+    Iterable<T> iterableB;
 
-    public ChainedIterator(Iterable<T> iterableA, Iterable<T> iterableB) {
-        this.iteratorA = iterableA.iterator();
-        this.iteratorB = iterableB.iterator();
-    }
-
-    public ChainedIterator(Iterator<T> iteratorA, Iterator<T> iteratorB) {
-        this.iteratorA = iteratorA;
-        this.iteratorB = iteratorB;
+    public ChainedIterable(Iterable<T> iterableA, Iterable<T> iterableB) {
+        this.iterableA = iterableA;
+        this.iterableB = iterableB;
     }
 
     @Override
-    public final boolean hasNext() {
-        return iteratorA.hasNext() || iteratorB.hasNext();
+    public final Iterator<T> iterator() {
+        return new ChainedIterator(iterableA.iterator(), iterableB.iterator());
     }
 
-    @Override
-    public final T next() {
-        if (iteratorA.hasNext()) {
-            return iteratorA.next();
+    /**
+     * Combines two iterators into a single iterator. The returned iterator iterates across the elements
+     * in {@code a}, followed by the elements in {@code b}. The source iterators are not polled until
+     * necessary.
+     * <p>
+     * The returned iterator does not support {@code remove()}.
+     */
+    public static class ChainedIterator<U extends Object> implements Iterator<U> {
+        Iterator<U> iteratorA;
+        Iterator<U> iteratorB; 
+
+        public ChainedIterator(Iterator<U> iteratorA, Iterator<U> iteratorB) {
+            this.iteratorA = iteratorA;
+            this.iteratorB = iteratorB;
         }
-        return iteratorB.next();
-    }
 
-    @Override
-    public final void remove() {
-        throw new UnsupportedOperationException();
-    }
+        @Override
+        public final boolean hasNext() {
+            return iteratorA.hasNext() || iteratorB.hasNext();
+        }
 
-    @Override
-    public final Iterator<T> iterator() {
-        return this;
+        @Override
+        public final U next() {
+            if (iteratorA.hasNext()) {
+                return iteratorA.next();
+            }
+            return iteratorB.next();
+        }
+
+        @Override
+        public final void remove() {
+            throw new UnsupportedOperationException();
+        }
     }
 }
diff --git a/dexlib2/src/main/java/com/android/tools/smali/util/CharSequenceUtils.java b/dexlib2/src/main/java/com/android/tools/smali/util/CharSequenceUtils.java
@@ -44,11 +44,11 @@ public String apply(Object o) {
     };
 
     public static int listHashCode(List<? extends CharSequence> list) {
-        return IteratorUtils.toList((Iterator)new TransformedIterator(list.iterator(), TO_STRING)).hashCode();
+        return IteratorUtils.toList(new TransformedIterable(list, TO_STRING)).hashCode();
     }
 
     public static boolean listEquals(List<? extends CharSequence> list1, List<? extends CharSequence> list2) {
-        return IteratorUtils.toList((Iterator)new TransformedIterator(list1.iterator(), TO_STRING)).equals(
-                IteratorUtils.toList((Iterator)new TransformedIterator(list2.iterator(), TO_STRING)));
+        return IteratorUtils.toList(new TransformedIterable(list1, TO_STRING)).equals(
+                IteratorUtils.toList(new TransformedIterable(list2, TO_STRING)));
     }
 }
diff --git a/...tools/smali/util/TransformedIterator.java → ...tools/smali/util/TransformedIterable.java b/...tools/smali/util/TransformedIterator.java → ...tools/smali/util/TransformedIterable.java
@@ -34,45 +34,63 @@
 import java.util.function.Function;
 
 /**
- * An iterator that will return the results of applying {@code transformFunction} to each element of
- * {@code backingIterator}.
+ * An wrapping instance that will return the {@code TransformedIterator} of {@code backingIterator}
+ * and {@code transformFunction}.
  * <p>
  * The returned iterator supports {@code remove()} if {@code backingIterator} does.
  */
-public class TransformedIterator<F extends Object, T extends Object>
-        implements Iterator<T>, Iterable<T> {
-    final Iterator<? extends F> backingIterator;
+public class TransformedIterable<F extends Object, T extends Object>
+        implements Iterable<T> {
+    final Iterable<? extends F> backingIterable;
     final Function<F, T> transformFunction;
 
-    public TransformedIterator(Iterable<? extends F> backingIterable,
+    public TransformedIterable(Iterable<? extends F> backingIterable,
             Function<F, T> transformFunction) {
-        this.backingIterator = backingIterable.iterator();
-        this.transformFunction = transformFunction;
-    }
-
-    public TransformedIterator(Iterator<? extends F> backingIterator,
-            Function<F, T> transformFunction) {
-        this.backingIterator = backingIterator;
+        this.backingIterable = backingIterable;
         this.transformFunction = transformFunction;
     }
 
     @Override
-    public final boolean hasNext() {
-        return backingIterator.hasNext();
+    public final Iterator<T> iterator() {
+        return new TransformedIterator<>(backingIterable.iterator(), transformFunction);
     }
 
-    @Override
-    public final T next() {
-        return transformFunction.apply(backingIterator.next());
-    }
+    /**
+     * An iterator that will return the results of applying {@code transformFunction} to each
+     * element of {@code backingIterator}.
+     * <p>
+     * The returned iterator supports {@code remove()} if {@code backingIterator} does.
+     */
+    public static class TransformedIterator<G extends Object, U extends Object>
+            implements Iterator<U> {
+        final Iterator<? extends G> backingIterator;
+        final Function<G, U> transformFunction;
 
-    @Override
-    public final void remove() {
-        backingIterator.remove();
-    }
+        public TransformedIterator(Iterable<? extends G> backingIterable,
+                Function<G, U> transformFunction) {
+            this.backingIterator = backingIterable.iterator();
+            this.transformFunction = transformFunction;
+        }
 
-    @Override
-    public final Iterator<T> iterator() {
-        return this;
+        public TransformedIterator(Iterator<? extends G> backingIterator,
+                Function<G, U> transformFunction) {
+            this.backingIterator = backingIterator;
+            this.transformFunction = transformFunction;
+        }
+
+        @Override
+        public final boolean hasNext() {
+            return backingIterator.hasNext();
+        }
+
+        @Override
+        public final U next() {
+            return transformFunction.apply(backingIterator.next());
+        }
+
+        @Override
+        public final void remove() {
+            backingIterator.remove();
+        }
     }
 }