1use crate::{from_ref::FromRef, mul_by_scalar::MulByScalar};
2use crypto_primitives::{ProjectElementWithConfig, SemiringConfig, Wrapper, boolean::Boolean};
3use num_traits::CheckedAdd;
4use thiserror::Error;
5
6pub trait InnerProduct<C, Lhs: ?Sized, Rhs, Output> {
11 fn inner_product<const CHECK: bool>(
14 cfg: &C,
15 lhs: &Lhs,
16 rhs: &[Rhs],
17 zero: Output,
18 ) -> Result<Output, InnerProductError>;
19}
20
21#[derive(Clone, Debug, PartialEq, Error)]
22pub enum InnerProductError {
23 #[error("The length of LHS and RHS does not match: LHS={lhs}, RHS={rhs}")]
24 LengthMismatch { lhs: usize, rhs: usize },
25 #[error("Arithmetic overflow")]
26 Overflow,
27}
28
29#[derive(Clone, Debug)]
34pub struct MBSInnerProduct;
35
36impl<C, Lhs, Rhs, Out> InnerProduct<C, [Lhs], Rhs, Out> for MBSInnerProduct
37where
38 Out: FromRef<Lhs> + CheckedAdd + MulByScalar<Rhs, Out>,
39{
40 #[allow(clippy::arithmetic_side_effects)] fn inner_product<const CHECK: bool>(
43 _cfg: &C,
44 lhs: &[Lhs],
45 rhs: &[Rhs],
46 zero: Out,
47 ) -> Result<Out, InnerProductError> {
48 if lhs.len() != rhs.len() {
49 return Err(InnerProductError::LengthMismatch {
50 lhs: lhs.len(),
51 rhs: rhs.len(),
52 });
53 }
54
55 lhs.iter().zip(rhs).try_fold(zero, |acc, (l, r)| {
56 let widened = Out::from_ref(l);
57 let product = widened
58 .mul_by_scalar::<CHECK>(r)
59 .ok_or(InnerProductError::Overflow)?;
60 if CHECK {
61 acc.checked_add(&product).ok_or(InnerProductError::Overflow)
62 } else {
63 Ok(acc + product)
64 }
65 })
66 }
67}
68
69#[derive(Clone, Debug)]
70pub struct NativeInnerProduct;
71
72impl<C> InnerProduct<C, [C::Element], C::Element, C::Element> for NativeInnerProduct
73where
74 C: SemiringConfig,
75{
76 fn inner_product<const CHECK: bool>(
77 cfg: &C,
78 lhs: &[C::Element],
79 rhs: &[C::Element],
80 zero: C::Element,
81 ) -> Result<C::Element, InnerProductError> {
82 if lhs.len() != rhs.len() {
83 return Err(InnerProductError::LengthMismatch {
84 lhs: lhs.len(),
85 rhs: rhs.len(),
86 });
87 }
88 lhs.iter().zip(rhs).try_fold(zero, |mut acc, (l, r)| {
89 let product = if CHECK {
90 cfg.checked_mul(l, r).ok_or(InnerProductError::Overflow)?
91 } else {
92 cfg.mul(l, r)
93 };
94 if CHECK {
95 acc = cfg
96 .checked_add(&acc, &product)
97 .ok_or(InnerProductError::Overflow)?;
98 } else {
99 cfg.add_assign(&mut acc, &product);
100 }
101 Ok(acc)
102 })
103 }
104}
105
106#[derive(Clone, Debug)]
111pub struct FieldInnerProduct;
112
113impl<C, Rhs> InnerProduct<C, [C::Element], Rhs, C::Element> for FieldInnerProduct
114where
115 C: SemiringConfig + ProjectElementWithConfig<Rhs>,
116{
117 fn inner_product<const CHECK: bool>(
118 cfg: &C,
119 lhs: &[C::Element],
120 rhs: &[Rhs],
121 zero: C::Element,
122 ) -> Result<C::Element, InnerProductError> {
123 if lhs.len() != rhs.len() {
124 return Err(InnerProductError::LengthMismatch {
125 lhs: lhs.len(),
126 rhs: rhs.len(),
127 });
128 }
129
130 Ok(lhs.iter().zip(rhs).fold(zero, |mut acc, (l, r)| {
131 let product = cfg.mul(l, &cfg.project(r));
132 cfg.add_assign(&mut acc, &product);
133 acc
134 }))
135 }
136}
137
138#[derive(Clone, Debug)]
142pub struct ScalarProduct;
143
144impl<C, Lhs, Rhs, Out> InnerProduct<C, Lhs, Rhs, Out> for ScalarProduct
145where
146 Out: FromRef<Lhs> + MulByScalar<Rhs, Out>,
147{
148 fn inner_product<const CHECK: bool>(
151 _cfg: &C,
152 lhs: &Lhs,
153 point: &[Rhs],
154 _zero: Out,
155 ) -> Result<Out, InnerProductError> {
156 if point.as_ref().len() != 1 {
157 Err(InnerProductError::LengthMismatch {
158 lhs: 1,
159 rhs: point.as_ref().len(),
160 })
161 } else {
162 Ok(Out::from_ref(lhs)
163 .mul_by_scalar::<CHECK>(&point[0])
164 .ok_or(InnerProductError::Overflow)?)
165 }
166 }
167}
168
169pub struct BooleanInnerProductAdd;
173
174impl<C, Rhs: Clone, Out: FromRef<Rhs> + CheckedAdd> InnerProduct<C, [Boolean], Rhs, Out>
175 for BooleanInnerProductAdd
176{
177 #[allow(clippy::arithmetic_side_effects)] fn inner_product<const CHECK: bool>(
180 _cfg: &C,
181 lhs: &[Boolean],
182 rhs: &[Rhs],
183 zero: Out,
184 ) -> Result<Out, InnerProductError> {
185 if lhs.len() != rhs.as_ref().len() {
186 return Err(InnerProductError::LengthMismatch {
187 lhs: lhs.len(),
188 rhs: rhs.as_ref().len(),
189 });
190 }
191
192 (0..lhs.len())
193 .filter(|&i| lhs[i].into_inner())
194 .try_fold(zero, |acc, i| {
195 let rhs = Out::from_ref(&rhs[i]);
196 if CHECK {
197 acc.checked_add(&rhs).ok_or(InnerProductError::Overflow)
198 } else {
199 Ok(acc + rhs)
200 }
201 })
202 }
203}
204
205#[cfg(test)]
206mod test {
207 use crate::{CHECKED, UNCHECKED};
208 use crypto_bigint::{U64, const_monty_params};
209 use crypto_primitives::crypto_bigint_const_monty::ConstMontyField;
210 use num_traits::ConstZero;
211
212 use super::*;
213
214 #[test]
215 fn test_inner_product_basic() {
216 let lhs = [1, 2, 3];
217 let rhs = [4, 5, 6];
218 assert_eq!(
219 MBSInnerProduct::inner_product::<CHECKED>(&(), &lhs, &rhs, 0),
220 Ok(4 + 2 * 5 + 3 * 6)
221 );
222 }
223
224 #[test]
225 fn scalar_product() {
226 let lhs = 42i32;
227 let rhs = 23i128;
228
229 assert_eq!(
230 ScalarProduct::inner_product::<CHECKED>(&(), &lhs, &[rhs], 0).unwrap(),
231 i128::from(lhs) * rhs
232 )
233 }
234
235 #[test]
236 fn boolean_checked_eq_mbs_inner_product() {
237 let lhs = [
238 Boolean::from(true),
239 Boolean::from(false),
240 Boolean::from(true),
241 Boolean::from(true),
242 ];
243 let rhs = [1i128, 2, 3, 4];
244
245 assert_eq!(
246 BooleanInnerProductAdd::inner_product::<CHECKED>(&(), &lhs, &rhs, 0),
247 MBSInnerProduct::inner_product::<CHECKED>(&(), &rhs, &lhs, 0i128)
248 );
249 }
250
251 const_monty_params!(Params, U64, "0000000000000007");
252
253 #[test]
254 fn boolean_unchecked_eq_boolean_checked() {
255 let lhs = [
256 Boolean::from(true),
257 Boolean::from(false),
258 Boolean::from(true),
259 Boolean::from(true),
260 ];
261 let rhs = [
262 ConstMontyField::<Params, 1>::from(1),
263 ConstMontyField::<Params, 1>::from(2),
264 ConstMontyField::<Params, 1>::from(3),
265 ConstMontyField::<Params, 1>::from(4),
266 ];
267
268 assert_eq!(
269 BooleanInnerProductAdd::inner_product::<CHECKED>(
270 &(),
271 &lhs,
272 &rhs,
273 ConstMontyField::ZERO
274 ),
275 BooleanInnerProductAdd::inner_product::<UNCHECKED>(
276 &(),
277 &lhs,
278 &rhs,
279 ConstMontyField::ZERO
280 )
281 );
282 }
283}