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

use super::{PyType, PyTypeRef};
use crate::{
    builtins::PyTupleRef,
    class::PyClassImpl,
    function::PosArgs,
    protocol::{PyIter, PyIterReturn},
    types::{Constructor, IterNext, Iterable, SelfIter},
    Context, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine,
};

#[pyclass(module = false, name = "map", traverse)]
#[derive(Debug)]
pub struct PyMap {
    mapper: PyObjectRef,
    iterators: Vec<PyIter>,
}

impl PyPayload for PyMap {
    fn class(ctx: &Context) -> &'static Py<PyType> {
        ctx.types.map_type
    }
}

impl Constructor for PyMap {
    type Args = (PyObjectRef, PosArgs<PyIter>);

    fn py_new(cls: PyTypeRef, (mapper, iterators): Self::Args, vm: &VirtualMachine) -> PyResult {
        let iterators = iterators.into_vec();
        PyMap { mapper, iterators }
            .into_ref_with_type(vm, cls)
            .map(Into::into)
    }
}

#[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE))]
impl PyMap {
    #[pymethod(magic)]
    fn length_hint(&self, vm: &VirtualMachine) -> PyResult<usize> {
        self.iterators.iter().try_fold(0, |prev, cur| {
            let cur = cur.as_ref().to_owned().length_hint(0, vm)?;
            let max = std::cmp::max(prev, cur);
            Ok(max)
        })
    }

    #[pymethod(magic)]
    fn reduce(&self, vm: &VirtualMachine) -> (PyTypeRef, PyTupleRef) {
        let mut vec = vec![self.mapper.clone()];
        vec.extend(self.iterators.iter().map(|o| o.clone().into()));
        (vm.ctx.types.map_type.to_owned(), vm.new_tuple(vec))
    }
}

impl SelfIter for PyMap {}
impl IterNext for PyMap {
    fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
        let mut next_objs = Vec::new();
        for iterator in &zelf.iterators {
            let item = match iterator.next(vm)? {
                PyIterReturn::Return(obj) => obj,
                PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)),
            };
            next_objs.push(item);
        }

        // the mapper itself can raise StopIteration which does stop the map iteration
        PyIterReturn::from_pyresult(zelf.mapper.call(next_objs, vm), vm)
    }
}

pub fn init(context: &Context) {
    PyMap::extend_class(context, context.types.map_type);
}