feat(zkvms_host_io): Implement convinence environment variable methods

1 files changed, 18 insertions, 1 deletions

diff --git a/zkvms_host_io/src/lib.rs b/zkvms_host_io/src/lib.rs
index a3bdb1b..c029975 100644
--- a/zkvms_host_io/src/lib.rs
+++ b/zkvms_host_io/src/lib.rs
@@ -2,7 +2,7 @@ use clap::{Parser, ValueEnum};
 use num_traits::NumCast;
 use serde::{ Serialize, Deserialize };
 use env_file_reader::read_str;
-use std::{ fs::read_to_string, collections::HashMap };
+use std::{env, option::Option, fs::read_to_string, collections::HashMap};
 pub use input_macros::foreach_input_field;
 
 static DEFAULT_PUBLIC_INPUT: &str = include_str!(concat!(env!("INPUTS_DIR"), "/default_public_input.toml"));
@@ -31,6 +31,23 @@ pub struct RunWith<T> {
     pub default_env: HashMap<String, String>,
 }
 
+impl<T> RunWith<T> {
+    pub fn env_then_or<U>(&self, variable_name: &str, then_apply: fn(String) -> Option<U>, else_const: U) -> U {
+        env::var(variable_name)
+            .ok()
+            .and_then(then_apply)
+            .unwrap_or(self
+                .default_env
+                .get(variable_name)
+                .and_then(|x| then_apply(x.clone()))
+                .unwrap_or(else_const))
+    }
+
+    pub fn env_or(&self, variable_name: &str, else_const: &str) -> String {
+        self.env_then_or(variable_name, |x| Some(x), else_const.to_string())
+    }
+}
+
 input_macros::generate_input_struct!();
 
 pub fn read_args() -> RunWith<Input> {