disc-v 0.1.2

// Copyright (c) 2026 Christopher Swenson
// Licensed under the Apache-2.0 license.

//! Extensions, and sets of them.

use core::fmt;

/// Defines [`Extension`], with each variant's name in an ISA string, and
/// [`Extension::ALL`], so that every variant is in `ALL`.
macro_rules! extensions {
    ($($(#[$doc:meta])* $variant:ident = $name:literal,)*) => {
        /// A RISC-V extension that defines instructions disc-v can decode.
        ///
        /// Extensions that define no instructions, such as Zkt or Sstc, are
        /// not listed; [`Decoder::from_march`](crate::Decoder::from_march)
        /// accepts and ignores them.
        #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
        #[non_exhaustive]
        pub enum Extension {
            $($(#[$doc])* $variant,)*
        }

        impl Extension {
            /// Every extension, in declaration order.
            const ALL: &[Extension] = &[$(Extension::$variant),*];

            /// The extension's name in an ISA string, such as `"zba"` or
            /// `"m"`.
            pub const fn name(self) -> &'static str {
                match self {
                    $(Extension::$variant => $name,)*
                }
            }
        }
    };
}

extensions! {
    /// The RV32E and RV64E bases: only registers x0-x15 exist, so
    /// instructions that use x16-x31 are reserved. This restricts the base
    /// ISA rather than adding instructions, so it is not in
    /// [`Extensions::DEFAULT`].
    E = "e",
    /// Integer multiplication and division.
    M = "m",
    /// Atomic instructions.
    A = "a",
    /// Single-precision floating point.
    F = "f",
    /// Double-precision floating point. Implies F.
    D = "d",
    /// Quad-precision floating point. Implies D.
    Q = "q",
    /// Compressed instructions: Zca, and Zcf and Zcd when F and D are
    /// enabled.
    C = "c",
    /// Compressed integer instructions.
    Zca = "zca",
    /// Compressed single-precision loads and stores (RV32 only). Implies Zca
    /// and F.
    Zcf = "zcf",
    /// Compressed double-precision loads and stores. Implies Zca and D.
    Zcd = "zcd",
    /// Compressed push, pop and register moves (`cm.push`, ...). Implies
    /// Zca, and conflicts with Zcd.
    Zcmp = "zcmp",
    /// Compressed table jumps (`cm.jt`, `cm.jalt`). Implies Zca and Zicsr,
    /// and conflicts with Zcd.
    Zcmt = "zcmt",
    /// Load and store register pairs on RV32 (`ld`, `sd`).
    Zilsd = "zilsd",
    /// Compressed load and store register pairs on RV32 (`c.ld`, `c.sd`,
    /// ...). Implies Zilsd and Zca, and conflicts with Zcf.
    Zclsd = "zclsd",
    /// Vectors.
    V = "v",
    /// The hypervisor extension.
    H = "h",
    /// Control and status register instructions.
    Zicsr = "zicsr",
    /// `fence.i`.
    Zifencei = "zifencei",
    /// Integer conditional operations (`czero.*`).
    Zicond = "zicond",
    /// Load-acquire and store-release (`lw.aq`, `sw.rl`, ...).
    Zalasr = "zalasr",
    /// Atomic compare-and-swap (`amocas.*`). Implies A.
    Zacas = "zacas",
    /// Byte and halfword atomics (`amoadd.b`, ...). Implies A.
    Zabha = "zabha",
    /// Wait-on-reservation-set (`wrs.*`).
    Zawrs = "zawrs",
    /// Cache-block management (`cbo.clean`, `cbo.flush`, `cbo.inval`).
    Zicbom = "zicbom",
    /// Cache-block zero (`cbo.zero`).
    Zicboz = "zicboz",
    /// Cache-block prefetch hints (`prefetch.*`).
    Zicbop = "zicbop",
    /// Non-temporal locality hints (`ntl.*`).
    Zihintntl = "zihintntl",
    /// The `pause` hint.
    Zihintpause = "zihintpause",
    /// May-be-operations (`mop.r.N`, `mop.rr.N`).
    Zimop = "zimop",
    /// Compressed may-be-operations (`c.mop.N`). Implies C.
    Zcmop = "zcmop",
    /// Shadow stacks (`sspush`, `sspopchk`, `ssrdp`, `ssamoswap.*`).
    /// Implies Zimop.
    Zicfiss = "zicfiss",
    /// Landing pads (`lpad`).
    Zicfilp = "zicfilp",
    /// Address generation (`sh1add`, ...).
    Zba = "zba",
    /// Basic bit manipulation (`andn`, `clz`, ...).
    Zbb = "zbb",
    /// Carry-less multiplication (`clmul`, ...).
    Zbc = "zbc",
    /// Single-bit instructions (`bset`, ...).
    Zbs = "zbs",
    /// Bit manipulation for cryptography (`pack`, `brev8`, ...). It also
    /// provides several Zbb instructions, such as `rol` and `rev8`.
    Zbkb = "zbkb",
    /// Carry-less multiplication for cryptography: Zbc's `clmul` and
    /// `clmulh`.
    Zbkc = "zbkc",
    /// Crossbar permutations (`xperm4`, `xperm8`).
    Zbkx = "zbkx",
    /// AES decryption.
    Zknd = "zknd",
    /// AES encryption.
    Zkne = "zkne",
    /// SHA-256 and SHA-512.
    Zknh = "zknh",
    /// SM4.
    Zksed = "zksed",
    /// SM3.
    Zksh = "zksh",
    /// Simple compressed instructions (`c.lbu`, `c.mul`, ...). Implies C.
    Zcb = "zcb",
    /// Half-precision floating point. Implies Zfhmin.
    Zfh = "zfh",
    /// Minimal half-precision floating point: loads, stores and
    /// conversions. Implies F.
    Zfhmin = "zfhmin",
    /// Additional floating-point instructions (`fli`, `fround`, ...).
    /// Implies F.
    Zfa = "zfa",
    /// Conversions between BFloat16 and single precision. Implies F.
    Zfbfmin = "zfbfmin",
    /// Single-precision floating point in the integer registers: F's
    /// instructions except its loads, stores and moves, with integer
    /// register operands. Conflicts with F (and so with the extensions that
    /// imply it, such as D, Zfa and C's Zcf).
    Zfinx = "zfinx",
    /// Double-precision floating point in the integer registers (register
    /// pairs on RV32). Implies Zfinx.
    Zdinx = "zdinx",
    /// Half-precision floating point in the integer registers. Implies
    /// Zhinxmin.
    Zhinx = "zhinx",
    /// Minimal half-precision floating point in the integer registers: the
    /// conversions. Implies Zfinx.
    Zhinxmin = "zhinxmin",
    /// Vector BFloat16 conversions. Implies V.
    Zvfbfmin = "zvfbfmin",
    /// Vector BFloat16 widening multiply-add. Implies Zvfbfmin.
    Zvfbfwma = "zvfbfwma",
    /// Vector bit manipulation. Implies V.
    Zvbb = "zvbb",
    /// Vector carry-less multiplication. Implies V.
    Zvbc = "zvbc",
    /// Vector GCM/GMAC. Implies V.
    Zvkg = "zvkg",
    /// Vector AES. Implies V.
    Zvkned = "zvkned",
    /// Vector SHA-256. Implies V.
    Zvknha = "zvknha",
    /// Vector SHA-256 and SHA-512. Implies Zvknha, whose instructions it
    /// shares.
    Zvknhb = "zvknhb",
    /// Vector SM4. Implies V.
    Zvksed = "zvksed",
    /// Vector SM3. Implies V.
    Zvksh = "zvksh",
    /// Fine-grained address-translation cache invalidation (`sinval.vma`,
    /// ...).
    Svinval = "svinval",
    /// Debug mode (`dret`).
    Sdext = "sdext",
    /// Resumable non-maskable interrupts (`mnret`).
    Smrnmi = "smrnmi",
    /// Control-transfer records (`sctrclr`).
    Ssctr = "ssctr",
}

// Extensions are a bit set in a u128.
const _: () = assert!(Extension::ALL.len() <= 128);

impl Extension {
    /// The extension named `name` in an ISA string, ignoring case.
    pub fn from_name(name: &str) -> Option<Extension> {
        Extension::ALL
            .iter()
            .copied()
            .find(|e| e.name().eq_ignore_ascii_case(name))
    }

    /// The extensions this one implies.
    const fn implies(self) -> &'static [Extension] {
        use Extension::*;
        match self {
            D | Zfhmin | Zfa | Zfbfmin => &[F],
            C | Zcb | Zcmop | Zcmp => &[Zca],
            Zcf => &[Zca, F],
            Zcd => &[Zca, D],
            Zcmt => &[Zca, Zicsr],
            Zclsd => &[Zilsd, Zca],
            Q => &[D],
            Zfh => &[Zfhmin],
            Zdinx | Zhinxmin => &[Zfinx],
            Zhinx => &[Zhinxmin],
            Zacas | Zabha => &[A],
            Zicfiss => &[Zimop],
            Zvbb | Zvbc | Zvkg | Zvkned | Zvknha | Zvksed | Zvksh | Zvfbfmin => &[V],
            Zvfbfwma => &[Zvfbfmin],
            Zvknhb => &[Zvknha],
            _ => &[],
        }
    }

    /// The extension that provides this one's instructions with their
    /// floating-point operands in the integer registers, such as Zfinx for
    /// F, or the extension itself if there is none.
    pub(crate) const fn in_x_registers(self) -> Extension {
        use Extension::*;
        match self {
            F => Zfinx,
            D => Zdinx,
            Zfh => Zhinx,
            Zfhmin => Zhinxmin,
            ext => ext,
        }
    }

    const fn bit(self) -> u128 {
        1 << self as u32
    }
}

impl fmt::Display for Extension {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.name())
    }
}

