@@ -82,19 +82,35 @@ impl FunctionLocations for VecFunctionLocations {
8282 }
8383}
8484
85- pub type LineNumber = u16 ; // TODO u32, newtype
85+ /// Either the line number, or the bytecode index needed to get it.
86+ #[ derive( Copy , Clone , Serialize , Deserialize , Hash , Eq , PartialEq , Debug ) ]
87+ pub enum LineNumberInfo {
88+ LineNumber ( u32 ) ,
89+ BytecodeIndex ( i32 ) ,
90+ }
91+
92+ impl LineNumberInfo {
93+ fn get_line_number ( & self ) -> u32 {
94+ if let LineNumberInfo :: LineNumber ( lineno) = self {
95+ * lineno
96+ } else {
97+ debug_assert ! ( false , "This should never happen." ) ;
98+ 0
99+ }
100+ }
101+ }
86102
87103/// A specific location: file + function + line number.
88104#[ derive( Clone , Debug , PartialEq , Eq , Copy , Hash , Serialize , Deserialize ) ]
89105pub struct CallSiteId {
90106 /// The function + filename. We use IDs for performance reasons (faster hashing).
91- function : FunctionId ,
107+ pub function : FunctionId ,
92108 /// Line number within the _file_, 1-indexed.
93- line_number : LineNumber ,
109+ pub line_number : LineNumberInfo ,
94110}
95111
96112impl CallSiteId {
97- pub fn new ( function : FunctionId , line_number : u16 ) -> CallSiteId {
113+ pub fn new ( function : FunctionId , line_number : LineNumberInfo ) -> CallSiteId {
98114 CallSiteId {
99115 function,
100116 line_number,
@@ -108,7 +124,7 @@ impl CallSiteId {
108124pub struct Callstack {
109125 calls : Vec < CallSiteId > ,
110126 #[ derivative( Hash = "ignore" , PartialEq = "ignore" ) ]
111- cached_callstack_id : Option < ( u16 , CallstackId ) > , // first bit is line number
127+ cached_callstack_id : Option < ( u32 , CallstackId ) > , // first bit is line number
112128}
113129
114130impl Callstack {
@@ -126,10 +142,14 @@ impl Callstack {
126142 }
127143 }
128144
129- pub fn start_call ( & mut self , parent_line_number : u16 , callsite_id : CallSiteId ) {
145+ pub fn to_vec ( & self ) -> Vec < CallSiteId > {
146+ self . calls . clone ( )
147+ }
148+
149+ pub fn start_call ( & mut self , parent_line_number : u32 , callsite_id : CallSiteId ) {
130150 if parent_line_number != 0 {
131151 if let Some ( mut call) = self . calls . last_mut ( ) {
132- call. line_number = parent_line_number;
152+ call. line_number = LineNumberInfo :: LineNumber ( parent_line_number) ;
133153 }
134154 }
135155 self . calls . push ( callsite_id) ;
@@ -141,7 +161,7 @@ impl Callstack {
141161 self . cached_callstack_id = None ;
142162 }
143163
144- pub fn id_for_new_allocation < F > ( & mut self , line_number : u16 , get_callstack_id : F ) -> CallstackId
164+ pub fn id_for_new_allocation < F > ( & mut self , line_number : u32 , get_callstack_id : F ) -> CallstackId
145165 where
146166 F : FnOnce ( & Callstack ) -> CallstackId ,
147167 {
@@ -156,7 +176,7 @@ impl Callstack {
156176 // Set the new line number:
157177 if line_number != 0 {
158178 if let Some ( call) = self . calls . last_mut ( ) {
159- call. line_number = line_number;
179+ call. line_number = LineNumberInfo :: LineNumber ( line_number) ;
160180 }
161181 }
162182
@@ -194,7 +214,8 @@ impl Callstack {
194214 . map ( |( id, ( function, filename) ) | {
195215 if to_be_post_processed {
196216 // Get Python code.
197- let code = linecache. get_source_line ( filename, id. line_number as usize ) ;
217+ let code = linecache
218+ . get_source_line ( filename, id. line_number . get_line_number ( ) as usize ) ;
198219 // Leading whitespace is dropped by SVG, so we'd like to
199220 // replace it with non-breaking space. However, inferno
200221 // trims whitespace
@@ -215,15 +236,15 @@ impl Callstack {
215236 format ! (
216237 "{filename}:{line} ({function});\u{2800} {code}" ,
217238 filename = filename,
218- line = id. line_number,
239+ line = id. line_number. get_line_number ( ) ,
219240 function = function,
220241 code = & code. trim_end( ) ,
221242 )
222243 } else {
223244 format ! (
224245 "{filename}:{line} ({function})" ,
225246 filename = filename,
226- line = id. line_number,
247+ line = id. line_number. get_line_number ( ) ,
227248 function = function,
228249 )
229250 }
@@ -335,6 +356,10 @@ impl Allocation {
335356 }
336357}
337358
359+ fn same_callstack ( callstack : & Callstack ) -> Cow < Callstack > {
360+ Cow :: Borrowed ( callstack)
361+ }
362+
338363/// The main data structure tracking everything.
339364pub struct AllocationTracker < FL : FunctionLocations > {
340365 // malloc()/calloc():
@@ -607,18 +632,22 @@ impl<FL: FunctionLocations> AllocationTracker<FL> {
607632 self . dump_to_flamegraph ( path, true , "peak-memory" , "Peak Tracked Memory Usage" , true ) ;
608633 }
609634
610- pub fn to_lines (
611- & self ,
635+ pub fn to_lines < ' a , F > (
636+ & ' a self ,
612637 peak : bool ,
613638 to_be_post_processed : bool ,
614- ) -> impl ExactSizeIterator < Item = String > + ' _ {
639+ update_callstack : F ,
640+ ) -> impl ExactSizeIterator < Item = String > + ' _
641+ where
642+ F : Fn ( & Callstack ) -> Cow < Callstack > + ' a ,
643+ {
615644 let by_call = self . combine_callstacks ( peak) . into_iter ( ) ;
616645 let id_to_callstack = self . interner . get_reverse_map ( ) ;
617646 let mut linecache = LineCacher :: default ( ) ;
618647 by_call. map ( move |( callstack_id, size) | {
619648 format ! (
620649 "{} {}" ,
621- id_to_callstack. get( & callstack_id) . unwrap( ) . as_string(
650+ update_callstack ( id_to_callstack. get( & callstack_id) . unwrap( ) ) . as_string(
622651 to_be_post_processed,
623652 & self . functions,
624653 ";" ,
@@ -675,7 +704,7 @@ impl<FL: FunctionLocations> AllocationTracker<FL> {
675704 subtitle,
676705 "bytes" ,
677706 to_be_post_processed,
678- |tbpp| self . to_lines ( peak, tbpp) ,
707+ |tbpp| self . to_lines ( peak, tbpp, same_callstack ) ,
679708 )
680709 }
681710
@@ -746,8 +775,9 @@ impl<FL: FunctionLocations> AllocationTracker<FL> {
746775
747776#[ cfg( test) ]
748777mod tests {
749- use crate :: memorytracking:: { ProcessUid , PARENT_PROCESS } ;
778+ use crate :: memorytracking:: { same_callstack , ProcessUid , PARENT_PROCESS } ;
750779
780+ use super :: LineNumberInfo :: LineNumber ;
751781 use super :: {
752782 Allocation , AllocationTracker , CallSiteId , Callstack , CallstackInterner , FunctionId ,
753783 FunctionLocations , VecFunctionLocations , HIGH_32BIT , MIB ,
@@ -815,7 +845,7 @@ mod tests {
815845 let ( process, allocation_size) = * allocated_sizes. get( i) . unwrap( ) ;
816846 let process = ProcessUid ( process) ;
817847 let mut cs = Callstack :: new( ) ;
818- cs. start_call( 0 , CallSiteId :: new( FunctionId :: new( i as u64 ) , 0 ) ) ;
848+ cs. start_call( 0 , CallSiteId :: new( FunctionId :: new( i as u64 ) , LineNumber ( 0 ) ) ) ;
819849 let cs_id = tracker. get_callstack_id( & cs) ;
820850 tracker. add_allocation( process, i as usize , allocation_size, cs_id) ;
821851 expected_memory_usage. push_back( allocation_size) ;
@@ -854,7 +884,7 @@ mod tests {
854884 let ( process, allocation_size) = * allocated_sizes. get( i) . unwrap( ) ;
855885 let process = ProcessUid ( process) ;
856886 let mut cs = Callstack :: new( ) ;
857- cs. start_call( 0 , CallSiteId :: new( FunctionId :: new( i as u64 ) , 0 ) ) ;
887+ cs. start_call( 0 , CallSiteId :: new( FunctionId :: new( i as u64 ) , LineNumber ( 0 ) ) ) ;
858888 let csid = tracker. get_callstack_id( & cs) ;
859889 tracker. add_anon_mmap( process, addresses[ i] as usize , allocation_size, csid) ;
860890 expected_memory_usage. push_back( allocation_size) ;
@@ -891,7 +921,7 @@ mod tests {
891921 let ( process, allocation_size) = * allocated_sizes. get( i) . unwrap( ) ;
892922 let process = ProcessUid ( process) ;
893923 let mut cs = Callstack :: new( ) ;
894- cs. start_call( 0 , CallSiteId :: new( FunctionId :: new( i as u64 ) , 0 ) ) ;
924+ cs. start_call( 0 , CallSiteId :: new( FunctionId :: new( i as u64 ) , LineNumber ( 0 ) ) ) ;
895925 let cs_id = tracker. get_callstack_id( & cs) ;
896926 tracker. add_allocation( process, i as usize , allocation_size, cs_id) ;
897927 expected_memory_usage += allocation_size;
@@ -900,7 +930,7 @@ mod tests {
900930 let ( process, allocation_size) = * allocated_mmaps. get( i) . unwrap( ) ;
901931 let process = ProcessUid ( process) ;
902932 let mut cs = Callstack :: new( ) ;
903- cs. start_call( 0 , CallSiteId :: new( FunctionId :: new( i as u64 ) , 0 ) ) ;
933+ cs. start_call( 0 , CallSiteId :: new( FunctionId :: new( i as u64 ) , LineNumber ( 0 ) ) ) ;
904934 let csid = tracker. get_callstack_id( & cs) ;
905935 tracker. add_anon_mmap( process, mmap_addresses[ i] as usize , allocation_size, csid) ;
906936 expected_memory_usage += allocation_size;
@@ -933,9 +963,10 @@ mod tests {
933963
934964 // Parent line number does nothing if it's first call:
935965 let mut cs1 = Callstack :: new ( ) ;
936- let id1 = CallSiteId :: new ( fid1, 2 ) ;
937- let id2 = CallSiteId :: new ( fid3, 45 ) ;
938- let id3 = CallSiteId :: new ( fid5, 6 ) ;
966+ // CallSiteId::new($a, $b) ==>> CallSiteId::new($a, LineNumber($b))
967+ let id1 = CallSiteId :: new ( fid1, LineNumber ( 2 ) ) ;
968+ let id2 = CallSiteId :: new ( fid3, LineNumber ( 45 ) ) ;
969+ let id3 = CallSiteId :: new ( fid5, LineNumber ( 6 ) ) ;
939970 cs1. start_call ( 123 , id1) ;
940971 assert_eq ! ( cs1. calls, vec![ id1] ) ;
941972
@@ -950,7 +981,11 @@ mod tests {
950981 cs2. start_call ( 12 , id3) ;
951982 assert_eq ! (
952983 cs2. calls,
953- vec![ CallSiteId :: new( fid1, 10 ) , CallSiteId :: new( fid3, 12 ) , id3]
984+ vec![
985+ CallSiteId :: new( fid1, LineNumber ( 10 ) ) ,
986+ CallSiteId :: new( fid3, LineNumber ( 12 ) ) ,
987+ id3
988+ ]
954989 ) ;
955990 }
956991
@@ -960,10 +995,10 @@ mod tests {
960995 let fid3 = FunctionId :: new ( 3u64 ) ;
961996
962997 let mut cs1 = Callstack :: new ( ) ;
963- cs1. start_call ( 0 , CallSiteId :: new ( fid1, 2 ) ) ;
998+ cs1. start_call ( 0 , CallSiteId :: new ( fid1, LineNumber ( 2 ) ) ) ;
964999 let cs1b = cs1. clone ( ) ;
9651000 let mut cs2 = Callstack :: new ( ) ;
966- cs2. start_call ( 0 , CallSiteId :: new ( fid3, 4 ) ) ;
1001+ cs2. start_call ( 0 , CallSiteId :: new ( fid3, LineNumber ( 4 ) ) ) ;
9671002 let cs3 = Callstack :: new ( ) ;
9681003 let cs3b = Callstack :: new ( ) ;
9691004
@@ -1014,7 +1049,7 @@ mod tests {
10141049
10151050 let fid1 = FunctionId :: new ( 1u64 ) ;
10161051
1017- cs1. start_call ( 0 , CallSiteId :: new ( fid1, 2 ) ) ;
1052+ cs1. start_call ( 0 , CallSiteId :: new ( fid1, LineNumber ( 2 ) ) ) ;
10181053 let id1 =
10191054 cs1. id_for_new_allocation ( 1 , |cs| interner. get_or_insert_id ( Cow :: Borrowed ( & cs) , || ( ) ) ) ;
10201055 let id2 =
@@ -1025,7 +1060,7 @@ mod tests {
10251060 assert_ne ! ( id2, id0) ;
10261061 assert_ne ! ( id2, id1) ;
10271062
1028- cs1. start_call ( 3 , CallSiteId :: new ( fid1, 2 ) ) ;
1063+ cs1. start_call ( 3 , CallSiteId :: new ( fid1, LineNumber ( 2 ) ) ) ;
10291064 let id3 =
10301065 cs1. id_for_new_allocation ( 4 , |cs| interner. get_or_insert_id ( Cow :: Borrowed ( & cs) , || ( ) ) ) ;
10311066 assert_ne ! ( id3, id0) ;
@@ -1041,7 +1076,7 @@ mod tests {
10411076 assert_eq ! ( id1, id1c) ;
10421077
10431078 // Check for cache invalidation in start_call:
1044- cs1. start_call ( 1 , CallSiteId :: new ( fid1, 1 ) ) ;
1079+ cs1. start_call ( 1 , CallSiteId :: new ( fid1, LineNumber ( 1 ) ) ) ;
10451080 let id4 =
10461081 cs1. id_for_new_allocation ( 1 , |cs| interner. get_or_insert_id ( Cow :: Borrowed ( & cs) , || ( ) ) ) ;
10471082 assert_ne ! ( id4, id0) ;
@@ -1063,9 +1098,9 @@ mod tests {
10631098
10641099 let mut tracker = new_tracker ( ) ;
10651100 let mut cs1 = Callstack :: new ( ) ;
1066- cs1. start_call ( 0 , CallSiteId :: new ( fid1, 2 ) ) ;
1101+ cs1. start_call ( 0 , CallSiteId :: new ( fid1, LineNumber ( 2 ) ) ) ;
10671102 let mut cs2 = Callstack :: new ( ) ;
1068- cs2. start_call ( 0 , CallSiteId :: new ( fid3, 4 ) ) ;
1103+ cs2. start_call ( 0 , CallSiteId :: new ( fid3, LineNumber ( 4 ) ) ) ;
10691104
10701105 let cs1_id = tracker. get_callstack_id ( & cs1) ;
10711106
@@ -1156,12 +1191,12 @@ mod tests {
11561191 . functions
11571192 . add_function ( "c" . to_string ( ) , "cf" . to_string ( ) ) ;
11581193
1159- let id1 = CallSiteId :: new ( fid1, 1 ) ;
1194+ let id1 = CallSiteId :: new ( fid1, LineNumber ( 1 ) ) ;
11601195 // Same function, different line number—should be different item:
1161- let id1_different = CallSiteId :: new ( fid1, 7 ) ;
1162- let id2 = CallSiteId :: new ( fid2, 2 ) ;
1196+ let id1_different = CallSiteId :: new ( fid1, LineNumber ( 7 ) ) ;
1197+ let id2 = CallSiteId :: new ( fid2, LineNumber ( 2 ) ) ;
11631198
1164- let id3 = CallSiteId :: new ( fid3, 3 ) ;
1199+ let id3 = CallSiteId :: new ( fid3, LineNumber ( 3 ) ) ;
11651200 let mut cs1 = Callstack :: new ( ) ;
11661201 cs1. start_call ( 0 , id1) ;
11671202 cs1. start_call ( 0 , id2. clone ( ) ) ;
@@ -1200,7 +1235,7 @@ mod tests {
12001235 "c:3 (cf) 234" ,
12011236 "a:7 (af);b:2 (bf) 6000" ,
12021237 ] ;
1203- let mut result2: Vec < String > = tracker. to_lines ( true , false ) . collect ( ) ;
1238+ let mut result2: Vec < String > = tracker. to_lines ( true , false , same_callstack ) . collect ( ) ;
12041239 result2. sort ( ) ;
12051240 expected2. sort ( ) ;
12061241 assert_eq ! ( expected2, result2) ;
0 commit comments