AST: Sort TypeRefinementContexts for better lookup performance

include/swift/AST/TypeRefinementContext.h

@@ -185,6 +185,8 @@ class TypeRefinementContext : public ASTAllocated<TypeRefinementContext> {
     unsigned needsExpansion : 1;
   } LazyInfo = {};
 
+  void verify(const TypeRefinementContext *parent, ASTContext &ctx) const;
+
   TypeRefinementContext(ASTContext &Ctx, IntroNode Node,
                         TypeRefinementContext *Parent, SourceRange SrcRange,
                         const AvailabilityRange &Info,
@@ -290,10 +292,7 @@ class TypeRefinementContext : public ASTAllocated<TypeRefinementContext> {
   }
 
   /// Adds a child refinement context.
-  void addChild(TypeRefinementContext *Child) {
-    assert(Child->getSourceRange().isValid());
-    Children.push_back(Child);
-  }
+  void addChild(TypeRefinementContext *Child, ASTContext &Ctx);
 
   /// Returns the inner-most TypeRefinementContext descendant of this context
   /// for the given source location.
@@ -306,6 +305,10 @@ class TypeRefinementContext : public ASTAllocated<TypeRefinementContext> {
     LazyInfo.needsExpansion = needsExpansion;
   }
 
+  /// Recursively check the tree for integrity. If any errors are found, emits
+  /// diagnosticts to stderr and aborts.
+  void verify(ASTContext &ctx);
+
   SWIFT_DEBUG_DUMPER(dump(SourceManager &SrcMgr));
   void dump(raw_ostream &OS, SourceManager &SrcMgr) const;
   void print(raw_ostream &OS, SourceManager &SrcMgr, unsigned Indent = 0) const;

include/swift/Basic/SourceManager.h

@@ -55,6 +55,7 @@ class GeneratedSourceInfo {
   enum Kind {
 #define MACRO_ROLE(Name, Description) Name##MacroExpansion,
 #include "swift/Basic/MacroRoles.def"
+#undef MACRO_ROLE
 
     /// A new function body that is replacing an existing function body.
     ReplacedFunctionBody,
@@ -66,6 +67,23 @@ class GeneratedSourceInfo {
     DefaultArgument,
   } kind;
 
+  static StringRef kindToString(GeneratedSourceInfo::Kind kind) {
+    switch (kind) {
+#define MACRO_ROLE(Name, Description)                                          \
+  case Name##MacroExpansion:                                                   \
+    return #Name "MacroExpansion";
+#include "swift/Basic/MacroRoles.def"
+#undef MACRO_ROLE
+    case ReplacedFunctionBody:
+      return "ReplacedFunctionBody";
+    case PrettyPrinted:
+      return "PrettyPrinted";
+    case DefaultArgument:
+      return "DefaultArgument";
+    }
+    llvm_unreachable("Invalid kind");
+  }
+
   /// The source range in the enclosing buffer where this source was generated.
   ///
   /// This could point at a macro expansion or at the implicit location at

lib/AST/ASTVerifier.cpp

@@ -36,6 +36,7 @@
 #include "swift/AST/SourceFile.h"
 #include "swift/AST/Stmt.h"
 #include "swift/AST/TypeCheckRequests.h"
+#include "swift/AST/TypeRefinementContext.h"
 #include "swift/AST/TypeRepr.h"
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/SourceManager.h"
@@ -3882,23 +3883,15 @@ void swift::verify(SourceFile &SF) {
     return;
   Verifier verifier(SF, &SF);
   SF.walk(verifier);
-}
 
-bool swift::shouldVerify(const Decl *D, const ASTContext &Context) {
-  if (!shouldVerifyGivenContext(Context))
-    return false;
-
-  if (const auto *ED = dyn_cast<ExtensionDecl>(D)) {
-    return shouldVerify(ED->getExtendedNominal(), Context);
-  }
-
-  const auto *VD = dyn_cast<ValueDecl>(D);
-  if (!VD) {
-    // Verify declarations without names everywhere.
-    return true;
+  // Verify the TypeRefinementContext hierarchy.
+  if (auto TRC = SF.getTypeRefinementContext()) {
+    TRC->verify(SF.getASTContext());
   }
+}
 
-  return true;
+bool swift::shouldVerify(const Decl *D, const ASTContext &Context) {
+  return shouldVerifyGivenContext(Context);
 }
 
 void swift::verify(Decl *D) {

lib/AST/TypeRefinementContext.cpp

@@ -14,14 +14,15 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "swift/AST/TypeRefinementContext.h"
+
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/Decl.h"
-#include "swift/AST/Module.h"
-#include "swift/AST/Stmt.h"
 #include "swift/AST/Expr.h"
+#include "swift/AST/Module.h"
 #include "swift/AST/SourceFile.h"
+#include "swift/AST/Stmt.h"
 #include "swift/AST/TypeCheckRequests.h"
-#include "swift/AST/TypeRefinementContext.h"
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/SourceManager.h"
 
@@ -35,7 +36,7 @@ TypeRefinementContext::TypeRefinementContext(
       ExplicitAvailabilityInfo(ExplicitInfo) {
   if (Parent) {
     assert(SrcRange.isValid());
-    Parent->addChild(this);
+    Parent->addChild(this, Ctx);
     assert(Info.isContainedIn(Parent->getAvailabilityInfo()));
   }
   Ctx.addDestructorCleanup(Children);
@@ -169,52 +170,64 @@ TypeRefinementContext::createForWhileStmtBody(ASTContext &Ctx, WhileStmt *S,
       Ctx, S, Parent, S->getBody()->getSourceRange(), Info, /* ExplicitInfo */Info);
 }
 
-/// Determine whether the child location is somewhere within the parent
-/// range.
-static bool rangeContainsTokenLocWithGeneratedSource(
-    SourceManager &sourceMgr, SourceRange parentRange, SourceLoc childLoc) {
-  if (sourceMgr.rangeContainsTokenLoc(parentRange, childLoc))
-    return true;
-
-  auto childBuffer = sourceMgr.findBufferContainingLoc(childLoc);
-  while (!sourceMgr.rangeContainsTokenLoc(
-      sourceMgr.getRangeForBuffer(childBuffer), parentRange.Start)) {
-    auto *info = sourceMgr.getGeneratedSourceInfo(childBuffer);
-    if (!info)
-      return false;
-
-    childLoc = info->originalSourceRange.getStart();
-    if (childLoc.isInvalid())
-      return false;
-
-    childBuffer = sourceMgr.findBufferContainingLoc(childLoc);
+void TypeRefinementContext::addChild(TypeRefinementContext *Child,
+                                     ASTContext &Ctx) {
+  assert(Child->getSourceRange().isValid());
+
+  // Handle the first child.
+  if (Children.empty()) {
+    Children.push_back(Child);
+    return;
   }
 
-  return sourceMgr.rangeContainsTokenLoc(parentRange, childLoc);
+  // Handle a child that is ordered after the existing children (this should be
+  // the common case).
+  auto &srcMgr = Ctx.SourceMgr;
+  if (srcMgr.isBefore(Children.back()->getSourceRange().Start,
+                      Child->getSourceRange().Start)) {
+    Children.push_back(Child);
+    return;
+  }
+
+  // Insert the child amongst the existing sorted children.
+  auto iter = std::upper_bound(
+      Children.begin(), Children.end(), Child,
+      [&srcMgr](TypeRefinementContext *lhs, TypeRefinementContext *rhs) {
+        return srcMgr.isBefore(lhs->getSourceRange().Start,
+                               rhs->getSourceRange().Start);
+      });
+
+  Children.insert(iter, Child);
 }
 
 TypeRefinementContext *
 TypeRefinementContext::findMostRefinedSubContext(SourceLoc Loc,
                                                  ASTContext &Ctx) {
   assert(Loc.isValid());
 
-  if (SrcRange.isValid() &&
-      !rangeContainsTokenLocWithGeneratedSource(Ctx.SourceMgr, SrcRange, Loc))
+  if (SrcRange.isValid() && !Ctx.SourceMgr.containsTokenLoc(SrcRange, Loc))
     return nullptr;
 
   auto expandedChildren = evaluateOrDefault(
       Ctx.evaluator, ExpandChildTypeRefinementContextsRequest{this}, {});
 
-  // For the moment, we perform a linear search here, but we can and should
-  // do something more efficient.
-  for (TypeRefinementContext *Child : expandedChildren) {
-    if (auto *Found = Child->findMostRefinedSubContext(Loc, Ctx)) {
-      return Found;
-    }
+  // Do a binary search to find the first child with a source range that
+  // ends after the given location.
+  auto iter = std::lower_bound(
+      expandedChildren.begin(), expandedChildren.end(), Loc,
+      [&Ctx](TypeRefinementContext *context, SourceLoc loc) {
+        return Ctx.SourceMgr.isBefore(context->getSourceRange().End, loc);
+      });
+
+  // Check whether the matching child or any of its descendants contain
+  // the given location.
+  if (iter != expandedChildren.end()) {
+    if (auto found = (*iter)->findMostRefinedSubContext(Loc, Ctx))
+      return found;
   }
 
-  // Loc is in this context's range but not in any child's, so this context
-  // must be the inner-most context.
+  // The location is in this context's range but not in any child's, so this
+  // context must be the innermost context.
   return this;
 }
 
@@ -391,7 +404,24 @@ void TypeRefinementContext::print(raw_ostream &OS, SourceManager &SrcMgr,
   auto R = getSourceRange();
   if (R.isValid()) {
     OS << " src_range=";
-    R.print(OS, SrcMgr, /*PrintText=*/false);
+
+    if (getReason() != Reason::Root) {
+      R.print(OS, SrcMgr, /*PrintText=*/false);
+    } else if (auto info = SrcMgr.getGeneratedSourceInfo(
+                   Node.getAsSourceFile()->getBufferID())) {
+      info->originalSourceRange.print(OS, SrcMgr, /*PrintText=*/false);
+    } else {
+      OS << "<unknown>";
+    }
+  }
+
+  if (getReason() == Reason::Root) {
+    if (auto info = SrcMgr.getGeneratedSourceInfo(
+            Node.getAsSourceFile()->getBufferID())) {
+      OS << " generated_kind=" << GeneratedSourceInfo::kindToString(info->kind);
+    } else {
+      OS << " file=" << Node.getAsSourceFile()->getFilename().str();
+    }
   }
 
   if (!ExplicitAvailabilityInfo.isAlwaysAvailable())
@@ -462,3 +492,78 @@ void ExpandChildTypeRefinementContextsRequest::cacheResult(
   TRC->Children = children;
   TRC->setNeedsExpansion(false);
 }
+
+/// Emit an error message, dump each context with its corresponding label, and
+/// abort.
+static void
+verificationError(ASTContext &ctx, llvm::StringRef msg,
+                  std::initializer_list<
+                      std::pair<const char *, const TypeRefinementContext *>>
+                      labelsAndNodes) {
+  llvm::errs() << msg << "\n";
+  for (auto pair : labelsAndNodes) {
+    auto label = std::get<0>(pair);
+    auto context = std::get<1>(pair);
+    llvm::errs() << label << ":\n";
+    context->print(llvm::errs(), ctx.SourceMgr);
+  }
+  abort();
+}
+
+void TypeRefinementContext::verify(const TypeRefinementContext *parent,
+                                   ASTContext &ctx) const {
+  // Verify the children first.
+  for (auto child : Children) {
+    child->verify(this, ctx);
+  }
+
+  // Verify that the children are in sorted order and that their source ranges
+  // do not overlap.
+  auto &srcMgr = ctx.SourceMgr;
+  if (Children.size() > 1) {
+    auto const *previous = Children.front();
+    for (auto const *current : ArrayRef(Children).drop_front()) {
+      if (!srcMgr.isAtOrBefore(previous->getSourceRange().Start,
+                               current->getSourceRange().Start))
+        verificationError(
+            ctx, "out of order children",
+            {{"child 1", previous}, {"child 2", current}, {"parent", this}});
+
+      if (srcMgr.containsLoc(previous->getSourceRange(),
+                             current->getSourceRange().Start))
+        verificationError(
+            ctx, "overlapping children",
+            {{"child 1", previous}, {"child 2", current}, {"parent", this}});
+
+      previous = current;
+    }
+  }
+
+  // Only root nodes are allowed to have no parent.
+  if (!parent) {
+    if (getReason() != Reason::Root)
+      verificationError(ctx, "interior node without parent", {{"node", this}});
+    return;
+  }
+
+  // All nodes with a parent must have a valid source range.
+  if (!SrcRange.isValid())
+    verificationError(ctx, "invalid source range", {{"node", this}});
+
+  if (getReason() != Reason::Root) {
+    auto parentRange = parent->SrcRange;
+    if (parentRange.isValid() &&
+        !(srcMgr.isAtOrBefore(parentRange.Start, SrcRange.Start) &&
+          srcMgr.isAtOrBefore(SrcRange.End, parentRange.End)))
+      verificationError(ctx, "child source range not contained",
+                        {{"child", this}, {"parent", this}});
+  }
+
+  if (!AvailabilityInfo.isContainedIn(parent->AvailabilityInfo))
+    verificationError(ctx, "child availability range not contained",
+                      {{"child", this}, {"parent", this}});
+}
+
+void TypeRefinementContext::verify(ASTContext &ctx) {
+  verify(nullptr, ctx);
+}