//! Mutate custom sections.

use super::Mutator;

use rand::prelude::*;

use std::borrow::Cow;

#[derive(Clone, Copy)]

pub struct CustomSectionMutator;

impl Mutator for CustomSectionMutator {

    fn can_mutate(&self, config: &crate::WasmMutate) -> bool {

        config.info().has_custom_section()

    }

    fn mutate<'a>(

        &self,

        config: &'a mut crate::WasmMutate,

    ) -> crate::Result<Box<dyn Iterator<Item = crate::Result<wasm_encoder::Module>> + 'a>> {

        let custom_section_indices: Vec<_> = config

            .info()

            .raw_sections

            .iter()

            .enumerate()

            .filter(|(_i, s)| s.id == wasm_encoder::SectionId::Custom as u8)

            .map(|(i, _s)| i)

            .collect();

        assert!(!custom_section_indices.is_empty());

        let custom_section_index = *custom_section_indices.choose(config.rng()).unwrap();

        let reader = config.info().get_binary_reader(custom_section_index);

        let old_custom_section = wasmparser::CustomSectionReader::new(reader).unwrap();

        let name_string;

        let data_vec;

        let mut name = old_custom_section.name();

        let mut data = old_custom_section.data();

        if config.rng().random_ratio(1, 20) {

            // Mutate the custom section's name.

            let mut new_name = name.to_string().into_bytes();

            config.raw_mutate(

                &mut new_name,

                if config.reduce {

                    name.len().saturating_sub(1)

                } else {

                    std::cmp::max(name.len() * 2, 32)

                },

            )?;

            name_string = String::from_utf8_lossy(&new_name).to_string();

            name = &name_string;

        } else {

            // Mutate the custom section's data.

            let mut new_data = data.to_vec();

            config.raw_mutate(

                &mut new_data,

                if config.reduce {

                    data.len().saturating_sub(1)

                } else {

                    std::cmp::max(data.len() * 2, 32)

                },

            )?;

            data_vec = new_data;

            data = &data_vec;

        };

        Ok(Box::new(std::iter::once(Ok(config

            .info()

            .replace_section(

                custom_section_index,

                &wasm_encoder::CustomSection {

                    name: name.into(),

                    data: Cow::Borrowed(data),

                },

            )))))

    }

}

#[derive(Clone, Copy)]

pub struct AddCustomSectionMutator;

const MAX_NEW_DATA_LEN: usize = 100;

const MAX_NEW_NAME_LEN: usize = 20;

impl Mutator for AddCustomSectionMutator {

    fn can_mutate(&self, config: &crate::WasmMutate) -> bool {

        !config.reduce

    }

    fn mutate<'a>(

        &self,

        config: &'a mut crate::WasmMutate,

    ) -> crate::Result<Box<dyn Iterator<Item = crate::Result<wasm_encoder::Module>> + 'a>> {

        let num_sections = config.info().raw_sections.len();

        let new_custom_section_idx = config.rng().random_range(0..=num_sections);

        let mut name = vec![];

        config.raw_mutate(&mut name, MAX_NEW_NAME_LEN)?;

        let name = String::from_utf8_lossy(&name);

        let mut data = vec![];

        config.raw_mutate(&mut data, MAX_NEW_DATA_LEN)?;

        Ok(Box::new(std::iter::once(Ok(config.info().insert_section(

            new_custom_section_idx,

            &wasm_encoder::CustomSection {

                name: name.into(),

                data: Cow::Borrowed(&data),

            },

        )))))

    }

}

#[derive(Copy, Clone)]

pub struct ReorderCustomSectionMutator;

impl Mutator for ReorderCustomSectionMutator {

    fn can_mutate(&self, config: &crate::WasmMutate) -> bool {

        config.info().has_custom_section() && config.info().raw_sections.len() > 1

    }

    fn mutate<'a>(

        &self,

        config: &'a mut crate::WasmMutate,

    ) -> crate::Result<Box<dyn Iterator<Item = crate::Result<wasm_encoder::Module>> + 'a>> {

        let custom_section_indices: Vec<_> = config

            .info()

            .raw_sections

            .iter()

            .enumerate()

            .filter(|(_i, s)| s.id == wasm_encoder::SectionId::Custom as u8)

            .map(|(i, _s)| i)

            .collect();

        assert!(!custom_section_indices.is_empty());

        let src_idx = *custom_section_indices.choose(config.rng()).unwrap();

        let num_sections = config.info().raw_sections.len();

        let mut dest_idx;

        loop {

            dest_idx = config.rng().random_range(0..num_sections);

            if dest_idx != src_idx {

                break;

            }

            config.consume_fuel(1)?;

        }

        Ok(Box::new(std::iter::once(Ok(config

            .info()

            .move_section(src_idx, dest_idx)))))

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn test_add_custom_section() {

        crate::mutators::match_mutation(

            r#"

                (module)

            "#,

            AddCustomSectionMutator,

            r#"

                (module

                    (@custom "" "")

                )

            "#,

        );

    }

    #[test]

    fn test_grow_custom_section() {

        crate::mutators::match_mutation(

            r#"

                (module

                    (@custom "name" "data")

                )

            "#,

            CustomSectionMutator,

            r#"

                (module

                    (@custom "name" "datadata")

                )

            "#,

        );

    }

    #[test]

    fn test_shrink_custom_section() {

        crate::mutators::match_mutation(

            r#"

                (module

                    (@custom "name" "data")

                )

            "#,

            CustomSectionMutator,

            r#"

                (module

                    (@custom "name" "d")

                )

            "#,

        );

    }

    #[test]

    fn test_mutate_custom_section() {

        crate::mutators::match_mutation(

            r#"

                (module

                    (@custom "name" "data")

                )

            "#,

            CustomSectionMutator,

            r#"

                (module

                    (@custom "name" "aaaa")

                )

            "#,

        );

    }

    #[test]

    fn test_mutate_custom_section_name() {

        crate::mutators::match_mutation(

            r#"

                (module

                    (@custom "custom-name" "data")

                )

            "#,

            CustomSectionMutator,

            r#"

                (module

                    (@custom "custom" "data")

                )

            "#,

        );

    }

    #[test]

    fn test_reorder_custom_section() {

        crate::mutators::match_mutation(

            r#"

            (module

                (@custom "name2" "data")

                (@custom "name3" "data")

            )

            "#,

            ReorderCustomSectionMutator,

            r#"

            (module

                (@custom "name3" "data")

                (@custom "name2" "data")

            )

            "#,

        )

    }

    #[test]

    fn test_reorder_custom_section2() {

        crate::mutators::match_mutation(

            r#"

            (module

                (func (export "x"))

                (@custom "name2" "data")

                (@custom "name3" "data")

            )

            "#,

            ReorderCustomSectionMutator,

            r#"

            (module

                (@custom "name3" "data")

                (func (export "x"))

                (@custom "name2" "data")

            )

            "#,

        )

    }

}