Add environment variable overrides for stack, heap

If you define a runtime, you can call `stack_size_env_override` to
define an optional environment variable checked by the runtime builder.
Same goes for the heap. A user can set these environment variables to
override the runtime's stack / heap size. You can use this package's
examples to try it out; see the updated build script.

There's no default environment variable for either memory region. The
package that defines the runtime needs to opt-in to this feature.

4 files changed, 215 insertions, 17 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index aa37d8c..17b4417 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,14 @@
 
 ## [Unreleased]
 
+Add configurations to `RuntimeBuilder`:
+
+- `stack_size_env_override`
+- `heap_size_env_override`
+
+Use these methods to define environment variables that can override the
+stack / heap sizes.
+
 ## [0.1.3] 2023-10-01
 
 Ensure that the runtime supports the GNU linker, `ld`.
diff --git a/board/build.rs b/board/build.rs
index 9ee64e8..8474e85 100644
--- a/board/build.rs
+++ b/board/build.rs
@@ -28,6 +28,9 @@ fn main() {
                     .data(imxrt_rt::Memory::Dtcm)
                     .bss(imxrt_rt::Memory::Dtcm)
                     .uninit(imxrt_rt::Memory::Dtcm)
+                    .stack_size_env_override("THIS_WONT_BE_CONSIDERED")
+                    .stack_size_env_override("BOARD_STACK")
+                    .heap_size_env_override("BOARD_HEAP")
                     .build()
                     .unwrap()
             }
@@ -37,6 +40,8 @@ fn main() {
             )
             .heap_size(1024)
             .rodata(imxrt_rt::Memory::Flash)
