1use std::{collections::HashMap, rc::Rc};
4
5use crate::{
6 function::{Arity, Function},
7 stdlib::{NativeError, NativeResult},
8 value::Value,
9};
10
11pub enum FunctionResult {
13 Exists { pure: bool },
15 NotFound,
17 WrongArity { min: usize, max: usize },
19}
20
21pub trait Environment {
24 fn variable(&self, name: &str) -> Option<Rc<Value>>;
26
27 fn call(&self, name: &str, params: &[Value]) -> NativeResult;
33
34 fn variable_exists(&self, name: &str) -> bool;
36
37 fn function_exists(&self, name: &str, arity: usize) -> FunctionResult;
39}
40
41#[allow(clippy::module_name_repetitions)]
44#[derive(Default)]
45pub struct StaticEnvironment {
46 variables: HashMap<String, Rc<Value>>,
47 functions: HashMap<String, Rc<Function>>,
48}
49
50fn get_env_key(name: &str) -> String {
52 name.to_lowercase()
53}
54
55impl StaticEnvironment {
56 pub fn add_variable(&mut self, name: &str, value: Value) {
58 self.variables.insert(get_env_key(name), Rc::new(value));
59 }
60
61 pub fn remove_variable(&mut self, name: &str) -> Option<Rc<Value>> {
63 self.variables.remove(&get_env_key(name))
64 }
65
66 pub fn clear_variables(&mut self) {
68 self.variables.clear();
69 }
70
71 pub fn add_function(&mut self, func: Function) {
73 self.functions
74 .insert(get_env_key(&func.name), Rc::new(func));
75 }
76
77 pub fn add_functions(&mut self, functions: Vec<Function>) {
79 for func in functions {
80 self.add_function(func);
81 }
82 }
83
84 pub fn remove_function(&mut self, name: &str) -> Option<Rc<Function>> {
87 self.functions.remove(&get_env_key(name))
88 }
89
90 #[must_use]
92 pub fn list_functions(&self) -> Vec<Rc<Function>> {
93 self.functions.values().cloned().collect()
94 }
95}
96
97impl Environment for StaticEnvironment {
98 fn variable(&self, name: &str) -> Option<Rc<Value>> {
99 self.variables.get(&get_env_key(name)).cloned()
100 }
101
102 fn call(&self, name: &str, params: &[Value]) -> NativeResult {
103 let function = self
104 .functions
105 .get(&get_env_key(name))
106 .ok_or(NativeError::FunctionNotFound(name.to_string()))?;
107
108 let call = function.func;
109 call(params)
110 }
111
112 fn variable_exists(&self, name: &str) -> bool {
113 self.variables.contains_key(&get_env_key(name))
114 }
115
116 fn function_exists(&self, name: &str, param_count: usize) -> FunctionResult {
117 if let Some(function) = self.functions.get(&get_env_key(name)) {
118 match function.arity {
119 Arity::Polyadic { required, optional } => {
120 let min = required;
121 let max = required + optional;
122
123 if param_count < min || param_count > max {
124 FunctionResult::WrongArity { min, max }
125 } else {
126 FunctionResult::Exists {
127 pure: function.pure,
128 }
129 }
130 }
131 Arity::Variadic if param_count > 0 => FunctionResult::Exists {
132 pure: function.pure,
133 },
134 Arity::Variadic => FunctionResult::WrongArity { min: 1, max: 99 }, Arity::None => FunctionResult::WrongArity { min: 0, max: 0 },
136 }
137 } else {
138 FunctionResult::NotFound
139 }
140 }
141}
142
143#[cfg(test)]
144mod test {
145
146 use super::*;
147 use crate::{compile, execute};
148
149 #[test]
150 fn static_variables() {
151 let mut env = StaticEnvironment::default();
152
153 env.add_variable("some_var", Value::Number(42.0));
154 let ast = compile("some_var = 42").unwrap();
155 assert_eq!(Ok(Value::Boolean(true)), execute(&env, &ast));
156
157 env.remove_variable("some_var");
158 assert_eq!(Ok(Value::Boolean(false)), execute(&env, &ast));
159
160 env.add_variable("some_var", Value::Number(42.0));
161 let ast = compile("some_var = 42").unwrap();
162 assert_eq!(Ok(Value::Boolean(true)), execute(&env, &ast));
163
164 env.clear_variables();
165 assert_eq!(Ok(Value::Boolean(false)), execute(&env, &ast));
166 }
167
168 #[test]
169 fn static_functions() {
170 fn test_func(_params: &[Value]) -> NativeResult {
171 unreachable!()
172 }
173 let mut env = StaticEnvironment::default();
174
175 env.add_function(Function::new(test_func, Arity::Variadic, "test(...)"));
176
177 let registered = env.list_functions();
178 assert_eq!(1, registered.len());
179 assert_eq!("test", registered.first().unwrap().name);
180 let removed = env.remove_function("test").unwrap();
181
182 assert_eq!(removed.name, registered.first().unwrap().name);
183 }
184}