start implementing visitors

androscalpel/src/apk.rs

@@ -490,37 +490,33 @@ impl Apk {
     pub fn get_method_handle_from_idx(idx: usize, dex: &DexFileReader) -> Result<MethodHandle> {
         let handle = dex.get_method_handle(idx)?;
         match handle.method_handle_type {
-            MethodHandleType::StaticPut => Ok(MethodHandle::StaticPut(StaticPut(
+            MethodHandleType::StaticPut => Ok(MethodHandle::StaticPut {
-                Self::get_id_field_from_idx(handle.field_or_method_id as usize, dex)?,
+                field: Self::get_id_field_from_idx(handle.field_or_method_id as usize, dex)?,
-            ))),
+            }),
-            MethodHandleType::StaticGet => Ok(MethodHandle::StaticGet(StaticGet(
+            MethodHandleType::StaticGet => Ok(MethodHandle::StaticGet {
-                Self::get_id_field_from_idx(handle.field_or_method_id as usize, dex)?,
+                field: Self::get_id_field_from_idx(handle.field_or_method_id as usize, dex)?,
-            ))),
+            }),
-            MethodHandleType::InstancePut => Ok(MethodHandle::InstancePut(InstancePut(
+            MethodHandleType::InstancePut => Ok(MethodHandle::InstancePut {
-                Self::get_id_field_from_idx(handle.field_or_method_id as usize, dex)?,
+                field: Self::get_id_field_from_idx(handle.field_or_method_id as usize, dex)?,
-            ))),
+            }),
-            MethodHandleType::InstanceGet => Ok(MethodHandle::InstanceGet(InstanceGet(
+            MethodHandleType::InstanceGet => Ok(MethodHandle::InstanceGet {
-                Self::get_id_field_from_idx(handle.field_or_method_id as usize, dex)?,
+                field: Self::get_id_field_from_idx(handle.field_or_method_id as usize, dex)?,
-            ))),
+            }),
-            MethodHandleType::InvokeStatic => Ok(MethodHandle::InvokeStatic(InvokeStatic(
+            MethodHandleType::InvokeStatic => Ok(MethodHandle::InvokeStatic {
-                Self::get_id_method_from_idx(handle.field_or_method_id as usize, dex)?,
+                method: Self::get_id_method_from_idx(handle.field_or_method_id as usize, dex)?,
-            ))),
+            }),
-            MethodHandleType::InvokeInstance => Ok(MethodHandle::InvokeInstance(InvokeInstance(
+            MethodHandleType::InvokeInstance => Ok(MethodHandle::InvokeInstance {
-                Self::get_id_method_from_idx(handle.field_or_method_id as usize, dex)?,
+                method: Self::get_id_method_from_idx(handle.field_or_method_id as usize, dex)?,
-            ))),
+            }),
-            MethodHandleType::InvokeConstructor => {
+            MethodHandleType::InvokeConstructor => Ok(MethodHandle::InvokeConstructor {
-                Ok(MethodHandle::InvokeConstructor(InvokeConstructor(
+                method: Self::get_id_method_from_idx(handle.field_or_method_id as usize, dex)?,
-                    Self::get_id_method_from_idx(handle.field_or_method_id as usize, dex)?,
+            }),
-                )))
+            MethodHandleType::InvokeDirect => Ok(MethodHandle::InvokeDirect {
-            }
+                method: Self::get_id_method_from_idx(handle.field_or_method_id as usize, dex)?,
-            MethodHandleType::InvokeDirect => Ok(MethodHandle::InvokeDirect(InvokeDirect(
+            }),
-                Self::get_id_method_from_idx(handle.field_or_method_id as usize, dex)?,
+            MethodHandleType::InvokeInterface => Ok(MethodHandle::InvokeInterface {
-            ))),
+                method: Self::get_id_method_from_idx(handle.field_or_method_id as usize, dex)?,
-            MethodHandleType::InvokeInterface => {
+            }),
-                Ok(MethodHandle::InvokeInterface(InvokeInterface(
-                    Self::get_id_method_from_idx(handle.field_or_method_id as usize, dex)?,
-                )))
-            }
         }
     }
 
@@ -793,7 +789,7 @@ impl Apk {
             labels.insert(addr, label.clone());
         }
         let ins = match format.clone() {
-            Format10X { op: 0x00 } => Instruction::Nop(Nop::new()),
+            Format10X { op: 0x00 } => Instruction::Nop {},
             Format12X { op: 0x01, va, vb } => Instruction::Move(Move::new(va as u16, vb as u16)),
             Format22X { op: 0x02, va, vb } => Instruction::Move(Move::new(va as u16, vb)),
             Format32X { op: 0x03, va, vb } => Instruction::Move(Move::new(va, vb)),

androscalpel/src/code.rs

@@ -472,7 +472,7 @@ impl Code {
                     }
                     last_ins_was_a_label = true;
                 }
-                Instruction::Nop(_) => (),
+                Instruction::Nop {} => (),
                 instr => {
                     last_ins_was_a_label = false;
                     new_insns.push(instr);

androscalpel/src/dex_id.rs

@@ -10,7 +10,7 @@ use anyhow::{anyhow, bail, Context};
 use pyo3::class::basic::CompareOp;
 use pyo3::prelude::*;
 
-use crate::{scalar::*, DexString, DexValue, Result};
+use crate::{scalar::*, DexString, DexValue, Result, Visitable, VisitableMut, Visitor, VisitorMut};
 use androscalpel_serializer::{StringDataItem, Uleb128};
 
 /// The type of a method. The shorty is formated as described in
@@ -25,6 +25,30 @@ pub struct IdMethodType {
     pub(crate) parameters: Vec<IdType>,
 }
 
+impl<V: Visitor> Visitable<V> for IdMethodType {
+    fn default_visit(&self, v: &mut V) -> Result<()> {
+        v.visit_string(&self.shorty)?;
+        v.visit_type(&self.return_type)?;
+        for ty in &self.parameters {
+            v.visit_type(&ty)?;
+        }
+        Ok(())
+    }
+}
+
+impl<V: VisitorMut> VisitableMut<V> for IdMethodType {
+    fn default_visit_mut(mut self, v: &mut V) -> Result<Self> {
+        self.shorty = v.visit_string(self.shorty)?;
+        self.return_type = v.visit_type(self.return_type)?;
+        self.parameters = self
+            .parameters
+            .into_iter()
+            .map(|ty| v.visit_type(ty))
+            .collect::<Result<_>>()?;
+        Ok(self)
+    }
+}
+
 impl Ord for IdMethodType {
     fn cmp(&self, other: &Self) -> Ordering {
         self.return_type
@@ -220,6 +244,18 @@ impl IdMethodType {
 #[pyclass]
 #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct IdType(pub(crate) DexString);
+
+impl<V: Visitor> Visitable<V> for IdType {
+    fn default_visit(&self, v: &mut V) -> Result<()> {
+        v.visit_string(&self.0)
+    }
+}
+
+impl<V: VisitorMut> VisitableMut<V> for IdType {
+    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
+        Ok(Self(v.visit_string(self.0)?))
+    }
+}
 #[pymethods]
 impl IdType {
     pub fn to_json(&self) -> Result<String> {
@@ -637,6 +673,24 @@ pub struct IdField {
     pub class_: IdType,
 }
 
+impl<V: Visitor> Visitable<V> for IdField {
+    fn default_visit(&self, v: &mut V) -> Result<()> {
+        v.visit_string(&self.name)?;
+        v.visit_type(&self.type_)?;
+        v.visit_type(&self.class_)?;
+        Ok(())
+    }
+}
+
+impl<V: VisitorMut> VisitableMut<V> for IdField {
+    fn default_visit_mut(mut self, v: &mut V) -> Result<Self> {
+        self.name = v.visit_string(self.name)?;
+        self.type_ = v.visit_type(self.type_)?;
+        self.class_ = v.visit_type(self.class_)?;
+        Ok(self)
+    }
+}
+
 #[pymethods]
 impl IdField {
     pub fn to_json(&self) -> Result<String> {
@@ -795,6 +849,24 @@ pub struct IdMethod {
     pub name: DexString,
 }
 
+impl<V: Visitor> Visitable<V> for IdMethod {
+    fn default_visit(&self, v: &mut V) -> Result<()> {
+        v.visit_string(&self.name)?;
+        v.visit_method_type(&self.proto)?;
+        v.visit_type(&self.class_)?;
+        Ok(())
+    }
+}
+
+impl<V: VisitorMut> VisitableMut<V> for IdMethod {
+    fn default_visit_mut(mut self, v: &mut V) -> Result<Self> {
+        self.name = v.visit_string(self.name)?;
+        self.proto = v.visit_method_type(self.proto)?;
+        self.class_ = v.visit_type(self.class_)?;
+        Ok(self)
+    }
+}
+
 #[pymethods]
 impl IdMethod {
     pub fn to_json(&self) -> Result<String> {
@@ -972,6 +1044,18 @@ impl SmaliName for IdMethod {
 #[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)]
 pub struct IdEnum(pub IdField);
 
+impl<V: Visitor> Visitable<V> for IdEnum {
+    fn default_visit(&self, v: &mut V) -> Result<()> {
+        v.visit_field_id(&self.0)
+    }
+}
+
+impl<V: VisitorMut> VisitableMut<V> for IdEnum {
+    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
+        Ok(Self(v.visit_field_id(self.0)?))
+    }
+}
+
 #[pymethods]
 impl IdEnum {
     pub fn to_json(&self) -> Result<String> {

androscalpel/src/dex_string.rs

@@ -1,4 +1,4 @@
-use crate::Result;
+use crate::{Result, Visitable, VisitableMut, Visitor, VisitorMut};
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use std::cmp::{Ord, PartialOrd};
 use std::collections::hash_map::DefaultHasher;
@@ -13,6 +13,19 @@ use pyo3::prelude::*;
 #[derive(Clone, PartialEq, Eq, Ord, PartialOrd)]
 pub struct DexString(pub androscalpel_serializer::StringDataItem);
 
+// Kinda useless
+impl<V: Visitor> Visitable<V> for DexString {
+    fn default_visit(&self, _: &mut V) -> Result<()> {
+        Ok(())
+    }
+}
+
+impl<V: VisitorMut> VisitableMut<V> for DexString {
+    fn default_visit_mut(self, _: &mut V) -> Result<Self> {
+        Ok(self)
+    }
+}
+
 impl std::fmt::Debug for DexString {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
         if let Ok(string) = TryInto::<String>::try_into(self) {

androscalpel/src/dex_writer.rs

@@ -581,7 +581,7 @@ impl DexWriter {
                         // https://cs.android.com/android/platform/superproject/main/+/main:art/runtime/verifier/method_verifier.cc;drc=e8c3e7be783937a340cd4f3280b69962d6f1ea0c;l=1347
                         // The ART check if the array data table is 4 bytes aligned (= 2 ins alligned)
                         // TODO: check how it is donne in android and other dex generation code.
-                        let nop = Instruction::Nop(Nop).get_raw_ins()?;
+                        let nop = (Instruction::Nop {}).get_raw_ins()?;
                         payload_addr += nop.size() / 2;
                         payloads.push(nop);
                     }
@@ -632,7 +632,7 @@ impl DexWriter {
                         // https://cs.android.com/android/platform/superproject/main/+/main:art/runtime/verifier/method_verifier.cc;drc=e8c3e7be783937a340cd4f3280b69962d6f1ea0c;l=1464
                         // The ART check if the switch table is 4 bytes aligned (= 2 ins alligned)
                         // TODO: check how it is donne in android and other dex generation code.
-                        let nop = Instruction::Nop(Nop).get_raw_ins()?;
+                        let nop = (Instruction::Nop {}).get_raw_ins()?;
                         payload_addr += nop.size() / 2;
                         payloads.push(nop);
                     }
@@ -1388,7 +1388,7 @@ impl DexWriter {
         }
         if addr % 2 != 0 {
             // make sure the payload section is 4 bytes aligned
-            let nop = Instruction::Nop(Nop).get_raw_ins()?;
+            let nop = (Instruction::Nop {}).get_raw_ins()?;
             //addr += nop.size() / 2;
             insns.push(nop);
         }
@@ -1582,66 +1582,66 @@ impl DexWriter {
     /// 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)) => (
+            MethodHandle::StaticPut { field } => (
                 *self.field_ids.get(field).ok_or(anyhow!(
                     "Field {} not found in dex writer",
                     field.__repr__()
                 ))? as u16,
                 MethodHandleType::StaticPut,
             ),
-            MethodHandle::StaticGet(StaticGet(field)) => (
+            MethodHandle::StaticGet { field } => (
                 *self.field_ids.get(field).ok_or(anyhow!(
                     "Field {} not found in dex writer",
                     field.__repr__()
                 ))? as u16,
                 MethodHandleType::StaticGet,
             ),
-            MethodHandle::InstancePut(InstancePut(field)) => (
+            MethodHandle::InstancePut { field } => (
                 *self.field_ids.get(field).ok_or(anyhow!(
                     "Field {} not found in dex writer",
                     field.__repr__()
                 ))? as u16,
                 MethodHandleType::InstancePut,
             ),
-            MethodHandle::InstanceGet(InstanceGet(field)) => (
+            MethodHandle::InstanceGet { field } => (
                 *self.field_ids.get(field).ok_or(anyhow!(
                     "Field {} not found in dex writer",
                     field.__repr__()
                 ))? as u16,
                 MethodHandleType::InstanceGet,
             ),
-            MethodHandle::InvokeStatic(InvokeStatic(meth)) => (
+            MethodHandle::InvokeStatic { method } => (
-                *self.method_ids.get(meth).ok_or(anyhow!(
+                *self.method_ids.get(method).ok_or(anyhow!(
                     "Method {} not found in dex writer",
-                    meth.__repr__()
+                    method.__repr__()
                 ))? as u16,
                 MethodHandleType::InvokeStatic,
             ),
-            MethodHandle::InvokeInstance(InvokeInstance(meth)) => (
+            MethodHandle::InvokeInstance { method } => (
-                *self.method_ids.get(meth).ok_or(anyhow!(
+                *self.method_ids.get(method).ok_or(anyhow!(
                     "Method {} not found in dex writer",
-                    meth.__repr__()
+                    method.__repr__()
                 ))? as u16,
                 MethodHandleType::InvokeInstance,
             ),
-            MethodHandle::InvokeConstructor(InvokeConstructor(meth)) => (
+            MethodHandle::InvokeConstructor { method } => (
-                *self.method_ids.get(meth).ok_or(anyhow!(
+                *self.method_ids.get(method).ok_or(anyhow!(
                     "Method {} not found in dex writer",
-                    meth.__repr__()
+                    method.__repr__()
                 ))? as u16,
                 MethodHandleType::InvokeConstructor,
             ),
-            MethodHandle::InvokeDirect(InvokeDirect(meth)) => (
+            MethodHandle::InvokeDirect { method } => (
-                *self.method_ids.get(meth).ok_or(anyhow!(
+                *self.method_ids.get(method).ok_or(anyhow!(
                     "Method {} not found in dex writer",
-                    meth.__repr__()
+                    method.__repr__()
                 ))? as u16,
                 MethodHandleType::InvokeDirect,
             ),
-            MethodHandle::InvokeInterface(InvokeInterface(meth)) => (
+            MethodHandle::InvokeInterface { method } => (
-                *self.method_ids.get(meth).ok_or(anyhow!(
+                *self.method_ids.get(method).ok_or(anyhow!(
                     "Method {} not found in dex writer",
-                    meth.__repr__()
+                    method.__repr__()
                 ))? as u16,
                 MethodHandleType::InvokeInterface,
             ),

androscalpel/src/lib.rs

@@ -21,6 +21,7 @@ pub mod method_handle;
 pub mod py_utils;
 pub mod scalar;
 pub mod value;
+pub mod visitor;
 
 pub use annotation::*;
 pub use apk::*;
@@ -36,6 +37,7 @@ pub use method::*;
 pub use method_handle::*;
 pub use scalar::*;
 pub use value::*;
+pub use visitor::*;
 
 #[cfg(test)]
 mod tests;
@@ -61,16 +63,6 @@ fn androscalpel(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<IdMethod>()?;
     m.add_class::<IdEnum>()?;
 
-    m.add_class::<StaticPut>()?;
-    m.add_class::<StaticGet>()?;
-    m.add_class::<InstancePut>()?;
-    m.add_class::<InstanceGet>()?;
-    m.add_class::<InvokeStatic>()?;
-    m.add_class::<InvokeInstance>()?;
-    m.add_class::<InvokeConstructor>()?;
-    m.add_class::<InvokeDirect>()?;
-    m.add_class::<InvokeInterface>()?;
-
     m.add_class::<HiddenApiData>()?;
     m.add_class::<HiddenApiPermission>()?;
     m.add_class::<HiddenApiDomain>()?;

androscalpel/src/method_handle.rs

@@ -4,790 +4,150 @@ use serde::{Deserialize, Serialize};
 use std::collections::HashSet;
 use std::hash::Hash;
 
-use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
 
 use crate::dex_id::*;
-use crate::DexString;
+use crate::{
-use crate::Result;
+    DexString, FieldIdCollector, MethodHandleCollector, MethodIdCollector, MethodTypeCollector,
+    Result, StringCollector, TypeCollector, Visitable, VisitableMut, Visitor, VisitorMut,
+};
 
 /// The structure use to reference a method invocation.
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize, Serialize)]
+#[pyclass]
 pub enum MethodHandle {
-    StaticPut(StaticPut),
+    StaticPut { field: IdField },
-    StaticGet(StaticGet),
+    StaticGet { field: IdField },
-    InstancePut(InstancePut),
+    InstancePut { field: IdField },
-    InstanceGet(InstanceGet),
+    InstanceGet { field: IdField },
-    InvokeStatic(InvokeStatic),
+    InvokeStatic { method: IdMethod },
-    InvokeInstance(InvokeInstance),
+    InvokeInstance { method: IdMethod },
-    InvokeConstructor(InvokeConstructor),
+    InvokeConstructor { method: IdMethod },
-    InvokeDirect(InvokeDirect),
+    InvokeDirect { method: IdMethod },
-    InvokeInterface(InvokeInterface),
+    InvokeInterface { method: IdMethod },
 }
 
-#[pyclass]
+impl<V: Visitor> Visitable<V> for MethodHandle {
-#[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize, Serialize)]
+    fn default_visit(&self, v: &mut V) -> Result<()> {
-pub struct StaticPut(pub IdField);
-
-#[pymethods]
-impl StaticPut {
-    pub fn to_json(&self) -> Result<String> {
-        Ok(serde_json::to_string(self)?)
-    }
-
-    #[staticmethod]
-    pub fn from_json(json: &str) -> Result<Self> {
-        Ok(serde_json::from_str(json)?)
-    }
-
-    #[new]
-    pub fn new(val: IdField) -> Self {
-        Self(val)
-    }
-
-    pub fn get_field(&self) -> IdField {
-        self.0.clone()
-    }
-
-    pub fn __str__(&self) -> String {
-        self.__repr__()
-    }
-
-    pub fn __repr__(&self) -> String {
-        format!("StaticPut({})", self.0.__str__())
-    }
-
-    /// Return all strings referenced in the handle.
-    pub fn get_all_strings(&self) -> HashSet<DexString> {
-        self.0.get_all_strings()
-    }
-
-    /// Return all types referenced in the handle.
-    pub fn get_all_types(&self) -> HashSet<IdType> {
-        self.0.get_all_types()
-    }
-
-    /// Return all prototypes referenced in the handle.
-    pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
-        HashSet::new()
-    }
-
-    /// Return all field ids referenced in the handle.
-    pub fn get_all_field_ids(&self) -> HashSet<IdField> {
-        let mut fields = HashSet::new();
-        fields.insert(self.0.clone());
-        fields
-    }
-
-    /// Return all method ids referenced in the handle.
-    pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
-        HashSet::new()
-    }
-
-    /// Return all method handles referenced in the handle.
-    pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
-        let mut methods = HashSet::new();
-        methods.insert(MethodHandle::StaticPut(self.clone()));
-        methods
-    }
-}
-#[pyclass]
-#[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize, Serialize)]
-pub struct StaticGet(pub IdField);
-
-#[pymethods]
-impl StaticGet {
-    pub fn to_json(&self) -> Result<String> {
-        Ok(serde_json::to_string(self)?)
-    }
-
-    #[staticmethod]
-    pub fn from_json(json: &str) -> Result<Self> {
-        Ok(serde_json::from_str(json)?)
-    }
-
-    #[new]
-    pub fn new(val: IdField) -> Self {
-        Self(val)
-    }
-
-    pub fn get_field(&self) -> IdField {
-        self.0.clone()
-    }
-
-    pub fn __str__(&self) -> String {
-        self.__repr__()
-    }
-
-    pub fn __repr__(&self) -> String {
-        format!("StaticGet({})", self.0.__str__())
-    }
-
-    /// Return all strings referenced in the handle.
-    pub fn get_all_strings(&self) -> HashSet<DexString> {
-        self.0.get_all_strings()
-    }
-
-    /// Return all types referenced in the handle.
-    pub fn get_all_types(&self) -> HashSet<IdType> {
-        self.0.get_all_types()
-    }
-
-    /// Return all prototypes referenced in the handle.
-    pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
-        HashSet::new()
-    }
-
-    /// Return all field ids referenced in the handle.
-    pub fn get_all_field_ids(&self) -> HashSet<IdField> {
-        let mut fields = HashSet::new();
-        fields.insert(self.0.clone());
-        fields
-    }
-
-    /// Return all method ids referenced in the handle.
-    pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
-        HashSet::new()
-    }
-
-    /// Return all method handles referenced in the handle.
-    pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
-        let mut methods = HashSet::new();
-        methods.insert(MethodHandle::StaticGet(self.clone()));
-        methods
-    }
-}
-#[pyclass]
-#[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize, Serialize)]
-pub struct InstancePut(pub IdField);
-
-#[pymethods]
-impl InstancePut {
-    pub fn to_json(&self) -> Result<String> {
-        Ok(serde_json::to_string(self)?)
-    }
-
-    #[staticmethod]
-    pub fn from_json(json: &str) -> Result<Self> {
-        Ok(serde_json::from_str(json)?)
-    }
-
-    #[new]
-    pub fn new(val: IdField) -> Self {
-        Self(val)
-    }
-
-    pub fn get_field(&self) -> IdField {
-        self.0.clone()
-    }
-
-    pub fn __str__(&self) -> String {
-        self.__repr__()
-    }
-
-    pub fn __repr__(&self) -> String {
-        format!("InstancePut({})", self.0.__str__())
-    }
-
-    /// Return all strings referenced in the handle.
-    pub fn get_all_strings(&self) -> HashSet<DexString> {
-        self.0.get_all_strings()
-    }
-
-    /// Return all types referenced in the handle.
-    pub fn get_all_types(&self) -> HashSet<IdType> {
-        self.0.get_all_types()
-    }
-
-    /// Return all prototypes referenced in the handle.
-    pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
-        HashSet::new()
-    }
-
-    /// Return all field ids referenced in the handle.
-    pub fn get_all_field_ids(&self) -> HashSet<IdField> {
-        let mut fields = HashSet::new();
-        fields.insert(self.0.clone());
-        fields
-    }
-
-    /// Return all method ids referenced in the handle.
-    pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
-        HashSet::new()
-    }
-
-    /// Return all method handles referenced in the handle.
-    pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
-        let mut methods = HashSet::new();
-        methods.insert(MethodHandle::InstancePut(self.clone()));
-        methods
-    }
-}
-#[pyclass]
-#[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize, Serialize)]
-pub struct InstanceGet(pub IdField);
-
-#[pymethods]
-impl InstanceGet {
-    pub fn to_json(&self) -> Result<String> {
-        Ok(serde_json::to_string(self)?)
-    }
-
-    #[staticmethod]
-    pub fn from_json(json: &str) -> Result<Self> {
-        Ok(serde_json::from_str(json)?)
-    }
-
-    #[new]
-    pub fn new(val: IdField) -> Self {
-        Self(val)
-    }
-
-    pub fn get_field(&self) -> IdField {
-        self.0.clone()
-    }
-
-    pub fn __str__(&self) -> String {
-        self.__repr__()
-    }
-
-    pub fn __repr__(&self) -> String {
-        format!("InstanceGet({})", self.0.__str__())
-    }
-
-    /// Return all strings referenced in the handle.
-    pub fn get_all_strings(&self) -> HashSet<DexString> {
-        self.0.get_all_strings()
-    }
-
-    /// Return all types referenced in the handle.
-    pub fn get_all_types(&self) -> HashSet<IdType> {
-        self.0.get_all_types()
-    }
-
-    /// Return all prototypes referenced in the handle.
-    pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
-        HashSet::new()
-    }
-
-    /// Return all field ids referenced in the handle.
-    pub fn get_all_field_ids(&self) -> HashSet<IdField> {
-        let mut fields = HashSet::new();
-        fields.insert(self.0.clone());
-        fields
-    }
-
-    /// Return all method ids referenced in the handle.
-    pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
-        HashSet::new()
-    }
-
-    /// Return all method handles referenced in the handle.
-    pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
-        let mut methods = HashSet::new();
-        methods.insert(MethodHandle::InstanceGet(self.clone()));
-        methods
-    }
-}
-
-#[pyclass]
-#[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize, Serialize)]
-pub struct InvokeStatic(pub IdMethod);
-
-#[pymethods]
-impl InvokeStatic {
-    pub fn to_json(&self) -> Result<String> {
-        Ok(serde_json::to_string(self)?)
-    }
-
-    #[staticmethod]
-    pub fn from_json(json: &str) -> Result<Self> {
-        Ok(serde_json::from_str(json)?)
-    }
-
-    #[new]
-    pub fn new(val: IdMethod) -> Self {
-        Self(val)
-    }
-
-    pub fn get_method(&self) -> IdMethod {
-        self.0.clone()
-    }
-
-    pub fn __str__(&self) -> String {
-        self.__repr__()
-    }
-
-    pub fn __repr__(&self) -> String {
-        format!("InvokeStatic({})", self.0.__str__())
-    }
-
-    /// Return all strings referenced in the handle.
-    pub fn get_all_strings(&self) -> HashSet<DexString> {
-        self.0.get_all_strings()
-    }
-
-    /// Return all types referenced in the handle.
-    pub fn get_all_types(&self) -> HashSet<IdType> {
-        self.0.get_all_types()
-    }
-
-    /// Return all prototypes referenced in the handle.
-    pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
-        self.0.get_all_protos()
-    }
-
-    /// Return all field ids referenced in the handle.
-    pub fn get_all_field_ids(&self) -> HashSet<IdField> {
-        HashSet::new()
-    }
-
-    /// Return all method ids referenced in the handle.
-    pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
-        let mut methods = HashSet::new();
-        methods.insert(self.0.clone());
-        methods
-    }
-
-    /// Return all method handles referenced in the handle.
-    pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
-        let mut methods = HashSet::new();
-        methods.insert(MethodHandle::InvokeStatic(self.clone()));
-        methods
-    }
-}
-
-#[pyclass]
-#[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize, Serialize)]
-pub struct InvokeInstance(pub IdMethod);
-
-#[pymethods]
-impl InvokeInstance {
-    pub fn to_json(&self) -> Result<String> {
-        Ok(serde_json::to_string(self)?)
-    }
-
-    #[staticmethod]
-    pub fn from_json(json: &str) -> Result<Self> {
-        Ok(serde_json::from_str(json)?)
-    }
-
-    #[new]
-    pub fn new(val: IdMethod) -> Self {
-        Self(val)
-    }
-
-    pub fn get_method(&self) -> IdMethod {
-        self.0.clone()
-    }
-
-    pub fn __str__(&self) -> String {
-        self.__repr__()
-    }
-
-    pub fn __repr__(&self) -> String {
-        format!("InvokeInstance({})", self.0.__str__())
-    }
-
-    /// Return all strings referenced in the handle.
-    pub fn get_all_strings(&self) -> HashSet<DexString> {
-        self.0.get_all_strings()
-    }
-
-    /// Return all types referenced in the handle.
-    pub fn get_all_types(&self) -> HashSet<IdType> {
-        self.0.get_all_types()
-    }
-
-    /// Return all prototypes referenced in the handle.
-    pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
-        self.0.get_all_protos()
-    }
-
-    /// Return all field ids referenced in the handle.
-    pub fn get_all_field_ids(&self) -> HashSet<IdField> {
-        HashSet::new()
-    }
-
-    /// Return all method ids referenced in the handle.
-    pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
-        let mut methods = HashSet::new();
-        methods.insert(self.0.clone());
-        methods
-    }
-
-    /// Return all method handles referenced in the handle.
-    pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
-        let mut methods = HashSet::new();
-        methods.insert(MethodHandle::InvokeInstance(self.clone()));
-        methods
-    }
-}
-
-#[pyclass]
-#[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize, Serialize)]
-pub struct InvokeConstructor(pub IdMethod);
-
-#[pymethods]
-impl InvokeConstructor {
-    pub fn to_json(&self) -> Result<String> {
-        Ok(serde_json::to_string(self)?)
-    }
-
-    #[staticmethod]
-    pub fn from_json(json: &str) -> Result<Self> {
-        Ok(serde_json::from_str(json)?)
-    }
-
-    #[new]
-    pub fn new(val: IdMethod) -> Self {
-        Self(val)
-    }
-
-    pub fn get_method(&self) -> IdMethod {
-        self.0.clone()
-    }
-
-    pub fn __str__(&self) -> String {
-        self.__repr__()
-    }
-
-    pub fn __repr__(&self) -> String {
-        format!("InvokeConstructor({})", self.0.__str__())
-    }
-
-    /// Return all strings referenced in the handle.
-    pub fn get_all_strings(&self) -> HashSet<DexString> {
-        self.0.get_all_strings()
-    }
-
-    /// Return all types referenced in the handle.
-    pub fn get_all_types(&self) -> HashSet<IdType> {
-        self.0.get_all_types()
-    }
-
-    /// Return all prototypes referenced in the handle.
-    pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
-        self.0.get_all_protos()
-    }
-
-    /// Return all field ids referenced in the handle.
-    pub fn get_all_field_ids(&self) -> HashSet<IdField> {
-        HashSet::new()
-    }
-
-    /// Return all method ids referenced in the handle.
-    pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
-        let mut methods = HashSet::new();
-        methods.insert(self.0.clone());
-        methods
-    }
-
-    /// Return all method handles referenced in the handle.
-    pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
-        let mut methods = HashSet::new();
-        methods.insert(MethodHandle::InvokeConstructor(self.clone()));
-        methods
-    }
-}
-
-#[pyclass]
-#[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize, Serialize)]
-pub struct InvokeDirect(pub IdMethod);
-
-#[pymethods]
-impl InvokeDirect {
-    pub fn to_json(&self) -> Result<String> {
-        Ok(serde_json::to_string(self)?)
-    }
-
-    #[staticmethod]
-    pub fn from_json(json: &str) -> Result<Self> {
-        Ok(serde_json::from_str(json)?)
-    }
-
-    #[new]
-    pub fn new(val: IdMethod) -> Self {
-        Self(val)
-    }
-
-    pub fn get_method(&self) -> IdMethod {
-        self.0.clone()
-    }
-
-    pub fn __str__(&self) -> String {
-        self.__repr__()
-    }
-
-    pub fn __repr__(&self) -> String {
-        format!("InvokeDirect({})", self.0.__str__())
-    }
-
-    /// Return all strings referenced in the handle.
-    pub fn get_all_strings(&self) -> HashSet<DexString> {
-        self.0.get_all_strings()
-    }
-
-    /// Return all types referenced in the handle.
-    pub fn get_all_types(&self) -> HashSet<IdType> {
-        self.0.get_all_types()
-    }
-
-    /// Return all prototypes referenced in the handle.
-    pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
-        self.0.get_all_protos()
-    }
-
-    /// Return all field ids referenced in the handle.
-    pub fn get_all_field_ids(&self) -> HashSet<IdField> {
-        HashSet::new()
-    }
-
-    /// Return all method ids referenced in the handle.
-    pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
-        let mut methods = HashSet::new();
-        methods.insert(self.0.clone());
-        methods
-    }
-
-    /// Return all method handles referenced in the handle.
-    pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
-        let mut methods = HashSet::new();
-        methods.insert(MethodHandle::InvokeDirect(self.clone()));
-        methods
-    }
-}
-
-#[pyclass]
-#[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize, Serialize)]
-pub struct InvokeInterface(pub IdMethod);
-
-#[pymethods]
-impl InvokeInterface {
-    pub fn to_json(&self) -> Result<String> {
-        Ok(serde_json::to_string(self)?)
-    }
-
-    #[staticmethod]
-    pub fn from_json(json: &str) -> Result<Self> {
-        Ok(serde_json::from_str(json)?)
-    }
-
-    #[new]
-    pub fn new(val: IdMethod) -> Self {
-        Self(val)
-    }
-
-    pub fn get_method(&self) -> IdMethod {
-        self.0.clone()
-    }
-
-    pub fn __str__(&self) -> String {
-        self.__repr__()
-    }
-
-    pub fn __repr__(&self) -> String {
-        format!("InvokeInterface({})", self.0.__str__())
-    }
-
-    /// Return all strings referenced in the handle.
-    pub fn get_all_strings(&self) -> HashSet<DexString> {
-        self.0.get_all_strings()
-    }
-
-    /// Return all types referenced in the handle.
-    pub fn get_all_types(&self) -> HashSet<IdType> {
-        self.0.get_all_types()
-    }
-
-    /// Return all prototypes referenced in the handle.
-    pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
-        self.0.get_all_protos()
-    }
-
-    /// Return all field ids referenced in the handle.
-    pub fn get_all_field_ids(&self) -> HashSet<IdField> {
-        HashSet::new()
-    }
-
-    /// Return all method ids referenced in the handle.
-    pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
-        let mut methods = HashSet::new();
-        methods.insert(self.0.clone());
-        methods
-    }
-
-    /// Return all method handles referenced in the handle.
-    pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
-        let mut methods = HashSet::new();
-        methods.insert(MethodHandle::InvokeInterface(self.clone()));
-        methods
-    }
-}
-
-impl<'source> FromPyObject<'source> for MethodHandle {
-    fn extract(ob: &'source PyAny) -> PyResult<Self> {
-        if let Ok(val) = StaticPut::extract(ob) {
-            Ok(Self::StaticPut(val))
-        } else if let Ok(val) = StaticGet::extract(ob) {
-            Ok(Self::StaticGet(val))
-        } else if let Ok(val) = InstancePut::extract(ob) {
-            Ok(Self::InstancePut(val))
-        } else if let Ok(val) = InstanceGet::extract(ob) {
-            Ok(Self::InstanceGet(val))
-        } else if let Ok(val) = InvokeStatic::extract(ob) {
-            Ok(Self::InvokeStatic(val))
-        } else if let Ok(val) = InvokeInstance::extract(ob) {
-            Ok(Self::InvokeInstance(val))
-        } else if let Ok(val) = InvokeConstructor::extract(ob) {
-            Ok(Self::InvokeConstructor(val))
-        } else if let Ok(val) = InvokeDirect::extract(ob) {
-            Ok(Self::InvokeDirect(val))
-        } else if let Ok(val) = InvokeInterface::extract(ob) {
-            Ok(Self::InvokeInterface(val))
-        } else {
-            Err(PyErr::new::<PyTypeError, _>(format!(
-                "{} is not a castable as a MethodHandle",
-                ob.repr()?
-            )))
-        }
-    }
-}
-
-impl IntoPy<PyObject> for MethodHandle {
-    fn into_py(self, py: Python<'_>) -> PyObject {
         match self {
-            Self::StaticPut(val) => val.into_py(py),
+            Self::StaticPut { field } => v.visit_field_id(field)?,
-            Self::StaticGet(val) => val.into_py(py),
+            Self::StaticGet { field } => v.visit_field_id(field)?,
-            Self::InstancePut(val) => val.into_py(py),
+            Self::InstancePut { field } => v.visit_field_id(field)?,
-            Self::InstanceGet(val) => val.into_py(py),
+            Self::InstanceGet { field } => v.visit_field_id(field)?,
-            Self::InvokeStatic(val) => val.into_py(py),
+            Self::InvokeStatic { method } => v.visit_method_id(method)?,
-            Self::InvokeInstance(val) => val.into_py(py),
+            Self::InvokeInstance { method } => v.visit_method_id(method)?,
-            Self::InvokeConstructor(val) => val.into_py(py),
+            Self::InvokeConstructor { method } => v.visit_method_id(method)?,
-            Self::InvokeDirect(val) => val.into_py(py),
+            Self::InvokeDirect { method } => v.visit_method_id(method)?,
-            Self::InvokeInterface(val) => val.into_py(py),
+            Self::InvokeInterface { method } => v.visit_method_id(method)?,
+        }
+        Ok(())
+    }
+}
+
+impl<V: VisitorMut> VisitableMut<V> for MethodHandle {
+    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
+        match self {
+            Self::StaticPut { field } => Ok(Self::StaticPut {
+                field: v.visit_field_id(field)?,
+            }),
+            Self::StaticGet { field } => Ok(Self::StaticGet {
+                field: v.visit_field_id(field)?,
+            }),
+            Self::InstancePut { field } => Ok(Self::InstancePut {
+                field: v.visit_field_id(field)?,
+            }),
+            Self::InstanceGet { field } => Ok(Self::InstanceGet {
+                field: v.visit_field_id(field)?,
+            }),
+            Self::InvokeStatic { method } => Ok(Self::InvokeStatic {
+                method: v.visit_method_id(method)?,
+            }),
+            Self::InvokeInstance { method } => Ok(Self::InvokeInstance {
+                method: v.visit_method_id(method)?,
+            }),
+            Self::InvokeConstructor { method } => Ok(Self::InvokeConstructor {
+                method: v.visit_method_id(method)?,
+            }),
+            Self::InvokeDirect { method } => Ok(Self::InvokeDirect {
+                method: v.visit_method_id(method)?,
+            }),
+            Self::InvokeInterface { method } => Ok(Self::InvokeInterface {
+                method: v.visit_method_id(method)?,
+            }),
         }
     }
 }
 
+#[pymethods]
 impl MethodHandle {
-    // Not exposed to python, but meh, let's keep it coherent
+    pub fn to_json(&self) -> Result<String> {
-    pub fn __str__(&self) -> String {
+        Ok(serde_json::to_string(self)?)
-        match self {
+    }
-            Self::StaticPut(val) => val.__str__(),
+
-            Self::StaticGet(val) => val.__str__(),
+    #[staticmethod]
-            Self::InstancePut(val) => val.__str__(),
+    pub fn from_json(json: &str) -> Result<Self> {
-            Self::InstanceGet(val) => val.__str__(),
+        Ok(serde_json::from_str(json)?)
-            Self::InvokeStatic(val) => val.__str__(),
+    }
-            Self::InvokeInstance(val) => val.__str__(),
+
-            Self::InvokeConstructor(val) => val.__str__(),
+    pub fn __str__(&self) -> String {
-            Self::InvokeDirect(val) => val.__str__(),
+        self.__repr__()
-            Self::InvokeInterface(val) => val.__str__(),
-        }
     }
 
-    // Not exposed to python, but meh, let's keep it coherent
     pub fn __repr__(&self) -> String {
         match self {
-            Self::StaticPut(val) => val.__repr__(),
+            Self::StaticPut { field } => format!("StaticPut({})", field.__str__()),
-            Self::StaticGet(val) => val.__repr__(),
+            Self::StaticGet { field } => format!("StaticGet({})", field.__str__()),
-            Self::InstancePut(val) => val.__repr__(),
+            Self::InstancePut { field } => format!("InstancePut({})", field.__str__()),
-            Self::InstanceGet(val) => val.__repr__(),
+            Self::InstanceGet { field } => format!("InstanceGet({})", field.__str__()),
-            Self::InvokeStatic(val) => val.__repr__(),
+            Self::InvokeStatic { method } => format!("InvokeStatic({})", method.__str__()),
-            Self::InvokeInstance(val) => val.__repr__(),
+            Self::InvokeInstance { method } => format!("InvokeInstance({})", method.__str__()),
-            Self::InvokeConstructor(val) => val.__str__(),
+            Self::InvokeConstructor { method } => {
-            Self::InvokeDirect(val) => val.__repr__(),
+                format!("InvokeConstructor({})", method.__str__())
-            Self::InvokeInterface(val) => val.__repr__(),
+            }
+            Self::InvokeDirect { method } => format!("InvokeDirect({})", method.__str__()),
+            Self::InvokeInterface { method } => format!("InvokeInterface({})", method.__str__()),
         }
     }
 
     /// Return all strings referenced in the Handle.
     pub fn get_all_strings(&self) -> HashSet<DexString> {
-        match self {
+        let mut visitor = StringCollector::default();
-            Self::StaticPut(val) => val.get_all_strings(),
+        visitor.visit_method_handle(self).unwrap();
-            Self::StaticGet(val) => val.get_all_strings(),
+        visitor.result()
-            Self::InstancePut(val) => val.get_all_strings(),
-            Self::InstanceGet(val) => val.get_all_strings(),
-            Self::InvokeStatic(val) => val.get_all_strings(),
-            Self::InvokeInstance(val) => val.get_all_strings(),
-            Self::InvokeConstructor(val) => val.get_all_strings(),
-            Self::InvokeDirect(val) => val.get_all_strings(),
-            Self::InvokeInterface(val) => val.get_all_strings(),
-        }
     }
 
     /// Return all types referenced in the Handle.
     pub fn get_all_types(&self) -> HashSet<IdType> {
-        match self {
+        let mut visitor = TypeCollector::default();
-            Self::StaticPut(val) => val.get_all_types(),
+        visitor.visit_method_handle(self).unwrap();
-            Self::StaticGet(val) => val.get_all_types(),
+        visitor.result()
-            Self::InstancePut(val) => val.get_all_types(),
-            Self::InstanceGet(val) => val.get_all_types(),
-            Self::InvokeStatic(val) => val.get_all_types(),
-            Self::InvokeInstance(val) => val.get_all_types(),
-            Self::InvokeConstructor(val) => val.get_all_types(),
-            Self::InvokeDirect(val) => val.get_all_types(),
-            Self::InvokeInterface(val) => val.get_all_types(),
-        }
     }
 
     /// Return all prototypes referenced in the handle.
     pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
-        match self {
+        let mut visitor = MethodTypeCollector::default();
-            Self::StaticPut(val) => val.get_all_protos(),
+        visitor.visit_method_handle(self).unwrap();
-            Self::StaticGet(val) => val.get_all_protos(),
+        visitor.result()
-            Self::InstancePut(val) => val.get_all_protos(),
-            Self::InstanceGet(val) => val.get_all_protos(),
-            Self::InvokeStatic(val) => val.get_all_protos(),
-            Self::InvokeInstance(val) => val.get_all_protos(),
-            Self::InvokeConstructor(val) => val.get_all_protos(),
-            Self::InvokeDirect(val) => val.get_all_protos(),
-            Self::InvokeInterface(val) => val.get_all_protos(),
-        }
     }
 
     /// Return all field ids referenced in the handle.
     pub fn get_all_field_ids(&self) -> HashSet<IdField> {
-        match self {
+        let mut visitor = FieldIdCollector::default();
-            Self::StaticPut(val) => val.get_all_field_ids(),
+        visitor.visit_method_handle(self).unwrap();
-            Self::StaticGet(val) => val.get_all_field_ids(),
+        visitor.result()
-            Self::InstancePut(val) => val.get_all_field_ids(),
-            Self::InstanceGet(val) => val.get_all_field_ids(),
-            Self::InvokeStatic(val) => val.get_all_field_ids(),
-            Self::InvokeInstance(val) => val.get_all_field_ids(),
-            Self::InvokeConstructor(val) => val.get_all_field_ids(),
-            Self::InvokeDirect(val) => val.get_all_field_ids(),
-            Self::InvokeInterface(val) => val.get_all_field_ids(),
-        }
     }
 
     /// Return all method ids referenced in the handle.
     pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
-        match self {
+        let mut visitor = MethodIdCollector::default();
-            Self::StaticPut(val) => val.get_all_method_ids(),
+        visitor.visit_method_handle(self).unwrap();
-            Self::StaticGet(val) => val.get_all_method_ids(),
+        visitor.result()
-            Self::InstancePut(val) => val.get_all_method_ids(),
-            Self::InstanceGet(val) => val.get_all_method_ids(),
-            Self::InvokeStatic(val) => val.get_all_method_ids(),
-            Self::InvokeInstance(val) => val.get_all_method_ids(),
-            Self::InvokeConstructor(val) => val.get_all_method_ids(),
-            Self::InvokeDirect(val) => val.get_all_method_ids(),
-            Self::InvokeInterface(val) => val.get_all_method_ids(),
-        }
     }
 
     /// Return all method handles referenced in the handle.
     pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
-        match self {
+        let mut visitor = MethodHandleCollector::default();
-            Self::StaticPut(val) => val.get_all_method_handles(),
+        visitor.visit_method_handle(self).unwrap();
-            Self::StaticGet(val) => val.get_all_method_handles(),
+        visitor.result()
-            Self::InstancePut(val) => val.get_all_method_handles(),
-            Self::InstanceGet(val) => val.get_all_method_handles(),
-            Self::InvokeStatic(val) => val.get_all_method_handles(),
-            Self::InvokeInstance(val) => val.get_all_method_handles(),
-            Self::InvokeConstructor(val) => val.get_all_method_handles(),
-            Self::InvokeDirect(val) => val.get_all_method_handles(),
-            Self::InvokeInterface(val) => val.get_all_method_handles(),
-        }
     }
 }

androscalpel/src/visitor.rs

@@ -0,0 +1,189 @@
+//! The visitor trait and common implementations.
+
+use crate::{
+    ins::Instruction, DexString, IdEnum, IdField, IdMethod, IdMethodType, IdType, MethodHandle,
+    Result,
+};
+use std::collections::HashSet;
+
+pub trait Visitor: Sized {
+    fn visit_instruction(&mut self, ins: &Instruction) -> Result<()> {
+        ins.default_visit(self)
+    }
+    fn visit_string(&mut self, string: &DexString) -> Result<()> {
+        string.default_visit(self)
+    }
+    fn visit_type(&mut self, ty: &IdType) -> Result<()> {
+        ty.default_visit(self)
+    }
+    fn visit_method_type(&mut self, mty: &IdMethodType) -> Result<()> {
+        mty.default_visit(self)
+    }
+    fn visit_field_id(&mut self, id: &IdField) -> Result<()> {
+        id.default_visit(self)
+    }
+    fn visit_method_id(&mut self, id: &IdMethod) -> Result<()> {
+        id.default_visit(self)
+    }
+    fn visit_enum_id(&mut self, id: &IdEnum) -> Result<()> {
+        id.default_visit(self)
+    }
+    fn visit_method_handle(&mut self, handle: &MethodHandle) -> Result<()> {
+        handle.default_visit(self)
+    }
+}
+
+pub trait VisitorMut: Sized {
+    fn visit_instruction(&mut self, ins: Instruction) -> Result<Instruction> {
+        ins.default_visit_mut(self)
+    }
+    fn visit_string(&mut self, string: DexString) -> Result<DexString> {
+        string.default_visit_mut(self)
+    }
+    fn visit_type(&mut self, ty: IdType) -> Result<IdType> {
+        ty.default_visit_mut(self)
+    }
+    fn visit_method_type(&mut self, mty: IdMethodType) -> Result<IdMethodType> {
+        mty.default_visit_mut(self)
+    }
+    fn visit_field_id(&mut self, id: IdField) -> Result<IdField> {
+        id.default_visit_mut(self)
+    }
+    fn visit_method_id(&mut self, id: IdMethod) -> Result<IdMethod> {
+        id.default_visit_mut(self)
+    }
+    fn visit_enum_id(&mut self, id: IdEnum) -> Result<IdEnum> {
+        id.default_visit_mut(self)
+    }
+    fn visit_method_handle(&mut self, handle: MethodHandle) -> Result<MethodHandle> {
+        handle.default_visit_mut(self)
+    }
+}
+
+/// Trait for structures that can be visited.
+/// In theorie, this is not required, as the [`Visitor`] trait implement
+/// all tree traversal, however if the visitor do not do a complex
+/// operation it can be usefull to have a default tree traversal accessible
+/// to re-implemented trait methods.
+pub trait Visitable<V: Visitor> {
+    fn default_visit(&self, visitor: &mut V) -> Result<()>;
+}
+
+/// Trait for structures that can be modified by a visitor.
+/// In theorie, this is not required, as the [`VisitorMut`] trait implement
+/// all tree traversal, however if the visitor do not do a complex
+/// operation it can be usefull to have a default tree traversal accessible
+/// to re-implemented trait methods.
+pub trait VisitableMut<V: VisitorMut> {
+    fn default_visit_mut(self, visitor: &mut V) -> Result<Self>
+    where
+        Self: Sized;
+}
+
+#[derive(Debug, Default)]
+pub struct StringCollector {
+    pub strings: HashSet<DexString>,
+}
+
+impl StringCollector {
+    pub fn result(self) -> HashSet<DexString> {
+        self.strings
+    }
+}
+
+impl Visitor for StringCollector {
+    fn visit_string(&mut self, string: &DexString) -> Result<()> {
+        self.strings.insert(string.clone());
+        Ok(())
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct TypeCollector {
+    pub types: HashSet<IdType>,
+}
+
+impl TypeCollector {
+    pub fn result(self) -> HashSet<IdType> {
+        self.types
+    }
+}
+
+impl Visitor for TypeCollector {
+    fn visit_type(&mut self, ty: &IdType) -> Result<()> {
+        self.types.insert(ty.clone());
+        Ok(())
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct MethodTypeCollector {
+    pub method_types: HashSet<IdMethodType>,
+}
+
+impl MethodTypeCollector {
+    pub fn result(self) -> HashSet<IdMethodType> {
+        self.method_types
+    }
+}
+
+impl Visitor for MethodTypeCollector {
+    fn visit_method_type(&mut self, mty: &IdMethodType) -> Result<()> {
+        self.method_types.insert(mty.clone());
+        Ok(())
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct FieldIdCollector {
+    pub field_ids: HashSet<IdField>,
+}
+
+impl FieldIdCollector {
+    pub fn result(self) -> HashSet<IdField> {
+        self.field_ids
+    }
+}
+
+impl Visitor for FieldIdCollector {
+    fn visit_field_id(&mut self, id: &IdField) -> Result<()> {
+        self.field_ids.insert(id.clone());
+        Ok(())
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct MethodIdCollector {
+    pub method_ids: HashSet<IdMethod>,
+}
+
+impl MethodIdCollector {
+    pub fn result(self) -> HashSet<IdMethod> {
+        self.method_ids
+    }
+}
+
+impl Visitor for MethodIdCollector {
+    fn visit_method_id(&mut self, id: &IdMethod) -> Result<()> {
+        self.method_ids.insert(id.clone());
+        Ok(())
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct MethodHandleCollector {
+    pub handles: HashSet<MethodHandle>,
+}
+
+impl MethodHandleCollector {
+    pub fn result(self) -> HashSet<MethodHandle> {
+        self.handles
+    }
+}
+
+impl Visitor for MethodHandleCollector {
+    fn visit_method_handle(&mut self, handle: &MethodHandle) -> Result<()> {
+        self.handles.insert(handle.clone());
+        Ok(())
+    }
+}

androscalpel_serializer/src/core/mod.rs

@@ -33,6 +33,7 @@ impl std::fmt::Display for Error {
             Self::DeserializationError(msg) => write!(f, "{}", msg),
             Self::InvalidStringEncoding(msg) => write!(f, "{}", msg),
             Self::InconsistantStruct(msg) => write!(f, "{}", msg),
+            Self::OutOfBound(msg) => write!(f, "{}", msg),
         }
     }
 }

androscalpel_serializer/src/debug.rs

@@ -351,25 +351,29 @@ impl<'a> DebugStateMachine<'a> {
         Ok(self.debug_info.bytecode[self.pc])
     }
 
-    pub fn tick(&mut self) -> Result<Option<DebugInfo>> {
+    pub fn tick(&mut self) -> Option<DebugInfo> {
-        let ins = self.get_ins()?;
+        let ins = if let Ok(ins) = self.get_ins() {
+            ins
+        } else {
+            return Some(DebugInfo::EndOfData);
+        };
         self.pc += 1;
         match ins {
             DbgBytecode::EndSequence => {
                 self.pc = self.debug_info.bytecode.len();
-                Ok(Some(DebugInfo::EndOfData))
+                Some(DebugInfo::EndOfData)
             }
             DbgBytecode::AdvancePC {
                 addr_diff: Uleb128(addr_diff),
             } => {
                 self.address += addr_diff;
-                Ok(None)
+                None
             }
             DbgBytecode::AdvanceLine {
                 line_diff: Sleb128(line_diff),
             } => {
                 self.line = (self.line as i32 + line_diff) as u32;
-                Ok(None)
+                None
             }
             DbgBytecode::StartLocal {
                 register_num: Uleb128(register_num),
@@ -398,11 +402,11 @@ impl<'a> DebugStateMachine<'a> {
                     sig_idx: None,
                     in_scope: true,
                 };
-                Ok(Some(DebugInfo::DefLocal {
+                Some(DebugInfo::DefLocal {
                     addr: self.address,
                     reg: register_num,
                     val: self.register_states[register_num as usize],
-                }))
+                })
             }
             DbgBytecode::StartLocalExtended {
                 register_num: Uleb128(register_num),
@@ -436,20 +440,20 @@ impl<'a> DebugStateMachine<'a> {
                     },
                     in_scope: true,
                 };
-                Ok(Some(DebugInfo::DefLocal {
+                Some(DebugInfo::DefLocal {
                     addr: self.address,
                     reg: register_num,
                     val: self.register_states[register_num as usize],
-                }))
+                })
             }
             DbgBytecode::EndLocal {
                 register_num: Uleb128(register_num),
             } => {
                 self.register_states[register_num as usize].in_scope = false;
-                Ok(Some(DebugInfo::EndLocal {
+                Some(DebugInfo::EndLocal {
                     addr: self.address,
                     reg: register_num,
-                }))
+                })
             }
             DbgBytecode::RestartLocal {
                 register_num: Uleb128(register_num),
@@ -463,35 +467,35 @@ impl<'a> DebugStateMachine<'a> {
                     })
                 }
                 self.register_states[register_num as usize].in_scope = true;
-                Ok(Some(DebugInfo::DefLocal {
+                Some(DebugInfo::DefLocal {
                     addr: self.address,
                     reg: register_num,
                     val: self.register_states[register_num as usize],
-                }))
+                })
             }
             DbgBytecode::SetPrologueEnd => {
                 //self.prologue_end = true;
-                Ok(Some(DebugInfo::PrologueEnd { addr: self.address }))
+                Some(DebugInfo::PrologueEnd { addr: self.address })
             }
             DbgBytecode::SetEpilogueBegin => {
                 //self.epilogue_begin = true;
-                Ok(Some(DebugInfo::EpilogueBegin { addr: self.address }))
+                Some(DebugInfo::EpilogueBegin { addr: self.address })
             }
             DbgBytecode::SetFile { name_idx: NO_INDEX } => {
                 //self.source_file_idx = None;
-                Ok(Some(DebugInfo::SetSourceFile {
+                Some(DebugInfo::SetSourceFile {
                     addr: self.address,
                     source_file_idx: None,
-                }))
+                })
             }
             DbgBytecode::SetFile {
                 name_idx: Uleb128p1(name_idx),
             } => {
                 //self.source_file_idx = Some(name_idx);
-                Ok(Some(DebugInfo::SetSourceFile {
+                Some(DebugInfo::SetSourceFile {
                     addr: self.address,
                     source_file_idx: Some(name_idx),
-                }))
+                })
             }
             DbgBytecode::SpecialOpcode(op) => {
                 //if op >= 0x0a {
@@ -502,10 +506,10 @@ impl<'a> DebugStateMachine<'a> {
                 let adjusted_opcode = op as u32 - 0x0a;
                 self.line = (self.line as i32 + (adjusted_opcode as i32 % 15) - 4) as u32;
                 self.address += adjusted_opcode / 15;
-                Ok(Some(DebugInfo::SetLineNumber {
+                Some(DebugInfo::SetLineNumber {
                     addr: self.address,
                     line_num: self.line,
-                }))
+                })
             }
         }
     }