[SourceKit] Use recored #if regions in "active regions" request

* Record each IfConfig clause location info in SourceFile
* Update SILProfiler to handle them
* Update SwiftLangSupport::findActiveRegionsInFile() to use the recorded
  regions instead of walking into AST to find #if regions

rdar://118082146

include/swift/AST/SourceFile.h

@@ -49,36 +49,41 @@ enum class RestrictedImportKind {
 using ImportAccessLevel = std::optional<AttributedImport<ImportedModule>>;
 
 /// Stores range information for a \c #if block in a SourceFile.
-class IfConfigRangeInfo final {
-  /// The range of the entire \c #if block, including \c #else and \c #endif.
-  CharSourceRange WholeRange;
+class IfConfigClauseRangeInfo final {
+public:
+  enum ClauseKind {
+    // Active '#if', '#elseif', or '#else' clause.
+    ActiveClause,
+    // Inactive '#if', '#elseif', or '#else' clause.
+    InactiveClause,
+    // '#endif' directive.
+    EndDirective,
+  };
 
-  /// The range of the active selected body, if there is one. This does not
-  /// include the outer syntax of the \c #if. This may be invalid, which
-  /// indicates there is no active body.
-  CharSourceRange ActiveBodyRange;
+private:
+  /// Source location of '#if', '#elseif', etc.
+  SourceLoc DirectiveLoc;
+  /// Character source location of body starts.
+  SourceLoc BodyLoc;
+  /// Location of the end of the body.
+  SourceLoc EndLoc;
+
+  ClauseKind Kind;
 
 public:
-  IfConfigRangeInfo(CharSourceRange wholeRange, CharSourceRange activeBodyRange)
-      : WholeRange(wholeRange), ActiveBodyRange(activeBodyRange) {
-    assert(wholeRange.getByteLength() > 0 && "Range must be non-empty");
-    assert(activeBodyRange.isInvalid() || wholeRange.contains(activeBodyRange));
+  IfConfigClauseRangeInfo(SourceLoc DirectiveLoc, SourceLoc BodyLoc,
+                          SourceLoc EndLoc, ClauseKind Kind)
+      : DirectiveLoc(DirectiveLoc), BodyLoc(BodyLoc), EndLoc(EndLoc),
+        Kind(Kind) {
+    assert(DirectiveLoc.isValid() && BodyLoc.isValid() && EndLoc.isValid());
   }
 
-  CharSourceRange getWholeRange() const { return WholeRange; }
-  SourceLoc getStartLoc() const { return WholeRange.getStart(); }
+  SourceLoc getStartLoc() const { return DirectiveLoc; }
+  CharSourceRange getDirectiveRange(const SourceManager &SM) const;
+  CharSourceRange getWholeRange(const SourceManager &SM) const;
+  CharSourceRange getBodyRange(const SourceManager &SM) const;
 
-  friend bool operator==(const IfConfigRangeInfo &lhs,
-                         const IfConfigRangeInfo &rhs) {
-    return lhs.WholeRange == rhs.WholeRange &&
-           lhs.ActiveBodyRange == rhs.ActiveBodyRange;
-  }
-
-  /// Retrieve the ranges produced by subtracting the active body range from
-  /// the whole range. This includes both inactive branches as well as the
-  /// other syntax of the \c #if.
-  SmallVector<CharSourceRange, 2>
-  getRangesWithoutActiveBody(const SourceManager &SM) const;
+  ClauseKind getKind() const { return Kind; }
 };
 
 /// A file containing Swift source code.
@@ -244,9 +249,9 @@ private:
   ParserStatePtr DelayedParserState =
       ParserStatePtr(/*ptr*/ nullptr, /*deleter*/ nullptr);
 
-  struct IfConfigRangesData {
+  struct IfConfigClauseRangesData {
     /// All the \c #if source ranges in this file.
-    std::vector<IfConfigRangeInfo> Ranges;
+    std::vector<IfConfigClauseRangeInfo> Ranges;
 
     /// Whether the elemnts in \c Ranges are sorted in source order within
     /// this file. We flip this to \c false any time a new range gets recorded,
@@ -255,7 +260,7 @@ private:
   };
 
   /// Stores all the \c #if source range info in this file.
-  mutable IfConfigRangesData IfConfigRanges;
+  mutable IfConfigClauseRangesData IfConfigClauseRanges;
 
   friend class HasImportsMatchingFlagRequest;
 
@@ -500,12 +505,16 @@ public:
      const_cast<SourceFile *>(this)->MissingImportedModules.insert(module);
   }
 
-  /// Record the source range info for a parsed \c #if block.
-  void recordIfConfigRangeInfo(IfConfigRangeInfo ranges);
+  /// Record the source range info for a parsed \c #if clause.
+  void recordIfConfigClauseRangeInfo(const IfConfigClauseRangeInfo &range);
 
-  /// Retrieve the source range infos for any \c #if blocks contained within a
+  /// Retrieve the source range info for any \c #if clauses in the file.
+  ArrayRef<IfConfigClauseRangeInfo> getIfConfigClauseRanges() const;
+
+  /// Retrieve the source range infos for any \c #if clauses contained within a
   /// given source range of this file.
-  ArrayRef<IfConfigRangeInfo> getIfConfigsWithin(SourceRange outer) const;
+  ArrayRef<IfConfigClauseRangeInfo>
+  getIfConfigClausesWithin(SourceRange outer) const;
 
   void getMissingImportedModules(
          SmallVectorImpl<ImportedModule> &imports) const override;

lib/AST/Module.cpp

@@ -2778,56 +2778,66 @@ SourceFile::getImportAccessLevel(const ModuleDecl *targetModule) const {
   return restrictiveImport;
 }
 
-SmallVector<CharSourceRange, 2>
-IfConfigRangeInfo::getRangesWithoutActiveBody(const SourceManager &SM) const {
-  SmallVector<CharSourceRange, 2> result;
-  if (ActiveBodyRange.isValid()) {
-    // Split the whole range by the active range.
-    result.emplace_back(SM, WholeRange.getStart(), ActiveBodyRange.getStart());
-    result.emplace_back(SM, ActiveBodyRange.getEnd(), WholeRange.getEnd());
-  } else {
-    // No active body, we just return the whole range.
-    result.push_back(WholeRange);
+CharSourceRange
+IfConfigClauseRangeInfo::getDirectiveRange(const SourceManager &SM) const {
+  return CharSourceRange(SM, DirectiveLoc, BodyLoc);
+}
+
+CharSourceRange
+IfConfigClauseRangeInfo::getBodyRange(const SourceManager &SM) const {
+  return CharSourceRange(SM, BodyLoc, EndLoc);
+}
+
+CharSourceRange
+IfConfigClauseRangeInfo::getWholeRange(const SourceManager &SM) const {
+  return CharSourceRange(SM, DirectiveLoc, EndLoc);
+}
+
+void SourceFile::recordIfConfigClauseRangeInfo(
+    const IfConfigClauseRangeInfo &range) {
+  IfConfigClauseRanges.Ranges.push_back(range);
+  IfConfigClauseRanges.IsSorted = false;
+}
+
+ArrayRef<IfConfigClauseRangeInfo> SourceFile::getIfConfigClauseRanges() const {
+  if (!IfConfigClauseRanges.IsSorted) {
+    auto &SM = getASTContext().SourceMgr;
+    // Sort the ranges if we need to.
+    llvm::sort(
+        IfConfigClauseRanges.Ranges, [&](const IfConfigClauseRangeInfo &lhs,
+                                         const IfConfigClauseRangeInfo &rhs) {
+          return SM.isBeforeInBuffer(lhs.getStartLoc(), rhs.getStartLoc());
+        });
+
+    // Be defensive and eliminate duplicates in case we've parsed twice.
+    auto newEnd = llvm::unique(
+        IfConfigClauseRanges.Ranges, [&](const IfConfigClauseRangeInfo &lhs,
+                                         const IfConfigClauseRangeInfo &rhs) {
+          if (lhs.getStartLoc() != rhs.getStartLoc())
+            return false;
+          assert(lhs.getBodyRange(SM) == rhs.getBodyRange(SM) &&
+                 "range changed on a re-parse?");
+          return true;
+        });
+    IfConfigClauseRanges.Ranges.erase(newEnd,
+                                      IfConfigClauseRanges.Ranges.end());
+    IfConfigClauseRanges.IsSorted = true;
   }
-  return result;
+
+  return IfConfigClauseRanges.Ranges;
 }
 
-void SourceFile::recordIfConfigRangeInfo(IfConfigRangeInfo ranges) {
-  IfConfigRanges.Ranges.push_back(ranges);
-  IfConfigRanges.IsSorted = false;
-}
-
-ArrayRef<IfConfigRangeInfo>
-SourceFile::getIfConfigsWithin(SourceRange outer) const {
+ArrayRef<IfConfigClauseRangeInfo>
+SourceFile::getIfConfigClausesWithin(SourceRange outer) const {
   auto &SM = getASTContext().SourceMgr;
   assert(SM.getRangeForBuffer(BufferID).contains(outer.Start) &&
          "Range not within this file?");
 
-  if (!IfConfigRanges.IsSorted) {
-    // Sort the ranges if we need to.
-    llvm::sort(IfConfigRanges.Ranges, [&](IfConfigRangeInfo lhs,
-                                          IfConfigRangeInfo rhs) {
-      return SM.isBeforeInBuffer(lhs.getStartLoc(), rhs.getStartLoc());
-    });
-
-    // Be defensive and eliminate duplicates in case we've parsed twice.
-    auto newEnd = std::unique(
-        IfConfigRanges.Ranges.begin(), IfConfigRanges.Ranges.end(),
-        [&](const IfConfigRangeInfo &lhs, const IfConfigRangeInfo &rhs) {
-          if (lhs.getWholeRange() != rhs.getWholeRange())
-            return false;
-
-          assert(lhs == rhs && "Active ranges changed on a re-parse?");
-          return true;
-        });
-    IfConfigRanges.Ranges.erase(newEnd, IfConfigRanges.Ranges.end());
-    IfConfigRanges.IsSorted = true;
-  }
-
   // First let's find the first #if that is after the outer start loc.
-  auto ranges = llvm::ArrayRef(IfConfigRanges.Ranges);
+  auto ranges = getIfConfigClauseRanges();
   auto lower = llvm::lower_bound(
-      ranges, outer.Start, [&](IfConfigRangeInfo range, SourceLoc loc) {
+      ranges, outer.Start,
+      [&](const IfConfigClauseRangeInfo &range, SourceLoc loc) {
         return SM.isBeforeInBuffer(range.getStartLoc(), loc);
       });
   if (lower == ranges.end() ||
@@ -2836,7 +2846,8 @@ SourceFile::getIfConfigsWithin(SourceRange outer) const {
   }
   // Next let's find the first #if that's after the outer end loc.
   auto upper = llvm::upper_bound(
-      ranges, outer.End, [&](SourceLoc loc, IfConfigRangeInfo range) {
+      ranges, outer.End,
+      [&](SourceLoc loc, const IfConfigClauseRangeInfo &range) {
         return SM.isBeforeInBuffer(loc, range.getStartLoc());
       });
   return llvm::ArrayRef(lower, upper - lower);

lib/Parse/ParseIfConfig.cpp

@@ -889,10 +889,12 @@ Result Parser::parseIfConfigRaw(
           ClauseLoc, Condition, isActive, IfConfigElementsRole::Skipped);
     }
 
-    // Record the active body range for the SourceManager.
-    if (shouldEvaluate && isActive) {
-      assert(!activeBodyRange.isValid() && "Multiple active regions?");
-      activeBodyRange = CharSourceRange(SourceMgr, bodyStart, Tok.getLoc());
+    // Record the clause range info in SourceFile.
+    if (shouldEvaluate) {
+      auto kind = isActive ? IfConfigClauseRangeInfo::ActiveClause
+                           : IfConfigClauseRangeInfo::InactiveClause;
+      SF.recordIfConfigClauseRangeInfo(
+          {ClauseLoc, bodyStart, Tok.getLoc(), kind});
     }
 
     if (Tok.isNot(tok::pound_elseif, tok::pound_else))
@@ -905,11 +907,11 @@ Result Parser::parseIfConfigRaw(
   SourceLoc EndLoc;
   bool HadMissingEnd = parseEndIfDirective(EndLoc);
 
-  // Record the #if ranges on the SourceManager.
+  // Record the '#end' ranges in SourceFile.
   if (!HadMissingEnd && shouldEvaluate) {
-    auto wholeRange = Lexer::getCharSourceRangeFromSourceRange(
-        SourceMgr, SourceRange(startLoc, EndLoc));
-    SF.recordIfConfigRangeInfo({wholeRange, activeBodyRange});
+    SourceLoc EndOfEndLoc = getEndOfPreviousLoc();
+    SF.recordIfConfigClauseRangeInfo({EndLoc, EndOfEndLoc, EndOfEndLoc,
+                                      IfConfigClauseRangeInfo::EndDirective});
   }
   return finish(EndLoc, HadMissingEnd);
 }

lib/SIL/IR/SILProfiler.cpp

@@ -1165,14 +1165,37 @@ public:
                                            Counter.getLLVMCounter()));
     }
     // Add any skipped regions present in the outer range.
-    for (auto IfConfig : SF->getIfConfigsWithin(OuterRange)) {
-      for (auto SkipRange : IfConfig.getRangesWithoutActiveBody(SM)) {
-        auto Start = SM.getLineAndColumnInBuffer(SkipRange.getStart());
-        auto End = SM.getLineAndColumnInBuffer(SkipRange.getEnd());
-        assert(Start.first <= End.first && "region start and end out of order");
-        Regions.push_back(MappedRegion::skipped(Start.first, Start.second,
-                                                End.first, End.second));
+    for (auto clause : SF->getIfConfigClausesWithin(OuterRange)) {
+      CharSourceRange SkipRange;
+      switch (clause.getKind()) {
+      case IfConfigClauseRangeInfo::ActiveClause:
+      case IfConfigClauseRangeInfo::EndDirective:
+        SkipRange = clause.getDirectiveRange(SM);
+        break;
+      case IfConfigClauseRangeInfo::InactiveClause:
+        SkipRange = clause.getWholeRange(SM);
+        break;
       }
+      if (SkipRange.getByteLength() == 0)
+        continue;
+
+      auto Start = SM.getLineAndColumnInBuffer(SkipRange.getStart());
+      auto End = SM.getLineAndColumnInBuffer(SkipRange.getEnd());
+      assert(Start.first <= End.first && "region start and end out of order");
+
+      // If this is consecutive with the last one, expand it.
+      if (!Regions.empty()) {
+        auto &last = Regions.back();
+        if (last.RegionKind == MappedRegion::Kind::Skipped &&
+            last.EndLine == Start.first && last.EndCol == Start.second) {
+          last.EndLine = End.first;
+          last.EndCol = End.second;
+          continue;
+        }
+      }
+
+      Regions.push_back(MappedRegion::skipped(Start.first, Start.second,
+                                              End.first, End.second));
     }
     return SILCoverageMap::create(M, SF, Filename, Name, PGOFuncName, Hash,
                                   Regions, CounterBuilder.getExpressions());

test/SourceKit/ActiveRegions/nested2.swift

@@ -0,0 +1,19 @@
+func foo() -> Int {
+  #if true
+  #if true
+  return 1
+  #else
+  return 2
+  #endif
+  #else
+  return 3
+  #endif
+}
+
+// RUN: %sourcekitd-test -req=active-regions %s -- -D FLAG_1 -module-name active_regions %s | %FileCheck -check-prefix=CHECK1 %s
+// CHECK1: START IF CONFIGS
+// CHECK1-NEXT: 2:3 - active
+// CHECK1-NEXT: 3:3 - active
+// CHECK1-NEXT: 5:3 - inactive
+// CHECK1-NEXT: 8:3 - inactive
+// CHECK1-NEXT: END IF CONFIGS

test/SourceKit/ActiveRegions/postfix.swift

@@ -0,0 +1,18 @@
+struct Value {
+  var prop1: Int = 0
+  var prop2: Int = 0
+}
+func test(value: Value) {
+  let _ = value
+#if false
+    .prop1
+#else
+    .prop2
+#endif
+}
+
+// RUN: %sourcekitd-test -req=active-regions %s -- -module-name active_regions %s | %FileCheck -check-prefix=CHECK1 %s
+// CHECK1: START IF CONFIGS
+// CHECK1-NEXT: 7:1 - inactive
+// CHECK1-NEXT: 9:1 - active
+// CHECK1-NEXT: END IF CONFIGS

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -2652,36 +2652,6 @@ void SwiftLangSupport::findRelatedIdentifiersInFile(
 // SwiftLangSupport::findActiveRegionsInFile
 //===----------------------------------------------------------------------===//
 
-namespace {
-class IfConfigScanner : public SourceEntityWalker {
-  unsigned BufferID = -1;
-  SmallVectorImpl<IfConfigInfo> &Infos;
-  bool Cancelled = false;
-
-public:
-  explicit IfConfigScanner(unsigned BufferID,
-                           SmallVectorImpl<IfConfigInfo> &Infos)
-      : BufferID(BufferID), Infos(Infos) {}
-
-private:
-  bool walkToDeclPre(Decl *D, CharSourceRange Range) override {
-    if (Cancelled)
-      return false;
-
-    if (auto *IfDecl = dyn_cast<IfConfigDecl>(D)) {
-      for (auto &Clause : IfDecl->getClauses()) {
-        unsigned Offset = D->getASTContext().SourceMgr.getLocOffsetInBuffer(
-            Clause.Loc, BufferID);
-        Infos.emplace_back(Offset, Clause.isActive);
-      }
-    }
-
-    return true;
-  }
-};
-
-} // end anonymous namespace
-
 void SwiftLangSupport::findActiveRegionsInFile(
     StringRef PrimaryFilePath, StringRef InputBufferName,
     ArrayRef<const char *> Args, SourceKitCancellationToken CancellationToken,
@@ -2717,16 +2687,25 @@ void SwiftLangSupport::findActiveRegionsInFile(
         return;
       }
 
-      SmallVector<IfConfigInfo> Configs;
-      IfConfigScanner Scanner(*SF->getBufferID(), Configs);
-      Scanner.walk(SF);
+      auto &SM = SF->getASTContext().SourceMgr;
+      auto BufferID = *SF->getBufferID();
 
-      // Sort by offset so nested decls are reported
-      // in source order (not tree order).
-      llvm::sort(Configs,
-                 [](const IfConfigInfo &LHS, const IfConfigInfo &RHS) -> bool {
-                   return LHS.Offset < RHS.Offset;
-                 });
+      SmallVector<IfConfigInfo> Configs;
+      for (auto &range : SF->getIfConfigClauseRanges()) {
+        bool isActive = false;
+        switch (range.getKind()) {
+        case IfConfigClauseRangeInfo::ActiveClause:
+          isActive = true;
+          break;
+        case IfConfigClauseRangeInfo::InactiveClause:
+          isActive = false;
+          break;
+        case IfConfigClauseRangeInfo::EndDirective:
+          continue;
+        }
+        auto offset = SM.getLocOffsetInBuffer(range.getStartLoc(), BufferID);
+        Configs.emplace_back(offset, isActive);
+      }
       ActiveRegionsInfo Info;
       Info.Configs = Configs;
       Receiver(RequestResult<ActiveRegionsInfo>::fromResult(Info));