redo the drain · rust-lang/rust@e3a2c23 (original) (raw)

1

use crate::assert_unsafe_precondition;

2

use crate:📑:Destruct;

3

use crate::mem::ManuallyDrop;

1

use crate:📑:{Destruct, PhantomData};

2

use crate::mem::{ManuallyDrop, SizedTypeProperties, conjure_zst};

3

use crate::ptr::{NonNull, drop_in_place, from_raw_parts_mut, null_mut};

4

5

impl<'l, 'f, T, U, const N: usize, F: FnMut(T) -> U> Drain<'l, 'f, T, N, F> {

6

/// This function returns a function that lets you index the given array in const.

7

/// As implemented it can optimize better than iterators, and can be constified.

8

/// It acts like a sort of guard (owns the array) and iterator combined, which can be implemented

9

/// as it is a struct that implements const fn;

10

`` +

/// in that regard it is somewhat similar to an array::Iter implementing UncheckedIterator.

11

`` +

/// The only method you're really allowed to call is next(),

12

/// anything else is more or less UB, hence this function being unsafe.

13

/// Moved elements will not be dropped.

14

/// This will also not actually store the array.

15

///

16

`` +

/// SAFETY: must only be called N times. Thou shalt not drop the array either.

17

`` +

// FIXME(const-hack): this is a hack for let guard = Guard(array); |i| f(guard[i]).

18

#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]

19

pub(super) const unsafe fn new(array: &'l mut ManuallyDrop<[T; N]>, f: &'f mut F) -> Self {

20

// dont drop the array, transfers "ownership" to Self

21

let ptr: NonNull = NonNull::from_mut(array).cast();

22

// SAFETY:

23

`` +

// Adding slice.len() to the starting pointer gives a pointer

24

`` +

// at the end of slice. end will never be dereferenced, only checked

25

`` +

// for direct pointer equality with ptr to check if the drainer is done.

26

unsafe {

27

let end = if T::IS_ZST { null_mut() } else { ptr.as_ptr().add(N) };

28

Self { ptr, end, f, l: PhantomData }

29

}

30

}

31

}

32

+

33

`` +

/// See [Drain::new]; this is our fake iterator.

5

34

`#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]

6

35

`#[unstable(feature = "array_try_map", issue = "79711")]

7

pub(super) struct Drain<'a, T, U, const N: usize, F: FnMut(T) -> U> {

8

array: ManuallyDrop<[T; N]>,

9

moved: usize,

10

f: &'a mut F,

36

