Enable calling procedures in contracts

Strata/Languages/Laurel/ConstrainedTypeElim.lean

@@ -75,7 +75,7 @@ def constraintCallFor (ptMap : ConstrainedTypeMap) (ty : HighType)
     For nested types, the function calls the parent's constraint function. -/
 def mkConstraintFunc (ptMap : ConstrainedTypeMap) (ct : ConstrainedType) : Procedure :=
   let baseType := resolveType ptMap ct.base
-  let bodyExpr := match ct.base.val with
+  let bodyExpr: StmtExprMd := match ct.base.val with
     | .UserDefined parent =>
       if ptMap.contains parent.text then
         let paramId := { ct.valueName with uniqueId := none }
@@ -89,7 +89,7 @@ def mkConstraintFunc (ptMap : ConstrainedTypeMap) (ct : ConstrainedType) : Proce
   { name := mkId s!"{ct.name.text}$constraint"
     inputs := [{ name := ct.valueName, type := baseType }]
     outputs := [{ name := mkId "result", type := { val := .TBool, source := none } }]
-    body := .Transparent { val := .Block [bodyExpr] none, source := none }
+    body := .Transparent { val := .Return bodyExpr, source := none }
     isFunctional := true
     decreases := none
     preconditions := [] }
@@ -244,11 +244,11 @@ public def constrainedTypeElim (model : SemanticModel) (program : Program)
    funcDiags)
 
 /-- Pipeline pass: constrained type elimination. -/
-public def constrainedTypeElimPass : LaurelPass where
+public def constrainedTypeElimPass : LoweringPass where
   name := "ConstrainedTypeElim"
   documentation := "Eliminates constrained types by replacing them with their base types and generating constraint-checking functions and witness procedures. Type tests against constrained types are rewritten to call the generated constraint function."
   needsResolves := true
