WIP

androscalpel/src/annotation.rs

@@ -111,6 +111,20 @@ impl DexAnnotationItem {
     }
 }
 
+impl<V: Visitor> Visitable<V> for DexAnnotationItem {
+    fn default_visit(&self, v: &mut V) -> Result<()> {
+        v.visit_annotation(&self.annotation)
+    }
+}
+impl<V: VisitorMut> VisitableMut<V> for DexAnnotationItem {
+    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
+        Ok(Self {
+            annotation: v.visit_annotation(self.annotation)?,
+            ..self
+        })
+    }
+}
+
 /// An annotation.
 #[pyclass]
 #[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]

androscalpel/src/class.rs

@@ -7,7 +7,7 @@ use pyo3::prelude::*;
 
 use crate::{
     DexAnnotationItem, DexString, Field, IdField, IdMethod, IdMethodType, IdType, Method,
-    MethodHandle, Result,
+    MethodHandle, Result, Visitable, VisitableMut, Visitor, VisitorMut,
 };
 use androscalpel_serializer::consts::*;
 
@@ -387,3 +387,92 @@ impl Class {
         self == other
     }
 }
+
+impl<V: Visitor> Visitable<V> for Class {
+    fn default_visit(&self, v: &mut V) -> Result<()> {
+        v.visit_type(&self.descriptor)?;
+        if let Some(superclass) = &self.superclass {
+            v.visit_type(&superclass)?;
+        }
+        for interface in &self.interfaces {
+            v.visit_type(&interface)?;
+        }
+        if let Some(source_file) = &self.source_file {
+            v.visit_string(&source_file)?;
+        }
+        for (id, field) in &self.static_fields {
+            v.visit_field_id(&id)?;
+            v.visit_field(&field)?;
+        }
+        for (id, field) in &self.instance_fields {
+            v.visit_field_id(&id)?;
+            v.visit_field(&field)?;
+        }
+        for (id, meth) in &self.direct_methods {
+            v.visit_method_id(&id)?;
+            v.visit_method(&meth)?;
+        }
+        for (id, meth) in &self.virtual_methods {
+            v.visit_method_id(&id)?;
+            v.visit_method(&meth)?;
+        }
+        for annot in &self.annotations {
+            v.visit_annotation_item(&annot)?;
+        }
+        Ok(())
+    }
+}
+impl<V: VisitorMut> VisitableMut<V> for Class {
+    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
+        Ok(Self {
+            descriptor: v.visit_type(self.descriptor)?,
+            superclass: self
+                .superclass
+                .map(|superclass| v.visit_type(superclass))
+                .transpose()?,
+            interfaces: self
+                .interfaces
+                .into_iter()
+                .map(|interface| v.visit_type(interface))
+                .collect::<Result<_>>()?,
+            source_file: self
+                .source_file
+                .map(|source_file| v.visit_string(source_file))
+                .transpose()?,
+            static_fields: {
+                let mut static_fields = HashMap::new();
+                for (id, field) in self.static_fields.into_iter() {
+                    static_fields.insert(v.visit_field_id(id)?, v.visit_field(field)?);
+                }
+                static_fields
+            },
+            instance_fields: {
+                let mut instance_fields = HashMap::new();
+                for (id, field) in self.instance_fields.into_iter() {
+                    instance_fields.insert(v.visit_field_id(id)?, v.visit_field(field)?);
+                }
+                instance_fields
+            },
+            direct_methods: {
+                let mut direct_methods = HashMap::new();
+                for (id, meth) in self.direct_methods.into_iter() {
+                    direct_methods.insert(v.visit_method_id(id)?, v.visit_method(meth)?);
+                }
+                direct_methods
+            },
+            virtual_methods: {
+                let mut virtual_methods = HashMap::new();
+                for (id, meth) in self.virtual_methods.into_iter() {
+                    virtual_methods.insert(v.visit_method_id(id)?, v.visit_method(meth)?);
+                }
+                virtual_methods
+            },
+            annotations: self
+                .annotations
+                .into_iter()
+                .map(|annotation| v.visit_annotation_item(annotation))
+                .collect::<Result<_>>()?,
+            ..self
+        })
+    }
+}

androscalpel/src/code.rs

@@ -8,6 +8,7 @@ use pyo3::prelude::*;
 
 use crate::{
     ins::Instruction, DexString, IdField, IdMethod, IdMethodType, IdType, MethodHandle, Result,
+    Visitable, VisitableMut, Visitor, VisitorMut,
 };
 
 // TODO: make this easy to edit/manipulate, maybe move to Method
@@ -493,3 +494,38 @@ impl Code {
         })
     }
 }
