From b313cd3a591728aa513ba04ff7f58dc56fe49d26 Mon Sep 17 00:00:00 2001 From: Nikhil Swamy Date: Wed, 15 Apr 2026 09:05:56 -0700 Subject: [PATCH] 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> --- lib/C11.ml | 2 + lib/CStar.ml | 2 + lib/CStarToC11.ml | 8 + lib/GotoForEarlyReturn.ml | 172 ++++++++++++++++++++++ lib/Options.ml | 2 + lib/PrintC.ml | 4 + src/Karamel.ml | 11 ++ test/Makefile | 6 +- test/goto_return/GotoReturn.expected.c | 195 +++++++++++++++++++++++++ test/goto_return/GotoReturn.fst | 85 +++++++++++ test/goto_return/Makefile | 54 +++++++ test/goto_return/krmlheader | 1 + 12 files changed, 541 insertions(+), 1 deletion(-) create mode 100644 lib/GotoForEarlyReturn.ml create mode 100644 test/goto_return/GotoReturn.expected.c create mode 100644 test/goto_return/GotoReturn.fst create mode 100644 test/goto_return/Makefile create mode 100644 test/goto_return/krmlheader diff --git a/lib/C11.ml b/lib/C11.ml index 31d8f89e3..5ab9ccd29 100644 --- a/lib/C11.ml +++ b/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 diff --git a/lib/CStar.ml b/lib/CStar.ml index edd99b86a..1387f496c 100644 --- a/lib/CStar.ml +++ b/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 = diff --git a/lib/CStarToC11.ml b/lib/CStarToC11.ml index 6d6f1fed6..2e2f51ed8 100644 --- a/lib/CStarToC11.ml +++ b/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 ] diff --git a/lib/GotoForEarlyReturn.ml b/lib/GotoForEarlyReturn.ml new file mode 100644 index 000000000..be8be6f0d --- /dev/null +++ b/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 diff --git a/lib/Options.ml b/lib/Options.ml index 9e8a71b74..2c68260a7 100644 --- a/lib/Options.ml +++ b/lib/Options.ml @@ -164,3 +164,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 diff --git a/lib/PrintC.ml b/lib/PrintC.ml index 41aaff055..d8ea3c5a6 100644 --- a/lib/PrintC.ml +++ b/lib/PrintC.ml @@ -474,6 +474,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 diff --git a/src/Karamel.ml b/src/Karamel.ml index a368bf840..689c22fa9 100644 --- a/src/Karamel.ml +++ b/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 *) @@ -790,6 +793,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 diff --git a/test/Makefile b/test/Makefile index 5f78bef66..5f0d082e2 100644 --- a/test/Makefile +++ b/test/Makefile @@ -85,7 +85,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) @@ -303,6 +303,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 diff --git a/test/goto_return/GotoReturn.expected.c b/test/goto_return/GotoReturn.expected.c new file mode 100644 index 000000000..b40da9703 --- /dev/null +++ b/test/goto_return/GotoReturn.expected.c @@ -0,0 +1,195 @@ +/* krml header omitted for test repeatability */ + + +#include "GotoReturn.h" + +int32_t GotoReturn_early_return_int(int32_t x) +{ + int32_t result = (int32_t)0; + if (x == (int32_t)0) + { + result = (int32_t)42; + goto exit; + } + if (x == (int32_t)1) + { + result = (int32_t)43; + goto exit; + } + result = (int32_t)44; + goto exit; + exit: ; + return result; +} + +int32_t GotoReturn_no_early_return(int32_t x) +{ + return x; +} + +int32_t GotoReturn_early_return_two_args(int32_t x, int32_t y) +{ + int32_t result = (int32_t)0; + if (x == (int32_t)0) + { + result = (int32_t)42; + goto exit; + } + result = y; + goto exit; + exit: ; + return result; +} + +int32_t GotoReturn_collision_test(int32_t x) +{ + int32_t result0 = (int32_t)0; + int32_t result = x; + if (result == (int32_t)0) + { + result0 = (int32_t)42; + goto exit; + } + if (result == (int32_t)1) + { + result0 = (int32_t)43; + goto exit; + } + result0 = (int32_t)44; + goto exit; + exit: ; + return result0; +} + +bool GotoReturn_early_return_bool(int32_t x) +{ + bool result = false; + if (x == (int32_t)0) + { + result = true; + goto exit; + } + if (x == (int32_t)1) + { + result = false; + goto exit; + } + result = true; + goto exit; + exit: ; + return result; +} + +uint32_t GotoReturn_early_return_uint32(uint32_t x) +{ + uint32_t result = 0U; + if (x == 0U) + { + result = 42U; + goto exit; + } + if (x == 1U) + { + result = 43U; + goto exit; + } + result = 44U; + goto exit; + exit: ; + return result; +} + +uint64_t GotoReturn_early_return_uint64(uint64_t x) +{ + uint64_t result = 0ULL; + if (x == 0ULL) + { + result = 42ULL; + goto exit; + } + if (x == 1ULL) + { + result = 43ULL; + goto exit; + } + result = 44ULL; + goto exit; + exit: ; + return result; +} + +int64_t GotoReturn_early_return_int64(int64_t x) +{ + int64_t result = (int64_t)0LL; + if (x == (int64_t)0LL) + { + result = (int64_t)42LL; + goto exit; + } + if (x == (int64_t)1LL) + { + result = (int64_t)43LL; + goto exit; + } + result = (int64_t)44LL; + goto exit; + exit: ; + return result; +} + +void GotoReturn_early_return_void(int32_t x) +{ + if (x == (int32_t)0) + { + GotoReturn_side_effect((int32_t)1); + goto exit; + } + if (x == (int32_t)1) + { + GotoReturn_side_effect((int32_t)2); + goto exit; + } + GotoReturn_side_effect((int32_t)3); + exit: ; +} + +GotoReturn_point GotoReturn_early_return_struct(int32_t x) +{ + GotoReturn_point result = { 0U }; + if (x == (int32_t)0) + { + result = ((GotoReturn_point){ .px = (int32_t)1, .py = (int32_t)2 }); + goto exit; + } + if (x == (int32_t)1) + { + result = ((GotoReturn_point){ .px = (int32_t)3, .py = (int32_t)4 }); + goto exit; + } + result = ((GotoReturn_point){ .px = (int32_t)5, .py = (int32_t)6 }); + goto exit; + exit: ; + return result; +} + +void GotoReturn_unit_return(void) +{ + +} + +int32_t GotoReturn_main(void) +{ + GotoReturn_early_return_int((int32_t)0); + GotoReturn_no_early_return((int32_t)1); + GotoReturn_early_return_two_args((int32_t)0, (int32_t)1); + GotoReturn_collision_test((int32_t)2); + GotoReturn_early_return_bool((int32_t)0); + GotoReturn_early_return_uint32(0U); + GotoReturn_early_return_uint64(0ULL); + GotoReturn_early_return_int64((int64_t)0LL); + GotoReturn_early_return_void((int32_t)0); + GotoReturn_early_return_struct((int32_t)0); + GotoReturn_unit_return(); + return (int32_t)0; +} + diff --git a/test/goto_return/GotoReturn.fst b/test/goto_return/GotoReturn.fst new file mode 100644 index 000000000..3e2206605 --- /dev/null +++ b/test/goto_return/GotoReturn.fst @@ -0,0 +1,85 @@ +module GotoReturn + +(* Early return with Int32 *) +let early_return_int (x: Int32.t) : Int32.t = + if x = 0l then 42l + else if x = 1l then 43l + else 44l + +(* No early return: single return at the end *) +let no_early_return (x: Int32.t) : Int32.t = + x + +(* Early return with two args *) +let early_return_two_args (x: Int32.t) (y: Int32.t) : Int32.t = + if x = 0l then 42l + else y + +(* Collision avoidance: local named result *) +let collision_test (x: Int32.t) : Int32.t = + let result = x in + if result = 0l then 42l + else if result = 1l then 43l + else 44l + +(* Early return with bool *) +let early_return_bool (x: Int32.t) : bool = + if x = 0l then true + else if x = 1l then false + else true + +(* Early return with UInt32 *) +let early_return_uint32 (x: FStar.UInt32.t) : FStar.UInt32.t = + let open FStar.UInt32 in + if x = 0ul then 42ul + else if x = 1ul then 43ul + else 44ul + +(* Early return with UInt64 *) +let early_return_uint64 (x: FStar.UInt64.t) : FStar.UInt64.t = + let open FStar.UInt64 in + if x = 0UL then 42UL + else if x = 1UL then 43UL + else 44UL + +(* Early return with Int64 *) +let early_return_int64 (x: FStar.Int64.t) : FStar.Int64.t = + let open FStar.Int64 in + if x = 0L then 42L + else if x = 1L then 43L + else 44L + +(* Early return with void: use assumed side-effecting function *) +assume val side_effect: Int32.t -> FStar.All.ML unit + +let early_return_void (x: Int32.t) : FStar.All.ML unit = + if x = 0l then side_effect 1l + else if x = 1l then side_effect 2l + else side_effect 3l + +(* Struct type for testing *) +type point = { px: Int32.t; py: Int32.t } + +(* Early return with struct *) +let early_return_struct (x: Int32.t) : point = + if x = 0l then { px = 1l; py = 2l } + else if x = 1l then { px = 3l; py = 4l } + else { px = 5l; py = 6l } + +(* Unit/void interaction: returning unit values *) +let unit_return (x: unit) : FStar.All.ML unit = + if true then x else () + +let main () : FStar.All.ML FStar.Int32.t = + let _ = early_return_int 0l in + let _ = no_early_return 1l in + let _ = early_return_two_args 0l 1l in + let _ = collision_test 2l in + let _ = early_return_bool 0l in + let _ = early_return_uint32 0ul in + let _ = early_return_uint64 0UL in + let _ = early_return_int64 0L in + early_return_void 0l; + let _ = early_return_struct 0l in + unit_return (); + 0l diff --git a/test/goto_return/Makefile b/test/goto_return/Makefile new file mode 100644 index 000000000..21c5bda0e --- /dev/null +++ b/test/goto_return/Makefile @@ -0,0 +1,54 @@ +# Test for -goto_for_early_return transformation. +# Extracts GotoReturn.fst, compiles with -goto_for_early_return -fnoreturn-else, +# verifies the output matches expected C, and runs the executable. + +KRML_EXE ?= $(abspath ../../krml) +FSTAR_EXE ?= fstar.exe +OUTPUT_DIR = _output +CACHE_DIR = _cache + +.DELETE_ON_ERROR: +.SECONDARY: +.DEFAULT_GOAL := all + +Q?=@ + +FSTAR = $(FSTAR_EXE) --cache_dir $(CACHE_DIR) --odir $(OUTPUT_DIR) --warn_error -321 + +$(CACHE_DIR)/GotoReturn.fst.checked: GotoReturn.fst + @mkdir -p $(CACHE_DIR) $(OUTPUT_DIR) + @printf " %-14s %s\n" "CHECK" "GotoReturn.fst" + $(Q)$(FSTAR) -c $< -o $@ + +$(OUTPUT_DIR)/GotoReturn.krml: $(CACHE_DIR)/GotoReturn.fst.checked + @mkdir -p $(OUTPUT_DIR) + @printf " %-14s %s\n" "EXTRACT" "GotoReturn" + $(Q)$(FSTAR) --codegen krml --extract_module GotoReturn GotoReturn.fst + +# Compile with -goto_for_early_return and -fnoreturn-else +$(OUTPUT_DIR)/GotoReturn.c: $(OUTPUT_DIR)/GotoReturn.krml + @printf " %-14s %s\n" "KRML" "GotoReturn" + $(Q)$(KRML_EXE) -minimal -skip-makefiles -skip-compilation -skip-linking \ + -bundle GotoReturn=* \ + -fnoreturn-else -goto_for_early_return \ + -tmpdir $(OUTPUT_DIR) \ + -header krmlheader \ + $< + +# Diff against expected output +$(OUTPUT_DIR)/GotoReturn.diff: $(OUTPUT_DIR)/GotoReturn.c GotoReturn.expected.c + @printf " %-14s %s\n" "DIFF" "$<" + $(Q)diff -u --strip-trailing-cr GotoReturn.expected.c $< + $(Q)touch $@ + +# Accept current output as expected +accept: $(OUTPUT_DIR)/GotoReturn.c + @printf " %-14s %s\n" "ACCEPT" "$<" + $(Q)cp $< GotoReturn.expected.c + +all: $(OUTPUT_DIR)/GotoReturn.diff + +clean: + rm -rf $(OUTPUT_DIR) $(CACHE_DIR) + +.PHONY: all clean accept diff --git a/test/goto_return/krmlheader b/test/goto_return/krmlheader new file mode 100644 index 000000000..5974791c5 --- /dev/null +++ b/test/goto_return/krmlheader @@ -0,0 +1 @@ +/* krml header omitted for test repeatability */