Merge pull request #67973 from zachary0kent/lazy-immediate-globals

[Immediate] JIT'ing Globals

Author: zachary-kent

include/swift/AST/ASTTypeIDZone.def

@@ -37,7 +37,6 @@ SWIFT_TYPEID(FragileFunctionKind)
 SWIFT_TYPEID(PolymorphicEffectKind)
 SWIFT_TYPEID(PolymorphicEffectRequirementList)
 SWIFT_TYPEID(TangentPropertyInfo)
-SWIFT_TYPEID(SymbolSourceMap)
 SWIFT_TYPEID(Type)
 SWIFT_TYPEID(TypePair)
 SWIFT_TYPEID(TypeWitnessAndDecl)

include/swift/AST/ASTTypeIDs.h

@@ -70,7 +70,6 @@ struct FragileFunctionKind;
 enum class PolymorphicEffectKind : uint8_t;
 class PolymorphicEffectRequirementList;
 class SourceFile;
-class SymbolSourceMap;
 struct TangentPropertyInfo;
 class Type;
 class TypeAliasDecl;

include/swift/AST/SILGenRequests.h

@@ -21,6 +21,7 @@
 #include "swift/AST/EvaluatorDependencies.h"
 #include "swift/AST/SimpleRequest.h"
 #include "swift/AST/SourceFile.h"
