byteor_pipeline_kernel/

step.rs

//! Step primitives.

use crate::lane::{
    LaneRx, LaneRxBorrow, LaneRxSlot, LaneSlot, LaneTx, LaneTxSlot, LaneTxSlotError, LaneTxWith,
    LaneTxWithResult,
};
use crate::lane::{LaneTxError, LaneTxWithError};

use byteor_pipeline_spec::OnFullV1;
use indexbus_core::WaitStrategy;

/// Outcome of a stage transform.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum TransformOutcome {
    /// Drop the message.
    Drop,
    /// Forward the original input.
    ForwardInput,
    /// Forward the first N bytes of the output buffer.
    ForwardOutput(usize),
}

/// Receive → transform → publish primitive.
pub fn recv_transform_publish_step<Rx, Tx, W, F>(
    rx: &mut Rx,
    tx: &mut Tx,
    input: &mut [u8],
    output: &mut [u8],
    on_full: OnFullV1,
    wait: &mut W,
    mut transform: F,
) -> core::result::Result<bool, LaneTxError>
where
    Rx: LaneRx,
    Tx: LaneTx,
    W: WaitStrategy,
    F: FnMut(&[u8], &mut [u8]) -> TransformOutcome,
{
    let Some(n) = rx.recv(input) else {
        return Ok(false);
    };
    let input = &input[..n];
    match transform(input, output) {
        TransformOutcome::Drop => {}
        TransformOutcome::ForwardInput => {
            publish_with_policy(tx, input, on_full, wait)?;
        }
        TransformOutcome::ForwardOutput(out_n) => {
            let out_n = out_n.min(output.len());
            publish_with_policy(tx, &output[..out_n], on_full, wait)?;
        }
    }

    Ok(true)
}

/// Receive-with (borrowed) → transform → publish primitive.
///
/// This is the same logical primitive as [`recv_transform_publish_step`], but avoids copying the
/// input bytes into an intermediate buffer.
pub fn recv_with_transform_publish_step<Rx, Tx, F>(
    rx: &mut Rx,
    tx: &mut Tx,
    output: &mut [u8],
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
    mut transform: F,
) -> core::result::Result<bool, LaneTxError>
where
    Rx: LaneRxBorrow,
    Tx: LaneTx,
    F: FnMut(&[u8], &mut [u8]) -> TransformOutcome,
{
    let res = rx
        .recv_with(|input| match transform(input, output) {
            TransformOutcome::Drop => Ok(()),
            TransformOutcome::ForwardInput => publish_with_policy(tx, input, on_full, wait),
            TransformOutcome::ForwardOutput(out_n) => {
                let out_n = out_n.min(output.len());
                publish_with_policy(tx, &output[..out_n], on_full, wait)
            }
        })
        .map_err(|_| LaneTxError::Failed)?;

    match res {
        None => Ok(false),
        Some(r) => {
            r?;
            Ok(true)
        }
    }
}

/// Receive-with (borrowed) → transform → publish, up to `batch` messages.
///
/// Returns `Ok(true)` if at least one message was processed.
pub fn recv_with_transform_publish_batch<Rx, Tx, F>(
    rx: &mut Rx,
    tx: &mut Tx,
    output: &mut [u8],
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
    batch: usize,
    mut transform: F,
) -> core::result::Result<bool, LaneTxError>
where
    Rx: LaneRxBorrow,
    Tx: LaneTx,
    F: FnMut(&[u8], &mut [u8]) -> TransformOutcome,
{
    let batch = batch.max(1);
    let mut did_any = false;

    for _ in 0..batch {
        let did = recv_with_transform_publish_step(rx, tx, output, on_full, wait, &mut transform)?;
        if !did {
            break;
        }
        did_any = true;
    }

    Ok(did_any)
}

/// Receive → (MapOk) transform → publish-with (slot write).
///
/// This avoids an intermediate output buffer copy by writing directly into the lane slot.
pub fn recv_map_ok_publish_with_step<Rx, Tx>(
    rx: &mut Rx,
    tx: &mut Tx,
    input: &mut [u8],
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
    map_ok: fn(&[u8], &mut [u8]) -> usize,
) -> core::result::Result<bool, LaneTxError>
where
    Rx: LaneRx,
    Tx: LaneTxWith,
{
    let Some(n) = rx.recv(input) else {
        return Ok(false);
    };
    let input = &input[..n];
    publish_with_policy_with(tx, on_full, wait, |slot| map_ok(input, slot))?;
    Ok(true)
}

