1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170

use crate::core_crypto::prelude::{SignedNumeric, UnsignedNumeric};
use crate::integer::block_decomposition::DecomposableInto;
use crate::integer::ciphertext::{CompactCiphertextList, RadixCiphertext};
use crate::integer::encryption::{create_clear_radix_block_iterator, encrypt_words_radix_impl};
use crate::integer::{ClientKey, SignedRadixCiphertext};
use crate::shortint::{
    CompactPublicKey as ShortintCompactPublicKey,
    CompressedCompactPublicKey as ShortintCompressedCompactPublicKey,
};
use serde::{Deserialize, Serialize};

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CompactPublicKey {
    pub(crate) key: ShortintCompactPublicKey,
}

impl CompactPublicKey {
    pub fn new(client_key: &ClientKey) -> Self {
        let key = ShortintCompactPublicKey::new(&client_key.key);
        Self { key }
    }

    pub fn try_new(client_key: &ClientKey) -> Option<Self> {
        let key = ShortintCompactPublicKey::try_new(&client_key.key)?;
        Some(Self { key })
    }

    /// Deconstruct a [`CompactPublicKey`] into its constituents.
    pub fn into_raw_parts(self) -> ShortintCompactPublicKey {
        self.key
    }

    /// Construct a [`CompactPublicKey`] from its constituents.
    pub fn from_raw_parts(key: ShortintCompactPublicKey) -> Self {
        Self { key }
    }

    pub fn encrypt_radix<T>(&self, message: T, num_blocks: usize) -> RadixCiphertext
    where
        T: DecomposableInto<u64> + UnsignedNumeric,
    {
        encrypt_words_radix_impl(
            &self.key,
            message,
            num_blocks,
            ShortintCompactPublicKey::encrypt,
        )
    }

    pub fn encrypt_signed_radix<T>(&self, message: T, num_blocks: usize) -> SignedRadixCiphertext
    where
        T: DecomposableInto<u64> + SignedNumeric,
    {
        encrypt_words_radix_impl(
            &self.key,
            message,
            num_blocks,
            ShortintCompactPublicKey::encrypt,
        )
    }

    pub fn encrypt_radix_compact<T: DecomposableInto<u64>>(
        &self,
        message: T,
        num_blocks_per_integer: usize,
    ) -> CompactCiphertextList {
        let clear_block_iter = create_clear_radix_block_iterator(
            message,
            self.key.parameters.message_modulus(),
            num_blocks_per_integer,
        );

        let ct_list = self.key.encrypt_iter(clear_block_iter);
        CompactCiphertextList {
            ct_list,
            num_blocks_per_integer,
        }
    }

    pub fn encrypt_slice_radix_compact<T: DecomposableInto<u64>>(
        &self,
        messages: &[T],
        num_blocks: usize,
    ) -> CompactCiphertextList {
        self.encrypt_iter_radix_compact(messages.iter().copied(), num_blocks)
    }

    pub fn encrypt_iter_radix_compact<T: DecomposableInto<u64>>(
        &self,
        mut message_iter: impl Iterator<Item = T>,
        num_blocks_per_integer: usize,
    ) -> CompactCiphertextList {
        let mut iterator_chain;
        match (message_iter.next(), message_iter.next()) {
            (None, None) => panic!("At least one message is required"),
            (None, Some(_)) => unreachable!(),
            (Some(first_message), None) => {
                // Cannot form a chain
                return self.encrypt_radix_compact(first_message, num_blocks_per_integer);
            }
            (Some(first_message), Some(second_message)) => {
                let first_iter = create_clear_radix_block_iterator(
                    first_message,
                    self.key.parameters.message_modulus(),
                    num_blocks_per_integer,
                );
                let second_iter = create_clear_radix_block_iterator(
                    second_message,
                    self.key.parameters.message_modulus(),
                    num_blocks_per_integer,
                );

                iterator_chain =
                    Box::new(first_iter.chain(second_iter)) as Box<dyn Iterator<Item = u64>>;
            }
        }

        for message in message_iter {
            let other_iter = create_clear_radix_block_iterator(
                message,
                self.key.parameters.message_modulus(),
                num_blocks_per_integer,
            );

            iterator_chain = Box::new(iterator_chain.chain(other_iter));
        }

        let ct_list = self.key.encrypt_iter(iterator_chain);
        CompactCiphertextList {
            ct_list,
            num_blocks_per_integer,
        }
    }

    pub fn size_elements(&self) -> usize {
        self.key.size_elements()
    }

    pub fn size_bytes(&self) -> usize {
        self.key.size_bytes()
    }
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CompressedCompactPublicKey {
    pub(crate) key: ShortintCompressedCompactPublicKey,
}

impl CompressedCompactPublicKey {
    pub fn new(client_key: &ClientKey) -> Self {
        let key = ShortintCompressedCompactPublicKey::new(&client_key.key);
        Self { key }
    }

    /// Deconstruct a [`CompressedCompactPublicKey`] into its constituents.
    pub fn into_raw_parts(self) -> ShortintCompressedCompactPublicKey {
        self.key
    }

    /// Construct a [`CompressedCompactPublicKey`] from its constituents.
    pub fn from_raw_parts(key: ShortintCompressedCompactPublicKey) -> Self {
        Self { key }
    }

    pub fn decompress(&self) -> CompactPublicKey {
        CompactPublicKey {
            key: self.key.decompress(),
        }
    }
}