[NFC] Propagate options for converting the stack trace to stack frames to ManagedStackTrace

PiperOrigin-RevId: 351844466
Change-Id: Ic11fa2ac74920005353e945ae9f97bc2e5577d9a

tensorflow/core/util/managed_stack_trace.h

@@ -22,21 +22,35 @@ limitations under the License.
 
 namespace tensorflow {
 
+// Maps filename/line_no combination into a stack frame.
+using StackTraceMap =
+    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&,
+                                                     const StackTraceFilter&,
+                                                     bool, int);
+
 // Language agnostic stack trace class. It only saves an id, and language
 // clients are responsible for managing the actual stack trace objects.
 class ManagedStackTrace {
  public:
-  ManagedStackTrace(int id, std::vector<StackFrame> (*to_stack_frames)(int))
+  ManagedStackTrace(int id, ToStackFramesFunctor* to_stack_frames)
       : id_(id), to_stack_frames_(to_stack_frames) {}
 
   // Returns stack trace as a vector of `StackFrame`s.
-  std::vector<StackFrame> ToStackFrames() const {
-    return to_stack_frames_(id_);
+  std::vector<StackFrame> ToStackFrames(const StackTraceMap& mapper = {},
+                                        const StackTraceFilter& filtered = {},
+                                        bool reverse_traversal = false,
+                                        int limit = -1) const {
+    return to_stack_frames_(id_, mapper, filtered, reverse_traversal, limit);
   }
 
  private:
   int id_;
-  std::vector<StackFrame> (*to_stack_frames_)(int);
+  ToStackFramesFunctor* to_stack_frames_;
 };
 
 }  // namespace tensorflow

tensorflow/python/util/stack_trace.h

@@ -42,13 +42,6 @@ inline void DCheckPyGilStateForStackTrace() {
 #endif
 }
 
-// Maps filename/line_no combination into a stack frame.
-using StackTraceMap =
-    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*)>;
-
 // A class for capturing Python stack trace.
 class StackTrace final {
  public:
@@ -152,19 +145,26 @@ class StackTraceManager {
 // Singleton StackTraceManager.
 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,
+    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);
+  PyGILState_Release(gstate);
+  return result;
+}
+
 // Returns Python stack trace object that can be converted to string.
 // Note that the actual stack trace is kept in a circular buffer for string
 // conversion could fail if it's evicted before.
 // Python GIL must be acquired beforehand.
 inline ManagedStackTrace GetStackTrace(int limit) {
   DCheckPyGilStateForStackTrace();
-  return ManagedStackTrace(stack_trace_manager->Capture(limit), [](int id) {
-    PyGILState_STATE gstate = PyGILState_Ensure();
-    std::vector<StackFrame> result =
-        stack_trace_manager->Get(id)->ToStackFrames();
-    PyGILState_Release(gstate);
-    return result;
-  });
+  return ManagedStackTrace(stack_trace_manager->Capture(limit),
+                           &ManagedStackTraceToStackFrames);
 }
 
 }  // namespace tensorflow