arcis-compiler 0.9.7

use crate::{
    core::{
        bounds::BoolBounds,
        circuits::boolean::byte::Byte,
        expressions::{expr::EvalFailure, field_expr::InputId, InputKind},
    },
    traits::Keccak,
    utils::{ignore_for_equality::IgnoreForEquality, unique_id::UniqueId},
};
use arcis_internal_expr_macro::Expr;
use serde::{Deserialize, Serialize};
use std::{cell::Cell, marker::PhantomData, rc::Rc};

#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct RandomBitId(UniqueId);

impl RandomBitId {
    pub fn new() -> Self {
        RandomBitId(UniqueId::new())
    }
}

impl Default for RandomBitId {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct BitInputInfo {
    pub kind: InputKind,
    pub name: String,
    pub has_already_been_found_unused: IgnoreForEquality<Cell<bool>>,
}

impl Default for BitInputInfo {
    fn default() -> Self {
        let kind = InputKind::Secret;
        let name = "_".to_owned();
        let has_already_been_found_unused = Default::default();
        BitInputInfo {
            kind,
            name,
            has_already_been_found_unused,
        }
    }
}

/// Expressions for boolean circuits.
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, Expr)]
pub enum BitExpr<B: Clone> {
    // B = Bit, used in boolean circuits
    /// A bit equal to the product of two bits.
    And(B, B),
    /// A bit equal to the negation of a bit.
    Not(B),
    /// A bit equal to the sum of two bits.
    Xor(B, B),
    /// Reveals a bit, making it plaintext.
    /// Revealing an already revealed or plaintext bit will work and not increase circuit size.
    Reveal(B),
    /// A bit equal to a constant.
    Val(bool),
    /// Arbitrary binop. Bools are f(false, false), f(false, true), f(true, false), f(true, true)
    Binop(B, B, [bool; 4]),
    /// Random bit
    Random(RandomBitId),
    /// The keccak-f1600 permutation.
    KeccakF1600(Vec<B>, usize),
    /// An input that will produce a bit.
    Input(InputId, Rc<BitInputInfo>),
}

impl BitExpr<bool> {
    pub fn eval_binop(b1: bool, b2: bool, arr: [bool; 4]) -> bool {
        arr[b1 as usize * 2 + b2 as usize]
    }
    pub fn eval(self) -> Result<bool, EvalFailure> {
        let res = match self {
            BitExpr::And(b1, b2) => b1 && b2,
            BitExpr::Not(b) => !b,
            BitExpr::Xor(b1, b2) => b1 ^ b2,
            BitExpr::Reveal(b) => b,
            BitExpr::Val(b) => b,
            BitExpr::Binop(b1, b2, arr) => Self::eval_binop(b1, b2, arr),
            BitExpr::Random(_) => EvalFailure::err_imp("Cannot evaluate random gate here.")?,
            BitExpr::KeccakF1600(bits, i) => {
                assert!(bits.len() == 1600 && i < 1600);
                let mut state = [Byte::<bool>::from(0); 200];
                bits.chunks(8).enumerate().for_each(|(i, chunk)| {
                    state[i] =
                        Byte::new(chunk.to_vec().try_into().unwrap_or_else(|v: Vec<bool>| {
                            panic!("Expected a Vec of length 8 (found {})", v.len())
                        }));
                });
                let res = Keccak::f1600(state);
                res[i / 8].get_bits()[i % 8]
            }

            BitExpr::Input(_, _) => EvalFailure::err_imp("Input not evaluable here")?,
        };
        Ok(res)
    }
    pub fn is_plaintext(&self) -> bool {
        match self {
            BitExpr::Reveal(_) => true,
            BitExpr::Val(_) => true,
            BitExpr::Random(_) => false,
            BitExpr::KeccakF1600(_, _) => true,
            BitExpr::Input(_, info) => info.kind.is_plaintext(),
            _ => self.get_deps().iter().all(|x| *x),
        }
    }
}

impl<T: Clone> BitExpr<T> {
    pub fn is_eval_deterministic_fn_from_deps(&self) -> bool {
        !matches!(self, BitExpr::Random(_) | BitExpr::Input(..))
    }
    pub fn get_input(&self) -> Option<InputId> {
        if let BitExpr::Input(id, _) = self {
            Some(*id)
        } else {
            None
        }
    }
    pub fn get_input_name(&self) -> &str {
        if let BitExpr::Input(_, info) = self {
            info.name.as_str()
        } else {
            ""
        }
    }
    pub fn get_is_input_already_optimized_out(&self) -> Option<&Cell<bool>> {
        if let BitExpr::Input(_, info) = self {
            Some(&info.has_already_been_found_unused.0)
        } else {
            None
        }
    }
}

impl BitExpr<BoolBounds> {
    pub fn bounds(self) -> BoolBounds {
        match self {
            BitExpr::And(b1, b2) => b1 & b2,
            BitExpr::Not(b) => !b,
            BitExpr::Xor(b1, b2) => b1 ^ b2,
            BitExpr::Reveal(b) => b,
            BitExpr::Val(b) => BoolBounds::from(b),
            BitExpr::Binop(b1, b2, arr) => {
                let mut can_be_false = false;
                let mut can_be_true = false;
                for bool_1 in b1 {
                    for bool_2 in b2 {
                        let res = BitExpr::eval_binop(bool_1, bool_2, arr);
                        if res {
                            can_be_true = true;
                        } else {
                            can_be_false = true;
                        }
                    }
                }
                BoolBounds::new(can_be_false, can_be_true)
            }
            BitExpr::Random(_) => BoolBounds::new(true, true),
            BitExpr::KeccakF1600(..) => BoolBounds::new(true, true),
            BitExpr::Input(_, _) => BoolBounds::new(true, true),
        }
    }
}