AlephProver prep helper scripts

LeanUtils/ExtractSorry.lean

@@ -36,14 +36,35 @@ where
 def extractSorries (T : InfoTree) : IO (List <| SorryData Format) :=
   traverseInfoTree ppGoalIfNoMVar T
 
+/-- One record per source token.
+
+A `sorry` *tactic* elaborates to a `sorry` *term*, so the info trees carry two
+nodes for the same token with the same goal.  They differ only in `kind`, which
+would otherwise defeat the plain deduplication; merge them and keep the
+`"tactic"` record, since that is the position the token occupies in the source
+(a replacement there needs no leading `by`). -/
+def dedupByToken (sorries : List ParsedSorry) : List ParsedSorry :=
+  sorries.foldl (init := []) fun acc ps =>
+    let sameToken (q : ParsedSorry) :=
+      q.startPos == ps.startPos && q.endPos == ps.endPos &&
+        q.parentDecl == ps.parentDecl && q.goal == ps.goal
+    match acc.find? sameToken with
+    | none => acc ++ [ps]
+    | some q =>
+      if q.kind != some "tactic" && ps.kind == some "tactic" then
+        acc.map fun r => if sameToken r then ps else r
+      else acc
+
 /-- `parseFile myLeanFile` extracts the sorries contained in the Lean file `myLeanFile`. -/
 def parseFile (path : System.FilePath) : IO (List ParsedSorry) := do
+  unsafe enableInitializersExecution
+  let projectSearchPath ← getProjectSearchPath path
+  searchPathRef.set projectSearchPath
   -- Throw an error if the oleans of the file can't be found...
   path.checkOLeans
   let (fileMap, trees) ← extractInfoTrees path
   -- TODO(Paul-Lez): here ideally we should filter `trees` so we only run
   -- `extractSorries` on infotrees that arise from theorems/lemmas/definitions/...
   let sorryLists  ← trees.mapM extractSorries
   let sorryLists : List ParsedSorry := sorryLists.flatten'.map (SorryData.toParsedSorry fileMap)
-  let sorryLists := sorryLists.dedup'
-  return sorryLists
+  return dedupByToken sorryLists

LeanUtils/TargetEnv.lean

