wdext 0.1.0

// SPDX-FileCopyrightText: 2026 takubokudori
// SPDX-License-Identifier: MIT OR Apache-2.0
use crate::{
    error::WdErrorKind, expr::PtrEvalExpr, object::traits::list::*,
    wd::dbgmodel::WdCallingConventionKind, *,
};
use std::{collections::HashSet, marker::PhantomData, num::ParseIntError};

/// Parses a MASM-formatted value into a [`DebuggeeOffset`].
///
/// ```
/// use wdext::util::parse_masm_value;
///
/// assert_eq!(parse_masm_value("0123456789abcdef"), Ok(0x0123456789abcdef));
/// assert_eq!(
///     parse_masm_value("0x0123456789abcdef"),
///     Ok(0x0123456789abcdef)
/// );
/// assert_eq!(parse_masm_value("1010h"), Ok(0x1010));
/// assert_eq!(parse_masm_value("0t1010"), Ok(0o1010));
/// assert_eq!(parse_masm_value("0y1010"), Ok(0b1010));
/// assert_eq!(parse_masm_value("0b1010"), Ok(0xb1010));
/// assert_eq!(parse_masm_value("0n1010"), Ok(1010));
/// assert_eq!(parse_masm_value("0x"), Ok(0));
/// assert_eq!(parse_masm_value("0n"), Ok(0));
/// assert_eq!(parse_masm_value("0t"), Ok(0));
/// assert_eq!(parse_masm_value("0y"), Ok(0));
/// assert_eq!(
///     parse_masm_value("012345678`9abcdef"),
///     Ok(0x0123456789abcdef)
/// );
/// assert_eq!(
///     parse_masm_value("0x012345678`9abcdef"),
///     Ok(0x0123456789abcdef)
/// );
/// assert_eq!(parse_masm_value("a`bcd"), Ok(0xabcd));
/// assert_eq!(parse_masm_value("0xa`bcd"), Ok(0xabcd));
/// assert_eq!(parse_masm_value("0x`bcd"), Ok(0xbcd));
/// assert_eq!(parse_masm_value("0xa`"), Ok(0xa));
/// assert_eq!(parse_masm_value("a`"), Ok(0xa));
/// assert_eq!(parse_masm_value("a`b``c``d"), Ok(0xabcd));
/// assert_eq!(parse_masm_value("--+-+1010h"), Ok(-4112i64 as u64));
/// parse_masm_value("").unwrap_err();
/// parse_masm_value("+").unwrap_err();
/// parse_masm_value("`bcd").unwrap_err();
/// parse_masm_value("bcdx").unwrap_err();
/// parse_masm_value("1`0h").unwrap_err();
/// parse_masm_value("h").unwrap_err();
/// ```
pub fn parse_masm_value(
    s: impl AsRef<str>,
) -> Result<DebuggeeOffset, ParseIntError> {
    let s = s.as_ref();
    let mut negative = false;
    let mut sign_len = 0;

    for &byte in s.as_bytes() {
        match byte {
            b'+' => sign_len += 1,
            b'-' => {
                sign_len += 1;
                negative = !negative;
            }
            _ => break,
        }
    }

    let value = parse_unsigned_masm_value(&s[sign_len..])?;

    Ok(if negative {
        value.wrapping_neg()
    } else {
        value
    })
}

fn parse_unsigned_masm_value(
    s: impl AsRef<str>,
) -> Result<DebuggeeOffset, ParseIntError> {
    let s = s.as_ref();
    if s.is_empty() {
        // Since we cannot create a ParseIntError ourselves, return this instead.
        return DebuggeeOffset::from_str_radix(s, 16);
    }

    if let Some(s) = s.strip_suffix("h") {
        return DebuggeeOffset::from_str_radix(s, 16);
    }
    if let Some(s) = s.strip_prefix("0n") {
        if s.is_empty() {
            return Ok(0);
        }
        return DebuggeeOffset::from_str_radix(s, 10);
    }
    if let Some(s) = s.strip_prefix("0t") {
        if s.is_empty() {
            return Ok(0);
        }
        return DebuggeeOffset::from_str_radix(s, 8);
    }
    if let Some(s) = s.strip_prefix("0y") {
        if s.is_empty() {
            return Ok(0);
        }
        return DebuggeeOffset::from_str_radix(s, 2);
    }
    if s.starts_with('`') {
        // Numbers starting with a backtick are invalid.
        // Since we cannot create a ParseIntError ourselves, return this instead.
        return DebuggeeOffset::from_str_radix(s, 16);
    }
    let s = s.strip_prefix("0x").unwrap_or(s);
    if s.is_empty() {
        return Ok(0);
    }
    if !s.as_bytes().contains(&b'`') {
        DebuggeeOffset::from_str_radix(s, 16)
    } else {
        let s = s.replace('`', "");
        DebuggeeOffset::from_str_radix(&s, 16)
    }
}

