Move data recorded with yk_promote into MTThread.

I've slowly come to realise that calls to `yk_promote` could (but don't
currently!) end up being in compiled in two different ways:

  1. A literal call to `yk_promote_*`.
  2. An inlined call to a side-band recording mechanism (e.g.
     `PTWRITE`).

Because I hadn't teased these two apart in my mind, my previous attempt
had kind-of tried to merge the two together in one API. It's now clear
that was entirely misguided on my part.

This commit can be seen as partly undoing my previous attempt: literal
calls to `yk_promote_*` now record the data in `MTThread` which is a
thread local. However, when tracing stops, the promotion data is now
extracted from `MTThread` so it can safely be moved to another thread.

That means we can now handle ykjit#1 above: but we don't yet have an API that
can handle ykjit#2.

ykrt/src/mt.rs

@@ -263,19 +263,29 @@ impl MT {
                 match Arc::clone(&tracer).start_recorder() {
                     Ok(tt) => THREAD_MTTHREAD.with(|mtt| {
                         *mtt.thread_tracer.borrow_mut() = Some(tt);
+                        *mtt.promotions.borrow_mut() = Some(Vec::new());
                     }),
                     Err(e) => todo!("{e:?}"),
                 }
             }
             TransitionControlPoint::StopTracing(hl_arc) => {
-                // Assuming no bugs elsewhere, the `unwrap` cannot fail, because `StartTracing`
+                // Assuming no bugs elsewhere, the `unwrap`s cannot fail, because `StartTracing`
                 // will have put a `Some` in the `Rc`.
-                let thrdtrcr = THREAD_MTTHREAD.with(|mtt| mtt.thread_tracer.take().unwrap());
+                let (thrdtrcr, promotions) = THREAD_MTTHREAD.with(|mtt| {
+                    (
+                        mtt.thread_tracer.take().unwrap(),
+                        mtt.promotions.take().unwrap(),
+                    )
+                });
                 match thrdtrcr.stop() {
                     Ok(utrace) => {
                         #[cfg(feature = "yk_jitstate_debug")]
                         print_jit_state("stop-tracing");
-                        self.queue_compile_job(utrace, hl_arc, None);
+                        self.queue_compile_job(
+                            (utrace, promotions.into_boxed_slice()),
+                            hl_arc,
+                            None,
+                        );
                     }
                     Err(_e) => {
                         #[cfg(feature = "yk_jitstate_debug")]
@@ -284,14 +294,23 @@ impl MT {
                 }
             }
             TransitionControlPoint::StopSideTracing(hl_arc, sti, parent) => {
-                // Assuming no bugs elsewhere, the `unwrap` cannot fail, because `StartTracing`
-                // will have put a `Some` in the `Rc`.
-                let thrdtrcr = THREAD_MTTHREAD.with(|mtt| mtt.thread_tracer.take().unwrap());
+                // Assuming no bugs elsewhere, the `unwrap`s cannot fail, because
+                // `StartSideTracing` will have put a `Some` in the `Rc`.
+                let (thrdtrcr, promotions) = THREAD_MTTHREAD.with(|mtt| {
+                    (
+                        mtt.thread_tracer.take().unwrap(),
+                        mtt.promotions.take().unwrap(),
+                    )
+                });
                 match thrdtrcr.stop() {
                     Ok(utrace) => {
                         #[cfg(feature = "yk_jitstate_debug")]
                         print_jit_state("stop-side-tracing");
-                        self.queue_compile_job(utrace, hl_arc, Some((sti, parent)));
+                        self.queue_compile_job(
+                            (utrace, promotions.into_boxed_slice()),
+                            hl_arc,
+                            Some((sti, parent)),
+                        );
                     }
                     Err(_e) => {
                         #[cfg(feature = "yk_jitstate_debug")]
@@ -556,6 +575,7 @@ impl MT {
                 match Arc::clone(&tracer).start_recorder() {
                     Ok(tt) => THREAD_MTTHREAD.with(|mtt| {
                         *mtt.thread_tracer.borrow_mut() = Some(tt);
+                        *mtt.promotions.borrow_mut() = Some(Vec::new());
                     }),
                     Err(e) => todo!("{e:?}"),
                 }
@@ -594,7 +614,10 @@ pub(crate) struct MTThread {
     /// When tracing is active, this will be `RefCell<Some(...)>`; when tracing is inactive
     /// `RefCell<None>`. We need to keep track of the [Tracer] used to start the [ThreadTracer], as
     /// trace mapping requires a reference to the [Tracer].
-    pub(crate) thread_tracer: RefCell<Option<Box<dyn TraceRecorder>>>,
+    thread_tracer: RefCell<Option<Box<dyn TraceRecorder>>>,
+    /// Records calls to `yk_promote`. When tracing is active, this will be `RefCell<Some(...)>`;
+    /// when tracing is inactive `RecCell<None>`.
+    promotions: RefCell<Option<Vec<usize>>>,
     // Raw pointers are neither send nor sync.
     _dont_send_or_sync_me: PhantomData<*mut ()>,
 }
@@ -604,9 +627,23 @@ impl MTThread {
         MTThread {
             tracing: RefCell::new(None),
             thread_tracer: RefCell::new(None),
+            promotions: RefCell::new(None),
             _dont_send_or_sync_me: PhantomData,
         }
     }
+
+    /// Records `val` as a value to be promoted. Returns `true` if either: no trace is being
+    /// recorded; or recording the promotion succeeded.
+    ///
+    /// If `false` is returned, the current trace is unable to record the promotion successfully
+    /// and further calls are probably pointless, though they will not cause the tracer to enter
+    /// undefined behaviour territory.
+    pub(crate) fn promote_usize(&self, val: usize) -> bool {
+        if let Some(promotions) = &mut *self.promotions.borrow_mut() {
+            promotions.push(val);
+        }
+        true
+    }
 }
 
 /// What action should a caller of [MT::transition_control_point] take?

ykrt/src/promote.rs

@@ -11,10 +11,8 @@ use crate::mt::THREAD_MTTHREAD;
 #[no_mangle]
 pub extern "C" fn __yk_promote_usize(val: usize) {
     THREAD_MTTHREAD.with(|mtt| {
-        if let Some(tt) = mtt.thread_tracer.borrow().as_ref() {
-            // We ignore the return value for `promote_usize` as we can't really cancel tracing from
-            // this function.
-            tt.promote_usize(val);
-        }
+        // We ignore the return value for `promote_usize` as we can't really cancel tracing from
+        // this function.
+        mtt.promote_usize(val);
     });
 }

ykrt/src/trace/hwt/mod.rs

@@ -1,7 +1,7 @@
 //! Hardware tracing via hwtracer.
 
 use super::{errors::InvalidTraceError, AOTTraceIterator, TraceRecorder, TracedAOTBlock};
-use std::{cell::RefCell, error::Error, sync::Arc};
+use std::{error::Error, sync::Arc};
 
 pub(crate) mod mapper;
 pub(crate) use mapper::HWTMapper;
@@ -23,21 +23,17 @@ impl super::Tracer for HWTracer {
     fn start_recorder(self: Arc<Self>) -> Result<Box<dyn TraceRecorder>, Box<dyn Error>> {
         Ok(Box::new(HWTTraceRecorder {
             thread_tracer: Arc::clone(&self.backend).start_collector()?,
-            promotions: RefCell::new(Vec::new()),
         }))
     }
 }
 
 /// Hardware thread tracer.
 struct HWTTraceRecorder {
     thread_tracer: Box<dyn hwtracer::ThreadTracer>,
-    promotions: RefCell<Vec<usize>>,
 }
 
 impl TraceRecorder for HWTTraceRecorder {
-    fn stop(
-        self: Box<Self>,
-    ) -> Result<(Box<dyn AOTTraceIterator>, Box<[usize]>), InvalidTraceError> {
+    fn stop(self: Box<Self>) -> Result<Box<dyn AOTTraceIterator>, InvalidTraceError> {
         let tr = self.thread_tracer.stop_collector().unwrap();
         let mut mt = HWTMapper::new();
         let mapped = mt
@@ -46,19 +42,11 @@ impl TraceRecorder for HWTTraceRecorder {
         if mapped.is_empty() {
             Err(InvalidTraceError::EmptyTrace)
         } else {
-            Ok((
-                Box::new(HWTTraceIterator {
-                    trace: mapped.into_iter(),
-                }),
-                self.promotions.into_inner().into_boxed_slice(),
-            ))
+            Ok(Box::new(HWTTraceIterator {
+                trace: mapped.into_iter(),
+            }))
         }
     }
-
-    fn promote_usize(&self, val: usize) -> bool {
-        self.promotions.borrow_mut().push(val);
-        true
-    }
 }
 
 struct HWTTraceIterator {

ykrt/src/trace/mod.rs

@@ -41,14 +41,7 @@ pub(crate) fn default_tracer() -> Result<Arc<dyn Tracer>, Box<dyn Error>> {
 pub(crate) trait TraceRecorder {
     /// Stop recording a trace of the current thread and return an iterator which successively
     /// produces the traced blocks.
-    fn stop(
-        self: Box<Self>,
-    ) -> Result<(Box<dyn AOTTraceIterator>, Box<[usize]>), InvalidTraceError>;
-    /// Records `val` as a value to be promoted at this point in the trace. Returns `true` if
-    /// recording succeeded or `false` otherwise. If `false` is returned, this `TraceRecorder` is
-    /// no longer able to trace successfully and further calls are probably pointless, though they
-    /// must not cause the tracer to enter undefined behaviour territory.
-    fn promote_usize(&self, val: usize) -> bool;
+    fn stop(self: Box<Self>) -> Result<Box<dyn AOTTraceIterator>, InvalidTraceError>;
 }
 
 /// An iterator which takes an underlying raw trace and successively produces [TracedAOTBlock]s.