share helpers between #model_check and #simulate, add widget support

Author: rnbguy

Veil/Frontend/DSL/Module/Elaborators.lean

@@ -550,6 +550,66 @@ def getModelCheckingMode (modeStx : Syntax) : ModelCheckingMode :=
     | `(modelCheckMode| compiled) => .compiled
     | _ => .default
 
+
+/-- Get all action label names for never-enabled action warnings. -/
+private def getActionLabelNames (mod : Module) : CommandElabM (List String) := do
+  let labelTypeName ← resolveGlobalConstNoOverload labelType
+  return mod.actions.map (fun a => s!"{labelTypeName}.{a.name}") |>.toList
+
+private def warnAboutTransitions (mod : Module) : CommandElabM Unit := do
+  let transitions := mod.procedures.filter (·.info.isTransition)
+  if transitions.isEmpty then return
+  let names := ", ".intercalate (transitions.map (·.info.name.toString) |>.toList)
+  logWarning m!"Explicit state model checking of transitions is SLOW!\n\n\
+    The current implementation enumerates all possible states and filters those satisfying \
+    the transition relation. Your specification has {transitions.size} \
+    transition{if transitions.size > 1 then "s" else ""}: {names}\n\n\
+    Consider encoding transitions as imperative actions where possible."
+
+private def resolveTheoryTerm (cmdName : String) (theoryTermOpt : Option Term)
+    (mod : Module) (instTerm : Term) : CommandElabM Term := do
+  match theoryTermOpt with
+  | some t => pure t
+  | none =>
+    unless mod.immutableComponents.isEmpty do
+      let fieldStrs := mod.immutableComponents.map (fun c => s!"{c.name} := ...")
+      let theoryExample := "{ " ++ ", ".intercalate fieldStrs.toList ++ " }"
+      throwError "This module has immutable fields, so you must specify the theory instantiation:\n\
+        {cmdName} {instTerm} {theoryExample}"
+    `({})
+
+/-- Prepend `name` with `mod.name`. -/
+private def mkIdentWithModName' (mod : Module) (name : Name) : Ident :=
+  Lean.mkIdent (mod.name ++ name)
+
+/-- Build search parameters for model checking / simulation. -/
+private def buildSearchParameters (mod : Module) (config : ModelCheckerConfig) : CommandElabM Term := do
+  let mkProp (sa : StateAssertion) : CommandElabM Term :=
+    `($(mkIdent ``Veil.ModelChecker.SafetyProperty.mk)
+        ($(mkIdent `name) := $(quote sa.name))
+        ($(mkIdent `property) := fun $(mkIdent `th) $(mkIdent `st) => $(mkIdentWithModName' mod sa.name) $(mkIdent `th) $(mkIdent `st)))
+  let safetyList ← `([$((← mod.invariants.mapM mkProp)),*])
+  let terminatingProp ← match mod.terminations[0]? with
+    | some t => mkProp t
+    | none => `($(mkIdent `default))
+  let earlyTermConds ← do
+    let base ← `([$(mkIdent ``Veil.ModelChecker.EarlyTerminationCondition.foundViolatingState),
+                  $(mkIdent ``Veil.ModelChecker.EarlyTerminationCondition.assertionFailed),
+                  $(mkIdent ``Veil.ModelChecker.EarlyTerminationCondition.deadlockOccurred)])
+    if config.maxDepth > 0 then `($base ++ [$(mkIdent ``Veil.ModelChecker.EarlyTerminationCondition.reachedDepthBound) $(quote config.maxDepth)])
+    else pure base
+  `({ $(mkIdent `invariants):ident := $safetyList, $(mkIdent `terminating):ident := $terminatingProp,
+      $(mkIdent `earlyTerminationConditions):ident := $earlyTermConds })
+
+/-- Display a TraceDisplayViewer widget with the given result JSON. -/
+private def displayResultWidget (stx : Syntax) (resultTerm : Term) : CommandElabM Unit := do
+  let widgetExpr ← `(open ProofWidgets.Jsx in
+    <ProofWidgets.TraceDisplayViewer result={$resultTerm} layout={"vertical"} />)
+  let html ← ← liftTermElabM <| ProofWidgets.HtmlCommand.evalCommandMHtml <| ← ``(ProofWidgets.HtmlEval.eval $widgetExpr)
+  liftCoreM <| Widget.savePanelWidgetInfo
+    (hash ProofWidgets.HtmlDisplayPanel.javascript)
+    (return json% { html: $(← Server.rpcEncode html) }) stx
+
 @[command_elab Veil.modelCheck]
 def elabModelCheck : CommandElab := fun stx => do
   -- Use dynamic trace class name for detailed profiling
@@ -561,19 +621,6 @@ def elabModelCheck : CommandElab := fun stx => do
     let cfg := stx[4]
     elabModelCheckCore stx mode instTerm theoryTermOpt cfg
 where
-  /-- Get the theory term, defaulting to `{}` if not provided and there are no theory fields.
-      Throws a helpful error if theory fields exist but no term was provided. -/
-  getTheoryTerm (theoryTermOpt : Option Term) (mod : Module) (instTerm : Term) : CommandElabM Term := do
-    match theoryTermOpt with
-    | some t => pure t
-    | none =>
-      unless mod.immutableComponents.isEmpty do
-        let fieldStrs := mod.immutableComponents.map (fun c => s!"{c.name} := ...")
-        let theoryExample := "{ " ++ ", ".intercalate fieldStrs.toList ++ " }"
-        throwError "This module has immutable fields, so you must specify the theory instantiation:\n\
-          #model_check {instTerm} {theoryExample}"
-      `({})
-
   /-- Generate the model source for compilation:
       1. Insert `set_option veil.__modelCheckCompileMode true` after imports
       2. Keep everything up to the point where the spec was finalized
