Associate PackageUnit with ModuleDecl

* Weakly reference ModuleDecl from PackageUnit
* Add PackageUnit decl context getter and use it for a package AccessScope
* Return module decl referenced by PackageUnit in getModuleScopeContext and getParentModule
* Handle package acl in access scope checkers
* Fix tests
Resolves rdar://104987295, rdar://105187216

Author: elsh

include/swift/AST/AccessScope.h

@@ -66,10 +66,6 @@ class AccessScope {
   bool operator==(AccessScope RHS) const { return Value == RHS.Value; }
   bool operator!=(AccessScope RHS) const { return !(*this == RHS); }
   bool hasEqualDeclContextWith(AccessScope RHS) const {
-    if (isPublic())
-      return RHS.isPublic();
-    if (isPackage())
-      return RHS.isPackage();
     return getDeclContext() == RHS.getDeclContext();
   }
 
@@ -81,26 +77,23 @@ class AccessScope {
   }
   bool isFileScope() const;
   bool isInternal() const;
-  bool isPackage() const {
-    return !Value.getPointer() && Value.getInt() == AccessLimitKind::Package;
-  }
+  bool isPackage() const;
 
   /// Returns true if the context of this (use site) is more restrictive than
   /// the argument context (decl site). This function does _not_ check the
   /// restrictiveness of the access level between this and the argument. \see
   /// AccessScope::isInContext
   bool isChildOf(AccessScope AS) const {
-    if (isInContext()) {
+    if (isPackage())
+      return AS.isPublic();
+    else if (isInContext()) {
       if (AS.isInContext())
         return allowsPrivateAccess(getDeclContext(), AS.getDeclContext());
       else
-        return AS.isPackage() || AS.isPublic();
-    }
-    if (isPackage())
-      return AS.isPublic();
-    // If this is public, it can't be less than access level of AS
-    // so return false
-    return false;
+        return AS.isPublic();
+    } else // If this is public, it can't be less than access level of AS so
+           // return false
+      return false;
   }
 
   /// Result depends on whether it's called at a use site or a decl site:

include/swift/AST/DeclContext.h

@@ -512,9 +512,11 @@ class alignas(1 << DeclContextAlignInBits) DeclContext
     return false;
   }
 
-  /// Returns the package context of the parent module.
+  /// Returns the package unit of this context.
+  /// \p lookupIfNotCurrent If the current decl context is not PackageUnit, look
+  /// it up via parent module
   LLVM_READONLY
-  PackageUnit *getParentModulePackage() const;
+  PackageUnit *getPackageContext(bool lookupIfNotCurrent = false) const;
 
   /// Returns the module context that contains this context.
   LLVM_READONLY

include/swift/AST/DiagnosticsSema.def

