chore: Incremental elaboration of Verso (first draft)

examples/website/DemoSite/Blog/Conditionals.lean

@@ -3,6 +3,8 @@ import DemoSite.Categories
 open Verso Genre Blog
 open DemoSite
 
+set_option trace.Elab.reuse true
+
 #doc (Post) "Conditional Expressions in Lean" =>
 
 %%%
@@ -11,6 +13,7 @@ date := {year := 2024, month := 1, day := 15}
 categories := [examples, other]
 %%%
 
+
 Finally started blogging!
 This post describes the syntax and semantics of conditional expressions in Lean.
 Here are some examples:
@@ -103,6 +106,12 @@ mutual
 end
 ```
 
+```lean demo
+--- foo
+example := 99
+```
+
+
 Here is a proof with some lambdas and big terms in it, to check highlighting:
 ```lean demo
 def grow : Nat → α → α
@@ -115,6 +124,7 @@ def grow : Nat → α → α
 theorem grow_10_id {α} : grow (α := α) 6 = id := by
   repeat unfold grow
   all_goals sorry
+
 ```
 
 Here is a proof with big terms in the context:

src/verso/Verso/Doc/Concrete.lean

@@ -168,6 +168,48 @@ macro "%docName" moduleName:ident : term =>
 def currentDocName [Monad m] [MonadEnv m] : m Name := do
   pure <| docName <| (← Lean.MonadEnv.getEnv).mainModule
 
+open Language
+
+structure DocElabSnapshotState where
+  docState : DocElabM.State
+  partState : PartElabM.State
+  termState : Term.SavedState
+deriving Nonempty
+
+open Lean Elab Term
+private def getSnapshotState : PartElabM DocElabSnapshotState := do
+  pure ⟨← getThe _, ← getThe _,  ← saveState (m := TermElabM)⟩
+
+structure DocElabSnapshotData extends Language.Snapshot where
+  stx      : Syntax
+  finished : Task DocElabSnapshotState
+deriving Nonempty
+
+inductive DocElabSnapshot where
+  | mk (data : DocElabSnapshotData) (next : Array (SnapshotTask DocElabSnapshot))
+deriving TypeName, Nonempty
+
+partial instance : ToSnapshotTree DocElabSnapshot where
+  toSnapshotTree := go where
+    go := fun ⟨s, next⟩ => ⟨s.toSnapshot, next.map (·.map (sync := true) go)⟩
+
+
+abbrev DocElabSnapshot.data : DocElabSnapshot → DocElabSnapshotData
+  | .mk data _ => data
+
+abbrev DocElabSnapshot.next : DocElabSnapshot → Array (SnapshotTask DocElabSnapshot)
+  | .mk _ next => next
+
+private def exnMessage (exn : Exception) : TermElabM Message := do
+  match exn with
+  | .error ref msg =>
+    let ref    := replaceRef ref (← MonadLog.getRef)
+    let pos    := ref.getPos?.getD 0
+    let endPos := ref.getTailPos?.getD pos
+    let fileMap ← getFileMap
+    let msgData ← addMessageContext msg
+    pure { fileName := (← getFileName), pos := fileMap.toPosition pos, endPos := fileMap.toPosition endPos, data := msgData, severity := .error }
+  | oops@(.internal _ _) => throw oops
 
 elab "#doc" "(" genre:term ")" title:inlineStr "=>" text:completeDocument eof:eoi : command => open Lean Elab Term Command PartElabM DocElabM in do
   findGenreCmd genre
@@ -180,20 +222,73 @@ elab "#doc" "(" genre:term ")" title:inlineStr "=>" text:completeDocument eof:eo
     let ⟨`<low| [~_ ~(titleName@(.node _ _ titleParts))]>⟩ := title
       | dbg_trace "nope {ppSyntax title}" throwUnsupportedSyntax
     let titleString := inlinesToString (← getEnv) titleParts
