openvm-cuda-common 2.0.0

/*
 * Source: https://github.com/risc0/risc0 (ref=093f46925a0366d72df505256d35c54a033ddad0)
 * Status: MODIFIED from risc0/sys/kernels/zkp/cuda/fpext.h
 * Imported: 2025-01-25 by @gaxiom
 * 
 * LOCAL CHANGES (high level):
 * - 2025-03-25: add operator-()
 * - 2025-09-19: use sppark's bb31_4_t impl for core arithmetic operations
 */

// Copyright 2024 RISC Zero, Inc.
//
// 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
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// 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.

#pragma once

/// \file
/// Defines FpExt, a finite field F_p^4, based on Fp via the irreducible polynomial x^4 - 11.

#include "fp.h"

/// Instances of FpExt are element of a finite field F_p^4.  They are represented as elements of
/// F_p[X] / (X^4 - 11). Basically, this is a 'big' finite field (about 2^128 elements), which is
/// used when the security of various operations depends on the size of the field.  It has the field
/// Fp as a subfield, which means operations by the two are compatible, which is important.  The
/// irreducible polynomial was chosen to be the simplest possible one, x^4 - B, where 11 is the
/// smallest B which makes the polynomial irreducible.
struct FpExt {
    /// The elements of FpExt, elems[0] + elems[1]*X + elems[2]*X^2 + elems[3]*x^4
    union {
        Fp elems[4];
        bb31_4_t rep;
    };
    
    /// Default constructor makes the zero elements
    __device__ FpExt() : rep(bb31_t{0u}) {}

    /// Initialize from uint32_t
    __device__ explicit FpExt(uint32_t x) : rep(bb31_t{x}) {}

    /// Convert from Fp to FpExt.
    __device__ explicit FpExt(Fp x) : rep(static_cast<bb31_t>(x)) {}

    /// Explicitly construct an FpExt from parts
    __device__ FpExt(Fp a, Fp b, Fp c, Fp d) {
        elems[0] = a;
        elems[1] = b;
        elems[2] = c;
        elems[3] = d;
    }

    // Implement the addition/subtraction overloads
    __device__ FpExt operator+=(FpExt rhs) {
        rep += rhs.rep;
        return *this;
    }

    __device__ FpExt operator-=(FpExt rhs) {
        rep -= rhs.rep;
        return *this;
    }

    __device__ FpExt operator+(FpExt rhs) const {
        FpExt result = *this;
        result += rhs;
        return result;
    }

    __device__ FpExt operator-(FpExt rhs) const {
        FpExt result = *this;
        result -= rhs;
        return result;
    }

    __device__ FpExt operator-(Fp rhs) const {
        FpExt result = *this;
        result.elems[0] -= rhs;
        return result;
    }

    __device__ FpExt operator-() const { return FpExt() - *this; }

    // Implement the simple multiplication case by the subfield Fp
    // Fp * FpExt is done as a free function due to C++'s operator overloading rules.
    __device__ FpExt operator*=(Fp rhs) {
        rep *= static_cast<bb31_t>(rhs);
        return *this;
    }

    __device__ FpExt operator*(Fp rhs) const {
        FpExt result = *this;
        result *= rhs;
        return result;
    }

    __device__ FpExt operator*=(FpExt rhs) {
        rep *= rhs.rep;
        return *this;
    }

    __device__ FpExt operator*(FpExt rhs) const {
        FpExt result = *this;
        result *= rhs;
        return result;
    }

    // Equality
    __device__ bool operator==(FpExt rhs) const { return rep == rhs.rep; }
    __device__ bool operator!=(FpExt rhs) const { return rep != rhs.rep; }

    __device__ Fp constPart() const { return elems[0]; }
};

/// Overload for case where LHS is Fp (RHS case is handled as a method)
__device__ inline FpExt operator*(Fp a, FpExt b) { return b * a; }

/// Raise an FpExt to a power
__device__ inline FpExt pow(FpExt x, uint32_t n) {
    FpExt r; r.rep = x.rep ^ n;
    return r;
}

template <class I, std::enable_if_t<std::is_integral_v<I>, int> = 0>
__device__ inline FpExt pow(FpExt x, I n) {
    return pow(x, static_cast<uint32_t>(n));
}

/// Compute the multiplicative inverse of an FpExt.
__device__ inline FpExt inv(FpExt in) {
    FpExt result;
    result.rep = in.rep.reciprocal();
    return result;
}

static_assert(sizeof(FpExt) == 16, "FpExt must be 16 bytes");
static_assert(sizeof(FpExt) == sizeof(bb31_4_t), "FpExt and bb31_4_t sizes must match");
static_assert(alignof(FpExt) == alignof(bb31_4_t), "FpExt and bb31_4_t align must match");