Polynomial.huff

/// @title Polynomial Arithmetic Utilities
/// @notice SPDX-License-Identifier: MIT
/// @author erhant <https://github.com/erhant>
/// @dev for every [offset, length] that these macros expect in memory,
/// it means that there are coefficients (a_0, a_1, ..., a_k) stored and
/// any a_i can be accessed by (offset + a_i * 0x20); the corresponding
/// polynomial is P(x) = a_0 + a_1*x + a_2*x^2 + ... + a_k*x^k

/// @dev multiply coefficients by some value
///
/// expects `length` many words stored in memory following `offset`.
///
/// args:    order of polynomial
///
/// takes:   [offset, length, val]
/// returns: [offset, length]
#define macro POLY_SCALE(order) = takes (3) returns (2) {
    // input:             [o, l, v] (offset, length, val)

    // ensure val is in the field
    <order> swap1 mod   // [o, l, v'] (v' := v % p)

    // multiply len by 32 since we will go over the memory word-by-word
    swap1               // [o, v, l]
    0x20 mul            // [o, v, l'] (l' = l * 32)

    // begin loop, starting from i := 0
    0x00                // [o, v, l, i] (index)
    loop_begin:
    // [o, v, l, i]

        // loop condition
        dup2 dup2       // [o, v, l, i, l, i]
        lt iszero       // [o, v, l, i, i >= l]
        loop_end jumpi

        // compute memory index
        dup1 dup5       // [o, v, l, i, i, o]
        add             // [o, v, l, i, k]          (k := o+i)
        dup1            // [o, v, l, i, k, k]

        // load coefficient from memory
        mload           // [o, v, l, i, k, c]       (c := mem[k])

        // scale
        <order> swap1   // [o, v, l, i, k, p, c]
        dup6            // [o, v, l, i, k, p, c, v]
        mulmod          // [o, v, l, i, k, c']      (c' := (c * v) % p)

        // load back result into memory
        swap1 mstore    // [o, v, l, i]

        // increment index one word
        0x20 add        // [o, v, l, i']            (i' = i' + 32)

    // [o, v, l, i]
    loop_begin jump
    loop_end:

    pop swap1 pop   // [o, l]

    // revert changes to length
    0x05 shr        // [o, l'] (l' := l >> 5 = l / 32)

    // output:         [o, l]
}

/// @dev add two polynomials
///
/// expects `length` many words stored in memory following `offset_P` and `offset_Q`.
///
/// args:    order of polynomial
///
/// takes:   [offset_P, offset_Q, length]
/// returns: [offset, length]
#define macro POLY_ADD(order) = takes (3) returns (2) {
    // input:      [oP, oQ, l] (offset_P, offset_Q, length)

    // multiply len by 32 since we will go over the memory word-by-word
    0x05 shl    // [oP, oQ, l'] (l' = l << 5 = l * 32)

    // begin loop, starting from i := 0
    0x00        // [oP, oQ, l, i] (index)
    loop_begin:
    // [oP, oQ, l, i]

        // loop condition
        dup2 dup2       // [oP, oQ, l, i, l, i]
        lt iszero       // [oP, oQ, l, i, i >= l]
        loop_end jumpi

        // compute memory index for the first polynomial
        dup1 dup5       // [oP, oQ, l, i, i, oP]
        add             // [oP, oQ, l, i, kP]           (kP := oP+i)
        dup1            // [oP, oQ, l, i, kP, kP]

        // compute memory index for the second polynomial
        dup3 dup6       // [oP, oQ, l, i, kP, kP, i, oQ]
        add             // [oP, oQ, l, i, kP, kP, kQ]   (kQ := oQ+i)

        // load coefficients
        mload           // [oP, oQ, l, i, kP, kP, cQ]   (cQ := mem[kQ])
        swap1 mload     // [oP, oQ, l, i, kP, cQ, cP]   (cP := mem[kP])

        // add
        <order> swap2   // [oP, oQ, l, i, kP, p, cP, cQ]
        addmod          // [oP, oQ, l, i, kP, c']       (c' := (cP + cQ) % p)

        // load back result into memory (at P's offset)
        swap1           // [oP, oQ, l, i, c, kP] 
        mstore          // [oP, oQ, l, i]

        // increment index one word
        0x20 add        // [oP, oQ, l, i']              (i' = i' + 32)

    // [oP, oQ, l, i]
    loop_begin jump
    loop_end:

    pop swap1 pop       // [oP, l]

    // div length by 32 to go back to normal indexing
    // which is equal to shifting to right 5 times
    0x05 shr            // [oP, l'] (l' = l >> 5 = l / 32)

    // output:             [oP, l]
}

