zisk-sm-binary 1.1.0-alpha

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

/* PIL Binary Operations Table used by Binary
                                                                                     Accumulated rows
    MINU       2^8 (A) x 2^8 (B) x 2^1 (POS_IND) x 2^1 (CIN) x FLAGS  = 2^18 + 2^17 |       393,216
    MIN        2^8 (A) x 2^8 (B) x 2^1 (POS_IND) x 2^1 (CIN) x FLAGS  = 2^18 + 2^17 |       786,432
    MAXU       2^8 (A) x 2^8 (B) x 2^1 (POS_IND) x 2^1 (CIN) x FLAGS  = 2^18 + 2^17 |     1,179,648
    MAX        2^8 (A) x 2^8 (B) x 2^1 (POS_IND) x 2^1 (CIN) x FLAGS  = 2^18 + 2^17 |     1,572,864
    LT_ABS_NP  2^8 (A) x 2^8 (B) x 2^2 (POS_IND) x 2^2 (CIN)          =    2^20     |     2,621,440
    LT_ABS_PN  2^8 (A) x 2^8 (B) x 2^2 (POS_IND) x 2^1 (CIN)          =    2^19     |     3,145,728
    LTU        2^8 (A) x 2^8 (B) x 2^1 (POS_IND) x 2^1 (CIN)          =    2^18     |     3,407,872
    LT         2^8 (A) x 2^8 (B) x 2^1 (POS_IND) x 2^1 (CIN)          =    2^18     |     3,670,016
    GT         2^8 (A) x 2^8 (B) x 2^1 (POS_IND) x 2^1 (CIN)          =    2^18     |     3,932,160
    EQ         2^8 (A) x 2^8 (B) x 2^1 (POS_IND) x 2^1 (CIN)          =    2^18     |     4,194,304
    ADD        2^8 (A) x 2^8 (B) x 2^1 (POS_IND) x 2^1 (CIN)          =    2^18     |     4,456,448
    SUB        2^8 (A) x 2^8 (B) x 2^1 (POS_IND) x 2^1 (CIN)          =    2^18     |     4,718,592
    LEU        2^8 (A) x 2^8 (B) x 2^1 (POS_IND) x 2^1 (CIN)          =    2^18     |     4,980,736
    LE         2^8 (A) x 2^8 (B) x 2^1 (POS_IND) x 2^1 (CIN)          =    2^18     |     5,242,880
    AND        2^8 (A) x 2^8 (B) x 2^1 (POS_IND)                      =    2^17     |     5,373,952
    OR         2^8 (A) x 2^8 (B) x 2^1 (POS_IND)                      =    2^17     |     5,505,024
    XOR        2^8 (A) x 2^8 (B) x 2^1 (POS_IND)                      =    2^17     |     5,636,096
    SEXT_00    2^8 (A) x 2^8 (B)                 x CIN x 2^1 (FLAGS)  = 2^17 + 2^16 |     5,832,704
    SEXT_FF    2^8 (A) x 2^8 (B)                 x CIN x 2^1 (FLAGS)  = 2^17 + 2^16 |     6,029,312
    ANDN       2^8 (A) x 2^8 (B) x 2^1 (POS_IND)                      =    2^17     |     6,160,384
    ORN        2^8 (A) x 2^8 (B) x 2^1 (POS_IND)                      =    2^17     |     6,291,456
    XNOR       2^8 (A) x 2^8 (B) x 2^1 (POS_IND)                      =    2^17     |     6,422,528
    BREV8      2^8 (A) x 2^8 (B) x 2^1 (POS_IND)                      =    2^17     |     6,553,600  => BINARY_TABLE_SIZE

    Total N = 6,553,600 rows, with 2^22 (4,194,304) < N < 2^23 (8,388,608)
    -------------------------------------------------------------------------------------------------------------------
    Notes:
    - MINU/MIN/MAXU/MAX: POS_IND=0 (2^18 rows with FLAGS), POS_IND=1 (2^17 rows without FLAGS)
    - SEXT_00/SEXT_FF: FLAGS=0 (2^17 rows with CIN), FLAGS=1 (2^16 rows without CIN)
    - Bitwise ops don't use POS_IND operationally, but include it for consistency with other ops
    - Sign extension ops don't use A,B,CIN operationally, but include them for consistency with other ops
    - SEXT_00, SEXT_FF are the only (unary) operation which are not a ZisK OP but are used to prove the rest
    - ANDN, ORN, XNOR, BREV8 are B-extension bitwise ops, single-byte like AND/OR/XOR (2^17 rows each)
*/

