Omit epilogue if a function does not return

src/cc/arch/aarch64/emit_code.c

@@ -553,36 +553,38 @@ static void emit_defun(Function *func) {
 
   emit_bb_irs(fnbe->bbcon);
 
-  // Epilogue
-  if (!no_stmt) {
-    if (func->flag & FUNCF_STACK_MODIFIED) {
-      // Stack pointer might be changed if alloca is used, so it need to be recalculated.
-      size_t size = frame_size + ALIGN(callee_saved_count * POINTER_SIZE, 16);
-      const char *value;
-      if (size <= 0x0fff)
-        value = IM(size);
-      else
-        mov_immediate(value = SP, size, true, false);
-      SUB(SP, FP, value);
-    }
-    pop_callee_save_regs(used_reg_bits, fnbe->ra->used_freg_bits);
-    if (fp_saved) {
-      const char *value;
-      if (frame_size <= 0x0fff) {
-        value = IM(frame_size);
-      } else {
-        // Break x17
-        mov_immediate(value = X17, frame_size, true, false);
+  if (!function_not_returned(fnbe)) {
+    // Epilogue
+    if (!no_stmt) {
+      if (func->flag & FUNCF_STACK_MODIFIED) {
+        // Stack pointer might be changed if alloca is used, so it need to be recalculated.
+        size_t size = frame_size + ALIGN(callee_saved_count * POINTER_SIZE, 16);
+        const char *value;
+        if (size <= 0x0fff)
+          value = IM(size);
+        else
+          mov_immediate(value = SP, size, true, false);
+        SUB(SP, FP, value);
       }
-      ADD(SP, SP, value);
+      pop_callee_save_regs(used_reg_bits, fnbe->ra->used_freg_bits);
+      if (fp_saved) {
+        const char *value;
+        if (frame_size <= 0x0fff) {
+          value = IM(frame_size);
+        } else {
+          // Break x17
+          mov_immediate(value = X17, frame_size, true, false);
+        }
+        ADD(SP, SP, value);
+      }
+
+      if (fp_saved || lr_saved)
+        LDP(FP, LR, POST_INDEX(SP, 16));
     }
 
-    if (fp_saved || lr_saved)
-      LDP(FP, LR, POST_INDEX(SP, 16));
+    RET();
   }
 
-  RET();
-
   // Static variables are emitted through global variables.
 }
 

src/cc/arch/x64/emit_code.c

@@ -543,31 +543,35 @@ static void emit_defun(Function *func) {
 
   emit_bb_irs(fnbe->bbcon);
 
-  // Epilogue
-  if (!no_stmt) {
-    if (func->flag & FUNCF_STACK_MODIFIED) {
-      // Stack pointer might be changed if alloca is used, so it need to be recalculated.
-      LEA(OFFSET_INDIRECT(callee_saved_count * -POINTER_SIZE - frame_size, RBP, NULL, 1),
-          RSP);
-    }
+  if (!function_not_returned(fnbe)) {
+    // Epilogue
+    if (!no_stmt) {
+      if (func->flag & FUNCF_STACK_MODIFIED) {
+        // Stack pointer might be changed if alloca is used, so it need to be recalculated.
+        LEA(OFFSET_INDIRECT(callee_saved_count * -POINTER_SIZE - frame_size, RBP, NULL, 1),
+            RSP);
+      }
 
-    pop_callee_save_regs(fnbe->ra->used_reg_bits, fnbe->ra->used_freg_bits);
+      pop_callee_save_regs(fnbe->ra->used_reg_bits, fnbe->ra->used_freg_bits);
 
-    if (rbp_saved) {
-      MOV(RBP, RSP);
-      stackpos -= frame_size;
-      POP(RBP); POP_STACK_POS();
-    } else if (frame_size > 0) {
-      ADD(IM(frame_size), RSP);
-      stackpos -= frame_size;
+      if (rbp_saved) {
+        MOV(RBP, RSP);
+        stackpos -= frame_size;
+        POP(RBP); POP_STACK_POS();
+      } else if (frame_size > 0) {
+        ADD(IM(frame_size), RSP);
+        stackpos -= frame_size;
+      }
     }
-  }
 
-  RET();
+    RET();
 
-  // Static variables are emitted through global variables.
+    assert(stackpos == 8);
+  } else {
+    stackpos = 8;
+  }
 
-  assert(stackpos == 8);
+  // Static variables are emitted through global variables.
 }
 
 void emit_code(Vector *decls) {

src/cc/backend/emit_util.c

@@ -7,6 +7,7 @@
 #include <stdint.h>  // int64_t
 #include <stdlib.h>  // realloc
 
+#include "ir.h"
 #include "table.h"
 #include "util.h"
 
@@ -147,3 +148,15 @@ void emit_bss(const char *label, size_t size, size_t align) {
 void init_emit(FILE *fp) {
   emit_fp = fp;
 }
+
+bool function_not_returned(FuncBackend *fnbe) {
+  BB *bb = fnbe->bbcon->bbs->data[fnbe->bbcon->bbs->len - 1];
+  if (bb->irs->len > 0) {
+    IR *ir = bb->irs->data[bb->irs->len - 1];
+    if (ir->kind == IR_JMP && ir->jmp.cond == COND_ANY && ir->jmp.bb != NULL) {
+      // No fallthrough exists: the function does not return.
+      return true;
+    }
+  }
+  return false;
+}

src/cc/backend/emit_util.h

@@ -2,9 +2,12 @@
 
 #pragma once
 
+#include <stdbool.h>
 #include <stdint.h>  // int64_t
 #include <stdio.h>
 
+typedef struct FuncBackend FuncBackend;
+
 #ifndef EMIT_LABEL
 #define EMIT_LABEL(label)  emit_label(label)
 #endif
@@ -36,3 +39,5 @@ void emit_asm4(const char *op, const char *a1, const char *a2, const char *a3, c
 void emit_align_p2(int align);
 void emit_comment(const char *comment, ...);
 void emit_bss(const char *label, size_t size, size_t align);
+
+bool function_not_returned(FuncBackend *fnbe);

src/cc/backend/optimize.c

@@ -66,14 +66,12 @@ static void remove_unnecessary_bb(BBContainer *bbcon) {
   for (;;) {
     table_init(&keeptbl);
     assert(bbs->len > 0);
-    BB *bb0 = bbs->data[0];
-    table_put(&keeptbl, bb0->label, bb0);
 
-    for (int i = 0; i < bbs->len - 1; ++i) {
+    for (int i = 0; i < bbs->len; ++i) {
       BB *bb = bbs->data[i];
       bool remove = false;
       IR *ir_jmp = is_last_jmp(bb);
-      if (bb->irs->len == 0) {  // Empty BB.
+      if (bb->irs->len == 0 && bb->next != NULL) {  // Empty BB.
         replace_jmp_destination(bbcon, bb, bb->next);
         remove = true;
       } else if (bb->irs->len == 1 && ir_jmp != NULL && ir_jmp->jmp.cond == COND_ANY &&
@@ -95,7 +93,7 @@ static void remove_unnecessary_bb(BBContainer *bbcon) {
 
       if (ir_jmp != NULL)
         table_put(&keeptbl, ir_jmp->jmp.bb->label, bb);
-      if (ir_jmp == NULL || ir_jmp->jmp.cond != COND_ANY)
+      if ((ir_jmp == NULL || ir_jmp->jmp.cond != COND_ANY) && bb->next != NULL)
         table_put(&keeptbl, bb->next->label, bb);
 
       IR *tjmp = is_last_jtable(bb);
@@ -110,7 +108,7 @@ static void remove_unnecessary_bb(BBContainer *bbcon) {
     }
 
     bool again = false;
-    for (int i = 0; i < bbs->len; ++i) {
+    for (int i = 1; i < bbs->len; ++i) {
       BB *bb = bbs->data[i];
       if (!table_try_get(&keeptbl, bb->label, NULL)) {
         if (i > 0) {

tests/test.sh

@@ -180,6 +180,7 @@ test_function() {
   compile_error 'conflict func' 'void main(); int main(int, char**){return 0;}'
   compile_error 'duplicate var & func' 'int main; int main(){return 0;}'
   compile_error 'duplicate func & var' 'int main(){return 0;} int main;'
+  try_direct 'infinite loop and exit' 77 '#include <stdlib.h>\nint main(){for (int i = 0; ; ++i) if (i == 10) exit(77);}'
 
   end_test_suite
 }