kiddo 6.0.0

use crate::kd_tree::ConstructionError;
use crate::traits::leaf_strategy::ConstructibleLeafStrategy;
use crate::{Axis, Content, KdTree, StemStrategy};

use super::construction::{validate_auto_generated_items, DefaultConstruction, SerialConstruction};
#[cfg(feature = "multi-threaded")]
use super::construction::{ParallelConstruction, DEFAULT_PARALLEL_CONSTRUCTION_THRESHOLD};

/// Configures how a [`KdTree`] is constructed.
///
/// With the default `multi-threaded` feature the builder starts from parallel
/// construction at or above `DEFAULT_PARALLEL_CONSTRUCTION_THRESHOLD` points,
/// using the current Rayon thread pool. Callers can use
/// `rayon::ThreadPool::install` to control its thread count, or
/// [`KdTreeBuilder::with_serial_construction`] to opt out.
///
/// Without that feature the parallel policy does not exist: the builder starts
/// from [`SerialConstruction`], and `with_parallel_construction` and its
/// threshold variant are not compiled. See [`DefaultConstruction`].
///
/// # Examples
///
/// ```rust
/// # #[cfg(feature = "multi-threaded")] {
/// use kiddo::leaf_strategy::FlatVec;
/// use kiddo::{Eytzinger, KdTree};
///
/// type Tree = KdTree<f64, u32, Eytzinger, FlatVec<f64, u32, 3, 32>, 3, 32>;
///
/// let points = vec![[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]];
/// let tree = Tree::builder()
///     .with_parallel_construction_threshold(1_000)
///     .build_from_slice(&points)
///     .unwrap();
///
/// assert_eq!(tree.size(), 2);
/// # }
/// ```
#[must_use = "a construction builder does nothing until a build method is called"]
pub struct KdTreeBuilder<A, T, SS, LS, const K: usize, const B: usize, P = DefaultConstruction> {
    // Only the parallel build path reads the policy; the serial one is a ZST
    // marker, so without `multi-threaded` this field is never loaded.
    #[cfg_attr(not(feature = "multi-threaded"), allow(dead_code))]
    pub(in crate::kd_tree) policy: P,
    pub(in crate::kd_tree) _phantom: std::marker::PhantomData<(A, T, SS, LS)>,
}

impl<A, T, SS, LS, const K: usize, const B: usize, P> KdTreeBuilder<A, T, SS, LS, K, B, P> {
    /// Forces serial construction.
    pub fn with_serial_construction(self) -> KdTreeBuilder<A, T, SS, LS, K, B, SerialConstruction> {
        KdTreeBuilder {
            policy: SerialConstruction,
            _phantom: std::marker::PhantomData,
        }
    }

    /// Forces parallel construction for soft-bucket leaf strategies.
    ///
    /// Hard-bucket strategies currently retain their serial construction
    /// algorithm.
    ///
    /// Requires the `multi-threaded` feature.
    #[cfg(feature = "multi-threaded")]
    pub fn with_parallel_construction(
        self,
    ) -> KdTreeBuilder<A, T, SS, LS, K, B, ParallelConstruction> {
        KdTreeBuilder {
            policy: ParallelConstruction::with_threshold(1),
            _phantom: std::marker::PhantomData,
        }
    }