/// Receive-with (borrowed) → (MapOk) transform → publish-with (slot write).
///
/// This avoids both the input copy (borrowed receive) and the output copy (slot write).
pub fn recv_with_map_ok_publish_with_step<Rx, Tx>(
    rx: &mut Rx,
    tx: &mut Tx,
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
    map_ok: fn(&[u8], &mut [u8]) -> usize,
) -> core::result::Result<bool, LaneTxError>
where
    Rx: LaneRxBorrow,
    Tx: LaneTxWith,
{
    let res = rx
        .recv_with(|input| publish_with_policy_with(tx, on_full, wait, |slot| map_ok(input, slot)))
        .map_err(|_| LaneTxError::Failed)?;
    match res {
        None => Ok(false),
        Some(r) => {
            r?;
            Ok(true)
        }
    }
}

/// Receive-with (borrowed) → (Map) transform → publish-with (slot write), fallible.
///
/// This avoids both the input copy (borrowed receive) and the output copy (slot write), while
/// allowing the stage to return an error without publishing.
pub fn recv_with_map_publish_with_result_step<Rx, Tx, E>(
    rx: &mut Rx,
    tx: &mut Tx,
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
    map: impl FnMut(&[u8], &mut [u8]) -> core::result::Result<usize, E>,
) -> core::result::Result<bool, LaneTxWithError<E>>
where
    Rx: LaneRxBorrow,
    Tx: LaneTxWithResult,
{
    let mut map = map;
    let res = rx
        .recv_with(|input| {
            publish_with_policy_with_result(tx, on_full, wait, |slot| map(input, slot))
        })
        .map_err(|_| LaneTxWithError::Failed)?;
    match res {
        None => Ok(false),
        Some(r) => {
            r?;
            Ok(true)
        }
    }
}

/// Receive-with (borrowed) → forward → publish-with (slot write).
///
/// This avoids the intermediate input buffer copy by borrowing the input bytes and copying them
/// directly into the publish slot.
pub fn recv_with_forward_publish_with_step<Rx, Tx>(
    rx: &mut Rx,
    tx: &mut Tx,
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
) -> core::result::Result<bool, LaneTxError>
where
    Rx: LaneRxBorrow,
    Tx: LaneTxWith,
{
    let res = rx
        .recv_with(|input| {
            publish_with_policy_with(tx, on_full, wait, |slot| {
                let n = input.len().min(slot.len());
                slot[..n].copy_from_slice(&input[..n]);
                n
            })
        })
        .map_err(|_| LaneTxError::Failed)?;
    match res {
        None => Ok(false),
        Some(r) => {
            r?;
            Ok(true)
        }
    }
}

/// Receive-slot → forward → publish-slot.
///
/// If the TX is incompatible with the slot's pool, this falls back to copying the slot bytes
/// into a normal `publish`.
pub fn recv_slot_forward_publish_slot_step<Rx, Tx>(
    rx: &mut Rx,
    tx: &mut Tx,
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
) -> core::result::Result<bool, LaneTxError>
where
    Rx: LaneRxSlot,
    Tx: LaneTxSlot<Rx::Slot> + LaneTx,
{
    let Some(slot) = rx.try_recv_slot().map_err(|_| LaneTxError::Failed)? else {
        return Ok(false);
    };

    publish_slot_with_policy_or_copy(tx, slot, on_full, wait)?;
    Ok(true)
}

/// Receive-slot → filter → (forward via slot if true).
///
/// The `filter` closure returns `true` to forward or `false` to drop.
pub fn recv_slot_filter_forward_publish_slot_step<Rx, Tx, F>(
    rx: &mut Rx,
    tx: &mut Tx,
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
    mut filter: F,
) -> core::result::Result<bool, LaneTxError>
where
    Rx: LaneRxSlot,
    Tx: LaneTxSlot<Rx::Slot> + LaneTx,
    Rx::Slot: LaneSlot,
    F: FnMut(&[u8]) -> bool,
{
    let Some(slot) = rx.try_recv_slot().map_err(|_| LaneTxError::Failed)? else {
        return Ok(false);
    };

    if !filter(slot.as_bytes()) {
        // Dropping releases the slot back to the pool.
        return Ok(true);
    }

    publish_slot_with_policy_or_copy(tx, slot, on_full, wait)?;
    Ok(true)
}

