feat: activate scopes when unpickling

README.md

@@ -139,3 +139,5 @@ can find needed imports.
 ## Future work
 
 * Replay tactic scripts from tactic mode back into the original `sorry`.
+* Currently if you create scoped environment extensions (e.g. scoped notations) in a session
+  these are not correctly pickled and unpickled in later sessions.

REPL/Snapshots.lean

@@ -8,6 +8,22 @@ import REPL.Util.Pickle
 
 open Lean Elab
 
+namespace Lean.Elab.Command
+
+@[inline] def CommandElabM.run (x : CommandElabM α) (ctx : Context) (s : State) : EIO Exception (α × State) :=
+  (x ctx).run s
+
+@[inline] def CommandElabM.run' (x : CommandElabM α) (ctx : Context) (s : State) : EIO Exception α :=
+  Prod.fst <$> x.run ctx s
+
+@[inline] def CommandElabM.toIO (x : CommandElabM α) (ctx : Context) (s : State) : IO (α × State) := do
+  match (← (x.run ctx s).toIO') with
+  | Except.error (Exception.error _ msg)   => throw <| IO.userError (← msg.toString)
+  | Except.error (Exception.internal id _) => throw <| IO.userError <| "internal exception #" ++ toString id.idx
+  | Except.ok a                            => return a
+
+end Lean.Elab.Command
+
 namespace REPL
 
 /--
@@ -20,7 +36,8 @@ structure CommandSnapshot where
 
 namespace CommandSnapshot
 
-open Lean.Elab.Command in
+open Lean.Elab.Command
+
 /-- A copy of `Command.State` with the `Environment`, caches, and logging omitted. -/
 structure CompactableCommandSnapshot where
   -- env         : Environment
@@ -35,6 +52,15 @@ structure CompactableCommandSnapshot where
 
 open System (FilePath)
 
+/--
+Run a `CommandElabM` monadic function in the current `ProofSnapshot`,
+updating the `Command.State`.
+-/
+def runCommandElabM (p : CommandSnapshot) (t : CommandElabM α) : IO (α × CommandSnapshot) := do
+  let (a, cmdState) ← (CommandElabM.toIO · p.cmdContext p.cmdState) do t
+  return (a, { p with cmdState })
+
+
 /--
 Pickle a `CommandSnapshot`, discarding closures and non-essential caches.
 
@@ -57,10 +83,14 @@ def unpickle (path : FilePath) : IO (CommandSnapshot × CompactedRegion) := unsa
     _root_.unpickle (Array Import × PHashMap Name ConstantInfo × CompactableCommandSnapshot ×
       Command.Context) path
   let env ← (← importModules imports {} 0).replay (HashMap.ofList map₂.toList)
-  let p' :=
+  let p' : CommandSnapshot :=
   { cmdState := { cmdState with env }
     cmdContext }
-  return (p', region)
+  let (_, p'') ← p'.runCommandElabM do
+    for o in ← getOpenDecls do
+      if let .simple ns _ := o then do
+        activateScoped ns
+  return (p'', region)
 
 end CommandSnapshot
 
@@ -228,7 +258,7 @@ def unpickle (path : FilePath) : IO (ProofSnapshot × CompactedRegion) := unsafe
       Core.Context × Meta.State × CompactableMetaContext × Term.State × CompactableTermContext ×
       Tactic.State × Tactic.Context) path
   let env ← (← importModules imports {} 0).replay (HashMap.ofList map₂.toList)
-  let p' :=
+  let p' : ProofSnapshot :=
   { coreState := { coreState with env }
     coreContext
     metaState
@@ -237,6 +267,10 @@ def unpickle (path : FilePath) : IO (ProofSnapshot × CompactedRegion) := unsafe
     termContext := { termContext with }
     tacticState
     tacticContext }
-  return (p', region)
+  let (_, p'') ← p'.runCoreM do
+    for o in ← getOpenDecls do
+      if let .simple ns _ := o then
+        activateScoped ns
+  return (p'', region)
 
 end ProofSnapshot

test/Mathlib/test/H20231215_2.expected.out

@@ -2,9 +2,9 @@
 
 {"sorries":
  [{"proofState": 0,
-   "pos": {"line": 2, "column": 29},
-   "goal": "⊢ factorial 9 % 10 = 0",
-   "endPos": {"line": 2, "column": 34}}],
+   "pos": {"line": 2, "column": 20},
+   "goal": "⊢ 9! % 10 = 0",
+   "endPos": {"line": 2, "column": 25}}],
  "messages":
  [{"severity": "warning",
    "pos": {"line": 1, "column": 8},

test/Mathlib/test/H20231215_2.in

@@ -1,3 +1,3 @@
 {"unpickleEnvFrom": "test/H20231215.olean"}
 
-{"cmd": "theorem mathd_numbertheory_739 :\n  factorial 9 % 10 = 0 := by sorry", "env": 0}
+{"cmd": "theorem mathd_numbertheory_739 :\n  9! % 10 = 0 := by sorry", "env": 0}

test/pickle_open_scoped.expected.out

@@ -0,0 +1,18 @@
+{"env": 0}
+
+{"env": 1}
+
+{"sorries":
+ [{"proofState": 0,
+   "pos": {"line": 1, "column": 22},
+   "goal": "⊢ ◾",
+   "endPos": {"line": 1, "column": 27}}],
+ "messages":
+ [{"severity": "warning",
+   "pos": {"line": 1, "column": 7},
+   "endPos": {"line": 1, "column": 14},
+   "data": "declaration uses 'sorry'"}],
+ "env": 2}
+
+{"env": 1}
+

test/pickle_open_scoped.in

@@ -1,13 +1,8 @@
-{"cmd": "namespace X"}
+{"cmd": "import Lean"}
 
-{"cmd": "def f (x : Nat) := x + 37", "env": 0}
+{"cmd": "open Lean.Compiler", "env": 0}
 
-{"cmd": "scoped notation:10000 n \"!\" => f n", "env": 1}
+{"cmd": "unsafe example : ◾ := sorry", "env": 1}
 
-{"cmd": "end X", "env": 2}
+{"pickleTo": "test/f.olean", "env": 1}
 
-{"cmd": "open X", "env": 3}
-
-{"cmd": "example : 39 = 2 ! := rfl", "env": 4}
-
-{"pickleTo": "test/f.olean", "env": 4}

test/pickle_open_scoped_2.expected.out

@@ -0,0 +1,14 @@
+{"env": 0}
+
+{"sorries":
+ [{"proofState": 0,
+   "pos": {"line": 1, "column": 22},
+   "goal": "⊢ ◾",
+   "endPos": {"line": 1, "column": 27}}],
+ "messages":
+ [{"severity": "warning",
+   "pos": {"line": 1, "column": 7},
+   "endPos": {"line": 1, "column": 14},
+   "data": "declaration uses 'sorry'"}],
+ "env": 1}
+

test/pickle_open_scoped_2.in

@@ -1,3 +1,3 @@
 {"unpickleEnvFrom": "test/f.olean"}
 
-{"cmd": "example : 39 = 2 ! := rfl", "env": 1}
+{"cmd": "unsafe example : ◾ := sorry", "env": 0}

test/pickle_scoped_notation.in

@@ -0,0 +1,20 @@
+{"cmd": "import Lean"}
+
+{"cmd": "namespace X", "env": 0}
+
+{"cmd": "def f (x : Nat) := x + 37", "env": 1}
+
+{"cmd": "scoped notation:10000 n \"!\" => f n", "env": 2}
+
+{"cmd": "end X", "env": 3}
+
+{"cmd": "open X", "env": 4}
+
+{"cmd": "example : 39 = 2 ! := rfl", "env": 5}
+
+{"cmd": "open Lean.Compiler", "env": 5}
+
+{"cmd": "unsafe example : ◾ := sorry", "env": 7}
+
+{"pickleTo": "test/f.olean", "env": 7}
+

test/pickle_scoped_notation_2.in

@@ -0,0 +1,3 @@
+{"unpickleEnvFrom": "test/f.olean"}
+
+{"cmd": "example : 39 = 2 ! := rfl", "env": 0}