Add annotation lookup functionality to AnnotationDirectoryItem

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

@@ -31,9 +31,6 @@ package org.jf.baksmali.Adaptors;
 import org.jf.dexlib.Util.Utf8Utils;
 import org.jf.util.IndentingWriter;
 import org.jf.dexlib.*;
-import static org.jf.dexlib.AnnotationDirectoryItem.FieldAnnotation;
-import static org.jf.dexlib.AnnotationDirectoryItem.MethodAnnotation;
-import static org.jf.dexlib.AnnotationDirectoryItem.ParameterAnnotation;
 import org.jf.dexlib.Code.Analysis.ValidationException;
 import org.jf.dexlib.Code.Format.Instruction21c;
 import org.jf.dexlib.Code.Format.Instruction41c;
@@ -49,10 +46,6 @@ public class ClassDefinition {
     private ClassDefItem classDefItem;
     private ClassDataItem classDataItem;
 
-    private SparseArray<AnnotationSetItem> methodAnnotationsMap;
-    private SparseArray<AnnotationSetItem> fieldAnnotationsMap;
-    private SparseArray<AnnotationSetRefList> parameterAnnotationsMap;
-
     private SparseArray<FieldIdItem> fieldsSetInStaticConstructor;
 
     protected boolean validationErrors;
@@ -60,7 +53,6 @@ public class ClassDefinition {
     public ClassDefinition(ClassDefItem classDefItem) {
         this.classDefItem = classDefItem;
         this.classDataItem = classDefItem.getClassData();
-        buildAnnotationMaps();
         findFieldsSetInStaticConstructor();
     }
 
@@ -68,40 +60,6 @@ public class ClassDefinition {
         return validationErrors;
     }
 
-    private void buildAnnotationMaps() {
-        AnnotationDirectoryItem annotationDirectory = classDefItem.getAnnotations();
-        if (annotationDirectory == null) {
-            methodAnnotationsMap = new SparseArray<AnnotationSetItem>(0);
-            fieldAnnotationsMap = new SparseArray<AnnotationSetItem>(0);
-            parameterAnnotationsMap = new SparseArray<AnnotationSetRefList>(0);
-            return;
-        }
-
-        int fieldAnnotationCount = annotationDirectory.getFieldAnnotationCount();
-        fieldAnnotationsMap = new SparseArray<AnnotationSetItem>(fieldAnnotationCount);
-        if (fieldAnnotationCount > 0) {
-            for (FieldAnnotation fieldAnnotation: annotationDirectory.getFieldAnnotations()) {
-                fieldAnnotationsMap.put(fieldAnnotation.field.getIndex(), fieldAnnotation.annotationSet);
-            }
-        }
-
-        int methodAnnotationCount = annotationDirectory.getMethodAnnotationCount();
-        methodAnnotationsMap = new SparseArray<AnnotationSetItem>(methodAnnotationCount);
-        if (methodAnnotationCount > 0) {
-            for (MethodAnnotation methodAnnotation: annotationDirectory.getMethodAnnotations()) {
-                methodAnnotationsMap.put(methodAnnotation.method.getIndex(), methodAnnotation.annotationSet);
-            }
-        }
-
-        int parameterAnnotationCount = annotationDirectory.getParameterAnnotationCount();
-        parameterAnnotationsMap = new SparseArray<AnnotationSetRefList>(parameterAnnotationCount);
-        if (parameterAnnotationCount > 0) {
-            for (ParameterAnnotation parameterAnnotation: annotationDirectory.getParameterAnnotations()) {
-                parameterAnnotationsMap.put(parameterAnnotation.method.getIndex(), parameterAnnotation.annotationSet);
-            }
-        }
-    }
-
     private void findFieldsSetInStaticConstructor() {
         fieldsSetInStaticConstructor = new SparseArray<FieldIdItem>();
 
@@ -154,7 +112,6 @@ public class ClassDefinition {
         writeInstanceFields(writer);
         writeDirectMethods(writer);
         writeVirtualMethods(writer);
-        return ;
     }
 
     private void writeClass(IndentingWriter writer) throws IOException {
@@ -261,12 +218,16 @@ public class ClassDefinition {
             if (i < staticInitializers.length) {
                 encodedValue = staticInitializers[i];
             }
-            AnnotationSetItem annotationSet = fieldAnnotationsMap.get(field.field.getIndex());
+            AnnotationSetItem fieldAnnotations = null;
+            AnnotationDirectoryItem annotations = classDefItem.getAnnotations();
+            if (annotations != null) {
+                fieldAnnotations = annotations.getFieldAnnotations(field.field);
+            }
 
             boolean setInStaticConstructor =
                     fieldsSetInStaticConstructor.get(field.field.getIndex()) != null;
 
-            FieldDefinition.writeTo(writer, field, encodedValue, annotationSet, setInStaticConstructor);
+            FieldDefinition.writeTo(writer, field, encodedValue, fieldAnnotations, setInStaticConstructor);
         }
     }
 
@@ -283,15 +244,19 @@ public class ClassDefinition {
         writer.write("\n\n");
         writer.write("# instance fields\n");
         boolean first = true;
-        for (ClassDataItem.EncodedField field: classDataItem.getInstanceFields()) {
+        for (ClassDataItem.EncodedField field: encodedFields) {
             if (!first) {
                 writer.write('\n');
             }
             first = false;
 
-            AnnotationSetItem annotationSet = fieldAnnotationsMap.get(field.field.getIndex());
+            AnnotationSetItem fieldAnnotations = null;
+            AnnotationDirectoryItem annotations = classDefItem.getAnnotations();
+            if (annotations != null) {
+                fieldAnnotations = annotations.getFieldAnnotations(field.field);
+            }
 
-            FieldDefinition.writeTo(writer, field, null, annotationSet, false);
+            FieldDefinition.writeTo(writer, field, null, fieldAnnotations, false);
         }
     }
 
@@ -335,11 +300,16 @@ public class ClassDefinition {
             }
             first = false;
 
-            AnnotationSetItem annotationSet = methodAnnotationsMap.get(method.method.getIndex());
+            AnnotationSetItem methodAnnotations = null;
-            AnnotationSetRefList parameterAnnotationList = parameterAnnotationsMap.get(method.method.getIndex());
+            AnnotationSetRefList parameterAnnotations = null;
+            AnnotationDirectoryItem annotations = classDefItem.getAnnotations();
+            if (annotations != null) {
+                methodAnnotations = annotations.getMethodAnnotations(method.method);
+                parameterAnnotations = annotations.getParameterAnnotations(method.method);
+            }
 
             MethodDefinition methodDefinition = new MethodDefinition(method);
-            methodDefinition.writeTo(writer, annotationSet, parameterAnnotationList);
+            methodDefinition.writeTo(writer, methodAnnotations, parameterAnnotations);
 
             ValidationException validationException = methodDefinition.getValidationException();
             if (validationException != null) {

dexlib/src/main/java/org/jf/dexlib/AnnotationDirectoryItem.java

@@ -36,9 +36,7 @@ import org.jf.dexlib.Util.ReadOnlyArrayList;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
-import java.util.Arrays;
+import java.util.*;
-import java.util.Collections;
-import java.util.List;
 
 public class AnnotationDirectoryItem extends Item<AnnotationDirectoryItem> {
     @Nullable
@@ -375,6 +373,58 @@ public class AnnotationDirectoryItem extends Item<AnnotationDirectoryItem> {
         return ReadOnlyArrayList.of(parameterAnnotations);
     }
 
+    /**
+     * Gets the field annotations for the given field, or null if no annotations are defined for that field
+     * @param fieldIdItem The field to get the annotations for
+     * @return An <code>AnnotationSetItem</code> containing the field annotations, or null if none are found
+     */
+    @Nullable
+    public AnnotationSetItem getFieldAnnotations(FieldIdItem fieldIdItem) {
+        if (fieldAnnotations == null) {
+            return null;
+        }
+        int index = Arrays.binarySearch(fieldAnnotations, fieldIdItem);
+        if (index < 0) {
+            return null;
+        }
+        return fieldAnnotations[index].annotationSet;
+    }
+
+    /**
+     * Gets the method annotations for the given method, or null if no annotations are defined for that method
+     * @param methodIdItem The method to get the annotations for
+     * @return An <code>AnnotationSetItem</code> containing the method annotations, or null if none are found
+     */
+    @Nullable
+    public AnnotationSetItem getMethodAnnotations(MethodIdItem methodIdItem) {
+        if (methodAnnotations == null) {
+            return null;
+        }
+        int index = Arrays.binarySearch(methodAnnotations, methodIdItem);
+        if (index < 0) {
+            return null;
+        }
+        return methodAnnotations[index].annotationSet;
+    }
+
+    /**
+     * Gets the parameter annotations for the given method, or null if no parameter annotations are defined for that
+     * method
+     * @param methodIdItem The method to get the parameter annotations for
+     * @return An <code>AnnotationSetRefList</code> containing the parameter annotations, or null if none are found
+     */
+    @Nullable
+    public AnnotationSetRefList getParameterAnnotations(MethodIdItem methodIdItem) {
+        if (parameterAnnotations == null) {
+            return null;
+        }
+        int index = Arrays.binarySearch(parameterAnnotations, methodIdItem);
+        if (index < 0) {
+            return null;
+        }
+        return parameterAnnotations[index].annotationSet;
+    }
+
     /**
      * @return The number of field annotations in this <code>AnnotationDirectoryItem</code>
      */
@@ -454,7 +504,7 @@ public class AnnotationDirectoryItem extends Item<AnnotationDirectoryItem> {
         return (this.compareTo(other) == 0);
     }
 
-    public static class FieldAnnotation implements Comparable<FieldAnnotation> {
+    public static class FieldAnnotation implements Comparable<Convertible<FieldIdItem>>, Convertible<FieldIdItem> {
         public final FieldIdItem field;
         public final AnnotationSetItem annotationSet;
 
@@ -463,8 +513,8 @@ public class AnnotationDirectoryItem extends Item<AnnotationDirectoryItem> {
             this.annotationSet = annotationSet;
         }
 
-        public int compareTo(FieldAnnotation other) {
+        public int compareTo(Convertible<FieldIdItem> other) {
-            return field.compareTo(other.field);
+            return field.compareTo(other.convert());
         }
 
         @Override
@@ -479,9 +529,13 @@ public class AnnotationDirectoryItem extends Item<AnnotationDirectoryItem> {
         public int hashCode() {
             return field.hashCode() + 31 * annotationSet.hashCode();
         }
+
+        public FieldIdItem convert() {
+            return field;
+        }
     }
 
-    public static class MethodAnnotation implements Comparable<MethodAnnotation> {
+    public static class MethodAnnotation implements Comparable<Convertible<MethodIdItem>>, Convertible<MethodIdItem> {
         public final MethodIdItem method;
         public final AnnotationSetItem annotationSet;
 
@@ -490,8 +544,8 @@ public class AnnotationDirectoryItem extends Item<AnnotationDirectoryItem> {
             this.annotationSet = annotationSet;
         }
 
-        public int compareTo(MethodAnnotation other) {
+        public int compareTo(Convertible<MethodIdItem> other) {
-            return method.compareTo(other.method);
+            return method.compareTo(other.convert());
         }
 
         @Override
@@ -506,9 +560,14 @@ public class AnnotationDirectoryItem extends Item<AnnotationDirectoryItem> {
         public int hashCode() {
             return method.hashCode() + 31 * annotationSet.hashCode();
         }
+
+        public MethodIdItem convert() {
+            return method;
+        }
     }
 
-    public static class ParameterAnnotation implements Comparable<ParameterAnnotation> {
+    public static class ParameterAnnotation implements Comparable<Convertible<MethodIdItem>>,
+            Convertible<MethodIdItem> {
         public final MethodIdItem method;
         public final AnnotationSetRefList annotationSet;
 
@@ -517,8 +576,8 @@ public class AnnotationDirectoryItem extends Item<AnnotationDirectoryItem> {
             this.annotationSet = annotationSet;
         }
 
-        public int compareTo(ParameterAnnotation other) {
+        public int compareTo(Convertible<MethodIdItem> other) {
-            return method.compareTo(other.method);
+            return method.compareTo(other.convert());
         }
 
         @Override
@@ -533,5 +592,9 @@ public class AnnotationDirectoryItem extends Item<AnnotationDirectoryItem> {
         public int hashCode() {
             return method.hashCode() + 31 * annotationSet.hashCode();
         }
+
+        public MethodIdItem convert() {
+            return method;
+        }
     }
 }

dexlib/src/main/java/org/jf/dexlib/Convertible.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2012, Google Inc.
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *     * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *     * Neither the name of Google Inc. nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package org.jf.dexlib;
+
+/**
+ * Describes an object that can be converted to a different type
+ */
+public interface Convertible<T> {
+    T convert();
+}

dexlib/src/main/java/org/jf/dexlib/FieldIdItem.java

@@ -31,7 +31,7 @@ package org.jf.dexlib;
 import org.jf.dexlib.Util.AnnotatedOutput;
 import org.jf.dexlib.Util.Input;
 
-public class FieldIdItem extends Item<FieldIdItem> {
+public class FieldIdItem extends Item<FieldIdItem> implements Convertible<FieldIdItem> {
     private int hashCode = 0;
 
     private TypeIdItem classType;
@@ -237,4 +237,8 @@ public class FieldIdItem extends Item<FieldIdItem> {
                 fieldType == other.fieldType &&
                 fieldName == other.fieldName);
     }
+
+    public FieldIdItem convert() {
+        return this;
+    }
 }

dexlib/src/main/java/org/jf/dexlib/MethodIdItem.java

@@ -31,7 +31,7 @@ package org.jf.dexlib;
 import org.jf.dexlib.Util.AnnotatedOutput;
 import org.jf.dexlib.Util.Input;
 
-public class MethodIdItem extends Item<MethodIdItem> {
+public class MethodIdItem extends Item<MethodIdItem> implements Convertible<MethodIdItem> {
     private int hashCode = 0;
 
     private TypeIdItem classType;
@@ -249,4 +249,8 @@ public class MethodIdItem extends Item<MethodIdItem> {
                 methodPrototype == other.methodPrototype &&
                 methodName == other.methodName);
     }
+
+    public MethodIdItem convert() {
+        return this;
+    }
 }