* Hook up PackageUnit to ModuleDecl

* Use PackageUnit decl context for a package AccessScope
* Return the module decl weakly referenced by PackageUnit in module decl context getters
* Handle package in access scope checkers
* Fix tests

Resolves rdar://104987295

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();
+        return 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;
+    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,9 @@ class alignas(1 << DeclContextAlignInBits) DeclContext
     return false;
   }
 
-  /// Returns the package context of the parent module.
+  /// Returns the package unit of this context.
   LLVM_READONLY
-  PackageUnit *getParentModulePackage() const;
+  PackageUnit *getPackageContext(bool lookupIfNotCurrent = false) const;
 
   /// Returns the module context that contains this context.
   LLVM_READONLY

include/swift/AST/Module.h

@@ -165,27 +165,36 @@ class ModuleSourceFileLocationMap;
 class PackageUnit: public DeclContext {
 
   Identifier PackageName;
+  ModuleDecl &SourceModule;
 
-  PackageUnit(Identifier name);
+  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);
+  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;
   }
 
   Identifier getName() const {
     return PackageName;
   }
+
+  ModuleDecl &getSourceModule() {
+    return SourceModule;
+  }
+
+  bool isSamePackageAs(PackageUnit *other) {
+    return getName() == other->getName();
+  }
 };
 
 /// The minimum unit of compilation.
@@ -202,8 +211,8 @@ 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;
+  /// The package this module belongs to
+  PackageUnit *Package;
 
   /// Module name to use when referenced in clients module interfaces.
   mutable Identifier ExportAsName;
@@ -440,18 +449,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,8 @@ static AccessLevel getAdjustedFormalAccess(const ValueDecl *VD,
   }
 
   if (useDC) {
+    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 +3927,11 @@ 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);
+    assert(pkg && "decl has a package access level but there's no associated package; pass -package-name");
+    return AccessScope(pkg);
+  }
   case AccessLevel::Public:
   case AccessLevel::Open:
     return AccessScope::getPublic();
@@ -3965,10 +3970,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 +4084,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,24 @@ 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));
-  return nullptr;
+PackageUnit *DeclContext::getPackageContext(bool lookupIfNotCurrent) const {
+  if (ParentAndKind.getInt() == ASTHierarchy::Package) {
+    return const_cast<PackageUnit *>(cast<PackageUnit>(this));
+  }
+
+  if (lookupIfNotCurrent) {
+    auto mdecl = getParentModule();
+    return mdecl->getPackage();
+  }
+   return nullptr;
 }
 
 ModuleDecl *DeclContext::getParentModule() const {
+  if (auto pkg = getPackageContext()) {
+    auto &mdecl = pkg->getSourceModule();
+    return &mdecl;
+  }
+
   const DeclContext *DC = this;
   while (!DC->isModuleContext())
     DC = DC->getParent();
@@ -292,7 +301,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 +332,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 +357,19 @@ SourceFile *DeclContext::getOutermostParentSourceFile() const {
 }
 
 DeclContext *DeclContext::getModuleScopeContext() const {
-  auto DC = const_cast<DeclContext*>(this);
+  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 +1251,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 +1275,11 @@ bool AccessScope::allowsPrivateAccess(const DeclContext *useDC, const DeclContex
   if (useDC->isChildContextOf(sourceDC))
     return true;
 
+  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,11 @@ SourceFile::getImportAccessLevel(const ModuleDecl *targetModule) const {
   return restrictiveImport;
 }
 
+void ModuleDecl::setPackageName(Identifier name) {
+  PackageName = 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 {}
   }