4 files changed, 49 insertions, 2 deletions

diff --git a/embassy-stm32-wpan/Cargo.toml b/embassy-stm32-wpan/Cargo.toml
index 4f53a400a..360ca5f4b 100644
--- a/embassy-stm32-wpan/Cargo.toml
+++ b/embassy-stm32-wpan/Cargo.toml
@@ -44,6 +44,8 @@ defmt = ["dep:defmt", "embassy-sync/defmt", "embassy-embedded-hal/defmt", "embas
 ble = ["dep:stm32wb-hci"]
 mac = ["dep:bitflags", "dep:embassy-net-driver" ]
 
+extended = []
+
 stm32wb10cc = [ "embassy-stm32/stm32wb10cc" ]
 stm32wb15cc = [ "embassy-stm32/stm32wb15cc" ]
 stm32wb30ce = [ "embassy-stm32/stm32wb30ce" ]
diff --git a/embassy-stm32-wpan/build.rs b/embassy-stm32-wpan/build.rs
index 94aac070d..7ab458bf2 100644
--- a/embassy-stm32-wpan/build.rs
+++ b/embassy-stm32-wpan/build.rs
@@ -18,9 +18,22 @@ fn main() {
     // stm32wb tl_mbox link sections
 
     let out_file = out_dir.join("tl_mbox.x").to_string_lossy().to_string();
-    fs::write(out_file, fs::read_to_string("tl_mbox.x.in").unwrap()).unwrap();
+    let in_file;
+    if env::var_os("CARGO_FEATURE_EXTENDED").is_some() {
+        if env::vars()
+            .map(|(a, _)| a)
+            .any(|x| x.starts_with("CARGO_FEATURE_STM32WB1"))
+        {
+            in_file = "tl_mbox_extended_wb1.x.in";
+        } else {
+            in_file = "tl_mbox_extended_wbx5.x.in";
+        }
+    } else {
+        in_file = "tl_mbox.x.in";
+    }
+    fs::write(out_file, fs::read_to_string(in_file).unwrap()).unwrap();
     println!("cargo:rustc-link-search={}", out_dir.display());
-    println!("cargo:rerun-if-changed=tl_mbox.x.in");
+    println!("cargo:rerun-if-changed={}", in_file);
 }
 
 enum GetOneError {
diff --git a/embassy-stm32-wpan/tl_mbox_extended_wb1.x.in b/embassy-stm32-wpan/tl_mbox_extended_wb1.x.in
new file mode 100644
index 000000000..4cffdaddd
--- /dev/null
+++ b/embassy-stm32-wpan/tl_mbox_extended_wb1.x.in
@@ -0,0 +1,16 @@
+MEMORY 
+{
+    RAM_SHARED (xrw)           : ORIGIN = 0x20030000, LENGTH = 4K
+    RAMB_SHARED (xrw)          : ORIGIN = 0x20030028, LENGTH = 4K
+}
+
+/*
+ * Scatter the mailbox interface memory sections in shared memory
+ */
+SECTIONS
+{
+    TL_REF_TABLE                     (NOLOAD) : { *(TL_REF_TABLE) } >RAM_SHARED
+
+    MB_MEM1 (NOLOAD)                          : { *(MB_MEM1) } >RAMB_SHARED
+    MB_MEM2 (NOLOAD)                          : { _sMB_MEM2 = . ; *(MB_MEM2) ; _eMB_MEM2 = . ; } >RAMB_SHARED
+}
diff --git a/embassy-stm32-wpan/tl_mbox_extended_wbx5.x.in b/embassy-stm32-wpan/tl_mbox_extended_wbx5.x.in
new file mode 100644
index 000000000..281d637a9
--- /dev/null
+++ b/embassy-stm32-wpan/tl_mbox_extended_wbx5.x.in
@@ -0,0 +1,16 @@
+MEMORY 
+{
+    RAM_SHARED (xrw)           : ORIGIN = 0x20030000, LENGTH = 2K
+    RAMB_SHARED (xrw)          : ORIGIN = 0x20038000, LENGTH = 10K
+}
+
+/*
+ * Scatter the mailbox interface memory sections in shared memory
+ */
+SECTIONS
+{
+    TL_REF_TABLE                     (NOLOAD) : { *(TL_REF_TABLE) } >RAM_SHARED
+
+    MB_MEM1 (NOLOAD)                          : { *(MB_MEM1) } >RAMB_SHARED
+    MB_MEM2 (NOLOAD)                          : { _sMB_MEM2 = . ; *(MB_MEM2) ; _eMB_MEM2 = . ; } >RAMB_SHARED
+}