Skip to main content

zinc_utils/
inner_product.rs

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
6/// A trait for inner product algorithms implementations.
7///
8/// `C` is the config needed to perform the operations: `()` for
9/// self-sufficient types, or a field config for dynamic field elements.
10pub trait InnerProduct<C, Lhs: ?Sized, Rhs, Output> {
11    /// The main entry point for the inner product.
12    /// `CHECK` determines whether the implementation should check for overflow.
13    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/// An implementation of inner product that piggies back
30/// on the `MulByScalar` and `CheckedAdd` traits.
31/// It does `mul_by_scalar` for products of terms
32/// and then combines the results using either `add` or `checked_add`.
33#[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    /// The mul-by-scalar inner product.
41    #[allow(clippy::arithmetic_side_effects)] // Used in unchecked mode
42    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/// An implementation of inner product over a dynamic field: projects the RHS
107/// entries into the field and folds with the field operations.
108///
109/// Field operations cannot overflow, so `CHECK` is ignored.
110#[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/// The inner product for vectors of length 1 (a.k.a. scalars).
139/// Uses `mul_by_scalar` to multiply the only components of vectors
140/// to get the result.
141#[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    /// A scalar inner product. Assumes `Lhs` is a scalar type
149    /// and always asserts that `point` has only one component.
150    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
169/// The inner product for slices containing `Boolean` elements.
170/// Uses `add` or `checked_add` to sum the elements of the RHS that
171/// correspond to `true` elements of the boolean slice.
172pub struct BooleanInnerProductAdd;
173
174impl<C, Rhs: Clone, Out: FromRef<Rhs> + CheckedAdd> InnerProduct<C, [Boolean], Rhs, Out>
175    for BooleanInnerProductAdd
176{
177    /// Boolean inner product.
178    #[allow(clippy::arithmetic_side_effects)] // Used in unchecked mode
179    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}