1- use std:: sync:: atomic:: { AtomicBool , Ordering } ;
1+ use std:: sync:: atomic:: { AtomicBool , AtomicUsize , Ordering } ;
22
33use crate :: parallel:: { num_threads, Reduce } ;
44
@@ -89,23 +89,25 @@ where
8989/// for file-io as it won't make use of sorted inputs well.
9090/// Note that `periodic` is not guaranteed to be called in case other threads come up first and finish too fast.
9191// TODO: better docs
92- pub fn in_parallel_with_slice < I , S , E > (
93- input : & [ I ] ,
92+ pub fn in_parallel_with_slice < I , S , R , E > (
93+ input : & mut [ I ] ,
9494 thread_limit : Option < usize > ,
9595 new_thread_state : impl FnMut ( usize ) -> S + Send + Clone ,
96- consume : impl FnMut ( & I , & mut S ) -> Result < ( ) , E > + Send + Clone ,
96+ consume : impl FnMut ( & mut I , & mut S ) -> Result < ( ) , E > + Send + Clone ,
9797 mut periodic : impl FnMut ( ) -> Option < std:: time:: Duration > + Send ,
98- ) -> Result < Vec < S > , E >
98+ state_to_rval : impl FnOnce ( S ) -> R + Send + Clone ,
99+ ) -> Result < Vec < R > , E >
99100where
100- I : Send + Sync ,
101- E : Send + Sync ,
102- S : Send ,
101+ I : Send ,
102+ E : Send ,
103+ R : Send ,
103104{
104105 let num_threads = num_threads ( thread_limit) ;
105- let num_items = input. len ( ) ;
106106 let mut results = Vec :: with_capacity ( num_threads) ;
107107 let stop_everything = & AtomicBool :: default ( ) ;
108+ let index = & AtomicUsize :: default ( ) ;
108109
110+ // TODO: use std::thread::scope() once Rust 1.63 is available.
109111 crossbeam_utils:: thread:: scope ( {
110112 move |s| {
111113 s. spawn ( {
@@ -124,24 +126,44 @@ where
124126 }
125127 } ) ;
126128
129+ let input_len = input. len ( ) ;
130+ struct Input < I > ( * mut [ I ] )
131+ where
132+ I : Send ;
133+
134+ // SAFETY: I is Send + Sync, so is a *mut [I]
135+ #[ allow( unsafe_code) ]
136+ unsafe impl < I > Send for Input < I > where I : Send { }
137+
127138 let threads: Vec < _ > = ( 0 ..num_threads)
128139 . map ( |thread_id| {
129140 s. spawn ( {
130141 let mut new_thread_state = new_thread_state. clone ( ) ;
142+ let state_to_rval = state_to_rval. clone ( ) ;
131143 let mut consume = consume. clone ( ) ;
144+ let input = Input ( input as * mut [ I ] ) ;
132145 move |_| {
133146 let mut state = new_thread_state ( thread_id) ;
134- for input_index in ( thread_id..num_items) . step_by ( num_threads) {
147+ while let Ok ( input_index) = index
148+ . fetch_update ( Ordering :: SeqCst , Ordering :: SeqCst , |x| ( x < input_len) . then ( || x + 1 ) )
149+ {
135150 if stop_everything. load ( Ordering :: Relaxed ) {
136151 break ;
137152 }
138- let item = & input[ input_index] ;
153+ // SAFETY: our atomic counter for `input_index` is only ever incremented, yielding
154+ // each item exactly once.
155+ let item = {
156+ #[ allow( unsafe_code) ]
157+ unsafe {
158+ & mut ( & mut * input. 0 ) [ input_index]
159+ }
160+ } ;
139161 if let Err ( err) = consume ( item, & mut state) {
140162 stop_everything. store ( true , Ordering :: Relaxed ) ;
141163 return Err ( err) ;
142164 }
143165 }
144- Ok ( state)
166+ Ok ( state_to_rval ( state) )
145167 }
146168 } )
147169 } )
0 commit comments