tlua/
rust_tables.rs

1use crate::{
2    ffi,
3    lua_tables::LuaTable,
4    tuples::TuplePushError::{self, First, Other},
5    AsLua, LuaRead, LuaState, Push, PushGuard, PushInto, PushOne, PushOneInto, ReadResult, Void,
6    WrongType,
7};
8
9use std::collections::{BTreeMap, HashMap, HashSet};
10use std::fmt::{self, Debug};
11use std::hash::Hash;
12use std::iter;
13use std::num::NonZeroI32;
14
15#[inline]
16pub(crate) fn push_iter<L, I>(lua: L, iterator: I) -> Result<PushGuard<L>, (PushIterErrorOf<I>, L)>
17where
18    L: AsLua,
19    I: Iterator,
20    <I as Iterator>::Item: PushInto<LuaState>,
21{
22    // creating empty table
23    unsafe { ffi::lua_newtable(lua.as_lua()) };
24
25    for (elem, index) in iterator.zip(1..) {
26        let size = match elem.push_into_lua(lua.as_lua()) {
27            Ok(pushed) => pushed.forget_internal(),
28            Err((err, _)) => unsafe {
29                // TODO(gmoshkin): return an error capturing this push guard
30                // drop the lua table
31                drop(PushGuard::new(lua.as_lua(), 1));
32                return Err((PushIterError::ValuePushError(err), lua));
33            },
34        };
35
36        match size {
37            0 => continue,
38            1 => {
39                lua.as_lua().push_one(index).forget_internal();
40                unsafe { ffi::lua_insert(lua.as_lua(), -2) }
41                unsafe { ffi::lua_settable(lua.as_lua(), -3) }
42            }
43            2 => unsafe { ffi::lua_settable(lua.as_lua(), -3) },
44            n => unsafe {
45                // TODO(gmoshkin): return an error capturing this push guard
46                // n + 1 == n values from the recent push + lua table
47                drop(PushGuard::new(lua.as_lua(), n + 1));
48                return Err((PushIterError::TooManyValues(n), lua));
49            },
50        }
51    }
52
53    unsafe { Ok(PushGuard::new(lua, 1)) }
54}
55
56pub type PushIterErrorOf<I> = PushIterError<<<I as Iterator>::Item as PushInto<LuaState>>::Err>;
57
58#[derive(Debug, PartialEq, Eq)]
59pub enum PushIterError<E> {
60    TooManyValues(i32),
61    ValuePushError(E),
62}
63
64impl<E> PushIterError<E> {
65    pub fn map<F, R>(self, f: F) -> PushIterError<R>
66    where
67        F: FnOnce(E) -> R,
68    {
69        match self {
70            Self::ValuePushError(e) => PushIterError::ValuePushError(f(e)),
71            Self::TooManyValues(n) => PushIterError::TooManyValues(n),
72        }
73    }
74}
75
76impl<E> fmt::Display for PushIterError<E>
77where
78    E: fmt::Display,
79{
80    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
81        match self {
82            Self::TooManyValues(n) => {
83                write!(
84                    fmt,
85                    "Can only push 1 or 2 values as lua table item, got {} instead",
86                    n,
87                )
88            }
89            Self::ValuePushError(e) => {
90                write!(fmt, "Pushing iterable item failed: {}", e)
91            }
92        }
93    }
94}
95
96// NOTE: only the following From<_> for Void implementations are correct,
97//       don't add other ones!
98
99// T::Err: Void => no error possible
100// NOTE: making this one generic would conflict with the below implementations.
101impl From<PushIterError<Void>> for Void {
102    fn from(_: PushIterError<Void>) -> Self {
103        unreachable!("no way to create instance of Void")
104    }
105}
106
107// T::Err: Void; (T,) => no error possible
108impl<T> From<PushIterError<TuplePushError<T, Void>>> for Void
109where
110    T: Into<Void>,
111{
112    fn from(_: PushIterError<TuplePushError<T, Void>>) -> Self {
113        unreachable!("no way to create instance of Void")
114    }
115}
116
117// K::Err: Void; V::Err: Void; (K, V) => no error possible
118impl<K, V> From<PushIterError<TuplePushError<K, TuplePushError<V, Void>>>> for Void
119where
120    K: Into<Void>,
121    V: Into<Void>,
122{
123    fn from(_: PushIterError<TuplePushError<K, TuplePushError<V, Void>>>) -> Self {
124        unreachable!("no way to create instance of Void")
125    }
126}
127
128////////////////////////////////////////////////////////////////////////////////
129// TableFromIter
130////////////////////////////////////////////////////////////////////////////////
131
132/// A wrapper struct for converting arbitrary iterators into lua tables. Use
133/// this instead of converting the iterator into a `Vec` to avoid unnecessary
134/// allocations
135/// # Example
136/// ```no_run
137/// use std::io::BufRead;
138/// let lua = tlua::Lua::new();
139/// lua.set(
140///     "foo",
141///     tlua::TableFromIter(std::io::stdin().lock().lines().flatten()),
142/// )
143/// // Global variable 'foo' now contains an array of lines read from stdin
144/// ```
145pub struct TableFromIter<I>(pub I);
146
147impl<L, I> PushInto<L> for TableFromIter<I>
148where
149    L: AsLua,
150    I: Iterator,
151    <I as Iterator>::Item: PushInto<LuaState>,
152{
153    type Err = PushIterError<<I::Item as PushInto<LuaState>>::Err>;
154
155    fn push_into_lua(self, lua: L) -> crate::PushIntoResult<L, Self> {
156        push_iter(lua, self.0)
157    }
158}
159
160impl<L, I> PushOneInto<L> for TableFromIter<I>
161where
162    L: AsLua,
163    I: Iterator,
164    <I as Iterator>::Item: PushInto<LuaState>,
165{
166}
167
168////////////////////////////////////////////////////////////////////////////////
169// Vec
170////////////////////////////////////////////////////////////////////////////////
171
172impl<L, T> Push<L> for Vec<T>
173where
174    L: AsLua,
175    T: Push<LuaState>,
176{
177    type Err = PushIterError<T::Err>;
178
179    #[inline]
180    fn push_to_lua(&self, lua: L) -> Result<PushGuard<L>, (Self::Err, L)> {
181        push_iter(lua, self.iter())
182    }
183}
184
185impl<L, T> PushOne<L> for Vec<T>
186where
187    L: AsLua,
188    T: Push<LuaState>,
189{
190}
191
192impl<L, T> PushInto<L> for Vec<T>
193where
194    L: AsLua,
195    T: PushInto<LuaState>,
196{
197    type Err = PushIterError<T::Err>;
198
199    #[inline]
200    fn push_into_lua(self, lua: L) -> Result<PushGuard<L>, (Self::Err, L)> {
201        push_iter(lua, self.into_iter())
202    }
203}
204
205impl<L, T> PushOneInto<L> for Vec<T>
206where
207    L: AsLua,
208    T: PushInto<LuaState>,
209{
210}
211
212impl<L, T> LuaRead<L> for Vec<T>
213where
214    L: AsLua,
215    T: for<'a> LuaRead<PushGuard<&'a LuaTable<L>>>,
216    T: 'static,
217{
218    fn lua_read_at_position(lua: L, index: NonZeroI32) -> ReadResult<Self, L> {
219        // We need this as iteration order isn't guaranteed to match order of
220        // keys, even if they're numeric
221        // https://www.lua.org/manual/5.2/manual.html#pdf-next
222        let table = LuaTable::lua_read_at_position(lua, index)?;
223        let mut dict: BTreeMap<i32, T> = BTreeMap::new();
224
225        let mut max_key = i32::MIN;
226        let mut min_key = i32::MAX;
227
228        {
229            let mut iter = table.iter::<i32, T>();
230            while let Some(maybe_kv) = iter.next() {
231                let (key, value) = crate::unwrap_ok_or! { maybe_kv,
232                    Err(e) => {
233                        drop(iter);
234                        let lua = table.into_inner();
235                        let e = e.when("converting Lua table to Vec<_>")
236                            .expected_type::<Self>();
237                        return Err((lua, e))
238                    }
239                };
240                max_key = max_key.max(key);
241                min_key = min_key.min(key);
242                dict.insert(key, value);
243            }
244        }
245
246        if dict.is_empty() {
247            return Ok(vec![]);
248        }
249
250        if min_key != 1 {
251            // Rust doesn't support sparse arrays or arrays with negative
252            // indices
253            let e = WrongType::info("converting Lua table to Vec<_>")
254                .expected("indexes in range 1..N")
255                .actual(format!("value with index {}", min_key));
256            return Err((table.into_inner(), e));
257        }
258
259        let mut result = Vec::with_capacity(max_key as _);
260
261        // We expect to start with first element of table and have this
262        // be smaller that first key by one
263        let mut previous_key = 0;
264
265        // By this point, we actually iterate the map to move values to Vec
266        // and check that table represented non-sparse 1-indexed array
267        for (k, v) in dict {
268            if previous_key + 1 != k {
269                let e = WrongType::info("converting Lua table to Vec<_>")
270                    .expected("indexes in range 1..N")
271                    .actual(format!("Lua table with missing index {}", previous_key + 1));
272                return Err((table.into_inner(), e));
273            } else {
274                // We just push, thus converting Lua 1-based indexing
275                // to Rust 0-based indexing
276                result.push(v);
277                previous_key = k;
278            }
279        }
280
281        Ok(result)
282    }
283}
284
285////////////////////////////////////////////////////////////////////////////////
286// \[T]
287////////////////////////////////////////////////////////////////////////////////
288
289impl<L, T> Push<L> for [T]
290where
291    L: AsLua,
292    T: Push<LuaState>,
293{
294    type Err = PushIterError<T::Err>;
295
296    #[inline]
297    fn push_to_lua(&self, lua: L) -> Result<PushGuard<L>, (Self::Err, L)> {
298        push_iter(lua, self.iter())
299    }
300}
301
302impl<L, T> PushOne<L> for [T]
303where
304    L: AsLua,
305    T: Push<LuaState>,
306{
307}
308
309////////////////////////////////////////////////////////////////////////////////
310// [T; N]
311////////////////////////////////////////////////////////////////////////////////
312
313impl<L, T, const N: usize> Push<L> for [T; N]
314where
315    L: AsLua,
316    T: Push<LuaState>,
317{
318    type Err = PushIterError<T::Err>;
319
320    #[inline]
321    fn push_to_lua(&self, lua: L) -> Result<PushGuard<L>, (Self::Err, L)> {
322        push_iter(lua, self.iter())
323    }
324}
325
326impl<L, T, const N: usize> PushOne<L> for [T; N]
327where
328    L: AsLua,
329    T: Push<LuaState>,
330{
331}
332
333impl<L, T, const N: usize> PushInto<L> for [T; N]
334where
335    L: AsLua,
336    T: PushInto<LuaState>,
337{
338    type Err = PushIterError<T::Err>;
339
340    #[inline]
341    fn push_into_lua(self, lua: L) -> Result<PushGuard<L>, (Self::Err, L)> {
342        push_iter(lua, IntoIterator::into_iter(self))
343    }
344}
345
346impl<L, T, const N: usize> PushOneInto<L> for [T; N]
347where
348    L: AsLua,
349    T: PushInto<LuaState>,
350{
351}
352
353impl<L, T, const N: usize> LuaRead<L> for [T; N]
354where
355    L: AsLua,
356    T: for<'a> LuaRead<PushGuard<&'a LuaTable<L>>>,
357    T: 'static,
358{
359    fn lua_read_at_position(lua: L, index: NonZeroI32) -> ReadResult<Self, L> {
360        let table = LuaTable::lua_read_at_position(lua, index)?;
361        let mut res = std::mem::MaybeUninit::uninit();
362        let ptr = &mut res as *mut _ as *mut [T; N] as *mut T;
363        let mut was_assigned = [false; N];
364        let mut err = None;
365
366        for maybe_kv in table.iter::<i32, T>() {
367            match maybe_kv {
368                Ok((key, value)) if 1 <= key && key as usize <= N => {
369                    let i = (key - 1) as usize;
370                    unsafe { std::ptr::write(ptr.add(i), value) }
371                    was_assigned[i] = true;
372                }
373                Err(e) => {
374                    err = Some(Error::Subtype(e));
375                    break;
376                }
377                Ok((index, _)) => {
378                    err = Some(Error::WrongIndex(index));
379                    break;
380                }
381            }
382        }
383
384        if err.is_none() {
385            err = was_assigned
386                .iter()
387                .zip(1..)
388                .find(|(&was_assigned, _)| !was_assigned)
389                .map(|(_, i)| Error::MissingIndex(i));
390        }
391
392        let err = crate::unwrap_or! { err,
393            return Ok(unsafe { res.assume_init() });
394        };
395
396        for i in IntoIterator::into_iter(was_assigned)
397            .enumerate()
398            .flat_map(|(i, was_assigned)| was_assigned.then_some(i))
399        {
400            unsafe { std::ptr::drop_in_place(ptr.add(i)) }
401        }
402
403        let when = "converting Lua table to array";
404        let e = match err {
405            Error::Subtype(err) => err.when(when).expected_type::<Self>(),
406            Error::WrongIndex(index) => WrongType::info(when)
407                .expected(format!("indexes in range 1..={}", N))
408                .actual(format!("value with index {}", index)),
409            Error::MissingIndex(index) => WrongType::info(when)
410                .expected(format!("indexes in range 1..={}", N))
411                .actual(format!("Lua table with missing index {}", index)),
412        };
413        return Err((table.into_inner(), e));
414
415        enum Error {
416            Subtype(WrongType),
417            WrongIndex(i32),
418            MissingIndex(i32),
419        }
420    }
421}
422
423////////////////////////////////////////////////////////////////////////////////
424// HashMap
425////////////////////////////////////////////////////////////////////////////////
426
427impl<L, K, V, S> LuaRead<L> for HashMap<K, V, S>
428where
429    L: AsLua,
430    K: 'static + Hash + Eq,
431    K: for<'k> LuaRead<&'k LuaTable<L>>,
432    V: 'static,
433    V: for<'v> LuaRead<PushGuard<&'v LuaTable<L>>>,
434    S: Default,
435    S: std::hash::BuildHasher,
436{
437    fn lua_read_at_position(lua: L, index: NonZeroI32) -> ReadResult<Self, L> {
438        let table = LuaTable::lua_read_at_position(lua, index)?;
439        let res: Result<_, _> = table.iter().collect();
440        res.map_err(|err| {
441            let l = table.into_inner();
442            let e = err
443                .when("converting Lua table to HashMap<_, _>")
444                .expected_type::<Self>();
445            (l, e)
446        })
447    }
448}
449
450macro_rules! push_hashmap_impl {
451    ($self:expr, $lua:expr) => {
452        push_iter($lua, $self.into_iter()).map_err(|(e, lua)| match e {
453            PushIterError::TooManyValues(_) => unreachable!("K and V implement PushOne"),
454            PushIterError::ValuePushError(First(e)) => (First(e), lua),
455            PushIterError::ValuePushError(Other(e)) => (Other(e.first()), lua),
456        })
457    };
458}
459
460impl<L, K, V, S> Push<L> for HashMap<K, V, S>
461where
462    L: AsLua,
463    K: PushOne<LuaState> + Eq + Hash + Debug,
464    V: PushOne<LuaState> + Debug,
465{
466    type Err = TuplePushError<K::Err, V::Err>;
467
468    #[inline]
469    fn push_to_lua(&self, lua: L) -> Result<PushGuard<L>, (Self::Err, L)> {
470        push_hashmap_impl!(self, lua)
471    }
472}
473
474impl<L, K, V, S> PushOne<L> for HashMap<K, V, S>
475where
476    L: AsLua,
477    K: PushOne<LuaState> + Eq + Hash + Debug,
478    V: PushOne<LuaState> + Debug,
479{
480}
481
482impl<L, K, V, S> PushInto<L> for HashMap<K, V, S>
483where
484    L: AsLua,
485    K: PushOneInto<LuaState> + Eq + Hash + Debug,
486    V: PushOneInto<LuaState> + Debug,
487{
488    type Err = TuplePushError<K::Err, V::Err>;
489
490    #[inline]
491    fn push_into_lua(self, lua: L) -> Result<PushGuard<L>, (Self::Err, L)> {
492        push_hashmap_impl!(self, lua)
493    }
494}
495
496impl<L, K, V, S> PushOneInto<L> for HashMap<K, V, S>
497where
498    L: AsLua,
499    K: PushOneInto<LuaState> + Eq + Hash + Debug,
500    V: PushOneInto<LuaState> + Debug,
501{
502}
503
504////////////////////////////////////////////////////////////////////////////////
505// HashSet
506////////////////////////////////////////////////////////////////////////////////
507
508macro_rules! push_hashset_impl {
509    ($self:expr, $lua:expr) => {
510        push_iter($lua, $self.into_iter().zip(iter::repeat(true))).map_err(|(e, lua)| match e {
511            PushIterError::TooManyValues(_) => unreachable!("K implements PushOne"),
512            PushIterError::ValuePushError(First(e)) => (e, lua),
513            PushIterError::ValuePushError(Other(_)) => {
514                unreachable!("no way to create instance of Void")
515            }
516        })
517    };
518}
519
520impl<L, K, S> Push<L> for HashSet<K, S>
521where
522    L: AsLua,
523    K: PushOne<LuaState> + Eq + Hash + Debug,
524{
525    type Err = K::Err;
526
527    #[inline]
528    fn push_to_lua(&self, lua: L) -> Result<PushGuard<L>, (K::Err, L)> {
529        push_hashset_impl!(self, lua)
530    }
531}
532
533impl<L, K, S> PushOne<L> for HashSet<K, S>
534where
535    L: AsLua,
536    K: PushOne<LuaState> + Eq + Hash + Debug,
537{
538}
539
540impl<L, K, S> PushInto<L> for HashSet<K, S>
541where
542    L: AsLua,
543    K: PushOneInto<LuaState> + Eq + Hash + Debug,
544{
545    type Err = K::Err;
546
547    #[inline]
548    fn push_into_lua(self, lua: L) -> Result<PushGuard<L>, (K::Err, L)> {
549        push_hashset_impl!(self, lua)
550    }
551}
552
553impl<L, K, S> PushOneInto<L> for HashSet<K, S>
554where
555    L: AsLua,
556    K: PushOneInto<LuaState> + Eq + Hash + Debug,
557{
558}