zisk-sm-binary 1.3.0-alpha

Binary operations state machine for the ZisK zkVM
require "std_lookup.pil"
require "operations.pil"
require "opids.pil"
require "binary_table.pil"

/*
    Operation table summary:
    ==========================

    List of 64-bit proven operations:
    ┼───────────────┼──────────┼──────────┼────────────────┼───────┼────────────────────┼─────────┼
    │     name      │    op    │   b_op   │ use_first_byte │ carry │ use_last_cout_as_c │ ZisK OP │
    ┼───────────────┼──────────┼──────────┼────────────────┼───────┼────────────────────┼─────────┼
    │ MINU          │   0x02   │   0x02   │                │   X   │                    │    X    │
    │ MIN           │   0x03   │   0x03   │                │   X   │                    │    X    │
    │ MAXU          │   0x04   │   0x04   │                │   X   │                    │    X    │
    │ MAX           │   0x05   │   0x05   │                │   X   │                    │    X    │
    │ LTU           │   0x06   │   0x06   │                │   X   │       X            │    X    │
    │ LT            │   0x07   │   0x07   │                │   X   │       X            │    X    │
    │ GT (*)        │   0x08   │   0x08   │                │   X   │       X            │         │
    │ EQ            │   0x09   │   0x09   │                │   X   │       X            │    X    │
    │ ADD           │   0x0a   │   0x0a   │                │   X   │                    │    X    │
    │ SUB           │   0x0b   │   0x0b   │                │   X   │                    │    X    │
    │ LEU           │   0x0c   │   0x0c   │                │   X   │       X            │    X    │
    │ LE            │   0x0d   │   0x0d   │                │   X   │       X            │    X    │
    │ AND           │   0x0e   │   0x0e   │                │       │                    │    X    │
    │ OR            │   0x0f   │   0x0f   │                │       │                    │    X    │
    │ XOR           │   0x10   │   0x10   │                │       │                    │    X    │
    │ LT_ABS_NP (*) │   0x50   │   0x50   │        X       │   X   │       X            │         │
    │ LT_ABS_PN (*) │   0x51   │   0x51   │        X       │   X   │       X            │         │
    │ BREV8         │   0x52   │   0x52   │                │       │                    │    X    │
    │ ANDN          │   0x53   │   0x53   │                │       │                    │    X    │
    │ ORN           │   0x54   │   0x54   │                │       │                    │    X    │
    │ XNOR          │   0x55   │   0x55   │                │       │                    │    X    │
    │ SH1ADD        │   0x56   │   0x56   │                │   X   │                    │    X    │
    │ SH2ADD        │   0x57   │   0x57   │                │   X   │                    │    X    │
    │ SH3ADD        │   0x58   │   0x58   │                │   X   │                    │    X    │
    ┼───────────────┼──────────┼──────────┼────────────────┼───────┼────────────────────┼─────────┼
    (*) This operation is used by the Arith component

    List of 32-bit proven operations:
    │───────────┼──────────┼──────────┼───────┼────────────────────┼─────────│
    │   name    │    op    │   b_op   │ carry │ use_last_cout_as_c │ ZisK OP │
    │───────────┼──────────┼──────────┼───────┼────────────────────┼─────────│
    │ MINU_W    │   0x12   │   0x02   │   X   │                    │    X    │
    │ MIN_W     │   0x13   │   0x03   │   X   │                    │    X    │
    │ MAXU_W    │   0x14   │   0x04   │   X   │                    │    X    │
    │ MAX_W     │   0x15   │   0x05   │   X   │                    │    X    │
    │ LTU_W     │   0x16   │   0x06   │   X   │       X            │    X    │
    │ LT_W      │   0x17   │   0x07   │   X   │       X            │    X    │
    │ GT_W      │   0x18   │   0x08   │   X   │       X            │         │
    │ EQ_W      │   0x19   │   0x09   │   X   │       X            │    X    │
    │ ADD_W     │   0x1a   │   0x0a   │   X   │                    │    X    │
    │ SUB_W     │   0x1b   │   0x0b   │   X   │                    │    X    │
    │ LEU_W     │   0x1c   │   0x0c   │   X   │       X            │    X    │
    │ LE_W      │   0x1d   │   0x0d   │   X   │       X            │    X    │
    │───────────┼──────────┼──────────┼───────┼────────────────────┼─────────│
    Note: op = b_op + 0x10*mode32

    mode32 is a free witness, so from any b_op this air can also prove b_op + 0x10. Every opcode
    in the first table therefore reserves opcode+0x10 as well, and the ops with no 32-bit variant
    (LT_ABS_*, BREV8, ANDN, ORN, XNOR, SHxADD) leave their shadows 0x60-0x68 permanently empty:
    handing one of those slots to a different operation would let this air satisfy it with the
    wrong semantics

    Lanes per row:
    ==============

    A "lane" is the full set of columns that proves ONE binary operation: b_op, the per-byte
    free_in_a / free_in_b / free_in_c / carry chains, and the mode32 / result_is_a / use_first_byte /
    c_is_signed flags. With `lanes_x_row` > 1 each row packs that many INDEPENDENT operations, side by
    side.

    Nothing in this air relates a row to its neighbour — every constraint and every lookup is local to
    one operation — so packing lanes is a pure widening: the constraints are the same ones, repeated
    per lane, and the air proves `lanes_x_row` operations per row instead of one. This is the same
    packing `BinaryAddHi` does with `lanes_x_row`, and the note there about preferring an ODD count
    (the operation-bus contributions amortize in pairs under the degree bound) applies here too.

    `padding_size` stays a single air value: it counts padding OPERATIONS, not rows, so it already
    speaks in lane units.
*/