/// A set of [`Extension`]s.
///
/// Adding an extension also adds the extensions it implies (D implies F),
/// and removing one also removes the extensions that imply it (removing F
/// removes D).
///
/// ```
/// use disc_v::{Extension, Extensions};
///
/// let exts = Extensions::GC.with(Extension::Zba).without(Extension::D);
/// assert!(exts.contains(Extension::F));
/// assert!(!exts.contains(Extension::Q));
/// ```
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct Extensions(u128);

impl Extensions {
    /// No extensions: only the base integer ISA.
    pub const EMPTY: Extensions = Extensions(0);

    /// Every supported extension that does not conflict with another. This
    /// is what [`decode`](fn@crate::decode) and the other free functions use.
    pub const DEFAULT: Extensions = Extensions::of(Extension::ALL)
        .without(Extension::E)
        .without(Extension::Zcmp)
        .without(Extension::Zcmt)
        .without(Extension::Zclsd)
        .without(Extension::Zfinx);

    /// G and C: IMAFD, Zicsr, Zifencei and C, as in RV32GC and RV64GC.
    pub const GC: Extensions = {
        use Extension::*;
        Extensions::of(&[M, A, F, D, C, Zicsr, Zifencei])
    };

    /// The extensions of the RVA23U64 profile that define instructions.
    pub const RVA23U64: Extensions = {
        use Extension::*;
        Extensions::of(&[
            M,
            A,
            F,
            D,
            C,
            V,
            Zicsr,
            Zicond,
            Zawrs,
            Zicbom,
            Zicboz,
            Zicbop,
            Zihintntl,
            Zihintpause,
            Zimop,
            Zcmop,
            Zba,
            Zbb,
            Zbs,
            Zcb,
            Zfhmin,
            Zfa,
            Zvbb,
        ])
    };

