miden-air 0.32.1

//! Message structs for LogUp bus interactions.
//!
//! Each struct represents a reduced denominator encoding: `α + Σ βⁱ · field_i`.
//! Fields are named for readability; the [`super::lookup::LookupMessage`] trait
//! (implemented further down in this file) provides the `encode` method that
//! produces the extension-field value.
//!
//! Chiplet messages are addressed by interaction-specific bus domains (one [`BusId`]
//! variant per semantic message kind). Constructors pick the interaction domain; payloads
//! start directly with the semantic fields (addr, ctx, etc.).
//!
//! All structs are generic over `E` (base-field expression type, typically `AB::Expr`).

use core::array;

use miden_core::{
    WORD_SIZE,
    field::{Algebra, PrimeCharacteristicRing},
};

use crate::{
    lookup::{Challenges, message::LookupMessage},
    trace::chiplets::hasher::{RATE_LEN, STATE_WIDTH},
};

// MESSAGE PAYLOAD ALIASES
// ================================================================================================

type SpongeState<E> = [E; STATE_WIDTH];
type Rate<E> = [E; RATE_LEN];
type WordFields<E> = [E; WORD_SIZE];

// BUS IDENTIFIERS
// ================================================================================================

/// Width of the `beta_powers` table `Challenges` precomputes for Miden's bus
/// messages, i.e. the exponent of `gamma = beta^MIDEN_MAX_MESSAGE_WIDTH` used in
/// `bus_prefix[i] = alpha + (i + 1) * gamma`.
///
/// Must match the recursive verifier's hardcoded `gamma = beta^16` computation in
/// `crates/lib/core/asm/sys/vm/public_inputs.masm` (4 squarings). The const assertion
/// below is a tripwire: anyone changing `MIDEN_MAX_MESSAGE_WIDTH` must also update the
/// MASM-side computation in lockstep, or the build fails here.
pub const MIDEN_MAX_MESSAGE_WIDTH: usize = 16;

// Tripwire for the MASM-side `gamma = beta^16` hardcoding in
// `crates/lib/core/asm/sys/vm/public_inputs.masm:239-251` (4 sequential squarings).
// If this width ever changes, that MASM must change in lockstep.
const _: () = assert!(
    MIDEN_MAX_MESSAGE_WIDTH == 16,
    "MIDEN_MAX_MESSAGE_WIDTH is hardcoded as 16 by the MASM recursive verifier (4 squarings to reach gamma = beta^16). Update `crates/lib/core/asm/sys/vm/public_inputs.masm` before changing this constant.",
);

/// Domain-separated bus interaction identifier.
///
/// Each variant identifies a distinct bus interaction type. When encoding a message,
/// the bus is cast to `usize` and indexes into
/// [`Challenges::bus_prefix`](crate::lookup::Challenges) to obtain the additive base
/// `bus_prefix[bus] = alpha + (bus + 1) * gamma`.
#[repr(usize)]
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum BusId {
    // --- Out-of-circuit (boundary correction / eval_external) ---
    /// Kernel ROM init: kernel procedure digests from variable-length public inputs.
    KernelRomInit = 0,
    /// Block hash table (decoder p2): root program hash boundary correction.
    BlockHashTable = 1,
    /// Log-deferred state: initial/final deferred-root boundary correction.
    LogDeferredRoot = 2,

    // --- In-circuit buses ---
    KernelRomCall = 3,
    HasherLinearHashInit = 4,
    HasherReturnState = 5,
    HasherAbsorption = 6,
    HasherReturnHash = 7,
    HasherMerkleVerifyInit = 8,
    HasherMerkleOldInit = 9,
    HasherMerkleNewInit = 10,
    MemoryReadElement = 11,
    MemoryWriteElement = 12,
    MemoryReadWord = 13,
    MemoryWriteWord = 14,
    Bitwise = 15,
    AceInit = 16,
    /// Block stack table (decoder p1): tracks control flow block nesting.
    BlockStackTable = 17,
    /// Op group table (decoder p3): tracks operation batch consumption.
    OpGroupTable = 18,
    /// Stack overflow table.
    StackOverflowTable = 19,
    /// Sibling table: shares Merkle tree sibling nodes between old/new root computations.
    SiblingTable = 20,
    /// Range checker bus (LogUp).
    RangeCheck = 21,
    /// ACE wiring bus (LogUp).
    AceWiring = 22,
    /// Hasher perm-link input bus: pairs controller-input rows with perm-cycle row 0.
    HasherPermLinkInput = 23,
    /// Hasher perm-link output bus: pairs controller-output rows with perm-cycle row 15.
    HasherPermLinkOutput = 24,
}

impl BusId {
    /// Last variant discriminant. Paired with the static assertion below, `COUNT` stays
    /// in lockstep with the enum: adding a new variant with a higher discriminant bumps
    /// `COUNT` automatically (and the assertion flags a missed update if the new variant's
    /// discriminant isn't contiguous).
    pub const COUNT: usize = Self::HasherPermLinkOutput as usize + 1;
}

