feat: API to avoid deadlocks from dropped promises (#6958)

This PR improves the `Promise` API by considering how dropped promises
can lead to never-finished tasks.

Author: Kha

Diff for: src/Init/System/IO.lean

@@ -238,7 +238,12 @@ protected def TaskState.toString : TaskState → String
 
 instance : ToString TaskState := ⟨TaskState.toString⟩
 
-/-- Returns current state of the `Task` in the Lean runtime's task manager. -/
+/--
+Returns current state of the `Task` in the Lean runtime's task manager.
+
+Note that for tasks derived from `Promise`s, `waiting` and `running` should be considered
+equivalent.
+-/
 @[extern "lean_io_get_task_state"] opaque getTaskState : @& Task α → BaseIO TaskState
 
 /-- Check if the task has finished execution, at which point calling `Task.get` will return immediately. -/

Diff for: src/Init/System/Promise.lean

@@ -21,14 +21,13 @@ private structure PromiseImpl (α : Type) : Type where
 
 Typical usage is as follows:
 1. `let promise ← Promise.new` creates a promise
-2. `promise.result : Task α` can now be passed around
-3. `promise.result.get` blocks until the promise is resolved
+2. `promise.result? : Task (Option α)` can now be passed around
+3. `promise.result?.get` blocks until the promise is resolved
 4. `promise.resolve a` resolves the promise
-5. `promise.result.get` now returns `a`
+5. `promise.result?.get` now returns `some a`
 
-Every promise must eventually be resolved.
-Otherwise the memory used for the promise will be leaked,
-and any tasks depending on the promise's result will wait forever.
+If the promise is dropped without ever being resolved, `promise.result?.get` will return `none`.
+See `Promise.result!/resultD` for other ways to handle this case.
 -/
 def Promise (α : Type) : Type := PromiseImpl α
 
@@ -47,12 +46,32 @@ Only the first call to this function has an effect.
 @[extern "lean_io_promise_resolve"]
 opaque Promise.resolve (value : α) (promise : @& Promise α) : BaseIO Unit
 
+/--
+Like `Promise.result`, but resolves to `none` if the promise is dropped without ever being resolved.
+-/
+@[extern "lean_io_promise_result_opt"]
+opaque Promise.result? (promise : @& Promise α) : Task (Option α)
+
+-- SU: not planning to make this public without a lot more thought and motivation
+@[extern "lean_option_get_or_block"]
+private opaque Option.getOrBlock! [Nonempty α] : Option α → α
+
 /--
 The result task of a `Promise`.
 
-The task blocks until `Promise.resolve` is called.
+The task blocks until `Promise.resolve` is called. If the promise is dropped without ever being
+resolved, evaluating the task will panic and, when not using fatal panics, block forever. Use
+`Promise.result?` to handle this case explicitly.
+-/
+def Promise.result! (promise : @& Promise α) : Task α :=
+  let _ : Nonempty α := promise.h
+  promise.result?.map (sync := true) Option.getOrBlock!
+
+@[inherit_doc Promise.result!, deprecated Promise.result! (since := "2025-02-05")]
+def Promise.result := @Promise.result!
+
+/--
+Like `Promise.result`, but resolves to `dflt` if the promise is dropped without ever being resolved.
 -/
-@[extern "lean_io_promise_result"]
-opaque Promise.result (promise : Promise α) : Task α :=
-  have : Nonempty α := promise.h
-  Classical.choice inferInstance
+def Promise.resultD (promise : Promise α) (dflt : α): Task α :=
+  promise.result?.map (sync := true) (·.getD dflt)

Diff for: src/Lean/AddDecl.lean

@@ -77,12 +77,9 @@ def addDecl (decl : Declaration) : CoreM Unit := do
   async.commitConst async.asyncEnv (some info)
   setEnv async.mainEnv
   let checkAct ← Core.wrapAsyncAsSnapshot fun _ => do
-    try
-      setEnv async.asyncEnv
-      doAdd
-      async.commitCheckEnv (← getEnv)
-    finally
-      async.commitFailure
+    setEnv async.asyncEnv
+    doAdd
+    async.commitCheckEnv (← getEnv)
   let t ← BaseIO.mapTask (fun _ => checkAct) env.checked
   let endRange? := (← getRef).getTailPos?.map fun pos => ⟨pos, pos⟩
   Core.logSnapshotTask { range? := endRange?, task := t }

Diff for: src/Lean/Elab/Command.lean

@@ -489,38 +489,38 @@ partial def elabCommand (stx : Syntax) : CommandElabM Unit := do
                   return oldSnap
                 let oldCmds? := oldSnap?.map fun old =>
                   if old.newStx.isOfKind nullKind then old.newStx.getArgs else #[old.newStx]
-                Language.withAlwaysResolvedPromises cmds.size fun cmdPromises => do
-                  snap.new.resolve <| .ofTyped {
-                    diagnostics := .empty
-                    macroDecl := decl
-                    newStx := stxNew
-                    newNextMacroScope := nextMacroScope
-                    hasTraces
-                    next := Array.zipWith (fun cmdPromise cmd =>
-                      { range? := cmd.getRange?, task := cmdPromise.result }) cmdPromises cmds
-                    : MacroExpandedSnapshot
-                  }
-                  -- After the first command whose syntax tree changed, we must disable
-                  -- incremental reuse
-                  let mut reusedCmds := true
-                  let opts ← getOptions
-                  -- For each command, associate it with new promise and old snapshot, if any, and
-                  -- elaborate recursively
-                  for cmd in cmds, cmdPromise in cmdPromises, i in [0:cmds.size] do
-                    let oldCmd? := oldCmds?.bind (·[i]?)
-                    withReader ({ · with snap? := some {
-                      new := cmdPromise
-                      old? := do
-                        guard reusedCmds
-                        let old ← oldSnap?
-                        return { stx := (← oldCmd?), val := (← old.next[i]?) }
-                    } }) do
-                      elabCommand cmd
-                      -- Resolve promise for commands not supporting incrementality; waiting for
-                      -- `withAlwaysResolvedPromises` to do this could block reporting by later
-                      -- commands
-                      cmdPromise.resolve default
-                    reusedCmds := reusedCmds && oldCmd?.any (·.eqWithInfoAndTraceReuse opts cmd)
+                let cmdPromises ← cmds.mapM fun _ => IO.Promise.new
+                snap.new.resolve <| .ofTyped {
+                  diagnostics := .empty
+                  macroDecl := decl
+                  newStx := stxNew
+                  newNextMacroScope := nextMacroScope
+                  hasTraces
+                  next := Array.zipWith (fun cmdPromise cmd =>
+                    { range? := cmd.getRange?, task := cmdPromise.resultD default }) cmdPromises cmds
+                  : MacroExpandedSnapshot
+                }
+                -- After the first command whose syntax tree changed, we must disable
+                -- incremental reuse
+                let mut reusedCmds := true
+                let opts ← getOptions
+                -- For each command, associate it with new promise and old snapshot, if any, and
+                -- elaborate recursively
+                for cmd in cmds, cmdPromise in cmdPromises, i in [0:cmds.size] do
+                  let oldCmd? := oldCmds?.bind (·[i]?)
+                  withReader ({ · with snap? := some {
+                    new := cmdPromise
+                    old? := do
+                      guard reusedCmds
+                      let old ← oldSnap?
+                      return { stx := (← oldCmd?), val := (← old.next[i]?) }
+                  } }) do
+                    elabCommand cmd
+                    -- Resolve promise for commands not supporting incrementality; waiting for
+                    -- `withAlwaysResolvedPromises` to do this could block reporting by later
+                    -- commands
+                    cmdPromise.resolve default
+                  reusedCmds := reusedCmds && oldCmd?.any (·.eqWithInfoAndTraceReuse opts cmd)
               else
                 elabCommand stxNew
         | _ =>

Diff for: src/Lean/Elab/MutualDef.lean

@@ -249,7 +249,7 @@ where
   mkBodyTask (body : Syntax) (new : IO.Promise (Option BodyProcessedSnapshot)) :
       Language.SnapshotTask (Option BodyProcessedSnapshot) :=
     let rangeStx := getBodyTerm? body |>.getD body
-    { range? := rangeStx.getRange?, task := new.result }
+    { range? := rangeStx.getRange?, task := new.resultD default }
 
   /--
   If `body` allows for incremental tactic reporting and reuse, creates a snapshot task out of the
@@ -261,7 +261,7 @@ where
    := do
     if let some e := getBodyTerm? body then
       if let `(by $tacs*) := e then
-        return (e, some { range? := mkNullNode tacs |>.getRange?, task := tacPromise.result })
+        return (e, some { range? := mkNullNode tacs |>.getRange?, task := tacPromise.resultD default })
     tacPromise.resolve default
     return (none, none)
 
@@ -1005,45 +1005,45 @@ def elabMutualDef (vars : Array Expr) (sc : Command.Scope) (views : Array DefVie
   else
     go
 where
-  go :=
-    withAlwaysResolvedPromises views.size fun bodyPromises =>
-    withAlwaysResolvedPromises views.size fun tacPromises => do
-      let scopeLevelNames ← getLevelNames
-      let headers ← elabHeaders views bodyPromises tacPromises
-      let headers ← levelMVarToParamHeaders views headers
-      let allUserLevelNames := getAllUserLevelNames headers
-      withFunLocalDecls headers fun funFVars => do
-        for view in views, funFVar in funFVars do
-          addLocalVarInfo view.declId funFVar
-        let values ←
-          try
-            let values ← elabFunValues headers vars sc
-            Term.synthesizeSyntheticMVarsNoPostponing
-            values.mapM (instantiateMVarsProfiling ·)
-          catch ex =>
-            logException ex
-            headers.mapM fun header => withRef header.declId <| mkLabeledSorry header.type (synthetic := true) (unique := true)
-        let headers ← headers.mapM instantiateMVarsAtHeader
-        let letRecsToLift ← getLetRecsToLift
-        let letRecsToLift ← letRecsToLift.mapM instantiateMVarsAtLetRecToLift
-        checkLetRecsToLiftTypes funFVars letRecsToLift
-        (if headers.all (·.kind.isTheorem) && !deprecated.oldSectionVars.get (← getOptions) then withHeaderSecVars vars sc headers else withUsed vars headers values letRecsToLift) fun vars => do
-          let preDefs ← MutualClosure.main vars headers funFVars values letRecsToLift
-          checkAllDeclNamesDistinct preDefs
-          for preDef in preDefs do
-            trace[Elab.definition] "{preDef.declName} : {preDef.type} :=\n{preDef.value}"
-          let preDefs ← withLevelNames allUserLevelNames <| levelMVarToParamTypesPreDecls preDefs
-          let preDefs ← instantiateMVarsAtPreDecls preDefs
-          let preDefs ← shareCommonPreDefs preDefs
-          let preDefs ← fixLevelParams preDefs scopeLevelNames allUserLevelNames
-          for preDef in preDefs do
-            trace[Elab.definition] "after eraseAuxDiscr, {preDef.declName} : {preDef.type} :=\n{preDef.value}"
-          addPreDefinitions preDefs
-          processDeriving headers
-      for view in views, header in headers do
-        -- NOTE: this should be the full `ref`, and thus needs to be done after any snapshotting
-        -- that depends only on a part of the ref
-        addDeclarationRangesForBuiltin header.declName view.modifiers.stx view.ref
+  go := do
+    let bodyPromises ← views.mapM fun _ => IO.Promise.new
+    let tacPromises ← views.mapM fun _ => IO.Promise.new
+    let scopeLevelNames ← getLevelNames
+    let headers ← elabHeaders views bodyPromises tacPromises
+    let headers ← levelMVarToParamHeaders views headers
+    let allUserLevelNames := getAllUserLevelNames headers
+    withFunLocalDecls headers fun funFVars => do
+      for view in views, funFVar in funFVars do
+        addLocalVarInfo view.declId funFVar
+      let values ←
+        try
+          let values ← elabFunValues headers vars sc
+          Term.synthesizeSyntheticMVarsNoPostponing
+          values.mapM (instantiateMVarsProfiling ·)
+        catch ex =>
+          logException ex
+          headers.mapM fun header => withRef header.declId <| mkLabeledSorry header.type (synthetic := true) (unique := true)
+      let headers ← headers.mapM instantiateMVarsAtHeader
+      let letRecsToLift ← getLetRecsToLift
+      let letRecsToLift ← letRecsToLift.mapM instantiateMVarsAtLetRecToLift
+      checkLetRecsToLiftTypes funFVars letRecsToLift
+      (if headers.all (·.kind.isTheorem) && !deprecated.oldSectionVars.get (← getOptions) then withHeaderSecVars vars sc headers else withUsed vars headers values letRecsToLift) fun vars => do
+        let preDefs ← MutualClosure.main vars headers funFVars values letRecsToLift
+        checkAllDeclNamesDistinct preDefs
+        for preDef in preDefs do
+          trace[Elab.definition] "{preDef.declName} : {preDef.type} :=\n{preDef.value}"
+        let preDefs ← withLevelNames allUserLevelNames <| levelMVarToParamTypesPreDecls preDefs
+        let preDefs ← instantiateMVarsAtPreDecls preDefs
+        let preDefs ← shareCommonPreDefs preDefs
+        let preDefs ← fixLevelParams preDefs scopeLevelNames allUserLevelNames
+        for preDef in preDefs do
+          trace[Elab.definition] "after eraseAuxDiscr, {preDef.declName} : {preDef.type} :=\n{preDef.value}"
+        addPreDefinitions preDefs
+        processDeriving headers
+    for view in views, header in headers do
+      -- NOTE: this should be the full `ref`, and thus needs to be done after any snapshotting
+      -- that depends only on a part of the ref
+      addDeclarationRangesForBuiltin header.declName view.modifiers.stx view.ref
 
 
   processDeriving (headers : Array DefViewElabHeader) := do
@@ -1060,46 +1060,46 @@ namespace Command
 
 def elabMutualDef (ds : Array Syntax) : CommandElabM Unit := do
   let opts ← getOptions
-  withAlwaysResolvedPromises ds.size fun headerPromises => do
-    let snap? := (← read).snap?
-    let mut views := #[]
-    let mut defs := #[]
-    let mut reusedAllHeaders := true
-    for h : i in [0:ds.size], headerPromise in headerPromises do
-      let d := ds[i]
-      let modifiers ← elabModifiers ⟨d[0]⟩
-      if ds.size > 1 && modifiers.isNonrec then
-        throwErrorAt d "invalid use of 'nonrec' modifier in 'mutual' block"
-      let mut view ← mkDefView modifiers d[1]
-      let fullHeaderRef := mkNullNode #[d[0], view.headerRef]
-      if let some snap := snap? then
-        view := { view with headerSnap? := some {
-          old? := do
-            -- transitioning from `Context.snap?` to `DefView.headerSnap?` invariant: if the
-            -- elaboration context and state are unchanged, and the syntax of this as well as all
-            -- previous headers is unchanged, then the elaboration result for this header (which
-            -- includes state from elaboration of previous headers!) should be unchanged.
-            guard reusedAllHeaders
-            let old ← snap.old?
-            -- blocking wait, `HeadersParsedSnapshot` (and hopefully others) should be quick
-            let old ← old.val.get.toTyped? DefsParsedSnapshot
-            let oldParsed ← old.defs[i]?
-            guard <| fullHeaderRef.eqWithInfoAndTraceReuse opts oldParsed.fullHeaderRef
-            -- no syntax guard to store, we already did the necessary checks
-            return ⟨.missing, oldParsed.headerProcessedSnap⟩
-          new := headerPromise
-        } }
-        defs := defs.push {
-          fullHeaderRef
-          headerProcessedSnap := { range? := d.getRange?, task := headerPromise.result }
-        }
-        reusedAllHeaders := reusedAllHeaders && view.headerSnap?.any (·.old?.isSome)
-      views := views.push view
+  let headerPromises ← ds.mapM fun _ => IO.Promise.new
+  let snap? := (← read).snap?
+  let mut views := #[]
+  let mut defs := #[]
+  let mut reusedAllHeaders := true
+  for h : i in [0:ds.size], headerPromise in headerPromises do
+    let d := ds[i]
+    let modifiers ← elabModifiers ⟨d[0]⟩
+    if ds.size > 1 && modifiers.isNonrec then
+      throwErrorAt d "invalid use of 'nonrec' modifier in 'mutual' block"
+    let mut view ← mkDefView modifiers d[1]
+    let fullHeaderRef := mkNullNode #[d[0], view.headerRef]
     if let some snap := snap? then
-      -- no non-fatal diagnostics at this point
-      snap.new.resolve <| .ofTyped { defs, diagnostics := .empty : DefsParsedSnapshot }
-    let sc ← getScope
-    runTermElabM fun vars => Term.elabMutualDef vars sc views
+      view := { view with headerSnap? := some {
+        old? := do
+          -- transitioning from `Context.snap?` to `DefView.headerSnap?` invariant: if the
+          -- elaboration context and state are unchanged, and the syntax of this as well as all
+          -- previous headers is unchanged, then the elaboration result for this header (which
+          -- includes state from elaboration of previous headers!) should be unchanged.
+          guard reusedAllHeaders
+          let old ← snap.old?
+          -- blocking wait, `HeadersParsedSnapshot` (and hopefully others) should be quick
+          let old ← old.val.get.toTyped? DefsParsedSnapshot
+          let oldParsed ← old.defs[i]?
+          guard <| fullHeaderRef.eqWithInfoAndTraceReuse opts oldParsed.fullHeaderRef
+          -- no syntax guard to store, we already did the necessary checks
+          return ⟨.missing, oldParsed.headerProcessedSnap⟩
+        new := headerPromise
+      } }
+      defs := defs.push {
+        fullHeaderRef
+        headerProcessedSnap := { range? := d.getRange?, task := headerPromise.resultD default }
+      }
+      reusedAllHeaders := reusedAllHeaders && view.headerSnap?.any (·.old?.isSome)
+    views := views.push view
+  if let some snap := snap? then
+    -- no non-fatal diagnostics at this point
+    snap.new.resolve <| .ofTyped { defs, diagnostics := .empty : DefsParsedSnapshot }
+  let sc ← getScope
+  runTermElabM fun vars => Term.elabMutualDef vars sc views
 
 builtin_initialize
   registerTraceClass `Elab.definition.mkClosure

Diff for: src/Lean/Elab/Tactic/Basic.lean

@@ -224,26 +224,26 @@ where
                     guard <| state.term.meta.core.traceState.traces.size == 0
                     guard <| traceState.traces.size == 0
                     return old.val.get
-                  Language.withAlwaysResolvedPromise fun promise => do
-                    -- Store new unfolding in the snapshot tree
-                    snap.new.resolve {
-                      stx := stx'
+                  let promise ← IO.Promise.new
+                  -- Store new unfolding in the snapshot tree
+                  snap.new.resolve {
+                    stx := stx'
+                    diagnostics := .empty
+                    inner? := none
+                    finished := .pure {
                       diagnostics := .empty
-                      inner? := none
-                      finished := .pure {
-                        diagnostics := .empty
-                        state? := (← Tactic.saveState)
-                      }
-                      next := #[{ range? := stx'.getRange?, task := promise.result }]
+                      state? := (← Tactic.saveState)
                     }
-                    -- Update `tacSnap?` to old unfolding
-                    withTheReader Term.Context ({ · with tacSnap? := some {
-                      new := promise
-                      old? := do
-                        let old ← old?
-                        return ⟨old.stx, (← old.next.get? 0)⟩
-                    } }) do
-                      evalTactic stx'
+                    next := #[{ range? := stx'.getRange?, task := promise.resultD default }]
+                  }
+                  -- Update `tacSnap?` to old unfolding
+                  withTheReader Term.Context ({ · with tacSnap? := some {
+                    new := promise
+                    old? := do
+                      let old ← old?
+                      return ⟨old.stx, (← old.next.get? 0)⟩
+                  } }) do
+                    evalTactic stx'
                   return
               evalTactic stx'
         catch ex => handleEx s failures ex (expandEval s ms evalFns)