From ca1d471646d39dc925faa00b8500691db4b5a9b6 Mon Sep 17 00:00:00 2001
From: mlokr <jdoka@crownsterling.io>
Date: Thu, 16 May 2024 13:29:16 +0200
Subject: [PATCH] adding imm operations is come cases

---
 hblang/src/codegen.rs | 77 +++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 74 insertions(+), 3 deletions(-)

diff --git a/hblang/src/codegen.rs b/hblang/src/codegen.rs
index c47cfa73..e86f2cca 100644
--- a/hblang/src/codegen.rs
+++ b/hblang/src/codegen.rs
@@ -1231,6 +1231,7 @@ impl<'a> Codegen<'a> {
 
                 let lsize = self.size_of(left.ty);
                 let ty = ctx.ty().unwrap_or(left.ty);
+
                 let lhs = match std::mem::take(&mut ctx).loc() {
                     Some(Loc::RegRef(reg)) if Loc::RegRef(reg) == left.loc => LinReg(reg, None),
                     Some(loc) => {
@@ -1241,12 +1242,30 @@ impl<'a> Codegen<'a> {
                 };
                 let right = self.expr_ctx(right, Ctx::Inferred(left.ty))?;
                 let rsize = self.size_of(right.ty);
-                let rhs = self.loc_to_reg(right.loc, rsize);
 
                 let ty = self.assert_ty(expr.pos(), left.ty, right.ty);
                 let size = self.size_of(ty);
                 let signed = bt::is_signed(ty);
 
+                if let Loc::Imm(mut imm) = right.loc
+                    && let Some(oper) = Self::imm_math_op(op, signed, size, false)
+                {
+                    if matches!(op, T::Add | T::Sub) {
+                        if bt::is_pointer(left.ty) {
+                            let size = self.size_of(self.pointers[ty as usize]);
+                            imm *= size;
+                        }
+                    }
+
+                    self.code.encode(oper(lhs.0, lhs.0, imm));
+                    break 'ops Some(Value {
+                        ty,
+                        loc: Loc::Reg(lhs),
+                    });
+                }
+
+                let rhs = self.loc_to_reg(right.loc, rsize);
+
                 log::dbg!(
                     "binop: {} {} {}",
                     self.display_ty(ty),
@@ -1358,6 +1377,45 @@ impl<'a> Codegen<'a> {
         Some(ops[size.ilog2() as usize])
     }
 
+    fn imm_math_op(
+        op: T,
+        signed: bool,
+        size: u64,
+        first_imm: bool,
+    ) -> Option<fn(u8, u8, u64) -> (usize, [u8; instrs::MAX_SIZE])> {
+        use instrs as i;
+
+        macro_rules! def_op {
+            ($name:ident |$a:ident, $b:ident, $c:ident| $($tt:tt)*) => {
+                macro_rules! $name {
+                    ($$($$op:ident),*) => {
+                        [$$(
+                            |$a, $b, $c: u64| i::$$op($($tt)*),
+                        )*]
+                    }
+                }
+            };
+        }
+
+        def_op!(basic_op | a, b, c | a, b, c as _);
+        def_op!(sub_op | a, b, c | b, a, c.wrapping_neg() as _);
+
+        let ops = match op {
+            T::Add => basic_op!(addi8, addi16, addi32, addi64),
+            T::Sub if !first_imm => sub_op!(addi8, addi16, addi32, addi64),
+            T::Mul => basic_op!(muli8, muli16, muli32, muli64),
+            T::Band => return Some(i::andi),
+            T::Bor => return Some(i::ori),
+            T::Xor => return Some(i::xori),
+            T::Shr if signed && !first_imm => basic_op!(srui8, srui16, srui32, srui64),
+            T::Shr if !first_imm => basic_op!(srui8, srui16, srui32, srui64),
+            T::Shl if !first_imm => basic_op!(slui8, slui16, slui32, slui64),
+            _ => return None,
+        };
+
+        Some(ops[size.ilog2() as usize])
+    }
+
     fn struct_op(&mut self, op: T, ty: Type, ctx: Ctx, left: Loc, right: Loc) -> Option<Value> {
         if let TypeKind::Struct(stuct) = TypeKind::from_ty(ty) {
             let dst = match ctx {
@@ -1380,9 +1438,22 @@ impl<'a> Codegen<'a> {
         }
 
         let size = self.size_of(ty);
-        let (lhs, owned) = self.loc_to_reg_ref(&left, size);
-        let rhs = self.loc_to_reg(right, size);
         let signed = bt::is_signed(ty);
+        let (lhs, owned) = self.loc_to_reg_ref(&left, size);
+
+        if let Loc::Imm(imm) = right
+            && let Some(op) = Self::imm_math_op(op, signed, size, false)
+        {
+            self.code.encode(op(lhs, lhs, imm));
+            return if let Ctx::Dest(dest) = ctx {
+                self.assign(dest.ty, dest.loc, owned.map_or(Loc::RegRef(lhs), Loc::Reg));
+                Some(Value::VOID)
+            } else {
+                Some(Value::new(ty, owned.map_or(Loc::RegRef(lhs), Loc::Reg)))
+            };
+        }
+
+        let rhs = self.loc_to_reg(right, size);
 
         if let Some(op) = Self::math_op(op, signed, size) {
             self.code.encode(op(lhs, lhs, rhs.0));