// Per-variant discriminant locks. `BusId::COUNT` only catches gaps — a *reorder* that
// kept the high watermark would silently swap which `bus_prefix[i]` each variant resolves
// to, breaking domain separation across every emitter and consumer. These per-variant
// asserts pin the entire layout so any reorder fails at compile time.
//
// If a new bus is added: append it after the current tail, bump `HasherPermLinkOutput`'s
// expected index here only if necessary, and add a matching assert for the new variant.
const _: () = assert!(BusId::KernelRomInit as usize == 0);
const _: () = assert!(BusId::BlockHashTable as usize == 1);
const _: () = assert!(BusId::LogDeferredRoot as usize == 2);
const _: () = assert!(BusId::KernelRomCall as usize == 3);
const _: () = assert!(BusId::HasherLinearHashInit as usize == 4);
const _: () = assert!(BusId::HasherReturnState as usize == 5);
const _: () = assert!(BusId::HasherAbsorption as usize == 6);
const _: () = assert!(BusId::HasherReturnHash as usize == 7);
const _: () = assert!(BusId::HasherMerkleVerifyInit as usize == 8);
const _: () = assert!(BusId::HasherMerkleOldInit as usize == 9);
const _: () = assert!(BusId::HasherMerkleNewInit as usize == 10);
const _: () = assert!(BusId::MemoryReadElement as usize == 11);
const _: () = assert!(BusId::MemoryWriteElement as usize == 12);
const _: () = assert!(BusId::MemoryReadWord as usize == 13);
const _: () = assert!(BusId::MemoryWriteWord as usize == 14);
const _: () = assert!(BusId::Bitwise as usize == 15);
const _: () = assert!(BusId::AceInit as usize == 16);
const _: () = assert!(BusId::BlockStackTable as usize == 17);
const _: () = assert!(BusId::OpGroupTable as usize == 18);
const _: () = assert!(BusId::StackOverflowTable as usize == 19);
const _: () = assert!(BusId::SiblingTable as usize == 20);
const _: () = assert!(BusId::RangeCheck as usize == 21);
const _: () = assert!(BusId::AceWiring as usize == 22);
const _: () = assert!(BusId::HasherPermLinkInput as usize == 23);
const _: () = assert!(BusId::HasherPermLinkOutput as usize == 24);

// HASHER MESSAGES
// ================================================================================================

/// Hasher chiplet message: a [`BusId`] tag plus a variable-width payload.
///
/// All hasher messages encode as `bus_prefix[kind] + [addr, node_index, ...payload]`; only
/// the payload width differs between variants.
#[derive(Clone, Debug)]
pub struct HasherMsg<E> {
    pub kind: BusId,
    pub addr: E,
    pub node_index: E,
    pub payload: HasherPayload<E>,
}

/// Payload for a [`HasherMsg`]; width varies per interaction kind.
#[derive(Clone, Debug)]
pub enum HasherPayload<E> {
    /// 12-lane sponge state.
    State(SpongeState<E>),
    /// 8-lane rate.
    Rate(Rate<E>),
    /// 4-element word/digest.
    Word(WordFields<E>),
}

/// AIR-side Merkle-init message selected from the controller sub-selectors and rate halves.
///
/// On controller rows, `s1` and `s2` are independently constrained to be boolean. The three
/// Merkle encodings are `(s1, s2) = (0, 1)` for MP, `(1, 0)` for MV, and `(1, 1)` for MU.
/// `direction_bit` selects the rate half containing the leaf word.
#[derive(Clone, Debug)]
pub(super) struct MerkleInitFromSelectorsMsg<E> {
    pub s1: E,
    pub s2: E,
    pub direction_bit: E,
    pub addr: E,
    pub node_index: E,
    pub rate_0: WordFields<E>,
    pub rate_1: WordFields<E>,
}

impl<E: PrimeCharacteristicRing + Clone> HasherMsg<E> {
    // --- State messages (14 payload elements: [addr, node_index, state[12]]) ---

    /// Linear hash / control block init: full 12-lane sponge state.
    ///
    /// Used by: HPERM input, LOGDEFERRED input.
    pub fn linear_hash_init(addr: E, state: SpongeState<E>) -> Self {
        Self {
            kind: BusId::HasherLinearHashInit,
            addr,
            node_index: E::ZERO,
            payload: HasherPayload::State(state),
        }
    }

    /// Control block init: 8 rate lanes + opcode at `capacity[1]`, zeros elsewhere.
    ///
    /// Used by: JOIN, SPLIT, LOOP, SPAN, CALL, SYSCALL, DYN, DYNCALL.
    pub fn control_block(addr: E, rate: &Rate<E>, opcode: u8) -> Self {
        let state = [
            rate[0].clone(),
            rate[1].clone(),
            rate[2].clone(),
            rate[3].clone(),
            rate[4].clone(),
            rate[5].clone(),
            rate[6].clone(),
            rate[7].clone(),
            E::ZERO,
            E::from_u16(opcode as u16),
            E::ZERO,
            E::ZERO,
        ];
        Self {
            kind: BusId::HasherLinearHashInit,
            addr,
            node_index: E::ZERO,
            payload: HasherPayload::State(state),
        }
    }

