1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxCopy, VecZnxSubScalarInplace, VecZnxSwithcDegree,
  4        VmpPMatAlloc, VmpPMatPrepare,
  5    },
  6    layouts::{Backend, Module, ScratchOwned},
  7    oep::{
  8        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
  9        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
 10    },
 11};
 12use sampling::source::Source;
 13
 14use crate::{
 15    layouts::{
 16        GGLWESwitchingKey, GLWESecret,
 17        compressed::GGLWESwitchingKeyCompressed,
 18        prepared::{GLWESecretPrepared, PrepareAlloc},
 19    },
 20    trait_families::{Decompress, GLWEDecryptFamily},
 21};
 22
 23use crate::trait_families::{GGLWEEncryptSkFamily, GLWESecretPreparedModuleFamily};
 24
 25pub fn test_gglwe_switching_key_encrypt_sk<B: Backend>(
 26    module: &Module<B>,
 27    basek: usize,
 28    k_ksk: usize,
 29    digits: usize,
 30    rank_in: usize,
 31    rank_out: usize,
 32    sigma: f64,
 33) where
 34    Module<B>: GGLWEEncryptSkFamily<B>
 35        + GLWESecretPreparedModuleFamily<B>
 36        + GLWEDecryptFamily<B>
 37        + VecZnxSwithcDegree
 38        + VecZnxAddScalarInplace
 39        + VecZnxSubScalarInplace
 40        + VecZnxCopy
 41        + VmpPMatAlloc<B>
 42        + VmpPMatPrepare<B>,
 43    B: TakeVecZnxDftImpl<B>
 44        + TakeVecZnxBigImpl<B>
 45        + TakeSvpPPolImpl<B>
 46        + ScratchOwnedAllocImpl<B>
 47        + ScratchOwnedBorrowImpl<B>
 48        + ScratchAvailableImpl<B>
 49        + TakeScalarZnxImpl<B>
 50        + TakeVecZnxImpl<B>
 51        + VecZnxDftAllocBytesImpl<B>
 52        + VecZnxBigAllocBytesImpl<B>
 53        + TakeSvpPPolImpl<B>,
 54{
 55    let n: usize = module.n();
 56    let rows: usize = (k_ksk - digits * basek) / (digits * basek);
 57
 58    let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_in, rank_out);
 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(GGLWESwitchingKey::encrypt_sk_scratch_space(
 65        module, n, basek, k_ksk, rank_in, rank_out,
 66    ));
 67
 68    let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
 69    sk_in.fill_ternary_prob(0.5, &mut source_xs);
 70
 71    let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
 72    sk_out.fill_ternary_prob(0.5, &mut source_xs);
 73    let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
 74
 75    ksk.encrypt_sk(
 76        module,
 77        &sk_in,
 78        &sk_out,
 79        &mut source_xa,
 80        &mut source_xe,
 81        sigma,
 82        scratch.borrow(),
 83    );
 84
 85    ksk.key
 86        .assert_noise(module, &sk_out_prepared, &sk_in.data, sigma);
 87}
 88
 89pub fn test_gglwe_switching_key_compressed_encrypt_sk<B: Backend>(
 90    module: &Module<B>,
 91    basek: usize,
 92    k_ksk: usize,
 93    digits: usize,
 94    rank_in: usize,
 95    rank_out: usize,
 96    sigma: f64,
 97) where
 98    Module<B>: GGLWEEncryptSkFamily<B>
 99        + GLWESecretPreparedModuleFamily<B>
100        + GLWEDecryptFamily<B>
101        + VecZnxSwithcDegree
102        + VecZnxAddScalarInplace
103        + VecZnxSubScalarInplace
104        + VecZnxCopy
105        + VmpPMatAlloc<B>
106        + VmpPMatPrepare<B>,
107    B: TakeVecZnxDftImpl<B>
108        + TakeVecZnxBigImpl<B>
109        + TakeSvpPPolImpl<B>
110        + ScratchOwnedAllocImpl<B>
111        + ScratchOwnedBorrowImpl<B>
112        + ScratchAvailableImpl<B>
113        + TakeScalarZnxImpl<B>
114        + TakeVecZnxImpl<B>
115        + VecZnxDftAllocBytesImpl<B>
116        + VecZnxBigAllocBytesImpl<B>
117        + TakeSvpPPolImpl<B>,
118{
119    let n: usize = module.n();
120    let rows: usize = (k_ksk - digits * basek) / (digits * basek);
121
122    let mut ksk_compressed: GGLWESwitchingKeyCompressed<Vec<u8>> =
123        GGLWESwitchingKeyCompressed::alloc(n, basek, k_ksk, rows, digits, rank_in, rank_out);
124
125    let mut source_xs: Source = Source::new([0u8; 32]);
126    let mut source_xe: Source = Source::new([0u8; 32]);
127
128    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKeyCompressed::encrypt_sk_scratch_space(
129        module, n, basek, k_ksk, rank_in, rank_out,
130    ));
131
132    let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
133    sk_in.fill_ternary_prob(0.5, &mut source_xs);
134
135    let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
136    sk_out.fill_ternary_prob(0.5, &mut source_xs);
137    let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
138
139    let seed_xa = [1u8; 32];
140
141    ksk_compressed.encrypt_sk(
142        module,
143        &sk_in,
144        &sk_out,
145        seed_xa,
146        &mut source_xe,
147        sigma,
148        scratch.borrow(),
149    );
150
151    let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_in, rank_out);
152    ksk.decompress(module, &ksk_compressed);
153
154    ksk.key
155        .assert_noise(module, &sk_out_prepared, &sk_in.data, sigma);
156}