fix: use non-contiguous allocations wherever possible (#364)

This change brings back non-contiguous physical allocations in places where it makes sense to have them. This helps to make better use of memory especially as we approach the maxmimum system memory.

kernel/src/vm/frame_alloc/arena.rs

···
         };
 
         for region in self.inner.by_ref() {
-            let pages_in_hole = region.start.checked_sub_addr(arena.end).unwrap() / arch::PAGE_SIZE;
+            tracing::debug!(arena.end=?arena.end,region=?region, "Attempting to add free region");
+
+            debug_assert!(!arena.is_overlapping(&region));
+
+            let pages_in_hole = if arena.end <= region.start {
+                // the region is higher than the current arena
+                region.start.checked_sub_addr(arena.end).unwrap() / arch::PAGE_SIZE
+            } else {
+                debug_assert!(region.end <= arena.start);
+                // the region is lower than the current arena
+                arena.start.checked_sub_addr(region.end).unwrap() / arch::PAGE_SIZE
+            };
+
             let waste_from_hole = ARENA_PAGE_BOOKKEEPING_SIZE * pages_in_hole;
 
             if self.wasted_bytes + waste_from_hole > MAX_WASTED_ARENA_BYTES {
···
             } else {
                 self.wasted_bytes += waste_from_hole;
                 arena.end = region.end;
+
+                if arena.end <= region.start {
+                    arena.end = region.end;
+                } else {
+                    arena.start = region.start;
+                }
             }
         }
 

loader/src/arch/riscv64.rs

···
     );
     debug_assert!(
         virt % PAGE_SIZE == 0,
-        "virtual address must be aligned to at least 4KiB page size {virt:?}"
+        "virtual address must be aligned to at least 4KiB page size ({virt:#x})"
     );
     debug_assert!(
         phys % PAGE_SIZE == 0,
-        "physical address must be aligned to at least 4KiB page size {phys:?}"
+        "physical address must be aligned to at least 4KiB page size ({phys:#x})"
     );
 
     // To map out contiguous chunk of physical memory into the virtual address space efficiently
···
                     // we need to allocate a new sub-table and retry.
                     // allocate a new physical frame to hold the next level table and
                     // mark this PTE as a valid internal node pointing to that sub-table.
-                    let frame = frame_alloc
-                        .allocate_one_zeroed(phys_off)
-                        .ok_or(Error::NoMemory)?; // we should always be able to map a single page
+                    let frame = frame_alloc.allocate_one_zeroed(phys_off)?; // we should always be able to map a single page
 
                     // TODO memory barrier
 
···
                 pgtable = pgtable_ptr_from_phys(pte.get_address_and_flags().0, phys_off);
             } else {
                 unreachable!(
-                    "Invalid state: PTE can't be valid leaf (this means {virt:?} is already mapped) {pte:?} {pte:p}"
+                    "Invalid state: PTE can't be valid leaf (this means {virt:#x} is already mapped) {pte:?} {pte:p}"
                 );
             }
         }

loader/src/boot_info.rs

···
 // copied, modified, or distributed except according to those terms.
 
 use crate::arch;
-use crate::error::Error;
 use crate::frame_alloc::FrameAllocator;
 use core::alloc::Layout;
 use core::mem::MaybeUninit;
···
     hart_mask: usize,
     rng_seed: [u8; 32],
 ) -> crate::Result<*mut BootInfo> {
-    let frame = frame_alloc
-        .allocate_contiguous_zeroed(
-            Layout::from_size_align(arch::PAGE_SIZE, arch::PAGE_SIZE).unwrap(),
-            arch::KERNEL_ASPACE_BASE,
-        )
-        .ok_or(Error::NoMemory)?;
+    let frame = frame_alloc.allocate_contiguous_zeroed(
+        Layout::from_size_align(arch::PAGE_SIZE, arch::PAGE_SIZE).unwrap(),
+        arch::KERNEL_ASPACE_BASE,
+    )?;
     let page = physical_address_offset.checked_add(frame).unwrap();
 
-    let memory_regions = init_boot_info_memory_regions(page, frame_alloc, fdt_phys, loader_phys);
+    let memory_regions =
+        init_boot_info_memory_regions(page, frame_alloc, fdt_phys, loader_phys, kernel_phys);
 
     let mut boot_info = BootInfo::new(memory_regions);
     boot_info.physical_address_offset = physical_address_offset;
···
     frame_alloc: FrameAllocator,
     fdt_phys: Range<usize>,
     loader_phys: Range<usize>,
+    kernel_phys: Range<usize>,
 ) -> MemoryRegions {
     // Safety: we just allocated a whole frame for the boot info
     let regions: &mut [MaybeUninit<MemoryRegion>] = unsafe {
···
 
     let mut len = 0;
     let mut push_region = |region: MemoryRegion| {
+        debug_assert!(!region.range.is_empty());
         regions[len].write(region);
         len += 1;
     };
···
         });
     }
 
