hegeltest 0.19.1

// Native backend DataSource implementation.
//
// Wraps a NativeTestCase behind interior mutability so it can implement
// the DataSource trait (which takes &self).  The handle is shared with
// the engine so it can read back the recorded nodes / spans and, via
// `mark_complete`, the test case outcome after the test body returns.

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};

use ciborium::Value;

use crate::backend::{DataSource, DataSourceError, TestCaseResult};
use crate::native::bignum::{BigInt, ToPrimitive};
use crate::native::core::{ChoiceNode, EngineError, ManyState, NativeTestCase, Span, Status};
use crate::native::schema;
use crate::test_case::invalid_argument;

/// Per-test-case state shared between `NativeDataSource` and the engine
/// that owns the handle.  The engine constructs both halves up-front,
/// hands the data source into the test body, and then reads back nodes,
/// spans, and the outcome (populated by `mark_complete`) through the
/// handle.
pub struct NativeTestCaseInner {
    pub ntc: NativeTestCase,
    /// Set by [`DataSource::mark_complete`] after the test body returns.
    /// `None` only if `mark_complete` was never called — which the lifecycle
    /// in `run_lifecycle::run_test_case` guarantees won't happen — so the
    /// engine can safely unwrap when reading the outcome back.
    pub outcome: Option<TestCaseResult>,
    /// `tc.target()` observations recorded during the test case, keyed by
    /// label. Populated by [`DataSource::target_observation`]; read back by
    /// the targeting phase via [`NativeDataSource::take_target_observations`].
    pub target_observations: HashMap<String, f64>,
}

/// Shared handle to the per-test-case inner state.
pub type NativeTestCaseHandle = Arc<Mutex<NativeTestCaseInner>>;

pub struct NativeDataSource {
    inner: NativeTestCaseHandle,
    aborted: AtomicBool,
}

impl NativeDataSource {
    /// Create a new `NativeDataSource` and return a shared handle.
    ///
    /// The handle is the only way the engine reads back per-test-case
    /// state: choice nodes, spans, and the outcome reported by
    /// [`DataSource::mark_complete`].
    pub fn new(ntc: NativeTestCase) -> (Self, NativeTestCaseHandle) {
        let inner = Arc::new(Mutex::new(NativeTestCaseInner {
            ntc,
            outcome: None,
            target_observations: HashMap::new(),
        }));
        let handle = Arc::clone(&inner);
        (
            NativeDataSource {
                inner,
                aborted: AtomicBool::new(false),
            },
            handle,
        )
    }

    /// Convenience: extract choice nodes from a handle after a test case.
    pub fn take_nodes(handle: &NativeTestCaseHandle) -> Vec<ChoiceNode> {
        handle
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .ntc
            .nodes
            .clone()
    }

    /// Convenience: extract spans from a handle after a test case.
    pub fn take_spans(handle: &NativeTestCaseHandle) -> Vec<Span> {
        handle
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .ntc
            .spans
            .clone()
            .into_vec()
    }

    /// Drain the `tc.target()` observations the test body recorded.
    ///
    /// Used by the targeting phase in `test_runner` to read back per-label
    /// scores after a test case completes.
    pub fn take_target_observations(handle: &NativeTestCaseHandle) -> HashMap<String, f64> {
        std::mem::take(
            &mut handle
                .lock()
                .unwrap_or_else(|e| e.into_inner())
                .target_observations,
        )
    }

    /// Read the outcome reported via [`DataSource::mark_complete`].
    ///
    /// Panics if `mark_complete` was never called; the cross-backend
    /// lifecycle in `run_lifecycle::run_test_case` guarantees it always is.
    pub fn take_outcome(handle: &NativeTestCaseHandle) -> TestCaseResult {
        handle
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .outcome
            .take()
            .expect("mark_complete must be called for every test case")
    }

