1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
use crate::{
arch::{
self,
word::{DoubleWord, Word},
},
math::{mul_add_2carry, mul_add_carry},
primitive::{double_word, split_dword},
};
pub fn square(b: &mut [Word], a: &[Word]) {
debug_assert!(b.len() == a.len() * 2);
let n = a.len();
// `b` is zero on entry (contract from sqr::sqr).
//
// a^2 = (diagonal) + 2 * (upper triangle)
// The upper triangle is accumulated first, then doubled in place while the
// diagonal a[i]^2 terms are added, fusing the shift into a single pass.
// ---- first step: upper (off-diagonal) triangle ----
//
// For each pair of source limbs (a[i], a[i+1]) we add their off-diagonal
// products against the shared suffix a[i+2..]:
// * the lone corner product a[i] * a[i+1] (column 2i+1), and
// * (a[i] + a[i+1]*B) * a[i+2..] via a two-word multiply-accumulate kernel
// (two independent mul-accumulate chains, one accumulator load/store
// per two multiplier limbs), mirroring the multiplication basecase.
// This halves the accumulator memory traffic of the triangle versus a
// one-limb-per-row sweep.
//
// Carries are propagated in-place: each pair's high product words land at
// b[n+i], b[n+i+1] and the next pair's kernel sweep re-reads b[n+i+1] and
// propagates further, so only a single pending carry word `cy` ever escapes
// a sweep. (The triangle T < B^(2n-1) for any realistic n, so `cy` resolves
// to 0; it is folded into the final aggregation exactly like the old code's
// trailing carry.)
let mut cy: Word = 0;
let mut i = 0;
while i + 2 < n {
let m0 = a[i];
let m1 = a[i + 1];
let rhs = &a[i + 2..];
let base = 2 * i + 2; // kernel covers columns 2i+2 .. n+i
// corner a[i]*a[i+1] at column 2i+1 (= base-1)
let (lo, hi) = mul_add_carry(m0, m1, 0);
let (v, c1) = b[base - 1].overflowing_add(lo);
b[base - 1] = v;
// hi + c1 fits in a Word (hi <= Word::MAX-1, c1 in {0,1}); seed the
// kernel's low carry chain with it so no separate ripple is needed.
let init_lo = hi.wrapping_add(Word::from(c1));
// two-word kernel: b[base..] += (m0 + m1*B) * rhs
let mut carry_lo = init_lo;
let mut carry_hi: Word = 0;
for (x, &y) in b[base..base + rhs.len()].iter_mut().zip(rhs.iter()) {
(*x, carry_lo) = mul_add_2carry(y, m0, *x, carry_lo);
(carry_lo, carry_hi) = mul_add_2carry(y, m1, carry_lo, carry_hi);
}
// add the high product words plus the pending carry at columns n+i, n+i+1
let p = base + rhs.len(); // = n + i
let (s, c) = arch::add::add_with_carry(b[p], carry_lo, cy);
b[p] = s;
let (s, c2) = arch::add::add_with_carry(b[p + 1], carry_hi, Word::from(c));
b[p + 1] = s;
cy = Word::from(c2);
i += 2;
}
// leftover single row when n is even: a[n-2]*a[n-1] at column 2n-3
if i == n - 2 {
let (lo, hi) = mul_add_carry(a[i], a[i + 1], 0);
let base = 2 * i + 1; // = 2n-3
let (v, c1) = b[base].overflowing_add(lo);
b[base] = v;
// column 2n-2 = base+1 also receives the pending carry `cy`
let (s, c) = arch::add::add_with_carry(b[base + 1], hi, Word::from(c1));
let (s, c2) = arch::add::add_with_carry(s, cy, Word::from(c));
b[base + 1] = s;
cy = Word::from(c2);
}
// ---- second step: double the triangle and add the diagonal a[i]^2 ----
let (mut c1, mut c2) = (false, false);
for (m, b01) in a.iter().zip(b.chunks_exact_mut(2)) {
let b0 = b01.first().unwrap();
let b1 = b01.last().unwrap();
// new [b0, b1] = m^2 + 2 * [b0, b1] + c1 + c2
let (s0, s1) = mul_add_2carry(*m, *m, *b0, *b0);
let s = double_word(s0, s1);
let wb1 = double_word(0, *b1);
let (s, oc1) = s.overflowing_add(wb1 + c1 as DoubleWord);
let (s, oc2) = s.overflowing_add(wb1 + c2 as DoubleWord);
let (s0, s1) = split_dword(s);
*b01.first_mut().unwrap() = s0;
*b01.last_mut().unwrap() = s1;
c1 = oc1;
c2 = oc2;
}
// aggregate carry bits (cy is the triangle's trailing carry, ~always 0)
*b.last_mut().unwrap() += cy + c1 as Word + c2 as Word;
}
#[cfg(test)]
mod tests {
use crate::arch::word::Word;
use crate::UBig;
use alloc::vec;
use alloc::vec::Vec;
fn lcg_words(seed: u64, words: usize) -> UBig {
let mut limbs: Vec<Word> = Vec::with_capacity(words);
let mut s = seed.wrapping_add(words as u64);
for _ in 0..words {
s = s
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
limbs.push(s as Word);
}
if let Some(top) = limbs.last_mut() {
*top |= 1; // ensure full width (no leading-zero stripping)
}
UBig::from_words(&limbs)
}
#[test]
fn sqr_matches_mul_simple_range() {
// sizes 2..=128 span simple directly (<=30) and as the basecase of karatsuba/toom
for words in 2..=128 {
for seed in 0..6u64 {
let a = lcg_words(seed, words);
assert_eq!(a.sqr(), &a * &a, "sqr != mul at words={words} seed={seed}");
}
}
}
#[test]
fn sqr_matches_mul_adversarial() {
// patterns that stress carry propagation: all-ones, single high bit, mixed
for words in [2, 3, 4, 5, 7, 8, 15, 16, 17, 30, 31, 32, 48] {
let ones = vec![Word::MAX; words];
let a = UBig::from_words(&ones);
assert_eq!(a.sqr(), &a * &a, "all-ones mismatch at words={words}");
let mut hi: Vec<Word> = vec![0; words];
hi[words - 1] = Word::MAX;
let a = UBig::from_words(&hi);
assert_eq!(a.sqr(), &a * &a, "single-high-bit mismatch at words={words}");
let mixed: Vec<Word> = (0..words)
.map(|k| {
if k % 2 == 0 {
Word::MAX
} else {
Word::from(1u8)
}
})
.collect();
let a = UBig::from_words(&mixed);
assert_eq!(a.sqr(), &a * &a, "mixed mismatch at words={words}");
}
}
}