Merge pull request #978 from ltratt/rework_trace_enum

Rework `TracedAOTBlock` enum.

ykrt/src/compile/jitc_llvm.rs

@@ -1,11 +1,11 @@
-//! An LLVM JIT backend. Currently a minimal wrapper around the fact that [MappedTrace]s are hardcoded
+//! An LLVM JIT backend. Currently a minimal wrapper around the fact that [MappedAOTBlockTrace]s are hardcoded
 //! to be compiled with LLVM.
 
 use crate::{
     compile::{CompilationError, CompiledTrace, Compiler},
     location::HotLocation,
     mt::{SideTraceInfo, MT},
-    trace::{AOTTraceIterator, TracedAOTBlock},
+    trace::{AOTTraceIterator, TraceAction},
 };
 use object::{Object, ObjectSection};
 use parking_lot::Mutex;
@@ -79,19 +79,23 @@ impl JITCLLVM {
         Arc::new(JITCLLVM)
     }
 
-    fn encode_trace(&self, irtrace: &Vec<TracedAOTBlock>) -> (Vec<*const i8>, Vec<usize>, usize) {
+    fn encode_trace(&self, irtrace: &Vec<TraceAction>) -> (Vec<*const i8>, Vec<usize>, usize) {
         let trace_len = irtrace.len();
         let mut func_names = Vec::with_capacity(trace_len);
         let mut bbs = Vec::with_capacity(trace_len);
         for blk in irtrace {
-            if blk.is_unmappable() {
-                // The block was unmappable. Indicate this with a null function name.
-                func_names.push(ptr::null());
-                // Block indices for unmappable blocks are irrelevant so we may pass anything here.
-                bbs.push(0);
-            } else {
-                func_names.push(blk.func_name().as_ptr());
-                bbs.push(blk.bb());
+            match blk {
+                TraceAction::MappedAOTBlock { func_name, bb } => {
+                    func_names.push(func_name.as_ptr());
+                    bbs.push(*bb);
+                }
+                TraceAction::UnmappableBlock => {
+                    // The block was unmappable. Indicate this with a null function name.
+                    func_names.push(ptr::null());
+                    // Block indices for unmappable blocks are irrelevant so we may pass anything here.
+                    bbs.push(0);
+                }
+                TraceAction::Promotion => todo!(),
             }
         }
         (func_names, bbs, trace_len)

ykrt/src/compile/jitc_yk/trace_builder.rs

@@ -3,7 +3,7 @@
 use super::aot_ir::{self, IRDisplay, Module};
 use super::jit_ir;
 use crate::compile::CompilationError;
-use crate::trace::TracedAOTBlock;
+use crate::trace::TraceAction;
 use std::collections::HashMap;
 
 /// The argument index of the trace inputs struct in the control point call.
@@ -17,7 +17,7 @@ struct TraceBuilder<'a> {
     /// The JIT IR this struct builds.
     jit_mod: jit_ir::Module,
     /// The mapped trace.
-    mtrace: &'a Vec<TracedAOTBlock>,
+    mtrace: &'a Vec<TraceAction>,
     // Maps an AOT instruction to a jit instruction via their index-based IDs.
     local_map: HashMap<aot_ir::InstructionID, jit_ir::InstrIdx>,
 }
@@ -29,7 +29,7 @@ impl<'a> TraceBuilder<'a> {
     ///  - `trace_name`: The eventual symbol name for the JITted code.
     ///  - `aot_mod`: The AOT IR module that the trace flows through.
     ///  - `mtrace`: The mapped trace.
-    fn new(trace_name: String, aot_mod: &'a Module, mtrace: &'a Vec<TracedAOTBlock>) -> Self {
+    fn new(trace_name: String, aot_mod: &'a Module, mtrace: &'a Vec<TraceAction>) -> Self {
         Self {
             aot_mod,
             mtrace,
@@ -39,14 +39,15 @@ impl<'a> TraceBuilder<'a> {
     }
 
     // Given a mapped block, find the AOT block ID, or return `None` if it is unmapped.
-    fn lookup_aot_block(&self, tb: &TracedAOTBlock) -> Option<aot_ir::BlockID> {
+    fn lookup_aot_block(&self, tb: &TraceAction) -> Option<aot_ir::BlockID> {
         match tb {
-            TracedAOTBlock::Mapped { func_name, bb } => {
+            TraceAction::MappedAOTBlock { func_name, bb } => {
                 let func_name = func_name.to_str().unwrap(); // safe: func names are valid UTF-8.
                 let func = self.aot_mod.func_idx(func_name);
                 Some(aot_ir::BlockID::new(func, aot_ir::BlockIdx::new(*bb)))
             }
-            TracedAOTBlock::Unmappable { .. } => None,
+            TraceAction::UnmappableBlock { .. } => None,
+            TraceAction::Promotion => todo!(),
         }
     }
 
@@ -227,15 +228,16 @@ impl<'a> TraceBuilder<'a> {
         // Find the block containing the control point call. This is the (sole) predecessor of the
         // first (guaranteed mappable) block in the trace.
         let prev = match first_blk {
-            TracedAOTBlock::Mapped { func_name, bb } => {
+            TraceAction::MappedAOTBlock { func_name, bb } => {
                 debug_assert!(*bb > 0);
                 // It's `- 1` due to the way the ykllvm block splitting pass works.
-                TracedAOTBlock::Mapped {
+                TraceAction::MappedAOTBlock {
                     func_name: func_name.clone(),
                     bb: bb - 1,
                 }
             }
-            TracedAOTBlock::Unmappable => panic!(),
+            TraceAction::UnmappableBlock => panic!(),
+            TraceAction::Promotion => todo!(),
         };
 
         let firstblk = self.lookup_aot_block(&prev);
@@ -244,11 +246,11 @@ impl<'a> TraceBuilder<'a> {
         for tblk in self.mtrace {
             match self.lookup_aot_block(tblk) {
                 Some(bid) => {
-                    // Mapped block
+                    // MappedAOTBlock block
                     self.process_block(bid)?;
                 }
                 None => {
-                    // Unmappable block
+                    // UnmappableBlock block
                     todo!();
                 }
             }
@@ -260,7 +262,7 @@ impl<'a> TraceBuilder<'a> {
 /// Given a mapped trace (through `aot_mod`), assemble and return a Yk IR trace.
 pub(super) fn build(
     aot_mod: &Module,
-    mtrace: &Vec<TracedAOTBlock>,
+    mtrace: &Vec<TraceAction>,
 ) -> Result<jit_ir::Module, CompilationError> {
     // FIXME: the XXX below should be a thread-safe monotonically incrementing integer.
     TraceBuilder::new("__yk_compiled_trace_XXX".into(), aot_mod, mtrace).build()

ykrt/src/trace/hwt/mapper.rs

@@ -1,6 +1,6 @@
 //! The mapper translates a hwtracer trace into an IR trace.
 
-use crate::trace::TracedAOTBlock;
+use crate::trace::TraceAction;
 use hwtracer::llvm_blockmap::LLVM_BLOCK_MAP;
 use hwtracer::Trace;
 use std::error::Error;
@@ -53,7 +53,7 @@ impl HWTMapper {
     /// During codegen LLVM may remove the unconditional jump and simply place bb1 and bb2
     /// consecutively, allowing bb1 to fall-thru to bb2. In the eyes of hwtracer, a fall-thru does
     /// not terminate a block, so whereas LLVM sees two blocks, hwtracer sees only one.
-    fn map_block(&mut self, block: &hwtracer::Block) -> Vec<Option<TracedAOTBlock>> {
+    fn map_block(&mut self, block: &hwtracer::Block) -> Vec<Option<TraceAction>> {
         let b_rng = block.vaddr_range();
         if b_rng.is_none() {
             // If the address range of the block isn't known, then it follows that we can't map
@@ -103,7 +103,7 @@ impl HWTMapper {
                 );
                 if let Some(sym_name) = sio.dli_sname() {
                     for bb in ent.value.corr_bbs() {
-                        ret.push(Some(TracedAOTBlock::new_mapped(
+                        ret.push(Some(TraceAction::new_mapped_aot_block(
                             sym_name.to_owned(),
                             usize::try_from(*bb).unwrap(),
                         )));
@@ -128,11 +128,8 @@ impl HWTMapper {
     ///
     /// The returned trace will always start with a mapped block (the unmappable prefix of the
     /// foreign "turn on tracing" routine is omitted).
-    pub fn map_trace(
-        &mut self,
-        trace: Box<dyn Trace>,
-    ) -> Result<Vec<TracedAOTBlock>, Box<dyn Error>> {
-        let mut ret: Vec<TracedAOTBlock> = Vec::new();
+    pub fn map_trace(&mut self, trace: Box<dyn Trace>) -> Result<Vec<TraceAction>, Box<dyn Error>> {
+        let mut ret: Vec<TraceAction> = Vec::new();
 
         let mut trace_iter = trace.iter_blocks();
         // The first block contains the control point so we need to remove that.
@@ -142,7 +139,7 @@ impl HWTMapper {
                 // we know that the control point will be contained in a single mappable block.
                 assert!(matches!(
                     self.map_block(&x).as_slice(),
-                    &[Some(TracedAOTBlock::Mapped { .. })]
+                    &[Some(TraceAction::MappedAOTBlock { .. })]
                 ));
             }
             _ => unreachable!(),
@@ -159,8 +156,12 @@ impl HWTMapper {
                 // trace isn't empty (we never report the leading unmappable code in a trace). We
                 // also take care to collapse consecutive unmappable blocks into one.
                 if let Some(last) = ret.last_mut() {
-                    if !last.is_unmappable() {
-                        ret.push(TracedAOTBlock::new_unmappable());
+                    match last {
+                        TraceAction::MappedAOTBlock { .. } => {
+                            ret.push(TraceAction::new_unmappable_block());
+                        }
+                        TraceAction::UnmappableBlock => (),
+                        TraceAction::Promotion => todo!(),
                     }
                 }
             } else {
@@ -202,7 +203,7 @@ impl HWTMapper {
             Some(x) => {
                 // This is a rough proxy for "check that we removed only the thing we want to
                 // remove".
-                assert!(matches!(x, TracedAOTBlock::Unmappable));
+                assert!(matches!(x, TraceAction::UnmappableBlock));
             }
             _ => unreachable!(),
         }

ykrt/src/trace/hwt/mod.rs

@@ -1,6 +1,6 @@
 //! Hardware tracing via hwtracer.
 
-use super::{errors::InvalidTraceError, AOTTraceIterator, TraceRecorder, TracedAOTBlock};
+use super::{errors::InvalidTraceError, AOTTraceIterator, TraceAction, TraceRecorder};
 use std::{error::Error, sync::Arc};
 
 pub(crate) mod mapper;
@@ -50,13 +50,13 @@ impl TraceRecorder for HWTTraceRecorder {
 }
 
 struct HWTTraceIterator {
-    trace: std::vec::IntoIter<TracedAOTBlock>,
+    trace: std::vec::IntoIter<TraceAction>,
 }
 
 impl AOTTraceIterator for HWTTraceIterator {}
 
 impl Iterator for HWTTraceIterator {
-    type Item = TracedAOTBlock;
+    type Item = TraceAction;
     fn next(&mut self) -> Option<Self::Item> {
         self.trace.next()
     }

ykrt/src/trace/mod.rs

@@ -1,15 +1,23 @@
 //! Record and process traces.
+//!
+//! "Tracing" is split into the following stages:
+//!
+//! 1. *Record* the trace with a [Tracer], which abstracts over a specific *tracer backend*. The
+//!    tracer backend may use one of several low-level tracing methods (e.g. a hardware tracer like
+//!    PT or a software tracer). The tracer backend stores the recorded low-level trace in an
+//!    internal format of its choosing.
+//! 2. *Process* the recorded trace. The tracer backend returns an iterator which produces
+//!    [TraceAction]s.
+//! 3. *Compile* the processed trace. That happens in [compile](crate::compile) module.
+//!
+//! This module thus contains tracing backends which can record and process traces.
 
 #![allow(clippy::len_without_is_empty)]
 #![allow(clippy::new_without_default)]
 #![allow(clippy::missing_safety_doc)]
 
 mod errors;
-use std::{
-    error::Error,
-    ffi::{CStr, CString},
-    sync::Arc,
-};
+use std::{error::Error, ffi::CString, sync::Arc};
 
 #[cfg(tracer_hwt)]
 pub(crate) mod hwt;
@@ -37,21 +45,21 @@ pub(crate) fn default_tracer() -> Result<Arc<dyn Tracer>, Box<dyn Error>> {
     Err("No tracing backend for this platform/configuration.".into())
 }
 
-/// A thread which is currently recording a trace.
+/// An instance of a [Tracer] which is currently recording a trace of the current thread.
 pub(crate) trait TraceRecorder {
     /// Stop recording a trace of the current thread and return an iterator which successively
-    /// produces the traced blocks.
+    /// produces [TraceAction]s.
     fn stop(self: Box<Self>) -> Result<Box<dyn AOTTraceIterator>, InvalidTraceError>;
 }
 
-/// An iterator which takes an underlying raw trace and successively produces [TracedAOTBlock]s.
-pub(crate) trait AOTTraceIterator: Iterator<Item = TracedAOTBlock> + Send {}
+/// An iterator which [TraceRecord]s use to process a trace into [TraceAction]s.
+pub(crate) trait AOTTraceIterator: Iterator<Item = TraceAction> + Send {}
 
-/// An AOT LLVM IR block that has been traced at JIT time.
+/// A processed item from a trace.
 #[derive(Debug, Eq, PartialEq)]
-pub enum TracedAOTBlock {
+pub enum TraceAction {
     /// A sucessfully mapped block.
-    Mapped {
+    MappedAOTBlock {
         /// The name of the function containing the block.
         ///
         /// PERF: Use a string pool to avoid duplicated function names in traces.
@@ -62,42 +70,26 @@ pub enum TracedAOTBlock {
     /// One or more machine blocks that could not be mapped.
     ///
     /// This usually means that the blocks were compiled outside of ykllvm.
-    Unmappable,
+    UnmappableBlock,
+    /// A value promoted and recorded within the low-level trace (e.g. `PTWRITE`). In essence these
+    /// are calls to `yk_promote` that have been inlined so that the tracer backend can handle them
+    /// rather than being handled by yk's generic run-time support for `yk_promote`.
+    ///
+    /// While no tracer backend currently uses this variant, it's present to remind us that this a
+    /// useful possibility.
+    Promotion,
 }
 
-impl TracedAOTBlock {
-    pub fn new_mapped(func_name: CString, bb: usize) -> Self {
+impl TraceAction {
+    pub fn new_mapped_aot_block(func_name: CString, bb: usize) -> Self {
         // At one point, `bb = usize::MAX` was a special value, but it no longer is. We believe
         // that no part of the code sets/checks for this value, but just in case there is a
         // laggardly part of the code which does so, we've left this `assert` behind to catch it.
         debug_assert_ne!(bb, usize::MAX);
-        Self::Mapped { func_name, bb }
-    }
-
-    pub fn new_unmappable() -> Self {
-        Self::Unmappable
-    }
-
-    /// If `self` is a mapped block, return the function name, otherwise panic.
-    pub fn func_name(&self) -> &CStr {
-        if let Self::Mapped { func_name, .. } = self {
-            func_name.as_c_str()
-        } else {
-            panic!();
-        }
-    }
-
-    /// If `self` is a mapped block, return the basic block index, otherwise panic.
-    pub fn bb(&self) -> usize {
-        if let Self::Mapped { bb, .. } = self {
-            *bb
-        } else {
-            panic!();
-        }
+        Self::MappedAOTBlock { func_name, bb }
     }
 
-    /// Determines whether `self` represents unmappable code.
-    pub fn is_unmappable(&self) -> bool {
-        matches!(self, Self::Unmappable)
+    pub fn new_unmappable_block() -> Self {
+        Self::UnmappableBlock
     }
 }