core: refactor typecheck/magic_methods.rs operators & add op symbol name

nac3core/src/codegen/expr.rs

@@ -21,7 +21,7 @@ use crate::{
         DefinitionId, TopLevelDef,
     },
     typecheck::{
-        magic_methods::{binop_assign_name, binop_name, unaryop_name},
+        magic_methods::{BinOpVariant, OpInfo},
         typedef::{FunSignature, FuncArg, Type, TypeEnum, TypeVarId, Unifier, VarMap},
     },
 };
@@ -1167,7 +1167,7 @@ pub fn gen_binop_expr_with_values<'ctx, G: CodeGenerator>(
     op: Operator,
     right: (&Option<Type>, BasicValueEnum<'ctx>),
     loc: Location,
-    is_aug_assign: bool,
+    variant: BinOpVariant,
 ) -> Result<Option<ValueEnum<'ctx>>, String> {
     let (left_ty, left_val) = left;
     let (right_ty, right_val) = right;
@@ -1222,7 +1222,10 @@ pub fn gen_binop_expr_with_values<'ctx, G: CodeGenerator>(
                     generator,
                     ctx,
                     ndarray_dtype1,
-                    if is_aug_assign { Some(left_val) } else { None },
+                    match variant {
+                        BinOpVariant::Normal => None,
+                        BinOpVariant::AugAssign => Some(left_val),
+                    },
                     left_val,
                     right_val,
                 )?
@@ -1231,7 +1234,10 @@ pub fn gen_binop_expr_with_values<'ctx, G: CodeGenerator>(
                     generator,
                     ctx,
                     ndarray_dtype1,
-                    if is_aug_assign { Some(left_val) } else { None },
+                    match variant {
+                        BinOpVariant::Normal => None,
+                        BinOpVariant::AugAssign => Some(left_val),
+                    },
                     (left_val.as_base_value().into(), false),
                     (right_val.as_base_value().into(), false),
                     |generator, ctx, (lhs, rhs)| {
@@ -1242,7 +1248,7 @@ pub fn gen_binop_expr_with_values<'ctx, G: CodeGenerator>(
                             op,
                             (&Some(ndarray_dtype2), rhs),
                             ctx.current_loc,
-                            is_aug_assign,
+                            variant,
                         )?
                         .unwrap()
                         .to_basic_value_enum(
@@ -1267,7 +1273,10 @@ pub fn gen_binop_expr_with_values<'ctx, G: CodeGenerator>(
                 generator,
                 ctx,
                 ndarray_dtype,
-                if is_aug_assign { Some(ndarray_val) } else { None },
+                match variant {
+                    BinOpVariant::Normal => None,
+                    BinOpVariant::AugAssign => Some(ndarray_val),
+                },
                 (left_val, !is_ndarray1),
                 (right_val, !is_ndarray2),
                 |generator, ctx, (lhs, rhs)| {
@@ -1278,7 +1287,7 @@ pub fn gen_binop_expr_with_values<'ctx, G: CodeGenerator>(
                         op,
                         (&Some(ndarray_dtype), rhs),
                         ctx.current_loc,
-                        is_aug_assign,
+                        variant,
                     )?
                     .unwrap()
                     .to_basic_value_enum(ctx, generator, ndarray_dtype)
@@ -1293,13 +1302,15 @@ pub fn gen_binop_expr_with_values<'ctx, G: CodeGenerator>(
             unreachable!("must be tobj")
         };
         let (op_name, id) = {
-            let (binop_name, binop_assign_name) =
+            let normal_method_name = OpInfo::from_binop(op, BinOpVariant::Normal).method_name;
-                (binop_name(op).into(), binop_assign_name(op).into());
+            let assign_method_name = OpInfo::from_binop(op, BinOpVariant::AugAssign).method_name;
+
             // if is aug_assign, try aug_assign operator first
-            if is_aug_assign && fields.contains_key(&binop_assign_name) {
+            if variant == BinOpVariant::AugAssign && fields.contains_key(&assign_method_name.into())
-                (binop_assign_name, *obj_id)
+            {
+                (assign_method_name.into(), *obj_id)
             } else {
-                (binop_name, *obj_id)
+                (normal_method_name.into(), *obj_id)
             }
         };
 
@@ -1349,7 +1360,7 @@ pub fn gen_binop_expr<'ctx, G: CodeGenerator>(
     op: Operator,
     right: &Expr<Option<Type>>,
     loc: Location,
-    is_aug_assign: bool,
+    variant: BinOpVariant,
 ) -> Result<Option<ValueEnum<'ctx>>, String> {
     let left_val = if let Some(v) = generator.gen_expr(ctx, left)? {
         v.to_basic_value_enum(ctx, generator, left.custom.unwrap())?
@@ -1369,7 +1380,7 @@ pub fn gen_binop_expr<'ctx, G: CodeGenerator>(
         op,
         (&right.custom, right_val),
         loc,
-        is_aug_assign,
+        variant,
     )
 }
 
