# ML-KEM Sender / Receiver

> Prove post-quantum ML-KEM (Kyber) KeyGen, Encaps and Decaps in zero knowledge at all FIPS 203 levels, with keys and shared secrets kept private.

- URL: https://oumuamua.dev/primitives/encryption/mlkem
- Published: 2026-10-03T10:49:16+07:00
- Updated: 2026-10-11T12:51:24+07:00

The ML-KEM chiplet proves FIPS 203 key generation, encapsulation and decapsulation at ML-KEM-512, ML-KEM-768 and
ML-KEM-1024, with seeds, keys and the shared secret kept in the witness. One pipeline proves up to 256 calls of any of
the three kinds in a single proof. It shares `hekate-pqc` with the [ML-DSA chiplet](https://oumuamua.dev/primitives/signatures/mldsa).

## What the proof states

Each call runs one of the FIPS 203 internal algorithms:

| Call   | FIPS 203                               | In               | Out                        |
|:-------|:---------------------------------------|:-----------------|:---------------------------|
| KeyGen | Algorithm 16, `ML-KEM.KeyGen_internal` | seed d           | ek, H(ek), dk_PKE          |
| Encaps | Algorithm 17, `ML-KEM.Encaps_internal` | ek, randomness m | shared key K, ciphertext c |
| Decaps | Algorithm 18, `ML-KEM.Decaps_internal` | dk, c            | shared key K, valid        |

KeyGen leaves z to the caller, who appends it to form dk = dk_PKE ‖ ek ‖ H(ek) ‖ z. Decaps re-encrypts and compares.
When c fails the comparison, it returns the implicit-rejection key J(z ‖ c) as FIPS 203 specifies and reports valid =
false to the host program.

The FIPS 203 input checks are part of the statement. The Codec table rejects an encapsulation-key coefficient at or
above q, the modulus check, and a decapsulation key whose stored H(ek) differs from the hash of its ek breaks the word
bus, the hash check. Trace generation refuses both inputs before any proving starts.

The host program decides what becomes public. Both [examples](https://oumuamua.dev/primitives/encryption/mlkem#examples) publish H(ek) and c and keep the seeds, the
keys and K in the witness.

## How the tables split the work

Encapsulation runs through the tables like this, per the hekate-pqc README:

```text
FIPS 203 Encaps_internal(ek, m)               table
────────────────────────────────────────────  ───────────────
h ← H(ek); (K, r) ← G(m ‖ h)                  ctrl, Keccak
t̂ ← ByteDecode12(ek); μ ← Decompress1(m)      Codec
 ← SampleNTT(ρ); y, e1, e2 ← CBD(r)          Sampler, Keccak
ŷ ← NTT(y); s ← e2 + μ                        NTT
û ← Âᵀ∘ŷ; v̂ ← t̂ᵀ∘ŷ                            PolyArith
u ← NTT⁻¹(û) + e1; v ← NTT⁻¹(v̂) + s           NTT
c1 ← Compress_du(u); c2 ← Compress_dv(v)      Codec
```

Decaps decrypts c on the same tables and re-encrypts, and a KemSelect table compares c with c′ and returns K′ or the
implicit-rejection key K̄. KeyGen samples from G(d ‖ k) and encodes t̂ and ŝ in the Codec. The host program trades
32-bit words with the ctrl table over one service bus, and the tables pass coefficients, words and Keccak lanes over
internal buses, each value emitted once by its producer and once by its consumer.

## Outside the proof

- **Randomness.** The seeds d and z and the encapsulation randomness m enter as witness, and the proof does not show
  they were drawn uniformly. Drawing them is the job of the external algorithms, ML-KEM.KeyGen (Algorithm 19) and
  ML-KEM.Encaps (Algorithm 20), which stay with the caller.

## How it is tested

The tests run KeyGen, Encaps and Decaps at all three parameter sets, match their keys and ciphertexts against
RustCrypto's `ml-kem` crate, and prove the calls. Forgery tests swap one input at a time, such as a key derived for another message, a
foreign dk_PKE or a different ρ, and check that each swap breaks the word or lane bus that guards it. A non-canonical
encapsulation key breaks the Codec's modulus check, and a decapsulation key with a foreign hash breaks the word bus.
NIST's KAT files are not part of the suite.

## Examples

Two programs model the two sides of a key exchange at the level `HEKATE_LEVEL` selects:

- [`mlkem_sender.rs`][sender] proves Encaps of K to the recipient's ek and an AES-256-CTR encryption of a message under
  K, joined by the `key_link` bus that carries K into the AES key. It publishes H(ek), c, the 12-byte nonce of each
  counter block, the message and its ciphertext; ek, m and K stay private. The program also checks that the recipient
  recovers K from c and decrypts the message.
- [`mlkem_receiver.rs`][receiver] proves KeyGen from the seed d chained into Decaps of a sender's c, through the
  `dk_link` bus. It publishes H(ek), c and valid; d, z and dk stay private. The program also checks that the K it
  decapsulates matches the sender's.

From a hekate checkout:

```bash
HEKATE_LEVEL=768 just example mlkem_sender public     # HEKATE_MESSAGE=<text> sets the payload
HEKATE_LEVEL=768 just example mlkem_receiver public   # 512 | 768 | 1024
```

`public` builds the variable-time `table-math` prover that the README's measurements use. Without it the recipe builds
the constant-time prover, the one to use when the witness is private.

The programs below are the two examples at ML-KEM-768, as two crates, `mlkem-sender` and `mlkem-receiver`. Each takes
three files: `program.rs` declares the host table, the tables it attaches and their buses, and publishes the
statement's values; `trace.rs` runs the host side and fills the tables; `main.rs` proves and verifies.

Both links are permutation buses without clocks, which Hekate accepts only with a clock waiver: the written reason
their keys are unique, here fixed indices that each side emits once. The engine checks that the
reason is present; review checks the argument.

Start from the crates in [Installation](https://oumuamua.dev/hekate/docs/getting-started/installation) and add `hekate-pqc`, `ml-kem` and
`rand`, plus `hekate-aes` for the sender.

### Sender

```rust src/program.rs:1
use hekate_aes::{Aes256Chiplet, AesRound256Air, host_key_selector_shape, host_selector_shape};
use hekate_core::config::Config;
use hekate_core::errors;
use hekate_math::{Block128, TowerField};
use hekate_pqc::mlkem::{self, MLKEM_DATA_BUS_ID, MlKemChiplet, MlKemParams};
use hekate_program::circuit::{Circuit, CircuitProgram, Col, ColRange};
use hekate_program::constraint::builder::ConstraintSystem;
use hekate_program::permutation::{PermutationCheckSpec, Source};
use hekate_program::{FixedShape, define_columns};

pub const PARAMS: MlKemParams = MlKemParams::ML_KEM_768;
pub const SEED_WORDS: usize = 8;
pub const BLOCK: usize = 16;
pub const NONCE_BYTES: usize = 12;

const KEY_LINK_BUS_ID: &str = "key_link";
const KEY_LINK_WAIVER: &str = "see mlkem-sender/src/program.rs: \
     link_idx is a fixed column over K's host rows and a constant \
     on the first AES row; each index is emitted once on each side";

pub type F = Block128;
pub type Published = (usize, usize, F);

define_columns! {
    pub SenderColumns {
        WORD: B32,
        SEL: Bit,

        LINK_IDX: B16,
        LINK_SEL: Bit,

        KEY: [B8; 32],
        KEY_SELECTOR: Bit,
        DATA: [B8; 16],
        SELECTOR: Bit,

        PACKED: [B32; 8],
        FIRST: Bit,

        CARRY: Bit,

        COUNTER: [B8; 16],
        PLAINTEXT: [B8; 16],
        CIPHERTEXT: [B8; 16],
    }
}

pub struct Layout {
    words: usize,
    pub h: usize,
    pub key: usize,
    pub c: usize,
    pub c_words: usize,
}

impl Layout {
    pub fn new(params: MlKemParams) -> Self {
        let m = 96 * params.k() + SEED_WORDS;
        let h = m + SEED_WORDS;
        let key = h + SEED_WORDS;
        let c = key + SEED_WORDS;
        let c_words = 8 * (params.du() as usize * params.k() + params.dv() as usize);

        Self {
            words: c + c_words,
            h,
            key,
            c,
            c_words,
        }
    }
}

pub fn aes_chiplet(blocks: usize, config: &Config) -> errors::Result<Aes256Chiplet> {
    let floor = config.min_table_rows();

    Aes256Chiplet::new(
        (blocks * AesRound256Air::BLOCK_ROWS)
            .next_power_of_two()
            .max(floor),
        (blocks * AesRound256Air::ACTIVE_ROWS)
            .next_power_of_two()
            .max(floor),
        blocks,
    )
}

pub fn build(
    pipeline: &MlKemChiplet<F>,
    aes: &Aes256Chiplet,
    layout: &Layout,
    message_len: usize,
    rows: usize,
    statement: &[Published],
) -> errors::Result<CircuitProgram<F>> {
    let blocks = message_len.div_ceil(BLOCK);

    let mut cx = Circuit::<F>::new("MlKemSender", rows)?;

    let cols = cx.schema(&SenderColumns::build_layout());

    let at = |col: usize| cols.at(col);

    let word = at(SenderColumns::WORD);
    let key: [Col; 32] = core::array::from_fn(|j| at(SenderColumns::KEY + j));
    let data: [Col; BLOCK] = core::array::from_fn(|j| at(SenderColumns::DATA + j));

    fix_shapes(&mut cx, cols, layout, blocks);

    cx.call(&mlkem::service(), &[word], at(SenderColumns::SEL))?;
    cx.call(
        &AesRound256Air::link_service(),
        &data,
        at(SenderColumns::SELECTOR),
    )?;
    cx.call(
        &AesRound256Air::key_service(),
        &key,
        at(SenderColumns::KEY_SELECTOR),
    )?;

    let link_idx = Source::Column(at(SenderColumns::LINK_IDX).index());

    cx.bus(
        KEY_LINK_BUS_ID,
        key_link(link_idx, word, at(SenderColumns::LINK_SEL)),
    );

    for j in 0..SEED_WORDS {
        let packed = at(SenderColumns::PACKED + j);

        cx.bus(
            KEY_LINK_BUS_ID,
            key_link(Source::Const(j as u128), packed, at(SenderColumns::FIRST)),
        );
    }

    add_constraints(cx.cs(), cols);

    for &(c, row, _) in statement {
        cx.publish(at(c), row);
    }

    let pad = (BLOCK - message_len % BLOCK) % BLOCK;

    for j in BLOCK - pad..BLOCK {
        cx.boundary(at(SenderColumns::PLAINTEXT + j), 2 * blocks - 1, F::ZERO);
    }

    cx.attach_namespaced("mlkem", pipeline.defs()?, &[MLKEM_DATA_BUS_ID])?;
    cx.attach_namespaced("aes256", aes.defs()?, &Aes256Chiplet::EXTERNAL_BUS_IDS)?;

    cx.compile()
}

fn fix_shapes(cx: &mut Circuit<F>, cols: ColRange, layout: &Layout, blocks: usize) {
    let sel = FixedShape::Cadence {
        stride: 1,
        count: layout.words,
        origin: 0,
        values: vec![F::ONE],
    };
    let link_idx = (1..SEED_WORDS).map(|j| (layout.key + j, F::from(j as u32)));
    let link_sel = (0..SEED_WORDS).map(|j| (layout.key + j, F::ONE));
    let carry = FixedShape::Cadence {
        stride: 1,
        count: 2 * blocks - 1,
        origin: 0,
        values: vec![F::ONE],
    };

    cx.fix(cols.at(SenderColumns::SEL), sel);
    cx.fix(
        cols.at(SenderColumns::LINK_IDX),
        FixedShape::Sparse(link_idx.collect()),
    );
    cx.fix(
        cols.at(SenderColumns::LINK_SEL),
        FixedShape::Sparse(link_sel.collect()),
    );
    cx.fix(
        cols.at(SenderColumns::SELECTOR),
        host_selector_shape(2, blocks),
    );
    cx.fix(
        cols.at(SenderColumns::KEY_SELECTOR),
        host_key_selector_shape(2, blocks),
    );
    cx.fix(
        cols.at(SenderColumns::FIRST),
        FixedShape::Sparse(vec![(0, F::ONE)]),
    );
    cx.fix(cols.at(SenderColumns::CARRY), carry);

    for i in 0..BLOCK - NONCE_BYTES {
        let index = (0..blocks).map(|b| (2 * b, F::from((b as u32).to_be_bytes()[i])));

        cx.fix(
            cols.at(SenderColumns::COUNTER + NONCE_BYTES + i),
            FixedShape::Sparse(index.collect()),
        );
    }
}

fn add_constraints(cs: &ConstraintSystem<F>, cols: ColRange) {
    let col = |c: usize| cs.col(cols.at(c).index());
    let next = |c: usize| cs.next(cols.at(c).index());

    let one = cs.one();
    let carry = col(SenderColumns::CARRY);
    let selector = col(SenderColumns::SELECTOR);
    let key_selector = col(SenderColumns::KEY_SELECTOR);
    let output = selector + key_selector;

    cs.constrain((one + col(SenderColumns::SEL)) * col(SenderColumns::WORD));

    for j in 0..SEED_WORDS {
        let packed = (0..4).fold(col(SenderColumns::PACKED + j), |acc, i| {
            let shift = cs.constant(F::from(1u128 << (8 * i)));

            acc + shift * col(SenderColumns::KEY + 4 * j + i)
        });

        cs.constrain(packed);
    }

    for j in 0..32 {
        let key = SenderColumns::KEY + j;

        cs.constrain(carry * (next(key) + col(key)));
        cs.constrain((one + selector) * col(key));
    }

    for j in 0..BLOCK {
        let data = col(SenderColumns::DATA + j);
        let counter = col(SenderColumns::COUNTER + j);
        let plaintext = col(SenderColumns::PLAINTEXT + j);
        let ciphertext = col(SenderColumns::CIPHERTEXT + j);

        let whitened = data + counter + col(SenderColumns::KEY + j);
        let encrypted = ciphertext + plaintext + data;

        cs.constrain(key_selector * whitened);
        cs.constrain(output * encrypted);

        cs.constrain((one + selector) * data);
        cs.constrain((one + key_selector) * counter);
        cs.constrain((one + output) * plaintext);
        cs.constrain((one + output) * ciphertext);
    }
}

fn key_link(idx: Source, word: Col, selector: Col) -> PermutationCheckSpec {
    PermutationCheckSpec::new(
        vec![
            (idx, b"kappa_key_idx" as &[u8]),
            (Source::Column(word.index()), b"kappa_key_word" as &[u8]),
        ],
        Some(selector.index()),
    )
    .with_clock_waiver(KEY_LINK_WAIVER)
}
```

```rust src/trace.rs:1
use crate::program::{BLOCK, F, Layout, NONCE_BYTES, Published, SEED_WORDS, SenderColumns};
use hekate_aes::Aes256Chiplet;
use hekate_aes::trace::{Aes256Call, aes256_encrypt_block, expand_key_256};
use hekate_core::errors;
use hekate_core::trace::{ColumnTrace, TraceBuilder};
use hekate_math::{Bit, Block8, Block16, Block32, TowerField};
use std::error::Error;

pub struct Ctr {
    pub counters: Vec<[u8; BLOCK]>,
    pub keystream: Vec<[u8; BLOCK]>,
    pub plaintext: Vec<u8>,
    pub ciphertext: Vec<u8>,
}

impl Ctr {
    pub fn encrypt(key: &[u8; 32], nonce: &[u8; NONCE_BYTES], message: &[u8]) -> Self {
        let round_keys = expand_key_256(key);
        let blocks = message.len().div_ceil(BLOCK);

        let counters: Vec<[u8; BLOCK]> = (0..blocks as u32)
            .map(|b| {
                let mut block = [0u8; BLOCK];
                block[..NONCE_BYTES].copy_from_slice(nonce);
                block[NONCE_BYTES..].copy_from_slice(&b.to_be_bytes());

                block
            })
            .collect();

        let keystream: Vec<[u8; BLOCK]> = counters
            .iter()
            .map(|counter| aes256_encrypt_block(&round_keys, counter))
            .collect();

        let mut plaintext = message.to_vec();
        plaintext.resize(blocks * BLOCK, 0);

        let ciphertext = plaintext
            .iter()
            .zip(keystream.iter().flatten())
            .map(|(p, k)| p ^ k)
            .collect();

        Self {
            counters,
            keystream,
            plaintext,
            ciphertext,
        }
    }

    pub fn blocks(&self) -> usize {
        self.counters.len()
    }
}

pub fn statement(
    layout: &Layout,
    words: &[u32],
    ctr: &Ctr,
    message_len: usize,
) -> Result<Vec<Published>, Box<dyn Error>> {
    let h = layout.h..layout.h + SEED_WORDS;
    let c = layout.c..layout.c + layout.c_words;

    let mut statement = Vec::new();

    for row in h.chain(c) {
        let word = words
            .get(row)
            .ok_or("the host words end before the published rows")?;

        statement.push((SenderColumns::WORD, row, F::from(u128::from(*word))));
    }

    for (b, counter) in ctr.counters.iter().enumerate() {
        for (j, &byte) in counter.iter().take(NONCE_BYTES).enumerate() {
            statement.push((SenderColumns::COUNTER + j, 2 * b, F::from(byte)));
        }
    }

    let bytes = ctr.plaintext.iter().zip(&ctr.ciphertext).take(message_len);

    for (i, (&plain, &cipher)) in bytes.enumerate() {
        let (row, j) = (2 * (i / BLOCK) + 1, i % BLOCK);

        statement.push((SenderColumns::PLAINTEXT + j, row, F::from(plain)));
        statement.push((SenderColumns::CIPHERTEXT + j, row, F::from(cipher)));
    }

    Ok(statement)
}

pub fn host_trace(
    words: &[u32],
    layout: &Layout,
    key: &[u8; 32],
    ctr: &Ctr,
    rows: usize,
) -> errors::Result<ColumnTrace> {
    let mut tb = TraceBuilder::new_secret(
        &SenderColumns::build_layout(),
        rows.trailing_zeros() as usize,
    )?;

    for (r, &word) in words.iter().enumerate() {
        tb.set_b32(SenderColumns::WORD, r, Block32::from(word))?;
        tb.set_bit(SenderColumns::SEL, r, Bit::ONE)?;
    }

    for j in 0..SEED_WORDS {
        tb.set_b16(SenderColumns::LINK_IDX, layout.key + j, Block16(j as u16))?;
        tb.set_bit(SenderColumns::LINK_SEL, layout.key + j, Bit::ONE)?;
    }

    let aes_rows = 2 * ctr.blocks();
    let (packed, _) = key.as_chunks::<4>();

    for r in 0..aes_rows {
        for (j, &byte) in key.iter().enumerate() {
            tb.set_b8(SenderColumns::KEY + j, r, Block8(byte))?;
        }

        for (j, &bytes) in packed.iter().enumerate() {
            let word = Block32::from(u32::from_le_bytes(bytes));

            tb.set_b32(SenderColumns::PACKED + j, r, word)?;
        }

        tb.set_bit(SenderColumns::SELECTOR, r, Bit::ONE)?;
        tb.set_bit(
            SenderColumns::CARRY,
            r,
            Bit::from(u8::from(r + 1 < aes_rows)),
        )?;
    }

    tb.set_bit(SenderColumns::FIRST, 0, Bit::ONE)?;

    for (b, counter) in ctr.counters.iter().enumerate() {
        tb.set_bit(SenderColumns::KEY_SELECTOR, 2 * b, Bit::ONE)?;

        for (j, (&byte, &k)) in counter.iter().zip(key).enumerate() {
            tb.set_b8(SenderColumns::COUNTER + j, 2 * b, Block8(byte))?;
            tb.set_b8(SenderColumns::DATA + j, 2 * b, Block8(byte ^ k))?;
        }
    }

    let output = ctr
        .keystream
        .iter()
        .flatten()
        .zip(ctr.plaintext.iter())
        .zip(&ctr.ciphertext);

    for (i, ((&keystream, &plain), &cipher)) in output.enumerate() {
        let (row, j) = (2 * (i / BLOCK) + 1, i % BLOCK);

        tb.set_b8(SenderColumns::DATA + j, row, Block8(keystream))?;
        tb.set_b8(SenderColumns::PLAINTEXT + j, row, Block8(plain))?;
        tb.set_b8(SenderColumns::CIPHERTEXT + j, row, Block8(cipher))?;
    }

    Ok(tb.build())
}

pub fn aes_traces(
    aes: &Aes256Chiplet,
    key: &[u8; 32],
    ctr: &Ctr,
) -> errors::Result<Vec<ColumnTrace>> {
    let round_keys = expand_key_256(key);

    let calls: Vec<Aes256Call> = ctr
        .counters
        .iter()
        .map(|counter| Aes256Call {
            key: *key,
            plaintext: *counter,
            round_keys,
        })
        .collect();

    aes.generate_traces(&calls)
}
```

```rust src/main.rs:1
mod program;
mod trace;

use hekate_core::config::Config;
use hekate_crypto::DefaultHasher;
use hekate_crypto::transcript::Transcript;
use hekate_pqc::mlkem::{MlKemCall, MlKemChiplet, MlKemInput, MlKemOutput};
use hekate_program::digest::program_id;
use hekate_program::{ProgramInstance, ProgramWitness};
use hekate_prover_sys::prove;
use hekate_sdk::preflight;
use hekate_verifier::HekateVerifier;
use ml_kem::{Decapsulate, DecapsulationKey768, Generate, KeyExport};
use program::{F, Layout, NONCE_BYTES, PARAMS, aes_chiplet, build};
use rand::{TryRng, rngs::SysRng};
use std::error::Error;
use trace::{Ctr, aes_traces, host_trace, statement};

const LABEL: &[u8] = b"mlkem-sender";
const MESSAGE: &str =
    "Hekate ML-KEM sender: this message travels under AES-256-CTR with the encapsulated key.";

fn main() -> Result<(), Box<dyn Error>> {
    let recipient = DecapsulationKey768::try_generate_from_rng(&mut SysRng)?;
    let ek = recipient.encapsulation_key().to_bytes();
    let message = MESSAGE.as_bytes();

    let mut m = [0; 32];
    let mut nonce = [0; NONCE_BYTES];

    SysRng.try_fill_bytes(&mut m)?;
    SysRng.try_fill_bytes(&mut nonce)?;

    let pipeline = MlKemChiplet::<F>::new(PARAMS, &[MlKemCall::Encaps])?;
    let traced = pipeline.trace(&[MlKemInput::Encaps { ek: &ek, m: &m }])?;

    let [MlKemOutput::Encaps { key, c }] = traced.outputs.as_slice() else {
        return Err("the pipeline returns one Encaps output".into());
    };

    let config = Config::default();
    let ctr = Ctr::encrypt(key, &nonce, message);
    let aes = aes_chiplet(ctr.blocks(), &config)?;

    let layout = Layout::new(PARAMS);
    let rows = traced.words.len().max(2 * ctr.blocks()).next_power_of_two();
    let statement = statement(&layout, &traced.words, &ctr, message.len())?;
    let program = build(&pipeline, &aes, &layout, message.len(), rows, &statement)?;

    let inputs = statement.iter().map(|&(_, _, value)| value);
    let instance = ProgramInstance::new(rows, inputs.collect());

    let mut chiplets = traced.traces;
    chiplets.extend(aes_traces(&aes, key, &ctr)?);

    let host = host_trace(&traced.words, &layout, key, &ctr, rows)?;
    let witness = ProgramWitness::new(host).with_chiplets(chiplets);

    let report = preflight(&program, &instance, &witness)?;

    if !report.is_clean() {
        eprintln!("{report}");

        return Err("the trace fails preflight".into());
    }

    let mut seed = [0; 32];
    SysRng.try_fill_bytes(&mut seed)?;

    let proof = prove(LABEL, &program, &instance, &witness, &config, seed, None)?;

    let id = program_id::<F, _>(&program)?;
    let mut transcript = Transcript::<DefaultHasher>::new(LABEL);

    let accepted = HekateVerifier::<F, DefaultHasher>::verify(
        &id,
        &program,
        &instance,
        &proof,
        &mut transcript,
        &config,
    )?;

    if !accepted {
        return Err("the verifier rejects the encapsulation and the encryption".into());
    }

    let received = recipient.decapsulate_slice(c)?;
    let decrypted = Ctr::encrypt(received.as_ref(), &nonce, &ctr.ciphertext[..message.len()]);

    if *key != received || !decrypted.ciphertext.starts_with(message) {
        return Err("the recipient misses K or the message".into());
    }

    println!("c          {} bytes, public", c.len());
    println!("recipient  recovers K and decrypts the message");
    println!("the verifier accepts the encapsulation and the encryption");

    Ok(())
}
```

### Receiver

```rust src/program.rs:1
use hekate_core::errors;
use hekate_math::{Block128, TowerField};
use hekate_pqc::mlkem::{self, MLKEM_DATA_BUS_ID, MlKemChiplet, MlKemParams};
use hekate_program::circuit::{Circuit, CircuitProgram, Col};
use hekate_program::permutation::{PermutationCheckSpec, Source};
use hekate_program::{FixedShape, define_columns};

pub const PARAMS: MlKemParams = MlKemParams::ML_KEM_768;

const SEED_WORDS: usize = 8;

const LINK_BUS_ID: &str = "dk_link";
const LINK_WAIVER: &str = "see mlkem-receiver/src/program.rs: \
     link_idx is a fixed column; every index is emitted once \
     on a KeyGen output row and once on a Decaps input row";

pub type F = Block128;

define_columns! {
    pub ReceiverColumns {
        WORD: B32,
        SEL: Bit,

        LINK_IDX: B16,
        LINK_SEL: Bit,
    }
}

pub struct Layout {
    words: usize,
    h: usize,
    c: usize,
    c_words: usize,
    valid: usize,
    pub links: Vec<(usize, u16)>,
}

impl Layout {
    pub fn new(params: MlKemParams) -> Self {
        let w = 96 * params.k();
        let c_words = 8 * (params.du() as usize * params.k() + params.dv() as usize);

        let ek_out = SEED_WORDS;
        let rho_out = ek_out + w;
        let h_out = rho_out + SEED_WORDS;
        let dk_out = h_out + SEED_WORDS;

        let dk_in = dk_out + w;
        let ek_in = dk_in + w;
        let rho_in = ek_in + w;
        let h_in = rho_in + SEED_WORDS;
        let c = h_in + 2 * SEED_WORDS;
        let valid = c + c_words + SEED_WORDS;

        let mut links = Vec::with_capacity(4 * w + 4 * SEED_WORDS);

        for (keygen, decaps, first, len) in [
            (dk_out, dk_in, 0, w),
            (ek_out, ek_in, w, w),
            (rho_out, rho_in, 2 * w, SEED_WORDS),
            (h_out, h_in, 2 * w + SEED_WORDS, SEED_WORDS),
        ] {
            for j in 0..len {
                let idx = (first + j) as u16;

                links.push((keygen + j, idx));
                links.push((decaps + j, idx));
            }
        }

        links.sort_unstable();

        Self {
            words: valid + 1,
            h: h_out,
            c,
            c_words,
            valid,
            links,
        }
    }

    pub fn published(&self) -> Vec<usize> {
        (self.h..self.h + SEED_WORDS)
            .chain(self.c..self.c + self.c_words)
            .chain([self.valid])
            .collect()
    }
}

pub fn build(
    pipeline: &MlKemChiplet<F>,
    layout: &Layout,
    rows: usize,
) -> errors::Result<CircuitProgram<F>> {
    let mut cx = Circuit::<F>::new("MlKemReceiver", rows)?;

    let cols = cx.schema(&ReceiverColumns::build_layout());

    let word = cols.at(ReceiverColumns::WORD);
    let sel = cols.at(ReceiverColumns::SEL);
    let link_idx = cols.at(ReceiverColumns::LINK_IDX);
    let link_sel = cols.at(ReceiverColumns::LINK_SEL);

    let indices = layout
        .links
        .iter()
        .map(|&(row, idx)| (row, F::from(u32::from(idx))))
        .filter(|&(_, v)| v != F::ZERO);

    cx.fix(
        sel,
        FixedShape::Cadence {
            stride: 1,
            count: layout.words,
            origin: 0,
            values: vec![F::ONE],
        },
    );
    cx.fix(link_idx, FixedShape::Sparse(indices.collect()));
    cx.fix(
        link_sel,
        FixedShape::Sparse(layout.links.iter().map(|&(row, _)| (row, F::ONE)).collect()),
    );

    let cs = cx.cs();
    let col = |c: Col| cs.col(c.index());

    cs.constrain((cs.one() + col(sel)) * col(word));

    cx.call(&mlkem::service(), &[word], sel)?;

    cx.bus(
        LINK_BUS_ID,
        PermutationCheckSpec::new(
            vec![
                (Source::Column(link_idx.index()), b"kappa_link_idx" as &[u8]),
                (Source::Column(word.index()), b"kappa_link_word" as &[u8]),
            ],
            Some(link_sel.index()),
        )
        .with_clock_waiver(LINK_WAIVER),
    );

    for row in layout.published() {
        cx.publish(word, row);
    }

    cx.attach_namespaced("mlkem", pipeline.defs()?, &[MLKEM_DATA_BUS_ID])?;

    cx.compile()
}
```

```rust src/trace.rs:1
use crate::program::{F, Layout, ReceiverColumns};
use hekate_core::errors;
use hekate_core::trace::{ColumnTrace, TraceBuilder};
use hekate_math::{Bit, Block16, Block32, TowerField};
use hekate_pqc::mlkem::{MlKemCall, MlKemChiplet, MlKemInput, MlKemOutput, MlKemParams};
use std::error::Error;

pub struct KeyPair {
    pub ek: Vec<u8>,
    pub dk: Vec<u8>,
}

impl KeyPair {
    pub fn generate(
        params: MlKemParams,
        d: &[u8; 32],
        z: &[u8; 32],
    ) -> Result<Self, Box<dyn Error>> {
        let keygen = MlKemChiplet::<F>::new(params, &[MlKemCall::KeyGen])?;
        let generated = keygen.trace(&[MlKemInput::KeyGen { d }])?;

        let [MlKemOutput::KeyGen { ek, dk_pke, h }] = generated.outputs.as_slice() else {
            return Err("the KeyGen pipeline returns one KeyGen output".into());
        };

        Ok(Self {
            ek: ek.clone(),
            dk: [dk_pke.as_slice(), ek, h, z].concat(),
        })
    }
}

pub fn host_trace(words: &[u32], layout: &Layout, rows: usize) -> errors::Result<ColumnTrace> {
    let mut tb = TraceBuilder::new_secret(
        &ReceiverColumns::build_layout(),
        rows.trailing_zeros() as usize,
    )?;

    for (r, &word) in words.iter().enumerate() {
        tb.set_b32(ReceiverColumns::WORD, r, Block32::from(word))?;
        tb.set_bit(ReceiverColumns::SEL, r, Bit::ONE)?;
    }

    for &(row, idx) in &layout.links {
        tb.set_b16(ReceiverColumns::LINK_IDX, row, Block16(idx))?;
        tb.set_bit(ReceiverColumns::LINK_SEL, row, Bit::ONE)?;
    }

    Ok(tb.build())
}
```

```rust src/main.rs:1
mod program;
mod trace;

use hekate_core::config::Config;
use hekate_crypto::DefaultHasher;
use hekate_crypto::transcript::Transcript;
use hekate_pqc::mlkem::{MlKemCall, MlKemChiplet, MlKemInput, MlKemOutput};
use hekate_program::digest::program_id;
use hekate_program::{ProgramInstance, ProgramWitness};
use hekate_prover_sys::prove;
use hekate_sdk::preflight;
use hekate_verifier::HekateVerifier;
use ml_kem::{Encapsulate, EncapsulationKey768, TryKeyInit};
use program::{F, Layout, PARAMS, build};
use rand::rngs::{StdRng, SysRng};
use rand::{SeedableRng, TryRng};
use std::error::Error;
use trace::{KeyPair, host_trace};

const LABEL: &[u8] = b"mlkem-receiver";

fn main() -> Result<(), Box<dyn Error>> {
    let mut d = [0; 32];
    let mut z = [0; 32];

    SysRng.try_fill_bytes(&mut d)?;
    SysRng.try_fill_bytes(&mut z)?;

    let keys = KeyPair::generate(PARAMS, &d, &z)?;

    let sender = EncapsulationKey768::new_from_slice(&keys.ek)?;
    let (c, sent) = sender.encapsulate_with_rng(&mut StdRng::try_from_rng(&mut SysRng)?);

    let pipeline = MlKemChiplet::<F>::new(PARAMS, &[MlKemCall::KeyGen, MlKemCall::Decaps])?;
    let traced = pipeline.trace(&[
        MlKemInput::KeyGen { d: &d },
        MlKemInput::Decaps {
            dk: &keys.dk,
            c: &c,
        },
    ])?;

    let [_, MlKemOutput::Decaps { key, valid }] = traced.outputs.as_slice() else {
        return Err("the pipeline returns KeyGen then Decaps outputs".into());
    };

    if *key != sent {
        return Err("the decapsulated key differs from the sender's".into());
    }

    let config = Config::default();
    let layout = Layout::new(PARAMS);
    let rows = traced.words.len().next_power_of_two();
    let program = build(&pipeline, &layout, rows)?;

    let public_inputs = layout
        .published()
        .iter()
        .map(|&row| traced.words.get(row).map(|&word| F::from(u128::from(word))))
        .collect::<Option<Vec<F>>>()
        .ok_or("the host words end before the published rows")?;
    let instance = ProgramInstance::new(rows, public_inputs);

    let host = host_trace(&traced.words, &layout, rows)?;
    let witness = ProgramWitness::new(host).with_chiplets(traced.traces);

    let report = preflight(&program, &instance, &witness)?;

    if !report.is_clean() {
        eprintln!("{report}");

        return Err("the trace fails preflight".into());
    }

    let mut seed = [0; 32];
    SysRng.try_fill_bytes(&mut seed)?;

    let proof = prove(LABEL, &program, &instance, &witness, &config, seed, None)?;

    let id = program_id::<F, _>(&program)?;
    let mut transcript = Transcript::<DefaultHasher>::new(LABEL);

    let accepted = HekateVerifier::<F, DefaultHasher>::verify(
        &id,
        &program,
        &instance,
        &proof,
        &mut transcript,
        &config,
    )?;

    if !accepted {
        return Err("the verifier rejects the key generation and the decapsulation".into());
    }

    println!("c          {} bytes, public", c.len());
    println!("valid      {valid}, K matches the sender's");
    println!("the verifier accepts the key generation and the decapsulation");

    Ok(())
}
```

## Performance

Zero-knowledge proofs, measured on Apple M3 Max, 16 cores. Build: `--release`, `std parallel blake3 table-math public`. Best of three on an otherwise idle machine. The [Benchmarks](https://oumuamua.dev/hekate/benchmarks) page has the base-protocol figures and the command for each run.

| Workload | Input | Proving | Verify | Proof size | Peak memory |
|:---|:---|:---|:---|:---|:---|
| ML-KEM-768 sender | FIPS 203 | 413 ms | 11.7 ms | 1,541 KiB | 190 MiB |
| ML-KEM-768 receiver | FIPS 203 | 558 ms | 13.9 ms | 1,884 KiB | 269 MiB |

## Security

The crate has not been audited and may contain bugs and security flaws, as its README states. Hekate proves 100 bits
of soundness at every parameter set, below the security level of each, which makes the proof the weaker link; the
proven operations still run the full FIPS 203 parameter sets.

## License

`hekate-pqc` is AGPL-3.0-only. Commercial licenses are available from Oumuamua Labs at
[info@oumuamua.dev](mailto:info@oumuamua.dev). The [README's Usage section][usage] builds an ML-KEM-768 encapsulation
host program, and the API reference is on docs.rs.

[sender]: https://github.com/oumuamua-labs/hekate/blob/main/hekate/examples/mlkem_sender.rs
[receiver]: https://github.com/oumuamua-labs/hekate/blob/main/hekate/examples/mlkem_receiver.rs
[usage]: https://github.com/oumuamua-labs/hekate/blob/main/hekate-pqc/README.md#usage