1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

//! Module containing structs with parallel and serial iteration methods for batch
//! running of `sim_drive` method
// crate local
use super::RustSimDrive;
use crate::imports::*;
use crate::proc_macros::add_pyo3_api;
#[cfg(feature = "pyo3")]
use crate::pyo3imports::*;
use rayon::prelude::*;

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
#[add_pyo3_api(
    #[pyo3(name = "sim_drive")]
    /// Calls `sim_drive` method for each simdrive instance in vec.
    /// # Arguments:
    /// * parallelize: whether to parallelize `sim_drive` calls, defaults to `true`
    fn sim_drive_py(&mut self, parallelize: Option<bool>) -> anyhow::Result<()> {
        self.sim_drive(parallelize)
    }

    #[pyo3(name = "push")]
    fn push_py(&mut self, sd: RustSimDrive) {
        self.push(sd)
    }

    #[pyo3(name = "pop")]
    fn pop_py(&mut self) -> Option<RustSimDrive> {
        self.pop()
    }

    #[pyo3(name = "remove")]
    fn remove_py(&mut self, idx: usize) {
        self.remove(idx);
    }

    #[pyo3(name = "insert")]
    fn insert_py(&mut self, idx: usize, sd: RustSimDrive) {
        self.insert(idx, sd);
    }
)]
pub struct SimDriveVec(pub Vec<RustSimDrive>);

impl SimDriveVec {
    /// Calls `sim_drive` method for each simdrive instance in vec.
    /// # Arguments:
    /// * parallelize: whether to parallelize `sim_drive` calls
    pub fn sim_drive(&mut self, parallelize: Option<bool>) -> anyhow::Result<()> {
        let parallelize = parallelize.unwrap_or(true);
        if parallelize {
            self.0.par_iter_mut().enumerate().try_for_each(|(i, sd)| {
                sd.sim_drive(None, None)
                    .with_context(|| format!("simdrive idx: {}", i))
            })?;
        } else {
            self.0.iter_mut().enumerate().try_for_each(|(i, sd)| {
                sd.sim_drive(None, None)
                    .with_context(|| format!("simdrive idx: {}", i))
            })?;
        }
        Ok(())
    }

    pub fn push(&mut self, sd: RustSimDrive) {
        self.0.push(sd);
    }

    pub fn pop(&mut self) -> Option<RustSimDrive> {
        self.0.pop()
    }

    pub fn remove(&mut self, idx: usize) {
        self.0.remove(idx);
    }

    pub fn insert(&mut self, idx: usize, sd: RustSimDrive) {
        self.0.insert(idx, sd);
    }
}

impl SerdeAPI for SimDriveVec {}