+
+impl<V: Visitor> Visitable<V> for Code {
+    fn default_visit(&self, v: &mut V) -> Result<()> {
+        for ins in &self.insns {
+            v.visit_instruction(&ins)?;
+        }
+        Ok(())
+    }
+}
+impl<V: VisitorMut> VisitableMut<V> for Code {
+    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
+        let parameter_names = if let Some(parameter_names) = self.parameter_names {
+            let mut new_param = vec![];
+            for param in parameter_names {
+                if let Some(param) = param {
+                    new_param.push(Some(v.visit_string(param)?));
+                } else {
+                    new_param.push(None);
+                }
+            }
+            Some(new_param)
+        } else {
+            None
+        };
+        Ok(Self {
+            parameter_names,
+            insns: self
+                .insns
+                .into_iter()
+                .map(|ins| v.visit_instruction(ins))
+                .collect::<Result<_>>()?,
+            ..self
+        })
+    }
+}

androscalpel/src/field.rs

@@ -7,7 +7,7 @@ use pyo3::prelude::*;
 
 use crate::{
     DexAnnotationItem, DexString, DexValue, HiddenApiData, IdField, IdMethod, IdMethodType, IdType,
-    MethodHandle, Result,
+    MethodHandle, Result, Visitable, VisitableMut, Visitor, VisitorMut,
 };
 use androscalpel_serializer::consts::*;
 
@@ -257,3 +257,40 @@ impl Field {
         self == other
     }
 }
+
+impl<V: Visitor> Visitable<V> for Field {
+    fn default_visit(&self, v: &mut V) -> Result<()> {
+        v.visit_field_id(&self.descriptor)?;
+        v.visit_field_visibility(&self.visibility)?;
+        if let Some(val) = &self.value {
+            v.visit_value(&val)?;
+        }
+        for annot in &self.annotations {
+            v.visit_annotation_item(&annot)?;
+        }
+        if let Some(hiddenapi) = &self.hiddenapi {
+            v.visit_hidden_api_data(&hiddenapi)?;
+        }
+
+        Ok(())
+    }
+}
+impl<V: VisitorMut> VisitableMut<V> for Field {
+    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
+        Ok(Self {
+            descriptor: v.visit_field_id(self.descriptor)?,
+            visibility: v.visit_field_visibility(self.visibility)?,
+            value: self.value.map(|val| v.visit_value(val)).transpose()?,
+            annotations: self
+                .annotations
+                .into_iter()
+                .map(|annot| v.visit_annotation_item(annot))
+                .collect::<Result<_>>()?,
+            hiddenapi: self
+                .hiddenapi
+                .map(|hiddenapi| v.visit_hidden_api_data(hiddenapi))
+                .transpose()?,
+            ..self
+        })
+    }
+}

androscalpel/src/hiddenapi.rs

@@ -1,5 +1,6 @@
 //! Representation of the hidden API data
 
+use crate::{Result, Visitable, VisitableMut, Visitor, VisitorMut};
 use log::warn;
 use pyo3::prelude::*;
 use serde::{Deserialize, Serialize};
@@ -39,6 +40,22 @@ impl From<HiddenApiData> for androscalpel_serializer::HiddenApiFlags {
     }
 }
 
+impl<V: Visitor> Visitable<V> for HiddenApiData {
+    fn default_visit(&self, v: &mut V) -> Result<()> {
+        v.visit_hidden_api_permission(&self.permission)?;
+        v.visit_hidden_api_domain(&self.domain)?;
+        Ok(())
+    }
+}
+impl<V: VisitorMut> VisitableMut<V> for HiddenApiData {
+    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
+        Ok(Self {
+            permission: v.visit_hidden_api_permission(self.permission)?,
+            domain: v.visit_hidden_api_domain(self.domain)?,
+        })
+    }
+}
+
 #[pyclass]
 #[derive(Debug, PartialEq, Eq, Hash, Clone, Deserialize, Serialize)]
 pub enum HiddenApiPermission {

androscalpel/src/method.rs

@@ -7,7 +7,7 @@ use pyo3::prelude::*;
 
 use crate::{
     Code, DexAnnotationItem, DexString, HiddenApiData, IdField, IdMethod, IdMethodType, IdType,
-    MethodHandle, Result,
+    MethodHandle, Result, Visitable, VisitableMut, Visitor, VisitorMut,
 };
 use androscalpel_serializer::consts::*;
 
@@ -298,3 +298,54 @@ impl Method {
         self == other
     }
 }
+
+impl<V: Visitor> Visitable<V> for Method {
+    fn default_visit(&self, v: &mut V) -> Result<()> {
+        v.visit_method_id(&self.descriptor)?;
+        v.visit_method_visibility(&self.visibility)?;
+        for annot in &self.annotations {
+            v.visit_annotation_item(&annot)?;
+        }
+        for parameter_annotations in &self.parameters_annotations {
+            for annot in parameter_annotations {
+                v.visit_annotation_item(&annot)?;
+            }
+        }
+        if let Some(hiddenapi) = &self.hiddenapi {
+            v.visit_hidden_api_data(&hiddenapi)?;
+        }
+        if let Some(code) = &self.code {
+            v.visit_code(&code)?;
+        }
+        Ok(())
+    }
+}
+impl<V: VisitorMut> VisitableMut<V> for Method {
+    fn default_visit_mut(self, v: &mut V) -> Result<Self> {
+        Ok(Self {
+            descriptor: v.visit_method_id(self.descriptor)?,
+            visibility: v.visit_method_visibility(self.visibility)?,
+            annotations: self
+                .annotations
+                .into_iter()
+                .map(|annot| v.visit_annotation_item(annot))
+                .collect::<Result<_>>()?,
+            parameters_annotations: self
+                .parameters_annotations
+                .into_iter()
+                .map(|parameter_annotations| {
+                    parameter_annotations
+                        .into_iter()
+                        .map(|annot| v.visit_annotation_item(annot))
+                        .collect::<Result<_>>()
+                })
+                .collect::<Result<_>>()?,
+            hiddenapi: self
+                .hiddenapi
+                .map(|hiddenapi| v.visit_hidden_api_data(hiddenapi))
+                .transpose()?,
+            code: self.code.map(|code| v.visit_code(code)).transpose()?,
+            ..self
+        })
+    }
+}

