From f85ef8f5f839a2072a7143fabac8a231a360632c Mon Sep 17 00:00:00 2001
From: NikVolf <nikvolf@gmail.com>
Date: Fri, 7 Apr 2017 20:00:30 +0300
Subject: [PATCH] also module extensions

---
 src/builder/module.rs   | 41 +++++++++++++++++++++++++++++++++++++++++
 src/elements/section.rs |  8 +++++++-
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/src/builder/module.rs b/src/builder/module.rs
index 7e81319..4f192bc 100644
--- a/src/builder/module.rs
+++ b/src/builder/module.rs
@@ -8,11 +8,20 @@ pub struct ModuleBuilder<F=Identity> {
     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>,
 }
 
@@ -21,6 +30,7 @@ impl From<elements::Module> for ModuleScaffold {
         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() {
@@ -28,6 +38,7 @@ impl From<elements::Module> for ModuleScaffold {
                 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); }
                 _ => {}
             }
         }
@@ -36,6 +47,7 @@ impl From<elements::Module> for ModuleScaffold {
             functions: funcs.unwrap_or_default(),
             types: types.unwrap_or_default(),
             import: import.unwrap_or_default(),
+            code: code.unwrap_or_default(),
             other: sections,
         }
     }
@@ -98,6 +110,35 @@ impl<F> ModuleBuilder<F> where F: Invoke<elements::Module> {
         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 module = &mut self.module;
+
+        let type_ref = match signature {
+            code::Signature::Inline(func_type) => {
+                module.types.types_mut().push(elements::Type::Function(func_type));
+                module.types.types().len() as u32 - 1
+            }
+            code::Signature::TypeReference(type_ref) => {
+                type_ref
+            }
+        };
+
+        module.functions.entries_mut().push(elements::Func::new(type_ref));
+        let signature_index = module.functions.entries_mut().len() as u32 - 1;
+        module.code.bodies_mut().push(body);
+        let body_index = module.code.bodies_mut().len() as u32 - 1;
+
+        CodeLocation {
+            signature: signature_index,
+            body: body_index,
+        }
+    }
+
     /// Push signatures in the module, returning corresponding indices of pushed signatures
     pub fn push_signatures(&mut self, signatures: code::SignatureBindings) -> Vec<u32> {
         let module = &mut self.module;
diff --git a/src/elements/section.rs b/src/elements/section.rs
index c2220b1..947eadc 100644
--- a/src/elements/section.rs
+++ b/src/elements/section.rs
@@ -457,6 +457,7 @@ impl Serialize for ExportSection {
 }
 
 /// Section with function bodies of the module.
+#[derive(Default)]
 pub struct CodeSection(Vec<FuncBody>);
 
 impl CodeSection {
@@ -465,10 +466,15 @@ impl CodeSection {
         CodeSection(bodies)
     }
 
-    /// All function bodies in the section
+    /// All function bodies in the section.
     pub fn bodies(&self) -> &[FuncBody] {
         &self.0
     }
+
+    /// All function bodies in the section, mutable.
+    pub fn bodies_mut(&mut self) -> &mut Vec<FuncBody> {
+        &mut self.0
+    }
 }
 
 impl Deserialize for CodeSection {