handling early returns inside the functions

KagatiFoundation · Oct 19, 2024 · 6ec546a · 6ec546a
1 parent dab4946
commit 6ec546a
Show file tree

Hide file tree

Showing 5 changed files with 132 additions and 30 deletions.
diff --git a/examples/input1.bic b/examples/input1.bic
@@ -1,15 +1,17 @@
-def extern print_int(x: integer) -> void;
-def extern print(x: str) -> void;
-
-def plus(x: integer, y: integer) -> integer {
-    if (x > y) {
-        return x - y;
-    }
-    return x + y;
-}
+def extern sock(port: integer) -> integer;
+def extern sock_listen(fd: integer, backlog: integer) -> void;
+def extern sock_accept(fd: integer) -> integer;
+def extern sock_write(fd: integer, msg: str, size: integer) -> void;
+def extern sock_close(fd: integer) -> void;
 
 def main() -> void {
-    let a: integer = 12;
-    let b: integer = plus(a, 2);
-    print_int(b);
+    let fd = sock(8080);
+    sock_listen(fd, 3);
+    let new_fd = sock_accept(fd);
+
+    let http_response = "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: 25\r\n\r\n<b>Hello from Bichara</b>";
+
+    sock_write(new_fd, http_response, 89);
+    sock_close(new_fd);
+    sock_close(fd);
 }
diff --git a/examples/sock.bich b/examples/sock.bich
@@ -0,0 +1,17 @@
+def extern sock(port: integer) -> integer;
+def extern sock_listen(fd: integer, backlog: integer) -> void;
+def extern sock_accept(fd: integer) -> integer;
+def extern sock_write(fd: integer, msg: str, size: integer) -> void;
+def extern sock_close(fd: integer) -> void;
+
+def main() -> void {
+    let fd = sock(8080);
+    sock_listen(fd, 3);
+    let new_fd = sock_accept(fd);
+
+    let http_response = "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: 25\r\n\r\n<b>Hello from Bichara</b>";
+
+    sock_write(new_fd, http_response, 89);
+    sock_close(new_fd);
+    sock_close(fd);
+}
diff --git a/examples/sock.c b/examples/sock.c
@@ -0,0 +1,65 @@
+// Server side C program to demonstrate Socket
+// programming
+#include <netinet/in.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/socket.h>
+#include <unistd.h>
+
+struct sockaddr_in address;
+socklen_t addrlen;
+
+void sock_close(int fd) {
+    close(fd);
+}
+
+void sock_write(int fd, const char* msg, int size) {
+    send(fd, msg, size, 0);
+}
+
+void sock_listen(int fd, int backlog) {
+    if (listen(fd, backlog) < 0) {
+        perror("listen");
+        exit(EXIT_FAILURE);
+    }
+    puts("Socket listen");
+}
+
+int sock_accept(int fd) {
+    int new_socket = accept(fd, (struct sockaddr*)&address, &addrlen);
+    if (new_socket < 0) {
+        perror("accept");
+        exit(EXIT_FAILURE);
+    }
+    puts("New socket connection accepted");
+    return new_socket;
+}
+
+int sock(int port) {
+    int opt = 1;
+    addrlen = sizeof(address);
+    char buffer[1024] = { 0 };
+
+    int server_fd = socket(AF_INET, SOCK_STREAM, 0);
+    if (server_fd < 0) {
+        perror("socket failed");
+        exit(EXIT_FAILURE);
+    }
+    puts("Socket open");
+
+    if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt))) {
+        perror("setsockopt");
+        exit(EXIT_FAILURE);
+    }
+    address.sin_family = AF_INET;
+    address.sin_addr.s_addr = INADDR_ANY;
+    address.sin_port = htons(port);
+
+    if (bind(server_fd, (struct sockaddr*)&address, sizeof(address)) < 0) {
+        perror("bind failed");
+        exit(EXIT_FAILURE);
+    }
+    puts("Socket bound");
+    return server_fd;
+}
diff --git a/src/code_gen/aarch64/aarch64_codegen.rs b/src/code_gen/aarch64/aarch64_codegen.rs
@@ -35,6 +35,7 @@ use crate::ast::{
     AST
 };
 use crate::code_gen::codegen::CodeGenResult;
