fast_ntfs 1.0.1

// Copyright 2021-2026 Colin Finck <colin@reactos.org>
// SPDX-License-Identifier: MIT OR Apache-2.0

use core::cell::RefCell;
use core::cmp::Ordering;
use core::marker::PhantomData;

use alloc::vec;
use alloc::vec::Vec;

use crate::attribute::{NtfsAttributeItem, NtfsAttributeType};
use crate::error::{NtfsError, Result};
use crate::index_entry::{
    IndexEntryRange, IndexNodeEntryIndex, IndexNodeEntryRanges, IndexNodeSearchResult,
    NtfsIndexEntry, NtfsIndexEntryFlags,
};
use crate::indexes::{NtfsIndexEntryKey, NtfsIndexEntryType};
use crate::io::{Read, Seek};
use crate::structured_values::{NtfsIndexAllocation, NtfsIndexAllocationDataRuns, NtfsIndexRoot};
use crate::types::{NtfsPosition, Vcn};

const FINDER_SUBNODE_CACHE_CAPACITY: usize = 256;

/// Helper structure to iterate over all entries of an index or find a specific one.
///
/// The `E` type parameter of [`NtfsIndexEntryType`] specifies the type of the index entries.
/// The most common one is [`NtfsFileNameIndex`] for file name indexes, commonly known as "directories".
/// Check out [`NtfsFile::directory_index`] to return an [`NtfsIndex`] object for a directory without
/// any hassles.
///
/// [`NtfsFile::directory_index`]: crate::NtfsFile::directory_index
/// [`NtfsFileNameIndex`]: crate::indexes::NtfsFileNameIndex
#[derive(Clone, Debug)]
pub struct NtfsIndex<'n, 'f, E>
where
    E: NtfsIndexEntryType,
{
    index_record_size: u32,
    index_root_entry_ranges: IndexNodeEntryRanges<E>,
    index_root_position: NtfsPosition,
    index_allocation_item: Option<NtfsAttributeItem<'n, 'f>>,
    index_allocation_data_runs: RefCell<Option<NtfsIndexAllocationDataRuns>>,
    entry_type: PhantomData<E>,
}