-  run := fun p m =>
+  run := fun p m _ =>
     let (p', diags) := constrainedTypeElim m p
     (p', diags, {})
 

Strata/Languages/Laurel/CoreDefinitionsForLaurel.lean

@@ -27,16 +27,21 @@ program Laurel;
 
 datatype LaurelResolutionErrorPlaceholder {}
 datatype Float64IsNotSupportedYet {}
+datatype LaurelUnit { MkLaurelUnit() }
 
 // The types for these Map functions are incorrect.
 // We'll fix them when Laurel supports polymorphism
-function select(map: int, key: int) : int
+// And then we can remove the datatype Box as well
+// And remove the hacky filter in HeapParameterization
+datatype Box { MkBox() }
+
+function select(map: int, key: int) : Box
   external;
 
-function update(map: int, key: int, value: int) : int
+function update(map: int, key: int, value: int) : Box
   external;
 
-function const(value: int) : int
+function const(value: int) : Box
   external;
 
 #end

Strata/Languages/Laurel/CoreGroupingAndOrdering.lean

@@ -5,7 +5,9 @@
 -/
 
 module
-public import Strata.Languages.Laurel.TransparencyPass
+public import Strata.Languages.Laurel.LaurelAST
+public import Strata.Languages.Laurel.UnorderedCore
+public import Strata.Languages.Laurel.LaurelPass
 import Strata.DL.Lambda.LExpr
 import StrataDDM.Util.Graph.Tarjan
 import Strata.Languages.Laurel.Grammar.AbstractToConcreteTreeTranslator
@@ -27,6 +29,7 @@ declarations before they are emitted as Strata Core declarations.
 namespace Strata.Laurel
 
 open Lambda (LMonoTy LExpr)
+open Std (Format ToFormat)
 
 /-- Collect all `UserDefined` type names referenced in a `HighType`, including nested ones. -/
 def collectTypeRefs : HighTypeMd → List String
@@ -54,7 +57,7 @@ def collectStaticCallNames (expr : StmtExprMd) : List String :=
   match val with
   | .StaticCall callee args =>
       callee.text :: args.flatMap (fun a => collectStaticCallNames a)
-  | .PrimitiveOp _ args => args.flatMap (fun a => collectStaticCallNames a)
+  | .PrimitiveOp _ args _ => args.flatMap (fun a => collectStaticCallNames a)
   | .IfThenElse cond t e =>
       collectStaticCallNames cond ++
       collectStaticCallNames t ++
@@ -113,18 +116,18 @@ Build the procedure call graph, run Tarjan's SCC algorithm, and return each SCC
 as a list of procedures paired with a flag indicating whether the SCC is recursive.
 Results are in reverse topological order: dependencies before dependents.
 
-Procedures with `invokeOn` are placed as early as possible — before
+Procedures with axioms are placed as early as possible — before
 unrelated procedures without them — by stably partitioning them first before building
 the graph. Tarjan then naturally assigns them lower indices, causing them to appear
 earlier in the output.
 -/
 public def computeSccDecls (program : UnorderedCoreWithLaurelTypes) : List (List Procedure × Bool) :=
-  -- Stable partition: procedures with invokeOn come first, preserving relative
+  -- Stable partition: procedures with axioms come first, preserving relative
   -- order within each group. Tarjan then places them earlier in the topological output.
   let allProcs := program.functions ++ program.coreProcedures
-  let (withInvokeOn, withoutInvokeOn) :=
-    allProcs.partition (fun p => p.invokeOn.isSome)
-  let orderedProcs : List Procedure := withInvokeOn ++ withoutInvokeOn
+  let (withAxioms, withoutAxioms) :=
+    allProcs.partition (fun p => !p.axioms.isEmpty)
+  let orderedProcs : List Procedure := withAxioms ++ withoutAxioms
 
   -- Build a call-graph over all procedures.
   -- An edge proc → callee means proc's body/contracts contain a StaticCall to callee.
@@ -142,7 +145,8 @@ public def computeSccDecls (program : UnorderedCoreWithLaurelTypes) : List (List
       | _ => []
     let contractExprs : List StmtExprMd :=
       proc.preconditions.map (·.condition) ++
-      proc.invokeOn.toList
+      proc.invokeOn.toList ++
+      proc.axioms
     (bodyExprs ++ contractExprs).flatMap collectStaticCallNames
 
   -- Build the OutGraph for Tarjan.
@@ -225,7 +229,7 @@ Functions are grouped into SCCs (for mutual recursion). Proofs are emitted
 as individual `procedure` decls. Both participate in the topological ordering
 so that axioms are available to functions that need them.
 -/
-public def orderFunctionsAndProcedures (program : UnorderedCoreWithLaurelTypes) : CoreWithLaurelTypes :=
+def orderFunctionsAndProcedures (program : UnorderedCoreWithLaurelTypes) : CoreWithLaurelTypes :=
   let datatypeDecls := (groupDatatypesByScc' program).map OrderedDecl.datatypes
   let constantDecls := program.constants.map OrderedDecl.constant
   let funcNames : Std.HashSet String :=
@@ -254,4 +258,16 @@ where
       let members := comp.toList.filterMap fun idx => dtsArr[idx]?
       if members.isEmpty then none else some members
 
+public def orderingPass : LaurelPass UnorderedCoreWithLaurelTypes CoreWithLaurelTypes where
+  name := "OrderingPass"
+  comesBefore := []
+  documentation := "Produce a `CoreWithLaurelTypes` from a `UnorderedCoreWithLaurelTypes` by
+computing a combined ordering of functions and proofs using the call graph,
+then collecting datatypes and constants.
+Functions are grouped into SCCs (for mutual recursion). Proofs are emitted
+as individual `procedure` decls. Both participate in the topological ordering
+so that axioms are available to functions that need them."
+  run := fun p _ _ =>
+    (orderFunctionsAndProcedures p, [], {})
+
 end Strata.Laurel

Strata/Languages/Laurel/DesugarShortCircuit.lean

@@ -25,42 +25,43 @@ namespace Strata.Laurel
 
 public section
 
-private def bare (v : StmtExpr) : StmtExprMd := ⟨v, none⟩
 
 /-- Local rewrite of a single short-circuit node. Recursion is handled by `mapStmtExpr`. -/
-private def desugarShortCircuitNode (model : SemanticModel) (expr : StmtExprMd) : StmtExprMd :=
+private def desugarShortCircuitNode (imperativeCallees : List String) (expr : StmtExprMd) : StmtExprMd :=
   let source := expr.source
+  let wrap (v : StmtExpr) : StmtExprMd := ⟨v, source⟩
   match expr.val with
   | .PrimitiveOp op args _ =>
     match op, args with
     -- With bottom-up traversal, `a` and `b` are already desugared (nested
     -- short-circuits converted to IfThenElse). The check still works because
     -- `containsAssignmentOrImperativeCall` recurses into IfThenElse.
     | .AndThen, [a, b] | .Implies, [a, b] =>
-      if containsAssignmentOrImperativeCall model b then
+      if containsAssignmentOrImperativeCall imperativeCallees b then
         let elseVal := match op with | .AndThen => false | _ => true
-        ⟨.IfThenElse a b (some (bare (.LiteralBool elseVal))), source⟩
+        ⟨.IfThenElse a b (some (wrap (.LiteralBool elseVal))), source⟩
       else expr
     | .OrElse, [a, b] =>
-      if containsAssignmentOrImperativeCall model b then
-        ⟨.IfThenElse a (bare (.LiteralBool true)) (some b), source⟩
+      if containsAssignmentOrImperativeCall imperativeCallees b then
+        ⟨.IfThenElse a (wrap (.LiteralBool true)) (some b), source⟩
       else expr
     | _, _ => expr
   | _ => expr
 
 /-- Desugar short-circuit operators in a program. -/
-def desugarShortCircuit (model : SemanticModel) (program : Program) : Program :=
-  mapProgram (mapStmtExpr (desugarShortCircuitNode model)) program
+def desugarShortCircuit (program : Program) : Program :=
+  let imperativeCallees := (program.staticProcedures.filter (!·.isFunctional)).map (·.name.text)
+  mapProgram (mapStmtExpr (desugarShortCircuitNode imperativeCallees)) program
 
 end -- public section
 
 /-- Pipeline pass: desugar short-circuit operators. -/
-public def desugarShortCircuitPass : LaurelPass where
+public def desugarShortCircuitPass : LoweringPass where
   name := "DesugarShortCircuit"
   documentation := "Rewrites short-circuit boolean operators (`&&` and `||`) into equivalent conditional expressions. This simplifies subsequent passes and the final translation to Core, which does not have short-circuit semantics built in."
-  run := fun p m =>
-    (desugarShortCircuit m p, [], {})
+  run := fun p _ _ =>
+    (desugarShortCircuit p, [], {})
   comesBefore := [
-      ⟨ liftExpressionAssignmentsPass, "The desugar short circuit pass introduces if-then-else expressions whose control-flow must be taken into account by the lifting pass."⟩]
+      ⟨ liftImperativeExpressionsPass.meta, "The desugar short circuit pass introduces if-then-else expressions whose control-flow must be taken into account by the lifting pass."⟩]
 
 end Strata.Laurel

Strata/Languages/Laurel/EliminateDeterministicHoles.lean

@@ -90,10 +90,10 @@ def eliminateDeterministicHoles (program : Program) : Program × Statistics :=
 end -- public section
 
 /-- Pipeline pass: eliminate deterministic holes. -/
-public def eliminateDeterministicHolesPass : LaurelPass where
+public def eliminateDeterministicHolesPass : LoweringPass where
   name := "EliminateDeterministicHoles"
   documentation := "Replaces every deterministic hole with a call to a freshly generated uninterpreted function. After this pass the program contains only non-deterministic holes. Assumes `InferHoleTypes` has already annotated holes with types."
-  run := fun p _m =>
+  run := fun p _m _ =>
     let (p', stats) := eliminateDeterministicHoles p
     (p', [], stats)
 

Strata/Languages/Laurel/EliminateDoWhile.lean

@@ -73,10 +73,10 @@ def eliminateDoWhile (program : Program) : Program :=
   (mapProgramProceduresM rewrite program |>.run {}).fst
 
 /-- Pipeline pass: eliminate post-test (`do … while`) loops. -/
-public def eliminateDoWhilePass : LaurelPass where
+public def eliminateDoWhilePass : LoweringPass where
   name := "EliminateDoWhile"
   documentation := "Lowers post-test `While` loops (the `do … while` form) into the pre-test loop `{ while(true) invariant I { BODY; if (!COND) exit L } } L`, with a fresh `$`-prefixed exit label `L`. Runs early so no later pass observes a post-test loop; the invariant is checked at the loop head, matching `while`."
-  run := fun p _m => (eliminateDoWhile p, [], {})
+  run := fun p _m _ => (eliminateDoWhile p, [], {})
 
 end -- public section
 end Strata.Laurel

Strata/Languages/Laurel/EliminateIncrDecr.lean

@@ -100,10 +100,10 @@ def eliminateIncrDecr (program : Program) : Program :=
   mapProgramProcedures lowerProcedure program
 
 /-- Pipeline pass: eliminate increment/decrement operators. -/
-public def eliminateIncrDecrPass : LaurelPass where
+public def eliminateIncrDecrPass : LoweringPass where
   name := "EliminateIncrDecr"
   documentation := "Lowers Java-style increment/decrement operators (`++x`, `x++`, `--x`, `x--`) into existing Laurel assignment and arithmetic constructs. Prefix forms yield the new value; postfix forms yield the old value. Runs early so that no later pass observes an `.IncrDecr` node."
-  run := fun p _m => (eliminateIncrDecr p, [], {})
+  run := fun p _m _ => (eliminateIncrDecr p, [], {})
 
 end -- public section
 end Strata.Laurel

Strata/Languages/Laurel/EliminateReturnStatements.lean

@@ -0,0 +1,83 @@
+/-
+  Copyright Strata Contributors
+
+  SPDX-License-Identifier: Apache-2.0 OR MIT
+-/
+module
+
+public import Strata.Languages.Laurel.MapStmtExpr
+public import Strata.Languages.Laurel.LaurelPass
+
+/-!
+# Eliminate Return Statements
+
+Replaces `return` statements in imperative procedure bodies with assignments
+to the output parameters followed by an `exit` to a labelled block that wraps
+the entire body. This ensures that code placed after the body block (e.g.,
+postcondition assertions inserted by the contract pass) is always reached.
+
+This pass should run after `EliminateReturnsInExpression` (which handles
+functional/expression-position returns) and before the contract pass.
+-/
+
+namespace Strata.Laurel
+
+public section
+
+private def returnLabel : String := "$return"
+
+
+
+
+/-- Transform a single procedure: wrap body in a labelled block and replace returns. -/
+private def eliminateReturnStmts (proc : Procedure) : Procedure :=
+  match proc.body with
+  | .Opaque postconds (some impl) mods =>
+    let impl' := replaceReturn proc.outputs impl
+    let wrapped := match impl'.val with
+      | .Block stmts none => ⟨.Block stmts (some returnLabel), impl'.source⟩
+      | _ => ⟨ .Block [impl'] (some returnLabel), proc.name.source ⟩
+    { proc with body := .Opaque postconds (some wrapped) mods }
+  | .Transparent body =>
+    let body' := replaceReturn proc.outputs body
+    let wrapped := match body'.val with
+      | .Block stmts none => ⟨.Block stmts (some returnLabel), body'.source⟩
+      | _ => ⟨ .Block [body'] (some returnLabel), proc.name.source ⟩
+    { proc with body := .Transparent wrapped }
+  | _ => proc
+where
+
+  /-- Replace `Return val` with `output := val; exit "$return"` (or just `exit`
+    for valueless returns). Uses `mapStmtExpr` for bottom-up traversal. -/
+  replaceReturn (outputs : List Parameter) (expr : StmtExprMd) : StmtExprMd :=
+    mapStmtExpr (fun e =>
+      match e.val with
+      | .Return (some val) =>
+        /- Handling valued return is required because the heap param pass introduces valued return in
+        Strata/Languages/Laurel/HeapParameterizationConstants.lean
+        We should change that so we can remove this case.
+        -/
+        match outputs with
+        | [out] =>
+          let assign := ⟨ .Assign [⟨ .Local out.name, expr.source ⟩] val, expr.source ⟩
+          let exit := ⟨ .Exit returnLabel, expr.source ⟩
+          ⟨.Block [assign, exit] none, e.source⟩
+        | _ => ⟨ .Exit returnLabel, expr.source ⟩
+      | .Return none => ⟨ .Exit returnLabel, expr.source ⟩
+      | _ => e) expr
+
+/-- Transform a program by eliminating return statements in all procedure bodies. -/
+def eliminateReturnStatements (program : Program) : Program :=
+  { program with staticProcedures := program.staticProcedures.map eliminateReturnStmts }
+
+public def eliminateReturnStatementsPass : LoweringPass where
+  name := "EliminateReturnStatements"
+  documentation := "Lower return statements to exit statements. Wrap each procedure body with a 'return' block"
+  run := fun p _m _ =>
+    let p' := eliminateReturnStatements p
+    (p', [], {})
+  -- comesBefore := [contractPass]
+
+end -- public section
+
+end Strata.Laurel