From 6369db94ab1d35052c6aad5950cedbb0a1df58f2 Mon Sep 17 00:00:00 2001
From: lyken <lyken@m-labs.hk>
Date: Tue, 30 Jul 2024 16:32:38 +0800
Subject: [PATCH] core/codegen: gen_assign to take in value_ty

---
 nac3core/src/codegen/expr.rs      |  2 +-
 nac3core/src/codegen/generator.rs |  3 ++-
 nac3core/src/codegen/stmt.rs      | 14 +++++++++-----
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/nac3core/src/codegen/expr.rs b/nac3core/src/codegen/expr.rs
index 432b358b..990ac83a 100644
--- a/nac3core/src/codegen/expr.rs
+++ b/nac3core/src/codegen/expr.rs
@@ -1128,7 +1128,7 @@ pub fn gen_comprehension<'ctx, G: CodeGenerator>(
             )
             .into_pointer_value();
         let val = ctx.build_gep_and_load(arr_ptr, &[tmp], Some("val"));
-        generator.gen_assign(ctx, target, val.into())?;
+        generator.gen_assign(ctx, target, val.into(), ctx.primitives.int32)?;
     }
 
     // Emits the content of `cont_bb`
diff --git a/nac3core/src/codegen/generator.rs b/nac3core/src/codegen/generator.rs
index bb822f19..5259d280 100644
--- a/nac3core/src/codegen/generator.rs
+++ b/nac3core/src/codegen/generator.rs
@@ -123,11 +123,12 @@ pub trait CodeGenerator {
         ctx: &mut CodeGenContext<'ctx, '_>,
         target: &Expr<Option<Type>>,
         value: ValueEnum<'ctx>,
+        value_ty: Type,
     ) -> Result<(), String>
     where
         Self: Sized,
     {
-        gen_assign(self, ctx, target, value)
+        gen_assign(self, ctx, target, value, value_ty)
     }
 
     /// Generate code for a while expression.
diff --git a/nac3core/src/codegen/stmt.rs b/nac3core/src/codegen/stmt.rs
index f7ab5ee7..393119de 100644
--- a/nac3core/src/codegen/stmt.rs
+++ b/nac3core/src/codegen/stmt.rs
@@ -209,6 +209,7 @@ pub fn gen_assign<'ctx, G: CodeGenerator>(
     ctx: &mut CodeGenContext<'ctx, '_>,
     target: &Expr<Option<Type>>,
     value: ValueEnum<'ctx>,
+    value_ty: Type,
 ) -> Result<(), String> {
     let llvm_usize = generator.get_size_type(ctx.ctx);
 
@@ -225,7 +226,7 @@ pub fn gen_assign<'ctx, G: CodeGenerator>(
                     .builder
                     .build_extract_value(v, u32::try_from(i).unwrap(), "struct_elem")
                     .unwrap();
-                generator.gen_assign(ctx, elt, v.into())?;
+                generator.gen_assign(ctx, elt, v.into(), value_ty)?;
             }
         }
         ExprKind::Subscript { value: ls, slice, .. }
@@ -434,7 +435,7 @@ pub fn gen_for<G: CodeGenerator>(
             .map(BasicValueEnum::into_int_value)
             .unwrap();
         let val = ctx.build_gep_and_load(arr_ptr, &[index], Some("val"));
-        generator.gen_assign(ctx, target, val.into())?;
+        generator.gen_assign(ctx, target, val.into(), ctx.primitives.int32)?;
         generator.gen_block(ctx, body.iter())?;
     }
 
@@ -1585,14 +1586,16 @@ pub fn gen_stmt<G: CodeGenerator>(
         }
         StmtKind::AnnAssign { target, value, .. } => {
             if let Some(value) = value {
+                let value_ty = value.custom.unwrap();
                 let Some(value) = generator.gen_expr(ctx, value)? else { return Ok(()) };
-                generator.gen_assign(ctx, target, value)?;
+                generator.gen_assign(ctx, target, value, value_ty)?;
             }
         }
         StmtKind::Assign { targets, value, .. } => {
+            let value_ty = value.custom.unwrap();
             let Some(value) = generator.gen_expr(ctx, value)? else { return Ok(()) };
             for target in targets {
-                generator.gen_assign(ctx, target, value.clone())?;
+                generator.gen_assign(ctx, target, value.clone(), value_ty)?;
             }
         }
         StmtKind::Continue { .. } => {
@@ -1606,6 +1609,7 @@ 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_ty = value.custom.unwrap();
             let value = gen_binop_expr(
                 generator,
                 ctx,
@@ -1614,7 +1618,7 @@ pub fn gen_stmt<G: CodeGenerator>(
                 value,
                 stmt.location,
             )?;
-            generator.gen_assign(ctx, target, value.unwrap())?;
+            generator.gen_assign(ctx, target, value.unwrap(), value_ty)?;
         }
         StmtKind::Try { .. } => gen_try(generator, ctx, stmt)?,
         StmtKind::Raise { exc, .. } => {