Cache store scoping when possible

[stephencelis]

Currently, store scoping can happen quite often: whenever a view that scopes a store is recomputed, the old scoped store is discarded and a brand new scoped store is created. This involves creating a new object and parent/child subscription every single time a store is scoped. While stores are generally lightweight objects, this is still unnecessary work.

This PR introduces the ability to cache child stores directly in the parent via new scope methods. These scope methods take a state key path instead of state transform function. The parent store uses this key path as a hashable identifier to store the child store in a dictionary. This means that after the first time a store is scoped, scoping to it again is a simple dictionary lookup.

Anyone using key path expression syntax for their scopes will automatically benefit from this change, while the older style of scoping (soft deprecated for now and probably soft deprecated for some time) will mostly go through the older code path.

Adapting an existing benchmark to scope a store X number of times before sending an action to the scoped store shows a pretty significant improvement:

Before:

name                       time             std        iterations
-----------------------------------------------------------------
Store.Scoped nested tap: 1   53291.000 ns ±   6.35 %      26188
Store.Scoped nested tap: 2 1818583.000 ns ±  31.98 %       2514
Store.Scoped nested tap: 3 3289500.000 ns ±   6.57 %        469
Store.Scoped nested tap: 4 4203916.500 ns ±   5.79 %        370
Store.Scoped nested tap: 5 5205125.000 ns ±   5.25 %        296

After:

name                       time             std        iterations
-----------------------------------------------------------------
Store.Scoped nested tap: 1   52375.000 ns ±   5.75 %      26657
Store.Scoped nested tap: 2  109833.000 ns ±   6.05 %      12640
Store.Scoped nested tap: 3  165834.000 ns ±   5.53 %       8371
Store.Scoped nested tap: 4  228583.000 ns ±   4.47 %       6103
Store.Scoped nested tap: 5  299500.000 ns ±   4.11 %       4645

[stephencelis]

These deletions in our integration suite represent scoped stores that are able to be reused rather than discarded and recreated each render.

[stephencelis]

Since stores are now being cached, we think this rescope machinery is no longer necessary in most cases. I also did not measure much of a performance difference in its removal in our very limited suite.

@iampatbrown Let me know your thoughts on this, especially if you have apprehensions. I think we could potentially keep rescope for stores that are not cached, though I'm not 100% sure what that would look like.

[iampatbrown]

No apprehensions :) I'm a fan of the cached approach.
Scoping via a key path is probably preferable. Computed properties might be an alternative for people using a closure as well.

[stephencelis]

We should probably explore using a dedicated reducer conformance for scoping (as was done before with rescope), though it'd be nice to measure the difference.

[iampatbrown]

It could potentially contain some of the strong/weak dance too.

[pyrtsa]

Could you explain how this affects the lifetimes of child stores? Will the parent retain them for as long as it lives, or if not then when will a child be removed from the children dictionary?

[stephencelis]

Could you explain how this affects the lifetimes of child stores? Will the parent retain them for as long as it lives, or if not then when will a child be removed from the children dictionary?

@pyrtsa Child stores live as long as they are "valid". Optional/presented children are removed when the associated state goes away. While this means some stores may live longer than previously, where they'd only live as long as a view is in the hierarchy, we think it won't ever be an inordinate number.

[acosmicflamingo]

Many of the custom UICollectionViewListCell subclasses I have in my codebase have a Store in them that are actually tuples, and I am wondering if/how I could leverage caching from this PR in that use-case. For example, here's how I register my supplementary views:

let footerRegistration = UICollectionView.SupplementaryRegistration<SupplementaryView>(
  elementKind: UICollectionView.elementKindSectionFooter
) { [weak store] footerView, elementKind, indexPath in
  guard let store else { return }

  store.scope(
    state: {
      let item = $0.sections[indexPath.section].footer
      return (
        item: item,
        elementKind: elementKind
      )
    },
    action: { $0 }
  ).ifLet { [weak footerView] store in
    footerView?.update(with: store)
  }
  .store(in: &footerView.cancellables)
}

Would there be a way I could benefit from cached scoping if I am using closures to transform a parent state to a tuple child state? I imagine the answer is no because I think 99% of the time the scoped state is a subset of original state, but that can't be said for what I'm doing in the example above (elementKind value does not come from store).

[stephencelis]

@acosmicflamingo You can always define computed properties and subscripts to make new key paths. In this case I think you could add a custom subscript that takes the element kind.

Something like:

extension Feature.State {
  subscript(
    section section: Int, elementKind elementKind: ElementKind
  ) -> (item: Item, elementKind: ElementKind) {
    (item: self.sections[section].footer, elementKind: elementKind)
  }
}

