Add support for function declaration statements

Strata/DL/Imperative/PureExpr.lean

@@ -21,6 +21,8 @@ structure PureExpr : Type 1 where
   Expr    : Type
   /-- Types -/
   Ty      : Type
+  /-- Expression metadata type (for use in function declarations, etc.) -/
+  ExprMetadata : Type
   /-- Typing environment, expected to contain a map of variables to their types,
   type substitution, etc.
   -/

Strata/DL/Imperative/Stmt.lean

@@ -7,6 +7,7 @@
 
 
 import Strata.DL.Imperative.Cmd
+import Strata.DL.Lambda.Factory
 
 namespace Imperative
 ---------------------------------------------------------------------
@@ -38,6 +39,8 @@ inductive Stmt (P : PureExpr) (Cmd : Type) : Type where
   This will likely be removed, in favor of an alternative view of imperative
   programs that is purely untructured. -/
   | goto     (label : String) (md : MetaData P := .empty)
+  /-- A function declaration within a statement block. -/
+  | funcDecl (decl : Lambda.PureFunc P) (md : MetaData P := .empty)
   deriving Inhabited
 
 /-- A block is simply an abbreviation for a list of commands. -/
@@ -69,19 +72,22 @@ def Stmt.inductionOn {P : PureExpr} {Cmd : Type}
       motive (Stmt.loop guard measure invariant body md))
     (goto_case : ∀ (label : String) (md : MetaData P),
       motive (Stmt.goto label md))
+    (funcDecl_case : ∀ (decl : Lambda.PureFunc P) (md : MetaData P),
+      motive (Stmt.funcDecl decl md))
     (s : Stmt P Cmd) : motive s :=
   match s with
   | Stmt.cmd c => cmd_case c
   | Stmt.block label b md =>
-    block_case label b md (fun s _ => inductionOn cmd_case block_case ite_case loop_case goto_case s)
+    block_case label b md (fun s _ => inductionOn cmd_case block_case ite_case loop_case goto_case funcDecl_case s)
   | Stmt.ite cond thenb elseb md =>
     ite_case cond thenb elseb md
-      (fun s _ => inductionOn cmd_case block_case ite_case loop_case goto_case s)
-      (fun s _ => inductionOn cmd_case block_case ite_case loop_case goto_case s)
+      (fun s _ => inductionOn cmd_case block_case ite_case loop_case goto_case funcDecl_case s)
+      (fun s _ => inductionOn cmd_case block_case ite_case loop_case goto_case funcDecl_case s)
   | Stmt.loop guard measure invariant body md =>
     loop_case guard measure invariant body md
-      (fun s _ => inductionOn cmd_case block_case ite_case loop_case goto_case s)
+      (fun s _ => inductionOn cmd_case block_case ite_case loop_case goto_case funcDecl_case s)
   | Stmt.goto label md => goto_case label md
+  | Stmt.funcDecl decl md => funcDecl_case decl md
   termination_by s
 
 ---------------------------------------------------------------------
@@ -97,6 +103,7 @@ def Stmt.sizeOf (s : Imperative.Stmt P C) : Nat :=
   | .ite c tss ess _ => 3 + sizeOf c + Block.sizeOf tss + Block.sizeOf ess
   | .loop g _ _ bss _ => 3 + sizeOf g + Block.sizeOf bss
   | .goto _ _ => 1
+  | .funcDecl _ _ => 1
 
 @[simp]
 def Block.sizeOf (ss : Imperative.Block P C) : Nat :=
@@ -159,6 +166,11 @@ def Stmt.getVars [HasVarsPure P P.Expr] [HasVarsPure P C] (s : Stmt P C) : List
   | .ite _ tbss ebss _ => Block.getVars tbss ++ Block.getVars ebss
   | .loop _ _ _ bss _ => Block.getVars bss
   | .goto _ _  => []
+  | .funcDecl decl _ =>
+    -- Get variables from function body (including parameters for simplicity)
+    match decl.body with
+    | none => []
+    | some body => HasVarsPure.getVars body
   termination_by (Stmt.sizeOf s)
 
 def Block.getVars [HasVarsPure P P.Expr] [HasVarsPure P C] (ss : Block P C) : List P.Ident :=
