arcis-compiler 0.9.4

A framework for writing secure multi-party computation (MPC) circuits to be executed on the Arcium network.
Documentation 
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

use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        circuits::{
            arithmetic::{PowCircuit, SqrtCircuit},
            traits::arithmetic_circuit::ArithmeticCircuit,
        },
        compile_passes::compilation_pass::LocalCompilationPass,
        expressions::{
            circuit::{ArithmeticCircuitId, BaseCircuitId, GeneralCircuitId},
            expr::Expr,
            field_expr::FieldExpr,
            other_expr::OtherExpr,
        },
        ir_builder::IRBuilder,
    },
    utils::field::{BaseField, ScalarField},
};
use ff::Field;
use rustc_hash::FxHashMap;

type ArithmeticCircuitRunCache = FxHashMap<(Vec<usize>, ArithmeticCircuitId), Vec<usize>>;
type BaseCircuitRunCache = FxHashMap<(Vec<usize>, BaseCircuitId), Vec<usize>>;
type GeneralCircuitRunCache =
    FxHashMap<(Vec<usize>, Vec<usize>, GeneralCircuitId), (Vec<usize>, Vec<usize>)>;

/// Builds unoptimized arithmetic circuits.
#[derive(Default)]
pub struct ArithmeticCircuitBuilder {
    expr_store: IRBuilder,
    arithmetic_circuit_cache: ArithmeticCircuitRunCache,
    base_circuit_cache: BaseCircuitRunCache,
    general_circuit_cache: GeneralCircuitRunCache,
}

impl ArithmeticCircuitBuilder {
    fn build_arithmetic_circuits<F: ActuallyUsedField>(
        &mut self,
        expr: FieldExpr<F, usize>,
    ) -> Expr<usize> {
        match expr {
            FieldExpr::SubCircuit(v, c, idx) => {
                let res = match self.arithmetic_circuit_cache.get(&(v.clone(), c)) {
                    None => {
                        let circuit_outputs =
                            c.to_circuit::<F>().run_usize(&v, &mut self.expr_store);
                        let res = circuit_outputs.get(idx).cloned();
                        self.arithmetic_circuit_cache
                            .insert((v.clone(), c), circuit_outputs);
                        res
                    }
                    Some(circuit_outputs) => circuit_outputs.get(idx).cloned(),
                };
                res.map(|expr_id| self.expr_store.get_expr(expr_id).clone())
                    .unwrap_or(F::field_expr_to_expr(FieldExpr::Val(F::ZERO)))
            }
            FieldExpr::Sqrt(v) if !self.expr_store.get_is_plaintext(v) => {
                let circuit_outputs = <SqrtCircuit as ArithmeticCircuit<F>>::run_usize(
                    &SqrtCircuit,
                    &[v],
                    &mut self.expr_store,
                );
                let res = circuit_outputs[0];
                self.expr_store.get_expr(res).clone()
            }
            FieldExpr::Pow(v, e, is_expected_non_zero) if !self.expr_store.get_is_plaintext(v) => {
                let circuit_outputs = <PowCircuit as ArithmeticCircuit<F>>::run_usize(
                    &PowCircuit {
                        exponent: e,
                        is_expected_non_zero,
                    },
                    &[v],
                    &mut self.expr_store,
                );
                let res = circuit_outputs[0];
                self.expr_store.get_expr(res).clone()
            }
            _ => F::field_expr_to_expr(expr),
        }
    }

    fn build_general_circuits(&mut self, expr: OtherExpr<usize>) -> Expr<usize> {
        match expr {
            OtherExpr::ScalarGeneralCircuit(s, b, c, idx) => {
                let res = match self.general_circuit_cache.get(&(s.clone(), b.clone(), c)) {
                    None => {
                        let result = c.to_circuit().run_usize(&s, &b, &mut self.expr_store);
                        self.general_circuit_cache
                            .insert((s.clone(), b.clone(), c), result.clone());
                        result.0.get(idx).cloned()
                    }
                    Some((scalar_res, _)) => scalar_res.get(idx).cloned(),
                };
                res.map(|expr_id| self.expr_store.get_expr(expr_id).clone())
                    .unwrap_or(Expr::Scalar(FieldExpr::Val(ScalarField::ZERO)))
            }
            OtherExpr::BaseGeneralCircuit(s, b, c, idx) => {
                let res = match self.general_circuit_cache.get(&(s.clone(), b.clone(), c)) {
                    None => {
                        let result = c.to_circuit().run_usize(&s, &b, &mut self.expr_store);
                        self.general_circuit_cache
                            .insert((s.clone(), b.clone(), c), result.clone());
                        result.1.get(idx).cloned()
                    }
                    Some((_, base_res)) => base_res.get(idx).cloned(),
                };
                res.map(|expr_id| self.expr_store.get_expr(expr_id).clone())
                    .unwrap_or(Expr::Base(FieldExpr::Val(BaseField::ZERO)))
            }
            OtherExpr::BaseArithmeticCircuit(v, c, idx) => {
                let res = match self.base_circuit_cache.get(&(v.clone(), c)) {
                    None => {
                        let circuit_outputs = c.to_circuit().run_usize(&v, &mut self.expr_store);
                        let res = circuit_outputs.get(idx).cloned();
                        self.base_circuit_cache
                            .insert((v.clone(), c), circuit_outputs);
                        res
                    }
                    Some(circuit_outputs) => circuit_outputs.get(idx).cloned(),
                };
                res.map(|expr_id| self.expr_store.get_expr(expr_id).clone())
                    .unwrap_or(Expr::Base(FieldExpr::Val(BaseField::ZERO)))
            }
            _ => Expr::Other(expr),
        }
    }
}

impl LocalCompilationPass for ArithmeticCircuitBuilder {
    fn expr_store(&mut self) -> &mut IRBuilder {
        &mut self.expr_store
    }

    fn transform(&mut self, expr: Expr<usize>, _is_plaintext: bool) -> Expr<usize> {
        match expr {
            Expr::Scalar(e) => self.build_arithmetic_circuits(e),
            Expr::Base(e) => self.build_arithmetic_circuits(e),
            Expr::Other(e) => self.build_general_circuits(e),
            _ => expr,
        }
    }
}