@@ -1679,6 +1679,7 @@ ERROR(access_control_open_bad_decl,none,
 WARNING(access_control_non_objc_open_member,none,
       "non-'@objc' %0 in extensions cannot be overridden; use 'public' instead",
       (DescriptiveDeclKind))
+ERROR(access_control_requires_package_name, none, "decl has a package access level but no -package-name was passed", ())
 
 ERROR(invalid_decl_attribute,none,
       "'%0' attribute cannot be applied to this declaration", (DeclAttribute))

include/swift/AST/Module.h

@@ -159,33 +159,48 @@ class OverlayFile;
 /// location.
 class ModuleSourceFileLocationMap;
 
-/// Package unit used to allow grouping of modules by a package name.
-/// ModuleDecl can set PackageUnit as a parent in its DeclContext
-/// See \c ModuleDecl
+/// A unit that allows grouping of modules by a package name.
+/// It's set as a property in ModuleDecl, instead of as its parent decl context,
+/// since otherwise it will break the existing decl context lookups that assume
+/// ModuleDecl as the top level context. See \c ModuleDecl
 class PackageUnit: public DeclContext {
-
+  /// Identifies this package and used for the equality check
   Identifier PackageName;
-
-  PackageUnit(Identifier name);
+  /// Weakly reference ModuleDecl that points to this package.
+  /// Instead of having multiple ModuleDecls pointing to one PackageUnit, we
+  /// create and set one PackageUnit per ModuleDecl, to make it easier to look
+  /// up the module pointing to this package; this look up is needed in existing
+  /// DeclContext functions, e.g. \c DeclContext::getModuleScopeContext,
+  /// and \c DeclContext::getParentModule
+  ModuleDecl &SourceModule;
+
+  PackageUnit(Identifier name, ModuleDecl &src)
+      : DeclContext(DeclContextKind::Package, nullptr), PackageName(name),
+        SourceModule(src) {}
 
 public:
-  static PackageUnit *
-  create(Identifier name, ASTContext &ctx) {
-    return new (ctx) PackageUnit(name);
+  static PackageUnit *create(Identifier name, ModuleDecl &src,
+                             ASTContext &ctx) {
+    return new (ctx) PackageUnit(name, src);
   }
 
   static bool classof(const DeclContext *DC) {
     return DC->getContextKind() == DeclContextKind::Package;
   }
 
-  static bool classof(const PackageUnit *PU) {
-    // FIXME: add a correct check
-    return true;
-  }
+  static bool classof(const PackageUnit *PU) { return true; }
 
   Identifier getName() const {
     return PackageName;
   }
+
+  ModuleDecl &getSourceModule() { return SourceModule; }
+
+  /// Equality check via package name instead of pointer comparison.
+  /// Returns false if the name is empty.
+  bool isSamePackageAs(PackageUnit *other) {
+    return !(getName().empty()) && getName() == other->getName();
+  }
 };
 
 /// The minimum unit of compilation.
@@ -202,8 +217,10 @@ class ModuleDecl
   /// The ABI name of the module, if it differs from the module name.
   mutable Identifier ModuleABIName;
 
-  /// The name of the package this module belongs to
-  mutable Identifier PackageName;
+  /// A package this module belongs to. It's set as a property instead of a
+  /// parent decl context; otherwise it will break the existing decl context
+  /// lookups that assume ModuleDecl as the top level context.
+  PackageUnit *Package;
 
   /// Module name to use when referenced in clients module interfaces.
   mutable Identifier ExportAsName;
@@ -440,18 +457,21 @@ class ModuleDecl
     ModuleABIName = name;
   }
 
-  /// Get the package name of the module
-  Identifier getPackageName() const { return PackageName; }
-
-  /// Set the name of the package this module belongs to
-  void setPackageName(Identifier name) {
-    PackageName = name;
-    // TODO: uncomment when PackageUnit gets passed to constructor
-//    PackageUnit *pkgUnit = PackageUnit::create(name, getASTContext());
-//    DeclContext newContext = DeclContext(DeclContextKind::Module, pkgUnit);
-//    setDeclContext(&newContext);
+  /// Get the package name of this module
+  /// FIXME: remove this and bump module version rdar://104723918
+  Identifier getPackageName() const {
+    if (auto pkg = getPackage())
+      return pkg->getName();
+    return Identifier();
   }
 
+  /// Get the package associated with this module
+  PackageUnit *getPackage() const { return Package; }
+
+  /// Set the package this module is associated with
+  /// FIXME: rename this with setPackage(name) rdar://104723918
+  void setPackageName(Identifier name);
+
   Identifier getExportAsName() const { return ExportAsName; }
 
   void setExportAsName(Identifier name) {

lib/AST/ASTVerifier.cpp

@@ -963,12 +963,6 @@ class Verifier : public ASTWalker {
       if (D->hasAccess()) {
         PrettyStackTraceDecl debugStack("verifying access", D);
         if (!D->getASTContext().isAccessControlDisabled()) {
-          if (D->getFormalAccessScope().isPackage() &&
-              D->getFormalAccess() < AccessLevel::Package) {
-            Out << "non-package decl has no formal access scope\n";
-            D->dump(Out);
-            abort();
-          }
           if (D->getFormalAccessScope().isPublic() &&
               D->getFormalAccess() < AccessLevel::Public) {
             Out << "non-public decl has no formal access scope\n";

lib/AST/Decl.cpp

@@ -3773,6 +3773,10 @@ static AccessLevel getAdjustedFormalAccess(const ValueDecl *VD,
   }
 
   if (useDC) {
+    // If the use site decl context is PackageUnit, just return
+    // the access level that's passed in
+    if (auto usePkg = useDC->getPackageContext())
+      return access;
     // Check whether we need to modify the access level based on
     // @testable/@_private import attributes.
     auto *useSF = dyn_cast<SourceFile>(useDC->getModuleScopeContext());
@@ -3925,8 +3929,14 @@ getAccessScopeForFormalAccess(const ValueDecl *VD,
                                      : AccessLimitKind::None);
   case AccessLevel::Internal:
     return AccessScope(resultDC->getParentModule());
-  case AccessLevel::Package:
-    return AccessScope::getPackage();
+  case AccessLevel::Package: {
+    auto pkg = resultDC->getPackageContext(true);
+    if (!pkg) {
+      auto &d = VD->getASTContext().Diags;
+      d.diagnose(VD->getLoc(), diag::access_control_requires_package_name);
+    }
+    return AccessScope(pkg);
+  }
   case AccessLevel::Public:
   case AccessLevel::Open:
     return AccessScope::getPublic();
@@ -3965,10 +3975,6 @@ static bool checkAccessUsingAccessScopes(const DeclContext *useDC,
   // useDC is null only when caller wants to skip non-public type checks.
   if (!useDC) return true;
 
-  // Check package access; accessing package decl should not be allowed if package names are different
-  if (accessScope.isPackage())
-    return VD->getDeclContext()->getParentModule()->getPackageName() == useDC->getParentModule()->getPackageName();
-
   // Check SPI access
   if (!VD->isSPI()) return true;
   auto useSF = dyn_cast<SourceFile>(useDC->getModuleScopeContext());
@@ -4083,7 +4089,11 @@ static bool checkAccess(const DeclContext *useDC, const ValueDecl *VD,
     auto *useSF = dyn_cast<SourceFile>(useFile);
     return useSF && useSF->hasTestableOrPrivateImport(access, sourceModule);
   }
-  case AccessLevel::Package:
+  case AccessLevel::Package: {
+    auto srcPkg = sourceDC->getPackageContext(true);
+    auto usePkg = useDC->getPackageContext(true);
+    return usePkg->isSamePackageAs(srcPkg);
+  }
   case AccessLevel::Public:
   case AccessLevel::Open:
     return true;

lib/AST/DeclContext.cpp

@@ -274,15 +274,32 @@ DeclContext *DeclContext::getParentForLookup() const {
   return getParent();
 }
 
-PackageUnit *DeclContext::getParentModulePackage() const {
-  auto parentModule = getParentModule();
-  auto pkg = parentModule->getParent();
-  if (pkg)
-    return const_cast<PackageUnit *>(cast<PackageUnit>(pkg));
+PackageUnit *DeclContext::getPackageContext(bool lookupIfNotCurrent) const {
+  // Current decl context might be PackageUnit, which is not in the
+  // DeclContext hierarchy, so check for it first
+  if (ParentAndKind.getInt() == ASTHierarchy::Package) {
+    return const_cast<PackageUnit *>(cast<PackageUnit>(this));
+  }
+
+  // If the current context is not PackageUnit, look it up via
+  // the parent module if needed
+  if (lookupIfNotCurrent) {
+    auto mdecl = getParentModule();
+    return mdecl->getPackage();
+  }
   return nullptr;
 }
 
 ModuleDecl *DeclContext::getParentModule() const {
+  // If the current context is PackageUnit, return the module
+  // decl context pointing to the current context. This check
+  // needs to be done first as PackageUnit is not in the DeclContext
+  // hierarchy.
+  if (auto pkg = getPackageContext()) {
+    auto &mdecl = pkg->getSourceModule();
+    return &mdecl;
+  }
+
   const DeclContext *DC = this;
   while (!DC->isModuleContext())
     DC = DC->getParent();
@@ -292,7 +309,7 @@ ModuleDecl *DeclContext::getParentModule() const {
 SourceFile *DeclContext::getParentSourceFile() const {
   const DeclContext *DC = this;
   SourceLoc loc;
-  while (!DC->isModuleScopeContext()) {
+  while (DC && !DC->isModuleScopeContext()) {
     // If we don't have a source location yet, try to grab one from this
     // context.
     if (loc.isInvalid()) {
@@ -323,6 +340,9 @@ SourceFile *DeclContext::getParentSourceFile() const {
     DC = DC->getParent();
   }
 
+  if (!DC)
+    return nullptr;
+
   auto fallbackSF = const_cast<SourceFile *>(dyn_cast<SourceFile>(DC));
   if (auto module = DC->getParentModule()) {
     if (auto sf = module->getSourceFileContainingLocation(loc))
@@ -345,15 +365,24 @@ SourceFile *DeclContext::getOutermostParentSourceFile() const {
 }
 
 DeclContext *DeclContext::getModuleScopeContext() const {
-  auto DC = const_cast<DeclContext*>(this);
+  // If the current context is PackageUnit, return the module
+  // decl context pointing to the current context. This check
+  // needs to be done first as PackageUnit is not in the DeclContext
+  // hierarchy.
+  if (auto pkg = getPackageContext()) {
+    auto &mdecl = pkg->getSourceModule();
+    return &mdecl;
+  }
 
+  auto DC = const_cast<DeclContext *>(this);
   while (true) {
     if (DC->ParentAndKind.getInt() == ASTHierarchy::FileUnit)
       return DC;
     if (auto NextDC = DC->getParent()) {
       DC = NextDC;
     } else {
-      assert(isa<ModuleDecl>(DC->getAsDecl()));
+      assert(DC && isa<ModuleDecl>(DC->getAsDecl()) &&
+             "no module decl context found!");
       return DC;
     }
   }
@@ -1235,6 +1264,11 @@ bool AccessScope::isInternal() const {
   return DC && isa<ModuleDecl>(DC);
 }
 
+bool AccessScope::isPackage() const {
+  auto DC = getDeclContext();
+  return DC && isa<PackageUnit>(DC);
+}
+
 AccessLevel AccessScope::accessLevelForDiagnostics() const {
   if (isPublic())
     return AccessLevel::Public;
@@ -1254,6 +1288,14 @@ bool AccessScope::allowsPrivateAccess(const DeclContext *useDC, const DeclContex
   if (useDC->isChildContextOf(sourceDC))
     return true;
 
+  // If the decl site has package acl, but the use site
+  // has internal or less acl, check if it belongs to
+  // the same package as the decl site's to allow access.
+  if (auto srcPkg = sourceDC->getPackageContext()) {
+    if (auto usePkg = sourceDC->getPackageContext(true)) {
+      return usePkg->isSamePackageAs(srcPkg);
+    }
+  }
   // Do not allow access if the sourceDC is in a different file
   auto useSF = useDC->getParentSourceFile();
   if (useSF != sourceDC->getParentSourceFile())

lib/AST/Module.cpp

@@ -543,11 +543,6 @@ void SourceLookupCache::invalidate() {
   (void)SameSizeSmallVector{std::move(AllVisibleValues)};
 }
 
-PackageUnit::PackageUnit(Identifier name)
-  : DeclContext(DeclContextKind::Package, nullptr) {
-    PackageName = name;
-}
-
 //===----------------------------------------------------------------------===//
 // Module Implementation
 //===----------------------------------------------------------------------===//
@@ -3068,6 +3063,10 @@ SourceFile::getImportAccessLevel(const ModuleDecl *targetModule) const {
   return restrictiveImport;
 }
 
+void ModuleDecl::setPackageName(Identifier name) {
+  Package = PackageUnit::create(name, *this, getASTContext());
+}
+
 bool ModuleDecl::isImportedImplementationOnly(const ModuleDecl *module) const {
   if (module == this) return false;
 

lib/Sema/TypeCheckAccess.cpp

@@ -2286,7 +2286,7 @@ static void checkExtensionGenericParamAccess(const ExtensionDecl *ED) {
     desiredAccessScope = AccessScope(ED->getModuleContext());
     break;
   case AccessLevel::Package:
-    desiredAccessScope = AccessScope::getPackage();
+    desiredAccessScope = AccessScope(ED->getPackageContext(true));
     break;
   case AccessLevel::Public:
   case AccessLevel::Open:

test/Sema/accessibility.swift

@@ -1,4 +1,4 @@
-// RUN: %target-typecheck-verify-swift -enable-objc-interop -disable-objc-attr-requires-foundation-module -swift-version 5
+// RUN: %target-typecheck-verify-swift -enable-objc-interop -disable-objc-attr-requires-foundation-module -swift-version 5 -package-name mylib
 
 /// Test structs with protocols and extensions
 public protocol PublicProto {
@@ -1021,8 +1021,8 @@ public struct PublicWithInternalSettersConformPublic : PublicMutationOperations
 }
 
 public struct PublicWithPackageSettersConformPublic : PublicMutationOperations {
-  public package(set) var size = 0 // FIXME: rdar://104987295 this should error
-  public package(set) subscript (_: Int) -> Int { // FIXME: rdar://104987295 this should error
+  public package(set) var size = 0 // expected-error {{setter for property 'size' must be declared public because it matches a requirement in public protocol 'PublicMutationOperations'}} {{none}} expected-note {{mark the property as 'public' to satisfy the requirement}} {{10-23=}}
+  public package(set) subscript (_: Int) -> Int { // expected-error {{subscript setter must be declared public because it matches a requirement in public protocol 'PublicMutationOperations'}} {{none}} expected-note {{mark the subscript as 'public' to satisfy the requirement}} {{10-23=}}
     get { return 42 }
     set {}
   }