From d2052557f7f7d0c81c9148f14b5ec3871d837473 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveTH@foxmail.com>
Date: Sun, 14 Apr 2024 21:39:10 +0800
Subject: [PATCH] fix `line_profiler`

---
 include/pocketpy/profiler.h |  7 ++++---
 src/ceval.cpp               |  2 +-
 src/profiler.cpp            | 21 +++++++--------------
 3 files changed, 12 insertions(+), 18 deletions(-)

diff --git a/include/pocketpy/profiler.h b/include/pocketpy/profiler.h
index 4ed36695..b102a5f6 100644
--- a/include/pocketpy/profiler.h
+++ b/include/pocketpy/profiler.h
@@ -14,7 +14,8 @@ struct _LineRecord{
 };
 
 struct _FrameRecord{
-    LinkedFrame* frame;
+    int callstack_size;
+    Frame* frame;
     clock_t prev_time;
     _LineRecord* prev_record;
 };
@@ -26,8 +27,8 @@ struct LineProfiler{
     std::set<FuncDecl*> functions;
 
     void begin();
-    void _step(LinkedFrame*);
-    void _step_end(LinkedFrame*, int);
+    void _step(int, Frame*);
+    void _step_end(int, Frame*, int);
     void end();
     Str stats();
 };
diff --git a/src/ceval.cpp b/src/ceval.cpp
index 93376389..2d4cf77b 100644
--- a/src/ceval.cpp
+++ b/src/ceval.cpp
@@ -74,7 +74,7 @@ PyObject* VM::_run_top_frame(){
 
 #define CEVAL_STEP_CALLBACK() \
     if(_ceval_on_step) _ceval_on_step(this, frame, byte);   \
-    if(_profiler) _profiler->_step(callstack._tail);        \
+    if(_profiler) _profiler->_step(callstack.size(), frame);        \
     if(!_next_breakpoint.empty()) { _next_breakpoint._step(this); }
 
 #define DISPATCH_OP_CALL() { frame = top_frame(); goto __NEXT_FRAME; }
diff --git a/src/profiler.cpp b/src/profiler.cpp
index 35eb60bf..db7e63a0 100644
--- a/src/profiler.cpp
+++ b/src/profiler.cpp
@@ -18,17 +18,16 @@ void LineProfiler::begin(){
     frames.clear();
 }
 
-void LineProfiler::_step(LinkedFrame* linked_frame){
-    Frame* frame = &linked_frame->frame;
+void LineProfiler::_step(int callstack_size, Frame* frame){
     auto line_info = frame->co->lines[frame->_ip];
     if(line_info.is_virtual) return;
     std::string_view filename = frame->co->src->filename.sv();
     int line = line_info.lineno;
 
     if(frames.empty()){
-        frames.push({linked_frame, clock(), nullptr});
+        frames.push({callstack_size, frame, clock(), nullptr});
     }else{
-        _step_end(linked_frame, line);
+        _step_end(callstack_size, frame, line);
     }
 
     auto& file_records = records[filename];
@@ -44,19 +43,13 @@ void LineProfiler::_step(LinkedFrame* linked_frame){
     frames.top().prev_record = &file_records.at(line);
 }
 
-void LineProfiler::_step_end(LinkedFrame* linked_frame, int line){
+void LineProfiler::_step_end(int callstack_size, Frame* frame, int line){
     clock_t now = clock();
     _FrameRecord& top_frame_record = frames.top();
     _LineRecord* prev_record = top_frame_record.prev_record;
 
-    int id_delta;
-    if(linked_frame == top_frame_record.frame){
-        id_delta = 0;
-    }else if(linked_frame->f_back == top_frame_record.frame){
-        id_delta = 1;
-    }else{
-        id_delta = -1;  // unsafe
-    }
+    int id_delta = callstack_size - top_frame_record.callstack_size;
+    PK_ASSERT(abs(id_delta) <= 1)
 
     // current line is about to change
     if(prev_record->line != line){
@@ -67,7 +60,7 @@ void LineProfiler::_step_end(LinkedFrame* linked_frame, int line){
     }
     
     if(id_delta == 1){
-        frames.push({linked_frame, now, nullptr});
+        frames.push({callstack_size, frame, now, nullptr});
     }else{
         if(id_delta == -1) frames.pop();
     }