And then:

store.scope(
  state: \.[section: indexPath.section, elementKind: elementKind],
  action: { $0 }
).ifLet { [weak footerView] store in
  footerView?.update(with: store)
}
.store(in: &footerView.cancellables)

[acosmicflamingo]

@stephencelis BRILLIANT! Thank you for the reply AND example too!

[acosmicflamingo]

@mbrandonw that is correct. The commit where the regression occurs only includes changes where Store.action is now being passed a casePath instead of a closure. I also verified that the scope where the issue occurs only exists in store-tree and not 1.4.2. I'll try to create a reproducer soon.

[acosmicflamingo]

@mbrandonw this might be quite a corner case that only impacts me because I am doing something a little different in order to work around an issue that has to do with how UICollectionView itself handles applying snapshots when the snapshot sectionIdentifiers are dynamic. There was a workaround that an Apple engineer in the UICollectionView team recommended years ago where to work around this issue I should apply an empty snapshot first and then the actual snapshot to a data source...so it might not be a coincidence that the instances where I'm getting these issues are the same ones that involve this workaround. Any time a UICollectionReusableView instance is brought into view, it contains the right data as well.

I'll have to do further investigation but I do not think that this is a blocker for the store-tree release. In all other cases in my app where I utilize diffable data source snapshot API, everything works perfectly.

[mbrandonw]

@acosmicflamingo Thanks for following up. Please do let us know if you uncover anything fishy happening.

[acosmicflamingo]

Hmm...how does a Store know when to clear out the Store cache whose key comprises the keyPath and casePath of the state and action, respectively? With some print statements, I have found that a store instance theoretically has the right value, but the scoped store does not reflect that:

  if
    let item = store.withState(\.snapshotData.sections[safeIndex: indexPath.section])
  {
    print("Store.withState:", item)
  }
  if let scopedState = store.scope(
    state: \.snapshotData.sections[safeIndex: indexPath.section],
    action: \.self
  ).withState({ $0 }) {
    print("ScopedStore.withState", scopedState)
  }
  
Prints:
Store.withState: string("A")
ScopedStore.withState string("Next Week")
Store.withState: string("B")
ScopedStore.withState string("Today")

When I asked the question earlier about the SupplementaryView displaying the wrong text, I thought that there was some race condition going on (and maybe there still is), but now I can definitely say that a Store instance does certainly have the value I expect, but scoping to that with the kePath/casePath scope function gives me the wrong value. I know you asked for a reproducer but it's going to take a whole lot more time than I anticipated to do that.

[acosmicflamingo]

I'm pretty confident I have found the root of the problem and it looks to be a race condition. In my scenario, I'm extracting some scoped state within a cell registration closure and it appears that by getting the cached Store, it's so absolutely fast that my SupplementaryView instance is extracting some state before the stateSubject publisher has even had a chance to emit a new stream and update itself with a new value: https://github.com/pointfreeco/swift-composable-architecture/blob/50b5cc6fffc09bf2c2a1ecca3058aa4b8e61b4da/Sources/ComposableArchitecture/Store.swift#L559C10-L559C10.

I can confirm this with some trusty print statements:

ChildState has updated in scoped subscription closure with: Optional(MediaLibraryFeature.MediaLibrarySection.string("A"))
Store.withState: string("A")
Returning childStore
ScopedStore.withState string("A")
MediaLibrarySectionCell updated with A
Store.withState: string("B")
Returning childStore
ScopedStore.withState string("Today")
MediaLibrarySectionCell updated with Today
ChildState has updated in scoped subscription closure with: Optional(MediaLibraryFeature.MediaLibrarySection.string("B"))
ChildState has updated in scoped subscription closure with: Optional(MediaLibraryFeature.MediaLibrarySection.string("A"))
ChildState has updated in scoped subscription closure with: Optional(MediaLibraryFeature.MediaLibrarySection.string("B"))

Notice that the first header view is correct, and the logs show that the subscription for the letter "A" is executed before any of the Store.scope functions within the cell registration closure have printed anything. Now notice that the order is flipped on the incorrect "Today" header at the bottom (which should be the letter "B"). It's very interesting that the logs show that MediaLibrarySectionCell have updated with the out-of-date state before "ChildState has updated in scoped subscription closure with: Optional(MediaLibraryFeature.MediaLibrarySection.string("B"))" ever appeared.

[mbrandonw]

Thanks for the details, but without a minimal reproducing project there really isn't much help we can offer. I think that would be the best path forward.

[acosmicflamingo]

