1use std::collections::HashMap;
2use std::rc::Rc;
3
4type BodyCache = HashMap<String, Rc<Vec<StmtEntry>>>;
6
7fn expand_tilde(path: &str) -> String {
12 if path == "~" || path.starts_with("~/") || path.starts_with("~\\") {
13 let home = std::env::var("HOME")
14 .or_else(|_| std::env::var("USERPROFILE"))
15 .unwrap_or_default();
16 if home.is_empty() {
17 return path.to_string();
18 }
19 if path == "~" {
20 home
21 } else {
22 format!("{}{}", home, &path[1..])
23 }
24 } else {
25 path.to_string()
26 }
27}
28
29use indexmap::IndexMap;
30use ndarray::Array2;
31use num_complex::Complex;
32
33use crate::env::{Env, Value};
34use crate::env::{load_workspace, save_workspace, save_workspace_vars};
35use crate::eval::{
36 Base, Expr, FormatMode, autoload_cache_insert, current_func_name, eval_with_io, format_complex,
37 format_scalar, format_value_full, get_display_base, get_display_compact, get_display_fmt,
38 global_declare, global_frame_pop, global_frame_push, global_get, global_init_if_absent,
39 global_refresh_into_env, global_set, is_global, is_persistent, persistent_declare,
40 persistent_frame_pop, persistent_frame_push, persistent_load, persistent_save,
41 set_autoload_hook, set_display_ctx, set_eval_str_hook, set_fn_call_hook, set_last_err,
42 set_nargout,
43};
44use crate::io::IoContext;
45use crate::parser::{Stmt, StmtEntry, parse_stmts};
46
47thread_local! {
48 static RUN_DEPTH: std::cell::Cell<u32> = const { std::cell::Cell::new(0) };
50
51 static SCRIPT_DIR_STACK: std::cell::RefCell<Vec<std::path::PathBuf>> =
57 const { std::cell::RefCell::new(Vec::new()) };
58
59 static SESSION_PATH: std::cell::RefCell<Vec<std::path::PathBuf>> =
64 const { std::cell::RefCell::new(Vec::new()) };
65
66 static BODY_CACHE: std::cell::RefCell<BodyCache> =
76 std::cell::RefCell::new(HashMap::new());
77}
78
79fn silence_all(stmts: Vec<StmtEntry>) -> Vec<StmtEntry> {
87 stmts
88 .into_iter()
89 .map(|(stmt, _, line)| {
90 let stmt = match stmt {
91 Stmt::If {
92 cond,
93 body,
94 elseif_branches,
95 else_body,
96 } => Stmt::If {
97 cond,
98 body: silence_all(body),
99 elseif_branches: elseif_branches
100 .into_iter()
101 .map(|(c, b)| (c, silence_all(b)))
102 .collect(),
103 else_body: else_body.map(silence_all),
104 },
105 Stmt::For {
106 var,
107 range_expr,
108 body,
109 } => Stmt::For {
110 var,
111 range_expr,
112 body: silence_all(body),
113 },
114 Stmt::While { cond, body } => Stmt::While {
115 cond,
116 body: silence_all(body),
117 },
118 Stmt::DoUntil { body, cond } => Stmt::DoUntil {
119 body: silence_all(body),
120 cond,
121 },
122 Stmt::Switch {
123 expr,
124 cases,
125 otherwise_body,
126 } => Stmt::Switch {
127 expr,
128 cases: cases
129 .into_iter()
130 .map(|(v, b)| (v, silence_all(b)))
131 .collect(),
132 otherwise_body: otherwise_body.map(silence_all),
133 },
134 Stmt::TryCatch {
135 try_body,
136 catch_var,
137 catch_body,
138 } => Stmt::TryCatch {
139 try_body: silence_all(try_body),
140 catch_var,
141 catch_body: silence_all(catch_body),
142 },
143 other => other,
144 };
145 (stmt, true, line)
146 })
147 .collect()
148}
149
150fn get_or_parse_body(body_source: &str) -> Result<Rc<Vec<StmtEntry>>, String> {
156 BODY_CACHE.with(|cache| {
157 let mut cache = cache.borrow_mut();
158 if let Some(body) = cache.get(body_source) {
159 return Ok(Rc::clone(body));
160 }
161 let stmts =
162 parse_stmts(body_source).map_err(|e| format!("function body parse error: {e}"))?;
163 let silent = silence_all(stmts);
164 let rc = Rc::new(silent);
165 cache.insert(body_source.to_string(), Rc::clone(&rc));
166 Ok(rc)
167 })
168}
169
170fn annotate_line(e: String, line: usize) -> String {
174 if line == 0 || e.contains("near line") {
175 e
176 } else {
177 format!("{e} near line {line}")
178 }
179}
180
181fn strip_near_line(s: String) -> String {
186 if let Some(pos) = s.rfind(" near line ") {
187 let after = &s[pos + " near line ".len()..];
188 if !after.is_empty() && after.bytes().all(|b| b.is_ascii_digit()) {
189 return s[..pos].to_string();
190 }
191 }
192 s
193}
194
195pub enum Signal {
201 Break,
203 Continue,
205 Return,
207}
208
209pub fn init() {
213 set_fn_call_hook(call_user_function);
214 set_autoload_hook(try_autoload);
215 set_eval_str_hook(eval_str_impl);
216}
217
218fn eval_str_impl(code: &str, env: &crate::env::Env) -> Result<crate::env::Value, String> {
224 let mut env_clone = env.clone();
225 env_clone
226 .entry("ans".to_string())
227 .or_insert(crate::env::Value::Scalar(0.0));
228 let mut tmp_io = crate::io::IoContext::new();
229 let fmt = get_display_fmt();
230 let base = get_display_base();
231 let compact = get_display_compact();
232 let stmts = parse_stmts(code).map_err(|e| format!("eval: parse error: {e}"))?;
233 exec_stmts(&stmts, &mut env_clone, &mut tmp_io, &fmt, base, compact)?;
234 Ok(env_clone
235 .get("ans")
236 .cloned()
237 .unwrap_or(crate::env::Value::Void))
238}
239
240fn try_autoload(name: &str) -> bool {
247 if name.contains('.') {
248 return try_autoload_pkg(name);
249 }
250 let candidates = [format!("{name}.calc"), format!("{name}.m")];
251 for candidate in &candidates {
252 let Some(path) = resolve_script_path(candidate) else {
253 continue;
254 };
255 let Ok(content) = std::fs::read_to_string(&path) else {
256 continue;
257 };
258 let Ok(stmts) = parse_stmts(&content) else {
259 continue;
260 };
261 if !matches!(stmts.first(), Some((Stmt::FunctionDef { .. }, _, _))) {
262 continue;
263 }
264 let primary_name = match &stmts[0].0 {
265 Stmt::FunctionDef { name, .. } => name.clone(),
266 _ => continue,
267 };
268 let mut locals: IndexMap<String, Value> = IndexMap::new();
269 for (stmt, _, _) in &stmts {
270 if let Stmt::FunctionDef {
271 name: n,
272 outputs,
273 params,
274 body_source,
275 doc,
276 } = stmt
277 && n != &primary_name
278 {
279 locals.insert(
280 n.clone(),
281 Value::Function {
282 outputs: outputs.clone(),
283 params: params.clone(),
284 body_source: body_source.clone(),
285 locals: IndexMap::new(),
286 doc: doc.clone(),
287 },
288 );
289 }
290 }
291 if let Stmt::FunctionDef {
292 outputs,
293 params,
294 body_source,
295 doc,
296 ..
297 } = &stmts[0].0
298 {
299 autoload_cache_insert(
300 primary_name,
301 Value::Function {
302 outputs: outputs.clone(),
303 params: params.clone(),
304 body_source: body_source.clone(),
305 locals,
306 doc: doc.clone(),
307 },
308 );
309 return true;
310 }
311 }
312 false
313}
314
315fn try_autoload_pkg(qualified: &str) -> bool {
321 let parts: Vec<&str> = qualified.split('.').collect();
322 if parts.len() < 2 {
323 return false;
324 }
325 let func_name = *parts.last().unwrap();
326 let pkg_parts = &parts[..parts.len() - 1];
327
328 let mut rel_prefix = std::path::PathBuf::new();
330 for pkg in pkg_parts {
331 rel_prefix.push(format!("+{pkg}"));
332 }
333
334 let candidates = [
335 rel_prefix.join(format!("{func_name}.calc")),
336 rel_prefix.join(format!("{func_name}.m")),
337 ];
338
339 let mut search_dirs: Vec<std::path::PathBuf> = Vec::new();
341 SCRIPT_DIR_STACK.with(|s| search_dirs.extend(s.borrow().iter().cloned()));
342 search_dirs.push(std::path::PathBuf::from("."));
343 SESSION_PATH.with(|s| search_dirs.extend(s.borrow().iter().cloned()));
344
345 for dir in &search_dirs {
346 for candidate in &candidates {
347 let full = dir.join(candidate);
348 let Ok(content) = std::fs::read_to_string(&full) else {
349 continue;
350 };
351 let Ok(stmts) = parse_stmts(&content) else {
352 continue;
353 };
354 if !matches!(stmts.first(), Some((Stmt::FunctionDef { .. }, _, _))) {
355 continue;
356 }
357 let primary_name = match &stmts[0].0 {
358 Stmt::FunctionDef { name, .. } => name.clone(),
359 _ => continue,
360 };
361 let mut locals: IndexMap<String, Value> = IndexMap::new();
362 for (stmt, _, _) in &stmts {
363 if let Stmt::FunctionDef {
364 name: n,
365 outputs,
366 params,
367 body_source,
368 doc,
369 } = stmt
370 && n != &primary_name
371 {
372 locals.insert(
373 n.clone(),
374 Value::Function {
375 outputs: outputs.clone(),
376 params: params.clone(),
377 body_source: body_source.clone(),
378 locals: IndexMap::new(),
379 doc: doc.clone(),
380 },
381 );
382 }
383 }
384 if let Stmt::FunctionDef {
385 outputs,
386 params,
387 body_source,
388 doc,
389 ..
390 } = &stmts[0].0
391 {
392 autoload_cache_insert(
393 qualified.to_string(),
394 Value::Function {
395 outputs: outputs.clone(),
396 params: params.clone(),
397 body_source: body_source.clone(),
398 locals,
399 doc: doc.clone(),
400 },
401 );
402 return true;
403 }
404 }
405 }
406 false
407}
408
409pub fn script_dir_push(dir: &std::path::Path) {
415 SCRIPT_DIR_STACK.with(|s| s.borrow_mut().push(dir.to_path_buf()));
416}
417
418pub fn script_dir_pop() {
420 SCRIPT_DIR_STACK.with(|s| s.borrow_mut().pop());
421}
422
423pub fn session_path_init(paths: Vec<std::path::PathBuf>) {
428 SESSION_PATH.with(|p| *p.borrow_mut() = paths);
429}
430
431pub fn session_path_add(path: std::path::PathBuf, append: bool) {
436 SESSION_PATH.with(|p| {
437 let mut v = p.borrow_mut();
438 v.retain(|e| e != &path);
439 if append {
440 v.push(path);
441 } else {
442 v.insert(0, path);
443 }
444 });
445}
446
447pub fn session_path_remove(path: &std::path::Path) {
449 SESSION_PATH.with(|p| p.borrow_mut().retain(|e| e.as_path() != path));
450}
451
452pub fn session_path_list() -> Vec<std::path::PathBuf> {
454 SESSION_PATH.with(|p| p.borrow().clone())
455}
456
457fn call_user_function(
464 name: &str,
465 func: &Value,
466 args: &[Value],
467 caller_env: &Env,
468 io: &mut IoContext,
469) -> Result<Value, String> {
470 let Value::Function {
471 outputs,
472 params,
473 body_source,
474 locals,
475 ..
476 } = func
477 else {
478 return Err("call_user_function: not a Function value".to_string());
479 };
480
481 global_frame_push();
483 persistent_frame_push(name); let mut local_env = Env::new();
489 local_env.insert("i".to_string(), Value::Complex(0.0, 1.0));
490 local_env.insert("j".to_string(), Value::Complex(0.0, 1.0));
491 local_env.insert("ans".to_string(), Value::Scalar(0.0));
492 for (fn_name, val) in locals.iter() {
495 local_env.insert(fn_name.clone(), val.clone());
496 }
497 for (var_name, val) in caller_env.iter() {
498 if matches!(val, Value::Function { .. } | Value::Lambda(_)) {
499 local_env.insert(var_name.clone(), val.clone());
500 }
501 }
502
503 let has_varargin = params.last().is_some_and(|p| p == "varargin");
505 let fixed_params = if has_varargin {
506 ¶ms[..params.len() - 1]
507 } else {
508 params.as_slice()
509 };
510
511 let effective_args = if args.len() > params.len() {
516 if !has_varargin && args.len() > params.len() + 1 {
517 return Err(format!(
518 "Too many arguments: expected at most {}, got {}",
519 params.len(),
520 args.len()
521 ));
522 }
523 if has_varargin {
524 args
525 } else {
526 &args[..params.len()]
528 }
529 } else {
530 args
531 };
532
533 for (p, a) in fixed_params.iter().zip(effective_args.iter()) {
535 local_env.insert(p.clone(), a.clone());
536 }
537
538 if has_varargin {
540 let extra: Vec<Value> = effective_args
544 .get(fixed_params.len()..)
545 .unwrap_or(&[])
546 .to_vec();
547 let varargin = Value::Cell(extra);
548 local_env.insert("varargin".to_string(), varargin);
549 }
550
551 let nargin = effective_args.len().min(params.len());
552 local_env.insert("nargin".to_string(), Value::Scalar(nargin as f64));
553 local_env.insert("nargout".to_string(), Value::Scalar(outputs.len() as f64));
554
555 let body = get_or_parse_body(body_source)?;
557 let fmt = get_display_fmt();
558 let base = get_display_base();
559 let compact = get_display_compact();
560 let exec_result = exec_stmts(&body, &mut local_env, io, &fmt, base, compact);
561
562 let (func_name_saved, persistent_names) = persistent_frame_pop();
564 for var_name in &persistent_names {
565 if let Some(val) = local_env.get(var_name) {
566 persistent_save(&func_name_saved, var_name, val.clone());
567 }
568 }
569 global_frame_pop();
571
572 match exec_result? {
574 None | Some(Signal::Return) => {}
575 Some(Signal::Break) => return Err("'break' outside loop".to_string()),
576 Some(Signal::Continue) => return Err("'continue' outside loop".to_string()),
577 }
578
579 if outputs.is_empty() {
581 return Ok(Value::Void);
582 }
583
584 if outputs.len() == 1 && outputs[0] == "varargout" {
586 let cell = local_env.remove("varargout").unwrap_or(Value::Cell(vec![]));
587 return match cell {
588 Value::Cell(mut v) => {
589 if v.is_empty() {
590 Ok(Value::Void)
591 } else if v.len() == 1 {
592 Ok(v.remove(0))
593 } else {
594 Ok(Value::Tuple(v))
595 }
596 }
597 other => Ok(other),
598 };
599 }
600
601 if outputs.len() == 1 {
602 return Ok(local_env.remove(&outputs[0]).unwrap_or(Value::Void));
603 }
604 let vals: Vec<Value> = outputs
605 .iter()
606 .map(|o| local_env.remove(o).unwrap_or(Value::Void))
607 .collect();
608 Ok(Value::Tuple(vals))
609}
610
611pub fn resolve_script_path(name: &str) -> Option<std::path::PathBuf> {
617 let p = std::path::Path::new(name);
626 let mut bases: Vec<std::path::PathBuf> = Vec::new();
627
628 SCRIPT_DIR_STACK.with(|stack| {
630 for dir in stack.borrow().iter().rev() {
631 bases.push(dir.join("private").join(p));
632 bases.push(dir.join(p));
633 }
634 });
635
636 bases.push(p.to_path_buf());
638
639 SESSION_PATH.with(|sp| {
641 for dir in sp.borrow().iter() {
642 bases.push(dir.join(p));
643 }
644 });
645
646 for base in &bases {
647 if base.extension().is_some() {
648 if base.exists() {
649 return Some(base.clone());
650 }
651 continue;
653 }
654 let with_calc = base.with_extension("calc");
655 if with_calc.exists() {
656 return Some(with_calc);
657 }
658 let with_m = base.with_extension("m");
659 if with_m.exists() {
660 return Some(with_m);
661 }
662 }
663 None
664}
665
666fn is_truthy(val: &Value) -> bool {
674 match val {
675 Value::Scalar(n) => *n != 0.0 && !n.is_nan(),
676 Value::Matrix(m) => m.iter().all(|&x| x != 0.0 && !x.is_nan()),
677 Value::Complex(re, im) => *re != 0.0 || *im != 0.0,
678 Value::ComplexMatrix(m) => m.iter().all(|c| c.re != 0.0 || c.im != 0.0),
679 Value::Str(s) | Value::StringObj(s) => !s.is_empty(),
680 Value::Void => false,
681 Value::Lambda(_) | Value::Function { .. } | Value::Tuple(_) => true,
684 Value::Cell(v) => !v.is_empty(),
686 Value::Struct(_) | Value::StructArray(_) => true,
688 Value::DateTime(ts) => !ts.is_nan(),
690 Value::Duration(s) => *s != 0.0,
691 Value::DateTimeArray(v) | Value::DurationArray(v) => !v.is_empty(),
692 }
693}
694
695fn print_value(label: Option<&str>, val: &Value, fmt: &FormatMode, base: Base, compact: bool) {
700 match val {
701 Value::Void => {}
702 Value::Scalar(n) => {
703 if let Some(name) = label {
704 println!("{name} = {}", format_scalar(*n, base, fmt));
705 } else {
706 println!("{}", format_scalar(*n, base, fmt));
707 }
708 }
709 Value::Matrix(_) => {
710 if let Some(full) = format_value_full(val, fmt) {
711 let prefix = label.unwrap_or("ans");
712 println!("{prefix} =");
713 println!("{full}");
714 if !compact {
715 println!();
716 }
717 }
718 }
719 Value::Complex(re, im) => {
720 if let Some(name) = label {
721 println!("{name} = {}", format_complex(*re, *im, fmt));
722 } else {
723 println!("{}", format_complex(*re, *im, fmt));
724 }
725 }
726 Value::Str(s) | Value::StringObj(s) => {
727 if let Some(name) = label {
728 println!("{name} = {s}");
729 } else {
730 println!("{s}");
731 }
732 }
733 Value::Lambda(_) => {
734 if let Some(name) = label {
735 println!("{name} = @<lambda>");
736 } else {
737 println!("@<lambda>");
738 }
739 }
740 Value::Function {
741 outputs, params, ..
742 } => {
743 let params_str = params.join(", ");
744 let out_str = match outputs.len() {
745 0 => String::new(),
746 1 => format!("{} = ", outputs[0]),
747 _ => format!("[{}] = ", outputs.join(", ")),
748 };
749 if let Some(name) = label {
750 println!("{name} = @function {out_str}{name}({params_str})");
751 } else {
752 println!("@function {out_str}f({params_str})");
753 }
754 }
755 Value::Tuple(vals) => {
756 for (i, v) in vals.iter().enumerate() {
759 print_value(label.map(|_| "ans").or(Some("ans")), v, fmt, base, compact);
760 let _ = i;
761 }
762 }
763 Value::ComplexMatrix(_) => {
764 if let Some(full) = format_value_full(val, fmt) {
765 let prefix = label.unwrap_or("ans");
766 println!("{prefix} =");
767 println!("{full}");
768 if !compact {
769 println!();
770 }
771 }
772 }
773 Value::Cell(_) | Value::Struct(_) | Value::StructArray(_) => {
774 if let Some(full) = format_value_full(val, fmt) {
775 let prefix = label.unwrap_or("ans");
776 println!("{prefix} =");
777 println!("{full}");
778 if !compact {
779 println!();
780 }
781 }
782 }
783 Value::DateTime(ts) => {
784 let s = crate::datetime::format_datetime(*ts);
785 if let Some(name) = label {
786 println!("{name} = {s}");
787 } else {
788 println!("{s}");
789 }
790 }
791 Value::Duration(secs) => {
792 let s = crate::datetime::format_duration(*secs);
793 if let Some(name) = label {
794 println!("{name} = {s}");
795 } else {
796 println!("{s}");
797 }
798 }
799 Value::DateTimeArray(_) | Value::DurationArray(_) => {
800 if let Some(full) = format_value_full(val, fmt) {
801 let prefix = label.unwrap_or("ans");
802 println!("{prefix} =");
803 println!("{full}");
804 if !compact {
805 println!();
806 }
807 }
808 }
809 }
810}
811
812fn set_nested(
817 mut map: IndexMap<String, Value>,
818 path: &[String],
819 val: Value,
820) -> Result<IndexMap<String, Value>, String> {
821 let (first, rest) = path.split_first().expect("set_nested: empty path");
822 if rest.is_empty() {
823 map.insert(first.clone(), val);
824 } else {
825 let inner = match map.shift_remove(first) {
826 Some(Value::Struct(m)) => m,
827 None => IndexMap::new(),
828 Some(other) => {
829 map.insert(first.clone(), other);
830 return Err(format!("'{first}' is not a struct"));
831 }
832 };
833 let updated = set_nested(inner, rest, val)?;
834 map.insert(first.clone(), Value::Struct(updated));
835 }
836 Ok(map)
837}
838
839pub fn exec_stmts(
850 stmts: &[StmtEntry],
851 env: &mut Env,
852 io: &mut IoContext,
853 fmt: &FormatMode,
854 base: Base,
855 compact: bool,
856) -> Result<Option<Signal>, String> {
857 set_display_ctx(fmt, base, compact);
859
860 for (stmt, silent, stmt_line) in stmts {
861 match stmt {
862 Stmt::Assign(name, expr) => {
863 set_nargout(1);
864 let val = eval_with_io(expr, env, io).map_err(|e| annotate_line(e, *stmt_line))?;
865 env.insert(name.clone(), val.clone());
866 if is_global(name) {
868 global_set(name, val.clone());
869 }
870 if is_persistent(name) {
872 persistent_save(¤t_func_name(), name, val.clone());
873 }
874 if !silent && !matches!(val, Value::Void) {
875 print_value(Some(name), &val, fmt, base, compact);
876 }
877 }
878
879 Stmt::Global(names) => {
880 for name in names {
881 global_declare(name);
882 global_init_if_absent(name);
883 if let Some(local_val) = env.remove(name) {
886 global_set(name, local_val.clone());
887 env.insert(name.clone(), local_val);
888 } else if let Some(global_val) = global_get(name) {
889 env.insert(name.clone(), global_val);
890 }
891 }
892 }
893
894 Stmt::Persistent(names) => {
895 let func = current_func_name();
898 for name in names {
899 persistent_declare(name);
900 if let Some(saved) = persistent_load(&func, name) {
901 env.insert(name.clone(), saved);
903 } else {
904 env.insert(name.clone(), Value::Matrix(ndarray::Array2::zeros((0, 0))));
908 }
909 }
910 }
911
912 Stmt::Expr(expr) => {
913 if let Expr::Call(fn_name, args) = expr
915 && matches!(fn_name.as_str(), "addpath" | "rmpath" | "path")
916 {
917 match fn_name.as_str() {
918 "addpath" => {
919 if args.is_empty() || args.len() > 2 {
920 return Err(
921 "addpath: expects 1 or 2 arguments: addpath(dir) or addpath(dir, '-end')".to_string()
922 );
923 }
924 let path_val = eval_with_io(&args[0], env, io)?;
925 let path_str = match &path_val {
926 Value::Str(s) | Value::StringObj(s) => s.clone(),
927 _ => {
928 return Err(
929 "addpath: argument must be a string (directory path)"
930 .to_string(),
931 );
932 }
933 };
934 let append = if args.len() == 2 {
935 let flag_val = eval_with_io(&args[1], env, io)?;
936 match &flag_val {
937 Value::Str(s) | Value::StringObj(s) if s == "-end" => true,
938 Value::Str(_) | Value::StringObj(_) => {
939 return Err(
940 "addpath: second argument must be '-end' (to append) or omitted (to prepend)".to_string()
941 );
942 }
943 _ => {
944 return Err(
945 "addpath: second argument must be a string '-end'"
946 .to_string(),
947 );
948 }
949 }
950 } else {
951 false
952 };
953 let expanded = expand_tilde(&path_str);
954 let pb = std::path::PathBuf::from(&expanded);
955 session_path_add(pb, append);
956 if !silent {
957 for p in session_path_list() {
958 println!("{}", p.display());
959 }
960 }
961 }
962 "rmpath" => {
963 if args.len() != 1 {
964 return Err("rmpath: expects exactly 1 argument".to_string());
965 }
966 let path_val = eval_with_io(&args[0], env, io)?;
967 let path_str = match &path_val {
968 Value::Str(s) | Value::StringObj(s) => s.clone(),
969 _ => {
970 return Err(
971 "rmpath: argument must be a string (directory path)"
972 .to_string(),
973 );
974 }
975 };
976 let expanded = expand_tilde(&path_str);
977 session_path_remove(std::path::Path::new(&expanded));
978 }
979 "path" => {
980 if !args.is_empty() {
981 return Err("path: takes no arguments".to_string());
982 }
983 if !silent {
984 let paths = session_path_list();
985 if paths.is_empty() {
986 println!("(search path is empty)");
987 } else {
988 for p in &paths {
989 println!("{}", p.display());
990 }
991 }
992 }
993 }
994 _ => unreachable!(),
995 }
996 continue;
997 }
998
999 if let Expr::Call(fn_name, args) = expr
1002 && matches!(fn_name.as_str(), "run" | "source")
1003 && args.len() == 1
1004 {
1005 let path_val = eval_with_io(&args[0], env, io)?;
1006 let filename = match &path_val {
1007 Value::Str(s) | Value::StringObj(s) => s.clone(),
1008 _ => {
1009 return Err(format!("{fn_name}: argument must be a string (filename)"));
1010 }
1011 };
1012 let script_path = resolve_script_path(&filename)
1013 .ok_or_else(|| format!("{fn_name}: script not found: '{filename}'"))?;
1014 let content = std::fs::read_to_string(&script_path).map_err(|e| {
1015 format!("{fn_name}: cannot read '{}': {e}", script_path.display())
1016 })?;
1017 let depth = RUN_DEPTH.with(|d| d.get());
1018 if depth >= 64 {
1019 return Err(format!(
1020 "{fn_name}: maximum script nesting depth (64) exceeded"
1021 ));
1022 }
1023 RUN_DEPTH.with(|d| d.set(depth + 1));
1024 if let Some(dir) = script_path.parent() {
1027 SCRIPT_DIR_STACK.with(|s| s.borrow_mut().push(dir.to_path_buf()));
1028 }
1029 let run_stmts = parse_stmts(&content).map_err(|e| {
1030 format!("{fn_name}: parse error in '{}': {e}", script_path.display())
1031 })?;
1032
1033 let is_fn_file = !run_stmts.is_empty()
1039 && run_stmts
1040 .iter()
1041 .all(|(s, _, _)| matches!(s, Stmt::FunctionDef { .. }));
1042 let result = if is_fn_file {
1043 let primary_name = match &run_stmts[0].0 {
1044 Stmt::FunctionDef { name, .. } => name.clone(),
1045 _ => unreachable!(),
1046 };
1047 let mut locals: IndexMap<String, Value> = IndexMap::new();
1048 for (stmt, _, _) in &run_stmts {
1049 if let Stmt::FunctionDef {
1050 name,
1051 outputs,
1052 params,
1053 body_source,
1054 doc,
1055 } = stmt
1056 && name != &primary_name
1057 {
1058 locals.insert(
1059 name.clone(),
1060 Value::Function {
1061 outputs: outputs.clone(),
1062 params: params.clone(),
1063 body_source: body_source.clone(),
1064 locals: IndexMap::new(),
1065 doc: doc.clone(),
1066 },
1067 );
1068 }
1069 }
1070 if let Stmt::FunctionDef {
1071 outputs,
1072 params,
1073 body_source,
1074 doc,
1075 ..
1076 } = &run_stmts[0].0
1077 {
1078 env.insert(
1079 primary_name,
1080 Value::Function {
1081 outputs: outputs.clone(),
1082 params: params.clone(),
1083 body_source: body_source.clone(),
1084 locals,
1085 doc: doc.clone(),
1086 },
1087 );
1088 }
1089 Ok(None)
1090 } else {
1091 for (stmt, _, _) in run_stmts.iter() {
1094 if let Stmt::FunctionDef {
1095 name,
1096 outputs,
1097 params,
1098 body_source,
1099 doc,
1100 } = stmt
1101 {
1102 env.insert(
1103 name.clone(),
1104 Value::Function {
1105 outputs: outputs.clone(),
1106 params: params.clone(),
1107 body_source: body_source.clone(),
1108 locals: IndexMap::new(),
1109 doc: doc.clone(),
1110 },
1111 );
1112 }
1113 }
1114 exec_stmts(&run_stmts, env, io, fmt, base, compact)
1115 };
1116 SCRIPT_DIR_STACK.with(|s| s.borrow_mut().pop());
1117 RUN_DEPTH.with(|d| d.set(depth));
1118 match result? {
1121 None => {}
1122 Some(sig) => return Ok(Some(sig)),
1123 }
1124 continue;
1125 }
1126
1127 if let Expr::Call(fn_name, args) = expr
1131 && fn_name == "eval"
1132 && (args.len() == 1 || args.len() == 2)
1133 {
1134 let code_val = eval_with_io(&args[0], env, io)?;
1135 let code_str = match code_val {
1136 Value::Str(s) | Value::StringObj(s) => s,
1137 _ => return Err("eval: argument must be a string".to_string()),
1138 };
1139 let depth = RUN_DEPTH.with(|d| d.get());
1140 if depth >= 64 {
1141 return Err("eval: maximum nesting depth (64) exceeded".to_string());
1142 }
1143 RUN_DEPTH.with(|d| d.set(depth + 1));
1144 let run_result = (|| -> Result<Option<Signal>, String> {
1145 let stmts = parse_stmts(&code_str)
1146 .map_err(|e| format!("eval: parse error: {e}"))?;
1147 exec_stmts(&stmts, env, io, fmt, base, compact)
1148 })();
1149 RUN_DEPTH.with(|d| d.set(depth));
1150 match run_result {
1151 Err(e) if args.len() == 2 => {
1152 set_last_err(&e);
1153 let catch_val = eval_with_io(&args[1], env, io)?;
1154 let catch_str = match catch_val {
1155 Value::Str(s) | Value::StringObj(s) => s,
1156 _ => {
1157 return Err("eval: catch argument must be a string".to_string());
1158 }
1159 };
1160 let catch_stmts = parse_stmts(&catch_str)
1161 .map_err(|e| format!("eval: catch parse error: {e}"))?;
1162 match exec_stmts(&catch_stmts, env, io, fmt, base, compact)? {
1163 None => {}
1164 Some(sig) => return Ok(Some(sig)),
1165 }
1166 }
1167 Err(e) => return Err(e),
1168 Ok(None) => {}
1169 Ok(Some(sig)) => return Ok(Some(sig)),
1170 }
1171 continue;
1172 }
1173
1174 if let Expr::Call(fn_name, args) = expr
1176 && fn_name == "clear"
1177 {
1178 if args.is_empty() {
1179 env.clear();
1180 } else {
1181 for arg in args {
1182 let key = match arg {
1183 Expr::StrLiteral(s) | Expr::StringObjLiteral(s) => s.clone(),
1184 other => match eval_with_io(other, env, io)? {
1185 Value::Str(s) | Value::StringObj(s) => s,
1186 _ => continue,
1187 },
1188 };
1189 env.remove(&key);
1190 }
1191 }
1192 continue;
1193 }
1194
1195 if let Expr::Call(fn_name, args) = expr
1197 && fn_name == "format"
1198 {
1199 let arg = match args.first() {
1200 Some(Expr::StrLiteral(s)) => s.as_str(),
1201 None => "",
1202 _ => "",
1203 };
1204 let new_fmt = match arg {
1205 "" | "short" => FormatMode::Short,
1206 "long" => FormatMode::Long,
1207 "shorte" | "shortE" => FormatMode::ShortE,
1208 "longe" | "longE" => FormatMode::LongE,
1209 "shortg" | "shortG" => FormatMode::ShortG,
1210 "longg" | "longG" => FormatMode::LongG,
1211 "bank" => FormatMode::Bank,
1212 "rat" => FormatMode::Rat,
1213 "hex" => FormatMode::Hex,
1214 "+" => FormatMode::Plus,
1215 "compact" | "loose" => get_display_fmt(),
1216 s => s
1217 .parse::<usize>()
1218 .map(FormatMode::Custom)
1219 .unwrap_or_else(|_| get_display_fmt()),
1220 };
1221 let new_compact = match arg {
1222 "compact" => true,
1223 "loose" => false,
1224 _ => get_display_compact(),
1225 };
1226 set_display_ctx(&new_fmt, base, new_compact);
1227 continue;
1228 }
1229
1230 if let Expr::Call(fn_name, args) = expr
1232 && matches!(fn_name.as_str(), "save" | "load" | "ws" | "wl")
1233 {
1234 let is_save = matches!(fn_name.as_str(), "save" | "ws");
1235 if is_save {
1236 let (path_opt, var_names) = if args.is_empty() {
1237 (None, vec![])
1238 } else {
1239 let path_val = eval_with_io(&args[0], env, io)?;
1240 let path_str = match path_val {
1241 Value::Str(s) | Value::StringObj(s) => s,
1242 _ => return Err("save: path argument must be a string".to_string()),
1243 };
1244 let mut vars: Vec<String> = Vec::new();
1245 for a in &args[1..] {
1246 let v = eval_with_io(a, env, io)?;
1247 match v {
1248 Value::Str(s) | Value::StringObj(s) => vars.push(s),
1249 _ => {
1250 return Err(
1251 "save: variable names must be strings".to_string()
1252 );
1253 }
1254 }
1255 }
1256 (Some(path_str), vars)
1257 };
1258 let result = match &path_opt {
1259 None => {
1260 let home = std::env::var("HOME")
1261 .or_else(|_| std::env::var("USERPROFILE"))
1262 .unwrap_or_default();
1263 let p = std::path::Path::new(&home)
1264 .join(".config")
1265 .join("ccalc")
1266 .join("workspace.toml");
1267 save_workspace(env, &p)
1268 }
1269 Some(p) if var_names.is_empty() => {
1270 save_workspace(env, std::path::Path::new(p))
1271 }
1272 Some(p) => {
1273 let refs: Vec<&str> =
1274 var_names.iter().map(String::as_str).collect();
1275 save_workspace_vars(env, std::path::Path::new(p), &refs)
1276 }
1277 };
1278 if let Err(e) = result {
1279 return Err(format!("save: {e}"));
1280 }
1281 } else {
1282 let loaded = if args.is_empty() {
1284 let home = std::env::var("HOME")
1285 .or_else(|_| std::env::var("USERPROFILE"))
1286 .unwrap_or_default();
1287 let p = std::path::Path::new(&home)
1288 .join(".config")
1289 .join("ccalc")
1290 .join("workspace.toml");
1291 load_workspace(&p)
1292 } else {
1293 let path_val = eval_with_io(&args[0], env, io)?;
1294 let path_str = match path_val {
1295 Value::Str(s) | Value::StringObj(s) => s,
1296 _ => return Err("load: path argument must be a string".to_string()),
1297 };
1298 load_workspace(std::path::Path::new(&path_str))
1299 };
1300 match loaded {
1301 Ok(ws) => env.extend(ws),
1302 Err(e) => return Err(format!("load: {e}")),
1303 }
1304 }
1305 continue;
1306 }
1307
1308 let val = eval_with_io(expr, env, io).map_err(|e| annotate_line(e, *stmt_line))?;
1309 env.insert("ans".to_string(), val.clone());
1310 if !silent && !matches!(val, Value::Void) {
1311 print_value(None, &val, fmt, base, compact);
1312 }
1313 }
1314
1315 Stmt::If {
1316 cond,
1317 body,
1318 elseif_branches,
1319 else_body,
1320 } => {
1321 let cond_val =
1322 eval_with_io(cond, env, io).map_err(|e| annotate_line(e, *stmt_line))?;
1323 let chosen: Option<&[StmtEntry]> = if is_truthy(&cond_val) {
1324 Some(body)
1325 } else {
1326 let mut found = None;
1327 for (ei_cond, ei_body) in elseif_branches {
1328 if is_truthy(
1329 &eval_with_io(ei_cond, env, io)
1330 .map_err(|e| annotate_line(e, *stmt_line))?,
1331 ) {
1332 found = Some(ei_body.as_slice());
1333 break;
1334 }
1335 }
1336 if found.is_none() {
1337 found = else_body.as_deref();
1338 }
1339 found
1340 };
1341 if let Some(body_stmts) = chosen
1342 && let Some(sig) = exec_stmts(body_stmts, env, io, fmt, base, compact)?
1343 {
1344 return Ok(Some(sig));
1345 }
1346 }
1347
1348 Stmt::For {
1349 var,
1350 range_expr,
1351 body,
1352 } => {
1353 let range_val =
1354 eval_with_io(range_expr, env, io).map_err(|e| annotate_line(e, *stmt_line))?;
1355 let iter_cols: Vec<Value> = match range_val {
1356 Value::Scalar(n) => vec![Value::Scalar(n)],
1357 Value::Matrix(m) => {
1358 let nrows = m.nrows();
1359 let ncols = m.ncols();
1360 (0..ncols)
1361 .map(|j| {
1362 if nrows == 1 {
1363 Value::Scalar(m[[0, j]])
1365 } else {
1366 let mut col = Array2::zeros((nrows, 1));
1368 for i in 0..nrows {
1369 col[[i, 0]] = m[[i, j]];
1370 }
1371 Value::Matrix(col)
1372 }
1373 })
1374 .collect()
1375 }
1376 _ => return Err("'for' range must evaluate to a scalar or matrix".to_string()),
1377 };
1378
1379 'for_loop: for col_val in iter_cols {
1380 env.insert(var.clone(), col_val);
1381 match exec_stmts(body, env, io, fmt, base, compact)? {
1382 None => {}
1383 Some(Signal::Break) => break 'for_loop,
1384 Some(Signal::Continue) => continue 'for_loop,
1385 Some(Signal::Return) => return Ok(Some(Signal::Return)),
1386 }
1387 }
1388 }
1389
1390 Stmt::While { cond, body } => loop {
1391 if !is_truthy(
1392 &eval_with_io(cond, env, io).map_err(|e| annotate_line(e, *stmt_line))?,
1393 ) {
1394 break;
1395 }
1396 match exec_stmts(body, env, io, fmt, base, compact)? {
1397 None => {}
1398 Some(Signal::Break) => break,
1399 Some(Signal::Continue) => continue,
1400 Some(Signal::Return) => return Ok(Some(Signal::Return)),
1401 }
1402 },
1403
1404 Stmt::Break => return Ok(Some(Signal::Break)),
1405 Stmt::Continue => return Ok(Some(Signal::Continue)),
1406
1407 Stmt::Switch {
1409 expr,
1410 cases,
1411 otherwise_body,
1412 } => {
1413 let switch_val =
1414 eval_with_io(expr, env, io).map_err(|e| annotate_line(e, *stmt_line))?;
1415 let mut matched = false;
1416 'switch_loop: for (case_exprs, case_body) in cases {
1417 for case_expr in case_exprs {
1418 let case_val = eval_with_io(case_expr, env, io)
1419 .map_err(|e| annotate_line(e, *stmt_line))?;
1420 let is_match = if let Value::Cell(cell_elems) = &case_val {
1423 cell_elems.iter().any(|elem| match (&switch_val, elem) {
1424 (Value::Scalar(a), Value::Scalar(b)) => a == b,
1425 _ => {
1426 let sv = match &switch_val {
1427 Value::Str(s) | Value::StringObj(s) => Some(s.as_str()),
1428 _ => None,
1429 };
1430 let cv = match elem {
1431 Value::Str(s) | Value::StringObj(s) => Some(s.as_str()),
1432 _ => None,
1433 };
1434 matches!((sv, cv), (Some(a), Some(b)) if a == b)
1435 }
1436 })
1437 } else {
1438 match (&switch_val, &case_val) {
1439 (Value::Scalar(a), Value::Scalar(b)) => a == b,
1440 _ => {
1441 let sv = match &switch_val {
1442 Value::Str(s) | Value::StringObj(s) => Some(s.as_str()),
1443 _ => None,
1444 };
1445 let cv = match &case_val {
1446 Value::Str(s) | Value::StringObj(s) => Some(s.as_str()),
1447 _ => None,
1448 };
1449 matches!((sv, cv), (Some(a), Some(b)) if a == b)
1450 }
1451 }
1452 };
1453 if is_match {
1454 if let Some(sig) = exec_stmts(case_body, env, io, fmt, base, compact)? {
1455 return Ok(Some(sig));
1456 }
1457 matched = true;
1458 break 'switch_loop;
1459 }
1460 }
1461 }
1462 if !matched
1463 && let Some(ob) = otherwise_body
1464 && let Some(sig) = exec_stmts(ob, env, io, fmt, base, compact)?
1465 {
1466 return Ok(Some(sig));
1467 }
1468 }
1469
1470 Stmt::DoUntil { body, cond } => loop {
1472 match exec_stmts(body, env, io, fmt, base, compact)? {
1473 Some(Signal::Break) => break,
1474 Some(Signal::Continue) | None => {}
1475 Some(Signal::Return) => return Ok(Some(Signal::Return)),
1476 }
1477 if is_truthy(
1478 &eval_with_io(cond, env, io).map_err(|e| annotate_line(e, *stmt_line))?,
1479 ) {
1480 break;
1481 }
1482 },
1483
1484 Stmt::TryCatch {
1486 try_body,
1487 catch_var,
1488 catch_body,
1489 } => match exec_stmts(try_body, env, io, fmt, base, compact) {
1490 Ok(None) => {}
1491 Ok(Some(sig)) => return Ok(Some(sig)),
1492 Err(msg) => {
1493 set_last_err(&msg);
1494 if let Some(var) = catch_var {
1495 let mut map = IndexMap::new();
1496 map.insert("message".to_string(), Value::Str(strip_near_line(msg)));
1497 env.insert(var.clone(), Value::Struct(map));
1498 }
1499 if let Some(sig) = exec_stmts(catch_body, env, io, fmt, base, compact)? {
1500 return Ok(Some(sig));
1501 }
1502 }
1503 },
1504
1505 Stmt::FunctionDef {
1507 name,
1508 outputs,
1509 params,
1510 body_source,
1511 doc,
1512 } => {
1513 env.insert(
1514 name.clone(),
1515 Value::Function {
1516 outputs: outputs.clone(),
1517 params: params.clone(),
1518 body_source: body_source.clone(),
1519 locals: IndexMap::new(),
1520 doc: doc.clone(),
1521 },
1522 );
1523 }
1524
1525 Stmt::Return => return Ok(Some(Signal::Return)),
1527
1528 Stmt::CellSet(cell_name, idx_expr, val_expr) => {
1530 let cell_len = match env.get(cell_name) {
1532 Some(Value::Cell(v)) => v.len(),
1533 _ => 0,
1534 };
1535 let _owned_end;
1536 let env_end: &Env = if crate::eval::contains_end(idx_expr) {
1537 _owned_end = write_env_with_end(env, cell_len);
1538 &_owned_end
1539 } else {
1540 env
1541 };
1542 let idx = eval_with_io(idx_expr, env_end, io)
1543 .map_err(|e| annotate_line(e, *stmt_line))?;
1544 let rhs =
1545 eval_with_io(val_expr, env, io).map_err(|e| annotate_line(e, *stmt_line))?;
1546 let i = match idx {
1547 Value::Scalar(n) => n as isize,
1548 _ => return Err(format!("{cell_name}{{}}: index must be a scalar integer")),
1549 };
1550 match env.get_mut(cell_name) {
1551 Some(Value::Cell(v)) => {
1552 if i < 1 {
1553 return Err(format!(
1554 "{cell_name}{{}}: index {i} out of range (1..{})",
1555 v.len()
1556 ));
1557 }
1558 let idx = (i - 1) as usize;
1559 if idx >= v.len() {
1561 v.resize(idx + 1, Value::Scalar(0.0));
1562 }
1563 v[idx] = rhs.clone();
1564 }
1565 Some(_) => {
1566 return Err(format!(
1567 "'{cell_name}' is not a cell array; use () for regular indexing"
1568 ));
1569 }
1570 None => {
1571 if i < 1 {
1573 return Err(format!("{cell_name}{{}}: index {i} must be >= 1"));
1574 }
1575 let idx = (i - 1) as usize;
1576 let mut v = vec![Value::Scalar(0.0); idx + 1];
1577 v[idx] = rhs.clone();
1578 env.insert(cell_name.clone(), Value::Cell(v));
1579 }
1580 }
1581 if !silent && let Some(val) = env.get(cell_name) {
1582 print_value(Some(cell_name), val, fmt, base, compact);
1583 }
1584 }
1585
1586 Stmt::FieldSet(base_name, path, rhs_expr) => {
1588 let rhs =
1589 eval_with_io(rhs_expr, env, io).map_err(|e| annotate_line(e, *stmt_line))?;
1590 let root = match env.remove(base_name) {
1591 Some(Value::Struct(m)) => m,
1592 None => IndexMap::new(),
1593 Some(other) => {
1594 env.insert(base_name.clone(), other);
1595 return Err(format!("'{base_name}' is not a struct"));
1596 }
1597 };
1598 let updated = set_nested(root, path, rhs)?;
1599 let struct_val = Value::Struct(updated);
1600 if !silent {
1601 print_value(Some(base_name), &struct_val, fmt, base, compact);
1602 }
1603 env.insert(base_name.clone(), struct_val);
1604 }
1605
1606 Stmt::StructArrayFieldSet(base_name, idx_expr, path, rhs_expr) => {
1608 let rhs =
1609 eval_with_io(rhs_expr, env, io).map_err(|e| annotate_line(e, *stmt_line))?;
1610 let idx_val =
1611 eval_with_io(idx_expr, env, io).map_err(|e| annotate_line(e, *stmt_line))?;
1612 let idx = match &idx_val {
1613 Value::Scalar(n) => {
1614 let i = *n as isize;
1615 if i < 1 {
1616 return Err(format!(
1617 "Struct array index must be a positive integer, got {n}"
1618 ));
1619 }
1620 i as usize
1621 }
1622 _ => return Err("Struct array index must be a scalar integer".to_string()),
1623 };
1624 let mut arr: Vec<IndexMap<String, Value>> = match env.remove(base_name) {
1626 Some(Value::StructArray(v)) => v,
1627 Some(Value::Struct(m)) => vec![m],
1629 None => Vec::new(),
1630 Some(other) => {
1631 env.insert(base_name.clone(), other);
1632 return Err(format!("'{base_name}' is not a struct array"));
1633 }
1634 };
1635 while arr.len() < idx {
1637 arr.push(IndexMap::new());
1638 }
1639 let elem = arr[idx - 1].clone();
1641 let updated_elem = set_nested(elem, path, rhs)?;
1642 arr[idx - 1] = updated_elem;
1643 let arr_val = Value::StructArray(arr);
1644 if !silent {
1645 print_value(Some(base_name), &arr_val, fmt, base, compact);
1646 }
1647 env.insert(base_name.clone(), arr_val);
1648 }
1649
1650 Stmt::IndexSet {
1652 name,
1653 indices,
1654 value,
1655 } => {
1656 let rhs = eval_with_io(value, env, io).map_err(|e| annotate_line(e, *stmt_line))?;
1657 if is_persistent(name) {
1660 let func = current_func_name();
1661 if let Some(fresh) = persistent_load(&func, name) {
1662 env.insert(name.clone(), fresh);
1663 }
1664 }
1665 exec_index_set(name, indices, rhs, env, io)?;
1666 if is_persistent(name)
1668 && let Some(val) = env.get(name)
1669 {
1670 persistent_save(¤t_func_name(), name, val.clone());
1671 }
1672 if !silent && let Some(val) = env.get(name) {
1673 print_value(Some(name), val, fmt, base, compact);
1674 }
1675 }
1676
1677 Stmt::MultiAssign { targets, expr } => {
1679 set_nargout(targets.len());
1680 let val = eval_with_io(expr, env, io).map_err(|e| annotate_line(e, *stmt_line))?;
1681 let vals: Vec<Value> = match val {
1682 Value::Tuple(v) => v,
1683 other => vec![other],
1684 };
1685 for (i, target) in targets.iter().enumerate() {
1686 if target == "~" {
1687 continue; }
1689 let v = vals.get(i).cloned().unwrap_or(Value::Void);
1690 env.insert(target.clone(), v.clone());
1691 if !silent && !matches!(v, Value::Void) {
1692 print_value(Some(target), &v, fmt, base, compact);
1693 }
1694 }
1695 }
1696 }
1697 }
1698 global_refresh_into_env(env);
1702 Ok(None)
1703}
1704
1705enum WriteIdx {
1709 All,
1711 Positions(Vec<usize>),
1713}
1714
1715fn resolve_write_dim(
1721 expr: &crate::eval::Expr,
1722 dim_size: usize,
1723 env: &Env,
1724 io: &mut IoContext,
1725) -> Result<WriteIdx, String> {
1726 if matches!(expr, crate::eval::Expr::Colon) {
1727 return Ok(WriteIdx::All);
1728 }
1729 let val = eval_with_io(expr, env, io)?;
1730 let floats: Vec<f64> = match val {
1731 Value::Scalar(n) => vec![n],
1732 Value::Complex(re, im) => {
1733 if im != 0.0 {
1734 return Err("Index must be real, not complex".to_string());
1735 }
1736 vec![re]
1737 }
1738 Value::Matrix(m) => {
1739 let total = m.nrows() * m.ncols();
1740 if m.nrows() > 1 && m.ncols() > 1 && total != dim_size {
1741 return Err("Index must be a scalar or vector, not a 2-D matrix".to_string());
1742 }
1743 if m.nrows() > 1 && m.ncols() > 1 {
1745 let mut v = Vec::with_capacity(total);
1746 for col in 0..m.ncols() {
1747 for row in 0..m.nrows() {
1748 v.push(m[[row, col]]);
1749 }
1750 }
1751 v
1752 } else {
1753 m.iter().copied().collect()
1754 }
1755 }
1756 _ => return Err("Index must be numeric".to_string()),
1757 };
1758 if dim_size > 0 && floats.len() == dim_size && floats.iter().all(|&f| f == 0.0 || f == 1.0) {
1760 let positions: Vec<usize> = floats
1761 .iter()
1762 .enumerate()
1763 .filter(|&(_, &f)| f == 1.0)
1764 .map(|(i, _)| i)
1765 .collect();
1766 return Ok(WriteIdx::Positions(positions));
1767 }
1768 let positions: Result<Vec<usize>, String> = floats
1770 .iter()
1771 .map(|&n| {
1772 let i = n.round() as i64;
1773 if i < 1 {
1774 return Err(format!("Index {i} must be >= 1"));
1775 }
1776 Ok(i as usize - 1)
1777 })
1778 .collect();
1779 Ok(WriteIdx::Positions(positions?))
1780}
1781
1782fn write_env_with_end(env: &Env, dim_size: usize) -> Env {
1784 let mut e = env.clone();
1785 e.insert("end".to_string(), Value::Scalar(dim_size as f64));
1786 e
1787}
1788
1789fn exec_index_set(
1797 name: &str,
1798 indices: &[crate::eval::Expr],
1799 rhs: Value,
1800 env: &mut Env,
1801 io: &mut IoContext,
1802) -> Result<(), String> {
1803 let needs_complex = matches!(env.get(name), Some(Value::ComplexMatrix(_)))
1805 || matches!(&rhs, Value::Complex(_, _) | Value::ComplexMatrix(_));
1806
1807 if needs_complex {
1808 return exec_index_set_complex(name, indices, rhs, env, io);
1809 }
1810
1811 match indices.len() {
1812 1 => {
1813 let (total, nrows_hint, ncols_hint) = match env.get(name) {
1815 Some(Value::Matrix(m)) => (m.nrows() * m.ncols(), m.nrows(), m.ncols()),
1816 Some(Value::Scalar(_)) => (1, 1, 1),
1817 None | Some(Value::Void) => (0, 0, 0),
1818 Some(_) => {
1819 return Err(format!(
1820 "'{name}' is not a matrix; cannot use () indexed assignment"
1821 ));
1822 }
1823 };
1824
1825 let widx = {
1828 let _owned_end;
1829 let env_end: &Env = if crate::eval::contains_end(&indices[0]) {
1830 _owned_end = write_env_with_end(env, total);
1831 &_owned_end
1832 } else {
1833 env
1834 };
1835 resolve_write_dim(&indices[0], total, env_end, io)?
1836 };
1837
1838 let positions: Vec<usize> = match widx {
1840 WriteIdx::All => (0..total).collect(),
1841 WriteIdx::Positions(p) => p,
1842 };
1843
1844 let rhs_vals: Vec<f64> = match &rhs {
1846 Value::Scalar(n) => vec![*n; positions.len()],
1847 Value::Matrix(m) => {
1848 let flat: Vec<f64> = m.iter().copied().collect();
1849 if flat.len() != positions.len() {
1850 return Err(format!(
1851 "Assignment dimension mismatch: {} positions but {} values",
1852 positions.len(),
1853 flat.len()
1854 ));
1855 }
1856 flat
1857 }
1858 _ => {
1859 return Err(
1860 "Indexed assignment: RHS must be a numeric scalar or matrix".to_string()
1861 );
1862 }
1863 };
1864
1865 let required = positions.iter().copied().max().map(|m| m + 1).unwrap_or(0);
1867 let required = required.max(total);
1868
1869 let (out_rows, out_cols) = if nrows_hint == 0 || ncols_hint == 0 {
1871 (1, required)
1873 } else if nrows_hint == 1 {
1874 (1, required)
1875 } else if ncols_hint == 1 {
1876 (required, 1)
1877 } else if required > total {
1878 return Err("Cannot grow a 2-D matrix with linear indexing".to_string());
1879 } else {
1880 (nrows_hint, ncols_hint)
1881 };
1882
1883 let mut mat = match env.remove(name) {
1886 Some(Value::Matrix(m)) => m,
1887 Some(Value::Scalar(n)) => Array2::from_elem((1, 1), n),
1888 None | Some(Value::Void) => Array2::zeros((0, 0)),
1889 Some(other) => {
1890 env.insert(name.to_string(), other);
1891 return Err(format!(
1892 "'{name}' is not a matrix; cannot use () indexed assignment"
1893 ));
1894 }
1895 };
1896
1897 if required > total || out_rows != mat.nrows() || out_cols != mat.ncols() {
1899 let mut new_mat = Array2::<f64>::zeros((out_rows, out_cols));
1900 for old_p in 0..total {
1901 let old_row = old_p % mat.nrows().max(1);
1902 let old_col = old_p / mat.nrows().max(1);
1903 let new_row = old_p % out_rows;
1904 let new_col = old_p / out_rows;
1905 if old_row < mat.nrows() && old_col < mat.ncols() {
1906 new_mat[[new_row, new_col]] = mat[[old_row, old_col]];
1907 }
1908 }
1909 mat = new_mat;
1910 }
1911
1912 for (&pos, &val) in positions.iter().zip(rhs_vals.iter()) {
1914 let row = pos % mat.nrows();
1915 let col = pos / mat.nrows();
1916 mat[[row, col]] = val;
1917 }
1918
1919 let result = if mat.nrows() == 1 && mat.ncols() == 1 {
1920 Value::Scalar(mat[[0, 0]])
1921 } else {
1922 Value::Matrix(mat)
1923 };
1924 env.insert(name.to_string(), result);
1925 }
1926 2 => {
1927 let (nrows, ncols) = match env.get(name) {
1929 Some(Value::Matrix(m)) => (m.nrows(), m.ncols()),
1930 Some(Value::Scalar(_)) => (1, 1),
1931 None | Some(Value::Void) => (0, 0),
1932 Some(_) => {
1933 return Err(format!(
1934 "'{name}' is not a matrix; cannot use () indexed assignment"
1935 ));
1936 }
1937 };
1938
1939 let (ridx, cidx) = {
1941 let _owned_r;
1942 let env_r: &Env = if crate::eval::contains_end(&indices[0]) {
1943 _owned_r = write_env_with_end(env, nrows);
1944 &_owned_r
1945 } else {
1946 env
1947 };
1948 let _owned_c;
1949 let env_c: &Env = if crate::eval::contains_end(&indices[1]) {
1950 _owned_c = write_env_with_end(env, ncols);
1951 &_owned_c
1952 } else {
1953 env
1954 };
1955 (
1956 resolve_write_dim(&indices[0], nrows, env_r, io)?,
1957 resolve_write_dim(&indices[1], ncols, env_c, io)?,
1958 )
1959 };
1960
1961 let rows: Vec<usize> = match ridx {
1962 WriteIdx::All => (0..nrows.max(1)).collect(),
1963 WriteIdx::Positions(p) => p,
1964 };
1965 let cols: Vec<usize> = match cidx {
1966 WriteIdx::All => (0..ncols.max(1)).collect(),
1967 WriteIdx::Positions(p) => p,
1968 };
1969
1970 let req_rows = rows
1971 .iter()
1972 .copied()
1973 .max()
1974 .map(|m| m + 1)
1975 .unwrap_or(0)
1976 .max(nrows);
1977 let req_cols = cols
1978 .iter()
1979 .copied()
1980 .max()
1981 .map(|m| m + 1)
1982 .unwrap_or(0)
1983 .max(ncols);
1984
1985 let n_sel = rows.len() * cols.len();
1987 let rhs_vals: Vec<f64> = match &rhs {
1988 Value::Scalar(n) => vec![*n; n_sel],
1989 Value::Matrix(m) => {
1990 let flat: Vec<f64> = m.iter().copied().collect();
1991 if flat.len() != n_sel {
1992 return Err(format!(
1993 "Assignment dimension mismatch: {}×{} = {} positions but {} values",
1994 rows.len(),
1995 cols.len(),
1996 n_sel,
1997 flat.len()
1998 ));
1999 }
2000 flat
2001 }
2002 _ => {
2003 return Err(
2004 "Indexed assignment: RHS must be a numeric scalar or matrix".to_string()
2005 );
2006 }
2007 };
2008
2009 let mut mat = match env.remove(name) {
2011 Some(Value::Matrix(m)) => m,
2012 Some(Value::Scalar(n)) => Array2::from_elem((1, 1), n),
2013 None | Some(Value::Void) => Array2::zeros((0, 0)),
2014 Some(other) => {
2015 env.insert(name.to_string(), other);
2016 return Err(format!(
2017 "'{name}' is not a matrix; cannot use () indexed assignment"
2018 ));
2019 }
2020 };
2021
2022 if req_rows != nrows || req_cols != ncols {
2024 let mut new_mat = Array2::<f64>::zeros((req_rows, req_cols));
2025 for r in 0..nrows {
2026 for c in 0..ncols {
2027 new_mat[[r, c]] = mat[[r, c]];
2028 }
2029 }
2030 mat = new_mat;
2031 }
2032
2033 let mut k = 0;
2035 for &r in &rows {
2036 for &c in &cols {
2037 mat[[r, c]] = rhs_vals[k];
2038 k += 1;
2039 }
2040 }
2041
2042 let result = if mat.nrows() == 1 && mat.ncols() == 1 {
2043 Value::Scalar(mat[[0, 0]])
2044 } else {
2045 Value::Matrix(mat)
2046 };
2047 env.insert(name.to_string(), result);
2048 }
2049 _ => return Err("Indexed assignment supports at most 2 indices".to_string()),
2050 }
2051 Ok(())
2052}
2053
2054fn exec_index_set_complex(
2059 name: &str,
2060 indices: &[crate::eval::Expr],
2061 rhs: Value,
2062 env: &mut Env,
2063 io: &mut IoContext,
2064) -> Result<(), String> {
2065 fn rhs_to_complex(rhs: &Value, n_slots: usize) -> Result<Vec<Complex<f64>>, String> {
2067 match rhs {
2068 Value::Scalar(n) => Ok(vec![Complex::new(*n, 0.0); n_slots]),
2069 Value::Complex(re, im) => Ok(vec![Complex::new(*re, *im); n_slots]),
2070 Value::Matrix(m) => {
2071 let flat: Vec<Complex<f64>> = m.iter().map(|&x| Complex::new(x, 0.0)).collect();
2072 if flat.len() != n_slots {
2073 return Err(format!(
2074 "Assignment dimension mismatch: {} positions but {} values",
2075 n_slots,
2076 flat.len()
2077 ));
2078 }
2079 Ok(flat)
2080 }
2081 Value::ComplexMatrix(m) => {
2082 let flat: Vec<Complex<f64>> = m.iter().copied().collect();
2083 if flat.len() != n_slots {
2084 return Err(format!(
2085 "Assignment dimension mismatch: {} positions but {} values",
2086 n_slots,
2087 flat.len()
2088 ));
2089 }
2090 Ok(flat)
2091 }
2092 _ => Err("Indexed assignment: RHS must be a numeric value".to_string()),
2093 }
2094 }
2095
2096 fn take_as_complex(name: &str, env: &mut Env) -> Result<Array2<Complex<f64>>, String> {
2098 match env.remove(name) {
2099 Some(Value::ComplexMatrix(m)) => Ok(m),
2100 Some(Value::Matrix(m)) => Ok(m.mapv(|x| Complex::new(x, 0.0))),
2101 Some(Value::Scalar(n)) => Ok(Array2::from_elem((1, 1), Complex::new(n, 0.0))),
2102 None | Some(Value::Void) => Ok(Array2::zeros((0, 0))),
2103 Some(other) => {
2104 env.insert(name.to_string(), other);
2105 Err(format!(
2106 "'{name}' is not a matrix; cannot use () indexed assignment"
2107 ))
2108 }
2109 }
2110 }
2111
2112 match indices.len() {
2113 1 => {
2114 let (total, nrows_hint, ncols_hint) = match env.get(name) {
2115 Some(Value::ComplexMatrix(m)) => (m.nrows() * m.ncols(), m.nrows(), m.ncols()),
2116 Some(Value::Matrix(m)) => (m.nrows() * m.ncols(), m.nrows(), m.ncols()),
2117 Some(Value::Scalar(_)) => (1, 1, 1),
2118 None | Some(Value::Void) => (0, 0, 0),
2119 Some(_) => {
2120 return Err(format!(
2121 "'{name}' is not a matrix; cannot use () indexed assignment"
2122 ));
2123 }
2124 };
2125
2126 let widx = {
2127 let _owned_end;
2128 let env_end: &Env = if crate::eval::contains_end(&indices[0]) {
2129 _owned_end = write_env_with_end(env, total);
2130 &_owned_end
2131 } else {
2132 env
2133 };
2134 resolve_write_dim(&indices[0], total, env_end, io)?
2135 };
2136
2137 let positions: Vec<usize> = match widx {
2138 WriteIdx::All => (0..total).collect(),
2139 WriteIdx::Positions(p) => p,
2140 };
2141
2142 let rhs_vals = rhs_to_complex(&rhs, positions.len())?;
2143
2144 let required = positions.iter().copied().max().map(|m| m + 1).unwrap_or(0);
2145 let required = required.max(total);
2146
2147 let (out_rows, out_cols) = if nrows_hint == 0 || ncols_hint == 0 || nrows_hint == 1 {
2148 (1, required)
2149 } else if ncols_hint == 1 {
2150 (required, 1)
2151 } else if required > total {
2152 return Err("Cannot grow a 2-D matrix with linear indexing".to_string());
2153 } else {
2154 (nrows_hint, ncols_hint)
2155 };
2156
2157 let mut mat = take_as_complex(name, env)?;
2158
2159 if required > total || out_rows != mat.nrows() || out_cols != mat.ncols() {
2160 let mut new_mat = Array2::<Complex<f64>>::zeros((out_rows, out_cols));
2161 for old_p in 0..total {
2162 let old_row = old_p % mat.nrows().max(1);
2163 let old_col = old_p / mat.nrows().max(1);
2164 let new_row = old_p % out_rows;
2165 let new_col = old_p / out_rows;
2166 if old_row < mat.nrows() && old_col < mat.ncols() {
2167 new_mat[[new_row, new_col]] = mat[[old_row, old_col]];
2168 }
2169 }
2170 mat = new_mat;
2171 }
2172
2173 for (&pos, &val) in positions.iter().zip(rhs_vals.iter()) {
2174 let row = pos % mat.nrows();
2175 let col = pos / mat.nrows();
2176 mat[[row, col]] = val;
2177 }
2178
2179 env.insert(name.to_string(), Value::ComplexMatrix(mat));
2180 }
2181 2 => {
2182 let (nrows, ncols) = match env.get(name) {
2183 Some(Value::ComplexMatrix(m)) => (m.nrows(), m.ncols()),
2184 Some(Value::Matrix(m)) => (m.nrows(), m.ncols()),
2185 Some(Value::Scalar(_)) => (1, 1),
2186 None | Some(Value::Void) => (0, 0),
2187 Some(_) => {
2188 return Err(format!(
2189 "'{name}' is not a matrix; cannot use () indexed assignment"
2190 ));
2191 }
2192 };
2193
2194 let (ridx, cidx) = {
2195 let _owned_r;
2196 let env_r: &Env = if crate::eval::contains_end(&indices[0]) {
2197 _owned_r = write_env_with_end(env, nrows);
2198 &_owned_r
2199 } else {
2200 env
2201 };
2202 let _owned_c;
2203 let env_c: &Env = if crate::eval::contains_end(&indices[1]) {
2204 _owned_c = write_env_with_end(env, ncols);
2205 &_owned_c
2206 } else {
2207 env
2208 };
2209 (
2210 resolve_write_dim(&indices[0], nrows, env_r, io)?,
2211 resolve_write_dim(&indices[1], ncols, env_c, io)?,
2212 )
2213 };
2214
2215 let rows: Vec<usize> = match ridx {
2216 WriteIdx::All => (0..nrows.max(1)).collect(),
2217 WriteIdx::Positions(p) => p,
2218 };
2219 let cols: Vec<usize> = match cidx {
2220 WriteIdx::All => (0..ncols.max(1)).collect(),
2221 WriteIdx::Positions(p) => p,
2222 };
2223
2224 let req_rows = rows
2225 .iter()
2226 .copied()
2227 .max()
2228 .map(|m| m + 1)
2229 .unwrap_or(0)
2230 .max(nrows);
2231 let req_cols = cols
2232 .iter()
2233 .copied()
2234 .max()
2235 .map(|m| m + 1)
2236 .unwrap_or(0)
2237 .max(ncols);
2238
2239 let n_sel = rows.len() * cols.len();
2240 let rhs_vals = rhs_to_complex(&rhs, n_sel)?;
2241
2242 let mut mat = take_as_complex(name, env)?;
2243
2244 if req_rows != nrows || req_cols != ncols {
2245 let mut new_mat = Array2::<Complex<f64>>::zeros((req_rows, req_cols));
2246 for r in 0..nrows {
2247 for c in 0..ncols {
2248 if r < mat.nrows() && c < mat.ncols() {
2249 new_mat[[r, c]] = mat[[r, c]];
2250 }
2251 }
2252 }
2253 mat = new_mat;
2254 }
2255
2256 let mut k = 0;
2257 for &r in &rows {
2258 for &c in &cols {
2259 mat[[r, c]] = rhs_vals[k];
2260 k += 1;
2261 }
2262 }
2263
2264 env.insert(name.to_string(), Value::ComplexMatrix(mat));
2265 }
2266 _ => return Err("Indexed assignment supports at most 2 indices".to_string()),
2267 }
2268 Ok(())
2269}