+            .stack_size_env_override("BOARD_STACK")
+            .heap_size_env_override("BOARD_HEAP")
             .build()
             .unwrap(),
             "imxrt1170evk_cm7" => imxrt_rt::RuntimeBuilder::from_flexspi(
@@ -44,6 +49,8 @@ fn main() {
                 16 * 1024 * 1024,
             )
             .rodata(imxrt_rt::Memory::Dtcm)
+            .stack_size_env_override("BOARD_STACK")
+            .heap_size_env_override("BOARD_HEAP")
             .build()
             .unwrap(),
             _ => continue,
diff --git a/src/host.rs b/src/host.rs
index 131457a..dbe023f 100644
--- a/src/host.rs
+++ b/src/host.rs
@@ -120,6 +120,63 @@ struct FlashOpts {
     flexspi: FlexSpi,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct EnvOverride {
+    default: usize,
+    env: Option<String>,
+}
+
+impl EnvOverride {
+    const fn new(default: usize) -> Self {
+        Self { default, env: None }
+    }
+    fn set_env_key(&mut self, key: String) {
+        self.env = Some(key);
+    }
+    fn read(&self) -> Result<usize, Box<dyn std::error::Error>> {
+        if let Some(env) = &self.env {
+            // If the user sets multiple environment variables for the same runtime
+            // property (like stack, heap), we will only re-run when the variable
+            // we care about changes. An example might help:
+            //
+            //      let mut bldr = RuntimeBuilder::from_flexspi(/* ... */);
+            //      bldr.stack_size_env_override("COMMON_STACK");
+            //      if special_condition() {
+            //          bldr.stack_size_env_override("SPECIAL_STACK");
+            //      }
+            //
+            // If we take the branch, then we re-run the build if SPECIAL_STACK
+            // changes. Otherwise, we re-run if COMMON_STACK changes.
+            //
+            // I previously put this into `set_env_key`. That would mean we re-run
+            // the build if _either_ enviroment variable changes. But then I thought
+            // about the user who writes their build script like
+            //
+            //      if special_condition() {
+            //          bldr.stack_size_env_override("SPECIAL_STACK");
+            //      } else {
+            //          bldr.stack_size_env_override("COMMON_STACK");
+            //      }
+            //
+            // and how that would introduce different re-run behavior than the first
+            // example, even though the variable selection is the same. Coupling the
+            // re-run behavior to the variable selection behavior seems less surprising.
+            println!("cargo:rerun-if-env-changed={env}");
+        }
+
+        if let Some(val) = self.env.as_ref().and_then(|key| env::var(key).ok()) {
+            let val = if val.ends_with('k') || val.ends_with('K') {
+                val[..val.len() - 1].parse::<usize>()? * 1024
+            } else {
+                val.parse::<usize>()?
+            };
+            Ok(val)
+        } else {
+            Ok(self.default)
+        }
+    }
+}
+
 /// Builder for the i.MX RT runtime.
 ///
 /// `RuntimeBuilder` let you assign sections to memory regions. It also lets
@@ -169,6 +226,71 @@ struct FlashOpts {
 ///
 /// assert_eq!(b, RuntimeBuilder::from_flexspi(family, FLASH_SIZE));
 /// ```
+///
+/// # Environment overrides
+///
+/// Certain memory regions, like the stack and heap, can be sized using environment
+/// variables. As the provider of the runtime, you can use `*_env_override` methods
+/// to select the environment variable(s) that others may use to set the size, in bytes,
+/// for these memory regions.
+///
+/// The rest of this section describes how environment variables interact with other
+/// methods on this builder. Although the examples use stack size, the concepts apply
+/// to all regions that can be sized with environment variables.
+///
+/// ```no_run
+/// # use imxrt_rt::{Family, RuntimeBuilder, Memory};
+/// # const FLASH_SIZE: usize = 16 * 1024;
+/// # let family = Family::Imxrt1060;
+/// RuntimeBuilder::from_flexspi(family, FLASH_SIZE)
+///     .stack_size_env_override("YOUR_STACK_SIZE")
+///     // ...
+///     # .build().unwrap();
+/// ```
+///
+/// In the above example, if a user set an environment variable `YOUR_STACK_SIZE=1024`, then
+/// the runtime's stack size is 1024. Otherwise, the stack size is the default stack size.
+///
+/// > As a convenience, a user can use a `k` or `K` suffix to specify multiples of 1024 bytes.
+/// > For example, the environment variables `YOUR_STACK_SIZE=4k` and `YOUR_STACK_SIZE=4K` are
+/// > each equivalent to `YOUR_STACK_SIZE=4096`.
+///
+/// ```no_run
+/// # use imxrt_rt::{Family, RuntimeBuilder, Memory};
+/// # const FLASH_SIZE: usize = 16 * 1024;
+/// # let family = Family::Imxrt1060;
+/// RuntimeBuilder::from_flexspi(family, FLASH_SIZE)
+///     .stack_size_env_override("YOUR_STACK_SIZE")
+///     .stack_size(2048)
+///     // ...
+///     # .build().unwrap();
+///
+/// RuntimeBuilder::from_flexspi(family, FLASH_SIZE)
+///     .stack_size(2048)
+///     .stack_size_env_override("YOUR_STACK_SIZE")
+///     // ...
+///     # .build().unwrap();
+/// ```
+///
+/// In the above example, the two builders produce the same runtime. The builder
+/// selects the stack size from the environment variable, if available. Otherwise,
+/// the stack size is 2048 bytes. The call order is irrelevant, since the builder
+/// doesn't consult the environment until you invoke [`build()`](Self::build).
+///
+/// ```no_run
+/// # use imxrt_rt::{Family, RuntimeBuilder, Memory};
+/// # const FLASH_SIZE: usize = 16 * 1024;
+/// # let family = Family::Imxrt1060;
+/// RuntimeBuilder::from_flexspi(family, FLASH_SIZE)
+///     .stack_size_env_override("INVALIDATED")
+///     .stack_size_env_override("YOUR_STACK_SIZE")
+///     // ...
+///     # .build().unwrap();
+/// ````
+///
+/// In the above example, `YOUR_STACK_SIZE` invalidates the call with `INVALIDATED`.
+/// Therefore, `YOUR_STACK_SIZE` controls the stack size, if set. Otherwise, the stack
+/// size is the default stack size.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct RuntimeBuilder {
     family: Family,
@@ -180,9 +302,9 @@ pub struct RuntimeBuilder {
     bss: Memory,
     uninit: Memory,
     stack: Memory,
-    stack_size: usize,
+    stack_size: EnvOverride,
     heap: Memory,
-    heap_size: usize,
+    heap_size: EnvOverride,
     flash_opts: Option<FlashOpts>,
     linker_script_name: String,
 }
@@ -205,9 +327,9 @@ impl RuntimeBuilder {
             bss: Memory::Ocram,
             uninit: Memory::Ocram,
             stack: Memory::Dtcm,
-            stack_size: 8 * 1024,
+            stack_size: EnvOverride::new(8 * 1024),
             heap: Memory::Dtcm,
-            heap_size: 0,
+            heap_size: EnvOverride::new(0),
             flash_opts: Some(FlashOpts {
                 size: flash_size,
                 flexspi: FlexSpi::family_default(family),
@@ -261,7 +383,15 @@ impl RuntimeBuilder {
     }
     /// Set the size, in bytes, of the stack.
     pub fn stack_size(&mut self, bytes: usize) -> &mut Self {
-        self.stack_size = bytes;
+        self.stack_size.default = bytes;
+        self
+    }
+    /// Let end users override the stack size using an environment variable.
+    ///
+    /// See the [environment overrides](Self#environment-overrides) documentation
+    /// for more information.
+    pub fn stack_size_env_override(&mut self, key: impl AsRef<str>) -> &mut Self {
+        self.stack_size.set_env_key(key.as_ref().into());
         self
     }
     /// Set the memory placement for the heap.
@@ -274,7 +404,15 @@ impl RuntimeBuilder {
     }
     /// Set the size, in bytes, of the heap.
     pub fn heap_size(&mut self, bytes: usize) -> &mut Self {
-        self.heap_size = bytes;
+        self.heap_size.default = bytes;
+        self
+    }
+    /// Let end users override the heap size using an environment variable.
+    ///
+    /// See the [environment overrides](Self#environment-overrides) documentation
+    /// for more information.
+    pub fn heap_size_env_override(&mut self, key: impl AsRef<str>) -> &mut Self {
+        self.heap_size.set_env_key(key.as_ref().into());
         self
     }
     /// Set the FlexSPI peripheral that interfaces flash.
@@ -384,8 +522,8 @@ impl RuntimeBuilder {
         region_alias(writer, "STACK", self.stack)?;
         region_alias(writer, "HEAP", self.heap)?;
         // Used in the linker script and / or target code.
-        writeln!(writer, "__stack_size = {:#010X};", self.stack_size)?;
-        writeln!(writer, "__heap_size = {:#010X};", self.heap_size)?;
+        writeln!(writer, "__stack_size = {:#010X};", self.stack_size.read()?)?;
+        writeln!(writer, "__heap_size = {:#010X};", self.heap_size.read()?)?;
 
         if self.flash_opts.is_some() {
             // Runtime will see different VMA and LMA, and copy the sections.
diff --git a/tests/inspect_elf.rs b/tests/inspect_elf.rs
index 34da382..8df4b6a 100644
--- a/tests/inspect_elf.rs
+++ b/tests/inspect_elf.rs
@@ -10,8 +10,8 @@ use std::{fs, path::PathBuf, process::Command};
 
 type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
 
-/// Build an example, returning a path to the ELF.
-fn cargo_build(board: &str) -> Result<PathBuf> {
+/// Build an example with optional environment variables, returning a path to the ELF.
+fn cargo_build_with_envs(board: &str, envs: &[(&str, &str)]) -> Result<PathBuf> {
     let status = Command::new("cargo")
         .arg("build")
         .arg("--example=blink-rtic")
@@ -19,6 +19,7 @@ fn cargo_build(board: &str) -> Result<PathBuf> {
         .arg("--target=thumbv7em-none-eabihf")
         .arg(format!("--target-dir=target/{}", board))
         .arg("--quiet")
+        .envs(envs.iter().copied())
         .spawn()?
         .wait()?;
 
@@ -38,6 +39,11 @@ fn cargo_build(board: &str) -> Result<PathBuf> {
     Ok(path)
 }
 
+/// Build an example, returning a path to the ELF.
+fn cargo_build(board: &str) -> Result<PathBuf> {
+    cargo_build_with_envs(board, &[])
+}
+
 struct ImxrtBinary<'a> {
     elf: &'a Elf<'a>,
     contents: &'a [u8],
@@ -294,7 +300,7 @@ fn imxrt1010evk() {
     assert_eq!(binary.section_lma(".heap"), heap.address, "Heap is NOLOAD");
 }
 
-fn baseline_teensy4(binary: &ImxrtBinary, dcd_at_runtime: u32) {
+fn baseline_teensy4(binary: &ImxrtBinary, dcd_at_runtime: u32, stack_size: u64, heap_size: u64) {
     assert_eq!(
         Fcb {
             address: 0x6000_0000,
@@ -320,10 +326,10 @@ fn baseline_teensy4(binary: &ImxrtBinary, dcd_at_runtime: u32) {
     assert_eq!(
         Section {
             address: DTCM,
-            size: 8 * 1024
+            size: stack_size
         },
         stack,
-        "stack not at ORIGIN(DTCM), or not 8 KiB large"
+        "stack not at ORIGIN(DTCM), or not {stack_size} bytes large"
     );
     assert_eq!(binary.section_lma(".stack"), stack.address);
 
@@ -402,10 +408,10 @@ fn baseline_teensy4(binary: &ImxrtBinary, dcd_at_runtime: u32) {
     assert_eq!(
         Section {
             address: uninit.address + aligned(uninit.size, 4),
-            size: 1024
+            size: heap_size
         },
         heap,
-        "1 KiB heap in DTCM behind uninit"
+        "{heap_size} byte heap in DTCM behind uninit"
     );
     assert_eq!(binary.section_lma(".heap"), heap.address, "Heap is NOLOAD");
 }
@@ -424,7 +430,7 @@ fn teensy4() {
         binary.symbol_value("__dcd_end")
     );
     assert_eq!(binary.symbol_value("__dcd"), Some(0));
-    baseline_teensy4(&binary, 0);
+    baseline_teensy4(&binary, 0, 8 * 1024, 1024);
 }
 
 #[test]
@@ -446,7 +452,7 @@ fn teensy4_fake_dcd() {
         binary.symbol_value("__dcd_start"),
     );
     assert_eq!(dcd.st_size % 4, 0);
-    baseline_teensy4(&binary, dcd_start as u32);
+    baseline_teensy4(&binary, dcd_start as u32, 8 * 1024, 1024);
 }
 
 #[test]
@@ -459,6 +465,45 @@ fn teensy4_fake_dcd_missize_fail() {
 
 #[test]
 #[ignore = "building an example can take time"]
+fn teensy4_env_overrides() {
+    let path = cargo_build_with_envs(
+        "teensy4",
+        &[
+            ("BOARD_STACK", "4096"),
+            ("THIS_WONT_BE_CONSIDERED", "12288"),
+            ("BOARD_HEAP", "8192"),
+        ],
+    )
+    .expect("Unable to build example");
+    let contents = fs::read(path).expect("Could not read ELF file");
+    let elf = Elf::parse(&contents).expect("Could not parse ELF");
+
+    let binary = ImxrtBinary::new(&elf, &contents);
+    baseline_teensy4(&binary, 0, 4 * 1024, 8 * 1024);
+}
+
+#[test]
+#[ignore = "building an example can take time"]
+fn teensy4_env_overrides_kib() {
+    let path = cargo_build_with_envs("teensy4", &[("BOARD_STACK", "5K"), ("BOARD_HEAP", "9k")])
+        .expect("Unable to build example");
+    let contents = fs::read(path).expect("Could not read ELF file");
+    let elf = Elf::parse(&contents).expect("Could not parse ELF");
+
+    let binary = ImxrtBinary::new(&elf, &contents);
+    baseline_teensy4(&binary, 0, 5 * 1024, 9 * 1024);
+}
+
+#[test]
+#[should_panic]
+#[ignore = "building an example can take time"]
+fn teensy4_env_override_fail() {
+    cargo_build_with_envs("teensy4", &[("BOARD_STACK", "1o24")])
+        .expect("Build should fail since BOARD_STACK can't be parsed");
+}
+
+#[test]
+#[ignore = "building an example can take time"]
 fn imxrt1170evk_cm7() {
     let path = cargo_build("imxrt1170evk-cm7").expect("Unable to build example");
     let contents = fs::read(path).expect("Could not read ELF file");