    /// Return full sponge state after permutation.
    ///
    /// Used by: HPERM output, LOGDEFERRED output.
    pub fn return_state(addr: E, state: SpongeState<E>) -> Self {
        Self {
            kind: BusId::HasherReturnState,
            addr,
            node_index: E::ZERO,
            payload: HasherPayload::State(state),
        }
    }

    // --- Rate messages (10 payload elements: [addr, node_index, rate[8]]) ---

    /// Absorb new rate into running hash.
    ///
    /// Used by: RESPAN.
    pub fn absorption(addr: E, rate: Rate<E>) -> Self {
        Self {
            kind: BusId::HasherAbsorption,
            addr,
            node_index: E::ZERO,
            payload: HasherPayload::Rate(rate),
        }
    }

    // --- Word messages (6 payload elements: [addr, node_index, word[4]]) ---

    /// Return digest only (node_index = 0).
    ///
    /// Used by: END, MPVERIFY output, MRUPDATE output.
    pub fn return_hash(addr: E, word: WordFields<E>) -> Self {
        Self {
            kind: BusId::HasherReturnHash,
            addr,
            node_index: E::ZERO,
            payload: HasherPayload::Word(word),
        }
    }

    /// Start Merkle path verification (with explicit node_index).
    ///
    /// Used by: MPVERIFY input.
    pub fn merkle_verify_init(addr: E, node_index: E, word: WordFields<E>) -> Self {
        Self {
            kind: BusId::HasherMerkleVerifyInit,
            addr,
            node_index,
            payload: HasherPayload::Word(word),
        }
    }

    /// Start Merkle update, old path (with explicit node_index).
    ///
    /// Used by: MRUPDATE old input.
    pub fn merkle_old_init(addr: E, node_index: E, word: WordFields<E>) -> Self {
        Self {
            kind: BusId::HasherMerkleOldInit,
            addr,
            node_index,
            payload: HasherPayload::Word(word),
        }
    }

    /// Start Merkle update, new path (with explicit node_index).
    ///
    /// Used by: MRUPDATE new input.
    pub fn merkle_new_init(addr: E, node_index: E, word: WordFields<E>) -> Self {
        Self {
            kind: BusId::HasherMerkleNewInit,
            addr,
            node_index,
            payload: HasherPayload::Word(word),
        }
    }
}

// MEMORY MESSAGES
// ================================================================================================

/// Memory chiplet message. Variants differ by payload size.
///
/// Encodes as `bus_prefix[bus] + [ctx, addr, clk, ...payload]`. Use the [`MemoryMsg`]
/// associated functions (`read_element`, `write_element`, `read_word`, `write_word`) to
/// build messages with the correct interaction kind.
#[derive(Clone, Debug)]
pub enum MemoryMsg<E> {
    /// 5-element message: `[ctx, addr, clk, element]`.
    ///
    /// `#[non_exhaustive]` forces external construction through the typed
    /// [`MemoryMsg::read_element`] / [`MemoryMsg::write_element`] helpers, which pin
    /// `bus` to `MemoryReadElement` / `MemoryWriteElement`. Direct external construction
    /// with an arbitrary `BusId` would silently break bus domain separation.
    #[non_exhaustive]
    Element {
        bus: BusId,
        ctx: E,
        addr: E,
        clk: E,
        element: E,
    },
    /// 8-element message: `[ctx, addr, clk, word[0..4]]`.
    ///
    /// `#[non_exhaustive]` forces external construction through the typed
    /// [`MemoryMsg::read_word`] / [`MemoryMsg::write_word`] helpers — see
    /// [`MemoryMsg::Element`] for rationale.
    #[non_exhaustive]
    Word {
        bus: BusId,
        ctx: E,
        addr: E,
        clk: E,
        word: WordFields<E>,
    },
}

impl<E> MemoryMsg<E> {
    /// Read a single element from memory.
    pub fn read_element(ctx: E, addr: E, clk: E, element: E) -> Self {
        Self::Element {
            bus: BusId::MemoryReadElement,
            ctx,
            addr,
            clk,
            element,
        }
    }

    /// Write a single element to memory.
    pub fn write_element(ctx: E, addr: E, clk: E, element: E) -> Self {
        Self::Element {
            bus: BusId::MemoryWriteElement,
            ctx,
            addr,
            clk,
            element,
        }
    }

    /// Read a 4-element word from memory.
    pub fn read_word(ctx: E, addr: E, clk: E, word: WordFields<E>) -> Self {
        Self::Word {
            bus: BusId::MemoryReadWord,
            ctx,
            addr,
            clk,
            word,
        }
    }

    /// Write a 4-element word to memory.
    pub fn write_word(ctx: E, addr: E, clk: E, word: WordFields<E>) -> Self {
        Self::Word {
            bus: BusId::MemoryWriteWord,
            ctx,
            addr,
            clk,
            word,
        }
    }
}

// BITWISE MESSAGE
// ================================================================================================

/// Bitwise chiplet message (4 elements): `[op, a, b, result]`.
#[derive(Clone, Debug)]
pub struct BitwiseMsg<E> {
    pub op: E,
    pub a: E,
    pub b: E,
    pub result: E,
}

