Skip to main content

tg_rcore_tutorial_driver/
devices.rs

1use alloc::{collections::btree_map::BTreeMap, sync::Arc};
2use tg_console::println;
3use virtio_drivers::{DeviceType, Hal, MmioTransport, Transport, VirtIOHeader};
4
5use crate::{GpuDevice, VirtIOGpuWrapper, buffer::OverflowStrategy, input::{InputDevice, VirtIOInputWrapper}, block::{BlockDevice, VirtIOBlockWrapper}, qemu::{claim_irq, complete_irq, setup_interrupt_for, enable_external_interrupts}, address::*};
6
7pub trait Device : Send + Sync {
8    fn handle_irq(&self);
9}
10
11pub struct VirtIOProbingIterator {
12    current: usize
13}
14
15impl VirtIOProbingIterator {
16    pub fn new() -> Self {
17        Self {
18            current: 0
19        }
20    }
21}
22
23impl Iterator for VirtIOProbingIterator {
24    type Item = (usize, &'static mut VirtIOHeader);
25
26    fn next(&mut self) -> Option<Self::Item> {
27        let addr = VIRTIO_START + self.current * VIRTIO_STEP;
28        if addr >= VIRTIO_END {
29            return None;
30        } else {
31            self.current += 1;
32            // Note that addr 9x1000_1000 has id 1, so add self.current first
33            Some((self.current, unsafe { &mut *(addr as *mut VirtIOHeader) }))
34        }
35    }
36}
37
38pub struct DeviceManager {
39    gpu: Option<Arc<dyn GpuDevice>>,
40    keyboard: Option<Arc<dyn InputDevice>>,
41    block: Option<Arc<dyn BlockDevice>>,
42    irq_map: BTreeMap<usize, Arc<dyn Device>>,
43}
44
45impl DeviceManager {
46    pub fn new<H: Hal + 'static>() -> Self {
47        let mut manager = Self {
48            gpu: None,
49            keyboard: None,
50            block: None,
51            irq_map: BTreeMap::new(),
52        };
53
54        let mut devices = VirtIOProbingIterator::new();
55
56        while let Some((slot, header)) = devices.next() {
57            let Ok(transport) = (unsafe { MmioTransport::new(core::ptr::NonNull::from(header)) }) else {
58                continue;
59            };
60            match transport.device_type() {
61                DeviceType::GPU if manager.gpu.is_none() => {
62                    if let Ok(gpu) = VirtIOGpuWrapper::<H>::new(transport) {
63                        let gpu = Arc::new(gpu);
64                        manager.gpu = Some(gpu.clone());
65                        manager.irq_map.insert(slot, gpu.clone() as _);
66                        // Interrupts for GPU aren't quite needed now
67                        // setup_interrupt_for(slot);
68                        println!("GPU device found at slot {slot}");
69                    }
70                }
71                DeviceType::Input if manager.keyboard.is_none() => {
72                    if let Ok(keyboard) = VirtIOInputWrapper::<H>::new(transport, OverflowStrategy::DropOldest) {
73                        let keyboard = Arc::new(keyboard);
74                        manager.keyboard = Some(keyboard.clone());
75                        manager.irq_map.insert(slot, keyboard.clone() as _);
76                        setup_interrupt_for(slot);
77                        println!("Keyboard device found at slot {slot}");
78                    }
79                }
80                DeviceType::Block if manager.block.is_none() => {
81                    if let Ok(block) = VirtIOBlockWrapper::<H>::new(transport) {
82                        let block = Arc::new(block);
83                        manager.block = Some(block.clone());
84                        manager.irq_map.insert(slot, block.clone() as _);
85                        setup_interrupt_for(slot);
86                        println!("Block device found at slot {slot}");
87                    }
88                }
89                _ => {}
90            }
91        }
92
93        enable_external_interrupts();
94
95        manager
96    }
97
98    fn get_device(&self, slot: u32) -> Option<Arc<dyn Device>> {
99        self.irq_map.get(&(slot as usize)).cloned()
100    }
101
102    pub fn get_gpu(&self) -> Option<Arc<dyn GpuDevice>> {
103        self.gpu.clone()
104    }
105
106    pub fn get_keyboard(&self) -> Option<Arc<dyn InputDevice>> {
107        self.keyboard.clone()
108    }
109
110    pub fn get_block(&self) -> Option<Arc<dyn BlockDevice>> {
111        self.block.clone()
112    }
113
114    pub fn handle_external_interrupt(&self) {
115        let Some(slot) = claim_irq() else { return; };
116        self.get_device(slot).unwrap().handle_irq();
117        complete_irq(slot);
118    }
119}