const int BINARY_TABLE_SIZE = 6_553_600;
const int BINARY_TABLE_SIZE_P2 = 2**23;
const int OP_SEXT_00 = 0x200;
const int OP_SEXT_FF = 0x201;

airtemplate BinaryTable(int N = 0) {
    if (N == 0) {
        N = VIRTUAL ? BINARY_TABLE_SIZE : BINARY_TABLE_SIZE_P2;
    } else if (N < BINARY_TABLE_SIZE) {
        error(`N must be at least ${BINARY_TABLE_SIZE}, but N=${N} was provided`);
    }

    const int SIGN_BYTE = 0x80;

    // Input A (8 bits)
    col fixed A = [0..255]...;                          

    // Input B (8 bits)
    col fixed B = [0:P2_8..255:P2_8]...;

    // Indicator of the byte position (<=2 bits)
    //  - 2 indicates the first byte
    //  - 1 indicates the last byte
    //  - 0 indicates a middle byte
    col fixed POS_IND = [[0:P2_18, 1:P2_17]:2,              // MINU,MIN
                         [0:P2_18, 1:P2_17]:2,              // MAXU,MAX
                         [0:P2_16..2:P2_16, 2:P2_16]:4,     // LT_ABS_NP
                         [0:P2_16..2:P2_16, 2:P2_16]:2,     // LT_ABS_PN
                         [0:P2_16, 1:P2_16]:(4*2),          // LTU,LT,GT,EQ
                         [0:P2_16, 1:P2_16]:(2*2),          // ADD,SUB
                         [0:P2_16, 1:P2_16]:(2*2),          // LEU,LE
                         [0:P2_16, 1:P2_16]:3,              // AND,OR,XOR
                          0:(P2_17 + P2_16),                // SEXT_00
                          0:(P2_17 + P2_16),                // SEXT_FF
                         [0:P2_16, 1:P2_16]:4]...;          // ANDN,ORN,XNOR,BREV8

    // Input carry (<=2 bits)
    col fixed CIN = [[0:P2_16, 1:P2_16]:(3*2),        // MINU,MIN
                     [0:P2_16, 1:P2_16]:(3*2),        // MAXU,MAX
                     [0:P2_18..3:P2_18],              // LT_ABS_NP
                     [0:P2_18, 1:P2_18],              // LT_ABS_PN
                     [0:P2_17, 1:P2_17]:4,            // LTU,LT,GT,EQ
                     [0:P2_17, 1:P2_17]:2,            // ADD,SUB
                     [0:P2_17, 1:P2_17]:2,            // LEU,LE
                      0:(P2_17*3),                    // AND,OR,XOR
                     [0:P2_16, 1:P2_16, 0:P2_16],     // SEXT_00
                     [0:P2_16, 1:P2_16, 0:P2_16],     // SEXT_FF
                      0:(P2_17*4)]...;                // ANDN,ORN,XNOR,BREV8

    // Operation opcode
    col fixed OP = [OP_MINU:(P2_18 + P2_17), OP_MIN:(P2_18 + P2_17),     // MINU,MIN
                    OP_MAXU:(P2_18 + P2_17), OP_MAX:(P2_18 + P2_17),     // MAXU,MAX
                    OP_LT_ABS_NP:P2_20,                                  // LT_ABS_NP
                    OP_LT_ABS_PN:P2_19,                                  // LT_ABS_PN
                    OP_LTU:P2_18, OP_LT:P2_18, OP_GT:P2_18, OP_EQ:P2_18, // LTU,LT,GT,EQ
                    OP_ADD:P2_18, OP_SUB:P2_18,                          // ADD,SUB
                    OP_LEU:P2_18, OP_LE:P2_18,                           // LEU,LE
                    OP_AND:P2_17, OP_OR:P2_17, OP_XOR:P2_17,             // AND,OR,XOR
                    OP_SEXT_00:(P2_17 + P2_16),                          // SEXT_00
                    OP_SEXT_FF:(P2_17 + P2_16),                          // SEXT_FF
                    OP_ANDN:P2_17, OP_ORN:P2_17,                         // ANDN, ORN
                    OP_XNOR:P2_17, OP_BREV8:P2_17]...;                   // XNOR, BREV8

    col fixed C;     // Output C (8 bits)
    col fixed FLAGS; // Combined flags (8 bits): cout + 2*result_is_a + 4*use_first_byte + 8*c_is_signed

    int relative_index = 0;
    #pragma transpile logfile:/tmp/binary_table.txt
    for (int i = 0; i < N; i++) {
        const int [pos_ind, op, a, b, cin] = [POS_IND[i], OP[i], A[i], B[i], CIN[i]];
        relative_index = (i == 0 || OP[i-1] != op) ? 0 : relative_index + 1;
        const int pfirst = (pos_ind == 2) ? 1 : 0;
        const int plast = (pos_ind == 1) ? 1 : 0;
        int c = 0;
        int cout = 0;
        int result_is_a = 0;
        int use_first_byte = 0;
        int c_is_signed = 0;
        switch (op) {
            case OP_MINU,OP_MIN: // MINU,MIN: Minimum operation
                // Returns min(a,b). MINU for unsigned, MIN for signed.
                // Example: MINU(5,3) = 3, MIN(-1,3) = -1
                if (a < b) {
                    cout = 1;
                } else if (a == b) {
                    cout = cin;
                } else {
                    cout = 0;
                }
                
                if (plast == 0) {
                    result_is_a = (relative_index >= P2_17) ? 1 : 0;
                    c = result_is_a ? a : b;
                } else {
                    result_is_a = cout;

                    // For signed MIN with different signs, the negative number is smaller
                    if (op == OP_MIN && (a & SIGN_BYTE) != (b & SIGN_BYTE)) {
                        result_is_a = (a & SIGN_BYTE) ? 1 : 0;
                    }

                    c = result_is_a ? a : b;
                    c_is_signed = (c & SIGN_BYTE) ? 1 : 0;
                    cout = 0;  // Clear carry for final result
                }

            case OP_MAXU,OP_MAX: // MAXU,MAX: Maximum operation
                // Returns max(a,b). MAXU for unsigned, MAX for signed.
                // Example: MAXU(5,3) = 5, MAX(-1,3) = 3
                if (a > b) {
                    cout = 1;
                } else if (a == b) {
                    cout = cin;
                } else {
                    cout = 0;
                }

                if (plast == 0) {
                    result_is_a = (relative_index >= P2_17) ? 1 : 0;
                    c = result_is_a ? a : b;
                } else {
                    result_is_a = cout;

                    // For signed MAX with different signs, the negative number is smaller
                    if (op == OP_MAX && (a & SIGN_BYTE) != (b & SIGN_BYTE)) {
                        result_is_a = (b & SIGN_BYTE) ? 1 : 0;
                    }

                    c = result_is_a ? a : b;
                    c_is_signed = (c & SIGN_BYTE) ? 1 : 0;
                    cout = 0;  // Clear carry for final result
                }

            case OP_LT_ABS_NP: // LT_ABS_NP: Absolute value comparison |a| < |b| where a < 0 and b > 0
                // Compares absolute values when a is negative and b is positive
                // Example: LT_ABS_NP(-3, 5) = 1 (|-3| < |5|), LT_ABS_NP(-7, 5) = 0 (|-7| > |5|)

                // Notice that |a| < |b| if and only if |a| - b < 0
                // Since a is negative, |a| = -a = (a ^ 0xFF) + 1
                // This produces two carries:
                //   - the carry of the LT operation which compares |a| and b
                //   - the carry of the addition a ^ 0xFF + 1
                //
                // Both necessary carries are encoded by cin as 
                //                  cin = 0bYX,
                // where X is the carry of the LT operation and Y is 
                // the carry of the operation a ^ 0xFF + cneg
                //
                // Since a < 0, cneg should be 1 at the first byte and (a ^ 0xFF + cneg) >> 8 at the rest

                // Decode the carries
                const int clt = pfirst ? 0 : cin & 0x01;
                const int cneg = pfirst ? 1 : (cin & 0x02) >> 1;

                const int _a = (a ^ 0xFF) + cneg;
                const int abs_a = _a & 0xFF;

                if (abs_a < b) {
                    cout = 1;
                } else if (abs_a == b) {
                    cout = clt;
                } else {
                    cout = 0;
                }

                // Encode the negation carry for the next byte
                cout += 2*(_a >> 8);
                
                c = 0;

                use_first_byte = 1;

            case OP_LT_ABS_PN: // LT_ABS_PN: Absolute value comparison |a| < |b| where a > 0 and b < 0
                // Compares absolute values when a is positive and b is negative
                // Example: LT_ABS_PN(3, -5) = 1 (|3| < |-5|), LT_ABS_PN(7, -5) = 0 (|7| > |-5|)

                // Notice that |a| < |b| if and only if a - |b| < 0
                // Since b is negative, |b| = -b = (b ^ 0xFF) + 1
                //
                // Unlike LT_ABS_NP, a single carry is enough here: the +1 of the
                // negation lives in the subtrahend, so when byte 0 of b is zero it
                // just makes sub = a - 256 < 0, i.e. a borrow. That borrow is exactly
                // the negation carry, and cin already propagates it upward, so it
                // never needs to be tracked separately
                const int _a = a;
                const int _b = b ^ 0xFF;
                const int sub = pfirst ? _a - (_b + 1) : _a - _b;

                if (sub < 0) {
                    cout = 1;
                } else if (sub == 0) {
                    cout = cin;
                } else {
                    cout = 0;
                }

                c = 0;

                use_first_byte = 1;

            case OP_LTU,OP_LT: // LTU,LT: Less than comparison
                // Returns 1 if a < b, 0 otherwise. LTU for unsigned, LT for signed.
                // Example: LTU(3,5) = 1, LT(-1,3) = 1, LTU(5,3) = 0
                if (a < b) {
                    cout = 1;
                } else if (a == b) {
                    cout = cin;
                } else {
                    cout = 0;
                }

                c = 0;

                // If the chunk is signed, then the result is the sign of a
                if (op == OP_LT && plast && (a & SIGN_BYTE) != (b & SIGN_BYTE)) {
                    cout = (a & SIGN_BYTE) ? 1 : 0;
                }

            case OP_GT: // GT: Greater than comparison (signed)
                // Returns 1 if a > b, 0 otherwise (signed comparison)
                // Example: GT(5,3) = 1, GT(-1,3) = 0, GT(3,3) = 0
                if (a > b) {
                    cout = 1;
                } else if (a == b) {
                    cout = cin;
                } else {
                    cout = 0;
                }

                c = 0;

                // The result is the sign of b
                if (plast && (a & SIGN_BYTE) != (b & SIGN_BYTE)) {
                    cout = (b & SIGN_BYTE) ? 1 : 0;
                }

            case OP_EQ: // EQ: Equality comparison
                // Returns 1 if a == b, 0 otherwise
                // Example: EQ(5,5) = 1, EQ(3,5) = 0
                if (cin == 0 && a == b) {
                    cout = 0; // is_eq
                } else {
                    cout = 1; // is_neq
                }

                c = 0;

                if (plast) {
                    // cout = 0 means a == b => change cout = 1
                    cout = 1 - cout;
                }

            case OP_ADD: // ADD: Addition with carry
                // Performs a + b + cin
                // Example: ADD(5,3,0) = 8, ADD(255,1,0) = 0
                const int sum = cin + a + b;
                c = sum & 0xFF;
                cout = plast ? 0 : sum >> 8;

                if (plast) {
                    c_is_signed = (c & SIGN_BYTE) ? 1 : 0;
                }

            case OP_SUB: // SUB: Subtraction with borrow
                // Performs a - b - cin (where cin is borrow)
                // Example: SUB(5,3,0) = 2, SUB(3,5,0) = 254 with borrow=1
                const int borrow = (a - cin) < b ? 1 : 0;
                c = P2_8 * borrow + a - cin - b;
                cout = plast ? 0 : borrow;

                if (plast) {
                    c_is_signed = (c & SIGN_BYTE) ? 1 : 0;
                }

            case OP_LEU,OP_LE: // LEU,LE: Less than or equal comparison
                // Returns 1 if a <= b, 0 otherwise. LEU for unsigned, LE for signed.
                // Example: LEU(3,5) = 1, LE(-1,3) = 1, LEU(5,3) = 0, LEU(3,3) = 1
                if (a < b) {
                    cout = 0; // is_le
                } else if (a == b) {
                    cout = cin;
                } else {
                    cout = 1; // is_gt
                }

                c = 0;

                if (plast) {
                    // cout = 0 means a <= b => change cout = 1
                    cout = 1 - cout;

                    // If the chunk is signed, then the result is the sign of a
                    if (op == OP_LE && (a & SIGN_BYTE) != (b & SIGN_BYTE)) {
                        cout = (a & SIGN_BYTE) ? 1 : 0;
                    }
                }

            case OP_AND: // AND: Bitwise AND
                // Performs bitwise AND operation: a & b
                // Example: AND(0b1100, 0b1010) = 0b1000 = 8
                c = a & b;

            case OP_OR: // OR: Bitwise OR
                // Performs bitwise OR operation: a | b
                // Example: OR(0b1100, 0b1010) = 0b1110 = 14
                c = a | b;

            case OP_XOR: // XOR: Bitwise XOR
                // Performs bitwise XOR operation: a ^ b
                // Example: XOR(0b1100, 0b1010) = 0b0110 = 6
                c = a ^ b;

            case OP_SEXT_00: // SEXT_00: Sign extension with 0
                // Example: SEXT_00(a) = 0x00
                cout = cin;
                c = 0x00;

                result_is_a = (relative_index >= P2_17) ? 1 : 0;
                c_is_signed = 0;

            case OP_SEXT_FF: // SEXT_FF: Sign extension with 1
                // Example: SEXT_FF(a) = 0xFF
                cout = cin;
                c = 0xFF;

                result_is_a = (relative_index >= P2_17) ? 1 : 0;
                c_is_signed = 1;

            case OP_ANDN: // ANDN: bitwise AND with negated b -> a & ~b
                c = a & (b ^ 0xFF);

            case OP_ORN: // ORN: bitwise OR with negated b -> a | ~b
                c = a | (b ^ 0xFF);

            case OP_XNOR: // XNOR: bitwise NOT of XOR -> ~(a ^ b)
                c = (a ^ b) ^ 0xFF;

            case OP_BREV8: // BREV8: reverse the bit order within the byte (operand in b)
                c = 0;
                for (int k = 0; k < 8; k++) {
                    c = c | (((b >> k) & 0x01) << (7 - k));
                }

            default:
                error(`Invalid operation opcode: ${op} at row ${i}`);
        }

        C[i] = c;
        const int flags = cout + 2*result_is_a + 4*use_first_byte + 8*c_is_signed;
        FLAGS[i] = flags;
        log(`T[${i}] = [${pos_ind}, ${op}, ${a}, ${b}, ${cin}, ${c}, ${flags}]`);
    }

    col witness multiplicity;
    lookup_proves(BINARY_TABLE_ID, [POS_IND, OP, A, B, CIN, C, FLAGS], multiplicity);
}