impl<E: PrimeCharacteristicRing> BitwiseMsg<E> {
    const AND_SELECTOR: u32 = 0;
    const XOR_SELECTOR: u32 = 1;

    /// Bitwise AND message (op selector = 0).
    pub fn and(a: E, b: E, result: E) -> Self {
        Self {
            op: E::from_u32(Self::AND_SELECTOR),
            a,
            b,
            result,
        }
    }

    /// Bitwise XOR message (op selector = 1).
    pub fn xor(a: E, b: E, result: E) -> Self {
        Self {
            op: E::from_u32(Self::XOR_SELECTOR),
            a,
            b,
            result,
        }
    }
}

// DECODER MESSAGES
// ================================================================================================

/// Block stack message: `[block_id, parent_id, is_loop, ctx, fmp, depth, fn_hash[4]]`.
///
/// `Simple` — for blocks that don't save context (JOIN/SPLIT/SPAN/DYN/LOOP/RESPAN/END-simple).
/// Context fields are encoded as zeros.
///
/// `Full` — for blocks that save/restore the caller's execution context
/// (CALL/SYSCALL/DYNCALL/END-call).
#[derive(Clone, Debug)]
pub enum BlockStackMsg<E> {
    Simple {
        block_id: E,
        parent_id: E,
        is_loop: E,
    },
    Full {
        block_id: E,
        parent_id: E,
        is_loop: E,
        ctx: E,
        fmp: E,
        depth: E,
        fn_hash: WordFields<E>,
    },
}

/// Block hash queue message (7 elements):
/// `[child_hash[4], parent, is_first_child, is_loop_body]`.
///
/// `FirstChild` — first child of a JOIN (is_first_child = 1, is_loop_body = 0).
/// `Child` — non-first, non-loop child (is_first_child = 0, is_loop_body = 0).
/// `LoopBody` — loop body entry (is_first_child = 0, is_loop_body = 1).
/// `End` — removal at END; both flags are computed expressions.
#[derive(Clone, Debug)]
pub enum BlockHashMsg<E> {
    FirstChild {
        parent: E,
        child_hash: WordFields<E>,
    },
    Child {
        parent: E,
        child_hash: WordFields<E>,
    },
    LoopBody {
        parent: E,
        child_hash: WordFields<E>,
    },
    End {
        parent: E,
        child_hash: WordFields<E>,
        is_first_child: E,
        is_loop_body: E,
    },
}

/// Op group table message (3 elements): `[batch_id, group_pos, group_value]`.
#[derive(Clone, Debug)]
pub struct OpGroupMsg<E> {
    pub batch_id: E,
    pub group_pos: E,
    pub group_value: E,
}

impl<E: PrimeCharacteristicRing + Clone> OpGroupMsg<E> {
    /// Create an op group message. Computes `group_pos = group_count - offset`.
    pub fn new<V>(batch_id: &E, group_count: V, offset: u16, group_value: E) -> Self
    where
        V: core::ops::Sub<E, Output = E> + Clone,
    {
        Self {
            batch_id: batch_id.clone(),
            group_pos: group_count - E::from_u16(offset),
            group_value,
        }
    }
}

// STACK MESSAGE
// ================================================================================================

/// Stack overflow table message (3 elements): `[clk, val, prev]`.
///
/// `clk` is the cycle at which the value spilled past `stack[15]`, `val` is the spilled element,
/// and `prev` links to the previous overflow entry (the prior `b1`).
#[derive(Clone, Debug)]
pub struct StackOverflowMsg<E> {
    pub clk: E,
    pub val: E,
    pub prev: E,
}

// HASHER PERM-LINK MESSAGE
// ================================================================================================

/// Beta-power offset at which the Poseidon2 state starts in a perm-link denominator.
///
/// Offset 2 aligns the state lanes with full-state hasher messages; beta^1 is unused.
const HASHER_PERM_LINK_STATE_OFFSET: usize = 2;
const _: () = assert!(HASHER_PERM_LINK_STATE_OFFSET + STATE_WIDTH <= MIDEN_MAX_MESSAGE_WIDTH);

/// Hasher perm-link message: `perm_id` plus the full Poseidon2 state.
///
/// The id ties a controller input/output row pair to the corresponding Poseidon2 permutation
/// instance. `state` is encoded at `HASHER_PERM_LINK_STATE_OFFSET` to match full-state hasher
/// messages.
#[derive(Clone, Debug)]
pub enum HasherPermLinkMsg<E> {
    Input { perm_id: E, state: SpongeState<E> },
    Output { perm_id: E, state: SpongeState<E> },
}

// KERNEL ROM MESSAGE
// ================================================================================================

/// Kernel ROM message (4 elements): `bus_prefix[bus] + [digest[4]]`.
///
/// Two bus domains: INIT (one remove per declared procedure, balanced by the boundary
/// correction from public inputs) and CALL (one insert per SYSCALL, carrying the
/// multiplicity from kernel ROM column 0; balanced by decoder-emitted SYSCALL removes).
#[derive(Clone, Debug)]
pub struct KernelRomMsg<E> {
    bus: BusId,
    pub digest: WordFields<E>,
}