/// Receive-slot → inspect → (forward via slot).
///
/// The `inspect` closure returns `true` to forward or `false` to drop.
pub fn recv_slot_inspect_forward_publish_slot_step<Rx, Tx, F>(
    rx: &mut Rx,
    tx: &mut Tx,
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
    mut inspect: F,
) -> core::result::Result<bool, LaneTxError>
where
    Rx: LaneRxSlot,
    Tx: LaneTxSlot<Rx::Slot> + LaneTx,
    Rx::Slot: LaneSlot,
    F: FnMut(&[u8]) -> bool,
{
    let Some(slot) = rx.try_recv_slot().map_err(|_| LaneTxError::Failed)? else {
        return Ok(false);
    };

    if !inspect(slot.as_bytes()) {
        return Ok(true);
    }

    publish_slot_with_policy_or_copy(tx, slot, on_full, wait)?;
    Ok(true)
}

/// Receive-slot → transform → (forward via slot on `ForwardInput`).
///
/// This enables a transform stage to do true zero-copy forwarding for no-op cases.
///
/// - `Drop` drops the slot (releasing it back to the pool).
/// - `ForwardInput` forwards the input slot to `tx` when compatible, otherwise falls back to copying.
/// - `ForwardOutput(n)` publishes the output bytes (copying) and drops the input slot.
pub fn recv_slot_transform_forward_publish_slot_step<Rx, Tx, F>(
    rx: &mut Rx,
    tx: &mut Tx,
    output: &mut [u8],
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
    mut transform: F,
) -> core::result::Result<bool, LaneTxError>
where
    Rx: LaneRxSlot,
    Tx: LaneTxSlot<Rx::Slot> + LaneTx,
    Rx::Slot: LaneSlot,
    F: FnMut(&[u8], &mut [u8]) -> TransformOutcome,
{
    let Some(slot) = rx.try_recv_slot().map_err(|_| LaneTxError::Failed)? else {
        return Ok(false);
    };

    match transform(slot.as_bytes(), output) {
        TransformOutcome::Drop => {
            // Dropping releases the slot back to the pool.
        }
        TransformOutcome::ForwardInput => {
            publish_slot_with_policy_or_copy(tx, slot, on_full, wait)?;
        }
        TransformOutcome::ForwardOutput(out_n) => {
            let out_n = out_n.min(output.len());
            publish_with_policy(tx, &output[..out_n], on_full, wait)?;
            // Dropping releases the slot back to the pool.
        }
    }

    Ok(true)
}

/// Publish an owned slot with `on_full` backpressure semantics.
///
/// On `Incompatible`, falls back to copying the slot bytes into `LaneTx::publish`.
#[inline]
pub fn publish_slot_with_policy_or_copy<Tx, S>(
    tx: &mut Tx,
    mut slot: S,
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
) -> core::result::Result<(), LaneTxError>
where
    Tx: LaneTxSlot<S> + LaneTx,
    S: LaneSlot,
{
    loop {
        match tx.publish_slot(slot) {
            Ok(()) => {
                wait.reset();
                return Ok(());
            }
            Err(LaneTxSlotError::Full(s)) => match on_full {
                OnFullV1::Drop => return Ok(()),
                OnFullV1::Block => {
                    slot = s;
                    wait.wait();
                    continue;
                }
            },
            Err(LaneTxSlotError::Incompatible(s)) => {
                // Fall back to copying into the destination.
                publish_with_policy_wait(tx, s.as_bytes(), on_full, wait)?;
                return Ok(());
            }
            Err(LaneTxSlotError::Failed(_s)) => return Err(LaneTxError::Failed),
        }
    }
}

/// Receive → transform → publish, up to `batch` messages.
///
/// Returns `Ok(true)` if at least one message was processed.
#[allow(clippy::too_many_arguments)]
pub fn recv_transform_publish_batch<Rx, Tx, F>(
    rx: &mut Rx,
    tx: &mut Tx,
    input: &mut [u8],
    output: &mut [u8],
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
    batch: usize,
    mut transform: F,
) -> core::result::Result<bool, LaneTxError>
where
    Rx: LaneRx,
    Tx: LaneTx,
    F: FnMut(&[u8], &mut [u8]) -> TransformOutcome,
{
    let batch = batch.max(1);
    let mut did_any = false;

    for _ in 0..batch {
        let Some(n) = rx.recv(input) else {
            break;
        };
        did_any = true;
        let input = &input[..n];
        match transform(input, output) {
            TransformOutcome::Drop => {}
            TransformOutcome::ForwardInput => {
                publish_with_policy_wait(tx, input, on_full, wait)?;
            }
            TransformOutcome::ForwardOutput(out_n) => {
                let out_n = out_n.min(output.len());
                publish_with_policy_wait(tx, &output[..out_n], on_full, wait)?;
            }
        }
    }

    Ok(did_any)
}

