Add with_ctx to allow attaching context to decoder in the middle of decoding

Author: branchseer

Cargo.toml

@@ -37,6 +37,7 @@ unty = "0.0.3"
 
 # Used for tests
 [dev-dependencies]
+ouroboros = "0.18.3"
 serde_derive = "1.0"
 serde_json = { version = "1.0", default-features = false }
 tempfile = "3.2"

src/de/decoder.rs

@@ -95,3 +95,45 @@ impl<R: Reader, C: Config, Ctx> Decoder for DecoderImpl<R, C, Ctx> {
         &mut self.ctx
     }
 }
+
+pub struct WithContext<'a, D: ?Sized, C> {
+    pub(crate) decoder: &'a mut D,
+    pub(crate) ctx: &'a mut C,
+}
+
+impl<'a, C, D: Decoder + ?Sized> Sealed for WithContext<'a, D, C> {}
+
+impl<'a, Ctx, D: Decoder + ?Sized> Decoder for WithContext<'a, D, Ctx> {
+    type R = D::R;
+
+    type C = D::C;
+
+    type Ctx = Ctx;
+
+    fn ctx(&mut self) -> &mut Self::Ctx {
+        &mut self.ctx
+    }
+
+    fn reader(&mut self) -> &mut Self::R {
+        self.decoder.reader()
+    }
+
+    fn config(&self) -> &Self::C {
+        self.decoder.config()
+    }
+
+    fn claim_bytes_read(&mut self, n: usize) -> Result<(), DecodeError> {
+        self.decoder.claim_bytes_read(n)
+    }
+
+    fn unclaim_bytes_read(&mut self, n: usize) {
+        self.decoder.unclaim_bytes_read(n)
+    }
+}
+
+impl<'de, 'a, C, D: BorrowDecoder<'de>> BorrowDecoder<'de> for WithContext<'a, D, C> {
+    type BR = D::BR;
+    fn borrow_reader(&mut self) -> &mut Self::BR {
+        self.decoder.borrow_reader()
+    }
+}

src/de/mod.rs

@@ -5,7 +5,10 @@ mod impl_core;
 mod impl_tuples;
 mod impls;
 
-use self::read::{BorrowReader, Reader};
+use self::{
+    decoder::WithContext,
+    read::{BorrowReader, Reader},
+};
 use crate::{
     config::{Config, InternalLimitConfig},
     error::DecodeError,
@@ -138,6 +141,10 @@ pub trait Decoder: Sealed {
 
     fn ctx(&mut self) -> &mut Self::Ctx;
 
+    fn with_ctx<'a, C>(&'a mut self, ctx: &'a mut C) -> WithContext<'a, Self, C> {
+        WithContext { decoder: self, ctx }
+    }
+
     /// Returns a mutable reference to the reader
     fn reader(&mut self) -> &mut Self::R;
 

tests/ctx.rs

@@ -1,6 +1,6 @@
 use bincode::{
-    config, de::BorrowDecoder, decode_from_slice_with_ctx, encode_to_vec, error::DecodeError,
-    BorrowDecode, Decode, Encode,
+    config, de::BorrowDecoder, decode_from_slice, decode_from_slice_with_ctx, encode_to_vec,
+    error::DecodeError, BorrowDecode, Decode, Encode,
 };
 use bumpalo::{collections::Vec, vec, Bump};
 
@@ -66,16 +66,36 @@ enum _EnumContainer<'bump> {
     Vec(CodableVec<'bump, u32>),
 }
 
+#[ouroboros::self_referencing]
+struct SelfReferencing {
+    bump: Bump,
+    #[borrows(bump)]
+    #[not_covariant]
+    container: Container<'this>,
+}
+
+impl<C> Decode<C> for SelfReferencing {
+    fn decode<D: bincode::de::Decoder<Ctx = C>>(decoder: &mut D) -> Result<Self, DecodeError> {
+        SelfReferencing::try_new(Bump::new(), |mut bump| {
+            Container::decode(&mut decoder.with_ctx(&mut bump))
+        })
+    }
+}
+
 #[test]
 fn decode_with_context() {
+    let config = config::standard();
     let bump = Bump::new();
     let container = Container {
         vec: CodableVec(vec![in &bump; 1, 2, 3]),
     };
 
-    let bytes = encode_to_vec(&container, config::standard()).unwrap();
+    let bytes = encode_to_vec(&container, config).unwrap();
     let (decoded_container, _) =
-        decode_from_slice_with_ctx::<_, Container, _>(&bytes, config::standard(), &bump).unwrap();
+        decode_from_slice_with_ctx::<_, Container, _>(&bytes, config, &bump).unwrap();
 
     assert_eq!(container, decoded_container);
+
+    let self_referencing: SelfReferencing = decode_from_slice(&bytes, config).unwrap().0;
+    self_referencing.with_container(|c| assert_eq!(&container, c))
 }