> For the complete documentation index, see [llms.txt](https://thias-organization.gitbook.io/p256-documentation/llms.txt). Markdown versions of documentation pages are available by appending `.md` to page URLs; this page is available as [Markdown](https://thias-organization.gitbook.io/p256-documentation/risc-zero-p256-accelerator/field-element-montgomery-form.md).

# Field Element (Montgomery Form)

There are 2 ways we can optimise `FieldElement`. One of it is to retain all field arithmetic in [Montgomery form](https://en.wikipedia.org/wiki/Montgomery_modular_multiplication) as per the original implementation. This was the initial optimisation where we try to preserve most of the functions and only modify things that are necessary as highlighted below. However, this optimisation only led to a cycle count of 7.7 million which is only a 2x speedup.

In the [original](https://github.com/automata-network/RustCrypto-elliptic-curves/blob/risczero/p256/src/arithmetic/field.rs) `field.rs` of RISC Zero's library, `FieldElement` was not defined for specific architecture like [what has been done for k256](/p256-documentation/risc-zero-k256-accelerator/field-element-implementation.md). Hence, the first step is to define new `FieldElement` for 32-bit, 64-bit and zkVM architectures and utilise conditional compilation.

```rust
#[cfg_attr(
    all(target_os = "zkvm", target_arch = "riscv32"),
    path = "field/field_risc0.rs"
)]
#[cfg_attr(
    all(
        not(all(target_os = "zkvm", target_arch = "riscv32")),
        target_pointer_width = "32"
    ),
    path = "field/field32.rs"
)]
#[cfg_attr(target_pointer_width = "64", path = "field/field64.rs")]
mod field_impl;
```

This allows us to abstract away most of the field arithmetic operations and define specific instructions and operations for each type of architecture. For example, the original `sub` function is defined only for the 32-bit architecture:

```rust
/// Returns self - rhs mod p
pub const fn sub(&self, rhs: &Self) -> Self {
    let a = u256_to_u64x4(self.0);
    let b = u256_to_u64x4(rhs.0);
    Self::sub_inner(a[0], a[1], a[2], a[3], 0, b[0], b[1], b[2], b[3], 0).0
}
```

But we can modify it to just

```rust
/// Returns self - rhs mod p
pub const fn sub(&self, rhs: &Self) -> Self {
    Self(field_impl::sub(self.0, rhs.0))
}
```

where we can define `field_impl::sub` differently for different architectures. This is similar for other field arithmetic operations like `add`, `multiply`, etc.

## Double operation

In the original implementation, `double` just involves adding a point to itself. In the modified operation, we make use a new helper function `mul_single` to multiply the point by a single limb integer, if we are in the zkVM architecture

```rust
/// Multiplies by a single-limb integer.
/// Multiplies the magnitude by the same value.
pub fn mul_single(&self, rhs: u32) -> Self {
    Self(field_impl::mul_single(self.0, rhs))
}

/// Returns 2*self.
pub fn double(&self) -> Self {
    if cfg!(all(target_os = "zkvm", target_arch = "riscv32")) {
        self.mul_single(2)
    } else {
        self.add(self)
    }
}
```

This is because we are able to use `modmul_u256_denormalized` function to accelerate the multiplication operation. In `field_risc0.rs`, `mul_single` is defined as follows:

```rust
pub(super) fn mul_single(a: U256, rhs: u32) -> U256 {
    risc0::modmul_u256_denormalized(
        &a,
        &U256::from_words([rhs, 0, 0, 0, 0, 0, 0, 0]),
        &MODULUS_256,
    )
}
```

Changing this single operation saw a slight reduction in cycle count from 16562421 to 15791519 cycles.

## Multiply operation

Modifying this operation is more complicated as it involves breaking down 2 large integers into `limbs` and multiplying them correctly to ensure that it gives the right result. 

The original implementation uses a school book multiplication using `mac` to get the individual limbs and their carries. Then it does a Montgomery reduction to get the final result.

For the 32-bit and 64-bit architecture, we can simplify this process by using `BigInt`'s `mul_wide` function

```rust
pub(super) fn mul(a: U256, b: U256) -> U256 {
    let (lo, hi): (U256, U256) = a.mul_wide(&b);
    montgomery_reduce(lo, hi)
}
```

It is more complicated for zkVM as we cannot just simply use `modmul_u256_denormalized` like what RISC Zero did for the k256 curve. This is because the p256 curve arithmetic is implemented for [montgomery form](https://en.wikipedia.org/wiki/Montgomery_modular_multiplication#Arithmetic_in_Montgomery_form) which is not the same as that in the k256 curve.

Hence, we needed a way to achieve a similar result to 

```rust
let (lo, hi): (U256, U256) = a.mul_wide(&b);
```

while using the accelerator from RISC Zero and then calling `montgomery_reduce`. 

`mul_wide_256` performs a wide multiplication of 2 256-bit integer values using another accelerator `mul_wide_u128`. Because this function only works for 128-bit integers, we need to decompose the original 258-bit integer values $$a$$ and $$b$$ into

$$
\begin{align\*}
a &= a\_1 \* 2^{128} + a\_0\\
b &= b\_1 \* 2^{128} + b\_0
\end{align\*}
$$

Then the product $$p = a \* b$$ can be expressed as

$$
\begin{align\*}
p = &\[a\_1 \times b\_1 ] \times 2^{256} +\\
&\[a\_1 \times b\_0 + a\_0 \times b\_1] \times 2^{128} +\\
&\[a\_0 \times b\_0] \
\end{align\*}
$$

and we can make use of `mul_wide_u128` to calculate the constants 

```rust
// Perform the four multiplications using RISC Zero Accelerator
let p0 = risc0::mul_wide_u128(&a0, &b0);
let p1 = risc0::mul_wide_u128(&a0, &b1);
let p2 = risc0::mul_wide_u128(&a1, &b0);
let p3 = risc0::mul_wide_u128(&a1, &b1);
```

From here, we need to initialise a 512-bit integer and be careful when we carry out multiplications for each limb and account for their carries. The details can be found in the definition of `mul_wide_256`.

### Downstream functions and constants

Because we use `mul_wide_u128`, our `multiply` function cannot be a `const` anymore and this affects many other operations downstream such as `square`, `to_montgomery`, `from_uint_unchecked`, `invert_unchecked`, etc.

This creates an issue when it comes to defining constants in the `PrimeField` implementation. This was the original implementation:

```rust
const TWO_INV: Self = Self::from_u64(2).invert_unchecked();
const MULTIPLICATIVE_GENERATOR: Self = Self::from_u64(6);
const S: u32 = 1;
const ROOT_OF_UNITY: Self =
    Self::from_hex("ffffffff00000001000000000000000000000000fffffffffffffffffffffffe");
const ROOT_OF_UNITY_INV: Self = Self::ROOT_OF_UNITY.invert_unchecked();
const DELTA: Self = Self::from_u64(36);
```

The issue is we cannot reference a non-const function in a constant declaration like `const ROOT_OF_UNITY_INV`. To mitigate this, we can just do a simple check of the hex value using the original implementation by using test cases. 

For example we can check the hex value of `root_of_unity_inv` in the test and use that hex value for directly instead of calling the non-const functions.

```rust
// ...
const ROOT_OF_UNITY: Self = Self(U256::from_be_hex(
    "FFFFFFFE00000002000000000000000000000001FFFFFFFFFFFFFFFFFFFFFFFE",
));
const ROOT_OF_UNITY_INV: Self = Self(U256::from_be_hex(
    "FFFFFFFE00000002000000000000000000000001FFFFFFFFFFFFFFFFFFFFFFFE",
));
// ...

#[test]
fn root_of_unity_constant() {
    let root_of_unity = FieldElement::from_hex(
        "ffffffff00000001000000000000000000000000fffffffffffffffffffffffe",
    );
    let root_of_unity_inv = root_of_unity.invert_unchecked();
    assert_eq!(root_of_unity, FieldElement::ROOT_OF_UNITY);
    assert_eq!(root_of_unity_inv, FieldElement::ROOT_OF_UNITY_INV);
    assert_eq!(
        (FieldElement::ROOT_OF_UNITY * FieldElement::ROOT_OF_UNITY_INV),
        FieldElement::ONE
    )
}
```

This was also done for other constants that were affected.