/// Receive from one of N RX lanes (round-robin) → publish.
///
/// This is a building block for fan-in roles (merge/join) without requiring allocs in the hot
/// loop. The caller owns the round-robin cursor.
pub fn recv_select_publish_step<Tx>(
    rxs: &mut [&mut dyn LaneRx],
    rr: &mut usize,
    tx: &mut Tx,
    buf: &mut [u8],
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
) -> core::result::Result<bool, LaneTxError>
where
    Tx: LaneTx,
{
    if rxs.is_empty() {
        return Ok(false);
    }

    let n = rxs.len();
    let start = (*rr) % n;

    for offset in 0..n {
        let idx = (start + offset) % n;
        let Some(msg_n) = rxs[idx].recv(buf) else {
            continue;
        };
        publish_with_policy(tx, &buf[..msg_n], on_full, wait)?;
        *rr = (idx + 1) % n;
        return Ok(true);
    }

    // No work: still advance the cursor so that the first-lane check doesn't stay pinned.
    *rr = (start + 1) % n;
    Ok(false)
}

/// Receive from one of N RX lanes (round-robin) → publish, up to `batch` messages.
///
/// Returns `Ok(true)` if at least one message was processed.
pub fn recv_select_publish_batch<Tx>(
    rxs: &mut [&mut dyn LaneRx],
    rr: &mut usize,
    tx: &mut Tx,
    buf: &mut [u8],
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
    batch: usize,
) -> core::result::Result<bool, LaneTxError>
where
    Tx: LaneTx,
{
    let batch = batch.max(1);
    let mut did_any = false;

    for _ in 0..batch {
        let did = recv_select_publish_step(rxs, rr, tx, buf, on_full, wait)?;
        if !did {
            break;
        }
        did_any = true;
    }

    Ok(did_any)
}

#[inline]
/// Publish a message to a lane, applying `on_full` backpressure semantics with an explicit wait
/// strategy.
///
/// - On success, resets `wait`.
/// - On `Full` + `OnFullV1::Block`, calls `wait.wait()` and retries.
/// - On `Full` + `OnFullV1::Drop`, returns `Ok(())` (drop the message).
pub fn publish_with_policy<Tx: LaneTx>(
    tx: &mut Tx,
    msg: &[u8],
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
) -> core::result::Result<(), LaneTxError> {
    publish_with_policy_wait(tx, msg, on_full, wait)
}

#[inline]
fn publish_with_policy_with<Tx: LaneTxWith>(
    tx: &mut Tx,
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
    mut fill: impl FnMut(&mut [u8]) -> usize,
) -> core::result::Result<(), LaneTxError> {
    loop {
        match tx.publish_with(|slot| fill(slot).min(slot.len())) {
            Ok(()) => {
                wait.reset();
                return Ok(());
            }
            Err(LaneTxError::Full) => match on_full {
                OnFullV1::Drop => return Ok(()),
                OnFullV1::Block => {
                    wait.wait();
                    continue;
                }
            },
            Err(e) => return Err(e),
        }
    }
}

#[inline]
fn publish_with_policy_with_result<Tx: LaneTxWithResult, E>(
    tx: &mut Tx,
    on_full: OnFullV1,
    wait: &mut impl WaitStrategy,
    mut fill: impl FnMut(&mut [u8]) -> core::result::Result<usize, E>,
) -> core::result::Result<(), LaneTxWithError<E>> {
    loop {
        match tx.publish_with_result(|slot| fill(slot).map(|n| n.min(slot.len()))) {
            Ok(()) => {
                wait.reset();
                return Ok(());
            }
            Err(LaneTxWithError::Full) => match on_full {
                OnFullV1::Drop => return Ok(()),
                OnFullV1::Block => {
                    wait.wait();
                    continue;
                }
            },
            Err(e @ LaneTxWithError::Encode(_)) => return Err(e),
            Err(e @ LaneTxWithError::Failed) => return Err(e),
        }
    }
}

