[NFC][Stack Traces] Do not build a functor where we can use the translation map directly

PiperOrigin-RevId: 356403200
Change-Id: I6ae0231a08bcc24b103f3fa8f95233aa2116a952

tensorflow/core/util/BUILD

@@ -482,6 +482,7 @@ cc_library(
     ],
     deps = [
         "//tensorflow/core/platform:stack_frame",
+        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
     ],

tensorflow/core/util/managed_stack_trace.h

@@ -19,20 +19,22 @@ limitations under the License.
 #include <string>
 #include <vector>
 
+#include "absl/container/flat_hash_map.h"
 #include "absl/strings/match.h"
 #include "absl/types/optional.h"
 #include "tensorflow/core/platform/stack_frame.h"
 
 namespace tensorflow {
 
+using SourceLoc = std::tuple<std::string, int>;
+
 // Maps filename/line_no combination into a stack frame.
-using StackTraceMap =
+using SourceMap = absl::flat_hash_map<SourceLoc, StackFrame>;
-    std::function<absl::optional<StackFrame>(std::pair<const char*, int>)>;
 
 // Returns "true" on filenames which should be skipped.
 using StackTraceFilter = std::function<bool(const char*)>;
 
-using ToStackFramesFunctor = std::vector<StackFrame>(int, const StackTraceMap&,
+using ToStackFramesFunctor = std::vector<StackFrame>(int, const SourceMap&,
                                                      const StackTraceFilter&,
                                                      bool, int);
 
@@ -54,11 +56,12 @@ class ManagedStackTrace {
       : id_(id), to_stack_frames_(to_stack_frames) {}
 
   // Returns stack trace as a vector of `StackFrame`s.
-  std::vector<StackFrame> ToStackFrames(const StackTraceMap& mapper = {},
+  std::vector<StackFrame> ToStackFrames(const SourceMap& source_map = {},
                                         const StackTraceFilter& filtered = {},
                                         bool reverse_traversal = false,
                                         int limit = -1) const {
-    return to_stack_frames_(id_, mapper, filtered, reverse_traversal, limit);
+    return to_stack_frames_(id_, source_map, filtered, reverse_traversal,
+                            limit);
   }
 
  private:

tensorflow/python/util/stack_trace.cc

@@ -42,7 +42,7 @@ const char* GetPythonString(PyObject* o) {
 namespace tensorflow {
 
 std::vector<StackFrame> StackTrace::ToStackFrames(
-    const StackTraceMap& mapper, const StackTraceFilter& filtered,
+    const SourceMap& source_map, const StackTraceFilter& filtered,
     bool reverse_traversal, int limit) const {
   DCheckPyGilStateForStackTrace();
   std::vector<StackFrame> result;
@@ -61,11 +61,9 @@ std::vector<StackFrame> StackTrace::ToStackFrames(
       continue;
     }
 
-    absl::optional<StackFrame> mapped =
+    const auto it = source_map.find(std::make_tuple(file_name, line_number));
-        mapper ? mapper(std::make_pair(file_name, line_number)) : absl::nullopt;
+    if (it != source_map.end()) {
-
+      result.push_back(it->second);
-    if (mapped) {
-      result.push_back(*mapped);
     } else {
       result.emplace_back(StackFrame{file_name, line_number,
                                      GetPythonString(code_obj.first->co_name)});

tensorflow/python/util/stack_trace.h

@@ -88,12 +88,12 @@ class StackTrace final {
   }
 
   // Returns a structured representation of the captured stack trace.
-  // `mapper` provides a custom mapping for translating stack frames, `filter`
+  // `source_map` provides a custom mapping for translating stack frames,
-  // returns `true` for the stack frames which should be omitted.
+  // `filter` returns `true` for the stack frames which should be omitted.
   //
   // `reverse_traversal` changes the traversal order of the stack trace, and
   // `limit` bounds the number of returned frames (after filtering).
-  std::vector<StackFrame> ToStackFrames(const StackTraceMap& mapper = {},
+  std::vector<StackFrame> ToStackFrames(const SourceMap& source_map = {},
                                         const StackTraceFilter& filtered = {},
                                         bool reverse_traversal = false,
                                         int limit = -1) const;
@@ -149,11 +149,11 @@ extern StackTraceManager* const stack_trace_manager;
 // Converts the ManagedStackTrace (identified by ID) to a vector of stack
 // frames.
 inline std::vector<StackFrame> ManagedStackTraceToStackFrames(
-    int id, const StackTraceMap& mapper, const StackTraceFilter& filtered,
+    int id, const SourceMap& source_map, const StackTraceFilter& filtered,
     bool reverse_traversal, int limit) {
   PyGILState_STATE gstate = PyGILState_Ensure();
   std::vector<StackFrame> result = stack_trace_manager->Get(id)->ToStackFrames(
-      mapper, filtered, reverse_traversal, limit);
+      source_map, filtered, reverse_traversal, limit);
   PyGILState_Release(gstate);
   return result;
 }

tensorflow/python/util/tf_stack.cc

@@ -54,10 +54,6 @@ namespace {
 
 namespace py = pybind11;
 
-using SourceLoc = std::tuple<std::string, int>;
-
-using SourceMap = absl::flat_hash_map<SourceLoc, StackFrame>;
-
 using StringSet = absl::flat_hash_set<std::string>;
 
 // Python wrapper for a SourceMap.
@@ -149,8 +145,7 @@ class StackTraceWrapper : public AbstractStackTrace {
     PyGILState_STATE state = PyGILState_Ensure();
 
     stack_frames_cache_ = captured_.ToStackFrames(
-        [&](std::pair<const char*, int> p) { return StackTraceMapping(p); },
+        *source_map_, [&](const char* f) { return StackTraceFiltering(f); });
-        [&](const char* f) { return StackTraceFiltering(f); });
     stack_frames_cache_->pop_back();  // Drop last stack frame.
     PyGILState_Release(state);
     return *stack_frames_cache_;
@@ -163,7 +158,7 @@ class StackTraceWrapper : public AbstractStackTrace {
 
     PyGILState_STATE state = PyGILState_Ensure();
     std::vector<StackFrame> last_frame = captured_.ToStackFrames(
-        [&](std::pair<const char*, int> p) { return StackTraceMapping(p); },
+        *source_map_,
         [&](const char* file_name) {
           return StackTraceFiltering(file_name) ||
                  IsInternalFrameForFilename(file_name);
@@ -226,14 +221,6 @@ class StackTraceWrapper : public AbstractStackTrace {
         });
   }
 
-  absl::optional<StackFrame> StackTraceMapping(SourceLoc loc) const {
-    if (source_map_->contains(loc)) {
-      return source_map_->at(loc);
-    }
-
-    return absl::nullopt;
-  }
-
   bool StackTraceFiltering(const char* file_name) const {
     return filter_->contains(file_name);
   }