use crate::{
add,
arch::word::{SignedWord, Word},
helper_macros::debug_assert_zero,
math,
memory::{self, Memory},
sqr,
Sign::*,
};
use alloc::alloc::Layout;
pub const MIN_LEN: usize = 3;
pub fn memory_requirement_up_to(n: usize) -> Layout {
let num_words = 3 * n + 2 * (math::ceil_log2(n) as usize);
memory::array_layout::<Word>(num_words)
}
pub fn square(b: &mut [Word], a: &[Word], memory: &mut Memory) {
let n = a.len();
debug_assert!(n >= MIN_LEN && b.len() == 2 * n);
let mid = (n + 1) / 2;
let (a_lo, a_hi) = a.split_at(mid);
let mut carry: SignedWord = 0;
let mut carry_c0: SignedWord = 0; let mut carry_c1: SignedWord = 0;
{
let (p0, mut mem) = memory.allocate_slice_fill::<Word>(2 * mid, 0);
sqr::sqr(p0, a_lo, &mut mem);
carry_c0 += add::add_signed_same_len_in_place(&mut b[..2 * mid], Positive, p0);
carry_c1 += add::add_signed_same_len_in_place(&mut b[mid..3 * mid], Positive, p0);
}
{
let p2_len = 2 * (n - mid);
let (p2, mut mem) = memory.allocate_slice_fill::<Word>(p2_len, 0);
sqr::sqr(p2, a_hi, &mut mem);
carry += add::add_signed_same_len_in_place(&mut b[2 * mid..], Positive, p2);
carry_c1 += add::add_signed_in_place(&mut b[mid..3 * mid], Positive, p2);
}
{
let (diff, mut mem) = memory.allocate_slice_copy(a_lo);
let diff_sign = add::sub_in_place_with_sign(diff, a_hi);
if diff_sign == Negative {
diff[..(n - mid)].copy_from_slice(a_hi);
diff[(n - mid)..].fill(0);
debug_assert_zero!(add::sub_in_place(diff, a_lo));
}
let (diff_sq, mut mem) = mem.allocate_slice_fill::<Word>(2 * mid, 0);
sqr::sqr(diff_sq, diff, &mut mem);
carry_c1 += add::add_signed_same_len_in_place(&mut b[mid..3 * mid], Negative, diff_sq);
}
carry_c1 += add::add_signed_word_in_place(&mut b[2 * mid..3 * mid], carry_c0);
carry += add::add_signed_word_in_place(&mut b[3 * mid..], carry_c1);
debug_assert!(carry.abs() <= 1);
}