[Parse] Store #if ranges on the SourceFile

Record any `#if`'s we encounter, and make a note
of their active region.

Author: hamishknight

include/swift/AST/SourceFile.h

@@ -48,6 +48,39 @@ enum class RestrictedImportKind {
 /// Import that limits the access level of imported entities.
 using ImportAccessLevel = llvm::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;
+
+  /// 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;
+
+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));
+  }
+
+  CharSourceRange getWholeRange() const { return WholeRange; }
+  SourceLoc getStartLoc() const { return WholeRange.getStart(); }
+
+  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;
+};
+
 /// A file containing Swift source code.
 ///
 /// This is a .swift or .sil file (or a virtual file, such as the contents of
@@ -209,6 +242,19 @@ class SourceFile final : public FileUnit {
   ParserStatePtr DelayedParserState =
       ParserStatePtr(/*ptr*/ nullptr, /*deleter*/ nullptr);
 
+  struct IfConfigRangesData {
+    /// All the \c #if source ranges in this file.
+    std::vector<IfConfigRangeInfo> 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,
+    /// and lazily do the sorting when doing a query.
+    bool IsSorted = false;
+  };
+
+  /// Stores all the \c #if source range info in this file.
+  mutable IfConfigRangesData IfConfigRanges;
+
   friend class HasImportsMatchingFlagRequest;
 
   /// Indicates which import options have valid caches. Storage for
@@ -448,6 +494,13 @@ class SourceFile final : public FileUnit {
      const_cast<SourceFile *>(this)->MissingImportedModules.insert(module);
   }
 
+  /// Record the source range info for a parsed \c #if block.
+  void recordIfConfigRangeInfo(IfConfigRangeInfo ranges);
+
+  /// Retrieve the source range infos for any \c #if blocks contained within a
+  /// given source range of this file.
+  ArrayRef<IfConfigRangeInfo> getIfConfigsWithin(SourceRange outer) const;
+
   void getMissingImportedModules(
          SmallVectorImpl<ImportedModule> &imports) const override;
 

lib/AST/Module.cpp

@@ -3326,6 +3326,70 @@ 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);
+  }
+  return result;
+}
+
+void SourceFile::recordIfConfigRangeInfo(IfConfigRangeInfo ranges) {
+  IfConfigRanges.Ranges.push_back(ranges);
+  IfConfigRanges.IsSorted = false;
+}
+
+ArrayRef<IfConfigRangeInfo>
+SourceFile::getIfConfigsWithin(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::makeArrayRef(IfConfigRanges.Ranges);
+  auto lower = llvm::lower_bound(
+      ranges, outer.Start, [&](IfConfigRangeInfo range, SourceLoc loc) {
+        return SM.isBeforeInBuffer(range.getStartLoc(), loc);
+      });
+  if (lower == ranges.end() ||
+      SM.isBeforeInBuffer(outer.End, lower->getStartLoc())) {
+    return {};
+  }
+  // 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) {
+        return SM.isBeforeInBuffer(loc, range.getStartLoc());
+      });
+  return llvm::makeArrayRef(lower, upper - lower);
+}
+
 void ModuleDecl::setPackageName(Identifier name) {
   Package = PackageUnit::create(name, *this, getASTContext());
 }

lib/Parse/ParseIfConfig.cpp

@@ -785,6 +785,7 @@ Result Parser::parseIfConfigRaw(
                             bool isActive, IfConfigElementsRole role)>
       parseElements,
     llvm::function_ref<Result(SourceLoc endLoc, bool hadMissingEnd)> finish) {
+  auto startLoc = Tok.getLoc();
   assert(Tok.is(tok::pound_if));
 
   Parser::StructureMarkerRAII ParsingDecl(
@@ -824,6 +825,7 @@ Result Parser::parseIfConfigRaw(
 
   bool foundActive = false;
   bool isVersionCondition = false;
+  CharSourceRange activeBodyRange;
   while (1) {
     bool isElse = Tok.is(tok::pound_else);
     SourceLoc ClauseLoc = consumeToken();
@@ -881,6 +883,7 @@ Result Parser::parseIfConfigRaw(
     }
 
     // Parse elements
+    auto bodyStart = Lexer::getLocForEndOfToken(SourceMgr, PreviousLoc);
     llvm::SaveAndRestore<bool> S(InInactiveClauseEnvironment,
                                  InInactiveClauseEnvironment || !isActive);
     // Disable updating the interface hash inside inactive blocks.
@@ -899,6 +902,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());
+    }
+
     if (Tok.isNot(tok::pound_elseif, tok::pound_else))
       break;
 
@@ -909,6 +918,12 @@ Result Parser::parseIfConfigRaw(
   SourceLoc EndLoc;
   bool HadMissingEnd = parseEndIfDirective(EndLoc);
 
+  // Record the #if ranges on the SourceManager.
+  if (!HadMissingEnd && shouldEvaluate) {
+    auto wholeRange = Lexer::getCharSourceRangeFromSourceRange(
+        SourceMgr, SourceRange(startLoc, EndLoc));
+    SF.recordIfConfigRangeInfo({wholeRange, activeBodyRange});
+  }
   return finish(EndLoc, HadMissingEnd);
 }