impl<'n, 'f, E> NtfsIndex<'n, 'f, E>
where
    E: NtfsIndexEntryType,
{
    /// Creates a new [`NtfsIndex`] object from a previously looked up [`NtfsIndexRoot`] attribute
    /// (contained in an [`NtfsAttributeItem`]) and, in case of a large index, a matching
    /// [`NtfsIndexAllocation`] attribute (also contained in an [`NtfsAttributeItem`]).
    ///
    /// If you just want to look up files in a directory, check out [`NtfsFile::directory_index`],
    /// which looks up the correct [`NtfsIndexRoot`] and [`NtfsIndexAllocation`] attributes for you.
    ///
    /// [`NtfsFile::directory_index`]: crate::NtfsFile::directory_index
    pub fn new(
        index_root_item: NtfsAttributeItem<'n, 'f>,
        index_allocation_item: Option<NtfsAttributeItem<'n, 'f>>,
    ) -> Result<Self> {
        let index_root_attribute = index_root_item.to_attribute()?;
        index_root_attribute.ensure_ty(NtfsAttributeType::IndexRoot)?;
        let index_root = index_root_attribute.resident_structured_value::<NtfsIndexRoot>()?;

        if let Some(item) = &index_allocation_item {
            let attribute = item.to_attribute()?;
            attribute.ensure_ty(NtfsAttributeType::IndexAllocation)?;
        }

        Self::from_index_root(index_root, index_allocation_item)
    }

    fn from_index_root(
        index_root: NtfsIndexRoot<'_>,
        index_allocation_item: Option<NtfsAttributeItem<'n, 'f>>,
    ) -> Result<Self> {
        if index_allocation_item.is_none() && index_root.is_large_index() {
            return Err(NtfsError::MissingIndexAllocation {
                position: index_root.position(),
            });
        }

        let index_record_size = index_root.index_record_size();
        let index_root_entry_ranges = index_root.entry_ranges();
        let index_root_position = index_root.position();
        let entry_type = PhantomData;

        Ok(Self {
            index_record_size,
            index_root_entry_ranges,
            index_root_position,
            index_allocation_item,
            index_allocation_data_runs: RefCell::new(None),
            entry_type,
        })
    }

    pub(crate) fn new_with_fs_from_index_root<T>(
        index_root: NtfsIndexRoot<'_>,
        index_allocation_item: Option<NtfsAttributeItem<'n, 'f>>,
        fs: &mut T,
    ) -> Result<Self>
    where
        T: Read + Seek,
    {
        let index = Self::from_index_root(index_root, index_allocation_item)?;
        if let Some(item) = &index.index_allocation_item {
            let attribute = item.to_attribute()?;
            let allocation = attribute.structured_value::<_, NtfsIndexAllocation>(fs)?;
            *index.index_allocation_data_runs.borrow_mut() = Some(allocation.into_data_runs());
        }
        Ok(index)
    }

    /// Returns an [`NtfsIndexEntries`] iterator to perform an in-order traversal of this index.
    pub fn entries<'i>(&'i self) -> NtfsIndexEntries<'n, 'f, 'i, E> {
        NtfsIndexEntries::new(self)
    }

    /// Returns an [`NtfsIndexFinder`] structure to efficiently find an entry in this index.
    pub fn finder<'i>(&'i self) -> NtfsIndexFinder<'n, 'f, 'i, E> {
        NtfsIndexFinder::new(self)
    }

    fn subnode_entry_ranges_with_buffer<T>(
        &self,
        fs: &mut T,
        vcn: Vcn,
        buffer: Vec<u8>,
    ) -> Result<IndexNodeEntryRanges<E>>
    where
        T: Read + Seek,
    {
        let item =
            self.index_allocation_item
                .as_ref()
                .ok_or(NtfsError::MissingIndexAllocation {
                    position: self.index_root_position,
                })?;
        if let Some(data_runs) = self.index_allocation_data_runs.borrow().as_ref() {
            let record =
                data_runs.record_from_vcn_with_buffer(fs, self.index_record_size, vcn, buffer)?;
            return Ok(record.into_entry_ranges());
        }

        let attribute = item.to_attribute()?;
        let allocation = attribute.structured_value::<_, NtfsIndexAllocation>(fs)?;
        let record = allocation.record_from_vcn(fs, self.index_record_size, vcn)?;
        *self.index_allocation_data_runs.borrow_mut() = Some(allocation.data_runs().clone());

        Ok(record.into_entry_ranges())
    }
}

/// Iterator over
///   all index entries of an index,
///   sorted ascending by the index key,
///   returning an [`NtfsIndexEntry`] for each entry.
///
/// This iterator is returned from the [`NtfsIndex::entries`] function.
#[derive(Clone, Debug)]
pub struct NtfsIndexEntries<'n, 'f, 'i, E>
where
    E: NtfsIndexEntryType,
{
    index: &'i NtfsIndex<'n, 'f, E>,
    inner_iterators: Vec<IndexNodeEntryRanges<E>>,
    following_entries: Vec<Option<IndexEntryRange<E>>>,
    buffer: Option<Vec<u8>>,
}

