codegen label entrypoint fix

Author: eldiablolives

src/lower/arm64/codegen.c

@@ -11,6 +11,12 @@ static uint64_t cg_now_ns(void) {
   return (t * (uint64_t)tb.numer) / (uint64_t)tb.denom;
 }
 
+static const char *g_entry_label_override = NULL;
+
+void cg_set_entry_label_override(const char *label) {
+  g_entry_label_override = label;
+}
+
 typedef struct {
   const char *key;
   symtab_entry *val;
@@ -433,6 +439,7 @@ static void emit_adrp_add(cg_blob_t *out, cg_profile_t *prof, uint32_t *w, size_
 
 int cg_emit_arm64(const ir_prog_t *ir, cg_blob_t *out) {
   if (!ir || !out) return -1;
+  const char *entry_override = (g_entry_label_override && g_entry_label_override[0]) ? g_entry_label_override : NULL;
   const int profile_enabled = out->profile_enabled;
   memset(out,0,sizeof(*out));
   out->profile_enabled = profile_enabled;
@@ -512,14 +519,20 @@ int cg_emit_arm64(const ir_prog_t *ir, cg_blob_t *out) {
 
   /* Determine function entry labels: first label, labels named "main", labels prefixed with "fn_", and labels that are CALL targets. */
   const uint64_t t_func0 = prof ? cg_now_ns() : 0;
+  int override_found = 0;
   for (size_t i = 0; i < n_labels; i++) {
     int is_func = 0;
-    if (!first_label_seen) { is_func = 1; first_label_seen = 1; }
+    if (entry_override && label_list[i].name && strcmp(label_list[i].name, entry_override) == 0) { is_func = 1; override_found = 1; }
+    if (!entry_override && !first_label_seen) { is_func = 1; first_label_seen = 1; }
     if (label_list[i].name && strcmp(label_list[i].name, "main") == 0) is_func = 1;
     if (label_list[i].name && strncmp(label_list[i].name, "fn_", 3) == 0) is_func = 1;
     if (label_list[i].name && seen_has(call_targets, n_calls, label_list[i].name)) is_func = 1;
     if (is_func) func_count++;
   }
+  if (entry_override && !override_found) {
+    if (prof) prof->func_detect_ns = cg_now_ns() - t_func0;
+    CG_FAIL(NULL, "--entry-label symbol not found in input labels");
+  }
   if (func_count == 0 && any_instr) func_count = 1;
   if (prof) prof->func_detect_ns = cg_now_ns() - t_func0;
 
@@ -556,7 +569,8 @@ int cg_emit_arm64(const ir_prog_t *ir, cg_blob_t *out) {
     if (e->kind == IR_ENTRY_LABEL) {
       if (e->u.label.name) {
         int is_func = 0;
-        if (!func_seen) is_func = 1;
+        if (entry_override && strcmp(e->u.label.name, entry_override) == 0) is_func = 1;
+        if (!entry_override && !func_seen) is_func = 1;
         if (strcmp(e->u.label.name, "main") == 0) is_func = 1;
         if (strncmp(e->u.label.name, "fn_", 3) == 0) is_func = 1;
         if (seen_has(call_targets, n_calls, e->u.label.name)) is_func = 1;
@@ -566,6 +580,18 @@ int cg_emit_arm64(const ir_prog_t *ir, cg_blob_t *out) {
           func_has_prologue = 0;
         }
         symtab_update(&ctx, out->syms, e->u.label.name, pcw*4);
+
+        /* Emit the function prologue at the function entry label.
+         * This ensures any immediately-following label is placed after the
+         * prologue (distinct address), avoiding accidental stack growth when
+         * branching to a basic-block label that would otherwise share the
+         * function entry address.
+         */
+        if (is_func && !func_has_prologue) {
+          w[pcw++] = enc_stp_fp_lr();
+          w[pcw++] = enc_mov_fp_sp();
+          func_has_prologue = 1;
+        }
       }
       if (prof) prof->pass2_labels_ns += cg_now_ns() - t_e0;
       continue;

src/lower/arm64/codegen.h

@@ -64,4 +64,10 @@ typedef struct {
 int cg_emit_arm64(const ir_prog_t *ir, cg_blob_t *out);
 void cg_free(cg_blob_t *out);
 
+/* Optional: override the implicit entry-function heuristic.
+ * When set, cg_emit_arm64 will treat this label as a function entry.
+ * Pass NULL to clear.
+ */
+void cg_set_entry_label_override(const char *label);
+
 #endif /* ZINGCC_CODEGEN_H */

src/lower/arm64/main.c

@@ -44,6 +44,7 @@ static void usage(const char *prog) {
     "  --emit-pc-map Write JSON mapping of code offsets to IR ids/lines\n"
     "  --emit-lldb   Write an LLDB trace script to the given path\n"
     "  --trace-func  Function symbol to trace (default: main)\n"
+    "  --entry-label Override implicit entry label (disables first-label heuristic)\n"
     "  --trace-syms  Comma list of symbols to dump at trace breakpoints\n"
     "  --trace-regs  Comma list of registers (e.g. w0,x0) to dump\n"
     "  --profile     Print per-phase timings to stderr\n"
@@ -480,6 +481,7 @@ int main(int argc, char **argv) {
   int json_dump = 0;
   const char *emit_lldb_path = NULL;
   const char *trace_func = "main";
+  const char *entry_label = NULL;
   const char *trace_syms = NULL;
   const char *trace_regs = NULL;
   int profile = 0;
@@ -510,6 +512,7 @@ int main(int argc, char **argv) {
     if (strcmp(a, "--emit-pc-map") == 0 && argi < argc) { emit_pc_map_path = argv[argi++]; continue; }
     if (strcmp(a, "--emit-lldb") == 0 && argi < argc) { emit_lldb_path = argv[argi++]; continue; }
     if (strcmp(a, "--trace-func") == 0 && argi < argc) { trace_func = argv[argi++]; continue; }
+    if (strcmp(a, "--entry-label") == 0 && argi < argc) { entry_label = argv[argi++]; continue; }
     if (strcmp(a, "--trace-syms") == 0 && argi < argc) { trace_syms = argv[argi++]; continue; }
     if (strcmp(a, "--trace-regs") == 0 && argi < argc) { trace_regs = argv[argi++]; continue; }
     if (strcmp(a, "--input") == 0 && argi < argc) { inputs[nin++] = argv[argi++]; continue; }
@@ -544,6 +547,9 @@ int main(int argc, char **argv) {
     ir_prog_t prog;
     ir_init(&prog);
 
+    /* Configure optional entry label override for codegen. */
+    cg_set_entry_label_override(entry_label);
+
     const uint64_t t_parse0 = profile ? now_ns() : 0;
     if (json_ir_read(fp, &prog) != 0) {
       fprintf(stderr, "[lower] failed to parse IR: %s\n", in_path);

src/lower/arm64/tests/codegen_cases.jsonl

@@ -1,3 +1,8 @@
+{"ir":"zasm-v1.0","k":"label","name":"main"}
+{"ir":"zasm-v1.0","k":"label","name":"loop"}
+{"ir":"zasm-v1.0","k":"instr","m":"LD","ops":[{"t":"sym","v":"HL"},{"t":"num","v":0}]}
+{"ir":"zasm-v1.0","k":"instr","m":"JR","ops":[{"t":"sym","v":"loop"}]}
+
 {"ir":"zasm-v1.0","k":"dir","d":"DB","name":"msg","args":[{"t":"str","v":"Hi"},{"t":"num","v":0}]}
 {"ir":"zasm-v1.0","k":"dir","d":"RESB","name":"buf","args":[{"t":"num","v":9000}]}
 {"ir":"zasm-v1.0","k":"instr","m":"LD","ops":[{"t":"sym","v":"HL"},{"t":"num","v":42}]}

src/lower/arm64/tests/codegen_test.c

@@ -4,6 +4,20 @@
 #include <stdlib.h>
 #include <string.h>
 
+static symtab_entry *find_sym(symtab_entry *root, const char *name) {
+  for (symtab_entry *s = root; s; s = s->next) {
+    if (s->name && strcmp(s->name, name) == 0) return s;
+  }
+  return NULL;
+}
+
+static int ir_has_label(const ir_prog_t *ir, const char *name) {
+  for (ir_entry_t *e = ir->head; e; e = e->next) {
+    if (e->kind == IR_ENTRY_LABEL && e->u.label.name && strcmp(e->u.label.name, name) == 0) return 1;
+  }
+  return 0;
+}
+
 static int run_file(const char *path){
   FILE *f = fopen(path,"r");
   if (!f) { perror("open"); return 1; }
@@ -12,6 +26,25 @@ static int run_file(const char *path){
   fclose(f);
   cg_blob_t blob={0};
   if (cg_emit_arm64(&ir,&blob)!=0){fprintf(stderr,"codegen failed\n"); ir_free(&ir); return 1;}
+
+  /* Regression guard: if a stream begins with adjacent labels (e.g. main then loop)
+   * and lowerer function prologue insertion happens at the first instruction rather
+   * than at the function entry label, both labels can wind up with the same address.
+   * Branching to the second label can then re-run the prologue and grow the stack.
+   */
+  if (ir_has_label(&ir, "main") && ir_has_label(&ir, "loop")) {
+    symtab_entry *s_main = find_sym(blob.syms, "main");
+    symtab_entry *s_loop = find_sym(blob.syms, "loop");
+    if (s_main && s_loop && s_main->off != (size_t)-1 && s_loop->off != (size_t)-1) {
+      if (s_loop->off <= s_main->off) {
+        fprintf(stderr, "codegen regression: loop label (off=%zu) not after main (off=%zu)\n", s_loop->off, s_main->off);
+        cg_free(&blob);
+        ir_free(&ir);
+        return 1;
+      }
+    }
+  }
+
   printf("code %zu bytes, data %zu bytes, relocs %u\n", blob.code_len, blob.data_len, blob.reloc_count);
   cg_free(&blob);
   ir_free(&ir);