chacha20_fused.ea

// chacha20_fused.ea — Fused ChaCha20 encrypt + plaintext statistics kernel
// 4-block interleaved ILP (identical round loop to chacha20.ea) plus
// per-plaintext-byte sum/count/min/max in a single pass.

// Rotation helpers — identical to chacha20.ea
func rotl16(v: i32x4) -> i32x4 {
    return (v .<< splat(16)) .| ((v .>> splat(16)) .& splat(65535))
}

func rotl12(v: i32x4) -> i32x4 {
    return (v .<< splat(12)) .| ((v .>> splat(20)) .& splat(4095))
}

func rotl8(v: i32x4) -> i32x4 {
    return (v .<< splat(8)) .| ((v .>> splat(24)) .& splat(255))
}

func rotl7(v: i32x4) -> i32x4 {
    return (v .<< splat(7)) .| ((v .>> splat(25)) .& splat(127))
}

// Internal single-block function for tail handling
func chacha20_block(key: *i32, nonce: *i32, counter: i32, out: *mut i32) {
    let row0_init: i32x4 = [1634760805, 857760878, 2036477234, 1797285236]i32x4
    let row1_init: i32x4 = load(key, 0)
    let row2_init: i32x4 = load(key, 4)
    let n0: i32 = nonce[0]
    let n1: i32 = nonce[1]
    let n2: i32 = nonce[2]
    let row3_init: i32x4 = [counter, n0, n1, n2]i32x4
    let mut r0: i32x4 = row0_init
    let mut r1: i32x4 = row1_init
    let mut r2: i32x4 = row2_init
    let mut r3: i32x4 = row3_init
    let mut round: i32 = 0
    while round < 10 {
        r0 = r0 .+ r1
        r3 = rotl16(r3 .^ r0)
        r2 = r2 .+ r3
        r1 = rotl12(r1 .^ r2)
        r0 = r0 .+ r1
        r3 = rotl8(r3 .^ r0)
        r2 = r2 .+ r3
        r1 = rotl7(r1 .^ r2)
        r1 = shuffle(r1, [1, 2, 3, 0])
        r2 = shuffle(r2, [2, 3, 0, 1])
        r3 = shuffle(r3, [3, 0, 1, 2])
        r0 = r0 .+ r1
        r3 = rotl16(r3 .^ r0)
        r2 = r2 .+ r3
        r1 = rotl12(r1 .^ r2)
        r0 = r0 .+ r1
        r3 = rotl8(r3 .^ r0)
        r2 = r2 .+ r3
        r1 = rotl7(r1 .^ r2)
        r1 = shuffle(r1, [3, 0, 1, 2])
        r2 = shuffle(r2, [2, 3, 0, 1])
        r3 = shuffle(r3, [1, 2, 3, 0])
        round = round + 1
    }
    store(out, 0, r0 .+ row0_init)
    store(out, 4, r1 .+ row1_init)
    store(out, 8, r2 .+ row2_init)
    store(out, 12, r3 .+ row3_init)
}

// Sum all 16 bytes of a u8x16 vector, returning i32
func sum_u8x16(chunk: u8x16) -> i32 {
    let w0: i32x4 = widen_u8_i32x4(chunk)
    let s1: u8x16 = shuffle(chunk, [4,5,6,7, 0,1,2,3, 8,9,10,11, 12,13,14,15])
    let w1: i32x4 = widen_u8_i32x4(s1)
    let s2: u8x16 = shuffle(chunk, [8,9,10,11, 0,1,2,3, 4,5,6,7, 12,13,14,15])
    let w2: i32x4 = widen_u8_i32x4(s2)
    let s3: u8x16 = shuffle(chunk, [12,13,14,15, 0,1,2,3, 4,5,6,7, 8,9,10,11])
    let w3: i32x4 = widen_u8_i32x4(s3)
    return reduce_add(w0) + reduce_add(w1) + reduce_add(w2) + reduce_add(w3)
}

