parasol_cpu 0.10.0

use std::{
    borrow::BorrowMut,
    collections::HashMap,
    sync::{Arc, OnceLock},
};

use fhe_processor::FheProcessor;
pub use fhe_processor::{RunProgramOptions, RunProgramOptionsBuilder};
use parasol_concurrency::AtomicRefCell;
use parasol_runtime::{
    CircuitProcessor, Encryption, Evaluation, FheCircuit, L0LweCiphertext, L1GgswCiphertext,
    L1GlweCiphertext, L1LweCiphertext, TrivialOne, TrivialZero,
    fluent::{
        DynamicGenericInt, FheCircuitCtx, GenericInt, PackedDynamicGenericInt, PackedGenericInt,
        Sign,
    },
};
use rayon::ThreadPool;
use serde::{Deserialize, Serialize};

use crate::{
    Byte, Error, Memory, Ptr32, Result, Word, proc::gas_model::GasModel,
    tomasulo::scoreboard::ScoreboardEntryRef,
};

use self::ops::trivially_encrypt_value_l1glwe;

/// Argument handling for Parasol programs.
mod args;
pub use args::*;

#[doc(hidden)]
pub mod assembly;
mod ops;

mod fhe_processor;
mod gas_model;

#[cfg(test)]
mod tests;

pub(crate) use assembly::*;

#[doc(hidden)]
pub enum Ciphertext {
    #[allow(unused)]
    L0Lwe {
        data: Vec<Arc<AtomicRefCell<L0LweCiphertext>>>,
    },
    #[allow(unused)]
    L1Lwe {
        data: Vec<Arc<AtomicRefCell<L1LweCiphertext>>>,
    },
    L1Glwe {
        data: Vec<Arc<AtomicRefCell<L1GlweCiphertext>>>,
    },
    #[allow(unused)]
    L1Ggsw {
        data: Vec<Arc<AtomicRefCell<L1GgswCiphertext>>>,
    },
}

impl Ciphertext {
    pub fn len(&self) -> usize {
        match self {
            Self::L0Lwe { data } => data.len(),
            Self::L1Lwe { data } => data.len(),
            Self::L1Glwe { data } => data.len(),
            Self::L1Ggsw { data } => data.len(),
        }
    }

    #[allow(unused)]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    #[allow(unused)]
    pub fn unwrap_l1glwe(&self) -> &[Arc<AtomicRefCell<L1GlweCiphertext>>] {
        match self {
            Self::L1Glwe { data } => data,
            _ => panic!("Ciphertext was not L1GlweCiphertext"),
        }
    }

    pub fn try_into_l1glwe(&self) -> Result<&[Arc<AtomicRefCell<L1GlweCiphertext>>]> {
        match self {
            Self::L1Glwe { data } => Ok(data),
            _ => Err(Error::EncryptionMismatch),
        }
    }
}

#[doc(hidden)]
/// The type of value stored in a register.
#[derive(Debug, Copy, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum RegisterValueType {
    Plaintext,
    L0LweCiphertext,
    L1LweCiphertext,
    L1GlweCiphertext,
    L1GgswCiphertext,
}

#[doc(hidden)]
pub enum Register {
    Plaintext { val: u128, width: u32 },

    Ciphertext(Ciphertext),
}

impl std::fmt::Debug for Register {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Plaintext { val, width } => {
                write!(f, "v={val}, w={width}")
            }
            Self::Ciphertext(c) => {
                write!(f, "v=<encrypted>, w={}", c.len())
            }
        }
    }
}

impl Register {
    /// How many bits is the value of the register?
    pub fn width(&self) -> usize {
        match self {
            Self::Plaintext { val: _, width } => *width as usize,
            Self::Ciphertext(x) => x.len(),
        }
    }

    /// Is the register a plaintext register?
    pub fn is_plaintext(&self) -> bool {
        matches!(self, Self::Plaintext { val: _, width: _ })
    }

    /// Is the register a ciphertext register?
    pub fn is_ciphertext(&self) -> bool {
        matches!(self, Self::Ciphertext(_))
    }

    /// What type of value is stored in the register?
    pub fn register_value_type(&self) -> RegisterValueType {
        match self {
            Self::Plaintext { val: _, width: _ } => RegisterValueType::Plaintext,
            Self::Ciphertext(Ciphertext::L0Lwe { data: _ }) => RegisterValueType::L0LweCiphertext,
            Self::Ciphertext(Ciphertext::L1Lwe { data: _ }) => RegisterValueType::L1LweCiphertext,
            Self::Ciphertext(Ciphertext::L1Glwe { data: _ }) => RegisterValueType::L1GlweCiphertext,
            Self::Ciphertext(Ciphertext::L1Ggsw { data: _ }) => RegisterValueType::L1GgswCiphertext,
        }
    }

