@@ -161,14 +161,21 @@ impl<T> TypedArena<T> {
161161 available_capacity_bytes >= at_least_bytes
162162 }
163163
164+ /// Ensures there's enough space in the current chunk to fit `len` objects.
165+ #[ inline]
166+ fn ensure_capacity ( & self , len : usize ) {
167+ if !self . can_allocate ( len) {
168+ self . grow ( len) ;
169+ debug_assert ! ( self . can_allocate( len) ) ;
170+ }
171+ }
172+
164173 #[ inline]
165174 unsafe fn alloc_raw_slice ( & self , len : usize ) -> * mut T {
166175 assert ! ( mem:: size_of:: <T >( ) != 0 ) ;
167176 assert ! ( len != 0 ) ;
168177
169- if !self . can_allocate ( len) {
170- self . grow ( len) ;
171- }
178+ self . ensure_capacity ( len) ;
172179
173180 let start_ptr = self . ptr . get ( ) ;
174181 self . ptr . set ( start_ptr. add ( len) ) ;
@@ -203,19 +210,20 @@ impl<T> TypedArena<T> {
203210
204211 match size_hint {
205212 ( min, Some ( max) ) if min == max => {
206- if min == 0 {
213+ // We know the exact number of elements the iterator will produce here
214+ let len = min;
215+
216+ if len == 0 {
207217 return & mut [ ] ;
208218 }
209219
210- if !self . can_allocate ( min) {
211- self . grow ( min) ;
212- }
220+ self . ensure_capacity ( len) ;
213221
214222 let slice = self . ptr . get ( ) ;
215223
216224 unsafe {
217225 let mut ptr = self . ptr . get ( ) ;
218- for _ in 0 ..min {
226+ for _ in 0 ..len {
219227 // Write into uninitialized memory.
220228 ptr:: write ( ptr, iter. next ( ) . unwrap ( ) ) ;
221229 // Advance the pointer.
@@ -224,7 +232,7 @@ impl<T> TypedArena<T> {
224232 // we destroy the correct amount
225233 self . ptr . set ( ptr) ;
226234 }
227- slice:: from_raw_parts_mut ( slice, min )
235+ slice:: from_raw_parts_mut ( slice, len )
228236 }
229237 }
230238 _ => {
@@ -239,7 +247,7 @@ impl<T> TypedArena<T> {
239247 let len = vec. len ( ) ;
240248 let start_ptr = self . alloc_raw_slice ( len) ;
241249 vec. as_ptr ( ) . copy_to_nonoverlapping ( start_ptr, len) ;
242- vev . set_len ( 0 ) ;
250+ vec . set_len ( 0 ) ;
243251 slice:: from_raw_parts_mut ( start_ptr, len)
244252 }
245253 } )
@@ -488,16 +496,19 @@ impl DroplessArena {
488496
489497 match size_hint {
490498 ( min, Some ( max) ) if min == max => {
491- if min == 0 {
499+ // We know the exact number of elements the iterator will produce here
500+ let len = min;
501+
502+ if len == 0 {
492503 return & mut [ ]
493504 }
494- let size = min . checked_mul ( mem:: size_of :: < T > ( ) ) . unwrap ( ) ;
505+ let size = len . checked_mul ( mem:: size_of :: < T > ( ) ) . unwrap ( ) ;
495506 let mem = self . alloc_raw ( size, mem:: align_of :: < T > ( ) ) as * mut _ as * mut T ;
496507 unsafe {
497- for i in 0 ..min {
508+ for i in 0 ..len {
498509 ptr:: write ( mem. offset ( i as isize ) , iter. next ( ) . unwrap ( ) )
499510 }
500- slice:: from_raw_parts_mut ( mem, min )
511+ slice:: from_raw_parts_mut ( mem, len )
501512 }
502513 }
503514 ( _, _) => {
@@ -515,7 +526,7 @@ impl DroplessArena {
515526 mem:: align_of :: < T > ( )
516527 ) as * mut _ as * mut T ;
517528 vec. as_ptr ( ) . copy_to_nonoverlapping ( start_ptr, len) ;
518- vev . set_len ( 0 ) ;
529+ vec . set_len ( 0 ) ;
519530 slice:: from_raw_parts_mut ( start_ptr, len)
520531 }
521532 } )
0 commit comments