/// @dev evaluate the polynomial at some point `x`
///
/// expects `length` many words stored in memory following `offset`.
///
/// uses Horner's method (https://zcash.github.io/halo2/background/polynomials.html#aside-horners-rule)
///
/// args:    order of polynomial
///
/// takes:   [offset, length, x]
/// returns: [offset, length, v]
#define macro POLY_EVAL(order) = takes (3) returns (3) {
    // input:      [o, l, x] (offset, length, x)

    // multiply len by 32 since we will go over the memory word-by-word
    swap1       // [o, x, l]
    0x05 shl    // [o, x, l'] (l' = l << 5 = l * 32)
    swap1       // [o, l, x]

    // sum accumulator starts at 0
    0x00        // [o, l, x, s] (s := 0)

    // index shall be (length - 1) * 32 = l - 32
    0x20        // [o, l, x, s, 32]
    dup4        // [o, l, x, s, 32, l]
    sub         // [o, l, x, s, i] (i := l - 32)

    // FIXME: change the name of loop label here to something other
    // than `loop_begin` and `loop_end` because it clashes with other
    // outer loops when this macro is called within.
    // however, when I change all labels to something more specific and long,
    // it get "byte index 24238 is out of bounds of ..." errors.
    evall:
    // [o, l, x, s, i]

        // loop condition
        // note that `i` is actually in range [l-1, 0] * 32, once it goes below 0
        // it shall underflow and this condition will fail nevertheless
        dup4 dup2       // [o, l, x, s, i, l, i]
        lt iszero       // [o, l, x, s, i, i >= l]
        evallend jumpi

        // compute memory index
        dup1 dup6       // [o, l, x, s, i, i, o]
        add             // [o, l, x, s, i, k]           (k := o+i)

        // mulmod accumulator with x
        <order>         // [o, l, x, s, i, k, p]        (p := order)
        dup4            // [o, l, x, s, i, k, p, s]
        dup6            // [o, l, x, s, i, k, p, s, x]
        mulmod          // [o, l, x, s, i, k, ss]       (ss := (s * x) % p)

        // addmod accumulator with coefficient
        <order>         // [o, l, x, s, i, k, ss, p]    (p : order)
        swap2           // [o, l, x, s, i, p, ss, k]
        mload           // [o, l, x, s, i, p, ss, c]    (c := mem[k])
        addmod          // [o, l, x, s, i, ss']         (ss' := (ss + c) % p)

        // update accumulator
        swap2           // [o, l, x, ss', i, s]
        pop             // [o, l, x, s',  i]             (s' := ss')

        // decrement index one word
        0x20            // [o, l, x, s, i, 32]
        swap1           // [o, l, x, s, 32, i]
        sub             // [o, l, x, s, i']             (i' := i - 32)

    // [o, l, x, s, i]
    evall jump
    evallend:
                        // [o, l, x, s, i]
    pop                 // [o, l, x, s]
    swap1               // [o, l, s, x]
    pop                 // [o, l, s]

    // divide length by 32 to revert the change at the start
    swap1               // [o, s, l]
    0x05 shr            // [o, s, l'] (l' = l >> 5 = l / 32)
    swap1               // [o, l, s]

    // output:             [o, l, s] (offset, length, evaluation)
}


/// @dev find the offset for a basis polynomial
///
/// works for both memory and code offsets
///
/// args: size of coefficients, and the number of basis polynomials
///
/// takes:   [basisId]
/// returns: [offset]
#define macro POLY_OFFSET(size, count) = takes (1) returns (1) {
    // input:      [b]   (basis id)
    <size>      // [b, s]
    <count>     // [b, s, c]
    mul mul     // [b*s*c]
    // output:     [o]   (offset for this basis)
}