#[inline]
fn publish_with_policy_wait<Tx: LaneTx, W: WaitStrategy>(
    tx: &mut Tx,
    msg: &[u8],
    on_full: OnFullV1,
    wait: &mut W,
) -> core::result::Result<(), LaneTxError> {
    loop {
        match tx.publish(msg) {
            Ok(()) => {
                wait.reset();
                return Ok(());
            }
            Err(LaneTxError::Full) => match on_full {
                OnFullV1::Drop => return Ok(()),
                OnFullV1::Block => {
                    wait.wait();
                    continue;
                }
            },
            Err(e) => return Err(e),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use indexbus_core::SpinWait;

    struct Rx {
        msgs: std::collections::VecDeque<Vec<u8>>,
    }

    impl LaneRx for Rx {
        fn recv(&mut self, buf: &mut [u8]) -> Option<usize> {
            let msg = self.msgs.pop_front()?;
            let n = msg.len().min(buf.len());
            buf[..n].copy_from_slice(&msg[..n]);
            Some(n)
        }
    }

    impl LaneRxBorrow for Rx {
        fn recv_with<F, R>(
            &mut self,
            f: F,
        ) -> core::result::Result<Option<R>, crate::lane::LaneRxError>
        where
            F: FnOnce(&[u8]) -> R,
        {
            let Some(msg) = self.msgs.pop_front() else {
                return Ok(None);
            };
            Ok(Some(f(&msg)))
        }
    }

    #[derive(Default)]
    struct Tx {
        out: Vec<Vec<u8>>,
    }

    impl LaneTx for Tx {
        fn publish(&mut self, msg: &[u8]) -> core::result::Result<(), LaneTxError> {
            self.out.push(msg.to_vec());
            Ok(())
        }
    }

    #[test]
    fn recv_select_publish_round_robin() {
        let mut rx0 = Rx {
            msgs: [b"a0".to_vec(), b"a1".to_vec()].into(),
        };
        let mut rx1 = Rx {
            msgs: [b"b0".to_vec(), b"b1".to_vec()].into(),
        };
        let mut tx = Tx::default();
        let mut rr = 0usize;
        let mut buf = [0u8; 16];

        let mut rxs: [&mut dyn LaneRx; 2] = [&mut rx0, &mut rx1];
        let mut wait = SpinWait::default();

        // Consume 4 messages.
        for _ in 0..4 {
            let did = recv_select_publish_step(
                &mut rxs,
                &mut rr,
                &mut tx,
                &mut buf,
                OnFullV1::Block,
                &mut wait,
            )
            .unwrap();
            assert!(did);
        }

        assert_eq!(
            tx.out,
            vec![
                b"a0".to_vec(),
                b"b0".to_vec(),
                b"a1".to_vec(),
                b"b1".to_vec()
            ]
        );
    }

    #[test]
    fn recv_transform_publish_batch_forwards_all_available_up_to_batch() {
        fn map_ok(input: &[u8], output: &mut [u8]) -> TransformOutcome {
            let n = input.len().min(output.len());
            output[..n].copy_from_slice(&input[..n]);
            TransformOutcome::ForwardOutput(n)
        }

        let mut rx = Rx {
            msgs: [b"a".to_vec(), b"b".to_vec(), b"c".to_vec()].into(),
        };
        let mut tx = Tx::default();
        let mut in_buf = [0u8; 16];
        let mut out_buf = [0u8; 16];
        let mut wait = SpinWait::default();

        let did = recv_transform_publish_batch(
            &mut rx,
            &mut tx,
            &mut in_buf,
            &mut out_buf,
            OnFullV1::Block,
            &mut wait,
            32,
            map_ok,
        )
        .unwrap();
        assert!(did);
        assert_eq!(tx.out, vec![b"a".to_vec(), b"b".to_vec(), b"c".to_vec()]);

        let did2 = recv_transform_publish_batch(
            &mut rx,
            &mut tx,
            &mut in_buf,
            &mut out_buf,
            OnFullV1::Block,
            &mut wait,
            32,
            map_ok,
        )
        .unwrap();
        assert!(!did2);
    }

    #[test]
    fn recv_select_publish_batch_consumes_multiple() {
        let mut rx0 = Rx {
            msgs: [b"a0".to_vec(), b"a1".to_vec()].into(),
        };
        let mut rx1 = Rx {
            msgs: [b"b0".to_vec(), b"b1".to_vec()].into(),
        };
        let mut tx = Tx::default();
        let mut rr = 0usize;
        let mut buf = [0u8; 16];
        let mut rxs: [&mut dyn LaneRx; 2] = [&mut rx0, &mut rx1];
        let mut wait = SpinWait::default();

        let did = recv_select_publish_batch(
            &mut rxs,
            &mut rr,
            &mut tx,
            &mut buf,
            OnFullV1::Block,
            &mut wait,
            4,
        )
        .unwrap();
        assert!(did);

        assert_eq!(
            tx.out,
            vec![
                b"a0".to_vec(),
                b"b0".to_vec(),
                b"a1".to_vec(),
                b"b1".to_vec()
            ]
        );
    }

    #[test]
    fn recv_with_transform_publish_step_borrows_input() {
        fn map_ok(input: &[u8], output: &mut [u8]) -> TransformOutcome {
            let n = input.len().min(output.len());
            output[..n].copy_from_slice(&input[..n]);
            TransformOutcome::ForwardOutput(n)
        }

        let mut rx = Rx {
            msgs: [b"abc".to_vec()].into(),
        };
        let mut tx = Tx::default();
        let mut out_buf = [0u8; 16];
        let mut wait = SpinWait::default();

        let did = recv_with_transform_publish_step(
            &mut rx,
            &mut tx,
            &mut out_buf,
            OnFullV1::Block,
            &mut wait,
            map_ok,
        )
        .unwrap();
        assert!(did);
        assert_eq!(tx.out, vec![b"abc".to_vec()]);
    }

    #[test]
    fn recv_with_map_publish_with_result_step_does_not_publish_on_error() {
        #[derive(Default)]
        struct TxWithResult {
            out: Vec<Vec<u8>>,
        }

        impl LaneTxWithResult for TxWithResult {
            fn publish_with_result<F, E>(
                &mut self,
                f: F,
            ) -> core::result::Result<(), crate::lane::LaneTxWithError<E>>
            where
                F: FnOnce(&mut [u8]) -> core::result::Result<usize, E>,
            {
                let mut slot = [0u8; 16];
                match f(&mut slot) {
                    Ok(n) => {
                        self.out.push(slot[..n.min(slot.len())].to_vec());
                        Ok(())
                    }
                    Err(e) => Err(crate::lane::LaneTxWithError::Encode(e)),
                }
            }
        }

        let mut rx = Rx {
            msgs: [b"abc".to_vec()].into(),
        };
        let mut tx = TxWithResult::default();
        let mut wait = SpinWait::default();

        let err = recv_with_map_publish_with_result_step(
            &mut rx,
            &mut tx,
            OnFullV1::Block,
            &mut wait,
            |_input, _slot| Err("map failed"),
        )
        .unwrap_err();

        assert!(tx.out.is_empty());
        assert_eq!(err, crate::lane::LaneTxWithError::Encode("map failed"));
    }

    #[test]
    fn recv_slot_transform_forward_publish_slot_forwards_slot_on_forward_input() {
        #[derive(Clone, Debug, PartialEq, Eq)]
        struct Slot(Vec<u8>);

        impl LaneSlot for Slot {
            fn as_bytes(&self) -> &[u8] {
                &self.0
            }
        }

        struct RxSlot {
            slots: std::collections::VecDeque<Slot>,
        }

        impl LaneRxSlot for RxSlot {
            type Slot = Slot;

            fn try_recv_slot(
                &mut self,
            ) -> core::result::Result<Option<Self::Slot>, crate::lane::LaneRxError> {
                Ok(self.slots.pop_front())
            }
        }

        #[derive(Default)]
        struct TxSlot {
            forwarded: Vec<Vec<u8>>,
            published: Vec<Vec<u8>>,
        }

        impl LaneTx for TxSlot {
            fn publish(&mut self, msg: &[u8]) -> core::result::Result<(), LaneTxError> {
                self.published.push(msg.to_vec());
                Ok(())
            }
        }

        impl LaneTxSlot<Slot> for TxSlot {
            fn publish_slot(
                &mut self,
                slot: Slot,
            ) -> core::result::Result<(), LaneTxSlotError<Slot>> {
                self.forwarded.push(slot.0);
                Ok(())
            }
        }

        let mut rx = RxSlot {
            slots: [Slot(b"abc".to_vec())].into(),
        };
        let mut tx = TxSlot::default();
        let mut out_buf = [0u8; 16];
        let mut wait = SpinWait::default();

        let did = recv_slot_transform_forward_publish_slot_step(
            &mut rx,
            &mut tx,
            &mut out_buf,
            OnFullV1::Block,
            &mut wait,
            |_input, _output| TransformOutcome::ForwardInput,
        )
        .unwrap();
        assert!(did);
        assert_eq!(tx.forwarded, vec![b"abc".to_vec()]);
        assert!(tx.published.is_empty());
    }

    #[test]
    fn recv_slot_transform_forward_publish_slot_publishes_output_on_forward_output() {
        #[derive(Clone, Debug, PartialEq, Eq)]
        struct Slot(Vec<u8>);

        impl LaneSlot for Slot {
            fn as_bytes(&self) -> &[u8] {
                &self.0
            }
        }

        struct RxSlot {
            slots: std::collections::VecDeque<Slot>,
        }

        impl LaneRxSlot for RxSlot {
            type Slot = Slot;

            fn try_recv_slot(
                &mut self,
            ) -> core::result::Result<Option<Self::Slot>, crate::lane::LaneRxError> {
                Ok(self.slots.pop_front())
            }
        }

        #[derive(Default)]
        struct TxSlot {
            forwarded: Vec<Vec<u8>>,
            published: Vec<Vec<u8>>,
        }

        impl LaneTx for TxSlot {
            fn publish(&mut self, msg: &[u8]) -> core::result::Result<(), LaneTxError> {
                self.published.push(msg.to_vec());
                Ok(())
            }
        }

        impl LaneTxSlot<Slot> for TxSlot {
            fn publish_slot(
                &mut self,
                slot: Slot,
            ) -> core::result::Result<(), LaneTxSlotError<Slot>> {
                self.forwarded.push(slot.0);
                Ok(())
            }
        }

        let mut rx = RxSlot {
            slots: [Slot(b"abc".to_vec())].into(),
        };
        let mut tx = TxSlot::default();
        let mut out_buf = [0u8; 16];
        let mut wait = SpinWait::default();

        let did = recv_slot_transform_forward_publish_slot_step(
            &mut rx,
            &mut tx,
            &mut out_buf,
            OnFullV1::Block,
            &mut wait,
            |_input, output| {
                output[0] = b'x';
                TransformOutcome::ForwardOutput(1)
            },
        )
        .unwrap();
        assert!(did);
        assert!(tx.forwarded.is_empty());
        assert_eq!(tx.published, vec![b"x".to_vec()]);
    }

    #[derive(Default)]
    struct CountingWait {
        resets: u64,
        waits: u64,
    }

    impl WaitStrategy for CountingWait {
        fn reset(&mut self) {
            self.resets += 1;
        }

        fn wait(&mut self) {
            self.waits += 1;
        }
    }

    struct FullOnceTx {
        calls: u32,
        published: Vec<Vec<u8>>,
    }

    impl LaneTx for FullOnceTx {
        fn publish(&mut self, msg: &[u8]) -> core::result::Result<(), LaneTxError> {
            self.calls = self.calls.saturating_add(1);
            if self.calls == 1 {
                return Err(LaneTxError::Full);
            }
            self.published.push(msg.to_vec());
            Ok(())
        }
    }

    #[test]
    fn publish_with_policy_drop_does_not_wait_on_full() {
        let mut tx = FullOnceTx {
            calls: 0,
            published: Vec::new(),
        };
        let mut wait = CountingWait::default();

        publish_with_policy(&mut tx, b"abc", OnFullV1::Drop, &mut wait).unwrap();

        // Full + Drop: message is dropped (never published) and no waiting occurs.
        assert!(tx.published.is_empty());
        assert_eq!(wait.waits, 0);
        assert_eq!(wait.resets, 0);
    }

    #[test]
    fn publish_with_policy_block_waits_and_retries_on_full() {
        let mut tx = FullOnceTx {
            calls: 0,
            published: Vec::new(),
        };
        let mut wait = CountingWait::default();

        publish_with_policy(&mut tx, b"abc", OnFullV1::Block, &mut wait).unwrap();

        // Full + Block: wait once, retry, then publish.
        assert_eq!(tx.published, vec![b"abc".to_vec()]);
        assert_eq!(wait.waits, 1);
        assert_eq!(wait.resets, 1);
    }
}