fix river's comment and add the test

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

mlir/include/mlir/IR/Diagnostics.h

public:
   SourceMgrDiagnosticHandler(llvm::SourceMgr &mgr, MLIRContext *ctx,
                              raw_ostream &os);
   SourceMgrDiagnosticHandler(llvm::SourceMgr &mgr, MLIRContext *ctx);
   ~SourceMgrDiagnosticHandler();
 
   /// Emit the given diagnostic information with the held source manager.
   void emitDiagnostic(Location loc, Twine message, DiagnosticSeverity kind);
 
-protected:
   /// Emit the given diagnostic with the held source manager.
   void emitDiagnostic(Diagnostic &diag);
 
+protected:
   /// Get a memory buffer for the given file, or nullptr if no file is
   /// available.
   const llvm::MemoryBuffer *getBufferForFile(StringRef filename);
 
   /// The source manager that we are wrapping.
   llvm::SourceMgr &mgr;
 
   /// The output stream to use when printing diagnostics.

mlir/lib/IR/Diagnostics.cpp

case StandardAttributes::CallSiteLocation:
     return getFileLineColLoc(loc.cast<CallSiteLoc>().getCallee());
   case StandardAttributes::FusedLocation:
     return getFileLineColLoc(loc.cast<FusedLoc>().getLocations().front());
   default:
     return llvm::None;
   }
 }
 
-/// Return a processable CallSiteLocation from the given location.
-static Optional<FileLineColLoc> getCallSiteLoc(Location loc) {
+/// Return a processable CallSiteLoc from the given location.
+static Optional<CallSiteLoc> getCallSiteLoc(Location loc) {
   switch (loc->getKind()) {
   case StandardAttributes::NameLocation:
     return getCallSiteLoc(loc.cast<NameLoc>().getChildLoc());
   case StandardAttributes::CallSiteLocation:
     return loc.cast<CallSiteLoc>();
   case StandardAttributes::FusedLocation:
     return getCallSiteLoc(loc.cast<FusedLoc>().getLocations().front());
   default:
void SourceMgrDiagnosticHandler::emitDiagnostic(Diagnostic &diag) {
 
   // If the diagnostic location was a call site location, then print the call
   // stack as well.
   if (auto callLoc = getCallSiteLoc(loc)) {
     // Print the call stack while valid, or until the limit is reached.
     Location callerLoc = callLoc->getCaller();
     for (unsigned curDepth = 0; curDepth < callStackLimit; ++curDepth) {
       emitDiagnostic(callerLoc, "called from", DiagnosticSeverity::Note);
-      if (auto subLoc = callerLoc.dyn_cast<CallSiteLoc>())
-        callerLoc = subLoc.getCaller();
-      else
+      callLoc = getCallSiteLoc(callerLoc);
+      if (!callLoc)
         break;
+      callerLoc = callLoc->getCaller();
     }
   }
 
   // Emit each of the notes.
   for (auto &note : diag.getNotes())
     emitDiagnostic(note.getLocation(), note.str(), note.getSeverity());
 }
 

mlir/test/IR/diagnostic-handler.mlir

+// RUN: mlir-opt %s -test-diagnostic-handler | FileCheck %s
+// This test verifies that diagnostic handler can emit the call stack successfully.
+
+func @call_site_loc_in_fused() -> i32 {
+  //CHECK: mysource1: note:
+  //CHECK: mysource2: note: called from
+  //CHECK: mysource3: note: called from
+  %3 = constant 3 : i32 loc(fused[callsite("foo"("mysource1":0:0) at callsite("mysource2":1:0 at "mysource3":2:0)), "bar"])
+  return %3 : i32
+}

mlir/test/lib/Transforms/CMakeLists.txt

 add_llvm_library(MLIRTestTransforms
   TestCallGraph.cpp
   TestConstantFold.cpp
+  TestDiagnosticHandler.cpp
   TestLoopFusion.cpp
   TestInlining.cpp
   TestLinalgTransforms.cpp
   TestLiveness.cpp
   TestLoopMapping.cpp
   TestLoopParametricTiling.cpp
   TestOpaqueLoc.cpp
   TestMemRefStrideCalculation.cpp

mlir/test/lib/Transforms/TestDiagnosticHandler.cpp

+//===- TestDiagnosticHandler.cpp - Pass to test diagnostic handler --------===//
+//
+// 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/Dialect/StandardOps/Ops.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+/// Pass that emit diagnostic for each op.
+struct TestDiagnosticHandler : public FunctionPass<TestDiagnosticHandler> {
+
+  void runOnFunction() override {
+    auto &os = llvm::outs();
+    llvm::SourceMgr fileSourceMgr;
+    SourceMgrDiagnosticHandler diagHandler(fileSourceMgr, &getContext(), os);
+
+    getFunction().walk([&](Operation *op) {
+      if (isa<FuncOp>(op) || op->isKnownTerminator())
+        return;
+      Diagnostic diag(op->getLoc(), DiagnosticSeverity::Note);
+      diagHandler.emitDiagnostic(diag);
+    });
+  }
+};
+
+} // end anonymous namespace
+
+static PassRegistration<TestDiagnosticHandler> pass("test-diagnostic-handler",
+                                                    "emit diagnostic for ops");