reinsert data sections

src/optimize.rs

@@ -8,14 +8,17 @@ pub fn optimize(m: &mut Module, keep_name_section: bool) {
     let temp_file_name = "temp.opt.wasm";
     let temp_path = PathBuf::from(temp_file_name);
 
-    // pull out the custom sections to preserve
+    // pull out a copy of the custom sections to preserve
+    let m_copy = parse_wasm(&m.emit_wasm(), keep_name_section).unwrap();
     let mut metadata_sections = Vec::new();
-    list_metadata(m).iter().for_each(|full_name| {
+    list_metadata(&m_copy).iter().for_each(|full_name| {
         match full_name.strip_prefix("icp:public ") {
-            Some(name) => metadata_sections.push(("public", name, get_metadata(m, name).unwrap())),
+            Some(name) => {
+                metadata_sections.push(("public", name, get_metadata(&m_copy, name).unwrap()))
+            }
             None => match full_name.strip_prefix("icp:private ") {
                 Some(name) => {
-                    metadata_sections.push(("private", name, get_metadata(m, name).unwrap()))
+                    metadata_sections.push(("private", name, get_metadata(&m_copy, name).unwrap()))
                 }
                 None => unreachable!(),
             },
@@ -30,8 +33,18 @@ pub fn optimize(m: &mut Module, keep_name_section: bool) {
         .run(temp_file_name, temp_file_name)
         .unwrap();
 
-    // FIXME re-insert the custom section before assigning back to m
-
-    // read back in from fs and assign back to m
+    // read optimized wasm back in from fs
     *m = parse_wasm_file(temp_path, keep_name_section).unwrap();
+
+    // re-insert the custom sections
+    metadata_sections
+        .iter()
+        .for_each(|(visibility, name, data)| {
+            let visibility = match *visibility {
+                "public" => Kind::Public,
+                "private" => Kind::Private,
+                _ => unreachable!(),
+            };
+            add_metadata(m, visibility, name, data.to_vec());
+        });
 }