feat: Goal state pickling

This commit is contained in:
Leni Aniva 2024-12-05 14:23:55 -08:00
parent 0f946880ae
commit 105fb7af4b
Signed by: aniva
GPG Key ID: 4D9B1C8D10EA4C50
2 changed files with 118 additions and 7 deletions

View File

@ -2,6 +2,7 @@ import Lean.Environment
import Lean.Replay import Lean.Replay
import Init.System.IOError import Init.System.IOError
import Std.Data.HashMap import Std.Data.HashMap
import Pantograph.Goal
/-! /-!
Input/Output functions Input/Output functions
@ -70,4 +71,92 @@ def environmentUnpickle (path : System.FilePath) : IO (Environment × CompactedR
let env ← importModules imports {} 0 let env ← importModules imports {} 0
return (← env.replay (Std.HashMap.ofList map₂.toList), region) return (← env.replay (Std.HashMap.ofList map₂.toList), region)
open Lean.Core in
structure CompactCoreState where
-- env : Environment
nextMacroScope : MacroScope := firstFrontendMacroScope + 1
ngen : NameGenerator := {}
-- traceState : TraceState := {}
-- cache : Cache := {}
-- messages : MessageLog := {}
-- infoState : Elab.InfoState := {}
@[export pantograph_goal_state_pickle_m]
def goalStatePickle (goalState : GoalState) (path : System.FilePath) : IO Unit :=
let {
savedState := {
term := {
meta := {
core,
meta,
}
«elab»,
},
tactic
}
root,
parentMVar?,
convMVar?,
calcPrevRhs?,
} := goalState
--let env := core.env
Pantograph.pickle path (
({ core with } : CompactCoreState),
meta,
«elab»,
tactic,
root,
parentMVar?,
convMVar?,
calcPrevRhs?,
)
@[export pantograph_goal_state_unpickle_m]
def goalStateUnpickle (path : System.FilePath) (env : Environment)
: IO (GoalState × CompactedRegion) := unsafe do
let ((
compactCore,
meta,
«elab»,
tactic,
root,
parentMVar?,
convMVar?,
calcPrevRhs?,
), region) ← Pantograph.unpickle (
CompactCoreState ×
Meta.State ×
Elab.Term.State ×
Elab.Tactic.State ×
MVarId ×
Option MVarId ×
Option (MVarId × MVarId × List MVarId) ×
Option (MVarId × Expr)
) path
let goalState := {
savedState := {
term := {
meta := {
core := {
compactCore with
passedHeartbeats := 0,
env,
},
meta,
},
«elab»,
},
tactic,
},
root,
parentMVar?,
convMVar?,
calcPrevRhs?,
}
return (goalState, region)
end Pantograph end Pantograph

View File

@ -31,12 +31,12 @@ def runCoreM { α } (state : Core.State) (testCoreM : TestT CoreM α) : TestM (
return (a, state') return (a, state')
def test_environment_pickling : TestM Unit := do def test_environment_pickling : TestM Unit := do
let stateSrc: Core.State := { env := ← getEnv } let coreSrc : Core.State := { env := ← getEnv }
let stateDst: Core.State := { env := ← getEnv } let coreDst : Core.State := { env := ← getEnv }
let name := `mystery let name := `mystery
let envPicklePath ← tempPath let envPicklePath ← tempPath
let ((), _) ← runCoreM stateSrc do let ((), _) ← runCoreM coreSrc do
let type: Expr := .forallE `p (.sort 0) (.forallE `h (.bvar 0) (.bvar 1) .default) .default let type: Expr := .forallE `p (.sort 0) (.forallE `h (.bvar 0) (.bvar 1) .default) .default
let value: Expr := .lam `p (.sort 0) (.lam `h (.bvar 0) (.bvar 0) .default) .default let value: Expr := .lam `p (.sort 0) (.lam `h (.bvar 0) (.bvar 0) .default) .default
let c := Lean.Declaration.defnDecl <| Lean.mkDefinitionValEx let c := Lean.Declaration.defnDecl <| Lean.mkDefinitionValEx
@ -54,7 +54,7 @@ def test_environment_pickling : TestM Unit := do
| .ok env' => pure env' | .ok env' => pure env'
environmentPickle env' envPicklePath environmentPickle env' envPicklePath
let _ ← runCoreM stateDst do let _ ← runCoreM coreDst do
let (env', _) ← environmentUnpickle envPicklePath let (env', _) ← environmentUnpickle envPicklePath
checkTrue s!"Has symbol {name}" (env'.find? name).isSome checkTrue s!"Has symbol {name}" (env'.find? name).isSome
let anotherName := `mystery2 let anotherName := `mystery2
@ -62,9 +62,30 @@ def test_environment_pickling : TestM Unit := do
IO.FS.removeFile envPicklePath IO.FS.removeFile envPicklePath
def test_goal_state_pickling_simple : TestM Unit := do
let coreSrc : Core.State := { env := ← getEnv }
let coreDst : Core.State := { env := ← getEnv }
let statePath ← tempPath
let type: Expr := .forallE `p (.sort 0) (.forallE `h (.bvar 0) (.bvar 1) .default) .default
let stateGenerate : MetaM GoalState := runTermElabMInMeta do
GoalState.create type
let ((), _) ← runCoreM coreSrc do
let state ← stateGenerate.run'
goalStatePickle state statePath
let ((), _) ← runCoreM coreDst do
let (goalState, _) ← goalStateUnpickle statePath (← getEnv)
let metaM : MetaM (List Expr) := do
goalState.goals.mapM λ goal => goalState.withContext goal goal.getType
let types ← metaM.run'
checkTrue "Goals" $ types[0]!.equal type
IO.FS.removeFile statePath
structure Test where structure Test where
name : String name : String
nInstances : Nat
routine: TestM Unit routine: TestM Unit
protected def Test.run (test: Test) (env: Lean.Environment) : IO LSpec.TestSeq := do protected def Test.run (test: Test) (env: Lean.Environment) : IO LSpec.TestSeq := do
@ -76,11 +97,12 @@ protected def Test.run (test: Test) (env: Lean.Environment) : IO LSpec.TestSeq :
match ← ((runTest $ test.routine).run' state).toBaseIO with match ← ((runTest $ test.routine).run' state).toBaseIO with
| .ok e => return e | .ok e => return e
| .error e => | .error e =>
return LSpec.check "Emitted exception" (e.toString == "") return LSpec.check s!"Emitted exception: {e.toString}" (e.toString == "")
def suite (env : Lean.Environment): List (String × IO LSpec.TestSeq) := def suite (env : Lean.Environment): List (String × IO LSpec.TestSeq) :=
let tests: List Test := [ let tests: List Test := [
{ name := "environment_pickling", nInstances := 2, routine := test_environment_pickling }, { name := "environment_pickling", routine := test_environment_pickling, },
{ name := "goal_state_pickling_simple", routine := test_goal_state_pickling_simple, },
] ]
tests.map (fun test => (test.name, test.run env)) tests.map (fun test => (test.name, test.run env))