assemblylift_core_io_guest/
lib.rs

1use std::future::Future;
2use std::io::BufReader;
3use std::marker::PhantomData;
4use std::pin::Pin;
5use std::task::{Context, Poll, Waker};
6
7use serde::{Deserialize, de::DeserializeOwned};
8
9use assemblylift_core_io_common::constants::{FUNCTION_INPUT_BUFFER_SIZE, IO_BUFFER_SIZE_BYTES};
10
11extern "C" {
12    // IO
13    fn __asml_abi_io_poll(id: u32) -> i32;
14    fn __asml_abi_io_len(id: u32) -> u32;
15    fn __asml_abi_io_load(id: u32) -> i32;
16    fn __asml_abi_io_next() -> i32;
17
18    // System clock
19    fn __asml_abi_clock_time_get() -> u64;
20
21    // Console
22    fn __asml_abi_console_log(ptr: *const u8, len: usize);
23
24    // Input
25    fn __asml_abi_input_start() -> i32;
26    fn __asml_abi_input_next() -> i32;
27    fn __asml_abi_input_length_get() -> u64;
28
29    // Z85
30    fn __asml_expabi_z85_encode(ptr: *const u8, len: usize, out_ptr: *const u8) -> i32;
31    fn __asml_expabi_z85_decode(ptr: *const u8, len: usize, out_ptr: *const u8) -> i32;
32}
33
34// Raw buffer holding serialized IO data
35pub static mut IO_BUFFER: [u8; IO_BUFFER_SIZE_BYTES] = [0; IO_BUFFER_SIZE_BYTES];
36
37#[no_mangle]
38pub fn __asml_guest_get_io_buffer_pointer() -> *const u8 {
39    unsafe { IO_BUFFER.as_ptr() }
40}
41
42fn console_log(message: String) {
43    unsafe { __asml_abi_console_log(message.as_ptr(), message.len()) }
44}
45
46pub fn get_time() -> u64 {
47    unsafe { __asml_abi_clock_time_get() }
48}
49
50pub struct IoDocument {
51    bytes_read: usize,
52    pages_read: usize,
53    length: usize,
54}
55
56impl IoDocument {
57    pub fn new(ioid: u32) -> Self {
58        unsafe { __asml_abi_io_load(ioid) };
59        Self {
60            bytes_read: 0,
61            pages_read: 0,
62            length: unsafe { __asml_abi_io_len(ioid) } as usize,
63        }
64    }
65
66    pub fn len(&self) -> usize {
67        self.length
68    }
69}
70
71impl std::io::Read for IoDocument {
72    fn read(&mut self, buf: &mut [u8]) -> Result<usize, std::io::Error> {
73        let mut bytes_read = 0usize;
74        if self.bytes_read < self.length {
75            for idx in 0..std::cmp::min(self.length, buf.len()) {
76                // unsafe: bytes_read is always positive, mod IO_BUFFER_SIZE_BYTES 
77                //         is always less than IO_BUFFER_SIZE_BYTES
78                buf[idx] = unsafe { 
79                    IO_BUFFER[self.bytes_read % IO_BUFFER_SIZE_BYTES]
80                };
81                bytes_read += 1;
82                self.bytes_read += 1;
83                if self.bytes_read % IO_BUFFER_SIZE_BYTES == 0 {
84                    unsafe { __asml_abi_io_next() };
85                    self.pages_read += 1;
86                }
87            }
88        }
89        Ok(bytes_read)
90    }
91}
92
93#[derive(Clone)]
94pub struct Io<'a, R> {
95    pub id: u32,
96    waker: Box<Option<Waker>>,
97    _phantom: PhantomData<&'a R>,
98}
99
100impl<'a, R: Deserialize<'a>> Io<'_, R> {
101    pub fn new(id: u32) -> Self {
102        Io {
103            id,
104            waker: Box::new(None),
105            _phantom: PhantomData,
106        }
107    }
108}
109
110impl<'a, R> Future for Io<'_, R> 
111where
112    R: DeserializeOwned,
113{
114    type Output = R;
115
116    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
117        match unsafe { __asml_abi_io_poll(self.id) } {
118            1 => Poll::Ready(read_response::<Self::Output>(self.id).unwrap()),
119            _ => {
120                self.waker = Box::new(Some(cx.waker().clone()));
121                Poll::Pending
122            }
123        }
124    }
125}
126
127fn read_response<'a, T>(id: u32) -> Option<T>
128where
129    T: DeserializeOwned,
130{
131    let doc = IoDocument::new(id);
132    let doc = BufReader::with_capacity(doc.len(), doc);
133    match serde_json::from_reader::<BufReader<IoDocument>, T>(doc) {
134        Ok(response) => Some(response),
135        Err(why) => {
136            console_log(format!("[ERROR] ioid={} {}", id, why.to_string()));
137            None
138        }
139    }
140}
141
142// Function Input Buffer
143
144pub static mut FUNCTION_INPUT_BUFFER: [u8; FUNCTION_INPUT_BUFFER_SIZE] =
145    [0; FUNCTION_INPUT_BUFFER_SIZE];
146
147// provided TO the wasm runtime (host)
148#[no_mangle]
149pub fn __asml_guest_get_function_input_buffer_pointer() -> *const u8 {
150    unsafe { FUNCTION_INPUT_BUFFER.as_ptr() }
151}
152
153pub struct FunctionInputBuffer {
154    bytes_read: usize,
155    pages_read: usize,
156    length: usize,
157}
158
159impl FunctionInputBuffer {
160    pub fn new() -> Self {
161        unsafe { __asml_abi_input_start() };
162        Self {
163            bytes_read: 0usize,
164            pages_read: 0usize,
165            length: unsafe { __asml_abi_input_length_get() as usize },
166        }
167    }
168}
169
170impl std::io::Read for FunctionInputBuffer {
171    fn read(&mut self, buf: &mut [u8]) -> Result<usize, std::io::Error> {
172        let mut bytes_read = 0usize;
173        if self.bytes_read < self.length {
174            for idx in 0..std::cmp::min(self.length, buf.len()) {
175                // unsafe: bytes_read is always positive, mod FUNCTION_INPUT_BUFFER_SIZE
176                //         is always less than FUNCTION_INPUT_BUFFER_SIZE
177                buf[idx] = unsafe {
178                    FUNCTION_INPUT_BUFFER[self.bytes_read % FUNCTION_INPUT_BUFFER_SIZE]
179                };
180                bytes_read += 1;
181                self.bytes_read += 1;
182                if self.bytes_read % FUNCTION_INPUT_BUFFER_SIZE == 0 {
183                    unsafe { __asml_abi_input_next() };
184                    self.pages_read += 1;
185                }
186            }
187        }
188        Ok(bytes_read)
189    }
190}
assemblylift_core_io_guest/lib.rs

assemblylift_core_io_guest/
lib.rs
