refactor(json): use UserData-based OrderedTable for json_encode/decode

Cargo.lock

@@ -3474,6 +3474,7 @@ dependencies = [
  "parking_lot",
  "paste",
  "ratatui",
+ "serde",
  "serde_json",
  "syntect",
  "tokio",

yazi-plugin/Cargo.toml

@@ -43,6 +43,7 @@ tokio-util    = { workspace = true }
 tracing       = { workspace = true }
 twox-hash     = { workspace = true }
 unicode-width = { workspace = true }
+serde        = { workspace = true }
 yazi-prebuilt = "0.1.0"
 
 [target."cfg(unix)".dependencies]

yazi-plugin/src/utils/json.rs

@@ -1,126 +1,240 @@
-use mlua::{Function, IntoLuaMulti, Lua, LuaSerdeExt, MetaMethod, Result, Table, Value};
-use yazi_binding::Error;
+use mlua::{
+	Function, IntoLuaMulti, Lua, LuaSerdeExt, MetaMethod, UserData, UserDataMethods, Value,
+};
 
-use std::cell::RefCell;
+use serde::Serialize;
+use serde_json::value::Value as JsonValue;
+
+use std::cell::{Ref, RefCell, RefMut};
 use std::rc::Rc;
 
+use yazi_binding::Error;
+
 use super::Utils;
 use crate::config::OPTS;
 
-struct OrderedTable<'lua> {
-	metatable: Table,
-	lua: &'lua Lua,
+#[derive(Clone)]
+enum OrderedTableIndex {
+	Key(String),
+	Index(usize),
 }
 
-impl<'lua> OrderedTable<'lua> {
-	const DATA_KEY: &'static str = "__data";
-	const ORDERED_KEY: &'static str = "__ordered";
+#[derive(Clone)]
+struct OrderedTable {
+	data: Rc<RefCell<JsonValue>>,
+	path: Vec<OrderedTableIndex>,
+}
 
-	fn new(lua: &'lua Lua) -> Result<Self> {
-		let mt = lua.create_table()?;
+impl Serialize for OrderedTable {
+	fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+	where
+		S: serde::Serializer,
+	{
+		use OrderedTableIndex::*;
+		let current = self.data.borrow();
 
-		mt.set(
-			MetaMethod::Index.name(),
-			lua.create_function(|_, (table, key): (Table, Value)| {
-				let data: Table = table.get(Self::DATA_KEY)?;
-				data.get::<Value>(key)
-			})?,
-		)?;
+		let mut value = &*current;
+		for index in &self.path {
+			match index {
+				Key(key) => {
+					value = value
+						.get(key)
+						.ok_or_else(|| serde::ser::Error::custom(format!("Key '{}' not found", key)))?;
+				}
+				Index(i) => {
+					value = value
+						.get(*i)
+						.ok_or_else(|| serde::ser::Error::custom(format!("Index {} out of bounds", i)))?;
+				}
+			}
+		}
 
-		mt.set(
-			MetaMethod::Pairs.name(),
-			lua.create_function(|lua, table: Table| {
-				let ordered: Value = table.get(Self::ORDERED_KEY)?;
-				let data: Table = table.get(Self::DATA_KEY)?;
+		value.serialize(serializer)
+	}
+}
 
-				let idx = Rc::new(RefCell::new(0));
-
-				let iter = lua.create_function_mut(move |_, ()| {
-					let mut index = idx.borrow_mut();
-
-					*index += 1;
-
-					let key: Value = match &ordered {
-						Value::Table(ordered) => ordered.get(*index)?,
-						_ => {
-							if data.len()? >= *index {
-								Value::Integer(*index)
-							} else {
-								Value::Nil
-							}
-						}
-					};
-
-					if key == Value::Nil {
-						Ok((Value::Nil, Value::Nil))
-					} else {
-						let val = data.get(key.clone())?;
-						Ok((key, val))
-					}
-				})?;
-
-				Ok((iter, Value::Nil, Value::Nil))
-			})?,
-		)?;
-
-		mt.set(
-			MetaMethod::Len.name(),
-			lua.create_function(|_, table: Table| {
-				let data: Table = table.get(Self::DATA_KEY)?;
-
-				data.len()
-			})?,
-		)?;
-
-		Ok(Self { metatable: mt, lua })
+impl OrderedTable {
+	fn new(data: Rc<RefCell<JsonValue>>, path: Vec<OrderedTableIndex>) -> Self {
+		Self { data, path }
 	}
 
-	fn wrap(&self, value: &serde_json::Value) -> Result<Value> {
-		let data_table = self.lua.create_table()?;
-		let ordered_table = self.lua.create_table()?;
-		let wrapper = self.lua.create_table()?;
+	fn get_current(&self) -> Option<Ref<JsonValue>> {
+		let root = self.data.borrow();
 
-		wrapper.set_metatable(Some(self.metatable.clone()));
-
-		match &value {
-			serde_json::Value::Object(obj) => {
-				let mut idx = 1;
-
-				for (k, v) in obj {
-					let key_str = self.lua.create_string(k)?;
-					let val = self.wrap(v)?;
-
-					data_table.set(key_str.clone(), val)?;
-					ordered_table.set(idx, key_str)?;
-					idx += 1;
-				}
-
-				wrapper.set(Self::DATA_KEY, data_table)?;
-				wrapper.set(Self::ORDERED_KEY, ordered_table)?;
-
-				Ok(Value::Table(wrapper))
-			}
-			serde_json::Value::Array(arr) => {
-				for (i, v) in arr.iter().enumerate() {
-					data_table.set(i + 1, self.wrap(v)?)?;
-				}
-
-				wrapper.set(Self::DATA_KEY, data_table)?;
-				wrapper.set(Self::ORDERED_KEY, Value::Nil)?;
-
-				Ok(Value::Table(wrapper))
-			}
-			_ => self.lua.to_value_with(value, OPTS),
+		let mut tmp = &*root;
+		for comp in &self.path {
+			tmp = match comp {
+				OrderedTableIndex::Key(k) => tmp.get(k)?,
+				OrderedTableIndex::Index(i) => tmp.get(*i)?,
+			};
 		}
+
+		Some(Ref::map(root, move |root_value| {
+			let mut current = root_value;
+			for comp in &self.path {
+				current = match comp {
+					OrderedTableIndex::Key(k) => current.get(k).unwrap(),
+					OrderedTableIndex::Index(i) => current.get(*i).unwrap(),
+				};
+			}
+			current
+		}))
+	}
+
+	fn get_current_mut(&self) -> Option<RefMut<JsonValue>> {
+		let root = self.data.borrow_mut();
+
+		let mut tmp = &*root;
+		for comp in &self.path {
+			tmp = match comp {
+				OrderedTableIndex::Key(k) => tmp.get(k)?,
+				OrderedTableIndex::Index(i) => tmp.get(*i)?,
+			};
+		}
+
+		Some(RefMut::map(root, move |root_value| {
+			let mut current = root_value;
+			for comp in &self.path {
+				current = match comp {
+					OrderedTableIndex::Key(k) => current.get_mut(k).unwrap(),
+					OrderedTableIndex::Index(i) => current.get_mut(*i).unwrap(),
+				};
+			}
+			current
+		}))
+	}
+}
+
+impl UserData for OrderedTable {
+	fn add_fields<F: mlua::UserDataFields<Self>>(fields: &mut F) {
+		fields.add_meta_field("__ordered", true);
+	}
+
+	fn add_methods<M: UserDataMethods<Self>>(methods: &mut M) {
+		methods.add_meta_method(MetaMethod::Index, |lua, this: &OrderedTable, key: Value| {
+			let new_path = match key {
+				Value::String(s) => OrderedTableIndex::Key(s.to_str()?.to_string()),
+				Value::Integer(i) => OrderedTableIndex::Index((i - 1) as usize),
+				_ => return Ok(Value::Nil),
+			};
+
+			let sub = OrderedTable::new(Rc::clone(&this.data), {
+				let mut p = this.path.clone();
+				p.push(new_path);
+				p
+			});
+
+			if let Some(v) = sub.clone().get_current() {
+				if matches!(*v, JsonValue::Object(_) | JsonValue::Array(_)) {
+					Ok(Value::UserData(lua.create_userdata(sub)?))
+				} else {
+					lua.to_value_with(&*v, OPTS)
+				}
+			} else {
+				Ok(Value::Nil)
+			}
+		});
+
+		methods.add_meta_method_mut(MetaMethod::Pairs, |lua, this: &mut OrderedTable, ()| {
+			let curr =
+				this.get_current_mut().ok_or_else(|| mlua::Error::RuntimeError("invalid path".into()))?;
+
+			let keys = match &*curr {
+				JsonValue::Object(obj) => {
+					obj.keys().map(|k| (Value::String(lua.create_string(k).unwrap()))).collect()
+				}
+				JsonValue::Array(arry) => (1..=arry.len()).map(|i| Value::Integer(i as i64)).collect(),
+				_ => vec![],
+			};
+			let path = this.path.clone();
+			let data = Rc::clone(&this.data);
+			let idx = Rc::new(RefCell::new(keys.into_iter()));
+			let iter = lua.create_function_mut(move |lua, ()| {
+				let mut index = idx.borrow_mut();
+				if let Some(key) = index.next() {
+					let new_path = match &key {
+						Value::String(s) => OrderedTableIndex::Key(s.to_str()?.to_string()),
+						Value::Integer(i) => OrderedTableIndex::Index((i - 1) as usize),
+						_ => return Ok((Value::Nil, Value::Nil)),
+					};
+
+					Ok((
+						key.clone(),
+						Value::UserData(lua.create_userdata(OrderedTable::new(Rc::clone(&data), {
+							let mut p = path.clone();
+							p.push(new_path);
+							p
+						}))?),
+					))
+				} else {
+					Ok((Value::Nil, Value::Nil))
+				}
+			})?;
+
+			Ok((iter, Value::Nil, Value::Nil))
+		});
+
+		methods.add_meta_method(MetaMethod::Len, |_, this: &OrderedTable, ()| {
+			Ok(if let Some(data) = this.get_current() {
+				Value::Integer(match &*data {
+					JsonValue::Array(arr) => arr.len() as i64,
+					_ => 0,
+				})
+			} else {
+				Value::Nil
+			})
+		});
+
+		methods.add_meta_method_mut(
+			MetaMethod::NewIndex,
+			|_, this: &mut OrderedTable, (key, value): (Value, Value)| {
+				let value = serde_json::to_value(value).map_err(mlua::Error::external)?;
+				let mut curr =
+					this.get_current_mut().ok_or_else(|| mlua::Error::RuntimeError("invalid path".into()))?;
+
+				match &mut *curr {
+					JsonValue::Object(obj) => {
+						let k = key
+							.as_str()
+							.ok_or_else(|| mlua::Error::RuntimeError("object key must be string".into()))?;
+						obj.insert(k.to_string(), value);
+					}
+					JsonValue::Array(arr) => {
+						let i = key
+							.as_integer()
+							.ok_or_else(|| mlua::Error::RuntimeError("array index must be integer".into()))?;
+						let idx = (i - 1) as usize;
+						if idx < arr.len() {
+							arr[idx] = value;
+						}
+					}
+					_ => {
+						return Err(mlua::Error::RuntimeError("not an object or array".into()));
+					}
+				}
+
+				Ok(())
+			},
+		);
 	}
 }
 
 impl Utils {
 	pub(super) fn json_encode(lua: &Lua) -> mlua::Result<Function> {
 		lua.create_async_function(|lua, value: Value| async move {
-			match serde_json::to_string(&value) {
-				Ok(s) => (s, Value::Nil).into_lua_multi(&lua),
-				Err(e) => (Value::Nil, Error::Serde(e)).into_lua_multi(&lua),
+			let result = match value {
+				Value::UserData(ud) => {
+					let table = ud
+						.borrow::<OrderedTable>()
+						.map_err(|_| mlua::Error::RuntimeError("Unknown userdata type".into()))?;
+					serde_json::to_string(&*table).map_err(Error::Serde)
+				}
+				_ => serde_json::to_string(&value).map_err(Error::Serde),
+			};
+
+			match result {
+				Ok(s) => Ok((s, Value::Nil).into_lua_multi(&lua)?),
+				Err(e) => Ok((Value::Nil, e).into_lua_multi(&lua)?),
 			}
 		})
 	}
@@ -129,11 +243,12 @@ impl Utils {
 		lua.create_async_function(|lua, (s, ordered): (mlua::String, Option<Value>)| async move {
 			let ordered = ordered.unwrap_or(Value::Boolean(false));
 
-			match serde_json::from_slice::<serde_json::Value>(&s.as_bytes()) {
+			match serde_json::from_slice::<JsonValue>(&s.as_bytes()) {
 				Ok(v) => (
 					if let Value::Boolean(true) = ordered {
-						let ordered_table = OrderedTable::new(&lua)?;
-						ordered_table.wrap(&v)?
+						let data = Rc::new(RefCell::new(v));
+
+						Value::UserData(lua.create_userdata(OrderedTable::new(data, vec![]))?)
 					} else {
 						lua.to_value_with(&v, OPTS)?
 					},