StoredPropertyIterable + CustomKeyPathIterable

property-key-paths.swift

//============================================================================//
// Part 1. StoredPropertyIterable
// This models the purely static layout of a struct.
//============================================================================//

// This is an implementation detail that is required before PAT existentials are
// possible.
protocol _StoredPropertyIterableBase {
  static var _allStoredPropertiesTypeErased: [AnyKeyPath] { get }
  static var _recursivelyAllStoredPropertiesTypeErased: [AnyKeyPath] { get }
}

protocol StoredPropertyIterable : _StoredPropertyIterableBase {
  associatedtype AllStoredProperties : Collection
    where AllStoredProperties.Element == PartialKeyPath<Self>
  static var allStoredProperties: AllStoredProperties { get }
}

extension StoredPropertyIterable {
  static var _allStoredPropertiesTypeErased: [AnyKeyPath] {
    return allStoredProperties.map { $0 as AnyKeyPath }
  }
  
  static var _recursivelyAllStoredPropertiesTypeErased: [AnyKeyPath] {
    return recursivelyAllStoredProperties.map { $0 as AnyKeyPath }
  }
}

extension StoredPropertyIterable {
  static var recursivelyAllStoredProperties: [PartialKeyPath<Self>] {
    var kps: [PartialKeyPath<Self>] = []
    for kp in allStoredProperties {
      if let t = type(of: kp).valueType as? _StoredPropertyIterableBase.Type {
        for nkp in t._recursivelyAllStoredPropertiesTypeErased {
          kps.append(kp.appending(path: nkp)!)
        }
      }
      kps.append(kp)
    }
    return kps
  }
}

// Type-specific key paths getters.
extension StoredPropertyIterable {
  static func allStoredProperties<T>(ofValueType _: T.Type) -> [KeyPath<Self, T>] {
    return allStoredProperties.compactMap { $0 as? KeyPath<Self, T> }
  }
  
  static func allMutableStoredProperties<T>(
    ofValueType _: T.Type
  ) -> [WritableKeyPath<Self, T>] {
    return allStoredProperties(ofValueType: T.self).compactMap {
      $0 as? WritableKeyPath<Self, T>
    }
  }
  
  static func recursivelyAllStoredProperties<T>(
    ofValueType _: T.Type
  ) -> [KeyPath<Self, T>] {
    return recursivelyAllStoredProperties.compactMap { $0 as? KeyPath<Self, T> }
  }
  
  static func recursivelyAllMutableStoredProperties<T>(
    ofValueType _: T.Type
  ) -> [WritableKeyPath<Self, T>] {
    return recursivelyAllStoredProperties(ofValueType: T.self).compactMap {
      $0 as? WritableKeyPath<Self, T>
    }
  }
}

//============================================================================//
// Testbed
//============================================================================//

struct Meow : StoredPropertyIterable {
  var a: Float
  var b: Float
  
  // NOTE: This is *all* the compiler needs to synthesize.
  static var allStoredProperties: [PartialKeyPath<Meow>] {
    return [\Meow.a, \Meow.b]
  }
}

struct Foo : StoredPropertyIterable {
  var x: Int
  var y: Int
  let z: Int
  var c: Meow
  var arr: [Meow]
  var dict: [String : Meow]
  
  // NOTE: This is *all* the compiler needs to synthesize.
  static var allStoredProperties: [PartialKeyPath<Foo>] {
    return [\Foo.x, \Foo.y, \Foo.z, \Foo.c, \Foo.arr, \Foo.dict]
  }
}

Foo.allStoredProperties
Foo.allStoredProperties(ofValueType: Int.self)
Foo.allMutableStoredProperties(ofValueType: Int.self)
Foo.recursivelyAllStoredProperties


//============================================================================//
// Part 2. CustomKeyPathIterable
// This models both static and dynamic structures. When the conforming type
// also conforms to StoredPropertyIterable, it has a default implementation.
//============================================================================//

protocol _CustomKeyPathIterableBase {
  var _allKeyPathsTypeErased: [AnyKeyPath] { get }
  var _recursivelyAllKeyPathsTypeErased: [AnyKeyPath] { get }
}

protocol CustomKeyPathIterable : _CustomKeyPathIterableBase {
  associatedtype AllKeyPaths : Collection
    where AllKeyPaths.Element == PartialKeyPath<Self>
  var allKeyPaths: [PartialKeyPath<Self>] { get }
}

extension CustomKeyPathIterable {
  var _allKeyPathsTypeErased: [AnyKeyPath] {
    return allKeyPaths.map { $0 as AnyKeyPath }
  }
  var _recursivelyAllKeyPathsTypeErased: [AnyKeyPath] {
    return recursivelyAllKeyPaths.map { $0 as AnyKeyPath }
  }
}

extension CustomKeyPathIterable where Self : StoredPropertyIterable, AllKeyPaths == AllStoredProperties {
  var allKeyPaths: AllKeyPaths {
    return Self.allStoredProperties
  }
}

extension CustomKeyPathIterable {
  var recursivelyAllKeyPaths: [PartialKeyPath<Self>] {
    var kps: [PartialKeyPath<Self>] = []
    for kp in allKeyPaths {
      if let nested = self[keyPath: kp] as? _CustomKeyPathIterableBase {
        for nkp in nested._recursivelyAllKeyPathsTypeErased {
          kps.append(kp.appending(path: nkp)!)
        }
      }
      kps.append(kp)
    }
    return kps
  }
}

extension CustomKeyPathIterable {
  func allKeyPaths<T>(ofValueType _: T.Type) -> [KeyPath<Self, T>] {
    return allKeyPaths.compactMap { $0 as? KeyPath<Self, T> }
  }
  
  func recursivelyAllKeyPaths<T>(ofValueType _: T.Type) -> [KeyPath<Self, T>] {
    return recursivelyAllKeyPaths.compactMap { $0 as? KeyPath<Self, T> }
  }
}

//============================================================================//
// Testbed
//============================================================================//

extension Foo : CustomKeyPathIterable {
  // Compiler needs to derive this typealias as well, give that Foo already
  // conforms to `StoredPropertyIterable`.
  typealias AllKeyPaths = AllStoredProperties
}

extension Array : CustomKeyPathIterable {
  typealias AllKeyPaths = [PartialKeyPath<Array>]
  var allKeyPaths: [PartialKeyPath<Array>] {
    return indices.map { \Array[$0] }
  }
}

extension Dictionary : CustomKeyPathIterable {
  typealias AllKeyPaths = [PartialKeyPath<Dictionary>]
  var allKeyPaths: [PartialKeyPath<Dictionary>] {
    return keys.map { \Dictionary[$0]! }
  }
}

let foo = Foo(x: 1, y: 2, z: 3, c: Meow(a: 4, b: 5),
              arr: [Meow(a: 6, b: 7), Meow(a: 8, b: 9)],
              dict: ["a" : Meow(a: 10, b: 11), "b" : Meow(a: 12, b: 13)])
// Yay!
assert(foo.recursivelyAllKeyPaths(ofValueType: Meow.self).count == 5)