1 // Copyright 2024, The Android Open Source Project 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 use core::{future::Future, pin::Pin}; 16 use gbl_async::poll; 17 18 /// A container abstraction that takes input of dynamically typed [Future]s and stores them at 19 /// pinned memory locations. 20 pub trait PinFutContainer<'a> { 21 /// Adds and pins a new [Future] of any type generated by the given closure `f`. 22 /// 23 /// If operation cannot be performed, such as due to no capacity, `f` should not be called. add_with<F: Future<Output = ()> + 'a>(&mut self, f: impl FnOnce() -> F)24 fn add_with<F: Future<Output = ()> + 'a>(&mut self, f: impl FnOnce() -> F); 25 26 /// Calls the closure on each element in the container. Removes the element if it returns true. for_each_remove_if( &mut self, cb: impl FnMut(&mut Pin<&mut (dyn Future<Output = ()> + 'a)>) -> bool, )27 fn for_each_remove_if( 28 &mut self, 29 cb: impl FnMut(&mut Pin<&mut (dyn Future<Output = ()> + 'a)>) -> bool, 30 ); 31 } 32 33 /// An internal container abstraction that takes input of a specific type of [Future] and stores 34 /// them at pinned memory locations. 35 pub(crate) trait PinFutContainerTyped<'a, F: Future + 'a> { 36 /// Adds and pins a new [Future] of type T into the container returned by `f`. 37 /// 38 /// If operation cannot be performed, such as due to no capacity, `f` should not be called. add_with(&mut self, f: impl FnOnce() -> F)39 fn add_with(&mut self, f: impl FnOnce() -> F); 40 41 /// Calls the closure on each element in the container. Removes the element if it returns true. for_each_remove_if( &mut self, cb: impl FnMut(&mut Pin<&mut (dyn Future<Output = ()> + 'a)>) -> bool, )42 fn for_each_remove_if( 43 &mut self, 44 cb: impl FnMut(&mut Pin<&mut (dyn Future<Output = ()> + 'a)>) -> bool, 45 ); 46 47 /// Returns the number of items 48 #[cfg(test)] size(&mut self) -> usize49 fn size(&mut self) -> usize { 50 let mut res = 0; 51 self.for_each_remove_if(|_| { 52 res += 1; 53 false 54 }); 55 res 56 } 57 58 /// Polls all the [Future] once. Returns the number or unfinished ones. poll_all(&mut self) -> usize59 fn poll_all(&mut self) -> usize { 60 let mut res = 0; 61 self.for_each_remove_if(|v| { 62 let finished = poll(v).is_some(); 63 res += usize::from(!finished); 64 finished 65 }); 66 res 67 } 68 69 /// Runs until all futures are finished 70 #[cfg(test)] run(&mut self)71 fn run(&mut self) { 72 while self.poll_all() > 0 {} 73 } 74 } 75 76 /// `PinFutContainer` can implement `PinFutContainerTyped` for any [Future] type. 77 impl<'a, F: Future<Output = ()> + 'a, T: PinFutContainer<'a>> PinFutContainerTyped<'a, F> for T { add_with(&mut self, f: impl FnOnce() -> F)78 fn add_with(&mut self, f: impl FnOnce() -> F) { 79 PinFutContainer::add_with(self, move || f()) 80 } 81 for_each_remove_if( &mut self, cb: impl FnMut(&mut Pin<&mut (dyn Future<Output = ()> + 'a)>) -> bool, )82 fn for_each_remove_if( 83 &mut self, 84 cb: impl FnMut(&mut Pin<&mut (dyn Future<Output = ()> + 'a)>) -> bool, 85 ) { 86 PinFutContainer::for_each_remove_if(self, cb) 87 } 88 } 89 90 /// An implementation of `PinFutContainerTyped` backed by a preallocated slice. 91 pub(crate) struct PinFutSlice<'a, F> { 92 arr: &'a mut [Pin<&'a mut F>], 93 used: usize, 94 } 95 96 impl<'a, F> PinFutSlice<'a, F> { 97 /// Creates a new instance new(arr: &'a mut [Pin<&'a mut F>]) -> Self98 pub fn new(arr: &'a mut [Pin<&'a mut F>]) -> Self { 99 Self { arr, used: 0 } 100 } 101 } 102 103 impl<'a, F: Future<Output = ()> + 'a> PinFutContainerTyped<'a, F> for PinFutSlice<'a, F> { add_with(&mut self, f: impl FnOnce() -> F)104 fn add_with(&mut self, f: impl FnOnce() -> F) { 105 if self.used < self.arr.len() { 106 self.arr[self.used].set(f()); 107 self.used += 1; 108 } 109 } 110 for_each_remove_if( &mut self, mut cb: impl FnMut(&mut Pin<&mut (dyn Future<Output = ()> + 'a)>) -> bool, )111 fn for_each_remove_if( 112 &mut self, 113 mut cb: impl FnMut(&mut Pin<&mut (dyn Future<Output = ()> + 'a)>) -> bool, 114 ) { 115 // Iterates from the end because we swap remove with the last. 116 for idx in (0..self.used).rev() { 117 if cb(&mut (self.arr[idx].as_mut() as _)) { 118 // Swaps remove with the last element 119 self.used -= 1; 120 self.arr.swap(idx, self.used); 121 } 122 } 123 } 124 } 125