arcis 0.10.4

use crate::ArcisType;
use arcis_compiler::{compile::read_ir, ArcisInstruction, DomainKind, EvalValue};

/// Used for testing an `#[instruction]`.
///
/// Access an `#[instruction]` via `get_instruction("instruction_name")`, then
/// test it via `instruction.eval(...)`.
///
/// Full example:
///
/// ```
/// use arcis::*;
/// #[encrypted]
/// mod my_module {
///     use arcis::*;
///     #[instruction]
///     fn decrypt_client(a: Enc<Shared, u8>) -> u8 {
///         a.to_arcis().reveal()
///     }
/// }
///
/// #[cfg(test)]
/// mod tests {
///     use arcis::{testing::*, *};
///     #[test]
///     fn test_decrypt_client() {
///         // Access the `#[instruction]` by its name.
///         let instruction = get_instruction("decrypt_client");
///         let pubkey = ArcisX25519Pubkey::from_base58(b"test"); // Does not need to be a real pubkey.
///
///         // Test normal values + values above `u8::MAX`.
///         // A safety modulo is performed during `.to_arcis()` so that all values are really their rust type.
///         // These values will behave exactly like the `u8` `value % 256`.
///         for value in (0u128..1000).chain((u128::MAX - 1000)..=u128::MAX) {
///             let expected = value as u8;
///             let pubkey = Shared::new(pubkey);
///             let input = pubkey.from_arcis(value);
///             // Note: `input` is an `Enc<Shared, u128>` and the `#[instruction]` requires an `Enc<Shared, u128>`.
///             // This will still work because `u8` and `u128` take the same amount of ciphertexts (1).
///             let res = instruction.eval::<u8>(&input);
///             assert_eq!(res, expected);
///             // We interpreted the result as a `u8`, but it would have been the same as a `u128`.
///             let res_128 = instruction.eval::<u128>(&input);
///             assert_eq!(res as u128, res_128);
///         }
///
///         // Test big negative values.
///         // Negative values are considered as `value + 2^255 - 19`,
///         // and will behave exactly like the `u8` `value + 2^255 - 19 % 256`.
///         for value in i128::MIN..(i128::MIN + 1000) {
///             let expected = (value + 256 - 19) as u8; // -19 is necessary
///             let pubkey = Shared::new(pubkey);
///             let input = pubkey.from_arcis(value);
///             let res = instruction.eval::<u8>(&input);
///             assert_eq!(res, expected);
///             let res_128 = instruction.eval::<u128>(&input);
///             assert_eq!(res as u128, res_128);
///         }
///
///         // Test -1.
///         // Statistically, -1 will 99.999...% of the time be considered as 255.
///         {
///             let pubkey = Shared::new(pubkey);
///             let input = pubkey.from_arcis(-1);
///             let res = instruction.eval::<u8>(&input);
///             assert_eq!(res, 255); // statistically, no 256-19=237
///         }
///     }
/// }
/// ```
pub struct ArcisInstructionWithInfo {
    instruction: ArcisInstruction,
    output_domains: Vec<DomainKind>,
}

trait ArcisTypeOrRaw {
    fn to_raw_inputs(self) -> Vec<EvalValue>;
    fn from_raw_outputs(outputs: Vec<EvalValue>) -> Self;
}

impl ArcisTypeOrRaw for Vec<EvalValue> {
    fn to_raw_inputs(self) -> Vec<EvalValue> {
        self
    }
    fn from_raw_outputs(outputs: Vec<EvalValue>) -> Self {
        outputs
    }
}

impl<T: ArcisType> ArcisTypeOrRaw for T {
    fn to_raw_inputs(self) -> Vec<EvalValue> {
        let mut input_vals = Vec::new();
        self.handle_outputs(&mut input_vals);
        input_vals
    }
    fn from_raw_outputs(outputs: Vec<EvalValue>) -> Self {
        T::from_values(outputs.as_slice())
    }
}

impl ArcisInstructionWithInfo {
    fn eval_raw(&self, inputs: Vec<EvalValue>) -> Vec<EvalValue> {
        let mut rng = rand::thread_rng();
        self.instruction
            .mock_eval_vec(inputs, self.output_domains.as_slice(), &[], &mut rng)
    }
    /// Evaluates an instruction on inputs, then interprets the instruction result as `R`.
    #[allow(private_bounds)]
    pub fn eval<R: ArcisTypeOrRaw>(&self, inputs: impl ArcisTypeOrRaw) -> R {
        let inputs = inputs.to_raw_inputs();

        let output_vals = self.eval_raw(inputs);
        R::from_raw_outputs(output_vals)
    }
}

/// Gets the instruction with the given name.
/// For `#[instruction] fn my_instruction(...) -> ... { ... }`, the name is `"my_instruction"`.
pub fn get_instruction(name: &str) -> ArcisInstructionWithInfo {
    let workspace_root = match find_workspace_root() {
        Ok(path) => path,
        Err(error) => panic!("Unable to find workspace root: {}", error),
    };
    let path = workspace_root.join(format!("build/{}.arcis.ir", name));
    let ir =
        read_ir(path.to_str().expect("Cannot convert path to string.")).expect("Cannot read IR.");
    let output_domains = ir.get_output_domains();
    let instruction = ir.optimize_into_circuitable().to_async_mpc_circuit();
    ArcisInstructionWithInfo {
        instruction,
        output_domains,
    }
}

/// Walks up from `CARGO_MANIFEST_DIR` to find the nearest `Cargo.toml`
/// containing a `[workspace]` section.
fn find_workspace_root() -> Result<std::path::PathBuf, &'static str> {
    let start_dir = std::path::PathBuf::from(
        std::env::var("CARGO_MANIFEST_DIR").map_err(|_| "CARGO_MANIFEST_DIR not set")?,
    );
    let mut dir = start_dir.clone();
    loop {
        let manifest = dir.join("Cargo.toml");
        if manifest.exists() {
            let content = std::fs::read_to_string(&manifest).map_err(|_| "cannot read manifest")?;
            if content.contains("[workspace]") {
                return Ok(dir);
            }
        }
        if !dir.pop() {
            return Ok(start_dir);
        }
    }
}