Add batch processing to LeanRepl

README.md

@@ -57,6 +57,26 @@ Example output:
 
 showing any messages generated, and sorries with their goal states.
 
+## Batch Command mode
+
+Multiple commands can be run in a single batch using
+
+```json
+{ "cmds": ["theorem womp : 2 + 2 = 4 := by rfl", "theorem womp1 : 2 + 4 = 6 := by rfl"]}
+```
+
+All the same options from Command can be used and will be applied to each command in the `cmds` array. Additionally, you can specify the parrallelism mode using `mode`
+
+```json
+{ "cmds": ["theorem womp : 2 + 2 = 4 := by rfl", "theorem womp1 : 2 + 4 = 6 := by rfl"], "mode": "sequential"}
+{ "cmds": ["theorem womp : 2 + 2 = 4 := by rfl", "theorem womp1 : 2 + 4 = 6 := by rfl"], "mode": "naive"}
+{ "cmds": ["theorem womp : 2 + 2 = 4 := by rfl", "theorem womp1 : 2 + 4 = 6 := by rfl"], "mode": "parrallel", "buckets": 10}
+```
+`sequential` runs all of the commands sequentially. `naive` runs each command as its own multithreading task. `parrallel` splits the commands into `buckets` number of buckets and runs each bucket as a multithreading task.
+By default the parrallelism mode will be `sequential`
+
+Note that for parrallelism, both Command and Proof snapshots will not be saved.
+
 ## File mode
 
 There is a simple wrapper around command mode that allows reading in an entire file.

REPL/Frontend.lean

@@ -45,3 +45,18 @@ def processInput (input : String) (cmdState? : Option Command.State)
   | some cmdState => do
     pure ({ : Parser.ModuleParserState }, cmdState)
   processCommandsWithInfoTrees inputCtx parserState commandState
+
+
+/-
+  asTask but with a timeout
+-/
+def withTimeout {α : Type} (act : IO α) (timeoutMs : Nat) (prio := Task.Priority.default) : IO (Except IO.Error α) := do
+  let task ← IO.asTask act prio
+  for _ in [0:timeoutMs / 1000] do
+    if ← IO.hasFinished task then
+      return task.get
+    else
+      IO.sleep 1000
+  IO.cancel task
+  -- I'm not sure what the actual error code should be
+  return .error <| IO.Error.timeExpired 0 "timeout exceeded"

REPL/JSON.lean

@@ -26,6 +26,21 @@ If `env = some n`, builds on the existing environment `n`.
 structure Command extends CommandOptions where
   env : Option Nat
   cmd : String
+  gc : Option Bool := false
+deriving ToJson, FromJson
+
+structure BatchCommandOptions extends CommandOptions where
+  /-
+    mode = "sequential", "naive", "parrallel"
+    buckets is unused if mode is "sequential" or "naive"
+  -/
+  mode : Option String
+  buckets : Option Nat
+  timeout : Option Nat
+
+structure BatchCommand extends BatchCommandOptions where
+  env : Option Nat
+  cmds : Array String
 deriving ToJson, FromJson
 
 /-- Process a Lean file in a fresh environment. -/
@@ -120,7 +135,7 @@ A response to a Lean command.
 `env` can be used in later calls, to build on the stored environment.
 -/
 structure CommandResponse where
-  env : Nat
+  env : Option Nat
   messages : List Message := []
   sorries : List Sorry := []
   tactics : List Tactic := []
