diff --git a/src/instrumentation.rs b/src/instrumentation.rs
index a803cbc..8db0541 100644
--- a/src/instrumentation.rs
+++ b/src/instrumentation.rs
@@ -120,7 +120,7 @@ pub fn instrument(m: &mut Module, config: Config) -> Result<(), String> {
             && id != dynamic_counter_func
             && id != dynamic_counter64_func
         {
-            let is_partial_tracing = trace_only_ids.get(&id).is_some();
+            let is_partial_tracing = trace_only_ids.contains(&id);
             inject_profiling_prints(&m.types, printer, id, func, is_partial_tracing, &vars);
         }
     }
diff --git a/src/optimize.rs b/src/optimize.rs
index 9203f10..6e2b5e3 100644
--- a/src/optimize.rs
+++ b/src/optimize.rs
@@ -90,6 +90,15 @@ pub fn optimize(
     if let Some(max_size) = always_inline_max_function_size {
         optimizations.always_inline_max_size(*max_size);
     }
+    // The feature set should be align with IC `wasmtime` validation config:
+    // https://github.com/dfinity/ic/blob/6a6470d705a0f36fb94743b12892280409f85688/rs/embedders/src/wasm_utils/validation.rs#L1385
+    optimizations.enable_feature(wasm_opt::Feature::MutableGlobals);
+    optimizations.enable_feature(wasm_opt::Feature::TruncSat);
+    optimizations.enable_feature(wasm_opt::Feature::Simd);
+    optimizations.enable_feature(wasm_opt::Feature::BulkMemory);
+    optimizations.enable_feature(wasm_opt::Feature::SignExt);
+    optimizations.enable_feature(wasm_opt::Feature::ReferenceTypes);
+    optimizations.enable_feature(wasm_opt::Feature::Memory64);
     optimizations.run(temp_file.path(), temp_file.path())?;
 
     // read optimized wasm back in from temp file