Skip to content

Commit

Permalink
feat(hlapi): add tag system
Browse files Browse the repository at this point in the history 
Tag

The `Tag` allows to store bytes alongside of entities (keys, and ciphertext)
the main purpose of this system is to `tag` / identify ciphertext with their keys.

* When encrypted, a ciphertext gets the tag of the key used to encrypt it.
* Ciphertexts resulting from operations (add, sub, etc.) get the tag from the ServerKey used
* PublicKey gets its tag from the ClientKey that was used to create it
* ServerKey gets its tag from the ClientKey that was used to create it

User can change the tag of any entities at any point.

BREAKING CHANGE: Many of the into_raw_parts and from_raw_parts changed
to accommodate the addition of the `tag``
  • Loading branch information
@tmontaigu
tmontaigu committed Aug 12, 2024
1 parent 663322c commit 35efd30
Show file tree
Hide file tree
Showing 39 changed files with 1,902 additions and 515 deletions.
4 changes: 2 additions & 2 deletions tfhe/src/high_level_api/array.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ pub fn fhe_uint_array_eq<Id: FheUintId>(lhs: &[FheUint<Id>], rhs: &[FheUint<Id>]
let result = cpu_keys
.pbs_key()
.all_eq_slices_parallelized(&tmp_lhs, &tmp_rhs);
FheBool::new(result)
FheBool::new(result, cpu_keys.tag.clone())
})
}

Expand All @@ -37,6 +37,6 @@ pub fn fhe_uint_array_contains_sub_slice<Id: FheUintId>(
let result = cpu_keys
.pbs_key()
.contains_sub_slice_parallelized(&tmp_lhs, &tmp_pattern);
FheBool::new(result)
FheBool::new(result, cpu_keys.tag.clone())
})
}
51 changes: 43 additions & 8 deletions tfhe/src/high_level_api/backward_compatibility/booleans.rs
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
#![allow(deprecated)]

use serde::{Deserialize, Serialize};
use tfhe_versionable::{Versionize, VersionsDispatch};
use tfhe_versionable::{Upgrade, Versionize, VersionsDispatch};

use crate::high_level_api::booleans::InnerBooleanVersionOwned;
use crate::high_level_api::booleans::{InnerBooleanVersionOwned, InnerCompressedFheBool};
use crate::integer::ciphertext::{CompactCiphertextList, DataKind};
use crate::{CompactCiphertextList as HlCompactCiphertextList, CompressedFheBool, Error, FheBool};
use crate::{
CompactCiphertextList as HlCompactCiphertextList, CompressedFheBool, Error, FheBool, Tag,
};
use std::convert::Infallible;

// Manual impl
#[derive(Serialize, Deserialize)]
Expand All @@ -24,9 +27,33 @@ pub enum CompactFheBoolVersions {
V0(CompactFheBool),
}

#[derive(VersionsDispatch)]
pub enum InnerCompressedFheBoolVersions {
V0(InnerCompressedFheBool),
}

// Before V1 where we added the Tag, the CompressedFheBool
// was simply the inner enum
type CompressedFheBoolV0 = InnerCompressedFheBool;

impl Upgrade<CompressedFheBool> for CompressedFheBoolV0 {
type Error = Infallible;

fn upgrade(self) -> Result<CompressedFheBool, Self::Error> {
Ok(CompressedFheBool {
inner: match self {
Self::Seeded(s) => Self::Seeded(s),
Self::ModulusSwitched(m) => Self::ModulusSwitched(m),
},
tag: Tag::default(),
})
}
}

#[derive(VersionsDispatch)]
pub enum CompressedFheBoolVersions {
V0(CompressedFheBool),
V0(CompressedFheBoolV0),
V1(CompressedFheBool),
}

#[derive(VersionsDispatch)]
Expand Down Expand Up @@ -56,14 +83,18 @@ impl CompactFheBool {
.iter_mut()
.for_each(|info| *info = DataKind::Boolean);

let hl_list = HlCompactCiphertextList(self.list);
let hl_list = HlCompactCiphertextList {
inner: self.list,
tag: Tag::default(),
};
let list = hl_list.expand()?;

let block = list
.inner
.get::<crate::integer::BooleanBlock>(0)
.ok_or_else(|| Error::new("Failed to expand compact list".to_string()))??;

let mut ciphertext = FheBool::new(block);
let mut ciphertext = FheBool::new(block, Tag::default());
ciphertext.ciphertext.move_to_device_of_server_key_if_set();
Ok(ciphertext)
}
Expand All @@ -86,17 +117,21 @@ impl CompactFheBoolList {
.iter_mut()
.for_each(|info| *info = DataKind::Boolean);

let hl_list = HlCompactCiphertextList(self.list);
let hl_list = HlCompactCiphertextList {
inner: self.list,
tag: Tag::default(),
};
let list = hl_list.expand()?;
let len = list.len();

(0..len)
.map(|idx| {
let block = list
.inner
.get::<crate::integer::BooleanBlock>(idx)
.ok_or_else(|| Error::new("Failed to expand compact list".to_string()))??;

let mut ciphertext = FheBool::new(block);
let mut ciphertext = FheBool::new(block, Tag::default());
ciphertext.ciphertext.move_to_device_of_server_key_if_set();
Ok(ciphertext)
})
Expand Down
86 changes: 71 additions & 15 deletions tfhe/src/high_level_api/backward_compatibility/integers.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,11 @@ use crate::integer::ciphertext::{
};
use crate::shortint::ciphertext::CompressedModulusSwitchedCiphertext;
use crate::shortint::{Ciphertext, ServerKey};
use crate::{CompactCiphertextList as HlCompactCiphertextList, Error};
use crate::{CompactCiphertextList as HlCompactCiphertextList, Error, Tag};
use serde::{Deserialize, Serialize};

use self::unsigned::RadixCiphertext;

// Manual impl
#[derive(Serialize, Deserialize)]
#[cfg_attr(tfhe_lints, allow(tfhe_lints::serialize_without_versionize))]
Expand Down Expand Up @@ -67,11 +69,11 @@ impl Upgrade<CompressedSignedRadixCiphertext> for CompressedSignedRadixCiphertex
let blocks = ct
.blocks
.par_iter()
.map(|a| old_sk_decompress(&sk.key.key, a))
.map(|a| old_sk_decompress(&sk.key.key.key, a))
.collect();

let radix = BaseSignedRadixCiphertext { blocks };
sk.key
sk.pbs_key()
.switch_modulus_and_compress_signed_parallelized(&radix)
});
Ok(CompressedSignedRadixCiphertext::ModulusSwitched(upgraded))
Expand Down Expand Up @@ -105,11 +107,12 @@ impl Upgrade<CompressedRadixCiphertext> for CompressedRadixCiphertextV0 {
let blocks = ct
.blocks
.par_iter()
.map(|a| old_sk_decompress(&sk.key.key, a))
.map(|a| old_sk_decompress(&sk.key.key.key, a))
.collect();

let radix = BaseRadixCiphertext { blocks };
sk.key.switch_modulus_and_compress_parallelized(&radix)
sk.pbs_key()
.switch_modulus_and_compress_parallelized(&radix)
});
Ok(CompressedRadixCiphertext::ModulusSwitched(upgraded))
}
Expand Down Expand Up @@ -143,19 +146,56 @@ pub enum CompactFheIntListVersions<Id: FheIntId> {
V0(CompactFheIntList<Id>),
}

#[derive(Version)]
pub struct FheUintV0<Id: FheUintId> {
pub(in crate::high_level_api) ciphertext: RadixCiphertext,
pub(in crate::high_level_api) id: Id,
}

impl<Id: FheUintId> Upgrade<FheUint<Id>> for FheUintV0<Id> {
type Error = Infallible;

fn upgrade(self) -> Result<FheUint<Id>, Self::Error> {
Ok(FheUint {
ciphertext: self.ciphertext,
id: self.id,
tag: Tag::default(),
})
}
}

#[derive(VersionsDispatch)]
pub enum FheUintVersions<Id: FheUintId> {
V0(FheUint<Id>),
V0(FheUintV0<Id>),
V1(FheUint<Id>),
}

#[derive(VersionsDispatch)]
pub enum CompactFheUintVersions<Id: FheUintId> {
V0(CompactFheUint<Id>),
}

#[derive(Version)]
pub struct CompressedFheUintV0<Id>
where
Id: FheUintId,
{
pub(in crate::high_level_api) ciphertext: CompressedRadixCiphertext,
pub(in crate::high_level_api) id: Id,
}

impl<Id: FheUintId> Upgrade<CompressedFheUint<Id>> for CompressedFheUintV0<Id> {
type Error = Infallible;

fn upgrade(self) -> Result<CompressedFheUint<Id>, Self::Error> {
Ok(CompressedFheUint::new(self.ciphertext, Tag::default()))
}
}

#[derive(VersionsDispatch)]
pub enum CompressedFheUintVersions<Id: FheUintId> {
V0(CompressedFheUint<Id>),
V0(CompressedFheUintV0<Id>),
V1(CompressedFheUint<Id>),
}

#[derive(VersionsDispatch)]
Expand Down Expand Up @@ -186,13 +226,17 @@ where
.info
.iter_mut()
.for_each(|info| *info = DataKind::Signed(info.num_blocks()));
let hl_list = HlCompactCiphertextList(self.list);
let hl_list = HlCompactCiphertextList {
inner: self.list,
tag: Tag::default(),
};
let list = hl_list.expand()?;

let ct = list
.inner
.get::<crate::integer::SignedRadixCiphertext>(0)
.ok_or_else(|| Error::new("Failed to expand compact list".to_string()))??;
Ok(FheInt::new(ct))
Ok(FheInt::new(ct, Tag::default()))
}
}

Expand All @@ -217,17 +261,21 @@ where
.iter_mut()
.for_each(|info| *info = DataKind::Signed(info.num_blocks()));

let hl_list = HlCompactCiphertextList(self.list);
let hl_list = HlCompactCiphertextList {
inner: self.list,
tag: Tag::default(),
};
let list = hl_list.expand()?;

let len = list.len();

(0..len)
.map(|idx| {
let ct = list
.inner
.get::<crate::integer::SignedRadixCiphertext>(idx)
.ok_or_else(|| Error::new("Failed to expand compact list".to_string()))??;
Ok(FheInt::new(ct))
Ok(FheInt::new(ct, Tag::default()))
})
.collect::<Result<Vec<_>, _>>()
}
Expand All @@ -254,13 +302,17 @@ where
.iter_mut()
.for_each(|info| *info = DataKind::Unsigned(info.num_blocks()));

let hl_list = HlCompactCiphertextList(self.list);
let hl_list = HlCompactCiphertextList {
inner: self.list,
tag: Tag::default(),
};
let list = hl_list.expand()?;

let ct = list
.inner
.get::<crate::integer::RadixCiphertext>(0)
.ok_or_else(|| Error::new("Failed to expand compact list".to_string()))??;
Ok(FheUint::new(ct))
Ok(FheUint::new(ct, Tag::default()))
}
}

Expand All @@ -285,17 +337,21 @@ where
.iter_mut()
.for_each(|info| *info = DataKind::Unsigned(info.num_blocks()));

let hl_list = HlCompactCiphertextList(self.list);
let hl_list = HlCompactCiphertextList {
inner: self.list,
tag: Tag::default(),
};
let list = hl_list.expand()?;

let len = list.len();

(0..len)
.map(|idx| {
let ct = list
.inner
.get::<crate::integer::RadixCiphertext>(idx)
.ok_or_else(|| Error::new("Failed to expand compact list".to_string()))??;
Ok(FheUint::new(ct))
Ok(FheUint::new(ct, Tag::default()))
})
.collect::<Result<Vec<_>, _>>()
}
Expand Down
Loading

0 comments on commit 35efd30

Please sign in to comment.