PyPantograph/src/Pantograph/Serial.lean

import Pantograph.Goal
import Pantograph.Environment

import Lean.Environment
import Lean.Replay
import Lean.Util.ShareCommon
import Std.Data.HashMap

open Lean

/-!
Input/Output functions borrowed from REPL

# Pickling and unpickling objects

By abusing `saveModuleData` and `readModuleData` we can pickle and unpickle objects to disk.
-/

namespace Pantograph

/--
Save an object to disk.
If you need to write multiple objects from within a single declaration,
you will need to provide a unique `key` for each.
-/
def pickle {α : Type} (path : System.FilePath) (x : α) (key : Name := by exact decl_name%) : IO Unit :=
  saveModuleData path key (unsafe unsafeCast x)

/--
Load an object from disk.
Note: The returned `CompactedRegion` can be used to free the memory behind the value
of type `α`, using `CompactedRegion.free` (which is only safe once all references to the `α` are
released). Ignoring the `CompactedRegion` results in the data being leaked.
Use `withUnpickle` to call `CompactedRegion.free` automatically.

This function is unsafe because the data being loaded may not actually have type `α`, and this
may cause crashes or other bad behavior.
-/
unsafe def unpickle (α : Type) (path : System.FilePath) : IO (α × CompactedRegion) := do
  let (x, region) ← readModuleData path
  pure (unsafeCast x, region)

/-- Load an object from disk and run some continuation on it, freeing memory afterwards. -/
unsafe def withUnpickle [Monad m] [MonadLiftT IO m] {α β : Type}
    (path : System.FilePath) (f : α → m β) : m β := do
  let (x, region) ← unpickle α path
  let r ← f x
  region.free
  pure r

/--
Pickle an `Environment` to disk relative to a background.

We only store:
* the list of imports
* the new constants from `Environment.constants`
and when unpickling, we build a fresh `Environment` from the imports,
and then add the new constants.
-/
@[export pantograph_env_pickle_m]
def environmentPickle (env : Environment) (path : System.FilePath) (background? : Option Environment := .none)
  : IO Unit :=
  let distilled := distilEnvironment env background?
  pickle path (ShareCommon.shareCommon distilled)

def resurrectEnvironment
  (distilled : DistilledEnvironment)
  (background? : Option Environment := .none)
  : IO Environment := do
  let (imports, constArray) := distilled
  let env ← match background? with
    | .none => importModules imports {} 0 (loadExts := true)
    | .some env =>
      assert! env.imports == imports
      pure env
  env.replay (Std.HashMap.ofList constArray.toList)

/--
Unpickle an `Environment` from disk.

We construct a fresh `Environment` with the relevant imports,
and then replace the new constants.

The `CompactedRegion` must be free'd after using the environment. Otherwise
there could be memory leaks.
-/
@[export pantograph_env_unpickle_m]
def environmentUnpickle (path : System.FilePath) (background? : Option Environment := .none)
  : IO (Environment × CompactedRegion) := unsafe do
  let (distilled, region) ← unpickle (Array Import × ConstArray) path
  return (← resurrectEnvironment distilled background?, region)

/-- `CoreM`'s state, with information irrelevant to pickling masked out -/
structure CompactCoreState where
  -- env             : Environment
  nextMacroScope  : MacroScope     := firstFrontendMacroScope + 1
  ngen            : NameGenerator  := {}
  auxDeclNGen     : DeclNameGenerator := {}
  -- traceState      : TraceState     := {}
  -- cache           : Cache     := {}
  -- messages        : MessageLog     := {}
  -- infoState       : Elab.InfoState := {}

structure CompactGoalState where
  env : DistilledEnvironment

  core : CompactCoreState
  «meta» : Meta.State
  «elab» : Elab.Term.State
  tactic : Elab.Tactic.State
  root : MVarId
  parentMVars : List MVarId
  fragments : FragmentMap

/-- Pickles a goal state by taking its diff relative to a background
environment. This function eliminates all `MessageData` from synthetic
metavariables, because these `MessageData` objects frequently carry closures,
which cannot be pickled. If there is no synthetic metavariable, this would not
cause a difference. -/
@[export pantograph_goal_state_pickle_m]
def goalStatePickle (goalState : GoalState) (path : System.FilePath) (background? : Option Environment := .none)
  : IO Unit :=
  let {
    savedState := {
      term := {
        «meta» := {
          core := {
            env, nextMacroScope, ngen, auxDeclNGen, ..
          },
          «meta»,
        }
        «elab» := «elab»@{
          syntheticMVars, ..
        },
      },
      tactic
    }
    root,
    parentMVars,
    fragments,
  } := goalState
  -- Delete `MessageData`s
  let syntheticMVars : MVarIdMap _ := syntheticMVars.foldl (init := .empty) λ acc key val =>
    let kind := match val.kind with
      | .typeClass _ => .typeClass .none
      | .coe header? expectedType e f? _ => .coe header? expectedType e f? .none
      | k => k
    acc.insert key { val with kind }
  let compacted : CompactGoalState := {
    env := distilEnvironment env background?,

    core := { nextMacroScope, ngen, auxDeclNGen },
    «meta»,
    «elab» := { «elab» with syntheticMVars },
    tactic,

    root,
    parentMVars,
    fragments,
  }
  pickle path (ShareCommon.shareCommon compacted)

/--
Loads a goal state  from disk. The background is the environment which the goal
state was cached relative to.

The `CompactedRegion` must be free'd after using the goal state. Otherwise there
will be memory leaks.
-/
@[export pantograph_goal_state_unpickle_m]
def goalStateUnpickle (path : System.FilePath) (background? : Option Environment := .none)
    : IO (GoalState × CompactedRegion) := unsafe do
  let ({
    env,

    core,
    «meta»,
    «elab»,
    tactic,

    root,
    parentMVars,
    fragments,
  }, region) ← Pantograph.unpickle CompactGoalState path
  let env ← resurrectEnvironment env background?
  let goalState := {
    savedState := {
      term := {
        «meta» := {
          core := {
            core with
            passedHeartbeats := 0,
            env,
          },
          «meta»,
        },
        «elab»,
      },
      tactic,
    },
    root,
    parentMVars,
    fragments,
  }
  return (goalState, region)

end Pantograph