use std::alloc::{Allocator, Global};
use std::ops::Range;
use std::sync::Arc;
use feanor_math::algorithms::eea::signed_gcd;
use feanor_math::algorithms::matmul::ComputeInnerProduct;
use feanor_math::divisibility::DivisibilityRingStore;
use feanor_math::group::AbelianGroupStore;
use feanor_math::iters::multi_cartesian_product;
use feanor_math::matrix::OwnedMatrix;
use feanor_math::primitive_int::*;
use feanor_math::ring::*;
use feanor_math::rings::extension::*;
use feanor_math::rings::zn::zn_64::*;
use feanor_math::seq::*;
use feanor_math::algorithms::linsolve::LinSolveRing;
use feanor_math::homomorphism::*;
use tracing::instrument;
use crate::circuit::{Coefficient, PlaintextCircuit};
use crate::number_ring::galois::*;
use crate::number_ring::hypercube::isomorphism::*;
use crate::number_ring::hypercube::structure::*;
use crate::number_ring::quotient_by_int::NumberRingQuotientByIntBase;
use crate::number_ring::*;
use super::trace::extract_linear_map;
use crate::{NiceZn};
pub struct MatmulTransform<R>
where R: NumberRingQuotient,
<<R as RingExtension>::BaseRing as RingStore>::Type: NiceZn
{
data: Vec<(
// a representation of the used Galois automorphism w.r.t. the hypercube structure;
// different choices for the same Galois automorphism will give the same transform,
// but may have an effect on how well the transform can be optimized;
// the first entry is for the frobenius
Box<[i64]>,
// the coefficient of the Galois automorphism, to be multiplied to the input element
// after applying the automorphism
R::Element
)>
}
impl<R> MatmulTransform<R>
where R: NumberRingQuotient,
<<R as RingExtension>::BaseRing as RingStore>::Type: NiceZn
{
pub fn eq<S>(&self, ring: S, H: &HypercubeStructure, other: &Self) -> bool
where S: Copy + RingStore<Type = R>
{
self.check_valid(ring, H);
other.check_valid(ring, H);
if self.data.len() != other.data.len() {
return false;
}
for (self_d, other_d) in self.data.iter().zip(other.data.iter()) {
if !H.galois_group().eq_el(&H.map_incl_frobenius(&self_d.0), &H.map_incl_frobenius(&other_d.0)) {
return false;
}
if !ring.eq_el(&self_d.1, &other_d.1) {
return false;
}
}
return true;
}
#[instrument(skip_all)]
pub fn matmul1d<G, S>(H: &HypercubeIsomorphism<S>, dim_index: usize, matrix: G) -> MatmulTransform<R>
where G: Fn(usize, usize, &[usize]) -> El<SlotRingOf<S>>,
S: RingStore<Type = R>
{
let m = H.hypercube().dim_length(dim_index) as i64;
let mut result = MatmulTransform {
data: ((1 - m)..m).into_iter().filter_map(|s| {
let coeffs = H.hypercube().hypercube_iter(|idxs: &[usize]| if idxs[dim_index] as i64 >= s && idxs[dim_index] as i64 - s < m {
Some(matrix(idxs[dim_index], (idxs[dim_index] as i64 - s) as usize, idxs))
} else {
None
}).collect::<Vec<_>>();
if coeffs.iter().all(|x| x.is_none() || H.slot_ring().is_zero(x.as_ref().unwrap())) {
None
} else {
Some((
[0].into_iter().chain((0..H.hypercube().dim_count()).map(|j| if j == dim_index { s } else { 0 })).collect::<Vec<_>>().into_boxed_slice(),
H.from_slot_values(coeffs.into_iter().map(|x| x.unwrap_or_else(|| H.slot_ring().zero()))),
))
}
}).collect()
};
result.canonicalize(H.ring(), H.hypercube());
return result;
}
fn slot_ring_frobenius<'a, S>(H: &'a HypercubeIsomorphism<S>) -> Arc<dyn 'a + Fn(&El<SlotRingOf<S>>, usize) -> El<SlotRingOf<S>>>
where S: 'a + RingStore<Type = R>,
R: 'a
{
let d = H.hypercube().d();
let mut generator_frobenius_conjugates = Vec::with_capacity(d);
let mut current = H.slot_ring().canonical_gen();
for _ in 0..d {
generator_frobenius_conjugates.push((0..d).map(|j| H.slot_ring().pow(H.slot_ring().clone_el(¤t), j)).collect::<Vec<_>>());
current = H.slot_ring().pow(current, H.galois_group().representative(H.hypercube().p()) as usize);
}
#[instrument(skip_all)]
fn apply_frobenius<S>(generator_frobenius_conjugates: &Vec<Vec<El<SlotRingOf<S>>>>, slot_ring: &SlotRingOf<S>, d: usize, x: &El<SlotRingOf<S>>, count: usize) -> El<SlotRingOf<S>>
where S: RingStore,
S::Type: NumberRingQuotient,
<<S::Type as RingExtension>::BaseRing as RingStore>::Type: NiceZn
{
let mut result = slot_ring.zero();
let x_wrt_basis = slot_ring.wrt_canonical_basis(x);
for i in 0..d {
result = slot_ring.inclusion().fma_map(&generator_frobenius_conjugates[count][i], &x_wrt_basis.at(i), result);
}
return result;
}
return Arc::new(move |x: &El<SlotRingOf<S>>, count: usize| apply_frobenius::<S>(&generator_frobenius_conjugates, H.slot_ring(), d, x, count));
}
#[instrument(skip_all)]
pub fn blockmatmul0d<G, S>(H: &HypercubeIsomorphism<S>, matrix: G) -> MatmulTransform<R>
where G: Fn(usize, usize, &[usize]) -> El<R::BaseRing>,
S: RingStore<Type = R>
{
let d = H.slot_ring().rank();
let extract_coeff_factors = (0..d).map(|j|
extract_linear_map(H.slot_ring(), |x| H.slot_ring().wrt_canonical_basis(&x).at(j))
).collect::<Vec<_>>();
let frobenius = Self::slot_ring_frobenius(H);
let extract_coeff_factor_conjugates = extract_coeff_factors.iter().map(|c|
(0..d).map(|i| frobenius(c, i)).collect::<Vec<_>>()
).collect::<Vec<_>>();
let mut result = MatmulTransform {
data: (0..d).filter_map(|frobenius_index| {
let coeffs = H.hypercube().hypercube_iter(|idxs| {
<_ as ComputeInnerProduct>::inner_product_ref_fst(H.slot_ring().get_ring(), (0..d).map(|l| (
&extract_coeff_factor_conjugates[l][frobenius_index],
H.slot_ring().from_canonical_basis((0..d).map(|k| matrix(k, l, idxs)))
)))
}).collect::<Vec<_>>();
if coeffs.iter().all(|x| H.slot_ring().is_zero(x)) {
None
} else {
Some((
[frobenius_index as i64].into_iter().chain((0..H.hypercube().dim_count()).map(|_| 0)).collect::<Vec<_>>().into_boxed_slice(),
H.from_slot_values(coeffs),
))
}
}).collect()
};
result.canonicalize(H.ring(), H.hypercube());
return result;
}
#[instrument(skip_all)]
pub fn blockmatmul1d<G, S>(H: &HypercubeIsomorphism<S>, dim_index: usize, matrix: G) -> MatmulTransform<R>
where G: Fn((usize, usize), (usize, usize), &[usize]) -> El<R::BaseRing>,
S: RingStore<Type = R>
{
let m = H.hypercube().dim_length(dim_index) as i64;
let d = H.slot_ring().rank();
let extract_coeff_factors = (0..d).map(|j|
extract_linear_map(H.slot_ring(), |x| H.slot_ring().wrt_canonical_basis(&x).at(j))
).collect::<Vec<_>>();
let frobenius = Self::slot_ring_frobenius(H);
let extract_coeff_factor_conjugates = extract_coeff_factors.iter().map(|c|
(0..d).map(|i| frobenius(c, i)).collect::<Vec<_>>()
).collect::<Vec<_>>();
let mut result = MatmulTransform {
data: ((1 - m)..m).flat_map(|s| (0..d).map(move |frobenius_index| (s, frobenius_index))).filter_map(|(s, frobenius_index)| {
let coeffs = H.hypercube().hypercube_iter(|idxs| if idxs[dim_index] as i64 >= s && idxs[dim_index] as i64 - s < m {
let i = idxs[dim_index];
let j = (idxs[dim_index] as i64 - s) as usize;
<_ as ComputeInnerProduct>::inner_product_ref_fst(H.slot_ring().get_ring(), (0..d).map(|l| (
&extract_coeff_factor_conjugates[l][frobenius_index],
H.slot_ring().from_canonical_basis((0..d).map(|k| matrix((i, k), (j, l), idxs)))
)))
} else {
H.slot_ring().zero()
}).collect::<Vec<_>>();
if coeffs.iter().all(|x| H.slot_ring().is_zero(x)) {
None
} else {
let value = H.from_slot_values(coeffs);
Some((
[frobenius_index as i64].into_iter().chain((0..H.hypercube().dim_count()).map(|j| if j == dim_index { s } else { 0 })).collect::<Vec<_>>().into_boxed_slice(),
value
))
}
}).collect()
};
result.canonicalize(H.ring(), H.hypercube());
return result;
}
pub fn switch_ring<H, S>(&self, hom: H) -> MatmulTransform<S>
where S: NumberRingQuotient,
<S::BaseRing as RingStore>::Type: NiceZn,
H: Homomorphism<R, S>
{
MatmulTransform::<S> {
data: self.data.iter().map(|(g, coeff)| (g.clone(), hom.map_ref(coeff))).collect::<Vec<(_, S::Element)>>()
}
}
#[instrument(skip_all)]
pub fn inverse<S>(&self, H: &HypercubeIsomorphism<S>) -> Self
where S: RingStore<Type = R>
{
self.check_valid(H.ring(), H.hypercube());
let Gal = H.galois_group().parent();
let original_automorphisms = self.data.iter().map(|(g, _)| g.clone());
let inverse_automorphisms = original_automorphisms.clone().map(|g|
g.iter().map(|i| -i).collect::<Vec<_>>().into_boxed_slice()
).collect::<Vec<_>>();
let mut composed_automorphisms = original_automorphisms.clone().flat_map(|g| inverse_automorphisms.iter().map(move |s|
g.iter().zip(s.iter()).map(|(i1, i2)| i1 + i2).collect::<Vec<_>>().into_boxed_slice()
)).collect::<Vec<_>>();
composed_automorphisms.sort_unstable_by_key(|g| Gal.representative(&H.hypercube().map_incl_frobenius(g)));
composed_automorphisms.dedup_by(|a, b| Gal.eq_el(&H.hypercube().map_incl_frobenius(a), &H.hypercube().map_incl_frobenius(b)));
let mut matrix: OwnedMatrix<_> = OwnedMatrix::zero(composed_automorphisms.len(), inverse_automorphisms.len(), H.ring());
for (i, g) in original_automorphisms.enumerate() {
for (j, s) in inverse_automorphisms.iter().enumerate() {
let row_index = composed_automorphisms.binary_search_by_key(
&Gal.representative(&H.hypercube().map_incl_frobenius(&g.iter().zip(s.iter()).map(|(i1, i2)| i1 + i2).collect::<Vec<_>>())),
|g| Gal.representative(&H.hypercube().map_incl_frobenius(g))
).unwrap();
let entry = H.ring().get_ring().apply_galois_action(&self.data[i].1, &H.hypercube().map_incl_frobenius(s));
*matrix.at_mut(row_index, j) = entry;
}
}
let mut matrix_by_slots = (0..matrix.row_count()).map(|i| (0..matrix.col_count()).map(|j| H.get_slot_values(matrix.at(i, j))).collect::<Vec<_>>()).collect::<Vec<_>>();
let mut result_by_slots = (0..inverse_automorphisms.len()).map(|_| Vec::new()).collect::<Vec<_>>();
let mut lhs: OwnedMatrix<_> = OwnedMatrix::zero(matrix.row_count(), matrix.col_count(), H.slot_ring());
let mut rhs: OwnedMatrix<_> = OwnedMatrix::zero(matrix.row_count(), 1, H.slot_ring());
let mut sol: OwnedMatrix<_> = OwnedMatrix::zero(matrix.col_count(), 1, H.slot_ring());
for _ in H.hypercube().element_iter() {
for i in 0..matrix.row_count() {
for j in 0..matrix.col_count() {
*lhs.at_mut(i, j) = matrix_by_slots[i][j].next().unwrap();
}
}
assert!(Gal.is_identity(&H.hypercube().map_incl_frobenius(&composed_automorphisms[0])));
*rhs.at_mut(0, 0) = H.slot_ring().one();
for j in 1..matrix.row_count() {
*rhs.at_mut(j, 0) = H.slot_ring().zero();
}
H.slot_ring().get_ring().solve_right(lhs.data_mut(), rhs.data_mut(), sol.data_mut(), Global).assert_solved();
for j in 0..matrix.col_count() {
result_by_slots[j].push(H.slot_ring().clone_el(sol.at(j, 0)));
}
}
let result = result_by_slots.into_iter().map(|coeff_by_slots| H.from_slot_values(coeff_by_slots.into_iter())).collect::<Vec<_>>();
let mut result = Self {
data: self.data.iter().zip(result.into_iter()).map(|((g, _), coeff)| (
g.iter().map(|i| -i).collect::<Vec<_>>().into_boxed_slice(),
coeff
)).collect()
};
result.canonicalize(H.ring(), H.hypercube());
return result;
}
fn check_valid<S>(&self, _ring: S, H: &HypercubeStructure)
where S: RingStore<Type = R>
{
let Gal = H.galois_group().parent();
assert!(self.data.is_sorted_by_key(|(g, _)| Gal.representative(&H.map_incl_frobenius(g))));
for (i, (g, _)) in self.data.iter().enumerate() {
assert_eq!(H.dim_count() + 1, g.len());
for (j, (s, _)) in self.data.iter().enumerate() {
assert!(i == j || !Gal.eq_el(&H.map_incl_frobenius(g), &H.map_incl_frobenius(s)));
}
}
}
#[instrument(skip_all)]
pub fn compose<S>(&self, ring: S, H: &HypercubeStructure, run_first: &MatmulTransform<R>) -> Self
where S: Copy + RingStore<Type = R>
{
self.check_valid(ring, H);
run_first.check_valid(ring, H);
let mut result = Self {
data: self.data.iter().flat_map(|(self_g, self_coeff)| run_first.data.iter().map(|(first_g, first_coeff)| (
self_g.iter().zip(first_g.iter()).map(|(i1, i2)| i1 + i2).collect::<Vec<_>>().into_boxed_slice(),
ring.mul_ref_snd(ring.get_ring().apply_galois_action(first_coeff, &H.map_incl_frobenius(self_g)), self_coeff)
))).collect()
};
result.canonicalize(ring, &H);
return result;
}
#[instrument(skip_all)]
pub fn mult_scalar_slots<S>(H: &HypercubeIsomorphism<S>, scalar: &El<SlotRingOf<S>>) -> MatmulTransform<R>
where S: RingStore<Type = R>
{
Self::mult_ring_element(H.ring(), H.hypercube(), &H.from_slot_values((0..H.slot_count()).map(|_| H.slot_ring().clone_el(scalar))))
}
#[instrument(skip_all)]
pub fn mult_ring_element<S>(ring: S, H: &HypercubeStructure, factor: &El<S>) -> MatmulTransform<R>
where S: RingStore<Type = R>
{
return MatmulTransform {
data: vec![(
[0].into_iter().chain((0..H.dim_count()).map(|_| 0)).collect::<Vec<_>>().into_boxed_slice(),
ring.clone_el(factor)
)]
};
}
#[instrument(skip_all)]
pub fn identity<S>(ring: S, H: &HypercubeStructure) -> Self
where S: Copy + RingStore<Type = R>
{
Self::mult_ring_element(ring, H, &ring.one())
}
#[instrument(skip_all)]
pub fn shift<S>(ring: S, H: &HypercubeStructure, positions: &[i64]) -> Self
where S: Copy + RingStore<Type = R>
{
assert_eq!(H.dim_count(), positions.len());
Self {
data: vec![(
[0].into_iter().chain(positions.iter().copied()).collect::<Vec<_>>().into_boxed_slice(),
ring.one()
)]
}
}
#[instrument(skip_all)]
pub fn linear_combine_shifts<V, I, S>(H: &HypercubeIsomorphism<S>, summands: I) -> Self
where I: Iterator<Item = (V, El<S>)>,
V: VectorFn<i64>,
S: RingStore<Type = R>
{
let mut result = Self {
data: summands
.inspect(|(positions, _)| assert_eq!(H.hypercube().dim_count(), positions.len()))
.map(|(positions, factor)| (
[0].into_iter().chain(positions.iter()).collect::<Vec<_>>().into_boxed_slice(),
factor
)).collect()
};
result.canonicalize(H.ring(), H.hypercube());
return result;
}
#[instrument(skip_all)]
fn canonicalize<S>(&mut self, ring: S, H: &HypercubeStructure)
where S: Copy + RingStore<Type = R>
{
self.data.sort_unstable_by_key(|(g, _)| H.galois_group().representative(&H.map_incl_frobenius(g)));
let mut normalize_all_shifts = false;
self.data.dedup_by(|second, first| {
if H.galois_group().eq_el(&H.map_incl_frobenius(&second.0), &H.map_incl_frobenius(&first.0)) {
ring.add_assign_ref(&mut first.1, &second.1);
normalize_all_shifts = true;
return true;
} else {
return false;
}
});
if normalize_all_shifts {
for (g, _) in self.data.iter_mut() {
*g = H.std_preimage(&H.map_incl_frobenius(g)).iter().map(|i| *i as i64).collect::<Vec<_>>().into_boxed_slice();
}
}
}
fn compute_automorphisms_per_dimension<S>(&self, ring: S, H: &HypercubeStructure) -> (Vec<i64>, Vec<i64>, Vec<usize>, Vec<usize>)
where S: Copy + RingStore<Type = R>
{
self.check_valid(ring, H);
let mut max_step: Vec<i64> = Vec::new();
let mut min_step: Vec<i64> = Vec::new();
let mut gcd_step: Vec<usize> = Vec::new();
let mut sizes: Vec<usize> = Vec::new();
for i in 0..=H.dim_count() {
max_step.push(self.data.iter().map(|(g, _)| g[i]).max().unwrap());
min_step.push(self.data.iter().map(|(g, _)| g[i]).min().unwrap());
let gcd = self.data.iter().map(|(g, _)| g[i]).fold(0, |a, b| signed_gcd(a, b, StaticRing::<i64>::RING)).try_into().unwrap();
if gcd == 0 {
gcd_step.push(usize::MAX);
} else {
gcd_step.push(gcd);
}
assert!(gcd_step[i] > 0);
if gcd_step[i] != usize::MAX {
sizes.push(StaticRing::<i64>::RING.checked_div(&(max_step[i] - min_step[i] + gcd_step[i] as i64), &gcd_step[i].try_into().unwrap()).unwrap().try_into().unwrap());
} else {
sizes.push(1);
}
}
return (
max_step,
min_step,
gcd_step,
sizes
);
}
#[instrument(skip_all)]
pub fn baby_step_giant_step_params<V>(automorphisms_per_dim: V, preferred_baby_steps: usize) -> BabyStepGiantStepParams
where V: VectorFn<usize>
{
let mut baby_step_dims = 0;
let mut baby_steps = 1;
for i in 0..automorphisms_per_dim.len() {
let new_steps = baby_steps * automorphisms_per_dim.at(i);
if new_steps >= preferred_baby_steps {
break;
}
baby_step_dims += 1;
baby_steps = new_steps;
}
let mixed_dim_i = baby_step_dims;
let giant_step_start_dim = mixed_dim_i + 1;
let mixed_dim_baby_steps = (preferred_baby_steps - 1) / baby_steps + 1;
let baby_steps = baby_steps * mixed_dim_baby_steps;
assert!(baby_steps >= preferred_baby_steps);
let giant_steps = (giant_step_start_dim..automorphisms_per_dim.len()).map(|i| automorphisms_per_dim.at(i)).product::<usize>() * ((automorphisms_per_dim.at(mixed_dim_i) - 1) / mixed_dim_baby_steps + 1);
return BabyStepGiantStepParams {
pure_baby_step_dimensions: 0..baby_step_dims,
pure_giant_step_dimensions: giant_step_start_dim..automorphisms_per_dim.len(),
mixed_step_dimension: mixed_dim_i,
mixed_step_dimension_baby_steps: mixed_dim_baby_steps,
hoisted_automorphism_count: baby_steps - 1,
unhoisted_automorphism_count: giant_steps - 1
};
}
#[instrument(skip_all)]
pub fn to_circuit<S>(self, ring: S, H: &HypercubeStructure) -> PlaintextCircuit<R>
where S: Copy + RingStore<Type = R>
{
self.check_valid(ring, H);
let (_, _, _, sizes) = self.compute_automorphisms_per_dimension(ring, H);
const UNHOISTED_AUTO_COUNT_OVERHEAD: usize = 3;
let preferred_baby_steps = (1..=(sizes.iter().copied().product::<usize>())).min_by_key(|preferred_baby_steps| {
let params = Self::baby_step_giant_step_params(sizes.as_fn().map_fn(|s| *s), *preferred_baby_steps);
return params.hoisted_automorphism_count + params.unhoisted_automorphism_count * UNHOISTED_AUTO_COUNT_OVERHEAD;
}).unwrap();
return self.to_circuit_with_baby_steps(ring, H, preferred_baby_steps);
}
#[instrument(skip_all)]
pub fn to_circuit_many<S>(ring: S, H: &HypercubeStructure, transforms: Vec<Self>) -> PlaintextCircuit<R>
where S: Copy + RingStore<Type = R>
{
transforms.into_iter().fold(PlaintextCircuit::identity(1, ring), |current, next| next.to_circuit(ring, H).compose(current, ring))
}
#[instrument(skip_all)]
pub fn to_circuit_with_baby_steps<S>(self, ring: S, H: &HypercubeStructure, preferred_baby_steps: usize) -> PlaintextCircuit<R>
where S: Copy + RingStore<Type = R>
{
self.check_valid(ring, H);
let (max_step, min_step, gcd_step, sizes) = self.compute_automorphisms_per_dimension(ring, H);
let params = Self::baby_step_giant_step_params((0..sizes.len()).map_fn(|i| sizes[i]), preferred_baby_steps);
let mixed_dim_i = params.mixed_step_dimension;
let mixed_dim_baby_steps = params.mixed_step_dimension_baby_steps;
let mixed_dim_steps_on_top = StaticRing::<i64>::RING.checked_div(&max_step[mixed_dim_i], &gcd_step[mixed_dim_i].try_into().unwrap()).unwrap();
let mixed_dim_steps_on_bottom = StaticRing::<i64>::RING.checked_div(&min_step[mixed_dim_i], &gcd_step[mixed_dim_i].try_into().unwrap()).unwrap();
let mixed_dim_baby_steps_on_top = mixed_dim_steps_on_top % mixed_dim_baby_steps as i64;
let mixed_dim_baby_steps_on_bottom = mixed_dim_baby_steps_on_top - mixed_dim_baby_steps as i64 + 1;
let mixed_dim_giant_steps_on_top = StaticRing::<i64>::RING.checked_div(&(mixed_dim_steps_on_top - mixed_dim_baby_steps_on_top), &(mixed_dim_baby_steps as i64)).unwrap();
let mixed_dim_giant_steps_on_bottom = (mixed_dim_steps_on_bottom - mixed_dim_baby_steps_on_bottom).div_floor(mixed_dim_baby_steps as i64);
let mixed_dim_giant_step_range = ((mixed_dim_giant_steps_on_bottom * mixed_dim_baby_steps as i64 * gcd_step[mixed_dim_i] as i64)..=(mixed_dim_giant_steps_on_top * mixed_dim_baby_steps as i64 * gcd_step[mixed_dim_i] as i64)).step_by(mixed_dim_baby_steps * gcd_step[mixed_dim_i]);
let mixed_dim_baby_step_range = ((mixed_dim_baby_steps_on_bottom * gcd_step[mixed_dim_i] as i64)..=(mixed_dim_baby_steps_on_top * gcd_step[mixed_dim_i] as i64)).step_by(gcd_step[mixed_dim_i]);
debug_assert!(mixed_dim_baby_steps_on_top + mixed_dim_giant_steps_on_top * mixed_dim_baby_steps as i64 == mixed_dim_steps_on_top);
debug_assert!(mixed_dim_baby_steps_on_bottom + mixed_dim_giant_steps_on_bottom * mixed_dim_baby_steps as i64 <= mixed_dim_steps_on_bottom);
let giant_step_range_iters = [mixed_dim_giant_step_range].into_iter()
.chain(params.pure_giant_step_dimensions.clone().map(|i| (min_step[i]..=max_step[i]).step_by(gcd_step[i])));
let baby_step_range_iters = params.pure_baby_step_dimensions.clone().map(|i| (min_step[i]..=max_step[i]).step_by(gcd_step[i]))
.chain([mixed_dim_baby_step_range]);
let shift_or_frobenius = |dim_or_frobenius: usize, steps: i64| if dim_or_frobenius == 0 {
H.frobenius(steps)
} else {
H.map_1d(dim_or_frobenius - 1, steps)
};
let identity = shift_or_frobenius(0, 0);
let giant_steps_galois_els = multi_cartesian_product(giant_step_range_iters, |indices| {
indices[1..].iter()
.enumerate()
.map(|(i, s)| shift_or_frobenius(i + params.pure_giant_step_dimensions.start, *s))
.fold(shift_or_frobenius(mixed_dim_i, indices[0]), |a, b| H.galois_group().op(a, b))
}, |_, x| *x)
.map(|g_el| if H.galois_group().is_identity(&g_el) { None } else { Some(g_el) })
.collect::<Vec<_>>();
let baby_steps_galois_els = multi_cartesian_product(baby_step_range_iters, move |indices| {
indices.iter()
.enumerate()
.map(|(i, s)| shift_or_frobenius(i, *s))
.fold(identity.clone(), |a, b| H.galois_group().op(a, b))
}, |_, x| *x)
.collect::<Vec<_>>();
debug_assert_eq!(params.hoisted_automorphism_count, baby_steps_galois_els.len() - 1);
debug_assert!(params.unhoisted_automorphism_count == giant_steps_galois_els.len() - 1 || params.unhoisted_automorphism_count == giant_steps_galois_els.len());
let mut lin_transform_data = self.data;
let compiled_coeffs: Vec<Vec<Coefficient<_>>> = giant_steps_galois_els.iter().map(|gs_el| baby_steps_galois_els.iter().map(|bs_el| {
let gs_el = gs_el.clone().unwrap_or(H.galois_group().identity());
let total_el = H.galois_group().op_ref(&gs_el, bs_el);
let mut coeff = None;
lin_transform_data.retain(|(g, c)| if H.galois_group().eq_el(&H.map_incl_frobenius(g), &total_el) {
debug_assert!(coeff.is_none());
coeff = Some(ring.clone_el(c));
false
} else {
true
});
if coeff.is_none() {
return Coefficient::Zero;
} else {
return Coefficient::Other(ring.get_ring().apply_galois_action(coeff.as_ref().unwrap(), &H.galois_group().inv(&gs_el)));
}
}).collect::<Vec<_>>()).collect::<Vec<_>>();
let baby_step_circuit = PlaintextCircuit::gal_many(&baby_steps_galois_els, H.galois_group(), ring);
let mut current = PlaintextCircuit::constant(ring.zero(), ring).tensor(baby_step_circuit, ring);
for (g, coeffs) in giant_steps_galois_els.iter().cloned().zip(compiled_coeffs) {
debug_assert_eq!(baby_steps_galois_els.len() + 1, current.output_count());
let summand = if let Some(g) = g {
let galois_of_lin_transform = PlaintextCircuit::gal(g, &H.galois_group(), ring).compose(PlaintextCircuit::linear_transform(&coeffs, ring), ring);
PlaintextCircuit::add(ring).compose(
PlaintextCircuit::identity(1, ring).tensor(galois_of_lin_transform, ring),
ring
)
} else {
PlaintextCircuit::add(ring).compose(
PlaintextCircuit::identity(1, ring).tensor(PlaintextCircuit::linear_transform(&coeffs, ring), ring),
ring
)
};
current = summand.tensor(PlaintextCircuit::drop(1), ring).tensor(PlaintextCircuit::identity(baby_steps_galois_els.len(), ring), ring).compose(
current.output_twice(ring), ring
);
debug_assert_eq!(baby_steps_galois_els.len() + 1, current.output_count());
}
return PlaintextCircuit::identity(1, ring).tensor(PlaintextCircuit::drop(baby_steps_galois_els.len()), ring).compose(current, ring);
}
}
impl<NumberRing, A> NumberRingQuotientByIntBase<NumberRing, Zn, A>
where NumberRing: AbstractNumberRing,
A: Allocator + Clone
{
#[instrument(skip_all)]
pub fn compute_linear_transform(&self, H: &HypercubeStructure, el: &<Self as RingBase>::Element, transform: &MatmulTransform<Self>) -> <Self as RingBase>::Element {
assert!(H.galois_group().get_group() == self.acting_galois_group().get_group());
<_ as RingBase>::sum(self, transform.data.iter().map(|(s, c)| self.mul_ref_fst(c, self.apply_galois_action(el, &H.map_incl_frobenius(s)))))
}
}
#[derive(Debug)]
pub struct BabyStepGiantStepParams {
pure_baby_step_dimensions: Range<usize>,
pure_giant_step_dimensions: Range<usize>,
mixed_step_dimension: usize,
mixed_step_dimension_baby_steps: usize,
hoisted_automorphism_count: usize,
unhoisted_automorphism_count: usize
}
#[cfg(test)]
use crate::number_ring::pow2_cyclotomic::*;
#[cfg(test)]
use crate::number_ring::composite_cyclotomic::CompositeCyclotomicNumberRing;
#[cfg(test)]
use crate::number_ring::general_cyclotomic::OddSquarefreeCyclotomicNumberRing;
#[cfg(test)]
use feanor_math::assert_el_eq;
#[cfg(test)]
use crate::{ZZi64, ZZbig};
#[cfg(test)]
use feanor_math::integer::*;
#[test]
fn test_to_circuit_single() {
let number_ring: Pow2CyclotomicNumberRing = Pow2CyclotomicNumberRing::new(64);
let ring = NumberRingQuotientByIntBase::new(number_ring, Zn::new(23));
let hypercube = HypercubeStructure::default_pow2_hypercube(ring.acting_galois_group(), int_cast(23, ZZbig, ZZi64));
assert_eq!(1, hypercube.dim_count());
assert_eq!(8, hypercube.d());
assert_eq!(4, hypercube.dim_length(0));
let H = HypercubeIsomorphism::new::<false>(&&ring, &hypercube, None);
let transform = MatmulTransform::blockmatmul1d(&H, 0, |(i, k), (j, l), _| if j == i + 1 && k == 0 && l == 0 {
H.slot_ring().base_ring().one()
} else {
H.slot_ring().base_ring().zero()
});
let input = H.from_slot_values([1, 2, 3, 4].into_iter().map(|i| H.slot_ring().int_hom().map(i)));
let expected = H.from_slot_values([2, 3, 4, 0].into_iter().map(|i| H.slot_ring().int_hom().map(i)));
let compiled_transform = transform.to_circuit(H.ring(), H.hypercube());
let actual = compiled_transform.evaluate(&[input], ring.identity()).pop().unwrap();
assert_el_eq!(&ring, &expected, &actual);
let transform = MatmulTransform::linear_combine_shifts(&H, [
([-1].copy_els(), H.from_slot_values([1, 1, 1, 0].into_iter().map(|x| H.slot_ring().int_hom().map(x)))),
([0].copy_els(), H.from_slot_values([2, 2, 2, 2].into_iter().map(|x| H.slot_ring().int_hom().map(x)))),
([1].copy_els(), H.from_slot_values([0, 3, 3, 3].into_iter().map(|x| H.slot_ring().int_hom().map(x)))),
].into_iter());
let compiled_transform = transform.to_circuit(H.ring(), H.hypercube());
let input = H.from_slot_values([1, 2, 3, 4].into_iter().map(|i| H.slot_ring().int_hom().map(i)));
let expected = H.from_slot_values([4, 10, 16, 17].into_iter().map(|i| H.slot_ring().int_hom().map(i)));
let actual = compiled_transform.evaluate(&[input], ring.identity()).pop().unwrap();
assert_el_eq!(&ring, &expected, &actual);
assert_eq!(2, compiled_transform.required_galois_keys(H.galois_group()).len());
}
#[test]
fn test_compute_automorphisms_per_dimension() {
let ring = NumberRingQuotientByIntBase::new(CompositeCyclotomicNumberRing::new(3, 19), Zn::new(7));
let hypercube = HypercubeStructure::halevi_shoup_hypercube(ring.acting_galois_group(), int_cast(7, ZZbig, ZZi64));
let H = HypercubeIsomorphism::new::<false>(&&ring, &hypercube, None);
assert_eq!(2, H.hypercube().dim_count());
assert_eq!(3, H.slot_ring().rank());
assert_eq!(6, H.hypercube().dim_length(0));
assert_eq!(2, H.hypercube().dim_length(1));
let transform = MatmulTransform::blockmatmul1d(&H, 0, |_, _, _| H.slot_ring().base_ring().one());
let (max, min, gcd, sizes) = transform.compute_automorphisms_per_dimension(H.ring(), H.hypercube());
assert_eq!(vec![0, 0, 0], min);
assert_eq!(vec![2, 5, 0], max);
assert_eq!(vec![1, 1, usize::MAX], gcd);
assert_eq!(vec![3, 6, 1], sizes);
}
#[test]
fn test_compose() {
let number_ring: Pow2CyclotomicNumberRing = Pow2CyclotomicNumberRing::new(64);
let ring = NumberRingQuotientByIntBase::new(number_ring, Zn::new(23));
let hypercube = HypercubeStructure::default_pow2_hypercube(ring.acting_galois_group(), int_cast(23, ZZbig, ZZi64));
assert_eq!(1, hypercube.dim_count());
assert_eq!(8, hypercube.d());
assert_eq!(4, hypercube.dim_length(0));
let H = HypercubeIsomorphism::new::<false>(&&ring, &hypercube, None);
let transform1 = MatmulTransform::matmul1d(&H, 0, |i, j, _| if i == j + 1 {
H.slot_ring().one()
} else {
H.slot_ring().zero()
});
let transform2 = MatmulTransform::matmul1d(&H, 0, |i, j, _| if i + 1 == j {
H.slot_ring().one()
} else {
H.slot_ring().zero()
});
let composed_transform = transform1.compose(H.ring(), H.hypercube(), &transform2);
let input = H.from_slot_values([1, 2, 3, 4].into_iter().map(|i| H.slot_ring().int_hom().map(i)));
let expected = H.from_slot_values([0, 2, 3, 4].into_iter().map(|i| H.slot_ring().int_hom().map(i)));
let actual = ring.get_ring().compute_linear_transform(H.hypercube(), &input, &composed_transform);
assert_el_eq!(&ring, &expected, &actual);
}
#[test]
fn test_invert() {
let number_ring: Pow2CyclotomicNumberRing = Pow2CyclotomicNumberRing::new(64);
let ring = NumberRingQuotientByIntBase::new(number_ring, Zn::new(23));
let hypercube = HypercubeStructure::default_pow2_hypercube(ring.acting_galois_group(), int_cast(23, ZZbig, ZZi64));
assert_eq!(1, hypercube.dim_count());
assert_eq!(8, hypercube.d());
assert_eq!(4, hypercube.dim_length(0));
let H = HypercubeIsomorphism::new::<false>(&&ring, &hypercube, None);
let transform = MatmulTransform::matmul1d(&H, 0, |i, j, _| H.slot_ring().int_hom().map(StaticRing::<i32>::RING.pow(i as i32 + 1, j)));
let inverse_transform = transform.inverse(&H);
let input = H.from_slot_values([1, 2, 3, 4].into_iter().map(|i| H.slot_ring().int_hom().map(i)));
let expected = H.from_slot_values([0, 1, 0, 0].into_iter().map(|i| H.slot_ring().int_hom().map(i)));
let actual = ring.get_ring().compute_linear_transform(H.hypercube(), &input, &inverse_transform);
assert_el_eq!(&ring, &expected, &actual);
}
#[test]
fn test_blockmatmul1d() {
let ring = NumberRingQuotientByIntBase::new(OddSquarefreeCyclotomicNumberRing::new(5), Zn::new(23));
let hypercube = HypercubeStructure::halevi_shoup_hypercube(ring.acting_galois_group(), int_cast(23, ZZbig, ZZi64));
let H = HypercubeIsomorphism::new::<false>(&&ring, &hypercube, None);
let matrix = [
[1, 0, 1, 0],
[0, 0, 0, 2],
[0, 0, 0, 0],
[5, 0, 8, 8]
];
let lin_transform = MatmulTransform::blockmatmul1d(&H, 0, |(i, k), (j, l), idxs| {
assert_eq!(0, i);
assert_eq!(0, j);
assert_eq!(&[0], idxs);
H.slot_ring().base_ring().int_hom().map(matrix[k][l])
});
for i in 0..4 {
let input = H.ring().pow(H.ring().canonical_gen(), i);
let expected = H.ring().from_canonical_basis((0..4).map(|j| H.ring().base_ring().int_hom().map(matrix[j][i])));
let actual = ring.get_ring().compute_linear_transform(H.hypercube(), &input, &lin_transform);
assert_el_eq!(H.ring(), &expected, &actual);
}
let ring = NumberRingQuotientByIntBase::new(OddSquarefreeCyclotomicNumberRing::new(7), Zn::new(23));
let hypercube = HypercubeStructure::halevi_shoup_hypercube(ring.acting_galois_group(), int_cast(23, ZZbig, ZZi64));
let H = HypercubeIsomorphism::new::<false>(&&ring, &hypercube, None);
let matrix = [
[1, 0, 0],
[2, 0, 0],
[3, 0, 0]
];
let lin_transform = MatmulTransform::blockmatmul1d(&H, 0, |(i, k), (j, l), _idxs| {
if i == 0 && j == 1 {
H.slot_ring().base_ring().int_hom().map(matrix[k][l])
} else {
H.slot_ring().base_ring().zero()
}
});
for i in 0..3 {
let input = H.from_slot_values([H.slot_ring().zero(), H.slot_ring().pow(H.slot_ring().canonical_gen(), i)]);
let expected = H.from_slot_values([H.slot_ring().from_canonical_basis((0..3).map(|j| H.slot_ring().base_ring().int_hom().map(matrix[j][i]))), H.slot_ring().zero()]);
let actual = ring.get_ring().compute_linear_transform(H.hypercube(), &input, &lin_transform);
assert_el_eq!(H.ring(), &expected, &actual);
}
for i in 0..3 {
let input = H.from_slot_values([H.slot_ring().pow(H.slot_ring().canonical_gen(), i), H.slot_ring().zero()]);
let expected = H.ring().zero();
let actual = ring.get_ring().compute_linear_transform(H.hypercube(), &input, &lin_transform);
assert_el_eq!(H.ring(), &expected, &actual);
}
}