//! The visitor trait and common implementations.

use crate::{
    ins::Instruction, scalar::*, Apk, CallSite, Class, Code, DexAnnotation, DexAnnotationItem,
    DexFile, DexString, DexValue, Field, FieldVisibility, HiddenApiData, HiddenApiDomain,
    HiddenApiPermission, IdEnum, IdField, IdMethod, IdMethodType, IdType, Method, MethodHandle,
    MethodVisibility, 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)
    }
    fn visit_call_site(&mut self, site: &CallSite) -> Result<()> {
        site.default_visit(self)
    }
    fn visit_value(&mut self, val: &DexValue) -> Result<()> {
        val.default_visit(self)
    }
    fn visit_byte(&mut self, _: &DexByte) -> Result<()> {
        Ok(())
    }
    fn visit_short(&mut self, _: &DexShort) -> Result<()> {
        Ok(())
    }
    fn visit_char(&mut self, _: &DexChar) -> Result<()> {
        Ok(())
    }
    fn visit_int(&mut self, _: &DexInt) -> Result<()> {
        Ok(())
    }
    fn visit_long(&mut self, _: &DexLong) -> Result<()> {
        Ok(())
    }
    fn visit_float(&mut self, _: &DexFloat) -> Result<()> {
        Ok(())
    }
    fn visit_double(&mut self, _: &DexDouble) -> Result<()> {
        Ok(())
    }
    fn visit_array(&mut self, val: &DexArray) -> Result<()> {
        val.default_visit(self)
    }
    fn visit_annotation(&mut self, val: &DexAnnotation) -> Result<()> {
        val.default_visit(self)
    }
    fn visit_null(&mut self, _: &DexNull) -> Result<()> {
        Ok(())
    }
    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)
    }
    fn visit_dex_file(&mut self, dex_file: &DexFile) -> Result<()> {
        dex_file.default_visit(self)
    }
    fn visit_apk(&mut self, apk: &Apk) -> Result<()> {
        apk.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)
    }
    fn visit_call_site(&mut self, site: CallSite) -> Result<CallSite> {
        site.default_visit_mut(self)
    }
    fn visit_value(&mut self, val: DexValue) -> Result<DexValue> {
        val.default_visit_mut(self)
    }
    fn visit_byte(&mut self, val: DexByte) -> Result<DexByte> {
        Ok(val)
    }
    fn visit_short(&mut self, val: DexShort) -> Result<DexShort> {
        Ok(val)
    }
    fn visit_char(&mut self, val: DexChar) -> Result<DexChar> {
        Ok(val)
    }
    fn visit_int(&mut self, val: DexInt) -> Result<DexInt> {
        Ok(val)
    }
    fn visit_long(&mut self, val: DexLong) -> Result<DexLong> {
        Ok(val)
    }
    fn visit_float(&mut self, val: DexFloat) -> Result<DexFloat> {
        Ok(val)
    }
    fn visit_double(&mut self, val: DexDouble) -> Result<DexDouble> {
        Ok(val)
    }
    fn visit_array(&mut self, val: DexArray) -> Result<DexArray> {
        val.default_visit_mut(self)
    }
    fn visit_annotation(&mut self, val: DexAnnotation) -> Result<DexAnnotation> {
        val.default_visit_mut(self)
    }
    fn visit_null(&mut self, val: DexNull) -> Result<DexNull> {
        Ok(val)
    }
    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)
    }
    fn visit_dex_file(&mut self, dex_file: DexFile) -> Result<DexFile> {
        dex_file.default_visit_mut(self)
    }
    fn visit_apk(&mut self, apk: Apk) -> Result<Apk> {
        apk.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(())
    }
}
