[Macros/Plugin] Make 'features' a set of enum values

Author: rintaro

lib/ASTGen/Sources/ASTGen/PluginHost.swift

@@ -46,7 +46,7 @@ func swift_ASTGen_pluginServerLoadLibraryPlugin(
   cxxDiagnosticEngine: UnsafeMutablePointer<UInt8>
 ) -> Bool {
   let plugin =  CompilerPlugin(opaqueHandle: opaqueHandle)
-  assert(plugin.capability.features?.contains("loadPluginLibrary") == true)
+  assert(plugin.capability?.features.contains(.loadPluginLibrary) == true)
   let libraryPath = String(cString: libraryPath)
   let moduleName = String(cString: moduleName)
   let diagEngine = PluginDiagnosticsEngine(cxxDiagnosticEngine: cxxDiagnosticEngine)
@@ -67,6 +67,24 @@ func swift_ASTGen_pluginServerLoadLibraryPlugin(
 }
 
 struct CompilerPlugin {
+  struct Capability {
+    enum Feature: String {
+    case loadPluginLibrary = "load-plugin-library"
+    }
+
+    var protocolVersion: Int
+    var features: Set<Feature>
+
+    init(_ message: PluginMessage.PluginCapability) {
+      self.protocolVersion = message.protocolVersion
+      if let features = message.features {
+        self.features = Set(features.compactMap(Feature.init(rawValue:)))
+      } else {
+        self.features = []
+      }
+    }
+  }
+
   let opaqueHandle: UnsafeMutableRawPointer
 
   private func withLock<R>(_ body: () throws -> R) rethrows -> R {
@@ -111,26 +129,26 @@ struct CompilerPlugin {
   }
 
   func initialize() {
-    self.withLock {
-      // Get capability.
-      let response: PluginToHostMessage
-      do {
+    // Don't use `sendMessageAndWait` because we want to keep the lock until
+    // setting the returned value.
+    do {
+      try self.withLock {
+        // Get capability.
         try self.sendMessage(.getCapability)
-        response = try self.waitForNextMessage()
-      } catch {
-        assertionFailure(String(describing: error))
-        return
-      }
-      switch response {
-      case .getCapabilityResult(capability: let capability):
-        let ptr = UnsafeMutablePointer<PluginMessage.PluginCapability>.allocate(capacity: 1)
-        ptr.initialize(to: capability)
+        let response = try self.waitForNextMessage()
+        guard case .getCapabilityResult(let capability) = response else {
+          throw PluginError.invalidReponseKind
+        }
+        let ptr = UnsafeMutablePointer<Capability>.allocate(capacity: 1)
+        ptr.initialize(to: .init(capability))
         Plugin_setCapability(opaqueHandle, UnsafeRawPointer(ptr))
-      default:
-        assertionFailure("invalid response")
       }
+    } catch {
+      assertionFailure(String(describing: error))
+      return
     }
   }
+
   func deinitialize() {
     self.withLock {
       if let ptr = Plugin_getCapability(opaqueHandle) {
@@ -142,11 +160,11 @@ struct CompilerPlugin {
     }
   }
 
-  var capability: PluginMessage.PluginCapability {
+  var capability: Capability? {
     if let ptr = Plugin_getCapability(opaqueHandle) {
-      return ptr.assumingMemoryBound(to: PluginMessage.PluginCapability.self).pointee
+      return ptr.assumingMemoryBound(to: Capability.self).pointee
     }
-    return PluginMessage.PluginCapability(protocolVersion: 0)
+    return nil
   }
 }
 

lib/ASTGen/Sources/ASTGen/PluginMessages.swift

@@ -71,7 +71,7 @@ internal enum PluginToHostMessage: Codable {
     var protocolVersion: Int
 
     /// Optional features this plugin provides.
-    ///  * loadPluginLibrary: 'loadPluginLibrary' message is implemented.
+    ///  * 'load-plugin-library': 'loadPluginLibrary' message is implemented.
     var features: [String]?
   }