vyre-conform 0.1.0

//! `workgroup.stack` conform specification.
//!
//! The workgroup stack is a bounded LIFO data structure resident in SRAM,
//! coordinated across workgroup lanes via atomic operations. Operations
//! are expressed as serialized command sequences (push/pop/peek/len/is_empty)
//! dispatched to specific lanes.
//!
//! CPU reference semantics: a `Vec`-backed LIFO with explicit capacity.
//! Overflow => status 1. Underflow => status 2. Success => status 0.
//!
//! The WGSL kernel uses `var<workgroup>` storage and atomic length tracking
//! to implement the same contract on GPU shared memory.

use crate::spec::primitive::EquivalenceClass;
use crate::OpSpec;
use vyre_spec::GoldenSample;

// ─── Constants mirroring WGSL kernel ───────────────────────────────────

const STACK_CAPACITY: u32 = 256;

const OP_PUSH: u32 = 1;
const OP_POP: u32 = 2;
const OP_PEEK: u32 = 3;
const OP_LEN: u32 = 4;
const OP_IS_EMPTY: u32 = 5;

const STATUS_OK: u32 = 0;
const STATUS_OVERFLOW: u32 = 1;
const STATUS_UNDERFLOW: u32 = 2;

/// Location-agnostic operation metadata.
pub const VYRE_OP_METADATA: vyre_spec::OpMetadata = vyre_spec::OpMetadata {
    id: "workgroup.stack",
    layer: vyre_spec::Layer::L2,
    category: vyre_spec::MetadataCategory::C,
    version: 1,
    description: "workgroup-local bounded LIFO stack with atomic coordination",
    signature: "(Bytes) -> Bytes",
    strictness: "strict",
    archetype_signature: "(Bytes) -> Bytes",
};

/// Golden samples for this op.
///
/// Each sample encodes a command sequence as little-endian u32 quadruplets:
/// `[op, lane, value, reserved]` per command, and expects a result sequence
/// of `[status, value, len, is_empty]` per command.
pub const GOLDEN: &[GoldenSample] = &[
    GoldenSample {
        op_id: "workgroup.stack",
        // push(lane=0, value=42), pop(lane=0)
        // cmd0: op=1(push), lane=0, value=42, reserved=0
        // cmd1: op=2(pop),  lane=0, value=0,  reserved=0
        input: &[
            0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2A, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00,
        ],
        // res0: status=0(ok), value=0xFFFFFFFF(unused), len=1, is_empty=0
        // res1: status=0(ok), value=42, len=0, is_empty=1
        expected: &[
            0x00, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2A, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x01, 0x00, 0x00, 0x00,
        ],
        reason: "push then pop returns the pushed value (LIFO identity)",
    },
    GoldenSample {
        op_id: "workgroup.stack",
        // pop on empty stack
        // cmd0: op=2(pop), lane=0, value=0, reserved=0
        input: &[
            0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00,
        ],
        // res0: status=2(underflow), value=0xFFFFFFFF, len=0, is_empty=1
        expected: &[
            0x02, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00,
            0x00, 0x00,
        ],
        reason: "pop on empty stack reports underflow status",
    },
];

/// Known-answer tests for this op.
pub const KAT: &[vyre_spec::KatVector] = &[vyre_spec::KatVector {
    input: &[
        // push(lane=0, value=42)
        0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2A, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, // pop(lane=0)
        0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00,
    ],
    expected: &[
        0x00, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x2A, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00,
        0x00, 0x00,
    ],
    source: "hand-verified: push(42) then pop == (ok, 42)",
}];

/// Adversarial inputs for this op.
pub const ADVERSARIAL: &[vyre_spec::AdversarialInput] = &[
    vyre_spec::AdversarialInput {
        input: &[],
        reason: "empty command sequence — zero-length dispatch must not crash",
    },
    vyre_spec::AdversarialInput {
        // pop on empty then check is_empty
        input: &[
            0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x05, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00,
        ],
        reason: "underflow then is_empty — must report underflow (2) then empty (1)",
    },
];

