Migrate VariableTypeCollector to SourceRange

- Add SourceRange::contains and SourceRange::overlaps
- Use SourceRange in VariableTypeCollector

Author: fwcd

include/swift/Basic/SourceLoc.h

@@ -123,6 +123,14 @@ class SourceRange {
   /// includes both this range and the other one.
   void widen(SourceRange Other);
 
+  /// Checks whether this range contains the given location. Note that the given
+  /// location should correspond to the start of a token, since locations inside
+  /// the last token may be considered outside the range by this function.
+  bool contains(SourceLoc Loc) const;
+
+  /// Checks whether this range overlaps with the given range.
+  bool overlaps(SourceRange Other) const;
+
   bool operator==(const SourceRange &other) const {
     return Start == other.Start && End == other.End;
   }

include/swift/Sema/IDETypeChecking.h

@@ -225,13 +225,13 @@ namespace swift {
   };
 
   /// Collect type information for every variable declaration in \c SF
-  /// within the given range, defined by \c Offset and \c Length.
+  /// within the given range.
   /// All types will be printed to \c OS and the type offsets of the
   /// \c VariableTypeInfos will index into the string that backs this
   /// stream.
-  void collectVariableType(
-      SourceFile &SF, Optional<unsigned> Offset, Optional<unsigned> Length,
-      std::vector<VariableTypeInfo> &Scratch, llvm::raw_ostream &OS);
+  void collectVariableType(SourceFile &SF, SourceRange Range,
+                           std::vector<VariableTypeInfo> &Scratch,
+                           llvm::raw_ostream &OS);
 
   /// FIXME: All of the below goes away once CallExpr directly stores its
   /// arguments.

lib/Basic/SourceLoc.cpp

@@ -270,6 +270,15 @@ void SourceRange::widen(SourceRange Other) {
     End = Other.End;
 }
 
+bool SourceRange::contains(SourceLoc Loc) const {
+  return Start.Value.getPointer() <= Loc.Value.getPointer() &&
+         Loc.Value.getPointer() <= End.Value.getPointer();
+}
+
+bool SourceRange::overlaps(SourceRange Other) const {
+  return contains(Other.Start) || Other.contains(Start);
+}
+
 void SourceLoc::printLineAndColumn(raw_ostream &OS, const SourceManager &SM,
                                    unsigned BufferID) const {
   if (isInvalid()) {

lib/IDE/IDETypeChecking.cpp

@@ -738,11 +738,8 @@ class VariableTypeCollector : public SourceEntityWalker {
   SourceManager &SM;
   unsigned int BufferId;
 
-  /// The start offset of the range in which variable types are to be collected.
-  Optional<unsigned> TotalOffset;
-
-  /// The length of the range in which variable types are to be collected.
-  Optional<unsigned> TotalLength;
+  /// The range in which variable types are to be collected.
+  SourceRange TotalRange;
 
   /// The output vector for VariableTypeInfos emitted during traversal.
   std::vector<VariableTypeInfo> &Results;
@@ -768,26 +765,31 @@ class VariableTypeCollector : public SourceEntityWalker {
     return {TypeOffsets[PrintedType], PrintedType.size()};
   }
 
+  /// Checks whether the given range overlaps the total range in which we
+  /// collect variable types.
+  bool overlapsTotalRange(SourceRange Range) {
+    return TotalRange.isInvalid() || Range.overlaps(TotalRange);
+  }
+
 public:
-  VariableTypeCollector(SourceFile &SF, Optional<unsigned> Offset,
-                        Optional<unsigned> Length,
+  VariableTypeCollector(SourceFile &SF, SourceRange Range,
                         std::vector<VariableTypeInfo> &Results,
                         llvm::raw_ostream &OS)
       : SM(SF.getASTContext().SourceMgr), BufferId(*SF.getBufferID()),
-        TotalOffset(Offset), TotalLength(Length), Results(Results), OS(OS) {}
+        TotalRange(Range), Results(Results), OS(OS) {}
 
-  bool walkToDeclPre(Decl *D, CharSourceRange Range) override {
-    if (Range.isInvalid()) {
+  bool walkToDeclPre(Decl *D, CharSourceRange DeclRange) override {
+    if (DeclRange.isInvalid()) {
       return false;
     }
-    unsigned Offset = SM.getLocOffsetInBuffer(Range.getStart(), BufferId);
-    unsigned Length = Range.getByteLength();
     // Skip this declaration and its subtree if outside the range
-    if (TotalOffset.hasValue() && TotalLength.hasValue() &&
-        (Offset < *TotalOffset || Offset >= (*TotalOffset + *TotalLength))) {
+    if (!overlapsTotalRange(D->getSourceRange())) {
       return false;
     }
     if (auto VD = dyn_cast<VarDecl>(D)) {
+      unsigned VarOffset =
+          SM.getLocOffsetInBuffer(DeclRange.getStart(), BufferId);
+      unsigned VarLength = DeclRange.getByteLength();
       // Print the type to a temporary buffer
       SmallString<64> Buffer;
       {
@@ -802,21 +804,36 @@ class VariableTypeCollector : public SourceEntityWalker {
       bool HasExplicitType =
           VD->getTypeReprOrParentPatternTypeRepr() != nullptr;
       // Add the type information to the result list.
-      Results.emplace_back(Offset, Length, HasExplicitType, Ty.first);
+      Results.emplace_back(VarOffset, VarLength, HasExplicitType, Ty.first);
     }
     return true;
   }
+
+  bool walkToStmtPre(Stmt *S) override {
+    // Skip this statement and its subtree if outside the range
+    return overlapsTotalRange(S->getSourceRange());
+  }
+
+  bool walkToExprPre(Expr *E) override {
+    // Skip this expression and its subtree if outside the range
+    return overlapsTotalRange(E->getSourceRange());
+  }
+
+  bool walkToPatternPre(Pattern *P) override {
+    // Skip this pattern and its subtree if outside the range
+    return overlapsTotalRange(P->getSourceRange());
+  }
 };
 
 VariableTypeInfo::VariableTypeInfo(uint32_t Offset, uint32_t Length,
                                    bool HasExplicitType, uint32_t TypeOffset)
     : Offset(Offset), Length(Length), HasExplicitType(HasExplicitType),
       TypeOffset(TypeOffset) {}
 
-void swift::collectVariableType(
-    SourceFile &SF, Optional<unsigned> Offset, Optional<unsigned> Length,
-    std::vector<VariableTypeInfo> &Scratch, llvm::raw_ostream &OS) {
-  VariableTypeCollector Walker(SF, Offset, Length, Scratch, OS);
+void swift::collectVariableType(SourceFile &SF, SourceRange Range,
+                                std::vector<VariableTypeInfo> &Scratch,
+                                llvm::raw_ostream &OS) {
+  VariableTypeCollector Walker(SF, Range, Scratch, OS);
   Walker.walk(SF);
 }
 

tools/SourceKit/lib/SwiftLang/SwiftSourceDocInfo.cpp

@@ -2396,12 +2396,29 @@ void SwiftLangSupport::collectVariableTypes(
         : Receiver(std::move(Receiver)), Offset(Offset), Length(Length) {}
 
     void handlePrimaryAST(ASTUnitRef AstUnit) override {
-      auto *SF = AstUnit->getCompilerInstance().getPrimarySourceFile();
+      auto &CompInst = AstUnit->getCompilerInstance();
+      auto *SF = CompInst.getPrimarySourceFile();
+
+      // Construct the range for which variable types are to be queried. If
+      // offset/length are unset, the (default) range will be used, which
+      // corresponds to the entire document.
+      SourceRange Range;
+      if (Offset.hasValue() && Length.hasValue()) {
+        auto &SM = CompInst.getSourceMgr();
+        unsigned BufferID = SF->getBufferID().getValue();
+        SourceLoc Start = Lexer::getLocForStartOfToken(SM, BufferID, *Offset);
+        SourceLoc End =
+            Lexer::getLocForStartOfToken(SM, BufferID, *Offset + *Length);
+        Range = SourceRange(Start, End);
+      }
+
       std::vector<VariableTypeInfo> Infos;
       std::string TypeBuffer;
       llvm::raw_string_ostream OS(TypeBuffer);
       VariableTypesInFile Result;
-      collectVariableType(*SF, Offset, Length, Infos, OS);
+
+      collectVariableType(*SF, Range, Infos, OS);
+
       for (auto Info : Infos) {
         Result.Results.push_back({Info.Offset, Info.Length, Info.TypeOffset, Info.HasExplicitType});
       }