Rollup merge of rust-lang#157620 - Daniel-B-Smith:smithdb3/fileencoder-strategy, r=cjgillot

Add a strategy FnMut to inject behavior into the flush cycle

This PR adds the interface needed to avoid forking FileEncoder in rust-lang#154724.

This adds a FnMut used as a strategy/observer to see the bytes being flushed to disk. For rust-lang#154724, that would be used to calculate the metadata hash on the raw metadata bytes.

The 'static hack is pretty ugly, but otherwise, the FileEncoder lifetime goes exceptionally viral and requires some very unfortunate workarounds to avoid self referential lifetimes. That's a bit of a timebomb if the wrong FileEncoder client tries to use the new API, but the worst case outcome is a wild goose chase of lifetime chasing.

Full disclosure: I used Claude Code to generate much of the PR. I have fully reviewed the code and stand behind it.

r? @cjgillot

Author: jhpratt

compiler/rustc_data_structures/src/marker.rs

@@ -64,6 +64,7 @@ already_send!(
         [std::io::Error][std::fs::File][std::panic::Location<'_>][rustc_arena::DroplessArena]
         [jobserver_crate::Client][jobserver_crate::HelperThread][crate::memmap::Mmap]
         [crate::profiling::SelfProfiler][crate::owned_slice::OwnedSlice]
+        [rustc_serialize::opaque::FileEncoder<'_>]
 );
 
 #[cfg(target_has_atomic = "64")]

compiler/rustc_incremental/src/persist/file_format.rs

@@ -28,7 +28,7 @@ const FILE_MAGIC: &[u8] = b"RSIC";
 /// Change this if the header format changes.
 const HEADER_FORMAT_VERSION: u16 = 0;
 
-pub(crate) fn write_file_header(stream: &mut FileEncoder, sess: &Session) {
+pub(crate) fn write_file_header(stream: &mut FileEncoder<'_>, sess: &Session) {
     stream.emit_raw_bytes(FILE_MAGIC);
     stream.emit_raw_bytes(&u16::to_le_bytes(HEADER_FORMAT_VERSION));
 
@@ -41,7 +41,7 @@ pub(crate) fn write_file_header(stream: &mut FileEncoder, sess: &Session) {
 
 pub(crate) fn save_in<F>(sess: &Session, path_buf: PathBuf, name: &str, encode: F)
 where
-    F: FnOnce(FileEncoder) -> FileEncodeResult,
+    F: FnOnce(FileEncoder<'static>) -> FileEncodeResult,
 {
     debug!("save: storing data in {}", path_buf.display());
 

compiler/rustc_incremental/src/persist/save.rs

@@ -131,7 +131,7 @@ pub fn save_work_product_index(
     });
 }
 
-fn encode_work_product_index(work_products: &WorkProductMap, encoder: &mut FileEncoder) {
+fn encode_work_product_index(work_products: &WorkProductMap, encoder: &mut FileEncoder<'_>) {
     let serialized_products: Vec<_> = work_products
         .to_sorted_stable_ord()
         .into_iter()

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -40,7 +40,7 @@ use crate::errors::{FailCreateFileEncoder, FailWriteFile};
 use crate::rmeta::*;
 
 pub(super) struct EncodeContext<'a, 'tcx> {
-    opaque: opaque::FileEncoder,
+    opaque: opaque::FileEncoder<'a>,
     tcx: TyCtxt<'tcx>,
     feat: &'tcx rustc_feature::Features,
     tables: TableBuilders,

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -371,7 +371,7 @@ macro_rules! define_tables {
         }
 
         impl TableBuilders {
-            fn encode(&self, buf: &mut FileEncoder) -> LazyTables {
+            fn encode(&self, buf: &mut FileEncoder<'_>) -> LazyTables {
                 LazyTables {
                     $($name1: self.$name1.encode(buf),)+
                     $($name2: self.$name2.encode(buf),)+

compiler/rustc_metadata/src/rmeta/table.rs

@@ -487,7 +487,7 @@ impl<I: Idx, const N: usize, T: FixedSizeEncoding<ByteArray = [u8; N]>> TableBui
         }
     }
 
-    pub(crate) fn encode(&self, buf: &mut FileEncoder) -> LazyTable<I, T> {
+    pub(crate) fn encode(&self, buf: &mut FileEncoder<'_>) -> LazyTable<I, T> {
         let pos = buf.position();
 
         let width = self.width;

compiler/rustc_middle/src/dep_graph/graph.rs

@@ -135,7 +135,7 @@ impl DepGraph {
         session: &Session,
         prev_graph: Arc<SerializedDepGraph>,
         prev_work_products: WorkProductMap,
-        encoder: FileEncoder,
+        encoder: FileEncoder<'static>,
     ) -> DepGraph {
         let prev_graph_node_count = prev_graph.node_count();
 
@@ -1137,7 +1137,7 @@ impl CurrentDepGraph {
     fn new(
         session: &Session,
         prev_graph_node_count: usize,
-        encoder: FileEncoder,
+        encoder: FileEncoder<'static>,
         previous: Arc<SerializedDepGraph>,
     ) -> Self {
         let mut stable_hasher = StableHasher::new();

compiler/rustc_middle/src/dep_graph/serialized.rs

@@ -564,13 +564,17 @@ struct LocalEncoderResult {
 struct EncoderState {
     next_node_index: AtomicU64,
     previous: Arc<SerializedDepGraph>,
-    file: Lock<Option<FileEncoder>>,
+    file: Lock<Option<FileEncoder<'static>>>,
     local: WorkerLocal<RefCell<LocalEncoderState>>,
     stats: Option<Lock<FxHashMap<DepKind, Stat>>>,
 }
 
 impl EncoderState {
-    fn new(encoder: FileEncoder, record_stats: bool, previous: Arc<SerializedDepGraph>) -> Self {
+    fn new(
+        encoder: FileEncoder<'static>,
+        record_stats: bool,
+        previous: Arc<SerializedDepGraph>,
+    ) -> Self {
         Self {
             previous,
             next_node_index: AtomicU64::new(0),
@@ -863,7 +867,7 @@ pub(crate) struct GraphEncoder {
 impl GraphEncoder {
     pub(crate) fn new(
         sess: &Session,
-        encoder: FileEncoder,
+        encoder: FileEncoder<'static>,
         prev_node_count: usize,
         previous: Arc<SerializedDepGraph>,
     ) -> Self {

compiler/rustc_middle/src/query/on_disk_cache.rs

@@ -201,7 +201,7 @@ impl OnDiskCache {
 
     /// Serialize the current-session data that will be loaded by [`OnDiskCache`]
     /// in a subsequent incremental compilation session.
-    pub fn serialize(tcx: TyCtxt<'_>, encoder: FileEncoder) -> FileEncodeResult {
+    pub fn serialize(tcx: TyCtxt<'_>, encoder: FileEncoder<'static>) -> FileEncodeResult {
         // Serializing the `DepGraph` should not modify it.
         tcx.dep_graph.with_ignore(|| {
             // Allocate `SourceFileIndex`es.
@@ -779,7 +779,7 @@ impl_ref_decoder! {<'tcx>
 /// An encoder that can write to the incremental compilation cache.
 pub struct CacheEncoder<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
-    encoder: FileEncoder,
+    encoder: FileEncoder<'static>,
     type_shorthands: FxHashMap<Ty<'tcx>, usize>,
     predicate_shorthands: FxHashMap<ty::PredicateKind<'tcx>, usize>,
     interpret_allocs: FxIndexSet<interpret::AllocId>,

compiler/rustc_serialize/src/opaque.rs

@@ -29,7 +29,10 @@ const BUF_SIZE: usize = 64 * 1024;
 /// `Vec`. `FileEncoder` is better because its memory use is determined by the
 /// size of the buffer, rather than the full length of the encoded data, and
 /// because it doesn't need to reallocate memory along the way.
-pub struct FileEncoder {
+///
+/// The `'a` lifetime is the borrow of the optional flush strategy (see
+/// `flush_strategy`); it is unused (`'static`) for encoders created without one.
+pub struct FileEncoder<'a> {
     // The input buffer. For adequate performance, we need to be able to write
     // directly to the unwritten region of the buffer, without calling copy_from_slice.
     // Note that our buffer is always initialized so that we can do that direct access
@@ -43,11 +46,12 @@ pub struct FileEncoder {
     // comment on `trait Encoder`.
     res: Result<(), io::Error>,
     path: PathBuf,
+    flush_strategy: Option<&'a mut (dyn FnMut(&[u8]) + Send)>,
     #[cfg(debug_assertions)]
     finished: bool,
 }
 
-impl FileEncoder {
+impl<'a> FileEncoder<'a> {
     pub fn new<P: AsRef<Path>>(path: P) -> io::Result<Self> {
         // File::create opens the file for writing only. When -Zmeta-stats is enabled, the metadata
         // encoder rewinds the file to inspect what was written. So we need to always open the file
@@ -62,11 +66,21 @@ impl FileEncoder {
             flushed: 0,
             file,
             res: Ok(()),
+            flush_strategy: None,
             #[cfg(debug_assertions)]
             finished: false,
         })
     }
 
+    pub fn with_flush_strategy<P: AsRef<Path>>(
+        path: P,
+        strategy: &'a mut (dyn FnMut(&[u8]) + Send),
+    ) -> io::Result<Self> {
+        let mut encoder = Self::new(path)?;
+        encoder.flush_strategy = Some(strategy);
+        Ok(encoder)
+    }
+
     #[inline]
     pub fn position(&self) -> usize {
         // Tracking position this way instead of having a `self.position` field
@@ -86,6 +100,9 @@ impl FileEncoder {
             self.res = self.file.write_all(&self.buf[..self.buffered]);
         }
         self.flushed += self.buffered;
+        if let Some(f) = &mut self.flush_strategy {
+            f(&self.buf[..self.buffered]);
+        }
         self.buffered = 0;
     }
 
@@ -115,6 +132,9 @@ impl FileEncoder {
         } else {
             if self.res.is_ok() {
                 self.res = self.file.write_all(buf);
+                if let Some(f) = &mut self.flush_strategy {
+                    f(buf);
+                }
             }
             self.flushed += buf.len();
         }
@@ -200,7 +220,7 @@ impl FileEncoder {
 }
 
 #[cfg(debug_assertions)]
-impl Drop for FileEncoder {
+impl Drop for FileEncoder<'_> {
     fn drop(&mut self) {
         if !std::thread::panicking() {
             assert!(self.finished);
@@ -217,7 +237,7 @@ macro_rules! write_leb128 {
     };
 }
 
-impl Encoder for FileEncoder {
+impl Encoder for FileEncoder<'_> {
     write_leb128!(emit_usize, usize, write_usize_leb128);
     write_leb128!(emit_u128, u128, write_u128_leb128);
     write_leb128!(emit_u64, u64, write_u64_leb128);
@@ -415,8 +435,8 @@ impl<'a> Decoder for MemDecoder<'a> {
 
 // Specialize encoding byte slices. This specialization also applies to encoding `Vec<u8>`s, etc.,
 // since the default implementations call `encode` on their slices internally.
-impl Encodable<FileEncoder> for [u8] {
-    fn encode(&self, e: &mut FileEncoder) {
+impl Encodable<FileEncoder<'_>> for [u8] {
+    fn encode(&self, e: &mut FileEncoder<'_>) {
         Encoder::emit_usize(e, self.len());
         e.emit_raw_bytes(self);
     }
@@ -438,9 +458,9 @@ impl IntEncodedWithFixedSize {
     pub const ENCODED_SIZE: usize = 8;
 }
 
-impl Encodable<FileEncoder> for IntEncodedWithFixedSize {
+impl Encodable<FileEncoder<'_>> for IntEncodedWithFixedSize {
     #[inline]
-    fn encode(&self, e: &mut FileEncoder) {
+    fn encode(&self, e: &mut FileEncoder<'_>) {
         let start_pos = e.position();
         e.write_array(self.0.to_le_bytes());
         let end_pos = e.position();