RequirementMachine: Speed up term simplification with a prefix trie

include/swift/Basic/LangOptions.h

@@ -464,6 +464,9 @@ namespace swift {
     /// Enables debugging output from the requirement machine.
     bool DebugRequirementMachine = false;
 
+    /// Enables statistics output from the requirement machine.
+    bool AnalyzeRequirementMachine = false;
+
     /// Maximum iteration count for requirement machine confluent completion
     /// algorithm.
     unsigned RequirementMachineStepLimit = 2000;

include/swift/Option/FrontendOptions.td

@@ -276,6 +276,10 @@ def disable_named_lazy_member_loading : Flag<["-"], "disable-named-lazy-member-l
 def debug_requirement_machine : Flag<["-"], "debug-requirement-machine">,
   HelpText<"Enables debugging output from the generics implementation">;
 
+def analyze_requirement_machine : Flag<["-"], "analyze-requirement-machine">,
+  Flags<[FrontendOption, HelpHidden, DoesNotAffectIncrementalBuild]>,
+  HelpText<"Print out requirement machine statistics at the end of the compilation job">;
+
 def requirement_machine_step_limit : Separate<["-"], "requirement-machine-step-limit">,
   Flags<[FrontendOption, HelpHidden, DoesNotAffectIncrementalBuild]>,
   HelpText<"Set the maximum steps before we give up on confluent completion">;

lib/AST/CMakeLists.txt

@@ -80,6 +80,8 @@ add_swift_host_library(swiftAST STATIC
   RequirementMachine/RewriteContext.cpp
   RequirementMachine/RewriteSystem.cpp
   RequirementMachine/RewriteSystemCompletion.cpp
+  RequirementMachine/Symbol.cpp
+  RequirementMachine/Term.cpp
   SILLayout.cpp
   Stmt.cpp
   SubstitutionMap.cpp

lib/AST/RequirementMachine/Histogram.h

@@ -0,0 +1,150 @@
+//===--- Histogram.h - Simple histogram for statistics ----------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2021 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/ArrayRef.h"
+
+#ifndef SWIFT_HISTOGRAM_H
+#define SWIFT_HISTOGRAM_H
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/raw_ostream.h"
+#include <algorithm>
+#include <vector>
+
+namespace swift {
+
+namespace rewriting {
+
+/// A simple histogram that supports two operations: recording an element, and
+/// printing a graphical representation.
+///
+/// Used by the rewrite system to record various statistics, which are dumped
+/// when the -analyze-requirement-machine frontend flag is used.
+class Histogram {
+  unsigned Size;
+  unsigned Start;
+  std::vector<unsigned> Buckets;
+  unsigned OverflowBucket;
+
+  const unsigned MaxWidth = 40;
+
+  /// This is not really the base-10 logarithm, but rather the number of digits
+  /// in \p x when printed in base 10.
+  ///
+  /// So log10(0) == 1, log10(1) == 1, log10(10) == 2, etc.
+  static unsigned log10(unsigned x) {
+    if (x == 0)
+      return 1;
+    unsigned result = 0;
+    while (x != 0) {
+      result += 1;
+      x /= 10;
+    }
+    return result;
+  }
+
+public:
+  /// Creates a histogram with \p size buckets, where the numbering begins at
+  /// \p start.
+  Histogram(unsigned size, unsigned start = 0)
+      : Size(size),
+        Start(start),
+        Buckets(size, 0),
+        OverflowBucket(0) {
+  }
+
+  /// Adds an entry to the histogram. Asserts if the \p value is smaller than
+  /// Start. The value is allowed to be greater than or equal to Start + Size,
+  /// in which case it's counted in the OverflowBucket.
+  void add(unsigned value) {
+    assert(value >= Start);
+    value -= Start;
+    if (value >= Size)
+      ++OverflowBucket;
+    else
+      ++Buckets[value];
+  }
+
+  /// Print a nice-looking graphical representation of the histogram.
+  void dump(llvm::raw_ostream &out, const StringRef labels[] = {}) {
+    unsigned sumValues = 0;
+    unsigned maxValue = 0;
+    for (unsigned i = 0; i < Size; ++i) {
+      sumValues += Buckets[i];
+      maxValue = std::max(maxValue, Buckets[i]);
+    }
+    sumValues += OverflowBucket;
+    maxValue = std::max(maxValue, OverflowBucket);
+
+    size_t maxLabelWidth = 3 + log10(Size + Start);
+    if (labels != nullptr) {
+      for (unsigned i = 0; i < Size; ++i)
+        maxLabelWidth = std::max(labels[i].size(), maxLabelWidth);
+    }
+
+    out << std::string(maxLabelWidth, ' ') << " |";
+    out << std::string(std::min(MaxWidth, maxValue), ' ');
+    out << maxValue << "\n";
+
+    out << std::string(maxLabelWidth, ' ') << " |";
+    out << std::string(std::min(MaxWidth, maxValue), ' ');
+    out << "*\n";
+
+    out << std::string(maxLabelWidth, '-') << "-+-";
+    out << std::string(std::min(MaxWidth, maxValue), '-') << "\n";
+
+    auto scaledValue = [&](unsigned value) {
+      if (maxValue > MaxWidth) {
+        if (value != 0) {
+          // If the value is non-zero, print at least one '#'.
+          return std::max(1U, (value * MaxWidth) / maxValue);
+        }
+      }
+      return value;
+    };
+
+    for (unsigned i = 0; i < Size; ++i) {
+      if (labels) {
+        out << std::string(maxLabelWidth - labels[i].size(), ' ');
+        out << labels[i];
+      } else {
+        unsigned key = i + Start;
+        out << std::string(maxLabelWidth - log10(key), ' ');
+        out << key;
+      }
+      out << " | ";
+
+      unsigned value = scaledValue(Buckets[i]);
+      out << std::string(value, '#') << "\n";
+    }
+
+    if (OverflowBucket > 0) {
+      out << ">= " << (Size + Start) << " | ";
+
+      unsigned value = scaledValue(OverflowBucket);
+      out << std::string(value, '#') << "\n";
+    }
+
+    out << std::string(maxLabelWidth, '-') << "-+-";
+    out << std::string(std::min(MaxWidth, maxValue), '-') << "\n";
+
+    out << std::string(maxLabelWidth, ' ') << " | ";
+    out << "Total: " << sumValues << "\n";
+  }
+};
+
+}  // end namespace rewriting
+
+}  // end namespace swift
+
+#endif

lib/AST/RequirementMachine/PropertyMap.cpp

@@ -291,9 +291,8 @@ static MutableTerm getRelativeTermForType(CanType typeWitness,
   }
 
   // Add the member type names.
-  std::reverse(symbols.begin(), symbols.end());
-  for (auto symbol : symbols)
-    result.add(symbol);
+  for (auto iter = symbols.rbegin(), end = symbols.rend(); iter != end; ++iter)
+    result.add(*iter);
 
   return result;
 }
@@ -449,8 +448,8 @@ namespace {
 ///
 /// These lower to the following two rules:
 ///
-///   T.[concrete: Foo<τ_0_0, τ_0_1, String> with {X.Y, Z}]
-///   T.[concrete: Foo<Int, τ_0_0, τ_0_1> with {A.B, W}]
+///   T.[concrete: Foo<τ_0_0, τ_0_1, String> with {X.Y, Z}] => T
+///   T.[concrete: Foo<Int, τ_0_0, τ_0_1> with {A.B, W}] => T
 ///
 /// The two concrete type symbols will be added to the property bag of 'T',
 /// and we will eventually end up in this method, where we will generate three
@@ -1063,17 +1062,22 @@ RewriteSystem::buildPropertyMap(PropertyMap &map,
     if (rule.isDeleted())
       continue;
 
-    const auto &lhs = rule.getLHS();
+    auto lhs = rule.getLHS();
+    auto rhs = rule.getRHS();
 
     // Collect all rules of the form T.[p] => T where T is canonical.
     auto property = lhs.back();
     if (!property.isProperty())
       continue;
 
-    MutableTerm key(lhs.begin(), lhs.end() - 1);
-    if (key != rule.getRHS())
+    if (lhs.size() - 1 != rhs.size())
       continue;
 
+    if (!std::equal(rhs.begin(), rhs.end(), lhs.begin()))
+      continue;
+
+    MutableTerm key(rhs);
+
 #ifndef NDEBUG
     assert(!simplify(key) &&
            "Right hand side of a property rule should already be reduced");

lib/AST/RequirementMachine/RewriteContext.cpp

@@ -19,6 +19,15 @@
 using namespace swift;
 using namespace rewriting;
 
+RewriteContext::RewriteContext(ASTContext &ctx)
+    : Context(ctx),
+      Stats(ctx.Stats),
+      SymbolHistogram(Symbol::NumKinds),
+      TermHistogram(4, /*Start=*/1),
+      RuleTrieHistogram(16, /*Start=*/1),
+      RuleTrieRootHistogram(16)  {
+}
+
 Term RewriteContext::getTermForType(CanType paramType,
                                     const ProtocolDecl *proto) {
   return Term::get(getMutableTermForType(paramType, proto), *this);
@@ -271,4 +280,20 @@ Type RewriteContext::getRelativeTypeForTerm(
   return getTypeForSymbolRange(
       term.begin() + prefix.size(), term.end(), genericParam,
       { }, protos, Context);
-}
+}
+
+/// We print stats in the destructor, which should get executed at the end of
+/// a compilation job.
+RewriteContext::~RewriteContext() {
+  if (Context.LangOpts.AnalyzeRequirementMachine) {
+    llvm::dbgs() << "--- Requirement Machine Statistics ---\n";
+    llvm::dbgs() << "\n* Symbol kind:\n";
+    SymbolHistogram.dump(llvm::dbgs(), Symbol::Kinds);
+    llvm::dbgs() << "\n* Term length:\n";
+    TermHistogram.dump(llvm::dbgs());
+    llvm::dbgs() << "\n* Rule trie fanout:\n";
+    RuleTrieHistogram.dump(llvm::dbgs());
+    llvm::dbgs() << "\n* Rule trie root fanout:\n";
+    RuleTrieRootHistogram.dump(llvm::dbgs());
+  }
+}

lib/AST/RequirementMachine/RewriteContext.h

@@ -18,7 +18,9 @@
 #include "swift/Basic/Statistic.h"
 #include "llvm/ADT/FoldingSet.h"
 #include "llvm/Support/Allocator.h"
-#include "RewriteSystem.h"
+#include "Histogram.h"
+#include "Symbol.h"
+#include "Term.h"
 
 namespace swift {
 
@@ -50,10 +52,16 @@ class RewriteContext final {
   ASTContext &Context;
 
 public:
-  /// Statistical counters.
+  /// Statistics.
   UnifiedStatsReporter *Stats;
 
-  RewriteContext(ASTContext &ctx) : Context(ctx), Stats(ctx.Stats) {}
+  /// Histograms.
+  Histogram SymbolHistogram;
+  Histogram TermHistogram;
+  Histogram RuleTrieHistogram;
+  Histogram RuleTrieRootHistogram;
+
+  explicit RewriteContext(ASTContext &ctx);
 
   Term getTermForType(CanType paramType, const ProtocolDecl *proto);
 
@@ -73,6 +81,8 @@ class RewriteContext final {
   Type getRelativeTypeForTerm(
                       const MutableTerm &term, const MutableTerm &prefix,
                       const ProtocolGraph &protos) const;
+
+  ~RewriteContext();
 };
 
 } // end namespace rewriting