@@ -183,6 +195,7 @@ def Stmt.definedVars [HasVarsImp P C] (s : Stmt P C) : List P.Ident :=
   | .cmd cmd => HasVarsImp.definedVars cmd
   | .block _ bss _ => Block.definedVars bss
   | .ite _ tbss ebss _ => Block.definedVars tbss ++ Block.definedVars ebss
+  | .funcDecl decl _ => [decl.name]  -- Function declaration defines the function name
   | _ => []
   termination_by (Stmt.sizeOf s)
 
@@ -202,6 +215,7 @@ def Stmt.modifiedVars [HasVarsImp P C] (s : Stmt P C) : List P.Ident :=
   | .block _ bss _ => Block.modifiedVars bss
   | .ite _ tbss ebss _ => Block.modifiedVars tbss ++ Block.modifiedVars ebss
   | .loop _ _ _ bss _ => Block.modifiedVars bss
+  | .funcDecl _ _ => []  -- Function declarations don't modify variables
   termination_by (Stmt.sizeOf s)
 
 def Block.modifiedVars [HasVarsImp P C] (ss : Block P C) : List P.Ident :=
@@ -262,6 +276,7 @@ def formatStmt (P : PureExpr) (s : Stmt P C)
   | .loop guard measure invariant body md => f!"{md}while ({guard}) ({measure}) ({invariant}) " ++
                         Format.bracket "{" f!"{formatBlock P body}" "}"
   | .goto label md => f!"{md}goto {label}"
+  | .funcDecl _ md => f!"{md}funcDecl <function>"
   termination_by s.sizeOf
 
 def formatBlock (P : PureExpr) (ss : List (Stmt P C))

Strata/DL/Lambda/Factory.lean

@@ -10,6 +10,8 @@ import Strata.DDM.AST
 import Strata.DDM.Util.Array
 import Strata.DL.Util.List
 import Strata.DL.Util.ListMap
+import Strata.DL.Imperative.PureExpr
+import Strata.DL.Imperative.MetaData
 
 /-!
 ## Lambda's Factory
@@ -57,9 +59,11 @@ abbrev LTySignature := Signature IDMeta LTy
 
 
 /--
-A Lambda factory function, where the body can be optional. Universally
-quantified type identifiers, if any, appear before this signature and can
-quantify over the type identifiers in it.
+A generic function structure, parameterized by identifier, expression, type, and metadata types.
+
+This structure can be instantiated for different expression languages.
+For Lambda expressions, use `LFunc`. For other expression systems, instantiate
+with appropriate types.
 
 A optional evaluation function can be provided in the `concreteEval` field for
 each factory function to allow the partial evaluator to do constant propagation
@@ -85,44 +89,74 @@ concrete/constants, this fails and it returns .none.
 (TODO) Use `.bvar`s in the body to correspond to the formals instead of using
 `.fvar`s.
 -/
-structure LFunc (T : LExprParams) where
-  name     : T.Identifier
+structure Func (IdentT : Type) (ExprT : Type) (TyT : Type) (MetadataT : Type) where
+  name     : IdentT
   typeArgs : List TyIdentifier := []
   isConstr : Bool := false --whether function is datatype constructor
-  inputs   : @LMonoTySignature T.IDMeta
-  output   : LMonoTy
-  body     : Option (LExpr T.mono) := .none
+  inputs   : ListMap IdentT TyT
+  output   : TyT
+  body     : Option ExprT := .none
   -- (TODO): Add support for a fixed set of attributes (e.g., whether to inline
   -- a function, etc.).
   attr     : Array String := #[]
-  -- The T.Metadata argument is the metadata that will be attached to the
+  -- The MetadataT argument is the metadata that will be attached to the
   -- resulting expression of concreteEval if evaluation was successful.
-  concreteEval : Option (T.Metadata → List (LExpr T.mono) → Option (LExpr T.mono)) := .none
-  axioms   : List (LExpr T.mono) := []  -- For axiomatic definitions
+  concreteEval : Option (MetadataT → List ExprT → Option ExprT) := .none
+  axioms   : List ExprT := []  -- For axiomatic definitions
+
+/--
+A Lambda factory function - instantiation of `Func` for Lambda expressions.
+
+Universally quantified type identifiers, if any, appear before this signature and can
+quantify over the type identifiers in it.
+-/
+abbrev LFunc (T : LExprParams) := Func (T.Identifier) (LExpr T.mono) LMonoTy T.Metadata
 
 /--
