Add -goto_for_early_return transformation

Adds a -goto_for_early_return CLI option. When enabled, functions with
early returns are transformed to use goto-based control flow:
- A result variable is allocated at the top (for non-void functions)
- Every early return is replaced by assignment + goto
- A label and final return are appended at the end

The pass operates on the CStar AST (before lowering to C11) for simpler
type access, natural block flattening, and use of existing helpers.

Detection uses a position-aware walk that distinguishes terminal returns
(in tail position) from early returns (in non-terminal position).

Includes tests covering int32, uint32, uint64, int64, bool, void, struct,
name collision avoidance, no-early-return passthrough, and unit/void
interaction.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: nikswamy

lib/C11.ml

@@ -139,6 +139,8 @@ and stmt =
   | Continue
   | Comment of string
     (** note: this is not in the C grammar *)
+  | Goto of ident
+  | Label of ident
 
 and program =
   declaration_or_function list

lib/CStar.ml

@@ -45,6 +45,8 @@ and stmt =
   | BufFill of typ * expr * expr * expr
   | BufFree of typ * expr
   | Block of block
+  | Goto of ident
+  | Label of ident
   | Comment of string
 
 and expr =

lib/CStarToC11.ml

@@ -109,6 +109,8 @@ and vars_of_stmt m = function
   | Return _
   | Break
   | Continue
+  | Goto _
+  | Label _
   | Comment _ ->
       S.empty
   | Ignore e
@@ -637,6 +639,12 @@ and mk_stmt m (stmt: stmt): C.stmt list =
   | Block stmts ->
       [ Compound (mk_stmts m stmts) ]
 
+  | Goto label ->
+      [ Goto label ]
+
+  | Label label ->
+      [ Label label ]
+
   | Break ->
       [ Break ]
 

lib/GotoForEarlyReturn.ml