-    // The memory occupied by the loader is not needed once the kernel is running.
-    // Mark it as usable.
+    // Most of the memory occupied by the loader is not needed once the kernel is running,
+    // but the kernel itself lies somewhere in the loader memory.
+    //
+    // We can still mark the range before and after the kernel as usable.
     push_region(MemoryRegion {
-        range: loader_phys,
+        range: Range::from(loader_phys.start..kernel_phys.start),
+        kind: MemoryRegionKind::Usable,
+    });
+    push_region(MemoryRegion {
+        range: Range::from(kernel_phys.end..loader_phys.end),
         kind: MemoryRegionKind::Usable,
     });
 

loader/src/frame_alloc.rs

···
 // copied, modified, or distributed except according to those terms.
 
 use crate::arch;
+use crate::error::Error;
 use core::alloc::Layout;
+use core::num::NonZeroUsize;
 use core::range::Range;
-use core::{iter, ptr, slice};
+use core::{cmp, iter, ptr, slice};
+use fallible_iterator::FallibleIterator;
 
 pub struct FrameAllocator<'a> {
     regions: &'a [Range<usize>],
     // offset from the top of memory regions
     offset: usize,
-    phys_offset: usize,
 }
 
 impl<'a> FrameAllocator<'a> {
     /// Create a new frame allocator over a given set of physical memory regions.
     #[must_use]
     pub fn new(regions: &'a [Range<usize>]) -> Self {
-        Self {
-            regions,
-            offset: 0,
-            phys_offset: 0,
-        }
-    }
-
-    pub fn set_phys_offset(&mut self, phys_offset: usize) {
-        self.phys_offset = phys_offset;
+        Self { regions, offset: 0 }
     }
 
     #[must_use]
···
         self.offset >> arch::PAGE_SHIFT
     }
 
-    pub fn allocate_one_zeroed(&mut self, phys_offset: usize) -> Option<usize> {
+    pub fn allocate_one_zeroed(&mut self, phys_offset: usize) -> Result<usize, Error> {
         self.allocate_contiguous_zeroed(
             // Safety: the layout is always valid
             unsafe { Layout::from_size_align_unchecked(arch::PAGE_SIZE, arch::PAGE_SIZE) },
···
         )
     }
 
-    pub fn allocate_contiguous(&mut self, layout: Layout) -> Option<usize> {
+    pub fn allocate(&mut self, layout: Layout) -> FrameIter<'a, '_> {
+        assert_eq!(
+            layout.align(),
+            arch::PAGE_SIZE,
+            "BootstrapAllocator only supports page-aligned allocations"
+        );
+
+        let remaining = layout.pad_to_align().size();
+
+        debug_assert!(remaining % arch::PAGE_SIZE == 0);
+        FrameIter {
+            alloc: self,
+            remaining,
+        }
+    }
+
+    pub fn allocate_zeroed(
+        &mut self,
+        layout: Layout,
+        phys_offset: usize,
+    ) -> FrameIterZeroed<'a, '_> {
+        FrameIterZeroed {
+            inner: self.allocate(layout),
+            phys_offset,
+        }
+    }
+
+    pub fn allocate_contiguous(&mut self, layout: Layout) -> Result<usize, Error> {
         let requested_size = layout.pad_to_align().size();
         assert_eq!(
             layout.align(),
···
 
                 let frame = region.end.checked_sub(offset + requested_size).unwrap();
                 self.offset += requested_size;
-
-                return Some(frame);
+                return Ok(frame);
             }
 
             offset -= region_size;
         }
 
-        None
+        Err(Error::NoMemory)
     }
 
     pub fn allocate_contiguous_zeroed(
         &mut self,
         layout: Layout,
         phys_offset: usize,
-    ) -> Option<usize> {
+    ) -> Result<usize, Error> {
         let requested_size = layout.pad_to_align().size();
         let addr = self.allocate_contiguous(layout)?;
         // Safety: we just allocated the frame
···
                 requested_size,
             );
         }
