use std::collections::HashMap;
#[derive(Debug, Clone, PartialEq)]
pub struct Variable {
pub value: String,
pub exported: bool,
pub readonly: bool,
}
impl Variable {
pub fn new(value: impl Into<String>) -> Self {
Variable {
value: value.into(),
exported: false,
readonly: false,
}
}
pub fn new_exported(value: impl Into<String>) -> Self {
Variable {
value: value.into(),
exported: true,
readonly: false,
}
}
}
#[derive(Debug, Clone)]
struct Scope {
vars: HashMap<String, Variable>,
positional_params: Vec<String>,
}
#[derive(Debug, Clone)]
pub struct VarStore {
scopes: Vec<Scope>,
environ_cache: Option<Vec<(String, String)>>,
}
impl VarStore {
pub fn new() -> Self {
VarStore {
scopes: vec![Scope {
vars: HashMap::new(),
positional_params: Vec::new(),
}],
environ_cache: None,
}
}
pub fn from_environ() -> Self {
let mut vars = HashMap::new();
for (key, value) in std::env::vars() {
vars.insert(key, Variable::new_exported(value));
}
VarStore {
scopes: vec![Scope {
vars,
positional_params: Vec::new(),
}],
environ_cache: None,
}
}
pub fn push_scope(&mut self, positional_params: Vec<String>) {
self.environ_cache = None;
self.scopes.push(Scope {
vars: HashMap::new(),
positional_params,
});
}
pub fn pop_scope(&mut self) {
self.environ_cache = None;
assert!(self.scopes.len() > 1, "cannot pop the global scope");
self.scopes.pop();
}
pub fn scope_depth(&self) -> usize {
self.scopes.len()
}
pub fn positional_params(&self) -> &[String] {
&self.scopes.last().unwrap().positional_params
}
pub fn set_positional_params(&mut self, params: Vec<String>) {
self.scopes.last_mut().unwrap().positional_params = params;
}
pub fn get(&self, name: &str) -> Option<&str> {
if self.scopes.len() == 1 {
return self.scopes[0].vars.get(name).map(|v| v.value.as_str());
}
for scope in self.scopes.iter().rev() {
if let Some(var) = scope.vars.get(name) {
return Some(var.value.as_str());
}
}
None
}
#[allow(dead_code)]
pub fn get_var(&self, name: &str) -> Option<&Variable> {
for scope in self.scopes.iter().rev() {
if let Some(var) = scope.vars.get(name) {
return Some(var);
}
}
None
}
pub fn set(&mut self, name: &str, value: impl Into<String>) -> Result<(), String> {
self.environ_cache = None;
let value = value.into();
if self.scopes.len() == 1 {
if let Some(existing) = self.scopes[0].vars.get(name) {
if existing.readonly {
return Err(format!("{}: readonly variable", name));
}
let exported = existing.exported;
self.scopes[0].vars.insert(
name.to_string(),
Variable {
value,
exported,
readonly: false,
},
);
} else {
self.scopes[0].vars.insert(name.to_string(), Variable::new(value));
}
return Ok(());
}
for scope in self.scopes.iter_mut().rev() {
if let Some(existing) = scope.vars.get(name) {
if existing.readonly {
return Err(format!("{}: readonly variable", name));
}
let exported = existing.exported;
scope.vars.insert(
name.to_string(),
Variable {
value,
exported,
readonly: false,
},
);
return Ok(());
}
}
self.scopes[0].vars.insert(name.to_string(), Variable::new(value));
Ok(())
}
pub fn set_with_options(
&mut self,
name: &str,
value: impl Into<String>,
allexport: bool,
) -> Result<(), String> {
self.environ_cache = None;
let value = value.into();
for scope in self.scopes.iter_mut().rev() {
if let Some(existing) = scope.vars.get(name) {
if existing.readonly {
return Err(format!("{}: readonly variable", name));
}
let exported = existing.exported || allexport;
scope.vars.insert(
name.to_string(),
Variable {
value,
exported,
readonly: false,
},
);
return Ok(());
}
}
let mut var = Variable::new(value);
if allexport {
var.exported = true;
}
self.scopes[0].vars.insert(name.to_string(), var);
Ok(())
}
pub fn unset(&mut self, name: &str) -> Result<(), String> {
self.environ_cache = None;
for scope in self.scopes.iter_mut().rev() {
if let Some(existing) = scope.vars.get(name) {
if existing.readonly {
return Err(format!("{}: readonly variable", name));
}
scope.vars.remove(name);
return Ok(());
}
}
Ok(())
}
pub fn export(&mut self, name: &str) {
self.environ_cache = None;
for scope in self.scopes.iter_mut().rev() {
if let Some(var) = scope.vars.get_mut(name) {
var.exported = true;
return;
}
}
self.scopes[0]
.vars
.insert(name.to_string(), Variable::new_exported(""));
}
pub fn set_readonly(&mut self, name: &str) {
self.environ_cache = None;
for scope in self.scopes.iter_mut().rev() {
if let Some(var) = scope.vars.get_mut(name) {
var.readonly = true;
return;
}
}
let mut var = Variable::new("");
var.readonly = true;
self.scopes[0].vars.insert(name.to_string(), var);
}
pub fn environ(&mut self) -> &[(String, String)] {
if self.environ_cache.is_none() {
self.environ_cache = Some(self.build_environ());
}
self.environ_cache.as_ref().unwrap()
}
fn build_environ(&self) -> Vec<(String, String)> {
let mut merged: HashMap<String, &Variable> = HashMap::new();
for scope in &self.scopes {
for (name, var) in &scope.vars {
merged.insert(name.clone(), var);
}
}
merged
.into_iter()
.filter(|(_, v)| v.exported)
.map(|(k, v)| (k, v.value.clone()))
.collect()
}
pub fn vars_iter(&self) -> impl Iterator<Item = (&str, &Variable)> {
let mut seen = std::collections::HashSet::new();
self.scopes
.iter()
.rev()
.flat_map(|s| s.vars.iter())
.filter_map(move |(k, v)| {
if seen.insert(k.as_str()) {
Some((k.as_str(), v))
} else {
None
}
})
}
}
impl Default for VarStore {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_get_set() {
let mut store = VarStore::new();
assert_eq!(store.get("FOO"), None);
store.set("FOO", "bar").unwrap();
assert_eq!(store.get("FOO"), Some("bar"));
}
#[test]
fn test_unset() {
let mut store = VarStore::new();
store.set("FOO", "bar").unwrap();
assert_eq!(store.get("FOO"), Some("bar"));
store.unset("FOO").unwrap();
assert_eq!(store.get("FOO"), None);
}
#[test]
fn test_readonly_prevents_set() {
let mut store = VarStore::new();
store.set("FOO", "bar").unwrap();
store.set_readonly("FOO");
let result = store.set("FOO", "baz");
assert!(result.is_err());
assert_eq!(store.get("FOO"), Some("bar"));
}
#[test]
fn test_readonly_prevents_unset() {
let mut store = VarStore::new();
store.set("FOO", "bar").unwrap();
store.set_readonly("FOO");
let result = store.unset("FOO");
assert!(result.is_err());
assert_eq!(store.get("FOO"), Some("bar"));
}
#[test]
fn test_export() {
let mut store = VarStore::new();
store.set("FOO", "bar").unwrap();
assert!(!store.get_var("FOO").unwrap().exported);
store.export("FOO");
assert!(store.get_var("FOO").unwrap().exported);
}
#[test]
fn test_environ_excludes_unexported() {
let mut store = VarStore::new();
store.set("FOO", "bar").unwrap();
store.set("BAZ", "qux").unwrap();
store.export("FOO");
let env = store.environ();
assert_eq!(env.len(), 1);
assert_eq!(env[0], ("FOO".to_string(), "bar".to_string()));
}
#[test]
fn test_from_environ() {
let store = VarStore::from_environ();
if let Some(var) = store.get_var("PATH") {
assert!(var.exported, "Variables from environ should be exported");
}
}
#[test]
fn test_push_pop_scope_positional_params() {
let mut store = VarStore::new();
store.set_positional_params(vec!["a".to_string(), "b".to_string()]);
assert_eq!(store.positional_params(), &["a", "b"]);
store.push_scope(vec!["x".to_string(), "y".to_string(), "z".to_string()]);
assert_eq!(store.positional_params(), &["x", "y", "z"]);
store.pop_scope();
assert_eq!(store.positional_params(), &["a", "b"]);
}
#[test]
fn test_scope_variable_lookup_walks_chain() {
let mut store = VarStore::new();
store.set("FOO", "global").unwrap();
store.push_scope(vec![]);
assert_eq!(store.get("FOO"), Some("global"));
store.set("FOO", "updated").unwrap();
store.pop_scope();
assert_eq!(store.get("FOO"), Some("updated"));
}
#[test]
fn test_scope_new_variable_goes_to_global() {
let mut store = VarStore::new();
store.push_scope(vec![]);
store.set("NEW_VAR", "value").unwrap();
store.pop_scope();
assert_eq!(store.get("NEW_VAR"), Some("value"));
}
#[test]
fn test_scope_readonly_across_scopes() {
let mut store = VarStore::new();
store.set("RO", "immutable").unwrap();
store.set_readonly("RO");
store.push_scope(vec![]);
let result = store.set("RO", "changed");
assert!(result.is_err());
assert_eq!(store.get("RO"), Some("immutable"));
store.pop_scope();
}
#[test]
fn test_scope_export_across_scopes() {
let mut store = VarStore::new();
store.set("EX", "value").unwrap();
store.push_scope(vec![]);
store.export("EX");
store.pop_scope();
assert!(store.get_var("EX").unwrap().exported);
}
#[test]
fn test_scope_unset_across_scopes() {
let mut store = VarStore::new();
store.set("DEL", "value").unwrap();
store.push_scope(vec![]);
store.unset("DEL").unwrap();
store.pop_scope();
assert_eq!(store.get("DEL"), None);
}
}