Skip to main content

wasefire_slice_cell/
lib.rs

1// Copyright 2025 Google LLC
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//! Slice with dynamic borrow checking.
16
17#![no_std]
18#![feature(slice_ptr_get)]
19
20extern crate alloc;
21
22use core::cell::RefCell;
23use core::marker::PhantomData;
24use core::ops::{Range, RangeBounds};
25
26#[cfg(feature = "_internal")]
27pub use internal::*;
28
29/// Slice with dynamic borrow checking.
30pub struct SliceCell<'a, T> {
31    _lifetime: PhantomData<&'a ()>,
32    data: *mut [T],
33    state: RefCell<internal::State>,
34}
35
36/// Access errors.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum Error {
39    /// Invalid range.
40    Range,
41
42    /// Invalid borrow.
43    Borrow,
44}
45
46impl<'a, T> SliceCell<'a, T> {
47    /// Returns the same exclusive slice but with dynamic borrow checking.
48    pub fn new(data: &'a mut [T]) -> Self {
49        Self { _lifetime: PhantomData, data, state: Default::default() }
50    }
51
52    /// Returns a shared reference to an element of the underlying slice.
53    pub fn get(&self, index: usize) -> Result<&T, Error> {
54        Ok(&self.get_range(index ..= index)?[0])
55    }
56
57    /// Returns an exclusive reference to an element of the underlying slice.
58    pub fn get_mut(&self, index: usize) -> Result<&mut T, Error> {
59        Ok(&mut self.get_range_mut(index ..= index)?[0])
60    }
61
62    /// Returns a shared slice reference of the underlying slice.
63    pub fn get_range(&self, range: impl RangeBounds<usize>) -> Result<&[T], Error> {
64        let range = self.range(range)?;
65        let false = range.is_empty() else { return Ok(&[]) };
66        // SAFETY: Checked by self.range() above.
67        let ptr = unsafe { self.data.as_mut_ptr().add(range.start) };
68        let len = range.len();
69        self.borrow(range)?;
70        // SAFETY: Checked by self.borrow() above.
71        Ok(unsafe { core::slice::from_raw_parts(ptr, len) })
72    }
73
74    /// Returns an exclusive slice reference of the underlying slice.
75    #[allow(clippy::mut_from_ref)]
76    pub fn get_range_mut(&self, range: impl RangeBounds<usize>) -> Result<&mut [T], Error> {
77        let range = self.range(range)?;
78        let false = range.is_empty() else { return Ok(&mut []) };
79        // SAFETY: Checked by self.range() above.
80        let ptr = unsafe { self.data.as_mut_ptr().add(range.start) };
81        let len = range.len();
82        self.borrow_mut(range)?;
83        // SAFETY: Checked by self.borrow_mut() above.
84        Ok(unsafe { core::slice::from_raw_parts_mut(ptr, len) })
85    }
86
87    /// Invalidates all references to the underlying slice.
88    pub fn reset(&mut self) {
89        self.state.take();
90    }
91
92    fn range(&self, range: impl RangeBounds<usize>) -> Result<Range<usize>, Error> {
93        internal::range_check(self.data.len(), range)
94    }
95
96    fn borrow(&self, range: Range<usize>) -> Result<(), Error> {
97        let access = internal::Access { exclusive: false, range };
98        internal::borrow_check(&mut self.state.borrow_mut(), access)
99    }
100
101    fn borrow_mut(&self, range: Range<usize>) -> Result<(), Error> {
102        let access = internal::Access { exclusive: true, range };
103        internal::borrow_check(&mut self.state.borrow_mut(), access)
104    }
105}
106
107#[cfg_attr(not(feature = "_internal"), allow(unreachable_pub))]
108mod internal {
109    use alloc::vec::Vec;
110    use core::ops::{Bound, Range, RangeBounds};
111
112    use crate::Error;
113
114    /// Makes sure the range is a sub-range of the underlying slice.
115    pub fn range_check(len: usize, range: impl RangeBounds<usize>) -> Result<Range<usize>, Error> {
116        let start = match range.start_bound() {
117            Bound::Included(x) => *x,
118            Bound::Excluded(x) => x.checked_add(1).ok_or(Error::Range)?,
119            Bound::Unbounded => 0,
120        };
121        let end = match range.end_bound() {
122            Bound::Included(x) => x.checked_add(1).ok_or(Error::Range)?,
123            Bound::Excluded(x) => *x,
124            Bound::Unbounded => len,
125        };
126        if start <= end && end <= len { Ok(start .. end) } else { Err(Error::Range) }
127    }
128
129    /// Sorted list of non-overlapping non-empty accesses since the last reset.
130    pub type State = Vec<Access>;
131
132    /// Describes an access into a slice.
133    #[cfg_attr(feature = "_internal", derive(Clone, PartialEq, Eq))]
134    pub struct Access {
135        /// Whether the access is shared or exclusive.
136        pub exclusive: bool,
137        /// The non-empty range that is accessed.
138        pub range: Range<usize>,
139    }
140
141    #[cfg(feature = "_internal")]
142    impl core::fmt::Debug for Access {
143        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
144            write!(f, "{}{:?}", if self.exclusive { "mut " } else { "" }, self.range)
145        }
146    }
147
148    /// Makes sure the non-empty range can be borrowed (shared or exclusively).
149    ///
150    /// If the range can be borrowed, the state is updated to reflect that access.
151    pub fn borrow_check(state: &mut State, new: Access) -> Result<(), Error> {
152        debug_assert!(!new.range.is_empty());
153        // Find the first existing access that ends after the new access starts.
154        let Some(i) = state.iter().position(|cur| new.range.start < cur.range.end) else {
155            // The new access does not overlap and is after all existing accesses.
156            state.push(new);
157            return Ok(());
158        };
159        // Find the first existing access that starts after the new access ends.
160        let j = match state[i ..].iter().position(|cur| new.range.end <= cur.range.start) {
161            None => state.len(),
162            Some(x) => i + x,
163        };
164        if i == j {
165            // The new access does not overlap and is before an existing access.
166            state.insert(i, new);
167            return Ok(());
168        }
169        // The new access overlaps with the existing accesses between i and j.
170        if new.exclusive || state[i .. j].iter().any(|x| x.exclusive) {
171            // Either the new or at least one of the existing accesses is exclusive.
172            return Err(Error::Borrow);
173        }
174        // Merge the new access with all the existing overlapping ones.
175        state[i].range.start = core::cmp::min(state[i].range.start, new.range.start);
176        state[i].range.end = core::cmp::max(state[j - 1].range.end, new.range.end);
177        state.drain(i + 1 .. j);
178        Ok(())
179    }
180}