[clangd] Support limiting down traversals from sources in FileDistance.

[ioeric]

Summary

This is useful for symbo scope proximity, where down traversals from
the global scope if not desired. The limitation in the current implementation
is that we may lose coverage of paths that when are otherwise reachable from a
more expensive node with larger down traversal limit, as the greedy algorithm
favors nodes with cheapest costs.

Another option is to simply special-case the root path in FileDistance.

[sammccall]

The slightly weird down-traversals semantics are indeed slightly weird. I'm happy to move forward with this approach though if we document it.

I did think of one alternative:

• instead of unsigned Source::MaxDownTraversals, have bool Source::AllowDownTraversals (obviously less flexible)
• instead of (min-cost, max-down-traversals) in the cache, store (first=min-cost, second=min-cost-allowing-downtraversals).

When populating tree based on sources, set both cache values if down-traversals are allowed, otherwise just the first.
When propagating scores from parent to child, only read from the second cache value, and write to both.
When looking up a result from the cache, use the first cached value.

I don't think there's a lot of practical difference for our use case, the advantage would be the semantics are easier to reason about. WDYT?

clangd/FileDistance.cpp

84

I don't understand this if condition

90

For consistency, tiebreak with down traversals: || Node.Cost == R.first->second.Cost && Node.MaxDownTraversals > R.first->second.MaxDownTraversals (note the opposite sign)

But probably clearer to write this by defining CacheNode::operator<

107

This fills in the cache if it wasn't set - it's probably harmless but not quite as clear. Does lookup not work here?

clangd/FileDistance.h

65

nit: "unlimited" might be a clearer name here.

68

We should document the slightly weird semantics.

unittests/clangd/FileDistanceTests.cpp

124

, because /a itself is reached more cheaply via the root.

[ioeric]

As discussed offline, the semantics is getting a bit complicated. As what we really want for scope proximity is just disallow down traversals from root, we can go with special-casing the root path for now: D53317

clangd/FileDistance.h

