Merge pull request #49 from audunhalland/mut_output

fix: Basic support for mutable output

CHANGELOG.md

@@ -5,6 +5,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## Unreleased
+### Added
+- Basic support for mutable outputs ([#49](https://github.com/audunhalland/unimock/pull/49))
 
 ## [0.6.2] - 2024-03-27
 ### Fixed

src/lib.rs

@@ -869,6 +869,13 @@ impl Unimock {
         self.value_chain.add(value)
     }
 
+    /// Convert the given value into a mutable reference.
+    ///
+    /// This can be useful when returning mutable references from `answers` functions.
+    pub fn make_mut<T: Send + Sync + 'static>(&mut self, value: T) -> &mut T {
+        self.value_chain.add_mut(value)
+    }
+
     #[track_caller]
     fn from_assembler(
         assembler_result: Result<MockAssembler, alloc::String>,

src/output.rs

@@ -4,6 +4,7 @@ use crate::alloc::Box;
 
 pub(crate) mod deep;
 pub(crate) mod lending;
+pub(crate) mod mut_lending;
 pub(crate) mod owning;
 pub(crate) mod shallow;
 pub(crate) mod static_ref;
@@ -59,8 +60,12 @@ pub trait ReturnDefault<K: Kind> {
     fn return_default() -> K::Return;
 }
 
+/// A "marker" for mutable types
+pub struct Mutable<T>(pub(crate) T);
+
 pub use deep::Deep;
 pub use lending::Lending;
+pub use mut_lending::MutLending;
 pub use owning::Owning;
 pub use shallow::Shallow;
 pub use static_ref::StaticRef;

src/output/mut_lending.rs

@@ -0,0 +1,23 @@
+use core::marker::PhantomData;
+
+use super::*;
+
+/// A type category for values lent out by Unimock.
+#[doc(hidden)]
+pub struct MutLending<T: ?Sized + 'static>(core::marker::PhantomData<fn() -> T>);
+
+impl<T: ?Sized + 'static> Kind for MutLending<T> {
+    type Return = MutLent<T>;
+}
+
+pub struct MutLent<T: ?Sized>(PhantomData<T>);
+
+impl<T: ?Sized + 'static> GetOutput for MutLent<T> {
+    type Output<'u> = &'u mut T
+        where
+            Self: 'u;
+
+    fn output(&self) -> Option<Self::Output<'_>> {
+        None
+    }
+}

src/output/shallow/option.rs

@@ -2,19 +2,25 @@ use core::borrow::Borrow;
 
 use crate::output::*;
 
-type Mix<T> = Shallow<Option<&'static T>>;
+type Ref<T> = Shallow<Option<&'static T>>;
+type Mut<T> = Shallow<Option<&'static mut T>>;
 
-type Response<T> = Option<Box<dyn Borrow<T> + Send + Sync>>;
+type RefResponse<T> = Option<Box<dyn Borrow<T> + Send + Sync>>;
+type MutResponse<T> = Option<Mutable<Box<dyn Borrow<T> + Send + Sync>>>;
 
-impl<T: ?Sized> Kind for Mix<T> {
-    type Return = Response<T>;
+impl<T: ?Sized> Kind for Ref<T> {
+    type Return = RefResponse<T>;
 }
 
-impl<T: ?Sized> Return for Mix<T> {
+impl<T: ?Sized> Kind for Mut<T> {
+    type Return = MutResponse<T>;
+}
+
+impl<T: ?Sized> Return for Ref<T> {
     type Type = Option<Box<dyn Borrow<T> + Send + Sync>>;
 }
 
-impl<T: ?Sized + 'static> GetOutput for Response<T> {
+impl<T: ?Sized + 'static> GetOutput for RefResponse<T> {
     type Output<'u> = Option<&'u T>
         where
             Self: 'u;
@@ -24,13 +30,23 @@ impl<T: ?Sized + 'static> GetOutput for Response<T> {
     }
 }
 