impl<E: PrimeCharacteristicRing + Clone> KernelRomMsg<E> {
    /// Kernel procedure call message (SYSCALL request side + chiplet CALL response).
    pub fn call(digest: WordFields<E>) -> Self {
        Self { bus: BusId::KernelRomCall, digest }
    }

    /// Kernel procedure init message (public-input boundary + chiplet INIT response).
    pub fn init(digest: WordFields<E>) -> Self {
        Self { bus: BusId::KernelRomInit, digest }
    }
}

// ACE MESSAGE
// ================================================================================================

/// ACE circuit evaluation init message (5 elements): `[clk, ctx, ptr, num_read, num_eval]`.
#[derive(Clone, Debug)]
pub struct AceInitMsg<E> {
    pub clk: E,
    pub ctx: E,
    pub ptr: E,
    pub num_read: E,
    pub num_eval: E,
}

// RANGE CHECK MESSAGE
// ================================================================================================

/// Range check message (1 element): `[value]`.
///
/// The denominator is `α + β⁰ · value`.
#[derive(Clone, Debug)]
pub struct RangeMsg<E> {
    pub value: E,
}

// LOG-DEFERRED STATE MESSAGE
// ================================================================================================

/// Log-deferred state message (4 elements): deferred root `state[4]`.
#[derive(Clone, Debug)]
pub struct LogDeferredMsg<E> {
    pub state: WordFields<E>,
}

// SIBLING TABLE MESSAGE
// ================================================================================================

// ACE WIRING MESSAGE
// ================================================================================================

/// ACE wiring bus message (5 elements): `[clk, ctx, id, v0, v1]`.
///
/// Encodes a single wire entry for the ACE wiring bus. Each wire carries
/// an identifier and a two-coefficient extension-field value.
#[derive(Clone, Debug)]
pub struct AceWireMsg<E> {
    pub clk: E,
    pub ctx: E,
    pub id: E,
    pub v0: E,
    pub v1: E,
}

// CHIPLET RESPONSE MESSAGES
// ================================================================================================

/// Memory chiplet response message with conditional element/word encoding.
///
/// The chiplet-side memory response must select between element access (4 payload
/// elements: `[ctx, addr, clk, element]`) and word access (7 payload elements:
/// `[ctx, addr, clk, word[4]]`) based on `is_word`. The label, address, and element are
/// all pre-computed from the chiplet columns (including the idx0/idx1 element mux).
#[derive(Clone, Debug)]
pub struct MemoryResponseMsg<E> {
    pub is_read: E,
    pub ctx: E,
    pub addr: E,
    pub clk: E,
    pub is_word: E,
    pub element: E,
    pub word: WordFields<E>,
}

// LOOKUP MESSAGE IMPLEMENTATIONS
// ================================================================================================

// --- HasherMsg (interaction-specific bus ids) ----------------------------------------------------

impl<E, EF> LookupMessage<E, EF> for HasherMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        let mut acc = challenges.bus_prefix[self.kind as usize].clone();
        acc += challenges.inner_product_at(0, &[self.addr.clone(), self.node_index.clone()]);
        let payload = match &self.payload {
            HasherPayload::State(state) => state.as_slice(),
            HasherPayload::Rate(rate) => rate.as_slice(),
            HasherPayload::Word(word) => word.as_slice(),
        };
        acc += challenges.inner_product_at(2, payload);
        acc
    }
}

impl<E, EF> LookupMessage<E, EF> for MerkleInitFromSelectorsMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        let s1 = self.s1.clone();
        let s2 = self.s2.clone();
        let not_s1 = E::ONE - s1.clone();
        let not_s2 = E::ONE - s2.clone();
        let f_mp = not_s1 * s2.clone();
        let f_mv = s1.clone() * not_s2;
        let f_mu = s1 * s2;

        let mut acc = challenges.bus_prefix[BusId::HasherMerkleVerifyInit as usize].clone() * f_mp
            + challenges.bus_prefix[BusId::HasherMerkleOldInit as usize].clone() * f_mv
            + challenges.bus_prefix[BusId::HasherMerkleNewInit as usize].clone() * f_mu;
        acc += challenges.inner_product_at(0, &[self.addr.clone(), self.node_index.clone()]);

        let bit = self.direction_bit.clone();
        let one_minus_bit = E::ONE - bit.clone();
        let word: WordFields<E> = array::from_fn(|i| {
            self.rate_0[i].clone() * one_minus_bit.clone() + self.rate_1[i].clone() * bit.clone()
        });
        acc += challenges.inner_product_at(2, &word);
        acc
    }
}

// --- MemoryMsg (interaction-specific bus ids) ----------------------------------------------------