struct FileDistanceOptions {
   unsigned UpCost = 2;      // |foo/bar.h -> foo|
   unsigned DownCost = 1;    // |foo -> foo/bar.h|
   unsigned IncludeCost = 2; // |foo.cc -> included_header.h|
 };
 
 struct SourceParams {
   // Base cost for paths starting at this source.
   unsigned Cost = 0;
-  // Limits the number of upwards traversals allowed from this source.
-  unsigned MaxUpTraversals = std::numeric_limits<unsigned>::max();
+  // Limits the number of upwards/downwards traversals allowed from this source.
+  static constexpr unsigned MaxTraversals =
+      std::numeric_limits<unsigned>::max();
+  unsigned MaxUpTraversals = MaxTraversals;
+  unsigned MaxDownTraversals = MaxTraversals;
 };
 
 // Supports lookups to find the minimum distance to a file from any source.
 // This object should be reused, it memoizes intermediate computations.
 class FileDistance {
 public:
   static constexpr unsigned Unreachable = std::numeric_limits<unsigned>::max();
 
   FileDistance(llvm::StringMap<SourceParams> Sources,
                const FileDistanceOptions &Opts = {});
 
   // Computes the minimum distance from any source to the file path.
   unsigned distance(llvm::StringRef Path);
 
 private:
   // Costs computed so far. Always contains sources and their ancestors.
   // We store hash codes only. Collisions are rare and consequences aren't dire.
-  llvm::DenseMap<llvm::hash_code, unsigned> Cache;
+  struct CacheNode {
+    unsigned Cost = Unreachable;
+    unsigned MaxDownTraversals = std::numeric_limits<unsigned>::max();
+  };
+  llvm::DenseMap<llvm::hash_code, CacheNode> Cache;
   FileDistanceOptions Opts;
 };
 
 // Supports lookups like FileDistance, but the lookup keys are URIs.
 // We convert each of the sources to the scheme of the URI and do a FileDistance
 // comparison on the bodies.
 class URIDistance {
 public:

clangd/FileDistance.cpp

 FileDistance::FileDistance(StringMap<SourceParams> Sources,
                            const FileDistanceOptions &Opts)
     : Opts(Opts) {
   llvm::DenseMap<hash_code, SmallVector<hash_code, 4>> DownEdges;
   // Compute the best distance following only up edges.
   // Keep track of down edges, in case we can use them to improve on this.
   for (const auto &S : Sources) {
     auto Canonical = canonicalize(S.getKey());
-    dlog("Source {0} = {1}, MaxUp = {2}", Canonical, S.second.Cost,
-         S.second.MaxUpTraversals);
+    dlog("Source {0} = {1}, MaxUp = {2}, MaxDown = {3}", Canonical,
+         S.second.Cost, S.second.MaxUpTraversals, S.second.MaxDownTraversals);
     // Walk up to ancestors of this source, assigning cost.
     StringRef Rest = Canonical;
     llvm::hash_code Hash = hash_value(Rest);
     for (unsigned I = 0; !Rest.empty(); ++I) {
       Rest = parent_path(Rest, sys::path::Style::posix);
       auto NextHash = hash_value(Rest);
       auto &Down = DownEdges[NextHash];
       if (!llvm::is_contained(Down, Hash))
         Down.push_back(Hash);
       // We can't just break after MaxUpTraversals, must still set DownEdges.
       if (I > S.getValue().MaxUpTraversals) {
         if (Cache.find(Hash) != Cache.end())
           break;
       } else {
-        unsigned Cost = S.getValue().Cost + I * Opts.UpCost;
-        auto R = Cache.try_emplace(Hash, Cost);
+        CacheNode Node;
+        Node.Cost = S.getValue().Cost + I * Opts.UpCost;
+        if (SourceParams::MaxTraversals - S.getValue().MaxDownTraversals >= I)
+          Node.MaxDownTraversals = S.getValue().MaxDownTraversals + I;
+        auto R = Cache.try_emplace(Hash, Node);
         if (!R.second) {
-          if (Cost < R.first->second) {
-            R.first->second = Cost;
+          // Note that we might lose coverage of some paths that are otherwise
+          // reachable with a higher cost but large down traversal limit.
+          if (Node.Cost < R.first->second.Cost) {
+            R.first->second = std::move(Node);
           } else {
             // If we're not the best way to get to this path, stop assigning.
             break;
           }
         }
       }
       Hash = NextHash;
     }
   }
   // Now propagate scores parent -> child if that's an improvement.
   // BFS ensures we propagate down chains (must visit parents before children).
   std::queue<hash_code> Next;
   for (auto Child : DownEdges.lookup(hash_value(llvm::StringRef(""))))
     Next.push(Child);
   while (!Next.empty()) {
-    auto ParentCost = Cache.lookup(Next.front());
-    for (auto Child : DownEdges.lookup(Next.front())) {
-      auto &ChildCost =
-          Cache.try_emplace(Child, Unreachable).first->getSecond();
-      if (ParentCost + Opts.DownCost < ChildCost)
-        ChildCost = ParentCost + Opts.DownCost;
-      Next.push(Child);
+    const auto &Parent = Cache[Next.front()];
+    if (Parent.MaxDownTraversals > 0) {
+      for (auto ChildHash : DownEdges.lookup(Next.front())) {
+        auto &Child =
+            Cache.try_emplace(ChildHash, CacheNode()).first->getSecond();
+        if (Parent.Cost + Opts.DownCost < Child.Cost) {
+          Child.Cost = Parent.Cost + Opts.DownCost;
+          Child.MaxDownTraversals = Parent.MaxDownTraversals - 1;
+        }
+        Next.push(ChildHash);
+      }
     }
     Next.pop();
   }
 }
 
 unsigned FileDistance::distance(StringRef Path) {
   auto Canonical = canonicalize(Path);
-  unsigned Cost = Unreachable;
+  CacheNode KnownNode;
   SmallVector<hash_code, 16> Ancestors;
   // Walk up ancestors until we find a path we know the distance for.
   for (StringRef Rest = Canonical; !Rest.empty();
        Rest = parent_path(Rest, sys::path::Style::posix)) {
     auto Hash = hash_value(Rest);
     auto It = Cache.find(Hash);
     if (It != Cache.end()) {
-      Cost = It->second;
+      KnownNode = It->second;
       break;
     }
     Ancestors.push_back(Hash);
   }
   // Now we know the costs for (known node, queried node].
   // Fill these in, walking down the directory tree.
-  for (hash_code Hash : reverse(Ancestors)) {
-    if (Cost != Unreachable)
-      Cost += Opts.DownCost;
-    Cache.try_emplace(Hash, Cost);
+  for (auto B = Ancestors.rbegin(), I = B, E = Ancestors.rend(); I != E; ++I) {
+    CacheNode Node;
+    unsigned DistanceToKnown = I - B + 1;
+    if (DistanceToKnown <= KnownNode.MaxDownTraversals) {
+      if (KnownNode.Cost != Unreachable)
+        Node.Cost = KnownNode.Cost + DistanceToKnown * Opts.DownCost;
+      if (KnownNode.MaxDownTraversals < SourceParams::MaxTraversals)
+        Node.MaxDownTraversals = KnownNode.MaxDownTraversals - (I - B) - 1;
+    }
+    Cache.try_emplace(*I, Node);
   }
-  dlog("distance({0} = {1})", Path, Cost);
-  return Cost;
+  dlog("distance({0} = {1})", Path, Cache[hash_value(Canonical)].Cost);
+  return Cache[hash_value(Canonical)].Cost;
 }
 
 unsigned URIDistance::distance(llvm::StringRef URI) {
   auto R = Cache.try_emplace(llvm::hash_value(URI), FileDistance::Unreachable);
   if (!R.second)
     return R.first->getSecond();
   if (auto U = clangd::URI::parse(URI)) {
     dlog("distance({0} = {1})", URI, U->body());

unittests/clangd/FileDistanceTests.cpp

TEST(FileDistance, LimitUpTraversals) {
 
   FileDistance D({{"/", CostLots}, {"/a/b/c", CheapButLimited}}, Opts);
   EXPECT_EQ(D.distance("/a"), 101u);
   EXPECT_EQ(D.distance("/a/z"), 102u);
   EXPECT_EQ(D.distance("/a/b"), 1u);
   EXPECT_EQ(D.distance("/a/b/z"), 2u);
 }
 
+TEST(FileDistance, LimitDownTraversals) {
+  FileDistanceOptions Opts;
+  Opts.UpCost = 1;
+  Opts.DownCost = 1;
+  SourceParams CheapButLimited, CostLots;
+  CheapButLimited.MaxDownTraversals = 1;
+  CostLots.Cost = 100;
+
+  FileDistance D({{"/", CheapButLimited}, {"/a/b", CostLots}}, Opts);
+  EXPECT_EQ(D.distance("/a"), 1u);
+  EXPECT_EQ(D.distance("/a/b"), 100u);
+  EXPECT_EQ(D.distance("/a/b/c"), 101u);
+}
+
+TEST(FileDistance, LimitDownTraversalsCorner) {
+  FileDistanceOptions Opts;
+  Opts.UpCost  = 1;
+  Opts.DownCost = 1;
+  SourceParams CheapButLimited, CostLots;
+  CheapButLimited.MaxDownTraversals = 1;
+  CostLots.Cost = 100;
+
+  FileDistance D({{"/", CheapButLimited}, {"/a", CostLots}}, Opts);
+  EXPECT_EQ(D.distance("/x"), 1u);
+  EXPECT_EQ(D.distance("/a"), 1u);
+  // These are actually reacahble from "/a" with high cost. But we don't support
+  // such case.
+  EXPECT_EQ(D.distance("/a/b"), FileDistance::Unreachable);
+  EXPECT_EQ(D.distance("/a/b/c"), FileDistance::Unreachable);
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang