mikros: Add a Future to get RPC function returns

library/std/src/os/mikros/ipc/rpc.rs

@@ -6,6 +6,9 @@ use crate::collections::HashMap;
 use crate::ops::{Deref, DerefMut};
 use crate::sync::{LazyLock, Mutex, RwLock};
 use crate::sys::buffers::KernelBufferAllocator;
+use crate::future::Future;
+use crate::pin::Pin;
+use crate::task::{Context, Poll};
 
 type MessageCallback = fn(IncomingCall);
 type ReturnCallback = Box<dyn FnOnce() + Sync + Send>;
@@ -84,6 +87,40 @@ impl CallId {
     pub fn set_return_callback(self, callback: ReturnCallback) {
         RETURN_CALLBACKS.lock().unwrap().insert(self, callback);
     }
+
+
+    #[stable(feature = "mikros", since = "1.80.0")]
+    pub fn get_return_async(self) -> RpcReturnFuture {
+        RpcReturnFuture {
+            call_id: self,
+            callback_registered: false,
+        }
+    }
+}
+
+#[stable(feature = "mikros", since = "1.80.0")]
+pub struct RpcReturnFuture {
+    call_id: CallId,
+    callback_registered: bool,
+}
+
+#[stable(feature = "mikros", since = "1.80.0")]
+impl Future for RpcReturnFuture {
+    type Output = RpcReturn;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let slf = self.get_mut();
+        if let Some(ret) = slf.call_id.try_get_return() {
+            return Poll::Ready(ret);
+        }
+        if !slf.callback_registered {
+            let waker = cx.waker().clone();
+            slf.call_id
+                .set_return_callback(Box::new(move || waker.wake()));
+            slf.callback_registered = true;
+        }
+        Poll::Pending
+    }
 }
 
 #[stable(feature = "mikros", since = "1.80.0")]