reorganise files

TrainingData.lean

@@ -1,5 +1,4 @@
 import TrainingData.Frontend
 import TrainingData.InfoTree.Basic
 import TrainingData.InfoTree.ToJson
-import TrainingData.InfoTree.TacticInvocation.Basic
-import TrainingData.Mathlib
+import TrainingData.Mathlib

TrainingData/InfoTree/Basic.lean

@@ -167,3 +167,61 @@ def findTacticNodes (t : InfoTree) : List (TacticInfo × ContextInfo × Persiste
   | _ => none
 
 end Lean.Elab.InfoTree
+
+
+namespace Lean.Elab.TacticInfo
+
+/-- Return the range of the tactic, as a pair of file positions. -/
+def range (info : TacticInfo) (ctx : ContextInfo) : Position × Position := ctx.fileMap.stxRange info.stx
+
+/-- Pretty print a tactic. -/
+def pp (info : TacticInfo) (ctx : ContextInfo) : IO Format :=
+  ctx.runMetaM {} try
+    Lean.PrettyPrinter.ppTactic ⟨info.stx⟩
+  catch _ =>
+    pure "<failed to pretty print>"
+
+open Meta
+
+/-- Run a tactic on the goals stored in a `TacticInfo`. -/
+def runMetaMGoalsBefore (info : TacticInfo) (ctx : ContextInfo) (x : List MVarId → MetaM α) : IO α := do
+  ctx.runMetaM {} <| Meta.withMCtx info.mctxBefore <| x info.goalsBefore
+
+/-- Run a tactic on the after goals stored in a `TacticInfo`. -/
+def runMetaMGoalsAfter (info : TacticInfo) (ctx : ContextInfo) (x : List MVarId → MetaM α) : IO α := do
+  ctx.runMetaM {} <| Meta.withMCtx info.mctxAfter <| x info.goalsAfter
+
+/-- Run a tactic on the main goal stored in a `TacticInfo`. -/
+def runMetaM (info : TacticInfo) (ctx : ContextInfo) (x : MVarId → MetaM α) : IO α := do
+  match info.goalsBefore.head? with
+  | none => throw <| IO.userError s!"No goals at {← info.pp ctx}"
+  | some g => info.runMetaMGoalsBefore ctx fun _ => do g.withContext <| x g
+
+def mainGoal (info : TacticInfo) (ctx : ContextInfo) : IO Expr :=
+  info.runMetaM ctx (fun g => do instantiateMVars (← g.getType))
+
+def formatMainGoal (info : TacticInfo) (ctx : ContextInfo) : IO Format :=
+  info.runMetaM ctx (fun g => do ppExpr (← instantiateMVars (← g.getType)))
+
+def goalState (info : TacticInfo) (ctx : ContextInfo) : IO (List Format) := do
+  info.runMetaMGoalsBefore ctx (fun gs => gs.mapM fun g => do Meta.ppGoal g)
+
+def goalStateAfter (info : TacticInfo) (ctx : ContextInfo) : IO (List Format) := do
+  info.runMetaMGoalsAfter ctx (fun gs => gs.mapM fun g => do Meta.ppGoal g)
+
+def ppExpr (info : TacticInfo) (ctx : ContextInfo) (e : Expr) : IO Format :=
+  info.runMetaM ctx (fun _ => do Meta.ppExpr (← instantiateMVars e))
+
+end Lean.Elab.TacticInfo
+
+namespace Lean.Elab.InfoTree
+
+/--
+Finds all tactic invocations in an `InfoTree`,
+ignoring structuring tactics (e.g. `by`, `;`, multiline tactics, parenthesized tactics).
+-/
+def tactics (t : InfoTree) : List (TacticInfo × ContextInfo) :=
+  t.findTacticNodes.map (fun ⟨i, ctx, _⟩ => ⟨i, ctx⟩)
+    |>.filter fun i => i.1.isSubstantive
+
+end Lean.Elab.InfoTree

TrainingData/InfoTree/ToJson.lean

@@ -1,4 +1,4 @@
-import TrainingData.InfoTree.TacticInvocation.Basic
+import TrainingData.InfoTree.Basic
 
 /-!
 # Exporting an `InfoTree` as Json
@@ -53,15 +53,14 @@ deriving ToJson
 
 -- Note: this is not responsible for converting the children to Json.
 def TacticInfo.toJson (i : TacticInfo) (ctx : ContextInfo) : IO TacticInfo.Json := do
-  let I : TacticInvocation := ⟨i, ctx, .empty⟩
   return {
     name := i.name?
     stx :=
-    { pp := Format.pretty (← I.pp),
+    { pp := Format.pretty (← i.pp ctx),
       -- raw := toString i.info.stx,
       range := i.stx.toRange ctx },
-    goalsBefore := (← I.goalState).map Format.pretty,
-    goalsAfter := (← I.goalStateAfter).map Format.pretty }
+    goalsBefore := (← i.goalState ctx).map Format.pretty,
+    goalsAfter := (← i.goalStateAfter ctx).map Format.pretty }
 
 structure CommandInfo.Json where
   elaborator : Option Name

scripts/comment_data.lean

@@ -1,6 +1,5 @@
 import TrainingData.Frontend
 import TrainingData.InfoTree.ToJson
-import TrainingData.InfoTree.TacticInvocation.Basic
 import Mathlib.Data.String.Defs
 import Mathlib.Lean.CoreM
 import Std.Lean.Util.Path

scripts/goal_comments.lean

@@ -1,6 +1,5 @@
 import TrainingData.Frontend
 import TrainingData.InfoTree.ToJson
-import TrainingData.InfoTree.TacticInvocation.Basic
 import Mathlib.Data.String.Defs
 import Mathlib.Lean.CoreM
 import Mathlib.Tactic.Change
@@ -33,12 +32,12 @@ instance : DecidableRel ((· : Range) < ·) := by
   change DecidableRel fun _ _ => _ ∧ _
   infer_instance
 
-namespace Lean.Elab.TacticInvocation
+namespace Lean.Elab.TacticInfo
 
-def rangesAndGoals (i : TacticInvocation) : IO (Range × String) := do
-  return ⟨i.range, (Format.joinSep (← i.goalStateAfter) "\n").pretty 1000000⟩
+def rangesAndGoals (i : TacticInfo) (ctx : ContextInfo) : IO (Range × String) := do
+  return ⟨i.range ctx, (Format.joinSep (← i.goalStateAfter ctx) "\n").pretty 1000000⟩
 
-end Lean.Elab.TacticInvocation
+end Lean.Elab.TacticInfo
 
 partial def dropEnclosed (L : List (Range × String)) : List (Range × String) :=
   let L' := L.filter fun ⟨r, _⟩ => ¬ L.any fun ⟨r', _⟩ => r < r'
@@ -67,7 +66,7 @@ def goalComments (args : Cli.Parsed) : IO UInt32 := do
     trees := trees.bind InfoTree.retainTacticInfo
     trees := trees.bind InfoTree.retainOriginal
     trees := trees.bind InfoTree.retainSubstantive
-    let L₁ ← (trees.bind InfoTree.tactics).mapM TacticInvocation.rangesAndGoals
+    let L₁ ← (trees.bind InfoTree.tactics).mapM fun ⟨i, c⟩ => i.rangesAndGoals c
     let L₂ := dropEnclosed L₁ |>.filter fun ⟨⟨⟨l₁, _⟩, ⟨l₂, _⟩⟩, _⟩  => l₁ = l₂
     let L₃ := (L₂.map fun ⟨r, s⟩ => (r, justTheGoal s)) |>.filter fun ⟨_, s⟩ => s != ""
     let mut src := (← moduleSource module).splitOn "\n"

scripts/tactic_benchmark.lean

@@ -100,6 +100,7 @@ def useAesop : TacticM Unit := do evalTactic (← `(tactic| aesop))
 def useExact? : TacticM Unit := do evalTactic (← `(tactic| exact?))
 def useRfl : TacticM Unit := do evalTactic (← `(tactic| intros; rfl))
 def useSimpAll : TacticM Unit := do evalTactic (← `(tactic| intros; simp_all))
+def useOmega : TacticM Unit := do evalTactic (← `(tactic| intros; omega))
 
 open Cli System
 
@@ -111,6 +112,7 @@ def tacticBenchmarkMain (args : Cli.Parsed) : IO UInt32 := do
     if args.hasFlag "exact" then pure useExact? else
     if args.hasFlag "rfl" then pure useRfl else
     if args.hasFlag "simp_all" then pure useSimpAll else
+    if args.hasFlag "omega" then pure useOmega else
     throw <| IO.userError "Specify a tactic, e.g. `--aesop`"
   let result := runTacticAtDecls module (fun _ => pure true) tac
   IO.println s!"{module}"
@@ -129,6 +131,7 @@ def tactic_benchmark : Cmd := `[Cli|
     "exact";       "Use `exact?`."
     "rfl";         "Use `intros; rfl`."
     "simp_all";    "Use `intros; simp_all`."
+    "omega";       "Use `intros; omega`."
 
   ARGS:
     module : ModuleName; "Lean module to compile and export InfoTrees."

scripts/tactic_benchmark.sh

@@ -4,10 +4,20 @@
 # or `scripts/tactic_benchmark.sh --simp_all Mathlib.Logic.Hydra` to run on just one file.
 # Results will go in `out/tactic_benchmark/simp_all/Mathlib.Logic.Hydra.bench`.
 
+FLAGS=()
+ARGS=()
 
-if [ "$#" -eq 1 ]; then
+for arg in "$@"; do
+    if [[ $arg == --* ]]; then
+        FLAGS+=("$arg")
+    else
+        ARGS+=("$arg")
+    fi
+done
+
+if [ ${#ARGS[@]} -eq 0 ]; then
   lake build tactic_benchmark
-  parallel -j32 ./scripts/tactic_benchmark.sh $1 -- ::: `cat .lake/packages/mathlib/Mathlib.lean | sed -e 's/import //'`
+  parallel -j32 ./scripts/tactic_benchmark.sh ${FLAGS[@]} ::: `cat .lake/packages/mathlib/Mathlib.lean | sed -e 's/import //'`
 else
   DIR=out/tactic_benchmark/${1#--}
   mkdir -p $DIR
@@ -17,6 +27,6 @@ else
     if [ ! -f .lake/build/bin/tactic_benchmark ]; then
       lake build tactic_benchmark
     fi
-    timeout 5m .lake/build/bin/tactic_benchmark $1 $mod > $DIR/$mod._bench && mv $DIR/$mod._bench $DIR/$mod.bench
+    timeout 5m .lake/build/bin/tactic_benchmark ${FLAGS[@]} $mod > $DIR/$mod._bench && mv $DIR/$mod._bench $DIR/$mod.bench
   fi
 fi

scripts/training_data.lean

@@ -1,6 +1,5 @@
 import TrainingData.Frontend
 import TrainingData.InfoTree.ToJson
-import TrainingData.InfoTree.TacticInvocation.Basic
 import Mathlib.Data.String.Defs
 import Mathlib.Lean.CoreM
 import Std.Lean.Util.Path
@@ -10,27 +9,27 @@ import Cli
 open Lean Elab IO Meta
 open Cli System
 
-namespace Lean.Elab.TacticInvocation
+namespace Lean.Elab.TacticInfo
 
-def verboseTrainingData (module : Name) (i : TacticInvocation) : IO String := do
+def verboseTrainingData (module : Name) (i : TacticInfo) (ctx : ContextInfo) : IO String := do
   let mut result := "===\n"
   result := result ++ s!"{module}\n---\n"
-  let sourceUpToTactic := Substring.mk (← moduleSource module) 0 (i.info.stx.getPos?.getD 0)
+  let sourceUpToTactic := Substring.mk (← moduleSource module) 0 (i.stx.getPos?.getD 0)
   let chunks := sourceUpToTactic.splitOn "\n\n"
   let declUpToTactic := chunks.getLast!.toString
   let offset := chunks.dropLast.foldl (init := 0) (fun t c => t + (c.toString.count '\n') + 2)
   result := result ++ s!"{offset}\n---\n{declUpToTactic}\n---\n"
-  result := result ++ (Format.joinSep (← i.goalState) "\n").pretty ++ "\n---\n"
-  let ⟨⟨startLine, startCol⟩, ⟨endLine, endCol⟩⟩ := i.range
+  result := result ++ (Format.joinSep (← i.goalState ctx) "\n").pretty ++ "\n---\n"
+  let ⟨⟨startLine, startCol⟩, ⟨endLine, endCol⟩⟩ := i.range ctx
   result := result ++ s!"{startLine}:{startCol}-{endLine}:{endCol}\n---\n"
-  result := result ++ (← i.pp).pretty ++ "\n---\n"
-  result := result ++ (Format.joinSep (← i.goalStateAfter) "\n").pretty ++ "\n---\n"
+  result := result ++ (← i.pp ctx).pretty ++ "\n---\n"
+  result := result ++ (Format.joinSep (← i.goalStateAfter ctx) "\n").pretty ++ "\n---\n"
   return result
 
-def proofStepData (i : TacticInvocation) : IO String := do
-  return "[GOAL]\n" ++ (Format.joinSep (← i.goalState) "\n").pretty ++ "\n[PROOFSTEP]\n" ++ (← i.pp).pretty
+def proofStepData (i : TacticInfo) (ctx : ContextInfo) : IO String := do
+  return "[GOAL]\n" ++ (Format.joinSep (← i.goalState ctx) "\n").pretty ++ "\n[PROOFSTEP]\n" ++ (← i.pp ctx).pretty
 
-end Lean.Elab.TacticInvocation
+end Lean.Elab.TacticInfo
 
 def trainingData (args : Cli.Parsed) : IO UInt32 := do
     searchPathRef.set compile_time_search_path%
@@ -40,11 +39,11 @@ def trainingData (args : Cli.Parsed) : IO UInt32 := do
     trees := trees.bind InfoTree.retainOriginal
     trees := trees.bind InfoTree.retainSubstantive
     for t in trees do
-      for t in t.tactics do
+      for (i, ctx) in t.tactics do
         if args.hasFlag "proofstep" then
-          IO.println (← t.proofStepData)
+          IO.println (← i.proofStepData ctx)
         else
-          IO.println (← t.verboseTrainingData module)
+          IO.println (← i.verboseTrainingData module ctx)
     return 0
 
 /-- Setting up command line options and help text for `lake exe training_data`. -/