rustler/types/
map.rs

1//! Utilities used to access and create Erlang maps.
2
3use super::atom;
4use crate::wrapper::map;
5use crate::{Decoder, Encoder, Env, Error, NifResult, Term};
6use std::ops::RangeInclusive;
7
8#[inline]
9pub fn map_new(env: Env) -> Term {
10    unsafe { Term::new(env, map::map_new(env.as_c_arg())) }
11}
12
13/// ## Map terms
14impl<'a> Term<'a> {
15    /// Constructs a new, empty map term.
16    ///
17    /// ### Elixir equivalent
18    /// ```elixir
19    /// %{}
20    /// ```
21    #[inline]
22    pub fn map_new(env: Env<'a>) -> Term<'a> {
23        map_new(env)
24    }
25
26    /// Construct a new map from two vectors
27    ///
28    /// ### Elixir equivalent
29    /// ```elixir
30    /// keys = ["foo", "bar"]
31    /// values = [1, 2]
32    /// Enum.zip(keys, values) |> Map.new()
33    /// ```
34    #[inline]
35    pub fn map_from_arrays(
36        env: Env<'a>,
37        keys: &[impl Encoder],
38        values: &[impl Encoder],
39    ) -> NifResult<Term<'a>> {
40        if keys.len() == values.len() {
41            let keys: Vec<_> = keys.iter().map(|k| k.encode(env).as_c_arg()).collect();
42            let values: Vec<_> = values.iter().map(|v| v.encode(env).as_c_arg()).collect();
43
44            unsafe {
45                map::make_map_from_arrays(env.as_c_arg(), &keys, &values)
46                    .map_or_else(|| Err(Error::BadArg), |map| Ok(Term::new(env, map)))
47            }
48        } else {
49            Err(Error::BadArg)
50        }
51    }
52
53    /// Construct a new map from two vectors of terms.
54    ///
55    /// It is identical to map_from_arrays, but requires the keys and values to
56    /// be encoded already - this is useful for constructing maps whose values
57    /// or keys are different Rust types, with the same performance as map_from_arrays.
58    pub fn map_from_term_arrays(
59        env: Env<'a>,
60        keys: &[Term<'a>],
61        values: &[Term<'a>],
62    ) -> NifResult<Term<'a>> {
63        if keys.len() == values.len() {
64            let keys: Vec<_> = keys.iter().map(|k| k.as_c_arg()).collect();
65            let values: Vec<_> = values.iter().map(|v| v.as_c_arg()).collect();
66
67            unsafe {
68                map::make_map_from_arrays(env.as_c_arg(), &keys, &values)
69                    .map_or_else(|| Err(Error::BadArg), |map| Ok(Term::new(env, map)))
70            }
71        } else {
72            Err(Error::BadArg)
73        }
74    }
75
76    /// Construct a new map from pairs of terms
77    ///
78    /// It is similar to `map_from_arrays` but
79    /// receives only one vector with the pairs
80    /// of `(key, value)`.
81    ///
82    /// ### Elixir equivalent
83    /// ```elixir
84    /// Map.new([{"foo", 1}, {"bar", 2}])
85    /// ```
86    #[inline]
87    pub fn map_from_pairs(
88        env: Env<'a>,
89        pairs: &[(impl Encoder, impl Encoder)],
90    ) -> NifResult<Term<'a>> {
91        let (keys, values): (Vec<_>, Vec<_>) = pairs
92            .iter()
93            .map(|(k, v)| (k.encode(env).as_c_arg(), v.encode(env).as_c_arg()))
94            .unzip();
95
96        unsafe {
97            map::make_map_from_arrays(env.as_c_arg(), &keys, &values)
98                .map_or_else(|| Err(Error::BadArg), |map| Ok(Term::new(env, map)))
99        }
100    }
101
102    /// Gets the value corresponding to a key in a map term.
103    ///
104    /// Returns Err(Error::BadArg) if the term is not a map or if
105    /// key doesn't exist in the map.
106    ///
107    /// ### Elixir equivalent
108    /// ```elixir
109    /// Map.get(self_term, key)
110    /// ```
111    #[inline]
112    pub fn map_get(self, key: impl Encoder) -> NifResult<Term<'a>> {
113        let env = self.get_env();
114        match unsafe {
115            map::get_map_value(env.as_c_arg(), self.as_c_arg(), key.encode(env).as_c_arg())
116        } {
117            Some(value) => Ok(unsafe { Term::new(env, value) }),
118            None => Err(Error::BadArg),
119        }
120    }
121
122    /// Gets the size of a map term.
123    ///
124    /// Returns Err(Error::BadArg) if the term is not a map.
125    ///
126    /// ### Elixir equivalent
127    /// ```elixir
128    /// map_size(self_term)
129    /// ```
130    #[inline]
131    pub fn map_size(self) -> NifResult<usize> {
132        let env = self.get_env();
133        unsafe { map::get_map_size(env.as_c_arg(), self.as_c_arg()).ok_or(Error::BadArg) }
134    }
135
136    /// Makes a copy of the self map term and sets key to value.
137    /// If the value already exists, it is overwritten.
138    ///
139    /// Returns Err(Error::BadArg) if the term is not a map.
140    ///
141    /// ### Elixir equivalent
142    /// ```elixir
143    /// Map.put(self_term, key, value)
144    /// ```
145    #[inline]
146    pub fn map_put(self, key: impl Encoder, value: impl Encoder) -> NifResult<Term<'a>> {
147        let env = self.get_env();
148
149        match unsafe {
150            map::map_put(
151                env.as_c_arg(),
152                self.as_c_arg(),
153                key.encode(env).as_c_arg(),
154                value.encode(env).as_c_arg(),
155            )
156        } {
157            Some(inner) => Ok(unsafe { Term::new(env, inner) }),
158            None => Err(Error::BadArg),
159        }
160    }
161
162    /// Makes a copy of the self map term and removes key. If the key
163    /// doesn't exist, the original map is returned.
164    ///
165    /// Returns Err(Error::BadArg) if the term is not a map.
166    ///
167    /// ### Elixir equivalent
168    /// ```elixir
169    /// Map.delete(self_term, key)
170    /// ```
171    #[inline]
172    pub fn map_remove(self, key: impl Encoder) -> NifResult<Term<'a>> {
173        let env = self.get_env();
174
175        match unsafe {
176            map::map_remove(env.as_c_arg(), self.as_c_arg(), key.encode(env).as_c_arg())
177        } {
178            Some(inner) => Ok(unsafe { Term::new(env, inner) }),
179            None => Err(Error::BadArg),
180        }
181    }
182
183    /// Makes a copy of the self map term where key is set to value.
184    ///
185    /// Returns Err(Error::BadArg) if the term is not a map of if key
186    /// doesn't exist.
187    #[inline]
188    pub fn map_update(self, key: impl Encoder, new_value: impl Encoder) -> NifResult<Term<'a>> {
189        let env = self.get_env();
190
191        match unsafe {
192            map::map_update(
193                env.as_c_arg(),
194                self.as_c_arg(),
195                key.encode(env).as_c_arg(),
196                new_value.encode(env).as_c_arg(),
197            )
198        } {
199            Some(inner) => Ok(unsafe { Term::new(env, inner) }),
200            None => Err(Error::BadArg),
201        }
202    }
203}
204
205struct SimpleMapIterator<'a> {
206    map: Term<'a>,
207    entry: map::MapIteratorEntry,
208    iter: Option<map::ErlNifMapIterator>,
209    last_key: Option<Term<'a>>,
210    done: bool,
211}
212
213impl<'a> SimpleMapIterator<'a> {
214    fn next(&mut self) -> Option<(Term<'a>, Term<'a>)> {
215        if self.done {
216            return None;
217        }
218
219        let iter = loop {
220            match self.iter.as_mut() {
221                None => {
222                    match unsafe {
223                        map::map_iterator_create(
224                            self.map.get_env().as_c_arg(),
225                            self.map.as_c_arg(),
226                            self.entry,
227                        )
228                    } {
229                        Some(iter) => {
230                            self.iter = Some(iter);
231                            continue;
232                        }
233                        None => {
234                            self.done = true;
235                            return None;
236                        }
237                    }
238                }
239                Some(iter) => {
240                    break iter;
241                }
242            }
243        };
244
245        let env = self.map.get_env();
246
247        unsafe {
248            match map::map_iterator_get_pair(env.as_c_arg(), iter) {
249                Some((key, value)) => {
250                    match self.entry {
251                        map::MapIteratorEntry::First => {
252                            map::map_iterator_next(env.as_c_arg(), iter);
253                        }
254                        map::MapIteratorEntry::Last => {
255                            map::map_iterator_prev(env.as_c_arg(), iter);
256                        }
257                    }
258                    let key = Term::new(env, key);
259                    self.last_key = Some(key);
260                    Some((key, Term::new(env, value)))
261                }
262                None => {
263                    self.done = true;
264                    None
265                }
266            }
267        }
268    }
269}
270
271impl Drop for SimpleMapIterator<'_> {
272    fn drop(&mut self) {
273        if let Some(iter) = self.iter.as_mut() {
274            unsafe {
275                map::map_iterator_destroy(self.map.get_env().as_c_arg(), iter);
276            }
277        }
278    }
279}
280
281pub struct MapIterator<'a> {
282    forward: SimpleMapIterator<'a>,
283    reverse: SimpleMapIterator<'a>,
284}
285
286impl<'a> MapIterator<'a> {
287    pub fn new(map: Term<'a>) -> Option<MapIterator<'a>> {
288        if map.is_map() {
289            Some(MapIterator {
290                forward: SimpleMapIterator {
291                    map,
292                    entry: map::MapIteratorEntry::First,
293                    iter: None,
294                    last_key: None,
295                    done: false,
296                },
297                reverse: SimpleMapIterator {
298                    map,
299                    entry: map::MapIteratorEntry::Last,
300                    iter: None,
301                    last_key: None,
302                    done: false,
303                },
304            })
305        } else {
306            None
307        }
308    }
309}
310
311impl<'a> Iterator for MapIterator<'a> {
312    type Item = (Term<'a>, Term<'a>);
313
314    fn next(&mut self) -> Option<Self::Item> {
315        self.forward.next().and_then(|(key, value)| {
316            if self.reverse.last_key == Some(key) {
317                self.forward.done = true;
318                self.reverse.done = true;
319                return None;
320            }
321            Some((key, value))
322        })
323    }
324}
325
326impl DoubleEndedIterator for MapIterator<'_> {
327    fn next_back(&mut self) -> Option<Self::Item> {
328        self.reverse.next().and_then(|(key, value)| {
329            if self.forward.last_key == Some(key) {
330                self.forward.done = true;
331                self.reverse.done = true;
332                return None;
333            }
334            Some((key, value))
335        })
336    }
337}
338
339impl<'a> Decoder<'a> for MapIterator<'a> {
340    fn decode(term: Term<'a>) -> NifResult<Self> {
341        match MapIterator::new(term) {
342            Some(iter) => Ok(iter),
343            None => Err(Error::BadArg),
344        }
345    }
346}
347
348impl<'a, T> Decoder<'a> for RangeInclusive<T>
349where
350    T: Decoder<'a>,
351{
352    fn decode(term: Term<'a>) -> NifResult<Self> {
353        let name = term.map_get(atom::__struct__())?;
354
355        match name.atom_to_string()?.as_ref() {
356            "Elixir.Range" => (),
357            _ => return Err(Error::BadArg),
358        }
359
360        let first = term.map_get(atom::first())?.decode::<T>()?;
361        let last = term.map_get(atom::last())?.decode::<T>()?;
362        if let Ok(step) = term.map_get(atom::step()) {
363            match step.decode::<i64>()? {
364                1 => (),
365                _ => return Err(Error::BadArg),
366            }
367        }
368
369        Ok(first..=last)
370    }
371}