fix: div/mod/shift operations on unsigned integers

3 files changed, 44 insertions, 11 deletions

diff --git a/src/ast.c b/src/ast.c
index 4d30a66..0dcec0d 100644
--- a/src/ast.c
+++ b/src/ast.c
@@ -126,6 +126,11 @@ bool type_is_unsized(Type* ty) {
     return ty->kind == TypeKind_void;
 }
 
+bool type_is_unsigned(Type* ty) {
+    return ty->kind == TypeKind_uchar || ty->kind == TypeKind_ushort || ty->kind == TypeKind_uint ||
+           ty->kind == TypeKind_ulong || ty->kind == TypeKind_ullong || ty->kind == TypeKind_bool;
+}
+
 int type_sizeof(Type* ty) {
     if (type_is_unsized(ty)) {
         fatal_error("type_sizeof: type size cannot be determined");
diff --git a/src/ast.h b/src/ast.h
index f12fd83..c6fd0e2 100644
--- a/src/ast.h
+++ b/src/ast.h
@@ -73,6 +73,7 @@ Type* type_new_static_string(int len);
 Type* type_array_to_ptr(Type* ty);
 Type* type_new_func(Type* result, AstNode* params);
 bool type_is_unsized(Type* ty);
+bool type_is_unsigned(Type* ty);
 
 int type_sizeof_struct(Type* ty);
 int type_sizeof_union(Type* ty);
diff --git a/src/codegen.c b/src/codegen.c
index 2da6851..10a0287 100644
--- a/src/codegen.c
+++ b/src/codegen.c
@@ -205,11 +205,21 @@ static void codegen_binary_expr(CodeGen* g, BinaryExprNode* expr, GenMode gen_mo
     } else if (expr->op == TokenKind_star) {
         fprintf(g->out, "  imul rax, rdi\n");
     } else if (expr->op == TokenKind_slash) {
-        fprintf(g->out, "  cqo\n");
-        fprintf(g->out, "  idiv rdi\n");
+        if (type_is_unsigned(expr->lhs->ty)) {
+            fprintf(g->out, "  xor rdx, rdx\n");
+            fprintf(g->out, "  div rdi\n");
+        } else {
+            fprintf(g->out, "  cqo\n");
+            fprintf(g->out, "  idiv rdi\n");
+        }
     } else if (expr->op == TokenKind_percent) {
-        fprintf(g->out, "  cqo\n");
-        fprintf(g->out, "  idiv rdi\n");
+        if (type_is_unsigned(expr->lhs->ty)) {
+            fprintf(g->out, "  xor rdx, rdx\n");
+            fprintf(g->out, "  div rdi\n");
+        } else {
+            fprintf(g->out, "  cqo\n");
+            fprintf(g->out, "  idiv rdi\n");
+        }
         fprintf(g->out, "  mov rax, rdx\n");
     } else if (expr->op == TokenKind_and) {
         fprintf(g->out, "  and rax, rdi\n");
@@ -221,9 +231,12 @@ static void codegen_binary_expr(CodeGen* g, BinaryExprNode* expr, GenMode gen_mo
         fprintf(g->out, "  mov rcx, rdi\n");
         fprintf(g->out, "  shl rax, cl\n");
     } else if (expr->op == TokenKind_rshift) {
-        // TODO: check if the operand is signed or unsigned
         fprintf(g->out, "  mov rcx, rdi\n");
-        fprintf(g->out, "  sar rax, cl\n");
+        if (type_is_unsigned(expr->lhs->ty)) {
+            fprintf(g->out, "  shr rax, cl\n");
+        } else {
+            fprintf(g->out, "  sar rax, cl\n");
+        }
     } else if (expr->op == TokenKind_eq) {
         fprintf(g->out, "  cmp rax, rdi\n");
         fprintf(g->out, "  sete al\n");
@@ -274,11 +287,21 @@ static void codegen_assign_expr_helper(CodeGen* g, AssignExprNode* expr) {
     } else if (expr->op == TokenKind_assign_mul) {
         fprintf(g->out, "  imul rax, rdi\n");
     } else if (expr->op == TokenKind_assign_div) {
-        fprintf(g->out, "  cqo\n");
-        fprintf(g->out, "  idiv rdi\n");
+        if (type_is_unsigned(expr->lhs->ty)) {
+            fprintf(g->out, "  xor rdx, rdx\n");
+            fprintf(g->out, "  div rdi\n");
+        } else {
+            fprintf(g->out, "  cqo\n");
+            fprintf(g->out, "  idiv rdi\n");
+        }
     } else if (expr->op == TokenKind_assign_mod) {
-        fprintf(g->out, "  cqo\n");
-        fprintf(g->out, "  idiv rdi\n");
+        if (type_is_unsigned(expr->lhs->ty)) {
+            fprintf(g->out, "  xor rdx, rdx\n");
+            fprintf(g->out, "  div rdi\n");
+        } else {
+            fprintf(g->out, "  cqo\n");
+            fprintf(g->out, "  idiv rdi\n");
+        }
         fprintf(g->out, "  mov rax, rdx\n");
     } else if (expr->op == TokenKind_assign_or) {
         fprintf(g->out, "  or rax, rdi\n");
@@ -291,7 +314,11 @@ static void codegen_assign_expr_helper(CodeGen* g, AssignExprNode* expr) {
         fprintf(g->out, "  shl rax, cl\n");
     } else if (expr->op == TokenKind_assign_rshift) {
         fprintf(g->out, "  mov rcx, rdi\n");
-        fprintf(g->out, "  sar rax, cl\n");
+        if (type_is_unsigned(expr->lhs->ty)) {
+            fprintf(g->out, "  shr rax, cl\n");
+        } else {
+            fprintf(g->out, "  sar rax, cl\n");
+        }
     } else {
         unreachable();
     }