darkfi/zk/gadget/
cond_select.rs

1/* This file is part of DarkFi (https://dark.fi)
2 *
3 * Copyright (C) 2020-2026 Dyne.org foundation
4 *
5 * Copyright (c) zkMove Authors
6 * SPDX-License-Identifier: Apache-2.0
7 *
8 * This program is free software: you can redistribute it and/or modify
9 * it under the terms of the GNU Affero General Public License as
10 * published by the Free Software Foundation, either version 3 of the
11 * License, or (at your option) any later version.
12 *
13 * This program is distributed in the hope that it will be useful,
14 * but WITHOUT ANY WARRANTY; without even the implied warranty of
15 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
16 * GNU Affero General Public License for more details.
17 *
18 * You should have received a copy of the GNU Affero General Public License
19 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
20 */
21
22use core::marker::PhantomData;
23
24use halo2_proofs::{
25    circuit::{AssignedCell, Chip, Layouter, Region},
26    pasta::group::ff::WithSmallOrderMulGroup,
27    plonk::{Advice, Column, ConstraintSystem, Error, Expression, Selector},
28    poly::Rotation,
29};
30
31pub const NUM_OF_UTILITY_ADVICE_COLUMNS: usize = 4;
32
33#[derive(Clone, Debug)]
34pub struct ConditionalSelectConfig<F: WithSmallOrderMulGroup<3> + Ord> {
35    advices: [Column<Advice>; NUM_OF_UTILITY_ADVICE_COLUMNS],
36    s_cs: Selector,
37    _marker: PhantomData<F>,
38}
39
40pub struct ConditionalSelectChip<F: WithSmallOrderMulGroup<3> + Ord> {
41    config: ConditionalSelectConfig<F>,
42}
43
44impl<F: WithSmallOrderMulGroup<3> + Ord> Chip<F> for ConditionalSelectChip<F> {
45    type Config = ConditionalSelectConfig<F>;
46    type Loaded = ();
47
48    fn config(&self) -> &Self::Config {
49        &self.config
50    }
51
52    fn loaded(&self) -> &Self::Loaded {
53        &()
54    }
55}
56
57impl<F: WithSmallOrderMulGroup<3> + Ord> ConditionalSelectChip<F> {
58    pub fn construct(config: <Self as Chip<F>>::Config) -> Self {
59        Self { config }
60    }
61
62    pub fn configure(
63        meta: &mut ConstraintSystem<F>,
64        advices: [Column<Advice>; NUM_OF_UTILITY_ADVICE_COLUMNS],
65    ) -> <Self as Chip<F>>::Config {
66        for column in &advices {
67            meta.enable_equality(*column);
68        }
69        let s_cs = meta.selector();
70
71        meta.create_gate("conditional_select", |meta| {
72            let lhs = meta.query_advice(advices[0], Rotation::cur());
73            let rhs = meta.query_advice(advices[1], Rotation::cur());
74            let out = meta.query_advice(advices[2], Rotation::cur());
75            let cond = meta.query_advice(advices[3], Rotation::cur());
76            let s_cs = meta.query_selector(s_cs);
77            let one = Expression::Constant(F::ONE);
78
79            vec![
80                // cond is 0 or 1
81                s_cs.clone() * (cond.clone() * (one - cond.clone())),
82                // lhs * cond + rhs * (1 - cond) = out
83                s_cs * ((lhs - rhs.clone()) * cond + rhs - out),
84            ]
85        });
86
87        ConditionalSelectConfig { advices, s_cs, _marker: PhantomData }
88    }
89
90    pub fn conditional_select(
91        &self,
92        layouter: &mut impl Layouter<F>,
93        a: AssignedCell<F, F>,
94        b: AssignedCell<F, F>,
95        cond: AssignedCell<F, F>,
96    ) -> Result<AssignedCell<F, F>, Error> {
97        let config = self.config();
98        let out = layouter.assign_region(
99            || "conditional_select",
100            |mut region: Region<'_, F>| {
101                config.s_cs.enable(&mut region, 0)?;
102
103                a.copy_advice(|| "copy a", &mut region, config.advices[0], 0)?;
104                b.copy_advice(|| "copy b", &mut region, config.advices[1], 0)?;
105
106                let cond = cond.copy_advice(|| "copy cond", &mut region, config.advices[3], 0)?;
107
108                let selected = cond
109                    .value()
110                    .copied()
111                    .to_field()
112                    .zip(a.value())
113                    .zip(b.value())
114                    .map(|((cond, a), b)| if cond == F::ONE.into() { a } else { b })
115                    .copied();
116
117                let cell =
118                    region.assign_advice(|| "select result", config.advices[2], 0, || selected)?;
119
120                Ok(cell)
121            },
122        )?;
123        Ok(out)
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use crate::zk::assign_free_advice;
131
132    use halo2_proofs::{
133        arithmetic::Field,
134        circuit::{floor_planner, Value},
135        dev::MockProver,
136        pasta::pallas,
137        plonk,
138        plonk::{Circuit, Instance},
139    };
140    use rand::rngs::OsRng;
141
142    #[derive(Clone)]
143    struct CondSelectConfig {
144        primary: Column<Instance>,
145        condselect_config: ConditionalSelectConfig<pallas::Base>,
146    }
147
148    #[derive(Default)]
149    struct CondSelectCircuit {
150        pub cond: Value<pallas::Base>,
151        pub a: Value<pallas::Base>,
152        pub b: Value<pallas::Base>,
153    }
154
155    impl Circuit<pallas::Base> for CondSelectCircuit {
156        type Config = CondSelectConfig;
157        type FloorPlanner = floor_planner::V1;
158        type Params = ();
159
160        fn without_witnesses(&self) -> Self {
161            Self::default()
162        }
163
164        fn configure(meta: &mut ConstraintSystem<pallas::Base>) -> Self::Config {
165            let advices: [Column<Advice>; NUM_OF_UTILITY_ADVICE_COLUMNS] = [
166                meta.advice_column(),
167                meta.advice_column(),
168                meta.advice_column(),
169                meta.advice_column(),
170            ];
171
172            let primary = meta.instance_column();
173            meta.enable_equality(primary);
174
175            for advice in advices.iter() {
176                meta.enable_equality(*advice);
177            }
178
179            let condselect_config = ConditionalSelectChip::configure(meta, advices);
180
181            Self::Config { primary, condselect_config }
182        }
183
184        fn synthesize(
185            &self,
186            config: Self::Config,
187            mut layouter: impl Layouter<pallas::Base>,
188        ) -> Result<(), plonk::Error> {
189            let condselect_chip =
190                ConditionalSelectChip::construct(config.condselect_config.clone());
191
192            let cond = assign_free_advice(
193                layouter.namespace(|| "Witness cond"),
194                config.condselect_config.advices[0],
195                self.cond,
196            )?;
197
198            let a = assign_free_advice(
199                layouter.namespace(|| "Witness a"),
200                config.condselect_config.advices[1],
201                self.a,
202            )?;
203
204            let b = assign_free_advice(
205                layouter.namespace(|| "Witness b"),
206                config.condselect_config.advices[2],
207                self.b,
208            )?;
209
210            let selection = condselect_chip.conditional_select(&mut layouter, a, b, cond)?;
211            layouter.constrain_instance(selection.cell(), config.primary, 0)?;
212
213            Ok(())
214        }
215    }
216
217    #[test]
218    fn cond_select_chip() -> crate::Result<()> {
219        // 1 should select A
220        let cond = pallas::Base::ONE;
221        let a = pallas::Base::random(&mut OsRng);
222        let b = pallas::Base::random(&mut OsRng);
223        let public_inputs = vec![a];
224
225        let circuit =
226            CondSelectCircuit { cond: Value::known(cond), a: Value::known(a), b: Value::known(b) };
227
228        let prover = MockProver::run(4, &circuit, vec![public_inputs])?;
229        prover.assert_satisfied();
230
231        // 0 should select B
232        let cond = pallas::Base::ZERO;
233        let a = pallas::Base::random(&mut OsRng);
234        let b = pallas::Base::random(&mut OsRng);
235        let public_inputs = vec![b];
236
237        let circuit =
238            CondSelectCircuit { cond: Value::known(cond), a: Value::known(a), b: Value::known(b) };
239
240        let prover = MockProver::run(4, &circuit, vec![public_inputs])?;
241        prover.assert_satisfied();
242
243        Ok(())
244    }
245}
darkfi/zk/gadget/cond_select.rs

darkfi/zk/gadget/
cond_select.rs