/// CPU reference oracle for the workgroup stack.
///
/// Input: concatenated little-endian `StackCommand` structs (16 bytes each):
///   `[op: u32, lane: u32, value: u32, reserved: u32]`
///
/// Output: concatenated little-endian `StackResult` structs (16 bytes each):
///   `[status: u32, value: u32, len: u32, is_empty: u32]`
///
/// Commands are processed sequentially. Only the command whose `lane` field
/// matches the "active" lane (0 for single-lane CPU reference) is executed.
#[inline]
pub(crate) fn cpu(input: &[u8]) -> Vec<u8> {
    const CMD_SIZE: usize = 16;
    const SENTINEL: u32 = 0xFFFF_FFFF;

    let command_count = input.len() / CMD_SIZE;
    let mut output = Vec::with_capacity(command_count * CMD_SIZE);
    let mut stack: Vec<u32> = Vec::new();

    for i in 0..command_count {
        let base = i * CMD_SIZE;
        let op = read_u32_le(input, base);
        let _lane = read_u32_le(input, base + 4);
        let value = read_u32_le(input, base + 8);

        let mut status = STATUS_OK;
        let mut result_value = SENTINEL;
        let len_after;
        let empty_after;

        match op {
            OP_PUSH => {
                if stack.len() >= STACK_CAPACITY as usize {
                    status = STATUS_OVERFLOW;
                    len_after = STACK_CAPACITY;
                    empty_after = 0;
                } else {
                    stack.push(value);
                    len_after = stack.len() as u32;
                    empty_after = 0;
                }
            }
            OP_POP => {
                if stack.is_empty() {
                    status = STATUS_UNDERFLOW;
                    len_after = 0;
                    empty_after = 1;
                } else {
                    result_value = stack.pop().unwrap_or(SENTINEL);
                    len_after = stack.len() as u32;
                    empty_after = u32::from(stack.is_empty());
                }
            }
            OP_PEEK => {
                if stack.is_empty() {
                    status = STATUS_UNDERFLOW;
                    len_after = 0;
                    empty_after = 1;
                } else {
                    result_value = stack.last().copied().unwrap_or(SENTINEL);
                    len_after = stack.len() as u32;
                    empty_after = 0;
                }
            }
            OP_LEN => {
                len_after = stack.len() as u32;
                empty_after = u32::from(stack.is_empty());
            }
            OP_IS_EMPTY => {
                len_after = stack.len() as u32;
                empty_after = u32::from(stack.is_empty());
                result_value = empty_after;
            }
            _ => {
                // Unknown op — no-op with current state
                len_after = stack.len() as u32;
                empty_after = u32::from(stack.is_empty());
            }
        }

        output.extend_from_slice(&status.to_le_bytes());
        output.extend_from_slice(&result_value.to_le_bytes());
        output.extend_from_slice(&len_after.to_le_bytes());
        output.extend_from_slice(&empty_after.to_le_bytes());
    }

    output
}

fn read_u32_le(data: &[u8], offset: usize) -> u32 {
    if offset + 4 > data.len() {
        return 0;
    }
    u32::from_le_bytes([
        data[offset],
        data[offset + 1],
        data[offset + 2],
        data[offset + 3],
    ])
}

fn wgsl() -> String {
    vyre::ops::workgroup::stack::lowering::WGSL.to_string()
}

/// Build the conformance specification for workgroup.stack.
#[inline]
pub fn vyre_op() -> OpSpec {
    use crate::enforce::category::Category;
    use crate::proof::comparator::ComparatorKind;
    use crate::spec::types::{BoundaryValue, DataType, OpSignature};

    let id = "workgroup.stack";

    OpSpec::builder(id)
        .signature(OpSignature {
            inputs: vec![DataType::Bytes],
            output: DataType::Bytes,
        })
        .cpu_fn(cpu)
        .wgsl_fn(wgsl)
        .laws(vec![])
        .alt_wgsl_fns(vec![])
        .declared_laws(Vec::<crate::spec::types::DeclaredLaw>::new())
        .spec_table(&[])
        .archetypes(&[])
        .mutation_sensitivity(&[])
        .oracle_override(None)
        .since_version(crate::spec::types::Version::V1_0)
        .docs_path("")
        .comparator(ComparatorKind::ExactMatch)
        .category(Category::C {
            hardware: "workgroup-sram-atomics",
            backend_availability: vyre_spec::BackendAvailabilityPredicate::new(|backend| {
                backend == "wgpu"
            }),
        })
        .overflow_contract(crate::spec::types::OverflowContract::Saturating)
        .strictness(crate::spec::types::Strictness::Strict)
        .version(1)
        .equivalence_classes(vec![
            EquivalenceClass::universal("single push-pop sequence"),
            EquivalenceClass::specific("overflow boundary", vec![OP_PUSH, 256]),
            EquivalenceClass::specific("underflow boundary", vec![OP_POP, OP_PEEK]),
        ])
        .boundary_values(vec![
            BoundaryValue::unary("empty stack pop", OP_POP),
            BoundaryValue::unary("empty stack peek", OP_PEEK),
            BoundaryValue::unary("empty stack len", OP_LEN),
            BoundaryValue::unary("empty stack is_empty", OP_IS_EMPTY),
            BoundaryValue::unary("push value zero", 0),
            BoundaryValue::unary("push value max", u32::MAX),
        ])
        .build()
        .expect("registry invariant violated")
}

