[mlir] Add support for walking locations similarly to Operations

[rriddle]

Summary

This allows for walking all nested locations of a given location, and is generally useful when processing locations.

Depends On D100436

[jpienaar]

mlir/lib/IR/Location.cpp

45

OOC Would a type switch be possible here?

[rriddle]

rebase

[rriddle]

rebase

llvm/include/llvm/ADT/TypeSwitch.h

public:
 
   /// As a default, invoke the given callable within the root value.
   template <typename CallableT>
   LLVM_NODISCARD ResultT Default(CallableT &&defaultFn) {
     if (result)
       return std::move(*result);
     return defaultFn(this->value);
   }
+  /// As a default, return the given value.
+  LLVM_NODISCARD ResultT Default(ResultT defaultResult) {
+    if (result)
+      return std::move(*result);
+    return defaultResult;
+  }
 
   LLVM_NODISCARD
   operator ResultT() {
     assert(result && "Fell off the end of a type-switch");
     return std::move(*result);
   }
 
 private:

llvm/unittests/ADT/TypeSwitchTest.cpp

TEST(TypeSwitchTest, CaseResult) {
   EXPECT_EQ(-1, translate(DerivedD()));
 }
 
 TEST(TypeSwitchTest, CasesResult) {
   auto translate = [](auto value) {
     return TypeSwitch<Base *, int>(&value)
         .Case<DerivedA, DerivedB, DerivedD>([](auto *) { return 0; })
         .Case([](DerivedC *) { return 1; })
-        .Default([](Base *) { return -1; });
+        .Default(-1);
   };
   EXPECT_EQ(0, translate(DerivedA()));
   EXPECT_EQ(0, translate(DerivedB()));
   EXPECT_EQ(1, translate(DerivedC()));
   EXPECT_EQ(0, translate(DerivedD()));
   EXPECT_EQ(-1, translate(DerivedE()));
 }
 

mlir/include/mlir/IR/Location.h

 #define MLIR_IR_LOCATION_H
 
 #include "mlir/IR/Attributes.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
 
 class Identifier;
+class Location;
+class WalkResult;
 
 //===----------------------------------------------------------------------===//
 // LocationAttr
 //===----------------------------------------------------------------------===//
 
 /// Location objects represent source locations information in MLIR.
 /// LocationAttr acts as the anchor for all Location based attributes.
 class LocationAttr : public Attribute {
 public:
   using Attribute::Attribute;
 
+  /// Walk all of the locations nested under, and including, the current.
+  WalkResult walk(function_ref<WalkResult(Location)> walkFn);
+
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
   static bool classof(Attribute attr);
 };
 
 //===----------------------------------------------------------------------===//
 // Location
 //===----------------------------------------------------------------------===//
 

mlir/lib/IR/Diagnostics.cpp

struct SourceMgrDiagnosticHandlerImpl {
   /// Mapping between file name and buffer ID's.
   llvm::StringMap<unsigned> filenameToBufId;
 };
 } // end namespace detail
 } // end namespace mlir
 
 /// Return a processable FileLineColLoc from the given location.
 static Optional<FileLineColLoc> getFileLineColLoc(Location loc) {
-  if (auto nameLoc = loc.dyn_cast<NameLoc>())
-    return getFileLineColLoc(loc.cast<NameLoc>().getChildLoc());
-  if (auto fileLoc = loc.dyn_cast<FileLineColLoc>())
-    return fileLoc;
-  if (auto callLoc = loc.dyn_cast<CallSiteLoc>())
-    return getFileLineColLoc(loc.cast<CallSiteLoc>().getCallee());
-  if (auto fusedLoc = loc.dyn_cast<FusedLoc>()) {
-    for (auto subLoc : loc.cast<FusedLoc>().getLocations()) {
-      if (auto callLoc = getFileLineColLoc(subLoc)) {
-        return callLoc;
-      }
-    }
-    return llvm::None;
-  }
-  return llvm::None;
+  Optional<FileLineColLoc> firstFileLoc;
+  loc->walk([&](Location loc) {
+    if (FileLineColLoc fileLoc = loc.dyn_cast<FileLineColLoc>()) {
+      firstFileLoc = fileLoc;
+      return WalkResult::interrupt();
+    return WalkResult::advance();
+  });
+  return firstFileLoc;
 }
 
 /// Return a processable CallSiteLoc from the given location.
 static Optional<CallSiteLoc> getCallSiteLoc(Location loc) {
   if (auto nameLoc = loc.dyn_cast<NameLoc>())
     return getCallSiteLoc(loc.cast<NameLoc>().getChildLoc());
   if (auto callLoc = loc.dyn_cast<CallSiteLoc>())
     return callLoc;

mlir/lib/IR/Location.cpp

 //===- Location.cpp - MLIR Location Classes -------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Location.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Identifier.h"
+#include "mlir/IR/Visitors.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
 /// Tablegen Attribute Definitions
 //===----------------------------------------------------------------------===//
 

 #include "mlir/IR/BuiltinLocationAttributes.cpp.inc"
       >();
 }
 
 //===----------------------------------------------------------------------===//
 // LocationAttr
 //===----------------------------------------------------------------------===//
 
+WalkResult LocationAttr::walk(function_ref<WalkResult(Location)> walkFn) {
+  if (walkFn(*this).wasInterrupted())
+    return WalkResult::interrupt();
+
+  return TypeSwitch<LocationAttr, WalkResult>(*this)
+      .Case([&](CallSiteLoc callLoc) -> WalkResult {
+        if (callLoc.getCallee()->walk(walkFn).wasInterrupted())
+          return WalkResult::interrupt();
+        return callLoc.getCaller()->walk(walkFn);
+      })
+      .Case([&](FusedLoc fusedLoc) -> WalkResult {
+        for (Location subLoc : fusedLoc.getLocations())
+          if (subLoc->walk(walkFn).wasInterrupted())
+            return WalkResult::interrupt();
+        return WalkResult::advance();
+      })
+      .Case([&](NameLoc nameLoc) -> WalkResult {
+        return nameLoc.getChildLoc()->walk(walkFn);
+      })
+      .Case([&](OpaqueLoc opaqueLoc) -> WalkResult {
+        return opaqueLoc.getFallbackLocation()->walk(walkFn);
+      })
+      .Default(WalkResult::advance());
+}
+
 /// Methods for support type inquiry through isa, cast, and dyn_cast.
 bool LocationAttr::classof(Attribute attr) {
   return attr.isa<CallSiteLoc, FileLineColLoc, FusedLoc, NameLoc, OpaqueLoc,
                   UnknownLoc>();
 }
 
 //===----------------------------------------------------------------------===//
 // CallSiteLoc