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