@mbrandonw Yep, I understand. Since I've narrowed the problem down, it'll be easier for me to create a reproducer. In addition, there are things I can do on my end to work around this (like slapping on a receive(on: DispatchQueue.main). Thanks!

[acosmicflamingo]

Here's a reproducer: https://github.com/acosmicflamingo/StoreScopingCacheReproducer.git. To reproduce, simply tap on the second segmented control button and you'll see that sometimes the header does not reflect what should be shown. These logs that I have also confirm what I'm seeing:

/*
Initial:
Store.withState: first("9")
ScopedStore.withState first("9")
MusicLibrarySectionCell updated with 9
Store.withState: first("A")
ScopedStore.withState first("A")
MusicLibrarySectionCell updated with A
Store.withState: first("B")
ScopedStore.withState first("B")
MusicLibrarySectionCell updated with B
Store.withState: first("C")
ScopedStore.withState first("C")
MusicLibrarySectionCell updated with C
Store.withState: first("D")
ScopedStore.withState first("D")
MusicLibrarySectionCell updated with D
*/

/*
After tapping second segment:
Store.withState: first("E")
ScopedStore.withState first("E")
MusicLibrarySectionCell updated with E
Store.withState: first("W")
ScopedStore.withState first("9")
MusicLibrarySectionCell updated with 9
Store.withState: first("S")
ScopedStore.withState first("S")
MusicLibrarySectionCell updated with S
Store.withState: first("H")
ScopedStore.withState first("H")
MusicLibrarySectionCell updated with H
Store.withState: first("D")
ScopedStore.withState first("D")
MusicLibrarySectionCell updated with D
Store.withState: first("E")
ScopedStore.withState first("E")
MusicLibrarySectionCell updated with E
Store.withState: first("S")
ScopedStore.withState first("S")
MusicLibrarySectionCell updated with S
Store.withState: first("H")
ScopedStore.withState first("H")
MusicLibrarySectionCell updated with H
Store.withState: first("W")
ScopedStore.withState first("9")
MusicLibrarySectionCell updated with 9
Store.withState: first("D")
ScopedStore.withState first("D")
MusicLibrarySectionCell updated with D
Store.withState: first("W")
ScopedStore.withState first("9")
MusicLibrarySectionCell updated with 9
Store.withState: first("H")
ScopedStore.withState first("H")
MusicLibrarySectionCell updated with H
Store.withState: first("S")
ScopedStore.withState first("S")
MusicLibrarySectionCell updated with S
Store.withState: first("E")
ScopedStore.withState first("E")
MusicLibrarySectionCell updated with E
*/

It is a race condition so sometimes it will work, in which case you just have to rebuild and try tapping again. Any of the headers could have the issue too; it's just a matter of looking in the log and seeing that Store.withState does not coincide with ScopedStore.withState. Sorry to dump cold water on the potential store-tree merge but maybe we'll find what I'm doing wrong :)

[acosmicflamingo]

It looks like the reproducer I created demonstrating this race condition still has the problem where the cached store is not guaranteed to have its sink function executed by the time it's returned, resulting in incorrect data sometimes being used to populate a UICollectionView's section headers. Is this a bug or a feature?

[mbrandonw]

@acosmicflamingo It's hard to say. In order for us to best look into the issue we think we need a smaller reproducing example. The current one is quite complex and it's hard to tell if it's an issue with the library or in UIKit.

Ideally the problem could be expressed as a simple unit test, but if that's not possible then an integration test with fewer moving parts would be best.

[acosmicflamingo]

@mbrandonw how is this latest commit acosmicflamingo/StoreScopingCacheReproducer@07930c8? If you check it out, you'll notice that I've removed all possible models, extensions, and view/controllers that makes the app more complex than it needs to be (and hopefully makes it easier to debug this issue). I've also added an integration test that will fail about half the time:

final class StoreScopingCacheReproducerUITests: XCTestCase {
  override func setUpWithError() throws {
    continueAfterFailure = false
  }

  func testSectionHeaderText() throws {
    let app = XCUIApplication()
    app.launch()

    // Always passes
    app.buttons["A to Z"].tap()
    let headerA = app.collectionViews.element(boundBy: 0).cells.element(boundBy: 0)
    XCTAssertEqual(headerA.label, "9")

    // Does not always pass
    app.buttons["Z to A"].tap()
    let headerZ = app.collectionViews.element(boundBy: 0).cells.element(boundBy: 0)
    XCTAssertEqual(headerZ.label, "W")
  }
}
 
Result: Executed 100 tests, with 58 failures (0 unexpected) in 527.966 (528.031) seconds