From 2f1a25c5699b8d3b787d19bf0c4ea1560fc25c11 Mon Sep 17 00:00:00 2001
From: Graham Kelly <gkgoat6700@gmail.com>
Date: Sat, 3 Feb 2024 13:47:08 -0500
Subject: [PATCH] custom sections

---
 Cargo.toml         |  1 +
 src/backend/mod.rs |  4 ++++
 src/frontend.rs    | 10 +++++++---
 src/ir/module.rs   |  4 ++++
 4 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 0fe3fe1..bdde85a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -23,6 +23,7 @@ addr2line = "0.19"
 # For fuzzing only. Versions must match those in fuzz/Cargo.toml.
 libfuzzer-sys = { version = "0.4", optional = true }
 wasm-smith = { version = "0.8", optional = true }
+indexmap = "2.2.2"
 
 [features]
 default = []
diff --git a/src/backend/mod.rs b/src/backend/mod.rs
index 4fc8632..e6ef716 100644
--- a/src/backend/mod.rs
+++ b/src/backend/mod.rs
@@ -7,6 +7,7 @@ use crate::Operator;
 use anyhow::Result;
 use rayon::prelude::*;
 use std::borrow::Cow;
+use wasm_encoder::CustomSection;
 
 pub mod stackify;
 use stackify::{Context as StackifyContext, WasmBlock};
@@ -787,6 +788,9 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result<Vec<u8>> {
     }
     names.functions(&func_names);
     into_mod.section(&names);
+    for (k, v) in module.custom_sections.iter() {
+        into_mod.section(&CustomSection { name: &k, data: &v });
+    }
 
     Ok(into_mod.finish())
 }
diff --git a/src/frontend.rs b/src/frontend.rs
index a7a86e8..ed462e3 100644
--- a/src/frontend.rs
+++ b/src/frontend.rs
@@ -297,7 +297,11 @@ fn handle_payload<'a>(
             extra_sections.debug_rnglists =
                 gimli::DebugRngLists::new(reader.data(), gimli::LittleEndian);
         }
-        Payload::CustomSection(_) => {}
+        Payload::CustomSection(reader) => {
+            module
+                .custom_sections
+                .insert(reader.name().to_owned(), reader.data().to_owned());
+        }
         Payload::Version { .. } => {}
         Payload::ElementSection(reader) => {
             for element in reader {
@@ -1151,9 +1155,9 @@ impl<'a, 'b> FunctionBodyBuilder<'a, 'b> {
             | wasmparser::Operator::TableGet { .. }
             | wasmparser::Operator::TableSet { .. }
             | wasmparser::Operator::TableGrow { .. }
-            | wasmparser::Operator::TableSize { .. } 
+            | wasmparser::Operator::TableSize { .. }
             | wasmparser::Operator::MemoryCopy { .. }
-            | wasmparser::Operator::MemoryFill { .. }=> {
+            | wasmparser::Operator::MemoryFill { .. } => {
                 self.emit(Operator::try_from(&op).unwrap(), loc)?
             }
 
diff --git a/src/ir/module.rs b/src/ir/module.rs
index 05795db..dcb1b9e 100644
--- a/src/ir/module.rs
+++ b/src/ir/module.rs
@@ -3,6 +3,7 @@ use crate::entity::{EntityRef, EntityVec};
 use crate::ir::{Debug, DebugMap, FunctionBody};
 use crate::{backend, frontend};
 use anyhow::Result;
+use indexmap::IndexMap;
 
 pub use crate::frontend::FrontendOptions;
 
@@ -19,6 +20,7 @@ pub struct Module<'a> {
     pub start_func: Option<Func>,
     pub debug: Debug,
     pub debug_map: DebugMap,
+    pub custom_sections: IndexMap<String,Vec<u8>>
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -143,6 +145,7 @@ impl<'a> Module<'a> {
             start_func: None,
             debug: Debug::default(),
             debug_map: DebugMap::default(),
+            custom_sections: IndexMap::new(),
         }
     }
 
@@ -165,6 +168,7 @@ impl<'a> Module<'a> {
             start_func: self.start_func,
             debug: self.debug,
             debug_map: self.debug_map,
+            custom_sections: self.custom_sections,
         }
     }
 }