@@ -1453,7 +1464,10 @@ pub fn gen_unaryop_expr_with_values<'ctx, G: CodeGenerator>(
             if op == ast::Unaryop::Invert {
                 ast::Unaryop::Not
             } else {
-                unreachable!("ufunc {} not supported for ndarray[bool, N]", unaryop_name(op))
+                unreachable!(
+                    "ufunc {} not supported for ndarray[bool, N]",
+                    OpInfo::from_unaryop(op).method_name
+                )
             }
         } else {
             op
@@ -2343,7 +2357,15 @@ pub fn gen_expr<'ctx, G: CodeGenerator>(
             }
         }
         ExprKind::BinOp { op, left, right } => {
-            return gen_binop_expr(generator, ctx, left, *op, right, expr.location, false);
+            return gen_binop_expr(
+                generator,
+                ctx,
+                left,
+                *op,
+                right,
+                expr.location,
+                BinOpVariant::Normal,
+            );
         }
         ExprKind::UnaryOp { op, operand } => return gen_unaryop_expr(generator, ctx, *op, operand),
         ExprKind::Compare { left, ops, comparators } => {

nac3core/src/codegen/numpy.rs

@@ -11,8 +11,7 @@ use crate::{
             call_ndarray_calc_broadcast_index, call_ndarray_calc_nd_indices,
             call_ndarray_calc_size,
         },
-        llvm_intrinsics,
+        llvm_intrinsics::{self, call_memcpy_generic},
-        llvm_intrinsics::call_memcpy_generic,
         stmt::{gen_for_callback_incrementing, gen_for_range_callback, gen_if_else_expr_callback},
         CodeGenContext, CodeGenerator,
     },
@@ -22,7 +21,10 @@ use crate::{
         numpy::{make_ndarray_ty, unpack_ndarray_var_tys},
         DefinitionId,
     },
-    typecheck::typedef::{FunSignature, Type, TypeEnum},
+    typecheck::{
+        magic_methods::BinOpVariant,
+        typedef::{FunSignature, Type, TypeEnum},
+    },
 };
 use inkwell::types::{AnyTypeEnum, BasicTypeEnum, PointerType};
 use inkwell::{
@@ -1632,7 +1634,7 @@ pub fn ndarray_matmul_2d<'ctx, G: CodeGenerator>(
                     Operator::Mult,
                     (&Some(elem_ty), b),
                     ctx.current_loc,
-                    false,
+                    BinOpVariant::Normal,
                 )?
                 .unwrap()
                 .to_basic_value_enum(ctx, generator, elem_ty)?;
@@ -1645,7 +1647,7 @@ pub fn ndarray_matmul_2d<'ctx, G: CodeGenerator>(
                     Operator::Add,
                     (&Some(elem_ty), a_mul_b),
                     ctx.current_loc,
-                    false,
+                    BinOpVariant::Normal,
                 )?
                 .unwrap()
                 .to_basic_value_enum(ctx, generator, elem_ty)?;

nac3core/src/codegen/stmt.rs

@@ -11,7 +11,10 @@ use crate::{
         gen_in_range_check,
     },
     toplevel::{helper::PrimDef, numpy::unpack_ndarray_var_tys, DefinitionId, TopLevelDef},