    /// Returns true if a previous request triggered a EngineError abort.
    /// Test-only helper — not part of the `DataSource` interface, so
    /// callers must hold a concrete `&NativeDataSource`.
    #[cfg(test)]
    pub(crate) fn test_aborted(&self) -> bool {
        self.aborted.load(Ordering::Relaxed)
    }

    /// Acquire the test-case state under the abort guard.  Returns
    /// `DataSourceError::StopTest` immediately if a previous call has already
    /// aborted the test case so subsequent draws short-circuit without
    /// touching `ntc`.
    fn with_ntc<R>(
        &self,
        f: impl FnOnce(&mut NativeTestCase) -> Result<R, EngineError>,
    ) -> Result<R, DataSourceError> {
        if self.aborted.load(Ordering::Relaxed) {
            return Err(self.aborted_error());
        }
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        f(&mut inner.ntc).map_err(|e| match e {
            EngineError::Overrun => {
                self.aborted.store(true, Ordering::Relaxed);
                DataSourceError::StopTest
            }
            EngineError::InvalidTestCase => {
                self.aborted.store(true, Ordering::Relaxed);
                DataSourceError::Assume
            }
            EngineError::InvalidArgument(msg) => DataSourceError::InvalidArgument(msg),
        })
    }

    fn aborted_error(&self) -> DataSourceError {
        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        match inner.ntc.status {
            Some(Status::Invalid) => DataSourceError::Assume,
            _ => DataSourceError::StopTest,
        }
    }
}

impl DataSource for NativeDataSource {
    fn generate(&self, schema: &Value) -> Result<Value, DataSourceError> {
        self.with_ntc(|ntc| schema::interpret_schema(ntc, schema))
    }

    fn start_span(&self, label: u64) -> Result<(), DataSourceError> {
        self.with_ntc(|ntc| {
            ntc.start_span(label);
            Ok(())
        })
    }

    fn stop_span(&self, discard: bool) -> Result<(), DataSourceError> {
        self.with_ntc(|ntc| {
            ntc.stop_span(discard);
            Ok(())
        })
    }

    fn new_collection(&self, min_size: u64, max_size: Option<u64>) -> Result<i64, DataSourceError> {
        self.with_ntc(|ntc| {
            let state = ManyState::new(min_size as usize, max_size.map(|n| n as usize));
            Ok(ntc.new_collection(state))
        })
    }

    fn collection_more(&self, collection_id: i64) -> Result<bool, DataSourceError> {
        self.with_ntc(|ntc| {
            let mut state = ntc
                .collections
                .remove(&collection_id)
                .expect("collection_more: unknown collection_id");
            let result = schema::many_more(ntc, &mut state);
            ntc.collections.insert(collection_id, state);
            result
        })
    }

    fn collection_reject(
        &self,
        collection_id: i64,
        _why: Option<&str>,
    ) -> Result<(), DataSourceError> {
        self.with_ntc(|ntc| {
            let mut state = ntc
                .collections
                .remove(&collection_id)
                .expect("collection_reject: unknown collection_id");
            let result = schema::many_reject(ntc, &mut state);
            ntc.collections.insert(collection_id, state);
            result
        })
    }

    fn new_state_machine(
        &self,
        rule_names: &[&str],
        invariant_names: &[&str],
    ) -> Result<i64, DataSourceError> {
        if rule_names.is_empty() {
            return Err(DataSourceError::InvalidArgument(
                "cannot run a state machine with no rules".to_string(),
            ));
        }
        let rules = rule_names.iter().map(|s| s.to_string()).collect();
        let invariants = invariant_names.iter().map(|s| s.to_string()).collect();
        self.with_ntc(|ntc| {
            let id = ntc.state_machines.len() as i64;
            ntc.state_machines
                .push(crate::native::core::NativeStateMachine::new(
                    rules, invariants,
                ));
            Ok(id)
        })
    }

