python-ast 1.1.0

use proc_macro2::TokenStream;
use pyo3::{Borrowed, FromPyObject, PyAny, PyResult, prelude::PyAnyMethods};
use quote::quote;
use serde::{Deserialize, Serialize};

use crate::{
    CodeGen, CodeGenContext, ExprType, Node, PythonOptions, SymbolTableScopes,
    BinOps, FromPythonString, PyAttributeExtractor,
};

/// Augmented assignment statement (e.g., x += 1, y -= 2, etc.)
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct AugAssign {
    /// The target being assigned to (left side)
    pub target: ExprType,
    /// The operator (+=, -=, *=, etc.)
    pub op: BinOps,
    /// The value being assigned (right side)
    pub value: ExprType,
    /// Position information
    pub lineno: Option<usize>,
    pub col_offset: Option<usize>,
    pub end_lineno: Option<usize>,
    pub end_col_offset: Option<usize>,
}

impl<'a, 'py> FromPyObject<'a, 'py> for AugAssign {
    type Error = pyo3::PyErr;
    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
        // Extract target
        let target = ob.extract_attr_with_context("target", "augmented assignment target")?;
        let target: ExprType = target.extract()?;
        
        // Extract operator
        let op = ob.extract_attr_with_context("op", "augmented assignment operator")?;
        let op_type_str = op.extract_type_name("augmented assignment operator")?;
        let op = BinOps::parse_or_unknown(&op_type_str);
        
        // Extract value
        let value = ob.extract_attr_with_context("value", "augmented assignment value")?;
        let value: ExprType = value.extract()?;
        
        Ok(AugAssign {
            target,
            op,
            value,
            lineno: ob.lineno(),
            col_offset: ob.col_offset(),
            end_lineno: ob.end_lineno(),
            end_col_offset: ob.end_col_offset(),
        })
    }
}

impl Node for AugAssign {
    fn lineno(&self) -> Option<usize> { self.lineno }
    fn col_offset(&self) -> Option<usize> { self.col_offset }
    fn end_lineno(&self) -> Option<usize> { self.end_lineno }
    fn end_col_offset(&self) -> Option<usize> { self.end_col_offset }
}

impl CodeGen for AugAssign {
    type Context = CodeGenContext;
    type Options = PythonOptions;
    type SymbolTable = SymbolTableScopes;

    fn find_symbols(self, symbols: Self::SymbolTable) -> Self::SymbolTable {
        // Process the value for symbols, but don't add new symbols for augmented assignment
        self.value.find_symbols(symbols)
    }