-Well-formedness properties of LFunc. These are split from LFunc because
-otherwise it becomes impossible to create a 'temporary' LFunc object whose
+A function declaration for use with `PureExpr` - instantiation of `Func` for
+any expression system that implements the `PureExpr` interface.
+-/
+abbrev PureFunc (P : Imperative.PureExpr) := Func P.Ident P.Expr P.Ty P.ExprMetadata
+
+/--
+Helper constructor for LFunc to maintain backward compatibility.
+-/
+def LFunc.mk {T : LExprParams} (name : T.Identifier) (typeArgs : List TyIdentifier := [])
+    (isConstr : Bool := false) (inputs : ListMap T.Identifier LMonoTy) (output : LMonoTy)
+    (body : Option (LExpr T.mono) := .none) (attr : Array String := #[])
+    (concreteEval : Option (T.Metadata → List (LExpr T.mono) → Option (LExpr T.mono)) := .none)
+    (axioms : List (LExpr T.mono) := []) : LFunc T :=
+  Func.mk name typeArgs isConstr inputs output body attr concreteEval axioms
+
+/--
+Well-formedness properties of Func. These are split from Func because
+otherwise it becomes impossible to create a 'temporary' Func object whose
 wellformedness might not hold yet.
 -/
-structure LFuncWF {T : LExprParams} (f : LFunc T) where
+structure FuncWF {IdentT ExprT TyT MetadataT : Type} (f : Func IdentT ExprT TyT MetadataT) where
   -- No args have same name.
   arg_nodup:
-    List.Nodup (f.inputs.map (·.1.name))
+    List.Nodup (f.inputs.map (·.1))
+  -- concreteEval does not succeed if the length of args is incorrect.
+  concreteEval_argmatch:
+    ∀ fn md args res, f.concreteEval = .some fn
+      → fn md args = .some res
+      → args.length = f.inputs.length
+
+/--
+Well-formedness properties of LFunc - extends FuncWF with Lambda-specific properties.
+-/
+structure LFuncWF {T : LExprParams} (f : LFunc T) extends FuncWF f where
   -- Free variables of body must be arguments.
   body_freevars:
     ∀ b freevars, f.body = .some b
       → freevars = LExpr.freeVars b
       → (∀ fv, fv ∈ freevars →
           ∃ arg, List.Mem arg f.inputs ∧ fv.1.name = arg.1.name)
-  -- concreteEval does not succeed if the length of args is incorrect.
-  concreteEval_argmatch:
-    ∀ fn md args res, f.concreteEval = .some fn
-      → fn md args = .some res
-      → args.length = f.inputs.length
 
-instance LFuncWF.arg_nodup_decidable {T : LExprParams} (f : LFunc T):
-    Decidable (List.Nodup (f.inputs.map (·.1.name))) := by
+instance FuncWF.arg_nodup_decidable {IdentT ExprT TyT MetadataT : Type} [DecidableEq IdentT]
+    (f : Func IdentT ExprT TyT MetadataT):
+    Decidable (List.Nodup (f.inputs.map (·.1))) := by
   apply List.nodupDecidable
 
 instance LFuncWF.body_freevars_decidable {T : LExprParams} (f : LFunc T):
@@ -151,24 +185,36 @@ instance LFuncWF.body_freevars_decidable {T : LExprParams} (f : LFunc T):
   | .none => by
     apply isTrue; grind
 
--- LFuncWF.concreteEval_argmatch is not decidable.
+-- FuncWF.concreteEval_argmatch and LFuncWF.concreteEval_argmatch are not decidable.-- FuncWF.body_freevars is commented out as it's expression-type specific
+-- FuncWF.concreteEval_argmatch is not decidable.
 
 instance [Inhabited T.Metadata] [Inhabited T.IDMeta] : Inhabited (LFunc T) where
   default := { name := Inhabited.default, inputs := [], output := LMonoTy.bool }
 
-instance : ToFormat (LFunc T) where
+instance {IdentT ExprT TyT MetadataT : Type} [ToFormat IdentT] [ToFormat ExprT] [ToFormat TyT] [Inhabited ExprT] : ToFormat (Func IdentT ExprT TyT MetadataT) where
   format f :=
     let attr := if f.attr.isEmpty then f!"" else f!"@[{f.attr}]{Format.line}"
     let typeArgs := if f.typeArgs.isEmpty
                     then f!""
                     else f!"∀{f.typeArgs}."
-    let type := f!"{typeArgs} ({Signature.format f.inputs}) → {f.output}"
+    -- Format inputs recursively like Signature.format
+    let rec formatInputs (inputs : List (IdentT × TyT)) : Format :=
+      match inputs with
+      | [] => f!""
+      | [(k, v)] => f!"({k} : {v})"
+      | (k, v) :: rest => f!"({k} : {v}) " ++ formatInputs rest
+    let type := f!"{typeArgs} ({formatInputs f.inputs}) → {f.output}"
     let sep := if f.body.isNone then f!";" else f!" :="
     let body := if f.body.isNone then f!"" else Std.Format.indentD f!"({f.body.get!})"
     f!"{attr}\
        func {f.name} : {type}{sep}\
        {body}"
 
+-- Provide explicit instance for LFunc to ensure proper resolution
+-- Requires ToFormat for T.IDMeta (for identifiers in expressions) and T.Metadata (for Inhabited LExpr)
+instance [ToFormat T.IDMeta] [Inhabited T.Metadata] : ToFormat (LFunc T) where
+  format := format
+
 def LFunc.type [DecidableEq T.IDMeta] (f : (LFunc T)) : Except Format LTy := do
   if !(decide f.inputs.keys.Nodup) then
     .error f!"[{f.name}] Duplicates found in the formals!\

Strata/Languages/C_Simp/C_Simp.lean

@@ -28,6 +28,7 @@ abbrev Expression : Imperative.PureExpr := {
   Ident := Lambda.Identifier Unit,
   Expr := Lambda.LExpr CSimpLParams.mono,
   Ty := Lambda.LTy,
+  ExprMetadata := CSimpLParams.Metadata,
   TyEnv := Lambda.TEnv Unit,
   TyContext := Lambda.LContext ⟨Unit, Unit⟩,
   EvalEnv := Lambda.LState ⟨Unit, String⟩,

Strata/Languages/C_Simp/Verify.lean

@@ -48,6 +48,7 @@ def translate_stmt (s: Imperative.Stmt C_Simp.Expression C_Simp.Command) : Core.
   | .block l b _md => .block l (b.map translate_stmt) {}
   | .ite cond thenb elseb _md => .ite (translate_expr cond) (thenb.map translate_stmt) (elseb.map translate_stmt) {}
   | .loop guard measure invariant body _md => .loop (translate_expr guard) (translate_opt_expr measure) (translate_opt_expr invariant) (body.map translate_stmt) {}
+  | .funcDecl _ _ => panic! "C_Simp does not support function declarations"
   | .goto label _md => .goto label {}
   termination_by s.sizeOf
   decreasing_by

Strata/Languages/Core/CallGraph.lean

@@ -102,6 +102,7 @@ partial def extractCallsFromStatement (stmt : Statement) : List String :=
     extractCallsFromStatements elseBody
   | .loop _ _ _ body _ => extractCallsFromStatements body
   | .goto _ _ => []
+  | .funcDecl _ _ => []
 
 /-- Extract procedure calls from a list of statements -/
 partial def extractCallsFromStatements (stmts : List Statement) : List String :=

Strata/Languages/Core/Expressions.lean

@@ -22,6 +22,7 @@ abbrev Expression : Imperative.PureExpr :=
      EqIdent := inferInstanceAs (DecidableEq (Lambda.Identifier _))
      Expr := Lambda.LExpr ⟨⟨ExpressionMetadata, Visibility⟩, Lambda.LMonoTy⟩,
      Ty := Lambda.LTy,
+     ExprMetadata := ExpressionMetadata,
      TyEnv := @Lambda.TEnv Visibility,
      TyContext := @Lambda.LContext ⟨ExpressionMetadata, Visibility⟩,
      EvalEnv := Lambda.LState ⟨ExpressionMetadata, Visibility⟩ }

Strata/Languages/Core/FunctionType.lean

@@ -40,12 +40,12 @@ def typeCheck (C: Core.Expression.TyContext) (Env : Core.Expression.TyEnv) (func
   | some body =>
     -- Temporarily add formals in the context.
     let Env := Env.pushEmptyContext
-    let Env := Env.addToContext func.inputPolyTypes
+    let Env := Env.addToContext (LFunc.inputPolyTypes func)
     -- Type check and annotate the body, and ensure that it unifies with the
     -- return type.
     let (bodya, Env) ← LExpr.resolve C Env body
     let bodyty := bodya.toLMonoTy
-    let (retty, Env) ← func.outputPolyType.instantiateWithCheck C Env
+    let (retty, Env) ← (LFunc.outputPolyType func).instantiateWithCheck C Env
     let S ← Constraints.unify [(retty, bodyty)] Env.stateSubstInfo |>.mapError format
     let Env := Env.updateSubst S
     let Env := Env.popContext

Strata/Languages/Core/Statement.lean

@@ -126,6 +126,11 @@ def Statement.eraseTypes (s : Statement) : Statement :=
     let body' := Statements.eraseTypes bss
     .loop guard measure invariant body' md
   | .goto l md => .goto l md
+  | .funcDecl decl md =>
+    let decl' := { decl with
+      body := decl.body.map Lambda.LExpr.eraseTypes,
+      axioms := decl.axioms.map Lambda.LExpr.eraseTypes }
+    .funcDecl decl' md
   termination_by (Stmt.sizeOf s)
   decreasing_by
   all_goals simp_wf <;> simp [sizeOf] <;> omega
@@ -203,6 +208,7 @@ def Statement.modifiedVarsTrans
     Statements.modifiedVarsTrans π tbss ++ Statements.modifiedVarsTrans π ebss
   | .loop _ _ _ bss _ =>
     Statements.modifiedVarsTrans π bss
+  | .funcDecl _ _ => []  -- Function declarations don't modify variables
   termination_by (Stmt.sizeOf s)
 
 def Statements.modifiedVarsTrans
@@ -241,6 +247,11 @@ def Statement.getVarsTrans
     Statements.getVarsTrans π tbss ++ Statements.getVarsTrans π ebss
   | .loop _ _ _ bss  _ =>
     Statements.getVarsTrans π bss
+  | .funcDecl decl _ =>
+    -- Get variables from function body (including parameters for simplicity)
+    match decl.body with
+    | none => []
+    | some body => HasVarsPure.getVars body
   termination_by (Stmt.sizeOf s)
 
 def Statements.getVarsTrans
@@ -284,6 +295,7 @@ def Statement.touchedVarsTrans
   | .block _ bss _ => Statements.touchedVarsTrans π bss
   | .ite _ tbss ebss _ => Statements.touchedVarsTrans π tbss ++ Statements.touchedVarsTrans π ebss
   | .loop _ _ _ bss _ => Statements.touchedVarsTrans π bss
+  | .funcDecl decl _ => [decl.name]  -- Function declaration touches (defines) the function name
   termination_by (Stmt.sizeOf s)
 
 def Statements.touchedVarsTrans
@@ -347,6 +359,12 @@ def Statement.substFvar (s : Core.Statement)
           (Block.substFvar body fr to)
           metadata
   | .goto _ _ => s
+  | .funcDecl decl md =>
+    -- Substitute in function body and axioms
+    let decl' := { decl with
+      body := decl.body.map (Lambda.LExpr.substFvar · fr to),
+      axioms := decl.axioms.map (Lambda.LExpr.substFvar · fr to) }
+    .funcDecl decl' md
   termination_by s.sizeOf
   decreasing_by all_goals(simp_wf; try omega)
 end
@@ -373,6 +391,10 @@ def Statement.renameLhs (s : Core.Statement)
       if l.name == fr  then to else l)) pname args metadata
   | .block lbl b metadata =>
     .block lbl (Block.renameLhs b fr to) metadata
+  | .funcDecl decl md =>
+    -- Rename function name if it matches
+    let decl' := if decl.name == fr then { decl with name := to } else decl
+    .funcDecl decl' md
   | .havoc _ _ | .assert _ _ _ | .assume _ _ _ | .ite _ _ _ _
   | .loop _ _ _ _ _ | .goto _ _ | .cover _ _ _ => s
   termination_by s.sizeOf