1use alloc::{collections::BTreeMap, vec::Vec};
2use core::fmt;
3
4use ax_memory_addr::{AddrRange, MemoryAddr};
5
6use crate::{MappingBackend, MappingError, MappingResult, MemoryArea};
7
8fn restore_overlapped_mappings<B: MappingBackend>(
14 old: &BTreeMap<B::Addr, MemoryArea<B>>,
15 range: AddrRange<B::Addr>,
16 context: &mut B::MutationContext,
17 page_table: &mut B::PageTable,
18) -> bool {
19 for area in old.values() {
20 if area.start() >= range.end {
21 break;
22 }
23 if area.end() <= range.start {
24 continue;
25 }
26 let start = area.start().max(range.start);
27 let end = area.end().min(range.end);
28 let Some(fragment) = AddrRange::try_new(start, end) else {
29 return false;
30 };
31 if !area.backend().map(
32 fragment.start,
33 fragment.size(),
34 area.flags(),
35 context,
36 page_table,
37 ) {
38 return false;
39 }
40 }
41 true
42}
43
44#[derive(Clone)]
46pub struct MemorySet<B: MappingBackend> {
47 areas: BTreeMap<B::Addr, MemoryArea<B>>,
48}
49
50impl<B: MappingBackend> MemorySet<B> {
51 pub const fn new() -> Self {
53 Self {
54 areas: BTreeMap::new(),
55 }
56 }
57
58 pub fn len(&self) -> usize {
60 self.areas.len()
61 }
62
63 pub fn is_empty(&self) -> bool {
65 self.areas.is_empty()
66 }
67
68 pub fn iter(&self) -> impl Iterator<Item = &MemoryArea<B>> {
70 self.areas.values()
71 }
72
73 pub fn restore_metadata_preimage(&mut self, preimage: Self) {
81 *self = preimage;
82 }
83
84 pub fn overlaps(&self, range: AddrRange<B::Addr>) -> bool {
86 if let Some((_, before)) = self.areas.range(..range.start).last()
87 && before.va_range().overlaps(range)
88 {
89 return true;
90 }
91 if let Some((_, after)) = self.areas.range(range.start..).next()
92 && after.va_range().overlaps(range)
93 {
94 return true;
95 }
96 false
97 }
98
99 pub fn find(&self, addr: B::Addr) -> Option<&MemoryArea<B>> {
101 let candidate = self.areas.range(..=addr).last().map(|(_, a)| a);
102 candidate.filter(|a| a.va_range().contains(addr))
103 }
104
105 pub fn find_free_area(
120 &self,
121 hint: B::Addr,
122 size: usize,
123 limit: AddrRange<B::Addr>,
124 align: usize,
125 ) -> Option<B::Addr> {
126 if align == 0 || !align.is_power_of_two() || size == 0 || limit.start >= limit.end {
131 return None;
132 }
133 if !size.is_multiple_of(align) {
134 return None;
136 }
137 let align_up = |address: B::Addr| {
139 address
140 .into()
141 .checked_add(align - 1)
142 .map(|value| B::Addr::from(value & !(align - 1)))
143 };
144 let mut last_end: <B as MappingBackend>::Addr = align_up(hint.max(limit.start))?;
145 if last_end < limit.start || last_end >= limit.end {
146 return None;
147 }
148 if let Some((_, area)) = self.areas.range(..last_end).last() {
149 last_end = align_up(last_end.max(area.end()))?;
150 if last_end >= limit.end {
151 return None;
152 }
153 }
154 for (&addr, area) in self.areas.range(last_end..) {
155 if addr >= limit.end {
156 break;
157 }
158 if last_end.checked_add(size).is_some_and(|end| end <= addr) {
159 if last_end
160 .checked_add(size)
161 .is_some_and(|end| end <= limit.end)
162 {
163 return Some(last_end);
164 }
165 return None;
166 }
167 last_end = align_up(area.end().max(limit.start))?;
168 if last_end >= limit.end {
169 return None;
170 }
171 }
172 if last_end
173 .checked_add(size)
174 .is_some_and(|end| end <= limit.end)
175 {
176 Some(last_end)
177 } else {
178 None
179 }
180 }
181
182 pub fn extend_area(
184 &mut self,
185 addr: B::Addr,
186 additional_size: usize,
187 context: &mut B::MutationContext,
188 page_table: &mut B::PageTable,
189 ) -> MappingResult {
190 if additional_size == 0 {
191 return Ok(());
192 }
193
194 let area_start = self
196 .areas
197 .range(..=addr)
198 .last()
199 .filter(|(_, a)| a.va_range().contains(addr))
200 .map(|(&start, _)| start)
201 .ok_or(MappingError::InvalidParam)?;
202
203 let area_end = self.areas[&area_start].end();
205 let new_end = area_end
206 .checked_add(additional_size)
207 .ok_or(MappingError::InvalidParam)?;
208 if let Some((_, next)) = self.areas.range(area_end..).next()
209 && new_end > next.start()
210 {
211 return Err(MappingError::AlreadyExists);
212 }
213
214 self.areas
215 .get_mut(&area_start)
216 .ok_or(MappingError::BadState)?
217 .grow_right(additional_size, context, page_table)?;
218 Ok(())
219 }
220
221 pub fn rollback_extend_area(
227 &mut self,
228 addr: B::Addr,
229 additional_size: usize,
230 context: &mut B::MutationContext,
231 page_table: &mut B::PageTable,
232 ) -> MappingResult {
233 if additional_size == 0 {
234 return Ok(());
235 }
236 let area_start = self
237 .areas
238 .range(..=addr)
239 .last()
240 .filter(|(_, area)| area.va_range().contains(addr))
241 .map(|(&start, _)| start)
242 .ok_or(MappingError::InvalidParam)?;
243 let area = self.areas.get(&area_start).ok_or(MappingError::BadState)?;
244 if additional_size >= area.size() {
245 return Err(MappingError::InvalidParam);
246 }
247 let suffix_start = area
248 .end()
249 .checked_sub(additional_size)
250 .ok_or(MappingError::InvalidParam)?;
251 let backend = area.backend().clone();
252 if !backend.validate_unmap(suffix_start, additional_size, page_table) {
253 return Err(MappingError::BadState);
254 }
255 if !backend.unmap(suffix_start, additional_size, context, page_table) {
256 return Err(MappingError::NeedsRepair);
257 }
258 let area = self
259 .areas
260 .get_mut(&area_start)
261 .ok_or(MappingError::BadState)?;
262 let old_size = area.size();
263 area.shrink_right_metadata(old_size - additional_size)
264 }
265
266 pub fn map(
275 &mut self,
276 area: MemoryArea<B>,
277 context: &mut B::MutationContext,
278 page_table: &mut B::PageTable,
279 unmap_overlap: bool,
280 ) -> MappingResult {
281 if area.va_range().is_empty() {
282 return Err(MappingError::InvalidParam);
283 }
284
285 let overlaps = self.overlaps(area.va_range());
286 let backup = overlaps.then(|| self.areas.clone());
287 if overlaps {
288 if unmap_overlap {
289 self.unmap(area.start(), area.size(), context, page_table)?;
290 } else {
291 return Err(MappingError::AlreadyExists);
292 }
293 } else {
294 area.validate_map(page_table)?;
299 }
300
301 let area_start = area.start();
302 let area_size = area.size();
303 let area_backend = area.backend().clone();
304 if let Err(error) = area.map_area(context, page_table) {
305 let reverted_new = area_backend.unmap(area_start, area_size, context, page_table);
310 let restored_old = backup.as_ref().is_none_or(|old| {
311 restore_overlapped_mappings(old, area.va_range(), context, page_table)
312 });
313 if !reverted_new || !restored_old {
314 if let Some(old) = backup {
315 self.areas = old;
316 }
317 return Err(MappingError::NeedsRepair);
318 }
319 if let Some(old) = backup {
320 self.areas = old;
321 }
322 return Err(error);
323 }
324
325 if self.areas.insert(area_start, area).is_some() {
326 let reverted_new = area_backend.unmap(area_start, area_size, context, page_table);
331 let restored_old = backup.as_ref().is_none_or(|old| {
332 restore_overlapped_mappings(
333 old,
334 AddrRange::from_start_size(area_start, area_size),
335 context,
336 page_table,
337 )
338 });
339 if let Some(old) = backup {
340 self.areas = old;
341 }
342 return Err(if reverted_new && restored_old {
343 MappingError::BadState
344 } else {
345 MappingError::NeedsRepair
346 });
347 }
348 Ok(())
349 }
350
351 pub fn insert_prepared_area(&mut self, area: MemoryArea<B>) -> MappingResult {
364 if area.va_range().is_empty() {
365 return Err(MappingError::InvalidParam);
366 }
367 if self.overlaps(area.va_range()) {
368 return Err(MappingError::AlreadyExists);
369 }
370 if self.areas.insert(area.start(), area).is_some() {
371 return Err(MappingError::BadState);
372 }
373 Ok(())
374 }
375
376 pub fn replace_exact_backend(
384 &mut self,
385 start: B::Addr,
386 size: usize,
387 backend: B,
388 ) -> MappingResult<B> {
389 let end = start.checked_add(size).ok_or(MappingError::InvalidParam)?;
390 let area = self
391 .areas
392 .get_mut(&start)
393 .filter(|area| area.end() == end)
394 .ok_or(MappingError::InvalidParam)?;
395 Ok(area.replace_backend(backend))
396 }
397
398 pub fn unmap_exact(
406 &mut self,
407 start: B::Addr,
408 size: usize,
409 context: &mut B::MutationContext,
410 page_table: &mut B::PageTable,
411 ) -> MappingResult<MemoryArea<B>> {
412 let end = start.checked_add(size).ok_or(MappingError::InvalidParam)?;
413 let area = self
414 .areas
415 .get(&start)
416 .filter(|area| area.end() == end)
417 .ok_or(MappingError::InvalidParam)?;
418 if area.validate_unmap_range(start, size, page_table).is_err() {
419 return Err(MappingError::BadState);
420 }
421 let backend = area.backend().clone();
422 if !backend.unmap(start, size, context, page_table) {
423 return Err(MappingError::BadState);
424 }
425 self.areas.remove(&start).ok_or(MappingError::BadState)
426 }
427
428 pub fn unmap(
435 &mut self,
436 start: B::Addr,
437 size: usize,
438 context: &mut B::MutationContext,
439 page_table: &mut B::PageTable,
440 ) -> MappingResult {
441 let range =
442 AddrRange::try_from_start_size(start, size).ok_or(MappingError::InvalidParam)?;
443 if range.is_empty() {
444 return Ok(());
445 }
446
447 self.validate_unmap(start, size, page_table)?;
448 let prepared = self.prepare_unmap_metadata(range)?;
449
450 self.for_each_intersecting_area(range, |area, unmap_start, unmap_size| {
456 area.unmap_range(unmap_start, unmap_size, context, page_table)
457 })?;
458
459 self.areas = prepared;
463 Ok(())
464 }
465
466 pub fn validate_unmap(
468 &self,
469 start: B::Addr,
470 size: usize,
471 page_table: &B::PageTable,
472 ) -> MappingResult {
473 let range =
474 AddrRange::try_from_start_size(start, size).ok_or(MappingError::InvalidParam)?;
475 if range.is_empty() {
476 return Ok(());
477 }
478
479 self.for_each_intersecting_area(range, |area, unmap_start, unmap_size| {
484 area.validate_unmap_range(unmap_start, unmap_size, page_table)
485 })
486 }
487
488 fn for_each_intersecting_area(
494 &self,
495 range: AddrRange<B::Addr>,
496 mut visit: impl FnMut(&MemoryArea<B>, B::Addr, usize) -> MappingResult,
497 ) -> MappingResult {
498 for area in self.areas.values() {
499 if area.start() >= range.end {
500 break;
501 }
502 if area.end() <= range.start {
503 continue;
504 }
505 let unmap_start = area.start().max(range.start);
506 let unmap_end = area.end().min(range.end);
507 visit(area, unmap_start, unmap_end.sub_addr(unmap_start))?;
508 }
509 Ok(())
510 }
511
512 pub fn unmap_metadata(&mut self, start: B::Addr, size: usize) -> MappingResult {
517 let range =
518 AddrRange::try_from_start_size(start, size).ok_or(MappingError::InvalidParam)?;
519 if range.is_empty() {
520 return Ok(());
521 }
522
523 self.areas = self.prepare_unmap_metadata(range)?;
524 Ok(())
525 }
526
527 fn prepare_unmap_metadata(
530 &self,
531 range: AddrRange<B::Addr>,
532 ) -> MappingResult<BTreeMap<B::Addr, MemoryArea<B>>> {
533 let mut areas = self.areas.clone();
534 let start = range.start;
535 let end = range.end;
536 areas.retain(|_, area| !area.va_range().contained_in(range));
537
538 if let Some((&before_start, before)) = areas.range_mut(..start).last() {
539 let before_end = before.end();
540 if before_end > start {
541 if before_end <= end {
542 before.shrink_right_metadata(start.sub_addr(before_start))?;
543 } else {
544 let right_part = before.split(end)?.ok_or(MappingError::BadState)?;
545 before.shrink_right_metadata(start.sub_addr(before_start))?;
546 if right_part.start() != end {
547 return Err(MappingError::BadState);
548 }
549 areas.insert(end, right_part);
550 }
551 }
552 }
553
554 if let Some((&after_start, _)) = areas.range(start..).next()
555 && after_start < end
556 {
557 let mut new_area = areas.remove(&after_start).ok_or(MappingError::BadState)?;
558 let after_end = new_area.end();
559 new_area.shrink_left_metadata(after_end.sub_addr(end))?;
560 if new_area.start() != end {
561 return Err(MappingError::BadState);
562 }
563 areas.insert(end, new_area);
564 }
565 Ok(areas)
566 }
567
568 fn containing_area_for_metadata_replacement(
570 &self,
571 area: &MemoryArea<B>,
572 ) -> MappingResult<B::Addr> {
573 if area.va_range().is_empty() {
574 return Err(MappingError::InvalidParam);
575 }
576
577 let start = area.start();
578 let end = area.end();
579
580 self.areas
581 .range(..=start)
582 .last()
583 .filter(|(_, old)| old.start() <= start && end <= old.end())
584 .map(|(&old_start, _)| old_start)
585 .ok_or(MappingError::InvalidParam)
586 }
587
588 pub fn validate_area_metadata_replacement(&self, area: &MemoryArea<B>) -> MappingResult {
590 self.containing_area_for_metadata_replacement(area)
591 .map(|_| ())
592 }
593
594 pub fn replace_area_metadata(&mut self, area: MemoryArea<B>) -> MappingResult {
596 let start = area.start();
597 let end = area.end();
598 let old_start = self.containing_area_for_metadata_replacement(&area)?;
599
600 let backup = self.areas.clone();
601 let result = (|| {
602 let Some(mut old_area) = self.areas.remove(&old_start) else {
603 return Err(MappingError::BadState);
604 };
605 if old_start < start {
606 let right_part = old_area.split(start)?.ok_or(MappingError::BadState)?;
607 self.areas.insert(old_start, old_area);
608 old_area = right_part;
609 }
610 if old_area.end() > end {
611 let right_part = old_area.split(end)?.ok_or(MappingError::BadState)?;
612 self.areas.insert(right_part.start(), right_part);
613 }
614 if self.areas.insert(start, area).is_some() {
615 return Err(MappingError::AlreadyExists);
616 }
617 Ok(())
618 })();
619 if result.is_err() {
620 self.areas = backup;
621 }
622 result
623 }
624
625 pub fn clear(
627 &mut self,
628 context: &mut B::MutationContext,
629 page_table: &mut B::PageTable,
630 ) -> MappingResult {
631 for area in self.areas.values() {
632 area.validate_unmap_range(area.start(), area.size(), page_table)?;
633 }
634 for (index, area) in self.areas.values().enumerate() {
635 if let Err(error) = area.unmap_area(context, page_table) {
636 return Err(if index == 0 {
637 error
638 } else {
639 MappingError::NeedsRepair
640 });
641 }
642 }
643 self.areas.clear();
644 Ok(())
645 }
646
647 pub fn protect(
657 &mut self,
658 start: B::Addr,
659 size: usize,
660 update_flags: impl Fn(B::Flags) -> Option<B::Flags>,
661 context: &mut B::MutationContext,
662 page_table: &mut B::PageTable,
663 ) -> MappingResult {
664 self.protect_with_reported_flags(
665 start,
666 size,
667 |flags, _reported_flags| update_flags(flags).map(|new_flags| (new_flags, new_flags)),
668 context,
669 page_table,
670 )
671 }
672
673 pub fn protect_with_reported_flags(
675 &mut self,
676 start: B::Addr,
677 size: usize,
678 update_flags: impl Fn(B::Flags, B::Flags) -> Option<(B::Flags, B::Flags)>,
679 context: &mut B::MutationContext,
680 page_table: &mut B::PageTable,
681 ) -> MappingResult {
682 let end = start.checked_add(size).ok_or(MappingError::InvalidParam)?;
683 if size == 0 {
684 return Ok(());
685 }
686 let mut operations = Vec::new();
687 for (&area_start, area) in &self.areas {
688 let area_end = area.end();
689 if area_start >= end {
690 break;
691 }
692 if area_end <= start {
693 continue;
694 }
695 if let Some((new_flags, new_reported_flags)) =
696 update_flags(area.flags(), area.reported_flags())
697 {
698 let protect_start = area_start.max(start);
699 let protect_end = area_end.min(end);
700 operations.push((
701 area_start,
702 protect_start,
703 protect_end,
704 area.flags(),
705 new_flags,
706 new_reported_flags,
707 ));
708 }
709 }
710
711 let mut prepared = self.areas.clone();
714 for &(area_start, protect_start, protect_end, _, new_flags, new_reported_flags) in
715 &operations
716 {
717 let original = &self.areas[&area_start];
718 if !original.backend().validate_protect(
719 protect_start,
720 protect_end.sub_addr(protect_start),
721 new_flags,
722 page_table,
723 ) {
724 return Err(MappingError::BadState);
725 }
726 let mut middle = prepared.remove(&area_start).ok_or(MappingError::BadState)?;
727 if area_start < protect_start {
728 let right = middle.split(protect_start)?.ok_or(MappingError::BadState)?;
729 prepared.insert(area_start, middle);
730 middle = right;
731 }
732 if protect_end < middle.end() {
733 let right = middle.split(protect_end)?.ok_or(MappingError::BadState)?;
734 prepared.insert(right.start(), right);
735 }
736 middle.set_flags_with_reported_flags(new_flags, new_reported_flags);
737 prepared.insert(middle.start(), middle);
738 }
739
740 for (index, &(area_start, protect_start, protect_end, _, new_flags, _)) in
741 operations.iter().enumerate()
742 {
743 let result = self.areas[&area_start].protect_range(
744 protect_start,
745 protect_end.sub_addr(protect_start),
746 new_flags,
747 context,
748 page_table,
749 );
750 if let Err(error) = result {
751 let mut restored = true;
752 for &(rollback_area_start, rollback_start, rollback_end, old_flags, ..) in
753 operations[..=index].iter().rev()
754 {
755 restored &= self.areas[&rollback_area_start]
756 .protect_range(
757 rollback_start,
758 rollback_end.sub_addr(rollback_start),
759 old_flags,
760 context,
761 page_table,
762 )
763 .is_ok();
764 }
765 return Err(if restored {
766 error
767 } else {
768 MappingError::NeedsRepair
769 });
770 }
771 }
772 self.areas = prepared;
773 Ok(())
774 }
775}
776
777impl<B: MappingBackend> Default for MemorySet<B> {
778 fn default() -> Self {
779 Self::new()
780 }
781}
782
783impl<B: MappingBackend> fmt::Debug for MemorySet<B>
784where
785 B::Addr: fmt::Debug,
786 B::Flags: fmt::Debug,
787{
788 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
789 f.debug_list().entries(self.areas.values()).finish()
790 }
791}