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.iter().zip(ranges.iter().skip(1)).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
331            .iter()
332            .zip(all_ranges.iter().skip(1))
333            .any(|(x, y)| x.overlaps(y))
334        {
335            return Err(Error::BadMemoryRanges);
336        }
337
338        let last_ram_entry = ram.last().ok_or(Error::BadMemoryRanges)?;
339        let end_of_ram = last_ram_entry.range.end();
340
341        if let Some(range) = vtl2_range {
342            if range.start() < end_of_ram {
343                return Err(Error::Vtl2RangeBeforeEndOfRam);
344            }
345        }
346
347        Ok(Self {
348            ram,
349            mmio,
350            pci_ecam,
351            pci_mmio,
352            vtl2_range,
353        })
354    }
355
356    /// The MMIO gap ranges.
357    pub fn mmio(&self) -> &[MemoryRange] {
358        &self.mmio
359    }
360
361    /// The populated RAM ranges. This does not include the vtl2_range.
362    pub fn ram(&self) -> &[MemoryRangeWithNode] {
363        &self.ram
364    }
365
366    /// A special memory range for VTL2, if any. This memory range is treated
367    /// like RAM, but is only used to hold VTL2 and is located above ram and
368    /// mmio.
369    pub fn vtl2_range(&self) -> Option<MemoryRange> {
370        self.vtl2_range
371    }
372
373    /// The total RAM size in bytes. This is not contiguous.
374    pub fn ram_size(&self) -> u64 {
375        self.ram.iter().map(|r| r.range.len()).sum()
376    }
377
378    /// One past the last byte of RAM.
379    pub fn end_of_ram(&self) -> u64 {
380        // always at least one RAM range
381        self.ram.last().expect("mmio set").range.end()
382    }
383
384    /// The ending RAM address below 4GB.
385    ///
386    /// Returns None if there is no RAM mapped below 4GB.
387    pub fn max_ram_below_4gb(&self) -> Option<u64> {
388        Some(
389            self.ram
390                .iter()
391                .rev()
392                .find(|r| r.range.end() < FOUR_GB)?
393                .range
394                .end(),
395        )
396    }
397
398    /// One past the last byte of RAM, MMIO, PCI ECAM, or PCI MMIO.
399    pub fn end_of_layout(&self) -> u64 {
400        [
401            self.mmio
402                .iter()
403                .filter(|r| !r.is_empty())
404                .map(|r| r.end())
405                .max()
406                .unwrap_or(0),
407            self.end_of_ram(),
408            self.pci_ecam.last().map(|r| r.end()).unwrap_or(0),
409            self.pci_mmio.last().map(|r| r.end()).unwrap_or(0),
410        ]
411        .into_iter()
412        .max()
413        .unwrap()
414    }
415
416    /// Probe a given address to see if it is in the memory layout described by
417    /// `self`. Returns the [`AddressType`] of the address if it is in the
418    /// layout.
419    ///
420    /// This does not check the vtl2_range.
421    pub fn probe_address(&self, address: u64) -> Option<AddressType> {
422        let ranges = self
423            .ram
424            .iter()
425            .map(|r| (&r.range, AddressType::Ram))
426            .chain(self.mmio.iter().map(|r| (r, AddressType::Mmio)))
427            .chain(self.pci_ecam.iter().map(|r| (r, AddressType::PciEcam)))
428            .chain(self.pci_mmio.iter().map(|r| (r, AddressType::PciMmio)));
429
430        for (range, address_type) in ranges {
431            if range.contains_addr(address) {
432                return Some(address_type);
433            }
434        }
435
436        None
437    }
438}
439
440#[cfg(test)]
441mod tests {
442    use super::*;
443
444    const KB: u64 = 1024;
445    const MB: u64 = 1024 * KB;
446    const GB: u64 = 1024 * MB;
447    const TB: u64 = 1024 * GB;
448
449    #[test]
450    fn layout() {
451        let mmio = &[
452            MemoryRange::new(GB..2 * GB),
453            MemoryRange::new(3 * GB..4 * GB),
454        ];
455        let ram = &[
456            MemoryRangeWithNode {
457                range: MemoryRange::new(0..GB),
458                vnode: 0,
459            },
460            MemoryRangeWithNode {
461                range: MemoryRange::new(2 * GB..3 * GB),
462                vnode: 0,
463            },
464            MemoryRangeWithNode {
465                range: MemoryRange::new(4 * GB..TB + 2 * GB),
466                vnode: 0,
467            },
468        ];
469
470        let layout = MemoryLayout::new(TB, mmio, &[], &[], None).unwrap();
471        assert_eq!(
472            layout.ram(),
473            &[
474                MemoryRangeWithNode {
475                    range: MemoryRange::new(0..GB),
476                    vnode: 0
477                },
478                MemoryRangeWithNode {
479                    range: MemoryRange::new(2 * GB..3 * GB),
480                    vnode: 0
481                },
482                MemoryRangeWithNode {
483                    range: MemoryRange::new(4 * GB..TB + 2 * GB),
484                    vnode: 0
485                },
486            ]
487        );
488        assert_eq!(layout.mmio(), mmio);
489        assert_eq!(layout.ram_size(), TB);
490        assert_eq!(layout.end_of_ram(), TB + 2 * GB);
491        assert_eq!(layout.end_of_layout(), TB + 2 * GB);
492
493        let layout = MemoryLayout::new_from_ranges(ram, mmio).unwrap();
494        assert_eq!(
495            layout.ram(),
496            &[
497                MemoryRangeWithNode {
498                    range: MemoryRange::new(0..GB),
499                    vnode: 0
500                },
501                MemoryRangeWithNode {
502                    range: MemoryRange::new(2 * GB..3 * GB),
503                    vnode: 0
504                },
505                MemoryRangeWithNode {
506                    range: MemoryRange::new(4 * GB..TB + 2 * GB),
507                    vnode: 0
508                },
509            ]
510        );
511        assert_eq!(layout.mmio(), mmio);
512        assert_eq!(layout.ram_size(), TB);
513        assert_eq!(layout.end_of_ram(), TB + 2 * GB);
514        assert_eq!(layout.end_of_layout(), TB + 2 * GB);
515    }
516
517    #[test]
518    fn bad_layout() {
519        MemoryLayout::new(TB + 1, &[], &[], &[], None).unwrap_err();
520        let mmio = &[
521            MemoryRange::new(3 * GB..4 * GB),
522            MemoryRange::new(GB..2 * GB),
523        ];
524        MemoryLayout::new(TB, mmio, &[], &[], None).unwrap_err();
525
526        MemoryLayout::new_from_ranges(&[], mmio).unwrap_err();
527
528        let ram = &[MemoryRangeWithNode {
529            range: MemoryRange::new(0..GB),
530            vnode: 0,
531        }];
532        MemoryLayout::new_from_ranges(ram, mmio).unwrap_err();
533
534        let ram = &[MemoryRangeWithNode {
535            range: MemoryRange::new(0..GB + MB),
536            vnode: 0,
537        }];
538        let mmio = &[
539            MemoryRange::new(GB..2 * GB),
540            MemoryRange::new(3 * GB..4 * GB),
541        ];
542        MemoryLayout::new_from_ranges(ram, mmio).unwrap_err();
543
544        let mmio = &[
545            MemoryRange::new(GB..2 * GB),
546            MemoryRange::new(3 * GB..4 * GB),
547        ];
548        let pci_ecam = &[MemoryRange::new(GB..GB + MB)];
549        MemoryLayout::new(TB, mmio, pci_ecam, &[], None).unwrap_err();
550
551        let mmio = &[
552            MemoryRange::new(GB..2 * GB),
553            MemoryRange::new(3 * GB..4 * GB),
554        ];
555        let pci_mmio = &[MemoryRange::new(GB..GB + MB)];
556        MemoryLayout::new(TB, mmio, &[], pci_mmio, None).unwrap_err();
557
558        let pci_ecam = &[MemoryRange::new(GB..GB + MB)];
559        let pci_mmio = &[MemoryRange::new(GB..GB + MB)];
560        MemoryLayout::new(TB, &[], pci_ecam, pci_mmio, None).unwrap_err();
561    }
562
563    #[test]
564    fn resolved_ranges_constructor() {
565        let ram = vec![
566            MemoryRangeWithNode {
567                range: MemoryRange::new(0..GB),
568                vnode: 0,
569            },
570            MemoryRangeWithNode {
571                range: MemoryRange::new(2 * GB..3 * GB),
572                vnode: 1,
573            },
574        ];
575        let mmio = vec![MemoryRange::new(GB..2 * GB)];
576        let pci_ecam = vec![MemoryRange::new(4 * GB..4 * GB + MB)];
577        let pci_mmio = vec![MemoryRange::new(5 * GB..6 * GB)];
578
579        let layout = MemoryLayout::new_from_resolved_ranges(
580            ram.clone(),
581            mmio.clone(),
582            pci_ecam.clone(),
583            pci_mmio.clone(),
584            None,
585        )
586        .unwrap();
587
588        assert_eq!(layout.ram(), ram);
589        assert_eq!(layout.mmio(), mmio);
590        assert_eq!(layout.probe_address(4 * GB), Some(AddressType::PciEcam));
591        assert_eq!(layout.probe_address(5 * GB), Some(AddressType::PciMmio));
592    }
593
594    #[test]
595    fn resolved_ranges_reject_overlap_with_fixed_ranges() {
596        let ram = vec![MemoryRangeWithNode {
597            range: MemoryRange::new(0..2 * GB),
598            vnode: 0,
599        }];
600        let mmio = vec![MemoryRange::new(GB..2 * GB)];
601
602        assert!(MemoryLayout::new_from_resolved_ranges(ram, mmio, vec![], vec![], None).is_err());
603    }
604
605    #[test]
606    fn resolved_ranges_validate_vtl2_against_ram_end() {
607        let ram = vec![
608            MemoryRangeWithNode {
609                range: MemoryRange::new(0..GB),
610                vnode: 0,
611            },
612            MemoryRangeWithNode {
613                range: MemoryRange::new(3 * GB..4 * GB),
614                vnode: 0,
615            },
616        ];
617        let mmio = vec![MemoryRange::new(GB..2 * GB)];
618        let vtl2_range = MemoryRange::new(2 * GB..2 * GB + MB);
619
620        assert!(matches!(
621            MemoryLayout::new_from_resolved_ranges(ram, mmio, vec![], vec![], Some(vtl2_range)),
622            Err(Error::Vtl2RangeBeforeEndOfRam)
623        ));
624    }
625
626    #[test]
627    fn pci_ranges() {
628        let mmio = &[MemoryRange::new(3 * GB..4 * GB)];
629        let pci_ecam = &[MemoryRange::new(2 * TB - GB..2 * TB)];
630        let pci_mmio = &[
631            MemoryRange::new(2 * GB..3 * GB),
632            MemoryRange::new(5 * GB..6 * GB),
633        ];
634
635        let layout = MemoryLayout::new(TB, mmio, pci_ecam, pci_mmio, None).unwrap();
636        assert_eq!(
637            layout.ram(),
638            &[
639                MemoryRangeWithNode {
640                    range: MemoryRange::new(0..2 * GB),
641                    vnode: 0,
642                },
643                MemoryRangeWithNode {
644                    range: MemoryRange::new(4 * GB..5 * GB),
645                    vnode: 0,
646                },
647                MemoryRangeWithNode {
648                    range: MemoryRange::new(6 * GB..TB + 3 * GB),
649                    vnode: 0,
650                },
651            ]
652        );
653        assert_eq!(layout.end_of_layout(), 2 * TB);
654
655        assert_eq!(layout.probe_address(2 * GB), Some(AddressType::PciMmio));
656        assert_eq!(
657            layout.probe_address(2 * GB + MB),
658            Some(AddressType::PciMmio)
659        );
660        assert_eq!(layout.probe_address(5 * GB), Some(AddressType::PciMmio));
661        assert_eq!(
662            layout.probe_address(5 * GB + MB),
663            Some(AddressType::PciMmio)
664        );
665        assert_eq!(
666            layout.probe_address(2 * TB - GB),
667            Some(AddressType::PciEcam)
668        );
669    }
670
671    #[test]
672    fn probe_address() {
673        let mmio = &[
674            MemoryRange::new(GB..2 * GB),
675            MemoryRange::new(3 * GB..4 * GB),
676        ];
677        let ram = &[
678            MemoryRangeWithNode {
679                range: MemoryRange::new(0..GB),
680                vnode: 0,
681            },
682            MemoryRangeWithNode {
683                range: MemoryRange::new(2 * GB..3 * GB),
684                vnode: 0,
685            },
686            MemoryRangeWithNode {
687                range: MemoryRange::new(4 * GB..TB + 2 * GB),
688                vnode: 0,
689            },
690        ];
691
692        let layout = MemoryLayout::new_from_ranges(ram, mmio).unwrap();
693
694        assert_eq!(layout.probe_address(0), Some(AddressType::Ram));
695        assert_eq!(layout.probe_address(256), Some(AddressType::Ram));
696        assert_eq!(layout.probe_address(2 * GB), Some(AddressType::Ram));
697        assert_eq!(layout.probe_address(4 * GB), Some(AddressType::Ram));
698        assert_eq!(layout.probe_address(TB), Some(AddressType::Ram));
699        assert_eq!(layout.probe_address(TB + 1), Some(AddressType::Ram));
700
701        assert_eq!(layout.probe_address(GB), Some(AddressType::Mmio));
702        assert_eq!(layout.probe_address(GB + 123), Some(AddressType::Mmio));
703        assert_eq!(layout.probe_address(3 * GB), Some(AddressType::Mmio));
704
705        assert_eq!(layout.probe_address(TB + 2 * GB), None);
706        assert_eq!(layout.probe_address(TB + 3 * GB), None);
707        assert_eq!(layout.probe_address(4 * TB), None);
708    }
709
710    #[test]
711    fn numa_two_nodes_even_split() {
712        // 4 GB total, 2 nodes of 2 GB each, MMIO gap at 2-3 GB.
713        let mmio = &[MemoryRange::new(2 * GB..3 * GB)];
714        let layout = MemoryLayout::new_with_numa(&[2 * GB, 2 * GB], mmio, &[], &[], None).unwrap();
715        assert_eq!(
716            layout.ram(),
717            &[
718                MemoryRangeWithNode {
719                    range: MemoryRange::new(0..2 * GB),
720                    vnode: 0,
721                },
722                MemoryRangeWithNode {
723                    range: MemoryRange::new(3 * GB..5 * GB),
724                    vnode: 1,
725                },
726            ]
727        );
728        assert_eq!(layout.ram_size(), 4 * GB);
729    }
730
731    #[test]
732    fn numa_two_nodes_mid_chunk_split() {
733        // 4 GB total, 2 nodes of 2 GB each, MMIO gap at 3-4 GB.
734        // Node 0's 2 GB fits entirely below the gap; node 1 continues above.
735        // But the first chunk is 3 GB, so node 0 takes 2 GB and node 1
736        // takes the remaining 1 GB of that chunk, plus 1 GB above the gap.
737        let mmio = &[MemoryRange::new(3 * GB..4 * GB)];
738        let layout = MemoryLayout::new_with_numa(&[2 * GB, 2 * GB], mmio, &[], &[], None).unwrap();
739        assert_eq!(
740            layout.ram(),
741            &[
742                MemoryRangeWithNode {
743                    range: MemoryRange::new(0..2 * GB),
744                    vnode: 0,
745                },
746                MemoryRangeWithNode {
747                    range: MemoryRange::new(2 * GB..3 * GB),
748                    vnode: 1,
749                },
750                MemoryRangeWithNode {
751                    range: MemoryRange::new(4 * GB..5 * GB),
752                    vnode: 1,
753                },
754            ]
755        );
756        assert_eq!(layout.ram_size(), 4 * GB);
757    }
758
759    #[test]
760    fn numa_three_nodes() {
761        // 3 GB total, 3 nodes of 1 GB each, no gaps.
762        let layout = MemoryLayout::new_with_numa(&[GB, GB, GB], &[], &[], &[], None).unwrap();
763        assert_eq!(
764            layout.ram(),
765            &[
766                MemoryRangeWithNode {
767                    range: MemoryRange::new(0..GB),
768                    vnode: 0,
769                },
770                MemoryRangeWithNode {
771                    range: MemoryRange::new(GB..2 * GB),
772                    vnode: 1,
773                },
774                MemoryRangeWithNode {
775                    range: MemoryRange::new(2 * GB..3 * GB),
776                    vnode: 2,
777                },
778            ]
779        );
780    }
781
782    #[test]
783    fn numa_single_node_matches_new() {
784        // Single node should produce the same layout as new().
785        let mmio = &[
786            MemoryRange::new(GB..2 * GB),
787            MemoryRange::new(3 * GB..4 * GB),
788        ];
789        let layout_new = MemoryLayout::new(TB, mmio, &[], &[], None).unwrap();
790        let layout_numa = MemoryLayout::new_with_numa(&[TB], mmio, &[], &[], None).unwrap();
791        assert_eq!(layout_new.ram(), layout_numa.ram());
792    }
793
794    #[test]
795    fn numa_bad_inputs() {
796        // Empty sizes.
797        MemoryLayout::new_with_numa(&[], &[], &[], &[], None).unwrap_err();
798        // Non-page-aligned size.
799        MemoryLayout::new_with_numa(&[GB + 1], &[], &[], &[], None).unwrap_err();
800        // Zero size.
801        MemoryLayout::new_with_numa(&[0], &[], &[], &[], None).unwrap_err();
802        // Mixed: one valid, one zero.
803        MemoryLayout::new_with_numa(&[GB, 0], &[], &[], &[], None).unwrap_err();
804    }
805}