impl<'n, 'f, 'i, E> NtfsIndexEntries<'n, 'f, 'i, E>
where
    E: NtfsIndexEntryType,
{
    fn new(index: &'i NtfsIndex<'n, 'f, E>) -> Self {
        let inner_iterators = vec![index.index_root_entry_ranges.clone()];
        let following_entries = Vec::new();
        let buffer = None;

        Self {
            index,
            inner_iterators,
            following_entries,
            buffer,
        }
    }

    /// See [`Iterator::next`].
    pub fn next<'a, T>(&'a mut self, fs: &mut T) -> Option<Result<NtfsIndexEntry<'a, E>>>
    where
        T: Read + Seek,
    {
        // NTFS B-tree indexes are composed out of nodes, with multiple entries per node.
        // Each entry may have a reference to a subnode.
        // If that is the case, the subnode entries comes before the parent entry lexicographically.
        //
        // An example for an unbalanced, but otherwise valid and sorted tree:
        //
        //                                   -------------
        // INDEX ROOT NODE:                  | 4 | 5 | 6 |
        //                                   -------------
        //                                     |
        //                                 ---------
        // INDEX ALLOCATION SUBNODE:       | 1 | 3 |
        //                                 ---------
        //                                       |
        //                                     -----
        // INDEX ALLOCATION SUBNODE:           | 2 |
        //                                     -----
        //
        let entry_range = loop {
            // Get the iterator from the current node level, if any.
            let iter = self.inner_iterators.last_mut()?;

            // Get the next `IndexEntryRange` from it.
            if let Some(entry_range) = iter.next() {
                let entry_range = iter_try!(entry_range);

                // Convert that `IndexEntryRange` to a (lifetime-bound) `NtfsIndexEntry`.
                let entry = iter_try!(entry_range.to_entry(iter.data()));
                let is_last_entry = entry.flags().contains(NtfsIndexEntryFlags::LAST_ENTRY);

                // Does this entry have a subnode that needs to be iterated first?
                if let Some(subnode_vcn) = entry.subnode_vcn() {
                    let subnode_vcn = iter_try!(subnode_vcn);

                    let buffer = self.buffer.take().unwrap_or_default();
                    let subnode_iter = iter_try!(self.index.subnode_entry_ranges_with_buffer(
                        fs,
                        subnode_vcn,
                        buffer
                    ));

                    let following_entry = if !is_last_entry {
                        // This entry comes after the subnode lexicographically, so save it.
                        // We'll pick it up again after the subnode iterator has been fully iterated.
                        Some(entry_range)
                    } else {
                        None
                    };

                    // Save this subnode's iterator and any following entry.
                    // We'll pick up the iterator through `self.inner_iterators.last_mut()` in the next loop iteration.
                    self.inner_iterators.push(subnode_iter);
                    self.following_entries.push(following_entry);
                } else if !is_last_entry {
                    // There is no subnode, and this is not the empty "last entry",
                    // so our entry comes next lexicographically.
                    break entry_range;
                }
            } else {
                // The iterator for this subnode level has been fully iterated.
                // Drop it.
                let is_root_iterator = self.inner_iterators.len() == 1;
                if let Some(iter) = self.inner_iterators.pop()
                    && !is_root_iterator
                {
                    let data = iter.into_data();
                    if self
                        .buffer
                        .as_ref()
                        .is_none_or(|buffer| buffer.capacity() < data.capacity())
                    {
                        self.buffer = Some(data);
                    }
                }

                // The entry, whose subnode we just fully iterated, may have been saved in `following_entries`.
                // This depends on its `is_last_entry` flag:
                //   * If it was not the last entry, it contains an entry that comes next lexicographically,
                //     and has therefore been saved in `following_entries`.
                //   * If it was the last entry, it contains no further information.
                //     `None` has been saved in `following_entries`, so that `following_entries.len()` always
                //     matches `inner_iterators.len() - 1`.
                //
                // If we just finished iterating the root-level node, `following_entries` is empty and we are done.
                // Otherwise, we can be sure that `inner_iterators.last()` is the matching iterator for converting
                // `IndexEntryRange` to a (lifetime-bound) `NtfsIndexEntry`.
                if let Some(entry_range) = self.following_entries.pop()? {
                    break entry_range;
                }
            }
        };

        let iter = self.inner_iterators.last().unwrap();
        let entry = iter_try!(entry_range.to_entry(iter.data()));

        Some(Ok(entry))
    }
}

/// Helper structure to efficiently find an entry in an index, created by [`NtfsIndex::finder`].
///
/// Up to 256 loaded index subnodes are cached for the lifetime of this finder, so reuse it when
/// looking up multiple entries in the same index. Cached keys are kept sorted by VCN for
/// logarithmic lookup, and the least recently used node is evicted when the cache is full.
///
/// This helper is required because a returned entry may borrow from a subnode buffer cached by the
/// finder. Copy the field(s) you need from the returned entry before reusing or dropping the finder.
pub struct NtfsIndexFinder<'n, 'f, 'i, E>
where
    E: NtfsIndexEntryType,
{
    index: &'i NtfsIndex<'n, 'f, E>,
    root_entry_index: Option<IndexNodeEntryIndex<E>>,
    subnodes: FinderSubnodeCache<E>,
}

