YJIT: Count the number of actually written bytes (#7658)

2 files changed, 31 insertions, 13 deletions

diff --git a/yjit/src/asm/mod.rs b/yjit/src/asm/mod.rs
index e0d4fb3cc0..9df246dcbb 100644
--- a/yjit/src/asm/mod.rs
+++ b/yjit/src/asm/mod.rs
@@ -62,6 +62,9 @@ pub struct CodeBlock {
     // Current writing position
     write_pos: usize,
 
+    // Total number of bytes written to past pages
+    past_page_bytes: usize,
+
     // Size reserved for writing a jump to the next page
     page_end_reserve: usize,
 
@@ -119,6 +122,7 @@ impl CodeBlock {
             mem_size,
             page_size,
             write_pos: 0,
+            past_page_bytes: 0,
             page_end_reserve: JMP_PTR_BYTES,
             label_addrs: Vec::new(),
             label_names: Vec::new(),
@@ -201,6 +205,9 @@ impl CodeBlock {
                 assert!(!cb.has_dropped_bytes());
             });
 
+            // Update past_page_bytes for code_size()
+            self.past_page_bytes += self.current_page_bytes();
+
             // Start the next code from dst_pos
             self.write_pos = dst_pos;
         }
@@ -327,20 +334,14 @@ impl CodeBlock {
         addrs
     }
 
-    /// Return the code size that has been used by this CodeBlock.
+    /// Return the number of bytes written by this CodeBlock.
     pub fn code_size(&self) -> usize {
-        let mut size = 0;
-        let current_page_idx = self.write_pos / self.page_size;
-        for page_idx in 0..self.num_mapped_pages() {
-            if page_idx == current_page_idx {
-                // Count only actually used bytes for the current page.
-                size += (self.write_pos % self.page_size).saturating_sub(self.page_start());
-            } else if !self.has_freed_page(page_idx) {
-                // Count an entire range for any non-freed pages that have been used.
-                size += self.page_end() - self.page_start() + self.page_end_reserve;
-            }
-        }
-        size
+        self.current_page_bytes() + self.past_page_bytes
+    }
+
+    /// Return the number of bytes written to the current page.
+    fn current_page_bytes(&self) -> usize {
+        (self.write_pos % self.page_size).saturating_sub(self.page_start())
     }
 
     /// Check if this code block has sufficient remaining capacity
@@ -636,6 +637,7 @@ impl CodeBlock {
         if let Some(&first_page) = freed_pages.first() {
             for cb in [&mut *self, ocb.unwrap()] {
                 cb.write_pos = cb.get_page_pos(first_page);
+                cb.past_page_bytes = 0;
                 cb.dropped_bytes = false;
                 cb.clear_comments();
             }
@@ -796,4 +798,18 @@ mod tests
         assert_eq!(uimm_num_bits((u32::MAX as u64) + 1), 64);
         assert_eq!(uimm_num_bits(u64::MAX), 64);
     }
+
+    #[test]
+    fn test_code_size() {
+        let mut cb = CodeBlock::new_dummy(CodeBlock::PREFERRED_CODE_PAGE_SIZE * 2);
+        cb.write_bytes(&[0, 0, 0, 0]);
+        assert_eq!(cb.code_size(), 4);
+
+        // Moving to the next page should not increase code_size
+        cb.next_page(cb.get_write_ptr(), |_, _| {});
+        assert_eq!(cb.code_size(), 4);
+
+        cb.write_bytes(&[0, 0, 0, 0]);
+        assert_eq!(cb.code_size(), 8);
+    }
 }
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 0e832dbe71..f937c6e625 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -480,6 +480,8 @@ fn rb_yjit_gen_stats_dict(context: bool) -> VALUE {
         #[cfg(feature="stats")]
         hash_aset_usize!(hash, "yjit_alloc_size", global_allocation_size());
 
+        // `context` is true at RubyVM::YJIT._print_stats for --yjit-stats. It's false by default
+        // for RubyVM::YJIT.runtime_stats because counting all Contexts could be expensive.
         if context {
             let live_context_count = get_live_context_count();
             let context_size = std::mem::size_of::<Context>();