github/Sunscreen
1
1
Fork 0
mirror of https://github.com/Sunscreen-tech/Sunscreen.git synced 2026-05-11 03:00:37 -04:00
Code Issues Packages Projects Releases Wiki Activity

changed uses of current fhe and zkp contexts to generic enum

Browse Source
This commit is contained in:
Matthew Liu
2023-08-01 13:53:26 -07:00
parent 6da0352a84
commit c575d24c77
7 changed files with 18 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
sunscreen/src/fhe/mod.rs
View File

@@ -179,9 +179,10 @@ where
let ctx = option
.as_mut()
.expect("Called Ciphertext::new() outside of a context.")
.unwrap_fhe_mut();
.unwrap_fhe_mut()
.unwrap();
f(ctx.unwrap())
f(ctx)
})
}

4
sunscreen/src/lib.rs
View File

@@ -65,6 +65,7 @@ pub mod types;
use fhe::{FheOperation, Literal, FheContext};
use petgraph::stable_graph::StableGraph;
use seal_fhe::Context;
use serde::{Deserialize, Serialize};
use sunscreen_runtime::{marker, Fhe, FheZkp, Zkp};
use sunscreen_zkp_backend::CompiledZkpProgram;
@@ -88,8 +89,7 @@ pub use sunscreen_runtime::{
pub use sunscreen_zkp_backend::{BackendField, Error as ZkpError, Result as ZkpResult, ZkpBackend};
pub use zkp::ZkpProgramFn;
pub use zkp::{
invoke_gadget, with_zkp_ctx, ZkpContext, ZkpContextOps, ZkpData, ZkpFrontendCompilation,
CURRENT_ZKP_CTX,
invoke_gadget, with_zkp_ctx, ZkpContext, ZkpContextOps, ZkpData, ZkpFrontendCompilation
};
pub use fhe::{
with_fhe_ctx, CURRENT_PROGRAM_CTX

2
sunscreen/src/types/intern/fhe_program_node.rs
View File

@@ -37,7 +37,7 @@ use std::ops::{Add, Div, Mul, Neg, Shl, Shr, Sub};
* construction.
*
* # Undefined behavior
* These types must be constructed while [`CURRENT_FHE_CTX`][crate::fhe::CURRENT_FHE_CTX] refers to a valid
* These types must be constructed while [`CURRENT_PROGRAM_CTX`][crate::fhe::CURRENT_PROGRAM_CTX] refers to a valid
* [`FheContext`](crate::fhe::FheContext). Furthermore, no [`FheProgramNode`] should outlive the said context.
* Violating any of these conditions may result in memory corruption or
* use-after-free.

4
sunscreen/src/types/intern/input.rs
View File

@@ -1,4 +1,4 @@
use crate::fhe::{with_fhe_ctx, FheContextOps};
use crate::{fhe::{with_fhe_ctx, FheContextOps}, ContextEnum};
pub use crate::types::{intern::FheProgramNode, Cipher, FheType, NumCiphertexts, TypeName};
/**
@@ -86,7 +86,7 @@ fn can_create_inputs() {
security_level: SecurityLevel::TC128,
});
ctx.swap(&RefCell::new(Some(unsafe { transmute(&mut context) })));
ctx.swap(&RefCell::new(Some(unsafe { transmute(&mut ContextEnum::Fhe(context.clone())) })));
let scalar_node: FheProgramNode<Rational> = FheProgramNode::input();
let mut offset = 0;

17
sunscreen/src/zkp/mod.rs
View File

@@ -4,7 +4,7 @@ use sunscreen_compiler_common::DebugData;
use sunscreen_runtime::CallSignature;
use sunscreen_zkp_backend::{BackendField, BigInt, Gadget, Operation as JitOperation};
use crate::Result;
use crate::{Result, CURRENT_PROGRAM_CTX};
use std::collections::HashMap;
use std::hash::Hash;
@@ -187,6 +187,7 @@ impl Operation {
* An implementation detail of a ZKP program. During compilation, it
* tracks how many public and private inputs have been added.
*/
#[derive(Clone)]
pub struct ZkpData {
next_public_input: usize,
next_private_input: usize,
@@ -367,14 +368,6 @@ impl Render for Operation {
}
}
thread_local! {
/**
* Contains the graph of a ZKP program during compilation. An
* implementation detail and not for public consumption.
*/
pub static CURRENT_ZKP_CTX: RefCell<Option<&'static mut ZkpContext>> = RefCell::new(None);
}
/**
* Runs the specified closure, injecting the current
* [`fhe_program`](crate::fhe_program) context.
@@ -383,11 +376,13 @@ pub fn with_zkp_ctx<F, R>(f: F) -> R
where
F: FnOnce(&mut ZkpContext) -> R,
{
CURRENT_ZKP_CTX.with(|ctx| {
CURRENT_PROGRAM_CTX.with(|ctx| {
let mut option = ctx.borrow_mut();
let ctx = option
.as_mut()
.expect("Called with_zkp_ctx() outside of a context.");
.expect("Called with_zkp_ctx() outside of a context.")
.unwrap_zkp_mut()
.unwrap();
f(ctx)
})

4
sunscreen_compiler_macros/src/fhe_program.rs
View File

@@ -134,7 +134,7 @@ pub fn fhe_program_impl(
fn build(&self, params: &sunscreen::Params) -> sunscreen::Result<sunscreen::fhe::FheFrontendCompilation> {
use std::cell::RefCell;
use std::mem::transmute;
use sunscreen::{fhe::{CURRENT_PROGRAM_CTX, FheContext}, Error, INDEX_ARENA, Result, Params, SchemeType, Value, types::{intern::{FheProgramNode, Input, Output}, NumCiphertexts, Type, TypeName, SwapRows, LaneCount, TypeNameInstance}};
use sunscreen::{fhe::{CURRENT_PROGRAM_CTX, FheContext}, ContextEnum, Error, INDEX_ARENA, Result, Params, SchemeType, Value, types::{intern::{FheProgramNode, Input, Output}, NumCiphertexts, Type, TypeName, SwapRows, LaneCount, TypeNameInstance}};
if SchemeType::Bfv != params.scheme_type {
return Err(Error::IncorrectScheme)
@@ -153,7 +153,7 @@ pub fn fhe_program_impl(
// Transmute away the lifetime to 'static. So long as we are careful with internal()
// panicing, this is safe because we set the context back to none before the funtion
// returns.
ctx.swap(&RefCell::new(Some(unsafe { transmute(&mut context) })));
ctx.swap(&RefCell::new(Some(unsafe { transmute(&mut ContextEnum::Fhe(context.clone())) })));
#(#var_decl)*

4
sunscreen_compiler_macros/src/zkp_program.rs
View File

@@ -171,7 +171,7 @@ fn parse_inner(_attr_params: ZkpProgramAttrs, input_fn: ItemFn) -> Result<TokenS
fn build(&self) -> sunscreen::Result<sunscreen::ZkpFrontendCompilation> {
use std::cell::RefCell;
use std::mem::transmute;
use sunscreen::{CURRENT_PROGRAM_CTX, ZkpContext, ZkpData, Error, INDEX_ARENA, Result, types::{zkp::{ProgramNode, CreateZkpProgramInput, ConstrainEq, IntoProgramNode}, TypeName}};
use sunscreen::{CURRENT_PROGRAM_CTX, ZkpContext, ContextEnum, ZkpData, Error, INDEX_ARENA, Result, types::{zkp::{ProgramNode, CreateZkpProgramInput, ConstrainEq, IntoProgramNode}, TypeName}};
let mut context = ZkpContext::new(ZkpData::new());
@@ -179,7 +179,7 @@ fn parse_inner(_attr_params: ZkpProgramAttrs, input_fn: ItemFn) -> Result<TokenS
// Transmute away the lifetime to 'static. So long as we are careful with internal()
// panicing, this is safe because we set the context back to none before the function
// returns.
ctx.swap(&RefCell::new(Some(unsafe { transmute(&mut context) })));
ctx.swap(&RefCell::new(Some(unsafe { transmute(&mut ContextEnum::Zkp(context.clone())) })));
#[allow(clippy::type_complexity)]
#[forbid(unused_variables)]