@@ -0,0 +1,184 @@
+import LeanUtils.ExtractSorry
+import Lean.Meta.Basic
+
+/-!
+# Locating one `sorry` in a file's info trees
+
+Given a `ParsedSorry` (position, parent declaration, goal text) this module
+re-elaborates the file and recovers the elaboration state at that token: the
+`ContextInfo`, the local context, and the goal type.  `KernelCheck` uses it to
+check a candidate proof term against that goal, and `ExtractGoal` uses it to
+restate the goal as a standalone theorem.
+
+Two situations need care:
+* a `sorry` tactic also elaborates to a `sorry` term, so one token yields two
+  nodes; they carry the same type and either will do;
+* one token can close several goals (`constructor <;> sorry`, `all_goals sorry`).
+  The recorded goal text then selects the intended one, compared with whitespace
+  collapsed because different tools wrap goals at different widths.
+-/
+
+open Lean Meta Elab Term Expr Meta Tactic
+
+/-- The elaboration state at one `sorry`. -/
+structure TargetEnvData where
+  ctx : ContextInfo
+  lctx : LocalContext
+  type : Expr
+  /-- Start positions of the enclosing commands, outermost first: the
+  declaration's own command and any `… in` wrappers around it. -/
+  commandPositions : List Position
+
+private structure MatchedTarget where
+  ctx : ContextInfo
+  type? : Option Expr
+  goal? : Option MVarId
+  lctx : LocalContext
+  commandPositions : List Position := []
+
+/-- Collapse all whitespace, so goals that differ only in line wrapping compare equal.
+
+The dataset's goal text was rendered by SorryDB's REPL at a different line width than
+`ExtractSorry` uses here; the content is identical, the line breaks are not. -/
+def goalConclusion (s : String) : String :=
+  -- everything after the LAST turnstile; resets the accumulator at each one
+  s.foldl (fun acc c => if c == '⊢' then "" else acc.push c) ""
+
+def normalizeGoalText (s : String) : String :=
+  -- character fold rather than `String.split`: the latter returns `List String` on
+  -- some toolchains and `Std.Iter String.Slice` on others, which breaks the build
+  -- for whichever Lean versions the fleet does not happen to be tested against.
+  (s.foldl (fun acc c =>
+      if c.isWhitespace then
+        if acc.endsWith " " then acc else acc.push ' '
+      else acc.push c) "").trim
+
+def findTargetEnv (tree : InfoTree) (targetSorry : ParsedSorry) : IO (List TargetEnvData) := do
+  let matched ← tree.visitM (m := IO) (postNode := fun ctx info _ children => do
+    let targets : List MatchedTarget := (children.flatMap' Option.toList).flatten'
+    match info with
+    | .ofTermInfo ti =>
+        if !targetSorry.acceptsKind "term" || !isSorryTerm ti.stx then
+          return targets
+        let some pos := ti.stx.getPos? | return targets
+        if targetSorry.startPos != ctx.fileMap.toPosition pos then
+          return targets
+        let some type := ti.expectedType? | return targets
+        return targets ++ [{
+          ctx
+          type? := some type
+          goal? := none
+          lctx := ti.lctx
+        }]
+    | .ofTacticInfo ti =>
+        if !targetSorry.acceptsKind "tactic" || !isSorryTactic ti.stx then
+          return targets
+        let some pos := ti.stx.getPos? | return targets
+        if targetSorry.startPos != ctx.fileMap.toPosition pos then
+          return targets
+        let goal ← match ti.goalsBefore with
+          | [goal] => pure goal
+          | goals => do
+            -- A single `sorry` token can close several goals (`constructor <;> sorry`).
+            -- SorryDB records one entry per goal, all sharing this source position, so
+            -- the entry's printed goal is what distinguishes them.  Select the matching
+            -- one instead of refusing; still refuse if it is ambiguous.
+            let mut matching : List MVarId := []
+            for candidate in goals do
+              if let some mdecl := ti.mctxBefore.decls.find? candidate then
+                let rendered ← ctx.runMetaM mdecl.lctx do
+                  return toString (← ppGoal candidate)
+                if normalizeGoalText rendered == normalizeGoalText targetSorry.goal then
+                  matching := matching ++ [candidate]
+            match matching with
+            | [goal] => pure goal
+            | _ =>
+              -- report how the selection went, so a rendering mismatch (0 matched)
+              -- can be told apart from a genuinely ambiguous token (>1 matched)
+              let rendered ← goals.mapM fun candidate => do
+                match ti.mctxBefore.decls.find? candidate with
+                | some mdecl => ctx.runMetaM mdecl.lctx do return toString (← ppGoal candidate)
+                | none => pure "<no mvar decl>"
+              throw (IO.userError s!"Found more than one goal ({goals.length} goals, \
+                {matching.length} matched target); target={(normalizeGoalText targetSorry.goal).take 300}; \
+                candidates={(rendered.map fun r => (normalizeGoalText r).take 300)}")
+        let some mdecl := ti.mctxBefore.decls.find? goal
+          | throw (IO.userError "Could not recover the target goal's local context")
+        return targets ++ [{
+          ctx
+          type? := none
+          goal? := some goal
+          lctx := mdecl.lctx
+        }]
+    | .ofCommandInfo ci =>
+        let some pos := ci.stx.getPos? | return targets
+        let commandPos := ctx.fileMap.toPosition pos
+        return targets.map fun target =>
+          { target with commandPositions := commandPos :: target.commandPositions }
+    | _ => return targets)
+
+  let matched := matched.get!
+  let targetDatas ← matched.mapM fun target => do
+    target.ctx.runMetaM target.lctx do
+      let type ← match target.type?, target.goal? with
+        | some type, none => pure type
+        | none, some goal => goal.getType
+        | _, _ => throwError "Bad case"
+      return [{
+        ctx := target.ctx
+        lctx := target.lctx
+        type
+        commandPositions :=
+          if target.commandPositions.isEmpty then [targetSorry.startPos]
+          else target.commandPositions
+      }]
+
+  return targetDatas.flatten'.filter fun data =>
+    data.ctx.parentDecl? == some targetSorry.parentDecl
+
+
+def findSorryTargetFromFile (path rawSorry : String) : IO (Except String (FileMap × TargetEnvData)) := do
+  unsafe enableInitializersExecution
+  let path : System.FilePath := { toString := path }
+  let path ← IO.FS.realPath path
+  let projectSearchPath ← getProjectSearchPath path
+  searchPathRef.set projectSearchPath
+  let a := Json.parse rawSorry
+  let json ← match a with
+    | .ok json => pure json
+    | .error e => return .error s!"Failed to parse input as valid JSON {e}"
+
+  let parsedSorry : ParsedSorry ← match (Lean.FromJson.fromJson? json) with
+  | .ok parsedSorry => pure parsedSorry
+  | .error e => return .error s!"Failed to deserialize ParsedSorry: {e}"
+
+  let (fileMap, trees) ← extractInfoTrees path
+
+  let targetEnvs ← trees.mapM (fun t => findTargetEnv t parsedSorry)
+
+  let targetEnvs := targetEnvs.flatten'
+  -- We might have both term-mode and tactic-mode info trees for the same source-level 'sorry'
+  -- (since the 'sorry' tactic will end up emitting a 'sorry' term)
+  -- We just pick the first one - as long as they all have the same type (which we check),
+  -- shouldn't matter
+  let some singleData := targetEnvs[0]? | return .error s!"Did not find any targetEnv"
+  if targetEnvs.all (fun d => d.type == singleData.type) then
+    return .ok (fileMap, singleData)
+  -- One source `sorry` can legitimately close SEVERAL goals (`all_goals sorry`,
+  -- `<;> sorry`), so the infotrees genuinely carry different types.  That is not an
+  -- ambiguity to give up on: the recorded goal says which one this task is.  Match on
+  -- the conclusion -- sibling goals share a local context and differ only in the target.
+  let wantConcl := normalizeGoalText (goalConclusion parsedSorry.goal)
+  let mut matching : List TargetEnvData := []
+  let mut rendered : List String := []
+  for d in targetEnvs do
+    let txt ← d.ctx.runMetaM d.lctx do
+      return toString (← ppExpr (← instantiateMVars d.type))
+    rendered := rendered ++ [normalizeGoalText txt]
+    if normalizeGoalText txt == wantConcl then
+      matching := matching ++ [d]
+  match matching with
+  | d :: _ => return .ok (fileMap, d)
+  | [] =>
+    return .error s!"Found different types for infotrees corresponding to same sorry; \
+      target={wantConcl.take 200}; candidates={rendered.map (·.take 200)}"

LeanUtils/Utils.lean

@@ -42,26 +42,45 @@ def visitSorryNode {Out} (ctx : ContextInfo) (node : Info)
     else return none
   | _ => return none
 
+/-- One `sorry` found in a file.
+
+`startByte`/`endByte` and `kind` are emitted by `ExtractSorry` and consumed by
+tools that rewrite the source (a byte range is what a splice needs; `kind`
+says whether the token is the `sorry` *tactic* or the `sorry` *term*, which
+decides whether a replacement needs a leading `by`).  They are optional on
+input so that a record with only line/column positions -- the shape stored in
+the SorryDB database -- still deserializes. -/
 structure ParsedSorry where
   goal : String
   startPos : Position
   endPos : Position
   parentDecl : Name
   hash : UInt64
+  startByte : Option Nat := none
+  endByte : Option Nat := none
+  /-- `"tactic"` or `"term"`; `none` accepts either. -/
+  kind : Option String := none
 deriving DecidableEq, FromJson
 
+/-- `true` unless `kind` is set and differs from `k`. -/
+def ParsedSorry.acceptsKind (ps : ParsedSorry) (k : String) : Bool :=
+  ps.kind.all (· == k)
+
 instance : ToJson ParsedSorry where
-  toJson ps := Json.mkObj [
-    ("goal", Json.str ps.goal),
-    ("location", Json.mkObj [
-      ("start_line", Json.num ps.startPos.line),
-      ("start_column", Json.num ps.startPos.column),
-      ("end_line", Json.num ps.endPos.line),
-      ("end_column", Json.num ps.endPos.column)
-    ]),
-    ("parentDecl", Json.str ps.parentDecl.toString),
-    ("hash", Json.num ps.hash.toNat)
-  ]
+  toJson ps :=
+    let location := [
+        ("start_line", Json.num ps.startPos.line),
+        ("start_column", Json.num ps.startPos.column),
+        ("end_line", Json.num ps.endPos.line),
+        ("end_column", Json.num ps.endPos.column)
+      ] ++ (ps.startByte.map fun b => ("start_byte", Json.num b)).toList
+        ++ (ps.endByte.map fun b => ("end_byte", Json.num b)).toList
+    Json.mkObj <| [
+      ("goal", Json.str ps.goal),
+      ("location", Json.mkObj location),
+      ("parentDecl", Json.str ps.parentDecl.toString),
+      ("hash", Json.num ps.hash.toNat)
+    ] ++ (ps.kind.map fun k => ("kind", Json.str k)).toList
 
 def SorryData.toParsedSorry {Out} [ToString Out] (fileMap : FileMap) :
     SorryData Out → ParsedSorry :=
@@ -72,6 +91,9 @@ def SorryData.toParsedSorry {Out} [ToString Out] (fileMap : FileMap) :
       endPos := fileMap.toPosition stx.getTailPos?.get!
       parentDecl
       hash := Hashable.hash <| ToString.toString out
+      startByte := some stx.getPos?.get!.byteIdx
+      endByte := some stx.getTailPos?.get!.byteIdx
+      kind := some (if isSorryTactic stx then "tactic" else "term")
     }
 
 instance : ToString ParsedSorry where
