66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
|
from typing import Sequence, List
|
||
|
|
|
||
|
|
from eth2spec.deneb.mainnet import BLSFieldElement
|
||
|
|
from eth2spec.utils import bls
|
||
|
|
|
||
|
|
from da.kzg_rs.common import G1
|
||
|
|
|
||
|
|
|
||
|
|
def fft_g1(vals: Sequence[G1], roots_of_unity: Sequence[BLSFieldElement], modulus: int) -> List[G1]:
|
||
|
|
if len(vals) == 1:
|
||
|
|
return vals
|
||
|
|
L = fft_g1(vals[::2], roots_of_unity[::2], modulus)
|
||
|
|
R = fft_g1(vals[1::2], roots_of_unity[::2], modulus)
|
||
|
|
o = [bls.Z1() for _ in vals]
|
||
|
|
for i, (x, y) in enumerate(zip(L, R)):
|
||
|
|
y_times_root = bls.multiply(y, roots_of_unity[i])
|
||
|
|
o[i] = (x + y_times_root)
|
||
|
|
o[i + len(L)] = x + -y_times_root
|
||
|
|
return o
|
||
|
|
|
||
|
|
|
||
|
|
def ifft_g1(vals: Sequence[G1], roots_of_unity: Sequence[BLSFieldElement], modulus: int) -> List[G1]:
|
||
|
|
assert len(vals) == len(roots_of_unity)
|
||
|
|
# modular inverse
|
||
|
|
invlen = pow(len(vals), modulus-2, modulus)
|
||
|
|
return [
|
||
|
|
bls.multiply(x, invlen)
|
||
|
|
for x in fft_g1(
|
||
|
|
vals, [roots_of_unity[0], *roots_of_unity[:0:-1]], modulus
|
||
|
|
)
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def _fft(
|
||
|
|
vals: Sequence[BLSFieldElement],
|
||
|
|
roots_of_unity: Sequence[BLSFieldElement],
|
||
|
|
modulus: int,
|
||
|
|
) -> Sequence[BLSFieldElement]:
|
||
|
|
if len(vals) == 1:
|
||
|
|
return vals
|
||
|
|
L = _fft(vals[::2], roots_of_unity[::2], modulus)
|
||
|
|
R = _fft(vals[1::2], roots_of_unity[::2], modulus)
|
||
|
|
o = [BLSFieldElement(0) for _ in vals]
|
||
|
|
for i, (x, y) in enumerate(zip(L, R)):
|
||
|
|
y_times_root = BLSFieldElement((int(y) * int(roots_of_unity[i])) % modulus)
|
||
|
|
o[i] = BLSFieldElement((int(x) + y_times_root) % modulus)
|
||
|
|
o[i + len(L)] = BLSFieldElement((int(x) - int(y_times_root) + modulus) % modulus)
|
||
|
|
return o
|
||
|
|
|
||
|
|
|
||
|
|
def fft(vals, root_of_unity, modulus):
|
||
|
|
assert len(vals) == len(root_of_unity)
|
||
|
|
return _fft(vals, root_of_unity, modulus)
|
||
|
|
|
||
|
|
|
||
|
|
def ifft(vals, roots_of_unity, modulus):
|
||
|
|
assert len(vals) == len(roots_of_unity)
|
||
|
|
# modular inverse
|
||
|
|
invlen = pow(len(vals), modulus-2, modulus)
|
||
|
|
return [
|
||
|
|
BLSFieldElement((int(x) * invlen) % modulus)
|
||
|
|
for x in _fft(
|
||
|
|
vals, [roots_of_unity[0], *roots_of_unity[:0:-1]], modulus
|
||
|
|
)
|
||
|
|
]
|