<commit_before>use super::invoke::{Invoke, Identity};
use super::code::{self, SignaturesBuilder, FunctionBuilder};
use super::import;
use elements;

/// Module builder
pub struct ModuleBuilder<F=Identity> {
    callback: F,
    module: ModuleScaffold,
}

/// Location of the internal module function
pub struct CodeLocation {
    /// Location (index in 'functions' section) of the signature
    pub signature: u32,
    /// Location (index in the 'code' section) of the body
    pub body: u32,
}

#[derive(Default)]
struct ModuleScaffold {
    pub functions: elements::FunctionsSection,
    pub types: elements::TypeSection,
    pub import: elements::ImportSection,
    pub code: elements::CodeSection,
    pub other: Vec<elements::Section>,
}

impl From<elements::Module> for ModuleScaffold {
    fn from(module: elements::Module) -> Self {
        let mut funcs: Option<elements::FunctionsSection> = None;
        let mut types: Option<elements::TypeSection> = None;
        let mut import: Option<elements::ImportSection> = None;
        let mut code: Option<elements::CodeSection> = None;

        let mut sections = module.into_sections();
        while let Some(section) = sections.pop() {
            match section {
                elements::Section::Type(sect) => { types = Some(sect); }
                elements::Section::Function(sect) => { funcs = Some(sect); }
                elements::Section::Import(sect) => { import = Some(sect); }
                elements::Section::Code(sect) => { code = Some(sect); }
                _ => {}
            }
        }

        ModuleScaffold {
            functions: funcs.unwrap_or_default(),
            types: types.unwrap_or_default(),
            import: import.unwrap_or_default(),
            code: code.unwrap_or_default(),
            other: sections,
        }
    }
}

impl From<ModuleScaffold> for elements::Module {
    fn from(module: ModuleScaffold) -> Self {
        let mut sections = Vec::new();

        let types = module.types;
        if types.types().len() > 0 {
            sections.push(elements::Section::Type(types));
        }
        let import = module.import;
        if import.entries().len() > 0 {
            sections.push(elements::Section::Import(import));
        }                
        let functions = module.functions;
        if functions.entries().len() > 0 {
            sections.push(elements::Section::Function(functions));
        }        
        let code = module.code;
        if code.bodies().len() > 0 {
            sections.push(elements::Section::Code(code));
        }
        sections.extend(module.other);
        elements::Module::new(sections)
    }
}

impl ModuleBuilder {
    /// New empty module builder
    pub fn new() -> Self {
        ModuleBuilder::with_callback(Identity)
    }
}

impl<F> ModuleBuilder<F> where F: Invoke<elements::Module> {
    /// New module builder with bound callback
    pub fn with_callback(callback: F) -> Self {
        ModuleBuilder {
            callback: callback,
            module: Default::default(),
        }
    }

    /// Builder from raw module
    pub fn with_module(mut self, module: elements::Module) -> Self {
        self.module = module.into();
        self
    }

    /// Fill module with sections from iterator
    pub fn with_sections<I>(mut self, sections: I) -> Self 
        where I: IntoIterator<Item=elements::Section>
    {
        self.module.other.extend(sections);
        self
    }

    /// Add additional section
    pub fn with_section(mut self, section: elements::Section) -> Self {
        self.module.other.push(section);
        self
    }

    /// Binds to the type section, creates additional types when required
    pub fn with_signatures(mut self, bindings: code::SignatureBindings) -> Self {
        self.push_signatures(bindings);
        self
    }

    /// Push stand-alone function definition, creating sections, signature and code blocks
    /// in corresponding sections.
    /// `FunctionDefinition` can be build using `builder::function` builder
    pub fn push_function(&mut self, func: code::FunctionDefinition) -> CodeLocation {
        let signature = func.signature;
        let body = func.code;

        let type_ref = self.resolve_type_ref(signature);

        self.module.functions.entries_mut().push(elements::Func::new(type_ref));
        let signature_index = self.module.functions.entries_mut().len() as u32 - 1;
        self.module.code.bodies_mut().push(body);
        let body_index = self.module.code.bodies_mut().len() as u32 - 1;

        CodeLocation {
            signature: signature_index,
            body: body_index,
        }
    }