-    typecheck::typedef::{FunSignature, Type, TypeEnum},
+    typecheck::{
+        magic_methods::BinOpVariant,
+        typedef::{FunSignature, Type, TypeEnum},
+    },
 };
 use inkwell::{
     attributes::{Attribute, AttributeLoc},
@@ -1574,7 +1577,15 @@ pub fn gen_stmt<G: CodeGenerator>(
         StmtKind::For { .. } => generator.gen_for(ctx, stmt)?,
         StmtKind::With { .. } => generator.gen_with(ctx, stmt)?,
         StmtKind::AugAssign { target, op, value, .. } => {
-            let value = gen_binop_expr(generator, ctx, target, *op, value, stmt.location, true)?;
+            let value = gen_binop_expr(
+                generator,
+                ctx,
+                target,
+                *op,
+                value,
+                stmt.location,
+                BinOpVariant::AugAssign,
+            )?;
             generator.gen_assign(ctx, target, value.unwrap())?;
         }
         StmtKind::Try { .. } => gen_try(generator, ctx, stmt)?,

nac3core/src/typecheck/magic_methods.rs

@@ -5,7 +5,7 @@ use crate::typecheck::{
     type_inferencer::*,
     typedef::{FunSignature, FuncArg, Type, TypeEnum, Unifier, VarMap},
 };
-use itertools::Itertools;
+use itertools::{iproduct, Itertools};
 use nac3parser::ast::StrRef;
 use nac3parser::ast::{Cmpop, Operator, Unaryop};
 use std::cmp::max;
@@ -13,64 +13,93 @@ use std::collections::HashMap;
 use std::rc::Rc;
 use strum::IntoEnumIterator;
 
-#[must_use]
+/// Details about an operator (unary, binary, etc...) in Python
-pub fn binop_name(op: Operator) -> &'static str {
+#[derive(Debug, Clone, Copy)]
-    match op {
+pub struct OpInfo {
-        Operator::Add => "__add__",
+    /// The method name of the binary operator.
-        Operator::Sub => "__sub__",
+    /// For addition, this would be `__add__`, and `__iadd__` if
-        Operator::Div => "__truediv__",
+    /// it is the augmented assigning variant.
-        Operator::Mod => "__mod__",
+    pub method_name: &'static str,
-        Operator::Mult => "__mul__",
+    /// The symbol of the binary operator.
-        Operator::Pow => "__pow__",
+    /// For addition, this would be `+`, and `+=` if
-        Operator::BitOr => "__or__",
+    /// it is the augmented assigning variant.
-        Operator::BitXor => "__xor__",
+    pub symbol: &'static str,
-        Operator::BitAnd => "__and__",
-        Operator::LShift => "__lshift__",
-        Operator::RShift => "__rshift__",
-        Operator::FloorDiv => "__floordiv__",
-        Operator::MatMult => "__matmul__",
-    }
 }
 
-#[must_use]
+/// Helper macro to conveniently build an [`OpInfo`].
-pub fn binop_assign_name(op: Operator) -> &'static str {
+///
-    match op {
+/// Example usage: `make_info("add", "+")` generates `OpInfo { name: "__add__", symbol: "+" }`
-        Operator::Add => "__iadd__",
+macro_rules! make_info {
-        Operator::Sub => "__isub__",
+    ($name:expr, $symbol:expr) => {
-        Operator::Div => "__itruediv__",
+        OpInfo { method_name: concat!("__", $name, "__"), symbol: $symbol }
-        Operator::Mod => "__imod__",
+    };
-        Operator::Mult => "__imul__",
-        Operator::Pow => "__ipow__",
-        Operator::BitOr => "__ior__",
-        Operator::BitXor => "__ixor__",
-        Operator::BitAnd => "__iand__",
-        Operator::LShift => "__ilshift__",
-        Operator::RShift => "__irshift__",
-        Operator::FloorDiv => "__ifloordiv__",
-        Operator::MatMult => "__imatmul__",
-    }
 }
 
-#[must_use]
+/// The variant of a binary operator.
-pub fn unaryop_name(op: Unaryop) -> &'static str {
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-    match op {
+pub enum BinOpVariant {
-        Unaryop::UAdd => "__pos__",
+    /// The normal variant.
-        Unaryop::USub => "__neg__",
+    /// For addition, it would be `+`.
-        Unaryop::Not => "__not__",
+    Normal,
-        Unaryop::Invert => "__inv__",
+    /// The "Augmented Assigning Operator" variant.
-    }
+    /// For addition, it would be `+=`.
+    AugAssign,
 }
 
-#[must_use]
+impl OpInfo {
-pub fn comparison_name(op: Cmpop) -> Option<&'static str> {
+    #[must_use]
-    match op {
+    pub fn from_binop(op: Operator, variant: BinOpVariant) -> Self {
-        Cmpop::Lt => Some("__lt__"),
+        // Helper macro to generate both the normal variant [`OpInfo`] and the
-        Cmpop::LtE => Some("__le__"),
+        // augmented assigning variant [`OpInfo`] for a binary operator conveniently.
-        Cmpop::Gt => Some("__gt__"),
+        macro_rules! info {
-        Cmpop::GtE => Some("__ge__"),
+            ($name:literal, $symbol:literal) => {
-        Cmpop::Eq => Some("__eq__"),
+                (make_info!($name, $symbol), make_info!(concat!("i", $name), concat!($symbol, "=")))
-        Cmpop::NotEq => Some("__ne__"),
+            };
-        _ => None,
+        }
+
+        let (normal_variant, aug_assign_variant) = match op {
+            Operator::Add => info!("add", "+"),
+            Operator::Sub => info!("sub", "-"),
+            Operator::Div => info!("truediv", "/"),
+            Operator::Mod => info!("mod", "%"),
+            Operator::Mult => info!("mul", "*"),
+            Operator::Pow => info!("pow", "**"),
+            Operator::BitOr => info!("or", "|"),
+            Operator::BitXor => info!("xor", "^"),
+            Operator::BitAnd => info!("and", "&"),
+            Operator::LShift => info!("lshift", "<<"),
+            Operator::RShift => info!("rshift", ">>"),
+            Operator::FloorDiv => info!("floordiv", "//"),
+            Operator::MatMult => info!("matmul", "@"),
+        };
+
+        match variant {
+            BinOpVariant::Normal => normal_variant,
+            BinOpVariant::AugAssign => aug_assign_variant,
+        }
+    }
+
+    #[must_use]
+    pub fn from_unaryop(op: Unaryop) -> Self {
+        match op {
+            Unaryop::UAdd => make_info!("pos", "+"),
+            Unaryop::USub => make_info!("neg", "-"),
+            Unaryop::Not => make_info!("not", "not"), // i.e., `not False`, so the symbol is just `not`.
+            Unaryop::Invert => make_info!("inv", "~"),
+        }
+    }
+
+    #[must_use]
+    pub fn from_cmpop(op: Cmpop) -> Option<Self> {
+        match op {
+            Cmpop::Lt => Some(make_info!("lt", "<")),
+            Cmpop::LtE => Some(make_info!("le", "<=")),
+            Cmpop::Gt => Some(make_info!("gt", ">")),
+            Cmpop::GtE => Some(make_info!("ge", ">=")),
+            Cmpop::Eq => Some(make_info!("eq", "==")),
+            Cmpop::NotEq => Some(make_info!("ne", "!=")),
+            _ => None,
+        }
     }
 }
 
@@ -115,23 +144,8 @@ pub fn impl_binop(
 
         let ret_ty = ret_ty.unwrap_or_else(|| unifier.get_fresh_var(None, None).ty);
 
-        for op in ops {
+        for (op, variant) in iproduct!(ops, [BinOpVariant::Normal, BinOpVariant::AugAssign]) {
-            fields.insert(binop_name(*op).into(), {
+            fields.insert(OpInfo::from_binop(*op, variant).method_name.into(), {
-                (
-                    unifier.add_ty(TypeEnum::TFunc(FunSignature {
-                        ret: ret_ty,
-                        vars: function_vars.clone(),
-                        args: vec![FuncArg {
-                            ty: other_ty,
-                            default_value: None,
-                            name: "other".into(),
-                        }],
-                    })),
-                    false,
-                )
-            });
-
-            fields.insert(binop_assign_name(*op).into(), {
                 (
                     unifier.add_ty(TypeEnum::TFunc(FunSignature {
                         ret: ret_ty,
@@ -155,7 +169,7 @@ pub fn impl_unaryop(unifier: &mut Unifier, ty: Type, ret_ty: Option<Type>, ops:
 
         for op in ops {
             fields.insert(
-                unaryop_name(*op).into(),
+                OpInfo::from_unaryop(*op).method_name.into(),
                 (
                     unifier.add_ty(TypeEnum::TFunc(FunSignature {
                         ret: ret_ty,
@@ -195,7 +209,7 @@ pub fn impl_cmpop(
 
         for op in ops {
             fields.insert(
-                comparison_name(*op).unwrap().into(),
+                OpInfo::from_cmpop(*op).unwrap().method_name.into(),
                 (
                     unifier.add_ty(TypeEnum::TFunc(FunSignature {
                         ret: ret_ty,

nac3core/src/typecheck/type_inferencer/mod.rs

@@ -466,7 +466,8 @@ impl<'a> Fold<()> for Inferencer<'a> {
                 (None, None) => {}
             },
             ast::StmtKind::AugAssign { target, op, value, .. } => {
-                let res_ty = self.infer_bin_ops(stmt.location, target, *op, value, true)?;
+                let res_ty =
+                    self.infer_bin_ops(stmt.location, target, *op, value, BinOpVariant::AugAssign)?;
                 self.unify(res_ty, target.custom.unwrap(), &stmt.location)?;
             }
             ast::StmtKind::Assert { test, msg, .. } => {
@@ -548,7 +549,7 @@ impl<'a> Fold<()> for Inferencer<'a> {
             }
             ExprKind::BoolOp { values, .. } => Some(self.infer_bool_ops(values)?),
             ExprKind::BinOp { left, op, right } => {
-                Some(self.infer_bin_ops(expr.location, left, *op, right, false)?)
+                Some(self.infer_bin_ops(expr.location, left, *op, right, BinOpVariant::Normal)?)
             }
             ExprKind::UnaryOp { op, operand } => {
                 Some(self.infer_unary_ops(expr.location, *op, operand)?)
@@ -1536,7 +1537,7 @@ impl<'a> Inferencer<'a> {
         left: &ast::Expr<Option<Type>>,
         op: ast::Operator,
         right: &ast::Expr<Option<Type>>,
-        is_aug_assign: bool,
+        variant: BinOpVariant,
     ) -> InferenceResult {
         let left_ty = left.custom.unwrap();
         let right_ty = right.custom.unwrap();
@@ -1544,27 +1545,32 @@ impl<'a> Inferencer<'a> {
         let method = if let TypeEnum::TObj { fields, .. } =
             self.unifier.get_ty_immutable(left_ty).as_ref()
         {
-            let (binop_name, binop_assign_name) =
+            let normal_method_name = OpInfo::from_binop(op, BinOpVariant::Normal).method_name;
-                (binop_name(op).into(), binop_assign_name(op).into());
+            let assign_method_name = OpInfo::from_binop(op, BinOpVariant::AugAssign).method_name;
+
             // if is aug_assign, try aug_assign operator first
-            if is_aug_assign && fields.contains_key(&binop_assign_name) {
+            if variant == BinOpVariant::AugAssign && fields.contains_key(&assign_method_name.into())
-                binop_assign_name
+            {
+                assign_method_name
             } else {
-                binop_name
+                normal_method_name
             }
         } else {
-            binop_name(op).into()
+            OpInfo::from_binop(op, variant).method_name
         };
 
-        let ret = if is_aug_assign {
+        let ret = match variant {
-            // The type of augmented assignment operator should never change
+            BinOpVariant::Normal => {
-            Some(left_ty)
+                typeof_binop(self.unifier, self.primitives, op, left_ty, right_ty)
-        } else {
+                    .map_err(|e| HashSet::from([format!("{e} (at {location})")]))?
-            typeof_binop(self.unifier, self.primitives, op, left_ty, right_ty)
+            }
-                .map_err(|e| HashSet::from([format!("{e} (at {location})")]))?
+            BinOpVariant::AugAssign => {
+                // The type of augmented assignment operator should never change
+                Some(left_ty)
+            }
         };
 
-        self.build_method_call(location, method, left_ty, vec![right_ty], ret)
+        self.build_method_call(location, method.into(), left_ty, vec![right_ty], ret)
     }
 
     fn infer_unary_ops(
@@ -1573,7 +1579,7 @@ impl<'a> Inferencer<'a> {
         op: ast::Unaryop,
         operand: &ast::Expr<Option<Type>>,
     ) -> InferenceResult {
-        let method = unaryop_name(op).into();
+        let method = OpInfo::from_unaryop(op).method_name.into();
 
         let ret = typeof_unaryop(self.unifier, self.primitives, op, operand.custom.unwrap())
             .map_err(|e| HashSet::from([format!("{e} (at {location})")]))?;
@@ -1603,8 +1609,9 @@ impl<'a> Inferencer<'a> {
 
         let mut res = None;
         for (a, b, c) in izip!(once(left).chain(comparators), comparators, ops) {
-            let method = comparison_name(*c)
+            let method = OpInfo::from_cmpop(*c)
                 .ok_or_else(|| HashSet::from(["unsupported comparator".to_string()]))?
+                .method_name
                 .into();
 
             let ret = typeof_cmpop(