Call main only once

Author: daxpedda

crates/cli-support/src/lib.rs

@@ -331,7 +331,8 @@ impl Bindgen {
                 .context("failed getting Wasm module")?,
         };
 
-        self.threads
+        let thread_counter_addr = self
+            .threads
             .run(&mut module)
             .with_context(|| "failed to prepare module for threading")?;
 
@@ -363,6 +364,7 @@ impl Bindgen {
             &mut module,
             self.externref,
             self.wasm_interface_types,
+            thread_counter_addr,
             self.emit_start,
         )?;
 

crates/cli-support/src/wit/mod.rs

@@ -8,6 +8,7 @@ use std::str;
 use walrus::MemoryId;
 use walrus::{ExportId, FunctionId, ImportId, Module};
 use wasm_bindgen_shared::struct_function_export_name;
+use wasm_bindgen_threads_xform::ThreadCounterAddr;
 
 mod incoming;
 mod nonstandard;
@@ -32,6 +33,7 @@ struct Context<'a> {
     descriptors: HashMap<String, Descriptor>,
     externref_enabled: bool,
     wasm_interface_types: bool,
+    thread_counter_addr: Option<ThreadCounterAddr>,
     support_start: bool,
 }
 
@@ -47,6 +49,7 @@ pub fn process(
     module: &mut Module,
     externref_enabled: bool,
     wasm_interface_types: bool,
+    thread_counter_addr: Option<ThreadCounterAddr>,
     support_start: bool,
 ) -> Result<(NonstandardWitSectionId, WasmBindgenAuxId), Error> {
     let mut storage = Vec::new();
@@ -66,6 +69,7 @@ pub fn process(
         start_found: false,
         externref_enabled,
         wasm_interface_types,
+        thread_counter_addr,
         support_start,
     };
     cx.init()?;
@@ -476,22 +480,38 @@ impl<'a> Context<'a> {
             return Ok(());
         }
 
-        let prev_start = match self.module.start {
-            Some(f) => f,
-            None => {
-                self.module.start = Some(id);
-                return Ok(());
-            }
+        let prev_start = if let Some(prev_start) = self.module.start {
+            // Note that we call the previous start function, if any, first. This is
+            // because the start function currently only shows up when it's injected
+            // through thread/externref transforms. These injected start functions
+            // need to happen before user code, so we always schedule them first.
+            let mut builder = walrus::FunctionBuilder::new(&mut self.module.types, &[], &[]);
+            builder.func_body().call(prev_start);
+            Some(builder)
+        } else {
+            None
         };
 
-        // Note that we call the previous start function, if any, first. This is
-        // because the start function currently only shows up when it's injected
-        // through thread/externref transforms. These injected start functions
-        // need to happen before user code, so we always schedule them first.
-        let mut builder = walrus::FunctionBuilder::new(&mut self.module.types, &[], &[]);
-        builder.func_body().call(prev_start).call(id);
-        let new_start = builder.finish(Vec::new(), &mut self.module.funcs);
-        self.module.start = Some(new_start);
+        if let Some(thread_counter_addr) = self.thread_counter_addr {
+            // If we support shared memory, we re-use the thread counter address to
+            // call the start function only once.
+            let mut builder = prev_start
+                .unwrap_or_else(|| walrus::FunctionBuilder::new(&mut self.module.types, &[], &[]));
+
+            thread_counter_addr.wrap_start(self.memory()?, builder.func_body(), id);
+
+            self.module.start = Some(builder.finish(Vec::new(), &mut self.module.funcs));
+        } else if let Some(mut builder) = prev_start {
+            // If we don't support shared memory, we still have to call the start
+            // function if we had a previous start function.
+            builder.func_body().call(id);
+            self.module.start = Some(builder.finish(Vec::new(), &mut self.module.funcs));
+        } else {
+            // If we didn't have a previous start function and don't support shared memory,
+            // just add the start function.
+            self.module.start = Some(id);
+        }
+
         Ok(())
     }
 

crates/threads-xform/src/lib.rs

@@ -1,6 +1,8 @@
 use anyhow::{anyhow, bail, Error};
 use std::cmp;
 use std::env;
+use walrus::ir::BinaryOp;
+use walrus::ir::LoadKind;
 use walrus::ir::Value;
 use walrus::{
     ir::MemArg, ExportItem, FunctionId, GlobalId, GlobalKind, InitExpr, InstrSeqBuilder, MemoryId,
@@ -23,6 +25,9 @@ pub struct Config {
     enabled: bool,
 }
 
+#[derive(Clone, Copy)]
+pub struct ThreadCounterAddr(i32);
+
 impl Config {
     /// Create a new configuration with default settings.
     pub fn new() -> Config {
@@ -103,9 +108,9 @@ impl Config {
     /// * Some stack space is prepared for each thread after the first one.
     ///
     /// More and/or less may happen here over time, stay tuned!
-    pub fn run(&self, module: &mut Module) -> Result<(), Error> {
+    pub fn run(&self, module: &mut Module) -> Result<Option<ThreadCounterAddr>, Error> {
         if !self.is_enabled(module) {
-            return Ok(());
+            return Ok(None);
         }
 
         let memory = wasm_conventions::get_memory(module)?;
@@ -165,7 +170,25 @@ impl Config {
         // should not be called from an agent that cannot block (e.g. the main document thread).
         inject_destroy(module, &tls, &stack, memory)?;
 
-        Ok(())
+        Ok(Some(ThreadCounterAddr(thread_counter_addr)))
+    }
+}
+
+impl ThreadCounterAddr {
+    pub fn wrap_start(self, memory: MemoryId, mut body: InstrSeqBuilder, id: FunctionId) {
+        // This happens after the thread counter was already increased, so "1"
+        // means we are in the first thread.
+        body.i32_const(self.0)
+            .load(memory, LoadKind::I32 { atomic: true }, ATOMIC_MEM_ARG)
+            .i32_const(1)
+            .binop(BinaryOp::I32Eq)
+            .if_else(
+                None,
+                |body| {
+                    body.call(id);
+                },
+                |_| {},
+            );
     }
 }