[flang] Faster implementation of FindScope()

[klausler]

Summary

The utility semantics::SemanticsContext::FindScope() maps a contiguous
range of cooked source characters to the innermost Scope containing
them. Its implementation is unacceptably slow on large (tens of
thousands of lines) source files with many program units; it traverses
each level of the scope tree linearly.

Replace this implementation with a single instance of std::multimap<>
used as an index from each Scope's source range back to the Scope.

Compilation time with "-fsyntax-only" on the 50,000-line test case
that motivated this change drops from 4.36s to 3.72s, and FindScope()
no longer stands out egregiously in the profile.

[PeteSteinfeld]

All builds and tests correctly and looks good.

flang/include/flang/Parser/provenance.h

public:
   void Register(CookedSource &);
 
 private:
   AllSources &allSources_;
   std::list<CookedSource> cooked_; // owns all CookedSource instances
   std::map<CharBlock, const CookedSource &, CharBlockPointerComparator> index_;
 };
 
-// For use as a Comparator for maps, sets, sorting, &c.
-class CharBlockComparator {
-public:
-  explicit CharBlockComparator(const AllCookedSources &all) : all_{all} {}
-  bool operator()(CharBlock x, CharBlock y) const {
-    return all_.Precedes(x, y);
-  }
-
-private:
-  const AllCookedSources &all_;
-};
-
 } // namespace Fortran::parser
 #endif // FORTRAN_PARSER_PROVENANCE_H_

flang/include/flang/Semantics/scope.h

public:
   }
 
   bool hasSAVE() const { return hasSAVE_; }
   void set_hasSAVE(bool yes = true) { hasSAVE_ = yes; }
 
   // The range of the source of this and nested scopes.
   const parser::CharBlock &sourceRange() const { return sourceRange_; }
   void AddSourceRange(parser::CharBlock);
