Mock-up of program transformation with value semantics and memoization

Memo.swift

// This file demonstrates how memoization might work in the transformation of a `ScopedProgram` into
// a `TypedProgram`, which additionally represents type information.

public struct TypedProgram {
  /// The prior representation of the program, without analysis of types.
  public let base: ScopedProgram

  /// A type.
  public struct Type_: Hashable {
    let id: Int
  }

  /// The types.
  public private(set) var types: Set<Type_>

  /// The type of each declaration.
  public private(set) var type: Memo<AST.Declaration, Type_>

  /// The supertype, if any, of each type.
  public private(set) var supertype: Memo<Type_, Type_?>

  /// Creates a partially-formed instance with empty `type` and `supertype` memos in preparation for
  /// type checking.
  private init(partiallyFormedFrom base: ScopedProgram) {
    self.base = base
    self.type = Memo()
    self.types = Set()
    self.supertype = Memo()
  }

  /// Creates a well-formed, typechecked representation of `base`, or returns `nil` indicating
  /// typechecking failed.
  ///
  /// In real source we would throw diagnostics.
  public init?(_ base: ScopedProgram) {
    var checker = TypeChecker(checking: base)
    guard let s = checker() else { return nil }
    self = s
  }

  /// Returns `true` iff `d`'s type is `t` or a subtype thereof.
  ///
  /// An example of how an algorithm that requires potential memo updating during typechecking can
  /// be re-used after typechecking is complete.
  public mutating func type(of d: AST.Declaration, isOrInherits t: Type_) -> Bool {
    // Create a mutable type-checker with the same memos as self so that the algorithm could in
    // principle write new memos.  Because a well-formed `TypedProgram` contains complete memos,
    // no actual mutation will occur, but this satisfies Swift.
    var c = TypeChecker(asAlgorithmContextFor: self)

    // Delegate the computation to the mutable context, which is then discarded.
    return c.type(of: d, isOrInherits: t)
  }

  /// The transformation from a ScopedProgram into either a TypedProgram or a type-checking
  /// diagnostic.
  private struct TypeChecker {
    /// The representation under construction.
    ///
    /// Note: this may be partially-formed.
    var program: TypedProgram

    /// Local state of the typechecker.
    ///
    /// This exists to demonstrate that the approach supports local typechecker state that is not
    /// part of the `program` under construction.  Note however that if the algorithm strategy of
    /// `type(_:isOrInherits:)` (below) is used, one must take care with respect to how this state
    /// is relied upon.  The safest kind of local state is just another Memo.
    var localState = 0

    /// Creates an instance that checks `programToCheck`.
    init(checking programToCheck: ScopedProgram) {
      program = .init(partiallyFormedFrom: programToCheck)
    }

    /// Creates an instance as a mutable context for algorithms on `p`, whose memos are already
    /// complete.
    init(asAlgorithmContextFor p: TypedProgram) {
      program = p
    }

    /// Returns the program passed to the initializer, fully typechecked, or `nil` to indicate
    /// type-checking failed.
    mutating func callAsFunction() -> TypedProgram? {
      // Compute/record types of declarations
      for d in program.base.base.declarations { _ = type(d) }
      // Compute/record supertypes of types.
      for t in program.types { _ = supertype(t) }
      return program
    }

    /// Returns `true` iff `d`'s type is `t` or a subtype thereof, memoizing this and other
    /// computations along the way.
    mutating func type(of d: AST.Declaration, isOrInherits t: Type_) -> Bool {
      var r = type(d)
      while r != t {
        guard let s = supertype(r) else { return false }
        r = s
      }
      return true
    }

    /// Returns the type of `d`, memoizing this and other computations along the way.
    mutating func type(_ d: AST.Declaration) -> Type_ {
      let r = program.type(d) { Type_(id: d.id) }

      // Record supertype relationships for as-yet-unseen types.
      //
      // Note the slightly awkward code arrangement. We can't do this inside of the closure above
      // because it would create overlapping mutable accesses.
      //
      // Given that construction has a supertype recording pass, this code is, strictly speaking,
      // needless; it was added to demonstrate the awkwardness, which could be overcome if `Memo`
      // had a more flexible, granular API.
      let t = Type_(id: d.id)
      if program.types.insert(t).inserted {
        _ = supertype(t)
      }

      return r
    }

    /// Returns the supertype of `t`, if any.
    mutating func supertype(_ t: Type_) -> Type_? {
      return program.supertype(t) {
        localState += 1
        return t.id % 2 == 0 ? nil : Type_(id: t.id - 1)
      }
    }
  }

}

/// A property map from `Key` to `Value`.
public struct Memo<Key: Hashable, Value> {

  private var storage: [Key: Value] = [:]

  subscript(k: Key) -> Value {
    storage[k]!
  }

  /// If `k` is in `self`, returns its value; otherwise returns `compute()`, storing it as the value
  /// for `k`.
  mutating func callAsFunction(_ k: Key, compute: ()->Value) -> Value {
    { (x: inout Value) in x }(&storage[k, default: compute()])
  }

}

// Here follows a mockup of the two representations that precede `TypedProgram`

/// An abstract syntax tree for a program.
public struct AST {

  /// A declaration in an AST.
  public struct Declaration: Hashable {

    /// The identity of this declaration
    let id: Int

  }

  /// A scope in an AST
  public struct Scope: Hashable { let id: Int }

  /// Creates an instance with `declCount` declarations
  init(declCount: Int) {
    declarations = (0..<declCount).map(Declaration.init)
    scopes = [Scope(id: 0)]
  }

  /// The declarations in this program.
  let declarations: [Declaration]
  let scopes: [Scope]
}


/// An AST where every declaration has been assigned to a scope.
public struct ScopedProgram {

  /// The representation sans scope resolution
  public let base: AST

  /// Creates an instance assigning a scope to each declaration in `base`.
  init(_ base: AST) {
    self.base = base
    self.scope = .init(uniqueKeysWithValues: base.declarations.enumerated().map { ($1, base.scopes[0])})
  }

  /// Scope resolution results
  let scope: [AST.Declaration: AST.Scope]

}