1use std::any::Any;
4use std::env;
5use std::panic;
6
7use runmat_builtins::{
8 BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
9 BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
10 CharArray, Value,
11};
12use runmat_macros::runtime_builtin;
13
14use crate::builtins::common::spec::{
15 BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
16 ReductionNaN, ResidencyPolicy, ShapeRequirements,
17};
18use crate::{build_runtime_error, gather_if_needed_async, BuiltinResult, RuntimeError};
19
20const MESSAGE_EMPTY_NAME: &str = "Environment variable name must not be empty.";
21const MESSAGE_NAME_HAS_EQUAL: &str = "Environment variable names must not contain '='.";
22const MESSAGE_NAME_HAS_NULL: &str = "Environment variable names must not contain null characters.";
23const MESSAGE_VALUE_HAS_NULL: &str =
24 "Environment variable values must not contain null characters.";
25const MESSAGE_OPERATION_FAILED: &str = "Unable to update environment variable: ";
26
27#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::io::repl_fs::setenv")]
28pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
29 name: "setenv",
30 op_kind: GpuOpKind::Custom("io"),
31 supported_precisions: &[],
32 broadcast: BroadcastSemantics::None,
33 provider_hooks: &[],
34 constant_strategy: ConstantStrategy::InlineLiteral,
35 residency: ResidencyPolicy::GatherImmediately,
36 nan_mode: ReductionNaN::Include,
37 two_pass_threshold: None,
38 workgroup_size: None,
39 accepts_nan_mode: false,
40 notes:
41 "Host-only environment mutation. GPU-resident arguments are gathered automatically before invoking the OS APIs.",
42};
43
44#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::io::repl_fs::setenv")]
45pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
46 name: "setenv",
47 shape: ShapeRequirements::Any,
48 constant_strategy: ConstantStrategy::InlineLiteral,
49 elementwise: None,
50 reduction: None,
51 emits_nan: false,
52 notes: "Environment updates terminate fusion; metadata registered for completeness.",
53};
54
55const BUILTIN_NAME: &str = "setenv";
56
57const SETENV_OUTPUT_STATUS: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
58 name: "status",
59 ty: BuiltinParamType::NumericScalar,
60 arity: BuiltinParamArity::Required,
61 default: None,
62 description: "0 on success, 1 on failure.",
63}];
64const SETENV_OUTPUT_STATUS_MESSAGE: [BuiltinParamDescriptor; 2] = [
65 BuiltinParamDescriptor {
66 name: "status",
67 ty: BuiltinParamType::NumericScalar,
68 arity: BuiltinParamArity::Required,
69 default: None,
70 description: "0 on success, 1 on failure.",
71 },
72 BuiltinParamDescriptor {
73 name: "message",
74 ty: BuiltinParamType::StringScalar,
75 arity: BuiltinParamArity::Required,
76 default: None,
77 description: "Failure message text, or empty on success.",
78 },
79];
80const SETENV_INPUTS_NAME_VALUE: [BuiltinParamDescriptor; 2] = [
81 BuiltinParamDescriptor {
82 name: "NAME",
83 ty: BuiltinParamType::StringScalar,
84 arity: BuiltinParamArity::Required,
85 default: None,
86 description: "Environment variable name.",
87 },
88 BuiltinParamDescriptor {
89 name: "VALUE",
90 ty: BuiltinParamType::StringScalar,
91 arity: BuiltinParamArity::Required,
92 default: None,
93 description: "Environment variable value (empty clears variable).",
94 },
95];
96const SETENV_SIGNATURES: [BuiltinSignatureDescriptor; 2] = [
97 BuiltinSignatureDescriptor {
98 label: "status = setenv(NAME, VALUE)",
99 inputs: &SETENV_INPUTS_NAME_VALUE,
100 outputs: &SETENV_OUTPUT_STATUS,
101 },
102 BuiltinSignatureDescriptor {
103 label: "[status, message] = setenv(NAME, VALUE)",
104 inputs: &SETENV_INPUTS_NAME_VALUE,
105 outputs: &SETENV_OUTPUT_STATUS_MESSAGE,
106 },
107];
108const SETENV_ERROR_TOO_FEW_INPUTS: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
109 code: "RM.SETENV.TOO_FEW_INPUTS",
110 identifier: None,
111 when: "Fewer than two positional inputs are supplied.",
112 message: "setenv: not enough input arguments",
113};
114const SETENV_ERROR_TOO_MANY_INPUTS: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
115 code: "RM.SETENV.TOO_MANY_INPUTS",
116 identifier: None,
117 when: "More than two positional inputs are supplied.",
118 message: "setenv: too many input arguments",
119};
120const SETENV_ERROR_NAME_TYPE: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
121 code: "RM.SETENV.NAME_TYPE",
122 identifier: None,
123 when: "NAME input is not a string scalar, char row, or string-array scalar.",
124 message: "setenv: NAME must be a string scalar or character vector",
125};
126const SETENV_ERROR_VALUE_TYPE: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
127 code: "RM.SETENV.VALUE_TYPE",
128 identifier: None,
129 when: "VALUE input is not a string scalar, char row, or string-array scalar.",
130 message: "setenv: VALUE must be a string scalar or character vector",
131};
132const SETENV_ERRORS: [BuiltinErrorDescriptor; 4] = [
133 SETENV_ERROR_TOO_FEW_INPUTS,
134 SETENV_ERROR_TOO_MANY_INPUTS,
135 SETENV_ERROR_NAME_TYPE,
136 SETENV_ERROR_VALUE_TYPE,
137];
138pub const SETENV_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
139 signatures: &SETENV_SIGNATURES,
140 output_mode: BuiltinOutputMode::ByRequestedOutputCount,
141 completion_policy: BuiltinCompletionPolicy::Public,
142 errors: &SETENV_ERRORS,
143};
144
145fn setenv_error(error: &'static BuiltinErrorDescriptor) -> RuntimeError {
146 let mut builder = build_runtime_error(error.message).with_builtin(BUILTIN_NAME);
147 if let Some(identifier) = error.identifier {
148 builder = builder.with_identifier(identifier);
149 }
150 builder.build()
151}
152
153fn map_control_flow(err: RuntimeError) -> RuntimeError {
154 let identifier = err.identifier().map(str::to_string);
155 let mut builder = build_runtime_error(format!("{BUILTIN_NAME}: {}", err.message()))
156 .with_builtin(BUILTIN_NAME)
157 .with_source(err);
158 if let Some(identifier) = identifier {
159 builder = builder.with_identifier(identifier);
160 }
161 builder.build()
162}
163
164#[runtime_builtin(
165 name = "setenv",
166 category = "io/repl_fs",
167 summary = "Set or clear environment variables with status outputs.",
168 keywords = "setenv,environment variable,status,message,unset",
169 accel = "cpu",
170 suppress_auto_output = true,
171 type_resolver(crate::builtins::io::type_resolvers::setenv_type),
172 descriptor(crate::builtins::io::repl_fs::setenv::SETENV_DESCRIPTOR),
173 builtin_path = "crate::builtins::io::repl_fs::setenv"
174)]
175async fn setenv_builtin(args: Vec<Value>) -> crate::BuiltinResult<Value> {
176 let eval = evaluate(&args).await?;
177 if let Some(out_count) = crate::output_count::current_output_count() {
178 if out_count == 0 {
179 return Ok(Value::OutputList(Vec::new()));
180 }
181 return Ok(crate::output_count::output_list_with_padding(
182 out_count,
183 eval.outputs(),
184 ));
185 }
186 Ok(eval.first_output())
187}
188
189pub async fn evaluate(args: &[Value]) -> BuiltinResult<SetenvResult> {
191 let gathered = gather_arguments(args).await?;
192 match gathered.len() {
193 0 | 1 => Err(setenv_error(&SETENV_ERROR_TOO_FEW_INPUTS)),
194 2 => apply(&gathered[0], &gathered[1]),
195 _ => Err(setenv_error(&SETENV_ERROR_TOO_MANY_INPUTS)),
196 }
197}
198
199#[derive(Debug, Clone)]
200pub struct SetenvResult {
201 status: f64,
202 message: String,
203}
204
205impl SetenvResult {
206 fn success() -> Self {
207 Self {
208 status: 0.0,
209 message: String::new(),
210 }
211 }
212
213 fn failure(message: String) -> Self {
214 Self {
215 status: 1.0,
216 message,
217 }
218 }
219
220 pub fn first_output(&self) -> Value {
221 Value::Num(self.status)
222 }
223
224 pub fn outputs(&self) -> Vec<Value> {
225 vec![Value::Num(self.status), char_array_value(&self.message)]
226 }
227
228 #[cfg(test)]
229 pub(crate) fn status(&self) -> f64 {
230 self.status
231 }
232
233 #[cfg(test)]
234 pub(crate) fn message(&self) -> &str {
235 &self.message
236 }
237}
238
239fn apply(name_value: &Value, value_value: &Value) -> BuiltinResult<SetenvResult> {
240 let name = extract_scalar_text(name_value, &SETENV_ERROR_NAME_TYPE)?;
241 let value = extract_scalar_text(value_value, &SETENV_ERROR_VALUE_TYPE)?;
242
243 if name.is_empty() {
244 return Ok(SetenvResult::failure(MESSAGE_EMPTY_NAME.to_string()));
245 }
246 if name.chars().any(|ch| ch == '=') {
247 return Ok(SetenvResult::failure(MESSAGE_NAME_HAS_EQUAL.to_string()));
248 }
249 if name.chars().any(|ch| ch == '\0') {
250 return Ok(SetenvResult::failure(MESSAGE_NAME_HAS_NULL.to_string()));
251 }
252 if value.chars().any(|ch| ch == '\0') {
253 return Ok(SetenvResult::failure(MESSAGE_VALUE_HAS_NULL.to_string()));
254 }
255
256 Ok(update_environment(&name, &value))
257}
258
259fn update_environment(name: &str, value: &str) -> SetenvResult {
260 if value.is_empty() {
261 match panic::catch_unwind(|| env::remove_var(name)) {
262 Ok(()) => SetenvResult::success(),
263 Err(payload) => SetenvResult::failure(format!(
264 "{}{}",
265 MESSAGE_OPERATION_FAILED,
266 panic_payload_to_string(payload)
267 )),
268 }
269 } else {
270 match panic::catch_unwind(|| env::set_var(name, value)) {
271 Ok(()) => SetenvResult::success(),
272 Err(payload) => SetenvResult::failure(format!(
273 "{}{}",
274 MESSAGE_OPERATION_FAILED,
275 panic_payload_to_string(payload)
276 )),
277 }
278 }
279}
280
281async fn gather_arguments(args: &[Value]) -> BuiltinResult<Vec<Value>> {
282 let mut out = Vec::with_capacity(args.len());
283 for value in args {
284 out.push(
285 gather_if_needed_async(value)
286 .await
287 .map_err(map_control_flow)?,
288 );
289 }
290 Ok(out)
291}
292
293fn extract_scalar_text(
294 value: &Value,
295 error_message: &'static BuiltinErrorDescriptor,
296) -> BuiltinResult<String> {
297 match value {
298 Value::String(text) => Ok(text.clone()),
299 Value::CharArray(array) => {
300 if array.rows != 1 {
301 return Err(setenv_error(error_message));
302 }
303 Ok(char_row_to_string(array))
304 }
305 Value::StringArray(array) => {
306 if array.data.len() == 1 {
307 Ok(array.data[0].clone())
308 } else {
309 Err(setenv_error(error_message))
310 }
311 }
312 _ => Err(setenv_error(error_message)),
313 }
314}
315
316fn char_row_to_string(array: &CharArray) -> String {
317 if array.cols == 0 {
318 return String::new();
319 }
320 let mut text = String::with_capacity(array.cols);
321 for col in 0..array.cols {
322 text.push(array.data[col]);
323 }
324 while text.ends_with(' ') {
325 text.pop();
326 }
327 text
328}
329
330fn char_array_value(text: &str) -> Value {
331 Value::CharArray(CharArray::new_row(text))
332}
333
334fn panic_payload_to_string(payload: Box<dyn Any + Send>) -> String {
335 match payload.downcast::<String>() {
336 Ok(msg) => *msg,
337 Err(payload) => match payload.downcast::<&'static str>() {
338 Ok(msg) => (*msg).to_string(),
339 Err(_) => "operation failed".to_string(),
340 },
341 }
342}
343
344#[cfg(test)]
345pub(crate) mod tests {
346 use super::*;
347 use crate::builtins::io::repl_fs::REPL_FS_TEST_LOCK;
348 use runmat_builtins::{CharArray, StringArray, Value};
349
350 fn setenv_builtin(args: Vec<Value>) -> BuiltinResult<Value> {
351 futures::executor::block_on(super::setenv_builtin(args))
352 }
353
354 fn evaluate(args: &[Value]) -> BuiltinResult<SetenvResult> {
355 futures::executor::block_on(super::evaluate(args))
356 }
357
358 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
359 #[test]
360 fn setenv_descriptor_signatures_cover_core_forms() {
361 let labels: Vec<&str> = SETENV_DESCRIPTOR
362 .signatures
363 .iter()
364 .map(|sig| sig.label)
365 .collect();
366 assert!(labels.contains(&"status = setenv(NAME, VALUE)"));
367 assert!(labels.contains(&"[status, message] = setenv(NAME, VALUE)"));
368 }
369
370 fn unique_name(suffix: &str) -> String {
371 format!("RUNMAT_TEST_SETENV_{}", suffix)
372 }
373
374 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
375 #[test]
376 fn setenv_sets_variable_and_returns_success() {
377 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
378 let name = unique_name("BASIC");
379 env::remove_var(&name);
380
381 let result = setenv_builtin(vec![
382 Value::String(name.clone()),
383 Value::String("value".to_string()),
384 ])
385 .expect("setenv");
386
387 assert_eq!(result, Value::Num(0.0));
388 assert_eq!(env::var(&name).unwrap(), "value");
389 env::remove_var(name);
390 }
391
392 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
393 #[test]
394 fn setenv_removes_variable_when_value_is_empty() {
395 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
396 let name = unique_name("REMOVE");
397 env::set_var(&name, "seed");
398
399 let result = setenv_builtin(vec![
400 Value::String(name.clone()),
401 Value::CharArray(CharArray::new_row("")),
402 ])
403 .expect("setenv");
404
405 assert_eq!(result, Value::Num(0.0));
406 assert!(env::var(&name).is_err());
407 env::remove_var(name);
408 }
409
410 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
411 #[test]
412 fn setenv_reports_failure_for_illegal_name() {
413 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
414 let eval = evaluate(&[
415 Value::String("INVALID=NAME".to_string()),
416 Value::String("value".to_string()),
417 ])
418 .expect("evaluate");
419
420 assert_eq!(eval.status(), 1.0);
421 assert_eq!(eval.message(), MESSAGE_NAME_HAS_EQUAL);
422 }
423
424 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
425 #[test]
426 fn setenv_reports_failure_for_empty_name() {
427 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
428 let eval = evaluate(&[
429 Value::String(String::new()),
430 Value::String("value".to_string()),
431 ])
432 .expect("evaluate");
433
434 assert_eq!(eval.status(), 1.0);
435 assert_eq!(eval.message(), MESSAGE_EMPTY_NAME);
436 }
437
438 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
439 #[test]
440 fn setenv_reports_failure_for_null_in_name() {
441 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
442 let eval = evaluate(&[
443 Value::String("BAD\0NAME".to_string()),
444 Value::String("value".to_string()),
445 ])
446 .expect("evaluate");
447
448 assert_eq!(eval.status(), 1.0);
449 assert_eq!(eval.message(), MESSAGE_NAME_HAS_NULL);
450 }
451
452 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
453 #[test]
454 fn setenv_reports_failure_for_null_in_value() {
455 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
456 let eval = evaluate(&[
457 Value::String("RUNMAT_NULL_VALUE".to_string()),
458 Value::String("abc\0def".to_string()),
459 ])
460 .expect("evaluate");
461
462 assert_eq!(eval.status(), 1.0);
463 assert_eq!(eval.message(), MESSAGE_VALUE_HAS_NULL);
464 }
465
466 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
467 #[test]
468 fn setenv_errors_when_name_is_not_text() {
469 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
470 let err =
471 setenv_builtin(vec![Value::Num(5.0), Value::String("value".to_string())]).unwrap_err();
472 assert_eq!(err.message(), SETENV_ERROR_NAME_TYPE.message);
473 }
474
475 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
476 #[test]
477 fn setenv_errors_when_value_is_not_text() {
478 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
479 let err = setenv_builtin(vec![
480 Value::String("RUNMAT_INVALID_VALUE".to_string()),
481 Value::Num(1.0),
482 ])
483 .unwrap_err();
484 assert_eq!(err.message(), SETENV_ERROR_VALUE_TYPE.message);
485 }
486
487 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
488 #[test]
489 fn setenv_accepts_scalar_string_array_arguments() {
490 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
491 let name = unique_name("STRING_ARRAY");
492 env::remove_var(&name);
493
494 let name_array =
495 StringArray::new(vec![name.clone()], vec![1]).expect("scalar string array name");
496 let value_array =
497 StringArray::new(vec!["VALUE".to_string()], vec![1]).expect("scalar string array");
498
499 let status = setenv_builtin(vec![
500 Value::StringArray(name_array),
501 Value::StringArray(value_array),
502 ])
503 .expect("setenv");
504
505 assert_eq!(status, Value::Num(0.0));
506 assert_eq!(env::var(&name).unwrap(), "VALUE");
507 env::remove_var(name);
508 }
509
510 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
511 #[test]
512 fn setenv_errors_for_string_array_with_multiple_elements() {
513 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
514 let array =
515 StringArray::new(vec!["A".to_string(), "B".to_string()], vec![2]).expect("array");
516 let err = setenv_builtin(vec![
517 Value::StringArray(array),
518 Value::String("value".to_string()),
519 ])
520 .unwrap_err();
521 assert_eq!(err.message(), SETENV_ERROR_NAME_TYPE.message);
522 }
523
524 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
525 #[test]
526 fn setenv_errors_for_char_array_with_multiple_rows() {
527 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
528 let array = CharArray::new(vec!['R', 'M'], 2, 1).expect("two-row char array");
529 let err = setenv_builtin(vec![
530 Value::CharArray(array),
531 Value::String("value".to_string()),
532 ])
533 .unwrap_err();
534 assert_eq!(err.message(), SETENV_ERROR_NAME_TYPE.message);
535 }
536
537 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
538 #[test]
539 fn setenv_char_array_input_trims_padding() {
540 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
541 let chars = vec!['F', 'O', 'O', ' '];
542 let array = CharArray::new(chars, 1, 4).unwrap();
543 let result =
544 extract_scalar_text(&Value::CharArray(array), &SETENV_ERROR_NAME_TYPE).unwrap();
545 assert_eq!(result, "FOO");
546 }
547
548 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
549 #[test]
550 fn setenv_outputs_success_message_is_empty_char_array() {
551 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
552 let name = unique_name("SUCCESS_MSG");
553 env::remove_var(&name);
554
555 let eval = evaluate(&[
556 Value::String(name.clone()),
557 Value::String("value".to_string()),
558 ])
559 .expect("evaluate");
560 let outputs = eval.outputs();
561 assert_eq!(outputs.len(), 2);
562 match &outputs[1] {
563 Value::CharArray(ca) => {
564 assert_eq!(ca.rows, 1);
565 assert_eq!(ca.cols, 0);
566 }
567 other => panic!("expected empty CharArray, got {other:?}"),
568 }
569
570 env::remove_var(name);
571 }
572
573 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
574 #[test]
575 fn setenv_outputs_return_status_and_message() {
576 let _guard = REPL_FS_TEST_LOCK.lock().unwrap();
577 let eval = evaluate(&[
578 Value::String("INVALID=NAME".to_string()),
579 Value::String("value".to_string()),
580 ])
581 .expect("evaluate");
582
583 let outputs = eval.outputs();
584 assert_eq!(outputs.len(), 2);
585 assert!(matches!(outputs[0], Value::Num(1.0)));
586 match &outputs[1] {
587 Value::CharArray(ca) => {
588 assert_eq!(ca.rows, 1);
589 let text: String = ca.data.iter().collect();
590 assert_eq!(text, MESSAGE_NAME_HAS_EQUAL);
591 }
592 other => panic!("expected CharArray message, got {other:?}"),
593 }
594 }
595}