r/interpreter/
assignment.rs

1//! Assignment and replacement semantics: `<-`, `<<-`, `->`, `x[i] <- v`,
2//! `x[[i]] <- v`, `x$name <- v`, and replacement functions like `names<-`.
3
4use derive_more::{Display, Error};
5
6use crate::interpreter::environment::Environment;
7use crate::interpreter::value::*;
8use crate::interpreter::Interpreter;
9use crate::parser::ast::{Arg, AssignOp, Expr};
10
11// region: AssignmentError
12
13/// Structured error type for assignment operations.
14#[derive(Debug, Display, Error)]
15pub enum AssignmentError {
16    #[display("invalid assignment target")]
17    InvalidTarget,
18
19    #[display("invalid index")]
20    InvalidIndex,
21
22    #[display("replacement value must be a vector")]
23    InvalidReplacementValue,
24
25    #[display("object is not subsettable")]
26    NotSubsettable,
27}
28
29impl From<AssignmentError> for RError {
30    fn from(e: AssignmentError) -> Self {
31        let kind = match &e {
32            AssignmentError::InvalidTarget => RErrorKind::Other,
33            AssignmentError::InvalidIndex | AssignmentError::NotSubsettable => RErrorKind::Index,
34            AssignmentError::InvalidReplacementValue => RErrorKind::Type,
35        };
36        RError::from_source(kind, e)
37    }
38}
39
40impl From<AssignmentError> for RFlow {
41    fn from(e: AssignmentError) -> Self {
42        RFlow::Error(RError::from(e))
43    }
44}
45
46// endregion
47
48// region: type-preserving replacement
49
50/// Replace elements in a vector at 1-based indices, preserving type when possible.
51/// If the replacement value's type differs, coerces both to the common type.
52fn replace_elements(
53    target: &Vector,
54    indices: &[Option<i64>],
55    replacement: &Vector,
56    max_idx: usize,
57) -> Vector {
58    macro_rules! replace_typed_option {
59        ($target_vals:expr, $repl_vals:expr, $variant:ident) => {{
60            let mut result = $target_vals.to_vec();
61            while result.len() < max_idx {
62                result.push(Default::default());
63            }
64            for (j, idx) in indices.iter().enumerate() {
65                if let Some(i) = idx {
66                    let i = usize::try_from(*i).unwrap_or(0);
67                    if i > 0 && i <= result.len() {
68                        result[i - 1] = $repl_vals
69                            .get(j % $repl_vals.len())
70                            .cloned()
71                            .unwrap_or_default();
72                    }
73                }
74            }
75            Vector::$variant(result.into())
76        }};
77    }
78
79    macro_rules! replace_typed_buffer {
80        ($target_vals:expr, $repl_vals:expr, $variant:ident) => {{
81            let mut result = $target_vals.clone();
82            // Extend to max_idx if needed
83            while result.len() < max_idx {
84                result.push(None);
85            }
86            for (j, idx) in indices.iter().enumerate() {
87                if let Some(i) = idx {
88                    let i = usize::try_from(*i).unwrap_or(0);
89                    if i > 0 && i <= result.len() {
90                        result.set(i - 1, $repl_vals.get_opt(j % $repl_vals.len()));
91                    }
92                }
93            }
94            Vector::$variant(result)
95        }};
96    }
97
98    // Same-type fast path
99    match (target, replacement) {
100        (Vector::Integer(tv), Vector::Integer(rv)) => replace_typed_buffer!(tv, rv, Integer),
101        (Vector::Double(tv), Vector::Double(rv)) => replace_typed_buffer!(tv, rv, Double),
102        (Vector::Character(tv), Vector::Character(rv)) => replace_typed_option!(tv, rv, Character),
103        (Vector::Logical(tv), Vector::Logical(rv)) => replace_typed_option!(tv, rv, Logical),
104        (Vector::Complex(tv), Vector::Complex(rv)) => replace_typed_option!(tv, rv, Complex),
105        (Vector::Raw(tv), Vector::Raw(rv)) => {
106            let mut result = tv.to_vec();
107            while result.len() < max_idx {
108                result.push(0);
109            }
110            for (j, idx) in indices.iter().enumerate() {
111                if let Some(i) = idx {
112                    let i = usize::try_from(*i).unwrap_or(0);
113                    if i > 0 && i <= result.len() {
114                        result[i - 1] = rv.get(j % rv.len()).copied().unwrap_or(0);
115                    }
116                }
117            }
118            Vector::Raw(result)
119        }
120        // Type mismatch — coerce both to doubles
121        _ => {
122            let mut result = target.to_doubles();
123            let repl = replacement.to_doubles();
124            while result.len() < max_idx {
125                result.push(None);
126            }
127            for (j, idx) in indices.iter().enumerate() {
128                if let Some(i) = idx {
129                    let i = usize::try_from(*i).unwrap_or(0);
130                    if i > 0 && i <= result.len() {
131                        result[i - 1] = repl
132                            .get(j % repl.len())
133                            .copied()
134                            .flatten()
135                            .map(Some)
136                            .unwrap_or(None);
137                    }
138                }
139            }
140            Vector::Double(result.into())
141        }
142    }
143}
144
145// endregion
146
147// region: Interpreter delegation
148
149impl Interpreter {
150    pub(super) fn eval_assign(
151        &self,
152        op: &AssignOp,
153        target: &Expr,
154        val: RValue,
155        env: &Environment,
156    ) -> Result<RValue, RFlow> {
157        eval_assign(self, op, target, val, env)
158    }
159}
160
161// endregion
162
163// region: eval_assign
164
165fn eval_assign(
166    interp: &Interpreter,
167    op: &AssignOp,
168    target: &Expr,
169    val: RValue,
170    env: &Environment,
171) -> Result<RValue, RFlow> {
172    match target {
173        Expr::Symbol(name) => {
174            match op {
175                AssignOp::SuperAssign | AssignOp::RightSuperAssign => {
176                    env.set_super(name.clone(), val.clone());
177                }
178                _ => {
179                    env.set(name.clone(), val.clone());
180                }
181            }
182            Ok(val)
183        }
184        // Assignment to index: x[i] <- val or x[i] <<- val
185        Expr::Index { object, indices } => eval_index_assign(interp, op, object, indices, val, env),
186        Expr::IndexDouble { object, indices } => {
187            eval_index_double_assign(interp, op, object, indices, val, env)
188        }
189        Expr::Dollar { object, member } => eval_dollar_assign(interp, op, object, member, val, env),
190        // Handle function calls on left side like names(x) <- val, attr(x, "which") <- val
191        Expr::Call {
192            func,
193            args: call_args,
194            ..
195        } => {
196            if let Expr::Symbol(fname) = func.as_ref() {
197                let replacement_fn = format!("{}<-", fname);
198                if let Some(first_arg) = call_args.first() {
199                    if let Some(ref val_expr) = first_arg.value {
200                        let obj = interp.eval_in(val_expr, env)?;
201                        if let Some(f) = env.get(&replacement_fn) {
202                            // Evaluate extra args (e.g. "which" in attr(x, "which") <- val)
203                            let mut positional = vec![obj];
204                            for arg in &call_args[1..] {
205                                if let Some(ref v) = arg.value {
206                                    positional.push(interp.eval_in(v, env)?);
207                                }
208                            }
209                            positional.push(val.clone());
210                            let result = interp.call_function(&f, &positional, &[], env)?;
211                            if let Expr::Symbol(var_name) = val_expr {
212                                env.set(var_name.clone(), result);
213                            }
214                            return Ok(val);
215                        }
216                    }
217                }
218            }
219            Err(AssignmentError::InvalidTarget.into())
220        }
221        // In R, "name" <- value creates a binding named "name"
222        Expr::String(name) => {
223            match op {
224                AssignOp::SuperAssign | AssignOp::RightSuperAssign => {
225                    env.set_super(name.clone(), val.clone());
226                }
227                _ => {
228                    env.set(name.clone(), val.clone());
229                }
230            }
231            Ok(val)
232        }
233        _ => Err(AssignmentError::InvalidTarget.into()),
234    }
235}
236
237// endregion
238
239// region: index assignment (x[i] <- val)
240
241// env_assign removed — all callers now use eval_assign for chained replacement support.
242
243fn eval_index_assign(
244    interp: &Interpreter,
245    op: &AssignOp,
246    object: &Expr,
247    indices: &[Arg],
248    val: RValue,
249    env: &Environment,
250) -> Result<RValue, RFlow> {
251    // Evaluate the target object expression to get its current value.
252    // This supports chained targets like `body(f)[2] <- val` where
253    // `object` is `Call(body, f)`, not just a symbol.
254    let mut obj = interp.eval_in(object, env)?;
255
256    if indices.is_empty() {
257        eval_assign(interp, op, object, val.clone(), env)?;
258        return Ok(val);
259    }
260
261    let idx_val = if let Some(val_expr) = &indices[0].value {
262        interp.eval_in(val_expr, env)?
263    } else {
264        return Ok(val);
265    };
266
267    match &mut obj {
268        RValue::Vector(v) => {
269            let idx_ints = match &idx_val {
270                RValue::Vector(iv) => iv.to_integers(),
271                _ => return Err(AssignmentError::InvalidIndex.into()),
272            };
273
274            let val_vec = match &val {
275                RValue::Vector(vv) => vv,
276                _ => {
277                    return Err(AssignmentError::InvalidReplacementValue.into());
278                }
279            };
280
281            // Determine max index to know if we need to extend
282            let max_idx = idx_ints
283                .iter()
284                .filter_map(|x| x.and_then(|i| usize::try_from(i).ok()))
285                .max()
286                .unwrap_or(0);
287
288            let new_vec = replace_elements(&v.inner, &idx_ints, &val_vec.inner, max_idx);
289
290            // Preserve attributes (dim, dimnames, class, names, etc.)
291            let mut rv = RVector::from(new_vec);
292            if let Some(attrs) = &v.attrs {
293                rv.attrs = Some(attrs.clone());
294            }
295            eval_assign(interp, op, object, RValue::Vector(rv), env)?;
296            Ok(val)
297        }
298        RValue::List(list) => {
299            match &idx_val {
300                RValue::Vector(rv) if matches!(rv.inner, Vector::Character(_)) => {
301                    let Vector::Character(names) = &rv.inner else {
302                        unreachable!()
303                    };
304                    if let Some(Some(name)) = names.first() {
305                        if let Some(entry) = list
306                            .values
307                            .iter_mut()
308                            .find(|(n, _)| n.as_ref() == Some(name))
309                        {
310                            entry.1 = val.clone();
311                        } else {
312                            list.values.push((Some(name.clone()), val.clone()));
313                        }
314                    }
315                }
316                RValue::Vector(iv) => {
317                    let i = usize::try_from(iv.as_integer_scalar().unwrap_or(0)).unwrap_or(0);
318                    if i > 0 && i <= list.values.len() {
319                        list.values[i - 1].1 = val.clone();
320                    }
321                }
322                _ => {}
323            }
324            eval_assign(interp, op, object, obj, env)?;
325            Ok(val)
326        }
327        RValue::Null => {
328            match &idx_val {
329                RValue::Vector(rv) if matches!(rv.inner, Vector::Character(_)) => {
330                    let Vector::Character(names) = &rv.inner else {
331                        unreachable!()
332                    };
333                    let mut list = RList::new(vec![]);
334                    if let Some(Some(name)) = names.first() {
335                        list.values.push((Some(name.clone()), val.clone()));
336                    }
337                    eval_assign(interp, op, object, RValue::List(list), env)?;
338                }
339                _ => {
340                    let idx = match &idx_val {
341                        RValue::Vector(iv) => {
342                            usize::try_from(iv.as_integer_scalar().unwrap_or(0)).unwrap_or(0)
343                        }
344                        _ => 0,
345                    };
346                    let mut doubles = vec![None; idx];
347                    if idx > 0 {
348                        if let RValue::Vector(vv) = &val {
349                            doubles[idx - 1] = vv.to_doubles().into_iter().next().flatten();
350                        }
351                    }
352                    eval_assign(
353                        interp,
354                        op,
355                        object,
356                        RValue::vec(Vector::Double(doubles.into())),
357                        env,
358                    )?;
359                }
360            }
361            Ok(val)
362        }
363        _ => Err(AssignmentError::NotSubsettable.into()),
364    }
365}
366
367// endregion
368
369// region: double-bracket assignment (x[[i]] <- val)
370
371fn eval_index_double_assign(
372    interp: &Interpreter,
373    op: &AssignOp,
374    object: &Expr,
375    indices: &[Arg],
376    val: RValue,
377    env: &Environment,
378) -> Result<RValue, RFlow> {
379    // Evaluate the target object — supports chained targets like body(f)[[2]] <- val
380    let mut obj = match interp.eval_in(object, env) {
381        Ok(v) => v,
382        Err(_) => RValue::List(RList::new(vec![])),
383    };
384    let idx_val = if let Some(val_expr) = &indices[0].value {
385        interp.eval_in(val_expr, env)?
386    } else {
387        return Ok(val);
388    };
389
390    match &mut obj {
391        RValue::List(list) => {
392            match &idx_val {
393                RValue::Vector(rv) if matches!(rv.inner, Vector::Character(_)) => {
394                    let Vector::Character(names) = &rv.inner else {
395                        unreachable!()
396                    };
397                    if let Some(Some(name)) = names.first() {
398                        if let Some(entry) = list
399                            .values
400                            .iter_mut()
401                            .find(|(n, _)| n.as_ref() == Some(name))
402                        {
403                            entry.1 = val.clone();
404                        } else {
405                            list.values.push((Some(name.clone()), val.clone()));
406                        }
407                    }
408                }
409                RValue::Vector(iv) => {
410                    let i = usize::try_from(iv.as_integer_scalar().unwrap_or(0)).unwrap_or(0);
411                    if i > 0 {
412                        while list.values.len() < i {
413                            list.values.push((None, RValue::Null));
414                        }
415                        list.values[i - 1].1 = val.clone();
416                    }
417                }
418                _ => {}
419            }
420            eval_assign(interp, op, object, obj, env)?;
421            Ok(val)
422        }
423        RValue::Environment(target_env) => {
424            // env[["key"]] <- value sets a variable in the environment
425            if let Some(name) = idx_val.as_vector().and_then(|v| v.as_character_scalar()) {
426                target_env.set(name, val.clone());
427            }
428            Ok(val)
429        }
430        RValue::Language(lang) => {
431            // lang[[i]] <- val  — replace the i-th element of the language object
432            let i = match &idx_val {
433                RValue::Vector(iv) => {
434                    usize::try_from(iv.as_integer_scalar().unwrap_or(0)).unwrap_or(0)
435                }
436                _ => 0,
437            };
438            if let Some(new_lang) = lang.set_element(i, &val) {
439                eval_assign(interp, op, object, RValue::Language(new_lang), env)?;
440                Ok(val)
441            } else {
442                Err(RError::new(
443                    RErrorKind::Index,
444                    format!(
445                        "attempt to select {} element{} from a language object of length {}",
446                        if i == 0 { "zero-th" } else { "beyond-end" },
447                        if i > 0 {
448                            format!(" ({})", i)
449                        } else {
450                            String::new()
451                        },
452                        lang.language_length()
453                    ),
454                )
455                .into())
456            }
457        }
458        _ => eval_index_assign(interp, op, object, indices, val, env),
459    }
460}
461
462// endregion
463
464// region: dollar assignment (x$name <- val)
465
466fn eval_dollar_assign(
467    interp: &Interpreter,
468    op: &AssignOp,
469    object: &Expr,
470    member: &str,
471    val: RValue,
472    env: &Environment,
473) -> Result<RValue, RFlow> {
474    // Evaluate the target object — supports chained targets
475    let mut obj = match interp.eval_in(object, env) {
476        Ok(v) => v,
477        Err(_) => RValue::List(RList::new(vec![])),
478    };
479    match &mut obj {
480        RValue::List(list) => {
481            if let Some(entry) = list
482                .values
483                .iter_mut()
484                .find(|(n, _)| n.as_deref() == Some(member))
485            {
486                entry.1 = val.clone();
487            } else {
488                list.values.push((Some(member.to_string()), val.clone()));
489            }
490            eval_assign(interp, op, object, obj, env)?;
491            Ok(val)
492        }
493        RValue::Null => {
494            let list = RList::new(vec![(Some(member.to_string()), val.clone())]);
495            eval_assign(interp, op, object, RValue::List(list), env)?;
496            Ok(val)
497        }
498        _ => {
499            let list = RList::new(vec![(Some(member.to_string()), val.clone())]);
500            eval_assign(interp, op, object, RValue::List(list), env)?;
501            Ok(val)
502        }
503    }
504}
505
506// endregion
r/interpreter/assignment.rs

r/interpreter/
assignment.rs