-        Some(addr)
+        Ok(addr)
+    }
+}
+
+pub struct FrameIter<'a, 'b> {
+    alloc: &'b mut FrameAllocator<'a>,
+    remaining: usize,
+}
+
+impl<'a> FrameIter<'a, '_> {
+    pub fn alloc(&mut self) -> &mut FrameAllocator<'a> {
+        self.alloc
+    }
+}
+
+impl FallibleIterator for FrameIter<'_, '_> {
+    type Item = (usize, NonZeroUsize);
+    type Error = Error;
+
+    fn next(&mut self) -> Result<Option<Self::Item>, Self::Error> {
+        if self.remaining > 0 {
+            let mut offset = self.alloc.offset;
+
+            for region in self.alloc.regions.iter().rev() {
+                let region_size = region.end.checked_sub(region.start).unwrap();
+
+                // only consider regions that we haven't already exhausted
+                if let Some(allocatable_size) = region_size.checked_sub(offset)
+                    && allocatable_size >= arch::PAGE_SIZE
+                {
+                    let allocation_size = cmp::min(self.remaining, allocatable_size)
+                        & 0usize.wrapping_sub(arch::PAGE_SIZE);
+                    debug_assert!(allocation_size % arch::PAGE_SIZE == 0);
+
+                    let frame = region.end.checked_sub(offset + allocation_size).unwrap();
+                    self.alloc.offset += allocation_size;
+                    self.remaining -= allocation_size;
+
+                    return Ok(Some((frame, NonZeroUsize::new(allocation_size).unwrap())));
+                }
+
+                offset -= region_size;
+            }
+
+            Err(Error::NoMemory)
+        } else {
+            Ok(None)
+        }
+    }
+}
+
+pub struct FrameIterZeroed<'a, 'b> {
+    inner: FrameIter<'a, 'b>,
+    phys_offset: usize,
+}
+
+impl<'a> FrameIterZeroed<'a, '_> {
+    pub fn alloc(&mut self) -> &mut FrameAllocator<'a> {
+        self.inner.alloc
+    }
+}
+
+impl FallibleIterator for FrameIterZeroed<'_, '_> {
+    type Item = (usize, NonZeroUsize);
+    type Error = Error;
+
+    fn next(&mut self) -> Result<Option<Self::Item>, Self::Error> {
+        let Some((base, len)) = self.inner.next()? else {
+            return Ok(None);
+        };
+
+        // Safety: we just allocated the frame
+        unsafe {
+            ptr::write_bytes::<u8>(
+                self.phys_offset.checked_add(base).unwrap() as *mut u8,
+                0,
+                len.get(),
+            );
+        }
+
+        Ok(Some((base, len)))
     }
 }
 

loader/src/main.rs

···
     map_physical_memory,
 };
 use arrayvec::ArrayVec;
-use core::alloc::Layout;
 use core::ffi::c_void;
 use core::range::Range;
-use core::{ptr, slice};
+use core::slice;
 use rand::SeedableRng;
 use rand_chacha::ChaCha20Rng;
 use spin::{Barrier, OnceLock};
···
     let self_regions = SelfRegions::collect(&minfo);
     log::debug!("{self_regions:#x?}");
 
+    let fdt_phys = {
+        let fdt = minfo.fdt.as_ptr_range();
+        Range::from(fdt.start as usize..fdt.end as usize)
+    };
+    let kernel_phys = {
+        let fdt = INLINED_KERNEL_BYTES.0.as_ptr_range();
+        Range::from(fdt.start as usize..fdt.end as usize)
+    };
+
     // Initialize the frame allocator
