sel4/state/
token.rs

1//
2// Copyright 2024, Colias Group, LLC
3//
4// SPDX-License-Identifier: MIT
5//
6
7use core::cell::{Ref, RefCell, RefMut, UnsafeCell};
8use core::fmt;
9use core::sync::atomic::{AtomicIsize, Ordering};
10
11pub(crate) struct TokenCell<K, A> {
12    token: K,
13    accessor: A,
14}
15
16pub(crate) trait Accessor<T> {
17    fn with<F, U>(&self, f: F) -> U
18    where
19        F: FnOnce(&UnsafeCell<T>) -> U;
20}
21
22impl<K: Token, A> TokenCell<K, A> {
23    pub(crate) const unsafe fn new(accessor: A) -> Self {
24        Self {
25            token: K::INIT,
26            accessor,
27        }
28    }
29
30    pub(crate) fn try_with<F, T, U>(&self, f: F) -> U
31    where
32        A: Accessor<T>,
33        F: FnOnce(Result<&T, BorrowError>) -> U,
34    {
35        let access = || {
36            self.accessor
37                .with(|cell| unsafe { cell.get().as_ref().unwrap() })
38        };
39        self.token.try_with(access, f)
40    }
41
42    pub(crate) fn try_with_mut<F, T, U>(&self, f: F) -> U
43    where
44        A: Accessor<T>,
45        F: FnOnce(Result<&mut T, BorrowMutError>) -> U,
46    {
47        let access = || {
48            self.accessor
49                .with(|cell| unsafe { cell.get().as_mut().unwrap() })
50        };
51        self.token.try_with_mut(access, f)
52    }
53}
54
55pub(crate) trait Token {
56    const INIT: Self;
57
58    type Borrow<'a>
59    where
60        Self: 'a;
61
62    type BorrowMut<'a>
63    where
64        Self: 'a;
65
66    fn try_borrow(&self) -> Result<Self::Borrow<'_>, BorrowError>;
67
68    fn try_borrow_mut(&self) -> Result<Self::BorrowMut<'_>, BorrowMutError>;
69
70    fn try_with<F, G, T, U>(&self, access_resource: G, f: F) -> T
71    where
72        F: FnOnce(Result<U, BorrowError>) -> T,
73        G: FnOnce() -> U,
74    {
75        let (_, r) = take_ok(self.try_borrow());
76        f(r.map(|_| access_resource()))
77    }
78
79    fn try_with_mut<F, G, T, U>(&self, access_resource: G, f: F) -> T
80    where
81        F: FnOnce(Result<U, BorrowMutError>) -> T,
82        G: FnOnce() -> U,
83    {
84        let (_, r) = take_ok(self.try_borrow_mut());
85        f(r.map(|_| access_resource()))
86    }
87}
88
89#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
90pub struct BorrowError(());
91
92impl BorrowError {
93    pub(crate) fn new() -> Self {
94        Self(())
95    }
96}
97
98impl fmt::Display for BorrowError {
99    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
100        write!(f, "already mutably borrowed")
101    }
102}
103
104#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
105pub struct BorrowMutError(());
106
107impl BorrowMutError {
108    pub(crate) fn new() -> Self {
109        Self(())
110    }
111}
112
113impl fmt::Display for BorrowMutError {
114    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
115        write!(f, "already borrowed")
116    }
117}
118
119fn take_ok<T, E>(r: Result<T, E>) -> (Option<T>, Result<(), E>) {
120    match r {
121        Ok(ok) => (Some(ok), Ok(())),
122        Err(err) => (None, Err(err)),
123    }
124}
125
126#[allow(dead_code)]
127pub(crate) struct UnsyncToken(RefCell<()>);
128
129impl Token for UnsyncToken {
130    #[allow(clippy::declare_interior_mutable_const)]
131    const INIT: Self = Self(RefCell::new(()));
132
133    type Borrow<'a> = Ref<'a, ()>;
134    type BorrowMut<'a> = RefMut<'a, ()>;
135
136    fn try_borrow(&self) -> Result<Self::Borrow<'_>, BorrowError> {
137        self.0.try_borrow().map_err(|_| BorrowError::new())
138    }
139
140    fn try_borrow_mut(&self) -> Result<Self::BorrowMut<'_>, BorrowMutError> {
141        self.0.try_borrow_mut().map_err(|_| BorrowMutError::new())
142    }
143}
144
145#[allow(dead_code)]
146pub(crate) struct SyncToken(BorrowFlag);
147
148type BorrowFlag = AtomicIsize;
149
150pub(crate) struct SyncTokenBorrow<'a>(&'a BorrowFlag);
151
152impl Drop for SyncTokenBorrow<'_> {
153    fn drop(&mut self) {
154        self.0.fetch_sub(1, Ordering::Release);
155    }
156}
157
158pub(crate) struct SyncTokenBorrowMut<'a>(&'a BorrowFlag);
159
160impl Drop for SyncTokenBorrowMut<'_> {
161    fn drop(&mut self) {
162        self.0.fetch_add(1, Ordering::Release);
163    }
164}
165
166impl Token for SyncToken {
167    #[allow(clippy::declare_interior_mutable_const)]
168    const INIT: Self = Self(AtomicIsize::new(0));
169
170    type Borrow<'a> = SyncTokenBorrow<'a>;
171    type BorrowMut<'a> = SyncTokenBorrowMut<'a>;
172
173    fn try_borrow(&self) -> Result<Self::Borrow<'_>, BorrowError> {
174        let mut current = self.0.load(Ordering::SeqCst);
175        loop {
176            if (0..isize::MAX).contains(&current) {
177                match self.0.compare_exchange(
178                    current,
179                    current + 1,
180                    Ordering::Acquire,
181                    Ordering::Relaxed,
182                ) {
183                    Ok(_) => return Ok(SyncTokenBorrow(&self.0)),
184                    Err(actual_current) => {
185                        current = actual_current;
186                    }
187                }
188            } else {
189                return Err(BorrowError::new());
190            }
191        }
192    }
193
194    fn try_borrow_mut(&self) -> Result<Self::BorrowMut<'_>, BorrowMutError> {
195        let mut current = self.0.load(Ordering::SeqCst);
196        loop {
197            if current == 0 {
198                match self.0.compare_exchange(
199                    current,
200                    current - 1,
201                    Ordering::Acquire,
202                    Ordering::Relaxed,
203                ) {
204                    Ok(_) => return Ok(SyncTokenBorrowMut(&self.0)),
205                    Err(actual_current) => {
206                        current = actual_current;
207                    }
208                }
209            } else {
210                return Err(BorrowMutError::new());
211            }
212        }
213    }
214}