FriTransform.sol

/*
  Copyright 2019-2022 StarkWare Industries Ltd.

  Licensed under the Apache License, Version 2.0 (the "License").
  You may not use this file except in compliance with the License.
  You may obtain a copy of the License at

  https://www.starkware.co/open-source-license/

  Unless required by applicable law or agreed to in writing,
  software distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions
  and limitations under the License.
*/
// SPDX-License-Identifier: Apache-2.0.
pragma solidity ^0.6.12;

import "./PrimeFieldElement0.sol";

/*
  The FRI transform for a coset of size 2 (x, -x) takes the inputs
  x, f(x), f(-x) and evalPoint
  and returns
  (f(x) + f(-x) + evalPoint*(f(x) - f(-x))/x) / 2.

  The implementation here modifies this transformation slightly:
  1. Since dividing by 2 does not affect the degree, it is omitted here (and in the prover).
  2. The division by x is replaced by multiplication by x^-1, x^-1 is passed as input to the
     transform and (x^-1)^2 is returned as it will be needed in the next layer.

  To apply the transformation on a larger coset the transformation above is used multiple times
  with the evaluation points: evalPoint, evalPoint^2, evalPoint^4, ...
*/
contract FriTransform is PrimeFieldElement0 {
    // The supported step sizes are 2, 3 and 4.
    uint256 internal constant FRI_MIN_STEP_SIZE = 2;
    uint256 internal constant FRI_MAX_STEP_SIZE = 4;

    // The largest power of 2 multiple of K_MODULUS that fits in a uint256.
    // The value is given as a constant because "Only direct number constants and references to such
    // constants are supported by inline assembly."
    // This constant is used in places where we delay the module operation to reduce gas usage.
    uint256 internal constant K_MODULUS_TIMES_16 = (
    0x8000000000000110000000000000000000000000000000000000000000000010
    );

    /*
      Performs a FRI transform for the coset of size friFoldedCosetSize that begins at index.

      Assumes the evaluations on the coset are stored at 'evaluationsOnCosetPtr'.
      See gatherCosetInputs for more detail.
    */
    function transformCoset(
        uint256 friHalfInvGroupPtr,
        uint256 evaluationsOnCosetPtr,
        uint256 cosetOffset,
        uint256 friEvalPoint,
        uint256 friCosetSize
    ) internal pure returns (uint256 nextLayerValue, uint256 nextXInv) {
        // Compare to expected FRI step sizes in order of likelihood, step size 3 being most common.
        if (friCosetSize == 8) {
            return
            transformCosetOfSize8(
                friHalfInvGroupPtr,
                evaluationsOnCosetPtr,
                cosetOffset,
                friEvalPoint
            );
        } else if (friCosetSize == 4) {
            return
            transformCosetOfSize4(
                friHalfInvGroupPtr,
                evaluationsOnCosetPtr,
                cosetOffset,
                friEvalPoint
            );
        } else if (friCosetSize == 16) {
            return
            transformCosetOfSize16(
                friHalfInvGroupPtr,
                evaluationsOnCosetPtr,
                cosetOffset,
                friEvalPoint
            );
        } else {
            require(false, "Only step sizes of 2, 3 or 4 are supported.");
        }
    }

    /*
      Applies 2 + 1 FRI transformations to a coset of size 2^2.

      evaluations on coset:                    f0 f1  f2 f3
      ----------------------------------------  \ / -- \ / -----------
                                                 f0    f2
      ------------------------------------------- \ -- / -------------
      nextLayerValue:                               f0

      For more detail, see description of the FRI transformations at the top of this file.
    */
    function transformCosetOfSize4(
        uint256 friHalfInvGroupPtr,
        uint256 evaluationsOnCosetPtr,
        uint256 cosetOffset_,
        uint256 friEvalPoint
    ) private pure returns (uint256 nextLayerValue, uint256 nextXInv) {
        assembly {
            let friEvalPointDivByX := mulmod(friEvalPoint, cosetOffset_, K_MODULUS)

            let f0 := mload(evaluationsOnCosetPtr)
            {
                let f1 := mload(add(evaluationsOnCosetPtr, 0x20))

            // f0 < 3P ( = 1 + 1 + 1).
                f0 := add(
                add(f0, f1),
                mulmod(
                friEvalPointDivByX,
                add(
                f0,
                // -fMinusX
                sub(K_MODULUS, f1)
                ),
                K_MODULUS
                )
                )
            }

            let f2 := mload(add(evaluationsOnCosetPtr, 0x40))
            {
                let f3 := mload(add(evaluationsOnCosetPtr, 0x60))
                f2 := addmod(
                add(f2, f3),
                mulmod(
                add(
                f2,
                // -fMinusX
                sub(K_MODULUS, f3)
                ),
                mulmod(mload(add(friHalfInvGroupPtr, 0x20)), friEvalPointDivByX, K_MODULUS),
                K_MODULUS
                ),
                K_MODULUS
                )
            }

            {
                let newXInv := mulmod(cosetOffset_, cosetOffset_, K_MODULUS)
                nextXInv := mulmod(newXInv, newXInv, K_MODULUS)
            }

        // f0 + f2 < 4P ( = 3 + 1).
            nextLayerValue := addmod(
            add(f0, f2),
            mulmod(
            mulmod(friEvalPointDivByX, friEvalPointDivByX, K_MODULUS),
            add(
            f0,
            // -fMinusX
            sub(K_MODULUS, f2)
            ),
            K_MODULUS
            ),
            K_MODULUS
            )
        }
    }

    /*
      Applies 4 + 2 + 1 FRI transformations to a coset of size 2^3.

      For more detail, see description of the FRI transformations at the top of this file.
    */
    function transformCosetOfSize8(
        uint256 friHalfInvGroupPtr,
        uint256 evaluationsOnCosetPtr,
        uint256 cosetOffset_,
        uint256 friEvalPoint
    ) private pure returns (uint256 nextLayerValue, uint256 nextXInv) {
        assembly {
            let f0 := mload(evaluationsOnCosetPtr)

            let friEvalPointDivByX := mulmod(friEvalPoint, cosetOffset_, K_MODULUS)
            let friEvalPointDivByXSquared := mulmod(
            friEvalPointDivByX,
            friEvalPointDivByX,
            K_MODULUS
            )
            let imaginaryUnit := mload(add(friHalfInvGroupPtr, 0x20))

            {
                let f1 := mload(add(evaluationsOnCosetPtr, 0x20))

            // f0 < 3P ( = 1 + 1 + 1).
                f0 := add(
                add(f0, f1),
                mulmod(
                friEvalPointDivByX,
                add(
                f0,
                // -fMinusX
                sub(K_MODULUS, f1)
                ),
                K_MODULUS
                )
                )
            }
            {
                let f2 := mload(add(evaluationsOnCosetPtr, 0x40))
                {
                    let f3 := mload(add(evaluationsOnCosetPtr, 0x60))

                // f2 < 3P ( = 1 + 1 + 1).
                    f2 := add(
                    add(f2, f3),
                    mulmod(
                    add(
                    f2,
                    // -fMinusX
                    sub(K_MODULUS, f3)
                    ),
                    mulmod(friEvalPointDivByX, imaginaryUnit, K_MODULUS),
                    K_MODULUS
                    )
                    )
                }

            // f0 < 7P ( = 3 + 3 + 1).
                f0 := add(
                add(f0, f2),
                mulmod(
                friEvalPointDivByXSquared,
                add(
                f0,
                // -fMinusX
                sub(K_MODULUS_TIMES_16, f2)
                ),
                K_MODULUS
                )
                )
            }
            {
                let f4 := mload(add(evaluationsOnCosetPtr, 0x80))
                {
                    let friEvalPointDivByX2 := mulmod(
                    friEvalPointDivByX,
                    mload(add(friHalfInvGroupPtr, 0x40)),
                    K_MODULUS
                    )
                    {
                        let f5 := mload(add(evaluationsOnCosetPtr, 0xa0))

                    // f4 < 3P ( = 1 + 1 + 1).
                        f4 := add(
                        add(f4, f5),
                        mulmod(
                        friEvalPointDivByX2,
                        add(
                        f4,
                        // -fMinusX
                        sub(K_MODULUS, f5)
                        ),
                        K_MODULUS
                        )
                        )
                    }

                    let f6 := mload(add(evaluationsOnCosetPtr, 0xc0))
                    {
                        let f7 := mload(add(evaluationsOnCosetPtr, 0xe0))

                    // f6 < 3P ( = 1 + 1 + 1).
                        f6 := add(
                        add(f6, f7),
                        mulmod(
                        add(
                        f6,
                        // -fMinusX
                        sub(K_MODULUS, f7)
                        ),
                        // friEvalPointDivByX2 * imaginaryUnit ==
                        // friEvalPointDivByX * mload(add(friHalfInvGroupPtr, 0x60)).
                        mulmod(friEvalPointDivByX2, imaginaryUnit, K_MODULUS),
                        K_MODULUS
                        )
                        )
                    }

                // f4 < 7P ( = 3 + 3 + 1).
                    f4 := add(
                    add(f4, f6),
                    mulmod(
                    mulmod(friEvalPointDivByX2, friEvalPointDivByX2, K_MODULUS),
                    add(
                    f4,
                    // -fMinusX
                    sub(K_MODULUS_TIMES_16, f6)
                    ),
                    K_MODULUS
                    )
                    )
                }

            // f0, f4 < 7P -> f0 + f4 < 14P && 9P < f0 + (K_MODULUS_TIMES_16 - f4) < 23P.
                nextLayerValue := addmod(
                add(f0, f4),
                mulmod(
                mulmod(friEvalPointDivByXSquared, friEvalPointDivByXSquared, K_MODULUS),
                add(
                f0,
                // -fMinusX
                sub(K_MODULUS_TIMES_16, f4)
                ),
                K_MODULUS
                ),
                K_MODULUS
                )
            }

            {
                let xInv2 := mulmod(cosetOffset_, cosetOffset_, K_MODULUS)
                let xInv4 := mulmod(xInv2, xInv2, K_MODULUS)
                nextXInv := mulmod(xInv4, xInv4, K_MODULUS)
            }
        }
    }

    /*
      Applies 8 + 4 + 2 + 1 FRI transformations to a coset of size 2^4.
      to obtain a single element.

      For more detail, see description of the FRI transformations at the top of this file.
    */
    function transformCosetOfSize16(
        uint256 friHalfInvGroupPtr,
        uint256 evaluationsOnCosetPtr,
        uint256 cosetOffset_,
        uint256 friEvalPoint
    ) private pure returns (uint256 nextLayerValue, uint256 nextXInv) {
        assembly {
            let friEvalPointDivByXTessed
            let f0 := mload(evaluationsOnCosetPtr)

            let friEvalPointDivByX := mulmod(friEvalPoint, cosetOffset_, K_MODULUS)
            let imaginaryUnit := mload(add(friHalfInvGroupPtr, 0x20))

            {
                let f1 := mload(add(evaluationsOnCosetPtr, 0x20))

            // f0 < 3P ( = 1 + 1 + 1).
                f0 := add(
                add(f0, f1),
                mulmod(
                friEvalPointDivByX,
                add(
                f0,
                // -fMinusX
                sub(K_MODULUS, f1)
                ),
                K_MODULUS
                )
                )
            }
            {
                let f2 := mload(add(evaluationsOnCosetPtr, 0x40))
                {
                    let f3 := mload(add(evaluationsOnCosetPtr, 0x60))

                // f2 < 3P ( = 1 + 1 + 1).
                    f2 := add(
                    add(f2, f3),
                    mulmod(
                    add(
                    f2,
                    // -fMinusX
                    sub(K_MODULUS, f3)
                    ),
                    mulmod(friEvalPointDivByX, imaginaryUnit, K_MODULUS),
                    K_MODULUS
                    )
                    )
                }
                {
                    let friEvalPointDivByXSquared := mulmod(
                    friEvalPointDivByX,
                    friEvalPointDivByX,
                    K_MODULUS
                    )
                    friEvalPointDivByXTessed := mulmod(
                    friEvalPointDivByXSquared,
                    friEvalPointDivByXSquared,
                    K_MODULUS
                    )

                // f0 < 7P ( = 3 + 3 + 1).
                    f0 := add(
                    add(f0, f2),
                    mulmod(
                    friEvalPointDivByXSquared,
                    add(
                    f0,
                    // -fMinusX
                    sub(K_MODULUS_TIMES_16, f2)
                    ),
                    K_MODULUS
                    )
                    )
                }
            }
            {
                let f4 := mload(add(evaluationsOnCosetPtr, 0x80))
                {
                    let friEvalPointDivByX2 := mulmod(
                    friEvalPointDivByX,
                    mload(add(friHalfInvGroupPtr, 0x40)),
                    K_MODULUS
                    )
                    {
                        let f5 := mload(add(evaluationsOnCosetPtr, 0xa0))

                    // f4 < 3P ( = 1 + 1 + 1).
                        f4 := add(
                        add(f4, f5),
                        mulmod(
                        friEvalPointDivByX2,
                        add(
                        f4,
                        // -fMinusX
                        sub(K_MODULUS, f5)
                        ),
                        K_MODULUS
                        )
                        )
                    }

                    let f6 := mload(add(evaluationsOnCosetPtr, 0xc0))
                    {
                        let f7 := mload(add(evaluationsOnCosetPtr, 0xe0))

                    // f6 < 3P ( = 1 + 1 + 1).
                        f6 := add(
                        add(f6, f7),
                        mulmod(
                        add(
                        f6,
                        // -fMinusX
                        sub(K_MODULUS, f7)
                        ),
                        // friEvalPointDivByX2 * imaginaryUnit ==
                        // friEvalPointDivByX * mload(add(friHalfInvGroupPtr, 0x60)).
                        mulmod(friEvalPointDivByX2, imaginaryUnit, K_MODULUS),
                        K_MODULUS
                        )
                        )
                    }

                // f4 < 7P ( = 3 + 3 + 1).
                    f4 := add(
                    add(f4, f6),
                    mulmod(
                    mulmod(friEvalPointDivByX2, friEvalPointDivByX2, K_MODULUS),
                    add(
                    f4,
                    // -fMinusX
                    sub(K_MODULUS_TIMES_16, f6)
                    ),
                    K_MODULUS
                    )
                    )
                }

            // f0 < 15P ( = 7 + 7 + 1).
                f0 := add(
                add(f0, f4),
                mulmod(
                friEvalPointDivByXTessed,
                add(
                f0,
                // -fMinusX
                sub(K_MODULUS_TIMES_16, f4)
                ),
                K_MODULUS
                )
                )
            }
            {
                let f8 := mload(add(evaluationsOnCosetPtr, 0x100))
                {
                    let friEvalPointDivByX4 := mulmod(
                    friEvalPointDivByX,
                    mload(add(friHalfInvGroupPtr, 0x80)),
                    K_MODULUS
                    )
                    {
                        let f9 := mload(add(evaluationsOnCosetPtr, 0x120))

                    // f8 < 3P ( = 1 + 1 + 1).
                        f8 := add(
                        add(f8, f9),
                        mulmod(
                        friEvalPointDivByX4,
                        add(
                        f8,
                        // -fMinusX
                        sub(K_MODULUS, f9)
                        ),
                        K_MODULUS
                        )
                        )
                    }

                    let f10 := mload(add(evaluationsOnCosetPtr, 0x140))
                    {
                        let f11 := mload(add(evaluationsOnCosetPtr, 0x160))
                    // f10 < 3P ( = 1 + 1 + 1).
                        f10 := add(
                        add(f10, f11),
                        mulmod(
                        add(
                        f10,
                        // -fMinusX
                        sub(K_MODULUS, f11)
                        ),
                        // friEvalPointDivByX4 * imaginaryUnit ==
                        // friEvalPointDivByX * mload(add(friHalfInvGroupPtr, 0xa0)).
                        mulmod(friEvalPointDivByX4, imaginaryUnit, K_MODULUS),
                        K_MODULUS
                        )
                        )
                    }

                // f8 < 7P ( = 3 + 3 + 1).
                    f8 := add(
                    add(f8, f10),
                    mulmod(
                    mulmod(friEvalPointDivByX4, friEvalPointDivByX4, K_MODULUS),
                    add(
                    f8,
                    // -fMinusX
                    sub(K_MODULUS_TIMES_16, f10)
                    ),
                    K_MODULUS
                    )
                    )
                }
                {
                    let f12 := mload(add(evaluationsOnCosetPtr, 0x180))
                    {
                        let friEvalPointDivByX6 := mulmod(
                        friEvalPointDivByX,
                        mload(add(friHalfInvGroupPtr, 0xc0)),
                        K_MODULUS
                        )
                        {
                            let f13 := mload(add(evaluationsOnCosetPtr, 0x1a0))

                        // f12 < 3P ( = 1 + 1 + 1).
                            f12 := add(
                            add(f12, f13),
                            mulmod(
                            friEvalPointDivByX6,
                            add(
                            f12,
                            // -fMinusX
                            sub(K_MODULUS, f13)
                            ),
                            K_MODULUS
                            )
                            )
                        }

                        let f14 := mload(add(evaluationsOnCosetPtr, 0x1c0))
                        {
                            let f15 := mload(add(evaluationsOnCosetPtr, 0x1e0))

                        // f14 < 3P ( = 1 + 1 + 1).
                            f14 := add(
                            add(f14, f15),
                            mulmod(
                            add(
                            f14,
                            // -fMinusX
                            sub(K_MODULUS, f15)
                            ),
                            // friEvalPointDivByX6 * imaginaryUnit ==
                            // friEvalPointDivByX * mload(add(friHalfInvGroupPtr, 0xe0)).
                            mulmod(friEvalPointDivByX6, imaginaryUnit, K_MODULUS),
                            K_MODULUS
                            )
                            )
                        }

                    // f12 < 7P ( = 3 + 3 + 1).
                        f12 := add(
                        add(f12, f14),
                        mulmod(
                        mulmod(friEvalPointDivByX6, friEvalPointDivByX6, K_MODULUS),
                        add(
                        f12,
                        // -fMinusX
                        sub(K_MODULUS_TIMES_16, f14)
                        ),
                        K_MODULUS
                        )
                        )
                    }

                // f8 < 15P ( = 7 + 7 + 1).
                    f8 := add(
                    add(f8, f12),
                    mulmod(
                    mulmod(friEvalPointDivByXTessed, imaginaryUnit, K_MODULUS),
                    add(
                    f8,
                    // -fMinusX
                    sub(K_MODULUS_TIMES_16, f12)
                    ),
                    K_MODULUS
                    )
                    )
                }

            // f0, f8 < 15P -> f0 + f8 < 30P && 16P < f0 + (K_MODULUS_TIMES_16 - f8) < 31P.
                nextLayerValue := addmod(
                add(f0, f8),
                mulmod(
                mulmod(friEvalPointDivByXTessed, friEvalPointDivByXTessed, K_MODULUS),
                add(
                f0,
                // -fMinusX
                sub(K_MODULUS_TIMES_16, f8)
                ),
                K_MODULUS
                ),
                K_MODULUS
                )
            }

            {
                let xInv2 := mulmod(cosetOffset_, cosetOffset_, K_MODULUS)
                let xInv4 := mulmod(xInv2, xInv2, K_MODULUS)
                let xInv8 := mulmod(xInv4, xInv4, K_MODULUS)
                nextXInv := mulmod(xInv8, xInv8, K_MODULUS)
            }
        }
    }
}