-  // Find the smallest scope under this one that contains source
-  const Scope *FindScope(parser::CharBlock) const;
-  Scope *FindScope(parser::CharBlock);
 
   // Attempts to find a match for a derived type instance
   const DeclTypeSpec *FindInstantiatedDerivedType(const DerivedTypeSpec &,
       DeclTypeSpec::Category = DeclTypeSpec::TypeDerived) const;
 
   bool IsModuleFile() const {
     return kind_ == Kind::Module && symbol_ &&
         symbol_->test(Symbol::Flag::ModFile);

flang/include/flang/Semantics/semantics.h

parser::Message &SayWithDecl(const Symbol &symbol,
       A &&...args) {
     auto &message{Say(at, std::move(msg), args...)};
     evaluate::AttachDeclaration(&message, symbol);
     return message;
   }
 
   const Scope &FindScope(parser::CharBlock) const;
   Scope &FindScope(parser::CharBlock);
+  void UpdateScopeIndex(Scope &, parser::CharBlock);
 
   bool IsInModuleFile(parser::CharBlock) const;
 
   const ConstructStack &constructStack() const { return constructStack_; }
   template <typename N> void PushConstruct(const N &node) {
     constructStack_.emplace_back(&node);
   }
   void PopConstruct();
public:
   // standard, if the initialization and biggest size appearances are in
   // different compilation units, the behavior will depend on the linker. The
   // linker may have the behavior described before, but it may also keep the
   // initialized common symbol without extending its size, or have some other
   // behavior.
   CommonBlockList GetCommonBlocks() const;
 
 private:
+  struct ScopeIndexComparator {
+    bool operator()(parser::CharBlock, parser::CharBlock) const;
+  };
+  using ScopeIndex =
+      std::multimap<parser::CharBlock, Scope &, ScopeIndexComparator>;
+  ScopeIndex::iterator SearchScopeIndex(parser::CharBlock);
+
   void CheckIndexVarRedefine(
       const parser::CharBlock &, const Symbol &, parser::MessageFixedText &&);
   void CheckError(const Symbol &);
 
   const common::IntrinsicTypeDefaultKinds &defaultKinds_;
   const common::LanguageFeatureControl &languageFeatures_;
   parser::AllCookedSources &allCookedSources_;
   std::optional<parser::CharBlock> location_;
   std::vector<std::string> searchDirectories_;
   std::vector<std::string> intrinsicModuleDirectories_;
   std::string moduleDirectory_{"."s};
   std::string moduleFileSuffix_{".mod"};
   bool warnOnNonstandardUsage_{false};
   bool warningsAreErrors_{false};
   bool debugModuleWriter_{false};
   const evaluate::IntrinsicProcTable intrinsics_;
   evaluate::TargetCharacteristics targetCharacteristics_;
   Scope globalScope_;
   Scope &intrinsicModulesScope_;
+  ScopeIndex scopeIndex_;
   parser::Messages messages_;
   evaluate::FoldingContext foldingContext_;
   ConstructStack constructStack_;
   struct IndexVarInfo {
     parser::CharBlock location;
     IndexVarKind kind;
   };
   std::map<SymbolRef, const IndexVarInfo, SymbolAddressCompare>

flang/lib/Semantics/check-io.cpp

if (auto expr{AnalyzeExpr(context_, *var)}) {
             .Attach(std::move(*whyNot));
       }
     }
   }
 }
 
 void IoChecker::CheckForPureSubprogram() const { // C1597
   CHECK(context_.location());
-  if (const Scope *
-      scope{context_.globalScope().FindScope(*context_.location())}) {
-    if (FindPureProcedureContaining(*scope)) {
-      context_.Say(
-          "External I/O is not allowed in a pure subprogram"_err_en_US);
-    }
+  const Scope &scope{context_.FindScope(*context_.location())};
+  if (FindPureProcedureContaining(scope)) {
+    context_.Say("External I/O is not allowed in a pure subprogram"_err_en_US);
   }
 }
 
 // Seeks out an allocatable or pointer ultimate component that is not
 // nested in a nonallocatable/nonpointer component with a specific
 // defined I/O procedure.
 static const Symbol *FindUnsafeIoDirectComponent(common::DefinedIo which,
     const DerivedTypeSpec &derived, const Scope &scope) {

flang/lib/Semantics/scope.cpp

bool Scope::CanImport(const SourceName &name) const {
   case ImportKind::All:
   case ImportKind::Default:
     return true;
   case ImportKind::Only:
     return importNames_.count(name) > 0;
   }
 }
 
-const Scope *Scope::FindScope(parser::CharBlock source) const {
-  return const_cast<Scope *>(this)->FindScope(source);
-}
-
-Scope *Scope::FindScope(parser::CharBlock source) {
-  bool isContained{sourceRange_.Contains(source)};
-  if (!isContained && !IsTopLevel() && kind_ != Kind::Module) {
-    return nullptr;
-  }
-  for (auto &child : children_) {
-    if (auto *scope{child.FindScope(source)}) {
-      return scope;
-    }
-  }
-  return isContained && !IsTopLevel() ? this : nullptr;
-}
-
 void Scope::AddSourceRange(parser::CharBlock source) {
   if (source.empty()) {
     return;
   }
   const parser::AllCookedSources &allCookedSources{context_.allCookedSources()};
   const parser::CookedSource *cooked{allCookedSources.Find(source)};
   if (!cooked) {
     CHECK(context_.IsTempName(source.ToString()));
     return;
   }
   for (auto *scope{this}; !scope->IsTopLevel(); scope = &scope->parent()) {
     CHECK(scope->sourceRange_.empty() == (scope->cookedSource_ == nullptr));
     if (!scope->cookedSource_) {
+      context_.UpdateScopeIndex(*scope, source);
       scope->cookedSource_ = cooked;
       scope->sourceRange_ = source;
     } else if (scope->cookedSource_ == cooked) {
-      scope->sourceRange_.ExtendToCover(source);
+      auto combined{scope->sourceRange()};
+      combined.ExtendToCover(source);
+      context_.UpdateScopeIndex(*scope, combined);
+      scope->sourceRange_ = combined;
     } else {
       // There's a bug that will be hard to fix; crash informatively
       const parser::AllSources &allSources{allCookedSources.allSources()};
       const auto describe{[&](parser::CharBlock src) {
         if (auto range{allCookedSources.GetProvenanceRange(src)}) {
           std::size_t offset;
           if (const parser::SourceFile *
               file{allSources.GetSourceFile(range->start(), &offset)}) {
             return "'"s + file->path() + "' at " + std::to_string(offset) +
                 " for " + std::to_string(range->size());
           } else {
             return "(GetSourceFile failed)"s;
           }
         } else {
           return "(GetProvenanceRange failed)"s;
         }
       }};
       std::string scopeDesc{describe(scope->sourceRange_)};
       std::string newDesc{describe(source)};
       common::die("AddSourceRange would have combined ranges from distinct "
                   "source files \"%s\" and \"%s\"",
           scopeDesc.c_str(), newDesc.c_str());
     }
+    // Note: If the "break;" here were unconditional (or, equivalently, if
+    // there were no loop at all) then the source ranges of parent scopes
+    // would not enclose the source ranges of their children.  Timing
+    // shows that it's cheap to maintain this property, with the exceptions
+    // of top-level scopes and for (sub)modules and their descendant
+    // submodules.
     if (scope->IsSubmodule()) {
       break; // Submodules are child scopes but not contained ranges
     }
   }
 }
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const Scope &scope) {
   os << Scope::EnumToString(scope.kind()) << " scope: ";

flang/lib/Semantics/semantics.cpp

if (!AnyFatalError()) {
     std::string buf;
     llvm::raw_string_ostream ss{buf};
     ss << symbol;
     common::die(
         "No error was reported but setting error on: %s", ss.str().c_str());
   }
 }
 
