Optimize sorted serde table traversal

Author: khvzak

benches/serde.rs

@@ -1,6 +1,7 @@
+use std::collections::BTreeMap;
 use std::time::Duration;
 
-use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
+use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main};
 
 use mlua::prelude::*;
 
@@ -76,6 +77,29 @@ fn decode_json(c: &mut Criterion) {
     });
 }
 
+fn sorted_map_traversal(c: &mut Criterion) {
+    let lua = Lua::new();
+    let mut group = c.benchmark_group("sorted map traversal");
+    for size in [0, 1, 2, 3, 4, 7, 8, 64, 4096] {
+        let table = lua
+            .create_table_from((0..size).map(|i| (format!("key_{i:08}"), i)))
+            .unwrap();
+        let value = LuaValue::Table(table);
+        group.bench_function(BenchmarkId::new("serialize", size), |b| {
+            b.iter(|| serde_json::to_string(&value.to_serializable().sort_keys(true)).unwrap());
+        });
+        group.bench_function(BenchmarkId::new("from_value", size), |b| {
+            b.iter(|| {
+                lua.from_value_with::<BTreeMap<String, i64>>(
+                    value.clone(),
+                    LuaDeserializeOptions::new().sort_keys(true),
+                )
+                .unwrap()
+            });
+        });
+    }
+}
+
 criterion_group! {
     name = benches;
     config = Criterion::default()
@@ -85,6 +109,7 @@ criterion_group! {
     targets =
         encode_json,
         decode_json,
+        sorted_map_traversal,
 }
 
 criterion_main!(benches);

src/serde/de.rs

@@ -518,7 +518,7 @@ pub(crate) enum MapPairs<'a> {
 impl<'a> MapPairs<'a> {
     pub(crate) fn new(t: &'a Table, sort_keys: bool) -> Result<Self> {
         if sort_keys {
-            let mut pairs = t.pairs::<Value, Value>().collect::<Result<Vec<_>>>()?;
+            let mut pairs = t.collect_pairs()?;
             pairs.sort_by(|(a, _), (b, _)| b.sort_cmp(a)); // reverse order as we pop values from the end
             Ok(MapPairs::Vec(pairs))
         } else {

src/table.rs

@@ -1068,6 +1068,43 @@ impl Table {
         None
     }
 
+    #[cfg(feature = "serde")]
+    pub(crate) fn collect_pairs(&self) -> Result<Vec<(Value, Value)>> {
+        let mut pairs = Vec::new();
+
+        #[cfg(not(feature = "luau"))]
+        unsafe {
+            const LIMIT: c_int = 8;
+            let lua = self.0.lua.lock();
+            let state = lua.state();
+            let _sg = StackGuard::new(state);
+            check_stack(state, 2 * LIMIT + 2)?;
+
+            lua.push_ref(&self.0);
+            let table_index = ffi::lua_gettop(state);
+            ffi::lua_pushnil(state);
+            // Finish small traversals before conversions or allocations can reenter Lua
+            for count in 0..LIMIT {
+                if ffi::lua_next(state, table_index) == 0 {
+                    pairs.reserve_exact(count as usize);
+                    for i in 0..count {
+                        let index = table_index + 1 + 2 * i;
+                        let key = lua.try_stack_value(index, None)?;
+                        pairs.push((key, lua.try_stack_value(index + 1, None)?));
+                    }
+                    return Ok(pairs);
+                }
+                ffi::lua_pushvalue(state, -2);
+            }
+        }
+
+        self.for_each(|key, value| {
+            pairs.push((key, value));
+            Ok(())
+        })?;
+        Ok(pairs)
+    }
+
     #[cfg(not(feature = "luau"))]
     #[inline]
     unsafe fn next(state: *mut ffi::lua_State) -> Result<bool> {