+use crate::code_gen::codegen::EARLY_RETURN;
 use crate::code_gen::codegen::NO_REG;
 use crate::code_gen::CodeGen;
 use crate::code_gen::register::RegManager;
@@ -62,7 +63,8 @@ pub struct Aarch64CodeGen<'aarch64> {
     reg_manager: RefCell<RegManager>,
     ctx: Option<Rc<RefCell<CompilerCtx<'aarch64>>>>,
     label_id: usize,
-    function_id: usize
+    function_id: usize,
+    early_return_label_id: usize
 }
 
 impl<'aarch64> CodeGen for Aarch64CodeGen<'aarch64> {
@@ -106,13 +108,13 @@ impl<'aarch64> CodeGen for Aarch64CodeGen<'aarch64> {
         let label_end: usize = self.get_next_label();
 
         // Evaluate the condition and store the result in a register
-        let _cond_result_reg: usize = self.gen_code_from_ast(ast.left.as_ref().unwrap(), label_if_false, ast.operation)?;
+        let _cond_result_reg: usize = self.gen_code_from_ast(ast.left.as_ref().unwrap(), label_if_false, ASTOperation::AST_IF)?;
 
         // Free registers to allow usage within the if-else body
         self.reg_manager.borrow_mut().deallocate_all();
 
         // Generate code for the 'if-true' block
-        _ = self.gen_code_from_ast(ast.mid.as_ref().unwrap(), reg, ast.operation)?;
+        _ = self.gen_code_from_ast(ast.mid.as_ref().unwrap(), reg, ASTOperation::AST_IF)?;
 
         // Free registers again to prepare for any subsequent operations
         self.reg_manager.borrow_mut().deallocate_all();
@@ -127,7 +129,7 @@ impl<'aarch64> CodeGen for Aarch64CodeGen<'aarch64> {
 
         // Generate code for the 'else' block, if it exists
         if let Some(ref right_ast) = ast.right {
-            _ = self.gen_code_from_ast(right_ast, reg, ast.operation)?;
+            _ = self.gen_code_from_ast(right_ast, reg, ASTOperation::AST_IF)?;
             self.reg_manager.borrow_mut().deallocate_all();
             _ = self.gen_label(label_end)?; // Mark the end of the if-else block
         }
@@ -228,6 +230,10 @@ impl<'aarch64> CodeGen for Aarch64CodeGen<'aarch64> {
 
         if let Some(ref body) = ast.left {
             _ = self.gen_code_from_ast(body, 0xFFFFFFFF, ast.operation)?;
+            if self.early_return_label_id != NO_REG {
+                println!("_L{}:", self.early_return_label_id);
+                self.early_return_label_id = NO_REG; // reset early return label after function code generation
+            }
         }
         // ldp -> Load Pair of Registers
         // Restore the saved frame pointer (x29) and link register (x30) from the stack.
@@ -283,7 +289,7 @@ impl<'aarch64> CodeGen for Aarch64CodeGen<'aarch64> {
             let ctx_borrow = ctx_rc.borrow();
             ctx_borrow.sym_table.get_symbol(id).unwrap().clone()
         } else {
-            panic!("Please provide a context to work on!");
+            return Err(CodeGenErr::NoContext);
         };
         if symbol.class == StorageClass::GLOBAL {
             let reg: usize = self.reg_manager.borrow_mut().allocate();
@@ -383,11 +389,23 @@ impl<'aarch64> CodeGen for Aarch64CodeGen<'aarch64> {
         Ok(r1)
     }
 
-    fn gen_return_stmt(&mut self, _result_reg: usize, _func_id: usize) -> CodeGenResult {
+    fn gen_return_stmt(&mut self, result_reg: usize, early_return: bool) -> CodeGenResult {
         // NOTE: Generate code depending on the function's type. i.e. use w0 for i32, x0 for i64 etc.
         // let func_ret_type: LitTypeVariant = self.sym_table.get_symbol(func_id).lit_type;
-
-        Ok(NO_REG)
+        // is it an early return? 
+        if self.early_return_label_id == NO_REG {
+            self.early_return_label_id = self.get_next_label();
+        }
+
+        if early_return {
+            println!("b _L{}", self.early_return_label_id);
+        }
+
+        if early_return {
+            Ok(EARLY_RETURN)
+        } else {
+            Ok(NO_REG)
+        }
     }
 
     fn gen_array_access(&mut self, id: usize, expr: &AST) -> CodeGenResult {
@@ -535,13 +553,9 @@ impl<'aarch64> CodeGen for Aarch64CodeGen<'aarch64> {
             // let param_reg_name: String = self.reg_manager.borrow().name(param_reg);
             // println!("mov {}, {}", param_reg_name, reg_name);
         }
-        let mut reg_mgr = self.reg_manager.borrow_mut();
-        let alloced_reg: usize = reg_mgr.allocate();
-        let alloced_reg_name: String = reg_mgr.name(alloced_reg);
         println!("bl _{}", func_info.name);
-        println!("mov {}, x0", alloced_reg_name);
         self.function_id = 0xFFFFFFFF;
-        Ok(alloced_reg)
+        Ok(0) // always return the 'x0' register's index after function calls
     }
 
     fn gen_var_assignment_stmt(&mut self, assign_stmt: &crate::ast::AssignStmt, expr_ast: &Expr) -> CodeGenResult {
@@ -559,7 +573,8 @@ impl<'aarch64> Aarch64CodeGen<'aarch64> {
             reg_manager,
             ctx: None,
             label_id: 0,
-            function_id: 0xFFFFFFFF
+            function_id: 0xFFFFFFFF,
+            early_return_label_id: 0xFFFFFFFF
         }
     }
 

diff --git a/src/code_gen/codegen.rs b/src/code_gen/codegen.rs
@@ -47,6 +47,8 @@ pub type CodeGenResult = Result<usize, CodeGenErr>;
 /// Indicating no register was produced from an code generation operation.
 pub const NO_REG: usize = 0xFFFFFFFF;
 
+pub const EARLY_RETURN: usize = 0xEEEEEEEE;
+
 pub trait CodeGen {
     /// Starts the code generation process.
     ///
@@ -131,11 +133,11 @@ pub trait CodeGen {
         }
         else if ast_node.operation == ASTOperation::AST_GLUE {
             if let Some(left) = ast_node.left.as_ref() {
-                _ = self.gen_code_from_ast(left, reg, ast_node.operation);
+                _ = self.gen_code_from_ast(left, reg, parent_ast_kind);
                 self.reg_manager().deallocate_all();
             }
             if let Some(right) = ast_node.right.as_ref() {
-                _ = self.gen_code_from_ast(right, reg, ast_node.operation);
+                _ = self.gen_code_from_ast(right, reg, parent_ast_kind);
                 self.reg_manager().deallocate_all();
             }
             Ok(NO_REG)
@@ -154,11 +156,12 @@ pub trait CodeGen {
             Ok(NO_REG)
         }
         else if ast_node.operation == ASTOperation::AST_RETURN {
+            let early_return = parent_ast_kind != ASTOperation::AST_FUNCTION;
             let possible_ret_stmt: Stmt = ast_node.kind.clone().unwrap_stmt();
             let return_expr: &AST = ast_node.left.as_ref().unwrap();
             let result_reg: usize = self.gen_expr(&return_expr.kind.clone().unwrap_expr(), ast_node.operation, reg, parent_ast_kind)?;
             return match possible_ret_stmt {
-                Stmt::Return(ret) => self.gen_return_stmt(result_reg, ret.func_id),
+                Stmt::Return(_) => self.gen_return_stmt(result_reg, early_return),
                 _ => Ok(NO_REG)
             };
         }
@@ -348,7 +351,7 @@ pub trait CodeGen {
 
     fn gen_array_access2(&mut self, symbol_id: usize, index: usize) -> CodeGenResult;
 
-    fn gen_return_stmt(&mut self, result_reg: usize, _func_id: usize) -> CodeGenResult;
+    fn gen_return_stmt(&mut self, result_reg: usize, early_return: bool) -> CodeGenResult;
 
     fn gen_func_call_stmt(&mut self, func_call_stmt: &FuncCallStmt) -> CodeGenResult;