@@ -71,7 +71,7 @@ internal struct StoreContext {
7171 let ( key, _) = lookupAtomKeyAndOverride ( of: atom)
7272
7373 if let cache = lookupCache ( of: atom, for: key) {
74- switchContext ( scope : cache. initializedScope )
74+ switchContext ( with : cache)
7575 . update ( atom: atom, for: key, cache: cache, newValue: value)
7676 }
7777 }
@@ -82,7 +82,7 @@ internal struct StoreContext {
8282
8383 if let cache = lookupCache ( of: atom, for: key) {
8484 let newValue = mutating ( cache. value, body)
85- switchContext ( scope : cache. initializedScope )
85+ switchContext ( with : cache)
8686 . update ( atom: atom, for: key, cache: cache, newValue: newValue)
8787 }
8888 }
@@ -131,7 +131,7 @@ internal struct StoreContext {
131131 func refresh< Node: AsyncAtom > ( _ atom: Node ) async -> Node . Produced {
132132 let ( key, override) = lookupAtomKeyAndOverride ( of: atom)
133133 let cache = lookupCache ( of: atom, for: key)
134- let localContext = switchContext ( scope : cache? . initializedScope )
134+ let localContext = cache. map ( switchContext ) ?? self
135135 let context = localContext. prepareForTransaction ( of: atom, for: key)
136136
137137 let value : Node . Produced
@@ -163,7 +163,7 @@ internal struct StoreContext {
163163 func refresh< Node: Refreshable > ( _ atom: Node ) async -> Node . Produced {
164164 let ( key, _) = lookupAtomKeyAndOverride ( of: atom)
165165 let cache = lookupCache ( of: atom, for: key)
166- let localContext = switchContext ( scope : cache? . initializedScope )
166+ let localContext = cache. map ( switchContext ) ?? self
167167 let state = localContext. getState ( of: atom, for: key)
168168 let currentContext = AtomCurrentContext ( store: localContext)
169169
@@ -194,7 +194,7 @@ internal struct StoreContext {
194194 let ( key, override) = lookupAtomKeyAndOverride ( of: atom)
195195
196196 if let cache = lookupCache ( of: atom, for: key) {
197- let localContext = switchContext ( scope : cache. initializedScope )
197+ let localContext = switchContext ( with : cache)
198198 let newValue = localContext. getValue ( of: atom, for: key, override: override)
199199 localContext. update ( atom: atom, for: key, cache: cache, newValue: newValue)
200200 }
@@ -204,10 +204,12 @@ internal struct StoreContext {
204204 @usableFromInline
205205 func reset( _ atom: some Resettable ) {
206206 let ( key, _) = lookupAtomKeyAndOverride ( of: atom)
207- let cache = lookupCache ( of: atom, for: key)
208- let localContext = switchContext ( scope: cache? . initializedScope)
209- let currentContext = AtomCurrentContext ( store: localContext)
210- atom. reset ( context: currentContext)
207+
208+ if let cache = lookupCache ( of: atom, for: key) {
209+ let localContext = switchContext ( with: cache)
210+ let currentContext = AtomCurrentContext ( store: localContext)
211+ atom. reset ( context: currentContext)
212+ }
211213 }
212214
213215 @usableFromInline
@@ -322,7 +324,7 @@ private extension StoreContext {
322324
323325 func updatePropagation( for key: AtomKey , cache: some AtomCacheProtocol ) {
324326 // Dependents must be updated with the scope at which they were initialised.
325- let localContext = switchContext ( scope : cache. initializedScope )
327+ let localContext = switchContext ( with : cache)
326328
327329 // Overridden atoms don't get updated transitively.
328330 let newValue = localContext. getValue ( of: cache. atom, for: key, override: nil )
@@ -422,7 +424,7 @@ private extension StoreContext {
422424
423425 if let state, let cache {
424426 // It must call release effect with the scope at which they were initialised.
425- let localContext = switchContext ( scope : cache. initializedScope )
427+ let localContext = switchContext ( with : cache)
426428 let currentContext = AtomCurrentContext ( store: localContext)
427429 state. effect. released ( context: currentContext)
428430 }
@@ -510,7 +512,7 @@ private extension StoreContext {
510512
511513 return AtomProducerContext ( store: self , transactionState: transactionState) { newValue in
512514 if let cache = lookupCache ( of: atom, for: key) {
513- switchContext ( scope : cache. initializedScope )
515+ switchContext ( with : cache)
514516 . update ( atom: atom, for: key, cache: cache, newValue: newValue)
515517 }
516518 }
@@ -647,11 +649,11 @@ private extension StoreContext {
647649 }
648650 }
649651
650- func switchContext( scope : Scope ? ) -> StoreContext {
652+ func switchContext( with cache : some AtomCacheProtocol ) -> StoreContext {
651653 StoreContext (
652654 store: store,
653655 rootScope: rootScope,
654- currentScope: scope
656+ currentScope: cache . initializedScope
655657 )
656658 }
657659}
0 commit comments