feat: context

Author: algebraic-dev

src/Std/Sync.lean

@@ -16,5 +16,6 @@ public import Std.Sync.Notify
 public import Std.Sync.Broadcast
 public import Std.Sync.StreamMap
 public import Std.Sync.CancellationToken
+public import Std.Sync.Context
 
 @[expose] public section

src/Std/Sync/CancellationToken.lean

@@ -23,6 +23,27 @@ that a cancellation has occurred.
 namespace Std
 open Std.Internal.IO.Async
 
+/--
+Reasons for cancellation.
+-/
+inductive CancellationReason where
+  /-- Cancelled due to a deadline or timeout -/
+  | deadline
+  /-- Cancelled due to shutdown -/
+  | shutdown
+  /-- Explicitly cancelled -/
+  | cancel
+  /-- Custom cancellation reason -/
+  | custom (msg : String)
+  deriving Repr, BEq
+
+instance : ToString CancellationReason where
+  toString
+    | .deadline => "deadline"
+    | .shutdown => "shutdown"
+    | .cancel => "cancel"
+    | .custom msg => s!"custom(\"{msg}\")"
+
 inductive CancellationToken.Consumer where
   | normal (promise : IO.Promise Unit)
   | select (finished : Waiter Unit)
@@ -44,9 +65,9 @@ The central state structure for a `CancellationToken`.
 -/
 structure CancellationToken.State where
   /--
-  Whether this token has been cancelled.
+  The cancellation reason if cancelled, none otherwise.
   -/
-  cancelled : Bool
+  reason : Option CancellationReason
 
   /--
   Consumers that are blocked waiting for cancellation.
@@ -66,21 +87,21 @@ namespace CancellationToken
 Create a new cancellation token.
 -/
 def new : BaseIO CancellationToken := do
-  return { state := ← Std.Mutex.new { cancelled := false, consumers := ∅ } }
+  return { state := ← Std.Mutex.new { reason := none, consumers := ∅ } }
 
 /--
-Cancel the token, notifying all currently waiting consumers with `true`.
+Cancel the token with the given reason, notifying all currently waiting consumers.
 Once cancelled, the token remains cancelled.
 -/
-def cancel (x : CancellationToken) : BaseIO Unit := do
+def cancel (x : CancellationToken) (reason : CancellationReason := .cancel) : BaseIO Unit := do
   x.state.atomically do
     let mut st ← get
 
-    if st.cancelled then
+    if st.reason.isSome then
       return
 
     let mut remainingConsumers := st.consumers
-    st := { cancelled := true, consumers := ∅ }
+    st := { reason := some reason, consumers := ∅ }
 
     while true do
       if let some (consumer, rest) := remainingConsumers.dequeue? then
@@ -97,7 +118,15 @@ Check if the token is cancelled.
 def isCancelled (x : CancellationToken) : BaseIO Bool := do
   x.state.atomically do
     let st ← get
-    return st.cancelled
+    return st.reason.isSome
+
+/--
+Get the cancellation reason if the token is cancelled.
+-/
+def getCancellationReason (x : CancellationToken) : BaseIO (Option CancellationReason) := do
+  x.state.atomically do
+    let st ← get
+    return st.reason
 
 /--
 Wait for cancellation. Returns a task that completes when cancelled,
@@ -106,7 +135,7 @@ def wait (x : CancellationToken) : IO (AsyncTask Unit) :=
   x.state.atomically do
     let st ← get
 
-    if st.cancelled then
+    if st.reason.isSome then
       return Task.pure (.ok ())
 
     let promise ← IO.Promise.new
@@ -131,7 +160,7 @@ def selector (token : CancellationToken) : Selector Unit := {
     token.state.atomically do
       let st ← get
 
-      if st.cancelled then
+      if st.reason.isSome then
         discard <| waiter.race (return false) (fun promise => do
           promise.resolve (.ok ())
           return true)

src/Std/Sync/Context.lean

@@ -0,0 +1,143 @@
+/-
+Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Sofia Rodrigues
+-/
+
+module
+
+prelude
+public import Std.Data
+public import Init.System.Promise
+public import Init.Data.Queue
+public import Std.Sync.Mutex
+public import Std.Sync.CancellationToken
+public import Std.Internal.Async.Select
+
+public section
+
+/-!
+Context interface for cancellation and deadline management
+-/
+
+namespace Std
+open Std.Internal.IO.Async
+
+/--
+The central state structure shared by all context types.cc
+-/
+structure ContextState where
+  /--
+  Map of token IDs to optional tokens and their children.
+  `none` represents a background context that cannot be cancelled.
+  -/
+  tokens : TreeMap UInt64 (Option CancellationToken × Array UInt64) := .empty
+
+  /--
+  Next available ID
+  -/
+  id : UInt64 := 1
+
+/--
+A cancellation context that allows multiple consumers to wait
+until cancellation is requested. Forms a tree structure where
+cancelling a parent cancels all children.
+-/
+structure Context where
+  state : Std.Mutex ContextState
+  token : CancellationToken
+  id : UInt64
+
+namespace Context
+
+/--
+Create a new root cancellation context.
+-/
+def new : BaseIO Context := do
+  let token ← Std.CancellationToken.new
+  return {
+    state := ← Std.Mutex.new { tokens := .empty |>.insert 0 (some token, #[]) },
+    token,
+    id := 0
+  }
+
+/--
+Fork a child context from a parent. If the parent is already cancelled,
+returns the parent context. Otherwise, creates a new child that will be
+cancelled when the parent is cancelled.
+-/
+def fork (root : Context) : BaseIO Context := do
+  if ← root.token.isCancelled then
+    return root
+
+  let token ← Std.CancellationToken.new
+
+  root.state.atomically do
+    let st ← get
+    let newId := st.id
+    set { st with
+      id := newId + 1,
+      tokens := st.tokens.insert newId (some token, #[])
+        |>.modify root.id (.map (·) (.push · newId))
+    }
+    return { state := root.state, token, id := newId }
+
+/--
+Recursively cancel a context and all its children with the given reason.
+-/
+private partial def cancelChildren (state : ContextState) (id : UInt64) (reason : CancellationReason) : BaseIO ContextState := do
+  let mut state := state
+
+  let some (tokenOpt, children) := state.tokens.get? id
+    | return state
+
+  for tokenId in children do
+    state ← cancelChildren state tokenId reason
+
+  if let some token := tokenOpt then
+    token.cancel reason
+
+  pure { state with tokens := state.tokens.erase id }
+
+/--
+Cancel this context and all child contexts with the given reason.
+-/
+def cancel (x : Context) (reason : CancellationReason) : BaseIO Unit := do
+  if ← x.token.isCancelled then
+    return
+
+  x.state.atomically do
+    let st ← get
+    let st ← cancelChildren st x.id reason
+    set st
+
+/--
+Check if the context is cancelled.
+-/
+@[inline]
+def isCancelled (x : Context) : BaseIO Bool := do
+  x.token.isCancelled
+
+/--
+Get the cancellation reason if the context is cancelled.
+-/
+@[inline]
+def getCancellationReason (x : Context) : BaseIO (Option CancellationReason) := do
+  x.token.getCancellationReason
+
+/--
+Wait for cancellation. Returns a task that completes when the context is cancelled.
+-/
+@[inline]
+def done (x : Context) : IO (AsyncTask Unit) :=
+  x.token.wait
+
+/--
+Creates a selector that waits for cancellation.
+-/
+@[inline]
+def doneSelector (x : Context) : Selector Unit :=
+  x.token.selector
+
+end Context
+end Std