From ca2597973e3f3bff43ec056e841ed1e3e71b3445 Mon Sep 17 00:00:00 2001
From: Brian Carroll <brian.carroll.ireland@gmail.com>
Date: Wed, 12 Jan 2022 08:45:24 +0000
Subject: [PATCH] Wasm: store function_count on the ImportSection

---
 compiler/gen_wasm/src/lib.rs                  |  3 +-
 compiler/gen_wasm/src/wasm_module/sections.rs | 28 +++++++++++++++----
 .../src/helpers/wasm32_test_result.rs         |  2 +-
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/compiler/gen_wasm/src/lib.rs b/compiler/gen_wasm/src/lib.rs
index d084e75260..d1f71b9619 100644
--- a/compiler/gen_wasm/src/lib.rs
+++ b/compiler/gen_wasm/src/lib.rs
@@ -96,8 +96,7 @@ pub fn build_module_help<'a>(
     let initial_module = WasmModule::preload(env.arena, preload_bytes);
 
     // Adjust Wasm function indices to account for functions from the object file
-    let runtime_import_fn_count: u32 = initial_module.import.function_count(); // to be imported at runtime (e.g. WASI)
-    let fn_index_offset: u32 = runtime_import_fn_count + initial_module.code.preloaded_count;
+    let fn_index_offset: u32 = initial_module.import.function_count + initial_module.code.preloaded_count;
 
     // Get a map of name to index for the preloaded functions
     // Assumes the preloaded object file has all symbols exported, as per `zig build-lib -dymamic`
diff --git a/compiler/gen_wasm/src/wasm_module/sections.rs b/compiler/gen_wasm/src/wasm_module/sections.rs
index fb19adc2ea..2366d62b09 100644
--- a/compiler/gen_wasm/src/wasm_module/sections.rs
+++ b/compiler/gen_wasm/src/wasm_module/sections.rs
@@ -56,7 +56,7 @@ pub trait Section<'a>: Sized {
 }
 
 macro_rules! section_impl {
-    ($structname: ident, $id: expr) => {
+    ($structname: ident, $id: expr, $from_count_and_bytes: expr) => {
         impl<'a> Section<'a> for $structname<'a> {
             const ID: SectionId = $id;
 
@@ -72,7 +72,7 @@ macro_rules! section_impl {
                 let (count, initial_bytes) = parse_section(Self::ID, module_bytes, cursor);
                 let mut bytes = Vec::with_capacity_in(initial_bytes.len() * 2, arena);
                 bytes.extend_from_slice(initial_bytes);
-                $structname { bytes, count }
+                $from_count_and_bytes(count, bytes)
             }
 
             fn size(&self) -> usize {
@@ -84,6 +84,13 @@ macro_rules! section_impl {
             }
         }
     };
+
+    ($structname: ident, $id: expr) => {
+        section_impl!($structname, $id, |count, bytes| $structname {
+            bytes,
+            count
+        });
+    };
 }
 
 impl<'a, Sec> Serialize for Sec
@@ -367,6 +374,7 @@ impl Serialize for Import {
 #[derive(Debug)]
 pub struct ImportSection<'a> {
     pub count: u32,
+    pub function_count: u32,
     pub bytes: Vec<'a, u8>,
 }
 
@@ -376,7 +384,7 @@ impl<'a> ImportSection<'a> {
         self.count += 1;
     }
 
-    pub fn function_count(&self) -> u32 {
+    fn update_function_count(&mut self) {
         let mut f_count = 0;
         let mut cursor = 0;
         while cursor < self.bytes.len() {
@@ -403,11 +411,21 @@ impl<'a> ImportSection<'a> {
             }
         }
 
-        f_count
+        self.function_count = f_count;
+    }
+
+    pub fn from_count_and_bytes(count: u32, bytes: Vec<'a, u8>) -> Self {
+        let mut created = ImportSection {
+            bytes,
+            count,
+            function_count: 0,
+        };
+        created.update_function_count();
+        created
     }
 }
 
-section_impl!(ImportSection, SectionId::Import);
+section_impl!(ImportSection, SectionId::Import, ImportSection::from_count_and_bytes);
 
 /*******************************************************************
  *
diff --git a/compiler/test_gen/src/helpers/wasm32_test_result.rs b/compiler/test_gen/src/helpers/wasm32_test_result.rs
index 34e2a86bea..b96f3e03fa 100644
--- a/compiler/test_gen/src/helpers/wasm32_test_result.rs
+++ b/compiler/test_gen/src/helpers/wasm32_test_result.rs
@@ -14,7 +14,7 @@ pub trait Wasm32TestResult {
         wrapper_name: &str,
         main_function_index: u32,
     ) {
-        let index = module.import.function_count()
+        let index = module.import.function_count
             + module.code.preloaded_count
             + module.code.code_builders.len() as u32;