An efficient storage foundation for type erasure.

EfficientAny.swift

// Copyright 2020 Penguin Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

/// A substitute for `T.self` that eliminates ambiguity about whether you mean the static or dynamic
/// type.
public struct Type<T> { init() {} }

fileprivate extension MemoryLayout {
  /// True iff `T` instances are stored in an existential's inline buffer space.
  static var fitsExistentialInlineBuffer: Bool {
    MemoryLayout<(T, Any.Type)>.size <= MemoryLayout<Any>.size
      && self.alignment <= MemoryLayout<Any>.alignment
      && _isBitwiseTakable(T.self)
  }
}

/// A wrapper for a value of any type, with efficient access and inline storage of bounded size for
/// small values.
///
/// Existential types such as `Any` provide the storage characteristics of `AnyValue`, but many
/// operations on existentials [optimize poorly](https://bugs.swift.org/browse/SR-13438), so
/// `AnyValue` may be needed to achieve efficient access.
///
/// Using `enums` and no existential, it is possible to build such storage for:
/// - types that are trivially copiable (such as Int), or
/// - collections of elements with the inline storage being (roughly) a multiple of the element
///   size.
/// You might consider using the `enum` approach where applicable.
public struct AnyValue {
  /// The underlying storage of the value.
  internal var storage: Any

  /// Creates an instance that stores `x`.
  ///
  /// - Postcondition: where `a` is the created instance, `a.storedType == T.self`, and `a[T.self]`
  ///   is equivalent to `x`.
  public init<T>(_ x: T) {
    if MemoryLayout<T>.fitsExistentialInlineBuffer { storage = x}
    else { storage = Box(x) as BoxBase }
  }

  /// The layout of an `Any`.
  private typealias AnyLayout = (Int, Int, Int, storedType: Any.Type)

  /// The type of the value stored directly in `storage`.
  private var directlyStoredType: Any.Type {    
    return Swift.withUnsafePointer(to: self) {
      UnsafeRawPointer($0).assumingMemoryBound(to: AnyLayout.self)[0].storedType
    }
  }

  /// The type of the stored value.
  public var storedType: Any.Type {
    let d = directlyStoredType
    if d != BoxBase.self { return d }
    return self[unsafelyAssuming: Type<BoxBase>()].storedType
  }

  /// Returns a pointer to the `T` which is assumed to be stored in `self`.
  private func pointer<T>(toStored _: Type<T>) -> UnsafePointer<T> {
    .init(Swift.withUnsafePointer(to: self) { Handle<T>($0) }.pointer)
  }

  // Existentials that store large types have their own built-in box, but we are unable to detect
  // uniqueness of an existential box, and they seem to be reallocated unpredictably in some cases.
  // Instead, we box manually and only use the `Any` for its inline storage.
  
  /// A base class for boxing types that don't fit the inline buffer.
  private class BoxBase {
    /// The type stored in this box.
    final let storedType: Any.Type

    /// Creates an instance with the given `storedType`
    fileprivate init(storedType: Any.Type) {
      self.storedType = storedType
    }

    /// Returns the boxed value.
    fileprivate var asAny: Any { fatalError("override me") }
  }
  
  private final class Box<T>: BoxBase {
    /// The boxed value
    var value: T

    /// Creates an instance storing `value`
    init(_ value: T) {
      assert(!(value is BoxBase), "unexpected double-boxing!")
      self.value = value
      super.init(storedType: T.self)
    }

    /// Returns the boxed value.
    fileprivate override var asAny: Any { value }
  }
  
  /// Returns a pointer to the `T` which is assumed to be stored in `self`.
  private mutating func mutablePointer<T>(toStored _: Type<T>) -> UnsafeMutablePointer<T> {
    if !MemoryLayout<T>.fitsExistentialInlineBuffer {
      if !isKnownUniquelyReferenced(&self[unsafelyAssuming: Type<BoxBase>()]) {
        storage = Box(self[unsafelyAssuming: Type<T>()]) as BoxBase
      }
    }
    return withUnsafeMutablePointer(to: &self) { Handle<T>($0) }.pointer
  }
  