    fn resolve_type_ref(&mut self, signature: code::Signature) -> u32 {
        match signature {
            code::Signature::Inline(func_type) => {
                self.module.types.types_mut().push(elements::Type::Function(func_type));
                self.module.types.types().len() as u32 - 1
            }
            code::Signature::TypeReference(type_ref) => {
                type_ref
            }
        }
    }

    /// Push one function signature, returning it's calling index.
    /// Can create corresponding type in type section.
    pub fn push_signature(&mut self, signature: code::Signature) -> u32 {
        self.resolve_type_ref(signature)
    }

    /// Push signatures in the module, returning corresponding indices of pushed signatures
    pub fn push_signatures(&mut self, signatures: code::SignatureBindings) -> Vec<u32> {
        let mut result = Vec::new();

        // todo: maybe reuse existing types with the equal signatures
        let raw_functions: Vec<u32> = signatures.into_iter().map(|binding|
            match binding {
                code::Signature::Inline(func_type) => {
                    self.module.types.types_mut().push(elements::Type::Function(func_type));
                    self.module.types.types().len() as u32 - 1
                }
                code::Signature::TypeReference(type_ref) => {
                    type_ref
                }
            }
        ).collect();

        for function in raw_functions {
            self.module.functions.entries_mut().push(elements::Func::new(function));
            result.push(self.module.functions.entries_mut().len() as u32 - 1);
        }

        result
    }

    pub fn function(self) -> FunctionBuilder<Self> {
        FunctionBuilder::with_callback(self)
    }

    /// Define functions section
    pub fn functions(self) -> SignaturesBuilder<Self> {
        SignaturesBuilder::with_callback(self)
    }

    /// With inserted import entry<commit_msg>fix also for signature collection<commit_after>use super::invoke::{Invoke, Identity};
use super::code::{self, SignaturesBuilder, FunctionBuilder};
use super::import;
use elements;

/// Module builder
pub struct ModuleBuilder<F=Identity> {
    callback: F,
    module: ModuleScaffold,
}

/// Location of the internal module function
pub struct CodeLocation {
    /// Location (index in 'functions' section) of the signature
    pub signature: u32,
    /// Location (index in the 'code' section) of the body
    pub body: u32,
}

#[derive(Default)]
struct ModuleScaffold {
    pub functions: elements::FunctionsSection,
    pub types: elements::TypeSection,
    pub import: elements::ImportSection,
    pub code: elements::CodeSection,
    pub other: Vec<elements::Section>,
}

impl From<elements::Module> for ModuleScaffold {
    fn from(module: elements::Module) -> Self {
        let mut funcs: Option<elements::FunctionsSection> = None;
        let mut types: Option<elements::TypeSection> = None;
        let mut import: Option<elements::ImportSection> = None;
        let mut code: Option<elements::CodeSection> = None;

        let mut sections = module.into_sections();
        while let Some(section) = sections.pop() {
            match section {
                elements::Section::Type(sect) => { types = Some(sect); }
                elements::Section::Function(sect) => { funcs = Some(sect); }
                elements::Section::Import(sect) => { import = Some(sect); }
                elements::Section::Code(sect) => { code = Some(sect); }
                _ => {}
            }
        }

        ModuleScaffold {
            functions: funcs.unwrap_or_default(),
            types: types.unwrap_or_default(),
            import: import.unwrap_or_default(),
            code: code.unwrap_or_default(),
            other: sections,
        }
    }
}

impl From<ModuleScaffold> for elements::Module {
    fn from(module: ModuleScaffold) -> Self {
        let mut sections = Vec::new();

        let types = module.types;
        if types.types().len() > 0 {
            sections.push(elements::Section::Type(types));
        }
        let import = module.import;
        if import.entries().len() > 0 {
            sections.push(elements::Section::Import(import));
        }                
        let functions = module.functions;
        if functions.entries().len() > 0 {
            sections.push(elements::Section::Function(functions));
        }        
        let code = module.code;
        if code.bodies().len() > 0 {
            sections.push(elements::Section::Code(code));
        }
        sections.extend(module.other);
        elements::Module::new(sections)
    }
}

impl ModuleBuilder {
    /// New empty module builder
    pub fn new() -> Self {
        ModuleBuilder::with_callback(Identity)
    }
}

impl<F> ModuleBuilder<F> where F: Invoke<elements::Module> {
    /// New module builder with bound callback
    pub fn with_callback(callback: F) -> Self {
        ModuleBuilder {
            callback: callback,
            module: Default::default(),
        }
    }