    /// The extensions of the RVA23S64 profile that define instructions:
    /// RVA23U64, Zifencei, Svinval and H.
    pub const RVA23S64: Extensions = {
        use Extension::*;
        Extensions::RVA23U64.union(Extensions::of(&[Zifencei, Svinval, H]))
    };

    /// The extensions in `exts`, and the extensions they imply.
    pub const fn of(exts: &[Extension]) -> Extensions {
        let mut set = Extensions::EMPTY;
        let mut i = 0;
        while i < exts.len() {
            set = set.with(exts[i]);
            i += 1;
        }
        set
    }

    /// Whether the set contains `ext`.
    pub const fn contains(self, ext: Extension) -> bool {
        self.0 & ext.bit() != 0
    }

    /// The set with `ext` and the extensions it implies added.
    pub const fn with(self, ext: Extension) -> Extensions {
        if self.contains(ext) {
            return self;
        }
        let mut set = Extensions(self.0 | ext.bit());
        let implied = ext.implies();
        let mut i = 0;
        while i < implied.len() {
            set = set.with(implied[i]);
            i += 1;
        }
        // C includes Zcf and Zcd when F and D are enabled.
        if set.contains(Extension::C) && set.contains(Extension::F) {
            set = set.with(Extension::Zcf);
        }
        if set.contains(Extension::C) && set.contains(Extension::D) {
            set = set.with(Extension::Zcd);
        }
        set
    }

