Skip to main content

fake_mmio/inject/
dispatch.rs

1// Copyright 2026 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use std::collections::BTreeMap;
6use std::marker::PhantomData;
7use std::ops::Range;
8use std::sync::Arc;
9
10use mmio::Register;
11
12use crate::inject::{
13    GenericVmoMemoryHandler, MaybeVmoMemoryHandler, VmoMemoryHandler, VmoOpHelper,
14};
15use crate::operand::MmioOperand;
16
17/// Dispatches MMIO operations to specific handlers based on memory range.
18pub struct OpDispatcher {
19    ranges: BTreeMap<usize, (usize, DispatchNode)>,
20}
21
22impl OpDispatcher {
23    fn get(&self, range: Range<usize>) -> Option<&DispatchNode> {
24        let (start, (len, node)) = self.ranges.range(..=range.start).next_back()?;
25        let end = start.checked_add(*len).expect("Overflow calculating range end");
26        if range.start >= *start && range.start < end {
27            assert!(
28                range.end <= end,
29                "Register access in {range:?} extends beyond range {start}..{end}"
30            );
31            Some(node)
32        } else {
33            None
34        }
35    }
36
37    fn ensure(&self, range: Range<usize>) -> &DispatchNode {
38        self.get(range.clone()).unwrap_or_else(|| panic!("{range:?} not installed"))
39    }
40}
41
42impl GenericVmoMemoryHandler for OpDispatcher {
43    fn load<T: MmioOperand>(&self, op: VmoOpHelper<'_, T>) -> T {
44        self.ensure(op.range()).load(op)
45    }
46
47    fn store<T: MmioOperand>(&self, op: VmoOpHelper<'_, T>, value: T) {
48        self.ensure(op.range()).store(op, value)
49    }
50}
51
52/// A trait marking a [`VmoMemoryHandler`] as suitable for use within
53/// [`Dispatch`].
54pub trait DispatchHandler: VmoMemoryHandler + Send + Sync + 'static {}
55
56impl<O: VmoMemoryHandler + Send + Sync + 'static> DispatchHandler for O {}
57
58/// A reference-counted MMIO handler node used in dispatch trees.
59#[derive(Clone)]
60pub struct DispatchNode(Arc<dyn DispatchHandler>);
61
62impl DispatchNode {
63    /// Creates a new [`Node`] wrapping `handler`.
64    pub fn new(handler: Arc<dyn DispatchHandler>) -> Self {
65        Self(handler)
66    }
67
68    /// Takes this [`DispatchNode`] and returns a type that implements
69    /// [`DispatchHandler`].
70    ///
71    /// This is useful to reuse `DispatchNode`s with other `inject` combinators.
72    pub fn into_handler(self) -> DispatchNodeHandler {
73        DispatchNodeHandler(self)
74    }
75
76    pub(crate) fn load<T: MmioOperand>(&self, op: VmoOpHelper<'_, T>) -> T {
77        T::load(&*self.0, op)
78    }
79
80    pub(crate) fn store<T: MmioOperand>(&self, op: VmoOpHelper<'_, T>, value: T) {
81        T::store(&*self.0, op, value)
82    }
83}
84
85impl<T: DispatchHandler> From<T> for DispatchNode {
86    fn from(value: T) -> Self {
87        Self(Arc::new(value))
88    }
89}
90
91/// A wrapper around [`DispatchNode`] that implements [`DispatchHandler`].
92///
93/// Use [`DispatchNodeHandler::into_node`] to retrieve a [`DispatchNode`] from
94/// this. Note that there exists a blanket `From` impl that wraps any
95/// [`DispatchHandler`] into an `Arc` for `DispatchNode`.
96#[derive(Clone)]
97pub struct DispatchNodeHandler(DispatchNode);
98
99impl DispatchNodeHandler {
100    /// Takes this [`DispatchNodeHandler`] transforming it back into a
101    /// `DispatchNode` for installation in a `Dispatch`.
102    ///
103    /// See [`DispatchNode::into_handler`].
104    pub fn into_node(self) -> DispatchNode {
105        self.0
106    }
107}
108
109impl GenericVmoMemoryHandler for DispatchNodeHandler {
110    fn load<T: MmioOperand>(&self, op: VmoOpHelper<'_, T>) -> T {
111        self.0.load(op)
112    }
113
114    fn store<T: MmioOperand>(&self, op: VmoOpHelper<'_, T>, value: T) {
115        self.0.store(op, value)
116    }
117}
118
119/// Builder for constructing a [`Dispatch`] handler.
120///
121/// See [`DispatchBuilder::add`] for how to construct dispatched ranges.
122pub struct DispatchBuilder {
123    inner: OpDispatcher,
124}
125
126impl Default for DispatchBuilder {
127    fn default() -> Self {
128        Self::new()
129    }
130}
131
132impl DispatchBuilder {
133    /// Creates a new empty [`DispatchBuilder`] with a no-op range.
134    pub fn new() -> Self {
135        Self { inner: OpDispatcher { ranges: Default::default() } }
136    }
137
138    /// Registers `entries` into this builder.
139    pub fn insert<T: DispatchBuilderEntries>(mut self, entries: T) -> Self {
140        for (range, node) in entries.into_entries() {
141            assert!(range.start < range.end, "invalid or empty range: {range:?}");
142            for (start, (len, _)) in self.inner.ranges.range(..range.end) {
143                let existing = *start..(*start + *len);
144                assert!(
145                    !(range.start < existing.end && existing.start < range.end),
146                    "{range:?} overlaps with {existing:?}"
147                );
148            }
149
150            assert!(
151                self.inner.ranges.insert(range.start, (range.end - range.start, node)).is_none()
152            );
153        }
154
155        self
156    }
157
158    /// A shorthand for [`DispatchBuilder::insert`] with a tuple of `(range,
159    /// node)`.
160    pub fn insert_node<T: Into<DispatchNode>>(self, range: Range<usize>, node: T) -> Self {
161        self.insert((range, node))
162    }
163
164    /// A shorthand for [`DispatchBuilder::insert`] with a tuple of
165    /// `(RegisterRange, node)`.
166    pub fn insert_reg<R: Register, T: Into<DispatchNode>>(self, node: T) -> Self {
167        self.insert((RegisterRange::<R>::new(), node))
168    }
169
170    /// Consumes the builder and returns the configured [`Dispatch`].
171    pub fn build(self) -> OpDispatcher {
172        self.inner
173    }
174}
175
176/// A trait abstractig addition of ranges to a [`DispatchBuilder`].
177pub trait DispatchBuilderEntries {
178    /// Consumes this type returning a series of `Range` and `DispatchNode` for
179    /// insertion in a [`DispatchBuilder`].
180    fn into_entries(self) -> impl Iterator<Item = (Range<usize>, DispatchNode)>;
181}
182
183impl<T: Into<DispatchNode>> DispatchBuilderEntries for (Range<usize>, T) {
184    fn into_entries(self) -> impl Iterator<Item = (Range<usize>, DispatchNode)> {
185        let (range, node) = self;
186        std::iter::once((range, node.into()))
187    }
188}
189
190/// A helper struct for producing `Range<usize>` from a [`Register`] definition.
191#[derive(Debug, Copy, Clone)]
192pub struct RegisterRange<R: Register>(PhantomData<R>);
193
194impl<R: Register> RegisterRange<R> {
195    /// Creates a new `RegisterRange`.
196    pub fn new() -> Self {
197        Self(PhantomData)
198    }
199
200    /// Returns the address range for register `R`.
201    pub fn into_range(self) -> Range<usize> {
202        R::OFFSET..(R::OFFSET + std::mem::size_of::<R::Value>())
203    }
204}
205
206impl<R: Register> Default for RegisterRange<R> {
207    fn default() -> Self {
208        Self::new()
209    }
210}
211
212impl<R: Register, T: Into<DispatchNode>> DispatchBuilderEntries for (RegisterRange<R>, T) {
213    fn into_entries(self) -> impl Iterator<Item = (Range<usize>, DispatchNode)> {
214        let (range, node) = self;
215        std::iter::once((range.into_range(), node.into()))
216    }
217}
218
219impl MaybeVmoMemoryHandler for OpDispatcher {
220    fn get_handler(&self, range: Range<usize>) -> Option<&(impl VmoMemoryHandler + ?Sized)> {
221        self.get(range).map(|node| &*node.0)
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228    use crate::CachedVmoMemory;
229    use crate::inject::VmoMemory;
230    use crate::inject::passthrough::Passthrough;
231    use mmio::MmioExt as _;
232
233    struct TestReg;
234    impl Register for TestReg {
235        type Value = u32;
236        const OFFSET: usize = 0x10;
237        fn from_raw(_value: u32) -> Self {
238            Self
239        }
240        fn to_raw(&self) -> u32 {
241            0
242        }
243    }
244
245    #[test]
246    fn test_dispatch_builder_and_node() {
247        let builder = DispatchBuilder::default();
248        let node = DispatchNode::from(Passthrough);
249        let handler = node.clone().into_handler();
250        let node_back = handler.into_node();
251
252        let reg_range = RegisterRange::<TestReg>::default();
253        assert_eq!(reg_range.into_range(), 0x10..0x14);
254
255        let dispatch =
256            builder.insert_node(0..4, node_back).insert_reg::<TestReg, _>(Passthrough).build();
257
258        let mut mmio = CachedVmoMemory::new_mapped(4096)
259            .expect("Failed to map VMO.")
260            .map(|vmo| VmoMemory::new(dispatch, vmo));
261
262        let val1 = 0x1122_3344_u32;
263        mmio.store::<u32>(0, val1);
264        assert_eq!(mmio.load::<u32>(0), val1);
265
266        let val2 = 0x5566_7788_u32;
267        mmio.store::<u32>(0x10, val2);
268        assert_eq!(mmio.load::<u32>(0x10), val2);
269    }
270
271    #[fuchsia::test(logging = false)]
272    #[should_panic(expected = "overlaps with")]
273    fn test_dispatch_overlapping_ranges_panics() {
274        let _ =
275            DispatchBuilder::new().insert_node(0..10, Passthrough).insert_node(5..15, Passthrough);
276    }
277
278    #[fuchsia::test(logging = false)]
279    #[should_panic(expected = "invalid or empty range")]
280    #[allow(clippy::reversed_empty_ranges)]
281    fn test_dispatch_invalid_range_panics() {
282        let _ = DispatchBuilder::new().insert_node(10..5, Passthrough);
283    }
284
285    #[fuchsia::test(logging = false)]
286    #[should_panic(expected = "invalid or empty range")]
287    fn test_dispatch_empty_range_panics() {
288        let _ = DispatchBuilder::new().insert_node(5..5, Passthrough);
289    }
290
291    #[fuchsia::test(logging = false)]
292    #[should_panic(expected = "not installed")]
293    fn test_dispatch_uninstalled_range_panics() {
294        let dispatch = DispatchBuilder::new().insert_node(0..4, Passthrough).build();
295        let mmio = CachedVmoMemory::new_mapped(4096)
296            .expect("Failed to map VMO.")
297            .map(|vmo| VmoMemory::new(dispatch, vmo));
298
299        let _ = mmio.load::<u32>(8);
300    }
301
302    #[fuchsia::test(logging = false)]
303    #[should_panic(expected = "extends beyond range")]
304    fn test_dispatch_extends_beyond_range_panics() {
305        let dispatch = DispatchBuilder::new().insert_node(0..4, Passthrough).build();
306        let mmio = CachedVmoMemory::new_mapped(4096)
307            .expect("Failed to map VMO.")
308            .map(|vmo| VmoMemory::new(dispatch, vmo));
309
310        let _ = mmio.load::<u64>(0);
311    }
312}