shell: keep variables in one table with an export attribute, and prefix assignments on their command

docs/runtime/shell.mdx

@@ -327,6 +327,26 @@ const foo = "bar123; rm -rf /tmp";
 await $`FOO=${foo} bun -e 'console.log(process.env.FOO)'`; // bar123; rm -rf /tmp\n
 ```
 
+A variable from the environment is exported, and so is each name you pass to `export`. A child process receives the current value of every exported variable. A plain assignment to any other name stays in the shell:
+
+```js
+import { $ } from "bun";
+
+// FOO gets the export attribute and keeps its value
+await $`FOO=foo; export FOO; bun -e 'console.log(process.env.FOO)'`; // foo\n
+
+// NODE_ENV comes from the environment, so the new value reaches `bun`
+await $`NODE_ENV=production; bun -e 'console.log(process.env.NODE_ENV)'`.env({
+  ...process.env,
+  NODE_ENV: "development",
+}); // production\n
+
+// BAR is not exported
+await $`BAR=bar; bun -e 'console.log(process.env.BAR)'`; // undefined\n
+```
+
+On Windows, variable names are ASCII-case-insensitive, as in `process.env`: `PATH`, `Path` and `path` are one variable, so an assignment under any spelling changes what child processes receive.
+
 ### Changing the environment variables
 
 By default, all commands use `process.env` as their environment variables.
@@ -466,6 +486,7 @@ For cross-platform compatibility, Bun Shell implements a set of builtin commands
 - `rm`: remove files and directories
 - `echo`: print text
 - `pwd`: print the working directory
+- `export`: export variables to child processes
 - `bun`: run bun in bun
 - `cat`
 - `touch`

src/runtime/shell/EnvMap.rs

@@ -1,13 +1,11 @@
-use bun_collections::array_hash_map::{self, ArrayHashContext, ArrayHashMap, Iter};
+use bun_collections::array_hash_map::{self, ArrayHashContext, ArrayHashMap};
 
 use crate::shell::EnvStr;
 
 pub(crate) struct EnvMap {
     map: EnvMapInner,
 }
 
-pub(crate) type Iterator<'a> = Iter<'a, EnvStr, EnvStr>;
-
 // Named `EnvMapInner` to avoid rustc confusing it with the unrelated mmap
 // `sys::c::MapType` / `sys::posix::MapType` in diagnostic suggestions.
 type EnvMapInner = ArrayHashMap<EnvStr, EnvStr, EnvMapContext>;
@@ -31,43 +29,27 @@ impl ArrayHashContext<EnvStr> for EnvMapContext {
     }
 
     fn eql(&self, a: &EnvStr, b: &EnvStr, _b_index: usize) -> bool {
-        #[cfg(windows)]
-        {
-            // Must be length-checked: "PATH" must NOT match "PATHEXT".
-            return bun_core::strings::eql_case_insensitive_asciii_check_length(
-                a.slice(),
-                b.slice(),
-            );
-        }
-        #[cfg(not(windows))]
-        {
-            a.slice() == b.slice()
-        }
+        name_eql(a.slice(), b.slice())
     }
 }
 
-impl EnvMap {
-    pub(crate) fn init() -> EnvMap {
-        EnvMap {
-            map: EnvMapInner::new(),
-        }
+/// Whether `a` and `b` name the same variable. Names are ASCII-case-insensitive on Windows.
+pub(crate) fn name_eql(a: &[u8], b: &[u8]) -> bool {
+    #[cfg(windows)]
+    {
+        // Must be length-checked: "PATH" must NOT match "PATHEXT".
+        return bun_core::strings::eql_case_insensitive_asciii_check_length(a, b);
     }
-
-    pub(crate) fn memory_cost(&self) -> usize {
-        let mut size: usize = core::mem::size_of::<EnvMap>();
-        size += core::mem::size_of_val(self.map.keys());
-        size += core::mem::size_of_val(self.map.values());
-        debug_assert_eq!(self.map.keys().len(), self.map.values().len());
-        for (key, value) in self.map.keys().iter().zip(self.map.values()) {
-            size += key.memory_cost();
-            size += value.memory_cost();
-        }
-        size
+    #[cfg(not(windows))]
+    {
+        a == b
     }
+}
 
-    pub(crate) fn init_with_capacity(cap: usize) -> EnvMap {
+impl EnvMap {
+    pub(crate) fn init() -> EnvMap {
         EnvMap {
-            map: EnvMapInner::with_capacity(cap),
+            map: EnvMapInner::new(),
         }
     }
 
@@ -84,51 +66,191 @@ impl EnvMap {
         *result.value_ptr = val;
     }
 
-    pub(crate) fn iterator(&mut self) -> Iterator<'_> {
-        self.map.iterator()
+    /// NOTE: Make sure you deref the string when done!
+    pub(crate) fn get(&self, key: EnvStr) -> Option<EnvStr> {
+        let val = *self.map.get(&key)?;
+        val.ref_();
+        Some(val)
+    }
+
+    fn deref_strings(&self) {
+        for (key, value) in self.map.keys().iter().zip(self.map.values()) {
+            key.deref();
+            value.deref();
+        }
+    }
+}
+
+impl Drop for EnvMap {
+    fn drop(&mut self) {
+        self.deref_strings();
+        // map storage freed by its own Drop
+    }
+}
+
+/// One shell variable.
+#[derive(Copy, Clone, Default)]
+enum Var {
+    /// Exported, with a value: a child process inherits it.
+    Exported(EnvStr),
+    /// Has a value, not exported.
+    Local(EnvStr),
+    /// Exported, no value yet.
+    #[default]
+    ExportedUnset,
+}
+
+impl Var {
+    fn value(self) -> Option<EnvStr> {
+        match self {
+            Var::Exported(value) | Var::Local(value) => Some(value),
+            Var::ExportedUnset => None,
+        }
+    }
+}
+
+type VarMap = ArrayHashMap<EnvStr, Var, EnvMapContext>;
+
+/// Every shell variable, one entry per name.
+pub(crate) struct VarTable {
+    map: VarMap,
+}
+
+impl VarTable {
+    pub(crate) fn init() -> VarTable {
+        VarTable { map: VarMap::new() }
     }
 
-    pub(crate) fn iter(&self) -> impl core::iter::Iterator<Item = (&EnvStr, &EnvStr)> {
-        self.map.keys().iter().zip(self.map.values())
+    pub(crate) fn init_with_capacity(cap: usize) -> VarTable {
+        VarTable {
+            map: VarMap::with_capacity(cap),
+        }
     }
 
     pub(crate) fn ensure_total_capacity(&mut self, new_capacity: usize) {
         self.map.ensure_total_capacity(new_capacity).expect("OOM");
     }
 
-    /// NOTE: Make sure you deref the string when done!
-    pub(crate) fn get(&self, key: EnvStr) -> Option<EnvStr> {
-        let val = *self.map.get(&key)?;
-        val.ref_();
-        Some(val)
+    pub(crate) fn memory_cost(&self) -> usize {
+        let mut size: usize = core::mem::size_of::<VarTable>();
+        size += core::mem::size_of_val(self.map.keys());
+        size += core::mem::size_of_val(self.map.values());
+        for (name, var) in self.map.iter() {
+            size += name.memory_cost();
+            size += var.value().map_or(0, EnvStr::memory_cost);
+        }
+        size
     }
 
-    pub(crate) fn clone(&self) -> EnvMap {
-        let new = EnvMap {
+    pub(crate) fn clone(&self) -> VarTable {
+        let new = VarTable {
             map: self.map.clone().expect("OOM"),
         };
-        new.ref_strings();
+        for (name, var) in new.map.iter() {
+            name.ref_();
+            if let Some(value) = var.value() {
+                value.ref_();
+            }
+        }
         new
     }
 
-    fn ref_strings(&self) {
-        for (key, value) in self.map.keys().iter().zip(self.map.values()) {
-            key.ref_();
-            value.ref_();
+    /// `NAME=value`: the name keeps its export attribute. Refs `value`, like `EnvMap::insert`.
+    pub(crate) fn set(&mut self, name: EnvStr, value: EnvStr) {
+        let slot = self.map.get_or_put(name).expect("OOM");
+        value.ref_();
+        if !slot.found_existing {
+            name.ref_();
+            *slot.value_ptr = Var::Local(value);
+            return;
         }
+        *slot.value_ptr = match *slot.value_ptr {
+            Var::Local(old) => {
+                old.deref();
+                Var::Local(value)
+            }
+            Var::Exported(old) => {
+                old.deref();
+                Var::Exported(value)
+            }
+            Var::ExportedUnset => Var::Exported(value),
+        };
     }
 
-    fn deref_strings(&self) {
-        for (key, value) in self.map.keys().iter().zip(self.map.values()) {
-            key.deref();
-            value.deref();
+    /// `export NAME=value`, or `export NAME` for `None`, which keeps the value. Refs `value`.
+    pub(crate) fn export(&mut self, name: EnvStr, value: Option<EnvStr>) {
+        let slot = self
+            .map
+            .get_or_put_value(name, Var::ExportedUnset)
+            .expect("OOM");
+        if !slot.found_existing {
+            name.ref_();
+        }
+        match (value, *slot.value_ptr) {
+            (Some(value), old) => {
+                value.ref_();
+                if let Some(old) = old.value() {
+                    old.deref();
+                }
+                *slot.value_ptr = Var::Exported(value);
+            }
+            (None, Var::Local(value)) => *slot.value_ptr = Var::Exported(value),
+            (None, Var::Exported(_) | Var::ExportedUnset) => {}
         }
     }
+
+    /// The value of `$NAME`. Deref it when done, like `EnvMap::get`.
+    pub(crate) fn get(&self, name: EnvStr) -> Option<EnvStr> {
+        let value = self.map.get(&name)?.value()?;
+        value.ref_();
+        Some(value)
+    }
+
+    /// The value a child process inherits for `name`. Deref it when done.
+    pub(crate) fn get_exported(&self, name: EnvStr) -> Option<EnvStr> {
+        let Var::Exported(value) = *self.map.get(&name)? else {
+            return None;
+        };
+        value.ref_();
+        Some(value)
+    }
+
+    /// The environment of a child process, one entry per name: each exported name that has a value, then `prefixes`, which replace an exported name.
+    pub(crate) fn child_env<'a>(
+        &'a self,
+        prefixes: Option<&'a EnvMap>,
+    ) -> impl core::iter::Iterator<Item = (&'a EnvStr, &'a EnvStr)> {
+        let (names, values): (&[EnvStr], &[EnvStr]) =
+            prefixes.map_or((&[], &[]), |p| (p.map.keys(), p.map.values()));
+        let replaced = move |name: &EnvStr| names.iter().any(|p| name_eql(p.slice(), name.slice()));
+        self.map
+            .iter()
+            .filter_map(move |(name, var)| match var {
+                Var::Exported(value) if !replaced(name) => Some((name, value)),
+                Var::Exported(_) | Var::Local(_) | Var::ExportedUnset => None,
+            })
+            .chain(names.iter().zip(values))
+    }
+
+    /// Every exported name. The value is `None` until the name is assigned.
+    pub(crate) fn exported_names(
+        &self,
+    ) -> impl core::iter::Iterator<Item = (EnvStr, Option<EnvStr>)> + '_ {
+        self.map.iter().filter_map(|(name, var)| match *var {
+            Var::Exported(value) => Some((*name, Some(value))),
+            Var::ExportedUnset => Some((*name, None)),
+            Var::Local(_) => None,
+        })
+    }
 }
 
-impl Drop for EnvMap {
+impl Drop for VarTable {
     fn drop(&mut self) {
-        self.deref_strings();
-        // map storage freed by its own Drop
+        for (name, var) in self.map.iter() {
+            name.deref();
+            if let Some(value) = var.value() {
+                value.deref();
+            }
+        }
     }
 }

src/runtime/shell/ParsedShellScript.rs

@@ -8,7 +8,7 @@ use bun_jsc::{
     JsResult, MarkedArgumentBuffer,
 };
 
-use super::env_map::EnvMap;
+use super::env_map::VarTable;
 use super::interpreter::ShellArgs;
 use super::shell_body::shell_cmd_from_js;
 use super::{EnvStr, Interpreter};
@@ -17,7 +17,7 @@ pub(crate) use crate::generated_bindgen::parsed_shell_script::js_ParsedShellScri
 
 pub(crate) struct ParsedShellScript {
     pub args: JsCell<Option<Box<ShellArgs>>>,
-    pub(crate) export_env: JsCell<Option<EnvMap>>,
+    pub(crate) export_env: JsCell<Option<VarTable>>,
     pub(crate) quiet: Cell<bool>,
     pub(crate) cwd: JsCell<Option<BunString>>,
     /// Read-only after construction (set once before the JS wrapper exists).
@@ -56,7 +56,7 @@ impl ParsedShellScript {
     }
 
     // Returns a tuple; callers destructure it.
-    pub(crate) fn take(&self) -> (Box<ShellArgs>, bool, Option<BunString>, Option<EnvMap>) {
+    pub(crate) fn take(&self) -> (Box<ShellArgs>, bool, Option<BunString>, Option<VarTable>) {
         let args = self.args.replace(None).expect("args already taken");
         let quiet = self.quiet.get();
         let cwd = self.cwd.replace(None);
@@ -91,7 +91,7 @@ impl ParsedShellScript {
         )?;
         // `defer object_iter.deinit()` — handled by Drop.
 
-        let mut env = EnvMap::init();
+        let mut env = VarTable::init();
         // errdefer env.deinit() — Drop on early-return handles this.
         env.ensure_total_capacity(object_iter.len);
 
@@ -109,11 +109,11 @@ impl ParsedShellScript {
             let keyref = EnvStr::init_ref_counted(keyslice);
             let valueref = EnvStr::init_ref_counted(slice);
 
-            env.insert(keyref, valueref);
+            env.export(keyref, Some(valueref));
             keyref.deref();
             valueref.deref();
         }
-        // Dropping the previous Option<EnvMap> deinits it.
+        // Dropping the previous Option<VarTable> deinits it.
         self.export_env.set(Some(env));
         Ok(())
     }