impl<E, EF> LookupMessage<E, EF> for MemoryMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        let bus = match self {
            Self::Element { bus, .. } | Self::Word { bus, .. } => *bus as usize,
        };
        let mut acc = challenges.bus_prefix[bus].clone();
        match self {
            Self::Element { ctx, addr, clk, element, .. } => {
                acc += challenges.inner_product_at(
                    0,
                    &[ctx.clone(), addr.clone(), clk.clone(), element.clone()],
                );
            },
            Self::Word { ctx, addr, clk, word, .. } => {
                acc += challenges.inner_product_at(0, &[ctx.clone(), addr.clone(), clk.clone()]);
                acc += challenges.inner_product_at(3, word.as_slice());
            },
        }
        acc
    }
}

// --- BitwiseMsg ----------------------------------------------------------------------------------

impl<E, EF> LookupMessage<E, EF> for BitwiseMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        challenges.encode(
            BusId::Bitwise as usize,
            [self.op.clone(), self.a.clone(), self.b.clone(), self.result.clone()],
        )
    }
}

// --- BlockStackMsg -------------------------------------------------------------------------------

impl<E, EF> LookupMessage<E, EF> for BlockStackMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        let mut acc = challenges.bus_prefix[BusId::BlockStackTable as usize].clone();
        match self {
            // `Simple` zero-pads to 10 slots; slots `3..10` contribute `β^k · 0 = 0` so
            // they are elided from the loop.
            Self::Simple { block_id, parent_id, is_loop } => {
                acc += challenges
                    .inner_product_at(0, &[block_id.clone(), parent_id.clone(), is_loop.clone()]);
            },
            Self::Full {
                block_id,
                parent_id,
                is_loop,
                ctx,
                fmp,
                depth,
                fn_hash,
            } => {
                acc += challenges.inner_product_at(
                    0,
                    &[
                        block_id.clone(),
                        parent_id.clone(),
                        is_loop.clone(),
                        ctx.clone(),
                        fmp.clone(),
                        depth.clone(),
                    ],
                );
                acc += challenges.inner_product_at(6, fn_hash.as_slice());
            },
        }
        acc
    }
}

// --- BlockHashMsg --------------------------------------------------------------------------------

impl<E, EF> LookupMessage<E, EF> for BlockHashMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        // Per-variant fan-in: produce the (parent, child_hash, is_first_child, is_loop_body)
        // tuple, then emit a flat 7-slot payload laid out as
        // `[child_hash[4], parent, is_first_child, is_loop_body]`.
        let (parent, child_hash, is_first_child, is_loop_body) = match self {
            Self::FirstChild { parent, child_hash } => (parent, child_hash, E::ONE, E::ZERO),
            Self::Child { parent, child_hash } => (parent, child_hash, E::ZERO, E::ZERO),
            Self::LoopBody { parent, child_hash } => (parent, child_hash, E::ZERO, E::ONE),
            Self::End {
                parent,
                child_hash,
                is_first_child,
                is_loop_body,
            } => (parent, child_hash, is_first_child.clone(), is_loop_body.clone()),
        };
        challenges.encode(
            BusId::BlockHashTable as usize,
            [
                child_hash[0].clone(),
                child_hash[1].clone(),
                child_hash[2].clone(),
                child_hash[3].clone(),
                parent.clone(),
                is_first_child,
                is_loop_body,
            ],
        )
    }
}

// --- OpGroupMsg ----------------------------------------------------------------------------------

impl<E, EF> LookupMessage<E, EF> for OpGroupMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        challenges.encode(
            BusId::OpGroupTable as usize,
            [self.batch_id.clone(), self.group_pos.clone(), self.group_value.clone()],
        )
    }
}

// --- StackOverflowMsg ----------------------------------------------------------------------------

impl<E, EF> LookupMessage<E, EF> for StackOverflowMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        challenges.encode(
            BusId::StackOverflowTable as usize,
            [self.clk.clone(), self.val.clone(), self.prev.clone()],
        )
    }
}

// --- KernelRomMsg --------------------------------------------------------------------------------

impl<E, EF> LookupMessage<E, EF> for KernelRomMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        challenges.encode(self.bus as usize, self.digest.clone())
    }
}

// --- AceInitMsg ----------------------------------------------------------------------------------

impl<E, EF> LookupMessage<E, EF> for AceInitMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        challenges.encode(
            BusId::AceInit as usize,
            [
                self.clk.clone(),
                self.ctx.clone(),
                self.ptr.clone(),
                self.num_read.clone(),
                self.num_eval.clone(),
            ],
        )
    }
}

// --- RangeMsg ------------------------------------------------------------------------------------

impl<E, EF> LookupMessage<E, EF> for RangeMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        challenges.encode(BusId::RangeCheck as usize, [self.value.clone()])
    }
}

// --- LogDeferredMsg ----------------------------------------------------------------------------

impl<E, EF> LookupMessage<E, EF> for LogDeferredMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        challenges.encode(BusId::LogDeferredRoot as usize, self.state.clone())
    }
}

// --- HasherPermLinkMsg ---------------------------------------------------------------------------