/// @dev load coefficients from code to memory.
///
/// expects `length` many coefficients (each `coeffSize` bytes) in code,
/// starting from the offset `codeOffset`.
///
/// there will be `length` many words stored in memory following `memOffset`.
///
/// takes:   [codeOffset, coeffSize, memOffset, length]
/// returns: [memOffset, length]
#define macro POLY_CODECOPY() = takes (4) returns (2) {
    // input:   [oC, s, oM, l]

    // compute right-shift amount (bits)
    dup3 0x08 mul   // [oC, s, oM, l, s'] (s' := s * 8) bytes to bits
    0x0100 sub      // [oC, s, oM, l, r]  (r  := 256 - s')
    swap1           // [oC, s, oM, r, l]

    // index starts with 0
    0x00            // [oC, s, oM, r, l, i] (i := 0)
    loop_begin:
    // [oC, s, oM, r, l, i]

        // loop condition
        dup2 dup2           // [oC, s, oM, r, l, i, l, i]
        lt iszero           // [oC, s, oM, r, l, i, i >= l]
        loop_end jumpi

        // compute memory location
        dup4 dup2           // [oC, s, oM, r, l, i, oM, i]
        0x20 mul            // [oC, s, oM, r, l, i, oM, i'] (i' := i * 32)
        add                 // [oC, s, oM, r, l, i, kM]     (kM := oM + i')

        // save one more for shifting later
        dup1                // [oC, s, oM, r, l, i, kM, kM] 

        // compute code location
        dup8 dup4 dup9      // [oC, s, oM, r, l, i, kM, kM, oC, i, s]
        mul                 // [oC, s, oM, r, l, i, kM, kM, oC, i'] (i' := i * s')
        add                 // [oC, s, oM, r, l, i, kM, kM, kC]     (kC := oC + i')

        // load coefficient to memory
        dup8 swap2          // [oC, s, oM, r, l, i, kM, s, kC, kM]
        codecopy            // [oC, s, oM, r, l, i, kM] (mem[kM] := code[kC:kC+s])

        // shift the value
        dup1 mload          // [oC, s, oM, r, l, i, kM, c]  (c  := mem[kM]) but is 0-rightpadded
        dup5 shr            // [oC, s, oM, r, l, i, kM, c'] (c' := c >> r)
        swap1 mstore        // [oC, s, oM, r, l, i]         (mem[kM] := c')
                
        // increment index
        0x01 add            // [oS, s, oM, r, l, i'] (i' := i + 1)

    // [oC, s, oM, r, l, i]
    loop_begin jump
    loop_end: 

    pop swap1       // [oC, s, oM, l, r]
    pop swap2       // [oC, l, oM, s]
    pop swap2       // [oM, l, oC]
    pop             // [oM, l]

    // output:         [oM, l] (memOffset, length)
}

/// @dev load coefficients from storage to memory.
///
/// expects `length` many words stored in storage, starting from slot `storageOffset`.
///
/// there will be `length` many words stored in memory following `memOffset`.
///
/// takes:   [storageOffset, memOffset, length]
/// returns: [memOffset, length]
#define macro POLY_SLOAD() = takes (3) returns (2) {
    // input:   [oS, oM, l]

    // index starts with 0
    0x00        // [oS, oM, l, i] (i := 0)
    loop_begin:
    // [oS, oM, l, i]

        // loop condition
        dup2 dup2           // [oS, oM, l, i, l, i]
        lt iszero           // [oS, oM, l, i, i >= l]
        loop_end jumpi
        
        // find storage location
        dup4 dup2           // [oS, oM, l, i, oS, i]
        add                 // [oS, oM, l, i, k] (k := oS + i)

        // load coefficient from storage
        sload               // [oS, oM, l, i, c] (c := $[k])

        // find memory location
        dup4 dup3           // [oS, oM, l, i, c, oM, i]
        0x20 mul            // [oS, oM, l, i, c, oM, i'] (i' := i * 32)
        add                 // [oS, oM, l, i, c, k] (k := oM + i')

        // store coefficient to memory
        mstore              // [oS, oM, l, i] (mem[k] := c)

        // increment index
        0x01 add            // [oS, oM, l, i'] (i' := i + 1)

    // [oS, oM, l, i]
    loop_begin jump
    loop_end:
    
    pop     // [oS, oM, l]
    swap1   // [oS, l, oM]
    swap2   // [oM, l, oS]
    pop     // [oM, l]

    // output: [oM, l] (memOffset, length)
}

/// @dev store coefficients in memory to storage.
///
/// expects `length` many words stored in memory, following `memOffset`.
///
/// there will be `length` many words stored in storage following slot `storageOffset`.
///
/// takes:   [memOffset, storageOffset, length]
/// returns: [storageOffset, length]
#define macro POLY_SSTORE() = takes (3) returns (2) {
    // input:   [oM, oS, l]

    // index starts with 0
    0x00        // [oM, oS, l, i] (i := 0)
    loop_begin:
    // [oM, oS, l, i]

        // loop condition
        dup2 dup2           // [oM, oS, l, i, l, i]
        lt iszero           // [oM, oS, l, i, i >= l]
        loop_end jumpi
        
        // find memory location
        dup4 dup2           // [oM, oS, l, i, oM, i]
        0x20 mul            // [oM, oS, l, i, oM, i'] (i' := i * 32)
        add                 // [oM, oS, l, i, k] (k := oM + i')

        // load coefficient from memory
        mload               // [oM, oS, l, i, c] (c := mem[k])

        // find storage location
        dup4 dup3           // [oM, oS, l, i, c, oS, i]
        add                 // [oM, oS, l, i, c, k] (k := oS + i')

        // store coefficient to storage
        sstore              // [oM, oS, l, i] ($[k] := c)

        // increment index
        0x01 add            // [oM, oS, l, i'] (i' := i + 1)

    // [oM, oS, l, i]
    loop_begin jump
    loop_end:
    
    pop     // [oM, oS, l, i]
    swap1   // [oM, l, oS]
    swap2   // [oS, l, oM]
    pop     // [oS, l]

    // output: [oS, l] (storageOffset, length)
}