1
- use core:: iter:: { InPlaceIterable , SourceIter } ;
1
+ use core:: iter:: { InPlaceIterable , SourceIter , TrustedRandomAccess } ;
2
2
use core:: mem:: { self , ManuallyDrop } ;
3
3
use core:: ptr:: { self } ;
4
4
52
52
)
53
53
} ;
54
54
55
- // use try-fold since
56
- // - it vectorizes better for some iterator adapters
57
- // - unlike most internal iteration methods, it only takes a &mut self
58
- // - it lets us thread the write pointer through its innards and get it back in the end
59
- let sink = InPlaceDrop { inner : dst_buf, dst : dst_buf } ;
60
- let sink = iterator
61
- . try_fold :: < _ , _ , Result < _ , !> > ( sink, write_in_place_with_drop ( dst_end) )
62
- . unwrap ( ) ;
63
- // iteration succeeded, don't drop head
64
- let dst = ManuallyDrop :: new ( sink) . dst ;
55
+ let len = SpecInPlaceCollect :: collect_in_place ( & mut iterator, dst_buf, dst_end) ;
65
56
66
57
let src = unsafe { iterator. as_inner ( ) . as_into_iter ( ) } ;
67
58
// check if SourceIter contract was upheld
72
63
// then the source pointer will stay in its initial position and we can't use it as reference
73
64
if src. ptr != src_ptr {
74
65
debug_assert ! (
75
- dst as * const _ <= src. ptr,
66
+ unsafe { dst_buf . add ( len ) as * const _ } <= src. ptr,
76
67
"InPlaceIterable contract violation, write pointer advanced beyond read pointer"
77
68
) ;
78
69
}
82
73
// but prevent drop of the allocation itself once IntoIter goes out of scope
83
74
src. forget_allocation ( ) ;
84
75
85
- let vec = unsafe {
86
- let len = dst. offset_from ( dst_buf) as usize ;
87
- Vec :: from_raw_parts ( dst_buf, len, cap)
88
- } ;
76
+ let vec = unsafe { Vec :: from_raw_parts ( dst_buf, len, cap) } ;
89
77
90
78
vec
91
79
}
@@ -106,3 +94,52 @@ fn write_in_place_with_drop<T>(
106
94
Ok ( sink)
107
95
}
108
96
}
97
+
98
+ /// Helper trait to hold specialized implementations of the in-place iterate-collect loop
99
+ trait SpecInPlaceCollect < T , I > : Iterator < Item = T > {
100
+ /// Collects an iterator (`self`) into the destination buffer (`dst`) and returns the number of items
101
+ /// collected. `end` is the last writable element of the allocation and used for bounds checks.
102
+ fn collect_in_place ( & mut self , dst : * mut T , end : * const T ) -> usize ;
103
+ }
104
+
105
+ impl < T , I > SpecInPlaceCollect < T , I > for I
106
+ where
107
+ I : Iterator < Item = T > ,
108
+ {
109
+ #[ inline]
110
+ default fn collect_in_place ( & mut self , dst_buf : * mut T , end : * const T ) -> usize {
111
+ // use try-fold since
112
+ // - it vectorizes better for some iterator adapters
113
+ // - unlike most internal iteration methods, it only takes a &mut self
114
+ // - it lets us thread the write pointer through its innards and get it back in the end
115
+ let sink = InPlaceDrop { inner : dst_buf, dst : dst_buf } ;
116
+ let sink =
117
+ self . try_fold :: < _ , _ , Result < _ , !> > ( sink, write_in_place_with_drop ( end) ) . unwrap ( ) ;
118
+ // iteration succeeded, don't drop head
119
+ unsafe { ManuallyDrop :: new ( sink) . dst . offset_from ( dst_buf) as usize }
120
+ }
121
+ }
122
+
123
+ impl < T , I > SpecInPlaceCollect < T , I > for I
124
+ where
125
+ I : Iterator < Item = T > + TrustedRandomAccess ,
126
+ {
127
+ #[ inline]
128
+ fn collect_in_place ( & mut self , dst_buf : * mut T , end : * const T ) -> usize {
129
+ let len = self . size ( ) ;
130
+ let mut drop_guard = InPlaceDrop { inner : dst_buf, dst : dst_buf } ;
131
+ for i in 0 ..len {
132
+ // Safety: InplaceIterable contract guarantees that for every element we read
133
+ // one slot in the underlying storage will have been freed up and we can immediately
134
+ // write back the result.
135
+ unsafe {
136
+ let dst = dst_buf. offset ( i as isize ) ;
137
+ debug_assert ! ( dst as * const _ <= end, "InPlaceIterable contract violation" ) ;
138
+ ptr:: write ( dst, self . __iterator_get_unchecked ( i) ) ;
139
+ drop_guard. dst = dst. add ( 1 ) ;
140
+ }
141
+ }
142
+ mem:: forget ( drop_guard) ;
143
+ len
144
+ }
145
+ }
0 commit comments