arcis-compiler 0.11.0

use crate::{
    core::{
        actually_used_field::FieldType,
        circuits::pre_compiled::{
            get_shared_rescue_base_field_key::get_shared_rescue_base_field_key_instruction,
            get_shared_rescue_scalar_field_key::get_shared_rescue_scalar_field_key_instruction,
        },
        expressions::other_expr::SharedRescueKeyData,
    },
    profile_summary::ProfileSummary,
    ArcisInput,
    ArcisInstruction,
    ArcisMetadata,
};
use rustc_hash::FxHashSet;
use std::{
    error::Error,
    ops::{Index, IndexMut},
    sync::LazyLock,
};

#[derive(Debug, Clone, Copy)]
enum AuxiliaryCircuit {
    SharedRescueKey(FieldType),
}

impl AuxiliaryCircuit {
    fn instruction(&self) -> ArcisInstruction {
        match self {
            AuxiliaryCircuit::SharedRescueKey(f) => match f {
                FieldType::BaseField => get_shared_rescue_base_field_key_instruction(),
                FieldType::ScalarField => get_shared_rescue_scalar_field_key_instruction(),
            },
        }
    }
}

const AUXILIARY_CIRCUITS: &[AuxiliaryCircuit] = &[
    AuxiliaryCircuit::SharedRescueKey(FieldType::BaseField),
    AuxiliaryCircuit::SharedRescueKey(FieldType::ScalarField),
];

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
struct AuxiliaryCircuitGlobalInfo {
    weight_outside_depth: AuxiliaryCircuitInfo<usize>,
    depth: AuxiliaryCircuitInfo<usize>,
}

static AUXILIARY_CIRCUIT_GLOBAL_INFO: LazyLock<AuxiliaryCircuitGlobalInfo> = LazyLock::new(|| {
    // If too slow, switch to using files, but freshness could be a problem?
    let mut res = AuxiliaryCircuitGlobalInfo::default();
    for circ in AUXILIARY_CIRCUITS {
        let instruction = circ.instruction();
        let mut info = instruction.profile_info();
        let depth = info.network_depth;
        info.network_depth = 0;
        res.weight_outside_depth[*circ] = info.weight();
        res.depth[*circ] = depth;
    }
    res
});

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct AuxiliaryCircuitInfo<T> {
    pub shared_rescue_base_field_key: T,
    pub shared_rescue_scalar_field_key: T,
}

impl<T> Index<AuxiliaryCircuit> for AuxiliaryCircuitInfo<T> {
    type Output = T;
    fn index(&self, index: AuxiliaryCircuit) -> &Self::Output {
        match index {
            AuxiliaryCircuit::SharedRescueKey(FieldType::BaseField) => {
                &self.shared_rescue_base_field_key
            }
            AuxiliaryCircuit::SharedRescueKey(FieldType::ScalarField) => {
                &self.shared_rescue_scalar_field_key
            }
        }
    }
}

impl<T> IndexMut<AuxiliaryCircuit> for AuxiliaryCircuitInfo<T> {
    fn index_mut(&mut self, index: AuxiliaryCircuit) -> &mut Self::Output {
        match index {
            AuxiliaryCircuit::SharedRescueKey(FieldType::BaseField) => {
                &mut self.shared_rescue_base_field_key
            }
            AuxiliaryCircuit::SharedRescueKey(FieldType::ScalarField) => {
                &mut self.shared_rescue_scalar_field_key
            }
        }
    }
}

impl AuxiliaryCircuitInfo<usize> {
    fn from_metadata(meta: &ArcisMetadata) -> (Self, Vec<usize>) {
        let mut res = Vec::new();
        let mut cum = AuxiliaryCircuitInfo::default();
        let mut used = AuxiliaryCircuitInfo::<FxHashSet<(usize, FieldType)>>::default();
        for input in &meta.input_order {
            res.push(cum.weight());
            if let ArcisInput::SharedRescueKey(SharedRescueKeyData {
                field_type,
                pubkey_input_id,
                ..
            }) = input
            {
                let circ = AuxiliaryCircuit::SharedRescueKey(*field_type);
                used[circ].insert((*pubkey_input_id, *field_type));
                cum[circ] = used[circ].len();
            }
        }
        res.push(cum.weight());
        (cum, res)
    }
    pub fn counts(meta: &ArcisMetadata) -> Self {
        Self::from_metadata(meta).0
    }
    pub fn weights_for_profiling(meta: &ArcisMetadata) -> Vec<usize> {
        Self::from_metadata(meta).1
    }
    pub fn weight(&self) -> usize {
        let mut weight_outside_depth = 0;
        let mut network_depth = 0;
        for circ in AUXILIARY_CIRCUITS {
            let n_circ = self[*circ];
            // Avoid infinite recursion.
            if n_circ != 0 {
                weight_outside_depth +=
                    n_circ * AUXILIARY_CIRCUIT_GLOBAL_INFO.weight_outside_depth[*circ];
                network_depth = AUXILIARY_CIRCUIT_GLOBAL_INFO.depth[*circ].max(network_depth);
            }
        }
        weight_outside_depth
            + ProfileSummary {
                network_depth,
                ..Default::default()
            }
            .weight()
    }
}

impl TryFrom<&[usize]> for AuxiliaryCircuitInfo<usize> {
    type Error = Box<dyn Error>;

    fn try_from(value: &[usize]) -> Result<Self, Self::Error> {
        let expected = AUXILIARY_CIRCUITS.len();
        if value.len() < expected {
            return Err(format!("Expected {} records, got {}", expected, value.len()).into());
        }
        let res = Self {
            shared_rescue_base_field_key: value[0],
            shared_rescue_scalar_field_key: value[1],
        };
        Ok(res)
    }
}