@@ -74,9 +74,8 @@ use crate::p2::bindings::{
7474} ;
7575use crate :: p2:: { FsError , IsATTY , WasiCtx , WasiImpl , WasiView } ;
7676use anyhow:: { Context , bail} ;
77- use std:: collections:: { BTreeMap , HashSet } ;
77+ use std:: collections:: { BTreeMap , BTreeSet , HashSet , btree_map } ;
7878use std:: mem:: { self , size_of, size_of_val} ;
79- use std:: ops:: { Deref , DerefMut } ;
8079use std:: slice;
8180use std:: sync:: Arc ;
8281use std:: sync:: atomic:: { AtomicU64 , Ordering } ;
@@ -286,21 +285,7 @@ struct WasiPreview1Adapter {
286285#[ derive( Debug , Default ) ]
287286struct Descriptors {
288287 used : BTreeMap < u32 , Descriptor > ,
289- free : Vec < u32 > ,
290- }
291-
292- impl Deref for Descriptors {
293- type Target = BTreeMap < u32 , Descriptor > ;
294-
295- fn deref ( & self ) -> & Self :: Target {
296- & self . used
297- }
298- }
299-
300- impl DerefMut for Descriptors {
301- fn deref_mut ( & mut self ) -> & mut Self :: Target {
302- & mut self . used
303- }
288+ free : BTreeSet < u32 > ,
304289}
305290
306291impl Descriptors {
@@ -377,42 +362,34 @@ impl Descriptors {
377362
378363 /// Returns next descriptor number, which was never assigned
379364 fn unused ( & self ) -> Result < u32 > {
380- match self . last_key_value ( ) {
365+ match self . used . last_key_value ( ) {
381366 Some ( ( fd, _) ) => {
382367 if let Some ( fd) = fd. checked_add ( 1 ) {
383368 return Ok ( fd) ;
384369 }
385- if self . len ( ) == u32:: MAX as usize {
370+ if self . used . len ( ) == u32:: MAX as usize {
386371 return Err ( types:: Errno :: Loop . into ( ) ) ;
387372 }
388373 // TODO: Optimize
389374 Ok ( ( 0 ..u32:: MAX )
390375 . rev ( )
391- . find ( |fd| !self . contains_key ( fd) )
376+ . find ( |fd| !self . used . contains_key ( fd) )
392377 . expect ( "failed to find an unused file descriptor" ) )
393378 }
394379 None => Ok ( 0 ) ,
395380 }
396381 }
397382
398- /// Removes the [Descriptor] corresponding to `fd`
399- fn remove ( & mut self , fd : types:: Fd ) -> Option < Descriptor > {
400- let fd = fd. into ( ) ;
401- let desc = self . used . remove ( & fd) ?;
402- self . free . push ( fd) ;
403- Some ( desc)
404- }
405-
406383 /// Pushes the [Descriptor] returning corresponding number.
407384 /// This operation will try to reuse numbers previously removed via [`Self::remove`]
408385 /// and rely on [`Self::unused`] if no free numbers are recorded
409386 fn push ( & mut self , desc : Descriptor ) -> Result < u32 > {
410- let fd = if let Some ( fd) = self . free . pop ( ) {
387+ let fd = if let Some ( fd) = self . free . pop_last ( ) {
411388 fd
412389 } else {
413390 self . unused ( ) ?
414391 } ;
415- assert ! ( self . insert( fd, desc) . is_none( ) ) ;
392+ assert ! ( self . used . insert( fd, desc) . is_none( ) ) ;
416393 Ok ( fd)
417394 }
418395}
@@ -453,15 +430,15 @@ impl Transaction<'_> {
453430 /// Returns [`types::Errno::Badf`] if no [`Descriptor`] is found
454431 fn get_descriptor ( & self , fd : types:: Fd ) -> Result < & Descriptor > {
455432 let fd = fd. into ( ) ;
456- let desc = self . descriptors . get ( & fd) . ok_or ( types:: Errno :: Badf ) ?;
433+ let desc = self . descriptors . used . get ( & fd) . ok_or ( types:: Errno :: Badf ) ?;
457434 Ok ( desc)
458435 }
459436
460437 /// Borrows [`File`] corresponding to `fd`
461438 /// if it describes a [`Descriptor::File`]
462439 fn get_file ( & self , fd : types:: Fd ) -> Result < & File > {
463440 let fd = fd. into ( ) ;
464- match self . descriptors . get ( & fd) {
441+ match self . descriptors . used . get ( & fd) {
465442 Some ( Descriptor :: File ( file) ) => Ok ( file) ,
466443 _ => Err ( types:: Errno :: Badf . into ( ) ) ,
467444 }
@@ -471,7 +448,7 @@ impl Transaction<'_> {
471448 /// if it describes a [`Descriptor::File`]
472449 fn get_file_mut ( & mut self , fd : types:: Fd ) -> Result < & mut File > {
473450 let fd = fd. into ( ) ;
474- match self . descriptors . get_mut ( & fd) {
451+ match self . descriptors . used . get_mut ( & fd) {
475452 Some ( Descriptor :: File ( file) ) => Ok ( file) ,
476453 _ => Err ( types:: Errno :: Badf . into ( ) ) ,
477454 }
@@ -485,7 +462,7 @@ impl Transaction<'_> {
485462 /// Returns [`types::Errno::Spipe`] if the descriptor corresponds to stdio
486463 fn get_seekable ( & self , fd : types:: Fd ) -> Result < & File > {
487464 let fd = fd. into ( ) ;
488- match self . descriptors . get ( & fd) {
465+ match self . descriptors . used . get ( & fd) {
489466 Some ( Descriptor :: File ( file) ) => Ok ( file) ,
490467 Some (
491468 Descriptor :: Stdin { .. } | Descriptor :: Stdout { .. } | Descriptor :: Stderr { .. } ,
@@ -518,7 +495,7 @@ impl Transaction<'_> {
518495 /// if it describes a [`Descriptor::Directory`]
519496 fn get_dir_fd ( & self , fd : types:: Fd ) -> Result < Resource < filesystem:: Descriptor > > {
520497 let fd = fd. into ( ) ;
521- match self . descriptors . get ( & fd) {
498+ match self . descriptors . used . get ( & fd) {
522499 Some ( Descriptor :: Directory { fd, .. } ) => Ok ( fd. borrowed ( ) ) ,
523500 _ => Err ( types:: Errno :: Badf . into ( ) ) ,
524501 }
@@ -1327,11 +1304,13 @@ impl wasi_snapshot_preview1::WasiSnapshotPreview1 for WasiP1Ctx {
13271304 _memory : & mut GuestMemory < ' _ > ,
13281305 fd : types:: Fd ,
13291306 ) -> Result < ( ) , types:: Error > {
1330- let desc = self
1331- . transact ( ) ?
1332- . descriptors
1333- . remove ( fd)
1334- . ok_or ( types:: Errno :: Badf ) ?;
1307+ let desc = {
1308+ let fd = fd. into ( ) ;
1309+ let mut st = self . transact ( ) ?;
1310+ let desc = st. descriptors . used . remove ( & fd) . ok_or ( types:: Errno :: Badf ) ?;
1311+ st. descriptors . free . insert ( fd) ;
1312+ desc
1313+ } ;
13351314 match desc {
13361315 Descriptor :: Stdin { stream, .. } => {
13371316 streams:: HostInputStream :: drop ( & mut self . as_io_impl ( ) , stream)
@@ -1830,8 +1809,17 @@ impl wasi_snapshot_preview1::WasiSnapshotPreview1 for WasiP1Ctx {
18301809 to : types:: Fd ,
18311810 ) -> Result < ( ) , types:: Error > {
18321811 let mut st = self . transact ( ) ?;
1833- let desc = st. descriptors . remove ( from) . ok_or ( types:: Errno :: Badf ) ?;
1834- st. descriptors . insert ( to. into ( ) , desc) ;
1812+ let from = from. into ( ) ;
1813+ let btree_map:: Entry :: Occupied ( desc) = st. descriptors . used . entry ( from) else {
1814+ return Err ( types:: Errno :: Badf . into ( ) ) ;
1815+ } ;
1816+ let to = to. into ( ) ;
1817+ if from != to {
1818+ let desc = desc. remove ( ) ;
1819+ st. descriptors . free . insert ( from) ;
1820+ st. descriptors . free . remove ( & to) ;
1821+ st. descriptors . used . insert ( to, desc) ;
1822+ }
18351823 Ok ( ( ) )
18361824 }
18371825
0 commit comments