From 70cf369d185263e014fd15cfb46586b18255b44c Mon Sep 17 00:00:00 2001
From: Fabian <copy@copy.sh>
Date: Thu, 10 Nov 2022 11:02:47 -0600
Subject: [PATCH] collect stats about segment offsets

---
 src/browser/print_stats.js |  4 ++++
 src/rust/modrm.rs          | 17 ++++++++++++++++-
 src/rust/profiler.rs       |  4 ++++
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/browser/print_stats.js b/src/browser/print_stats.js
index db3e27d1..f4ea525d 100644
--- a/src/browser/print_stats.js
+++ b/src/browser/print_stats.js
@@ -111,6 +111,10 @@ const print_stats = {
             "MODRM_COMPLEX",
             "SEG_OFFSET_OPTIMISED",
             "SEG_OFFSET_NOT_OPTIMISED",
+            "SEG_OFFSET_NOT_OPTIMISED_ES",
+            "SEG_OFFSET_NOT_OPTIMISED_FS",
+            "SEG_OFFSET_NOT_OPTIMISED_GS",
+            "SEG_OFFSET_NOT_OPTIMISED_NOT_FLAT",
         ];
 
         let j = 0;
diff --git a/src/rust/modrm.rs b/src/rust/modrm.rs
index c8d5e9fa..8fabb9f8 100644
--- a/src/rust/modrm.rs
+++ b/src/rust/modrm.rs
@@ -5,7 +5,7 @@ use jit::JitContext;
 use prefix::{PREFIX_MASK_SEGMENT, SEG_PREFIX_ZERO};
 use profiler;
 use regs::{BP, BX, DI, SI};
-use regs::{CS, DS, SS};
+use regs::{CS, DS, ES, FS, GS, SS};
 use regs::{EAX, EBP, EBX, ECX, EDI, EDX, ESI, ESP};
 
 pub struct ModrmByte {
@@ -285,6 +285,21 @@ pub fn jit_add_seg_offset_no_override(ctx: &mut JitContext, seg: u32) {
         return;
     }
     codegen::gen_profiler_stat_increment(ctx.builder, profiler::stat::SEG_OFFSET_NOT_OPTIMISED);
+    codegen::gen_profiler_stat_increment(
+        ctx.builder,
+        if seg == ES {
+            profiler::stat::SEG_OFFSET_NOT_OPTIMISED_ES
+        }
+        else if seg == FS {
+            profiler::stat::SEG_OFFSET_NOT_OPTIMISED_FS
+        }
+        else if seg == GS {
+            profiler::stat::SEG_OFFSET_NOT_OPTIMISED_GS
+        }
+        else {
+            profiler::stat::SEG_OFFSET_NOT_OPTIMISED_NOT_FLAT
+        },
+    );
 
     if seg != CS && seg != SS {
         if cfg!(feature = "profiler") {
diff --git a/src/rust/profiler.rs b/src/rust/profiler.rs
index b0ece211..e494c7fe 100644
--- a/src/rust/profiler.rs
+++ b/src/rust/profiler.rs
@@ -112,6 +112,10 @@ pub enum stat {
 
     SEG_OFFSET_OPTIMISED,
     SEG_OFFSET_NOT_OPTIMISED,
+    SEG_OFFSET_NOT_OPTIMISED_ES,
+    SEG_OFFSET_NOT_OPTIMISED_FS,
+    SEG_OFFSET_NOT_OPTIMISED_GS,
+    SEG_OFFSET_NOT_OPTIMISED_NOT_FLAT,
 }
 
 #[allow(non_upper_case_globals)]