solverforge-scoring 0.19.12

use super::{ExplainRow, Operator, RowChanges};
use crate::api::analysis::EntityRef;
use crate::stream::relational::{DenseRowStore, HandleMap, Leaf, RowHandle};

/// A borrowed child with its complete owning row and child position.
pub struct FlattenView<'a, V, T> {
    pub input: V,
    pub value: &'a T,
    pub child: usize,
}
impl<V: Copy, T> Copy for FlattenView<'_, V, T> {}
impl<V: Copy, T> Clone for FlattenView<'_, V, T> {
    fn clone(&self) -> Self {
        *self
    }
}
impl<V: ExplainRow, T> ExplainRow for FlattenView<'_, V, T> {
    fn explain(&self, entities: &mut Vec<EntityRef>) {
        self.input.explain(entities);
    }
}

/* Named source of borrowed children for one input operator.

`FlattenNode` borrows children out of the solution and ties the result to the
solution borrow. Spelled as a raw `for<'a> Fn(&'a S, O::View<'a>) -> &'a [T]`
that bound cannot appear in a struct field or a return type, so a stream that
owns a flattened operator could not name its own type. `FlattenSource<S, O>`
is that bound as a name: `O` carries the borrowed view, so both an ordinary
closure (blanket impl) and a concrete adapter (`ParentFlatten`) satisfy it.
*/
pub trait FlattenSource<S: 'static, O: Operator<S>> {
    type Item: 'static;

    fn flatten<'a>(&self, solution: &'a S, view: O::View<'a>) -> &'a [Self::Item];
}

/* Ordinary closures keep the current shape. */
impl<S: 'static, O: Operator<S>, F, T: 'static> FlattenSource<S, O> for F
where
    F: for<'a> Fn(&'a S, O::View<'a>) -> &'a [T],
{
    type Item = T;

    #[inline]
    fn flatten<'a>(&self, solution: &'a S, view: O::View<'a>) -> &'a [T] {
        (self)(solution, view)
    }
}

/* A nameable flattened extractor: any `FlattenExtract<P>` lifted onto a
collection operator over `P`. */
pub struct ParentFlatten<Inner> {
    inner: Inner,
}

impl<Inner> ParentFlatten<Inner> {
    pub fn new(inner: Inner) -> Self {
        Self { inner }
    }
}

impl<S, E, Inner, P, B> FlattenSource<S, super::CollectionNode<S, E>> for ParentFlatten<Inner>
where
    S: 'static,
    E: crate::stream::collection_extract::CollectionExtract<S, Item = P> + 'static,
    Inner: crate::stream::collection_extract::FlattenExtract<P, Item = B> + Send + Sync,
    P: 'static,
    B: 'static,
{
    type Item = B;

    #[inline]
    fn flatten<'a>(&self, _solution: &'a S, view: Leaf<'a, P>) -> &'a [B] {
        self.inner.extract(view.entity)
    }
}

struct Child {
    input: RowHandle,
    child: usize,
}
/// Flatten borrowed child slices from arbitrary input rows; retention owns only identities.
pub struct FlattenNode<O, F> {
    input: O,
    extractor: F,
    rows: DenseRowStore<Child>,
    outputs: HandleMap<Vec<RowHandle>>,
    /// Child-handle vectors retired by a retraction, kept for the next insert.
    ///
    /// A variable change retracts and re-inserts its owning row, so the same
    /// parent is rebuilt with a fresh handle on every event. Handing the retired
    /// buffer to that rebuild keeps the child handles in one allocation instead
    /// of re-growing a vector per event, which is the dominant heap traffic of a
    /// large local-search step.
    retired: Vec<Vec<RowHandle>>,
}
impl<O, F> FlattenNode<O, F> {
    pub fn new(input: O, extractor: F) -> Self {
        Self {
            input,
            extractor,
            rows: DenseRowStore::new(),
            outputs: HandleMap::new(),
            retired: Vec::new(),
        }
    }
    fn apply_changes<S: 'static>(&mut self, solution: &S, changes: RowChanges) -> RowChanges
    where
        O: Operator<S>,
        F: FlattenSource<S, O>,
    {
        let mut removed = Vec::new();
        for input in changes.removed {
            let Some(mut handles) = self.outputs.remove(input) else {
                continue;
            };
            removed.extend_from_slice(&handles);
            for handle in handles.drain(..) {
                self.rows.retract(handle);
            }
            handles.clear();
            self.retired.push(handles);
        }
        let mut inserted = Vec::new();
        for input in changes.inserted {
            let row = self
                .input
                .resolve(solution, input)
                .expect("inserted flattened input");
            let count = self.extractor.flatten(solution, row).len();
            let mut outputs = match self.outputs.remove(input) {
                Some(mut handles) => {
                    handles.clear();
                    handles
                }
                None => self.retired.pop().unwrap_or_default(),
            };
            outputs.reserve(count);
            inserted.reserve(count);
            for child in 0..count {
                let h = self.rows.insert(Child { input, child });
                outputs.push(h);
                inserted.push(h);
            }
            self.outputs.insert(input, outputs);
        }
        RowChanges { removed, inserted }
    }
}
impl<S: 'static, O: Operator<S>, F: FlattenSource<S, O> + 'static> Operator<S>
    for FlattenNode<O, F>
{
    type View<'a> = FlattenView<'a, O::View<'a>, F::Item>;
    type Evaluation = O::Evaluation;
    fn prepare_evaluation(&self, solution: &S) -> Self::Evaluation {
        self.input.prepare_evaluation(solution)
    }
    fn visit_evaluation<'a>(
        &'a self,
        solution: &'a S,
        evaluation: &'a Self::Evaluation,
        visitor: &mut impl FnMut(Self::View<'a>),
    ) {
        self.input
            .visit_evaluation(solution, evaluation, &mut |input| {
                for (child, value) in self.extractor.flatten(solution, input).iter().enumerate() {
                    visitor(FlattenView {
                        input,
                        value,
                        child,
                    });
                }
            });
    }
    fn clear(&mut self) {
        self.input.clear();
        self.rows.clear();
        self.outputs.clear();
        self.retired.clear();
    }
    fn initialize(&mut self, solution: &S) {
        self.clear();
        self.input.initialize(solution);
        self.apply_changes(
            solution,
            RowChanges {
                removed: Vec::new(),
                inserted: self.input.handles(),
            },
        );
    }
    fn handles(&self) -> Vec<RowHandle> {
        self.rows.iter().map(|(h, _)| h).collect()
    }
    fn resolve<'a>(&'a self, solution: &'a S, h: RowHandle) -> Option<Self::View<'a>> {
        let child = self.rows.get(h)?;
        let input = self.input.resolve(solution, child.input)?;
        let value = self.extractor.flatten(solution, input).get(child.child)?;
        Some(FlattenView {
            input,
            value,
            child: child.child,
        })
    }
    fn visit_provenance(&self, h: RowHandle, visitor: &mut impl FnMut(u32, usize, usize)) {
        if let Some(child) = self.rows.get(h) {
            self.input.visit_provenance(child.input, visitor);
        }
    }
    fn retract(&mut self, solution: &S, descriptor: usize, index: usize) -> RowChanges {
        let changes = self.input.retract(solution, descriptor, index);
        self.apply_changes(solution, changes)
    }
    fn insert(&mut self, solution: &S, descriptor: usize, index: usize) -> RowChanges {
        let changes = self.input.insert(solution, descriptor, index);
        self.apply_changes(solution, changes)
    }
}