arcis-compiler 0.9.7

use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        bounds::{BoolBounds, Bounds, FieldBounds, IsBounds},
        expressions::{
            bit_expr::BitExpr,
            conversion_expr::ConversionExpr,
            curve_expr::CurveExpr,
            expr::Expr,
            field_expr::FieldExpr,
            macro_uses::BoundUnFold,
            other_expr::OtherExpr,
        },
        ir::IntermediateRepresentation,
        tracking::Tracking,
    },
    utils::{
        field::{BaseField, ScalarField},
        hash_map_vec::HashMapVec,
    },
};
use ff::Field;

#[derive(Debug, Default, Clone)]
pub struct IRBuilder {
    exprs: HashMapVec<Expr<usize>>,
    bounds: Vec<Bounds>,
    is_plaintext: Vec<bool>,
    /// Whether we allow subcircuits that have not yet been built in the IR.
    /// Our compilation has multiple phases. A lot of phases are building expressions,
    /// for instance `ArithmeticCircuitBuilder` transforms `Sqrt` and `Pow` into their circuits.
    /// Therefore during or past this phase, we should not introduce these expressions anymore,
    /// and we should build them instead.
    /// This will be `true` when building the initial IR, and `false` during compilation phases.
    allow_unbuilt_circuits: bool,
}

/// Transforms empty bound to non-empty bound.
/// Avoids Internal Compiler Error from user-created UB.
fn safe_bound(bound: Bounds) -> Bounds {
    match bound {
        Bounds::Bit(b) => {
            if b.is_empty() {
                Bounds::Bit(BoolBounds::new(true, false))
            } else {
                bound
            }
        }
        Bounds::Scalar(b) => {
            if b.is_empty() {
                Bounds::Scalar(FieldBounds::new(ScalarField::ZERO, ScalarField::ONE))
            } else {
                bound
            }
        }
        Bounds::Base(b) => {
            if b.is_empty() {
                Bounds::Base(FieldBounds::new(BaseField::ZERO, BaseField::ONE))
            } else {
                bound
            }
        }
        Bounds::Curve(_) => bound,
    }
}

impl IRBuilder {
    pub fn new(allow_unbuilt_circuits: bool) -> IRBuilder {
        IRBuilder {
            exprs: HashMapVec::new(),
            bounds: Vec::new(),
            is_plaintext: Vec::new(),
            allow_unbuilt_circuits,
        }
    }
    pub fn are_unbuilt_circuits_allowed(&self) -> bool {
        self.allow_unbuilt_circuits
    }
    pub fn new_expr_with_info(
        &mut self,
        expr: Expr<usize>,
        old_bounds: Option<Bounds>,
        is_plaintext: bool,
    ) -> usize {
        let n = self.exprs.push(expr.clone());
        unsafe {
            // for performance reasons
            let bounds: Bounds = expr
                .clone()
                .apply(|i| *self.bounds.get_unchecked(i))
                .apply_2(&mut BoundUnFold)
                .bounds();
            let bounds = match old_bounds {
                None => bounds,
                Some(old_bounds) => old_bounds.inter(bounds),
            };
            let is_plaintext = is_plaintext
                || expr
                    .apply(|t| *self.is_plaintext.get_unchecked(t))
                    .is_plaintext();
            if n < self.bounds.len() {
                self.bounds[n] = safe_bound(bounds.inter(self.bounds[n]));
                self.is_plaintext[n] = *self.is_plaintext.get_unchecked(n) || is_plaintext;
            } else {
                self.bounds.push(safe_bound(bounds));
                self.is_plaintext.push(is_plaintext);
            };
        }
        n
    }
    pub fn new_expr(&mut self, expr: Expr<usize>) -> usize {
        self.new_expr_with_info(expr, None, false)
    }

    pub fn get_expr(&self, expr_id: usize) -> &Expr<usize> {
        self.exprs.get_val(expr_id)
    }

    pub fn contains_expr_id(&self, expr_id: usize) -> bool {
        self.exprs.contains_id(expr_id)
    }

    pub fn len(&self) -> usize {
        self.exprs.len()
    }

    #[allow(dead_code)]
    pub fn is_empty(&self) -> bool {
        self.exprs.is_empty()
    }

    pub fn get_bounds(&self, expr_id: usize) -> &Bounds {
        &self.bounds[expr_id]
    }

    pub fn get_is_plaintext(&self, expr_id: usize) -> bool {
        self.is_plaintext[expr_id]
    }
    pub fn reveal(&mut self, expr_id: usize) {
        self.is_plaintext[expr_id] = true;
    }

    pub fn into_ir(self, outputs: Vec<usize>) -> IntermediateRepresentation {
        IntermediateRepresentation::new(
            self.exprs.move_vec(),
            outputs,
            self.bounds,
            self.is_plaintext,
        )
    }
    pub fn into_ir_with_tracking(
        self,
        outputs: Vec<usize>,
        mut tracking: Tracking,
        old_id_to_new_id: Vec<usize>,
    ) -> IntermediateRepresentation {
        tracking.combine(old_id_to_new_id, self.exprs.len());
        IntermediateRepresentation::new_with_tracking(
            self.exprs.move_vec(),
            outputs,
            self.bounds,
            self.is_plaintext,
            tracking,
        )
    }
    pub fn into_ir_with_new_tracking(
        self,
        outputs: Vec<usize>,
        tracking_sizes: Vec<usize>,
    ) -> IntermediateRepresentation {
        IntermediateRepresentation::new_with_tracking(
            self.exprs.move_vec(),
            outputs,
            self.bounds,
            self.is_plaintext,
            Tracking::new(tracking_sizes),
        )
    }
}
pub trait ExprStore<F: ActuallyUsedField> {
    fn push_conversion(&mut self, expr: ConversionExpr<F, usize>) -> usize;
    fn push_field(&mut self, expr: FieldExpr<F, usize>) -> usize;
    fn push_bit(&mut self, expr: BitExpr<usize>) -> usize;
    fn push_curve(&mut self, expr: CurveExpr<usize>) -> usize;
    #[allow(dead_code)]
    fn push_other(&mut self, expr: OtherExpr<usize>) -> usize;
    fn get_expr(&self, expr_id: usize) -> Expr<usize>;
    fn bounds(&self, id: usize) -> FieldBounds<F>;
}

impl<F: ActuallyUsedField> ExprStore<F> for IRBuilder {
    fn push_conversion(&mut self, expr: ConversionExpr<F, usize>) -> usize {
        self.new_expr(F::conversion_expr_to_expr(expr))
    }

    fn push_field(&mut self, expr: FieldExpr<F, usize>) -> usize {
        self.new_expr(F::field_expr_to_expr(expr))
    }

    fn push_bit(&mut self, expr: BitExpr<usize>) -> usize {
        self.new_expr(Expr::Bit(expr))
    }

    fn push_curve(&mut self, expr: CurveExpr<usize>) -> usize {
        self.new_expr(Expr::Curve(expr))
    }

    fn push_other(&mut self, expr: OtherExpr<usize>) -> usize {
        self.new_expr(Expr::Other(expr))
    }

    fn get_expr(&self, expr_id: usize) -> Expr<usize> {
        self.get_expr(expr_id).clone()
    }

    fn bounds(&self, id: usize) -> FieldBounds<F> {
        F::bounds_to_field_bounds(*self.get_bounds(id))
    }
}