RequirementMachine: Split off ProtocolGraph into its own file

Author: slavapestov

include/swift/AST/LayoutConstraintKind.h

@@ -14,6 +14,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/Support/DataTypes.h"
+
 #ifndef SWIFT_LAYOUT_CONSTRAINTKIND_H
 #define SWIFT_LAYOUT_CONSTRAINTKIND_H
 

include/swift/AST/ProtocolGraph.h

@@ -0,0 +1,87 @@
+//===--- ProtocolGraph.h - Collects information about protocols -*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_PROTOCOLGRAPH_H
+#define SWIFT_PROTOCOLGRAPH_H
+
+#include "swift/AST/Requirement.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/TinyPtrVector.h"
+
+namespace swift {
+
+class ProtocolDecl;
+class AssociatedTypeDecl;
+
+namespace rewriting {
+
+struct ProtocolInfo {
+  ArrayRef<ProtocolDecl *> Inherited;
+  llvm::TinyPtrVector<AssociatedTypeDecl *> AssociatedTypes;
+  ArrayRef<Requirement> Requirements;
+
+  // Used by computeDepth() to detect circularity.
+  unsigned Mark : 1;
+
+  // Longest chain of protocol refinements, including this one.
+  // Greater than zero on valid code, might be zero if there's
+  // a cycle.
+  unsigned Depth : 31;
+
+  // Index of the protocol in the linear order.
+  unsigned Index : 32;
+
+  ProtocolInfo() {
+    Mark = 0;
+    Depth = 0;
+    Index = 0;
+  }
+
+  ProtocolInfo(ArrayRef<ProtocolDecl *> inherited,
+               llvm::TinyPtrVector<AssociatedTypeDecl *> &&types,
+               ArrayRef<Requirement> reqs)
+    : Inherited(inherited),
+      AssociatedTypes(types),
+      Requirements(reqs) {
+    Mark = 0;
+    Depth = 0;
+    Index = 0;
+  }
+};
+
+struct ProtocolGraph {
+  llvm::DenseMap<const ProtocolDecl *, ProtocolInfo> Info;
+  std::vector<const ProtocolDecl *> Protocols;
+  bool Debug = false;
+
+  void visitRequirements(ArrayRef<Requirement> reqs);
+
+  const ProtocolInfo &getProtocolInfo(
+      const ProtocolDecl *proto) const;
+
+  void addProtocol(const ProtocolDecl *proto);
+
+  void computeTransitiveClosure();
+
+  void computeLinearOrder();
+
+  void computeInheritedAssociatedTypes();
+
+private:
+  unsigned computeProtocolDepth(const ProtocolDecl *proto);
+};
+
+} // end namespace rewriting
+
+} // end namespace swift
+#endif

lib/AST/CMakeLists.txt

@@ -71,6 +71,7 @@ add_swift_host_library(swiftAST STATIC
   PlatformKind.cpp
   PrettyStackTrace.cpp
   ProtocolConformance.cpp
+  ProtocolGraph.cpp
   RawComment.cpp
   RequirementEnvironment.cpp
   RequirementMachine.cpp

lib/AST/ProtocolGraph.cpp

@@ -0,0 +1,133 @@
+//===--- ProtocolGraph.cpp - Collect information about protocols ----------===//
+//
+// 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 "swift/AST/ProtocolGraph.h"
+
+#include "swift/AST/Decl.h"
+#include "swift/AST/Requirement.h"
+
+using namespace swift;
+using namespace rewriting;
+
+void ProtocolGraph::visitRequirements(ArrayRef<Requirement> reqs) {
+  for (auto req : reqs) {
+    if (req.getKind() == RequirementKind::Conformance) {
+      addProtocol(req.getProtocolDecl());
+    }
+  }
+}
+
+const ProtocolInfo &ProtocolGraph::getProtocolInfo(
+    const ProtocolDecl *proto) const {
+  auto found = Info.find(proto);
+  assert(found != Info.end());
+  return found->second;
+}
+
+void ProtocolGraph::addProtocol(const ProtocolDecl *proto) {
+  if (Info.count(proto) > 0)
+    return;
+
+  Info[proto] = {proto->getInheritedProtocols(),
+                 proto->getAssociatedTypeMembers(),
+                 proto->getRequirementSignature()};
+  Protocols.push_back(proto);
+}
+
+void ProtocolGraph::computeTransitiveClosure() {
+  unsigned i = 0;
+  while (i < Protocols.size()) {
+    auto *proto = Protocols[i++];
+    visitRequirements(getProtocolInfo(proto).Requirements);
+  }
+}
+
+void ProtocolGraph::computeLinearOrder() {
+  for (const auto *proto : Protocols) {
+    (void) computeProtocolDepth(proto);
+  }
+
+  std::sort(
+      Protocols.begin(), Protocols.end(),
+      [&](const ProtocolDecl *lhs,
+          const ProtocolDecl *rhs) -> bool {
+        const auto &lhsInfo = getProtocolInfo(lhs);
+        const auto &rhsInfo = getProtocolInfo(rhs);
+
+        // protocol Base {} // depth 1
+        // protocol Derived : Base {} // depth 2
+        //
+        // Derived < Base in the linear order.
+        if (lhsInfo.Depth != rhsInfo.Depth)
+          return lhsInfo.Depth > rhsInfo.Depth;
+
+        return TypeDecl::compare(lhs, rhs) < 0;
+      });
+
+  for (unsigned i : indices(Protocols)) {
+    Info[Protocols[i]].Index = i;
+  }
+
+  if (Debug) {
+    for (const auto *proto : Protocols) {
+      const auto &info = getProtocolInfo(proto);
+      llvm::dbgs() << "@ Protocol " << proto->getName()
+                   << " Depth=" << info.Depth
+                   << " Index=" << info.Index << "\n";
+    }
+  }
+}
+
+void ProtocolGraph::computeInheritedAssociatedTypes() {
+  for (const auto *proto : Protocols) {
+    auto &info = Info[proto];
+
+    llvm::SmallDenseSet<const AssociatedTypeDecl *, 4> visited;
+    for (const auto *inherited : info.Inherited) {
+      if (inherited == proto)
+        continue;
+
+      for (auto *inheritedType : getProtocolInfo(inherited).AssociatedTypes) {
+        if (!visited.insert(inheritedType).second)
+          continue;
+
+        // The 'if (inherited == proto)' above avoids a potential
+        // iterator invalidation here.
+        info.AssociatedTypes.push_back(inheritedType);
+      }
+    }
+  }
+}
+
+unsigned ProtocolGraph::computeProtocolDepth(const ProtocolDecl *proto) {
+  auto &info = Info[proto];
+
+  if (info.Mark) {
+    // Already computed, or we have a cycle. Cycles are diagnosed
+    // elsewhere in the type checker, so we don't have to do
+    // anything here.
+    return info.Depth;
+  }
+
+  info.Mark = true;
+  unsigned depth = 0;
+
+  for (auto *inherited : info.Inherited) {
+    unsigned inheritedDepth = computeProtocolDepth(inherited);
+    depth = std::max(inheritedDepth, depth);
+  }
+
+  depth++;
+
+  info.Depth = depth;
+  return depth;
+}