airtemplate Binary(const int N = 2**21, const int RC = 2, const int bits = 64,
                   const int lanes_x_row = 1) {
    if (RC != 2 || bits != 64) {
        error(`Currently only RC=2 and bits=64 are supported, got RC=${RC}, bits=${bits}`);
    }

    assert(lanes_x_row > 0, "lanes_x_row must be greater than 0");

    // Default values
    const int BYTES = bits / 8;
    const int HALF_BYTES = BYTES / 2;
    const int BYTE_BASE = (1 << BYTES);
    const int BYTES_PER_CHUNK = BYTES / RC;

    // Primary columns, one set per lane
    col witness bits(7) b_op[lanes_x_row];                     // binary operation code
    col witness bits(8) free_in_a[lanes_x_row][BYTES];         // input A per byte
    col witness bits(8) free_in_b[lanes_x_row][BYTES];         // input B per byte
    col witness bits(8) free_in_c[lanes_x_row][BYTES];         // output C per byte
    col witness bits(4) carry[lanes_x_row][BYTES];             // per-byte carry chain:
                                            //      (cin[i]=carry[i-1], cout[i]=carry[i])
                                            // 4 bits to hold SH3ADD's {carry + bits shifted out} in [0,8] and
                                            // LT_ABS_NP's packed {clt + 2*cneg}, all other ops use only bit 0

    // Flags and helpers, one set per lane
    col witness bits(1) mode32[lanes_x_row];          // 1 if operation is in 32-bit mode; 0 if in 64-bit mode
    col witness bits(1) result_is_a[lanes_x_row];     // 1 if operation result is A (MIN/MAX family); 0 otherwise
    col witness bits(1) use_first_byte[lanes_x_row];  // 1 if operation uses first byte (LT_ABS_*); 0 otherwise
    col witness bits(1) c_is_signed[lanes_x_row];     // 1 if output C is signed; 0 otherwise

    // Auxiliary columns (primarily used to lower the lookup call complexity)
    col witness bits(10) b_op_or_sext[lanes_x_row];
    col witness bits(1) mode32_and_c_is_signed[lanes_x_row];

    const expr mode64[lanes_x_row];
    const expr cout[lanes_x_row];

    /*
    Perform, at the byte level, lookups against the binary table on inputs:
                  [pos_ind, opid, a, b, cin, c, cout + flags]
    where pos_ind indicates relevant byte positions

    Example for 8 bytes:
        │──────────────────│───────────│────│────│──────────│────│──────────────────-----------------------------------------------────────│
        │     pos_ind      │    opid   │ a  │ b  │ cin      │ c  │ cout + flags                                                           │
        │──────────────────│───────────│────│────│──────────│────│──────────────────-----------------------------------------------────────│
        │ 2*use_first_byte │   b_op    │ a0 │ b0 │ 0        │ c0 │ carry[0] + 16*result_is_a + 32*use_first_byte                          │
        │        0         │   b_op    │ a1 │ b1 │ carry[0] │ c1 │ carry[1] + 16*result_is_a + 32*use_first_byte                          │
        │        0         │   b_op    │ a2 │ b2 │ carry[1] │ c2 │ carry[2] + 16*result_is_a + 32*use_first_byte                          │
        │      mode32      │   b_op    │ a3 │ b3 │ carry[2] │ c3 │ carry[3] + 16*result_is_a + 32*use_first_byte + 64*(mode32*c_is_signed)│
        │        0         │ b_op|SEXT │ a4 │ b4 │ carry[3] │ c4 │ carry[4] + 16*result_is_a + 32*use_first_byte + 64*(mode32*c_is_signed)│
        │        0         │ b_op|SEXT │ a5 │ b5 │ carry[4] │ c5 │ carry[5] + 16*result_is_a + 32*use_first_byte + 64*(mode32*c_is_signed)│
        │        0         │ b_op|SEXT │ a6 │ b6 │ carry[5] │ c6 │ carry[6] + 16*result_is_a + 32*use_first_byte + 64*(mode32*c_is_signed)│
        │      mode64      │ b_op|SEXT │ a7 │ b7 │ carry[6] │ c7 │ carry[7] + 16*result_is_a + 32*use_first_byte + 64*c_is_signed         │
        │──────────────────│───────────│────│────│──────────│────│──────────────────-----------------------------------------------────────│

    The carry keeps the 4 low bits of the flags for itself: the SHxADD family transports in it both the
    addition carry and the bits shifted out of the previous byte (up to 8 for SH3ADD), so packing the
    remaining flags above it is what makes the decoding of the looked up FLAGS value unique
    */

    // Constraints to make sure that this component is called from the main component
    expr a[lanes_x_row][RC];
    expr b[lanes_x_row][RC];
    expr c[lanes_x_row][RC];

    for (int lane = 0; lane < lanes_x_row; lane++) {

        mode64[lane] = 1 - mode32[lane];
        cout[lane] = carry[lane][BYTES-1];

        // Selectors not constrained by the table should be constrained
        mode32[lane] * (1 - mode32[lane]) === 0; // Resilent constraint
        cout[lane] * (1 - cout[lane]) === 0;
        result_is_a[lane] * (1 - result_is_a[lane]) === 0;
        use_first_byte[lane] * (1 - use_first_byte[lane]) === 0;
        c_is_signed[lane] * (1 - c_is_signed[lane]) === 0;

        b_op_or_sext[lane] <== mode32[lane] * (c_is_signed[lane] * (OP_SEXT_FF - OP_SEXT_00) + OP_SEXT_00 - b_op[lane]) + b_op[lane];
        mode32_and_c_is_signed[lane] <== mode32[lane] * c_is_signed[lane];

        // Lookups
        lookup_assumes(BINARY_TABLE_ID, [2*use_first_byte[lane], b_op[lane], free_in_a[lane][0], free_in_b[lane][0], 0, free_in_c[lane][0], carry[lane][0] + 16*result_is_a[lane] + 32*use_first_byte[lane]]);
        for (int i = 1; i < BYTES; i++) {
            if (i < HALF_BYTES - 1) {
                lookup_assumes(BINARY_TABLE_ID, [0, b_op[lane], free_in_a[lane][i], free_in_b[lane][i], carry[lane][i-1], free_in_c[lane][i], carry[lane][i] + 16*result_is_a[lane] + 32*use_first_byte[lane]]);
            } else if (i == HALF_BYTES - 1) {
                lookup_assumes(BINARY_TABLE_ID, [mode32[lane], b_op[lane], free_in_a[lane][i], free_in_b[lane][i], carry[lane][i-1], free_in_c[lane][i], carry[lane][i] + 16*result_is_a[lane] + 32*use_first_byte[lane] + 64*mode32_and_c_is_signed[lane]]);
            } else if (i < BYTES - 1) {
                lookup_assumes(BINARY_TABLE_ID, [0, b_op_or_sext[lane], free_in_a[lane][i], free_in_b[lane][i], carry[lane][i-1], free_in_c[lane][i], carry[lane][i] + 16*result_is_a[lane] + 32*use_first_byte[lane] + 64*mode32_and_c_is_signed[lane]]);
            } else  {
                lookup_assumes(BINARY_TABLE_ID, [mode64[lane], b_op_or_sext[lane], free_in_a[lane][i], free_in_b[lane][i], carry[lane][i-1], free_in_c[lane][i], carry[lane][i] + 16*result_is_a[lane] + 32*use_first_byte[lane] + 64*c_is_signed[lane]]);
            }
        }

        for (int i = 0; i < RC; i++) {
            a[lane][i] = 0;
            b[lane][i] = 0;
            c[lane][i] = 0;
        }

        int byte = 0;
        int chunk = 0;
        for (int i = 0; i < BYTES; i++) {
            const int byte_weight = BYTE_BASE ** byte;
            a[lane][chunk] += (byte_weight * free_in_a[lane][i]);
            b[lane][chunk] += (byte_weight * free_in_b[lane][i]);
            c[lane][chunk] += (byte_weight * free_in_c[lane][i]);
            byte++;
            if (byte == BYTES_PER_CHUNK) {
                byte = 0;
                chunk++;
            }
        }

        // In comparison operations, we set c = 0 and so the result is cout
        // Otherwise, the result is simply c
        c[lane][0] += cout[lane];

        proves_operation(op: b_op[lane] + 0x10 * mode32[lane], a: a[lane], b: b[lane], c: c[lane], flag: cout[lane]);
    }

    airval padding_size;
    assumes_padding_operation(op: OP_ADD, padding_size:);
}