+#include "swift/AST/TBDGenRequests.h"
 #include "swift/SIL/SILDeclRef.h"
 
 namespace swift {
@@ -43,6 +44,8 @@ void reportEvaluatedRequest(UnifiedStatsReporter &stats,
 
 using SILRefsToEmit = llvm::SmallVector<SILDeclRef, 1>;
 
+using SymbolSources = llvm::SmallVector<SymbolSource, 1>;
+
 /// Describes a file or module to be lowered to SIL.
 struct ASTLoweringDescriptor {
   llvm::PointerUnion<FileUnit *, ModuleDecl *> context;
@@ -52,20 +55,17 @@ struct ASTLoweringDescriptor {
 
   /// A specific set of SILDeclRefs to emit. If set, only these refs will be
   /// emitted. Otherwise the entire \c context will be emitted.
-  llvm::Optional<SILRefsToEmit> refsToEmit;
+  llvm::Optional<SymbolSources> SourcesToEmit;
 
   friend llvm::hash_code hash_value(const ASTLoweringDescriptor &owner) {
     return llvm::hash_combine(owner.context, (void *)&owner.conv,
-                              (void *)&owner.opts,
-                              owner.refsToEmit);
+                              (void *)&owner.opts, owner.SourcesToEmit);
   }
 
   friend bool operator==(const ASTLoweringDescriptor &lhs,
                          const ASTLoweringDescriptor &rhs) {
-    return lhs.context == rhs.context &&
-           &lhs.conv == &rhs.conv &&
-           &lhs.opts == &rhs.opts &&
-           lhs.refsToEmit == rhs.refsToEmit;
+    return lhs.context == rhs.context && &lhs.conv == &rhs.conv &&
+           &lhs.opts == &rhs.opts && lhs.SourcesToEmit == rhs.SourcesToEmit;
   }
 
   friend bool operator!=(const ASTLoweringDescriptor &lhs,
@@ -76,17 +76,19 @@ struct ASTLoweringDescriptor {
 public:
   static ASTLoweringDescriptor
   forFile(FileUnit &sf, Lowering::TypeConverter &conv, const SILOptions &opts,
-          llvm::Optional<SILRefsToEmit> refsToEmit = llvm::None,
+          llvm::Optional<SymbolSources> SourcesToEmit = llvm::None,
           const IRGenOptions *irgenOptions = nullptr) {
-    return ASTLoweringDescriptor{&sf, conv, opts, irgenOptions, refsToEmit};
+    return ASTLoweringDescriptor{&sf, conv, opts, irgenOptions,
+                                 std::move(SourcesToEmit)};
   }
 
   static ASTLoweringDescriptor
   forWholeModule(ModuleDecl *mod, Lowering::TypeConverter &conv,
                  const SILOptions &opts,
-                 llvm::Optional<SILRefsToEmit> refsToEmit = llvm::None,
+                 llvm::Optional<SymbolSources> SourcesToEmit = llvm::None,
                  const IRGenOptions *irgenOptions = nullptr) {
-    return ASTLoweringDescriptor{mod, conv, opts, irgenOptions, refsToEmit};
+    return ASTLoweringDescriptor{mod, conv, opts, irgenOptions,
+                                 std::move(SourcesToEmit)};
   }
 
   /// Retrieves the files to generate SIL for. If the descriptor is configured

include/swift/AST/SourceFile.h

@@ -166,6 +166,14 @@ class SourceFile final : public FileUnit {
   /// if it were done from OpaqueResultTypeRequest.
   llvm::SetVector<OpaqueTypeDecl *> UnvalidatedOpaqueReturnTypes;
 
+  /// The list of functions defined in this file whose bodies have yet to be
+  /// typechecked. They must be held in this list instead of eagerly validated
+  /// because their bodies may force us to perform semantic checks of arbitrary
+  /// complexity, and we currently cannot handle those checks in isolation. E.g.
+  /// we cannot, in general, perform witness matching on singular requirements
+  /// unless the entire conformance has been evaluated.
+  std::vector<AbstractFunctionDecl *> DelayedFunctions;
+
   /// The list of top-level items in the source file. This is \c None if
   /// they have not yet been parsed.
   /// FIXME: Once addTopLevelDecl/prependTopLevelDecl
@@ -259,13 +267,11 @@ class SourceFile final : public FileUnit {
   /// The list of local type declarations in the source file.
   llvm::SetVector<TypeDecl *> LocalTypeDecls;
 
-  /// The list of functions defined in this file whose bodies have yet to be
-  /// typechecked. They must be held in this list instead of eagerly validated
-  /// because their bodies may force us to perform semantic checks of arbitrary
-  /// complexity, and we currently cannot handle those checks in isolation. E.g.
-  /// we cannot, in general, perform witness matching on singular requirements
-  /// unless the entire conformance has been evaluated.
-  std::vector<AbstractFunctionDecl *> DelayedFunctions;
+  /// Defer type checking of `AFD` to the end of `Sema`
+  void addDelayedFunction(AbstractFunctionDecl *AFD);
+
+  /// Typecheck the bodies of all lazily checked functions
+  void typeCheckDelayedFunctions();
 
   /// A mapping from Objective-C selectors to the methods that have
   /// those selectors.

include/swift/AST/TBDGenRequests.h

@@ -17,6 +17,8 @@
 #ifndef SWIFT_TBDGEN_REQUESTS_H
 #define SWIFT_TBDGEN_REQUESTS_H
 
+#include "llvm/ExecutionEngine/JITSymbol.h"
+
 #include "swift/AST/ASTTypeIDs.h"
 #include "swift/AST/SimpleRequest.h"
 #include "swift/IRGen/Linking.h"
@@ -148,6 +150,10 @@ class SymbolSource {
     /// A symbol introduced when emitting a SIL decl.
     SIL,
 
+    /// A symbol introduced when emitting a SIL global variable.
+    /// Only used in piecewise compilation for Immediate mode.
+    Global,
+
     /// A symbol introduced when emitting LLVM IR.
     IR,
 
@@ -162,13 +168,17 @@ class SymbolSource {
 
 private:
   union {
+    VarDecl *Global;
     SILDeclRef silDeclRef;
     irgen::LinkEntity irEntity;
   };
 
   explicit SymbolSource(SILDeclRef ref) : kind(Kind::SIL) {
     silDeclRef = ref;
   }
+
+  SymbolSource(VarDecl *Global) : kind(Kind::Global), Global(Global) {}
+
   explicit SymbolSource(irgen::LinkEntity entity) : kind(Kind::IR) {
     irEntity = entity;
   }
@@ -180,6 +190,9 @@ class SymbolSource {
   static SymbolSource forSILDeclRef(SILDeclRef ref) {
     return SymbolSource{ref};
   }
+
+  static SymbolSource forGlobal(VarDecl *Global) { return Global; }
+
   static SymbolSource forIRLinkEntity(irgen::LinkEntity entity) {
     return SymbolSource{entity};
   }
@@ -198,49 +211,93 @@ class SymbolSource {
     assert(kind == Kind::SIL);
     return silDeclRef;
   }
+
+  VarDecl *getGlobal() const {
+    assert(kind == Kind::Global);
+    return Global;
+  }
+
   irgen::LinkEntity getIRLinkEntity() const {
     assert(kind == Kind::IR);
     return irEntity;
   }
-};
-
-/// Maps a symbol back to its source for lazy compilation.
-class SymbolSourceMap {
-  friend class SymbolSourceMapRequest;
-
-  using Storage = llvm::StringMap<SymbolSource>;
-  const Storage *storage;
 
-  explicit SymbolSourceMap(const Storage *storage) : storage(storage) {
-    assert(storage);
+  /// Typecheck the entity wrapped by this `SymbolSource`
+  void typecheck() const {
+    switch (kind) {
+    case Kind::SIL: {
+      if (auto *AFD = silDeclRef.getAbstractFunctionDecl())
+        // If this entity is a `SILFunction`, check its body
+        AFD->getTypecheckedBody();
+      break;
+    }
+    case Kind::Global:
+    case Kind::IR:
+    case Kind::LinkerDirective:
+    case Kind::Unknown:
+      // Nothing to do
+      break;
+    }
   }
 
-public:
-  llvm::Optional<SymbolSource> find(StringRef symbol) const {
-    auto result = storage->find(symbol);
-    if (result == storage->end())
-      return llvm::None;
-    return result->second;
+  friend llvm::hash_code hash_value(const SymbolSource &S) {
+    auto Kind = S.kind;
+    switch (Kind) {
+    case Kind::SIL:
+      return llvm::hash_combine(Kind, S.silDeclRef);
+    case Kind::Global:
+      return llvm::hash_combine(Kind, S.Global);
+    case Kind::IR:
+      return llvm::hash_combine(Kind, S.irEntity);
+    case Kind::LinkerDirective:
+    case Kind::Unknown:
+      return llvm::hash_value(Kind);
+    }
   }
 
-  friend bool operator==(const SymbolSourceMap &lhs,
-                         const SymbolSourceMap &rhs) {
-    return lhs.storage == rhs.storage;
+  friend bool operator==(const SymbolSource &LHS, const SymbolSource &RHS) {
+    if (LHS.kind != RHS.kind)
+      return false;
+    switch (LHS.kind) {
+    case Kind::SIL:
+      return LHS.silDeclRef == RHS.silDeclRef;
+    case Kind::Global:
+      return LHS.Global == RHS.Global;
+    case Kind::IR:
+      return LHS.irEntity == RHS.irEntity;
+    case Kind::LinkerDirective:
+    case Kind::Unknown:
+      return true;
+    }
   }
-  friend bool operator!=(const SymbolSourceMap &lhs,
-                         const SymbolSourceMap &rhs) {
-    return !(lhs == rhs);
+
+  friend bool operator!=(const SymbolSource &LHS, const SymbolSource &RHS) {
+    return !(LHS == RHS);
   }
 
-  friend void simple_display(llvm::raw_ostream &out, const SymbolSourceMap &) {
-    out << "(symbol storage map)";
+  llvm::JITSymbolFlags getJITSymbolFlags() const {
+    switch (kind) {
+    case Kind::SIL:
+      return llvm::JITSymbolFlags::Callable | llvm::JITSymbolFlags::Exported;
+    case Kind::Global:
+      return llvm::JITSymbolFlags::Exported;
+    case Kind::IR:
+      llvm_unreachable("Unimplemented: Symbol flags for LinkEntities");
+    case Kind::LinkerDirective:
+      llvm_unreachable("Unsupported: Symbol flags for linker directives");
+    case Kind::Unknown:
+      llvm_unreachable("Unsupported: Symbol flags for unknown source");
+    }
   }
 };
 
+/// Maps a symbol back to its source for lazy compilation.
+using SymbolSourceMap = llvm::StringMap<SymbolSource>;
+
 /// Computes a map of symbols to their SymbolSource for a file or module.
 class SymbolSourceMapRequest
     : public SimpleRequest<SymbolSourceMapRequest,
-                           SymbolSourceMap(TBDGenDescriptor),
+                           const SymbolSourceMap *(TBDGenDescriptor),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -249,7 +306,8 @@ class SymbolSourceMapRequest
   friend SimpleRequest;
 
   // Evaluation.
-  SymbolSourceMap evaluate(Evaluator &evaluator, TBDGenDescriptor desc) const;
+  const SymbolSourceMap *evaluate(Evaluator &evaluator,
+                                  TBDGenDescriptor desc) const;
 
 public:
   // Cached.

include/swift/AST/TBDGenTypeIDZone.def

@@ -20,7 +20,7 @@ SWIFT_REQUEST(TBDGen, PublicSymbolsRequest,
               std::vector<std::string>(TBDGenDescriptor),
               Uncached, NoLocationInfo)
 SWIFT_REQUEST(TBDGen, SymbolSourceMapRequest,
-              SymbolSourceMap(TBDGenDescriptor),
+              const SymbolSourceMap *(TBDGenDescriptor),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TBDGen, APIGenRequest, apigen::API(TBDGenDescriptor), Uncached,
               NoLocationInfo)

include/swift/AST/TypeCheckRequests.h

@@ -2997,8 +2997,7 @@ class TypeCheckSourceFileRequest
   friend SimpleRequest;
 
   // Evaluation.
-  evaluator::SideEffect
-  evaluate(Evaluator &evaluator, SourceFile *SF) const;
+  evaluator::SideEffect evaluate(Evaluator &evaluator, SourceFile *SF) const;
 
 public:
   // Separate caching.

include/swift/Basic/Features.def

@@ -129,6 +129,9 @@ SUPPRESSIBLE_LANGUAGE_FEATURE(ExtensionMacroAttr, 0, "@attached(extension)", tru
 // Whether to enable @_used and @_section attributes
 EXPERIMENTAL_FEATURE(SymbolLinkageMarkers, true)
 
+// Whether to compile scripts lazily in immediate mode
+EXPERIMENTAL_FEATURE(LazyImmediate, false)
+
 // FIXME: MoveOnlyClasses is not intended to be in production,
 // but our tests currently rely on it, and we want to run those
 // tests in non-asserts builds too.

include/swift/Basic/LangOptions.h

@@ -799,6 +799,9 @@ namespace swift {
 
     /// See \ref FrontendOptions.PrintFullConvention
     bool PrintFullConvention = false;
+
+    /// Defer typechecking of declarations to their use at runtime
+    bool DeferToRuntime = false;
   };
 
   /// Options for controlling the behavior of the Clang importer.

include/swift/IRGen/Linking.h

@@ -515,6 +515,20 @@ class LinkEntity {
     return Kind(LINKENTITY_GET_FIELD(Data, Kind));
   }
 
+  friend llvm::hash_code hash_value(const LinkEntity &Entity) {
+    return llvm::hash_combine(Entity.Pointer, Entity.SecondaryPointer,
+                              Entity.Data);
+  }
+
+  friend bool operator==(const LinkEntity &LHS, const LinkEntity &RHS) {
+    return LHS.Pointer == RHS.Pointer &&
+           LHS.SecondaryPointer == RHS.SecondaryPointer && LHS.Data == RHS.Data;
+  }
+
+  friend bool operator!=(const LinkEntity &LHS, const LinkEntity &RHS) {
+    return !(LHS == RHS);
+  }
+
   static bool isDeclKind(Kind k) {
     return k <= Kind::AsyncFunctionPointerAST;
   }

include/swift/Immediate/Immediate.h

@@ -41,6 +41,9 @@ namespace swift {
   int RunImmediately(CompilerInstance &CI, const ProcessCmdLine &CmdLine,
                      const IRGenOptions &IRGenOpts, const SILOptions &SILOpts,
                      std::unique_ptr<SILModule> &&SM);
+
+  int RunImmediatelyFromAST(CompilerInstance &CI);
+
 } // end namespace swift
 
 #endif // SWIFT_IMMEDIATE_IMMEDIATE_H