[Compile Time Constant Extraction] Extract KeyPath Expressions

Author: Venkatesh Sriram

include/swift/AST/ConstTypeInfo.h

@@ -37,6 +37,7 @@ class CompileTimeValue {
     Tuple,
     Enum,
     Type,
+    KeyPath,
     Runtime
   };
 
@@ -199,6 +200,34 @@ class TypeValue : public CompileTimeValue {
   swift::Type Type;
 };
 
+/// A representation of a Keypath
+class KeyPathValue : public CompileTimeValue {
+public:
+  struct Component {
+    std::string Label;
+    swift::Type Type;
+  };
+  KeyPathValue(std::string Path,
+               swift::Type RootType,
+               std::vector<Component> Components)
+  : CompileTimeValue(ValueKind::KeyPath), Path(Path), RootType(RootType), Components(Components) {}
+
+  std::string getPath() const { return Path; }
+  swift::Type getRootType() const { return RootType; }
+  std::vector<Component> getComponents() const {
+    return Components;
+  }
+
+  static bool classof(const CompileTimeValue *T) {
+    return T->getKind() == ValueKind::KeyPath;
+  }
+
+private:
+  std::string Path;
+  swift::Type RootType;
+  std::vector<Component> Components;
+};
+
 /// A representation of an arbitrary value that does not fall under
 /// any of the above categories.
 class RuntimeValue : public CompileTimeValue {

lib/ConstExtract/ConstExtract.cpp

@@ -376,6 +376,33 @@ static std::shared_ptr<CompileTimeValue> extractCompileTimeValue(Expr *expr) {
       }
     } break;
 
+    case ExprKind::KeyPath: {
+        auto keyPathExpr = cast<KeyPathExpr>(expr);
+
+        auto rootType = keyPathExpr->getRootType();
+        std::vector<KeyPathValue::Component> components;
+
+        for (auto component: keyPathExpr->getComponents()) {
+            if (component.isResolved()) {
+                auto declRef = component.getDeclRef();
+                auto identifier = declRef.getDecl()->getBaseIdentifier().str();
+                auto type = component.getComponentType()->getRValueType();
+                components.push_back({identifier.str(), type});
+            }
+        }
+
+        std::string path = "";
+        auto numberOfComponents = static_cast<int>(components.size());
+        for (int i = 0; i < numberOfComponents; i++) {
+            if (i != 0) {
+                path += ".";
+            }
+            path += components[i].Label;
+        }
+
+        return std::make_shared<KeyPathValue>(path, rootType, components);
+    }
+
     case ExprKind::InjectIntoOptional: {
       auto injectIntoOptionalExpr = cast<InjectIntoOptionalExpr>(expr);
       return extractCompileTimeValue(injectIntoOptionalExpr->getSubExpr());
@@ -703,6 +730,25 @@ void writeValue(llvm::json::OStream &JSON,
     break;
   }
 
+  case CompileTimeValue::KeyPath: {
+      auto keyPathValue = cast<KeyPathValue>(value);
+      JSON.attribute("valueKind", "KeyPath");
+      JSON.attributeObject("value", [&]() {
+        JSON.attribute("path", keyPathValue->getPath());
+        JSON.attribute("rootType", toFullyQualifiedTypeNameString(keyPathValue->getRootType()));
+        JSON.attributeArray("components", [&] {
+          auto components = keyPathValue->getComponents();
+          for (auto c : components) {
+            JSON.object([&] {
+              JSON.attribute("label", c.Label);
+              JSON.attribute("type", toFullyQualifiedTypeNameString(c.Type));
+            });
+          }
+        });
+      });
+      break;
+  }
+
   case CompileTimeValue::ValueKind::Runtime: {
     JSON.attribute("valueKind", "Runtime");
     break;

test/ConstExtraction/ExtractKeyPaths.swift

@@ -0,0 +1,75 @@
+// RUN: %empty-directory(%t)
+// RUN: echo "[MyProto]" > %t/protocols.json
+
+// RUN: %target-swift-frontend -typecheck -emit-const-values-path %t/ExtractKeyPaths.swiftconstvalues -const-gather-protocols-file %t/protocols.json -primary-file %s
+// RUN: cat %t/ExtractKeyPaths.swiftconstvalues 2>&1 | %FileCheck %s
+
+protocol MyProto {}
+
+public struct MyType {
+    var nested: NestedOne
+
+    struct NestedOne {
+        var foo: NestedTwo
+    }
+
+    struct NestedTwo {
+        var bar: NestedThree
+    }
+
+    struct NestedThree {
+        var baz: String
+    }
+}
+
+public struct KeyPaths: MyProto {
+    static let nestedVariable = \MyType.nested.foo.bar.baz
+}
+
+// CHECK: [
+// CHECK-NEXT:   {
+// CHECK-NEXT:     "typeName": "ExtractKeyPaths.KeyPaths",
+// CHECK-NEXT:     "mangledTypeName": "15ExtractKeyPaths0bC0V",
+// CHECK-NEXT:     "kind": "struct",
+// CHECK-NEXT:     "file": "{{.*}}test{{/|\\\\}}ConstExtraction{{/|\\\\}}ExtractKeyPaths.swift",
+// CHECK-NEXT:    "line": 25,
+// CHECK-NEXT:    "conformances": [
+// CHECK-NEXT:      "ExtractKeyPaths.MyProto"
+// CHECK-NEXT:     ],
+// CHECK-NEXT:     "associatedTypeAliases": [],
+// CHECK-NEXT:    "properties": [
+// CHECK-NEXT:       {
+// CHECK-NEXT:        "label": "nestedVariable",
+// CHECK-NEXT:        "type": "Swift.WritableKeyPath<ExtractKeyPaths.MyType, Swift.String>",
+// CHECK-NEXT:        "mangledTypeName": "n/a - deprecated",
+// CHECK-NEXT:        "isStatic": "true",
+// CHECK-NEXT:        "isComputed": "false",
+// CHECK-NEXT:        "file": "{{.*}}test{{/|\\\\}}ConstExtraction{{/|\\\\}}ExtractKeyPaths.swift",
+// CHECK-NEXT:        "line": 26,
+// CHECK-NEXT:        "valueKind": "KeyPath",
+// CHECK-NEXT:        "value": {
+// CHECK-NEXT:         "path": "nested.foo.bar.baz",
+// CHECK-NEXT:         "rootType": "ExtractKeyPaths.MyType",
+// CHECK-NEXT:         "components": [
+// CHECK-NEXT:             {
+// CHECK-NEXT:               "label": "nested",
+// CHECK-NEXT:               "type": "ExtractKeyPaths.MyType.NestedOne"
+// CHECK-NEXT:             },
+// CHECK-NEXT:             {
+// CHECK-NEXT:               "label": "foo",
+// CHECK-NEXT:               "type": "ExtractKeyPaths.MyType.NestedTwo"
+// CHECK-NEXT:            },
+// CHECK-NEXT:            {
+// CHECK-NEXT:              "label": "bar",
+// CHECK-NEXT:              "type": "ExtractKeyPaths.MyType.NestedThree"
+// CHECK-NEXT:             },
+// CHECK-NEXT:             {
+// CHECK-NEXT:               "label": "baz",
+// CHECK-NEXT:               "type": "Swift.String"
+// CHECK-NEXT:             }
+// CHECK-NEXT:           ]
+// CHECK-NEXT:         }
+// CHECK-NEXT:       }
+// CHECK-NEXT:     ]
+// CHECK-NEXT:   }
+// CHECK-NEXT: ]