pub(super) struct Drain<'l, 'f, T, const N: usize, F> {

37

// FIXME(const-hack): This is essentially a slice::IterMut<'static>, replace when possible.

38

/// The pointer to the next element to return, or the past-the-end location

39

/// if the drainer is empty.

40

///

41

/// This address will be used for all ZST elements, never changed.

42

/// As we "own" this array, we dont need to store any lifetime.

43

ptr: NonNull,

44

/// For non-ZSTs, the non-null pointer to the past-the-end element.

45

/// For ZSTs, this is null.

46

end: *mut T,

47

+

48

f: &'f mut F,

49

l: PhantomData<&'l mut [T; N]>,

11

50

51

+

12

52

`#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]

13

53

`#[unstable(feature = "array_try_map", issue = "79711")]

14

impl<T, U, const N: usize, F> const FnOnce<(usize,)> for &mut Drain<'_, T, U, N, F>

54

impl<T, U, const N: usize, F> const FnOnce<(usize,)> for &mut Drain<'_, '_, T, N, F>

15

55

`where

16

56

`F: [const] FnMut(T) -> U,

17

57

18

58

`type Output = U;

19

59

60

/// This implementation is useless.

20

61

`extern "rust-call" fn call_once(mut self, args: (usize,)) -> Self::Output {

21

62

`self.call_mut(args)

22

63

23

64

24

65

`#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]

25

66

`#[unstable(feature = "array_try_map", issue = "79711")]

26

impl<T, U, const N: usize, F> const FnMut<(usize,)> for &mut Drain<'_, T, U, N, F>

67

impl<T, U, const N: usize, F> const FnMut<(usize,)> for &mut Drain<'_, '_, T, N, F>

27

68

`where

28

69

`F: [const] FnMut(T) -> U,

29

70

30

extern "rust-call" fn call_mut(&mut self, (i,): (usize,)) -> Self::Output {

31

`` -

// SAFETY: increment moved before moving. if f panics, we drop the rest.

32

self.moved += 1;

33

assert_unsafe_precondition!(

34

check_library_ub,

35

"musnt index array out of bounds", (i: usize = i, size: usize = N) => i < size

36

);

37

`` -

// SAFETY: the i should also always go up, and musnt skip any, else some things will be leaked.

38

// SAFETY: if it goes down, we will drop freed elements. not good.

39

`` -

// SAFETY: caller guarantees never called with number >= N (see Drain::new)

40

(self.f)(unsafe { self.array.as_ptr().add(i).read() })

71

`` +

// FIXME(const-hack): ideally this would be an unsafe fn next(), and to use it you would instead |_| unsafe { drain.next() }.

72

extern "rust-call" fn call_mut(

73

&mut self,

74

(_ /* ignore argument */,): (usize,),

75

) -> Self::Output {

76

if T::IS_ZST {

77

// its UB to call this more than N times, so returning more ZSTs is valid.

78

// SAFETY: its a ZST? we conjur.

79

(self.f)(unsafe { conjure_zst::() })

80

} else {

81

`` +

// increment before moving; if f panics, we drop the rest.

82

let p = self.ptr;

83

`` +

// SAFETY: caller guarantees never called more than N times (see Drain::new)

84

self.ptr = unsafe { self.ptr.add(1) };

85

// SAFETY: we are allowed to move this.

86

(self.f)(unsafe { p.read() })

87

}

41

88

42

89

43

90

`#[rustc_const_unstable(feature = "array_try_map", issue = "79711")]

44

91

`#[unstable(feature = "array_try_map", issue = "79711")]

45

impl<T: [const] Destruct, U, const N: usize, F: FnMut(T) -> U> const Drop

46

for Drain<'_, T, U, N, F>

47

{

92

impl<T: [const] Destruct, const N: usize, F> const Drop for Drain<'_, '_, T, N, F> {

48

93

`fn drop(&mut self) {

49

let mut n = self.moved;

50

while n != N {

51

// SAFETY: moved must always be < N

52

unsafe { self.array.as_mut_ptr().add(n).drop_in_place() };

53

n += 1;

94

if !T::IS_ZST {

95

// SAFETY: we cant read more than N elements

96

let slice = unsafe {

97

from_raw_parts_mut::<[T]>(

98

self.ptr.as_ptr(),

99

`` +

// SAFETY: start <= end

100

self.end.offset_from_unsigned(self.ptr.as_ptr()),

101

)

102

};

103

+

104

// SAFETY: By the type invariant, we're allowed to drop all these. (we own it, after all)

105

unsafe { drop_in_place(slice) }

54

106

55

107

56

108

57

impl<'a, T, U, const N: usize, F: FnMut(T) -> U> Drain<'a, T, U, N, F> {

58

/// This function returns a function that lets you index the given array in const.

59

/// As implemented it can optimize better than iterators, and can be constified.

60

/// It acts like a sort of guard and iterator combined, which can be implemented

61

/// as it is a struct that implements const fn;

62

`` -

/// in that regard it is somewhat similar to an array::Iter implementing UncheckedIterator.

63

`` -

/// The only method you're really allowed to call is next(),

64

/// anything else is more or less UB, hence this function being unsafe.

65

/// Moved elements will not be dropped.

66

///

67

`` -

/// Previously this was implemented as a wrapper around a slice::Iter, which

68

`` -

/// called read() on the returned &T; gnarly stuff.

69

///

70

`` -

/// SAFETY: must be called in order of 0..N, without indexing out of bounds. (see Drain::call_mut)

71

/// Potentially the function could completely disregard the supplied argument, however i think that behaviour would be unintuitive.

72

`` -

// FIXME(const-hack): this is a hack for let guard = Guard(array); |i| f(guard[i]).

73

pub(super) const unsafe fn new(array: [T; N], f: &'a mut F) -> Self {

74

Self { array: ManuallyDrop::new(array), moved: 0, f }

75

}

76

}