-    let allocatable_memories = allocatable_memory_regions(&minfo, &self_regions);
+    let allocatable_memories = allocatable_memory_regions(&minfo, &self_regions, fdt_phys);
+    log::debug!("allocatable memory regions {allocatable_memories:#x?}");
     let mut frame_alloc = FrameAllocator::new(&allocatable_memories);
 
     // initialize the random number generator
···
 
     // Initialize the page allocator
     let mut page_alloc = page_alloc::init(rng);
-
-    let fdt_phys = allocate_and_copy(&mut frame_alloc, minfo.fdt).unwrap();
-    let kernel_phys = allocate_and_copy(&mut frame_alloc, &INLINED_KERNEL_BYTES.0).unwrap();
 
     let root_pgtable = frame_alloc
         .allocate_one_zeroed(
···
         arch::activate_aspace(root_pgtable);
         log::trace!("activated.");
     }
-    frame_alloc.set_phys_offset(phys_off);
 
     // Safety: The kernel elf file is inlined into the loader executable as part of the build setup
     // which means we just need to parse it here.
     let kernel = parse_kernel(unsafe {
         let base = phys_off.checked_add(kernel_phys.start).unwrap();
-        let len = kernel_phys.end.checked_sub(kernel_phys.start).unwrap();
 
-        slice::from_raw_parts(base as *mut u8, len)
+        slice::from_raw_parts(base as *mut u8, INLINED_KERNEL_BYTES.0.len())
     })
     .unwrap();
     // print the elf sections for debugging purposes
