use alloc::collections::{BTreeMap, BTreeSet};
use alloc::vec::Vec;
use core::any::Any;
use core::cell::RefCell;
use std::boxed::Box;
use std::thread_local;
use crate::encode::BinaryDecode;
use crate::id_allocator::{BorrowIds, HeapIds, IdSlab, InstallIdBatch};
use crate::ipc::DecodedData;
use crate::ipc::{EncodedData, EncodedParts, IPCMessage, MessageType};
use crate::object_store::ObjectHandle;
use crate::runtime::WryIPC;
use crate::type_cache::TypeCache;
use crate::wire::BinaryEncode as AbiBinaryEncode;
use crate::wire::{JsFunctionSpec, JsRef, TypeDef};
#[derive(Default)]
pub(crate) struct OperationFreeFrame {
heap_ids: Vec<u64>,
object_handles: Vec<u32>,
}
pub struct Runtime {
encoder: EncodedParts,
encoder_has_pending_ops: bool,
heap_ids: HeapIds,
borrow_ids: BorrowIds,
object_handles: IdSlab<u32>,
is_batching: bool,
type_cache: TypeCache,
objects: BTreeMap<u32, Box<dyn Any>>,
pending_object_drops: BTreeSet<u32>,
op_free_stack: Vec<OperationFreeFrame>,
heap_ids_to_recycle_after_flush: Vec<JsRef>,
pending_type_ids: Vec<u32>,
ipc: WryIPC,
thread_locals: BTreeMap<*const (), Box<dyn Any>>,
}
impl Runtime {
pub(crate) fn new(ipc: WryIPC) -> Self {
Self {
encoder: EncodedParts::default(),
encoder_has_pending_ops: false,
heap_ids: HeapIds::new(),
borrow_ids: BorrowIds::new(),
object_handles: IdSlab::new(1),
is_batching: false,
type_cache: TypeCache::new(),
objects: BTreeMap::new(),
pending_object_drops: BTreeSet::new(),
op_free_stack: Vec::new(),
heap_ids_to_recycle_after_flush: Vec::new(),
pending_type_ids: Vec::new(),
ipc,
thread_locals: BTreeMap::new(),
}
}
pub(crate) fn ipc(&self) -> &WryIPC {
&self.ipc
}
pub(crate) fn get_next_placeholder_id(&mut self) -> u64 {
self.heap_ids.next_placeholder_id()
}
pub(crate) fn get_next_inbound_js_heap_id(&mut self) -> u64 {
self.heap_ids.next_inbound_js_heap_id()
}
pub(crate) fn get_next_borrow_id(&mut self) -> u64 {
self.borrow_ids.next_borrow_id()
}
pub(crate) fn push_borrow_frame(&mut self) {
self.borrow_ids.push_frame();
}
pub(crate) fn pop_borrow_frame(&mut self) {
self.borrow_ids.pop_frame();
}
pub(crate) fn release_heap_id(&mut self, id: u64) -> Option<u64> {
self.heap_ids.release_heap_slot(id);
match self.op_free_stack.last_mut() {
Some(frame) => {
frame.heap_ids.push(id);
None
}
None => Some(id),
}
}
pub(crate) fn recycle_heap_id(&mut self, id: u64) {
self.heap_ids.recycle_heap_id(id);
}
pub(crate) fn recycle_heap_id_if_released(&mut self, id: u64) -> bool {
self.heap_ids.recycle_heap_id_if_released(id)
}
pub(crate) fn defer_heap_id_recycle_until_flush(&mut self, id: u64) {
self.heap_ids_to_recycle_after_flush
.push(JsRef::from_raw(id));
}
pub(crate) fn take_message(&mut self) -> (IPCMessage, Vec<JsRef>) {
let reserved_ids = self
.take_reserved_placeholder_ids()
.into_iter()
.map(JsRef::from_raw)
.collect::<Vec<_>>();
let encoder = self.take_encoder();
let heap_ids_to_recycle_after_flush =
core::mem::take(&mut self.heap_ids_to_recycle_after_flush);
(
self.finish_rust_to_js_message(MessageType::Evaluate, encoder, Some(&reserved_ids)),
heap_ids_to_recycle_after_flush,
)
}
pub(crate) fn finish_respond_message(&mut self, encoder: EncodedData) -> IPCMessage {
self.finish_rust_to_js_message(
MessageType::Respond,
EncodedParts::from_encoded(encoder),
None,
)
}
fn finish_rust_to_js_message(
&mut self,
message_type: MessageType,
encoder: EncodedParts,
reserved_ids: Option<&[JsRef]>,
) -> IPCMessage {
let install_ids = self
.take_pending_install_ids()
.into_iter()
.map(JsRef::from_raw)
.collect::<Vec<_>>();
let mut prelude = Vec::new();
push_ref_list(&mut prelude, &install_ids);
if let Some(reserved_ids) = reserved_ids {
push_ref_list(&mut prelude, reserved_ids);
}
let pending_type_ids = core::mem::take(&mut self.pending_type_ids);
if reserved_ids.is_some() {
self.type_cache.push_pending_frame(pending_type_ids);
} else {
self.type_cache.ack_type_ids(&pending_type_ids);
}
encoder.into_message(message_type, &prelude)
}
pub(crate) fn is_empty(&self) -> bool {
!self.encoder_has_pending_ops
}
pub(crate) fn push_operation_frame(&mut self) {
self.op_free_stack.push(OperationFreeFrame::default());
}
pub(crate) fn release_object_handle(&mut self, handle: ObjectHandle) -> Option<Box<dyn Any>> {
match self.op_free_stack.last_mut() {
Some(frame) => {
frame.object_handles.push(handle.raw());
None
}
None => self.remove_object_untyped(handle),
}
}
pub(crate) fn pop_operation_frame(&mut self) -> OperationFreeFrame {
let frame = self
.op_free_stack
.pop()
.expect("pop_operation_frame called with empty frame stack");
if let Some(parent) = self.op_free_stack.last_mut() {
parent.heap_ids.extend(frame.heap_ids);
parent.object_handles.extend(frame.object_handles);
OperationFreeFrame::default()
} else {
frame
}
}
pub(crate) fn set_batching(&mut self, batching: bool) {
self.is_batching = batching;
}
pub(crate) fn is_batching(&self) -> bool {
self.is_batching
}
pub(crate) fn take_pending_install_ids(&mut self) -> InstallIdBatch {
self.heap_ids.take_pending_install_ids()
}
pub(crate) fn take_reserved_placeholder_ids(&mut self) -> Vec<u64> {
self.heap_ids.take_reserved_placeholder_ids()
}
pub(crate) fn take_encoder(&mut self) -> EncodedParts {
let next = EncodedParts::default();
self.encoder_has_pending_ops = false;
core::mem::replace(&mut self.encoder, next)
}
pub(crate) fn extend_encoder(&mut self, other: EncodedData) {
self.encoder.append_encoded(other);
self.encoder_has_pending_ops = true;
}
pub fn get_or_create_type_id(&mut self, type_def: &TypeDef) -> (u32, bool) {
let (id, can_use_cached) = self.type_cache.get_or_create_type_id(type_def);
if !can_use_cached {
self.pending_type_ids.push(id);
}
(id, can_use_cached)
}
pub(crate) fn pop_and_ack_type_cache_frame(&mut self) {
self.type_cache.pop_and_ack_pending_frame();
}
pub(crate) fn get_object<T: 'static>(&self, handle: u32) -> &T {
let boxed = self.objects.get(&handle).expect("invalid handle");
boxed.downcast_ref::<T>().expect("type mismatch")
}
pub fn remove_object_untyped(&mut self, handle: ObjectHandle) -> Option<Box<dyn Any>> {
let handle = handle.raw();
let object = self.objects.remove(&handle);
if object.is_some() {
self.object_handles.free(handle);
} else if self.object_handles.contains(handle) {
self.pending_object_drops.insert(handle);
}
object
}
}
fn push_ref_list(buf: &mut Vec<u32>, refs: &[JsRef]) {
buf.push(refs.len() as u32);
for js_ref in refs {
let id = js_ref.raw();
buf.push((id & 0xFFFF_FFFF) as u32);
buf.push((id >> 32) as u32);
}
}
impl Runtime {
pub fn resolve_function(&mut self, spec: JsFunctionSpec) -> u32 {
crate::function_registry::FUNCTION_REGISTRY
.resolve_function(spec)
.unwrap_or_else(|| panic!("Function not found for code: {}", spec.render_js_code()))
}
pub fn insert_object_box(&mut self, obj: Box<dyn Any>) -> ObjectHandle {
let handle = self.object_handles.alloc();
self.objects.insert(handle, obj);
ObjectHandle::from_raw(handle)
}
pub fn take_object_box(&mut self, handle: ObjectHandle) -> Option<Box<dyn Any>> {
self.objects.remove(&handle.raw())
}
pub fn reinsert_object_box(&mut self, handle: ObjectHandle, obj: Box<dyn Any>) {
if self.pending_object_drops.remove(&handle.raw()) {
self.object_handles.free(handle.raw());
drop(obj);
return;
}
assert!(
self.objects.insert(handle.raw(), obj).is_none(),
"object handle {} was reinserted while occupied",
handle.raw()
);
}
pub fn take_thread_local_box<K>(&mut self, key: &'static K) -> Option<Box<dyn Any>> {
self.thread_locals
.remove(&core::ptr::from_ref(key).cast::<()>())
}
pub fn insert_thread_local_box<K>(&mut self, key: &'static K, value: Box<dyn Any>) {
self.thread_locals
.insert(core::ptr::from_ref(key).cast::<()>(), value);
}
pub(crate) fn get_next_inbound_js_ref(&mut self) -> JsRef {
JsRef::from_raw(self.get_next_inbound_js_heap_id())
}
pub fn next_placeholder_ref(&mut self) -> JsRef {
JsRef::from_raw(self.get_next_placeholder_id())
}
pub fn next_borrowed_ref(&mut self) -> JsRef {
JsRef::from_raw(self.get_next_borrow_id())
}
}
thread_local! {
static RUNTIME: RefCell<Vec<Runtime>> = const { RefCell::new(Vec::new()) };
}
pub(crate) fn in_runtime<O>(runtime: Runtime, run: impl FnOnce() -> O) -> (Runtime, O) {
RUNTIME.with(|state| state.borrow_mut().push(runtime));
let out = run();
let runtime = RUNTIME.with(|state| {
state
.borrow_mut()
.pop()
.expect("No runtime available to pop")
});
(runtime, out)
}
pub fn with_runtime<R>(f: impl FnOnce(&mut Runtime) -> R) -> R {
RUNTIME.with(|state| {
let mut state = state.borrow_mut();
f(state.last_mut().expect("No runtime available"))
})
}
fn try_with_runtime<R>(f: impl FnOnce(&mut Runtime) -> R) -> Option<R> {
RUNTIME
.try_with(|state| {
let mut state = state.try_borrow_mut().ok()?;
Some(f(state.last_mut()?))
})
.ok()
.flatten()
}
fn runtime_installed() -> bool {
RUNTIME
.try_with(|state| {
state
.try_borrow()
.map(|state| !state.is_empty())
.unwrap_or(false)
})
.unwrap_or(false)
}
pub(crate) fn drop_js_object(js_ref: JsRef) {
let Some(Some(id)) = try_with_runtime(|runtime| runtime.release_heap_id(js_ref.raw())) else {
return;
};
let js_ref = JsRef::from_raw(id);
crate::js_helpers::js_drop_heap_ref(js_ref.raw());
recycle_heap_id_after_js_drop(js_ref);
}
pub(crate) fn dispose_js_rust_function(js_ref: JsRef) {
if runtime_installed() {
crate::js_helpers::js_dispose_rust_function(js_ref.raw());
}
}
pub(crate) fn drop_rust_object(handle: ObjectHandle) {
let Some(object) = try_with_runtime(|runtime| runtime.release_object_handle(handle)) else {
return;
};
drop(object);
}
pub(crate) fn is_batching() -> bool {
with_runtime(|state| state.is_batching())
}
fn recycle_heap_id_after_js_drop(js_ref: JsRef) {
with_runtime(|runtime| {
if runtime.is_batching() {
runtime.defer_heap_id_recycle_until_flush(js_ref.raw());
} else {
runtime.recycle_heap_id(js_ref.raw());
}
});
}
pub(crate) fn add_operation(
encoder: &mut EncodedData,
fn_id: u32,
add_args: impl FnOnce(&mut EncodedData),
) -> bool {
AbiBinaryEncode::encode(fn_id, encoder);
add_args(encoder);
encoder.take_needs_flush()
}
pub fn run_js_sync<R>(
fn_id: u32,
add_args: impl FnOnce(&mut EncodedData),
reserve_placeholder: impl FnOnce(&mut Runtime) -> Option<R>,
mut decode_result: impl for<'a> FnMut(DecodedData<'a>) -> R,
) -> R {
with_runtime(|state| state.push_operation_frame());
let mut batch = EncodedData::default();
let needs_flush = add_operation(&mut batch, fn_id, add_args);
with_runtime(|state| state.extend_encoder(batch));
let mut placeholder = with_runtime(|state| reserve_placeholder(state));
let result = if !is_batching() || needs_flush {
flush_and_then(move |data| placeholder.take().unwrap_or_else(|| decode_result(data)))
} else {
placeholder.unwrap_or_else(|| flush_and_then(decode_result))
};
let frame = with_runtime(|state| state.pop_operation_frame());
for id in frame.heap_ids {
let js_ref = JsRef::from_raw(id);
crate::js_helpers::js_drop_heap_ref(js_ref.raw());
recycle_heap_id_after_js_drop(js_ref);
}
for handle in frame.object_handles {
let object =
with_runtime(|state| state.remove_object_untyped(ObjectHandle::from_raw(handle)));
drop(object);
}
result
}
pub(crate) fn flush_and_return<R: BinaryDecode>() -> R {
flush_and_then(|mut data| R::decode(&mut data).expect("Failed to decode return value"))
}
pub(crate) fn flush_and_then<R>(mut then: impl for<'a> FnMut(DecodedData<'a>) -> R) -> R {
let (batch_msg, heap_ids_to_recycle_after_flush) = with_runtime(|state| state.take_message());
with_runtime(|runtime| runtime.ipc().send_ipc(batch_msg));
let mut heap_ids_to_recycle_after_flush = Some(heap_ids_to_recycle_after_flush);
loop {
if let Some(result) = crate::runtime::progress_js_with(&mut then) {
recycle_heap_ids_after_flush(
heap_ids_to_recycle_after_flush
.take()
.expect("heap IDs should only be recycled once per flush"),
);
return result;
}
}
}
fn recycle_heap_ids_after_flush(ids: Vec<JsRef>) {
for id in ids {
with_runtime(|state| {
state.recycle_heap_id_if_released(id.raw());
});
}
}
pub fn batch<R, F: FnOnce() -> R>(f: F) -> R {
let previous = with_runtime(|runtime| {
let previous = runtime.is_batching();
runtime.set_batching(true);
previous
});
let result = f();
if !previous {
force_flush();
}
with_runtime(|runtime| runtime.set_batching(previous));
result
}
pub fn batch_async<'a, R, F: core::future::Future<Output = R> + 'a>(
f: F,
) -> impl core::future::Future<Output = R> + 'a {
let mut f = Box::pin(f);
std::future::poll_fn(move |ctx| batch(|| f.as_mut().poll(ctx)))
}
pub fn force_flush() {
let has_pending = with_runtime(|state| !state.is_empty());
if has_pending {
flush_and_return::<()>();
}
}
#[cfg(test)]
mod take_encoder_tests {
use super::*;
use crate::ipc::DecodedVariant;
use crate::runtime::WryIPC;
fn test_runtime() -> Runtime {
let (ipc, _senders, _driver_commands) = WryIPC::new();
Runtime::new(ipc)
}
#[test]
fn take_encoder_yields_an_evaluate_message_with_no_request_id() {
let mut runtime = test_runtime();
assert!(runtime.is_empty());
let (message, _) = runtime.take_message();
let DecodedVariant::Evaluate { .. } = message.decoded().unwrap() else {
panic!("expected Evaluate message");
};
}
#[test]
fn object_drop_during_checkout_is_deferred_then_honored() {
use std::cell::Cell;
use std::rc::Rc;
struct DropFlag(Rc<Cell<bool>>);
impl Drop for DropFlag {
fn drop(&mut self) {
self.0.set(true);
}
}
let mut runtime = test_runtime();
let dropped = Rc::new(Cell::new(false));
let handle = runtime.insert_object_box(Box::new(DropFlag(dropped.clone())));
let checked_out = runtime
.take_object_box(handle)
.expect("invalid handle")
.downcast::<DropFlag>()
.expect("type mismatch");
assert!(runtime.remove_object_untyped(handle).is_none());
assert!(
!dropped.get(),
"object must not be dropped while it is checked out"
);
runtime.reinsert_object_box(handle, checked_out);
assert!(
dropped.get(),
"deferred drop must run once the checkout completes"
);
let reused = runtime.insert_object_box(Box::new(DropFlag(Rc::new(Cell::new(false)))));
assert_eq!(reused.raw(), handle.raw());
}
#[test]
fn remove_frees_handle_unlike_checkout_take() {
let mut runtime = test_runtime();
let checked_out = runtime.insert_object_box(Box::new(5_u32));
let checked_out_raw = checked_out.raw();
let value = runtime
.take_object_box(checked_out)
.expect("invalid handle");
assert_eq!(value.downcast_ref::<u32>().copied(), Some(5));
let removable = runtime.insert_object_box(Box::new(7_u32));
let removable_raw = removable.raw();
assert_ne!(removable_raw, checked_out_raw);
assert_eq!(
*runtime
.remove_object_untyped(removable)
.expect("invalid handle")
.downcast::<u32>()
.expect("type mismatch"),
7
);
let reused = runtime.insert_object_box(Box::new(9_u32));
assert_eq!(reused.raw(), removable_raw);
runtime.reinsert_object_box(checked_out, value);
}
}