Next Generation WASM Microkernel Operating System
wasm
os
rust
microkernel
1// Copyright 2025 Jonas Kruckenberg
2//
3// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
4// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5// http://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8use alloc::boxed::Box;
9use core::fmt::Formatter;
10use core::iter::{FlatMap, Flatten, FusedIterator};
11use core::mem::offset_of;
12use core::pin::Pin;
13use core::ptr::NonNull;
14use core::{array, fmt};
15
16use pin_project::pin_project;
17use wavltree::WAVLTree;
18
19use crate::arch;
20use crate::mem::frame_alloc::Frame;
21
22const FRAME_LIST_NODE_FANOUT: usize = 16;
23
24pub struct FrameList {
25 pub nodes: WAVLTree<FrameListNode>,
26 size: usize,
27}
28
29#[pin_project]
30#[derive(Debug)]
31pub struct FrameListNode {
32 links: wavltree::Links<FrameListNode>,
33 offset: usize,
34 frames: [Option<Frame>; FRAME_LIST_NODE_FANOUT],
35}
36
37pub struct Cursor<'a> {
38 cursor: wavltree::Cursor<'a, FrameListNode>,
39 index_in_node: usize,
40 offset: usize,
41}
42
43pub struct CursorMut<'a> {
44 cursor: wavltree::CursorMut<'a, FrameListNode>,
45 index_in_node: usize,
46 offset: usize,
47}
48
49pub enum Entry<'a> {
50 Occupied(OccupiedEntry<'a>),
51 Vacant(VacantEntry<'a>),
52}
53pub struct OccupiedEntry<'a> {
54 entry: wavltree::OccupiedEntry<'a, FrameListNode>,
55 index_in_node: usize,
56}
57pub struct VacantEntry<'a> {
58 entry: wavltree::Entry<'a, FrameListNode>,
59 index_in_node: usize,
60 offset: usize,
61}
62
63// === FrameList ===
64
65impl fmt::Debug for FrameList {
66 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
67 f.debug_struct("FrameList")
68 .field("size", &self.size)
69 .field_with("nodes", |f| {
70 let mut f = f.debug_list();
71 self.nodes.iter().for_each(|node| {
72 f.entry(node);
73 });
74 f.finish()
75 })
76 .finish()
77 }
78}
79
80impl FrameList {
81 pub fn new() -> Self {
82 Self {
83 nodes: WAVLTree::new(),
84 size: 0,
85 }
86 }
87
88 pub fn is_empty(&self) -> bool {
89 self.nodes.is_empty()
90 }
91
92 pub fn size(&self) -> usize {
93 self.size
94 }
95
96 pub fn get(&self, offset: usize) -> Option<&Frame> {
97 let node_offset = offset_to_node_offset(offset);
98 let node = self.nodes.find(&node_offset).get()?;
99
100 let frame = node.frames.get(offset_to_node_index(offset))?;
101 frame.as_ref()
102 }
103
104 pub fn get_mut(&mut self, offset: usize) -> Option<&mut Frame> {
105 let node_offset = offset_to_node_offset(offset);
106 let node = Pin::into_inner(self.nodes.find_mut(&node_offset).get_mut()?);
107
108 let frame = node.frames.get_mut(offset_to_node_index(offset))?;
109 frame.as_mut()
110 }
111
112 pub fn take(&mut self, offset: usize) -> Option<Frame> {
113 let node_offset = offset_to_node_offset(offset);
114 let node = Pin::into_inner(self.nodes.find_mut(&node_offset).get_mut()?);
115
116 let frame = node.frames.get_mut(offset_to_node_index(offset))?;
117 frame.take()
118 }
119
120 pub fn replace(&mut self, offset: usize, new: Frame) -> Option<Frame> {
121 let node_offset = offset_to_node_offset(offset);
122 let node = self.nodes.entry(&node_offset).or_insert_with(|| {
123 Box::pin(FrameListNode {
124 links: wavltree::Links::default(),
125 offset,
126 frames: [const { None }; FRAME_LIST_NODE_FANOUT],
127 })
128 });
129
130 // Safety: we'll not move out of the node, just manipulate its fields
131 let frame = unsafe { Pin::into_inner_unchecked(node) }
132 .frames
133 .get_mut(offset_to_node_index(offset))?;
134
135 frame.replace(new)
136 }
137
138 pub fn insert(&mut self, offset: usize, new: Frame) -> &mut Frame {
139 let node_offset = offset_to_node_offset(offset);
140 let node = self.nodes.entry(&node_offset).or_insert_with(|| {
141 Box::pin(FrameListNode {
142 links: wavltree::Links::default(),
143 offset,
144 frames: [const { None }; FRAME_LIST_NODE_FANOUT],
145 })
146 });
147
148 // Safety: we'll not move out of the node, just manipulate its fields
149 let frame = unsafe { Pin::into_inner_unchecked(node) }
150 .frames
151 .get_mut(offset_to_node_index(offset))
152 .unwrap();
153
154 frame.insert(new)
155 }
156
157 pub fn first(&self) -> Option<&Frame> {
158 let node = self.nodes.front().get()?;
159 node.frames.iter().find(|f| f.is_some())?.as_ref()
160 }
161
162 pub fn last(&self) -> Option<&Frame> {
163 let node = self.nodes.back().get()?;
164 node.frames.iter().rfind(|f| f.is_some())?.as_ref()
165 }
166
167 pub fn cursor(&self, offset: usize) -> Cursor<'_> {
168 let node_offset = offset_to_node_offset(offset);
169 let cursor = self.nodes.find(&node_offset);
170
171 Cursor {
172 cursor,
173 index_in_node: offset_to_node_index(offset),
174 offset,
175 }
176 }
177
178 pub fn cursor_mut(&mut self, offset: usize) -> CursorMut<'_> {
179 let node_offset = offset_to_node_offset(offset);
180 let cursor = self.nodes.find_mut(&node_offset);
181
182 CursorMut {
183 cursor,
184 index_in_node: offset_to_node_index(offset),
185 offset,
186 }
187 }
188
189 pub(crate) fn entry(&mut self, offset: usize) -> Entry<'_> {
190 let node_offset = offset_to_node_offset(offset);
191 let index_in_node = offset_to_node_index(offset);
192 let entry = self.nodes.entry(&node_offset);
193
194 match entry {
195 wavltree::Entry::Occupied(entry) if entry.get().frames[index_in_node].is_some() => {
196 Entry::Occupied(OccupiedEntry {
197 entry,
198 index_in_node,
199 })
200 }
201 entry => Entry::Vacant(VacantEntry {
202 entry,
203 index_in_node,
204 offset: node_offset,
205 }),
206 }
207 }
208
209 pub fn iter(&self) -> impl Iterator<Item = &Frame> {
210 self.nodes
211 .iter()
212 .flat_map(|node| node.frames.iter().filter_map(|f| f.as_ref()))
213 }
214
215 pub fn clear(&mut self) {
216 self.nodes.clear();
217 }
218
219 /// Asserts the frame list is in a valid state.
220 pub fn assert_valid(&self, ctx: &str) {
221 self.nodes.assert_valid(ctx);
222 self.iter().for_each(|frame| frame.assert_valid(ctx));
223 }
224}
225
226// === FrameList IntoIterator ===
227
228type FramesWithoutHoles = Flatten<array::IntoIter<Option<Frame>, FRAME_LIST_NODE_FANOUT>>;
229type IntoIterInner = FlatMap<
230 wavltree::IntoIter<FrameListNode>,
231 FramesWithoutHoles,
232 fn(Pin<Box<FrameListNode>>) -> FramesWithoutHoles,
233>;
234
235pub struct IntoIter(IntoIterInner);
236impl Iterator for IntoIter {
237 type Item = Frame;
238
239 fn next(&mut self) -> Option<Self::Item> {
240 self.0.next()
241 }
242}
243impl DoubleEndedIterator for IntoIter {
244 fn next_back(&mut self) -> Option<Self::Item> {
245 self.0.next_back()
246 }
247}
248impl FusedIterator for IntoIter {}
249
250impl IntoIterator for FrameList {
251 type Item = Frame;
252 type IntoIter = IntoIter;
253
254 fn into_iter(mut self) -> Self::IntoIter {
255 let inner: IntoIterInner = self
256 .nodes
257 .take()
258 .into_iter()
259 .flat_map(|node| Pin::into_inner(node).frames.into_iter().flatten());
260
261 IntoIter(inner)
262 }
263}
264
265impl FromIterator<Frame> for FrameList {
266 fn from_iter<T: IntoIterator<Item = Frame>>(iter: T) -> Self {
267 let mut nodes: WAVLTree<FrameListNode> = WAVLTree::new();
268
269 let mut offset = 0;
270 for frame in iter.into_iter() {
271 let node = nodes
272 .entry(&offset_to_node_offset(offset))
273 .or_insert_with(|| {
274 Box::pin(FrameListNode {
275 links: wavltree::Links::default(),
276 offset,
277 frames: [const { None }; FRAME_LIST_NODE_FANOUT],
278 })
279 });
280
281 node.project().frames[offset_to_node_index(offset)] = Some(frame);
282 offset += arch::PAGE_SIZE;
283 }
284
285 Self {
286 nodes,
287 size: offset,
288 }
289 }
290}
291
292// === FrameListNode ===
293
294// Safety: unsafe trait
295unsafe impl wavltree::Linked for FrameListNode {
296 type Handle = Pin<Box<FrameListNode>>;
297 type Key = usize;
298
299 fn into_ptr(handle: Self::Handle) -> NonNull<Self> {
300 // Safety: wavltree treats the ptr as pinned
301 unsafe { NonNull::from(Box::leak(Pin::into_inner_unchecked(handle))) }
302 }
303
304 unsafe fn from_ptr(ptr: NonNull<Self>) -> Self::Handle {
305 // Safety: `NonNull` *must* be constructed from a pinned reference
306 // which the tree implementation upholds.
307 unsafe { Pin::new_unchecked(Box::from_raw(ptr.as_ptr())) }
308 }
309
310 unsafe fn links(ptr: NonNull<Self>) -> NonNull<wavltree::Links<Self>> {
311 ptr.map_addr(|addr| {
312 let offset = offset_of!(Self, links);
313 addr.checked_add(offset).unwrap()
314 })
315 .cast()
316 }
317
318 fn get_key(&self) -> &Self::Key {
319 &self.offset
320 }
321}
322
323fn offset_to_node_offset(offset: usize) -> usize {
324 (offset) & 0usize.wrapping_sub(arch::PAGE_SIZE * FRAME_LIST_NODE_FANOUT)
325}
326
327fn offset_to_node_index(offset: usize) -> usize {
328 (offset >> arch::PAGE_SHIFT) % FRAME_LIST_NODE_FANOUT
329}
330
331// === Cursor ===
332
333impl<'a> Cursor<'a> {
334 /// Moves the cursor to the next [`Frame`] in the list
335 pub fn move_next(&mut self) {
336 self.offset += arch::PAGE_SIZE;
337
338 // if there is a current node AND the node still has unseen frames in it
339 // advance the offset
340 if let Some(node) = self.cursor.get() {
341 self.index_in_node += 1;
342 if node.frames.len() > self.index_in_node {
343 return;
344 }
345 }
346
347 // otherwise advance the cursor and reset the offset
348 self.cursor.move_next();
349 self.index_in_node = 0;
350 }
351
352 /// Returns the offset of the [`Frame`] in this list, will always be a multiple
353 /// of [`arch::PAGE_SIZE`].
354 pub fn offset(&self) -> usize {
355 self.offset
356 }
357
358 /// Returns a reference to the current [`Frame`] if any.
359 pub fn get(&self) -> Option<&'a Frame> {
360 let node = self.cursor.get()?;
361 node.frames.get(self.index_in_node)?.as_ref()
362 }
363}
364
365// === CursorMut ===
366
367impl<'a> CursorMut<'a> {
368 /// Moves the cursor to the next [`Frame`] in the list
369 pub fn move_next(&mut self) {
370 self.offset += arch::PAGE_SIZE;
371
372 // if there is a current node AND the node still has unseen frames in it
373 // advance the index
374 if let Some(node) = self.cursor.get() {
375 self.index_in_node += 1;
376 if node.frames.len() > self.index_in_node {
377 return;
378 }
379 }
380
381 // otherwise advance the cursor and reset the index
382 self.cursor.move_next();
383 self.index_in_node = 0;
384 }
385
386 /// Returns the offset of the [`Frame`] in this list, will always be a multiple
387 /// of [`arch::PAGE_SIZE`].
388 pub fn offset(&self) -> usize {
389 self.offset
390 }
391
392 pub fn remove(&mut self) -> Option<Frame> {
393 let node = Pin::into_inner(self.cursor.get_mut()?);
394 let frame = node.frames.get_mut(self.index_in_node)?.take()?;
395
396 // if the node has become empty remove it too
397 if node.frames.iter().all(Option::is_none) {
398 let _node = self.cursor.remove();
399 self.index_in_node = 0;
400 }
401
402 Some(frame)
403 }
404
405 /// Returns a reference to the current [`Frame`] if any.
406 pub fn get(&self) -> Option<&'a Frame> {
407 let node = self.cursor.get()?;
408 node.frames.get(self.index_in_node)?.as_ref()
409 }
410
411 /// Returns a mutable reference to the current [`Frame`] if any.
412 pub fn get_mut(&mut self) -> Option<&mut Frame> {
413 let node = Pin::into_inner(self.cursor.get_mut()?);
414 node.frames.get_mut(self.index_in_node)?.as_mut()
415 }
416}
417
418// === Entry ===
419
420impl<'a> Entry<'a> {
421 #[inline]
422 pub fn and_modify<F>(self, f: F) -> Self
423 where
424 F: FnOnce(&mut Frame),
425 {
426 match self {
427 Entry::Occupied(mut entry) => {
428 f(entry.get_mut());
429 Entry::Occupied(entry)
430 }
431 Entry::Vacant(entry) => Entry::Vacant(entry),
432 }
433 }
434
435 #[inline]
436 pub fn or_insert(self, default: Frame) -> &'a mut Frame {
437 match self {
438 Entry::Occupied(entry) => entry.into_mut(),
439 Entry::Vacant(entry) => entry.insert(default),
440 }
441 }
442
443 #[inline]
444 pub fn or_insert_with<F: FnOnce() -> Frame>(self, default: F) -> &'a mut Frame {
445 match self {
446 Entry::Occupied(entry) => entry.into_mut(),
447 Entry::Vacant(entry) => entry.insert(default()),
448 }
449 }
450}
451
452impl<'a> OccupiedEntry<'a> {
453 pub fn into_mut(mut self) -> &'a mut Frame {
454 // Safety: guaranteed by `FrameList::entry`
455 unsafe {
456 Pin::into_inner_unchecked(self.entry.get_mut())
457 .frames
458 .get_unchecked_mut(self.index_in_node)
459 .as_mut()
460 .unwrap_unchecked()
461 }
462 }
463
464 pub fn get(&self) -> &Frame {
465 // Safety: guaranteed by `FrameList::entry`
466 unsafe {
467 self.entry
468 .get()
469 .frames
470 .get_unchecked(self.index_in_node)
471 .as_ref()
472 .unwrap_unchecked()
473 }
474 }
475 pub fn get_mut(&mut self) -> &mut Frame {
476 // Safety: guaranteed by `FrameList::entry`
477 unsafe {
478 Pin::into_inner_unchecked(self.entry.get_mut())
479 .frames
480 .get_unchecked_mut(self.index_in_node)
481 .as_mut()
482 .unwrap_unchecked()
483 }
484 }
485 pub fn insert(&mut self, frame: Frame) -> Frame {
486 // Safety: guaranteed by `FrameList::entry`
487 unsafe {
488 self.entry
489 .get_mut()
490 .frames
491 .get_unchecked_mut(self.index_in_node)
492 .replace(frame)
493 .unwrap_unchecked()
494 }
495 }
496
497 pub fn remove(&mut self) -> Frame {
498 // Safety: guaranteed by `FrameList::entry`
499 unsafe {
500 self.entry
501 .get_mut()
502 .frames
503 .get_unchecked_mut(self.index_in_node)
504 .take()
505 .unwrap_unchecked()
506 }
507 }
508}
509impl<'a> VacantEntry<'a> {
510 pub fn insert(self, value: Frame) -> &'a mut Frame {
511 let mut node = self.entry.or_insert_with(|| {
512 Box::pin(FrameListNode {
513 links: wavltree::Links::default(),
514 offset: self.offset,
515 frames: [const { None }; FRAME_LIST_NODE_FANOUT],
516 })
517 });
518 let old = node.frames[self.index_in_node].replace(value);
519 debug_assert!(old.is_none());
520
521 // Safety: guaranteed by `FrameList::entry`
522 unsafe {
523 Pin::into_inner_unchecked(node)
524 .frames
525 .get_unchecked_mut(self.index_in_node)
526 .as_mut()
527 .unwrap_unchecked()
528 }
529 }
530
531 pub fn insert_entry(self, frame: Frame) -> OccupiedEntry<'a> {
532 let mut entry = match self.entry {
533 wavltree::Entry::Occupied(entry) => entry,
534 wavltree::Entry::Vacant(entry) => entry.insert_entry(Box::pin(FrameListNode {
535 links: wavltree::Links::default(),
536 offset: self.offset,
537 frames: [const { None }; FRAME_LIST_NODE_FANOUT],
538 })),
539 };
540
541 // Safety: guaranteed by `FrameList::entry`
542 unsafe {
543 *entry.get_mut().frames.get_unchecked_mut(self.index_in_node) = Some(frame);
544 }
545
546 OccupiedEntry {
547 entry,
548 index_in_node: self.index_in_node,
549 }
550 }
551}