ModuleInterface: Avoid printing @_spi enum elements in public swiftinterfaces

include/swift/AST/Decl.h

@@ -7479,7 +7479,16 @@ class EnumCaseDecl final : public Decl,
             Bits.EnumCaseDecl.NumElements};
   }
   SourceRange getSourceRange() const;
-
+
+  /// Returns the first of the member elements or null if there are no elements.
+  /// The attributes written with an EnumCaseDecl will be attached to each of
+  /// the elements instead so inspecting the attributes of the first element is
+  /// often useful.
+  EnumElementDecl *getFirstElement() const {
+    auto elements = getElements();
+    return elements.empty() ? nullptr : elements.front();
+  }
+
   static bool classof(const Decl *D) {
     return D->getKind() == DeclKind::EnumCase;
   }

include/swift/AST/DiagnosticsSema.def

@@ -1970,6 +1970,9 @@ ERROR(spi_attribute_on_protocol_requirement,none,
 ERROR(spi_attribute_on_frozen_stored_properties,none,
       "stored property %0 cannot be declared '@_spi' in a '@frozen' struct",
       (DeclName))
+ERROR(spi_attribute_on_frozen_enum_case,none,
+      "%0 %1 cannot be declared '@_spi' in a '@frozen' enum",
+      (DescriptiveDeclKind, DeclName))
 WARNING(spi_attribute_on_import_of_public_module,none,
         "'@_spi' import of %0 will not include any SPI symbols; "
         "%0 was built from the public interface at %1",

lib/AST/ASTPrinter.cpp

@@ -276,6 +276,18 @@ PrintOptions PrintOptions::printSwiftInterfaceFile(ModuleDecl *ModuleToPrint,
           return false;
       }
 
+      // Skip enum cases containing enum elements we wouldn't print.
+      if (auto *ECD = dyn_cast<EnumCaseDecl>(D)) {
+        if (auto *element = ECD->getFirstElement()) {
+          // Enum elements are usually not printed, so we have to override the
+          // print option controlling that.
+          PrintOptions optionsCopy = options;
+          optionsCopy.ExplodeEnumCaseDecls = true;
+          if (!shouldPrint(element, optionsCopy))
+            return false;
+        }
+      }
+
       return ShouldPrintChecker::shouldPrint(D, options);
     }
   };
