[ruby/mmtk] Fix environment variable parsing

Ues more idiomatic rust approaches. https://github.com/ruby/mmtk/commit/ef125f9eae
author: Kunshan Wang <[email protected]> 2025-05-28 16:21:09 +0800
committer: git <[email protected]> 2025-05-30 14:55:42 +0000
commit: d2a1ad00cbba41e22c11abf2948c23cd8d68f565 (patch)
tree: c76b3ae2e8679b9fea1621282e28ab24a9e60e2e /gc/mmtk/src
parent: 94688bdc7da3ac82f712d3cdc11859ace3bfa8ac (diff)
2 files changed, 72 insertions, 93 deletions
diff --git a/gc/mmtk/src/api.rs b/gc/mmtk/src/api.rs
index fff7a4a37c..a1b94d520d 100644
--- a/gc/mmtk/src/api.rs
+++ b/gc/mmtk/src/api.rs
@@ -3,6 +3,7 @@
 #![allow(clippy::missing_safety_doc)]
 
 use mmtk::util::options::PlanSelector;
+use std::str::FromStr;
 use std::sync::atomic::Ordering;
 
 use crate::abi::RawVecOfObjRef;
@@ -41,66 +42,61 @@ pub extern "C" fn mmtk_is_reachable(object: ObjectReference) -> bool {
 
 // =============== Bootup ===============
 
-fn mmtk_builder_default_parse_threads() -> usize {
-    let threads_str = std::env::var("MMTK_THREADS").unwrap_or("0".to_string());
+fn parse_env_var_with<T, F: FnOnce(&str) -> Option<T>>(key: &str, parse: F) -> Option<T> {
+    let val = match std::env::var(key) {
+        Ok(val) => val,
+        Err(std::env::VarError::NotPresent) => return None,
+        Err(std::env::VarError::NotUnicode(os_string)) => {
+            eprintln!("[FATAL] Invalid {key} {os_string:?}");
+            std::process::exit(1);
+        }
+    };
 
-    threads_str.parse::<usize>().unwrap_or_else(|_err| {
-        eprintln!("[FATAL] Invalid MMTK_THREADS {}", threads_str);
+    let parsed = parse(&val).unwrap_or_else(|| {
+        eprintln!("[FATAL] Invalid {key} {val}");
         std::process::exit(1);
-    })
-}
+    });
 
-fn mmtk_builder_default_parse_heap_min() -> usize {
-    const DEFAULT_HEAP_MIN: usize = 1 << 20;
+    Some(parsed)
+}
 
-    let heap_min_str = std::env::var("MMTK_HEAP_MIN").unwrap_or(DEFAULT_HEAP_MIN.to_string());
+fn parse_env_var<T: FromStr>(key: &str) -> Option<T> {
+    parse_env_var_with(key, |s| s.parse().ok())
+}
 
-    let size = parse_capacity(&heap_min_str, 0);
-    if size == 0 {
-        eprintln!("[FATAL] Invalid MMTK_HEAP_MIN {}", heap_min_str);
-        std::process::exit(1);
-    }
+fn mmtk_builder_default_parse_threads() -> Option<usize> {
+    parse_env_var("MMTK_THREADS")
+}
 
-    size
+fn mmtk_builder_default_parse_heap_min() -> usize {
+    const DEFAULT_HEAP_MIN: usize = 1 << 20;
+    parse_env_var_with("MMTK_HEAP_MIN", parse_capacity).unwrap_or(DEFAULT_HEAP_MIN)
 }
 
 fn mmtk_builder_default_parse_heap_max() -> usize {
-    let heap_max_str = std::env::var("MMTK_HEAP_MAX").unwrap_or(default_heap_max().to_string());
-
-    let size = parse_capacity(&heap_max_str, 0);
-    if size == 0 {
-        eprintln!("[FATAL] Invalid MMTK_HEAP_MAX {}", heap_max_str);
-        std::process::exit(1);
-    }
-
-    size
+    parse_env_var_with("MMTK_HEAP_MAX", parse_capacity).unwrap_or_else(default_heap_max)
 }
 
 fn mmtk_builder_default_parse_heap_mode(heap_min: usize, heap_max: usize) -> GCTriggerSelector {
-    let heap_mode_str = std::env::var("MMTK_HEAP_MODE").unwrap_or("dynamic".to_string());
+    let make_fixed = || GCTriggerSelector::FixedHeapSize(heap_max);
+    let make_dynamic = || GCTriggerSelector::DynamicHeapSize(heap_min, heap_max);
 
-    match heap_mode_str.as_str() {
-        "fixed" => GCTriggerSelector::FixedHeapSize(heap_max),
-        "dynamic" => GCTriggerSelector::DynamicHeapSize(heap_min, heap_max),
-        _ => {
-            eprintln!("[FATAL] Invalid MMTK_HEAP_MODE {}", heap_mode_str);
-            std::process::exit(1);
-        }
-    }
+    parse_env_var_with("MMTK_HEAP_MODE", |s| match s {
+        "fixed" => Some(make_fixed()),
+        "dynamic" => Some(make_dynamic()),
+        _ => None,
+    })
+    .unwrap_or_else(make_dynamic)
 }
 
 fn mmtk_builder_default_parse_plan() -> PlanSelector {
-    let plan_str = std::env::var("MMTK_PLAN").unwrap_or("Immix".to_string());
-
-    match plan_str.as_str() {
-        "NoGC" => PlanSelector::NoGC,
-        "MarkSweep" => PlanSelector::MarkSweep,
-        "Immix" => PlanSelector::Immix,
-        _ => {
-            eprintln!("[FATAL] Invalid MMTK_PLAN {}", plan_str);
-            std::process::exit(1);
-        }
-    }
+    parse_env_var_with("MMTK_PLAN", |s| match s {
+        "NoGC" => Some(PlanSelector::NoGC),
+        "MarkSweep" => Some(PlanSelector::MarkSweep),
+        "Immix" => Some(PlanSelector::Immix),
+        _ => None,
+    })
+    .unwrap_or(PlanSelector::Immix)
 }
 
 #[no_mangle]
@@ -108,9 +104,15 @@ pub extern "C" fn mmtk_builder_default() -> *mut MMTKBuilder {
     let mut builder = MMTKBuilder::new_no_env_vars();
     builder.options.no_finalizer.set(true);
 
-    let threads = mmtk_builder_default_parse_threads();
-    if threads > 0 {
-        builder.options.threads.set(threads);
+    if let Some(threads) = mmtk_builder_default_parse_threads() {
+        if !builder.options.threads.set(threads) {
+            // MMTk will validate it and reject 0.
+            eprintln!(
+                "[FATAL] Failed to set the number of MMTk threads to {}",
+                threads
+            );
+            std::process::exit(1);
+        }
     }
 
     let heap_min = mmtk_builder_default_parse_heap_min();
diff --git a/gc/mmtk/src/utils.rs b/gc/mmtk/src/utils.rs
index 3f149c66dd..71a7ae8dd2 100644
--- a/gc/mmtk/src/utils.rs
+++ b/gc/mmtk/src/utils.rs
@@ -97,35 +97,29 @@ pub fn default_heap_max() -> usize {
         .expect("Invalid Memory size") as usize
 }
 
-pub fn parse_capacity(input: &str, default: usize) -> usize {
+pub fn parse_capacity(input: &str) -> Option<usize> {
     let trimmed = input.trim();
 
     const KIBIBYTE: usize = 1024;
     const MEBIBYTE: usize = 1024 * KIBIBYTE;
     const GIBIBYTE: usize = 1024 * MEBIBYTE;
 
-    let (val, suffix) = if let Some(pos) = trimmed.find(|c: char| !c.is_numeric()) {
-        (&trimmed[..pos], &trimmed[pos..])
+    let (number, suffix) = if let Some(pos) = trimmed.find(|c: char| !c.is_numeric()) {
+        trimmed.split_at(pos)
     } else {
         (trimmed, "")
     };
 
-    // 1MiB is the default heap size
-    match (val, suffix) {
-        (number, "GiB") => number
-            .parse::<usize>()
-            .map(|v| v * GIBIBYTE)
-            .unwrap_or(default),
-        (number, "MiB") => number
-            .parse::<usize>()
-            .map(|v| v * MEBIBYTE)
-            .unwrap_or(default),
-        (number, "KiB") => number
-            .parse::<usize>()
-            .map(|v| v * KIBIBYTE)
-            .unwrap_or(default),
-        (number, "") => number.parse::<usize>().unwrap_or(default),
-        (_, _) => default,
+    let Ok(v) = number.parse::<usize>() else {
+        return None;
+    };
+
+    match suffix {
+        "GiB" => Some(v * GIBIBYTE),
+        "MiB" => Some(v * MEBIBYTE),
+        "KiB" => Some(v * KIBIBYTE),
+        "" => Some(v),
+        _ => None,
     }
 }
 
@@ -135,47 +129,30 @@ mod tests {
 
     #[test]
     fn test_parse_capacity_parses_bare_bytes() {
-        assert_eq!(1234, parse_capacity(&String::from("1234"), 0));
+        assert_eq!(Some(1234), parse_capacity("1234"));
     }
 
     #[test]
     fn test_parse_capacity_parses_kibibytes() {
-        assert_eq!(10240, parse_capacity(&String::from("10KiB"), 0))
+        assert_eq!(Some(10240), parse_capacity("10KiB"));
     }
 
     #[test]
     fn test_parse_capacity_parses_mebibytes() {
-        assert_eq!(10485760, parse_capacity(&String::from("10MiB"), 0))
+        assert_eq!(Some(10485760), parse_capacity("10MiB"))
     }
 
     #[test]
     fn test_parse_capacity_parses_gibibytes() {
-        assert_eq!(10737418240, parse_capacity(&String::from("10GiB"), 0))
+        assert_eq!(Some(10737418240), parse_capacity("10GiB"))
     }
 
     #[test]
-    fn test_parses_nonsense_value_as_default_max() {
-        let default = 100;
-
-        assert_eq!(
-            default,
-            parse_capacity(&String::from("notanumber"), default)
-        );
-        assert_eq!(
-            default,
-            parse_capacity(&String::from("5tartswithanumber"), default)
-        );
-        assert_eq!(
-            default,
-            parse_capacity(&String::from("number1nthemiddle"), default)
-        );
-        assert_eq!(
-            default,
-            parse_capacity(&String::from("numberattheend111"), default)
-        );
-        assert_eq!(
-            default,
-            parse_capacity(&String::from("mult1pl3numb3r5"), default)
-        );
+    fn test_parse_capacity_parses_nonsense_values() {
+        assert_eq!(None, parse_capacity("notanumber"));
+        assert_eq!(None, parse_capacity("5tartswithanumber"));
+        assert_eq!(None, parse_capacity("number1nthemiddle"));
+        assert_eq!(None, parse_capacity("numberattheend111"));
+        assert_eq!(None, parse_capacity("mult1pl3numb3r5"));
     }
 }
author	Kunshan Wang <[email protected]>	2025-05-28 16:21:09 +0800
committer	git <[email protected]>	2025-05-30 14:55:42 +0000
commit	d2a1ad00cbba41e22c11abf2948c23cd8d68f565 (patch)
tree	c76b3ae2e8679b9fea1621282e28ab24a9e60e2e /gc/mmtk/src
parent	94688bdc7da3ac82f712d3cdc11859ace3bfa8ac (diff)