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