impl<E, EF> LookupMessage<E, EF> for HasherPermLinkMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        let (bus, perm_id, state) = match self {
            Self::Input { perm_id, state } => (BusId::HasherPermLinkInput, perm_id, state),
            Self::Output { perm_id, state } => (BusId::HasherPermLinkOutput, perm_id, state),
        };
        let mut acc = challenges.bus_prefix[bus as usize].clone();
        acc += perm_id.clone();
        acc += challenges.inner_product_at(HASHER_PERM_LINK_STATE_OFFSET, state.as_slice());
        acc
    }
}

// --- AceWireMsg ----------------------------------------------------------------------------------

impl<E, EF> LookupMessage<E, EF> for AceWireMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        challenges.encode(
            BusId::AceWiring as usize,
            [
                self.clk.clone(),
                self.ctx.clone(),
                self.id.clone(),
                self.v0.clone(),
                self.v1.clone(),
            ],
        )
    }
}

// LookupMessage impls for the response + sibling structs
// ================================================================================================
//
// The `*ResponseMsg` structs below carry `LookupMessage<E, EF>` impls consumed by
// `lookup/buses/chiplet_responses.rs`. The runtime-muxed encoding (bus prefix muxed
// by `is_read`/`is_word` flags) keeps the response-column transition at degree 8.

impl<E, EF> LookupMessage<E, EF> for MemoryResponseMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        let bp = &challenges.beta_powers;
        let is_read = self.is_read.clone();
        let is_write: E = E::ONE - is_read.clone();
        let is_word = self.is_word.clone();
        let is_element: E = E::ONE - is_word.clone();

        // Mux only the bus prefix; the payload (ctx, addr, clk, ...) is shared. Factored
        // as a read/write select per access width so the four (read/write × element/word)
        // cases stay audit-visible without blowing the polynomial degree.
        let prefix_element = challenges.bus_prefix[BusId::MemoryReadElement as usize].clone()
            * is_read.clone()
            + challenges.bus_prefix[BusId::MemoryWriteElement as usize].clone() * is_write.clone();
        let prefix_word = challenges.bus_prefix[BusId::MemoryReadWord as usize].clone() * is_read
            + challenges.bus_prefix[BusId::MemoryWriteWord as usize].clone() * is_write;
        let prefix = prefix_element * is_element.clone() + prefix_word * is_word.clone();

        let mut acc = prefix;
        acc += bp[0].clone() * self.ctx.clone();
        acc += bp[1].clone() * self.addr.clone();
        acc += bp[2].clone() * self.clk.clone();

        // Element payload (gated by is_element) vs word payload (gated by is_word).
        acc += bp[3].clone() * self.element.clone() * is_element;
        acc += challenges.inner_product_at(3, self.word.as_slice()) * is_word;
        acc
    }
}

// SIBLING MESSAGES
// ================================================================================================
//
// [`SiblingMsg<E>`] carries an already selected rate half and a [`SiblingBit`] tag.
// [`SiblingFromRatesMsg<E>`] receives both rate halves and makes the same selection inside the AIR
// encoding. Both use the sparse β layout expected by the hasher chiplet.

/// Sibling-table message for the Merkle sibling bus.
///
/// The Merkle direction bit picks which half of the hasher rate block holds the sibling:
/// `bit = 0` → sibling at `h[4..8]`, payload lands in β positions `[1, 2, 7, 8, 9, 10]`
/// (mrupdate_id at β¹, node_index at β², rate1 at β⁷..β¹⁰); `bit = 1` → sibling at
/// `h[0..4]`, payload lands in β positions `[1, 2, 3, 4, 5, 6]`.
#[derive(Clone, Debug)]
pub struct SiblingMsg<E> {
    pub bit: SiblingBit,
    pub mrupdate_id: E,
    pub node_index: E,
    pub h: WordFields<E>,
}

/// Which half of the hasher rate block holds the sibling word for this row.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum SiblingBit {
    /// `bit = 0` — sibling lives in the high rate half (`h[4..8]`).
    Zero,
    /// `bit = 1` — sibling lives in the low rate half (`h[0..4]`).
    One,
}

impl<E, EF> LookupMessage<E, EF> for SiblingMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        let mut acc = challenges.bus_prefix[BusId::SiblingTable as usize].clone();
        acc += challenges.inner_product_at(1, &[self.mrupdate_id.clone(), self.node_index.clone()]);
        let base = match self.bit {
            SiblingBit::Zero => 7,
            SiblingBit::One => 3,
        };
        acc += challenges.inner_product_at(base, self.h.as_slice());
        acc
    }
}

/// AIR-side sibling message that selects the sibling from the two hasher rate halves.
///
/// For a boolean `direction_bit`, this encodes the same value as [`SiblingMsg`]. The Merkle input
/// constraints enforce booleanity. Selecting the rate here lets one signed interaction handle both
/// MRUPDATE legs and both directions.
#[derive(Clone, Debug)]
pub(super) struct SiblingFromRatesMsg<E> {
    pub direction_bit: E,
    pub mrupdate_id: E,
    pub node_index: E,
    pub rate_0: WordFields<E>,
    pub rate_1: WordFields<E>,
}