-    let ((), st, st') ← liftTermElabM <| PartElabM.run {} (.init titleName) <| do
-      let mut errors := #[]
-      setTitle titleString (← liftDocElabM <| titleParts.mapM elabInline)
-      for b in blocks do
-        try partCommand b
-        catch e => errors := errors.push e
-      closePartsUntil 0 endPos
-      for e in errors do
-        match e with
-        | .error stx msg => logErrorAt stx msg
-        | oops@(.internal _ _) => throw oops
-      pure ()
-    let finished := st'.partContext.toPartFrame.close endPos
-    pushInfoLeaf <| .ofCustomInfo {stx := (← getRef) , value := Dynamic.mk finished.toTOC}
-    saveRefs st st'
-    let docName ← mkIdentFrom title <$> currentDocName
-    elabCommand (← `(def $docName : Part $genre := $(← finished.toSyntax genre st'.linkDefs st'.footnoteDefs)))
+    let initState : PartElabM.State := .init titleName
+    let ctxt ← read
+    let some snap := ctxt.snap?
+      | throwErrorAt title "nope dawg"
+    withReader (fun ρ => {ρ with snap? := none}) do
+      let ((nextPromise, snapshotState), st, st') ← liftTermElabM <| PartElabM.run {} initState <| do
+        let mut oldSnap? := snap.old?.bind (·.get.toTyped? (α := DocElabSnapshot))
+
+        let mut nextPromise ← IO.Promise.new
+        snap.new.resolve <| DynamicSnapshot.ofTyped <| DocElabSnapshot.mk {
+            stx := ← `("I like to eat apples and bananas"),
+            finished := .pure <| ← getSnapshotState,
+            diagnostics := .empty
+          } #[{range? := text.raw.getRange?, task := nextPromise.result}]
+
+        setTitle titleString (← liftDocElabM <| titleParts.mapM elabInline)
+        let mut errors : MessageLog := .empty
+        for b in blocks do
+          let mut reuseState := false
+          if let some oldSnap := oldSnap? then
+            if let some next := oldSnap.next[0]? then
+              if next.get.data.stx.structRangeEqWithTraceReuse (← getOptions) b then
+                -- If they match, do nothing and advance the snapshot
+                let ⟨docState, partState, termState⟩ := next.get.data.finished.get
+                set docState
+                set partState
+                termState.restoreFull
+                errors := next.get.data.diagnostics.msgLog
+                oldSnap? := next.get
+                reuseState := true
+          let nextNextPromise ← IO.Promise.new
+          let updatedState ← IO.Promise.new
+          nextPromise.resolve <| DocElabSnapshot.mk {
+              stx := b,
+              finished := updatedState.result,
+              diagnostics := ⟨errors, none⟩
+            } #[{range? := b.getRange?, task := nextNextPromise.result}]
+
+          try
+            if not reuseState then
+              oldSnap? := none
+              try partCommand b
+              catch e => errors := errors.add <| ← exnMessage e
+
+            updatedState.resolve <| ← getSnapshotState
+            nextPromise := nextNextPromise
+          finally
+            -- In case of a fatal exception in partCommand, we need to make sure that the promise actually
+            -- gets resolved to avoid hanging forever. And the first resolve wins, so the second one is a
+            -- no-op the rest of the time.
+            nextPromise.resolve <| DocElabSnapshot.mk {
+                stx := b,
+                finished := .pure <| ← getSnapshotState,
+                diagnostics := ⟨errors, none⟩
+              } #[]
+            updatedState.resolve <| ← getSnapshotState -- some old state goes here
+
+        closePartsUntil 0 endPos
+
+        pure (nextPromise, ← getSnapshotState)
+      let finished := st'.partContext.toPartFrame.close endPos
+      pushInfoLeaf <| .ofCustomInfo {stx := (← getRef) , value := Dynamic.mk finished.toTOC}
+      saveRefs st st'
+      let docName ← mkIdentFrom title <$> currentDocName
+      elabCommand (← `(def $docName : Part $genre := $(← finished.toSyntax genre st'.linkDefs st'.footnoteDefs)))
+      nextPromise.resolve <| DocElabSnapshot.mk {
+        stx := ← `(55.8),
+        finished := .pure snapshotState
+        diagnostics := .empty
+      } #[]

src/verso/Verso/Doc/Elab/Monad.lean

@@ -69,6 +69,7 @@ def DocUses.add (uses : DocUses) (loc : Syntax) : DocUses := {uses with useSites
 structure DocElabM.State where
   linkRefs : HashMap String DocUses := {}
   footnoteRefs : HashMap String DocUses := {}
+deriving Nonempty
 
 /-- Custom info tree data to save footnote and reflink cross-references -/
 structure DocRefInfo where
@@ -182,15 +183,15 @@ structure PartFrame where
   metadata : Option (TSyntax `term)
   blocks : Array (TSyntax `term)
   priorParts : Array FinishedPart
-deriving Repr
+deriving Repr, Nonempty
 
 def PartFrame.close (fr : PartFrame) (endPos : String.Pos) : FinishedPart :=
   let (titlePreview, titleInlines) := fr.expandedTitle.getD ("<anonymous>", #[])
   .mk fr.titleSyntax titleInlines titlePreview fr.metadata fr.blocks fr.priorParts endPos
 
 structure PartContext extends PartFrame where
   parents : Array PartFrame
-deriving Repr
+deriving Repr, Nonempty
 
 def PartContext.level (ctxt : PartContext) : Nat := ctxt.parents.size
 def PartContext.close (ctxt : PartContext) (endPos : String.Pos) : Option PartContext := do
@@ -210,6 +211,7 @@ structure PartElabM.State where
   partContext : PartContext
   linkDefs : HashMap String (DocDef String) := {}
   footnoteDefs : HashMap String (DocDef (Array (TSyntax `term))) := {}
+deriving Nonempty
 
 
 def PartElabM.State.init (title : Syntax) (expandedTitle : Option (String × Array (TSyntax `term)) := none) : PartElabM.State where
@@ -223,6 +225,8 @@ def PartElabM.run (st : DocElabM.State) (st' : PartElabM.State) (act : PartElabM
 
 instance : MonadRef PartElabM := inferInstanceAs <| MonadRef (StateT DocElabM.State (StateT PartElabM.State TermElabM))
 
+instance : MonadAlwaysExcept Exception PartElabM := inferInstanceAs <| MonadAlwaysExcept Exception (StateT DocElabM.State (StateT PartElabM.State TermElabM))
+
 instance : AddErrorMessageContext PartElabM := inferInstanceAs <| AddErrorMessageContext (StateT DocElabM.State (StateT PartElabM.State TermElabM))
 
 instance : MonadQuotation PartElabM := inferInstanceAs <| MonadQuotation (StateT DocElabM.State (StateT PartElabM.State TermElabM))