    /// Builder from raw module
    pub fn with_module(mut self, module: elements::Module) -> Self {
        self.module = module.into();
        self
    }

    /// Fill module with sections from iterator
    pub fn with_sections<I>(mut self, sections: I) -> Self 
        where I: IntoIterator<Item=elements::Section>
    {
        self.module.other.extend(sections);
        self
    }

    /// Add additional section
    pub fn with_section(mut self, section: elements::Section) -> Self {
        self.module.other.push(section);
        self
    }

    /// Binds to the type section, creates additional types when required
    pub fn with_signatures(mut self, bindings: code::SignatureBindings) -> Self {
        self.push_signatures(bindings);
        self
    }

    /// Push stand-alone function definition, creating sections, signature and code blocks
    /// in corresponding sections.
    /// `FunctionDefinition` can be build using `builder::function` builder
    pub fn push_function(&mut self, func: code::FunctionDefinition) -> CodeLocation {
        let signature = func.signature;
        let body = func.code;

        let type_ref = self.resolve_type_ref(signature);

        self.module.functions.entries_mut().push(elements::Func::new(type_ref));
        let signature_index = self.module.functions.entries_mut().len() as u32 - 1;
        self.module.code.bodies_mut().push(body);
        let body_index = self.module.code.bodies_mut().len() as u32 - 1;

        CodeLocation {
            signature: signature_index,
            body: body_index,
        }
    }

    fn resolve_type_ref(&mut self, signature: code::Signature) -> u32 {
        match signature {
            code::Signature::Inline(func_type) => {
                // todo: maybe search for existing type
                self.module.types.types_mut().push(elements::Type::Function(func_type));
                self.module.types.types().len() as u32 - 1
            }
            code::Signature::TypeReference(type_ref) => {
                type_ref
            }
        }
    }

    /// Push one function signature, returning it's calling index.
    /// Can create corresponding type in type section.
    pub fn push_signature(&mut self, signature: code::Signature) -> u32 {
        self.resolve_type_ref(signature)
    }

    /// Push signatures in the module, returning corresponding indices of pushed signatures
    pub fn push_signatures(&mut self, signatures: code::SignatureBindings) -> Vec<u32> {
        signatures.into_iter().map(|binding|
            self.resolve_type_ref(binding)
        ).collect()
    }

    pub fn function(self) -> FunctionBuilder<Self> {
        FunctionBuilder::with_callback(self)
    }

    /// Define functions section
    pub fn functions(self) -> SignaturesBuilder<Self> {
        SignaturesBuilder::with_callback(self)
    }

    /// With inserted import entry
{'fraction_non_alphanumeric': 0.1195118085275004, 'fraction_numerical': 0.0042796005706134095, 'mean_word_length': 2.353441403135796, 'pattern_counts': {'":': 0, '<': 36, '<?xml version=': 0, '>': 88, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2319965', 'n_tokens_mistral': 3439, 'n_tokens_neox': 3334, 'n_words': 844}