    pub fn from_word(word: &Word) -> Self {
        if word.0[0].is_plaintext() {
            let mut val = 0u128;

            for (i, b) in word.0.iter().enumerate() {
                val |= (b.clone().unwrap_plaintext() as u128) << (8 * i)
            }

            Self::Plaintext { val, width: 32 }
        } else {
            let data = word
                .0
                .iter()
                .flat_map(|x| x.clone().unwrap_ciphertext())
                .collect::<Vec<_>>();

            Self::Ciphertext(Ciphertext::L1Glwe { data })
        }
    }
}

impl Default for Register {
    fn default() -> Self {
        Register::Plaintext { val: 0, width: 32 }
    }
}

/// Checks if the width of two registers is the same.
/// Used inside an instruction implementation.
pub fn check_register_width(
    a: &Register,
    b: &Register,
    instruction_id: usize,
    pc: u32,
) -> Result<()> {
    if a.width() != b.width() {
        return Err(Error::WidthMismatch {
            inst_id: instruction_id,
            pc,
        });
    }

    // TODO, relax the 128-bit limitation.
    if a.width() < 1 || a.width() > 128 {
        return Err(Error::unsupported_width(instruction_id, pc));
    }

    Ok(())
}

/// Convert a plaintext register to a L1 GLWE ciphertext register, or copy
/// the existing ciphertext register if it's already in that form.
///
/// Returns `Err` if the register is not a plaintext or L1 GLWE ciphertext
/// register.
pub fn register_to_l1glwe_by_trivial_lift(
    register: &Register,
    zero: &L1GlweCiphertext,
    one: &L1GlweCiphertext,
) -> Result<Vec<Arc<AtomicRefCell<L1GlweCiphertext>>>> {
    match register {
        Register::Plaintext { val, width } => {
            Ok(trivially_encrypt_value_l1glwe(*val, *width, zero, one))
        }
        Register::Ciphertext(Ciphertext::L1Glwe { data }) => Ok(data.clone()),
        _ => Err(Error::EncryptionMismatch),
    }
}

pub(crate) type Fault = Arc<OnceLock<Error>>;

pub(crate) struct FheProcessorAuxData {
    uop_processor: CircuitProcessor,
    flow: std::sync::mpsc::Receiver<()>,
    memory: Option<Arc<Memory>>,
    inflight_memory_ops: HashMap<Ptr32, ScoreboardEntryRef<DispatchIsaOp>>,
    l1glwe_zero: L1GlweCiphertext,
    l1glwe_one: L1GlweCiphertext,
    enc: Encryption,
    gas_model: GasModel,

    /// A sync or async error that can occur in the processor. When set, all previous in-flight
    /// instructions that haven't started become no-ops that immediately retire and notify
    /// their dependencies. This ensures that all outstanding scoreboard entries get dropped
    /// correctly and we don't leak memory.
    fault: Fault,
}

impl FheProcessorAuxData {
    pub fn new(enc: &Encryption, eval: &Evaluation, thread_pool: Option<Arc<ThreadPool>>) -> Self {
        let (uop_processor, flow) = CircuitProcessor::new(1024, thread_pool, eval, enc);

        let l1glwe_zero = L1GlweCiphertext::trivial_zero(enc);
        let l1glwe_one = L1GlweCiphertext::trivial_one(enc);

        Self {
            uop_processor,
            flow,
            memory: None,
            inflight_memory_ops: HashMap::new(),
            l1glwe_zero,
            l1glwe_one,
            enc: enc.clone(),
            gas_model: GasModel::new(),
            fault: Arc::new(OnceLock::new()),
        }
    }
}

/// The Parasol processor that can run programs over encrypted and plaintext data.
pub struct FheComputer {
    processor: FheProcessor,
}

impl FheComputer {
    /// Create a new [`FheComputer`]. Tasks will run on the global [`rayon::ThreadPool`].
    pub fn new(enc: &Encryption, eval: &Evaluation) -> Self {
        let aux_data = FheProcessorAuxData::new(enc, eval, None);

        let processor = FheProcessor::new(aux_data);

        Self { processor }
    }

    /// Create a new [`FheComputer`]. Tasks will run on the given [`rayon::ThreadPool`].
    pub fn new_with_threadpool(
        enc: &Encryption,
        eval: &Evaluation,
        thread_pool: Arc<ThreadPool>,
    ) -> Self {
        let aux_data = FheProcessorAuxData::new(enc, eval, Some(thread_pool));

        let processor = FheProcessor::new(aux_data);

        Self { processor }
    }