    fn to_rust(
        self,
        ctx: Self::Context,
        options: Self::Options,
        symbols: Self::SymbolTable,
    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
        // Compound assignment to a subscript (`counts[k] += 1`) is a
        // read-modify-write: the Load lowering of the target is a cloned
        // temporary (py_index), not a place, so read via py_index, combine,
        // and store back via py_set_index. The index is evaluated once.
        if let ExprType::Subscript(sub) = &self.target {
            // The receiver must be a place (see subscript_receiver_place);
            // a cloned receiver would silently drop the write-back.
            let receiver = crate::subscript_receiver_place(
                &sub.value,
                ctx.clone(),
                options.clone(),
                symbols.clone(),
            )?;
            let index = match &sub.kind {
                crate::SubscriptKind::Index(index) => index
                    .clone()
                    .to_rust(ctx.clone(), options.clone(), symbols.clone())?,
                crate::SubscriptKind::Slice { .. } => {
                    return Err(
                        "augmented assignment to a slice (`x[a:b] += ...`) is not supported"
                            .to_string()
                            .into(),
                    )
                }
            };
            let value = self.value.to_rust(ctx, options, symbols)?;
            let elem = quote!(__rython_elem);
            let combined = combine_op(&self.op, &elem, &value)?;
            // The receiver place is bound once so a nested chain
            // (`grid[i][j] += 1`) evaluates its intermediate lookups — and
            // any side effects in their indices — exactly once.
            return Ok(quote! {
                {
                    let __rython_recv = &mut (#receiver);
                    let __rython_idx = #index;
                    let __rython_elem = (__rython_recv).py_index(__rython_idx.clone())?;
                    (__rython_recv).py_set_index(__rython_idx, #combined)?;
                }
            });
        }

        let target = self.target.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
        let value = self.value.to_rust(ctx, options, symbols)?;

        // Generate the appropriate augmented assignment operator
        match self.op {
            // `+=` mirrors Python's `+` (string concat, list concat,
            // numeric promotion) via PyAdd.
            BinOps::Add => Ok(quote!(#target = (#target).py_add(&(#value)))),
            BinOps::Sub => Ok(quote!(#target -= #value)),
            BinOps::Mult => Ok(quote!(#target *= #value)),
            // Python's `/` is TRUE division: `x /= 2` on an int yields a
            // float. Rust's `/=` on an integer truncates silently, so
            // mirror the BinOp lowering instead (an int target then fails
            // to compile, which is loud rather than quietly wrong).
            BinOps::Div => Ok(quote!(#target = (#target) as f64 / (#value) as f64)),
            // Python // and % floor toward negative infinity / take the
            // divisor's sign; use the stdpython helpers instead of Rust's
            // truncating operators.
            BinOps::FloorDiv => Ok(quote!(#target = py_floordiv(#target, #value))),
            BinOps::Mod => Ok(quote!(#target = py_mod(#target, #value))),
            BinOps::BitAnd => Ok(quote!(#target &= #value)),
            BinOps::BitOr => Ok(quote!(#target |= #value)),
            BinOps::BitXor => Ok(quote!(#target ^= #value)),
            BinOps::LShift => Ok(quote!(#target <<= #value)),
            BinOps::RShift => Ok(quote!(#target >>= #value)),
            BinOps::Pow => {
                // Rust doesn't have **= operator, so we need to expand it
                Ok(quote!(#target = py_pow(#target, #value)))
            },
            BinOps::MatMult => {
                // Matrix multiplication assignment - not directly supported in Rust
                // Would need specific matrix library support
                Err(format!("Matrix multiplication assignment not supported in Rust").into())
            },
            BinOps::Unknown => {
                Err(format!("Unknown augmented assignment operator").into())
            },
        }
    }
}

/// The read-modify-write combination for a compound assignment: how the
/// current element and the operand produce the stored value.
fn combine_op(
    op: &BinOps,
    elem: &TokenStream,
    value: &TokenStream,
) -> Result<TokenStream, Box<dyn std::error::Error>> {
    Ok(match op {
        BinOps::Add => quote!((#elem).py_add(&(#value))),
        BinOps::Sub => quote!(#elem - #value),
        BinOps::Mult => quote!(#elem * #value),
        BinOps::Div => quote!((#elem) as f64 / (#value) as f64),
        BinOps::FloorDiv => quote!(py_floordiv(#elem, #value)),
        BinOps::Mod => quote!(py_mod(#elem, #value)),
        BinOps::Pow => quote!(py_pow(#elem, #value)),
        BinOps::BitAnd => quote!(#elem & #value),
        BinOps::BitOr => quote!(#elem | #value),
        BinOps::BitXor => quote!(#elem ^ #value),
        BinOps::LShift => quote!(#elem << #value),
        BinOps::RShift => quote!(#elem >> #value),
        other => {
            return Err(format!(
                "augmented assignment operator {:?} not supported on subscripts",
                other
            )
            .into())
        }
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::create_parse_test;

    create_parse_test!(test_add_assign, "x += 1", "test.py");
    create_parse_test!(test_sub_assign, "x -= 1", "test.py");
    create_parse_test!(test_mul_assign, "x *= 2", "test.py");
    create_parse_test!(test_div_assign, "x /= 3", "test.py");
    create_parse_test!(test_mod_assign, "x %= 4", "test.py");
    create_parse_test!(test_pow_assign, "x **= 2", "test.py");
    create_parse_test!(test_bitand_assign, "x &= 5", "test.py");
    create_parse_test!(test_bitor_assign, "x |= 6", "test.py");
    create_parse_test!(test_bitxor_assign, "x ^= 7", "test.py");
    create_parse_test!(test_lshift_assign, "x <<= 2", "test.py");
    create_parse_test!(test_rshift_assign, "x >>= 3", "test.py");
}