@@ -4204,14 +4216,14 @@ void PrintAST::printEnumElement(EnumElementDecl *elt) {
 }
 
 void PrintAST::visitEnumCaseDecl(EnumCaseDecl *decl) {
-  auto elems = decl->getElements();
-  if (!elems.empty()) {
+  if (auto *element = decl->getFirstElement()) {
     // Documentation comments over the case are attached to the enum elements.
-    printDocumentationComment(elems[0]);
-    printAttributes(elems[0]);
+    printDocumentationComment(element);
+    printAttributes(element);
   }
   Printer.printIntroducerKeyword("case", Options, " ");
 
+  auto elems = decl->getElements();
   llvm::interleave(elems.begin(), elems.end(),
     [&](EnumElementDecl *elt) {
       printEnumElement(elt);

lib/IDE/SyntaxModel.cpp

@@ -1077,8 +1077,8 @@ ASTWalker::PreWalkAction ModelASTWalker::walkToDeclPre(Decl *D) {
 
     // We need to handle the special case where attributes semantically
     // attach to enum element decls while syntactically locate before enum case decl.
-    if (!EnumCaseD->getElements().empty()) {
-      if (!handleAttrs(EnumCaseD->getElements().front()->getAttrs()))
+    if (auto *element = EnumCaseD->getFirstElement()) {
+      if (!handleAttrs(element->getAttrs()))
         return Action::SkipChildren();
     }
     if (pushStructureNode(SN, D)) {

lib/Sema/TypeCheckAttr.cpp

@@ -1150,6 +1150,17 @@ void AttributeChecker::visitSPIAccessControlAttr(SPIAccessControlAttr *attr) {
         }
       }
     }
+
+    // Forbid enum elements marked SPI in frozen types.
+    if (auto elt = dyn_cast<EnumElementDecl>(VD)) {
+      if (auto ED = dyn_cast<EnumDecl>(D->getDeclContext())) {
+        if (ED->getAttrs().hasAttribute<FrozenAttr>(/*allowInvalid*/ true) &&
+            !ED->isSPI()) {
+          diagnoseAndRemoveAttr(attr, diag::spi_attribute_on_frozen_enum_case,
+                                VD->getDescriptiveKind(), VD->getName());
+        }
+      }
+    }
   }
 
   if (auto ID = dyn_cast<ImportDecl>(D)) {

lib/Sema/TypeCheckAvailability.cpp

@@ -1609,9 +1609,8 @@ abstractSyntaxDeclForAvailableAttribute(const Decl *ConcreteSyntaxDecl) {
   } else if (auto *ECD = dyn_cast<EnumCaseDecl>(ConcreteSyntaxDecl)) {
     // Similar to the PatternBindingDecl case above, we return the
     // first EnumElementDecl.
-    ArrayRef<EnumElementDecl *> Elems = ECD->getElements();
-    if (!Elems.empty()) {
-      return Elems.front();
+    if (auto *Elem = ECD->getFirstElement()) {
+      return Elem;
     }
   }
 

test/SPI/Inputs/spi_helper.swift

@@ -89,3 +89,8 @@ public struct PublicStruct {
 @_spi(HelperSPI) public struct IntLike: ExpressibleByIntegerLiteral, Equatable {
   @_spi(HelperSPI) public init(integerLiteral: Int) {}
 }
+
+public enum PublicEnum {
+  case publicCase
+  @_spi(HelperSPI) case spiCase
+}

test/SPI/export_spi_from_spi_module.swift

@@ -11,15 +11,32 @@
   public init() {}
 }
 @_spi(S) public protocol SPIProtocol {}
+public enum PublicEnum {
+  case publicCase
+  @_spi(S) case spiCase
+}
 
 public func useOfSPITypeOk(_ p0: SPIProtocol, p1: SPIClass) -> SPIClass { fatalError() }
 
 @inlinable
-func inlinable() -> SPIClass {
+public func inlinable() -> SPIClass {
   spiFunc()
   _ = SPIClass()
 }
 
+@inlinable
+public func inlinable(_ e: PublicEnum) {
+  switch e {
+  case .publicCase: break
+  case .spiCase: break
+  @unknown default: break
+  }
+
+  if case .spiCase = e {}
+
+  _ = PublicEnum.spiCase
+}
+
 @frozen public struct FrozenStruct {
   public var spiInFrozen = SPIStruct()
   var spiTypeInFrozen = SPIStruct()

test/SPI/private_swiftinterface.swift

@@ -34,14 +34,6 @@
 // RUN: %target-swift-frontend -typecheck-module-from-interface -I %t %t/Main.swiftinterface
 // RUN: %target-swift-frontend -typecheck-module-from-interface -I %t %t/Main.private.swiftinterface -module-name Main
 
-/// Serialize and deserialize this module, then print.
-// RUN: %target-swift-frontend -emit-module %s -emit-module-path %t/Merged-partial.swiftmodule -swift-version 5 -I %t -module-name Merged -enable-library-evolution
-// RUN: %target-swift-frontend -merge-modules %t/Merged-partial.swiftmodule -module-name Merged -emit-module -emit-module-path %t/Merged.swiftmodule -I %t -emit-module-interface-path %t/Merged.swiftinterface -emit-private-module-interface-path %t/Merged.private.swiftinterface -enable-library-evolution -swift-version 5 -I %t
-// RUN: %FileCheck -check-prefix=CHECK-PUBLIC %s < %t/Merged.swiftinterface
-// RUN: %FileCheck -check-prefix=CHECK-PRIVATE %s < %t/Merged.private.swiftinterface
-// RUN: %target-swift-frontend -typecheck-module-from-interface -I %t %t/Merged.swiftinterface
-// RUN: %target-swift-frontend -typecheck-module-from-interface -I %t %t/Merged.private.swiftinterface -module-name Merged
-
 /// Both the public and private textual interfaces should have
 /// SPI information with `-library-level spi`.
 // RUN: %target-swift-frontend -typecheck %s -emit-module-interface-path %t/SPIModule.swiftinterface -emit-private-module-interface-path %t/SPIModule.private.swiftinterface -enable-library-evolution -swift-version 5 -I %t -module-name SPIModule -library-level spi
@@ -157,6 +149,34 @@ public struct PublicStruct {
   // CHECK-PUBLIC-NOT: spiWrappedDefault
 }
 
+@_spi(S) public enum SPIEnum {
+// CHECK-PRIVATE: @_spi(S) public enum SPIEnum
+// CHECK-PUBLIC-NOT: SPIEnum
+
+  case spiEnumCase
+  // CHECK-PRIVATE: case spiEnumCase
+  // CHECK-PUBLIC-NOT: spiEnumCase
+}
+
+public enum PublicEnum {
+  case publicCase
+  // CHECK-PUBLIC: case publicCase
+  // CHECK-PRIVATE: case publicCase
+
+  @_spi(S) case spiCase
+  // CHECK-PRIVATE: @_spi(S) case spiCase
+  // CHECK-PUBLIC-NOT: spiCase
+
+  @_spi(S) case spiCaseA, spiCaseB
+  // CHECK-PRIVATE: @_spi(S) case spiCaseA, spiCaseB
+  // CHECK-PUBLIC-NOT: spiCaseA
+  // CHECK-PUBLIC-NOT: spiCaseB
+
+  @_spi(S) case spiCaseWithPayload(_ c: SomeClass)
+  // CHECK-PRIVATE: @_spi(S) case spiCaseWithPayload(_: {{.*}}.SomeClass)
+  // CHECK-PUBLIC-NOT: spiCaseWithPayload
+}
+
 @_spi(LocalSPI) public protocol SPIProto3 {
 // CHECK-PRIVATE: @_spi(LocalSPI) public protocol SPIProto3
 // CHECK-PUBLIC-NOT: SPIProto3

test/SPI/public_client.swift

@@ -27,6 +27,7 @@ internalFunc() // expected-error {{cannot find 'internalFunc' in scope}}
 let c = SPIClass() // expected-error {{cannot find 'SPIClass' in scope}}
 let s = SPIStruct() // expected-error {{cannot find 'SPIStruct' in scope}}
 SPIEnum().spiMethod() // expected-error {{cannot find 'SPIEnum' in scope}}
+_ = PublicEnum.spiCase // expected-error {{'spiCase' is inaccessible due to '@_spi' protection level}}
 
 var ps = PublicStruct()
 let _ = PublicStruct(alt_init: 1) // expected-error {{argument passed to call that takes no arguments}}

test/SPI/spi_client.swift

@@ -40,6 +40,7 @@ print(s.spiVar)
 
 SPIEnum().spiMethod()
 SPIEnum.A.spiMethod()
+_ = PublicEnum.spiCase
 
 var ps = PublicStruct()
 let _ = PublicStruct(alt_init: 1)

test/SPI/spi_enum_element.swift

@@ -0,0 +1,48 @@
+// RUN: %target-typecheck-verify-swift -swift-version 5
+// RUN: %target-typecheck-verify-swift -swift-version 5 -enable-library-evolution
+
+public enum PublicEnum {
+  case publicCase
+  @_spi(S) case spiCase
+}
+
+@inlinable public func publicInlinableFunc(_ e: PublicEnum) {
+  switch e {
+  case .publicCase: break
+  case .spiCase: break   // FIXME: this should be diagnosed with "cannot use enum case 'spiCase' here; it is SPI"
+  @unknown default: break
+  }
+
+  if case .spiCase = e {} // FIXME: this should be diagnosed with "cannot use enum case 'spiCase' here; it is SPI"
+
+  _ = PublicEnum.spiCase // expected-error {{enum case 'spiCase' cannot be used in an '@inlinable' function because it is SPI}}
+}
+
+@_spi(S)
+@inlinable public func spiInlinableFunc(_ e: PublicEnum) {
+  switch e {
+  case .publicCase: break
+  case .spiCase: break
+  @unknown default: break
+  }
+
+  if case .spiCase = e {}
+
+  _ = PublicEnum.spiCase
+}
+
+public struct PublicStruct {}
+@_spi(S) public struct SPIStruct {} // expected-note {{type declared here}}
+
+public enum PublicEnumWithPayloads {
+  case publicCasePublicPayload(_ s: PublicStruct)
+  case publicCaseSPIPayload(_ s: SPIStruct) // expected-error {{cannot use struct 'SPIStruct' here; it is SPI}}
+  @_spi(S) case spiCasePublicPayload(_ s: PublicStruct)
+  @_spi(S) case spiCaseSPIPayload(_ s: SPIStruct)
+}
+
+@_spi(S)
+public enum SPIEnumWithPayloads {
+  case publicPayloadCase(_ s: PublicStruct)
+  case spiPayloadCase(_ s: SPIStruct)
+}

test/SPI/spi_members.swift

@@ -14,7 +14,7 @@ public class Bar {
   public init() {}
 }
 
-public struct Resilient {
+public struct ResilientStructSPIMembers {
   public init() {}
 
   @_spi(Foo) public func method(_: Bar) {}
@@ -28,7 +28,7 @@ public struct Resilient {
   @_spi(Foo) @Wrapper public var wrappedProperty2 = Bar()
 }
 
-@frozen public struct Good {
+@frozen public struct FrozenStructSPIMembers {
   public init() {}
 
   @_spi(Foo) public func method(_: Bar) {}
@@ -55,7 +55,7 @@ public struct Resilient {
   // expected-error@-1 {{stored property 'wrappedProperty2' cannot be declared '@_spi' in a '@frozen' struct}}
 }
 
-@frozen public struct Bad {
+@frozen public struct FrozenStructPublicMembers {
   public init() {}
 
   public func method(_: Bar) {} // expected-error {{cannot use class 'Bar' here; it is SPI}}
@@ -83,7 +83,7 @@ public struct Resilient {
   // expected-error@-3 {{initializer 'init()' cannot be used in a property initializer in a '@frozen' type because it is SPI}}
 }
 
-@frozen public struct BadPrivate {
+@frozen public struct FrozenStructPrivateMembers {
   private init() {}
 
   private func method(_: Bar) {}
@@ -110,3 +110,24 @@ public struct Resilient {
   // expected-error@-2 {{class 'Bar' cannot be used in a property initializer in a '@frozen' type because it is SPI}}
   // expected-error@-3 {{initializer 'init()' cannot be used in a property initializer in a '@frozen' type because it is SPI}}
 }
+
+public enum ResilientEnum {
+  @_spi(S)
+  case okSpiCase
+
+  @_spi(S)
+  case okSpiCaseWithPayload(_: Int)
+}
+
+@frozen
+public enum FrozenEnum {
+  case okCase
+
+  @_spi(S) // expected-error {{enum case 'spiCase' cannot be declared '@_spi' in a '@frozen' enum}}
+  case spiCase
+
+  case okCaseWithPayload(_: Int)
+
+  @_spi(S) // expected-error {{enum case 'spiCaseWithPayload' cannot be declared '@_spi' in a '@frozen' enum}}
+  case spiCaseWithPayload(_: Int)
+}