From a6a792c7407502c87eb7e20f9aa6c01ae023bedb Mon Sep 17 00:00:00 2001
From: Joe Groff <jgroff@apple.com>
Date: Mon, 26 Mar 2018 12:22:41 -0700
Subject: [PATCH] Add a `MemoryLayout<T>.offset(of:)` method for getting the
 offset of inline storage.

If a key path refers to inline storage of the root type, this produces the offset in bytes between a pointer to the root and a pointer to the projected storage. Otherwise, returns nil.
---
 stdlib/public/core/KeyPath.swift      | 21 +++++++++++++++
 stdlib/public/core/MemoryLayout.swift | 37 +++++++++++++++++++++++++++
 test/stdlib/KeyPath.swift             | 26 +++++++++++++++++++
 3 files changed, 84 insertions(+)

diff --git a/stdlib/public/core/KeyPath.swift b/stdlib/public/core/KeyPath.swift
index 73cf77745e33b..28845c9c2e17f 100644
--- a/stdlib/public/core/KeyPath.swift
+++ b/stdlib/public/core/KeyPath.swift
@@ -159,6 +159,27 @@ public class AnyKeyPath: Hashable, _AppendKeyPath {
     let base = UnsafeRawPointer(Builtin.projectTailElems(self, Int32.self))
     return try f(KeyPathBuffer(base: base))
   }
+
+  @_inlineable // FIXME(sil-serialize-all)
+  @_versioned // FIXME(sil-serialize-all)
+  internal var _storedInlineOffset: Int? {
+    return withBuffer {
+      var buffer = $0
+      var offset = 0
+      while true {
+        let (rawComponent, optNextType) = buffer.next()
+        switch rawComponent.header.kind {
+        case .struct:
+          offset += rawComponent._structOrClassOffset
+
+        case .class, .computed, .optionalChain, .optionalForce, .optionalWrap:
+          return .none
+        }
+
+        if optNextType == nil { return .some(offset) }
+      }
+    }
+  }
 }
 
 /// A partially type-erased key path, from a concrete root type to any
diff --git a/stdlib/public/core/MemoryLayout.swift b/stdlib/public/core/MemoryLayout.swift
index 3f09df4b33716..01982a21daf3f 100644
--- a/stdlib/public/core/MemoryLayout.swift
+++ b/stdlib/public/core/MemoryLayout.swift
@@ -162,4 +162,41 @@ extension MemoryLayout {
   public static func alignment(ofValue value: T) -> Int {
     return MemoryLayout.alignment
   }
+
+  /// Returns the offset of an inline stored property of `T` within the
+  /// in-memory representation of `T`.
+  ///
+  /// If the given `key` refers to inline storage within the
+  /// in-memory representation of `T`, and the storage is directly
+  /// addressable (meaning that accessing it does not need to trigger any
+  /// `didSet` or `willSet` accessors, perform any representation changes
+  /// such as bridging or closure reabstraction, or mask the value out of
+  /// overlapping storage as for packed bitfields), then the return value
+  /// is a distance in bytes that can be added to a pointer of type `T` to
+  /// get a pointer to the storage accessed by `key`. If the return value is
+  /// non-nil, then these formulations are equivalent:
+  ///
+  ///   var root: T, value: U
+  ///   var key: WritableKeyPath<T, U>
+  ///   // Mutation through the key path...
+  ///   root[keyPath: \.key] = value
+  ///   // ...is exactly equivalent to mutation through the offset pointer...
+  ///   withUnsafePointer(to: &root) {
+  ///     (UnsafeMutableRawPointer($0) + MemoryLayout<T>.offset(of: \.key))
+  ///       // ...which can be assumed to be bound to the target type
+  ///       .assumingMemoryBound(to: U.self).pointee = value
+  ///   }
+  ///
+  /// - Parameter key: A key path referring to storage that can be accessed
+  ///   through a value of type `T`.
+  /// - Returns: The offset in bytes from a pointer to a value of type `T`
+  ///   to a pointer to the storage referenced by `key`, or `nil` if no
+  ///   such offset is available for the storage referenced by `key`, such as
+  ///   because `key` is computed, has observers, requires reabstraction, or
+  ///   overlaps storage with other properties.
+  @_inlineable // FIXME(sil-serialize-all)
+  @_transparent
+  public static func offset(of key: PartialKeyPath<T>) -> Int? {
+    return key._storedInlineOffset
+  }
 }
diff --git a/test/stdlib/KeyPath.swift b/test/stdlib/KeyPath.swift
index 08b27cad7baae..072c72be0b9df 100644
--- a/test/stdlib/KeyPath.swift
+++ b/test/stdlib/KeyPath.swift
@@ -666,6 +666,32 @@ keyPath.test("subscripts") {
   expectEqual(base[keyPath: ints_be], (17 + 38).bigEndian)
 }
 
+struct NonOffsetableProperties {
+  // observers
+  var x: Int { didSet {} }
+  // reabstracted
+  var y: () -> ()
+  // computed
+  var z: Int { return 0 }
+}
+
+keyPath.test("offsets") {
+  let SLayout = MemoryLayout<S<Int>>.self
+  expectNotNil(SLayout.offset(of: \S<Int>.x))
+  expectNotNil(SLayout.offset(of: \S<Int>.y))
+  expectNotNil(SLayout.offset(of: \S<Int>.z))
+  expectNotNil(SLayout.offset(of: \S<Int>.p))
+  expectNotNil(SLayout.offset(of: \S<Int>.p.x))
+  expectNotNil(SLayout.offset(of: \S<Int>.p.y))
+  expectNotNil(SLayout.offset(of: \S<Int>.c))
+  expectNil(SLayout.offset(of: \S<Int>.c.x))
+
+  let NOPLayout = MemoryLayout<NonOffsetableProperties>.self
+  expectNil(NOPLayout.offset(of: \NonOffsetableProperties.x))
+  expectNil(NOPLayout.offset(of: \NonOffsetableProperties.y))
+  expectNil(NOPLayout.offset(of: \NonOffsetableProperties.z))
+}
+
 // SR-6096
 
 protocol Protocol6096 {}