+impl<T: ?Sized + 'static> GetOutput for MutResponse<T> {
+    type Output<'u> = Option<&'u mut T>
+        where
+            Self: 'u;
+
+    fn output(&self) -> Option<Self::Output<'_>> {
+        None
+    }
+}
+
 macro_rules! into {
     ($trait:ident, $method:ident) => {
-        impl<T0, T: ?Sized + 'static> $trait<Mix<T>> for Option<T0>
+        impl<T0, T: ?Sized + 'static> $trait<Ref<T>> for Option<T0>
         where
             T0: Borrow<T> + Send + Sync + 'static,
         {
-            fn $method(self) -> OutputResult<Response<T>> {
+            fn $method(self) -> OutputResult<RefResponse<T>> {
                 Ok(match self {
                     Some(val) => Some(Box::new(val)),
                     None => None,

src/output/shallow/result.rs

@@ -4,22 +4,28 @@ use core::borrow::Borrow;
 
 use crate::output::{owning::Owned, *};
 
-type Mix<T, E> = Shallow<Result<&'static T, E>>;
+type Ref<T, E> = Shallow<Result<&'static T, E>>;
+type Mut<T, E> = Shallow<Result<&'static mut T, E>>;
 
-type Response<T, E> = Result<Box<dyn Borrow<T> + Send + Sync>, Owned<E>>;
+type RefResponse<T, E> = Result<Box<dyn Borrow<T> + Send + Sync>, Owned<E>>;
+type MutResponse<T, E> = Result<Mutable<Box<dyn Borrow<T> + Send + Sync>>, Owned<E>>;
 
-impl<T: ?Sized, E: 'static> Kind for Mix<T, E> {
-    type Return = Response<T, E>;
+impl<T: ?Sized, E: 'static> Kind for Ref<T, E> {
+    type Return = RefResponse<T, E>;
 }
 
-impl<T: ?Sized, E: 'static> Return for Mix<T, E>
+impl<T: ?Sized, E: 'static> Kind for Mut<T, E> {
+    type Return = MutResponse<T, E>;
+}
+
+impl<T: ?Sized, E: 'static> Return for Ref<T, E>
 where
     Owning<E>: Return,
 {
     type Type = Result<Box<dyn Borrow<T> + Send + Sync>, <Owning<E> as Return>::Type>;
 }
 
-impl<T: ?Sized + 'static, E: 'static> GetOutput for Response<T, E> {
+impl<T: ?Sized + 'static, E: 'static> GetOutput for RefResponse<T, E> {
     type Output<'u> = Result<&'u T, E>
         where
             Self: 'u;
@@ -32,14 +38,27 @@ impl<T: ?Sized + 'static, E: 'static> GetOutput for Response<T, E> {
     }
 }
 
+impl<T: ?Sized + 'static, E: 'static> GetOutput for MutResponse<T, E> {
+    type Output<'u> = Result<&'u mut T, E>
+        where
+            Self: 'u;
+
+    fn output(&self) -> Option<Self::Output<'_>> {
+        match self {
+            Ok(_) => None,
+            Err(err) => Some(Err(err.output()?)),
+        }
+    }
+}
+
 macro_rules! into {
     ($trait:ident, $method:ident) => {
-        impl<T0, T: ?Sized + 'static, E0, E: 'static> $trait<Mix<T, E>> for Result<T0, E0>
+        impl<T0, T: ?Sized + 'static, E0, E: 'static> $trait<Ref<T, E>> for Result<T0, E0>
         where
             T0: Borrow<T> + Send + Sync + 'static,
             E0: $trait<Owning<E>>,
         {
-            fn $method(self) -> OutputResult<Response<T, E>> {
+            fn $method(self) -> OutputResult<RefResponse<T, E>> {
                 Ok(match self {
                     Ok(val) => Ok(Box::new(val)),
                     Err(err) => Err(err.$method()?),

src/value_chain.rs

@@ -8,7 +8,7 @@ use crate::alloc::Box;
 /// This allows retaining shared references to inserted values.
 #[derive(Default)]
 pub struct ValueChain {
-    root_node: OnceCell<Node>,
+    root: OnceCell<Node>,
 }
 
 impl ValueChain {
@@ -18,8 +18,21 @@ impl ValueChain {
         node.value.as_ref().downcast_ref::<T>().unwrap()
     }
 
+    pub fn add_mut<T: Any + Send + Sync>(&mut self, value: T) -> &mut T {
+        // note: There is no need for keeping the old chain.
+        // All those references are out of scope when add_mut is called.
+        self.root = Node::new(value).into();
+
+        self.root
+            .get_mut()
+            .unwrap()
+            .value
+            .downcast_mut::<T>()
+            .unwrap()
+    }
+
     fn push_node(&self, mut new_node: Node) -> &Node {
-        let mut cell = &self.root_node;
+        let mut cell = &self.root;
         loop {
             match cell.try_insert(new_node) {
                 Ok(new_node) => {
@@ -36,7 +49,7 @@ impl ValueChain {
 
 impl Drop for ValueChain {
     fn drop(&mut self) {
-        if let Some(node) = self.root_node.take() {
+        if let Some(node) = self.root.take() {
             drop(node.value);
             let mut cell = node.next;
 
@@ -73,3 +86,10 @@ fn it_works() {
     assert_eq!(&"", second);
     assert_eq!(&42.0, third);
 }
+
+#[test]
+fn it_works_mut() {
+    let mut value_chain = ValueChain::default();
+    let first = value_chain.add_mut(1);
+    *first += 1;
+}

tests/it/basic.rs

@@ -1193,3 +1193,39 @@ mod apply_fn {
         assert_eq!(1, b);
     }
 }
+
+mod mut_output {
+    use unimock::*;
+
+    #[unimock(api = TraitMock)]
+    trait Trait {
+        fn ref_mut(&mut self) -> &mut i32;
+        fn option_mut(&mut self) -> Option<&mut i32>;
+        fn result_mut(&mut self) -> Result<&mut i32, ()>;
+    }
+
+    #[test]
+    fn test_ref_mut() {
+        let mut u = Unimock::new(
+            TraitMock::ref_mut
+                .next_call(matching!())
+                .answers(&|u| u.make_mut(42)),
+        );
+
+        let num = u.ref_mut();
+        *num += 1;
+    }
+
+    #[test]
+    fn test_option_mut() {
+        let mut u = Unimock::new(
+            TraitMock::option_mut
+                .next_call(matching!())
+                .answers(&|u| Some(u.make_mut(42))),
+        );
+
+        if let Some(num) = u.option_mut() {
+            *num += 1;
+        }
+    }
+}

unimock_macros/src/unimock/output.rs

@@ -66,6 +66,7 @@ pub enum OutputWrapping {
 pub enum OutputKind {
     Owning,
     SelfReference,
+    MutSelfReference,
     ParamReference,
     StaticReference,
     Shallow,
@@ -81,6 +82,7 @@ impl OutputKind {
         match self {
             Self::Owning => "Owning",
             Self::SelfReference => "Lending",
+            Self::MutSelfReference => "MutLending",
             Self::ParamReference => "StaticRef",
             Self::StaticReference => "StaticRef",
             Self::Shallow => "Shallow",
@@ -208,7 +210,13 @@ pub fn determine_owned_or_deep_output_structure(
 
 fn make_generic_kind(ty: syn::Type, attr: &Attr) -> (OutputKind, syn::Type) {
     match ty {
-        syn::Type::Reference(reference) => (OutputKind::SelfReference, *reference.elem),
+        syn::Type::Reference(reference) => {
+            if reference.mutability.is_some() {
+                (OutputKind::MutSelfReference, *reference.elem)
+            } else {
+                (OutputKind::SelfReference, *reference.elem)
+            }
+        }
         syn::Type::Path(mut path) => {
             let mut kind = OutputKind::Owning;
 
@@ -295,12 +303,20 @@ fn determine_reference_ownership(
             "static" => OutputKind::StaticReference,
             _ => match find_param_lifetime(sig, &lifetime.ident) {
                 Some(index) => match index {
-                    0 => OutputKind::SelfReference,
+                    0 => {
+                        if type_reference.mutability.is_some() {
+                            OutputKind::MutSelfReference
+                        } else {
+                            OutputKind::SelfReference
+                        }
+                    }
                     _ => OutputKind::ParamReference,
                 },
                 None => OutputKind::SelfReference,
             },
         }
+    } else if type_reference.mutability.is_some() {
+        OutputKind::MutSelfReference
     } else {
         OutputKind::SelfReference
     }