1use super::{TypedResource, TypedResourceIndex};
2use alloc::vec::Vec;
3use anyhow::{Result, bail};
4use core::mem;
5use wasmtime_environ::component::{TypeFutureTableIndex, TypeStreamTableIndex};
6
7const MAX_HANDLE: u32 = 1 << 28;
11
12#[derive(Debug, Eq, PartialEq)]
15pub enum TransmitLocalState {
16 Write {
18 done: bool,
22 },
23 Read {
25 done: bool,
29 },
30 Busy,
32}
33
34pub enum RemovedResource {
36 Own { rep: u32 },
38 Borrow { scope: usize },
40}
41
42pub enum Waitable {
44 Subtask { is_host: bool },
45 Future,
46 Stream,
47}
48
49enum Slot {
50 Free {
51 next: u32,
52 },
53
54 ResourceOwn {
59 resource: TypedResource,
60 lend_count: u32,
61 },
62
63 ResourceBorrow {
69 resource: TypedResource,
70 scope: usize,
71 },
72
73 HostTask {
75 rep: u32,
76 },
77
78 GuestTask {
80 rep: u32,
81 },
82
83 Stream {
85 ty: TypeStreamTableIndex,
86 rep: u32,
87 state: TransmitLocalState,
88 },
89
90 Future {
92 ty: TypeFutureTableIndex,
93 rep: u32,
94 state: TransmitLocalState,
95 },
96
97 WaitableSet {
99 rep: u32,
100 },
101
102 ErrorContext {
104 rep: u32,
105 },
106}
107
108pub struct HandleTable {
109 next: u32,
110 slots: Vec<Slot>,
111}
112
113impl Default for HandleTable {
114 fn default() -> Self {
115 Self {
116 next: 0,
117 slots: Vec::new(),
118 }
119 }
120}
121
122impl HandleTable {
123 pub fn is_empty(&self) -> bool {
125 self.slots
126 .iter()
127 .all(|slot| matches!(slot, Slot::Free { .. }))
128 }
129
130 fn insert(&mut self, slot: Slot) -> Result<u32> {
131 let next = self.next as usize;
132 if next == self.slots.len() {
133 self.slots.push(Slot::Free {
134 next: self.next.checked_add(1).unwrap(),
135 });
136 }
137 let ret = self.next;
138 self.next = match mem::replace(&mut self.slots[next], slot) {
139 Slot::Free { next } => next,
140 _ => unreachable!(),
141 };
142 let ret = ret + 1;
147 if ret >= MAX_HANDLE {
148 bail!("cannot allocate another handle: index overflow");
149 }
150
151 Ok(ret)
152 }
153
154 fn remove(&mut self, idx: u32) -> Result<()> {
155 let to_fill = Slot::Free { next: self.next };
156 let slot = self.get_mut(idx)?;
157 *slot = to_fill;
158 self.next = idx - 1;
159 Ok(())
160 }
161
162 fn handle_index_to_table_index(&self, idx: u32) -> Option<usize> {
163 let idx = idx.checked_sub(1)?;
166 usize::try_from(idx).ok()
167 }
168
169 fn get_mut(&mut self, idx: u32) -> Result<&mut Slot> {
170 let slot = self
171 .handle_index_to_table_index(idx)
172 .and_then(|i| self.slots.get_mut(i));
173 match slot {
174 None | Some(Slot::Free { .. }) => bail!("unknown handle index {idx}"),
175 Some(slot) => Ok(slot),
176 }
177 }
178
179 pub fn resource_own_insert(&mut self, resource: TypedResource) -> Result<u32> {
182 self.insert(Slot::ResourceOwn {
183 resource,
184 lend_count: 0,
185 })
186 }
187
188 pub fn resource_borrow_insert(&mut self, resource: TypedResource, scope: usize) -> Result<u32> {
192 self.insert(Slot::ResourceBorrow { resource, scope })
193 }
194
195 pub fn resource_rep(&mut self, idx: TypedResourceIndex) -> Result<u32> {
200 match self.get_mut(idx.raw_index())? {
201 Slot::ResourceOwn { resource, .. } | Slot::ResourceBorrow { resource, .. } => {
202 resource.rep(&idx)
203 }
204 _ => bail!("index is not a resource"),
205 }
206 }
207
208 pub fn resource_lend(&mut self, idx: TypedResourceIndex) -> Result<(u32, bool)> {
217 match self.get_mut(idx.raw_index())? {
218 Slot::ResourceOwn {
219 resource,
220 lend_count,
221 } => {
222 let rep = resource.rep(&idx)?;
223 *lend_count = lend_count.checked_add(1).unwrap();
224 Ok((rep, true))
225 }
226 Slot::ResourceBorrow { resource, .. } => Ok((resource.rep(&idx)?, false)),
227 _ => bail!("index {} is not a resource", idx.raw_index()),
228 }
229 }
230
231 pub fn resource_undo_lend(&mut self, idx: TypedResourceIndex) -> Result<()> {
234 match self.get_mut(idx.raw_index())? {
235 Slot::ResourceOwn { lend_count, .. } => {
236 *lend_count -= 1;
237 Ok(())
238 }
239 _ => bail!("index {} is not an own resource", idx.raw_index()),
240 }
241 }
242
243 pub fn remove_resource(&mut self, idx: TypedResourceIndex) -> Result<RemovedResource> {
248 let ret = match self.get_mut(idx.raw_index())? {
249 Slot::ResourceOwn {
250 resource,
251 lend_count,
252 } => {
253 if *lend_count != 0 {
254 bail!("cannot remove owned resource while borrowed")
255 }
256 RemovedResource::Own {
257 rep: resource.rep(&idx)?,
258 }
259 }
260 Slot::ResourceBorrow { resource, scope } => {
261 resource.rep(&idx)?;
263 RemovedResource::Borrow { scope: *scope }
264 }
265 _ => bail!("index {} is not a resource", idx.raw_index()),
266 };
267 self.remove(idx.raw_index())?;
268 Ok(ret)
269 }
270
271 pub fn stream_insert_read(&mut self, ty: TypeStreamTableIndex, rep: u32) -> Result<u32> {
276 self.insert(Slot::Stream {
277 rep,
278 ty,
279 state: TransmitLocalState::Read { done: false },
280 })
281 }
282
283 pub fn stream_insert_write(&mut self, ty: TypeStreamTableIndex, rep: u32) -> Result<u32> {
288 self.insert(Slot::Stream {
289 rep,
290 ty,
291 state: TransmitLocalState::Write { done: false },
292 })
293 }
294
295 pub fn stream_rep(
301 &mut self,
302 expected_ty: TypeStreamTableIndex,
303 idx: u32,
304 ) -> Result<(u32, &mut TransmitLocalState)> {
305 match self.get_mut(idx)? {
306 Slot::Stream { rep, ty, state } => {
307 if *ty != expected_ty {
308 bail!("handle is a stream of a different type");
309 }
310 Ok((*rep, state))
311 }
312 _ => bail!("handle is not a stream"),
313 }
314 }
315
316 pub fn stream_remove_readable(
324 &mut self,
325 expected_ty: TypeStreamTableIndex,
326 idx: u32,
327 ) -> Result<(u32, bool)> {
328 let ret = match self.get_mut(idx)? {
329 Slot::Stream { rep, ty, state } => {
330 if *ty != expected_ty {
331 bail!("handle is a stream of a different type");
332 }
333 let is_done = match state {
334 TransmitLocalState::Read { done } => *done,
335 TransmitLocalState::Write { .. } => {
336 bail!("handle is not a readable end of a stream")
337 }
338 TransmitLocalState::Busy => bail!("cannot remove busy stream"),
339 };
340 (*rep, is_done)
341 }
342 _ => bail!("handle is not a stream"),
343 };
344 self.remove(idx)?;
345 Ok(ret)
346 }
347
348 pub fn stream_remove_writable(
350 &mut self,
351 expected_ty: TypeStreamTableIndex,
352 idx: u32,
353 ) -> Result<u32> {
354 let ret = match self.get_mut(idx)? {
355 Slot::Stream { rep, ty, state } => {
356 if *ty != expected_ty {
357 bail!("handle is a stream of a different type");
358 }
359 match state {
360 TransmitLocalState::Write { .. } => {}
361 TransmitLocalState::Read { .. } => {
362 bail!("passed read end to `stream.drop-writable`")
363 }
364 TransmitLocalState::Busy => bail!("cannot drop busy stream"),
365 }
366 *rep
367 }
368 _ => bail!("handle is not a stream"),
369 };
370 self.remove(idx)?;
371 Ok(ret)
372 }
373
374 pub fn future_insert_read(&mut self, ty: TypeFutureTableIndex, rep: u32) -> Result<u32> {
379 self.insert(Slot::Future {
380 rep,
381 ty,
382 state: TransmitLocalState::Read { done: false },
383 })
384 }
385
386 pub fn future_insert_write(&mut self, ty: TypeFutureTableIndex, rep: u32) -> Result<u32> {
391 self.insert(Slot::Future {
392 rep,
393 ty,
394 state: TransmitLocalState::Write { done: false },
395 })
396 }
397
398 pub fn future_rep(
404 &mut self,
405 expected_ty: TypeFutureTableIndex,
406 idx: u32,
407 ) -> Result<(u32, &mut TransmitLocalState)> {
408 match self.get_mut(idx)? {
409 Slot::Future { rep, ty, state } => {
410 if *ty != expected_ty {
411 bail!("handle is a future of a different type");
412 }
413 Ok((*rep, state))
414 }
415 _ => bail!("handle is not a future"),
416 }
417 }
418
419 pub fn future_remove_readable(
427 &mut self,
428 expected_ty: TypeFutureTableIndex,
429 idx: u32,
430 ) -> Result<(u32, bool)> {
431 let ret = match self.get_mut(idx)? {
432 Slot::Future { rep, ty, state } => {
433 if *ty != expected_ty {
434 bail!("handle is a future of a different type");
435 }
436 let is_done = match state {
437 TransmitLocalState::Read { done } => *done,
438 TransmitLocalState::Write { .. } => {
439 bail!("handle is not a readable end of a future")
440 }
441 TransmitLocalState::Busy => bail!("cannot remove busy future"),
442 };
443 (*rep, is_done)
444 }
445 _ => bail!("handle is not a future"),
446 };
447 self.remove(idx)?;
448 Ok(ret)
449 }
450
451 pub fn future_remove_writable(
453 &mut self,
454 expected_ty: TypeFutureTableIndex,
455 idx: u32,
456 ) -> Result<u32> {
457 let ret = match self.get_mut(idx)? {
458 Slot::Future { rep, ty, state } => {
459 if *ty != expected_ty {
460 bail!("handle is a future of a different type");
461 }
462 match state {
463 TransmitLocalState::Write { .. } => {}
464 TransmitLocalState::Read { .. } => {
465 bail!("passed read end to `future.drop-writable`")
466 }
467 TransmitLocalState::Busy => bail!("cannot drop busy future"),
468 }
469 *rep
470 }
471 _ => bail!("handle is not a future"),
472 };
473 self.remove(idx)?;
474 Ok(ret)
475 }
476
477 pub fn error_context_insert(&mut self, rep: u32) -> Result<u32> {
480 self.insert(Slot::ErrorContext { rep })
481 }
482
483 pub fn error_context_rep(&mut self, idx: u32) -> Result<u32> {
485 match self.get_mut(idx)? {
486 Slot::ErrorContext { rep } => Ok(*rep),
487 _ => bail!("handle is not an error-context"),
488 }
489 }
490
491 pub fn error_context_drop(&mut self, idx: u32) -> Result<u32> {
495 let rep = match self.get_mut(idx)? {
496 Slot::ErrorContext { rep } => *rep,
497 _ => bail!("handle is not an error-context"),
498 };
499 self.remove(idx)?;
500 Ok(rep)
501 }
502
503 pub fn subtask_insert_guest(&mut self, rep: u32) -> Result<u32> {
505 self.insert(Slot::GuestTask { rep })
506 }
507
508 pub fn subtask_insert_host(&mut self, rep: u32) -> Result<u32> {
510 self.insert(Slot::HostTask { rep })
511 }
512
513 pub fn subtask_rep(&mut self, idx: u32) -> Result<(u32, bool)> {
516 match self.get_mut(idx)? {
517 Slot::GuestTask { rep } => Ok((*rep, false)),
518 Slot::HostTask { rep } => Ok((*rep, true)),
519 _ => bail!("handle is not a subtask"),
520 }
521 }
522
523 pub fn subtask_remove(&mut self, idx: u32) -> Result<(u32, bool)> {
525 let ret = match self.get_mut(idx)? {
526 Slot::GuestTask { rep } => (*rep, false),
527 Slot::HostTask { rep } => (*rep, true),
528 _ => bail!("handle is not a subtask"),
529 };
530 self.remove(idx)?;
531 Ok(ret)
532 }
533
534 pub fn waitable_set_insert(&mut self, rep: u32) -> Result<u32> {
536 self.insert(Slot::WaitableSet { rep })
537 }
538
539 pub fn waitable_set_rep(&mut self, idx: u32) -> Result<u32> {
541 match self.get_mut(idx)? {
542 Slot::WaitableSet { rep, .. } => Ok(*rep),
543 _ => bail!("handle is not an waitable-set"),
544 }
545 }
546
547 pub fn waitable_set_remove(&mut self, idx: u32) -> Result<u32> {
549 let ret = match self.get_mut(idx)? {
550 Slot::WaitableSet { rep } => *rep,
551 _ => bail!("handle is not a waitable-set"),
552 };
553 self.remove(idx)?;
554 Ok(ret)
555 }
556
557 pub fn waitable_rep(&mut self, idx: u32) -> Result<(u32, Waitable)> {
560 match self.get_mut(idx)? {
561 Slot::GuestTask { rep } => Ok((*rep, Waitable::Subtask { is_host: false })),
562 Slot::HostTask { rep } => Ok((*rep, Waitable::Subtask { is_host: true })),
563 Slot::Future { rep, .. } => Ok((*rep, Waitable::Future)),
564 Slot::Stream { rep, .. } => Ok((*rep, Waitable::Stream)),
565 _ => bail!("handle is not a waitable"),
566 }
567 }
568}