  /// A tool for accessing the `T` stored in the `Any`.
  private struct Handle<T> {
    // Warning: this type may look like needless complication but efficient optimization for
    // AnyValue is extremely sensitive to the shape of its implementation and I've been unable to
    // eliminate it without hurting the generated code.
    
    /// Creates an instance for accessing the `Any` whose address is `address`.
    init(_ address: UnsafeRawPointer) {
      self.address = .init(mutating: address)
    }
    /// The address of the `Any` accessed by `self`
    private let address: UnsafeMutableRawPointer

    /// A pointer to the stored instance of type `T`.
    public var pointer: UnsafeMutablePointer<T> {
      if MemoryLayout<T>.fitsExistentialInlineBuffer {
        return address.assumingMemoryBound(to: T.self)
      }
      else {
        let boxBase = address.assumingMemoryBound(to: BoxBase.self).pointee
        let box = unsafeDowncast(boxBase, to: Box<T>.self)
        return withUnsafeMutablePointer(to: &box.value) { $0 }
      }
    }
  }

  /// Iff `storedType != T.self`, traps with an appropriate error message.
  private func typeCheck<T>(_: Type<T>) {
    if storedType != T.self { typeCheckFailure(T.self) }
  }

  /// Traps with an appropriate error message assuming that `storedType != T.self`.
  @inline(never)
  private func typeCheckFailure(_ expectedType: Any.Type) {
    fatalError("stored type \(storedType) != \(expectedType)")
  }

  /// Accesses the `T` stored in `self`.
  ///
  /// - Requires: `storedType == T.self`.
  @inline(__always) // Compiler likes to skip inlining this otherwise.
  public subscript<T>(_: Type<T>) -> T {
    get {
      defer { _fixLifetime(self) }
      typeCheck(Type<T>())
      return pointer(toStored: Type<T>()).pointee
    }
    _modify {
      typeCheck(Type<T>())
      yield &mutablePointer(toStored: Type<T>()).pointee
      _fixLifetime(self)
    }
  }

  /// Unsafely accesses the `T` stored in `self`.
  ///
  /// - Requires: `storedType == T.self`.
  @inline(__always) // Compiler likes to skip inlining this otherwise.
  public subscript<T>(unsafelyAssuming _: Type<T>) -> T {
    get {
      defer { _fixLifetime(self) }
      return pointer(toStored: Type<T>()).pointee
    }
    _modify {
      yield &mutablePointer(toStored: Type<T>()).pointee
      _fixLifetime(self)
    }
  }

  /// Stores `x` in `self`.
  ///
  /// This may be more efficient than `self = AnyValue(x)` because it uses the same allocated buffer
  /// when possible for large types.
  public mutating func store<T>(_ x: T) {
    defer { _fixLifetime(self) }
    if storedType == T.self { mutablePointer(toStored: Type<T>()).pointee = x }
    else { self = .init(x) }
  }

  /// The stored value.
  ///
  /// This property can be useful for interoperability with the rest of Swift, especially when you
  /// don't know the full dynamic type of the stored value.
  public var asAny: Any {
    if directlyStoredType != BoxBase.self { return storage }
    return self[unsafelyAssuming: Type<BoxBase>()].asAny
  }

  /// If the stored value is boxed or is an object, returns the ID of the box or object; returns
  /// `nil` otherwise.
  ///
  /// Used by tests.
  internal var boxOrObjectID_: ObjectIdentifier? {
    if !(directlyStoredType is AnyObject.Type) { return nil }
    return .init(self[unsafelyAssuming: Type<AnyObject>()])
  }
}

extension AnyValue: CustomStringConvertible, CustomDebugStringConvertible {
  /// A textual representation of this instance.
  public var description: String { String(describing: storage) }

  /// A string, suitable for debugging, that represents the instance.
  public var debugDescription: String { "AnyValue(\(String(reflecting: storage)))" }
}

// =========== Testing ======

import XCTest

fileprivate protocol P {}
extension Double: P {}

fileprivate class Base {}
fileprivate class Derived: Base {}

final class AnyValueTests: XCTestCase {
  // A type that will not fit in existential inline storage.
  typealias BufferOverflow = (Int, Int, Int, Int)