struct FinderCachedSubnode<E>
where
    E: NtfsIndexEntryType,
{
    vcn: Vcn,
    entry_ranges: IndexNodeEntryRanges<E>,
    entry_index: IndexNodeEntryIndex<E>,
    last_used: u64,
}

struct FinderSubnodeCache<E>
where
    E: NtfsIndexEntryType,
{
    capacity: usize,
    generation: u64,
    keys: Vec<(Vcn, usize)>,
    nodes: Vec<Option<FinderCachedSubnode<E>>>,
}

impl<E> FinderSubnodeCache<E>
where
    E: NtfsIndexEntryType,
{
    fn new(capacity: usize) -> Self {
        debug_assert!(capacity > 0);
        Self {
            capacity,
            generation: 0,
            keys: Vec::new(),
            nodes: Vec::new(),
        }
    }

    fn find_slot(&self, vcn: Vcn) -> Option<usize> {
        self.keys
            .binary_search_by_key(&vcn, |(cached_vcn, _)| *cached_vcn)
            .ok()
            .map(|key_index| self.keys[key_index].1)
    }

    fn get(&mut self, vcn: Vcn) -> Option<usize> {
        let slot = self.find_slot(vcn)?;
        let generation = self.next_generation();
        self.nodes[slot].as_mut().unwrap().last_used = generation;
        Some(slot)
    }

    fn node(&self, slot: usize) -> &FinderCachedSubnode<E> {
        self.nodes[slot].as_ref().unwrap()
    }

    fn prepare_insert(&mut self) -> (usize, Vec<u8>) {
        if self.keys.len() == self.capacity {
            let slot = self
                .nodes
                .iter()
                .enumerate()
                .filter_map(|(slot, node)| node.as_ref().map(|node| (slot, node.last_used)))
                .min_by_key(|(_, last_used)| *last_used)
                .unwrap()
                .0;
            let node = self.nodes[slot].take().unwrap();
            let key_index = self
                .keys
                .binary_search_by_key(&node.vcn, |(cached_vcn, _)| *cached_vcn)
                .unwrap();
            self.keys.remove(key_index);
            return (slot, node.entry_ranges.into_data());
        }

        if let Some(slot) = self.nodes.iter().position(Option::is_none) {
            return (slot, Vec::new());
        }

        let slot = self.nodes.len();
        self.nodes.push(None);
        (slot, Vec::new())
    }

    fn insert(
        &mut self,
        vcn: Vcn,
        slot: usize,
        entry_ranges: IndexNodeEntryRanges<E>,
        entry_index: IndexNodeEntryIndex<E>,
    ) {
        debug_assert!(self.keys.len() < self.capacity);
        debug_assert!(self.nodes[slot].is_none());
        let key_index = self
            .keys
            .binary_search_by_key(&vcn, |(cached_vcn, _)| *cached_vcn)
            .unwrap_err();
        self.keys.insert(key_index, (vcn, slot));
        let last_used = self.next_generation();
        self.nodes[slot] = Some(FinderCachedSubnode {
            vcn,
            entry_ranges,
            entry_index,
            last_used,
        });
    }

    fn next_generation(&mut self) -> u64 {
        if self.generation == u64::MAX {
            for node in self.nodes.iter_mut().flatten() {
                node.last_used /= 2;
            }
            self.generation /= 2;
        }
        self.generation += 1;
        self.generation
    }
}

