dbgscope 0.3.0

use std::collections::{HashMap, HashSet};
use std::sync::Mutex;

use super::{PoolDiagnostics, PoolSpan, PoolState, snapshot::PoolSnapshot, snapshot::WalkStalls};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Direction {
    Underflow,
    Overflow,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct Hole {
    pub span_index: usize,
    pub address: u64,
    pub size: u64,
    pub state: PoolState,
    pub distance: u64,
    pub immediate: bool,
}

#[derive(Debug, Clone, Default)]
pub(crate) struct PoolIndex {
    pub layout: crate::allocator::LayoutProvenance,
    pub spans: Vec<PoolSpan>,
    pub postings: HashMap<u32, Vec<usize>>,
    pub diagnostics: PoolDiagnostics,
    /// Whether the walk covered everything it set out to. Carried through from
    /// [`PoolSnapshot`] because a walk can end incomplete *without* emitting a
    /// diagnostic — `walk_vs` clears it when a readable region stops mid-chunk — and a
    /// caller that sees only counts and diagnostics would read that as a healthy result.
    pub complete: bool,
    /// Carried through from [`PoolSnapshot::budget_expired`] for the same reason `complete` is:
    /// an index that outlives the walk is the only thing a caller ever sees, so a reason left
    /// behind here is a reason nobody can recover.
    pub budget_expired: bool,
    /// The match threshold that intentionally stopped the walk, when one fired.
    pub stopped_after_matches: Option<usize>,
    /// Carried through from [`PoolSnapshot::stalls`], for the same reason the two flags above
    /// are: the index is all a caller ever sees of the walk.
    pub stalls: WalkStalls,
    /// Carried through from [`PoolSnapshot::refused_chunks`].
    pub refused_chunks: u64,
    /// Carried through from [`PoolSnapshot::unplaced_bytes`].
    pub unplaced_bytes: u64,
    row_postings: HashMap<RowIdentity, Vec<usize>>,
    span_rows: Vec<RowIdentity>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct RowIdentity {
    pool_kind: super::PoolKind,
    numa_node: u16,
    heap: super::HeapIdentity,
    backend: super::PoolBackend,
    subsegment: Option<u64>,
    size_class: u32,
}

impl From<&PoolSpan> for RowIdentity {
    fn from(span: &PoolSpan) -> Self {
        Self {
            pool_kind: span.pool_kind,
            numa_node: span.numa_node,
            heap: span.heap,
            backend: span.backend,
            subsegment: span.subsegment,
            size_class: span.size_class,
        }
    }
}

impl PoolIndex {
    pub(crate) fn build(mut snapshot: PoolSnapshot) -> Self {
        snapshot
            .spans
            .sort_by_key(|span| (span.heap, span.usable_address));
        let mut postings: HashMap<u32, Vec<usize>> = HashMap::new();
        let mut row_postings: HashMap<RowIdentity, Vec<usize>> = HashMap::new();
        let mut span_rows = Vec::with_capacity(snapshot.spans.len());
        for (index, span) in snapshot.spans.iter().enumerate() {
            if span.state == PoolState::Allocated {
                postings.entry(span.raw_tag).or_default().push(index);
            }
            let row = RowIdentity::from(span);
            row_postings.entry(row).or_default().push(index);
            span_rows.push(row);
        }
        Self {
            layout: snapshot.layout,
            spans: snapshot.spans,
            postings,
            diagnostics: snapshot.diagnostics,
            complete: snapshot.complete,
            budget_expired: snapshot.budget_expired,
            stopped_after_matches: snapshot.stopped_after_matches,
            stalls: snapshot.stalls,
            refused_chunks: snapshot.refused_chunks,
            unplaced_bytes: snapshot.unplaced_bytes,
            row_postings,
            span_rows,
        }
    }

    fn same_boundary(left: &PoolSpan, right: &PoolSpan) -> bool {
        left.pool_kind == right.pool_kind
            && left.numa_node == right.numa_node
            && left.heap == right.heap
            && left.backend == right.backend
            && left.subsegment == right.subsegment
            // Both have to be *chunks*, not merely readable ones: a gap has no header and no
            // successor, so letting one be a neighbour puts an allocation next to memory
            // nothing was decoded out of. Asked through `is_chunk` rather than named here, so
            // that a second kind of gap cannot arrive and pass this test by default — which is
            // exactly what `PoolState::Uncommitted` would have done.
            && left.state.is_chunk()
            && right.state.is_chunk()
    }

    pub(crate) fn predecessor(&self, index: usize) -> Option<usize> {
        let previous = index.checked_sub(1)?;
        Self::same_boundary(self.spans.get(previous)?, self.spans.get(index)?).then_some(previous)
    }

    pub(crate) fn successor(&self, index: usize) -> Option<usize> {
        let next = index.checked_add(1)?;
        Self::same_boundary(self.spans.get(index)?, self.spans.get(next)?).then_some(next)
    }

    pub(crate) fn holes(&self, allocation: usize, direction: Direction) -> Vec<Hole> {
        let Some(origin) = self.spans.get(allocation) else {
            return Vec::new();
        };
        let mut holes = Vec::new();
        let mut cursor = allocation;
        loop {
            let next = match direction {
                Direction::Underflow => self.predecessor(cursor),
                Direction::Overflow => self.successor(cursor),
            };
            let Some(index) = next else { break };
            let span = &self.spans[index];
            if span.state == PoolState::Allocated {
                break;
            }
            let distance = match direction {
                Direction::Underflow => origin.usable_address.saturating_sub(span.end()),
                Direction::Overflow => span.usable_address.saturating_sub(origin.end()),
            };
            holes.push(Hole {
                span_index: index,
                address: span.usable_address,
                size: span.size,
                state: span.state,
                distance,
                immediate: holes.is_empty() && distance == 0,
            });
            cursor = index;
        }
        holes
    }

    /// Matching allocations plus their local geometry. Unrelated allocations and
    /// free spans are deliberately retained so tag filtering never erases a hole.
    pub(crate) fn context_for_tag(&self, tag: u32) -> Vec<usize> {
        let mut selected_rows = HashSet::new();
        for &index in self.postings.get(&tag).into_iter().flatten() {
            if let Some(row) = self.span_rows.get(index) {
                selected_rows.insert(*row);
            }
        }
        let mut result: Vec<_> = selected_rows
            .into_iter()
            .flat_map(|row| self.row_postings.get(&row).into_iter().flatten().copied())
            .collect();
        result.sort_unstable();
        result
    }

    pub(crate) fn ranked_holes(&self, allocation: usize, direction: Direction) -> Vec<Hole> {
        let Some(origin) = self.spans.get(allocation) else {
            return Vec::new();
        };
        let mut holes = self.holes(allocation, direction);
        holes.sort_by_key(|hole| {
            let span = &self.spans[hole.span_index];
            (
                !hole.immediate,
                hole.state != PoolState::ReusableFree,
                span.backend != origin.backend || span.size_class != origin.size_class,
                hole.distance,
                u64::MAX - hole.size,
                hole.address,
            )
        });
        holes
    }
}

#[derive(Debug, Clone)]
struct CachedSnapshot {
    /// The full target identity, not just the generation counter. The generation alone
    /// is only bumped by debugger session notifications, which a purely programmatic
    /// host never receives — so a snapshot taken against one target could otherwise be
    /// handed back for another, silently describing memory that no longer exists.
    key: super::layout::SessionKey,
    index: PoolIndex,
}

#[derive(Default)]
pub(crate) struct SnapshotCache {
    entry: Mutex<Option<CachedSnapshot>>,
}

impl SnapshotCache {
    /// Generic over the builder's error so a walk's failure reaches the caller as itself.
    /// Pinning it to `String` here is what flattened an *interrupt* into "walking the pool
    /// failed" — a cache has no business deciding how a failure reads.
    pub(crate) fn get_or_refresh<F, E>(
        &self,
        key: super::layout::SessionKey,
        refresh: bool,
        build: F,
    ) -> Result<PoolIndex, E>
    where
        F: FnOnce() -> Result<PoolSnapshot, E>,
    {
        // A forced refresh supersedes the cached view even if rebuilding fails or
        // is interrupted. Never fall back to geometry the user explicitly replaced.
        if refresh {
            self.invalidate();
        }
        if !refresh
            && let Some(cached) = self.entry.lock().unwrap().as_ref()
            && cached.key == key
        {
            return Ok(cached.index.clone());
        }
        let snapshot = build()?;
        // Cached unless the walk was **cut short**. A budget that ran out or a match threshold
        // that fired left pool unwalked that the next query may need, and a snapshot that stops
        // where the target does not is the one kind a cache must not serve. A walk that
        // reached the end of the pool and found some of it unreadable is not that: a page a
        // live kernel has trimmed is as unreadable on the next walk as on this one, and on a
        // live kernel some always are — measured on two guests (2026-10-07), every walk ended
        // partial, so nothing was ever cached and a test asking twenty questions paid twenty
        // walks of 136 s each for one halted target that had not changed. The index carries
        // `complete`, `budget_expired` and `stopped_after_matches` with it, so a caller can
        // still read what the cached walk did not reach.
        let cacheable = !snapshot.budget_expired && snapshot.stopped_after_matches.is_none();
        let index = PoolIndex::build(snapshot);
        if cacheable {
            *self.entry.lock().unwrap() = Some(CachedSnapshot {
                key,
                index: index.clone(),
            });
        } else {
            self.invalidate();
        }
        Ok(index)
    }

    pub(crate) fn invalidate(&self) {
        *self.entry.lock().unwrap() = None;
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicUsize, Ordering};

    use super::*;
    use crate::pool::{HeapIdentity, PoolBackend, PoolKind};

    fn span(address: u64, tag: u32, state: PoolState, heap: u64) -> PoolSpan {
        let mut span = PoolSpan::allocation(
            address,
            0x20,
            tag,
            PoolKind::NonPagedNx,
            HeapIdentity {
                pool_state: 1,
                heap,
                special: false,
            },
            PoolBackend::Lfh,
        );
        span.state = state;
        span
    }

    #[test]
    fn test_pool_snapshot_index_and_cache_invalidation() {
        let tag = u32::from_le_bytes(*b"TEST");
        let other = u32::from_le_bytes(*b"OTHR");
        let snapshot = PoolSnapshot {
            spans: vec![
                span(0x1000, other, PoolState::Allocated, 1),
                span(0x1020, tag, PoolState::Allocated, 1),
                span(0x1040, 0, PoolState::ReusableFree, 1),
                span(0x1060, 0, PoolState::CachedFree, 1),
                span(0x1080, tag, PoolState::Allocated, 2),
            ],
            complete: true,
            ..PoolSnapshot::default()
        };
        let index = PoolIndex::build(snapshot.clone());
        assert_eq!(index.postings[&tag], vec![1, 4]);
        assert_eq!(index.successor(1), Some(2));
        assert_eq!(index.successor(3), None); // heap boundary
        let holes = index.ranked_holes(1, Direction::Overflow);
        assert_eq!(holes[0].address, 0x1040);
        assert!(holes[0].immediate);
        assert!(index.context_for_tag(tag).contains(&0)); // unrelated neighbor preserved

        let mut unreadable = span(0x1060, 0, PoolState::Unreadable, 1);
        unreadable.size_class = 0x20;
        let contextual = PoolIndex::build(PoolSnapshot {
            spans: vec![
                span(0x1000, tag, PoolState::Allocated, 1),
                span(0x1020, other, PoolState::Allocated, 1),
                span(0x1040, 0, PoolState::ReusableFree, 1),
                unreadable,
                span(0x1080, other, PoolState::Allocated, 1),
            ],
            complete: true,
            ..PoolSnapshot::default()
        });
        assert_eq!(contextual.context_for_tag(tag), vec![0, 1, 2, 3, 4]);
        assert_eq!(contextual.successor(2), None);

        // The other kind of gap, and the same answer. A span with nothing behind it has no
        // header and no successor either, so it cannot be an allocation's neighbour — and
        // arriving after this test was written is exactly why the check asks `is_chunk`
        // rather than naming `Unreadable`.
        let mut uncommitted = span(0x1060, 0, PoolState::Uncommitted, 1);
        uncommitted.size_class = 0x20;
        let across_a_gap = PoolIndex::build(PoolSnapshot {
            spans: vec![
                span(0x1040, 0, PoolState::ReusableFree, 1),
                uncommitted,
                span(0x1080, other, PoolState::Allocated, 1),
            ],
            complete: true,
            ..PoolSnapshot::default()
        });
        assert_eq!(across_a_gap.successor(0), None);
        assert_eq!(across_a_gap.predecessor(2), None);

        // Distinct targets, not just distinct generations: the cache keys on the whole
        // SessionKey so a snapshot cannot outlive the target it described.
        let session = |value: u64| super::super::layout::SessionKey {
            image: crate::dbgeng::KernelImage {
                base: 0x1000,
                ..Default::default()
            },
            session: value,
            target: 1,
        };
        let cache = SnapshotCache::default();
        let builds = AtomicUsize::new(0);
        // The error type is never produced here, only named: `get_or_refresh` is generic
        // over it now, so a builder that cannot fail still has to say what failing would look
        // like.
        let make = || -> Result<PoolSnapshot, String> {
            builds.fetch_add(1, Ordering::SeqCst);
            Ok(snapshot.clone())
        };
        cache.get_or_refresh(session(7), false, make).unwrap();
        cache.get_or_refresh(session(7), false, make).unwrap();
        assert_eq!(builds.load(Ordering::SeqCst), 1);
        cache.get_or_refresh(session(7), true, make).unwrap();
        assert_eq!(builds.load(Ordering::SeqCst), 2);
        cache.invalidate();
        cache.get_or_refresh(session(7), false, make).unwrap();
        cache.get_or_refresh(session(8), false, make).unwrap();
        assert_eq!(builds.load(Ordering::SeqCst), 4);

        // A walk that reached the end of the pool and found some of it unreadable is cached:
        // on a live kernel every walk ends that way, and the next question is about the same
        // halted target.
        let incomplete_builds = AtomicUsize::new(0);
        let make_incomplete = || -> Result<PoolSnapshot, String> {
            incomplete_builds.fetch_add(1, Ordering::SeqCst);
            let mut value = snapshot.clone();
            value.complete = false;
            Ok(value)
        };
        cache.invalidate();
        cache
            .get_or_refresh(session(9), false, make_incomplete)
            .unwrap();
        cache
            .get_or_refresh(session(9), false, make_incomplete)
            .unwrap();
        assert_eq!(incomplete_builds.load(Ordering::SeqCst), 1);

        // A walk cut short by its budget is not: what it did not reach is unwalked rather than
        // unreadable, and the next query may be the one that needs it.
        let expired_builds = AtomicUsize::new(0);
        let make_expired = || -> Result<PoolSnapshot, String> {
            expired_builds.fetch_add(1, Ordering::SeqCst);
            let mut value = snapshot.clone();
            value.complete = false;
            value.budget_expired = true;
            Ok(value)
        };
        cache.invalidate();
        cache
            .get_or_refresh(session(9), false, make_expired)
            .unwrap();
        cache
            .get_or_refresh(session(9), false, make_expired)
            .unwrap();
        assert_eq!(expired_builds.load(Ordering::SeqCst), 2);

        cache.invalidate();
        cache.get_or_refresh(session(10), false, make).unwrap();
        cache
            .get_or_refresh(session(10), true, make_expired)
            .unwrap();
        cache.get_or_refresh(session(10), false, make).unwrap();
        assert_eq!(builds.load(Ordering::SeqCst), 6);
        assert_eq!(expired_builds.load(Ordering::SeqCst), 3);

        let failed_refresh: Result<PoolIndex, String> =
            cache.get_or_refresh(session(10), true, || Err("interrupted".into()));
        assert_eq!(failed_refresh.unwrap_err(), "interrupted");
        cache.get_or_refresh(session(10), false, make).unwrap();
        assert_eq!(builds.load(Ordering::SeqCst), 7);
        // Two targets sharing a kernel base *and* generation — two dumps from one boot,
        // or a restarted session — must not share a cached snapshot. Nothing bumps the
        // generation for a purely programmatic host, so the target identity is the only
        // thing that tells them apart.
        cache.invalidate();
        let same_base = |target: u64| super::super::layout::SessionKey {
            image: crate::dbgeng::KernelImage {
                base: 0x1000,
                ..Default::default()
            },
            session: 1,
            target,
        };
        cache.get_or_refresh(same_base(100), false, make).unwrap();
        cache.get_or_refresh(same_base(101), false, make).unwrap();
        assert_eq!(builds.load(Ordering::SeqCst), 9);
    }

    #[test]
    fn test_match_limited_snapshot_is_never_cached_as_exhaustive() {
        let session = super::super::layout::SessionKey {
            image: crate::dbgeng::KernelImage {
                base: 0x1000,
                ..Default::default()
            },
            session: 1,
            target: 1,
        };
        let cache = SnapshotCache::default();
        let builds = AtomicUsize::new(0);
        let make_stopped = || -> Result<PoolSnapshot, String> {
            builds.fetch_add(1, Ordering::SeqCst);
            Ok(PoolSnapshot {
                complete: false,
                stopped_after_matches: Some(1),
                spans: vec![span(
                    0x1000,
                    u32::from_le_bytes(*b"Tgsm"),
                    PoolState::Allocated,
                    1,
                )],
                ..PoolSnapshot::default()
            })
        };

        let first = cache.get_or_refresh(session, false, make_stopped).unwrap();
        let second = cache.get_or_refresh(session, false, make_stopped).unwrap();
        assert_eq!(first.stopped_after_matches, Some(1));
        assert_eq!(second.stopped_after_matches, Some(1));
        assert_eq!(builds.load(Ordering::SeqCst), 2);

        let exhaustive_builds = AtomicUsize::new(0);
        let make_complete = || -> Result<PoolSnapshot, String> {
            exhaustive_builds.fetch_add(1, Ordering::SeqCst);
            Ok(PoolSnapshot {
                complete: true,
                spans: vec![
                    span(
                        0x1000,
                        u32::from_le_bytes(*b"Tgsm"),
                        PoolState::Allocated,
                        1,
                    ),
                    span(
                        0x1020,
                        u32::from_le_bytes(*b"Tgsm"),
                        PoolState::Allocated,
                        1,
                    ),
                ],
                ..PoolSnapshot::default()
            })
        };
        let complete = cache.get_or_refresh(session, false, make_complete).unwrap();
        let reused = cache
            .get_or_refresh(session, false, || -> Result<PoolSnapshot, String> {
                panic!("a complete cached snapshot should satisfy a limited query")
            })
            .unwrap();
        assert_eq!(complete.spans.len(), 2);
        assert_eq!(reused.spans.len(), 2);
        assert_eq!(reused.stopped_after_matches, None);
        assert_eq!(exhaustive_builds.load(Ordering::SeqCst), 1);
    }
}