impl<E, EF> LookupMessage<E, EF> for SiblingFromRatesMsg<E>
where
    E: PrimeCharacteristicRing + Clone,
    EF: PrimeCharacteristicRing + Clone + Algebra<E>,
{
    fn encode(&self, challenges: &Challenges<EF>) -> EF {
        let mut acc = challenges.bus_prefix[BusId::SiblingTable as usize].clone();
        acc += challenges.inner_product_at(1, &[self.mrupdate_id.clone(), self.node_index.clone()]);
        let bit = self.direction_bit.clone();
        let one_minus_bit = E::ONE - bit.clone();
        let selected_rate_0: WordFields<E> =
            array::from_fn(|i| self.rate_0[i].clone() * bit.clone());
        let selected_rate_1: WordFields<E> =
            array::from_fn(|i| self.rate_1[i].clone() * one_minus_bit.clone());
        acc += challenges.inner_product_at(3, &selected_rate_0);
        acc += challenges.inner_product_at(7, &selected_rate_1);
        acc
    }
}

#[cfg(test)]
mod tests {
    use miden_core::Felt;

    use super::{
        BusId, HasherMsg, MIDEN_MAX_MESSAGE_WIDTH, MerkleInitFromSelectorsMsg, SiblingBit,
        SiblingFromRatesMsg, SiblingMsg,
    };
    use crate::lookup::{Challenges, message::LookupMessage};

    /// The tests compare each combined encoder with the simpler messages it replaces. Several
    /// deterministic challenge pairs exercise different coefficients; boolean selector
    /// constraints provide the algebraic equivalence.
    const CHALLENGE_POINTS: [(u64, u64); 3] = [
        (29, 31),
        (0x0123_4567_89ab_cdef, 0xa5a5_5a5a_0f0f_1111),
        (1 << 40, (1 << 50) + 33),
    ];

    fn challenge_points() -> impl Iterator<Item = Challenges<Felt>> {
        CHALLENGE_POINTS.into_iter().map(|(alpha, beta)| {
            Challenges::new(
                Felt::new_unchecked(alpha),
                Felt::new_unchecked(beta),
                MIDEN_MAX_MESSAGE_WIDTH,
                BusId::COUNT,
            )
        })
    }

    #[test]
    fn sibling_from_rates_matches_selected_sibling_for_boolean_directions() {
        let mrupdate_id = Felt::from_u32(37);
        let node_index = Felt::from_u32(41);
        let rate_0 = [43, 47, 53, 59].map(Felt::from_u32);
        let rate_1 = [61, 67, 71, 73].map(Felt::from_u32);

        for challenges in challenge_points() {
            for (direction_bit, bit, sibling) in
                [(Felt::ZERO, SiblingBit::Zero, rate_1), (Felt::ONE, SiblingBit::One, rate_0)]
            {
                let from_rates = SiblingFromRatesMsg {
                    direction_bit,
                    mrupdate_id,
                    node_index,
                    rate_0,
                    rate_1,
                };
                let selected = SiblingMsg { bit, mrupdate_id, node_index, h: sibling };

                assert_eq!(
                    <SiblingFromRatesMsg<Felt> as LookupMessage<Felt, Felt>>::encode(
                        &from_rates,
                        &challenges,
                    ),
                    <SiblingMsg<Felt> as LookupMessage<Felt, Felt>>::encode(&selected, &challenges,),
                );
            }
        }
    }

    #[test]
    fn merkle_init_from_selectors_matches_typed_messages() {
        let addr = Felt::from_u32(37);
        let node_index = Felt::from_u32(41);
        let rate_0 = [43, 47, 53, 59].map(Felt::from_u32);
        let rate_1 = [61, 67, 71, 73].map(Felt::from_u32);

        for challenges in challenge_points() {
            for (s1, s2, kind) in [
                (Felt::ZERO, Felt::ONE, BusId::HasherMerkleVerifyInit),
                (Felt::ONE, Felt::ZERO, BusId::HasherMerkleOldInit),
                (Felt::ONE, Felt::ONE, BusId::HasherMerkleNewInit),
            ] {
                for direction_bit in [Felt::ZERO, Felt::ONE] {
                    let from_selectors = MerkleInitFromSelectorsMsg {
                        s1,
                        s2,
                        direction_bit,
                        addr,
                        node_index,
                        rate_0,
                        rate_1,
                    };
                    let word = if direction_bit == Felt::ZERO { rate_0 } else { rate_1 };
                    let typed = match kind {
                        BusId::HasherMerkleVerifyInit => {
                            HasherMsg::merkle_verify_init(addr, node_index, word)
                        },
                        BusId::HasherMerkleOldInit => {
                            HasherMsg::merkle_old_init(addr, node_index, word)
                        },
                        BusId::HasherMerkleNewInit => {
                            HasherMsg::merkle_new_init(addr, node_index, word)
                        },
                        _ => unreachable!("test cases contain only Merkle-init bus ids"),
                    };

                    assert_eq!(
                        <MerkleInitFromSelectorsMsg<Felt> as LookupMessage<Felt, Felt>>::encode(
                            &from_selectors,
                            &challenges,
                        ),
                        <HasherMsg<Felt> as LookupMessage<Felt, Felt>>::encode(&typed, &challenges,),
                    );
                }
            }
        }
    }
}