From 0f9bd2dc4c29f5e495f755ff365aac1326bcb5d8 Mon Sep 17 00:00:00 2001
From: John-John Tedro <udoprog@tedro.se>
Date: Wed, 8 Jan 2025 11:49:43 +0100
Subject: [PATCH] Clean up Context impl

---
 crates/musli/src/context/capture.rs         | 206 ++++++++++++++++++++
 crates/musli/src/context/context_error.rs   |   2 +
 crates/musli/src/context/default_context.rs | 178 +----------------
 crates/musli/src/context/mod.rs             |   6 +-
 crates/musli/src/context/trace.rs           |  62 +++---
 crates/musli/src/json/parser/into_parser.rs |   5 +-
 crates/musli/src/reader.rs                  |  10 +-
 crates/musli/src/writer.rs                  |   5 +-
 8 files changed, 274 insertions(+), 200 deletions(-)
 create mode 100644 crates/musli/src/context/capture.rs

diff --git a/crates/musli/src/context/capture.rs b/crates/musli/src/context/capture.rs
new file mode 100644
index 000000000..9d1482aac
--- /dev/null
+++ b/crates/musli/src/context/capture.rs
@@ -0,0 +1,206 @@
+use core::cell::UnsafeCell;
+use core::error::Error;
+use core::fmt;
+use core::marker::PhantomData;
+
+use super::{ContextError, ErrorMarker};
+
+mod sealed {
+    pub trait Sealed {}
+    impl Sealed for super::NoCapture {}
+    impl<E, A> Sealed for super::CaptureError<E, A> {}
+    impl<E, A> Sealed for super::SameError<E, A> {}
+}
+
+/// The trait governing how error capture is implemented.
+///
+/// See [`DefaultContext::with_capture`] or [`DefaultContext::with_same`] for
+/// more information.
+///
+/// [`DefaultContext::with_capture`]: super::DefaultContext::with_capture
+/// [`DefaultContext::with_same`]: super::DefaultContext::with_same
+pub trait Capture<A>
+where
+    Self: self::sealed::Sealed,
+{
+    #[doc(hidden)]
+    type Error;
+
+    #[doc(hidden)]
+    fn clear(&self);
+
+    #[doc(hidden)]
+    fn message<T>(&self, alloc: A, message: T) -> Self::Error
+    where
+        T: fmt::Display;
+
+    #[doc(hidden)]
+    fn custom<T>(&self, alloc: A, error: T) -> Self::Error
+    where
+        T: 'static + Send + Sync + Error;
+}
+
+/// Disable error capture.
+#[non_exhaustive]
+pub struct NoCapture;
+
+impl<A> Capture<A> for NoCapture {
+    type Error = ErrorMarker;
+
+    #[inline]
+    fn clear(&self) {}
+
+    #[inline]
+    fn message<T>(&self, alloc: A, message: T) -> Self::Error
+    where
+        T: fmt::Display,
+    {
+        _ = alloc;
+        _ = message;
+        ErrorMarker
+    }
+
+    #[inline]
+    fn custom<T>(&self, alloc: A, error: T) -> Self::Error
+    where
+        T: 'static + Send + Sync + Error,
+    {
+        _ = alloc;
+        _ = error;
+        ErrorMarker
+    }
+}
+
+/// Capture an error of the specified type.
+///
+/// See [`DefaultContext::with_same`] for more information.
+///
+/// [`DefaultContext::with_same`]: super::DefaultContext::with_same
+pub struct SameError<E, A> {
+    _marker: PhantomData<(E, A)>,
+}
+
+impl<E, A> SameError<E, A> {
+    #[inline]
+    pub(super) fn new() -> Self {
+        Self {
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<E, A> Capture<A> for SameError<E, A>
+where
+    E: ContextError<A>,
+{
+    type Error = E;
+
+    #[inline]
+    fn clear(&self) {}
+
+    #[inline]
+    fn message<T>(&self, alloc: A, message: T) -> Self::Error
+    where
+        T: fmt::Display,
+    {
+        E::message(alloc, message)
+    }
+
+    #[inline]
+    fn custom<T>(&self, alloc: A, error: T) -> Self::Error
+    where
+        T: 'static + Send + Sync + Error,
+    {
+        E::custom(alloc, error)
+    }
+}
+
+/// Capture an error of the specified type.
+///
+/// See [`DefaultContext::with_capture`] for more information.
+///
+/// [`DefaultContext::with_capture`]: super::DefaultContext::with_capture
+pub struct CaptureError<E, A> {
+    error: UnsafeCell<Option<E>>,
+    _marker: PhantomData<A>,
+}
+
+impl<E, A> CaptureError<E, A> {
+    #[inline]
+    pub(super) fn new() -> Self {
+        Self {
+            error: UnsafeCell::new(None),
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<E, A> CaptureError<E, A> {
+    #[inline]
+    pub(super) fn unwrap(&self) -> E {
+        // SAFETY: We're restricting access to the context, so that this is
+        // safe.
+        unsafe {
+            match (*self.error.get()).take() {
+                Some(error) => error,
+                None => panic!("no error captured"),
+            }
+        }
+    }
+
+    #[inline]
+    pub(super) fn result(&self) -> Result<(), E> {
+        // SAFETY: We're restricting access to the context, so that this is
+        // safe.
+        unsafe {
+            match (*self.error.get()).take() {
+                Some(error) => Err(error),
+                None => Ok(()),
+            }
+        }
+    }
+}
+
+impl<E, A> Capture<A> for CaptureError<E, A>
+where
+    E: ContextError<A>,
+{
+    type Error = ErrorMarker;
+
+    #[inline]
+    fn clear(&self) {
+        // SAFETY: We're restricting access to the context, so that this is
+        // safe.
+        unsafe {
+            (*self.error.get()) = None;
+        }
+    }
+
+    #[inline]
+    fn message<T>(&self, alloc: A, message: T) -> Self::Error
+    where
+        T: fmt::Display,
+    {
+        // SAFETY: We're restricting access to the context, so that this is
+        // safe.
+        unsafe {
+            (*self.error.get()) = Some(E::message(alloc, message));
+        }
+
+        ErrorMarker
+    }
+
+    #[inline]
+    fn custom<T>(&self, alloc: A, error: T) -> Self::Error
+    where
+        T: 'static + Send + Sync + Error,
+    {
+        // SAFETY: We're restricting access to the context, so that this is
+        // safe.
+        unsafe {
+            (*self.error.get()) = Some(E::custom(alloc, error));
+        }
+
+        ErrorMarker
+    }
+}
diff --git a/crates/musli/src/context/context_error.rs b/crates/musli/src/context/context_error.rs
index d58610675..b43f4c6b6 100644
--- a/crates/musli/src/context/context_error.rs
+++ b/crates/musli/src/context/context_error.rs
@@ -34,6 +34,7 @@ pub trait ContextError<A> {
 
 #[cfg(feature = "std")]
 impl<A> ContextError<A> for std::io::Error {
+    #[inline]
     fn custom<T>(_: A, message: T) -> Self
     where
         T: 'static + Send + Sync + Error,
@@ -41,6 +42,7 @@ impl<A> ContextError<A> for std::io::Error {
         std::io::Error::new(std::io::ErrorKind::Other, message)
     }
 
+    #[inline]
     fn message<T>(_: A, message: T) -> Self
     where
         T: fmt::Display,
diff --git a/crates/musli/src/context/default_context.rs b/crates/musli/src/context/default_context.rs
index 77f8ca85b..69cc67fb4 100644
--- a/crates/musli/src/context/default_context.rs
+++ b/crates/musli/src/context/default_context.rs
@@ -1,155 +1,14 @@
-use core::cell::UnsafeCell;
 use core::error::Error;
 use core::fmt;
-use core::marker::PhantomData;
 
 #[cfg(feature = "alloc")]
 use crate::alloc::System;
 use crate::{Allocator, Context};
 
-use super::{ContextError, ErrorMarker, Errors, NoTrace, Report, Trace, TraceConfig, TraceImpl};
-
-/// The trait implementing how an error is captured.
-pub trait Capture<A>
-where
-    A: Allocator,
-{
-    #[doc(hidden)]
-    type Error;
-
-    #[doc(hidden)]
-    fn clear(&self);
-
-    #[doc(hidden)]
-    fn message<T>(&self, alloc: A, message: T) -> Self::Error
-    where
-        T: fmt::Display;
-
-    #[doc(hidden)]
-    fn custom<T>(&self, alloc: A, error: T) -> Self::Error
-    where
-        T: 'static + Send + Sync + Error;
-}
-
-/// Disable error capture.
-#[non_exhaustive]
-pub struct NoCapture;
-
-impl<A> Capture<A> for NoCapture
-where
-    A: Allocator,
-{
-    type Error = ErrorMarker;
-
-    #[inline]
-    fn clear(&self) {}
-
-    #[inline]
-    fn message<T>(&self, alloc: A, message: T) -> Self::Error
-    where
-        T: fmt::Display,
-    {
-        _ = alloc;
-        _ = message;
-        ErrorMarker
-    }
-
-    #[inline]
-    fn custom<T>(&self, alloc: A, error: T) -> Self::Error
-    where
-        T: 'static + Send + Sync + Error,
-    {
-        _ = alloc;
-        _ = error;
-        ErrorMarker
-    }
-}
-
-/// Capture an error of the specified type.
-pub struct SameError<E, A>
-where
-    A: Allocator,
-{
-    _marker: PhantomData<(E, A)>,
-}
-
-impl<E, A> Capture<A> for SameError<E, A>
-where
-    E: ContextError<A>,
-    A: Allocator,
-{
-    type Error = E;
-
-    #[inline]
-    fn clear(&self) {}
-
-    #[inline]
-    fn message<T>(&self, alloc: A, message: T) -> Self::Error
-    where
-        T: fmt::Display,
-    {
-        E::message(alloc, message)
-    }
-
-    #[inline]
-    fn custom<T>(&self, alloc: A, error: T) -> Self::Error
-    where
-        T: 'static + Send + Sync + Error,
-    {
-        E::custom(alloc, error)
-    }
-}
-
-/// Capture an error of the specified type.
-pub struct CaptureError<E, A> {
-    error: UnsafeCell<Option<E>>,
-    _marker: PhantomData<A>,
-}
-
-impl<E, A> Capture<A> for CaptureError<E, A>
-where
-    E: ContextError<A>,
-    A: Allocator,
-{
-    type Error = ErrorMarker;
-
-    #[inline]
-    fn clear(&self) {
-        // SAFETY: We're restricting access to the context, so that this is
-        // safe.
-        unsafe {
-            (*self.error.get()) = None;
-        }
-    }
-
-    #[inline]
-    fn message<T>(&self, alloc: A, message: T) -> Self::Error
-    where
-        T: fmt::Display,
-    {
-        // SAFETY: We're restricting access to the context, so that this is
-        // safe.
-        unsafe {
-            (*self.error.get()) = Some(E::message(alloc, message));
-        }
-
-        ErrorMarker
-    }
-
-    #[inline]
-    fn custom<T>(&self, alloc: A, error: T) -> Self::Error
-    where
-        T: 'static + Send + Sync + Error,
-    {
-        // SAFETY: We're restricting access to the context, so that this is
-        // safe.
-        unsafe {
-            (*self.error.get()) = Some(E::custom(alloc, error));
-        }
-
-        ErrorMarker
-    }
-}
+use super::{
+    Capture, CaptureError, ContextError, Errors, NoCapture, NoTrace, Report, SameError, Trace,
+    TraceConfig, TraceImpl,
+};
 
 /// The default context which uses an allocator to track the location of errors.
 ///
@@ -318,10 +177,7 @@ where
         DefaultContext {
             alloc: self.alloc,
             trace: self.trace,
-            capture: CaptureError {
-                error: UnsafeCell::new(None),
-                _marker: PhantomData,
-            },
+            capture: CaptureError::new(),
         }
     }
 
@@ -368,9 +224,7 @@ where
         DefaultContext {
             alloc: self.alloc,
             trace: self.trace,
-            capture: SameError {
-                _marker: PhantomData,
-            },
+            capture: SameError::new(),
         }
     }
 }
@@ -384,27 +238,13 @@ where
     /// Unwrap the error marker or panic if there is no error.
     #[inline]
     pub fn unwrap(&self) -> E {
-        // SAFETY: We're restricting access to the context, so that this is
-        // safe.
-        unsafe {
-            match (*self.capture.error.get()).take() {
-                Some(error) => error,
-                None => panic!("no error captured"),
-            }
-        }
+        self.capture.unwrap()
     }
 
     /// Coerce a captured error into a result.
     #[inline]
     pub fn result(&self) -> Result<(), E> {
-        // SAFETY: We're restricting access to the context, so that this is
-        // safe.
-        unsafe {
-            match (*self.capture.error.get()).take() {
-                Some(error) => Err(error),
-                None => Ok(()),
-            }
-        }
+        self.capture.result()
     }
 }
 
@@ -415,7 +255,7 @@ where
     C: Capture<A>,
 {
     type Error = C::Error;
-    type Mark = <<B as TraceConfig>::Impl<A> as TraceImpl>::Mark;
+    type Mark = <<B as TraceConfig>::Impl<A> as TraceImpl<A>>::Mark;
     type Allocator = A;
 
     #[inline]
diff --git a/crates/musli/src/context/mod.rs b/crates/musli/src/context/mod.rs
index e3c37afcd..a8b0b2e49 100644
--- a/crates/musli/src/context/mod.rs
+++ b/crates/musli/src/context/mod.rs
@@ -9,13 +9,17 @@ mod trace;
 #[doc(inline)]
 pub use self::trace::{Error, Errors, NoTrace, Report, Trace, TraceConfig, TraceImpl};
 
+mod capture;
+#[doc(inline)]
+pub use self::capture::{Capture, CaptureError, NoCapture, SameError};
+
 mod error_marker;
 #[doc(inline)]
 pub use self::error_marker::ErrorMarker;
 
 mod default_context;
 #[doc(inline)]
-pub use self::default_context::{DefaultContext, NoCapture};
+pub use self::default_context::DefaultContext;
 
 mod context_error;
 #[doc(inline)]
diff --git a/crates/musli/src/context/trace.rs b/crates/musli/src/context/trace.rs
index a7bbab067..0364155bf 100644
--- a/crates/musli/src/context/trace.rs
+++ b/crates/musli/src/context/trace.rs
@@ -22,9 +22,12 @@ mod sealed {
 }
 
 /// Trait for marker types indicating the tracing configuration.
-pub trait TraceConfig: self::sealed::Sealed {
+pub trait TraceConfig
+where
+    Self: self::sealed::Sealed,
+{
     #[doc(hidden)]
-    type Impl<A>: TraceImpl<Allocator = A>
+    type Impl<A>: TraceImpl<A>
     where
         A: Allocator;
 
@@ -34,15 +37,15 @@ pub trait TraceConfig: self::sealed::Sealed {
         A: Allocator;
 }
 
-/// The trait governing tracing in a default context.
-pub trait TraceImpl
+/// The trait governing how tracing works in a [`DefaultContext`].
+///
+/// [`DefaultContext`]: super::DefaultContext
+pub trait TraceImpl<A>
 where
-    Self: Sized + self::sealed::Sealed,
+    Self: self::sealed::Sealed,
 {
     #[doc(hidden)]
     type Mark;
-    #[doc(hidden)]
-    type Allocator: Allocator;
 
     #[doc(hidden)]
     fn clear(&self);
@@ -54,22 +57,22 @@ where
     fn mark(&self) -> Self::Mark;
 
     #[doc(hidden)]
-    fn custom<T>(&self, alloc: Self::Allocator, message: &T)
+    fn custom<T>(&self, alloc: A, message: &T)
     where
         T: 'static + Send + Sync + fmt::Display + fmt::Debug;
 
     #[doc(hidden)]
-    fn message<T>(&self, alloc: Self::Allocator, message: &T)
+    fn message<T>(&self, alloc: A, message: &T)
     where
         T: fmt::Display;
 
     #[doc(hidden)]
-    fn marked_message<T>(&self, alloc: Self::Allocator, mark: &Self::Mark, message: &T)
+    fn marked_message<T>(&self, alloc: A, mark: &Self::Mark, message: &T)
     where
         T: fmt::Display;
 
     #[doc(hidden)]
-    fn marked_custom<T>(&self, alloc: Self::Allocator, mark: &Self::Mark, message: &T)
+    fn marked_custom<T>(&self, alloc: A, mark: &Self::Mark, message: &T)
     where
         T: 'static + Send + Sync + fmt::Display + fmt::Debug;
 
@@ -113,7 +116,7 @@ where
     fn leave_sequence_index(&self);
 
     #[doc(hidden)]
-    fn enter_map_key<T>(&self, alloc: Self::Allocator, field: &T)
+    fn enter_map_key<T>(&self, alloc: A, field: &T)
     where
         T: fmt::Display;
 
@@ -122,6 +125,10 @@ where
 }
 
 /// Marker type indicating that tracing is enabled.
+///
+/// See [`DefaultContext::with_trace`] for more information.
+///
+/// [`DefaultContext::with_trace`]: super::DefaultContext::with_trace
 #[non_exhaustive]
 pub struct Trace;
 
@@ -255,12 +262,11 @@ where
     }
 }
 
-impl<A> TraceImpl for WithTraceImpl<A>
+impl<A> TraceImpl<A> for WithTraceImpl<A>
 where
     A: Allocator,
 {
     type Mark = usize;
-    type Allocator = A;
 
     #[inline]
     fn clear(&self) {
@@ -285,7 +291,7 @@ where
     }
 
     #[inline]
-    fn custom<T>(&self, alloc: Self::Allocator, message: &T)
+    fn custom<T>(&self, alloc: A, message: &T)
     where
         T: 'static + Send + Sync + fmt::Display + fmt::Debug,
     {
@@ -295,7 +301,7 @@ where
     }
 
     #[inline]
-    fn message<T>(&self, alloc: Self::Allocator, message: &T)
+    fn message<T>(&self, alloc: A, message: &T)
     where
         T: fmt::Display,
     {
@@ -305,7 +311,7 @@ where
     }
 
     #[inline]
-    fn marked_message<T>(&self, alloc: Self::Allocator, mark: &Self::Mark, message: &T)
+    fn marked_message<T>(&self, alloc: A, mark: &Self::Mark, message: &T)
     where
         T: fmt::Display,
     {
@@ -315,7 +321,7 @@ where
     }
 
     #[inline]
-    fn marked_custom<T>(&self, alloc: Self::Allocator, mark: &Self::Mark, message: &T)
+    fn marked_custom<T>(&self, alloc: A, mark: &Self::Mark, message: &T)
     where
         T: 'static + Send + Sync + fmt::Display + fmt::Debug,
     {
@@ -394,7 +400,7 @@ where
     }
 
     #[inline]
-    fn enter_map_key<T>(&self, alloc: Self::Allocator, field: &T)
+    fn enter_map_key<T>(&self, alloc: A, field: &T)
     where
         T: fmt::Display,
     {
@@ -434,6 +440,11 @@ where
 }
 
 /// Trace configuration indicating that tracing is fully disabled.
+///
+/// This is the default behavior you get when calling [`new`] or [`new_in`].
+///
+/// [`new`]: super::new
+/// [`new_in`]: super::new_in
 #[non_exhaustive]
 pub struct NoTrace;
 
@@ -452,12 +463,11 @@ impl TraceConfig for NoTrace {
     }
 }
 
-impl<A> TraceImpl for NoTraceImpl<A>
+impl<A> TraceImpl<A> for NoTraceImpl<A>
 where
     A: Allocator,
 {
     type Mark = usize;
-    type Allocator = A;
 
     #[inline]
     fn clear(&self) {}
@@ -473,7 +483,7 @@ where
     }
 
     #[inline]
-    fn custom<T>(&self, alloc: Self::Allocator, message: &T)
+    fn custom<T>(&self, alloc: A, message: &T)
     where
         T: 'static + Send + Sync + fmt::Display + fmt::Debug,
     {
@@ -482,7 +492,7 @@ where
     }
 
     #[inline]
-    fn message<T>(&self, alloc: Self::Allocator, message: &T)
+    fn message<T>(&self, alloc: A, message: &T)
     where
         T: fmt::Display,
     {
@@ -491,7 +501,7 @@ where
     }
 
     #[inline]
-    fn marked_message<T>(&self, alloc: Self::Allocator, mark: &Self::Mark, message: &T)
+    fn marked_message<T>(&self, alloc: A, mark: &Self::Mark, message: &T)
     where
         T: fmt::Display,
     {
@@ -501,7 +511,7 @@ where
     }
 
     #[inline]
-    fn marked_custom<T>(&self, alloc: Self::Allocator, mark: &Self::Mark, message: &T)
+    fn marked_custom<T>(&self, alloc: A, mark: &Self::Mark, message: &T)
     where
         T: 'static + Send + Sync + fmt::Display + fmt::Debug,
     {
@@ -568,7 +578,7 @@ where
     fn leave_sequence_index(&self) {}
 
     #[inline]
-    fn enter_map_key<T>(&self, alloc: Self::Allocator, field: &T)
+    fn enter_map_key<T>(&self, alloc: A, field: &T)
     where
         T: fmt::Display,
     {
diff --git a/crates/musli/src/json/parser/into_parser.rs b/crates/musli/src/json/parser/into_parser.rs
index f365ea246..92c079f49 100644
--- a/crates/musli/src/json/parser/into_parser.rs
+++ b/crates/musli/src/json/parser/into_parser.rs
@@ -11,7 +11,10 @@ mod sealed {
 }
 
 /// Trait for types which can be converted into a [`Parser`].
-pub trait IntoParser<'de>: self::sealed::Sealed {
+pub trait IntoParser<'de>
+where
+    Self: self::sealed::Sealed,
+{
     /// The parser type being converted into.
     type Parser: Parser<'de>;
 
diff --git a/crates/musli/src/reader.rs b/crates/musli/src/reader.rs
index 526d63cd0..ddb92bb8f 100644
--- a/crates/musli/src/reader.rs
+++ b/crates/musli/src/reader.rs
@@ -22,7 +22,10 @@ mod sealed {
 }
 
 /// Coerce a type into a [`Reader`].
-pub trait IntoReader<'de>: self::sealed::Sealed {
+pub trait IntoReader<'de>
+where
+    Self: self::sealed::Sealed,
+{
     /// The reader type.
     type Reader: Reader<'de>;
 
@@ -34,7 +37,10 @@ pub trait IntoReader<'de>: self::sealed::Sealed {
 ///
 /// This requires the reader to be able to hand out contiguous references to the
 /// byte source through [`Reader::read_bytes`].
-pub trait Reader<'de>: self::sealed::Sealed {
+pub trait Reader<'de>
+where
+    Self: self::sealed::Sealed,
+{
     /// Type borrowed from self.
     ///
     /// Why oh why would we want to do this over having a simple `&'this mut T`?
diff --git a/crates/musli/src/writer.rs b/crates/musli/src/writer.rs
index fb387b1b9..f9ac023f5 100644
--- a/crates/musli/src/writer.rs
+++ b/crates/musli/src/writer.rs
@@ -23,7 +23,10 @@ mod sealed {
 }
 
 /// Coerce a type into a [`Writer`].
-pub trait IntoWriter: self::sealed::Sealed {
+pub trait IntoWriter
+where
+    Self: self::sealed::Sealed,
+{
     /// The output of the writer which will be returned after writing.
     type Ok;