/// Iterator over [`ListEntry`] nodes.
pub struct ListEntryIterator<'a, LE: ListEntry> {
    cur: LE::Ptr,
    ds: &'a WdDataSpaces,
    memory: &'a Memory,
    check_sentinel: bool,
    sentinel: LE::Ptr,
    visited: HashSet<LE::Ptr>,
    forward: bool,
    detect_cycles: bool,
    infinite_loop: bool,
    finished: bool,
}

impl<'a, LE: ListEntry> ListEntryIterator<'a, LE> {
    /// Creates a new list entry iterator.
    ///
    /// - `addr_head`: Address of the `_LIST_ENTRY` used as the list head.
    /// - `forward`: If `true`, follows `Flink`; otherwise follows `Blink`.
    /// - `include_initial_entrry`: If `true`, yields `addr_head` as the first item. If `false`, starts from `head->Flink` or `head->Blink`.
    /// - `detect_cycles`: If `true`, encountering an already visited entry yields an error whose kind is [`WdErrorKind::ListEntryCycleDetected`].
    /// - `sentinel`: Address to use as the sentinel. If `None`, the address specified by `addr_head` is used.
    /// - `infinite_loop`: The iterator does not stop at the sentinel, and `detect_cycles` is disabled. It loops indefinitely unless a read error occurs.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        ds: &'a WdDataSpaces,
        memory: &'a Memory,
        addr_head: impl Into<LE::Ptr>,
        sentinel: Option<DebuggeeOffset>,
        forward: bool,
        include_initial_entry: bool,
        detect_cycles: bool,
        infinite_loop: bool,
    ) -> WdResult<Self> {
        let initial_offset = addr_head.into();
        let cur = if include_initial_entry {
            initial_offset
        } else {
            let l = initial_offset.read(ds, memory)?;
            if forward { l.flink() } else { l.blink() }
        };
        Ok(Self {
            cur,
            ds,
            memory,
            check_sentinel: !include_initial_entry,
            sentinel: sentinel.map_or(initial_offset, Into::into),
            visited: HashSet::new(),
            forward,
            detect_cycles: detect_cycles && !infinite_loop,
            infinite_loop,
            finished: false,
        })
    }
}

impl<'a, LE: ListEntry> Iterator for ListEntryIterator<'a, LE> {
    type Item = WdResult<TObject<LE::Ptr, LE>>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.finished {
            return None;
        }
        if !self.infinite_loop {
            if self.check_sentinel && self.sentinel == self.cur {
                self.finished = true;
                return None;
            }
            self.check_sentinel = true;
        }
        let cur = self.cur;
        if self.detect_cycles && !self.visited.insert(self.cur) {
            self.finished = true;
            return Some(Err(WdErrorKind::ListEntryCycleDetected.into_error()));
        }
        let l = match self.cur.read(self.ds, self.memory) {
            Ok(x) => x,
            Err(e) => {
                self.finished = true;
                return Some(Err(e));
            }
        };

        let next_addr = if self.forward { l.flink() } else { l.blink() };
        self.cur = next_addr;
        Some(Ok(TObject::new(cur, l)))
    }
}

impl<'a, LE: ListEntry> std::iter::FusedIterator for ListEntryIterator<'a, LE> {}

pub struct TypedListEntryIterator<'a, LE, T>
where
    LE: TypedListEntry<T>,
{
    it: ListEntryIterator<'a, LE>,
    phantom: PhantomData<T>,
}