  // A value that will not fit in existential inline storage.
  let bufferOverflow = (1, 2, 3, 4)
  
  func test_Int() {
    var a = AnyValue(3)
    XCTAssert(a.storedType == Int.self)
    XCTAssertEqual(a[Type<Int>()], 3)
    XCTAssertEqual(a[unsafelyAssuming: Type<Int>()], 3)
    
    a[Type<Int>()] += 1
    XCTAssertEqual(a[Type<Int>()], 4)
    XCTAssertEqual(a.boxOrObjectID, nil)
    
    a[unsafelyAssuming: Type<Int>()] += 1
    XCTAssertEqual(a[Type<Int>()], 5)
  }
  
  func test_String() {
    var a = AnyValue("Three")
    XCTAssert(a.storedType == String.self)
    XCTAssertEqual(a[Type<String>()], "Three")
    XCTAssertEqual(a[unsafelyAssuming: Type<String>()], "Three")

    a[Type<String>()].append("D")
    XCTAssertEqual(a[Type<String>()], "ThreeD")

    // I don't know how big String is off the top of my head.
    let inlineStorageExpected = MemoryLayout<String>.size < 3 * MemoryLayout<Int>.size
      && MemoryLayout<String>.alignment <= MemoryLayout<Any>.alignment
    
    if inlineStorageExpected {
      XCTAssertEqual(a.boxOrObjectID, nil)
    }
    else {
      XCTAssertNotEqual(a.boxOrObjectID, nil)
    }
  }

  func test_BufferOverflow() {
    var a = AnyValue(bufferOverflow)
    XCTAssert(a.storedType == type(of: bufferOverflow))
    XCTAssert(a[Type<BufferOverflow>()] == bufferOverflow)
    XCTAssert(a[unsafelyAssuming: Type<BufferOverflow>()] == bufferOverflow)
    
    let boxID = a.boxOrObjectID
    XCTAssertNotEqual(boxID, nil)
    
    a[Type<BufferOverflow>()].0 += 42
    XCTAssert(a[Type<BufferOverflow>()] == (43, 2, 3, 4))
    XCTAssertEqual(a.boxOrObjectID, boxID, "Unexpected reallocation of box")
  }

  func test_class() {
    let base = Base()
    var a = AnyValue(base)
    XCTAssert(a.storedType == type(of: base))
    XCTAssert(a[Type<Base>()] === base)
    XCTAssert(a[unsafelyAssuming: Type<Base>()] === base)
    XCTAssertEqual(a.boxOrObjectID, ObjectIdentifier(base))
    
    let otherBase = Base()
    XCTAssert(otherBase !== base)
    a[Type<Base>()] = otherBase
    XCTAssert(a[Type<Base>()] === otherBase)
    XCTAssertEqual(a.boxOrObjectID, ObjectIdentifier(otherBase))
  }
  
  func testUpcastedSource() {
    // test postconditions storing a derived class instance as base.
    let derived = Derived() as Base
    var a = AnyValue(derived)
    XCTAssert(a.storedType == Base.self)
    XCTAssert(a[Type<Base>()] === derived)

    // test postconditions when storing an existential.
    a = AnyValue(Double.pi as P)
    XCTAssert(a.storedType == P.self)
    XCTAssertEqual(a[Type<P>()] as? Double, .pi)
  }

  func test_store() {
    var a = AnyValue(3)
    XCTAssertEqual(a.boxOrObjectID, nil)
    a.store(4)
    XCTAssertEqual(a[Type<Int>()], 4, "store didn't update the stored int")
    XCTAssertEqual(a.boxOrObjectID, nil)
    
    var bufferOverflow = self.bufferOverflow
    a.store(bufferOverflow)
    let boxID = a.boxOrObjectID
    XCTAssertNotEqual(boxID, nil, "BufferOverflow is too big, but apparently is stored inline.")
    XCTAssert(
      a[Type<BufferOverflow>()] == bufferOverflow, "store didn't update stored BufferOverflow")

    bufferOverflow.0 += 43
    a.store(bufferOverflow)
    XCTAssert(
      a[Type<BufferOverflow>()] == bufferOverflow,
      "store didn't correctly update stored BufferOverflow")
    
    XCTAssertEqual(a.boxOrObjectID, boxID, "store didn't reuse the box")
    let b = a
    XCTAssertEqual(
      b.boxOrObjectID, a.boxOrObjectID,
      "copies of boxed values unexpectedly don't share a box using CoW")
    a = AnyValue(bufferOverflow)
    XCTAssertNotEqual(
      a.boxOrObjectID, boxID,
      "Using store is no better than assigning over it with a new Any?")
    withExtendedLifetime(b) {}
  }

