From ac5712c84c543773e9e2e72755d775cdb3c091a5 Mon Sep 17 00:00:00 2001 From: Nathan Royer <61582713+NathanRoyer@users.noreply.github.com> Date: Fri, 22 Dec 2023 21:29:27 +0100 Subject: [PATCH 1/3] Restructure init of PCI legacy interrupts (INTx); add `lspci` app (#1081) * INTX interrupt handler registration for PCI devices is now done lazily on both architectures. * On aarch64, the interrupt numbers are known statically. * On x86_64, we expect that the driver or module using the PCI device knows which interrupt number to use. Technically this information can be discovered through ACPI AML, but we do not yet support that because it is tedious and difficult. * Clarify PCI interrupt information functions, e.g., MSI/MSI-x support, etc. Co-authored-by: Kevin Boos <1139460+kevinaboos@users.noreply.github.com> --- Cargo.lock | 11 ++ applications/lspci/Cargo.toml | 12 ++ applications/lspci/src/lib.rs | 75 ++++++++++ kernel/device_manager/src/lib.rs | 2 - kernel/e1000/src/lib.rs | 2 +- kernel/interrupts/src/aarch64/mod.rs | 22 --- kernel/pci/Cargo.toml | 1 + kernel/pci/src/lib.rs | 199 ++++++++++++++++++++------- 8 files changed, 249 insertions(+), 75 deletions(-) create mode 100644 applications/lspci/Cargo.toml create mode 100644 applications/lspci/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index c05769a28f..5bc9205de2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1925,6 +1925,16 @@ dependencies = [ "task", ] +[[package]] +name = "lspci" +version = "0.1.0" +dependencies = [ + "app_io", + "getopts", + "memory", + "pci", +] + [[package]] name = "lz4_flex" version = "0.9.3" @@ -2595,6 +2605,7 @@ dependencies = [ "arm_boards", "bit_field 0.7.0", "cpu", + "interrupt_controller", "interrupts", "log", "memory", diff --git a/applications/lspci/Cargo.toml b/applications/lspci/Cargo.toml new file mode 100644 index 0000000000..142434cc11 --- /dev/null +++ b/applications/lspci/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "lspci" +version = "0.1.0" +description = "An application which lists currently connected PCI devices." +authors = ["Nathan Royer "] +edition = "2021" + +[dependencies] +getopts = "0.2.21" +pci = { path = "../../kernel/pci" } +memory = { path = "../../kernel/memory" } +app_io = { path = "../../kernel/app_io" } diff --git a/applications/lspci/src/lib.rs b/applications/lspci/src/lib.rs new file mode 100644 index 0000000000..9ede0f89bf --- /dev/null +++ b/applications/lspci/src/lib.rs @@ -0,0 +1,75 @@ +//! This application lists currently connected PCI devices. + +#![no_std] + +extern crate alloc; +#[macro_use] extern crate app_io; +extern crate getopts; + +use alloc::vec::Vec; +use alloc::string::String; +use getopts::Options; +use pci::pci_device_iter; +use memory::PhysicalAddress; + +pub fn main(args: Vec) -> isize { + let mut opts = Options::new(); + opts.optflag("h", "help", "print this help menu"); + + let matches = match opts.parse(args) { + Ok(m) => m, + Err(_f) => { + println!("{}", _f); + print_usage(opts); + return -1; + } + }; + + if matches.opt_present("h") { + print_usage(opts); + return 0; + } + + if let Err(msg) = list_pci_devices() { + println!("Error: {}", msg); + } + + 0 +} + +fn list_pci_devices() -> Result<(), &'static str> { + for dev in pci_device_iter()? { + println!("{} -- {:04x}:{:04x}", dev.location, dev.vendor_id, dev.device_id); + println!("- class, subclass, prog_if: {:x}, {:x}, {:x}", dev.class, dev.subclass, dev.prog_if); + + for bar_idx in 0..6 { + let base = dev.determine_mem_base(bar_idx)?; + if base != PhysicalAddress::zero() { + let size = dev.determine_mem_size(bar_idx); + println!("- BAR {}: base = 0x{:x}, size = 0x{:x}", bar_idx, base, size); + } + } + + let support = dev.modern_interrupt_support(); + let supports = |b| match b { + true => "supported", + false => "not supported", + }; + + println!("- MSI interrupts: {}", supports(support.msi)); + println!("- MSI-X interrupts: {}", supports(support.msix)); + println!("- INTx enabled: {}", dev.pci_intx_enabled()); + println!("- INTx status: {}", dev.pci_get_intx_status(false)); + } + + Ok(()) +} + + +fn print_usage(opts: Options) { + println!("{}", opts.usage(USAGE)); +} + + +const USAGE: &str = "Usage: lspci +An application which lists currently connected PCI devices."; diff --git a/kernel/device_manager/src/lib.rs b/kernel/device_manager/src/lib.rs index 98581b8d1a..49089c69f2 100644 --- a/kernel/device_manager/src/lib.rs +++ b/kernel/device_manager/src/lib.rs @@ -84,8 +84,6 @@ pub fn init( } } - pci::init()?; - // Initialize/scan the PCI bus to discover PCI devices for dev in pci::pci_device_iter()? { debug!("Found PCI device: {:X?}", dev); diff --git a/kernel/e1000/src/lib.rs b/kernel/e1000/src/lib.rs index 715e470d00..c6c272c74d 100644 --- a/kernel/e1000/src/lib.rs +++ b/kernel/e1000/src/lib.rs @@ -156,7 +156,7 @@ impl E1000Nic { //debug!("e1000_nc bar_type: {0}, mem_base: {1}, io_base: {2}", e1000_nc.bar_type, e1000_nc.mem_base, e1000_nc.io_base); // Get interrupt number - let interrupt_num = match e1000_pci_dev.pci_get_interrupt_info() { + let interrupt_num = match e1000_pci_dev.pci_get_intx_info() { Ok((Some(irq), _pin)) => (irq + IRQ_BASE_OFFSET) as InterruptNumber, _ => return Err("e1000: PCI device had no interrupt number (IRQ vector)"), }; diff --git a/kernel/interrupts/src/aarch64/mod.rs b/kernel/interrupts/src/aarch64/mod.rs index e423391771..29590df929 100644 --- a/kernel/interrupts/src/aarch64/mod.rs +++ b/kernel/interrupts/src/aarch64/mod.rs @@ -217,28 +217,6 @@ pub fn init_pl011_rx_interrupt() -> Result<(), &'static str> { int_ctrl.set_destination(PL011_RX_SPI, Some(current_cpu()), u8::MAX) } -/// Sets an interrupt handler for legacy PCI interrupts: INTA, INTB, INTC, INTD -pub fn init_pci_interrupts(handlers: [InterruptHandler; 4]) -> Result<(), &'static str> { - let int_ctrl = SystemInterruptController::get() - .ok_or("SystemInterruptController was not yet initialized")?; - let dst = Some(cpu::bootstrap_cpu().unwrap()); - - let pci_intx_nums = BOARD_CONFIG.pci_intx.into_iter(); - let pci_intx_handlers = handlers.into_iter(); - - for (int_num, handler) in pci_intx_nums.zip(pci_intx_handlers) { - if let Err(existing_handler) = register_interrupt(int_num, handler) { - if handler as InterruptHandler != existing_handler { - return Err("A different interrupt handler has already been setup for that PCI interrupt"); - } - } - - int_ctrl.set_destination(int_num, dst, u8::MAX)?; - } - - Ok(()) -} - /// Registers an interrupt handler at the given IRQ interrupt number. /// /// The function fails if the interrupt number is reserved or is already in use. diff --git a/kernel/pci/Cargo.toml b/kernel/pci/Cargo.toml index 869d260769..f0b57682cc 100644 --- a/kernel/pci/Cargo.toml +++ b/kernel/pci/Cargo.toml @@ -21,6 +21,7 @@ port_io = { path = "../../libs/port_io" } [target.'cfg(target_arch = "aarch64")'.dependencies] arm_boards = { path = "../arm_boards" } +interrupt_controller = { path = "../interrupt_controller" } [lib] crate-type = ["rlib"] diff --git a/kernel/pci/src/lib.rs b/kernel/pci/src/lib.rs index d8e8dbe92d..0406a70217 100644 --- a/kernel/pci/src/lib.rs +++ b/kernel/pci/src/lib.rs @@ -1,5 +1,20 @@ //! PCI Configuration Space Access //! +//! ## Terminology +//! +//! This crate deals with multiple types of interrupts: +//! * Legacy (INTx) interrupts are the oldest PCI interrupt representation. +//! * Four interrupt pins are shared among all devices. +//! * This crate refers to these legacy interrupts as "intx". +//! +//! * MSI (Message Signaled Interrupts) appeared with PCI express. +//! * They allow devices to allocate up to 32 interrupt numbers. +//! * This crate refers to these interrupts as "msi". +//! +//! * MSI-X messages appeared with PCIe 3.0. +//! * They allow devices to allocate up to 2048 interrupt numbers. +//! * This crate refers to these interrupts as "msix". +//! //! Note: while pci currently uses port-io on x86 and mmio on aarch64, //! x86 may also support memory-based PCI configuration in the future; //! port-io is the legacy way to access the config space. @@ -10,6 +25,7 @@ #![no_std] #![allow(dead_code)] +#![feature(abi_x86_interrupt)] extern crate alloc; @@ -22,15 +38,18 @@ use bit_field::BitField; use volatile::Volatile; use zerocopy::FromBytes; use cpu::CpuId; -use interrupts::InterruptNumber; +use interrupts::{InterruptNumber, InterruptHandler, interrupt_handler, register_interrupt, EoiBehaviour}; #[cfg(target_arch = "x86_64")] -use port_io::Port; +use { + port_io::Port, + interrupts::IRQ_BASE_OFFSET, +}; #[cfg(target_arch = "aarch64")] use { arm_boards::BOARD_CONFIG, - interrupts::{EoiBehaviour, interrupt_handler, init_pci_interrupts}, + interrupt_controller::{SystemInterruptController, SystemInterruptControllerApi}, }; #[derive(Debug, Copy, Clone)] @@ -217,33 +236,74 @@ pub fn pci_device_iter() -> Result, &'s } static INTX_DEVICES: Mutex> = Mutex::new(Vec::new()); - -// Architecture-independent PCI interrupt handler -// Currently aarch64-only, because legacy interrupts aren't supported on x86 yet. -#[cfg(target_arch = "aarch64")] -interrupt_handler!(pci_int_handler, None, _stack_frame, { - let devices = INTX_DEVICES.lock(); - - for device in &*devices { - if device.pci_get_interrupt_status(true) { - device.pci_enable_interrupts(false); - log::info!("Device {} triggered an interrupt", device.location); - - let reader = device.interrupt_waker.lock(); - match &*reader { - Some(waker) => waker.wake_by_ref(), - None => log::error!("Device doesn't have an interrupt waker!"), +static INTX_NUMBERS: Mutex<[Option; 4]> = Mutex::new([None; 4]); + +// Architecture-independent INTx handlers +macro_rules! intx_handler { + ($name:ident, $num:literal) => { + interrupt_handler!($name, { + let intx_numbers = INTX_NUMBERS.lock(); + intx_numbers[$num].expect("uninitialized x86 PCI INTx handler") + }, _stack_frame, { + let devices = INTX_DEVICES.lock(); + + for device in &*devices { + if device.pci_get_intx_status(true) { + device.pci_enable_intx(false); + log::info!("Device {} triggered a legacy interrupt", device.location); + + let reader = device.intx_waker.lock(); + match &*reader { + Some(waker) => waker.wake_by_ref(), + None => log::error!("Device doesn't have an interrupt waker!"), + } + } } + + EoiBehaviour::HandlerDidNotSendEoi + }); + }; +} + +intx_handler!(pci_intx_handler_1, 0); +intx_handler!(pci_intx_handler_2, 1); +intx_handler!(pci_intx_handler_3, 2); +intx_handler!(pci_intx_handler_4, 3); + +static HANDLERS: [InterruptHandler; 4] = [pci_intx_handler_1, pci_intx_handler_2, pci_intx_handler_3, pci_intx_handler_4]; + +// Ensures that a PCI INTx interrupt handler is registered for this interrupt number +fn init_intx_handler(int_num: InterruptNumber) -> Result<(), &'static str> { + let mut intx_numbers = INTX_NUMBERS.lock(); + let mut slot = None; + + for i in 0..4 { + if intx_numbers[i] == Some(int_num) { + return Ok((/* this interrupt number was already initialized. */)); + } else if intx_numbers[i].is_none() { + slot = Some(i); } } - EoiBehaviour::HandlerDidNotSendEoi -}); + let slot = slot.ok_or("More than four different INTx numbers were encountered")?; + intx_numbers[slot] = Some(int_num); + core::mem::drop(intx_numbers); -/// Initializes the PCI interrupt handler -pub fn init() -> Result<(), &'static str> { - #[cfg(target_arch = "aarch64")] - init_pci_interrupts([pci_int_handler; 4])?; + let pci_intx_handler = HANDLERS[slot]; + if let Err(existing_handler) = register_interrupt(int_num, pci_intx_handler) { + if existing_handler != (pci_intx_handler as _) { + return Err("Couldn't lazily set PCI INTx handler"); + } + } + + #[cfg(target_arch = "aarch64")] { + let int_dst = Some(cpu::bootstrap_cpu().unwrap()); + let int_ctrl = SystemInterruptController::get() + .ok_or("SystemInterruptController was not yet initialized")?; + + // route interrupt to bootstrap processor + int_ctrl.set_destination(int_num, int_dst, u8::MAX)?; + } Ok(()) } @@ -326,11 +386,11 @@ fn scan_pci() -> Result, &'static str> { int_pin: location.pci_read_8(PCI_INTERRUPT_PIN), int_line: location.pci_read_8(PCI_INTERRUPT_LINE), location, - interrupt_waker: Mutex::new(None), + intx_waker: Mutex::new(None), }; // disable legacy interrupts initially - device.pci_enable_interrupts(false); + device.pci_enable_intx(false); device_list.push(device); } @@ -550,20 +610,15 @@ impl PciLocation { } /// Sets the PCI device's command bit 10 to disable legacy interrupts - pub fn pci_set_interrupt_disable_bit(&self, bit: bool) { + pub fn pci_set_intx_disable_bit(&self, bit: bool) { let command = self.pci_read_16(PCI_COMMAND); - // trace!("pci_set_interrupt_disable_bit: PciDevice: {}, read value: {:#x}", self, command); let new_value = match bit { true => command | PCI_COMMAND_INT_DISABLED, false => command & !PCI_COMMAND_INT_DISABLED, }; - self.pci_write_16(PCI_COMMAND, new_value); - /*trace!("pci_set_interrupt_disable_bit: PciDevice: {} read value AFTER WRITE CMD: {:#x}", - self, - self.pci_read_16(PCI_COMMAND), - );*/ + self.pci_write_16(PCI_COMMAND, new_value); } /// Explores the PCI config space and returns address of requested capability, if present. @@ -625,6 +680,13 @@ impl fmt::Debug for PciLocation { } } +/// Returned by [`PciDevice::modern_interrupt_support`] +pub struct ModernInterruptSupport { + /// `true` if this device supports MSI (Message Signaled Interrupts) + pub msi: bool, + /// `true` if this device supports MSI-X + pub msix: bool, +} /// Contains information common to every type of PCI Device, /// and offers functions for reading/writing to the PCI configuration space. @@ -637,7 +699,7 @@ pub struct PciDevice { pub location: PciLocation, /// The handling task for legacy PCI interrupts - pub interrupt_waker: Mutex>, + pub intx_waker: Mutex>, /// The class code, used to determine device type. pub class: u8, @@ -731,6 +793,14 @@ impl PciDevice { mem_size } + /// Queries and returns whether this PCI device supports MSI and MSI-X interrupts. + pub fn modern_interrupt_support(&self) -> ModernInterruptSupport { + ModernInterruptSupport { + msi: self.find_pci_capability(PciCapability::Msi).is_some(), + msix: self.find_pci_capability(PciCapability::Msix).is_some(), + } + } + /// Enable MSI interrupts for a PCI device. /// We assume the device only supports one MSI vector /// and set the interrupt number and core id for that vector. @@ -859,10 +929,10 @@ impl PciDevice { map_frame_range(mem_base, mem_size as usize, MMIO_FLAGS) } - /// Reads and returns this PCI device's interrupt line and interrupt pin registers. + /// Reads and returns this PCI device's INTx line and INTx pin registers. /// - /// Returns an error if this PCI device's interrupt pin value is invalid (greater than 4). - pub fn pci_get_interrupt_info(&self) -> Result<(Option, Option), &'static str> { + /// Returns an error if this PCI device's INTx pin value is invalid (greater than 4). + pub fn pci_get_intx_info(&self) -> Result<(Option, Option), &'static str> { let int_line = match self.pci_read_8(PCI_INTERRUPT_LINE) { 0xff => None, other => Some(other), @@ -881,25 +951,54 @@ impl PciDevice { } /// Enables/Disables legacy (INTx) interrupts for this device - pub fn pci_enable_interrupts(&self, enable: bool) { - self.pci_set_interrupt_disable_bit(!enable); + pub fn pci_enable_intx(&self, enable: bool) { + self.pci_set_intx_disable_bit(!enable); } - /// Reads and returns this PCI device's interrupt status flag. - pub fn pci_get_interrupt_status(&self, check_enabled: bool) -> bool { - const PCI_STATUS_INT: u16 = 1 << 3; + /// Checks that legacy interrupts are enabled in the command register. + pub fn pci_intx_enabled(&self) -> bool { + (self.pci_read_16(PCI_COMMAND) & PCI_COMMAND_INT_DISABLED) == 0 + } - let interrupt_enabled = || (self.pci_read_16(PCI_COMMAND) & PCI_COMMAND_INT_DISABLED) == 0; - let pending_interrupt = || (self.pci_read_16(PCI_STATUS) & PCI_STATUS_INT ) != 0; + /// Reads and returns this PCI device's legacy interrupt status flag. + pub fn pci_get_intx_status(&self, check_enabled: bool) -> bool { + const PCI_STATUS_INT: u16 = 1 << 3; + let pending_interrupt = || (self.pci_read_16(PCI_STATUS) & PCI_STATUS_INT) != 0; - ((!check_enabled) || interrupt_enabled()) && pending_interrupt() + ((!check_enabled) || self.pci_intx_enabled()) && pending_interrupt() } - /// Sets a task waker to be used when this device triggers an interrupt + /// Sets up the given `waker` to be woken when this PCI device triggers a legacy interrupt (INTx). /// /// Returns the previous interrupt waker for this device, if there was one. - pub fn set_interrupt_waker(&'static self, waker: Waker) -> Option { - let mut handle = self.interrupt_waker.lock(); + pub fn set_intx_waker(&'static self, waker: Waker) -> Result, &'static str> { + + // On x86, we don't yet support querying the ACPI tables to properly determine + // which interrupt numbers are used for legacy PCI interrupts. + // As a workaround, we lazily register these handlers when a driver + // calls this function, as by that time we're sure that + // the interrupt number in the device's config space is correct. + #[cfg(target_arch = "x86_64")] { + let int_num = match self.pci_get_intx_info() { + Ok((Some(irq), _pin)) => (irq + IRQ_BASE_OFFSET) as InterruptNumber, + _ => { + log::error!("Failed to get INTx info for PCI device {:?}", self); + return Err("PciDevice::set_intx_waker() failed to get INTx info"); + } + }; + + init_intx_handler(int_num)?; + } + + // On aarch64, we *do* know the interrupt numbers statically, + // but we do it lazily anyway for the sake of code clarity. + #[cfg(target_arch = "aarch64")] { + for int_num in BOARD_CONFIG.pci_intx { + init_intx_handler(int_num)?; + } + } + + let mut handle = self.intx_waker.lock(); let prev_value = handle.replace(waker); if prev_value.is_none() { @@ -907,7 +1006,7 @@ impl PciDevice { intx_devices.push(self) } - prev_value + Ok(prev_value) } } From 813015bed3665a9262b29bfbf4c51f94b1754d98 Mon Sep 17 00:00:00 2001 From: Niklas Date: Sat, 23 Dec 2023 19:54:27 +0100 Subject: [PATCH 2/3] Makefile: don't hardcode the path for `bash` (#1086) Some systems (e.g., NixOS) don't have `/bin/bash`, so use the `env` utility to locate `bash` instead. --- Makefile | 2 +- cfg/Config.mk | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index d6d1d219f1..1cf747dffd 100644 --- a/Makefile +++ b/Makefile @@ -1,7 +1,7 @@ ### This makefile is the top-level build script that builds all the crates in subdirectories ### and combines them into the final OS .iso image. ### It also provides convenient targets for running and debugging Theseus and using GDB on your host computer. -SHELL := /bin/bash +SHELL := /usr/bin/env bash ## Disable parallelism for this Makefile since it breaks the build, ## as our dependencies aren't perfectly specified for each target. diff --git a/cfg/Config.mk b/cfg/Config.mk index f01a672d54..eee4008f1b 100644 --- a/cfg/Config.mk +++ b/cfg/Config.mk @@ -3,7 +3,7 @@ ### So, to access the directory containing this file, you would use "../" .DEFAULT_GOAL := all -SHELL := /bin/bash +SHELL := /usr/bin/env bash ## specifies which architecture we're building for ARCH ?= x86_64 From 50c045ff6611e114a39be92115a1be1bb893b18e Mon Sep 17 00:00:00 2001 From: Klimenty Tsoutsman Date: Wed, 27 Dec 2023 16:55:48 +1100 Subject: [PATCH 3/3] Add `bit_set` library Necessary for upcoming `epoch_scheduler` improvements. Signed-off-by: Klimenty Tsoutsman --- libs/bit_set/Cargo.toml | 6 ++ libs/bit_set/src/iter.rs | 67 ++++++++++++++++ libs/bit_set/src/lib.rs | 165 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 238 insertions(+) create mode 100644 libs/bit_set/Cargo.toml create mode 100644 libs/bit_set/src/iter.rs create mode 100644 libs/bit_set/src/lib.rs diff --git a/libs/bit_set/Cargo.toml b/libs/bit_set/Cargo.toml new file mode 100644 index 0000000000..af52627a32 --- /dev/null +++ b/libs/bit_set/Cargo.toml @@ -0,0 +1,6 @@ +[package] +name = "bit_set" +version = "0.1.0" +authors = ["Klim Tsoutsman "] +description = "A bit set storing integers less than 64" +edition = "2021" diff --git a/libs/bit_set/src/iter.rs b/libs/bit_set/src/iter.rs new file mode 100644 index 0000000000..135f36c3c2 --- /dev/null +++ b/libs/bit_set/src/iter.rs @@ -0,0 +1,67 @@ +use core::intrinsics::unlikely; + +/// An iterator over a [`BitSet`]. +/// +/// [`BitSet`]: crate::BitSet +pub struct Iter { + set: u64, + current_mask: u64, +} + +impl Iter { + pub(crate) const fn new(set: u64) -> Self { + Self { + set, + current_mask: u64::MAX, + } + } +} + +impl Iterator for Iter { + type Item = usize; + + fn next(&mut self) -> Option { + let next_index = (self.set & self.current_mask).trailing_zeros(); + + if unlikely(next_index == 64) { + None + } else { + // https://users.rust-lang.org/t/how-to-make-an-integer-with-n-bits-set-without-overflow/63078 + self.current_mask = u64::MAX.checked_shl(next_index + 1).unwrap_or(0); + Some(next_index as usize) + } + } +} + +#[cfg(test)] +mod tests { + extern crate alloc; + + use alloc::vec::Vec; + + use crate::BitSet; + + #[test] + fn test_iter() { + let mut set = BitSet::new(); + set.insert(57); + set.insert(58); + set.insert(61); + set.insert(63); + assert_eq!(set.iter().collect::>(), [57, 58, 61, 63]); + + let mut set = BitSet::new(); + set.insert(0); + set.insert(8); + set.insert(16); + set.insert(24); + set.insert(32); + set.insert(40); + set.insert(48); + set.insert(56); + assert_eq!( + set.iter().collect::>(), + [0, 8, 16, 24, 32, 40, 48, 56] + ); + } +} diff --git a/libs/bit_set/src/lib.rs b/libs/bit_set/src/lib.rs new file mode 100644 index 0000000000..1724f46fc7 --- /dev/null +++ b/libs/bit_set/src/lib.rs @@ -0,0 +1,165 @@ +//! A bit set backed by a [`u64`]. +//! +//! See [`BitSet`] for more details. + +#![no_std] +#![feature(const_likely, core_intrinsics)] + +mod iter; + +use core::intrinsics::likely; + +pub use iter::Iter; + +/// A bit set backed by a [`u64`]. +/// +/// This is equivalent to a `HashSet` storing integers in the range `[0, +/// 64)`. +#[derive(Debug, Clone)] +pub struct BitSet { + inner: u64, +} + +impl BitSet { + /// Constructs a new, empty `BitSet`. + pub const fn new() -> Self { + Self { inner: 0 } + } + + /// Returns an iterator over the elements of the set. + pub const fn iter(&self) -> Iter { + Iter::new(self.inner) + } + + /// Returns `true` if the set contains the given element. + /// + /// # Panics + /// + /// Panics if `element` is greater than 63. + #[must_use] + pub const fn contains(&self, element: u8) -> bool { + assert!(element < 64); + self.inner & (1 << element) != 0 + } + + /// Adds an element to the set. + /// + /// # Panics + /// + /// Panics if `element` is greater than 63. + pub fn insert(&mut self, element: u8) { + assert!(element < 64); + self.inner |= 1 << element; + } + + /// Removes an element from the set. + /// + /// # Panics + /// + /// Panics if `element` is greater than 63. + pub fn remove(&mut self, element: u8) { + assert!(element < 64); + self.inner &= !(1 << element); + } + + /// Returns the smallest element in the set. + /// + /// Returns `None` if the set is empty. + #[must_use] + pub const fn min(&self) -> Option { + if likely(self.inner != 0) { + Some(self.inner.trailing_zeros() as u8) + } else { + None + } + } + + /// Returns the largest element in the set. + /// + /// Returns `None` if the set is empty. + #[must_use] + pub const fn max(&self) -> Option { + if likely(self.inner != 0) { + // self.inner.leading_zeros() <= 63 + Some(63 - self.inner.leading_zeros() as u8) + } else { + None + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_contains() { + let mut set = BitSet::new(); + + for i in 0..64 { + assert!(!set.contains(i)); + } + + set.insert(3); + + for i in 0..64 { + if i != 3 { + assert!(!set.contains(i)); + } else { + assert!(set.contains(i)); + } + } + + set.insert(0); + + for i in 0..64 { + if i != 0 && i != 3 { + assert!(!set.contains(i)); + } else { + assert!(set.contains(i)); + } + } + + set.insert(63); + + for i in 0..64 { + if i != 0 && i != 3 && i != 63 { + assert!(!set.contains(i)); + } else { + assert!(set.contains(i)); + } + } + } + + #[test] + fn test_remove() { + let mut set = BitSet::new(); + + set.insert(3); + set.insert(63); + set.remove(3); + + for i in 0..64 { + if i != 63 { + assert!(!set.contains(i)); + } else { + assert!(set.contains(i)); + } + } + } + + #[test] + fn test_min_max() { + let mut set = BitSet::new(); + assert_eq!(set.min(), None); + assert_eq!(set.max(), None); + + set.insert(5); + assert_eq!(set.min(), Some(5)); + assert_eq!(set.max(), Some(5)); + + set.insert(3); + assert_eq!(set.min(), Some(3)); + assert_eq!(set.max(), Some(5)); + } +}