impl<'a, LE, T> TypedListEntryIterator<'a, LE, T>
where
    LE: TypedListEntry<T>,
{
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        ds: &'a WdDataSpaces,
        memory: &'a Memory,
        addr_head: impl Into<LE::Ptr>,
        sentinel: Option<DebuggeeOffset>,
        forward: bool,
        include_initial_entry: bool,
        detect_cycles: bool,
        infinite_loop: bool,
    ) -> WdResult<Self> {
        Ok(Self {
            it: ListEntryIterator::new(
                ds,
                memory,
                addr_head,
                sentinel,
                forward,
                include_initial_entry,
                detect_cycles,
                infinite_loop,
            )?,
            phantom: Default::default(),
        })
    }
}

impl<'a, LE, T> Iterator for TypedListEntryIterator<'a, LE, T>
where
    LE: TypedListEntry<T>,
{
    type Item = WdResult<LE::TypedPtr>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.it.next()? {
            Ok(x) => Some(Ok(x
                .ptr()
                .offset()
                .wrapping_sub(LE::OFFSET_OF_ENTRY)
                .into())),
            Err(e) => Some(Err(e)),
        }
    }
}

impl<'a, LE, T> std::iter::FusedIterator for TypedListEntryIterator<'a, LE, T> where
    LE: TypedListEntry<T>
{
}

/// Iterator over [`SListEntry`] nodes.
pub struct SListEntryIterator<'a, LE: SListEntry> {
    cur: LE::Ptr,
    ds: &'a WdDataSpaces,
    memory: &'a Memory,
    check_sentinel: bool,
    sentinel: LE::Ptr,
    visited: HashSet<LE::Ptr>,
    detect_cycles: bool,
    finished: bool,
}

impl<'a, LE: SListEntry> SListEntryIterator<'a, LE> {
    /// Creates a new list entry iterator.
    ///
    /// - `addr_head`: Address of the `_LIST_ENTRY` used as the list head.
    /// - `detect_cycles`: If `true`, encountering an already visited entry yields an error whose kind is [`WdErrorKind::ListEntryCycleDetected`].
    /// - `sentinel`: Address to use as the sentinel. If `None`, `NULL` is used as the sentinel.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        ds: &'a WdDataSpaces,
        memory: &'a Memory,
        addr_head: impl Into<LE::Ptr>,
        sentinel: Option<DebuggeeOffset>,
        detect_cycles: bool,
    ) -> WdResult<Self> {
        let initial_offset = addr_head.into();
        Ok(Self {
            cur: initial_offset,
            ds,
            memory,
            check_sentinel: false,
            sentinel: sentinel.unwrap_or_default().into(),
            visited: HashSet::new(),
            detect_cycles,
            finished: false,
        })
    }
}

impl<'a, LE: SListEntry> Iterator for SListEntryIterator<'a, LE> {
    type Item = WdResult<TObject<LE::Ptr, LE>>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.finished {
            return None;
        }
        if self.check_sentinel && self.sentinel == self.cur {
            self.finished = true;
            return None;
        }
        self.check_sentinel = true;
        let cur = self.cur;
        if self.detect_cycles && !self.visited.insert(self.cur) {
            self.finished = true;
            return Some(Err(WdErrorKind::ListEntryCycleDetected.into_error()));
        }
        let l = match self.cur.read(self.ds, self.memory) {
            Ok(x) => x,
            Err(e) => {
                self.finished = true;
                return Some(Err(e));
            }
        };

        self.cur = l.next();
        if self.cur.is_null() {
            self.finished = true;
        }
        Some(Ok(TObject::new(cur, l)))
    }
}

impl<'a, LE: SListEntry> std::iter::FusedIterator
    for SListEntryIterator<'a, LE>
{
}

pub struct TypedSListEntryIterator<'a, LE, T>
where
    LE: TypedSListEntry<T>,
{
    it: SListEntryIterator<'a, LE>,
    phantom: PhantomData<T>,
}

