From c238c264e7253e18b4b72ebdd926e404730cead7 Mon Sep 17 00:00:00 2001
From: ychenfo <ychenfo@connect.ust.hk>
Date: Thu, 19 Aug 2021 11:18:58 +0800
Subject: [PATCH] add type vars to the primitive binop function def

---
 nac3core/src/typecheck/magic_methods.rs | 33 +++++++++++++------------
 1 file changed, 17 insertions(+), 16 deletions(-)

diff --git a/nac3core/src/typecheck/magic_methods.rs b/nac3core/src/typecheck/magic_methods.rs
index 7a9fbd6db..30e3c7531 100644
--- a/nac3core/src/typecheck/magic_methods.rs
+++ b/nac3core/src/typecheck/magic_methods.rs
@@ -66,26 +66,32 @@ pub fn comparison_name(op: &Cmpop) -> Option<&'static str> {
 
 pub fn impl_binop(
     unifier: &mut Unifier,
-    _store: &PrimitiveStore,
+    store: &PrimitiveStore,
     ty: Type,
     other_ty: &[Type],
     ret_ty: Type,
     ops: &[ast::Operator],
 ) {
     if let TypeEnum::TObj { fields, .. } = unifier.get_ty(ty).borrow() {
+        let (other_ty, other_var_id) = if other_ty.len() == 1 {
+            (other_ty[0], None)
+        } else {
+            let (ty, var_id) = unifier.get_fresh_var_with_range(other_ty);
+            (ty, Some(var_id))
+        };
+        let function_vars = if let Some(var_id) = other_var_id {
+            vec![(var_id, other_ty)].into_iter().collect::<HashMap<_, _>>()
+        } else {
+            HashMap::new()
+        };
         for op in ops {
             fields.borrow_mut().insert(binop_name(op).into(), {
-                let other = if other_ty.len() == 1 {
-                    other_ty[0]
-                } else {
-                    unifier.get_fresh_var_with_range(other_ty).0
-                };
                 unifier.add_ty(TypeEnum::TFunc(
                     FunSignature {
                         ret: ret_ty,
-                        vars: HashMap::new(),
+                        vars: function_vars.clone(),
                         args: vec![FuncArg {
-                            ty: other,
+                            ty: other_ty,
                             default_value: None,
                             name: "other".into(),
                         }],
@@ -95,17 +101,12 @@ pub fn impl_binop(
             });
 
             fields.borrow_mut().insert(binop_assign_name(op).into(), {
-                let other = if other_ty.len() == 1 {
-                    other_ty[0]
-                } else {
-                    unifier.get_fresh_var_with_range(other_ty).0
-                };
                 unifier.add_ty(TypeEnum::TFunc(
                     FunSignature {
-                        ret: ret_ty,
-                        vars: HashMap::new(),
+                        ret: store.none,
+                        vars: function_vars.clone(),
                         args: vec![FuncArg {
-                            ty: other,
+                            ty: other_ty,
                             default_value: None,
                             name: "other".into(),
                         }],