add primitive magic methods

nac3core/src/typecheck/magic_methods.rs

@@ -1,3 +1,4 @@
+use std::{collections::HashMap, rc::Rc};
 use rustpython_parser::ast::{Cmpop, Operator, Unaryop};
 
 pub fn binop_name(op: &Operator) -> &'static str {
@@ -56,3 +57,366 @@ pub fn comparison_name(op: &Cmpop) -> Option<&'static str> {
         _ => None,
     }
 }
+
+use crate::typecheck::{type_inferencer::*, typedef::{FunSignature, FuncArg, TypeEnum, Unifier}};
+use rustpython_parser::ast;
+pub fn set_primirives_magic_methods(store: &PrimitiveStore, unifier: &mut Unifier) {
+    // int32 --------
+    if let Some(TypeEnum::TObj {fields, .. }) = Rc::get_mut(&mut unifier.get_ty(store.int32)) {
+        for op in &[
+        ast::Operator::Add, 
+        ast::Operator::Sub, 
+        ast::Operator::Mult,
+        ast::Operator::Mod, 
+        ast::Operator::Pow, 
+        ast::Operator::LShift, 
+        ast::Operator::RShift, 
+        ast::Operator::BitOr, 
+        ast::Operator::BitXor, 
+        ast::Operator::BitAnd, 
+        ast::Operator::FloorDiv
+        ] {
+            fields.insert(
+                binop_name(op).to_string(), 
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.int32,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.int32,
+                        is_optional: false,
+                        name: "other".into() // the name does not matter here
+                    }],
+                }))
+            );
+
+            fields.insert(
+                binop_assign_name(op).to_string(), 
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.none,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.int32,
+                        is_optional: false,
+                        name: "other".into()
+                    }]
+                }))
+            );
+        };
+        // int div int gets float
+        fields.insert(
+            binop_assign_name(&ast::Operator::Div).into(),
+            unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                ret: store.float,
+                vars: HashMap::new(),
+                args: vec![FuncArg {
+                    ty: store.int32,
+                    is_optional: false,
+                    name: "other".into()
+                }]
+            }))
+        );
+
+        for op in &[
+            ast::Cmpop::Eq,
+            ast::Cmpop::NotEq,
+            ast::Cmpop::Lt,
+            ast::Cmpop::LtE,
+            ast::Cmpop::Gt,
+            ast::Cmpop::GtE,
+        ] {
+            fields.insert(
+                comparison_name(op).unwrap().to_string(), 
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.bool,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.int32,
+                        is_optional: false,
+                        name: "other".into()
+                    }],
+                }))
+            );
+        }
+
+        for op in &[
+            ast::Unaryop::UAdd,
+            ast::Unaryop::USub,
+            ast::Unaryop::Not,
+            ast::Unaryop::Invert,
+        ] {
+            fields.insert(
+                unaryop_name(op).into(), 
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.int32,
+                    vars: HashMap::new(),
+                    args: vec![]
+                }))
+            );
+        }
+    } else { unreachable!() }
+    // int32 --------
+    // int64 --------
+    if let Some(TypeEnum::TObj {fields, .. }) = Rc::get_mut(&mut unifier.get_ty(store.int64)) {
+        for op in &[
+        ast::Operator::Add, 
+        ast::Operator::Sub, 
+        ast::Operator::Mult,
+        ast::Operator::Mod, 
+        ast::Operator::Pow, 
+        ast::Operator::LShift, 
+        ast::Operator::RShift, 
+        ast::Operator::BitOr, 
+        ast::Operator::BitXor, 
+        ast::Operator::BitAnd, 
+        ast::Operator::FloorDiv
+        ] {
+            fields.insert(
+                binop_name(op).to_string(), 
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.int64,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.int64,
+                        is_optional: false,
+                        name: "other".into() // the name does not matter here
+                    }],
+                }))
+            );
+
+            fields.insert(
+                binop_assign_name(op).to_string(), 
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.none,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.int64,
+                        is_optional: false,
+                        name: "other".into()
+                    }]
+                }))
+            );
+        };
+        fields.insert(
+            binop_assign_name(&ast::Operator::Div).into(),
+            unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                ret: store.float,
+                vars: HashMap::new(),
+                args: vec![FuncArg {
+                    ty: store.int64,
+                    is_optional: false,
+                    name: "other".into()
+                }]
+            }))
+        );
+
+        for op in &[
+            ast::Cmpop::Eq,
+            ast::Cmpop::NotEq,
+            ast::Cmpop::Lt,
+            ast::Cmpop::LtE,
+            ast::Cmpop::Gt,
+            ast::Cmpop::GtE,
+        ] {
+            fields.insert(
+                comparison_name(op).unwrap().to_string(), 
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.bool,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.int64,
+                        is_optional: false,
+                        name: "other".into()
+                    }],
+                }))
+            );
+        }
+
+        for op in &[
+            ast::Unaryop::UAdd,
+            ast::Unaryop::USub,
+            ast::Unaryop::Not,
+            ast::Unaryop::Invert,
+        ] {
+            fields.insert(
+                unaryop_name(op).into(), 
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.int64,
+                    vars: HashMap::new(),
+                    args: vec![]
+                }))
+            );
+        }
+    } else { unreachable!() }
+    // int64 --------
+    // float --------
+    if let Some(TypeEnum::TObj {fields, .. }) = Rc::get_mut(&mut unifier.get_ty(store.float)) {
+        for op in &[
+        ast::Operator::Add, 
+        ast::Operator::Sub, 
+        ast::Operator::Mult, 
+        ast::Operator::Div, 
+        ast::Operator::Mod, 
+        ast::Operator::Pow,
+        ast::Operator::FloorDiv,
+        ] {
+            fields.insert(
+                binop_name(op).to_string(), 
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.float,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.float,
+                        is_optional: false,
+                        name: "other".into() // the name does not matter here
+                    }],
+                }))
+            );
+
+            fields.insert(
+                binop_assign_name(op).to_string(), 
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.none,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.float,
+                        is_optional: false,
+                        name: "other".into()
+                    }]
+                }))
+            );
+        };
+
+        for op in &[
+            ast::Cmpop::Eq,
+            ast::Cmpop::NotEq,
+            ast::Cmpop::Lt,
+            ast::Cmpop::LtE,
+            ast::Cmpop::Gt,
+            ast::Cmpop::GtE,
+        ] {
+            fields.insert(
+                comparison_name(op).unwrap().to_string(), 
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.bool,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.float,
+                        is_optional: false,
+                        name: "other".into()
+                    }],
+                }))
+            );
+        }
+
+        for op in &[
+            ast::Unaryop::UAdd,
+            ast::Unaryop::USub,
+            ast::Unaryop::Not,
+        ] {
+            fields.insert(
+                unaryop_name(op).into(), 
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.int64,
+                    vars: HashMap::new(),
+                    args: vec![]
+                }))
+            );
+        }
+    } else { unreachable!() }
+    // float --------
+    // bool ---------
+    if let Some(TypeEnum::TObj {fields, .. }) = Rc::get_mut(&mut unifier.get_ty(store.bool)) {
+        for op in &[
+        ast::Operator::Add, 
+        ast::Operator::Sub, 
+        ast::Operator::Mult, 
+        ast::Operator::Mod, 
+        ast::Operator::Pow, 
+        ast::Operator::LShift, 
+        ast::Operator::RShift, 
+        ast::Operator::FloorDiv
+        ] {
+            fields.insert(
+                binop_name(op).into(),
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.int32,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.bool,
+                        is_optional: false,
+                        name: "other".into()
+                    }]
+                }))
+            );
+
+            fields.insert(
+                binop_name(op).into(),
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.int32,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.int32,
+                        is_optional: false,
+                        name: "other".into()
+                    }]
+                }))
+            );
+
+            // binop_assignment will change type?
+            /* fields.insert(
+                binop_assignment_name(op).into(),
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.none,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.bool,
+                        is_optional: false,
+                        name: "other".into()
+                    }]
+                }))
+            ); */
+        };
+
+        for op in &[
+            ast::Operator::BitOr,
+            ast::Operator::BitXor,
+            ast::Operator::BitAnd
+        ] {
+            fields.insert(
+                binop_name(op).into(),
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.bool,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.int32,
+                        is_optional: false,
+                        name: "other".into()
+                    }]
+                }))
+            );
+        };
+
+        for op in &[
+            ast::Cmpop::Eq,
+            ast::Cmpop::NotEq,
+            ast::Cmpop::Lt,
+            ast::Cmpop::LtE,
+            ast::Cmpop::Gt,
+            ast::Cmpop::GtE,
+        ] {
+            fields.insert(
+                comparison_name(op).unwrap().into(), 
+                unifier.add_ty(TypeEnum::TFunc(FunSignature {
+                    ret: store.int32,
+                    vars: HashMap::new(),
+                    args: vec![FuncArg {
+                        ty: store.bool,
+                        is_optional: false,
+                        name: "other".into()
+                    }]
+                }))
+            );
+        }
+    }
+    // bool --------
+}

