Value semantic memoization in Swift

ValueMemoize.swift

// Just to give a sense of reality.
struct AST {}

/// The wrapper over a dictionary that maps keys onto non-optional
/// values, expecting them to have already been computed.
struct ImmutablePropertyMap<Key: Hashable, Value> {
  let storage: Dictionary<Key, Value>
  subscript(k: Key) -> Value { storage[k]! }
}

/// The wrapper over a type checker and property map storage ID that
/// maps keys onto values, lazily updating the property map storage in
/// the type checker.
struct MemoizedPropertyMap<Key: Hashable, Value> {
  var typeChecker: TypeChecker
  let mapID: WritableKeyPath<TypedProgram.State, [Key: Value]>
  let compute: (inout TypeChecker, Key) -> Value

  subscript(k: Key) -> Value {
    mutating get {
      // FIXME: there's a more efficient way to do this with subscript default values
      if let r = typeChecker.program.state[keyPath: mapID][k] { return r }
      let r = compute(&typeChecker, k)
      typeChecker.program.state[keyPath: mapID][k] = r
      return r
    }
  }
}

@dynamicMemberLookup
struct TypedProgram {
  let ast: AST
  /// The storage of property maps that need to be built as the program is typechecked.
  var state: State = State()

  /// The storage of property maps that need to be built as the program is typechecked.
  struct State {
    /// Maps an integer onto its square.
    var square: [Int: Int] = [:]

    /// Maps an integer onto its name.
    var name: [Int: String] = [:]
  }

  /// Accesses the property map corresponding to the storage indicated by mapID.
  subscript<K, V>(dynamicMember mapID: KeyPath<State, [K:V]>) -> ImmutablePropertyMap<K,V> {
    .init(storage: state[keyPath: mapID])
  }
}

/// Type erasure for `Computer` (below).
protocol AnyComputer {
  var key: PartialKeyPath<TypedProgram.State> { get }
  var computation: Any { get }
}

/// Notionally, a
struct Computer<Key: Hashable, Value>: AnyComputer {
  init(_ mapID:  KeyPath<TypedProgram.State, [Key: Value]>, _ f: @escaping (inout TypeChecker, Key)->Value) {
    self.mapID = mapID
    self.f = f
  }

  let mapID: KeyPath<TypedProgram.State, [Key: Value]>
  let f: (inout TypeChecker, Key)->Value

  var key: PartialKeyPath<TypedProgram.State> { mapID }
  var computation: Any { f }
}

let emptyTypeChecker = TypeChecker(program: TypedProgram(ast: AST()))

@dynamicMemberLookup
struct TypeChecker {
  var program: TypedProgram

  let computers: [PartialKeyPath<TypedProgram.State>: AnyComputer]
    = Dictionary(
      uniqueKeysWithValues: ([
                               Computer(\.square) { t, i in i * i },
                               Computer(\.name) { t, i in t.computeName(i) }
                             ] as [AnyComputer]).map { ($0.key, $0) })

  var counter = 0

  mutating func computeName(_ i: Int) -> String {
    counter += 1
    switch i {
    case 0: return "zero"
    case 1: return "one"
    case 2: return "two"
    default: return "<some-number>"
    }
  }

  subscript<K,V>(dynamicMember mapID: WritableKeyPath<TypedProgram.State, [K:V]>) -> MemoizedPropertyMap<K,V> {
    mutating get {
      fatalError("unreachable") // _modify will always be chosen *because* the getter is a mutating operation!
    }
    _modify {
      var y = MemoizedPropertyMap<K, V>(
        typeChecker: self, mapID: mapID, compute: computers[mapID]!.computation as! (inout TypeChecker, K)->V)
      self = emptyTypeChecker
      yield &y
      self = y.typeChecker
    }
  }

}

func main() {
  var t = TypeChecker(program: TypedProgram(ast: AST()))
  print(t.square[3])
  print(t.square[2])
  print(t.name[0])
  print(t.name[99])
  print(t.program)
}

main()