ahrightos_memory/
lib.rs

1#![cfg_attr(not(test), no_std)]
2
3pub const PAGE_SIZE: usize = 4096;
4pub const PAGE_TABLE_ENTRIES: usize = 512;
5pub const VM_REGION_MAX: usize = 32;
6pub const VM_REGION_MAX_PAGES: usize = 128;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
9#[repr(transparent)]
10pub struct VirtAddr(usize);
11
12impl VirtAddr {
13    pub const fn new(addr: usize) -> Self {
14        Self(addr)
15    }
16
17    pub const fn as_usize(self) -> usize {
18        self.0
19    }
20
21    pub fn as_mut_ptr<T>(self) -> *mut T {
22        self.0 as *mut T
23    }
24
25    pub const fn page_offset(self) -> usize {
26        self.0 & 0xFFF
27    }
28
29    pub const fn vpn0(self) -> usize {
30        (self.0 >> 12) & 0x1FF
31    }
32
33    pub const fn vpn1(self) -> usize {
34        (self.0 >> 21) & 0x1FF
35    }
36
37    pub const fn vpn2(self) -> usize {
38        (self.0 >> 30) & 0x1FF
39    }
40}
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
43#[repr(transparent)]
44pub struct PhysAddr(usize);
45
46impl PhysAddr {
47    pub const fn new(addr: usize) -> Self {
48        Self(addr)
49    }
50
51    pub const fn as_usize(self) -> usize {
52        self.0
53    }
54}
55
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub struct Page {
58    start_address: VirtAddr,
59}
60
61impl Page {
62    pub fn containing_address(addr: VirtAddr) -> Self {
63        Self {
64            start_address: VirtAddr::new(addr.as_usize() & !(PAGE_SIZE - 1)),
65        }
66    }
67
68    pub fn start_address(self) -> VirtAddr {
69        self.start_address
70    }
71}
72
73#[derive(Debug, Clone, Copy, PartialEq, Eq)]
74pub struct PhysFrame {
75    start_address: PhysAddr,
76}
77
78impl PhysFrame {
79    pub fn containing_address(addr: PhysAddr) -> Self {
80        Self {
81            start_address: PhysAddr::new(addr.as_usize() & !(PAGE_SIZE - 1)),
82        }
83    }
84
85    pub fn start_address(self) -> PhysAddr {
86        self.start_address
87    }
88}
89
90#[repr(u64)]
91#[derive(Debug, Clone, Copy, PartialEq, Eq)]
92pub enum PageTableFlags {
93    Valid = 1 << 0,
94    Read = 1 << 1,
95    Write = 1 << 2,
96    Execute = 1 << 3,
97    User = 1 << 4,
98    Global = 1 << 5,
99    Accessed = 1 << 6,
100    Dirty = 1 << 7,
101}
102
103#[repr(transparent)]
104#[derive(Debug, Clone, Copy, PartialEq, Eq)]
105pub struct PageTableEntry {
106    entry: usize,
107}
108
109impl PageTableEntry {
110    pub const fn new() -> Self {
111        Self { entry: 0 }
112    }
113
114    pub fn is_valid(&self) -> bool {
115        (self.entry & PageTableFlags::Valid as usize) != 0
116    }
117
118    pub fn is_leaf(&self) -> bool {
119        (self.entry & 0xE) != 0
120    }
121
122    pub fn ppn(&self) -> usize {
123        (self.entry >> 10) & 0xFFF_FFFF_FFFF
124    }
125
126    pub fn phys_addr(&self) -> PhysAddr {
127        PhysAddr::new(self.ppn() << 12)
128    }
129
130    pub fn set(&mut self, ppn: usize, flags: usize) {
131        self.entry = (ppn << 10) | flags;
132    }
133
134    pub fn flags(&self) -> usize {
135        self.entry & 0xFF
136    }
137}
138
139impl Default for PageTableEntry {
140    fn default() -> Self {
141        Self::new()
142    }
143}
144
145#[repr(align(4096))]
146pub struct PageTable {
147    pub entries: [PageTableEntry; PAGE_TABLE_ENTRIES],
148}
149
150impl PageTable {
151    pub const fn new() -> Self {
152        Self {
153            entries: [PageTableEntry::new(); PAGE_TABLE_ENTRIES],
154        }
155    }
156
157    pub fn get_entry(&self, index: usize) -> &PageTableEntry {
158        &self.entries[index]
159    }
160
161    pub fn get_entry_mut(&mut self, index: usize) -> &mut PageTableEntry {
162        &mut self.entries[index]
163    }
164
165    pub fn zero(&mut self) {
166        for entry in self.entries.iter_mut() {
167            *entry = PageTableEntry::new();
168        }
169    }
170}
171
172impl Default for PageTable {
173    fn default() -> Self {
174        Self::new()
175    }
176}
177
178pub struct SimpleFrameAllocator {
179    start_frame: usize,
180    next_frame: usize,
181    end_frame: usize,
182    free_stack: [usize; 256],
183    free_top: usize,
184}
185
186impl SimpleFrameAllocator {
187    pub fn new(kernel_end: usize, memory_end: usize) -> Self {
188        let next_frame = kernel_end.div_ceil(PAGE_SIZE);
189        let end_frame = memory_end / PAGE_SIZE;
190        Self {
191            start_frame: next_frame,
192            next_frame,
193            end_frame,
194            free_stack: [0; 256],
195            free_top: 0,
196        }
197    }
198
199    pub fn allocate(&mut self) -> Option<PhysFrame> {
200        if self.free_top > 0 {
201            self.free_top -= 1;
202            let frame_num = self.free_stack[self.free_top];
203            return Some(PhysFrame::containing_address(PhysAddr::new(
204                frame_num * PAGE_SIZE,
205            )));
206        }
207
208        if self.next_frame >= self.end_frame {
209            return None;
210        }
211
212        let frame = PhysFrame::containing_address(PhysAddr::new(self.next_frame * PAGE_SIZE));
213        self.next_frame += 1;
214        Some(frame)
215    }
216
217    pub fn deallocate(&mut self, frame: PhysFrame) {
218        let frame_num = frame.start_address().as_usize() / PAGE_SIZE;
219        if frame_num < self.start_frame
220            || frame_num >= self.next_frame
221            || self.free_top >= self.free_stack.len()
222            || self.free_stack[..self.free_top].contains(&frame_num)
223        {
224            return;
225        }
226        self.free_stack[self.free_top] = frame_num;
227        self.free_top += 1;
228    }
229
230    pub fn next_frame(&self) -> usize {
231        self.next_frame
232    }
233
234    pub fn set_next_frame(&mut self, value: usize) {
235        self.next_frame = value;
236    }
237}
238
239pub struct MemoryManager {
240    pub frame_allocator: SimpleFrameAllocator,
241}
242
243impl MemoryManager {
244    pub fn new(kernel_end: usize, memory_end: usize) -> Self {
245        Self {
246            frame_allocator: SimpleFrameAllocator::new(kernel_end, memory_end),
247        }
248    }
249}
250
251#[derive(Debug, Clone, Copy, PartialEq, Eq)]
252pub struct VmPermissions {
253    bits: u8,
254}
255
256impl VmPermissions {
257    pub const READ: Self = Self { bits: 1 << 0 };
258    pub const WRITE: Self = Self { bits: 1 << 1 };
259    pub const EXECUTE: Self = Self { bits: 1 << 2 };
260    pub const USER: Self = Self { bits: 1 << 3 };
261
262    pub const fn empty() -> Self {
263        Self { bits: 0 }
264    }
265
266    pub const fn from_bits(bits: u8) -> Self {
267        Self { bits }
268    }
269
270    pub const fn bits(self) -> u8 {
271        self.bits
272    }
273
274    pub const fn union(self, other: Self) -> Self {
275        Self {
276            bits: self.bits | other.bits,
277        }
278    }
279
280    pub const fn contains(self, other: Self) -> bool {
281        (self.bits & other.bits) == other.bits
282    }
283}
284
285#[derive(Debug, Clone, Copy, PartialEq, Eq)]
286pub enum VmRegionKind {
287    /// Anonymous mapping, suitable for mmap-like heap/scratch regions.
288    Anonymous,
289    /// File-backed mapping bookkeeping. Actual file page-in is a later kernel step.
290    FileBacked,
291    /// Grow-down user stack reservation with a bounded maximum size.
292    Stack { max_pages: usize },
293    /// Swap-backed anonymous reservation. Swap I/O is a later kernel step.
294    SwapBacked,
295}
296
297#[derive(Debug, Clone, Copy, PartialEq, Eq)]
298pub struct VmRegion {
299    start: VirtAddr,
300    end: VirtAddr,
301    kind: VmRegionKind,
302    permissions: VmPermissions,
303    committed_pages: u128,
304}
305
306impl VmRegion {
307    pub fn new(
308        start: VirtAddr,
309        len: usize,
310        kind: VmRegionKind,
311        permissions: VmPermissions,
312    ) -> Result<Self, &'static str> {
313        if len == 0 {
314            return Err("VM region length must be non-zero");
315        }
316        if start.page_offset() != 0 || !len.is_multiple_of(PAGE_SIZE) {
317            return Err("VM region must be page aligned");
318        }
319        let page_count = len / PAGE_SIZE;
320        if page_count == 0 || page_count > VM_REGION_MAX_PAGES {
321            return Err("VM region page count is outside teaching bounds");
322        }
323        if let VmRegionKind::Stack { max_pages } = kind {
324            if max_pages == 0 {
325                return Err("Stack maximum pages must be non-zero");
326            }
327            if max_pages > VM_REGION_MAX_PAGES {
328                return Err("Stack maximum pages exceed teaching bounds");
329            }
330            if page_count > max_pages {
331                return Err("Stack region length exceeds maximum pages");
332            }
333        }
334        let end = start
335            .as_usize()
336            .checked_add(len)
337            .ok_or("VM region end overflows")?;
338        Ok(Self {
339            start,
340            end: VirtAddr::new(end),
341            kind,
342            permissions,
343            committed_pages: 0,
344        })
345    }
346
347    pub const fn start(&self) -> VirtAddr {
348        self.start
349    }
350
351    pub const fn end(&self) -> VirtAddr {
352        self.end
353    }
354
355    pub const fn kind(&self) -> VmRegionKind {
356        self.kind
357    }
358
359    pub const fn permissions(&self) -> VmPermissions {
360        self.permissions
361    }
362
363    pub fn len(&self) -> usize {
364        self.end.as_usize() - self.start.as_usize()
365    }
366
367    pub fn is_empty(&self) -> bool {
368        self.len() == 0
369    }
370
371    pub fn page_count(&self) -> usize {
372        self.len() / PAGE_SIZE
373    }
374
375    pub fn contains(&self, addr: VirtAddr) -> bool {
376        addr.as_usize() >= self.start.as_usize() && addr.as_usize() < self.end.as_usize()
377    }
378
379    pub fn overlaps(&self, other: &VmRegion) -> bool {
380        self.start.as_usize() < other.end.as_usize() && other.start.as_usize() < self.end.as_usize()
381    }
382
383    pub fn committed_page_count(&self) -> usize {
384        self.committed_pages.count_ones() as usize
385    }
386
387    pub fn is_page_committed(&self, addr: VirtAddr) -> bool {
388        let Some(index) = self.page_index(addr) else {
389            return false;
390        };
391        (self.committed_pages & (1u128 << index)) != 0
392    }
393
394    pub fn commit_page(&mut self, addr: VirtAddr) -> Result<bool, &'static str> {
395        let index = self
396            .page_index(addr)
397            .ok_or("VM page is outside the region")?;
398        let mask = 1u128 << index;
399        let was_new = (self.committed_pages & mask) == 0;
400        self.committed_pages |= mask;
401        Ok(was_new)
402    }
403
404    pub fn grow_down_to(&mut self, fault_addr: VirtAddr) -> Result<bool, &'static str> {
405        let VmRegionKind::Stack { max_pages } = self.kind else {
406            return Err("Only stack VM regions can grow down");
407        };
408        let fault_page = align_down(fault_addr.as_usize());
409        if fault_page >= self.start.as_usize() {
410            return Ok(false);
411        }
412        let added_pages = (self.start.as_usize() - fault_page) / PAGE_SIZE;
413        let max_len = max_pages
414            .checked_mul(PAGE_SIZE)
415            .ok_or("Stack maximum size overflows")?;
416        let new_len = self.end.as_usize() - fault_page;
417        if new_len > max_len || new_len / PAGE_SIZE > VM_REGION_MAX_PAGES {
418            return Err("Stack growth exceeds reserved teaching bound");
419        }
420        self.start = VirtAddr::new(fault_page);
421        self.committed_pages <<= added_pages;
422        Ok(true)
423    }
424
425    fn page_index(&self, addr: VirtAddr) -> Option<usize> {
426        if !self.contains(addr) {
427            return None;
428        }
429        Some((align_down(addr.as_usize()) - self.start.as_usize()) / PAGE_SIZE)
430    }
431}
432
433#[derive(Debug, Clone)]
434pub struct VmRegionTable {
435    regions: [Option<VmRegion>; VM_REGION_MAX],
436}
437
438impl VmRegionTable {
439    pub const fn new() -> Self {
440        Self {
441            regions: [None; VM_REGION_MAX],
442        }
443    }
444
445    pub fn insert(&mut self, region: VmRegion) -> Result<(), &'static str> {
446        if self
447            .regions
448            .iter()
449            .flatten()
450            .any(|existing| existing.overlaps(&region))
451        {
452            return Err("VM region overlaps an existing mapping");
453        }
454        let slot = self
455            .regions
456            .iter_mut()
457            .find(|slot| slot.is_none())
458            .ok_or("VM region table is full")?;
459        *slot = Some(region);
460        Ok(())
461    }
462
463    pub fn find(&self, addr: VirtAddr) -> Option<&VmRegion> {
464        self.regions
465            .iter()
466            .flatten()
467            .find(|region| region.contains(addr))
468    }
469
470    pub fn find_mut(&mut self, addr: VirtAddr) -> Option<&mut VmRegion> {
471        self.regions
472            .iter_mut()
473            .flatten()
474            .find(|region| region.contains(addr))
475    }
476
477    pub fn remove(&mut self, start: VirtAddr) -> Option<VmRegion> {
478        let slot = self
479            .regions
480            .iter_mut()
481            .find(|region| region.map(|r| r.start == start).unwrap_or(false))?;
482        slot.take()
483    }
484
485    pub fn len(&self) -> usize {
486        self.regions.iter().flatten().count()
487    }
488
489    pub fn is_empty(&self) -> bool {
490        self.len() == 0
491    }
492
493    pub fn commit_page(&mut self, addr: VirtAddr) -> Result<bool, &'static str> {
494        let region = self.find_mut(addr).ok_or("No VM region contains page")?;
495        region.commit_page(addr)
496    }
497
498    pub fn grow_stack_for_fault(&mut self, fault_addr: VirtAddr) -> Result<bool, &'static str> {
499        let fault_page = align_down(fault_addr.as_usize());
500        let mut stack_slot = None;
501        let mut proposed_start = 0usize;
502        let mut proposed_end = 0usize;
503
504        for (index, region) in self.regions.iter().enumerate() {
505            let Some(region) = region else {
506                continue;
507            };
508            let VmRegionKind::Stack { max_pages } = region.kind else {
509                continue;
510            };
511            let lower_bound = region
512                .end
513                .as_usize()
514                .saturating_sub(max_pages.saturating_mul(PAGE_SIZE));
515            if fault_page >= lower_bound && fault_page < region.start.as_usize() {
516                stack_slot = Some(index);
517                proposed_start = fault_page;
518                proposed_end = region.end.as_usize();
519                break;
520            }
521        }
522
523        let stack_slot =
524            stack_slot.ok_or("Fault address is not inside a growable stack reservation")?;
525        if self.regions.iter().enumerate().any(|(index, region)| {
526            index != stack_slot
527                && region
528                    .as_ref()
529                    .is_some_and(|existing| ranges_overlap(proposed_start, proposed_end, existing))
530        }) {
531            return Err("Stack growth would overlap an existing VM region");
532        }
533
534        self.regions[stack_slot]
535            .as_mut()
536            .expect("stack slot was selected")
537            .grow_down_to(fault_addr)
538    }
539}
540
541impl Default for VmRegionTable {
542    fn default() -> Self {
543        Self::new()
544    }
545}
546
547const fn align_down(addr: usize) -> usize {
548    addr & !(PAGE_SIZE - 1)
549}
550
551fn ranges_overlap(start: usize, end: usize, region: &VmRegion) -> bool {
552    start < region.end.as_usize() && region.start.as_usize() < end
553}
554
555#[cfg(test)]
556mod tests {
557    use super::*;
558
559    #[test]
560    fn virt_addr_decomposes_sv39_parts() {
561        let addr = VirtAddr::new(0x1234_5678);
562        assert_eq!(addr.page_offset(), 0x678);
563        assert_eq!(addr.vpn0(), 0x145);
564    }
565
566    #[test]
567    fn frame_allocator_reuses_deallocated_frames() {
568        let mut allocator = SimpleFrameAllocator::new(0x8020_0000, 0x8030_0000);
569        let f1 = allocator.allocate().expect("first frame");
570        let f2 = allocator.allocate().expect("second frame");
571        assert_ne!(f1.start_address(), f2.start_address());
572
573        let f1_addr = f1.start_address();
574        allocator.deallocate(f1);
575        let f3 = allocator.allocate().expect("reused frame");
576        assert_eq!(f3.start_address(), f1_addr);
577    }
578
579    #[test]
580    fn page_table_entries_encode_ppn_and_flags() {
581        let mut entry = PageTableEntry::new();
582        let flags = PageTableFlags::Valid as usize | PageTableFlags::Read as usize;
583        entry.set(0x80200, flags);
584        assert!(entry.is_valid());
585        assert!(entry.is_leaf());
586        assert_eq!(entry.ppn(), 0x80200);
587        assert_eq!(entry.phys_addr(), PhysAddr::new(0x8020_0000));
588        assert_eq!(entry.flags(), flags);
589    }
590
591    #[test]
592    fn vm_region_table_tracks_lazy_mmap_commitment() {
593        let permissions = VmPermissions::READ
594            .union(VmPermissions::WRITE)
595            .union(VmPermissions::USER);
596        let region = VmRegion::new(
597            VirtAddr::new(0x4000_0000),
598            PAGE_SIZE * 4,
599            VmRegionKind::Anonymous,
600            permissions,
601        )
602        .expect("anonymous region");
603
604        let mut table = VmRegionTable::new();
605        table.insert(region).expect("insert region");
606        assert_eq!(table.len(), 1);
607        assert!(table.find(VirtAddr::new(0x4000_0123)).is_some());
608
609        assert!(table
610            .commit_page(VirtAddr::new(0x4000_0123))
611            .expect("first commit"));
612        assert!(!table
613            .commit_page(VirtAddr::new(0x4000_0abc))
614            .expect("duplicate commit"));
615        let region = table.find(VirtAddr::new(0x4000_0123)).unwrap();
616        assert_eq!(region.committed_page_count(), 1);
617        assert!(region.is_page_committed(VirtAddr::new(0x4000_0000)));
618        assert!(!region.is_page_committed(VirtAddr::new(0x4000_1000)));
619    }
620
621    #[test]
622    fn vm_region_table_rejects_overlap_and_bounds_stack_growth() {
623        let permissions = VmPermissions::READ
624            .union(VmPermissions::WRITE)
625            .union(VmPermissions::USER);
626        let mut table = VmRegionTable::new();
627        table
628            .insert(
629                VmRegion::new(
630                    VirtAddr::new(0x7000_0000),
631                    PAGE_SIZE,
632                    VmRegionKind::Stack { max_pages: 4 },
633                    permissions,
634                )
635                .unwrap(),
636            )
637            .unwrap();
638
639        let overlapping = VmRegion::new(
640            VirtAddr::new(0x7000_0000),
641            PAGE_SIZE,
642            VmRegionKind::Anonymous,
643            permissions,
644        )
645        .unwrap();
646        assert!(table.insert(overlapping).is_err());
647
648        assert!(table
649            .grow_stack_for_fault(VirtAddr::new(0x6fff_f000))
650            .expect("grow one page"));
651        assert!(table.find(VirtAddr::new(0x6fff_f000)).is_some());
652        assert!(table
653            .grow_stack_for_fault(VirtAddr::new(0x6fff_b000))
654            .is_err());
655    }
656
657    #[test]
658    fn vm_region_rejects_invalid_stack_bounds_at_construction() {
659        let permissions = VmPermissions::READ
660            .union(VmPermissions::WRITE)
661            .union(VmPermissions::USER);
662
663        assert_eq!(
664            VmRegion::new(
665                VirtAddr::new(0x7000_0000),
666                PAGE_SIZE,
667                VmRegionKind::Stack { max_pages: 0 },
668                permissions,
669            ),
670            Err("Stack maximum pages must be non-zero")
671        );
672        assert_eq!(
673            VmRegion::new(
674                VirtAddr::new(0x7000_0000),
675                PAGE_SIZE * 2,
676                VmRegionKind::Stack { max_pages: 1 },
677                permissions,
678            ),
679            Err("Stack region length exceeds maximum pages")
680        );
681        assert_eq!(
682            VmRegion::new(
683                VirtAddr::new(0x7000_0000),
684                PAGE_SIZE,
685                VmRegionKind::Stack {
686                    max_pages: VM_REGION_MAX_PAGES + 1,
687                },
688                permissions,
689            ),
690            Err("Stack maximum pages exceed teaching bounds")
691        );
692    }
693
694    #[test]
695    fn vm_region_stack_growth_preserves_overlap_invariant() {
696        let permissions = VmPermissions::READ
697            .union(VmPermissions::WRITE)
698            .union(VmPermissions::USER);
699        let mut table = VmRegionTable::new();
700        table
701            .insert(
702                VmRegion::new(
703                    VirtAddr::new(0x7fff_e000),
704                    PAGE_SIZE,
705                    VmRegionKind::Anonymous,
706                    permissions,
707                )
708                .unwrap(),
709            )
710            .unwrap();
711        table
712            .insert(
713                VmRegion::new(
714                    VirtAddr::new(0x7fff_f000),
715                    PAGE_SIZE,
716                    VmRegionKind::Stack { max_pages: 4 },
717                    permissions,
718                )
719                .unwrap(),
720            )
721            .unwrap();
722
723        assert_eq!(
724            table.grow_stack_for_fault(VirtAddr::new(0x7fff_e800)),
725            Err("Stack growth would overlap an existing VM region")
726        );
727        let stack = table.find(VirtAddr::new(0x7fff_f000)).unwrap();
728        assert_eq!(stack.start(), VirtAddr::new(0x7fff_f000));
729        assert!(matches!(stack.kind(), VmRegionKind::Stack { max_pages: 4 }));
730    }
731
732    #[test]
733    fn vm_region_stack_growth_keeps_committed_pages_at_original_addresses() {
734        let permissions = VmPermissions::READ
735            .union(VmPermissions::WRITE)
736            .union(VmPermissions::USER);
737        let mut table = VmRegionTable::new();
738        let stack_start = VirtAddr::new(0x7000_0000);
739        table
740            .insert(
741                VmRegion::new(
742                    stack_start,
743                    PAGE_SIZE,
744                    VmRegionKind::Stack { max_pages: 4 },
745                    permissions,
746                )
747                .unwrap(),
748            )
749            .unwrap();
750        assert!(table.commit_page(stack_start).unwrap());
751        assert!(table
752            .grow_stack_for_fault(VirtAddr::new(0x6fff_f800))
753            .unwrap());
754
755        let grown = table.find(stack_start).unwrap();
756        assert_eq!(grown.start(), VirtAddr::new(0x6fff_f000));
757        assert!(grown.is_page_committed(stack_start));
758        assert!(!grown.is_page_committed(VirtAddr::new(0x6fff_f000)));
759    }
760}
ahrightos_memory/lib.rs

ahrightos_memory/
lib.rs