···
 fn allocatable_memory_regions(
     minfo: &MachineInfo,
     self_regions: &SelfRegions,
+    fdt: Range<usize>,
 ) -> ArrayVec<Range<usize>, 16> {
-    let mut out = ArrayVec::new();
-    let to_exclude = Range::from(self_regions.executable.start..self_regions.read_write.end);
+    let mut temp: ArrayVec<Range<usize>, 16> = minfo.memories.clone();
 
-    for mut region in minfo.memories.clone() {
-        if to_exclude.contains(&region.start) && to_exclude.contains(&region.end) {
-            // remove region
-            continue;
-        } else if region.contains(&to_exclude.start) && region.contains(&to_exclude.end) {
-            out.push(Range::from(region.start..to_exclude.start));
-            out.push(Range::from(to_exclude.end..region.end));
-        } else if to_exclude.contains(&region.start) {
-            region.start = to_exclude.end;
-            out.push(region);
-        } else if to_exclude.contains(&region.end) {
-            region.end = to_exclude.start;
-            out.push(region);
-        } else {
-            out.push(region);
+    let mut exclude = |to_exclude: Range<usize>| {
+        for mut region in temp.take() {
+            if to_exclude.contains(&region.start) && to_exclude.contains(&region.end) {
+                // remove region
+                continue;
+            } else if region.contains(&to_exclude.start) && region.contains(&to_exclude.end) {
+                temp.push(Range::from(region.start..to_exclude.start));
+                temp.push(Range::from(to_exclude.end..region.end));
+            } else if to_exclude.contains(&region.start) {
+                region.start = to_exclude.end;
+                temp.push(region);
+            } else if to_exclude.contains(&region.end) {
+                region.end = to_exclude.start;
+                temp.push(region);
+            } else {
+                temp.push(region);
+            }
         }
-    }
+    };
 
-    out
-}
+    exclude(Range::from(
+        self_regions.executable.start..self_regions.read_write.end,
+    ));
 
-fn allocate_and_copy(frame_alloc: &mut FrameAllocator, src: &[u8]) -> Result<Range<usize>> {
-    let layout = Layout::from_size_align(src.len(), arch::PAGE_SIZE).unwrap();
-    let base = frame_alloc
-        .allocate_contiguous(layout)
-        .ok_or(Error::NoMemory)?;
+    exclude(fdt);
 
-    // Safety: we just allocated the frame
-    unsafe {
-        let dst = slice::from_raw_parts_mut(base as *mut u8, src.len());
+    // // merge adjacent regions
+    // let mut out: ArrayVec<Range<usize>, 16> = ArrayVec::new();
+    // 'outer: for region in temp {
+    //     for other in &mut out {
+    //         if region.start == other.end {
+    //             other.end = region.end;
+    //             continue 'outer;
+    //         }
+    //         if region.end == other.start {
+    //             other.start = region.start;
+    //             continue 'outer;
+    //         }
+    //     }
+    //
+    //     out.push(region);
+    // }
 
-        ptr::copy_nonoverlapping(src.as_ptr(), dst.as_mut_ptr(), dst.len());
+    temp.sort_unstable_by_key(|region| region.start);
+
+    #[cfg(debug_assertions)]
+    for (i, region) in temp.iter().enumerate() {
+        for (j, other) in temp.iter().enumerate() {
+            if i == j {
+                continue;
+            }
+
+            assert!(!other.contains(&region.start) && !other.contains(&region.end));
+        }
     }
 
-    Ok(Range::from(base..base.checked_add(layout.size()).unwrap()))
+    temp
 }

loader/src/mapping.rs

···
 use crate::{SelfRegions, arch};
 use bitflags::bitflags;
 use core::alloc::Layout;
-use core::num::{NonZero, NonZeroUsize};
+use core::num::NonZeroUsize;
 use core::range::Range;
 use core::{cmp, ptr, slice};
+use fallible_iterator::FallibleIterator;
 use loader_api::TlsTemplate;
 use xmas_elf::P64;
 use xmas_elf::dynamic::Tag;
···
             ph.align,
         );
 
-        let new_frame = frame_alloc
-            .allocate_contiguous_zeroed(
-                Layout::from_size_align(arch::PAGE_SIZE, arch::PAGE_SIZE).unwrap(),
-                arch::KERNEL_ASPACE_BASE,
-            )
-            .ok_or(Error::NoMemory)?;
+        let new_frame = frame_alloc.allocate_one_zeroed(arch::KERNEL_ASPACE_BASE)?;
 
         // Safety: we just allocated the frame
         unsafe {
···
     }
 
     log::trace!("zero_start {zero_start:#x} zero_end {zero_end:#x}");
-    let (additional_virt, additional_len) = {
+    let (mut virt, len) = {
         // zero_start either lies at a page boundary OR somewhere within the first page
         // by aligning up, we move it to the beginning of the *next* page.
         let start = checked_align_up(zero_start, ph.align).unwrap();
···
         (start, end.checked_sub(start).unwrap())
     };
 
-    if additional_len > 0 {
-        let additional_phys = frame_alloc
-            .allocate_contiguous_zeroed(
-                Layout::from_size_align(additional_len, arch::PAGE_SIZE).unwrap(),
-                arch::KERNEL_ASPACE_BASE,
-            )
-            .unwrap();
-
-        log::trace!(
-            "mapping additional zeros {additional_virt:#x}..{:#x}",
-            additional_virt.checked_add(additional_len).unwrap()
+    if len > 0 {
+        let mut phys_iter = frame_alloc.allocate_zeroed(
+            Layout::from_size_align(len, arch::PAGE_SIZE).unwrap(),
+            arch::KERNEL_ASPACE_BASE,
         );
-        // Safety: Leaving the address space in an invalid state here is fine since on panic we'll
-        // abort startup anyway
-        unsafe {
-            arch::map_contiguous(
-                root_pgtable,
-                frame_alloc,
-                additional_virt,
-                additional_phys,
-                NonZeroUsize::new(additional_len).unwrap(),
-                flags,
-                arch::KERNEL_ASPACE_BASE,
-            )?;
+
+        while let Some((phys, len)) = phys_iter.next()? {
+            log::trace!(
+                "mapping additional zeros {virt:#x}..{:#x}",
+                virt.checked_add(len.get()).unwrap()
+            );
+
+            // Safety: Leaving the address space in an invalid state here is fine since on panic we'll
+            // abort startup anyway
+            unsafe {
+                arch::map_contiguous(
+                    root_pgtable,
+                    phys_iter.alloc(),
+                    virt,
+                    phys,
+                    len,
+                    flags,
+                    arch::KERNEL_ASPACE_BASE,
+                )?;
+            }
+
+            virt += len.get();
         }
     }
 
···
         .unwrap()
         .0
         .pad_to_align();
+    log::trace!("allocating TLS segment {layout:?}...");
 
-    let phys = frame_alloc
-        .allocate_contiguous_zeroed(layout, phys_off)
-        .unwrap();
     let virt = page_alloc.allocate(layout);
+    let mut virt_start = virt.start;
 
-    log::trace!("Mapping TLS region {virt:#x?}...");
-    // Safety: Leaving the address space in an invalid state here is fine since on panic we'll
-    // abort startup anyway
-    unsafe {
-        arch::map_contiguous(
-            root_pgtable,
-            frame_alloc,
-            virt.start,
-            phys,
-            NonZero::new(layout.size()).unwrap(),
-            Flags::READ | Flags::WRITE,
-            phys_off,
-        )?;
+    let mut phys_iter = frame_alloc.allocate_zeroed(layout, phys_off);
+    while let Some((phys, len)) = phys_iter.next()? {
+        log::trace!(
+            "Mapping TLS region {virt_start:#x}..{:#x} {len} ...",
+            virt_start.checked_add(len.get()).unwrap()
+        );
+
+        // Safety: Leaving the address space in an invalid state here is fine since on panic we'll
+        // abort startup anyway
+        unsafe {
+            arch::map_contiguous(
+                root_pgtable,
+                phys_iter.alloc(),
+                virt_start,
+                phys,
+                len,
+                Flags::READ | Flags::WRITE,
+                phys_off,
+            )?;
+        }
+
+        virt_start += len.get();
     }
 
     Ok(TlsAllocation {
···
     let virt = page_alloc.allocate(layout_with_guard);
     log::trace!("Mapping stacks region {virt:#x?}...");
 
-    // Mapping stacks region 0xffffffc0c0000000..0xffffffc0c0101000...
-
     for hart in 0..minfo.hart_mask.count_ones() {
         let layout = Layout::from_size_align(per_cpu_size, arch::PAGE_SIZE).unwrap();
 
-        log::trace!("Allocating stack {layout:?}...");
-        // The stacks region doesn't need to be zeroed, since we will be filling it with
-        // the canary pattern anyway
-        let phys = frame_alloc
-            .allocate_contiguous(layout)
-            .ok_or(Error::NoMemory)?;
-
-        let virt = virt
+        let mut virt = virt
             .end
             .checked_sub(per_cpu_size_with_guard * hart as usize)
             .and_then(|a| a.checked_sub(per_cpu_size))
             .unwrap();
 
-        log::trace!(
-            "mapping stack for hart {hart} {virt:#x}..{:#x} => {phys:#x}..{:#x}",
-            virt.checked_add(per_cpu_size).unwrap(),
-            phys.checked_add(per_cpu_size).unwrap()
-        );
+        log::trace!("Allocating stack {layout:?}...");
+        // The stacks region doesn't need to be zeroed, since we will be filling it with
+        // the canary pattern anyway
+        let mut phys_iter = frame_alloc.allocate(layout);
 
-        // Safety: Leaving the address space in an invalid state here is fine since on panic we'll
-        // abort startup anyway
-        unsafe {
-            arch::map_contiguous(
-                root_pgtable,
-                frame_alloc,
-                virt,
-                phys,
-                NonZero::new(layout.size()).unwrap(),
-                Flags::READ | Flags::WRITE | Flags::USER,
-                phys_off,
-            )?;
+        while let Some((phys, len)) = phys_iter.next()? {
+            log::trace!(
+                "mapping stack for hart {hart} {virt:#x}..{:#x} => {phys:#x}..{:#x}",
+                virt.checked_add(len.get()).unwrap(),
+                phys.checked_add(len.get()).unwrap()
+            );
+
+            // Safety: Leaving the address space in an invalid state here is fine since on panic we'll
+            // abort startup anyway
+            unsafe {
+                arch::map_contiguous(
+                    root_pgtable,
+                    phys_iter.alloc(),
+                    virt,
+                    phys,
+                    len,
+                    Flags::READ | Flags::WRITE | Flags::USER,
+                    phys_off,
+                )?;
+            }
+
+            virt += len.get();
         }
     }