impl<'n, 'f, 'i, E> NtfsIndexFinder<'n, 'f, 'i, E>
where
    E: NtfsIndexEntryType,
{
    fn new(index: &'i NtfsIndex<'n, 'f, E>) -> Self {
        let root_entry_index = None;
        let subnodes = FinderSubnodeCache::new(FINDER_SUBNODE_CACHE_CAPACITY);

        Self {
            index,
            root_entry_index,
            subnodes,
        }
    }

    /// Finds an entry in this index using the given comparison function and returns an [`NtfsIndexEntry`]
    /// (if there is one).
    pub fn find<'a, T, F>(&'a mut self, fs: &mut T, cmp: F) -> Option<Result<NtfsIndexEntry<'a, E>>>
    where
        T: Read + Seek,
        F: Fn(&E::KeyType) -> Ordering,
    {
        self.find_by(fs, |slice, position| {
            let key = E::KeyType::key_from_slice(slice, position)?;
            Ok(cmp(&key))
        })
    }

    pub(crate) fn find_by<'a, T, F>(
        &'a mut self,
        fs: &mut T,
        cmp: F,
    ) -> Option<Result<NtfsIndexEntry<'a, E>>>
    where
        T: Read + Seek,
        F: Fn(&[u8], NtfsPosition) -> Result<Ordering>,
    {
        if self.root_entry_index.is_none() {
            self.root_entry_index =
                Some(iter_try!(self.index.index_root_entry_ranges.entry_index()));
        }

        // Start at the Index Root. Subnodes loaded by earlier searches stay cached in this finder.
        let mut subnode_index: Option<usize> = None;

        loop {
            let search_result = if let Some(index) = subnode_index {
                let subnode = self.subnodes.node(index);
                subnode
                    .entry_index
                    .find_by(subnode.entry_ranges.data(), &cmp)
            } else {
                self.root_entry_index
                    .as_ref()
                    .unwrap()
                    .find_by(self.index.index_root_entry_ranges.data(), &cmp)
            };

            match iter_try!(search_result?) {
                IndexNodeSearchResult::Found(entry_range) => {
                    let data = if let Some(index) = subnode_index {
                        self.subnodes.node(index).entry_ranges.data()
                    } else {
                        self.index.index_root_entry_ranges.data()
                    };
                    return Some(entry_range.to_entry(data));
                }
                IndexNodeSearchResult::Subnode(vcn) => {
                    if let Some(slot) = self.subnodes.get(vcn) {
                        subnode_index = Some(slot);
                    } else {
                        let (slot, buffer) = self.subnodes.prepare_insert();
                        let subnode =
                            iter_try!(self.index.subnode_entry_ranges_with_buffer(fs, vcn, buffer));
                        let entry_index = iter_try!(subnode.entry_index());
                        self.subnodes.insert(vcn, slot, subnode, entry_index);
                        subnode_index = Some(slot);
                    }
                }
            }
        }
    }
}

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

    fn empty_subnode(
        data: Vec<u8>,
    ) -> (
        IndexNodeEntryRanges<NtfsFileNameIndex>,
        IndexNodeEntryIndex<NtfsFileNameIndex>,
    ) {
        let entry_ranges = IndexNodeEntryRanges::new(data, 0..0, NtfsPosition::none());
        let entry_index = entry_ranges.entry_index().unwrap();
        (entry_ranges, entry_index)
    }

    #[test]
    fn finder_cache_evicts_lru_node_and_reuses_its_buffer() {
        let mut cache = FinderSubnodeCache::new(3);
        for vcn in 0..3 {
            let (entry_ranges, entry_index) = empty_subnode(Vec::with_capacity(64 + vcn as usize));
            let (slot, buffer) = cache.prepare_insert();
            assert!(buffer.is_empty());
            cache.insert(Vcn::from(vcn), slot, entry_ranges, entry_index);
        }

        // VCN 0 represents a hot root-adjacent node. Touching it makes VCN 1 the LRU node,
        // even though it is neither end of the VCN-sorted cache after the next insertion.
        assert!(cache.get(Vcn::from(0)).is_some());
        let victim_slot = cache.find_slot(Vcn::from(1)).unwrap();
        let victim_pointer = cache.node(victim_slot).entry_ranges.data().as_ptr();
        let (slot, buffer) = cache.prepare_insert();

        assert_eq!(slot, victim_slot);
        assert_eq!(buffer.as_ptr(), victim_pointer);
        assert!(cache.find_slot(Vcn::from(0)).is_some());
        assert!(cache.find_slot(Vcn::from(1)).is_none());

        let (entry_ranges, entry_index) = empty_subnode(buffer);
        cache.insert(Vcn::from(3), slot, entry_ranges, entry_index);
        assert_eq!(
            cache
                .node(cache.find_slot(Vcn::from(3)).unwrap())
                .entry_ranges
                .data()
                .as_ptr(),
            victim_pointer
        );
    }
}