run-rs 0.6.35

//! Plain and compound assignment.

use anyhow::{Result, bail};
use syn::spanned::Spanned;
use syn::{Expr, UnOp};

use crate::interpreter::bytecode::{BinKind, FieldName, Member, Op, Reg};

use super::place;
use super::{Compiler, NameLoc, idx16, int_literal};

impl Compiler<'_> {
    /// `*seq` for a `seq: &mut usize` parameter. A cell promoted or captured name keeps the strict op.
    pub(super) fn deref_param_reg(&self, expr: &Expr) -> Option<Reg> {
        let Expr::Path(p) = expr else { return None };
        if p.qself.is_some() || p.path.segments.len() != 1 {
            return None;
        }
        let name = p.path.segments[0].ident.to_string();
        let frame = self.frames.last()?;
        let reg = frame.local_reg(&name)?;
        if frame.mutable_locals.contains(&reg) {
            return None;
        }
        ((reg as usize) < frame.num_params).then_some(reg)
    }

    /// `deref_param_reg` for a `&mut` parameter a closure also writes. It lives in its capture
    /// cell, so `*x = v` stores into the cell and the return hands the cell's value back.
    pub(super) fn deref_param_cell(&self, expr: &Expr) -> Option<Reg> {
        let name = place::single_path_name(expr)?;
        let frame = self.frames.last()?;
        if frame.aliases.contains_key(&name) {
            return None;
        }
        let reg = frame.local_reg(&name)?;
        (frame.mutable_locals.contains(&reg) && (reg as usize) < frame.num_params).then_some(reg)
    }

    /// The variable a `*name` store writes directly. A `&mut variable` alias names the
    /// variable, which may live in an enclosing frame, and a `&mut` parameter of an enclosing
    /// function names its capture cell, see `deref_param_upvalue`.
    fn deref_write_target(&mut self, name: &str) -> Option<String> {
        let target = self.unalias(name);
        if target != name {
            return Some(target);
        }
        if let Some(target) = self.enclosing_alias_target(name) {
            return Some(target);
        }
        self.deref_param_upvalue(name).then(|| name.to_string())
    }

    /// The value stored into a local, owned. Its literals carry the local's width from the
    /// inference pass.
    pub(super) fn compile_stored_value(&mut self, value: &Expr) -> Result<Reg> {
        self.compile_owned_expr(value)
    }

    /// An owned value a store is about to take. A panic before the store drops it.
    fn hold_for_unwind(&mut self, val: Reg) {
        if self.ctx.has_drop {
            self.cur().unwind_temps.push(val);
        }
    }

    /// The store cloned the value into its place, so the register must not hold it any more.
    fn release_from_unwind(&mut self, val: Reg) {
        if self.ctx.has_drop {
            self.emit(Op::LoadUnit { dst: val });
        }
    }

    pub(super) fn compile_assign(&mut self, target: &Expr, value: &Expr) -> Result<()> {
        match target {
            Expr::Path(p) if p.path.segments.len() == 1 => {
                let name = p.path.segments[0].ident.to_string();
                let location = self.resolve_for_write(&name);
                let value = self.compile_stored_value(value)?;
                self.emit_name_store(location, value, &name)?;
            }
            // The value is evaluated before the place, so a panic in the index or the base
            // drops it like an argument the call never took. The store clones it into the
            // container, so the register is cleared after, or a later panic would drop it twice.
            Expr::Index(idx) => {
                let val = self.compile_stored_value(value)?;
                self.hold_for_unwind(val);
                let base = self.compile_place_base(&idx.expr)?;
                let key = self.compile_expr(&idx.index)?;
                self.set_line(idx.bracket_token.span.open());
                self.emit(Op::SetIndex { base, key, val });
                self.release_from_unwind(val);
            }
            Expr::Field(f) => {
                let val = self.compile_stored_value(value)?;
                self.hold_for_unwind(val);
                let base = self.compile_place_base(&f.base)?;
                let member = self.member_of(&f.member);
                self.emit(Op::SetField { base, member, val });
                self.release_from_unwind(val);
            }
            Expr::Unary(u) if matches!(u.op, UnOp::Deref(_)) => {
                // `*r = v` on a `&mut variable` alias writes the variable, which may live in an
                // enclosing frame
                if let Some(name) = place::single_path_name(&u.expr)
                    && let Some(target) = self.deref_write_target(&name)
                {
                    let location = self.resolve_for_write(&target);
                    let val = self.compile_stored_value(value)?;
                    self.emit_name_store(location, val, &target)?;
                    return Ok(());
                }
                let val = self.compile_owned_expr(value)?;
                if let Some(cell) = self.deref_param_cell(&u.expr) {
                    self.emit(Op::StoreCell { cell, src: val });
                } else if let Some(target) = self.deref_param_reg(&u.expr) {
                    self.emit(Op::SetDerefParam { target, val });
                } else {
                    let target = self.compile_expr(&u.expr)?;
                    self.emit(Op::SetDeref { target, val });
                }
            }
            Expr::Paren(p) => self.compile_assign(&p.expr, value)?,
            _ => bail!("invalid assignment target"),
        }
        Ok(())
    }

    /// The right operand evaluates before the place, so its panic fires first.
    pub(super) fn compile_compound_assign(
        &mut self,
        target: &Expr,
        op: BinKind,
        rhs: &Expr,
    ) -> Result<()> {
        match target {
            Expr::Path(p) if p.path.segments.len() == 1 => {
                let name = p.path.segments[0].ident.to_string();
                let location = self.resolve_for_write(&name);
                let typed = self.typed_arith(target, rhs, op);
                let rhs_reg = if int_literal(rhs).is_some() {
                    None
                } else {
                    Some(self.compile_expr(rhs)?)
                };
                let current = self.load_name_location(location, &name)?;
                // a plain local takes the result in place, every op reads before it writes
                let result = match location {
                    NameLoc::Local(reg) => reg,
                    _ => self.alloc(),
                };
                self.set_line(target.span());
                match (int_literal(rhs), rhs_reg) {
                    (Some(imm), _) => self.emit_bin_imm(result, current, imm, op, typed),
                    (None, Some(b)) => self.emit_bin(result, current, b, op, typed),
                    (None, None) => unreachable!("the right side compiled above"),
                }
                self.emit_name_store(location, result, &name)?;
            }
            Expr::Index(idx) => {
                let b = self.compile_expr(rhs)?;
                let base = self.compile_place_base(&idx.expr)?;
                let key = self.compile_expr(&idx.index)?;
                let cur = self.alloc();
                self.emit(Op::Index {
                    dst: cur,
                    base,
                    key,
                });
                let res = self.alloc();
                self.set_line(target.span());
                let typed = self.typed_arith(target, rhs, op);
                self.emit_bin(res, cur, b, op, typed);
                self.emit(Op::SetIndex {
                    base,
                    key,
                    val: res,
                });
            }
            Expr::Field(f) => {
                let b = self.compile_expr(rhs)?;
                let base = self.compile_place_base(&f.base)?;
                let member = self.member_of(&f.member);
                let cur = self.alloc();
                self.emit(Op::GetField {
                    dst: cur,
                    base,
                    member,
                });
                let res = self.alloc();
                self.set_line(target.span());
                let typed = self.typed_arith(target, rhs, op);
                self.emit_bin(res, cur, b, op, typed);
                self.emit(Op::SetField {
                    base,
                    member,
                    val: res,
                });
            }
            Expr::Unary(u) if matches!(u.op, UnOp::Deref(_)) => {
                self.compile_compound_deref_assign(u, op, rhs)?;
            }
            _ => bail!("invalid compound assignment target"),
        }
        Ok(())
    }

    /// The deref arm of `compile_compound_assign`.
    pub(super) fn compile_compound_deref_assign(
        &mut self,
        u: &syn::ExprUnary,
        op: BinKind,
        rhs: &Expr,
    ) -> Result<()> {
        // a `&mut variable` alias reads and writes the variable itself
        if let Some(name) = place::single_path_name(&u.expr)
            && let Some(target) = self.deref_write_target(&name)
        {
            let b = self.compile_expr(rhs)?;
            let location = self.resolve_for_write(&target);
            let current = self.load_name_location(location, &target)?;
            let result = self.alloc();
            self.set_line(u.span());
            self.emit(Op::Bin {
                dst: result,
                a: current,
                b,
                op,
            });
            self.emit_name_store(location, result, &target)?;
            return Ok(());
        }
        let b = self.compile_expr(rhs)?;
        if let Some(cell) = self.deref_param_cell(&u.expr) {
            let current = self.load_name_location(NameLoc::Cell(cell), "")?;
            let result = self.alloc();
            self.set_line(u.span());
            self.emit(Op::Bin {
                dst: result,
                a: current,
                b,
                op,
            });
            self.emit(Op::StoreCell { cell, src: result });
            return Ok(());
        }
        let param = self.deref_param_reg(&u.expr);
        let target = self.compile_expr(&u.expr)?;
        self.set_line(u.span());
        let Some(target) = param else {
            // the fused op holds the lock across the read-modify-write, so concurrent tasks can't
            // lose updates
            self.emit(Op::DerefBinAssign { target, val: b, op });
            return Ok(());
        };
        let current = self.alloc();
        self.emit(Op::Deref {
            dst: current,
            src: target,
        });
        let result = self.alloc();
        self.emit(Op::Bin {
            dst: result,
            a: current,
            b,
            op,
        });
        self.emit(Op::SetDerefParam {
            target,
            val: result,
        });
        Ok(())
    }

    pub(super) fn load_name_location(&mut self, location: NameLoc, name: &str) -> Result<Reg> {
        match location {
            NameLoc::Local(reg) => Ok(reg),
            NameLoc::Cell(cell) => {
                let reg = self.alloc();
                self.emit(Op::LoadCell { dst: reg, cell });
                Ok(reg)
            }
            NameLoc::Upvalue(idx) => {
                let reg = self.alloc();
                self.emit(Op::LoadUpvalue { dst: reg, idx });
                Ok(reg)
            }
            NameLoc::None => bail!("assignment to unknown variable `{name}`"),
        }
    }

    pub(super) fn emit_name_store(
        &mut self,
        location: NameLoc,
        src: Reg,
        name: &str,
    ) -> Result<()> {
        match location {
            NameLoc::Local(dst) if dst != src => {
                // the old value drops before the store, unless the binding only lent it
                let has_drop = self.ctx.has_drop;
                let f = self.cur();
                if has_drop && !f.shares_only(dst) && !f.drop_exempt.contains(&dst) {
                    f.drop_lists.push(vec![dst].into());
                    let list = idx16(f.drop_lists.len() - 1);
                    self.emit(Op::DropScope { list });
                }
                self.emit(Op::Move { dst, src });
            }
            NameLoc::Local(_) => {}
            NameLoc::Cell(cell) => self.emit(Op::StoreCell { cell, src }),
            NameLoc::Upvalue(idx) => self.emit(Op::StoreUpvalue { idx, src }),
            NameLoc::None => bail!("assignment to unknown variable `{name}`"),
        }
        Ok(())
    }

    pub(super) fn member_of(&mut self, member: &syn::Member) -> u16 {
        match member {
            syn::Member::Named(n) => {
                self.add_member(Member::Named(FieldName::new(n.to_string().into())))
            }
            syn::Member::Unnamed(i) => self.add_member(Member::Indexed(i.index as usize)),
        }
    }
}