insert annotation in frag

TODO.md

@@ -5,4 +5,7 @@
 - no nop when no payload
 - option to get label at every code addresses
 - name register / parameters
-- ord in python
+
+# frag meth
+- insert_class_data_item
+- insert_class_static_values

androscalpel/src/dex_fragment.rs

@@ -151,17 +151,25 @@ impl DexFragment {
             .map(|(idx, method)| (method.clone(), idx))
             .collect::<HashMap<IdMethod, usize>>();
 
+        let index = FragIndex {
+            strings: strings_index,
+            types: types_index,
+            protos: protos_index,
+            fields: fields_index,
+            methods: methods_index,
+        };
+
         frag.section_manager
             .add_elt(FragSection::ClassDefItem, None);
 
-        frag.class_def.class_idx = *types_index.get(&class.descriptor).ok_or(anyhow!(
+        frag.class_def.class_idx = *index.types.get(&class.descriptor).ok_or(anyhow!(
             "Type {} (type of class {}) not found in type list",
             class.descriptor.__str__(),
             class.__str__()
         ))? as u32;
         frag.class_def.access_flags = class.get_raw_access_flags();
         frag.class_def.superclass_idx = if let Some(sup) = &class.superclass {
-            *types_index.get(sup).ok_or(anyhow!(
+            *index.types.get(sup).ok_or(anyhow!(
                 "Type {} (superclass of class {}) not found in dex builder",
                 sup.__str__(),
                 class.__str__()
@@ -171,7 +179,7 @@ impl DexFragment {
         };
         frag.class_def.interfaces_off = 0; // set when linking
         frag.class_def.source_file_idx = if let Some(file) = &class.source_file {
-            *strings_index.get(file).ok_or(anyhow!(
+            *index.strings.get(file).ok_or(anyhow!(
                 "String {} (source file of class {}) not found in dex builder",
                 file.__str__(),
                 class.__str__()
@@ -180,7 +188,7 @@ impl DexFragment {
             NO_INDEX.0
         };
         frag.class_def.annotations_off = if class.has_annotations() {
-            frag.insert_annotations(class)?;
+            frag.insert_annotations(class, &index)?;
             1
         } else {
             0
@@ -214,12 +222,13 @@ impl DexFragment {
     /// in the whole file or 0), their value is set to the actual value prelink value + 1. This allow
     /// to distinguish the offset of the first item (equal to zero before linking) and the value
     /// 0 used to indicate an abscence of item.
-    fn insert_annotations(&mut self, class: &Class) -> Result<()> {
+    fn insert_annotations(&mut self, class: &Class, index: &FragIndex) -> Result<()> {
         let class_annotations_off = if !class.annotations.is_empty() {
             let class_annotations_off = self
                 .section_manager
                 .get_aligned_size(FragSection::AnnotationSetItem);
-            self.insert_class_annotation_set(class).with_context(|| {
+            self.insert_class_annotation_set(class, index)
+                .with_context(|| {
                     format!(
                         "Failed to insert class annotation for class {}",
                         class.__str__()
@@ -231,13 +240,11 @@ impl DexFragment {
         };
 
         let mut field_ids = vec![];
-        let (class, _) = self.class_defs.get(class_id).unwrap();
         field_ids.extend(class.static_fields.keys().cloned());
         field_ids.extend(class.instance_fields.keys().cloned());
         field_ids.sort();
         let mut field_annotations = vec![];
         for field_id in field_ids {
-            let (class, _) = self.class_defs.get(class_id).unwrap();
             let static_field = class.static_fields.get(&field_id);
             let instance_field = class.instance_fields.get(&field_id);
             let (is_static, field) = match (static_field, instance_field) {
@@ -245,33 +252,31 @@ impl DexFragment {
                 (None, Some(field)) => (false, field),
                 _ => bail!(
                     "Unexpected configuration: field {} is both a static and a instance field in {}", 
-                    field_id.__str__(), class_id.__str__()),
+                    field_id.__str__(), class.__str__()),
             };
             if !field.annotations.is_empty() {
                 let annotations_off = self
                     .section_manager
-                    .get_aligned_size(Section::AnnotationSetItem)
+                    .get_aligned_size(FragSection::AnnotationSetItem)
                     + 1;
-                self.insert_field_annotation_set(&field_id, is_static)?;
+                self.insert_field_annotation_set(field, index)?;
                 field_annotations.push(FieldAnnotation {
-                    field_idx: *self.field_ids.get(&field_id).ok_or(anyhow!(
+                    field_idx: *index.fields.get(&field_id).ok_or(anyhow!(
                         "Field {} in {} not found in dex builder",
                         field_id.__str__(),
-                        class_id.__str__(),
+                        class.__str__(),
                     ))? as u32,
-                    annotations_off, // linked in link_annotations()
+                    annotations_off,
                 });
             }
         }
 
         let mut method_ids = vec![];
-        let (class, _) = self.class_defs.get(class_id).unwrap();
         method_ids.extend(class.direct_methods.keys().cloned());
         method_ids.extend(class.virtual_methods.keys().cloned());
         method_ids.sort();
         let mut method_annotations = vec![];
         for method_id in &method_ids {
-            let (class, _) = self.class_defs.get(class_id).unwrap();
             let direct_method = class.direct_methods.get(method_id);
             let virtual_method = class.virtual_methods.get(method_id);
             let (is_direct, method) = match (direct_method, virtual_method) {
@@ -279,28 +284,27 @@ impl DexFragment {
                 (None, Some(method)) => (false, method),
                 _ => bail!(
                     "Unexpected configuration: method {} is both a direct and a virtual method in {}", 
-                    method_id.__str__(), class_id.__str__()),
+                    method_id.__str__(), class.__str__()),
             };
             if !method.annotations.is_empty() {
                 let annotations_off = self
                     .section_manager
-                    .get_aligned_size(Section::AnnotationSetItem)
+                    .get_aligned_size(FragSection::AnnotationSetItem)
                     + 1;
-                self.insert_method_annotation_set(method_id, is_direct)?;
+                self.insert_method_annotation_set(method, index)?;
                 method_annotations.push(MethodAnnotation {
-                    method_idx: *self.method_ids.get(method_id).ok_or(anyhow!(
+                    method_idx: *index.methods.get(method_id).ok_or(anyhow!(
                         "Method {} in {} not found in dex builder",
                         method_id.__str__(),
-                        class_id.__str__(),
+                        class.__str__(),
                     ))? as u32,
-                    annotations_off, // linked in link_annotations()
+                    annotations_off,
                 });
             }
         }
 
         let mut parameter_annotations = vec![];
         for method_id in method_ids {
-            let (class, _) = self.class_defs.get(class_id).unwrap();
             let direct_method = class.direct_methods.get(&method_id);
             let virtual_method = class.virtual_methods.get(&method_id);
             let (is_direct, method) = match (direct_method, virtual_method) {
@@ -308,39 +312,39 @@ impl DexFragment {
                 (None, Some(method)) => (false, method),
                 _ => bail!(
                     "Unexpected configuration: method {} is both a direct and a virtual method in {}", 
-                    method_id.__str__(), class_id.__str__()),
+                    method_id.__str__(), class.__str__()),
             };
             if !method.parameters_annotations.is_empty() {
                 let annotations_off = self
                     .section_manager
-                    .get_aligned_size(Section::AnnotationSetRefList)
+                    .get_aligned_size(FragSection::AnnotationSetRefList)
                     + 1;
-                self.insert_parameters_annotation_set_list(&method_id, is_direct)?;
+                self.insert_parameters_annotation_set_list(method, index)?;
                 parameter_annotations.push(ParameterAnnotation {
-                    method_idx: *self.method_ids.get(&method_id).ok_or(anyhow!(
+                    method_idx: *index.methods.get(&method_id).ok_or(anyhow!(
                         "Method {} in {} not found in dex builder",
                         method_id.__str__(),
-                        class_id.__str__(),
+                        class.__str__(),
                     ))? as u32,
-                    annotations_off, // linked in link_annotations()
+                    annotations_off,
                 });
             }
         }
 
         let item = AnnotationsDirectoryItem {
-            class_annotations_off, // linked in link_annotations()
+            class_annotations_off,
             field_annotations,
             method_annotations,
             parameter_annotations,
         };
         self.section_manager
             .add_elt(FragSection::AnnotationsDirectoryItem, Some(item.size()));
-        self.annotations_directory_items.push(item);
+        self.annotations_directory_items = Some(item);
         Ok(())
     }
 
     /// Insert the annnotations set for a class.
-    fn insert_class_annotation_set(&mut self, class: &Class) -> Result<()> {
+    fn insert_class_annotation_set(&mut self, class: &Class, index: &FragIndex) -> Result<()> {
         let mut annotations = class.annotations.clone();
 
         let mut annotation_set = AnnotationSetItem { entries: vec![] };
@@ -367,7 +371,7 @@ impl DexFragment {
                         runtime or system"
                     ), // TODO: check if this is true
                 },
-                annotation: self.dex_annotation_to_encoded_annotation(annot.annotation)?,
+                annotation: self.dex_annotation_to_encoded_annotation(annot.annotation, index)?,
             };
             self.section_manager
                 .add_elt(FragSection::AnnotationItem, Some(item.size()));
@@ -379,6 +383,170 @@ impl DexFragment {
         Ok(())
     }
 
+    /// Insert the annnotations set for a class.
+    fn insert_field_annotation_set(&mut self, field: &Field, index: &FragIndex) -> Result<()> {
+        let mut annotations = field.annotations.clone();
+
+        let mut annotation_set = AnnotationSetItem { entries: vec![] };
+        annotations.sort_by_key(|annot| annot.annotation.type_.clone());
+        for annot in annotations {
+            annotation_set.entries.push(AnnotationOffItem {
+                annotation_off: self
+                    .section_manager
+                    .get_aligned_size(FragSection::AnnotationItem),
+            });
+
+            let item = AnnotationItem {
+                visibility: match (
+                    annot.visibility_build,
+                    annot.visibility_runtime,
+                    annot.visibility_system,
+                ) {
+                    (true, false, false) => AnnotationVisibility::Build,
+                    (false, true, false) => AnnotationVisibility::Runtime,
+                    (false, false, true) => AnnotationVisibility::System,
+                    _ => bail!(
+                        "Annotation need visibility set to one and only one of build, \
+                        runtime or system"
+                    ), // TODO: check if this is true
+                },
+                annotation: self.dex_annotation_to_encoded_annotation(annot.annotation, index)?,
+            };
+            self.section_manager
+                .add_elt(FragSection::AnnotationItem, Some(item.size()));
+            self.annotation_items.push(item);
+        }
+        self.section_manager
+            .add_elt(FragSection::AnnotationSetItem, Some(annotation_set.size()));
+        self.annotation_set_items.push(annotation_set);
+        Ok(())
+    }
+
+    /// Insert the annnotations set for a method (but not the parameters annotations).
+    fn insert_method_annotation_set(&mut self, method: &Method, index: &FragIndex) -> Result<()> {
+        let mut annotations = method.annotations.clone();
+
+        let mut annotation_set = AnnotationSetItem { entries: vec![] };
+        annotations.sort_by_key(|annot| annot.annotation.type_.clone());
+        for annot in annotations {
+            annotation_set.entries.push(AnnotationOffItem {
+                annotation_off: self
+                    .section_manager
+                    .get_aligned_size(FragSection::AnnotationItem)
+                    + 1,
+            });
+
+            let item = AnnotationItem {
+                visibility: match (
+                    annot.visibility_build,
+                    annot.visibility_runtime,
+                    annot.visibility_system,
+                ) {
+                    (true, false, false) => AnnotationVisibility::Build,
+                    (false, true, false) => AnnotationVisibility::Runtime,
+                    (false, false, true) => AnnotationVisibility::System,
+                    _ => bail!(
+                        "Annotation need visibility set to one and only one of build, \
+                        runtime or system"
+                    ), // TODO: check if this is true
+                },
+                annotation: self.dex_annotation_to_encoded_annotation(annot.annotation, index)?,
+            };
+            self.section_manager
+                .add_elt(FragSection::AnnotationItem, Some(item.size()));
+            self.annotation_items.push(item);
+        }
+        self.section_manager
+            .add_elt(FragSection::AnnotationSetItem, Some(annotation_set.size()));
+        self.annotation_set_items.push(annotation_set);
+        Ok(())
+    }
+
+    /// Insert a [`MethodHandle`].
+    pub fn insert_method_handle(&mut self, handle: &MethodHandle, index: &FragIndex) -> Result<()> {
+        let (field_or_method_id, method_handle_type) = match handle {
+            MethodHandle::StaticPut(StaticPut(field)) => (
+                *index
+                    .fields
+                    .get(field)
+                    .ok_or(anyhow!("Field {} not found in dex writer", field.__str__()))?
+                    as u16,
+                MethodHandleType::StaticPut,
+            ),
+            MethodHandle::StaticGet(StaticGet(field)) => (
+                *index
+                    .fields
+                    .get(field)
+                    .ok_or(anyhow!("Field {} not found in dex writer", field.__str__()))?
+                    as u16,
+                MethodHandleType::StaticGet,
+            ),
+            MethodHandle::InstancePut(InstancePut(field)) => (
+                *index
+                    .fields
+                    .get(field)
+                    .ok_or(anyhow!("Field {} not found in dex writer", field.__str__()))?
+                    as u16,
+                MethodHandleType::InstancePut,
+            ),
+            MethodHandle::InstanceGet(InstanceGet(field)) => (
+                *index
+                    .fields
+                    .get(field)
+                    .ok_or(anyhow!("Field {} not found in dex writer", field.__str__()))?
+                    as u16,
+                MethodHandleType::InstanceGet,
+            ),
+            MethodHandle::InvokeStatic(InvokeStatic(meth)) => (
+                *index
+                    .methods
+                    .get(meth)
+                    .ok_or(anyhow!("Method {} not found in dex writer", meth.__str__()))?
+                    as u16,
+                MethodHandleType::InvokeStatic,
+            ),
+            MethodHandle::InvokeInstance(InvokeInstance(meth)) => (
+                *index
+                    .methods
+                    .get(meth)
+                    .ok_or(anyhow!("Method {} not found in dex writer", meth.__str__()))?
+                    as u16,
+                MethodHandleType::InvokeInstance,
+            ),
+            MethodHandle::InvokeConstructor(InvokeConstructor(meth)) => (
+                *index
+                    .methods
+                    .get(meth)
+                    .ok_or(anyhow!("Method {} not found in dex writer", meth.__str__()))?
+                    as u16,
+                MethodHandleType::InvokeConstructor,
+            ),
+            MethodHandle::InvokeDirect(InvokeDirect(meth)) => (
+                *index
+                    .methods
+                    .get(meth)
+                    .ok_or(anyhow!("Method {} not found in dex writer", meth.__str__()))?
+                    as u16,
+                MethodHandleType::InvokeDirect,
+            ),
+            MethodHandle::InvokeInterface(InvokeInterface(meth)) => (
+                *index
+                    .methods
+                    .get(meth)
+                    .ok_or(anyhow!("Method {} not found in dex writer", meth.__str__()))?
+                    as u16,
+                MethodHandleType::InvokeInterface,
+            ),
+        };
+        self.method_handles.push(MethodHandleItem {
+            method_handle_type,
+            field_or_method_id,
+            unused1: 0,
+            unused2: 0,
+        });
+        Ok(())
+    }
+
     /// Insert a code_item.
     ///
     /// # Warning
@@ -1283,7 +1451,7 @@ impl DexFragment {
                 }
                 Instruction::ConstMethodHandle(ins) => {
                     let method_handle_idx = self.method_handles.len();
-                    self.insert_method_handle(&ins.handle)?;
+                    self.insert_method_handle(&ins.handle, index)?;
                     let ins = ins.get_raw_ins(method_handle_idx);
                     addr += ins.size() / 2;
                     insns.push(ins);
@@ -1558,97 +1726,16 @@ impl DexFragment {
         Ok(())
     }
 
-    /// Insert a [`MethodHandle`].
-    pub fn insert_method_handle(&mut self, handle: &MethodHandle) -> Result<()> {
-        let (field_or_method_id, method_handle_type) = match handle {
-            MethodHandle::StaticPut(StaticPut(field)) => (
-                *self
-                    .field_ids
-                    .get(field)
-                    .ok_or(anyhow!("Field {} not found in dex writer", field.__str__()))?
-                    as u16,
-                MethodHandleType::StaticPut,
-            ),
-            MethodHandle::StaticGet(StaticGet(field)) => (
-                *self
-                    .field_ids
-                    .get(field)
-                    .ok_or(anyhow!("Field {} not found in dex writer", field.__str__()))?
-                    as u16,
-                MethodHandleType::StaticGet,
-            ),
-            MethodHandle::InstancePut(InstancePut(field)) => (
-                *self
-                    .field_ids
-                    .get(field)
-                    .ok_or(anyhow!("Field {} not found in dex writer", field.__str__()))?
-                    as u16,
-                MethodHandleType::InstancePut,
-            ),
-            MethodHandle::InstanceGet(InstanceGet(field)) => (
-                *self
-                    .field_ids
-                    .get(field)
-                    .ok_or(anyhow!("Field {} not found in dex writer", field.__str__()))?
-                    as u16,
-                MethodHandleType::InstanceGet,
-            ),
-            MethodHandle::InvokeStatic(InvokeStatic(meth)) => (
-                *self
-                    .method_ids
-                    .get(meth)
-                    .ok_or(anyhow!("Method {} not found in dex writer", meth.__str__()))?
-                    as u16,
-                MethodHandleType::InvokeStatic,
-            ),
-            MethodHandle::InvokeInstance(InvokeInstance(meth)) => (
-                *self
-                    .method_ids
-                    .get(meth)
-                    .ok_or(anyhow!("Method {} not found in dex writer", meth.__str__()))?
-                    as u16,
-                MethodHandleType::InvokeInstance,
-            ),
-            MethodHandle::InvokeConstructor(InvokeConstructor(meth)) => (
-                *self
-                    .method_ids
-                    .get(meth)
-                    .ok_or(anyhow!("Method {} not found in dex writer", meth.__str__()))?
-                    as u16,
-                MethodHandleType::InvokeConstructor,
-            ),
-            MethodHandle::InvokeDirect(InvokeDirect(meth)) => (
-                *self
-                    .method_ids
-                    .get(meth)
-                    .ok_or(anyhow!("Method {} not found in dex writer", meth.__str__()))?
-                    as u16,
-                MethodHandleType::InvokeDirect,
-            ),
-            MethodHandle::InvokeInterface(InvokeInterface(meth)) => (
-                *self
-                    .method_ids
-                    .get(meth)
-                    .ok_or(anyhow!("Method {} not found in dex writer", meth.__str__()))?
-                    as u16,
-                MethodHandleType::InvokeInterface,
-            ),
-        };
-        self.method_handles.push(MethodHandleItem {
-            method_handle_type,
-            field_or_method_id,
-            unused1: 0,
-            unused2: 0,
-        });
-        Ok(())
-    }
-
     /// Convert a [`DexValue`] to an [`EncodedValue`].
     ///
     /// # Warning
     ///
     /// This method can insert element in the dex file like method_handles.
-    pub fn dex_value_to_encoded_value(&mut self, value: &DexValue) -> Result<EncodedValue> {
+    pub fn dex_value_to_encoded_value(
+        &mut self,
+        value: &DexValue,
+        index: &FragIndex,
+    ) -> Result<EncodedValue> {
         match value {
             DexValue::Byte(DexByte(val)) => Ok(EncodedValue::Byte(*val)),
             DexValue::Short(DexShort(val)) => Ok(EncodedValue::Short(*val)),
@@ -1658,48 +1745,47 @@ impl DexFragment {
             DexValue::Float(DexFloat(val)) => Ok(EncodedValue::Float(*val)),
             DexValue::Double(DexDouble(val)) => Ok(EncodedValue::Double(*val)),
             DexValue::MethodType(val) => Ok(EncodedValue::MethodType(
-                *self.proto_ids.get(val).ok_or(anyhow!(
+                *index.protos.get(val).ok_or(anyhow!(
                     "Prototype {} not found in dex writer",
                     val.__str__()
                 ))? as u32,
             )),
             DexValue::MethodHandle(val) => {
-                // TODO: move to a method
                 let idx = self.method_handles.len() as u32;
-                self.insert_method_handle(val)?;
+                self.insert_method_handle(val, index)?;
                 Ok(EncodedValue::MethodHandle(idx))
             }
             DexValue::String(val) => Ok(EncodedValue::String(
-                *self
+                *index
                     .strings
                     .get(val)
                     .ok_or(anyhow!("String {} not found in dex writer", val.__str__()))?
                     as u32,
             )),
             DexValue::Type(val) => Ok(EncodedValue::Type(
-                *self
+                *index
-                    .type_ids
+                    .types
                     .get(val)
                     .ok_or(anyhow!("Type {} not found in dex writer", val.__str__()))?
                     as u32,
             )),
             DexValue::Field(val) => Ok(EncodedValue::Field(
-                *self
+                *index
-                    .field_ids
+                    .fields
                     .get(val)
                     .ok_or(anyhow!("Field {} not found in dex writer", val.__str__()))?
                     as u32,
             )),
             DexValue::Method(val) => Ok(EncodedValue::Method(
-                *self
+                *index
-                    .method_ids
+                    .methods
                     .get(val)
                     .ok_or(anyhow!("Method {} not found in dex writer", val.__str__()))?
                     as u32,
             )),
             DexValue::Enum(IdEnum(val)) => Ok(EncodedValue::Enum(
-                *self
+                *index
-                    .field_ids
+                    .fields
                     .get(val)
                     .ok_or(anyhow!("Field {} not found in dex writer", val.__str__()))?
                     as u32,
@@ -1708,14 +1794,14 @@ impl DexFragment {
                 let mut values = vec![];
                 for val in arr {
                     values.push(
-                        self.dex_value_to_encoded_value(val)
+                        self.dex_value_to_encoded_value(val, index)
                             .context("Error while serializing a array")?,
                     );
                 }
                 Ok(EncodedValue::Array(EncodedArray { values }))
             }
             DexValue::Annotation(val) => Ok(EncodedValue::Annotation(
-                self.dex_annotation_to_encoded_annotation(val.clone())?,
+                self.dex_annotation_to_encoded_annotation(val.clone(), index)?,
             )),
             DexValue::Null(DexNull) => Ok(EncodedValue::Null),
             DexValue::Boolean(DexBoolean(val)) => Ok(EncodedValue::Boolean(*val)),
@@ -1723,16 +1809,20 @@ impl DexFragment {
     }
 
     /// Insert an encoded_array in the encoded_array_item section.
-    fn insert_encoded_array_item(&mut self, DexArray(array): DexArray) -> Result<()> {
+    fn insert_encoded_array_item(
+        &mut self,
+        DexArray(array): DexArray,
+        index: &FragIndex,
+    ) -> Result<()> {
         let mut values = vec![];
         for value in array {
-            values.push(self.dex_value_to_encoded_value(&value)?);
+            values.push(self.dex_value_to_encoded_value(&value, index)?);
         }
         let item = EncodedArrayItem {
             value: EncodedArray { values },
         };
         self.section_manager
-            .add_elt(Section::EncodedArrayItem, Some(item.size()));
+            .add_elt(FragSection::EncodedArrayItem, Some(item.size()));
         self.encoded_array_items.push(item);
         Ok(())
     }
@@ -1794,6 +1884,7 @@ impl DexFragment {
     fn dex_annotation_to_encoded_annotation(
         &mut self,
         DexAnnotation { type_, elements }: DexAnnotation,
+        index: &FragIndex,
     ) -> Result<EncodedAnnotation> {
         let mut encoded_elements = vec![];
 
@@ -1803,7 +1894,7 @@ impl DexFragment {
         for name in elements_names {
             let elt = elements.get(name).unwrap();
             encoded_elements.push(AnnotationElement {
-                name_idx: Uleb128(*self.strings.get(name).ok_or(anyhow!(
+                name_idx: Uleb128(*index.strings.get(name).ok_or(anyhow!(
                     "{} (annotation element name) not found in dex builder",
                     name.__str__()
                 ))? as u32),
@@ -1811,7 +1902,7 @@ impl DexFragment {
             });
         }
         Ok(EncodedAnnotation {
-            type_idx: Uleb128(*self.type_ids.get(&type_).ok_or(anyhow!(
+            type_idx: Uleb128(*index.types.get(&type_).ok_or(anyhow!(
                 "Annotation type {} not found in dex builder",
                 type_.__str__(),
             ))? as u32),
@@ -1819,117 +1910,13 @@ impl DexFragment {
         })
     }
 
-    /// Insert the annnotations set for a class.
-    fn insert_field_annotation_set(
-        &mut self,
-        field_id: &IdField,
-        is_static_field: bool,
-    ) -> Result<()> {
-        let (class, _) = self.class_defs.get(&field_id.class_).unwrap();
-        let field = if is_static_field {
-            class.static_fields.get(field_id).unwrap()
-        } else {
-            class.instance_fields.get(field_id).unwrap()
-        };
-        let mut annotations = field.annotations.clone();
-
-        let mut annotation_set = AnnotationSetItem { entries: vec![] };
-        annotations.sort_by_key(|annot| annot.annotation.type_.clone());
-        for annot in annotations {
-            annotation_set.entries.push(AnnotationOffItem {
-                annotation_off: self
-                    .section_manager
-                    .get_aligned_size(Section::AnnotationItem),
-            }); // linked in link_annotations()
-
-            let item = AnnotationItem {
-                visibility: match (
-                    annot.visibility_build,
-                    annot.visibility_runtime,
-                    annot.visibility_system,
-                ) {
-                    (true, false, false) => AnnotationVisibility::Build,
-                    (false, true, false) => AnnotationVisibility::Runtime,
-                    (false, false, true) => AnnotationVisibility::System,
-                    _ => bail!(
-                        "Annotation need visibility set to one and only one of build, \
-                        runtime or system"
-                    ), // TODO: check if this is true
-                },
-                annotation: self.dex_annotation_to_encoded_annotation(annot.annotation)?,
-            };
-            self.section_manager
-                .add_elt(Section::AnnotationItem, Some(item.size()));
-            self.annotation_items.push(item);
-        }
-        self.section_manager
-            .add_elt(Section::AnnotationSetItem, Some(annotation_set.size()));
-        self.annotation_set_items.push(annotation_set);
-        Ok(())
-    }
-
-    /// Insert the annnotations set for a method (but not the parameters annotations).
-    fn insert_method_annotation_set(
-        &mut self,
-        method_id: &IdMethod,
-        is_direct_method: bool,
-    ) -> Result<()> {
-        let (class, _) = self.class_defs.get(&method_id.class_).unwrap();
-        let method = if is_direct_method {
-            class.direct_methods.get(method_id).unwrap()
-        } else {
-            class.virtual_methods.get(method_id).unwrap()
-        };
-        let mut annotations = method.annotations.clone();
-
-        let mut annotation_set = AnnotationSetItem { entries: vec![] };
-        annotations.sort_by_key(|annot| annot.annotation.type_.clone());
-        for annot in annotations {
-            annotation_set.entries.push(AnnotationOffItem {
-                annotation_off: self
-                    .section_manager
-                    .get_aligned_size(Section::AnnotationItem),
-            }); // linked in link_annotations()
-
-            let item = AnnotationItem {
-                visibility: match (
-                    annot.visibility_build,
-                    annot.visibility_runtime,
-                    annot.visibility_system,
-                ) {
-                    (true, false, false) => AnnotationVisibility::Build,
-                    (false, true, false) => AnnotationVisibility::Runtime,
-                    (false, false, true) => AnnotationVisibility::System,
-                    _ => bail!(
-                        "Annotation need visibility set to one and only one of build, \
-                        runtime or system"
-                    ), // TODO: check if this is true
-                },
-                annotation: self.dex_annotation_to_encoded_annotation(annot.annotation)?,
-            };
-            self.section_manager
-                .add_elt(Section::AnnotationItem, Some(item.size()));
-            self.annotation_items.push(item);
-        }
-        self.section_manager
-            .add_elt(Section::AnnotationSetItem, Some(annotation_set.size()));
-        self.annotation_set_items.push(annotation_set);
-        Ok(())
-    }
-
     /// Insert the annotations set for a method parameter.
     fn insert_parameters_annotation_set(
         &mut self,
-        method_id: &IdMethod,
+        method: &Method,
-        is_direct_method: bool,
         parameter_idx: usize,
+        index: &FragIndex,
     ) -> Result<()> {
-        let (class, _) = self.class_defs.get(&method_id.class_).unwrap();
-        let method = if is_direct_method {
-            class.direct_methods.get(method_id).unwrap()
-        } else {
-            class.virtual_methods.get(method_id).unwrap()
-        };
         let mut annotations = method.parameters_annotations[parameter_idx].clone();
 
         let mut annotation_set = AnnotationSetItem { entries: vec![] };
@@ -1938,8 +1925,9 @@ impl DexFragment {
             annotation_set.entries.push(AnnotationOffItem {
                 annotation_off: self
                     .section_manager
-                    .get_aligned_size(Section::AnnotationItem),
+                    .get_aligned_size(FragSection::AnnotationItem)
-            }); // linked in link_annotations()
+                    + 1,
+            });
 
             let item = AnnotationItem {
                 visibility: match (
@@ -1955,14 +1943,14 @@ impl DexFragment {
                            runtime or system"
                     ), // TODO: check if this is true
                 },
-                annotation: self.dex_annotation_to_encoded_annotation(annot.annotation)?,
+                annotation: self.dex_annotation_to_encoded_annotation(annot.annotation, index)?,
             };
             self.section_manager
-                .add_elt(Section::AnnotationItem, Some(item.size()));
+                .add_elt(FragSection::AnnotationItem, Some(item.size()));
             self.annotation_items.push(item);
         }
         self.section_manager
-            .add_elt(Section::AnnotationSetItem, Some(annotation_set.size()));
+            .add_elt(FragSection::AnnotationSetItem, Some(annotation_set.size()));
         self.annotation_set_items.push(annotation_set);
         Ok(())
     }
@@ -1970,16 +1958,10 @@ impl DexFragment {
     /// Insert the annotations set list for a method parameters.
     fn insert_parameters_annotation_set_list(
         &mut self,
-        method_id: &IdMethod,
+        method: &Method,
-        is_direct_method: bool,
+        index: &FragIndex,
     ) -> Result<()> {
         let mut list = AnnotationSetRefList { list: vec![] };
-        let (class, _) = self.class_defs.get(&method_id.class_).unwrap();
-        let method = if is_direct_method {
-            class.direct_methods.get(method_id).unwrap()
-        } else {
-            class.virtual_methods.get(method_id).unwrap()
-        };
         let param_has_annotation: Vec<_> = method
             .parameters_annotations
             .iter()
@@ -1990,17 +1972,17 @@ impl DexFragment {
                 annotations_off: if has_annotation {
                     let annotation_off = self
                         .section_manager
-                        .get_aligned_size(Section::AnnotationSetItem);
+                        .get_aligned_size(FragSection::AnnotationSetItem);
-                    self.insert_parameters_annotation_set(method_id, is_direct_method, param_idx)?;
+                    self.insert_parameters_annotation_set(method, param_idx, index)?;
                     annotation_off + 1
                 } else {
                     0
-                }, // linked in link_annotations()
+                },
             });
         }
 
         self.section_manager
-            .add_elt(Section::AnnotationSetRefList, Some(list.size()));
+            .add_elt(FragSection::AnnotationSetRefList, Some(list.size()));
         self.annotation_set_lists.push(list);
         Ok(())
     }
@@ -2879,3 +2861,12 @@ impl FragSectionManager {
         }
     }
 }
+
+/// Index that associate a type to its local id in a fragment.
+struct FragIndex {
+    pub strings: HashMap<DexString, usize>,
+    pub types: HashMap<IdType, usize>,
+    pub protos: HashMap<IdMethodType, usize>,
+    pub fields: HashMap<IdField, usize>,
+    pub methods: HashMap<IdMethod, usize>,
+}