Skip to main content

diskann_utils/
reborrow.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! A collection of tools for working with generalized references and scrounging through types.
7
8use sealed::{BoundTo, Sealed};
9
10/// An hybrid combination of reference covariance and borrowing for generalized reference
11/// types.
12/// ```
13/// use diskann_utils::Reborrow;
14/// let mut base = vec![1usize, 3usize];
15///
16/// let borrowed = base.reborrow();
17/// assert_eq!(borrowed[0], 1);
18/// assert_eq!(borrowed[1], 3);
19/// ```
20///
21/// # Notes
22///
23/// The extra hidden generic parameter is an implementation of
24/// <https://sabrinajewson.orgblog/the-better-alternative-to-lifetime-gats/> and allows
25/// [HRTB](https://doc.rust-lang.org/nomicon/hrtb.html) to work properly with shortened
26/// borrows.
27pub trait Reborrow<'this, Lifetime: Sealed = BoundTo<&'this Self>> {
28    type Target;
29
30    /// Borrow `self` into a generalized reference type and reborrow
31    fn reborrow(&'this self) -> Self::Target;
32}
33
34/// An hybrid combination of reference covariance and borrowing for generalized reference
35/// types.
36/// ```ignore
37/// use diskann_utils::ReborrowMut;
38/// let mut base = vec[0, 0];
39/// {
40///     let mut borrowed = base.reborrow_mut();
41///     borrowed[0] = 1;
42///     borrowed[1] = 3;
43/// }
44///
45/// assert_eq!(&*base, &*[1, 3]);
46/// ```
47pub trait ReborrowMut<'this, Lifetime: Sealed = BoundTo<&'this Self>> {
48    type Target;
49
50    /// Mutably borrow `self` into a generalized reference type and reborrow.
51    fn reborrow_mut(&'this mut self) -> Self::Target;
52}
53
54// Custom Impls
55impl<'short, T> Reborrow<'short> for &T
56where
57    T: ?Sized,
58{
59    type Target = &'short T;
60    fn reborrow(&'short self) -> Self::Target {
61        self
62    }
63}
64
65impl<'short, T> ReborrowMut<'short> for &mut T
66where
67    T: ?Sized,
68{
69    type Target = &'short mut T;
70    fn reborrow_mut(&'short mut self) -> Self::Target {
71        self
72    }
73}
74
75impl<'this, T> Reborrow<'this> for Vec<T> {
76    type Target = &'this [T];
77    fn reborrow(&'this self) -> Self::Target {
78        self
79    }
80}
81
82impl<'this, T> ReborrowMut<'this> for Vec<T> {
83    type Target = &'this mut [T];
84    fn reborrow_mut(&'this mut self) -> Self::Target {
85        self
86    }
87}
88
89impl<'this, T> Reborrow<'this> for Box<T>
90where
91    T: ?Sized,
92{
93    type Target = &'this T;
94    fn reborrow(&'this self) -> Self::Target {
95        self
96    }
97}
98
99impl<'this, T> ReborrowMut<'this> for Box<T>
100where
101    T: ?Sized,
102{
103    type Target = &'this mut T;
104    fn reborrow_mut(&'this mut self) -> Self::Target {
105        self
106    }
107}
108
109impl<'short, T> Reborrow<'short> for std::borrow::Cow<'_, T>
110where
111    T: std::borrow::ToOwned + ?Sized,
112{
113    type Target = &'short T;
114    fn reborrow(&'short self) -> Self::Target {
115        self
116    }
117}
118
119impl<'short> Reborrow<'short> for String {
120    type Target = &'short str;
121    fn reborrow(&'short self) -> Self::Target {
122        self
123    }
124}
125
126impl<'short> ReborrowMut<'short> for String {
127    type Target = &'short mut str;
128    fn reborrow_mut(&'short mut self) -> Self::Target {
129        self
130    }
131}
132
133////////////
134// Helper //
135////////////
136
137/// A container for types `T` providing an implementation `Deref<Target = T>` as well as
138/// reborrowing functionality mapped to `T as Deref` and `T as DerefMut`.
139#[derive(Debug, Clone, Copy, PartialEq)]
140#[repr(transparent)]
141pub struct Place<T>(pub T);
142
143impl<T> std::ops::Deref for Place<T> {
144    type Target = T;
145    fn deref(&self) -> &T {
146        &self.0
147    }
148}
149
150impl<T> std::ops::DerefMut for Place<T> {
151    fn deref_mut(&mut self) -> &mut T {
152        &mut self.0
153    }
154}
155
156impl<'this, T> Reborrow<'this> for Place<T>
157where
158    T: std::ops::Deref,
159{
160    type Target = &'this T::Target;
161    fn reborrow(&'this self) -> Self::Target {
162        self
163    }
164}
165
166impl<'this, T> ReborrowMut<'this> for Place<T>
167where
168    T: std::ops::DerefMut,
169{
170    type Target = &'this mut T::Target;
171    fn reborrow_mut(&'this mut self) -> Self::Target {
172        self
173    }
174}
175
176/// A container that reborrows by cloning the contained value.
177///
178/// Note that [`ReborrowMut`] is not implemented for this type.
179#[derive(Debug, Clone, Copy, PartialEq)]
180#[repr(transparent)]
181pub struct Cloned<T>(pub T)
182where
183    T: Clone;
184
185impl<T> std::ops::Deref for Cloned<T>
186where
187    T: Clone,
188{
189    type Target = T;
190    fn deref(&self) -> &Self::Target {
191        &self.0
192    }
193}
194
195impl<T> std::ops::DerefMut for Cloned<T>
196where
197    T: Clone,
198{
199    fn deref_mut(&mut self) -> &mut Self::Target {
200        &mut self.0
201    }
202}
203
204impl<'this, T> Reborrow<'this> for Cloned<T>
205where
206    T: Clone,
207{
208    type Target = Self;
209
210    fn reborrow(&'this self) -> Self::Target {
211        self.clone()
212    }
213}
214
215/// A container that reborrows by copying the contained value.
216///
217/// Note that [`ReborrowMut`] is not implemented for this type.
218#[derive(Debug, Clone, Copy, PartialEq)]
219#[repr(transparent)]
220pub struct Copied<T>(pub T)
221where
222    T: Copy;
223
224impl<T> std::ops::Deref for Copied<T>
225where
226    T: Copy,
227{
228    type Target = T;
229    fn deref(&self) -> &Self::Target {
230        &self.0
231    }
232}
233
234impl<T> std::ops::DerefMut for Copied<T>
235where
236    T: Copy,
237{
238    fn deref_mut(&mut self) -> &mut Self::Target {
239        &mut self.0
240    }
241}
242
243impl<'this, T> Reborrow<'this> for Copied<T>
244where
245    T: Copy,
246{
247    type Target = Self;
248
249    fn reborrow(&'this self) -> Self::Target {
250        *self
251    }
252}
253
254macro_rules! trivial_reborrow {
255    ($T:ty) => {
256        impl<'a> Reborrow<'a> for $T {
257            type Target = Self;
258
259            fn reborrow(&'a self) -> Self {
260                *self
261            }
262        }
263    };
264    ($($T:ty),* $(,)?) => {
265        $(trivial_reborrow!($T);)*
266    }
267}
268
269trivial_reborrow!(half::f16, f32, f64, u8, u16, u32, u64, i8, i16, i32, i64);
270
271/// Helper traits for the [HRTB]() trick.
272mod sealed {
273    pub trait Sealed: Sized {}
274    pub struct BoundTo<T>(T);
275    impl<T> Sealed for BoundTo<T> {}
276}
277
278///////////
279// Tests //
280///////////
281
282#[cfg(test)]
283mod tests {
284    use super::*;
285
286    fn test_hrtb_reborrow<T>(_x: T)
287    where
288        T: for<'a> Reborrow<'a>,
289    {
290    }
291
292    fn test_hrtb_reborrow_mut<T>(_x: T)
293    where
294        T: for<'a> ReborrowMut<'a>,
295    {
296    }
297
298    fn test_reborrow_constrained<T>(x: T) -> String
299    where
300        T: for<'a> Reborrow<'a, Target: std::fmt::Debug>,
301    {
302        format!("{:?}", x.reborrow())
303    }
304
305    #[test]
306    fn hrbt_reborrow() {
307        let x: &[usize] = &[10];
308        test_hrtb_reborrow(x);
309    }
310
311    #[test]
312    fn hrbt_reborrow_mut() {
313        let x: &mut [usize] = &mut [10];
314        test_hrtb_reborrow_mut(x);
315    }
316
317    #[test]
318    fn reborrow_constrained() {
319        let x: &[usize] = &[10];
320        let s = test_reborrow_constrained(x);
321        assert_eq!(s, "[10]");
322    }
323
324    ////////////////////////
325    // Reborrow built-ins //
326    ////////////////////////
327
328    #[test]
329    fn test_slice() {
330        let x: &[usize] = &[1, 2, 3];
331        let ptr = x.as_ptr();
332        let len = x.len();
333
334        let y: &[usize] = x.reborrow();
335        assert_eq!(ptr, y.as_ptr());
336        assert_eq!(len, y.len());
337    }
338
339    #[test]
340    fn test_mut_slice() {
341        let mut x: &mut [usize] = &mut [0, 0, 0];
342        let ptr = x.as_ptr();
343        let len = x.len();
344
345        let y: &mut [usize] = x.reborrow_mut();
346        assert_eq!(ptr, y.as_ptr());
347        assert_eq!(len, y.len());
348        y[0] = 1;
349        y[1] = 2;
350        y[2] = 3;
351
352        assert_eq!(x, [1, 2, 3]);
353    }
354
355    #[test]
356    fn test_vec() {
357        let x: Vec<usize> = vec![1, 2, 3];
358        let ptr = x.as_ptr();
359        let len = x.len();
360
361        let y: &[usize] = x.reborrow();
362        assert_eq!(ptr, y.as_ptr());
363        assert_eq!(len, y.len());
364    }
365
366    #[test]
367    fn test_vec_mut() {
368        let mut x: Vec<usize> = vec![0, 0, 0];
369        let ptr = x.as_ptr();
370        let len = x.len();
371
372        let y: &mut [usize] = x.reborrow_mut();
373        assert_eq!(ptr, y.as_ptr());
374        assert_eq!(len, y.len());
375        y[0] = 1;
376        y[1] = 2;
377        y[2] = 3;
378
379        assert_eq!(x, [1, 2, 3]);
380    }
381
382    #[test]
383    fn test_box() {
384        let x: Box<[usize]> = Box::new([1, 2, 3]);
385        let ptr = x.as_ptr();
386        let len = x.len();
387
388        let y: &[usize] = x.reborrow();
389        assert_eq!(ptr, y.as_ptr());
390        assert_eq!(len, y.len());
391    }
392
393    #[test]
394    fn test_box_mut() {
395        let mut x: Box<[usize]> = Box::new([0, 0, 0]);
396        let ptr = x.as_ptr();
397        let len = x.len();
398
399        let y: &mut [usize] = x.reborrow_mut();
400        assert_eq!(ptr, y.as_ptr());
401        assert_eq!(len, y.len());
402        y[0] = 1;
403        y[1] = 2;
404        y[2] = 3;
405
406        assert_eq!(&*x, [1, 2, 3]);
407    }
408
409    #[test]
410    fn test_cow() {
411        let x = &[1, 2, 3];
412        let ptr = x.as_ptr();
413        let len = x.len();
414        let cow = std::borrow::Cow::<[usize]>::Borrowed(x);
415
416        let y: &[usize] = cow.reborrow();
417        assert_eq!(ptr, y.as_ptr());
418        assert_eq!(len, y.len());
419
420        let cow = cow.into_owned();
421        let ptr = cow.as_ptr();
422        let len = cow.len();
423
424        let y: &[usize] = cow.reborrow();
425        assert_eq!(ptr, y.as_ptr());
426        assert_eq!(len, y.len());
427    }
428
429    #[test]
430    fn test_string() {
431        let mut x = String::from("hello world");
432        let ptr = x.as_ptr();
433        let len = x.len();
434
435        let y: &str = x.reborrow();
436        assert_eq!(y, x);
437        assert_eq!(ptr, y.as_ptr());
438        assert_eq!(len, y.len());
439
440        let y: &mut str = x.reborrow_mut();
441        assert_eq!(ptr, y.as_ptr());
442        assert_eq!(len, y.len());
443        y.make_ascii_uppercase();
444
445        assert_eq!(x, "HELLO WORLD");
446    }
447
448    ///////////
449    // Place //
450    ///////////
451
452    #[test]
453    fn test_place() {
454        let mut x: Place<Box<[usize]>> = Place(Box::new([0, 0, 0]));
455        // DerefMut through `Box`.
456        x[0] = 1;
457        x[1] = 2;
458        x[2] = 3;
459
460        assert_eq!(&**x, [1, 2, 3]);
461        assert_eq!(x.reborrow(), [1, 2, 3]);
462
463        *x = Box::new([2, 3, 4]);
464        assert_eq!(&**x, [2, 3, 4]);
465        assert_eq!(x.reborrow(), [2, 3, 4]);
466
467        let y = x.reborrow_mut();
468        y[0] = 10;
469        y[1] = 20;
470        y[2] = 30;
471        assert_eq!(&**x, [10, 20, 30]);
472        assert_eq!(x.reborrow(), [10, 20, 30]);
473    }
474
475    /////////////
476    // Helpers //
477    /////////////
478
479    #[test]
480    fn test_cloned() {
481        let mut x = Cloned(10);
482        assert_eq!(*x, 10);
483
484        let y: Cloned<usize> = x.reborrow();
485        assert_eq!(x, y);
486
487        // Test derive copy;
488        let z = x;
489        assert_eq!(*z, 10);
490
491        // DerefMut
492        *x = 50;
493        assert_eq!(*x, 50);
494        assert_ne!(x, y);
495    }
496
497    #[test]
498    fn test_copied() {
499        let mut x = Copied(10);
500        assert_eq!(*x, 10);
501
502        let y: Copied<usize> = x.reborrow();
503        assert_eq!(x, y);
504
505        // Test derive copy;
506        let z = x;
507        assert_eq!(*z, 10);
508
509        // DerefMut
510        *x = 50;
511        assert_eq!(*x, 50);
512        assert_ne!(x, y);
513    }
514
515    //////////////////////
516    // Trivial Reborrow //
517    //////////////////////
518
519    fn reborrow_to_self<T>(x: T) -> T
520    where
521        for<'a> T: Reborrow<'a, Target = T>,
522    {
523        x.reborrow()
524    }
525
526    #[test]
527    fn trivial_reborrows() {
528        assert_eq!(
529            reborrow_to_self::<half::f16>(Default::default()),
530            Default::default()
531        );
532        assert_eq!(reborrow_to_self::<f32>(1.0f32), 1.0);
533        assert_eq!(reborrow_to_self::<f64>(1.0f64), 1.0);
534
535        assert_eq!(reborrow_to_self::<i8>(1), 1);
536        assert_eq!(reborrow_to_self::<i16>(1), 1);
537        assert_eq!(reborrow_to_self::<i32>(1), 1);
538        assert_eq!(reborrow_to_self::<i64>(1), 1);
539
540        assert_eq!(reborrow_to_self::<u8>(1), 1);
541        assert_eq!(reborrow_to_self::<u16>(1), 1);
542        assert_eq!(reborrow_to_self::<u32>(1), 1);
543        assert_eq!(reborrow_to_self::<u64>(1), 1);
544    }
545}