Skip to main content

vm_topology/
memory.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Tools to the compute guest memory layout.
5
6use memory_range::MemoryRange;
7use memory_range::subtract_ranges;
8use thiserror::Error;
9
10const PAGE_SIZE: u64 = 4096;
11const FOUR_GB: u64 = 0x1_0000_0000;
12
13/// Represents a page-aligned byte range of memory, with additional metadata.
14#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
15#[cfg_attr(feature = "mesh", derive(mesh_protobuf::Protobuf))]
16#[cfg_attr(feature = "inspect", derive(inspect::Inspect))]
17pub struct MemoryRangeWithNode {
18    /// The memory range.
19    pub range: MemoryRange,
20    /// The virtual NUMA node the range belongs to.
21    pub vnode: u32,
22}
23
24impl core::fmt::Display for MemoryRangeWithNode {
25    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
26        write!(f, "{}({})", self.range, self.vnode)
27    }
28}
29
30/// Describes the memory layout of a guest.
31#[derive(Debug, Clone)]
32#[cfg_attr(feature = "inspect", derive(inspect::Inspect))]
33pub struct MemoryLayout {
34    #[cfg_attr(feature = "inspect", inspect(with = "inspect_ranges_with_metadata"))]
35    ram: Vec<MemoryRangeWithNode>,
36    #[cfg_attr(feature = "inspect", inspect(with = "inspect_ranges"))]
37    mmio: Vec<MemoryRange>,
38    #[cfg_attr(feature = "inspect", inspect(with = "inspect_ranges"))]
39    pci_ecam: Vec<MemoryRange>,
40    #[cfg_attr(feature = "inspect", inspect(with = "inspect_ranges"))]
41    pci_mmio: Vec<MemoryRange>,
42    /// The RAM range used by VTL2. This is not present in any of the stats
43    /// above.
44    vtl2_range: Option<MemoryRange>,
45}
46
47#[cfg(feature = "inspect")]
48fn inspect_ranges(ranges: &[MemoryRange]) -> impl '_ + inspect::Inspect {
49    inspect::iter_by_key(ranges.iter().map(|range| {
50        (
51            range.to_string(),
52            inspect::adhoc(|i| {
53                i.respond().hex("length", range.len());
54            }),
55        )
56    }))
57}
58
59#[cfg(feature = "inspect")]
60fn inspect_ranges_with_metadata(ranges: &[MemoryRangeWithNode]) -> impl '_ + inspect::Inspect {
61    inspect::iter_by_key(ranges.iter().map(|range| {
62        (
63            range.range.to_string(),
64            inspect::adhoc(|i| {
65                i.respond()
66                    .hex("length", range.range.len())
67                    .hex("vnode", range.vnode);
68            }),
69        )
70    }))
71}
72
73/// Memory layout creation error.
74#[derive(Debug, Error)]
75pub enum Error {
76    /// Invalid memory size.
77    #[error("invalid memory size")]
78    BadSize,
79    /// Invalid per-NUMA-node memory size.
80    #[error("invalid NUMA node memory size")]
81    BadNumaSize,
82    /// Empty NUMA memory sizes.
83    #[error("empty NUMA memory sizes")]
84    EmptyNumaSizes,
85    /// Invalid MMIO gap configuration.
86    #[error("invalid MMIO gap configuration")]
87    BadMmioGaps,
88    /// Invalid memory ranges.
89    #[error("invalid memory or MMIO ranges")]
90    BadMemoryRanges,
91    /// VTL2 range is below the end of ram, and overlaps.
92    #[error("vtl2 range is below end of ram")]
93    Vtl2RangeBeforeEndOfRam,
94}
95
96fn validate_ranges(ranges: &[MemoryRange]) -> Result<(), Error> {
97    validate_ranges_core(ranges, |x| x)
98}
99
100fn validate_ranges_with_metadata(ranges: &[MemoryRangeWithNode]) -> Result<(), Error> {
101    validate_ranges_core(ranges, |x| &x.range)
102}
103
104/// Ensures everything in a list of ranges is non-empty, in order, and
105/// non-overlapping.
106fn validate_ranges_core<T>(ranges: &[T], getter: impl Fn(&T) -> &MemoryRange) -> Result<(), Error> {
107    if ranges.iter().any(|x| getter(x).is_empty())
108        || !ranges.array_windows().all(|[x, y]| {
109            let x = getter(x);
110            let y = getter(y);
111            x <= y && !x.overlaps(y)
112        })
113    {
114        return Err(Error::BadMemoryRanges);
115    }
116
117    Ok(())
118}
119
120/// The type backing an address.
121#[derive(Debug, Copy, Clone, PartialEq, Eq)]
122pub enum AddressType {
123    /// The address describes ram.
124    Ram,
125    /// The address describes mmio.
126    Mmio,
127    /// The address describes PCI ECAM.
128    PciEcam,
129    /// The address describes PCI MMIO.
130    PciMmio,
131}
132
133impl MemoryLayout {
134    /// Makes a new memory layout for a guest with `ram_size` bytes of memory
135    /// and MMIO gaps at the locations specified by `gaps`.
136    ///
137    /// `ram_size` must be a multiple of the page size. Each mmio and device
138    /// reserved gap must be non-empty, and the gaps must be in order and
139    /// non-overlapping.
140    ///
141    /// `vtl2_range` describes a range of memory reserved for VTL2.
142    /// It is not reported in ram.
143    ///
144    /// All RAM is assigned to NUMA node 0.
145    pub fn new(
146        ram_size: u64,
147        mmio_gaps: &[MemoryRange],
148        pci_ecam_gaps: &[MemoryRange],
149        pci_mmio_gaps: &[MemoryRange],
150        vtl2_range: Option<MemoryRange>,
151    ) -> Result<Self, Error> {
152        if ram_size == 0 || ram_size & (PAGE_SIZE - 1) != 0 {
153            return Err(Error::BadSize);
154        }
155        Self::new_with_numa(
156            &[ram_size],
157            mmio_gaps,
158            pci_ecam_gaps,
159            pci_mmio_gaps,
160            vtl2_range,
161        )
162    }
163
164    /// Like [`Self::new()`], but distributes RAM across NUMA nodes according
165    /// to the per-node sizes in `numa_mem_sizes`.
166    ///
167    /// `numa_mem_sizes[i]` is the number of RAM bytes for vnode `i`.
168    /// Each size must be page-aligned and non-zero. The sum of all sizes
169    /// is the total guest RAM.
170    ///
171    /// RAM is placed sequentially around MMIO gaps, filling each node's
172    /// budget in order. When a node's budget is exhausted mid-chunk,
173    /// the chunk is split and the next node continues from that address.
174    pub fn new_with_numa(
175        numa_mem_sizes: &[u64],
176        mmio_gaps: &[MemoryRange],
177        pci_ecam_gaps: &[MemoryRange],
178        pci_mmio_gaps: &[MemoryRange],
179        vtl2_range: Option<MemoryRange>,
180    ) -> Result<Self, Error> {
181        if numa_mem_sizes.is_empty() {
182            return Err(Error::EmptyNumaSizes);
183        }
184
185        for &size in numa_mem_sizes {
186            if size == 0 || size & (PAGE_SIZE - 1) != 0 {
187                return Err(Error::BadNumaSize);
188            }
189        }
190
191        let ram_size: u64 = numa_mem_sizes
192            .iter()
193            .try_fold(0u64, |acc, &s| acc.checked_add(s))
194            .ok_or(Error::BadSize)?;
195
196        validate_ranges(mmio_gaps)?;
197        validate_ranges(pci_ecam_gaps)?;
198        validate_ranges(pci_mmio_gaps)?;
199
200        let mut combined_gaps = mmio_gaps
201            .iter()
202            .chain(pci_ecam_gaps)
203            .chain(pci_mmio_gaps)
204            .copied()
205            .collect::<Vec<_>>();
206        combined_gaps.sort();
207        validate_ranges(&combined_gaps)?;
208
209        let available = subtract_ranges(
210            [MemoryRange::new(0..MemoryRange::MAX_ADDRESS)],
211            combined_gaps,
212        );
213
214        // Distribute RAM across NUMA nodes, filling available ranges in order.
215        let mut ram = Vec::new();
216        let mut remaining = ram_size;
217        let mut node_idx = 0;
218        let mut node_remaining = numa_mem_sizes[0];
219
220        for range in available {
221            let range_size = remaining.min(range.len());
222            let mut offset = 0;
223
224            while offset < range_size {
225                if node_remaining == 0 {
226                    node_idx += 1;
227                    node_remaining = *numa_mem_sizes
228                        .get(node_idx)
229                        .expect("node budget exhausted before all RAM placed");
230                }
231
232                let piece = (range_size - offset).min(node_remaining);
233                let start = range.start() + offset;
234                ram.push(MemoryRangeWithNode {
235                    range: MemoryRange::new(start..start + piece),
236                    vnode: node_idx as u32,
237                });
238                offset += piece;
239                node_remaining -= piece;
240            }
241
242            remaining -= range_size;
243
244            if remaining == 0 {
245                break;
246            }
247        }
248
249        Self::build(
250            ram,
251            mmio_gaps.to_vec(),
252            pci_ecam_gaps.to_vec(),
253            pci_mmio_gaps.to_vec(),
254            vtl2_range,
255        )
256    }
257
258    /// Makes a new memory layout for a guest with the given mmio gaps and
259    /// memory ranges.
260    ///
261    /// `memory` and `gaps` ranges must be in sorted order and non-overlapping,
262    /// and describe page aligned ranges.
263    pub fn new_from_ranges(
264        memory: &[MemoryRangeWithNode],
265        gaps: &[MemoryRange],
266    ) -> Result<Self, Error> {
267        validate_ranges_with_metadata(memory)?;
268        validate_ranges(gaps)?;
269        Self::build(memory.to_vec(), gaps.to_vec(), vec![], vec![], None)
270    }
271
272    /// Makes a new memory layout from already-resolved RAM and fixed ranges.
273    ///
274    /// Each individual range must be non-empty, but the lists themselves may
275    /// be empty (e.g. no PCIe root complexes means empty PCI ECAM/MMIO
276    /// vectors). Ranges within each list must be sorted and non-overlapping.
277    /// MMIO gaps may contain empty placeholder ranges to preserve positional
278    /// indexing (e.g. `mmio()[0]` = low, `mmio()[1]` = high); empty entries
279    /// are ignored during validation. The combined layout is also validated
280    /// for overlaps, including the optional VTL2 range.
281    pub fn new_from_resolved_ranges(
282        ram: Vec<MemoryRangeWithNode>,
283        mmio_gaps: Vec<MemoryRange>,
284        pci_ecam_gaps: Vec<MemoryRange>,
285        pci_mmio_gaps: Vec<MemoryRange>,
286        vtl2_range: Option<MemoryRange>,
287    ) -> Result<Self, Error> {
288        validate_ranges_with_metadata(&ram)?;
289        // MMIO gaps may include empty placeholders for positional indexing;
290        // validate only the non-empty entries.
291        let non_empty_mmio: Vec<_> = mmio_gaps
292            .iter()
293            .copied()
294            .filter(|r| !r.is_empty())
295            .collect();
296        validate_ranges(&non_empty_mmio)?;
297        validate_ranges(&pci_ecam_gaps)?;
298        validate_ranges(&pci_mmio_gaps)?;
299
300        Self::build(ram, mmio_gaps, pci_ecam_gaps, pci_mmio_gaps, vtl2_range)
301    }
302
303    /// Builds the memory layout.
304    ///
305    /// `ram` must already be known to be sorted.
306    fn build(
307        ram: Vec<MemoryRangeWithNode>,
308        mmio: Vec<MemoryRange>,
309        pci_ecam: Vec<MemoryRange>,
310        pci_mmio: Vec<MemoryRange>,
311        vtl2_range: Option<MemoryRange>,
312    ) -> Result<Self, Error> {
313        // Filter out empty placeholder ranges before validation and overlap
314        // checks — they carry no physical meaning and exist only for
315        // positional indexing in the stored mmio vector.
316        let mut all_ranges = ram
317            .iter()
318            .map(|x| &x.range)
319            .chain(&mmio)
320            .chain(&vtl2_range)
321            .chain(&pci_ecam)
322            .chain(&pci_mmio)
323            .copied()
324            .filter(|r| !r.is_empty())
325            .collect::<Vec<_>>();
326
327        all_ranges.sort();
328        validate_ranges(&all_ranges)?;
329
330        if all_ranges.array_windows().any(|[x, y]| x.overlaps(y)) {
331            return Err(Error::BadMemoryRanges);
332        }
333
334        let last_ram_entry = ram.last().ok_or(Error::BadMemoryRanges)?;
335        let end_of_ram = last_ram_entry.range.end();
336
337        if let Some(range) = vtl2_range {
338            if range.start() < end_of_ram {
339                return Err(Error::Vtl2RangeBeforeEndOfRam);
340            }
341        }
342
343        Ok(Self {
344            ram,
345            mmio,
346            pci_ecam,
347            pci_mmio,
348            vtl2_range,
349        })
350    }
351
352    /// The MMIO gap ranges.
353    pub fn mmio(&self) -> &[MemoryRange] {
354        &self.mmio
355    }
356
357    /// The populated RAM ranges. This does not include the vtl2_range.
358    pub fn ram(&self) -> &[MemoryRangeWithNode] {
359        &self.ram
360    }
361
362    /// A special memory range for VTL2, if any. This memory range is treated
363    /// like RAM, but is only used to hold VTL2 and is located above ram and
364    /// mmio.
365    pub fn vtl2_range(&self) -> Option<MemoryRange> {
366        self.vtl2_range
367    }
368
369    /// The total RAM size in bytes. This is not contiguous.
370    pub fn ram_size(&self) -> u64 {
371        self.ram.iter().map(|r| r.range.len()).sum()
372    }
373
374    /// One past the last byte of RAM.
375    pub fn end_of_ram(&self) -> u64 {
376        // always at least one RAM range
377        self.ram.last().expect("mmio set").range.end()
378    }
379
380    /// The ending RAM address below 4GB.
381    ///
382    /// Returns None if there is no RAM mapped below 4GB.
383    pub fn max_ram_below_4gb(&self) -> Option<u64> {
384        Some(
385            self.ram
386                .iter()
387                .rev()
388                .find(|r| r.range.end() < FOUR_GB)?
389                .range
390                .end(),
391        )
392    }
393
394    /// One past the last byte of RAM, MMIO, PCI ECAM, or PCI MMIO.
395    pub fn end_of_layout(&self) -> u64 {
396        [
397            self.mmio
398                .iter()
399                .filter(|r| !r.is_empty())
400                .map(|r| r.end())
401                .max()
402                .unwrap_or(0),
403            self.end_of_ram(),
404            self.pci_ecam.last().map(|r| r.end()).unwrap_or(0),
405            self.pci_mmio.last().map(|r| r.end()).unwrap_or(0),
406        ]
407        .into_iter()
408        .max()
409        .unwrap()
410    }
411
412    /// Probe a given address to see if it is in the memory layout described by
413    /// `self`. Returns the [`AddressType`] of the address if it is in the
414    /// layout.
415    ///
416    /// This does not check the vtl2_range.
417    pub fn probe_address(&self, address: u64) -> Option<AddressType> {
418        let ranges = self
419            .ram
420            .iter()
421            .map(|r| (&r.range, AddressType::Ram))
422            .chain(self.mmio.iter().map(|r| (r, AddressType::Mmio)))
423            .chain(self.pci_ecam.iter().map(|r| (r, AddressType::PciEcam)))
424            .chain(self.pci_mmio.iter().map(|r| (r, AddressType::PciMmio)));
425
426        for (range, address_type) in ranges {
427            if range.contains_addr(address) {
428                return Some(address_type);
429            }
430        }
431
432        None
433    }
434}
435
436#[cfg(test)]
437mod tests {
438    use super::*;
439
440    const KB: u64 = 1024;
441    const MB: u64 = 1024 * KB;
442    const GB: u64 = 1024 * MB;
443    const TB: u64 = 1024 * GB;
444
445    #[test]
446    fn layout() {
447        let mmio = &[
448            MemoryRange::new(GB..2 * GB),
449            MemoryRange::new(3 * GB..4 * GB),
450        ];
451        let ram = &[
452            MemoryRangeWithNode {
453                range: MemoryRange::new(0..GB),
454                vnode: 0,
455            },
456            MemoryRangeWithNode {
457                range: MemoryRange::new(2 * GB..3 * GB),
458                vnode: 0,
459            },
460            MemoryRangeWithNode {
461                range: MemoryRange::new(4 * GB..TB + 2 * GB),
462                vnode: 0,
463            },
464        ];
465
466        let layout = MemoryLayout::new(TB, mmio, &[], &[], None).unwrap();
467        assert_eq!(
468            layout.ram(),
469            &[
470                MemoryRangeWithNode {
471                    range: MemoryRange::new(0..GB),
472                    vnode: 0
473                },
474                MemoryRangeWithNode {
475                    range: MemoryRange::new(2 * GB..3 * GB),
476                    vnode: 0
477                },
478                MemoryRangeWithNode {
479                    range: MemoryRange::new(4 * GB..TB + 2 * GB),
480                    vnode: 0
481                },
482            ]
483        );
484        assert_eq!(layout.mmio(), mmio);
485        assert_eq!(layout.ram_size(), TB);
486        assert_eq!(layout.end_of_ram(), TB + 2 * GB);
487        assert_eq!(layout.end_of_layout(), TB + 2 * GB);
488
489        let layout = MemoryLayout::new_from_ranges(ram, mmio).unwrap();
490        assert_eq!(
491            layout.ram(),
492            &[
493                MemoryRangeWithNode {
494                    range: MemoryRange::new(0..GB),
495                    vnode: 0
496                },
497                MemoryRangeWithNode {
498                    range: MemoryRange::new(2 * GB..3 * GB),
499                    vnode: 0
500                },
501                MemoryRangeWithNode {
502                    range: MemoryRange::new(4 * GB..TB + 2 * GB),
503                    vnode: 0
504                },
505            ]
506        );
507        assert_eq!(layout.mmio(), mmio);
508        assert_eq!(layout.ram_size(), TB);
509        assert_eq!(layout.end_of_ram(), TB + 2 * GB);
510        assert_eq!(layout.end_of_layout(), TB + 2 * GB);
511    }
512
513    #[test]
514    fn bad_layout() {
515        MemoryLayout::new(TB + 1, &[], &[], &[], None).unwrap_err();
516        let mmio = &[
517            MemoryRange::new(3 * GB..4 * GB),
518            MemoryRange::new(GB..2 * GB),
519        ];
520        MemoryLayout::new(TB, mmio, &[], &[], None).unwrap_err();
521
522        MemoryLayout::new_from_ranges(&[], mmio).unwrap_err();
523
524        let ram = &[MemoryRangeWithNode {
525            range: MemoryRange::new(0..GB),
526            vnode: 0,
527        }];
528        MemoryLayout::new_from_ranges(ram, mmio).unwrap_err();
529
530        let ram = &[MemoryRangeWithNode {
531            range: MemoryRange::new(0..GB + MB),
532            vnode: 0,
533        }];
534        let mmio = &[
535            MemoryRange::new(GB..2 * GB),
536            MemoryRange::new(3 * GB..4 * GB),
537        ];
538        MemoryLayout::new_from_ranges(ram, mmio).unwrap_err();
539
540        let mmio = &[
541            MemoryRange::new(GB..2 * GB),
542            MemoryRange::new(3 * GB..4 * GB),
543        ];
544        let pci_ecam = &[MemoryRange::new(GB..GB + MB)];
545        MemoryLayout::new(TB, mmio, pci_ecam, &[], None).unwrap_err();
546
547        let mmio = &[
548            MemoryRange::new(GB..2 * GB),
549            MemoryRange::new(3 * GB..4 * GB),
550        ];
551        let pci_mmio = &[MemoryRange::new(GB..GB + MB)];
552        MemoryLayout::new(TB, mmio, &[], pci_mmio, None).unwrap_err();
553
554        let pci_ecam = &[MemoryRange::new(GB..GB + MB)];
555        let pci_mmio = &[MemoryRange::new(GB..GB + MB)];
556        MemoryLayout::new(TB, &[], pci_ecam, pci_mmio, None).unwrap_err();
557    }
558
559    #[test]
560    fn resolved_ranges_constructor() {
561        let ram = vec![
562            MemoryRangeWithNode {
563                range: MemoryRange::new(0..GB),
564                vnode: 0,
565            },
566            MemoryRangeWithNode {
567                range: MemoryRange::new(2 * GB..3 * GB),
568                vnode: 1,
569            },
570        ];
571        let mmio = vec![MemoryRange::new(GB..2 * GB)];
572        let pci_ecam = vec![MemoryRange::new(4 * GB..4 * GB + MB)];
573        let pci_mmio = vec![MemoryRange::new(5 * GB..6 * GB)];
574
575        let layout = MemoryLayout::new_from_resolved_ranges(
576            ram.clone(),
577            mmio.clone(),
578            pci_ecam.clone(),
579            pci_mmio.clone(),
580            None,
581        )
582        .unwrap();
583
584        assert_eq!(layout.ram(), ram);
585        assert_eq!(layout.mmio(), mmio);
586        assert_eq!(layout.probe_address(4 * GB), Some(AddressType::PciEcam));
587        assert_eq!(layout.probe_address(5 * GB), Some(AddressType::PciMmio));
588    }
589
590    #[test]
591    fn resolved_ranges_reject_overlap_with_fixed_ranges() {
592        let ram = vec![MemoryRangeWithNode {
593            range: MemoryRange::new(0..2 * GB),
594            vnode: 0,
595        }];
596        let mmio = vec![MemoryRange::new(GB..2 * GB)];
597
598        assert!(MemoryLayout::new_from_resolved_ranges(ram, mmio, vec![], vec![], None).is_err());
599    }
600
601    #[test]
602    fn resolved_ranges_validate_vtl2_against_ram_end() {
603        let ram = vec![
604            MemoryRangeWithNode {
605                range: MemoryRange::new(0..GB),
606                vnode: 0,
607            },
608            MemoryRangeWithNode {
609                range: MemoryRange::new(3 * GB..4 * GB),
610                vnode: 0,
611            },
612        ];
613        let mmio = vec![MemoryRange::new(GB..2 * GB)];
614        let vtl2_range = MemoryRange::new(2 * GB..2 * GB + MB);
615
616        assert!(matches!(
617            MemoryLayout::new_from_resolved_ranges(ram, mmio, vec![], vec![], Some(vtl2_range)),
618            Err(Error::Vtl2RangeBeforeEndOfRam)
619        ));
620    }
621
622    #[test]
623    fn pci_ranges() {
624        let mmio = &[MemoryRange::new(3 * GB..4 * GB)];
625        let pci_ecam = &[MemoryRange::new(2 * TB - GB..2 * TB)];
626        let pci_mmio = &[
627            MemoryRange::new(2 * GB..3 * GB),
628            MemoryRange::new(5 * GB..6 * GB),
629        ];
630
631        let layout = MemoryLayout::new(TB, mmio, pci_ecam, pci_mmio, None).unwrap();
632        assert_eq!(
633            layout.ram(),
634            &[
635                MemoryRangeWithNode {
636                    range: MemoryRange::new(0..2 * GB),
637                    vnode: 0,
638                },
639                MemoryRangeWithNode {
640                    range: MemoryRange::new(4 * GB..5 * GB),
641                    vnode: 0,
642                },
643                MemoryRangeWithNode {
644                    range: MemoryRange::new(6 * GB..TB + 3 * GB),
645                    vnode: 0,
646                },
647            ]
648        );
649        assert_eq!(layout.end_of_layout(), 2 * TB);
650
651        assert_eq!(layout.probe_address(2 * GB), Some(AddressType::PciMmio));
652        assert_eq!(
653            layout.probe_address(2 * GB + MB),
654            Some(AddressType::PciMmio)
655        );
656        assert_eq!(layout.probe_address(5 * GB), Some(AddressType::PciMmio));
657        assert_eq!(
658            layout.probe_address(5 * GB + MB),
659            Some(AddressType::PciMmio)
660        );
661        assert_eq!(
662            layout.probe_address(2 * TB - GB),
663            Some(AddressType::PciEcam)
664        );
665    }
666
667    #[test]
668    fn probe_address() {
669        let mmio = &[
670            MemoryRange::new(GB..2 * GB),
671            MemoryRange::new(3 * GB..4 * GB),
672        ];
673        let ram = &[
674            MemoryRangeWithNode {
675                range: MemoryRange::new(0..GB),
676                vnode: 0,
677            },
678            MemoryRangeWithNode {
679                range: MemoryRange::new(2 * GB..3 * GB),
680                vnode: 0,
681            },
682            MemoryRangeWithNode {
683                range: MemoryRange::new(4 * GB..TB + 2 * GB),
684                vnode: 0,
685            },
686        ];
687
688        let layout = MemoryLayout::new_from_ranges(ram, mmio).unwrap();
689
690        assert_eq!(layout.probe_address(0), Some(AddressType::Ram));
691        assert_eq!(layout.probe_address(256), Some(AddressType::Ram));
692        assert_eq!(layout.probe_address(2 * GB), Some(AddressType::Ram));
693        assert_eq!(layout.probe_address(4 * GB), Some(AddressType::Ram));
694        assert_eq!(layout.probe_address(TB), Some(AddressType::Ram));
695        assert_eq!(layout.probe_address(TB + 1), Some(AddressType::Ram));
696
697        assert_eq!(layout.probe_address(GB), Some(AddressType::Mmio));
698        assert_eq!(layout.probe_address(GB + 123), Some(AddressType::Mmio));
699        assert_eq!(layout.probe_address(3 * GB), Some(AddressType::Mmio));
700
701        assert_eq!(layout.probe_address(TB + 2 * GB), None);
702        assert_eq!(layout.probe_address(TB + 3 * GB), None);
703        assert_eq!(layout.probe_address(4 * TB), None);
704    }
705
706    #[test]
707    fn numa_two_nodes_even_split() {
708        // 4 GB total, 2 nodes of 2 GB each, MMIO gap at 2-3 GB.
709        let mmio = &[MemoryRange::new(2 * GB..3 * GB)];
710        let layout = MemoryLayout::new_with_numa(&[2 * GB, 2 * GB], mmio, &[], &[], None).unwrap();
711        assert_eq!(
712            layout.ram(),
713            &[
714                MemoryRangeWithNode {
715                    range: MemoryRange::new(0..2 * GB),
716                    vnode: 0,
717                },
718                MemoryRangeWithNode {
719                    range: MemoryRange::new(3 * GB..5 * GB),
720                    vnode: 1,
721                },
722            ]
723        );
724        assert_eq!(layout.ram_size(), 4 * GB);
725    }
726
727    #[test]
728    fn numa_two_nodes_mid_chunk_split() {
729        // 4 GB total, 2 nodes of 2 GB each, MMIO gap at 3-4 GB.
730        // Node 0's 2 GB fits entirely below the gap; node 1 continues above.
731        // But the first chunk is 3 GB, so node 0 takes 2 GB and node 1
732        // takes the remaining 1 GB of that chunk, plus 1 GB above the gap.
733        let mmio = &[MemoryRange::new(3 * GB..4 * GB)];
734        let layout = MemoryLayout::new_with_numa(&[2 * GB, 2 * GB], mmio, &[], &[], None).unwrap();
735        assert_eq!(
736            layout.ram(),
737            &[
738                MemoryRangeWithNode {
739                    range: MemoryRange::new(0..2 * GB),
740                    vnode: 0,
741                },
742                MemoryRangeWithNode {
743                    range: MemoryRange::new(2 * GB..3 * GB),
744                    vnode: 1,
745                },
746                MemoryRangeWithNode {
747                    range: MemoryRange::new(4 * GB..5 * GB),
748                    vnode: 1,
749                },
750            ]
751        );
752        assert_eq!(layout.ram_size(), 4 * GB);
753    }
754
755    #[test]
756    fn numa_three_nodes() {
757        // 3 GB total, 3 nodes of 1 GB each, no gaps.
758        let layout = MemoryLayout::new_with_numa(&[GB, GB, GB], &[], &[], &[], None).unwrap();
759        assert_eq!(
760            layout.ram(),
761            &[
762                MemoryRangeWithNode {
763                    range: MemoryRange::new(0..GB),
764                    vnode: 0,
765                },
766                MemoryRangeWithNode {
767                    range: MemoryRange::new(GB..2 * GB),
768                    vnode: 1,
769                },
770                MemoryRangeWithNode {
771                    range: MemoryRange::new(2 * GB..3 * GB),
772                    vnode: 2,
773                },
774            ]
775        );
776    }
777
778    #[test]
779    fn numa_single_node_matches_new() {
780        // Single node should produce the same layout as new().
781        let mmio = &[
782            MemoryRange::new(GB..2 * GB),
783            MemoryRange::new(3 * GB..4 * GB),
784        ];
785        let layout_new = MemoryLayout::new(TB, mmio, &[], &[], None).unwrap();
786        let layout_numa = MemoryLayout::new_with_numa(&[TB], mmio, &[], &[], None).unwrap();
787        assert_eq!(layout_new.ram(), layout_numa.ram());
788    }
789
790    #[test]
791    fn numa_bad_inputs() {
792        // Empty sizes.
793        MemoryLayout::new_with_numa(&[], &[], &[], &[], None).unwrap_err();
794        // Non-page-aligned size.
795        MemoryLayout::new_with_numa(&[GB + 1], &[], &[], &[], None).unwrap_err();
796        // Zero size.
797        MemoryLayout::new_with_numa(&[0], &[], &[], &[], None).unwrap_err();
798        // Mixed: one valid, one zero.
799        MemoryLayout::new_with_numa(&[GB, 0], &[], &[], &[], None).unwrap_err();
800    }
801}