impl<'a, LE, T> TypedSListEntryIterator<'a, LE, T>
where
    LE: TypedSListEntry<T>,
{
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        ds: &'a WdDataSpaces,
        memory: &'a Memory,
        addr_head: impl Into<LE::Ptr>,
        sentinel: Option<DebuggeeOffset>,
        detect_cycles: bool,
    ) -> WdResult<Self> {
        Ok(Self {
            it: SListEntryIterator::new(
                ds,
                memory,
                addr_head,
                sentinel,
                detect_cycles,
            )?,
            phantom: Default::default(),
        })
    }
}

impl<'a, LE, T> Iterator for TypedSListEntryIterator<'a, LE, T>
where
    LE: TypedSListEntry<T>,
{
    type Item = WdResult<LE::TypedPtr>;

    fn next(&mut self) -> Option<Self::Item> {
        match self.it.next()? {
            Ok(x) => Some(Ok(x
                .ptr()
                .offset()
                .wrapping_sub(LE::OFFSET_OF_ENTRY)
                .into())),
            Err(e) => Some(Err(e)),
        }
    }
}

impl<'a, LE, T> std::iter::FusedIterator for TypedSListEntryIterator<'a, LE, T> where
    LE: TypedSListEntry<T>
{
}

pub trait BitMapWord: Copy {
    const BITS: u32;

    fn is_bit_set(self, bit_index: usize) -> bool;
}

macro_rules! impl_bitmap_word {
    ($($ty:ty),* $(,)?) => {
        $(
            impl BitMapWord for $ty {
                const BITS: u32 = <$ty>::BITS;

                #[inline]
                fn is_bit_set(self, bit_index: usize) -> bool {
                    self & ((1 as $ty) << bit_index) != 0
                }
            }
        )*
    };
}

impl_bitmap_word!(u8, u16, u32, u64);

#[derive(Debug, Clone)]
pub struct BitMapIterator<'a, W> {
    bitmap: &'a [W],
    bit_idx: usize,
    bit_size: usize,
}

impl<'a, W: BitMapWord> BitMapIterator<'a, W> {
    pub fn new(
        bitmap: &'a [W],
        bit_size: impl Into<Option<usize>>,
    ) -> Option<Self> {
        let bm_len = bitmap.len();

        let capacity_bits = bm_len.checked_mul(W::BITS as usize)?;

        let bit_size = bit_size.into().unwrap_or(capacity_bits);
        if bit_size > capacity_bits {
            return None;
        }

        Some(Self {
            bitmap,
            bit_idx: 0,
            bit_size,
        })
    }

    pub fn bit_size(&self) -> usize { self.bit_size }

    pub fn remaining_bits(&self) -> usize { self.bit_size - self.bit_idx }

    #[inline]
    pub fn bit_at(&self, bit_idx: usize) -> Option<bool> {
        if bit_idx >= self.bit_size {
            return None;
        }

        let word_bits = W::BITS as usize;
        let word_idx = bit_idx / word_bits;
        let bit_in_word = bit_idx % word_bits;

        Some(self.bitmap[word_idx].is_bit_set(bit_in_word))
    }
}

impl<'a, W: BitMapWord> Iterator for BitMapIterator<'a, W> {
    type Item = bool;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.bit_idx >= self.bit_size {
            return None;
        }

        let word_bits = W::BITS as usize;
        let word_idx = self.bit_idx / word_bits;
        let bit_in_word = self.bit_idx % word_bits;

        let result = self.bitmap[word_idx].is_bit_set(bit_in_word);

        self.bit_idx += 1;

        Some(result)
    }

    #[inline]
    fn nth(&mut self, n: usize) -> Option<Self::Item> {
        let target = self.bit_idx.checked_add(n)?;

        if target >= self.bit_size {
            self.bit_idx = self.bit_size;
            return None;
        }

        let result = self.bit_at(target);
        self.bit_idx = target + 1;
        result
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.remaining_bits();
        (remaining, Some(remaining))
    }
}

impl<'a, W: BitMapWord> std::iter::FusedIterator for BitMapIterator<'a, W> {}