    /// Uses parallel construction at or above `item_count`, and serial
    /// construction below it.
    ///
    /// The threshold also controls recursive Rayon join granularity.
    /// A threshold of zero is treated as forced parallel construction.
    ///
    /// Requires the `multi-threaded` feature.
    #[cfg(feature = "multi-threaded")]
    pub fn with_parallel_construction_threshold(
        self,
        item_count: usize,
    ) -> KdTreeBuilder<A, T, SS, LS, K, B, ParallelConstruction> {
        KdTreeBuilder {
            policy: ParallelConstruction::with_threshold(item_count),
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<A, T, SS, LS, const K: usize, const B: usize> Default
    for KdTreeBuilder<A, T, SS, LS, K, B, SerialConstruction>
{
    fn default() -> Self {
        Self {
            policy: SerialConstruction,
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<A, T, SS, LS, const K: usize, const B: usize>
    KdTreeBuilder<A, T, SS, LS, K, B, SerialConstruction>
where
    A: Axis<Coord = A>,
    T: Content,
    SS: StemStrategy,
    LS: ConstructibleLeafStrategy<A, T, SS, K, B>,
{
    /// Builds a tree from points, using source indices as items.
    pub fn build_from_slice(
        self,
        source: &[[A; K]],
    ) -> Result<KdTree<A, T, SS, LS, K, B>, ConstructionError>
    where
        T: TryFrom<usize>,
    {
        validate_auto_generated_items::<T>(source.len())?;
        KdTree::new_from_source_with(
            source,
            |point: &[A; K], dim| point[dim],
            |src_idx: usize, _point: &[A; K]| {
                T::try_from(src_idx).map_err(|_| {
                    ConstructionError::AutoGeneratedItemIndexOverflow {
                        item_count: source.len(),
                        item_type: core::any::type_name::<T>(),
                    }
                })
            },
        )
    }

    /// Builds a tree from a generic source and coordinate/item accessors.
    pub fn build_from_source<X, FA, FI>(
        self,
        source: &[X],
        axis_at: FA,
        item_at: FI,
    ) -> Result<KdTree<A, T, SS, LS, K, B>, ConstructionError>
    where
        FA: Fn(&X, usize) -> A,
        FI: Fn(usize, &X) -> T,
    {
        KdTree::new_from_source_with(source, axis_at, |src_idx, src| Ok(item_at(src_idx, src)))
    }

    /// Builds a tree from explicit item/point pairs.
    pub fn build_from_entries(
        self,
        source: &[(T, [A; K])],
    ) -> Result<KdTree<A, T, SS, LS, K, B>, ConstructionError> {
        self.build_from_source(
            source,
            |entry: &(T, [A; K]), dim| entry.1[dim],
            |_src_idx, entry: &(T, [A; K])| entry.0,
        )
    }
}

impl<A, SS, LS, const K: usize, const B: usize>
    KdTreeBuilder<A, (), SS, LS, K, B, SerialConstruction>
where
    A: Axis<Coord = A>,
    SS: StemStrategy,
    LS: ConstructibleLeafStrategy<A, (), SS, K, B>,
{
    /// Builds a tree with no stored items.
    pub fn build_from_slice_no_items(
        self,
        source: &[[A; K]],
    ) -> Result<KdTree<A, (), SS, LS, K, B>, ConstructionError> {
        self.build_from_source(
            source,
            |point: &[A; K], dim| point[dim],
            |_src_idx, _point| (),
        )
    }
}

#[cfg(feature = "multi-threaded")]
impl<A, T, SS, LS, const K: usize, const B: usize> Default
    for KdTreeBuilder<A, T, SS, LS, K, B, ParallelConstruction>
{
    fn default() -> Self {
        Self {
            policy: ParallelConstruction::with_threshold(DEFAULT_PARALLEL_CONSTRUCTION_THRESHOLD),
            _phantom: std::marker::PhantomData,
        }
    }
}

#[cfg(feature = "multi-threaded")]
impl<A, T, SS, LS, const K: usize, const B: usize>
    KdTreeBuilder<A, T, SS, LS, K, B, ParallelConstruction>
where
    A: Axis<Coord = A> + Send + Sync,
    T: Content,
    SS: StemStrategy,
    LS: ConstructibleLeafStrategy<A, T, SS, K, B>,
{
    /// Builds a tree from points, using source indices as items.
    pub fn build_from_slice(
        self,
        source: &[[A; K]],
    ) -> Result<KdTree<A, T, SS, LS, K, B>, ConstructionError>
    where
        T: TryFrom<usize>,
    {
        validate_auto_generated_items::<T>(source.len())?;
        KdTree::new_from_source_with_parallel_policy(
            source,
            |point: &[A; K], dim| point[dim],
            |src_idx: usize, _point: &[A; K]| {
                T::try_from(src_idx).map_err(|_| {
                    ConstructionError::AutoGeneratedItemIndexOverflow {
                        item_count: source.len(),
                        item_type: core::any::type_name::<T>(),
                    }
                })
            },
            self.policy,
        )
    }

    /// Builds a tree from a generic source and coordinate/item accessors.
    pub fn build_from_source<X, FA, FI>(
        self,
        source: &[X],
        axis_at: FA,
        item_at: FI,
    ) -> Result<KdTree<A, T, SS, LS, K, B>, ConstructionError>
    where
        X: Sync,
        FA: Fn(&X, usize) -> A + Sync,
        FI: Fn(usize, &X) -> T,
    {
        KdTree::new_from_source_with_parallel_policy(
            source,
            axis_at,
            |src_idx, src| Ok(item_at(src_idx, src)),
            self.policy,
        )
    }

    /// Builds a tree from explicit item/point pairs.
    pub fn build_from_entries(
        self,
        source: &[(T, [A; K])],
    ) -> Result<KdTree<A, T, SS, LS, K, B>, ConstructionError>
    where
        T: Sync,
    {
        self.build_from_source(
            source,
            |entry: &(T, [A; K]), dim| entry.1[dim],
            |_src_idx, entry: &(T, [A; K])| entry.0,
        )
    }
}

#[cfg(feature = "multi-threaded")]
impl<A, SS, LS, const K: usize, const B: usize>
    KdTreeBuilder<A, (), SS, LS, K, B, ParallelConstruction>
where
    A: Axis<Coord = A> + Send + Sync,
    SS: StemStrategy,
    LS: ConstructibleLeafStrategy<A, (), SS, K, B>,
{
    /// Builds a tree with no stored items.
    pub fn build_from_slice_no_items(
        self,
        source: &[[A; K]],
    ) -> Result<KdTree<A, (), SS, LS, K, B>, ConstructionError> {
        self.build_from_source(
            source,
            |point: &[A; K], dim| point[dim],
            |_src_idx, _point| (),
        )
    }
}