Added multicast to ViewStore publisher to reduce equality checks (#624)

* Added multicast to ViewStore publisher to reduce equality checks

* Replaced multicast with ViewStore CurrentValueSubject

* Added guard check to weak self

* Clean up and add test.

Co-authored-by: Brandon Williams <[email protected]>

Author: iampatbrown

Sources/ComposableArchitecture/Store.swift

@@ -416,27 +416,24 @@ public struct StorePublisher<State>: Publisher {
   public typealias Output = State
   public typealias Failure = Never
 
-  private let isDuplicate: (State, State) -> Bool
   public let upstream: AnyPublisher<State, Never>
 
   public func receive<S>(subscriber: S)
   where S: Subscriber, Failure == S.Failure, Output == S.Input {
-    self.upstream.removeDuplicates(by: isDuplicate).subscribe(subscriber)
+    self.upstream.subscribe(subscriber)
   }
 
   init<P>(
-    _ upstream: P,
-    removeDuplicates isDuplicate: @escaping (State, State) -> Bool
+    _ upstream: P
   ) where P: Publisher, Failure == P.Failure, Output == P.Output {
     self.upstream = upstream.eraseToAnyPublisher()
-    self.isDuplicate = isDuplicate
   }
 
   /// Returns the resulting publisher of a given key path.
   public subscript<LocalState>(
     dynamicMember keyPath: KeyPath<State, LocalState>
   ) -> StorePublisher<LocalState>
   where LocalState: Equatable {
-    .init(self.upstream.map(keyPath), removeDuplicates: ==)
+    .init(self.upstream.map(keyPath).removeDuplicates())
   }
 }

Sources/ComposableArchitecture/ViewStore.swift

@@ -55,16 +55,13 @@ public final class ViewStore<State, Action>: ObservableObject {
   /// A publisher of state.
   public let publisher: StorePublisher<State>
 
-  private var viewCancellable: AnyCancellable?
-
   // N.B. `ViewStore` does not use a `@Published` property, so `objectWillChange`
   // won't be synthesized automatically. To work around issues on iOS 13 we explicitly declare it.
   public private(set) lazy var objectWillChange = ObservableObjectPublisher()
 
-  /// The current state.
-  public var state: State { self.store.state.value }
-
-  private let store: Store<State, Action>
+  private let _send: (Action) -> Void
+  private let _state: CurrentValueSubject<State, Never>
+  private var viewCancellable: AnyCancellable?
 
   /// Initializes a view store from a store.
   ///
@@ -76,17 +73,27 @@ public final class ViewStore<State, Action>: ObservableObject {
     _ store: Store<State, Action>,
     removeDuplicates isDuplicate: @escaping (State, State) -> Bool
   ) {
-    self.publisher = StorePublisher(store.state, removeDuplicates: isDuplicate)
-    self.store = store
+    self._send = store.send
+    self._state = CurrentValueSubject(store.state.value)
+
+    self.publisher = StorePublisher(self._state)
     self.viewCancellable = store.state
-      .dropFirst()
       .removeDuplicates(by: isDuplicate)
-      .sink { [weak self] _ in self?.objectWillChange.send() }
+      .sink { [weak self] in
+        guard let self = self else { return }
+        self.objectWillChange.send()
+        self._state.send($0)
+      }
+  }
+
+  /// The current state.
+  public var state: State {
+    self._state.value
   }
 
   /// Returns the resulting value of a given key path.
   public subscript<LocalState>(dynamicMember keyPath: KeyPath<State, LocalState>) -> LocalState {
-    self.state[keyPath: keyPath]
+    self._state.value[keyPath: keyPath]
   }
 
   /// Sends an action to the store.
@@ -99,7 +106,7 @@ public final class ViewStore<State, Action>: ObservableObject {
   ///
   /// - Parameter action: An action.
   public func send(_ action: Action) {
-    self.store.send(action)
+    self._send(action)
   }
 
   /// Derives a binding from the store that prevents direct writes to state and instead sends
@@ -135,7 +142,7 @@ public final class ViewStore<State, Action>: ObservableObject {
     send localStateToViewAction: @escaping (LocalState) -> Action
   ) -> Binding<LocalState> {
     Binding(
-      get: { get(self.state) },
+      get: { get(self._state.value) },
       set: { newLocalState, transaction in
         if transaction.animation != nil {
           withTransaction(transaction) {

Tests/ComposableArchitectureTests/ViewStoreTests.swift

@@ -64,16 +64,16 @@ final class ViewStoreTests: XCTestCase {
     XCTAssertEqual(0, subEqualityChecks)
     viewStore4.send(())
     XCTAssertEqual(4, equalityChecks)
+    XCTAssertEqual(4, subEqualityChecks)
+    viewStore4.send(())
+    XCTAssertEqual(8, equalityChecks)
     XCTAssertEqual(8, subEqualityChecks)
     viewStore4.send(())
     XCTAssertEqual(12, equalityChecks)
-    XCTAssertEqual(20, subEqualityChecks)
-    viewStore4.send(())
-    XCTAssertEqual(20, equalityChecks)
-    XCTAssertEqual(32, subEqualityChecks)
+    XCTAssertEqual(12, subEqualityChecks)
     viewStore4.send(())
-    XCTAssertEqual(28, equalityChecks)
-    XCTAssertEqual(44, subEqualityChecks)
+    XCTAssertEqual(16, equalityChecks)
+    XCTAssertEqual(16, subEqualityChecks)
   }
 
   func testAccessViewStoreStateInPublisherSink() {
@@ -97,6 +97,28 @@ final class ViewStoreTests: XCTestCase {
 
     XCTAssertEqual([0, 1, 2, 3], results)
   }
+
+  func testWillSet() {
+    let reducer = Reducer<Int, Void, Void> { count, _, _ in
+      count += 1
+      return .none
+    }
+
+    let store = Store(initialState: 0, reducer: reducer, environment: ())
+    let viewStore = ViewStore(store)
+
+    var results: [Int] = []
+
+    viewStore.objectWillChange
+      .sink { _ in results.append(viewStore.state) }
+      .store(in: &self.cancellables)
+
+    viewStore.send(())
+    viewStore.send(())
+    viewStore.send(())
+
+    XCTAssertEqual([0, 1, 2], results)
+  }
 }
 
 private struct State: Equatable {