// Fused ChaCha20 encrypt + statistics — 4-block ILP hot loop
export func chacha20_encrypt_stats(
    key: *restrict i32, nonce: *restrict i32, counter: i32,
    plaintext: *restrict u8, ciphertext: *restrict mut u8, len: i32,
    ks_i32: *restrict mut i32, ks_u8: *restrict u8,
    pt_i32: *restrict i32, ct_i32: *restrict mut i32,
    out_sum: *mut i64, out_count: *mut i32,
    out_min: *mut u8, out_max: *mut u8
) {
    let row0_init: i32x4 = [1634760805, 857760878, 2036477234, 1797285236]i32x4
    let row1_init: i32x4 = load(key, 0)
    let row2_init: i32x4 = load(key, 4)
    let n0: i32 = nonce[0]
    let n1: i32 = nonce[1]
    let n2: i32 = nonce[2]
    let mut offset: i32 = 0
    let mut elem_off: i32 = 0
    let mut ctr: i32 = counter
    let mut running_sum: i64 = to_i64(0)
    let mut running_count: i32 = 0
    let mut running_min: u8 = 255
    let mut running_max: u8 = 0

    // ---- 4-block ILP hot loop (256 bytes per iteration) ----
    while offset + 256 <= len {
        let r3a_init: i32x4 = [ctr, n0, n1, n2]i32x4
        let r3b_init: i32x4 = [ctr + 1, n0, n1, n2]i32x4
        let r3c_init: i32x4 = [ctr + 2, n0, n1, n2]i32x4
        let r3d_init: i32x4 = [ctr + 3, n0, n1, n2]i32x4
        let mut a0: i32x4 = row0_init
        let mut a1: i32x4 = row1_init
        let mut a2: i32x4 = row2_init
        let mut a3: i32x4 = r3a_init
        let mut b0: i32x4 = row0_init
        let mut b1: i32x4 = row1_init
        let mut b2: i32x4 = row2_init
        let mut b3: i32x4 = r3b_init
        let mut c0: i32x4 = row0_init
        let mut c1: i32x4 = row1_init
        let mut c2: i32x4 = row2_init
        let mut c3: i32x4 = r3c_init
        let mut d0: i32x4 = row0_init
        let mut d1: i32x4 = row1_init
        let mut d2: i32x4 = row2_init
        let mut d3: i32x4 = r3d_init
        let mut round: i32 = 0
        while round < 10 {
            a0 = a0 .+ a1
            b0 = b0 .+ b1
            c0 = c0 .+ c1
            d0 = d0 .+ d1
            a3 = rotl16(a3 .^ a0)
            b3 = rotl16(b3 .^ b0)
            c3 = rotl16(c3 .^ c0)
            d3 = rotl16(d3 .^ d0)
            a2 = a2 .+ a3
            b2 = b2 .+ b3
            c2 = c2 .+ c3
            d2 = d2 .+ d3
            a1 = rotl12(a1 .^ a2)
            b1 = rotl12(b1 .^ b2)
            c1 = rotl12(c1 .^ c2)
            d1 = rotl12(d1 .^ d2)
            a0 = a0 .+ a1
            b0 = b0 .+ b1
            c0 = c0 .+ c1
            d0 = d0 .+ d1
            a3 = rotl8(a3 .^ a0)
            b3 = rotl8(b3 .^ b0)
            c3 = rotl8(c3 .^ c0)
            d3 = rotl8(d3 .^ d0)
            a2 = a2 .+ a3
            b2 = b2 .+ b3
            c2 = c2 .+ c3
            d2 = d2 .+ d3
            a1 = rotl7(a1 .^ a2)
            b1 = rotl7(b1 .^ b2)
            c1 = rotl7(c1 .^ c2)
            d1 = rotl7(d1 .^ d2)
            a1 = shuffle(a1, [1, 2, 3, 0])
            b1 = shuffle(b1, [1, 2, 3, 0])
            c1 = shuffle(c1, [1, 2, 3, 0])
            d1 = shuffle(d1, [1, 2, 3, 0])
            a2 = shuffle(a2, [2, 3, 0, 1])
            b2 = shuffle(b2, [2, 3, 0, 1])
            c2 = shuffle(c2, [2, 3, 0, 1])
            d2 = shuffle(d2, [2, 3, 0, 1])
            a3 = shuffle(a3, [3, 0, 1, 2])
            b3 = shuffle(b3, [3, 0, 1, 2])
            c3 = shuffle(c3, [3, 0, 1, 2])
            d3 = shuffle(d3, [3, 0, 1, 2])
            a0 = a0 .+ a1
            b0 = b0 .+ b1
            c0 = c0 .+ c1
            d0 = d0 .+ d1
            a3 = rotl16(a3 .^ a0)
            b3 = rotl16(b3 .^ b0)
            c3 = rotl16(c3 .^ c0)
            d3 = rotl16(d3 .^ d0)
            a2 = a2 .+ a3
            b2 = b2 .+ b3
            c2 = c2 .+ c3
            d2 = d2 .+ d3
            a1 = rotl12(a1 .^ a2)
            b1 = rotl12(b1 .^ b2)
            c1 = rotl12(c1 .^ c2)
            d1 = rotl12(d1 .^ d2)
            a0 = a0 .+ a1
            b0 = b0 .+ b1
            c0 = c0 .+ c1
            d0 = d0 .+ d1
            a3 = rotl8(a3 .^ a0)
            b3 = rotl8(b3 .^ b0)
            c3 = rotl8(c3 .^ c0)
            d3 = rotl8(d3 .^ d0)
            a2 = a2 .+ a3
            b2 = b2 .+ b3
            c2 = c2 .+ c3
            d2 = d2 .+ d3
            a1 = rotl7(a1 .^ a2)
            b1 = rotl7(b1 .^ b2)
            c1 = rotl7(c1 .^ c2)
            d1 = rotl7(d1 .^ d2)
            a1 = shuffle(a1, [3, 0, 1, 2])
            b1 = shuffle(b1, [3, 0, 1, 2])
            c1 = shuffle(c1, [3, 0, 1, 2])
            d1 = shuffle(d1, [3, 0, 1, 2])
            a2 = shuffle(a2, [2, 3, 0, 1])
            b2 = shuffle(b2, [2, 3, 0, 1])
            c2 = shuffle(c2, [2, 3, 0, 1])
            d2 = shuffle(d2, [2, 3, 0, 1])
            a3 = shuffle(a3, [1, 2, 3, 0])
            b3 = shuffle(b3, [1, 2, 3, 0])
            c3 = shuffle(c3, [1, 2, 3, 0])
            d3 = shuffle(d3, [1, 2, 3, 0])
            round = round + 1
        }

        // XOR plaintext with keystream and store ciphertext (via i32 pointers)
        store(ct_i32, elem_off, load(pt_i32, elem_off) .^ (a0 .+ row0_init))
        store(ct_i32, elem_off + 4, load(pt_i32, elem_off + 4) .^ (a1 .+ row1_init))
        store(ct_i32, elem_off + 8, load(pt_i32, elem_off + 8) .^ (a2 .+ row2_init))
        store(ct_i32, elem_off + 12, load(pt_i32, elem_off + 12) .^ (a3 .+ r3a_init))
        store(ct_i32, elem_off + 16, load(pt_i32, elem_off + 16) .^ (b0 .+ row0_init))
        store(ct_i32, elem_off + 20, load(pt_i32, elem_off + 20) .^ (b1 .+ row1_init))
        store(ct_i32, elem_off + 24, load(pt_i32, elem_off + 24) .^ (b2 .+ row2_init))
        store(ct_i32, elem_off + 28, load(pt_i32, elem_off + 28) .^ (b3 .+ r3b_init))
        store(ct_i32, elem_off + 32, load(pt_i32, elem_off + 32) .^ (c0 .+ row0_init))
        store(ct_i32, elem_off + 36, load(pt_i32, elem_off + 36) .^ (c1 .+ row1_init))
        store(ct_i32, elem_off + 40, load(pt_i32, elem_off + 40) .^ (c2 .+ row2_init))
        store(ct_i32, elem_off + 44, load(pt_i32, elem_off + 44) .^ (c3 .+ r3c_init))
        store(ct_i32, elem_off + 48, load(pt_i32, elem_off + 48) .^ (d0 .+ row0_init))
        store(ct_i32, elem_off + 52, load(pt_i32, elem_off + 52) .^ (d1 .+ row1_init))
        store(ct_i32, elem_off + 56, load(pt_i32, elem_off + 56) .^ (d2 .+ row2_init))
        store(ct_i32, elem_off + 60, load(pt_i32, elem_off + 60) .^ (d3 .+ r3d_init))

        // Stats on plaintext bytes (16 u8x16 chunks = 256 bytes)
        let mut si: i32 = 0
        while si < 16 {
            let pt_chunk: u8x16 = load(plaintext, offset + si * 16)
            let cmin: u8 = reduce_min(pt_chunk)
            let cmax: u8 = reduce_max(pt_chunk)
            if cmin < running_min {
                running_min = cmin
            }
            if cmax > running_max {
                running_max = cmax
            }
            running_sum = running_sum + to_i64(sum_u8x16(pt_chunk))
            si = si + 1
        }
        running_count = running_count + 256

        offset = offset + 256
        elem_off = elem_off + 64
        ctr = ctr + 4
    }

    // ---- Single-block loop for 64-byte remainders ----
    while offset + 64 <= len {
        let row3_init: i32x4 = [ctr, n0, n1, n2]i32x4
        let mut r0: i32x4 = row0_init
        let mut r1: i32x4 = row1_init
        let mut r2: i32x4 = row2_init
        let mut r3: i32x4 = row3_init
        let mut round: i32 = 0
        while round < 10 {
            r0 = r0 .+ r1
            r3 = rotl16(r3 .^ r0)
            r2 = r2 .+ r3
            r1 = rotl12(r1 .^ r2)
            r0 = r0 .+ r1
            r3 = rotl8(r3 .^ r0)
            r2 = r2 .+ r3
            r1 = rotl7(r1 .^ r2)
            r1 = shuffle(r1, [1, 2, 3, 0])
            r2 = shuffle(r2, [2, 3, 0, 1])
            r3 = shuffle(r3, [3, 0, 1, 2])
            r0 = r0 .+ r1
            r3 = rotl16(r3 .^ r0)
            r2 = r2 .+ r3
            r1 = rotl12(r1 .^ r2)
            r0 = r0 .+ r1
            r3 = rotl8(r3 .^ r0)
            r2 = r2 .+ r3
            r1 = rotl7(r1 .^ r2)
            r1 = shuffle(r1, [3, 0, 1, 2])
            r2 = shuffle(r2, [2, 3, 0, 1])
            r3 = shuffle(r3, [1, 2, 3, 0])
            round = round + 1
        }
        store(ct_i32, elem_off, load(pt_i32, elem_off) .^ (r0 .+ row0_init))
        store(ct_i32, elem_off + 4, load(pt_i32, elem_off + 4) .^ (r1 .+ row1_init))
        store(ct_i32, elem_off + 8, load(pt_i32, elem_off + 8) .^ (r2 .+ row2_init))
        store(ct_i32, elem_off + 12, load(pt_i32, elem_off + 12) .^ (r3 .+ row3_init))

        // Stats on 4 u8x16 chunks (64 bytes)
        let mut si: i32 = 0
        while si < 4 {
            let pt_chunk: u8x16 = load(plaintext, offset + si * 16)
            let cmin: u8 = reduce_min(pt_chunk)
            let cmax: u8 = reduce_max(pt_chunk)
            if cmin < running_min {
                running_min = cmin
            }
            if cmax > running_max {
                running_max = cmax
            }
            running_sum = running_sum + to_i64(sum_u8x16(pt_chunk))
            si = si + 1
        }
        running_count = running_count + 64

        offset = offset + 64
        elem_off = elem_off + 16
        ctr = ctr + 1
    }

    // ---- Sub-block tail (< 64 bytes) ----
    if offset < len {
        chacha20_block(key, nonce, ctr, ks_i32)
        let remaining: i32 = len - offset
        let mut chunk_off: i32 = 0

        // Full 16-byte chunks within the tail block
        while chunk_off + 16 <= remaining {
            let pt: u8x16 = load(plaintext, offset + chunk_off)
            let ks: u8x16 = load(ks_u8, chunk_off)
            store(ciphertext, offset + chunk_off, pt .^ ks)

            // Stats
            let cmin: u8 = reduce_min(pt)
            let cmax: u8 = reduce_max(pt)
            if cmin < running_min {
                running_min = cmin
            }
            if cmax > running_max {
                running_max = cmax
            }
            running_sum = running_sum + to_i64(sum_u8x16(pt))
            running_count = running_count + 16

            chunk_off = chunk_off + 16
        }

        // Byte-level tail (< 16 bytes)
        let tail_count: i32 = remaining - chunk_off
        if tail_count > 0 {
            let tail_off: i32 = offset + chunk_off
            let ks_tail: u8x16 = load_masked(ks_u8, chunk_off, tail_count)
            let pt_tail: u8x16 = load_masked(plaintext, tail_off, tail_count)
            store_masked(ciphertext, tail_off, pt_tail .^ ks_tail, tail_count)

            // Stats on tail bytes — byte-by-byte to avoid zeroed-lane corruption
            let mut ti: i32 = 0
            let mut tail_sum: i32 = 0
            while ti < tail_count {
                let b: u8 = plaintext[tail_off + ti]
                tail_sum = tail_sum + to_i32(b)
                if b < running_min {
                    running_min = b
                }
                if b > running_max {
                    running_max = b
                }
                ti = ti + 1
            }
            running_sum = running_sum + to_i64(tail_sum)
            running_count = running_count + tail_count
        }
    }

    // Write outputs
    out_sum[0] = running_sum
    out_count[0] = running_count
    out_min[0] = running_min
    out_max[0] = running_max
}