@@ -133,7 +148,9 @@ def Json.nonemptyList [ToJson α] (k : String) : List α → List (String × Jso
 
 instance : ToJson CommandResponse where
   toJson r := Json.mkObj <| .flatten [
-    [("env", r.env)],
+    match r.env with
+    | some x => [("env", x)]
+    | none => [],
     Json.nonemptyList "messages" r.messages,
     Json.nonemptyList "sorries" r.sorries,
     Json.nonemptyList "tactics" r.tactics,

REPL/Main.lean

@@ -110,6 +110,21 @@ def sorries (trees : List InfoTree) (env? : Option Environment) : M m (List Sorr
         let proofStateId ← proofState.mapM recordProofSnapshot
         return Sorry.of goal pos endPos proofStateId
 
+def sorriesGC (trees : List InfoTree) (env? : Option Environment) : IO (List Sorry) :=
+  trees.flatMap InfoTree.sorries |>.filter (fun t => match t.2.1 with
+    | .term _ none => false
+    | _ => true ) |>.mapM
+      fun ⟨ctx, g, pos, endPos⟩ => do
+        let (goal, _) ← match g with
+        | .tactic g => do
+           let s ← ProofSnapshot.create ctx none env? [g]
+           pure ("\n".intercalate <| (← s.ppGoals).map fun s => s!"{s}", some s)
+        | .term lctx (some t) => do
+           let s ← ProofSnapshot.create ctx lctx env? [] [t]
+           pure ("\n".intercalate <| (← s.ppGoals).map fun s => s!"{s}", some s)
+        | .term _ none => unreachable!
+        return Sorry.of goal pos endPos none
+
 def ppTactic (ctx : ContextInfo) (stx : Syntax) : IO Format :=
   ctx.runMetaM {} try
     Lean.PrettyPrinter.ppTactic ⟨stx⟩
@@ -125,6 +140,13 @@ def tactics (trees : List InfoTree) : M m (List Tactic) :=
       let proofStateId ← proofState.mapM recordProofSnapshot
       return Tactic.of goals tactic pos endPos proofStateId ns
 
+def tacticsGC (trees : List InfoTree) : IO (List Tactic) :=
+  trees.flatMap InfoTree.tactics |>.mapM
+    fun ⟨ctx, stx, goals, pos, endPos, ns⟩ => do
+      let goals := s!"{(← ctx.ppGoals goals)}".trim
+      let tactic := Format.pretty (← ppTactic ctx stx)
+      return Tactic.of goals tactic pos endPos none ns
+
 /-- Record a `ProofSnapshot` and generate a JSON response for it. -/
 def createProofStepReponse (proofState : ProofSnapshot) (old? : Option ProofSnapshot := none) :
     M m ProofStepResponse := do
@@ -184,18 +206,48 @@ def unpickleProofSnapshot (n : UnpickleProofState) : M IO (ProofStepResponse ⊕
   let (proofState, _) ← ProofSnapshot.unpickle n.unpickleProofStateFrom cmdSnapshot?
   Sum.inl <$> createProofStepReponse proofState
 
-/--
-Run a command, returning the id of the new environment, and any messages and sorries.
--/
-def runCommand (s : Command) : M IO (CommandResponse ⊕ Error) := do
-  let (cmdSnapshot?, notFound) ← do match s.env with
+
+def getCommandSnapshot (env : Option Nat) :  M IO (Option CommandSnapshot × Bool) :=  do match env with
   | none => pure (none, false)
   | some i => do match (← get).cmdStates[i]? with
     | some env => pure (some env, false)
     | none => pure (none, true)
-  if notFound then
-    return .inr ⟨"Unknown environment."⟩
-  let initialCmdState? := cmdSnapshot?.map fun c => c.cmdState
+
+def runCommandGCAux (initialCmdState? : Option Command.State) (s : Command) : IO (CommandResponse ⊕ Error) := do
+  let (_, messages, trees) ← try
+    IO.processInput s.cmd initialCmdState?
+  catch ex =>
+    return .inr ⟨ex.toString⟩
+  let messages ← messages.mapM fun m => Message.of m
+  -- For debugging purposes, sometimes we print out the trees here:
+  -- trees.forM fun t => do IO.println (← t.format)
+  let sorries ← sorriesGC trees (initialCmdState?.map (·.env))
+  let tactics ← match s.allTactics with
+  | some true => tacticsGC trees
+  | _ => pure []
+  let jsonTrees := match s.infotree with
+  | some "full" => trees
+  | some "tactics" => trees.flatMap InfoTree.retainTacticInfo
+  | some "original" => trees.flatMap InfoTree.retainTacticInfo |>.flatMap InfoTree.retainOriginal
+  | some "substantive" => trees.flatMap InfoTree.retainTacticInfo |>.flatMap InfoTree.retainSubstantive
+  | _ => []
+  let infotree ← if jsonTrees.isEmpty then
+    pure none
+  else
+    pure <| some <| Json.arr (← jsonTrees.toArray.mapM fun t => t.toJson none)
+  return .inl
+    { env := none,
+      messages,
+      sorries,
+      tactics
+      infotree }
+
+def runCommandGCAuxTimeout (initialCmdState? : Option Command.State) (s : Command) (timeout : Nat) : IO (CommandResponse ⊕ Error) := do
+  match ← IO.withTimeout (runCommandGCAux initialCmdState? s) timeout with
+  | .ok res => return res
+  | .error e => return .inr ⟨e.toString⟩
+
+def runCommandAux (cmdSnapshot? : Option CommandSnapshot) (initialCmdState? : Option Command.State) (s : Command) : M IO (CommandResponse ⊕ Error) := do
   let (cmdState, messages, trees) ← try
     IO.processInput s.cmd initialCmdState?
   catch ex =>
@@ -207,14 +259,6 @@ def runCommand (s : Command) : M IO (CommandResponse ⊕ Error) := do
   let tactics ← match s.allTactics with
   | some true => tactics trees
   | _ => pure []
-  let cmdSnapshot :=
-  { cmdState
-    cmdContext := (cmdSnapshot?.map fun c => c.cmdContext).getD
-      { fileName := "",
-        fileMap := default,
-        snap? := none,
-        cancelTk? := none } }
-  let env ← recordCommandSnapshot cmdSnapshot
   let jsonTrees := match s.infotree with
   | some "full" => trees
   | some "tactics" => trees.flatMap InfoTree.retainTacticInfo
@@ -225,20 +269,123 @@ def runCommand (s : Command) : M IO (CommandResponse ⊕ Error) := do
     pure none
   else
     pure <| some <| Json.arr (← jsonTrees.toArray.mapM fun t => t.toJson none)
+
+  let cmdSnapshot :=
+  { cmdState
+    cmdContext := (cmdSnapshot?.map fun c => c.cmdContext).getD
+      { fileName := "",
+        fileMap := default,
+        snap? := none,
+        cancelTk? := none } }
+  let env ← recordCommandSnapshot cmdSnapshot
   return .inl
-    { env,
+    { env := some env,
       messages,
       sorries,
       tactics
       infotree }
 
+/--
+Run a command, returning the id of the new environment, and any messages and sorries.
+-/
+def runCommand (s : Command) : M IO (CommandResponse ⊕ Error) := do
+  let (cmdSnapshot?, notFound) ← getCommandSnapshot s.env
+  if notFound then
+    return .inr ⟨"Unknown environment."⟩
+  let initialCmdState? := cmdSnapshot?.map fun c => c.cmdState
+  if s.gc = some true then
+    runCommandGCAux initialCmdState? s
+  else
+    runCommandAux cmdSnapshot? initialCmdState? s
+
+def splitArray {α : Type} (arr : Array α) (n : Nat) : Array (Array α) := Id.run do
+  if n ≤ 0 then #[]
+  else if n = 1 then #[arr]
+  else if arr.size = 0 then Array.replicate n #[]
+  else
+    let baseSize := arr.size / n
+    let remainder := arr.size % n
+
+    let mut result : Array (Array α) := #[]
+    let mut start : Nat := 0
+
+    for i in List.range n do
+      let extraElem := if i < remainder then 1 else 0
+      let endPos := start + baseSize + extraElem
+      let subArray := arr.extract start endPos
+      result := result.push subArray
+      start := start + baseSize + extraElem
+    result
+
+unsafe def batchVerifySequential (initialCmdState : Command.State) (cmds : Array Command) (timeout : Option Nat): IO (Array (CommandResponse ⊕ Error)) :=
+  let commandRunner :=
+    match timeout with
+    | some t => fun cmd => runCommandGCAuxTimeout initialCmdState cmd t
+    | none => runCommandGCAux initialCmdState
+  cmds.mapM commandRunner
+
+unsafe def batchVerifyParrallelNaive (initialCmdState : Command.State) (cmds : Array Command) (timeout : Option Nat) : IO (Array (CommandResponse ⊕ Error)) := do
+  let commandRunner :=
+    match timeout with
+    | some t => fun cmd => runCommandGCAuxTimeout initialCmdState cmd t
+    | none => runCommandGCAux initialCmdState
+  let tasks : Array (Task (Except IO.Error (CommandResponse ⊕ Error))) ← (cmds.mapM <| fun cmd => IO.asTask (commandRunner cmd)
+  )
+  tasks.mapM fun task => do
+    try
+      match task.get with
+      | .ok cmdres => return cmdres
+      | .error e => return .inr ⟨e.toString⟩
+    catch e =>
+      return .inr ⟨e.toString⟩
+
+unsafe def batchVerifyParrallel (commandState : Command.State) (cmds : Array Command) (buckets : Option Nat) (timeout : Option Nat) : IO (Array (CommandResponse ⊕ Error)) := do
+  let buckets :=
+    match buckets with
+    | some x => x
+    | none => max 50 cmds.size
+  let tasks ← (splitArray cmds buckets |>.mapM <| fun bucket => IO.asTask ( (batchVerifySequential commandState bucket timeout)))
+  tasks.flatMapM <|
+    fun task => do
+      try
+        match task.get with
+        | .ok cmdres => return cmdres
+        | .error e => return Array.replicate buckets (.inr ⟨e.toString⟩)
+      catch e =>
+        return Array.replicate buckets (.inr ⟨e.toString⟩)
+
 def processFile (s : File) : M IO (CommandResponse ⊕ Error) := do
   try
     let cmd ← IO.FS.readFile s.path
     runCommand { s with env := none, cmd }
   catch e =>
     pure <| .inr ⟨e.toString⟩
 
+unsafe def runBatchVerify (batch : BatchCommand) : M IO (Array (CommandResponse ⊕ Error) ⊕ Error) := do
+  let (cmdSnapshot?, notFound) ← getCommandSnapshot batch.env
+  if notFound then
+    return .inr ⟨"Unknown environment."⟩
+  let cmdState? := cmdSnapshot?.map fun c => c.cmdState
+  let commandState ← match cmdState? with
+    | none => do
+      let inputCtx   := Parser.mkInputContext "" "<input>"
+      let (header, _, messages) ← Parser.parseHeader inputCtx
+      let (env, messages) ← processHeader header {} messages inputCtx
+      pure (Command.mkState env messages {})
+    | some cmdState => do
+      pure cmdState
+  let cmds : Array Command := (batch.cmds.map fun cmd => { toCommandOptions := batch.toCommandOptions, env := none, cmd := cmd})
+  match batch.mode with
+  | some x =>
+    if x = "naive" then do
+      return .inl <| ← batchVerifyParrallelNaive commandState cmds batch.timeout
+    if x = "parrallel" then do
+      return .inl <| ← batchVerifyParrallel commandState cmds batch.buckets batch.timeout
+  | none =>
+    pure ()
+
+  return .inl <| ← batchVerifySequential commandState cmds batch.timeout
+
 /--
 Run a single tactic, returning the id of the new proof statement, and the new goals.
 -/
@@ -270,6 +417,12 @@ instance [ToJson α] [ToJson β] : ToJson (α ⊕ β) where
   | .inl a => toJson a
   | .inr b => toJson b
 
+
+structure Batch where
+  header : String
+  proofs : Array String
+deriving FromJson, ToJson
+
 /-- Commands accepted by the REPL. -/
 inductive Input
 | command : REPL.Command → Input
@@ -279,6 +432,7 @@ inductive Input
 | unpickleEnvironment : REPL.UnpickleEnvironment → Input
 | pickleProofSnapshot : REPL.PickleProofState → Input
 | unpickleProofSnapshot : REPL.UnpickleProofState → Input
+| batchVerify : REPL.BatchCommand → Input
 
 /-- Parse a user input string to an input command. -/
 def parse (query : String) : IO Input := do
@@ -299,6 +453,8 @@ def parse (query : String) : IO Input := do
     | .error _ => match fromJson? j with
     | .ok (r : REPL.Command) => return .command r
     | .error _ => match fromJson? j with
+    | .ok (r : REPL.BatchCommand) => return .batchVerify r
+    | .error _ => match fromJson? j with
     | .ok (r : REPL.File) => return .file r
     | .error e => throw <| IO.userError <| toString <| toJson <|
         (⟨"Could not parse as a valid JSON command:\n" ++ e⟩ : Error)
@@ -319,6 +475,7 @@ where loop : M IO Unit := do
   if query.startsWith "#" || query.startsWith "--" then loop else
   IO.println <| toString <| ← match ← parse query with
   | .command r => return toJson (← runCommand r)
+  | .batchVerify r => return toJson (← runBatchVerify r)
   | .file r => return toJson (← processFile r)
   | .proofStep r => return toJson (← runProofStep r)
   | .pickleEnvironment r => return toJson (← pickleCommandSnapshot r)

stats.txt

@@ -0,0 +1,24 @@
+# all profiling is one using meaure-command
+
+# test run on 1000 samples of
+# theorem womp (a b c : ℕ) : a + b + c = c + (b + a) := by sorry
+
+batch verify repl only, without env sharing: 78.0058819 seconds
+batch verify repl only, with env sharing: 2.0209996 seconds
+
+
+# first 20 from lean Goedel LM
+
+batch verify sequenctial: 472 seconds
+batch verify with 2 parrallel tasks: 444.80 seconds
+
+# first 1000 from lean Goedel LM
+
+batch verify sequenctial: 8930.02 seconds
+batch verify parrallel naive (1 proof per task): 1618.84 seconds
+batch verify with 2 parrallel tasks: 
+batch verify with 5 parrallel tasks:
+batch verify with 25 parrallel tasks: 1273.84 
+batch verify with 50 parrallel tasks: 1253.51 
+batch verify with 100 parrallel tasks: 1554.14
+batch verify with 200 parrallel tasks: 1613.46