+bool SemanticsContext::ScopeIndexComparator::operator()(
+    parser::CharBlock x, parser::CharBlock y) const {
+  return x.begin() < y.begin() ||
+      (x.begin() == y.begin() && x.size() > y.size());
+}
+
+auto SemanticsContext::SearchScopeIndex(parser::CharBlock source)
+    -> ScopeIndex::iterator {
+  if (!scopeIndex_.empty()) {
+    auto iter{scopeIndex_.upper_bound(source)};
+    auto begin{scopeIndex_.begin()};
+    do {
+      --iter;
+      if (iter->first.Contains(source)) {
+        return iter;
+      }
+    } while (iter != begin);
+  }
+  return scopeIndex_.end();
+}
+
 const Scope &SemanticsContext::FindScope(parser::CharBlock source) const {
   return const_cast<SemanticsContext *>(this)->FindScope(source);
 }
 
 Scope &SemanticsContext::FindScope(parser::CharBlock source) {
-  if (auto *scope{globalScope_.FindScope(source)}) {
-    return *scope;
+  if (auto iter{SearchScopeIndex(source)}; iter != scopeIndex_.end()) {
+    return iter->second;
   } else {
     common::die(
         "SemanticsContext::FindScope(): invalid source location for '%s'",
         source.ToString().c_str());
   }
 }
 
+void SemanticsContext::UpdateScopeIndex(
+    Scope &scope, parser::CharBlock newSource) {
+  if (scope.sourceRange().empty()) {
+    scopeIndex_.emplace(newSource, scope);
+  } else if (!scope.sourceRange().Contains(newSource)) {
+    auto iter{SearchScopeIndex(scope.sourceRange())};
+    CHECK(iter != scopeIndex_.end());
+    while (&iter->second != &scope) {
+      CHECK(iter != scopeIndex_.begin());
+      --iter;
+    }
+    scopeIndex_.erase(iter);
+    scopeIndex_.emplace(newSource, scope);
+  }
+}
+
 bool SemanticsContext::IsInModuleFile(parser::CharBlock source) const {
   for (const Scope *scope{&FindScope(source)}; !scope->IsGlobal();
        scope = &scope->parent()) {
     if (scope->IsModuleFile()) {
       return true;
     }
   }
   return false;