yazi/yazi-plugin/src/utils/sync.rs

use anyhow::Context;
use futures::future::join_all;
use mlua::{ExternalError, ExternalResult, Function, IntoLuaMulti, Lua, MultiValue, Table, Value, Variadic};
use tokio::sync::mpsc;
use yazi_binding::{Handle, MpscRx, MpscTx, MpscUnboundedRx, MpscUnboundedTx, OneshotRx, OneshotTx, runtime, runtime_mut};
use yazi_dds::Sendable;
use yazi_parser::app::PluginOpt;
use yazi_proxy::AppProxy;
use yazi_shared::{LOCAL_SET, data::Data};

use super::Utils;
use crate::loader::LOADER;

impl Utils {
	pub(super) fn co(lua: &Lua) -> mlua::Result<Function> {
		lua.create_function(|lua, f: Function| {
			let thread = lua.create_thread(f)?;
			lua.create_async_function(move |lua, mut args: MultiValue| {
				let thread = thread.clone();
				async move {
					loop {
						let values: MultiValue = thread.resume(args)?;
						if let Some(Value::LightUserData(ud)) = values.get(0)
							&& *ud == Lua::poll_pending()
						{
							args = lua.yield_with(values).await?;
						} else {
							return Ok(values);
						}
					}
				}
			})
		})
	}

	pub(super) fn sync(lua: &Lua) -> mlua::Result<Function> {
		lua.create_function(|lua, f: Function| {
			let mut rt = runtime_mut!(lua)?;
			let Some(block) = rt.put_block(&f) else {
				return Err("`ya.sync()` must be called in a plugin").into_lua_err();
			};

			let current = rt.current_owned()?;
			lua.create_async_function(move |lua, mut args: MultiValue| {
				let (f, current) = (f.clone(), current.clone());
				async move {
					let blocking = runtime!(lua)?.blocking;
					if blocking {
						args.push_front(Value::Table(LOADER.try_load(&lua, &current)?));
						f.call::<MultiValue>(args)
					} else {
						Self::retrieve(&lua, &current, block, args)
							.await
							.and_then(|data| Sendable::list_to_values(&lua, data))
							.with_context(|| {
								format!("Failed to execute sync block-{block} in `{current}` plugin")
							})
							.into_lua_err()
					}
				}
			})
		})
	}

	pub(super) fn r#async(lua: &Lua, isolate: bool) -> mlua::Result<Function> {
		if isolate {
			lua.create_function(|_, _: Function| {
				Err::<(), _>("`ya.async()` can only be used in sync context at the moment".into_lua_err())
			})
		} else {
			lua.create_function(|lua, (f, args): (Function, MultiValue)| {
				let name = runtime!(lua)?.current_owned()?;
				let lua = lua.clone();

				Ok(Handle::AsyncFn(LOCAL_SET.spawn_local(async move {
					let blocking = runtime_mut!(lua)?.critical_push(&name, false);
					let result = f.call_async::<MultiValue>(args).await;
					runtime_mut!(lua)?.critical_pop(blocking)?;

					if let Err(ref e) = result {
						match name.as_str() {
							"init" => tracing::error!("Failed to execute async block in `init.lua`: {e}"),
							s => tracing::error!("Failed to execute async block in `{s}` plugin: {e}"),
						}
					}
					result
				})))
			})
		}
	}

	pub(super) fn chan(lua: &Lua) -> mlua::Result<Function> {
		lua.create_function(|lua, (r#type, buffer): (mlua::String, Option<usize>)| {
			match (&*r#type.as_bytes(), buffer) {
				(b"mpsc", Some(buffer)) if buffer < 1 => {
					Err("Buffer size must be greater than 0".into_lua_err())
				}
				(b"mpsc", Some(buffer)) => {
					let (tx, rx) = tokio::sync::mpsc::channel::<Value>(buffer);
					(MpscTx(tx), MpscRx(rx)).into_lua_multi(lua)
				}
				(b"mpsc", None) => {
					let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Value>();
					(MpscUnboundedTx(tx), MpscUnboundedRx(rx)).into_lua_multi(lua)
				}
				(b"oneshot", _) => {
					let (tx, rx) = tokio::sync::oneshot::channel::<Value>();
					(OneshotTx(Some(tx)), OneshotRx(Some(rx))).into_lua_multi(lua)
				}
				_ => Err("Channel type must be `mpsc` or `oneshot`".into_lua_err()),
			}
		})
	}

	pub(super) fn join(lua: &Lua) -> mlua::Result<Function> {
		lua.create_async_function(|_, fns: Variadic<Function>| async move {
			let mut results = MultiValue::with_capacity(fns.len());
			for r in join_all(fns.into_iter().map(|f| f.call_async::<MultiValue>(()))).await {
				results.extend(r?);
			}
			Ok(results)
		})
	}

	// TODO
	pub(super) fn select(lua: &Lua) -> mlua::Result<Function> {
		lua.create_async_function(|_lua, _futs: MultiValue| async move { Ok(()) })
	}

	async fn retrieve(
		lua: &Lua,
		id: &str,
		calls: usize,
		args: MultiValue,
	) -> mlua::Result<Vec<Data>> {
		let args = Sendable::values_to_list(lua, args)?;
		let (tx, mut rx) = mpsc::channel::<Vec<Data>>(1);

		let id_ = id.to_owned();
		let callback = move |lua: &Lua, plugin: Table| {
			let Some(block) = runtime!(lua)?.get_block(&id_, calls) else {
				return Err("sync block not found".into_lua_err());
			};

			let args = [Ok(Value::Table(plugin))]
				.into_iter()
				.chain(args.into_iter().map(|d| Sendable::data_to_value(lua, d)))
				.collect::<mlua::Result<MultiValue>>()?;

			let values = Sendable::values_to_list(lua, block.call(args)?)?;
			tx.try_send(values).map_err(|_| "send failed".into_lua_err())
		};

		AppProxy::plugin(PluginOpt::new_callback(id.to_owned(), callback));

		rx.recv().await.ok_or("recv failed").into_lua_err()
	}
}