crates/sel4-initialize-tls: Refactor interface

Signed-off-by: Nick Spinale <nick@nickspinale.com>

crates/sel4-initialize-tls/src/lib.rs

@@ -66,16 +66,22 @@ impl UncheckedTlsImage {
         if self.memsz >= self.filesz && self.align.is_power_of_two() && self.align > 0 {
             Ok(TlsImage { checked: *self })
         } else {
-            Err(InvalidTlsImageError(()))
+            Err(InvalidTlsImageError::new())
         }
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct InvalidTlsImageError(());
 
+impl InvalidTlsImageError {
+    fn new() -> Self {
+        Self(())
+    }
+}
+
 #[repr(C)]
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct TlsImage {
     checked: UncheckedTlsImage,
 }
@@ -94,7 +100,7 @@ impl TlsImage {
     }
 
     #[allow(clippy::missing_safety_doc)]
-    pub unsafe fn initialize_tls_reservation(&self, reservation_start: *mut u8) {
+    pub unsafe fn initialize_reservation(&self, reservation_start: *mut u8) -> usize {
         let reservation_layout = self.reservation_layout();
         let reservation =
             slice::from_raw_parts_mut(reservation_start, reservation_layout.footprint().size());
@@ -103,18 +109,54 @@ impl TlsImage {
             .split_at_mut(self.checked.filesz);
         tdata.copy_from_slice(self.image_data().as_ref().unwrap());
         tbss.fill(0);
+        let thread_pointer = (reservation_start as usize)
+            .checked_add(reservation_layout.thread_pointer_offset())
+            .unwrap(); // TODO return error
         if cfg!(target_arch = "x86_64") {
-            let thread_pointer = (reservation_start as usize)
-                .checked_add(reservation_layout.thread_pointer_offset())
-                .unwrap();
             let thread_pointer_slice = &mut reservation
                 [reservation_layout.thread_pointer_offset()..][..mem::size_of::<usize>()];
             thread_pointer_slice.copy_from_slice(&thread_pointer.to_ne_bytes());
         }
+        thread_pointer
+    }
+
+    #[allow(clippy::missing_safety_doc)]
+    pub unsafe fn initialize_exact_reservation_region(
+        &self,
+        exact_reservation: &Region,
+    ) -> Result<usize, RegionLayoutError> {
+        if exact_reservation.fits_exactly(self.reservation_layout().footprint()) {
+            Ok(self.initialize_reservation(exact_reservation.start()))
+        } else {
+            Err(RegionLayoutError::new())
+        }
+    }
+
+    #[allow(clippy::missing_safety_doc)]
+    pub unsafe fn initialize_inexact_reservation_region(
+        &self,
+        inexact_reservation: &Region,
+    ) -> Result<usize, RegionLayoutError> {
+        if let Ok(TrimmedRegion { trimmed, .. }) =
+            inexact_reservation.trim(self.reservation_layout().footprint())
+        {
+            Ok(self.initialize_exact_reservation_region(&trimmed).unwrap())
+        } else {
+            Err(RegionLayoutError::new())
+        }
     }
 }
 
-#[derive(Copy, Clone, Debug)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct RegionLayoutError(());
+
+impl RegionLayoutError {
+    fn new() -> Self {
+        Self(())
+    }
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct TlsReservationLayout {
     footprint: Layout,
     segment_offset: usize,
@@ -165,3 +207,71 @@ impl TlsReservationLayout {
         self.thread_pointer_offset
     }
 }
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct Region {
+    start: *mut u8,
+    size: usize,
+}
+
+impl Region {
+    pub const fn new(start: *mut u8, size: usize) -> Self {
+        Self { start, size }
+    }
+
+    pub const fn start(&self) -> *mut u8 {
+        self.start
+    }
+
+    pub const fn size(&self) -> usize {
+        self.size
+    }
+
+    fn fits_exactly(&self, layout: Layout) -> bool {
+        self.size() == layout.size() && self.start().align_offset(layout.align()) == 0
+    }
+
+    fn trim(&self, layout: Layout) -> Result<TrimmedRegion, TrimRegionError> {
+        let start_addr = self.start() as usize;
+        let trimmed_start_addr = start_addr
+            .checked_next_multiple_of(layout.align())
+            .ok_or(TrimRegionError::new())?;
+        let remainder_start_addr = trimmed_start_addr
+            .checked_add(layout.size())
+            .ok_or(TrimRegionError::new())?;
+        let remainder_end_addr = start_addr
+            .checked_add(self.size())
+            .ok_or(TrimRegionError::new())?;
+        if remainder_start_addr > remainder_end_addr {
+            return Err(TrimRegionError::new());
+        }
+        Ok(TrimmedRegion {
+            padding: Region::new(start_addr as *mut u8, trimmed_start_addr - start_addr),
+            trimmed: Region::new(
+                trimmed_start_addr as *mut u8,
+                remainder_start_addr - trimmed_start_addr,
+            ),
+            remainder: Region::new(
+                remainder_start_addr as *mut u8,
+                remainder_end_addr - remainder_start_addr,
+            ),
+        })
+    }
+}
+
+struct TrimmedRegion {
+    #[allow(dead_code)]
+    padding: Region,
+    trimmed: Region,
+    #[allow(dead_code)]
+    remainder: Region,
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+struct TrimRegionError(());
+
+impl TrimRegionError {
+    fn new() -> Self {
+        Self(())
+    }
+}

crates/sel4-initialize-tls/src/on_heap.rs

@@ -4,32 +4,39 @@
 // SPDX-License-Identifier: BSD-2-Clause
 //
 
-use crate::{TlsImage, TlsReservationLayout};
+use core::alloc::Layout;
+
+use crate::{Region, TlsImage};
 
 pub struct HeapTlsReservation {
     start: *mut u8,
-    layout: TlsReservationLayout,
+    layout: Layout,
+    thread_pointer: usize,
 }
 
 impl HeapTlsReservation {
     fn initialize(tls_image: &TlsImage) -> Self {
-        let layout = tls_image.reservation_layout();
-        let start = unsafe { ::alloc::alloc::alloc(layout.footprint()) };
-        unsafe {
-            tls_image.initialize_tls_reservation(start);
-        };
-        Self { start, layout }
+        let layout = tls_image.reservation_layout().footprint();
+        let start = unsafe { ::alloc::alloc::alloc(layout) };
+        let region = Region::new(start, layout.size());
+        let thread_pointer =
+            unsafe { tls_image.initialize_exact_reservation_region(&region) }.unwrap();
+        Self {
+            start,
+            layout,
+            thread_pointer,
+        }
     }
 
     pub fn thread_pointer(&self) -> usize {
-        (self.start as usize) + self.layout.thread_pointer_offset()
+        self.thread_pointer
     }
 }
 
 impl Drop for HeapTlsReservation {
     fn drop(&mut self) {
         unsafe {
-            ::alloc::alloc::dealloc(self.start, self.layout.footprint());
+            ::alloc::alloc::dealloc(self.start, self.layout);
         }
     }
 }

crates/sel4-initialize-tls/src/on_stack/mod.rs

@@ -44,7 +44,7 @@ impl TlsImage {
         cont_fn: ContFn,
         cont_arg: *mut ContArg,
     ) -> ! {
-        self.initialize_tls_reservation(tls_reservation_start);
+        self.initialize_reservation(tls_reservation_start);
 
         let thread_pointer = tls_reservation_start
             .wrapping_byte_add(self.reservation_layout().thread_pointer_offset());

crates/sel4-initialize-tls/src/static_allocation.rs

@@ -4,10 +4,9 @@
 // SPDX-License-Identifier: BSD-2-Clause
 //
 
-use core::alloc::Layout;
 use core::cell::UnsafeCell;
 
-use crate::TlsImage;
+use crate::Region;
 
 #[repr(C)]
 pub struct StaticTlsAllocation<const N: usize, A = ()> {
@@ -32,40 +31,14 @@ impl<const N: usize, A> StaticTlsAllocation<N, A> {
     const fn start(&self) -> *mut u8 {
         self.space.get().cast()
     }
+
+    pub const fn region(&self) -> Region {
+        Region::new(self.start(), self.size())
+    }
 }
 
 impl<const N: usize, A> Default for StaticTlsAllocation<N, A> {
     fn default() -> Self {
         Self::new()
     }
 }
-
-impl TlsImage {
-    pub fn initialize_static_allocation<const N: usize, A>(
-        &self,
-        allocation: &StaticTlsAllocation<N, A>,
-    ) -> Result<usize, InitializeStaticTlsAllocationError> {
-        let layout = self.reservation_layout();
-        let footprint = layout.footprint();
-        let align_offset = allocation.start().align_offset(layout.footprint().align());
-        if align_offset + footprint.size() > allocation.size() {
-            return Err(InitializeStaticTlsAllocationError::AllocationTooSmall {
-                requested_footprint: footprint,
-                align_offset,
-            });
-        }
-        let start = allocation.start().wrapping_byte_add(align_offset);
-        unsafe {
-            self.initialize_tls_reservation(start);
-        };
-        Ok((allocation.start() as usize) + layout.thread_pointer_offset())
-    }
-}
-
-#[derive(Debug, Copy, Clone)]
-pub enum InitializeStaticTlsAllocationError {
-    AllocationTooSmall {
-        requested_footprint: Layout,
-        align_offset: usize,
-    },
-}