From 469ce3ce77c5dbf1e48ed264ebbdf4a5cbd19835 Mon Sep 17 00:00:00 2001
From: sofia <sofia@teascade.net>
Date: Mon, 28 Jul 2025 13:02:02 +0300
Subject: [PATCH] Improve type inferrence for binops

---
 examples/loops.reid                     | 11 -------
 examples/or_bitwise.reid                |  4 +--
 reid/src/mir/typecheck/typecheck.rs     |  2 +-
 reid/src/mir/typecheck/typeinference.rs |  3 ++
 reid/src/mir/typecheck/typerefs.rs      | 41 +++++++++++++++++++++++--
 5 files changed, 44 insertions(+), 17 deletions(-)

diff --git a/examples/loops.reid b/examples/loops.reid
index 65cef42..ae3e54b 100644
--- a/examples/loops.reid
+++ b/examples/loops.reid
@@ -1,17 +1,6 @@
 // Arithmetic, function calls and imports!
 
-import std::print;
-import std::from_str;
-import std::add_num_to_str;
-import std::free_string;
-
 fn main() -> u32 {
-  for i in 0 .. 15 {
-    let mut text = from_str("num: ");
-    add_num_to_str(&mut text, i);
-    print(text);
-    free_string(&text);
-  }
 
   let mut num = 0;
   while num < 10 {
diff --git a/examples/or_bitwise.reid b/examples/or_bitwise.reid
index e7b7487..bae3821 100644
--- a/examples/or_bitwise.reid
+++ b/examples/or_bitwise.reid
@@ -2,9 +2,9 @@
 
 
 fn main() -> bool {
-  let bwand = (0xffu32 & 0xf0) >> 4;
+  let bwand = (0xff & 0xf0) >> 4;
   let bwor = (0x0fu32 | 0x00) << 4;
-  let bwxor = (0xf0u32 | 0x0f);
+  let bwxor = (0xf0 | 0x0f);
 
   return (bwxor == 255) && ((bwand == 15) || false) && (bwor == 240);
 }
\ No newline at end of file
diff --git a/reid/src/mir/typecheck/typecheck.rs b/reid/src/mir/typecheck/typecheck.rs
index 59afd4a..dca27bb 100644
--- a/reid/src/mir/typecheck/typecheck.rs
+++ b/reid/src/mir/typecheck/typecheck.rs
@@ -431,7 +431,7 @@ impl Expression {
                     params: (lhs_type.clone(), rhs_type.clone()),
                     operator: *op,
                 });
-                dbg!(&binops, &expected_return_ty);
+                // dbg!(&lhs_type, &rhs_type, &binops, &ret_ty, &expected_return_ty);
                 if let Some(binop) = binops
                     .iter()
                     .filter(|f| f.1.return_ty.narrow_into(&expected_return_ty).is_ok())
diff --git a/reid/src/mir/typecheck/typeinference.rs b/reid/src/mir/typecheck/typeinference.rs
index 698b1d5..8adb0c6 100644
--- a/reid/src/mir/typecheck/typeinference.rs
+++ b/reid/src/mir/typecheck/typeinference.rs
@@ -390,6 +390,9 @@ impl Expression {
                         widened_rhs = widened_rhs.widen_into(&binop.hands.1);
                     }
                     let binop_res = type_refs.from_binop(*op, &lhs_ref, &rhs_ref);
+                    // dbg!(&return_ty);
+                    // dbg!(&binop_res);
+                    // dbg!(&lhs_ref, &rhs_ref, &binops, &widened_lhs, &widened_rhs);
                     lhs_ref.narrow(&type_refs.from_type(&widened_lhs).unwrap());
                     rhs_ref.narrow(&type_refs.from_type(&widened_rhs).unwrap());
                     *return_ty = binop_res.as_type();
diff --git a/reid/src/mir/typecheck/typerefs.rs b/reid/src/mir/typecheck/typerefs.rs
index fe91702..bdca213 100644
--- a/reid/src/mir/typecheck/typerefs.rs
+++ b/reid/src/mir/typecheck/typerefs.rs
@@ -332,21 +332,56 @@ impl<'outer> ScopeTypeRefs<'outer> {
                 .widen(self.types);
             self.narrow_to_type(&hint1, &ty)?;
             let hint1_typeref = self.types.retrieve_typeref(*hint1.0.borrow()).unwrap();
+            let hint2_typeref = self.types.retrieve_typeref(*hint2.0.borrow()).unwrap();
+
+            match (&hint1_typeref, &hint2_typeref) {
+                (TypeRefKind::Direct(ret_ty), TypeRefKind::BinOp(op, lhs, rhs)) => {
+                    let mut lhs_ref = self.from_type(&lhs).unwrap();
+                    let mut rhs_ref = self.from_type(&rhs).unwrap();
+                    let binops = self.available_binops(op, &mut lhs_ref, &mut rhs_ref);
+                    let mut binops = binops
+                        .iter()
+                        .filter(|b| b.return_ty.narrow_into(ret_ty).is_ok())
+                        .into_iter();
+
+                    if let Some(binop) = binops.next() {
+                        let mut lhs_widened = binop.hands.0.clone();
+                        let mut rhs_widened = binop.hands.1.clone();
+                        while let Some(binop) = binops.next() {
+                            lhs_widened = lhs_widened.widen_into(&binop.hands.0);
+                            rhs_widened = rhs_widened.widen_into(&binop.hands.1);
+                        }
+
+                        lhs_ref.narrow(&self.from_type(&lhs_widened).unwrap());
+                        rhs_ref.narrow(&self.from_type(&rhs_widened).unwrap());
+                    }
+                }
+                _ => {}
+            }
+
             for idx in self.types.type_refs.borrow_mut().iter_mut() {
-                match hint1_typeref {
-                    TypeRefKind::Direct(_) => {
+                match (&hint1_typeref, &hint2_typeref) {
+                    (TypeRefKind::Direct(_), TypeRefKind::Direct(_)) => {
                         if *idx == hint2.0 && idx != &hint1.0 {
                             *idx.borrow_mut() = *hint1.0.borrow();
                         }
                     }
-                    TypeRefKind::BinOp(_, _, _) => {
+                    (TypeRefKind::Direct(_), TypeRefKind::BinOp(..)) => {}
+                    (TypeRefKind::BinOp(..), TypeRefKind::Direct(..)) => {
                         // TODO may not be good ?
                         // if *idx == hint2.0 && idx != &hint1.0 {
                         //     *idx.borrow_mut() = *hint1.0.borrow();
                         // }
                     }
+                    (TypeRefKind::BinOp(..), TypeRefKind::BinOp(..)) => {
+                        // TODO may not be good ?
+                        if *idx == hint2.0 && idx != &hint1.0 {
+                            *idx.borrow_mut() = *hint1.0.borrow();
+                        }
+                    }
                 }
             }
+
             Some(TypeRef(hint1.0.clone(), self))
         }
     }