@@ -0,0 +1,172 @@
+(* Copyright (c) INRIA and Microsoft Corporation. All rights reserved. *)
+(* Licensed under the Apache 2.0 and MIT Licenses. *)
+
+(** Transform functions with early returns to use goto. When activated via
+    -goto_for_early_return, any function whose body contains a return statement
+    in non-tail position is rewritten so that:
+    - a return variable is allocated at the top (for non-void functions),
+    - every return is replaced by an assignment + goto,
+    - a label and final return are appended at the end.
+
+    This pass operates on the CStar AST, before lowering to C11. *)
+
+open CStar
+
+(* Collect all identifier names appearing in a block so we can pick fresh
+   names that don't collide. We walk expressions and statements collecting
+   variable references, declaration names, binder names, etc. *)
+let collect_names (body: block): (string, unit) Hashtbl.t =
+  let names = Hashtbl.create 32 in
+  let add n = Hashtbl.replace names n () in
+  let rec collect_expr (e: expr) =
+    match e with
+    | Var v -> add v
+    | Qualified (_, n) | Macro (_, n) -> add n
+    | Constant _ | Bool _ | StringLiteral _ | Any | EAbort _ | Type _
+    | BufNull | Op _ -> ()
+    | Cast (e, _) | Field (e, _) | AddrOf e | InlineComment (_, e, _) ->
+        collect_expr e
+    | BufRead (e1, e2) | BufSub (e1, e2) | Comma (e1, e2)
+    | BufCreate (_, e1, e2) ->
+        collect_expr e1; collect_expr e2
+    | Call (e, es) -> collect_expr e; List.iter collect_expr es
+    | BufCreateL (_, es) -> List.iter collect_expr es
+    | Struct (_, fes) -> List.iter (fun (_, e) -> collect_expr e) fes
+    | Stmt ss -> List.iter collect_stmt ss
+  and collect_stmt (s: stmt) =
+    match s with
+    | Abort _ | Break | Continue | Comment _ | Goto _ | Label _ -> ()
+    | Return (Some e) -> collect_expr e
+    | Return None -> ()
+    | Ignore e | BufFree (_, e) -> collect_expr e
+    | Decl (b, e) -> add b.name; collect_expr e
+    | Assign (e1, _, e2) -> collect_expr e1; collect_expr e2
+    | BufWrite (e1, e2, e3) ->
+        collect_expr e1; collect_expr e2; collect_expr e3
+    | BufBlit (_, e1, e2, e3, e4, e5) ->
+        List.iter collect_expr [e1; e2; e3; e4; e5]
+    | BufFill (_, e1, e2, e3) ->
+        List.iter collect_expr [e1; e2; e3]
+    | IfThenElse (_, e, b1, b2) ->
+        collect_expr e; collect_block b1; collect_block b2
+    | While (e, b) -> collect_expr e; collect_block b
+    | For (init, e, iter, b) ->
+        (match init with
+         | `Decl (bi, e') -> add bi.name; collect_expr e'
+         | `Stmt s -> collect_stmt s
+         | `Skip -> ());
+        collect_expr e; collect_stmt iter; collect_block b
+    | Switch (e, branches, default) ->
+        collect_expr e;
+        List.iter (fun (_, b) -> collect_block b) branches;
+        Option.iter collect_block default
+    | Block b -> collect_block b
+  and collect_block b = List.iter collect_stmt b
+  in
+  collect_block body;
+  names
+
+(* Position-aware early-return detection. A return in "terminal" position
+   (i.e., last statement in the function body, or in a branch of an
+   if-then-else/switch that is itself in terminal position) is NOT early.
+   Any return in non-terminal position IS early. *)
+let has_early_return (body: block): bool =
+  let found = ref false in
+  (* Walk statements. [terminal] tracks whether the current position could
+     be the "last thing" before the function returns naturally. *)
+  let rec walk_block ~terminal (stmts: block) =
+    match stmts with
+    | [] -> ()
+    | [s] -> walk_stmt ~terminal s
+    | s :: rest -> walk_stmt ~terminal:false s; walk_block ~terminal rest
+  and walk_stmt ~terminal (s: stmt) =
+    match s with
+    | Return _ ->
+        if not terminal then found := true
+    | IfThenElse (_, _, b1, b2) ->
+        walk_block ~terminal b1;
+        walk_block ~terminal b2
+    | Switch (_, branches, default) ->
+        List.iter (fun (_, b) -> walk_block ~terminal b) branches;
+        Option.iter (walk_block ~terminal) default
+    | While (_, b) ->
+        (* Loop body is not terminal — the loop may exit without returning *)
+        walk_block ~terminal:false b
+    | For (_, _, _, b) ->
+        walk_block ~terminal:false b
+    | Block b ->
+        walk_block ~terminal b
+    | Abort _ | Break | Continue | Comment _ | Goto _ | Label _
+    | Ignore _ | Decl _ | Assign _ | BufWrite _ | BufBlit _ | BufFill _
+    | BufFree _ ->
+        ()
+  in
+  walk_block ~terminal:true body;
+  !found
+
+(* Generate a type-specific zero initializer expression for the return
+   variable. *)
+let zero_initializer (t: typ): expr =
+  match t with
+  | Int w -> Constant (w, "0")
+  | Bool -> Bool false
+  | Pointer _ -> BufNull
+  | Void -> failwith "zero_initializer called on Void"
+  | Qualified _ | Struct _ | Enum _ | Union _ | Array _ | Function _
+  | Const _ ->
+      (* For aggregate/named types, produce a struct literal with a single
+         zero field. CStarToC11 translates this to { 0U }. *)
+      Struct (None, [(None, Constant (Constant.UInt8, "0"))])
+
+(* Rewrite a block, replacing Return statements with assignment + goto. *)
+let rewrite_block (ret_var: ident) (ret_typ: typ) (label: ident) (is_void: bool) (body: block): block =
+  let rec rewrite_stmts (stmts: block): block =
+    List.concat_map rewrite_one stmts
+  and rewrite_one (s: stmt): block =
+    match s with
+    | Return (Some e) when not is_void ->
+        [Assign (Var ret_var, ret_typ, e); Goto label]
+    | Return (Some _e) when is_void ->
+        Warn.fatal_error "void function returning a value"
+    | Return None ->
+        [Goto label]
+    | IfThenElse (ifdef, e, b1, b2) ->
+        [IfThenElse (ifdef, e, rewrite_stmts b1, rewrite_stmts b2)]
+    | Switch (e, branches, default) ->
+        [Switch (e,
+          List.map (fun (c, b) -> (c, rewrite_stmts b)) branches,
+          Option.map rewrite_stmts default)]
+    | While (e, b) ->
+        [While (e, rewrite_stmts b)]
+    | For (init, e, iter, b) ->
+        [For (init, e, iter, rewrite_stmts b)]
+    | Block b ->
+        [Block (rewrite_stmts b)]
+    | s -> [s]
+  in
+  rewrite_stmts body
+
+(* Rewrite a single CStar declaration if it has early returns. *)
+let rewrite_decl (d: decl): decl =
+  match d with
+  | Function (cc, flags, ret_typ, lid, binders, body) when has_early_return body ->
+      let is_void = (ret_typ = Void) in
+      let used = collect_names body in
+      let is_used n = Hashtbl.mem used n in
+      let ret_var = Idents.mk_fresh "result" is_used in
+      let label = Idents.mk_fresh "exit" is_used in
+      let rewritten = rewrite_block ret_var ret_typ label is_void body in
+      let new_body =
+        if is_void then
+          rewritten @ [Label label]
+        else
+          let init = zero_initializer ret_typ in
+          [Decl ({ name = ret_var; typ = ret_typ }, init)]
+          @ rewritten
+          @ [Label label; Return (Some (Var ret_var))]
+      in
+      Function (cc, flags, ret_typ, lid, binders, new_body)
+  | d -> d
+
+let rewrite_file (decls: decl list): decl list =
+  List.map rewrite_decl decls

lib/Options.ml

@@ -166,3 +166,5 @@ let allow_tapps = ref false
 (* Rust only: foo_bar_baz.fst gets emitted as foo/bar_baz.fst with depth=1 and
    foo/bar/baz.fst with depth = 2 (you get the idea). *)
 let depth = ref 1
+
+let goto_for_early_return = ref false

lib/PrintC.ml

@@ -493,6 +493,10 @@ and p_stmt (s: stmt) =
      string "break" ^^ semi
   | Continue ->
      string "continue" ^^ semi
+  | Goto label ->
+     group (string "goto" ^^ space ^^ string label ^^ semi)
+  | Label label ->
+     group (string label ^^ colon ^^ space ^^ semi)
 
 and p_stmts stmts = separate_map hardline p_stmt stmts
 

src/Karamel.ml

@@ -357,6 +357,9 @@ Supported options:|}
       one of the included system headers";
     "-faggressive-inlining", Arg.Set Options.aggressive_inlining, " attempt to inline \
       every variable for more compact code-generation";
+    "-goto_for_early_return", Arg.Set Options.goto_for_early_return, " replace early \
+      returns with assignments to a return variable and gotos to a label at the end \
+      of the function";
     "", Arg.Unit (fun _ -> ()), " ";
 
     (* For developers *)
@@ -806,6 +809,14 @@ Supported options:|}
 
     let files = List.filter (fun (_, decls) -> List.length decls > 0) files in
 
+    (* Apply goto-for-early-return on C* before lowering to C11 *)
+    let files =
+      if !Options.goto_for_early_return then
+        List.map (fun (name, decls) -> name, GotoForEarlyReturn.rewrite_file decls) files
+      else
+        files
+    in
+
     (* ... then to C *)
     let headers = CStarToC11.mk_headers c_name_map files in
     let deps = CStarToC11.drop_empty_headers deps headers in

test/Makefile

@@ -90,7 +90,7 @@ FSTAR		= $(FSTAR_EXE) \
   --trivial_pre_for_unannotated_effectful_fns false \
   --cmi --warn_error -274
 
-all: $(FILES) rust $(WASM_FILES) $(CUSTOM) ctypes-test sepcomp-test rust-val-test bundle-test rust-propererasure-bundle-test
+all: $(FILES) rust $(WASM_FILES) $(CUSTOM) ctypes-test sepcomp-test rust-val-test bundle-test rust-propererasure-bundle-test goto-return-test
 
 # Needs node
 wasm: $(WASM_FILES)
@@ -314,6 +314,10 @@ rust-propererasure-bundle-test:
 	+$(MAKE) -C rust-propererasure-bundle
 .PHONY: rust-propererasure-bundle-test
 
+goto-return-test:
+	+$(MAKE) -C goto_return
+.PHONY: goto-return-test
+
 CTYPES_HAND_WRITTEN_FILES=Client.ml Makefile
 
 .PHONY: $(LOWLEVEL_DIR)/Client.native