Implement timeout option

Author: augustepoiroux

REPL/JSON.lean

@@ -11,7 +11,14 @@ open Lean Elab InfoTree
 
 namespace REPL
 
-structure CommandOptions where
+/-- Base structure with timeout that all command types can inherit from -/
+structure BaseOptions where
+  /--
+  Optional timeout in milliseconds. If none, no timeout will be applied.
+  -/
+  timeout : Option Nat := none
+
+structure CommandOptions extends BaseOptions where
   allTactics : Option Bool := none
   rootGoals : Option Bool := none
   /--
@@ -37,7 +44,7 @@ deriving FromJson
 /--
 Run a tactic in a proof state.
 -/
-structure ProofStep where
+structure ProofStep extends BaseOptions where
   proofState : Nat
   tactic : String
 deriving ToJson, FromJson
@@ -169,21 +176,21 @@ structure Error where
   message : String
 deriving ToJson, FromJson
 
-structure PickleEnvironment where
+structure PickleEnvironment extends BaseOptions where
   env : Nat
   pickleTo : System.FilePath
 deriving ToJson, FromJson
 
-structure UnpickleEnvironment where
+structure UnpickleEnvironment extends BaseOptions where
   unpickleEnvFrom : System.FilePath
 deriving ToJson, FromJson
 
-structure PickleProofState where
+structure PickleProofState extends BaseOptions where
   proofState : Nat
   pickleTo : System.FilePath
 deriving ToJson, FromJson
 
-structure UnpickleProofState where
+structure UnpickleProofState extends BaseOptions where
   unpickleProofStateFrom : System.FilePath
   env : Option Nat
 deriving ToJson, FromJson

REPL/Main.lean

@@ -361,6 +361,23 @@ def runProofStep (s : ProofStep) : M IO (ProofStepResponse ⊕ Error) := do
     catch ex =>
       return .inr ⟨"Lean error:\n" ++ ex.toString⟩
 
+/--
+Run a task with an optional timeout in milliseconds.
+-/
+def runWithTimeout (timeout? : Option Nat) (task : M IO α) : M IO α :=
+  match timeout? with
+  | none => task
+  | some timeout => do
+    let jobTask ← IO.asTask (prio := .dedicated) (StateT.run task (← get))
+    let timeoutTask : IO (α × State) := do
+      IO.sleep timeout.toUInt32
+      throw <| IO.userError s!"Operation timed out after {timeout}ms"
+    match ← IO.waitAny [jobTask, ← IO.asTask timeoutTask] with
+    | .ok (a, state) => do
+      modify fun _ => state
+      return a
+    | .error e => throw e
+
 end REPL
 
 open REPL
@@ -425,14 +442,16 @@ where loop : M IO Unit := do
   if query = "" then
     return ()
   if query.startsWith "#" || query.startsWith "--" then loop else
-  IO.println <| toString <| ← match ← parse query with
-  | .command r => return toJson (← runCommand r)
-  | .file r => return toJson (← processFile r)
-  | .proofStep r => return toJson (← runProofStep r)
-  | .pickleEnvironment r => return toJson (← pickleCommandSnapshot r)
-  | .unpickleEnvironment r => return toJson (← unpickleCommandSnapshot r)
-  | .pickleProofSnapshot r => return toJson (← pickleProofSnapshot r)
-  | .unpickleProofSnapshot r => return toJson (← unpickleProofSnapshot r)
+  IO.println <| toString <| ← try
+      match ← parse query with
+      | .command r => return toJson (← runWithTimeout r.timeout (runCommand r))
+      | .file r => return toJson (← runWithTimeout r.timeout (processFile r))
+      | .proofStep r => return toJson (← runWithTimeout r.timeout (runProofStep r))
+      | .pickleEnvironment r => return toJson (← runWithTimeout r.timeout (pickleCommandSnapshot r))
+      | .unpickleEnvironment r => return toJson (← runWithTimeout r.timeout (unpickleCommandSnapshot r))
+      | .pickleProofSnapshot r => return toJson (← runWithTimeout r.timeout (pickleProofSnapshot r))
+      | .unpickleProofSnapshot r => return toJson (← runWithTimeout r.timeout (unpickleProofSnapshot r))
+    catch e => return toJson (⟨e.toString⟩ : Error)
   printFlush "\n" -- easier to parse the output if there are blank lines
   loop
 

test/timeout.expected.out

@@ -0,0 +1,46 @@
+{"sorries":
+ [{"proofState": 0,
+   "pos": {"line": 1, "column": 49},
+   "goal": "x : Nat\nh1 : x = 2\n⊢ x = 2",
+   "endPos": {"line": 1, "column": 54}}],
+ "messages":
+ [{"severity": "warning",
+   "pos": {"line": 1, "column": 8},
+   "endPos": {"line": 1, "column": 10},
+   "data": "declaration uses 'sorry'"}],
+ "env": 0}
+
+{"proofStatus": "Completed", "proofState": 1, "goals": []}
+
+{"sorries":
+ [{"proofState": 2,
+   "pos": {"line": 1, "column": 49},
+   "goal": "x : Nat\nh1 : x = 2\n⊢ x = 2",
+   "endPos": {"line": 1, "column": 54}}],
+ "messages":
+ [{"severity": "warning",
+   "pos": {"line": 1, "column": 8},
+   "endPos": {"line": 1, "column": 10},
+   "data": "declaration uses 'sorry'"}],
+ "env": 1}
+
+{"proofStatus": "Completed", "proofState": 3, "goals": []}
+
+{"message": "Operation timed out after 1ms"}
+
+{"message": "Unknown proof state."}
+
+{"sorries":
+ [{"proofState": 4,
+   "pos": {"line": 1, "column": 49},
+   "goal": "x : Nat\nh1 : x = 2\n⊢ x = 2",
+   "endPos": {"line": 1, "column": 54}}],
+ "messages":
+ [{"severity": "warning",
+   "pos": {"line": 1, "column": 8},
+   "endPos": {"line": 1, "column": 10},
+   "data": "declaration uses 'sorry'"}],
+ "env": 2}
+
+{"proofStatus": "Completed", "proofState": 5, "goals": []}
+

test/timeout.in

@@ -0,0 +1,15 @@
+{"cmd": "theorem aa (x : Nat) (h1 : x  = 2) : x = 2 := by sorry"}
+
+{"tactic": "assumption", "proofState": 0}
+
+{"cmd": "theorem aa (x : Nat) (h1 : x  = 2) : x = 2 := by sorry", "timeout": 1000}
+
+{"tactic": "assumption", "proofState": 2}
+
+{"cmd": "theorem aa (x : Nat) (h1 : x  = 2) : x = 2 := by sorry", "timeout": 1}
+
+{"tactic": "assumption", "proofState": 4}
+
+{"cmd": "theorem aa (x : Nat) (h1 : x  = 2) : x = 2 := by sorry"}
+
+{"tactic": "assumption", "proofState": 4}