Replace "getStates" by more efficient entities subscription (compressed states)

CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Changed: `HAData` now includes a `primitive` case to express non-array/dictionary values that aren't `null`.
 - Changed: WebSocket connection will now enable compression.
 - Fixed: Calling `HAConnection.connect()` and `HAConnection.disconnect()` off the main thread no longer occasionally crashes.
+- Removed: Usage of "get_states"
+- Added: More efficient API "subscribe_entities" replacing "get_states"
 
 ## [0.3] - 2021-07-08
 - Added: Subscriptions will now retry (when their request `shouldRetry`) when the HA config changes or components are loaded.

HAKit.podspec

@@ -1,6 +1,6 @@
 Pod::Spec.new do |s|
   s.name = 'HAKit'
-  s.version = '0.3'
+  s.version = '0.4'
   s.summary = 'Communicate with a Home Assistant instance.'
   s.author = 'Home Assistant'
 

Source/Caches/HACache.swift

@@ -44,6 +44,7 @@ public class HACache<ValueType> {
         self.connection = connection
         self.populateInfo = populate
         self.subscribeInfo = subscribe
+        self.subscribeOnlyInfo = nil
 
         self.start = { connection, cache, _ in
             Self.startPopulate(for: populate, on: connection, cache: cache) { cacheResult in
@@ -82,6 +83,7 @@ public class HACache<ValueType> {
         self.connection = connection
         self.populateInfo = nil
         self.subscribeInfo = nil
+        self.subscribeOnlyInfo = subscribe
 
         self.start = { connection, cache, state in
             state.isWaitingForPopulate = false
@@ -107,6 +109,7 @@ public class HACache<ValueType> {
         self.connection = incomingCache.connection
         self.populateInfo = nil
         self.subscribeInfo = nil
+        self.subscribeOnlyInfo = nil
         self.start = { _, someCache, _ in
             // unfortunately, using this value directly crashes the swift compiler, so we call into it with this
             let cache: HACache<ValueType> = someCache
@@ -136,6 +139,7 @@ public class HACache<ValueType> {
         }
         self.populateInfo = nil
         self.subscribeInfo = nil
+        self.subscribeOnlyInfo = nil
         state.mutate { state in
             state.current = constantValue
         }
@@ -321,6 +325,9 @@ public class HACache<ValueType> {
     /// If this cache was created with subscribe info, this contains that info
     /// This is largely intended for tests and is not used internally.
     public let subscribeInfo: [HACacheSubscribeInfo<ValueType>]?
+    /// If this cache was created with subscribe info, this contains that info
+    /// This is largely intended for tests and is not used internally.
+    public let subscribeOnlyInfo: HACacheSubscribeInfo<ValueType?>?
 
     /// Do the underlying populate send
     /// - Parameters:

Source/Caches/HACacheKeyStates.swift

@@ -0,0 +1,77 @@
+import Foundation
+
+/// Key for the cache
+internal struct HACacheKeyStates: HACacheKey {
+    static func create(connection: HAConnection) -> HACache<HACachedStates> {
+        .init(
+            connection: connection,
+            subscribe:
+            .init(
+                subscription: .subscribeEntities(),
+                transform: { info in
+                    .replace(processUpdates(
+                        info: info,
+                        shouldResetEntities: info.subscriptionPhase == .initial
+                    ))
+                }
+            )
+        )
+    }
+
+    /// Process updates from the compressed state to HAEntity
+    /// - Parameters:
+    ///   - info: The compressed state update and the current cached states
+    ///   - shouldResetEntities: True if current state needs to be ignored (e.g. re-connection)
+    /// - Returns: HAEntity cached states
+    /// Logic from: https://github.com/home-assistant/home-assistant-js-websocket/blob/master/lib/entities.ts
+    static func processUpdates(
+        info: HACacheTransformInfo<HACompressedStatesUpdates, HACachedStates?>,
+        shouldResetEntities: Bool
+    )
+        -> HACachedStates {
+        var states: HACachedStates
+        if shouldResetEntities {
+            states = .init(entities: [])
+        } else {
+            states = info.current ?? .init(entities: [])
+        }
+        if let additions = info.incoming.add {
+            for (entityId, updates) in additions {
+                if var currentState = states[entityId] {
+                    currentState.update(from: updates)
+                    states[entityId] = currentState
+                } else {
+                    do {
+                        states[entityId] = try updates.asEntity(entityId: entityId)
+                    } catch {
+                        HAGlobal.log(.error, "[Update-To-Entity-Error] Failed adding new entity: \(error)")
+                    }
+                }
+            }
+        }
+
+        if let subtractions = info.incoming.remove {
+            for entityId in subtractions {
+                states[entityId] = nil
+            }
+        }
+
+        if let changes = info.incoming.change {
+            changes.forEach { entityId, diff in
+                guard var entityState = states[entityId] else { return }
+
+                if let toAdd = diff.additions {
+                    entityState.add(toAdd)
+                    states[entityId] = entityState
+                }
+
+                if let toRemove = diff.subtractions {
+                    entityState.subtract(toRemove)
+                    states[entityId] = entityState
+                }
+            }
+        }
+
+        return states
+    }
+}

Source/Caches/HACacheSubscribeInfo.swift

@@ -36,9 +36,16 @@ public struct HACacheSubscribeInfo<OutgoingType> {
 
                 return transform(value)
             }, start: { connection, perform in
-                connection.subscribe(to: nonRetrySubscription, handler: { _, incoming in
+                var operationType: HACacheSubscriptionPhase = .initial
+                return connection.subscribe(to: nonRetrySubscription, handler: { _, incoming in
                     perform { current in
-                        transform(.init(incoming: incoming, current: current))
+                        let transform = transform(.init(
+                            incoming: incoming,
+                            current: current,
+                            subscriptionPhase: operationType
+                        ))
+                        operationType = .iteration
+                        return transform
                     }
                 })
             }
@@ -63,10 +70,19 @@ public struct HACacheSubscribeInfo<OutgoingType> {
     /// - Parameters:
     ///   - incoming: The incoming value, of some given type -- intended to be the IncomingType that created this
     ///   - current: The current value part of the transform info
+    ///   - subscriptionPhase: The phase in which the subscription is, initial iteration or subsequent
     /// - Throws: If the type of incoming does not match the original IncomingType
     /// - Returns: The response from the transform block
-    public func transform<IncomingType>(incoming: IncomingType, current: OutgoingType) throws -> Response {
-        try anyTransform(HACacheTransformInfo<IncomingType, OutgoingType>(incoming: incoming, current: current))
+    public func transform<IncomingType>(
+        incoming: IncomingType,
+        current: OutgoingType,
+        subscriptionPhase: HACacheSubscriptionPhase
+    ) throws -> Response {
+        try anyTransform(HACacheTransformInfo<IncomingType, OutgoingType>(
+            incoming: incoming,
+            current: current,
+            subscriptionPhase: subscriptionPhase
+        ))
     }
 
     /// The start handler

Source/Caches/HACacheTransformInfo.swift

@@ -9,4 +9,15 @@ public struct HACacheTransformInfo<IncomingType, OutgoingType> {
     /// For populate transforms, this is nil if an initial request hasn't been sent yet and the cache not reset.
     /// For subscribe transforms, this is nil if the populate did not produce results (or does not exist).
     public var current: OutgoingType
+
+    /// The current phase of the subscription
+    public var subscriptionPhase: HACacheSubscriptionPhase = .initial
+}
+
+/// The subscription phases
+public enum HACacheSubscriptionPhase {
+    /// `Initial` means it's the first time a value is returned
+    case initial
+    /// `Iteration` means subsequent iterations
+    case iteration
 }

Source/Caches/HACachedStates.swift

@@ -3,21 +3,6 @@ public extension HACachesContainer {
     var states: HACache<HACachedStates> { self[HACacheKeyStates.self] }
 }
 
-/// Key for the cache
-private struct HACacheKeyStates: HACacheKey {
-    static func create(connection: HAConnection) -> HACache<HACachedStates> {
-        .init(
-            connection: connection,
-            populate: .init(request: .getStates(), transform: { .init(entities: $0.incoming) }),
-            subscribe: .init(subscription: .stateChanged(), transform: { info in
-                var updated = info.current
-                updated[info.incoming.entityId] = info.incoming.newState
-                return .replace(updated)
-            })
-        )
-    }
-}
-
 /// Cached version of all entity states
 public struct HACachedStates {
     /// All entities

Source/Convenience/States.swift

@@ -9,13 +9,11 @@ public extension HATypedSubscription {
             "event_type": HAEventType.stateChanged.rawValue!,
         ]))
     }
-}
 
-public extension HATypedRequest {
-    /// Get the state of all entities
-    /// - Returns: A typed request that can be sent via `HAConnection`
-    static func getStates() -> HATypedRequest<[HAEntity]> {
-        .init(request: .init(type: .getStates, data: [:]))
+    /// Listen for compressed state changes of all entities
+    /// - Returns: A typed subscriptions that can be sent via `HAConnection`
+    static func subscribeEntities() -> HATypedSubscription<HACompressedStatesUpdates> {
+        .init(request: .init(type: .subscribeEntities, data: [:]))
     }
 }
 

Source/Data/HACompressedEntity.swift

@@ -0,0 +1,58 @@
+import Foundation
+
+public struct HACompressedStatesUpdates: HADataDecodable {
+    public var add: [String: HACompressedEntityState]?
+    public var remove: [String]?
+    public var change: [String: HACompressedEntityDiff]?
+
+    public init(data: HAData) throws {
+        self.add = try? data.decode("a")
+        self.remove = try? data.decode("r")
+        self.change = try? data.decode("c")
+    }
+}
+
+public struct HACompressedEntityState: HADataDecodable {
+    public var state: String
+    public var attributes: [String: Any]?
+    public var context: String?
+    public var lastChanged: Date?
+    public var lastUpdated: Date?
+
+    public init(data: HAData) throws {
+        self.state = try data.decode("s")
+        self.attributes = try? data.decode("a")
+        self.context = try? data.decode("c")
+        self.lastChanged = try? data.decode("lc")
+        self.lastUpdated = try? data.decode("lu")
+    }
+
+    func asEntity(entityId: String) throws -> HAEntity {
+        try HAEntity(
+            entityId: entityId,
+            state: state,
+            lastChanged: lastChanged ?? Date(),
+            lastUpdated: lastUpdated ?? Date(),
+            attributes: attributes ?? [:],
+            context: .init(id: context ?? "", userId: nil, parentId: nil)
+        )
+    }
+}
+
+public struct HACompressedEntityStateRemove: HADataDecodable {
+    public var attributes: [String]?
+
+    public init(data: HAData) throws {
+        self.attributes = try? data.decode("a")
+    }
+}
+
+public struct HACompressedEntityDiff: HADataDecodable {
+    public var additions: HACompressedEntityState?
+    public var subtractions: HACompressedEntityStateRemove?
+
+    public init(data: HAData) throws {
+        self.additions = try? data.decode("+")
+        self.subtractions = try? data.decode("-")
+    }
+}

Source/Data/HADecodeTransformable.swift

@@ -76,12 +76,18 @@ extension Date: HADecodeTransformable {
     /// - Parameter value: A string value to convert
     /// - Returns: The value converted to a Date, or nil if not possible
     public static func decode(unknown value: Any) -> Self? {
-        guard let value = value as? String else { return nil }
-        for formatter in Self.formatters {
-            if let string = formatter.date(from: value) {
-                return string
+        if let timestamp = value as? Double {
+            return Date(timeIntervalSince1970: timestamp)
+        }
+
+        if let value = value as? String {
+            for formatter in Self.formatters {
+                if let string = formatter.date(from: value) {
+                    return string
+                }
             }
         }
+
         return nil
     }
 }

Source/Data/HAEntity+CompressedEntity.swift

@@ -0,0 +1,23 @@
+import Foundation
+
+extension HAEntity {
+    mutating func update(from state: HACompressedEntityState) {
+        self.state = state.state
+        lastChanged = state.lastChanged ?? lastChanged
+        lastUpdated = state.lastUpdated ?? lastUpdated
+        attributes.dictionary = state.attributes ?? attributes.dictionary
+        context = .init(id: state.context ?? "", userId: nil, parentId: nil)
+    }
+
+    mutating func add(_ state: HACompressedEntityState) {
+        self.state = state.state
+        lastChanged = state.lastChanged ?? lastChanged
+        lastUpdated = state.lastUpdated ?? lastUpdated
+        attributes.dictionary.merge(state.attributes ?? [:]) { current, _ in current }
+        context = .init(id: state.context ?? "", userId: nil, parentId: nil)
+    }
+
+    mutating func subtract(_ state: HACompressedEntityStateRemove) {
+        attributes.dictionary = attributes.dictionary.filter { !(state.attributes?.contains($0.key) ?? false) }
+    }
+}