1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxCopy,
  4        VecZnxStd, VecZnxSubScalarInplace, VecZnxSwithcDegree, VmpPMatAlloc, VmpPMatPrepare,
  5    },
  6    layouts::{Backend, Module, ScratchOwned},
  7    oep::{
  8        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
  9        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
 10    },
 11};
 12use sampling::source::Source;
 13
 14use crate::{
 15    layouts::{GGLWEAutomorphismKey, GLWESecret, compressed::GGLWEAutomorphismKeyCompressed, prepared::GLWESecretExec},
 16    trait_families::{Decompress, GLWEDecryptFamily, GLWEKeyswitchFamily},
 17};
 18
 19use crate::trait_families::{GGLWEAutomorphismKeyEncryptSkFamily, GLWESecretExecModuleFamily};
 20
 21pub fn test_gglwe_automorphisk_key_encrypt_sk<B: Backend>(
 22    module: &Module<B>,
 23    basek: usize,
 24    k_ksk: usize,
 25    digits: usize,
 26    rank: usize,
 27    sigma: f64,
 28) where
 29    Module<B>: GGLWEAutomorphismKeyEncryptSkFamily<B>
 30        + GLWESecretExecModuleFamily<B>
 31        + GLWEKeyswitchFamily<B>
 32        + VecZnxAutomorphism
 33        + VecZnxSwithcDegree
 34        + VecZnxAddScalarInplace
 35        + VecZnxAutomorphism
 36        + VecZnxAutomorphismInplace
 37        + GLWEDecryptFamily<B>
 38        + VecZnxSubScalarInplace
 39        + VecZnxStd
 40        + VecZnxCopy
 41        + VmpPMatAlloc<B>
 42        + VmpPMatPrepare<B>,
 43    B: ScratchOwnedAllocImpl<B>
 44        + ScratchOwnedBorrowImpl<B>
 45        + ScratchAvailableImpl<B>
 46        + TakeScalarZnxImpl<B>
 47        + TakeVecZnxDftImpl<B>
 48        + TakeVecZnxImpl<B>
 49        + TakeSvpPPolImpl<B>
 50        + TakeVecZnxBigImpl<B>,
 51{
 52    let n: usize = module.n();
 53    let rows: usize = (k_ksk - digits * basek) / (digits * basek);
 54
 55    let mut atk: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
 56
 57    let mut source_xs: Source = Source::new([0u8; 32]);
 58    let mut source_xe: Source = Source::new([0u8; 32]);
 59    let mut source_xa: Source = Source::new([0u8; 32]);
 60
 61    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWEAutomorphismKey::encrypt_sk_scratch_space(
 62        module, n, basek, k_ksk, rank,
 63    ));
 64
 65    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
 66    sk.fill_ternary_prob(0.5, &mut source_xs);
 67
 68    let p = -5;
 69
 70    atk.encrypt_sk(
 71        module,
 72        p,
 73        &sk,
 74        &mut source_xa,
 75        &mut source_xe,
 76        sigma,
 77        scratch.borrow(),
 78    );
 79
 80    let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
 81    (0..atk.rank()).for_each(|i| {
 82        module.vec_znx_automorphism(
 83            module.galois_element_inv(p),
 84            &mut sk_out.data.as_vec_znx_mut(),
 85            i,
 86            &sk.data.as_vec_znx(),
 87            i,
 88        );
 89    });
 90    let sk_out_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_out);
 91
 92    atk.key
 93        .key
 94        .assert_noise(module, &sk_out_exec, &sk.data, sigma);
 95}
 96
 97pub fn test_gglwe_automorphisk_key_compressed_encrypt_sk<B: Backend>(
 98    module: &Module<B>,
 99    basek: usize,
100    k_ksk: usize,
101    digits: usize,
102    rank: usize,
103    sigma: f64,
104) where
105    Module<B>: GGLWEAutomorphismKeyEncryptSkFamily<B>
106        + GLWESecretExecModuleFamily<B>
107        + GLWEKeyswitchFamily<B>
108        + VecZnxAutomorphism
109        + VecZnxSwithcDegree
110        + VecZnxAddScalarInplace
111        + VecZnxAutomorphism
112        + VecZnxAutomorphismInplace
113        + GLWEDecryptFamily<B>
114        + VecZnxSubScalarInplace
115        + VecZnxStd
116        + VecZnxCopy
117        + VmpPMatAlloc<B>
118        + VmpPMatPrepare<B>,
119    B: ScratchOwnedAllocImpl<B>
120        + ScratchOwnedBorrowImpl<B>
121        + ScratchAvailableImpl<B>
122        + TakeScalarZnxImpl<B>
123        + TakeVecZnxDftImpl<B>
124        + TakeVecZnxImpl<B>
125        + TakeSvpPPolImpl<B>
126        + TakeVecZnxBigImpl<B>,
127{
128    let n: usize = module.n();
129    let rows: usize = (k_ksk - digits * basek) / (digits * basek);
130
131    let mut atk_compressed: GGLWEAutomorphismKeyCompressed<Vec<u8>> =
132        GGLWEAutomorphismKeyCompressed::alloc(n, basek, k_ksk, rows, digits, rank);
133
134    let mut source_xs: Source = Source::new([0u8; 32]);
135    let mut source_xe: Source = Source::new([0u8; 32]);
136
137    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWEAutomorphismKey::encrypt_sk_scratch_space(
138        module, n, basek, k_ksk, rank,
139    ));
140
141    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
142    sk.fill_ternary_prob(0.5, &mut source_xs);
143
144    let p = -5;
145
146    let seed_xa: [u8; 32] = [1u8; 32];
147
148    atk_compressed.encrypt_sk(
149        module,
150        p,
151        &sk,
152        seed_xa,
153        &mut source_xe,
154        sigma,
155        scratch.borrow(),
156    );
157
158    let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
159    (0..atk_compressed.rank()).for_each(|i| {
160        module.vec_znx_automorphism(
161            module.galois_element_inv(p),
162            &mut sk_out.data.as_vec_znx_mut(),
163            i,
164            &sk.data.as_vec_znx(),
165            i,
166        );
167    });
168    let sk_out_exec = GLWESecretExec::from(module, &sk_out);
169
170    let mut atk: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
171    atk.decompress(module, &atk_compressed);
172
173    atk.key
174        .key
175        .assert_noise(module, &sk_out_exec, &sk.data, sigma);
176}