portalis_transpiler/
threading_translator.rs

1//! Python Threading/Multiprocessing to Rust Translation
2//!
3//! Translates Python threading and multiprocessing patterns to Rust using:
4//! - std::thread for basic threading
5//! - wasi_threading primitives for cross-platform support
6//! - rayon for data parallelism
7//! - crossbeam for channels and synchronization
8
9use crate::python_ast::{PyExpr, PyStmt, TypeAnnotation};
10use std::collections::HashMap;
11
12/// Main threading translator
13pub struct ThreadingTranslator {
14    /// Track used threading patterns
15    used_patterns: Vec<ThreadingPattern>,
16    /// Required imports
17    required_imports: HashMap<String, Vec<String>>,
18}
19
20/// Threading patterns detected
21#[derive(Debug, Clone)]
22pub enum ThreadingPattern {
23    BasicThread,
24    ThreadWithArgs,
25    ThreadPool,
26    Lock,
27    RLock,
28    Semaphore,
29    Event,
30    Condition,
31    Queue,
32    Process,
33    Pool,
34    Barrier,
35    ThreadLocal,
36}
37
38/// Translation strategy for threading constructs
39#[derive(Debug, Clone)]
40pub enum ThreadingStrategy {
41    /// Use std::thread
42    StdThread,
43    /// Use wasi_threading
44    WasiThreading,
45    /// Use rayon for data parallelism
46    Rayon,
47    /// Use crossbeam for channels
48    Crossbeam,
49}
50
51impl Default for ThreadingTranslator {
52    fn default() -> Self {
53        Self::new()
54    }
55}
56
57impl ThreadingTranslator {
58    pub fn new() -> Self {
59        Self {
60            used_patterns: Vec::new(),
61            required_imports: HashMap::new(),
62        }
63    }
64
65    /// Translate Python threading.Thread to Rust
66    pub fn translate_thread_creation(
67        &mut self,
68        target_func: &str,
69        args: &[String],
70        daemon: bool,
71    ) -> String {
72        self.used_patterns.push(ThreadingPattern::BasicThread);
73        self.add_import("std::thread", vec![]);
74
75        if args.is_empty() {
76            format!(
77                r#"let handle = std::thread::spawn(|| {{
78    {}()
79}});"#,
80                target_func
81            )
82        } else {
83            let args_capture = args
84                .iter()
85                .map(|arg| format!("let {} = {}.clone();", arg, arg))
86                .collect::<Vec<_>>()
87                .join("\n    ");
88
89            let args_list = args.join(", ");
90
91            format!(
92                r#"{{
93    {}
94    let handle = std::thread::spawn(move || {{
95        {}({})
96    }});
97}}"#,
98                args_capture, target_func, args_list
99            )
100        }
101    }
102
103    /// Translate Python threading.Lock to Rust Mutex
104    pub fn translate_lock(&mut self) -> String {
105        self.used_patterns.push(ThreadingPattern::Lock);
106        self.add_import("std::sync", vec!["std::sync::Mutex".to_string()]);
107
108        "Mutex::new(())".to_string()
109    }
110
111    /// Translate Python lock.acquire() / lock.release() to RAII
112    pub fn translate_lock_usage(&mut self, lock_var: &str, body: &str) -> String {
113        format!(
114            r#"{{
115    let _guard = {}.lock().unwrap();
116{}
117}}"#,
118            lock_var, body
119        )
120    }
121
122    /// Translate Python threading.RLock to Rust RwLock
123    pub fn translate_rlock(&mut self) -> String {
124        self.used_patterns.push(ThreadingPattern::RLock);
125        self.add_import("std::sync", vec!["std::sync::RwLock".to_string()]);
126
127        "RwLock::new(())".to_string()
128    }
129
130    /// Translate Python threading.Semaphore
131    pub fn translate_semaphore(&mut self, value: usize) -> String {
132        self.used_patterns.push(ThreadingPattern::Semaphore);
133        self.add_import("tokio::sync", vec!["tokio::sync::Semaphore".to_string()]);
134
135        format!("Arc::new(Semaphore::new({}))", value)
136    }
137
138    /// Translate Python threading.Event
139    pub fn translate_event(&mut self) -> String {
140        self.used_patterns.push(ThreadingPattern::Event);
141        self.add_import("std::sync", vec![
142            "std::sync::Arc".to_string(),
143            "std::sync::Condvar".to_string(),
144            "std::sync::Mutex".to_string(),
145        ]);
146
147        r#"Arc::new((Mutex::new(false), Condvar::new()))"#.to_string()
148    }
149
150    /// Translate event.wait()
151    pub fn translate_event_wait(&mut self, event_var: &str) -> String {
152        format!(
153            r#"{{
154    let (lock, cvar) = &*{};
155    let mut started = lock.lock().unwrap();
156    while !*started {{
157        started = cvar.wait(started).unwrap();
158    }}
159}}"#,
160            event_var
161        )
162    }
163
164    /// Translate event.set()
165    pub fn translate_event_set(&mut self, event_var: &str) -> String {
166        format!(
167            r#"{{
168    let (lock, cvar) = &*{};
169    let mut started = lock.lock().unwrap();
170    *started = true;
171    cvar.notify_all();
172}}"#,
173            event_var
174        )
175    }
176
177    /// Translate Python queue.Queue to Rust crossbeam channel
178    pub fn translate_queue(&mut self) -> String {
179        self.used_patterns.push(ThreadingPattern::Queue);
180        self.add_import("crossbeam::channel", vec!["crossbeam::channel::unbounded".to_string()]);
181
182        "let (tx, rx) = unbounded();".to_string()
183    }
184
185    /// Translate queue.put()
186    pub fn translate_queue_put(&mut self, queue_var: &str, value: &str) -> String {
187        format!("{}.send({}).unwrap();", queue_var, value)
188    }
189
190    /// Translate queue.get()
191    pub fn translate_queue_get(&mut self, queue_var: &str) -> String {
192        format!("{}.recv().unwrap()", queue_var)
193    }
194
195    /// Translate Python multiprocessing.Process to Rust thread
196    pub fn translate_process(&mut self, target_func: &str, args: &[String]) -> String {
197        self.used_patterns.push(ThreadingPattern::Process);
198        // In Rust, we use threads for processes (true multiprocessing requires different approach)
199        self.translate_thread_creation(target_func, args, false)
200    }
201
202    /// Translate Python multiprocessing.Pool
203    pub fn translate_pool(&mut self, num_workers: usize) -> String {
204        self.used_patterns.push(ThreadingPattern::Pool);
205        self.add_import("rayon", vec!["rayon::ThreadPoolBuilder".to_string()]);
206
207        format!(
208            r#"ThreadPoolBuilder::new()
209    .num_threads({})
210    .build()
211    .unwrap()"#,
212            num_workers
213        )
214    }
215
216    /// Translate pool.map()
217    pub fn translate_pool_map(&mut self, pool_var: &str, func: &str, iterable: &str) -> String {
218        self.add_import("rayon::prelude", vec!["rayon::prelude::*".to_string()]);
219
220        format!(
221            r#"{}.install(|| {{
222    {}.par_iter()
223        .map(|x| {}(x))
224        .collect::<Vec<_>>()
225}})"#,
226            pool_var, iterable, func
227        )
228    }
229
230    /// Translate Python threading.Barrier
231    pub fn translate_barrier(&mut self, parties: usize) -> String {
232        self.used_patterns.push(ThreadingPattern::Barrier);
233        self.add_import("std::sync", vec!["std::sync::Barrier".to_string()]);
234
235        format!("Arc::new(Barrier::new({}))", parties)
236    }
237
238    /// Translate barrier.wait()
239    pub fn translate_barrier_wait(&mut self, barrier_var: &str) -> String {
240        format!("{}.wait();", barrier_var)
241    }
242
243    /// Translate Python threading.local()
244    pub fn translate_thread_local(&mut self) -> String {
245        self.used_patterns.push(ThreadingPattern::ThreadLocal);
246
247        r#"thread_local! {
248    static LOCAL_DATA: RefCell<HashMap<String, String>> = RefCell::new(HashMap::new());
249}"#
250        .to_string()
251    }
252
253    /// Generate comprehensive threading patterns
254    pub fn generate_threading_patterns() -> String {
255        r#"// Threading pattern helpers
256
257/// Producer-Consumer pattern
258pub fn producer_consumer<T: Send + 'static>(
259    num_producers: usize,
260    num_consumers: usize,
261    producer_fn: impl Fn(usize) -> T + Send + Sync + 'static,
262    consumer_fn: impl Fn(T) + Send + Sync + 'static,
263) where
264    T: Clone,
265{
266    use crossbeam::channel::unbounded;
267    use std::sync::Arc;
268
269    let (tx, rx) = unbounded();
270
271    // Spawn producers
272    let producer_fn = Arc::new(producer_fn);
273    for i in 0..num_producers {
274        let tx = tx.clone();
275        let producer_fn = producer_fn.clone();
276        std::thread::spawn(move || {
277            let item = producer_fn(i);
278            tx.send(item).unwrap();
279        });
280    }
281    drop(tx);
282
283    // Spawn consumers
284    let consumer_fn = Arc::new(consumer_fn);
285    let mut handles = vec![];
286    for _ in 0..num_consumers {
287        let rx = rx.clone();
288        let consumer_fn = consumer_fn.clone();
289        let handle = std::thread::spawn(move || {
290            while let Ok(item) = rx.recv() {
291                consumer_fn(item);
292            }
293        });
294        handles.push(handle);
295    }
296
297    for handle in handles {
298        handle.join().unwrap();
299    }
300}
301
302/// Work stealing pattern using rayon
303pub fn work_stealing_map<T, R>(items: Vec<T>, f: impl Fn(&T) -> R + Sync) -> Vec<R>
304where
305    T: Send + Sync,
306    R: Send,
307{
308    use rayon::prelude::*;
309    items.par_iter().map(f).collect()
310}
311
312/// Pipeline pattern with channels
313pub fn pipeline<T: Send + 'static, U: Send + 'static, V: Send + 'static>(
314    input: Vec<T>,
315    stage1: impl Fn(T) -> U + Send + 'static,
316    stage2: impl Fn(U) -> V + Send + 'static,
317) -> Vec<V> {
318    use crossbeam::channel::unbounded;
319
320    let (tx1, rx1) = unbounded();
321    let (tx2, rx2) = unbounded();
322
323    // Stage 1
324    std::thread::spawn(move || {
325        for item in input {
326            tx1.send(stage1(item)).unwrap();
327        }
328    });
329
330    // Stage 2
331    std::thread::spawn(move || {
332        while let Ok(item) = rx1.recv() {
333            tx2.send(stage2(item)).unwrap();
334        }
335    });
336
337    // Collect results
338    rx2.iter().collect()
339}
340
341/// Scatter-gather pattern
342pub fn scatter_gather<T, R>(
343    items: Vec<T>,
344    worker_fn: impl Fn(T) -> R + Send + Sync + 'static,
345    num_workers: usize,
346) -> Vec<R>
347where
348    T: Send + 'static,
349    R: Send + 'static,
350{
351    use crossbeam::channel::unbounded;
352    use std::sync::Arc;
353
354    let (work_tx, work_rx) = unbounded();
355    let (result_tx, result_rx) = unbounded();
356
357    // Send work
358    for item in items {
359        work_tx.send(item).unwrap();
360    }
361    drop(work_tx);
362
363    // Spawn workers
364    let worker_fn = Arc::new(worker_fn);
365    let mut handles = vec![];
366    for _ in 0..num_workers {
367        let work_rx = work_rx.clone();
368        let result_tx = result_tx.clone();
369        let worker_fn = worker_fn.clone();
370
371        let handle = std::thread::spawn(move || {
372            while let Ok(item) = work_rx.recv() {
373                let result = worker_fn(item);
374                result_tx.send(result).unwrap();
375            }
376        });
377        handles.push(handle);
378    }
379    drop(result_tx);
380
381    // Gather results
382    let results = result_rx.iter().collect();
383
384    for handle in handles {
385        handle.join().unwrap();
386    }
387
388    results
389}
390"#
391        .to_string()
392    }
393
394    // Helper methods
395
396    fn add_import(&mut self, crate_name: &str, items: Vec<String>) {
397        self.required_imports
398            .entry(crate_name.to_string())
399            .or_insert_with(Vec::new)
400            .extend(items);
401    }
402
403    pub fn generate_imports(&self) -> String {
404        let mut imports = Vec::new();
405
406        for (crate_name, items) in &self.required_imports {
407            if items.is_empty() {
408                imports.push(format!("use {};", crate_name));
409            } else {
410                for item in items {
411                    imports.push(format!("use {};", item));
412                }
413            }
414        }
415
416        // Add Arc for thread safety
417        if !imports.is_empty() {
418            imports.insert(0, "use std::sync::Arc;".to_string());
419        }
420
421        imports.join("\n")
422    }
423
424    pub fn get_used_patterns(&self) -> &[ThreadingPattern] {
425        &self.used_patterns
426    }
427}
428
429/// Multiprocessing mapper for Python multiprocessing module
430pub struct MultiprocessingMapper;
431
432impl MultiprocessingMapper {
433    /// Translate multiprocessing.Pool to rayon ThreadPool
434    pub fn translate_pool(num_processes: Option<usize>) -> String {
435        let workers = num_processes
436            .map(|n| n.to_string())
437            .unwrap_or_else(|| "num_cpus::get()".to_string());
438
439        format!(
440            r#"use rayon::ThreadPoolBuilder;
441
442let pool = ThreadPoolBuilder::new()
443    .num_threads({})
444    .build()
445    .unwrap();"#,
446            workers
447        )
448    }
449
450    /// Translate Pool.map to parallel iterator
451    pub fn translate_pool_map(func: &str, iterable: &str) -> String {
452        format!(
453            r#"use rayon::prelude::*;
454
455let results: Vec<_> = {}
456    .par_iter()
457    .map(|x| {}(x))
458    .collect();"#,
459            iterable, func
460        )
461    }
462
463    /// Translate Pool.starmap
464    pub fn translate_starmap(func: &str, iterable: &str) -> String {
465        format!(
466            r#"use rayon::prelude::*;
467
468let results: Vec<_> = {}
469    .par_iter()
470    .map(|(args)| {{
471        {}(args)
472    }})
473    .collect();"#,
474            iterable, func
475        )
476    }
477
478    /// Translate multiprocessing.Queue to crossbeam channel
479    pub fn translate_queue() -> String {
480        r#"use crossbeam::channel::unbounded;
481
482let (tx, rx) = unbounded();"#
483            .to_string()
484    }
485
486    /// Translate multiprocessing.Pipe
487    pub fn translate_pipe() -> String {
488        r#"use crossbeam::channel::unbounded;
489
490let (parent_tx, child_rx) = unbounded();
491let (child_tx, parent_rx) = unbounded();"#
492            .to_string()
493    }
494
495    /// Translate multiprocessing.Manager (shared memory)
496    pub fn translate_manager() -> String {
497        r#"use std::sync::{Arc, Mutex};
498
499// Shared data structure
500let shared_data = Arc::new(Mutex::new(HashMap::new()));"#
501            .to_string()
502    }
503
504    /// Translate multiprocessing.Lock
505    pub fn translate_lock() -> String {
506        r#"use std::sync::{Arc, Mutex};
507
508let lock = Arc::new(Mutex::new(()));"#
509            .to_string()
510    }
511
512    /// Translate multiprocessing.Value (shared value)
513    pub fn translate_value(type_hint: &str, initial: &str) -> String {
514        format!(
515            r#"use std::sync::{{Arc, Mutex}};
516
517let value = Arc::new(Mutex::new({} as {}));"#,
518            initial, type_hint
519        )
520    }
521
522    /// Translate multiprocessing.Array (shared array)
523    pub fn translate_array(type_hint: &str, size: usize) -> String {
524        format!(
525            r#"use std::sync::{{Arc, Mutex}};
526
527let array = Arc::new(Mutex::new(vec![{} as {}; {}]));"#,
528            "Default::default()", type_hint, size
529        )
530    }
531}
532
533#[cfg(test)]
534mod tests {
535    use super::*;
536
537    #[test]
538    fn test_thread_creation() {
539        let mut translator = ThreadingTranslator::new();
540        let code = translator.translate_thread_creation("worker_func", &[], false);
541        assert!(code.contains("std::thread::spawn"));
542        assert!(code.contains("worker_func()"));
543    }
544
545    #[test]
546    fn test_lock_translation() {
547        let mut translator = ThreadingTranslator::new();
548        let lock = translator.translate_lock();
549        assert!(lock.contains("Mutex::new"));
550    }
551
552    #[test]
553    fn test_queue_translation() {
554        let mut translator = ThreadingTranslator::new();
555        let queue = translator.translate_queue();
556        assert!(queue.contains("unbounded"));
557    }
558
559    #[test]
560    fn test_pool_translation() {
561        let mut translator = ThreadingTranslator::new();
562        let pool = translator.translate_pool(4);
563        assert!(pool.contains("ThreadPoolBuilder"));
564        assert!(pool.contains("num_threads(4)"));
565    }
566
567    #[test]
568    fn test_import_generation() {
569        let mut translator = ThreadingTranslator::new();
570        translator.translate_lock();
571        translator.translate_queue();
572
573        let imports = translator.generate_imports();
574        assert!(imports.contains("std::sync::Mutex"));
575        assert!(imports.contains("crossbeam::channel"));
576    }
577}
portalis_transpiler/threading_translator.rs

portalis_transpiler/
threading_translator.rs