nac3core/src/typecheck/type_inferencer/test.rs

@@ -45,10 +45,8 @@ struct TestEnvironment {
 }
 
 impl TestEnvironment {
-    pub fn basic_test_env() -> Option<TestEnvironment> {
+    pub fn basic_test_env() -> TestEnvironment {
-        use rustpython_parser::ast::Operator::*;
         let mut unifier = Unifier::new();
-        // let mut identifier_mapping = HashMap::new();
         
         let int32 = unifier.add_ty(TypeEnum::TObj {
             obj_id: 0,
@@ -77,15 +75,41 @@ impl TestEnvironment {
         });
         // identifier_mapping.insert("None".into(), none);
         let primitives = PrimitiveStore { int32, int64, float, bool, none };
+        set_primirives_magic_methods(&primitives, &mut unifier);
         
-        // if let TypeEnum::TObj {ref fields, ref params, .. } = *unifier.get_ty(int32) {
+        let id_to_name = [
-        //     for op in [Add, Sub, Mult, MatMult, Div, Mod, Pow, LShift, RShift, BitOr, BitXor, BitAnd, FloorDiv].into_iter() {
+            (0, "int32".to_string()),
-        //         let call = Rc::new(Call {posargs: vec![int32], kwargs: HashMap::new(), ret: int32, fun: RefCell::new(None)});
+            (1, "int64".to_string()),
-        //     };
+            (2, "float".to_string()),
-        //     None
+            (3, "bool".to_string()),
-        // } else {
+            (4, "none".to_string()),
-        //     None
+            (5, "Foo".to_string()),
-        // }
+            (6, "Bar".to_string()),
+            (7, "Bar2".to_string()),
+        ]
+        .iter()
+        .cloned()
+        .collect();
+
+        let mut identifier_mapping = HashMap::new();
+        identifier_mapping.insert("None".into(), none);
+       
+        let resolver =
+            Box::new(Resolver { identifier_mapping: identifier_mapping.clone(), class_names: Default::default() })
+                as Box<dyn SymbolResolver>;
+
+        TestEnvironment {
+            unifier,
+            function_data: FunctionData {
+                resolver,
+                bound_variables: Vec::new(),
+                return_type: None
+            },
+            primitives,
+            id_to_name,
+            identifier_mapping,
+            virtual_checks: Vec::new(),
+        }
     }
 
     fn new() -> TestEnvironment {