@@ -599,65 +646,22 @@ where
     let modelCheckCmd := String.Pos.Raw.extract src modelCheckStart modelCheckEnd
     return beforeImports ++ compileModePreamble ++ afterImportsToSpecFinalized ++ "\n" ++ modelCheckCmd ++ "\n"
 
-  /-- Prepend `name` with `mod.name`. Useful when expressions are printed out for debugging. -/
-  mkIdentWithModName (mod : Module) (name : Name) : Ident :=
-    Lean.mkIdent (mod.name ++ name)
-
-  /-- Display a TraceDisplayViewer widget with the given result term. -/
-  displayResultWidget (stx : Syntax) (resultTerm : Term) : CommandElabM Unit := do
-    let widgetExpr ← `(open ProofWidgets.Jsx in
-      <ProofWidgets.TraceDisplayViewer result={$resultTerm} layout={"vertical"} />)
-    let html ← ← liftTermElabM <| ProofWidgets.HtmlCommand.evalCommandMHtml <| ← ``(ProofWidgets.HtmlEval.eval $widgetExpr)
-    liftCoreM <| Widget.savePanelWidgetInfo
-      (hash ProofWidgets.HtmlDisplayPanel.javascript)
-      (return json% { html: $(← Server.rpcEncode html) }) stx
-
-  mkSearchParameters (mod : Module) (config : ModelCheckerConfig) : CommandElabM Term := do
-    -- Build SafetyProperty.mk syntax for a StateAssertion
-    let mkProp (sa : StateAssertion) : CommandElabM Term :=
-      `($(mkIdent ``Veil.ModelChecker.SafetyProperty.mk)
-          ($(mkIdent `name) := $(quote sa.name))
-          ($(mkIdent `property) := fun $(mkIdent `th) $(mkIdent `st) => $(mkIdentWithModName mod sa.name) $(mkIdent `th) $(mkIdent `st)))
-    let safetyList ← `([$((← mod.invariants.mapM mkProp)),*])
-    let terminatingProp ← match mod.terminations[0]? with
-      | some t => mkProp t
-      | none => `($(mkIdent `default))
-    let earlyTermConds ← do
-      let base ← `([$(mkIdent ``Veil.ModelChecker.EarlyTerminationCondition.foundViolatingState),
-                    $(mkIdent ``Veil.ModelChecker.EarlyTerminationCondition.assertionFailed),
-                    $(mkIdent ``Veil.ModelChecker.EarlyTerminationCondition.deadlockOccurred)])
-      if config.maxDepth > 0 then `($base ++ [$(mkIdent ``Veil.ModelChecker.EarlyTerminationCondition.reachedDepthBound) $(quote config.maxDepth)])
-      else pure base
-    `({ $(mkIdent `invariants):ident := $safetyList, $(mkIdent `terminating):ident := $terminatingProp,
-        $(mkIdent `earlyTerminationConditions):ident := $earlyTermConds })
-
   /-- Build the core model checker call syntax (without parallel config). -/
   mkModelCheckerCall (mod : Module) (config : ModelCheckerConfig)
       (instTerm theoryTerm : Term) : CommandElabM Term := do
     let inst := mkVeilImplementationDetailIdent `inst
     let th := mkVeilImplementationDetailIdent `th
     let instSortArgs ← (← mod.sortIdents).mapM fun sortIdent => `($inst.$(sortIdent))
-    let sp ← mkSearchParameters mod config
+    let sp ← buildSearchParameters mod config
     -- Model checker call with type annotation to help inference
     -- Note: findReachable takes parallelCfg, progressInstanceId, and cancelToken as the last three args
     `((let $inst : $instantiationType := $instTerm
        let $th : $theoryIdent $instSortArgs* := $theoryTerm
        $(mkIdent ``Veil.ModelChecker.Concrete.findReachable)
          ($(mkIdent `inhabσ) := $instInhabitedStateFieldConcreteType $instSortArgs*)
-         ($(mkIdentWithModName mod `enumerableTransitionSystem) $instSortArgs* $th)
+         ($(mkIdentWithModName' mod `enumerableTransitionSystem) $instSortArgs* $th)
          $sp : _ → _ → _ → IO _))
 
-  /-- Warn if the module contains transitions (which are slow to model check). -/
-  warnAboutTransitions (mod : Module) : CommandElabM Unit := do
-    let transitions := mod.procedures.filter (·.info.isTransition)
-    if transitions.isEmpty then return
-    let names := ", ".intercalate (transitions.map (·.info.name.toString) |>.toList)
-    logWarning m!"Explicit state model checking of transitions is SLOW!\n\n\
-      The current implementation enumerates all possible states and filters those satisfying \
-      the transition relation. Your specification has {transitions.size} \
-      transition{if transitions.size > 1 then "s" else ""}: {names}\n\n\
-      Consider encoding transitions as imperative actions where possible."
-
   /-- Create an error JSON object. -/
   errorJson (msg : String) : Json := Json.mkObj [("error", msg)]
 
@@ -743,11 +747,6 @@ where
       Term.synthesizeSyntheticMVarsNoPostponing
       unsafe Meta.evalExpr (IO Lean.Json) (mkApp (mkConst ``IO) (mkConst ``Lean.Json)) (← instantiateMVars expr)
 
-  /-- Get all action label names for never-enabled action warnings. -/
-  getActionLabelNames (mod : Module) : CommandElabM (List String) := do
-    let labelTypeName ← resolveGlobalConstNoOverload labelType
-    return mod.actions.map (fun a => s!"{labelTypeName}.{a.name}") |>.toList
-
   /-- Log model checking result. -/
   logModelCheckResult (stx : Syntax) (resultJson : Json) : CommandElabM Unit := do
     let msg := TraceDisplay.formatModelCheckingResult resultJson
@@ -805,7 +804,7 @@ where
     let mod ← getCurrentModule (errMsg := "You cannot #model_check outside of a Veil module!")
     mod.throwIfSpecNotFinalized
 
-    let theoryTerm ← getTheoryTerm theoryTermOpt mod instTerm
+    let theoryTerm ← resolveTheoryTerm "#model_check" theoryTermOpt mod instTerm
 
     warnAboutTransitions mod
     let config ← elabModelCheckerConfig cfg
@@ -926,7 +925,7 @@ private def mkSimulatorCall (mod : Module) (instTerm theoryTerm : Term)
   `((let $inst : $instantiationType := $instTerm
      let $th : $theoryIdent $instSortArgs* := $theoryTerm
      $(mkIdent ``Veil.ModelChecker.Simulation.simulate)
-       ($(Lean.mkIdent (mod.name ++ `enumerableTransitionSystem)) $instSortArgs* $th)
+       ($(mkIdentWithModName' mod `enumerableTransitionSystem) $instSortArgs* $th)
        $sp $th $cfgTerm))
 
 @[command_elab Veil.simulate]