    fn state_machine_next_rule(&self, state_machine_id: i64) -> Result<i64, DataSourceError> {
        self.with_ntc(|ntc| {
            let idx = state_machine_id as usize;
            // Move the machine out of `ntc` so `next_rule` can borrow `ntc`
            // mutably (same idea as the remove/insert dance in
            // `collection_more`).
            let mut machine = std::mem::take(&mut ntc.state_machines[idx]);
            let result = machine.next_rule(ntc);
            ntc.state_machines[idx] = machine;
            result
        })
    }

    fn primitive_boolean(&self, p: f64, forced: Option<bool>) -> Result<bool, DataSourceError> {
        self.with_ntc(|ntc| {
            if !(0.0..=1.0).contains(&p) {
                return Err(EngineError::InvalidArgument(format!(
                    "primitive_boolean(p = {p}) requires a probability in [0.0, 1.0]"
                )));
            }
            if forced == Some(true) && p == 0.0 {
                return Err(EngineError::InvalidArgument(
                    "primitive_boolean: cannot force true when p = 0.0".to_string(),
                ));
            }
            if forced == Some(false) && p == 1.0 {
                return Err(EngineError::InvalidArgument(
                    "primitive_boolean: cannot force false when p = 1.0".to_string(),
                ));
            }
            ntc.weighted(p, forced)
        })
    }

    fn new_pool(&self) -> Result<i64, DataSourceError> {
        self.with_ntc(|ntc| {
            let pool_id = ntc.variable_pools.len() as i64;
            ntc.variable_pools
                .push(crate::native::core::NativeVariables::new());
            Ok(pool_id)
        })
    }

    fn pool_add(&self, pool_id: i64) -> Result<i64, DataSourceError> {
        self.with_ntc(|ntc| Ok(ntc.variable_pools[pool_id as usize].next()))
    }

    fn pool_generate(&self, pool_id: i64, consume: bool) -> Result<i64, DataSourceError> {
        self.with_ntc(|ntc| {
            let pool_idx = pool_id as usize;
            let active = ntc.variable_pools[pool_idx].active();
            if active.is_empty() {
                ntc.status = Some(Status::Invalid);
                return Err(EngineError::InvalidTestCase);
            }
            // The variable ids drawn out of `active` are `i64`.
            let n = active.len();
            // Draw index from `[0, n-1]`.  Shrink towards `n-1`
            // (last added = most recent) by drawing `k` from
            // `[0, n-1]` and using `index = n-1-k`.
            let k = ntc
                .draw_integer(BigInt::from(0), BigInt::from(n as i64 - 1))?
                .to_i128()
                .unwrap() as usize;
            let variable_id = active[n - 1 - k];
            if consume {
                ntc.variable_pools[pool_idx].consume(variable_id);
            }
            Ok(variable_id)
        })
    }

    fn target_observation(&self, score: f64, label: &str) {
        // Mirror upstream `hypothesis.control.target`
        // (`control.py:354-356,372-376`): the
        // observation must be finite and each label may be observed at
        // most once per test case. These are usage errors, not discovered
        // counterexamples, so raise them via `invalid_argument!` for a clean
        // abort instead of letting the lifecycle shrink them as a "failure".
        if !score.is_finite() {
            invalid_argument!(
                "tc.target({score}, label={label:?}) requires a finite score; \
                 got non-finite value"
            );
        }
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        if inner.target_observations.contains_key(label) {
            invalid_argument!(
                "tc.target({score}, label={label:?}) would overwrite previous \
                 tc.target(_, label={label:?}); each label can be observed at \
                 most once per test case"
            );
        }
        inner.target_observations.insert(label.to_string(), score);
    }

    fn mark_complete(&self, result: &TestCaseResult) {
        // Record the outcome on the shared handle so the engine can read
        // it via `take_outcome` after the test body returns.  This is the
        // channel for per-test-case results: the engine consumes
        // `mark_complete` through the `DataSource` interface.
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner.outcome = Some(result.clone());
    }
}

#[cfg(test)]
#[path = "../../tests/embedded/native/data_source_tests.rs"]
mod tests;