[Macros] Add executable plugin support

Executable compiler plugins are programs invoked by the host compiler
and communicate with the host with IPC via standard IO (stdin/stdout.)
Each message is serialized in JSON, prefixed with a header which is a
64bit little-endian integer indicating the size of the message.

* Basic/ExecuteWithPipe: External program invocation. Lik
  llvm::sys::ExecuteNoWait() but establishing pipes to the child's
  stdin/stdout
* Basic/Sandbox: Sandboxed execution helper. Create command line
  arguments to be executed in sandbox environment (similar to SwiftPM's
  pluging sandbox)
* AST/PluginRepository: ASTContext independent plugin manager
* ASTGen/PluginHost: Communication with the plugin. Messages are
  serialized by ASTGen/LLVMJSON

rdar://101508815

Author: rintaro

include/swift/AST/ASTContext.h

@@ -79,6 +79,7 @@ namespace swift {
   class DifferentiableAttr;
   class ExtensionDecl;
   struct ExternalSourceLocs;
+  class LoadedExecutablePlugin;
   class ForeignRepresentationInfo;
   class FuncDecl;
   class GenericContext;
@@ -91,6 +92,7 @@ namespace swift {
   class ModuleDependencyInfo;
   class PatternBindingDecl;
   class PatternBindingInitializer;
+  class PluginRegistry;
   class SourceFile;
   class SourceLoc;
   class Type;
@@ -1458,6 +1460,16 @@ class ASTContext final {
 
   Type getNamedSwiftType(ModuleDecl *module, StringRef name);
 
+  LoadedExecutablePlugin *
+  lookupExecutablePluginByModuleName(Identifier moduleName);
+
+  /// Get the plugin registry this ASTContext is using.
+  PluginRegistry *getPluginRegistry() const;
+
+  /// Set the plugin registory this ASTContext should use.
+  /// This should be called before any plugin is loaded.
+  void setPluginRegistry(PluginRegistry *newValue);
+
 private:
   friend Decl;
 

include/swift/AST/CASTBridging.h

@@ -13,8 +13,13 @@
 #ifndef SWIFT_C_AST_ASTBRIDGING_H
 #define SWIFT_C_AST_ASTBRIDGING_H
 
+#include "swift/Basic/CBasicBridging.h"
 #include "swift/Basic/Compiler.h"
+
 #include <inttypes.h>
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
 
 #if __clang__
 // Provide macros to temporarily suppress warning about the use of
@@ -280,6 +285,31 @@ void Decl_dump(void *);
 void Stmt_dump(void *);
 void Type_dump(void *);
 
+//===----------------------------------------------------------------------===//
+// Plugins
+//===----------------------------------------------------------------------===//
+
+/// Set a capability data to the plugin object. Since the data is just a opaque
+/// pointer, it's not used in AST at all.
+void Plugin_setCapability(void *handle, const void *data);
+
+/// Get a coapability data set by \c Plugin_setCapability .
+const void *_Nullable Plugin_getCapability(void *handle);
+
+/// Lock the plugin. Clients should lock it during sending and recving the
+/// response.
+void Plugin_lock(void *handle);
+
+/// Unlock the plugin.
+void Plugin_unlock(void *handle);
+
+/// Sends the message to the plugin, returns true if there was an error.
+/// Clients should receive the response  by \c Plugin_waitForNextMessage .
+_Bool Plugin_sendMessage(void *handle, const BridgedData data);
+
+/// Receive a message from the plugin.
+_Bool Plugin_waitForNextMessage(void *handle, BridgedData *data);
+
 #ifdef __cplusplus
 }
 #endif

include/swift/AST/MacroDefinition.h

@@ -24,6 +24,11 @@ namespace swift {
 
 /// A reference to an external macro definition that is understood by ASTGen.
 struct ExternalMacroDefinition {
+  enum class PluginKind {
+    InProcess = 0,
+    Executable = 1,
+  };
+  PluginKind kind;
   /// ASTGen's notion of an macro definition, which is opaque to the C++ part
   /// of the compiler.
   void *opaqueHandle = nullptr;

include/swift/AST/PluginRegistry.h

@@ -0,0 +1,89 @@
+//===--- PluginRegistry.h ---------------------------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Chrono.h"
+#include "llvm/Support/Program.h"
+
+#include <mutex>
+#include <vector>
+
+namespace swift {
+
+class LoadedExecutablePlugin {
+  const llvm::sys::procid_t pid;
+  const llvm::sys::TimePoint<> LastModificationTime;
+  const int inputFileDescriptor;
+  const int outputFileDescriptor;
+
+  /// Opaque value of the protocol capability of the pluugin. This is a
+  /// value from ASTGen.
+  const void *capability = nullptr;
+
+  /// Cleanup function to call ASTGen.
+  std::function<void(void)> cleanup;
+
+  std::mutex mtx;
+
+  ssize_t write(const void *buf, size_t nbyte) const;
+  ssize_t read(void *buf, size_t nbyte) const;
+
+public:
+  LoadedExecutablePlugin(llvm::sys::procid_t pid,
+                         llvm::sys::TimePoint<> LastModificationTime,
+                         int inputFileDescriptor, int outputFileDescriptor);
+  ~LoadedExecutablePlugin();
+  llvm::sys::TimePoint<> getLastModificationTime() const {
+    return LastModificationTime;
+  }
+
+  void lock() { mtx.lock(); }
+  void unlock() { mtx.unlock(); }
+
+  /// Send a message to the plugin.
+  bool sendMessage(llvm::StringRef message) const;
+
+  /// Wait for a message from plugin and returns it.
+  std::string waitForNextMessage() const;
+
+  bool isInitialized() const { return bool(cleanup); }
+  void setCleanup(std::function<void(void)> cleanup) {
+    this->cleanup = cleanup;
+  }
+
+  llvm::sys::procid_t getPid() { return pid; }
+
+  const void *getCapability() { return capability; };
+  void setCapability(const void *newValue) { capability = newValue; };
+};
+
+class PluginRegistry {
+  /// Record of loaded plugin library modules.
+  llvm::StringMap<void *> LoadedPluginLibraries;
+
+  /// Record of loaded plugin executables.
+  llvm::StringMap<std::unique_ptr<LoadedExecutablePlugin>>
+      LoadedPluginExecutables;
+
+public:
+  bool loadLibraryPlugin(llvm::StringRef path, const char *&errorMsg);
+  LoadedExecutablePlugin *loadExecutablePlugin(llvm::StringRef path,
+                                               const char *&errorMsg);
+
+  const llvm::StringMap<void *> &getLoadedLibraryPlugins() const {
+    return LoadedPluginLibraries;
+  }
+};
+
+} // namespace swift

include/swift/AST/TypeCheckRequests.h

@@ -3931,7 +3931,7 @@ class ExpandPeerMacroRequest
 /// Resolve an external macro given its module and type name.
 class ExternalMacroDefinitionRequest
     : public SimpleRequest<ExternalMacroDefinitionRequest,
-                           ExternalMacroDefinition(
+                           llvm::Optional<ExternalMacroDefinition>(
                                ASTContext *, Identifier, Identifier),
                            RequestFlags::Cached> {
 public:
@@ -3940,9 +3940,9 @@ class ExternalMacroDefinitionRequest
 private:
   friend SimpleRequest;
 
-  ExternalMacroDefinition evaluate(
-      Evaluator &evaluator, ASTContext *ctx, Identifier moduleName,
-      Identifier typeName
+  llvm::Optional<ExternalMacroDefinition> evaluate(
+      Evaluator &evaluator, ASTContext *ctx,
+      Identifier moduleName, Identifier typeName
   ) const;
 
 public:

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -426,7 +426,7 @@ SWIFT_REQUEST(TypeChecker, CompilerPluginLoadRequest,
               void *(ASTContext *, Identifier),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ExternalMacroDefinitionRequest,
-              ExternalMacroDefinition(ASTContext *, Identifier, Identifier),
+              Optional<ExternalMacroDefinition>(ASTContext *, Identifier, Identifier),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ExpandMacroExpansionDeclRequest,
               ArrayRef<Decl *>(MacroExpansionDecl *),

include/swift/Basic/Program.h

@@ -13,6 +13,12 @@
 #ifndef SWIFT_BASIC_PROGRAM_H
 #define SWIFT_BASIC_PROGRAM_H
 
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/ErrorOr.h"
+#include "llvm/Support/Program.h"
+
 namespace swift {
 
 /// This function executes the program using the arguments provided,
@@ -33,6 +39,31 @@ namespace swift {
 int ExecuteInPlace(const char *Program, const char **args,
                    const char **env = nullptr);
 
+struct ChildProcessInfo {
+  llvm::sys::procid_t Pid;
+  int WriteFileDescriptor;
+  int ReadFileDescriptor;
+
+  ChildProcessInfo(llvm::sys::procid_t Pid, int WriteFileDescriptor,
+                   int ReadFileDescriptor)
+      : Pid(Pid), WriteFileDescriptor(WriteFileDescriptor),
+        ReadFileDescriptor(ReadFileDescriptor) {}
+};
+
+/// This function executes the program using the argument provided.
+/// Establish pipes between the current process and the spawned process.
+/// Returned a \c ChildProcessInfo where WriteFileDescriptor is piped to
+/// child's STDIN, and ReadFileDescriptor is piped from child's STDOUT.
+///
+/// \param program Path of the program to be executed
+/// \param args An array of strings that are passed to the program. The first
+/// element should be the name of the program.
+/// \param env An optional array of strings to use for the program's
+/// environment.
+llvm::ErrorOr<swift::ChildProcessInfo> ExecuteWithPipe(
+    llvm::StringRef program, llvm::ArrayRef<llvm::StringRef> args,
+    llvm::Optional<llvm::ArrayRef<llvm::StringRef>> env = llvm::None);
+
 } // end namespace swift
 
 #endif // SWIFT_BASIC_PROGRAM_H

include/swift/Basic/Sandbox.h

@@ -0,0 +1,44 @@
+//===--- Sandbox.h ----------------------------------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_BASIC_SANDBOX_H
+#define SWIFT_BASIC_SANDBOX_H
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Allocator.h"
+
+namespace swift {
+namespace Sandbox {
+
+/// Applies a sandbox invocation to the given command line (if the platform
+/// supports it), and returns the modified command line. On platforms that don't
+/// support sandboxing, the command line is returned unmodified.
+///
+/// - Parameters:
+///   - command: The command line to sandbox (including executable as first
+///              argument)
+///   - strictness: The basic strictness level of the standbox.
+///   - writableDirectories: Paths under which writing should be allowed, even
+///     if they would otherwise be read-only based on the strictness or paths in
+///     `readOnlyDirectories`.
+///   - readOnlyDirectories: Paths under which writing should be denied, even if
+///     they would have otherwise been allowed by the rules implied by the
+///     strictness level.
+bool apply(llvm::SmallVectorImpl<llvm::StringRef> &command,
+           llvm::BumpPtrAllocator &Alloc);
+
+} // namespace Sandbox
+} // namespace swift
+
+#endif // SWIFT_BASIC_SANDBOX_H