1 /*
2  * Copyright (C) 2021 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 use alloc::alloc::{AllocError, Allocator};
18 use alloc::boxed::Box;
19 use alloc::vec::Vec;
20 
21 /// Trait for fallible duplication of types that can be cloned.
22 ///
23 /// See the [`Clone`] trait for more details. This trait is identical except
24 /// that duplication may fail, e.g. if the allocator could not allocate more
25 /// space for the result.
26 pub trait TryClone: Sized {
27     /// Error type when the clone fails
28     type Error;
29 
30     /// Attempt to duplicate the value.
31     ///
32     /// See [`Clone::clone()`]. This method may fail with `Self::Error`, which
33     /// for a heap allocation, generally indicates that space for a duplicate
34     /// value could not be allocated from the heap.
35     ///
36     /// # Examples
37     ///
38     /// ```
39     /// let value = Box::new("Hello");
40     ///
41     /// let cloned_value = value.try_clone().expect("Box could not be cloned");
42     /// ```
try_clone(&self) -> Result<Self, Self::Error>43     fn try_clone(&self) -> Result<Self, Self::Error>;
44 }
45 
46 impl<T: Clone, A: Allocator + Clone> TryClone for Box<T, A> {
47     type Error = AllocError;
48 
49     #[inline]
try_clone(&self) -> Result<Self, Self::Error>50     fn try_clone(&self) -> Result<Self, Self::Error> {
51         let mut boxed = Self::try_new_uninit_in(Box::allocator(self).clone())?;
52         unsafe {
53             boxed.as_mut_ptr().write((**self).clone());
54             Ok(boxed.assume_init())
55         }
56     }
57 }
58 
59 #[inline]
try_to_vec<T: TryConvertVec, A: Allocator>(s: &[T], alloc: A) -> Result<Vec<T, A>, AllocError>60 fn try_to_vec<T: TryConvertVec, A: Allocator>(s: &[T], alloc: A) -> Result<Vec<T, A>, AllocError> {
61     T::try_to_vec(s, alloc)
62 }
63 
64 trait TryConvertVec {
try_to_vec<A: Allocator>(s: &[Self], alloc: A) -> Result<Vec<Self, A>, AllocError> where Self: Sized65     fn try_to_vec<A: Allocator>(s: &[Self], alloc: A) -> Result<Vec<Self, A>, AllocError>
66     where
67         Self: Sized;
68 }
69 
70 impl<T: Clone> TryConvertVec for T {
71     #[inline]
try_to_vec<A: Allocator>(s: &[Self], alloc: A) -> Result<Vec<Self, A>, AllocError>72     default fn try_to_vec<A: Allocator>(s: &[Self], alloc: A) -> Result<Vec<Self, A>, AllocError> {
73         struct DropGuard<'a, T, A: Allocator> {
74             vec: &'a mut Vec<T, A>,
75             num_init: usize,
76         }
77         impl<'a, T, A: Allocator> Drop for DropGuard<'a, T, A> {
78             #[inline]
79             fn drop(&mut self) {
80                 // SAFETY:
81                 // items were marked initialized in the loop below
82                 unsafe {
83                     self.vec.set_len(self.num_init);
84                 }
85             }
86         }
87         let mut vec = Vec::new_in(alloc);
88         // TODO: replace with try_with_capacity_in when
89         // https://github.com/rust-lang/rust/pull/86938 lands
90         vec.try_reserve_exact(s.len()).or(Err(AllocError))?;
91         let mut guard = DropGuard { vec: &mut vec, num_init: 0 };
92         let slots = guard.vec.spare_capacity_mut();
93         // .take(slots.len()) is necessary for LLVM to remove bounds checks
94         // and has better codegen than zip.
95         for (i, b) in s.iter().enumerate().take(slots.len()) {
96             guard.num_init = i;
97             slots[i].write(b.clone());
98         }
99         core::mem::forget(guard);
100         // SAFETY:
101         // the vec was allocated and initialized above to at least this length.
102         unsafe {
103             vec.set_len(s.len());
104         }
105         Ok(vec)
106     }
107 }
108 
109 #[cfg(not(no_global_oom_handling))]
110 impl<T: Copy> TryConvertVec for T {
111     #[inline]
try_to_vec<A: Allocator>(s: &[Self], alloc: A) -> Result<Vec<Self, A>, AllocError>112     fn try_to_vec<A: Allocator>(s: &[Self], alloc: A) -> Result<Vec<Self, A>, AllocError> {
113         let mut v = Vec::new_in(alloc);
114         // TODO: replace with try_with_capacity_in when
115         // https://github.com/rust-lang/rust/pull/86938 lands
116         v.try_reserve_exact(s.len()).or(Err(AllocError))?;
117         // SAFETY:
118         // allocated above with the capacity of `s`, and initialize to `s.len()` in
119         // ptr::copy_to_non_overlapping below.
120         unsafe {
121             s.as_ptr().copy_to_nonoverlapping(v.as_mut_ptr(), s.len());
122             v.set_len(s.len());
123         }
124         Ok(v)
125     }
126 }
127 
128 impl<T: Clone, A: Allocator + Clone> TryClone for Box<[T], A> {
129     type Error = AllocError;
130 
131     #[inline]
try_clone(&self) -> Result<Self, Self::Error>132     fn try_clone(&self) -> Result<Self, Self::Error> {
133         let alloc = Box::allocator(self).clone();
134         try_to_vec(&*self, alloc).map(Vec::into_boxed_slice)
135     }
136 }
137