    /// The set with `ext` and the extensions that imply it removed.
    ///
    /// C is a bundle of Zca and (with F and D) Zcf and Zcd: removing C
    /// removes every compressed instruction, and removing Zcf or Zcd removes
    /// C but keeps Zca.
    pub const fn without(self, ext: Extension) -> Extensions {
        let mut set = self.remove(ext);
        match ext {
            Extension::C => set = set.remove(Extension::Zca),
            Extension::Zcf | Extension::Zcd => set.0 &= !Extension::C.bit(),
            _ => {}
        }
        set
    }

    /// The set with `ext` and the extensions that imply it removed.
    const fn remove(self, ext: Extension) -> Extensions {
        let mut set = Extensions(self.0 & !ext.bit());
        let mut i = 0;
        while i < Extension::ALL.len() {
            let other = Extension::ALL[i];
            let implied = other.implies();
            let mut j = 0;
            while j < implied.len() {
                if implied[j] as u32 == ext as u32 && set.contains(other) {
                    set = set.remove(other);
                }
                j += 1;
            }
            i += 1;
        }
        set
    }

    /// The extensions in either set (and the extensions they imply
    /// together, such as Zcd for C and D).
    pub const fn union(self, other: Extensions) -> Extensions {
        let mut set = self;
        let mut i = 0;
        while i < Extension::ALL.len() {
            if other.contains(Extension::ALL[i]) {
                set = set.with(Extension::ALL[i]);
            }
            i += 1;
        }
        set
    }

    /// The extensions in the set, in declaration order.
    pub fn iter(self) -> impl Iterator<Item = Extension> {
        Extension::ALL
            .iter()
            .copied()
            .filter(move |&e| self.contains(e))
    }
}

impl Default for Extensions {
    /// [`Extensions::DEFAULT`].
    fn default() -> Self {
        Extensions::DEFAULT
    }
}

impl fmt::Debug for Extensions {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_set().entries(self.iter()).finish()
    }
}

impl From<Extension> for Extensions {
    fn from(ext: Extension) -> Self {
        Extensions::EMPTY.with(ext)
    }
}

impl<const N: usize> From<[Extension; N]> for Extensions {
    fn from(exts: [Extension; N]) -> Self {
        Extensions::of(&exts)
    }
}

impl From<&[Extension]> for Extensions {
    fn from(exts: &[Extension]) -> Self {
        Extensions::of(exts)
    }
}

impl FromIterator<Extension> for Extensions {
    fn from_iter<I: IntoIterator<Item = Extension>>(iter: I) -> Self {
        iter.into_iter().fold(Extensions::EMPTY, Extensions::with)
    }
}

impl core::ops::BitOr for Extensions {
    type Output = Extensions;

    fn bitor(self, other: Extensions) -> Extensions {
        self.union(other)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn names_round_trip() {
        for &ext in Extension::ALL {
            assert_eq!(Extension::from_name(ext.name()), Some(ext));
        }
        assert_eq!(Extension::from_name("ZBA"), Some(Extension::Zba));
        assert_eq!(Extension::from_name("xtheadba"), None);
    }

    #[test]
    fn all_lists_every_extension_once() {
        let mut bits = 0u128;
        for (i, &ext) in Extension::ALL.iter().enumerate() {
            assert_eq!(ext as usize, i);
            bits |= ext.bit();
        }
        assert_eq!(bits.count_ones() as usize, Extension::ALL.len());
    }

    #[test]
    fn implications() {
        use Extension::*;
        let set = Extensions::EMPTY.with(Q);
        assert_eq!(set, Extensions::of(&[F, D, Q]));
        assert_eq!(set.without(F), Extensions::EMPTY);
        assert_eq!(set.without(Q), Extensions::of(&[F, D]));
        assert!(Extensions::from(Zvknhb).contains(V));
    }
}