[Macros] Freestanding macro expansions in code blocks

include/swift/AST/ASTScope.h

@@ -278,10 +278,12 @@ class ASTScopeImpl : public ASTAllocated<ASTScopeImpl> {
   static const ASTScopeImpl *findStartingScopeForLookup(SourceFile *,
                                                         const SourceLoc where);
 
-protected:
+public:
   /// Not const because may reexpand some scopes.
   ASTScopeImpl *findInnermostEnclosingScope(SourceLoc,
                                             NullablePtr<raw_ostream>);
+
+protected:
   ASTScopeImpl *findInnermostEnclosingScopeImpl(SourceLoc,
                                                 NullablePtr<raw_ostream>,
                                                 SourceManager &,
@@ -1216,6 +1218,29 @@ class MacroDeclScope final : public ASTScopeImpl {
   NullablePtr<const GenericParamList> genericParams() const override;
 };
 
+class MacroExpansionDeclScope final : public ASTScopeImpl {
+public:
+  MacroExpansionDecl *const decl;
+
+  MacroExpansionDeclScope(MacroExpansionDecl *d) : decl(d) {}
+  virtual ~MacroExpansionDeclScope() {}
+
+protected:
+  ASTScopeImpl *expandSpecifically(ScopeCreator &scopeCreator) override;
+
+private:
+  void expandAScopeThatDoesNotCreateANewInsertionPoint(ScopeCreator &);
+
+public:
+  std::string getClassName() const override;
+  SourceRange
+  getSourceRangeOfThisASTNode(bool omitAssertions = false) const override;
+
+public:
+  virtual NullablePtr<Decl> getDeclIfAny() const override { return decl; }
+  Decl *getDecl() const { return decl; }
+};
+
 class AbstractStmtScope : public ASTScopeImpl {
 public:
   SourceRange

include/swift/AST/Decl.h

@@ -110,6 +110,7 @@ namespace swift {
   class VarDecl;
   class OpaqueReturnTypeRepr;
   class Witness;
+  class MacroExpansionExpr;
 
   namespace ast_scope {
   class AbstractPatternEntryScope;
@@ -8452,7 +8453,8 @@ class MacroExpansionDecl : public Decl {
   SourceLoc LeftAngleLoc, RightAngleLoc;
   ArrayRef<TypeRepr *> GenericArgs;
   ArgumentList *ArgList;
-  ArrayRef<Decl *> Rewritten;
+  // An implicit brace statement serving as the container for the expansion.
+  BraceStmt *Rewritten = nullptr;
 
   /// The referenced macro.
   ConcreteDeclRef macroRef;
@@ -8469,7 +8471,7 @@ class MacroExpansionDecl : public Decl {
       : Decl(DeclKind::MacroExpansion, dc), PoundLoc(poundLoc),
         Macro(macro), MacroLoc(macroLoc),
         LeftAngleLoc(leftAngleLoc), RightAngleLoc(rightAngleLoc),
-        GenericArgs(genericArgs), ArgList(args), Rewritten({}) {
+        GenericArgs(genericArgs), ArgList(args) {
     Bits.MacroExpansionDecl.Discriminator = InvalidDiscriminator;
   }
 
@@ -8485,10 +8487,12 @@ class MacroExpansionDecl : public Decl {
   DeclNameLoc getMacroLoc() const { return MacroLoc; }
   DeclNameRef getMacro() const { return Macro; }
   ArgumentList *getArgs() const { return ArgList; }
-  ArrayRef<Decl *> getRewritten() const { return Rewritten; }
-  void setRewritten(ArrayRef<Decl *> rewritten) { Rewritten = rewritten; }
+  void setArgs(ArgumentList *args) { ArgList = args; }
+  BraceStmt *getRewritten() const { return Rewritten; }
+  void setRewritten(BraceStmt *rewritten) { Rewritten = rewritten; }
   ConcreteDeclRef getMacroRef() const { return macroRef; }
   void setMacroRef(ConcreteDeclRef ref) { macroRef = ref; }
+  MacroExpansionExpr *createExpr() const;
 
   /// Returns a discriminator which determines this macro expansion's index
   /// in the sequence of macro expansions within the current function.

include/swift/AST/FileUnit.h

@@ -64,9 +64,15 @@ class FileUnit : public DeclContext, public ASTAllocated<FileUnit> {
   /// within this file.
   ///
   /// This does a simple local lookup, not recursively looking through imports.
-  virtual void lookupValue(DeclName name, NLKind lookupKind,
+  virtual void lookupValue(DeclName name, NLKind lookupKind, NLOptions options,
                            SmallVectorImpl<ValueDecl*> &result) const = 0;
 
+  void lookupValue(DeclName name, NLKind lookupKind,
+                   SmallVectorImpl<ValueDecl*> &result) const {
+    lookupValue(name, lookupKind, getDefaultNLOptions(lookupKind), result);
+  }
+
+
   /// Look up a local type declaration by its mangled name.
   ///
   /// This does a simple local lookup, not recursively looking through imports.
@@ -366,7 +372,7 @@ class BuiltinUnit final : public FileUnit {
 public:
   explicit BuiltinUnit(ModuleDecl &M);
 
-  virtual void lookupValue(DeclName name, NLKind lookupKind,
+  virtual void lookupValue(DeclName name, NLKind lookupKind, NLOptions options,
                            SmallVectorImpl<ValueDecl*> &result) const override;
 
   /// Find all Objective-C methods with the given selector.

include/swift/AST/LookupKinds.h

@@ -53,6 +53,9 @@ enum NLOptions : unsigned {
   // Include @usableFromInline and @inlinable
   NL_IncludeUsableFromInline = 1 << 6,
 
+  /// Don't expand macros.
+  NL_DisableMacroExpansions = 1 << 7,
+
   /// The default set of options used for qualified name lookup.
   ///
   /// FIXME: Eventually, add NL_ProtocolMembers to this, once all of the
@@ -79,6 +82,13 @@ static inline NLOptions operator~(NLOptions value) {
   return NLOptions(~(unsigned)value);
 }
 
+static inline NLOptions getDefaultNLOptions(NLKind kind) {
+  switch (kind) {
+  case NLKind::QualifiedLookup: return NL_QualifiedDefault;
+  case NLKind::UnqualifiedLookup: return NL_UnqualifiedDefault;
+  }
+}
+
 void simple_display(llvm::raw_ostream &out, NLOptions options);
 
 } // end namespace swift

include/swift/AST/MacroDeclaration.h

@@ -49,6 +49,8 @@ enum class MacroRole: uint32_t {
   /// An attached macro that generates synthesized members
   /// inside the declaration.
   SynthesizedMembers = 0x10,
+  /// A freestanding code item macro.
+  CodeItem = 0x20,
 };
 
 /// The contexts in which a particular macro declaration can be used.

include/swift/AST/Module.h

@@ -637,9 +637,19 @@ class ModuleDecl
   /// within the current module.
   ///
   /// This does a simple local lookup, not recursively looking through imports.
-  void lookupValue(DeclName Name, NLKind LookupKind,
+  void lookupValue(DeclName Name, NLKind LookupKind, NLOptions Options,
                    SmallVectorImpl<ValueDecl*> &Result) const;
 
+  /// Look up a (possibly overloaded) value set at top-level scope
+  /// (but with the specified access path, which may come from an import decl)
+  /// within the current module.
+  ///
+  /// This does a simple local lookup, not recursively looking through imports.
+  void lookupValue(DeclName Name, NLKind LookupKind,
+                   SmallVectorImpl<ValueDecl*> &Result) const {
+    lookupValue(Name, LookupKind, getDefaultNLOptions(LookupKind), Result);
+  }
+
   /// Look up a local type declaration by its mangled name.
   ///
   /// This does a simple local lookup, not recursively looking through imports.

include/swift/AST/NameLookup.h

@@ -239,6 +239,10 @@ enum class UnqualifiedLookupFlags {
   // This lookup should include results that are @inlinable or
   // @usableFromInline.
   IncludeUsableFromInline = 1 << 5,
+  // This lookup should not expand macros.
+  DisableMacroExpansions  = 1 << 6,
+  // This lookup should only return macros.
+  MacroLookup             = 1 << 7,
 };
 
 using UnqualifiedLookupOptions = OptionSet<UnqualifiedLookupFlags>;
@@ -794,6 +798,11 @@ class ASTScope : public ASTAllocated<ASTScope> {
   static std::pair<CaseStmt *, CaseStmt *>
   lookupFallthroughSourceAndDest(SourceFile *sourceFile, SourceLoc loc);
 
+  /// Returns true if the given location is within a \c MacroDecl or a top-level
+  /// \c MacroExpansionDecl.
+  static bool isLocWithinMacroDeclOrTopLevelMacroExpansionDeclScope(
+      SourceFile *sf, SourceLoc loc);
+
   SWIFT_DEBUG_DUMP;
   void print(llvm::raw_ostream &) const;
   void dumpOneScopeMapLocation(std::pair<unsigned, unsigned>);

include/swift/AST/SourceFile.h

@@ -427,7 +427,7 @@ class SourceFile final : public FileUnit {
   void cacheVisibleDecls(SmallVectorImpl<ValueDecl *> &&globals) const;
   const SmallVectorImpl<ValueDecl *> &getCachedVisibleDecls() const;
 
-  virtual void lookupValue(DeclName name, NLKind lookupKind,
+  virtual void lookupValue(DeclName name, NLKind lookupKind, NLOptions options,
                            SmallVectorImpl<ValueDecl*> &result) const override;
 
   virtual void lookupVisibleDecls(ImportPath::Access accessPath,

include/swift/AST/SynthesizedFileUnit.h

@@ -43,7 +43,7 @@ class SynthesizedFileUnit final : public FileUnit {
   /// Add a synthesized top-level declaration.
   void addTopLevelDecl(Decl *D) { TopLevelDecls.push_back(D); }
 
-  virtual void lookupValue(DeclName name, NLKind lookupKind,
+  virtual void lookupValue(DeclName name, NLKind lookupKind, NLOptions options,
                            SmallVectorImpl<ValueDecl *> &result) const override;
 
   void lookupObjCMethods(

include/swift/AST/TypeCheckRequests.h

@@ -3806,16 +3806,15 @@ class MacroDefinitionRequest
 /// Find the definition of a given macro.
 class ExpandMacroExpansionDeclRequest
     : public SimpleRequest<ExpandMacroExpansionDeclRequest,
-                           ArrayRef<Decl *>(MacroExpansionDecl *),
+                           BraceStmt *(MacroExpansionDecl *),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
 
 private:
   friend SimpleRequest;
 
-  ArrayRef<Decl *> evaluate(Evaluator &evaluator,
-                            MacroExpansionDecl *med) const;
+  BraceStmt *evaluate(Evaluator &evaluator, MacroExpansionDecl *med) const;
 
 public:
   bool isCached() const { return true; }

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -65,8 +65,8 @@ SWIFT_REQUEST(TypeChecker, DefaultArgumentInitContextRequest,
 SWIFT_REQUEST(TypeChecker, DefaultDefinitionTypeRequest,
               Type(AssociatedTypeDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, DefaultTypeRequest,
-              Type(KnownProtocolKind, const DeclContext *), SeparatelyCached,
-              NoLocationInfo)
+              Type(KnownProtocolKind, const DeclContext *),
+              SeparatelyCached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, DifferentiableAttributeTypeCheckRequest,
               IndexSubset *(DifferentiableAttr *),
               SeparatelyCached, NoLocationInfo)
@@ -453,7 +453,7 @@ SWIFT_REQUEST(TypeChecker, ExternalMacroDefinitionRequest,
               ExternalMacroDefinition(ASTContext *, Identifier, Identifier),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ExpandMacroExpansionDeclRequest,
-              ArrayRef<Decl *>(MacroExpansionDecl *),
+              BraceStmt *(MacroExpansionDecl *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ExpandMemberAttributeMacros,
               bool(Decl *),

include/swift/AST/TypeMemberVisitor.h

@@ -81,8 +81,10 @@ class TypeMemberVisitor : public DeclVisitor<ImplClass, RetTy> {
 
   /// Visit expanded macros.
   void visitMacroExpansionDecl(MacroExpansionDecl *D) {
-    for (auto *decl : D->getRewritten())
-      asImpl().visit(decl);
+    if (auto *rewritten = D->getRewritten())
+      for (auto node : rewritten->getElements())
+        if (auto *decl = node.dyn_cast<Decl *>())
+          asImpl().visit(decl);
   }
 };
 

include/swift/ClangImporter/ClangModule.h

@@ -65,7 +65,7 @@ class ClangModuleUnit final : public LoadedFile {
 
   virtual bool isSystemModule() const override;
 
-  virtual void lookupValue(DeclName name, NLKind lookupKind,
+  virtual void lookupValue(DeclName name, NLKind lookupKind, NLOptions options,
                            SmallVectorImpl<ValueDecl*> &results) const override;
 
   virtual TypeDecl *

include/swift/Serialization/SerializedModuleLoader.h

@@ -352,7 +352,7 @@ class SerializedASTFile final : public LoadedFile {
 
   virtual bool isSystemModule() const override;
 
-  virtual void lookupValue(DeclName name, NLKind lookupKind,
+  virtual void lookupValue(DeclName name, NLKind lookupKind, NLOptions options,
                            SmallVectorImpl<ValueDecl*> &results) const override;
 
   virtual StringRef

lib/AST/ASTContext.cpp

@@ -764,7 +764,8 @@ void ASTContext::lookupInModule(
 
   // Find all of the declarations with this name in the Swift module.
   auto identifier = getIdentifier(name);
-  M->lookupValue(identifier, NLKind::UnqualifiedLookup, results);
+  M->lookupValue(identifier, NLKind::UnqualifiedLookup, NL_UnqualifiedDefault,
+                 results);
 }
 
 void ASTContext::lookupInSwiftModule(
@@ -1138,7 +1139,7 @@ ProtocolDecl *ASTContext::getProtocol(KnownProtocolKind kind) const {
   if (!M)
     return nullptr;
   M->lookupValue(getIdentifier(getProtocolName(kind)),
-                 NLKind::UnqualifiedLookup, results);
+                 NLKind::UnqualifiedLookup, NL_UnqualifiedDefault, results);
 
   for (auto result : results) {
     if (auto protocol = dyn_cast<ProtocolDecl>(result)) {
@@ -5132,7 +5133,8 @@ static NominalTypeDecl *findUnderlyingTypeInModule(ASTContext &ctx,
                                                    ModuleDecl *module) {
   // Find all of the declarations with this name in the Swift module.
   SmallVector<ValueDecl *, 1> results;
-  module->lookupValue(name, NLKind::UnqualifiedLookup, results);
+  module->lookupValue(name, NLKind::UnqualifiedLookup, NL_UnqualifiedDefault,
+                      results);
   for (auto result : results) {
     if (auto nominal = dyn_cast<NominalTypeDecl>(result))
       return nominal;

lib/AST/ASTDemangler.cpp

@@ -107,6 +107,7 @@ ASTBuilder::createBuiltinType(StringRef builtinName,
         builtinName.drop_front(BUILTIN_TYPE_NAME_PREFIX.size());
     Ctx.TheBuiltinModule->lookupValue(Ctx.getIdentifier(strippedName),
                                       NLKind::QualifiedLookup,
+                                      NL_QualifiedDefault,
                                       decls);
 
     if (decls.size() == 1 && isa<TypeDecl>(decls[0]))

lib/AST/ASTDumper.cpp

@@ -1324,10 +1324,13 @@ namespace {
       OS << MED->getMacro();
       if (MED->getArgs()) {
         OS << '\n';
-        OS.indent(Indent + 2);
         printArgumentList(OS, MED->getArgs(), Indent,
                           [&](Expr *E) { printRec(E); });
       }
+      if (auto *rewritten = MED->getRewritten()) {
+        OS << " rewritten=\n";
+        printRec(rewritten, MED->getASTContext());
+      }
       PrintWithColorRAII(OS, ParenthesisColor) << ')';
     }
   };
@@ -3028,7 +3031,7 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
       printArgumentList(E->getArgs());
     }
     if (auto rewritten = E->getRewritten()) {
-      OS << '\n';
+      OS << " rewritten=\n";
       printRec(rewritten);
     }
     PrintWithColorRAII(OS, ParenthesisColor) << ')';

lib/AST/ASTScope.cpp

@@ -76,6 +76,34 @@ void ASTScope::dumpOneScopeMapLocation(std::pair<unsigned, unsigned> lineCol) {
   impl->dumpOneScopeMapLocation(lineCol);
 }
 
+bool ASTScope::isLocWithinMacroDeclOrTopLevelMacroExpansionDeclScope(
+    SourceFile *sf, SourceLoc loc) {
+  // The innermost scope containing the loc.
+  auto *innermostScope = sf->getScope().impl->findInnermostEnclosingScope(
+      loc, nullptr);
+  auto isMDOrTopLevelMED = [&](const ASTScopeImpl *scope) -> bool {
+    if (auto *decl = scope->getDeclIfAny().getPtrOrNull()) {
+      if (auto *med = dyn_cast<MacroExpansionDecl>(decl)) {
+        if (med->getDeclContext()->getContextKind() ==
+            DeclContextKind::FileUnit)
+          return true;
+      } else if (auto *md = dyn_cast<MacroDecl>(decl)) {
+        return true;
+      }
+    }
+    return false;
+  };
+  // Climb up the tree until we find a matching scope.
+  const ASTScopeImpl *scope = innermostScope;
+  while (!isMDOrTopLevelMED(scope)) {
+    auto *parentScope = scope->getLookupParent().getPtrOrNull();
+    if (!parentScope || parentScope == scope)
+      return false;
+    scope = parentScope;
+  }
+  return true;
+}
+
 #pragma mark ASTScopeImpl
 
 
@@ -151,6 +179,7 @@ DEFINE_GET_CLASS_NAME(DifferentiableAttributeScope)
 DEFINE_GET_CLASS_NAME(SubscriptDeclScope)
 DEFINE_GET_CLASS_NAME(EnumElementScope)
 DEFINE_GET_CLASS_NAME(MacroDeclScope)
+DEFINE_GET_CLASS_NAME(MacroExpansionDeclScope)
 DEFINE_GET_CLASS_NAME(IfStmtScope)
 DEFINE_GET_CLASS_NAME(WhileStmtScope)
 DEFINE_GET_CLASS_NAME(GuardStmtScope)