#[cfg(test)]
mod proptests {
    use super::*;

    #[test]
    fn coverage_artifacts_are_registered() {
        assert!(!KAT.is_empty());
        assert!(!ADVERSARIAL.is_empty());
    }

    #[test]
    fn push_pop_identity() {
        // push(42) then pop => value=42
        let mut input = Vec::new();
        // push command
        input.extend_from_slice(&OP_PUSH.to_le_bytes());
        input.extend_from_slice(&0u32.to_le_bytes()); // lane
        input.extend_from_slice(&42u32.to_le_bytes()); // value
        input.extend_from_slice(&0u32.to_le_bytes()); // reserved
                                                      // pop command
        input.extend_from_slice(&OP_POP.to_le_bytes());
        input.extend_from_slice(&0u32.to_le_bytes());
        input.extend_from_slice(&0u32.to_le_bytes());
        input.extend_from_slice(&0u32.to_le_bytes());

        let result = cpu(&input);
        assert_eq!(result.len(), 32); // 2 results × 16 bytes

        // push result
        let push_status = u32::from_le_bytes([result[0], result[1], result[2], result[3]]);
        assert_eq!(push_status, STATUS_OK);

        // pop result
        let pop_status = u32::from_le_bytes([result[16], result[17], result[18], result[19]]);
        let pop_value = u32::from_le_bytes([result[20], result[21], result[22], result[23]]);
        assert_eq!(pop_status, STATUS_OK);
        assert_eq!(pop_value, 42);
    }

    #[test]
    fn lifo_ordering() {
        let mut input = Vec::new();
        // push 10, 20, 30
        for val in [10u32, 20, 30] {
            input.extend_from_slice(&OP_PUSH.to_le_bytes());
            input.extend_from_slice(&0u32.to_le_bytes());
            input.extend_from_slice(&val.to_le_bytes());
            input.extend_from_slice(&0u32.to_le_bytes());
        }
        // pop 3 times
        for _ in 0..3 {
            input.extend_from_slice(&OP_POP.to_le_bytes());
            input.extend_from_slice(&0u32.to_le_bytes());
            input.extend_from_slice(&0u32.to_le_bytes());
            input.extend_from_slice(&0u32.to_le_bytes());
        }

        let result = cpu(&input);
        // 6 results × 16 bytes each
        assert_eq!(result.len(), 96);

        // Pop results at positions 3, 4, 5 (index × 16 + 4 for value field)
        let pop_vals: Vec<u32> = (3..6)
            .map(|i| {
                let base = i * 16 + 4;
                u32::from_le_bytes([
                    result[base],
                    result[base + 1],
                    result[base + 2],
                    result[base + 3],
                ])
            })
            .collect();
        assert_eq!(
            pop_vals,
            vec![30, 20, 10],
            "LIFO: last pushed = first popped"
        );
    }

    #[test]
    fn underflow_reported() {
        let mut input = Vec::new();
        input.extend_from_slice(&OP_POP.to_le_bytes());
        input.extend_from_slice(&0u32.to_le_bytes());
        input.extend_from_slice(&0u32.to_le_bytes());
        input.extend_from_slice(&0u32.to_le_bytes());

        let result = cpu(&input);
        let status = u32::from_le_bytes([result[0], result[1], result[2], result[3]]);
        assert_eq!(status, STATUS_UNDERFLOW);
    }

    #[test]
    fn overflow_reported() {
        let mut input = Vec::new();
        // push STACK_CAPACITY + 1 values
        for i in 0..=STACK_CAPACITY {
            input.extend_from_slice(&OP_PUSH.to_le_bytes());
            input.extend_from_slice(&0u32.to_le_bytes());
            input.extend_from_slice(&i.to_le_bytes());
            input.extend_from_slice(&0u32.to_le_bytes());
        }

        let result = cpu(&input);
        // Last push should overflow
        let last_base = STACK_CAPACITY as usize * 16;
        let status = u32::from_le_bytes([
            result[last_base],
            result[last_base + 1],
            result[last_base + 2],
            result[last_base + 3],
        ]);
        assert_eq!(status, STATUS_OVERFLOW);
    }

    #[test]
    fn empty_input_produces_empty_output() {
        let result = cpu(&[]);
        assert!(result.is_empty());
    }
}