[Compile Time Constant Extraction] Extract from extensions of types not declared in the same file/module (#70081)

Resolves rdar://118904022

Author: ApolloZhu

include/swift/ConstExtract/ConstExtractRequests.h

@@ -16,15 +16,16 @@
 #ifndef SWIFT_CONST_EXTRACT_REQUESTS_H
 #define SWIFT_CONST_EXTRACT_REQUESTS_H
 
-#include "swift/AST/SimpleRequest.h"
 #include "swift/AST/ASTTypeIDs.h"
 #include "swift/AST/ConstTypeInfo.h"
 #include "swift/AST/EvaluatorDependencies.h"
 #include "swift/AST/FileUnit.h"
 #include "swift/AST/Identifier.h"
 #include "swift/AST/NameLookup.h"
+#include "swift/AST/SimpleRequest.h"
 #include "swift/Basic/Statistic.h"
 #include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/TinyPtrVector.h"
 
 namespace swift {
@@ -35,9 +36,12 @@ class EnumDecl;
 
 /// Retrieve information about compile-time-known values
 class ConstantValueInfoRequest
-  : public SimpleRequest<ConstantValueInfoRequest,
-                         ConstValueTypeInfo(NominalTypeDecl *),
-                         RequestFlags::Cached> {
+    : public SimpleRequest<
+          ConstantValueInfoRequest,
+          ConstValueTypeInfo(
+              NominalTypeDecl *,
+              llvm::PointerUnion<const SourceFile *, ModuleDecl *>),
+          RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
 
@@ -46,7 +50,9 @@ class ConstantValueInfoRequest
 
   // Evaluation.
   ConstValueTypeInfo
-  evaluate(Evaluator &eval, NominalTypeDecl *nominal) const;
+  evaluate(Evaluator &eval, NominalTypeDecl *nominal,
+           llvm::PointerUnion<const SourceFile *, ModuleDecl *> extractionScope)
+      const;
 
 public:
   // Caching

include/swift/ConstExtract/ConstExtractTypeIDZone.def

@@ -16,5 +16,7 @@
 //===----------------------------------------------------------------------===//
 
 SWIFT_REQUEST(ConstExtract, ConstantValueInfoRequest,
-              ConstValueTypeInfo(NominalTypeDecl *), Cached,
-              NoLocationInfo)
+              ConstValueTypeInfo(
+                  NominalTypeDecl *,
+                  llvm::PointerUnion<const SourceFile *, ModuleDecl *>),
+              Cached, NoLocationInfo)

lib/ConstExtract/ConstExtract.cpp

@@ -53,14 +53,21 @@ class NominalTypeConformanceCollector : public ASTWalker {
   }
 
   PreWalkAction walkToDeclPre(Decl *D) override {
-    if (auto *NTD = llvm::dyn_cast<NominalTypeDecl>(D))
-      if (!isa<ProtocolDecl>(NTD))
+    auto *NTD = llvm::dyn_cast<NominalTypeDecl>(D);
+    if (!NTD)
+      if (auto *ETD = dyn_cast<ExtensionDecl>(D))
+        NTD = ETD->getExtendedNominal();
+    if (NTD)
+      if (!isa<ProtocolDecl>(NTD) && CheckedDecls.insert(NTD).second)
         for (auto &Protocol : NTD->getAllProtocols())
           if (Protocol->getAttrs().hasAttribute<ExtractConstantsFromMembersAttr>() ||
               Protocols.count(Protocol->getName().str().str()) != 0)
             ConformanceTypeDecls.push_back(NTD);
     return Action::Continue();
   }
+
+private:
+  std::unordered_set<NominalTypeDecl *> CheckedDecls;
 };
 
 std::string toFullyQualifiedTypeNameString(const swift::Type &Type) {
@@ -463,37 +470,52 @@ extractEnumCases(NominalTypeDecl *Decl) {
   return llvm::None;
 }
 
-ConstValueTypeInfo
-ConstantValueInfoRequest::evaluate(Evaluator &Evaluator,
-                                   NominalTypeDecl *Decl) const {
-  // Use 'getStoredProperties' to get lowered lazy and wrapped properties
-  auto StoredProperties = Decl->getStoredProperties();
-  std::unordered_set<VarDecl *> StoredPropertiesSet(StoredProperties.begin(),
-                                                    StoredProperties.end());
+ConstValueTypeInfo ConstantValueInfoRequest::evaluate(
+    Evaluator &Evaluator, NominalTypeDecl *Decl,
+    llvm::PointerUnion<const SourceFile *, ModuleDecl *> extractionScope)
+    const {
+
+  auto shouldExtract = [&](DeclContext *decl) {
+    if (auto SF = extractionScope.dyn_cast<const SourceFile *>())
+      return decl->getOutermostParentSourceFile() == SF;
+    return decl->getParentModule() == extractionScope.get<ModuleDecl *>();
+  };
 
   std::vector<ConstValueTypePropertyInfo> Properties;
-  for (auto Property : StoredProperties) {
-    Properties.push_back(extractTypePropertyInfo(Property));
-  }
+  llvm::Optional<std::vector<EnumElementDeclValue>> EnumCases;
 
-  for (auto Member : Decl->getMembers()) {
-    auto *VD = dyn_cast<VarDecl>(Member);
-    // Ignore plain stored properties collected above,
-    // instead gather up remaining static and computed properties.
-    if (!VD || StoredPropertiesSet.count(VD))
-      continue;
-    Properties.push_back(extractTypePropertyInfo(VD));
+  if (shouldExtract(Decl)) {
+    // Use 'getStoredProperties' to get lowered lazy and wrapped properties
+    auto StoredProperties = Decl->getStoredProperties();
+    std::unordered_set<VarDecl *> StoredPropertiesSet(StoredProperties.begin(),
+                                                      StoredProperties.end());
+    for (auto Property : StoredProperties) {
+      Properties.push_back(extractTypePropertyInfo(Property));
+    }
+
+    for (auto Member : Decl->getMembers()) {
+      auto *VD = dyn_cast<VarDecl>(Member);
+      // Ignore plain stored properties collected above,
+      // instead gather up remaining static and computed properties.
+      if (!VD || StoredPropertiesSet.count(VD))
+        continue;
+      Properties.push_back(extractTypePropertyInfo(VD));
+    }
+
+    EnumCases = extractEnumCases(Decl);
   }
 
   for (auto Extension: Decl->getExtensions()) {
-    for (auto Member : Extension->getMembers()) {
-      if (auto *VD = dyn_cast<VarDecl>(Member)) {
-        Properties.push_back(extractTypePropertyInfo(VD));
+    if (shouldExtract(Extension)) {
+      for (auto Member : Extension->getMembers()) {
+        if (auto *VD = dyn_cast<VarDecl>(Member)) {
+          Properties.push_back(extractTypePropertyInfo(VD));
+        }
       }
     }
   }
 
-  return ConstValueTypeInfo{Decl, Properties, extractEnumCases(Decl)};
+  return ConstValueTypeInfo{Decl, Properties, EnumCases};
 }
 
 std::vector<ConstValueTypeInfo>
@@ -507,7 +529,8 @@ gatherConstValuesForModule(const std::unordered_set<std::string> &Protocols,
   Module->walk(ConformanceCollector);
   for (auto *CD : ConformanceDecls)
     Result.emplace_back(evaluateOrDefault(CD->getASTContext().evaluator,
-                                          ConstantValueInfoRequest{CD}, {}));
+                                          ConstantValueInfoRequest{CD, Module},
+                                          {}));
   return Result;
 }
 
@@ -523,8 +546,8 @@ gatherConstValuesForPrimary(const std::unordered_set<std::string> &Protocols,
     D->walk(ConformanceCollector);
 
   for (auto *CD : ConformanceDecls)
-    Result.emplace_back(evaluateOrDefault(CD->getASTContext().evaluator,
-                                          ConstantValueInfoRequest{CD}, {}));
+    Result.emplace_back(evaluateOrDefault(
+        CD->getASTContext().evaluator, ConstantValueInfoRequest{CD, SF}, {}));
   return Result;
 }
 

test/ConstExtraction/ExtractFromExtension.swift

@@ -0,0 +1,33 @@
+// RUN: %empty-directory(%t)
+// RUN: echo "[MyProto]" > %t/protocols.json
+
+// RUN: %target-swift-frontend -typecheck -emit-const-values-path %t/ExtractFromExtension.swiftconstvalues -const-gather-protocols-file %t/protocols.json %S/Inputs/ProtocolConformances.swift -primary-file %s
+// RUN: cat %t/ExtractFromExtension.swiftconstvalues 2>&1 | %FileCheck %s
+
+extension MyType {
+    static let myValue = MyType("it is doable")
+}
+
+// CHECK:       "typeName": "ProtocolConformances.MyType",
+// CHECK:       "kind": "struct",
+// CHECK:       "conformances": [
+// CHECK-NEXT:    "ProtocolConformances.MyProto"
+// CHECK-NEXT:  ],
+// CHECK:       "properties": [
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "label": "myValue",
+// CHECK:           "type": "ProtocolConformances.MyType",
+// CHECK:           "valueKind": "InitCall",
+// CHECK-NEXT:      "value": {
+// CHECK-NEXT:        "type": "ProtocolConformances.MyType",
+// CHECK-NEXT:        "arguments": [
+// CHECK-NEXT:          {
+// CHECK-NEXT:            "label": "",
+// CHECK-NEXT:            "type": "Swift.String",
+// CHECK-NEXT:            "valueKind": "RawLiteral",
+// CHECK-NEXT:            "value": "it is doable"
+// CHECK-NEXT:          }
+// CHECK-NEXT:        ]
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  ]

test/ConstExtraction/Inputs/ProtocolConformances.swift

@@ -0,0 +1,15 @@
+public protocol MyProto { }
+
+public struct MyType {
+    private let value: String
+
+    public init(_ value: String) {
+        self.value = value
+    }
+    
+    private static let defaultValue = MyType("not me")
+}
+
+extension MyType: MyProto {
+    static let anotherValue = MyType("also not me")
+}