From 37a9cf0400aa2cc6615221ad0549bd10a21c65d8 Mon Sep 17 00:00:00 2001
From: sBubshait <s.bobshit@gmail.com>
Date: Fri, 14 Jun 2024 20:20:54 +0100
Subject: [PATCH] Update encode to handle branches with sign extend

---
 src/encode.c | 43 +++++++++++++++++++++++++++++++------------
 1 file changed, 31 insertions(+), 12 deletions(-)

diff --git a/src/encode.c b/src/encode.c
index 1f7ff7a..2f7cc0f 100644
--- a/src/encode.c
+++ b/src/encode.c
@@ -25,22 +25,41 @@ void setBits(word* wrd, uint8_t lsb, uint8_t msb, word value) {
     *wrd |= (value << lsb) & mask;    
 }
 
-/*
-static int getLabelOffset(symbol_table* table, char* label, int currentIndex) {
-    address target = st_get(table, label);
-    return target - (currentIndex * 4);
-}
-*/
+// Sign extend a given value to a 64-bit signed integer given the number of bits
+int64_t signExtend(dword value, unsigned int n) {
+    if (n == 0 || n >= 64) {
+        // If n_bits is 0 or greater than or equal to 64, return the value as is
+        return (int64_t)value;
+    }
+    
+    uint64_t sign_bit_mask = (uint64_t)1 << (n - 1);
+    
+    // Mask to isolate the n-bit value
+    uint64_t n_bit_mask = (sign_bit_mask << 1) - 1;
 
+    // Check if the sign bit is set
+    if (value & sign_bit_mask) {
+        // Sign bit is set, extend the sign
+        return (int64_t)(value | ~n_bit_mask);
+    } else {
+        // Sign bit is not set, return the value as is
+        return (int64_t)(value & n_bit_mask);
+    }
+}
+
+static int getLabelOffset(symbol_table* table, char* label, int currentIndex, int n_bits) {
+    address target = st_get(table, label);
+    return signExtend((unsigned int) (target - (currentIndex * 4)), n_bits);
+}
 
 // Generates assembled code based on the two-pass assembly method
-word encodeBranch(a64inst_instruction *instr) {
+word encodeBranch(a64inst_instruction *instr, int index, symbol_table *st) {
     word wrd = 0;
 
     switch (instr->data.BranchData.BranchType) {
     case a64inst_UNCONDITIONAL:
-        setBits(&wrd, 26, 30, 0x5);
-        setBits(&wrd, 0, 25, instr->data.BranchData.processOpData.unconditionalData.unconditionalOffset);
+        setBits(&wrd, 26, 29, 0x5);
+        setBits(&wrd, 0, 25, getLabelOffset(st, instr->data.BranchData.processOpData.unconditionalData.label, index, 26));
         break;
 
     case a64inst_REGISTER:
@@ -50,7 +69,7 @@ word encodeBranch(a64inst_instruction *instr) {
 
     case a64inst_CONDITIONAL:
         setBits(&wrd, 26, 32, 0x15);
-        setBits(&wrd, 5, 24, instr->data.BranchData.processOpData.conditionalData.offset);
+        setBits(&wrd, 5, 24, getLabelOffset(st, instr->data.BranchData.processOpData.conditionalData.label, index, 19));
         setBits(&wrd, 0, 4, instr->data.BranchData.processOpData.conditionalData.cond);
         break;
     }
@@ -168,7 +187,7 @@ word encodeLoadLiteral(a64inst_instruction cI) {
     return wrd;
 }
 
-word *encode(a64inst_instruction insts[], int instCount, symbol_table* table) {
+word *encode(a64inst_instruction insts[], int instCount, symbol_table* st) {
     word *arr = (word*)malloc(sizeof(word) * instCount);
     int index = 0;
     for (int i = 0; i < instCount; i++) {
@@ -202,7 +221,7 @@ word *encode(a64inst_instruction insts[], int instCount, symbol_table* table) {
                 // Labels are handled in the first pass and used for addressing.
                 break;
             case a64inst_BRANCH:
-                arr[index] = encodeBranch(&inst);
+                arr[index] = encodeBranch(&inst, index, st);
                 index++;
             default:
                 break;