/// Call-site function-argument iterator.
///
/// Indefinitely returns pointer expressions for function arguments.
pub struct CallSiteArgsIterator<'a> {
    index: usize,
    reg_args: &'a [&'a str],
    sp_name: &'a str,
    stack_offset: u64,
    ptr_width: PointerWidth,
    fmt_width: usize,
}

impl CallSiteArgsIterator<'static> {
    pub fn new_with_calling_convention(
        ptr_width: PointerWidth,
        cconv: WdCallingConventionKind,
    ) -> WdResult<Self> {
        use PointerWidth::*;
        use WdCallingConventionKind::*;

        let (reg_args, sp_name, stack_offset): (&[&str], &str, u64) =
            match (ptr_width, cconv) {
                (_, Unknown) => {
                    return Err(WdErrorKind::InvalidArgument.into_error());
                }

                (Ptr32, CDecl | StdCall | SysCall) => (&[], "@esp", 0),
                (Ptr32, FastCall) => (&["@ecx", "@edx"], "@esp", 0),
                (Ptr32, ThisCall) => (&["@ecx"], "@esp", 0),
                (Ptr64, CDecl | StdCall | FastCall | ThisCall) => {
                    (&["@rcx", "@rdx", "@r8", "@r9"], "@rsp", 4)
                }
                (Ptr64, SysCall) => {
                    (&["@r10", "@rdx", "@r8", "@r9"], "@rsp", 5)
                }
            };
        Ok(Self::new(ptr_width, reg_args, sp_name, stack_offset))
    }

    /// Specifies the formatting width for stack offsets.
    ///
    /// For example, setting this to 1 changes `@rsp+0x8` to `@rsp+0x08`.
    pub fn fmt_width(mut self, fmt_width: usize) -> Self {
        self.fmt_width = fmt_width;

        self
    }

    pub fn fit_fmt_width(self, num: usize) -> Self {
        let stack_args =
            num.saturating_sub(1).saturating_sub(self.reg_args.len());
        let width = format!(
            "{:x}",
            stack_args.saturating_mul(self.ptr_width.size() as usize)
        )
        .len();
        self.fmt_width(width - 1)
    }
}

#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub enum ArgLocation {
    Register,
    Stack,
}

impl ArgLocation {
    pub fn is_register(self) -> bool { matches!(self, ArgLocation::Register) }

    pub fn is_stack(self) -> bool { matches!(self, ArgLocation::Stack) }
}

#[derive(Debug, Clone)]
pub struct ArgInfo {
    pub expression: PtrEvalExpr,
    pub location: ArgLocation,
}

impl ArgInfo {
    pub fn is_register(&self) -> bool { self.location.is_register() }

    pub fn is_stack(&self) -> bool { self.location.is_stack() }
}

impl<'a> CallSiteArgsIterator<'a> {
    pub fn new(
        ptr_width: PointerWidth,
        reg_args: &'a [&'a str],
        sp_name: &'a str,
        stack_offset: u64,
    ) -> Self {
        Self {
            index: 0,
            reg_args,
            sp_name,
            stack_offset,
            ptr_width,
            fmt_width: 0,
        }
    }

    fn expression_at(&self, index: usize) -> ArgInfo {
        if let Some(&reg) = self.reg_args.get(index) {
            return ArgInfo {
                expression: reg.into(),
                location: ArgLocation::Register,
            };
        }

        let stack_index = (index - self.reg_args.len()) as u64;
        let byte_offset =
            (stack_index + self.stack_offset) * self.ptr_width.size();

        ArgInfo {
            expression: format!(
                "{}{byte_offset:+#0width$x}",
                self.sp_name,
                width = self.fmt_width + 4
            )
            .into(),
            location: ArgLocation::Stack,
        }
    }
}

impl<'a> Iterator for CallSiteArgsIterator<'a> {
    type Item = ArgInfo;

    fn next(&mut self) -> Option<Self::Item> {
        let result = self.expression_at(self.index);
        self.index += 1;
        Some(result)
    }

    fn size_hint(&self) -> (usize, Option<usize>) { (usize::MAX, None) }

    fn nth(&mut self, n: usize) -> Option<Self::Item> {
        self.index = self.index.checked_add(n)?;
        self.next()
    }
}