Support return statements in Laurel functions

.github/CODEOWNERS

@@ -3,7 +3,8 @@
 *       @strata-org/reviewers
 
 # Laurel language
-Strata/Languages/Laurel/     @strata-org/reviewers @keyboardDrummer @fabiomadge
-StrataTest/Languages/Laurel/ @strata-org/reviewers @keyboardDrummer @fabiomadge
-Strata/Languages/Python/     @strata-org/reviewers @keyboardDrummer @fabiomadge
-StrataTest/Languages/Python/ @strata-org/reviewers @keyboardDrummer @fabiomadge
+/Strata/Languages/Laurel        @strata-org/reviewers @keyboardDrummer @fabiomadge
+/StrataTest/Languages/Laurel    @strata-org/reviewers @keyboardDrummer @fabiomadge
+/Strata/Languages/Python        @strata-org/reviewers @keyboardDrummer @fabiomadge
+/StrataTest/Languages/Python    @strata-org/reviewers @keyboardDrummer @fabiomadge
+/docs/verso                     @strata-org/reviewers @keyboardDrummer @fabiomadge

Strata/Languages/Laurel/EliminateReturnsInExpression.lean

@@ -0,0 +1,116 @@
+/-
+  Copyright Strata Contributors
+
+  SPDX-License-Identifier: Apache-2.0 OR MIT
+-/
+
+import Strata.Languages.Laurel.Laurel
+import Strata.Util.Tactics
+
+/-!
+# Eliminate Returns in Expression Position
+
+Rewrites functional procedure bodies so that `return` statements are removed
+and early-return guard patterns become if-then-else expressions. This makes
+the body a pure expression tree suitable for translation to a Core function.
+
+The algorithm walks a block backwards (from last statement to first),
+accumulating a result expression:
+
+- The last statement is converted via `lastStmtToExpr` which strips `return`,
+  recurses into blocks, and handles if-then-else.
+- Each preceding statement wraps around the accumulated result via `stmtsToExpr`:
+  - `if (cond) { body }` (no else) becomes `if cond then lastStmtToExpr(body) else acc`
+  - Other statements are kept in a two-element block with the accumulator.
+
+-/
+
+namespace Strata.Laurel
+
+/-- Appending a singleton strictly increases `sizeOf`. -/
+private theorem List.sizeOf_lt_append_singleton [SizeOf α] (xs : List α) (y : α) :
+    sizeOf xs < sizeOf (xs ++ [y]) := by
+  induction xs with
+  | nil => simp_all; omega
+  | cons hd tl ih => simp_all [List.cons_append]
+
+/-- `dropLast` of a non-empty list has strictly smaller `sizeOf`. -/
+private theorem List.sizeOf_dropLast_lt [SizeOf α] {l : List α} (h_ne : l ≠ []) :
+    sizeOf l.dropLast < sizeOf l := by
+  have h_concat := List.dropLast_concat_getLast h_ne
+  have : sizeOf l = sizeOf (l.dropLast ++ [l.getLast h_ne]) := by rw [h_concat]
+  rw [this]
+  exact List.sizeOf_lt_append_singleton l.dropLast (l.getLast h_ne)
+
+mutual
+
+/--
+Fold a list of preceding statements (right-to-left) around an accumulator
+expression. Each `if-then` (no else) guard wraps as
+`if cond then lastStmtToExpr(body) else acc`; other statements produce
+`Block [stmt, acc]`.
+-/
+def stmtsToExpr (stmts : List StmtExprMd) (acc : StmtExprMd) (blockMd : MetaData)
+    : StmtExprMd :=
+  match stmts with
+  | [] => acc
+  | s :: rest =>
+    let acc' := stmtsToExpr rest acc blockMd
+    match s with
+    | ⟨.IfThenElse cond thenBr none, smd⟩ =>
+      ⟨.IfThenElse cond (lastStmtToExpr thenBr) (some acc'), smd⟩
+    | _ =>
+      ⟨.Block [s, acc'] none, blockMd⟩
+  termination_by (sizeOf stmts, 1)
+
+/--
+Convert the last statement of a block into an expression.
+- `return expr` → `expr`
+- A non-empty block → process last element, fold preceding statements
+- `if cond then A else B` → recurse into both branches
+- Anything else → kept as-is
+-/
+def lastStmtToExpr (stmt : StmtExprMd) : StmtExprMd :=
+  match stmt with
+  | ⟨.Return (some val), _⟩ => val
+  | ⟨.Block stmts _, md⟩ =>
+    match h_last : stmts.getLast? with
+    | some last =>
+      have := List.mem_of_getLast? h_last
+      let lastExpr := lastStmtToExpr last
+      let dropped := stmts.dropLast
+      -- have := List.dropLast_subset stmts
+      have h : sizeOf stmts.dropLast < sizeOf stmts :=
+        List.sizeOf_dropLast_lt (by intro h; simp [h] at h_last)
+      stmtsToExpr dropped lastExpr md
+    | none => stmt
+  | ⟨.IfThenElse cond thenBr (some elseBr), md⟩ =>
+    ⟨.IfThenElse cond (lastStmtToExpr thenBr) (some (lastStmtToExpr elseBr)), md⟩
+  | _ => stmt
+  termination_by (sizeOf stmt, 0)
+  decreasing_by
+    all_goals (simp_all; term_by_mem)
+
+end
+
+/--
+Apply return elimination to a functional procedure's body.
+The entire body is treated as an expression to be converted.
+-/
+def eliminateReturnsInExpression (proc : Procedure) : Procedure :=
+  if !proc.isFunctional then proc
+  else
+    match proc.body with
+    | .Transparent bodyExpr =>
+      { proc with body := .Transparent (lastStmtToExpr bodyExpr) }
+    | .Opaque postconds (some impl) modif =>
+      { proc with body := .Opaque postconds (some (lastStmtToExpr impl)) modif }
+    | _ => proc
+
+/--
+Transform a program by eliminating returns in all functional procedure bodies.
+-/
+def eliminateReturnsInExpressionTransform (program : Program) : Program :=
+  { program with staticProcedures := program.staticProcedures.map eliminateReturnsInExpression }
+
+end Laurel

Strata/Languages/Laurel/Grammar/ConcreteToAbstractTreeTranslator.lean

@@ -76,8 +76,8 @@ def translateBool (arg : Arg) : TransM Bool := do
 instance : Inhabited Parameter where
   default := { name := "" , type := ⟨.TVoid, #[]⟩ }
 
-def mkHighTypeMd (t : HighType) (md : MetaData Core.Expression) : HighTypeMd := ⟨t, md⟩
-def mkStmtExprMd (e : StmtExpr) (md : MetaData Core.Expression) : StmtExprMd := ⟨e, md⟩
+def mkHighTypeMd (t : HighType) (md : MetaData) : HighTypeMd := ⟨t, md⟩
+def mkStmtExprMd (e : StmtExpr) (md : MetaData) : StmtExprMd := ⟨e, md⟩
 def mkStmtExprMdEmpty (e : StmtExpr) : StmtExprMd := ⟨e, #[]⟩
 
 partial def translateHighType (arg : Arg) : TransM HighTypeMd := do

Strata/Languages/Laurel/Laurel.lean

@@ -91,8 +91,9 @@ inductive Operation : Type where
   | StrConcat
   deriving Repr
 
+abbrev MetaData := Imperative.MetaData Core.Expression
 -- Explicit instance needed for deriving Repr in the mutual block
-instance : Repr (Imperative.MetaData Core.Expression) := inferInstance
+instance : Repr MetaData := inferInstance
 
 /--
 A wrapper that pairs a value with source-level metadata such as source
@@ -104,7 +105,7 @@ structure WithMetadata (t : Type) : Type where
   /-- The wrapped value. -/
   val : t
   /-- Source-level metadata (locations, annotations). -/
-  md : Imperative.MetaData Core.Expression
+  md : MetaData
   deriving Repr
 
 /--

Strata/Languages/Laurel/LaurelToCoreTranslator.lean

@@ -11,6 +11,7 @@ import Strata.Languages.Core.Procedure
 import Strata.Languages.Core.Options
 import Strata.Languages.Laurel.Laurel
 import Strata.Languages.Laurel.LiftImperativeExpressions
+import Strata.Languages.Laurel.EliminateReturnsInExpression
 import Strata.Languages.Laurel.HeapParameterization
 import Strata.Languages.Laurel.TypeHierarchy
 import Strata.Languages.Laurel.LaurelTypes
@@ -115,12 +116,12 @@ def translateExpr (expr : StmtExprMd)
   -- Dummy expression used as placeholder when an error is emitted in pure context
   let dummy := .fvar () (⟨s!"DUMMY_VAR_{← freshId}", ()⟩) none
   -- Emit an error in pure context; panic in impure context (lifting invariant violated)
-  let disallowed (e : StmtExprMd) (msg : String) : TranslateM Core.Expression.Expr := do
+  let disallowed (md : MetaData) (msg : String) : TranslateM Core.Expression.Expr := do
     if isPureContext then
-      emitDiagnostic (e.md.toDiagnostic msg)
+      emitDiagnostic (md.toDiagnostic msg)
       return dummy
     else
-      panic! s!"translateExpr: {msg} (should have been lifted): {Std.Format.pretty (Std.ToFormat.format e)}"
+      panic! s!"translateExpr: {msg} (should have been lifted): {Std.Format.pretty (Std.ToFormat.format md)}"
   match h: expr.val with
   | .LiteralBool b => return .const () (.boolConst b)
   | .LiteralInt i => return .const () (.intConst i)
@@ -186,7 +187,7 @@ def translateExpr (expr : StmtExprMd)
   | .StaticCall callee args =>
       -- In a pure context, only Core functions (not procedures) are allowed
       if isPureContext && !model.isFunction callee then
-        disallowed expr "calls to procedures are not supported in functions or contracts"
+        disallowed expr.md "calls to procedures are not supported in functions or contracts"
       else
         let fnOp : Core.Expression.Expr := .op () ⟨callee.text, ()⟩ none
         args.attach.foldlM (fun acc ⟨arg, _⟩ => do
@@ -207,21 +208,38 @@ def translateExpr (expr : StmtExprMd)
       let re2 ← translateExpr e2 boundVars isPureContext
       return .eq () re1 re2
   | .Assign _ _ =>
-      disallowed expr "destructive assignments are not supported in functions or contracts"
+      disallowed expr.md "destructive assignments are not supported in functions or contracts"
   | .While _ _ _ _ =>
-      disallowed expr "loops are not supported in functions or contracts"
-  | .Exit _ => disallowed expr "exit is not supported in expression position"
+      disallowed expr.md "loops are not supported in functions or contracts"
+  | .Exit _ => disallowed expr.md "exit is not supported in expression position"
+
+  | .Block (⟨ .Assert _, md⟩ :: rest) label => do
+    _ ← disallowed md "asserts are not YET supported in functions or contracts"
+    translateExpr ⟨ StmtExpr.Block rest label, md ⟩ boundVars isPureContext
+  | .Block (⟨ .Assume _, md⟩ :: rest) label =>
+    _ ← disallowed md "assumes are not YET supported in functions or contracts"
+    translateExpr ⟨ StmtExpr.Block rest label, md ⟩ boundVars isPureContext
+  | .Block (⟨ .LocalVariable name ty (some initializer), md⟩ :: rest) label => do
+      let valueExpr ← translateExpr  initializer boundVars isPureContext
+      let bodyExpr ← translateExpr ⟨ StmtExpr.Block rest label, md ⟩ (name :: boundVars) isPureContext
+      disallowed md "local variables in functions are not YET supported"
+      -- This doesn't work because of a limitation in Core.
+      -- let coreMonoType := translateType ty
+      -- return .app () (.abs () (some coreMonoType) bodyExpr) valueExpr
+  | .Block (⟨ .LocalVariable name ty none, md⟩ :: rest) label =>
+    disallowed md "local variables in functions must have initializers"
+  | .Block (⟨ .IfThenElse cond thenBranch (some elseBranch), md⟩ :: rest) label =>
+    disallowed md "if-then-else only supported as the last statement in a block"
 
   | .IsType _ _ => panic "IsType should have been lowered"
   | .New _ => panic! s!"New should have been eliminated by typeHierarchyTransform"
   | .FieldSelect target fieldId =>
       -- Field selects should have been eliminated by heap parameterization
       -- If we see one here, it's an error in the pipeline
       panic! s!"FieldSelect should have been eliminated by heap parameterization: {Std.ToFormat.format target}#{fieldId.text}"
-
-  | .Block _ _ => panic "block expression not yet implemented (should be lowered in a separate pass)"
+  | .Block _ _ => panic! "block expression should have been lowered in a separate pass"
   | .LocalVariable _ _ _ => panic "local variable expression not yet implemented (should be lowered in a separate pass)"
-  | .Return _ => disallowed expr "return expression not yet implemented (should be lowered in a separate pass)"
+  | .Return _ => disallowed expr.md "return expression not yet implemented (should be lowered in a separate pass)"
 
   | .AsType target _ => panic "AsType expression not implemented"
   | .Assigned _ => panic "assigned expression not implemented"
@@ -423,68 +441,6 @@ def translateProcedure (proc : Procedure) : TranslateM Core.Procedure := do
   let spec : Core.Procedure.Spec := { modifies, preconditions, postconditions }
   return { header, spec, body }
 
-/--
-Check if a Laurel expression is pure (contains no side effects).
-Used to determine if a procedure can be translated as a Core function.
--/
-private def isPureExpr(expr: StmtExprMd): Bool :=
-  match _h : expr.val with
-  | .LiteralBool _ => true
-  | .LiteralInt _ => true
-  | .LiteralString _ => true
-  | .Identifier _ => true
-  | .PrimitiveOp _ args => args.attach.all (fun ⟨a, _⟩ => isPureExpr a)
-  | .IfThenElse c t none => isPureExpr c && isPureExpr t
-  | .IfThenElse c t (some e) => isPureExpr c && isPureExpr t && isPureExpr e
-  | .StaticCall _ args => args.attach.all (fun ⟨a, _⟩ => isPureExpr a)
-  | .New _ => false
-  | .ReferenceEquals e1 e2 => isPureExpr e1 && isPureExpr e2
-  | .Block [single] _ => isPureExpr single
-  | .Block _ _ => false
-  -- Statement-like
-  | .LocalVariable .. => true
-  | .While .. => false
-  | .Exit .. => false
-  | .Return .. => false
-  -- Expression-like
-  | .Assign .. => false
-  | .FieldSelect .. => true
-  | .PureFieldUpdate .. => true
-  -- Instance related
-  | .This => panic s!"isPureExpr not implemented for This"
-  | .AsType .. => panic s!"isPureExpr not supported for AsType"
-  | .IsType .. => panic s!"isPureExpr not supported for IsType"
-  | .InstanceCall .. => panic s!"isPureExpr not implemented for InstanceCall"
-  -- Verification specific
-  | .Forall .. => panic s!"isPureExpr not implemented for Forall"
-  | .Exists .. => panic s!"isPureExpr not implemented for Exists"
-  | .Assigned .. => panic s!"isPureExpr not supported for AsType"
-  | .Old .. => panic s!"isPureExpr not supported for AsType"
-  | .Fresh .. => panic s!"isPureExpr not supported for AsType"
-  | .Assert .. => panic s!"isPureExpr not implemented for Assert"
-  | .Assume .. => panic s!"isPureExpr not implemented for Assume"
-  | .ProveBy .. => panic s!"isPureExpr not implemented for ProveBy"
-  | .ContractOf .. => panic s!"isPureExpr not implemented for ContractOf"
-  | .Abstract => panic s!"isPureExpr not implemented for Abstract"
-  | .All => panic s!"isPureExpr not implemented for All"
-  -- Dynamic / closures
-  | .Hole => true
-  termination_by sizeOf expr
-  decreasing_by all_goals (have := WithMetadata.sizeOf_val_lt expr; term_by_mem)
-
-/-- Check if a pure-marked procedure can actually be represented as a Core function:
-    transparent body that is a pure expression and has exactly one output. -/
-private def canBeCoreFunctionBody (proc : Procedure) : Bool :=
-  match proc.body with
-  | .Transparent bodyExpr =>
-    isPureExpr bodyExpr &&
-    proc.outputs.length == 1
-  | .Opaque _ bodyExprOption _ =>
-    (bodyExprOption.map isPureExpr).getD true &&
-    proc.outputs.length == 1
-  | .External => false
-  | _ => false
-
 /--
 Translate a Laurel Procedure to a Core Function (when applicable) using `TranslateM`.
 Diagnostics for disallowed constructs in the function body are emitted into the monad state.
@@ -579,6 +535,7 @@ def translate (program : Program) : Except (Array DiagnosticModel) (Core.Program
   -- dbg_trace (toString (Std.Format.pretty (Std.ToFormat.format program)))
   -- dbg_trace "================================="
   let program := liftExpressionAssignments model program
+  let program := eliminateReturnsInExpressionTransform program
   let result := resolve program (some model)
   let (program, model) := (result.program, result.model)
   _resolutionDiags := _resolutionDiags ++ result.errors

Strata/Languages/Laurel/LiftImperativeExpressions.lean

@@ -321,7 +321,7 @@ def transformStmt (stmt : StmtExprMd) : LiftM (List StmtExprMd) := do
           -- If the initializer is a direct imperative StaticCall, don't lift it —
           -- translateStmt handles LocalVariable + StaticCall directly as a call statement.
           match _: initExprMd with
-          | WithMetadata.mk initExpr md =>
+          | WithMetadata.mk initExpr _ =>
           match _: initExpr with
           | .StaticCall callee args =>
               let model := (← get).model

StrataTest/Languages/Laurel/DivisionByZeroCheckTest.lean

@@ -26,9 +26,11 @@ procedure safeDivision() {
   assert z == 5;
 }
 
+// Error ranges are too wide because Core does not use expression locations
 procedure unsafeDivision(x: int) {
   var z: int := 10 / x;
-//              ^^^^^^ error: assertion does not hold
+//^^^^^^^^^^^^^^^^^^^^^ error: assertion does not hold
+// Error ranges are too wide because Core does not use expression locations
 }
 
 function pureDiv(x: int, y: int): int
@@ -44,7 +46,8 @@ procedure callPureDivSafe() {
 
 procedure callPureDivUnsafe(x: int) {
   var z: int := pureDiv(10, x);
-//              ^^^^^^^^^^^^^^ error: assertion does not hold
+//^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ error: assertion does not hold
+// Error ranges are too wide because Core does not use expression locations
 }
 "
 

StrataTest/Languages/Laurel/Examples/Fundamentals/T3_ControlFlow.lean

@@ -10,10 +10,46 @@ import StrataTest.Languages.Laurel.TestExamples
 open StrataTest.Util
 open Strata
 
-namespace Strata
-namespace Laurel
+namespace Strata.Laurel
 
 def program := r"
+function returnAtEnd(x: int) returns (r: int) {
+  if (x > 0) {
+    if (x == 1) {
+      return 1;
+    } else {
+      return 2;
+    }
+  } else {
+    return 3;
+  }
+}
+
+function guardInFunction(x: int) returns (r: int) {
+  if (x > 0) {
+    if (x == 1) {
+      return 1;
+    } else {
+      return 2;
+    }
+  }
+
+  return 3;
+}
+
+procedure testFunctions() {
+  // assert letsInFunction() == 2;
+  // assert letsInFunction() == 3; error: assertion does not hold
+
+  assert returnAtEnd(1) == 1;
+  assert returnAtEnd(1) == 2;
+//^^^^^^^^^^^^^^^^^^^^^^^^^^^ error: assertion does not hold
+
+  assert guardInFunction(1) == 1;
+  assert guardInFunction(1) == 2;
+//^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ error: assertion does not hold
+}
+
 procedure guards(a: int) returns (r: int)
 {
   var b: int := a + 2;