@@ -936,14 +935,15 @@ def elabSimulate : CommandElab := fun stx => do
     let theoryTermOpt : Option Term := if stx[2].isNone then none else some ⟨stx[2][0]⟩
     let mod ← getCurrentModule (errMsg := "You cannot #simulate outside of a Veil module!")
     mod.throwIfSpecNotFinalized
-    let theoryTerm ← getTheoryTerm theoryTermOpt mod instTerm
+    let theoryTerm ← resolveTheoryTerm "#simulate" theoryTermOpt mod instTerm
+    warnAboutTransitions mod
     let cfg0 ← elabSimulateConfig stx[3]
     let opts ← getOptions
     let maxTraces := if cfg0.maxTraces == 10000 then veil.simulate.maxTraces.get opts else cfg0.maxTraces
     let maxSteps := if cfg0.maxSteps == 100 then veil.simulate.maxSteps.get opts else cfg0.maxSteps
     let cfg : ModelChecker.Simulation.SimulateConfig := { cfg0 with maxTraces, maxSteps }
     let mcCfg : ModelCheckerConfig := { maxDepth := 0, sequential := false, parallelCfg := none }
-    let sp ← mkSearchParameters mod mcCfg
+    let sp ← buildSearchParameters mod mcCfg
     let callExpr ← mkSimulatorCall mod instTerm theoryTerm sp cfg
     let wrappedCallExpr ← `(Functor.map (fun r => Lean.Json.mkObj [
         ("result",      Lean.toJson (Veil.ModelChecker.Simulation.SimulateResult.result r)),
@@ -972,49 +972,16 @@ def elabSimulate : CommandElab := fun stx => do
     let stepsPerSec := if elapsedMs > 0 then totalSteps * 1000 / elapsedMs else 0
     let isViolation := resultJson.getObjValD "result" == Json.str "found_violation" ||
       resultJson.getObjValD "error" != .null
+    -- Log simulation-specific summary
     let summary := if isViolation then
       s!"simulation: found violation at depth {depth} (trace #{tracesRun}, {elapsedMs}ms, seed := {seed}). A shorter violation may exist at depth < {depth}."
     else
       s!"simulation: no violation in {tracesRun} traces ({totalSteps} steps, {elapsedMs}ms, {stepsPerSec} steps/s, seed := {seed}). Not exhaustive -- use #model_check for full coverage."
-    let details := TraceDisplay.formatModelCheckingResult resultJson
-    let msg := summary ++ "\n" ++ details
-    let violationIsError := veil.violationIsError.get opts
-    if isViolation && violationIsError then logErrorAt stx msg else logInfoAt stx msg
-where
-  /-- Get the theory term, defaulting to `{}` if not provided and there are no theory fields.
-      Throws a helpful error if theory fields exist but no term was provided. -/
-  getTheoryTerm (theoryTermOpt : Option Term) (mod : Module) (instTerm : Term) : CommandElabM Term := do
-    match theoryTermOpt with
-    | some t => pure t
-    | none =>
-      unless mod.immutableComponents.isEmpty do
-        let fieldStrs := mod.immutableComponents.map (fun c => s!"{c.name} := ...")
-        let theoryExample := "{ " ++ ", ".intercalate fieldStrs.toList ++ " }"
-        throwError "This module has immutable fields, so you must specify the theory instantiation:\n\
-          #simulate {instTerm} {theoryExample}"
-      `({})
-
-  /-- Prepend `name` with `mod.name`. Useful when expressions are printed out for debugging. -/
-  mkIdentWithModName (mod : Module) (name : Name) : Ident :=
-    Lean.mkIdent (mod.name ++ name)
-
-  /-- Build search parameters reused by simulator execution. -/
-  mkSearchParameters (mod : Module) (config : ModelCheckerConfig) : CommandElabM Term := do
-    let mkProp (sa : StateAssertion) : CommandElabM Term :=
-      `($(mkIdent ``Veil.ModelChecker.SafetyProperty.mk)
-          ($(mkIdent `name) := $(quote sa.name))
-          ($(mkIdent `property) := fun $(mkIdent `th) $(mkIdent `st) => $(mkIdentWithModName mod sa.name) $(mkIdent `th) $(mkIdent `st)))
-    let safetyList ← `([$((← mod.invariants.mapM mkProp)),*])
-    let terminatingProp ← match mod.terminations[0]? with
-      | some t => mkProp t
-      | none => `($(mkIdent `default))
-    let earlyTermConds ← do
-      let base ← `([$(mkIdent ``Veil.ModelChecker.EarlyTerminationCondition.foundViolatingState),
-                    $(mkIdent ``Veil.ModelChecker.EarlyTerminationCondition.assertionFailed),
-                    $(mkIdent ``Veil.ModelChecker.EarlyTerminationCondition.deadlockOccurred)])
-      if config.maxDepth > 0 then `($base ++ [$(mkIdent ``Veil.ModelChecker.EarlyTerminationCondition.reachedDepthBound) $(quote config.maxDepth)])
-      else pure base
-    `({ $(mkIdent `invariants):ident := $safetyList, $(mkIdent `terminating):ident := $terminatingProp,
-        $(mkIdent `earlyTerminationConditions):ident := $earlyTermConds })
-
+    logInfoAt stx summary
+    -- Log the same trace display as #model_check
+    elabModelCheck.logModelCheckResult stx resultJson
+    -- Display the same TraceDisplayViewer widget as #model_check
+    let (instanceId, _) ← ModelChecker.Concrete.allocProgressInstance (← getActionLabelNames mod)
+    ModelChecker.Concrete.finishProgress instanceId resultJson
+    ModelChecker.displayStreamingProgress stx instanceId
 end Veil