@@ -131,7 +153,17 @@ partial def getAllLakePaths (path : System.FilePath) : IO (Array System.FilePath
   unless ← path.pathExists do return #[]
   let dirEntries := (← path.readDir).map IO.FS.DirEntry.path
   if dirEntries.contains (path / ".lake") then
-    return (← getAllLakePaths <| path / ".lake/packages").push (path / ".lake/build/lib/lean")
+    -- A built package.  Recurse into its own dependencies, and ALSO into any
+    -- sub-packages sitting directly inside it: one git dependency can ship
+    -- several packages side by side (e.g. `packages/Hammer/HammerCore`), which
+    -- lake puts on LEAN_PATH but which this short-circuit would otherwise skip.
+    let nested ← getAllLakePaths <| path / ".lake/packages"
+    let subPkgs ← dirEntries.filterM fun entry => do
+      if entry == path / ".lake" then return false
+      if !(← entry.isDir) then return false
+      (entry / ".lake").pathExists
+    let fromSubPkgs ← subPkgs.mapM getAllLakePaths
+    return (nested ++ fromSubPkgs.flatten).push (path / ".lake/build/lib/lean")
   else
     let dirEntries ← dirEntries.filterM fun path ↦ path.isDir
     return (← dirEntries.mapM getAllLakePaths).flatten
@@ -144,7 +176,13 @@ def getProjectSearchPath (path : System.FilePath) : IO (System.SearchPath) := do
   let rootDir ← getProjectRootDirPath path
   let paths ← getAllLakePaths rootDir
   let originalSearchPath ← getBuiltinSearchPath (← findSysroot)
-  return originalSearchPath.append paths.toList
+  -- Honour LEAN_PATH when it is set: `lake env` derives it from the manifest,
+  -- which is authoritative for layouts a directory walk cannot infer.
+  let envPaths : List System.FilePath ← do
+    match ← IO.getEnv "LEAN_PATH" with
+    | some raw => pure (System.SearchPath.parse raw)
+    | none => pure []
+  return originalSearchPath.append (paths.toList ++ envPaths)
 
 def System.FilePath.checkOLeans (path : System.FilePath) : IO Unit := do
   discard <| Lean.findOLean (← moduleNameOfFileName path none)

LeanUtilsTest/LeanFileWithMultiGoalSorry.lean

@@ -0,0 +1,5 @@
+import Lean
+
+-- One `sorry` token closes two goals; SorryDB records one entry per goal.
+theorem both : (1 : Nat) = 1 ∧ (2 : Nat) = 2 := by
+  constructor <;> sorry

LeanUtilsTest/TestExtractGoal.lean

@@ -0,0 +1,52 @@
+import bins.ExtractGoal
+
+/-! `ExtractGoal <file> <ParsedSorry json> [flags]` restates the goal at one `sorry`
+as a standalone theorem, and prints the term that closes the original goal with
+it. Only positions, parent declaration, goal and hash are required in the record
+(hash as a decimal string). The payload has three parts, separated by markers:
+the source prefix up to the enclosing command, the helper theorem, and the
+application. -/
+
+-- A tactic `sorry` with one hypothesis in scope.
+/--
+info: {"ok":
+ "import Lean\n\n\n-- sorrydb-helper-start\ntheorem mytheorem (someLemma : True) : 1 + 1 = 2 := sorry\n-- sorrydb-application: mytheorem someLemma"}
+---
+info: 0
+-/
+#guard_msgs in
+#eval main ["LeanUtilsTest/LeanFileWithSorries.lean",
+  "{\"goal\": \"someLemma : True\\n⊢ 1 + 1 = 2\", \"startPos\": {\"line\": 6, \"column\": 2}, \"endPos\": {\"line\": 6, \"column\": 7}, \"parentDecl\": \"test\", \"hash\": \"3234805056349482567\"}"]
+
+-- A term-mode `sorry`: the prefix now contains the whole preceding declaration.
+/--
+info: {"ok":
+ "import Lean\n\ntheorem test : 1 + 1 = 2 := by\n  have someLemma : True := by\n    sorry\n  sorry\n\n\n\n-- sorrydb-helper-start\ntheorem mytheorem : 1 + 1 = 2 := sorry\n-- sorrydb-application: mytheorem"}
+---
+info: 0
+-/
+#guard_msgs in
+#eval main ["LeanUtilsTest/LeanFileWithSorries.lean",
+  "{\"goal\": \"⊢ 1 + 1 = 2\", \"startPos\": {\"line\": 10, \"column\": 2}, \"endPos\": {\"line\": 10, \"column\": 7}, \"parentDecl\": \"test'\", \"hash\": \"11128604812966687648\"}"]
+
+-- One token, two goals: the recorded goal text selects which one is meant.
+/--
+info: {"ok":
+ "import Lean\n\n-- One `sorry` token closes two goals; SorryDB records one entry per goal.\n\n-- sorrydb-helper-start\ntheorem mytheorem : 2 = 2 := sorry\n-- sorrydb-application: mytheorem"}
+---
+info: 0
+-/
+#guard_msgs in
+#eval main ["LeanUtilsTest/LeanFileWithMultiGoalSorry.lean",
+  "{\"goal\": \"⊢ 2 = 2\", \"startPos\": {\"line\": 5, \"column\": 18}, \"endPos\": {\"line\": 5, \"column\": 23}, \"parentDecl\": \"both\", \"hash\": \"0\", \"kind\": \"tactic\"}"]
+
+-- ... and a goal text matching neither is refused, listing the candidates.
+/--
+info: {"error":
+ "Found different types for infotrees corresponding to same sorry; target=3 = 3; candidates=[1 = 1, 2 = 2]"}
+---
+info: 0
+-/
+#guard_msgs in
+#eval main ["LeanUtilsTest/LeanFileWithMultiGoalSorry.lean",
+  "{\"goal\": \"⊢ 3 = 3\", \"startPos\": {\"line\": 5, \"column\": 18}, \"endPos\": {\"line\": 5, \"column\": 23}, \"parentDecl\": \"both\", \"hash\": \"0\", \"kind\": \"tactic\"}"]