    /// Run the given FHE program with user specified data and options including gas limit, etc,
    /// return the used gas and ([`Vec<Byte>`]) return value to be interpreted by the caller
    pub fn run_program_with_options_and_dynamic_return(
        &mut self,
        initial_pc: Ptr32,
        memory: &Arc<Memory>,
        args: CallData<Vec<Byte>>,
        options: &RunProgramOptions,
    ) -> Result<(u32, Vec<Byte>)> {
        self.processor.run_program_with_options_and_dynamic_return(
            memory,
            initial_pc,
            &args.to_dyn(),
            options,
        )
    }

    /// Run the given FHE program with user specified data and options including gas limit, etc,
    /// return the used gas and program return value
    pub fn run_program_with_options<T: ToArg>(
        &mut self,
        initial_pc: Ptr32,
        memory: &Arc<Memory>,
        args: CallData<T>,
        options: &RunProgramOptions,
    ) -> Result<(u32, T)> {
        self.processor
            .run_program_with_options(memory, initial_pc, &args, options)
    }

    /// Run the given FHE program with user specified data.
    pub fn run_program<T: ToArg>(
        &mut self,
        initial_pc: Ptr32,
        memory: &Arc<Memory>,
        args: CallData<T>,
    ) -> Result<T> {
        self.processor.run_program(memory, initial_pc, &args)
    }

    /// Run a graph in blocking mode.
    pub(crate) fn run_graph_blocking(&mut self, circuit: &FheCircuit) -> Result<()> {
        let uproc = self.processor.aux_data.uop_processor.borrow_mut();
        let fc = &self.processor.aux_data.flow;

        uproc.run_graph_blocking(circuit, fc)?;

        Ok(())
    }

    /// Packs a `GenericInt<N, L1GlweCiphertext, U>` into a `PackedGenericInt<N, L1GlweCiphertext, U>`.
    pub fn pack_int<const N: usize, U: Sign>(
        &mut self,
        input: GenericInt<N, L1GlweCiphertext, U>,
    ) -> Result<PackedGenericInt<N, L1GlweCiphertext, U>> {
        let ctx = FheCircuitCtx::new();

        let packed_ct = input
            .graph_inputs(&ctx)
            .pack(&ctx, &self.processor.aux_data.enc)
            .collect_output(&ctx, &self.processor.aux_data.enc);

        self.run_graph_blocking(&ctx.circuit.borrow())?;

        Ok(PackedGenericInt::from(packed_ct))
    }

    /// Similar to [`FheComputer::pack_int`] but works on [`DynamicGenericInt`]
    pub fn pack_int_dyn<U: Sign>(
        &mut self,
        input: DynamicGenericInt<L1GlweCiphertext, U>,
    ) -> Result<PackedDynamicGenericInt<L1GlweCiphertext, U>> {
        let ctx = FheCircuitCtx::new();

        let packed_ct = input
            .graph_inputs(&ctx)
            .pack(&ctx, &self.processor.aux_data.enc)
            .collect_output(&ctx, &self.processor.aux_data.enc);

        self.run_graph_blocking(&ctx.circuit.borrow())?;
        Ok(packed_ct)
    }

    /// Unpacks a `PackedGenericInt<N, L1GlweCiphertext, U>` into a `GenericInt<N, L1GlweCiphertext, U>`.
    pub fn unpack_int<const N: usize, U: Sign>(
        &mut self,
        input: PackedGenericInt<N, L1GlweCiphertext, U>,
    ) -> Result<GenericInt<N, L1GlweCiphertext, U>> {
        let ctx = FheCircuitCtx::new();

        let unpacked_ct = input
            .graph_input(&ctx)
            .unpack(&ctx)
            .convert(&ctx)
            .collect_outputs(&ctx, &self.processor.aux_data.enc);

        self.run_graph_blocking(&ctx.circuit.borrow())?;

        Ok(GenericInt::from(unpacked_ct))
    }

    /// Similar to [`FheComputer::unpack_int`] but works on [`PackedDynamicGenericInt`]
    pub fn unpack_int_dyn<U: Sign>(
        &mut self,
        input: PackedDynamicGenericInt<L1GlweCiphertext, U>,
    ) -> Result<DynamicGenericInt<L1GlweCiphertext, U>> {
        let ctx = FheCircuitCtx::new();

        let unpacked_ct = input
            .graph_input(&ctx)
            .unpack(&ctx)
            .convert(&ctx)
            .collect_outputs(&ctx, &self.processor.aux_data.enc);

        self.run_graph_blocking(&ctx.circuit.borrow())?;
        Ok(unpacked_ct)
    }
}