androscalpel/src/visitor.rs

@@ -1,8 +1,10 @@
 //! The visitor trait and common implementations.
 
 use crate::{
-    ins::Instruction, scalar::*, CallSite, DexAnnotation, DexString, DexValue, IdEnum, IdField,
-    IdMethod, IdMethodType, IdType, MethodHandle, Result,
+    ins::Instruction, scalar::*, CallSite, Class, Code, DexAnnotation, DexAnnotationItem,
+    DexString, DexValue, Field, FieldVisibility, HiddenApiData, HiddenApiDomain,
+    HiddenApiPermission, IdEnum, IdField, IdMethod, IdMethodType, IdType, Method, MethodHandle,
+    MethodVisibility, Result,
 };
 use std::collections::HashSet;
 
@@ -70,6 +72,36 @@ pub trait Visitor: Sized {
     fn visit_bool(&mut self, _: &DexBoolean) -> Result<()> {
         Ok(())
     }
+    fn visit_class(&mut self, class: &Class) -> Result<()> {
+        class.default_visit(self)
+    }
+    fn visit_field(&mut self, field: &Field) -> Result<()> {
+        field.default_visit(self)
+    }
+    fn visit_method(&mut self, method: &Method) -> Result<()> {
+        method.default_visit(self)
+    }
+    fn visit_annotation_item(&mut self, annotation: &DexAnnotationItem) -> Result<()> {
+        annotation.default_visit(self)
+    }
+    fn visit_field_visibility(&mut self, _: &FieldVisibility) -> Result<()> {
+        Ok(())
+    }
+    fn visit_method_visibility(&mut self, _: &MethodVisibility) -> Result<()> {
+        Ok(())
+    }
+    fn visit_hidden_api_data(&mut self, hidden_api: &HiddenApiData) -> Result<()> {
+        hidden_api.default_visit(self)
+    }
+    fn visit_hidden_api_permission(&mut self, _: &HiddenApiPermission) -> Result<()> {
+        Ok(())
+    }
+    fn visit_hidden_api_domain(&mut self, _: &HiddenApiDomain) -> Result<()> {
+        Ok(())
+    }
+    fn visit_code(&mut self, code: &Code) -> Result<()> {
+        code.default_visit(self)
+    }
 }
 
 pub trait VisitorMut: Sized {
@@ -136,6 +168,45 @@ pub trait VisitorMut: Sized {
     fn visit_bool(&mut self, val: DexBoolean) -> Result<DexBoolean> {
         Ok(val)
     }
+    fn visit_class(&mut self, class: Class) -> Result<Class> {
+        class.default_visit_mut(self)
+    }
+    fn visit_field(&mut self, field: Field) -> Result<Field> {
+        field.default_visit_mut(self)
+    }
+    fn visit_method(&mut self, method: Method) -> Result<Method> {
+        method.default_visit_mut(self)
+    }
+    fn visit_annotation_item(
+        &mut self,
+        annotation: DexAnnotationItem,
+    ) -> Result<DexAnnotationItem> {
+        annotation.default_visit_mut(self)
+    }
+    fn visit_field_visibility(&mut self, visibility: FieldVisibility) -> Result<FieldVisibility> {
+        Ok(visibility)
+    }
+    fn visit_method_visibility(
+        &mut self,
+        visibility: MethodVisibility,
+    ) -> Result<MethodVisibility> {
+        Ok(visibility)
+    }
+    fn visit_hidden_api_data(&mut self, hidden_api: HiddenApiData) -> Result<HiddenApiData> {
+        hidden_api.default_visit_mut(self)
+    }
+    fn visit_hidden_api_permission(
+        &mut self,
+        perm: HiddenApiPermission,
+    ) -> Result<HiddenApiPermission> {
+        Ok(perm)
+    }
+    fn visit_hidden_api_domain(&mut self, domain: HiddenApiDomain) -> Result<HiddenApiDomain> {
+        Ok(domain)
+    }
+    fn visit_code(&mut self, code: Code) -> Result<Code> {
+        code.default_visit_mut(self)
+    }
 }
 
 /// Trait for structures that can be visited.