  func staticType<T>(of _: T) -> Any.Type { T.self }
  
  func test_asAny() {
    var a = AnyValue(3)
    let aInt = a.asAny
    XCTAssert(staticType(of: aInt) == Any.self)
    XCTAssertEqual(aInt as? Int, 3)

    a.store(bufferOverflow)
    let aBufferOverflow = a.asAny
    XCTAssert(aBufferOverflow is BufferOverflow)
    if aBufferOverflow is BufferOverflow {
      XCTAssert(aBufferOverflow as! BufferOverflow == bufferOverflow)
    }

    let base = Base()
    a.store(base)
    let aBase = a.asAny
    XCTAssert(aBase as? Base === base)
  }

  func test_inlineValueSemantics() {
    var a = AnyValue(3)
    var b = a
    a[Type<Int>()] += 1
    XCTAssertEqual(a[Type<Int>()], 4)
    XCTAssertEqual(b[Type<Int>()], 3, "value semantics of stored Int not preserved by subscript.")
    b = a
    a.store(9)
    XCTAssertEqual(a[Type<Int>()], 9)
    XCTAssertEqual(b[Type<Int>()], 4, "value semantics of stored Int not preserved by store()")
  }

  func test_boxedValueSemantics() {
    var a = AnyValue(bufferOverflow)
    var b = a
    a[Type<BufferOverflow>()].0 += 41
    XCTAssertEqual(a[Type<BufferOverflow>()].0, 42)
    XCTAssertEqual(
      b[Type<BufferOverflow>()].0, 1, "value semantics of boxed value not preserved by subscript.")
    b = a
    a.store((4, 3, 2, 1))
    XCTAssert(a[Type<BufferOverflow>()] == (4, 3, 2, 1))
    XCTAssert(
      b[Type<BufferOverflow>()] == (42, 2, 3, 4),
      "value semantics of boxed value not preserved by store()")
  }

  static let allTests = [
    ("test_Int", test_Int),
    ("test_String", test_String),
    ("test_BufferOverflow", test_BufferOverflow),
    ("test_class", test_class),
    ("testUpcastedSource", testUpcastedSource),
    ("test_store", test_store),
    ("test_asAny", test_asAny),
    ("test_inlineValueSemantics", test_inlineValueSemantics),
    ("test_boxedValueSemantics", test_boxedValueSemantics),
  ]
}

// ==== Some Functions of which to inspect the disassembly === 

struct Int4 { var a, b, c, d: Int }

@inline(never)
func disassembleMe0(_ x: AnyValue) -> Int {
  x[unsafelyAssuming: Type<Int>()]
}

struct Int2 { var a, b: Int }

func disassembleMe1(_ x: AnyValue) -> Int2 {
  let a = x[unsafelyAssuming: Type<Int>()]
  var y = x
  y.store(Int2(a: 3, b: 4))
  let b = y[Type<Int2>()].a
  return Int2(a: a, b: b)
}

func disassembleMe1a(_ x: inout AnyValue)  {
  x[Type<Int4>()].a += 1
}


@inline(never)
func disassembleMe2(_ x: inout AnyValue) -> Int {
  x[unsafelyAssuming: Type<Int>()] += 1
  return x[Type<Int>()]
}

@inline(never)
func disassembleMe4(_ x: AnyValue) -> Int {
  x[Type<(Int,Int,Int,Int)>()].3
}

@inline(never)
func disassembleMe00(_ x: Any) -> Int {
  // Would be nice if this was efficient, but…
  (x as? Int).unsafelyUnwrapped
}


// Local Variables:
// fill-column: 100
// End: