feat(hlapi): bind string [r]split_once

tfhe/src/high_level_api/strings/ascii/mod.rs

@@ -3,6 +3,7 @@ mod contains;
 mod find;
 mod no_pattern;
 mod replace;
+mod split;
 mod strip;
 mod trim;
 

tfhe/src/high_level_api/strings/ascii/split.rs

@@ -0,0 +1,199 @@
+use crate::high_level_api::global_state::with_internal_keys;
+use crate::high_level_api::keys::InternalServerKey;
+use crate::prelude::FheStringSplitOnce;
+use crate::{ClearString, FheAsciiString, FheBool};
+
+impl FheStringSplitOnce<&Self> for FheAsciiString {
+    /// Splits the encrypted string into two substrings at the first occurrence of the pattern
+    /// and returns a tuple of the two substrings along with a boolean
+    /// indicating if the split occurred.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use tfhe::prelude::*;
+    /// use tfhe::{
+    ///     generate_keys, set_server_key, ClearString, ConfigBuilder, FheAsciiString,
+    ///     FheStringIsEmpty, FheStringLen,
+    /// };
+    ///
+    /// let (client_key, server_key) = generate_keys(ConfigBuilder::default());
+    /// set_server_key(server_key);
+    ///
+    /// let string = FheAsciiString::try_encrypt("tfhe is an fhe scheme", &client_key).unwrap();
+    /// let separator = FheAsciiString::try_encrypt(" ", &client_key).unwrap();
+    ///
+    /// // separator is found
+    /// let (lhs, rhs, split_occurred) = string.split_once(&separator);
+    /// assert!(split_occurred.decrypt(&client_key));
+    /// let lhs_decrypted = lhs.decrypt(&client_key);
+    /// assert_eq!(&lhs_decrypted, "tfhe");
+    /// let rhs_decrypted = rhs.decrypt(&client_key);
+    /// assert_eq!(&rhs_decrypted, "is an fhe scheme");
+    ///
+    /// // separator is not found
+    /// let separator = ClearString::new("_".to_string());
+    /// let (lhs, rhs, split_occurred) = string.split_once(&separator);
+    /// assert!(!split_occurred.decrypt(&client_key));
+    /// let lhs_decrypted = lhs.decrypt(&client_key);
+    /// assert_eq!(&lhs_decrypted, "");
+    /// let rhs_decrypted = rhs.decrypt(&client_key);
+    /// assert_eq!(&rhs_decrypted, "fhe is an fhe scheme");
+    /// ```
+    fn split_once(&self, pat: &Self) -> (Self, Self, FheBool) {
+        with_internal_keys(|keys| match keys {
+            InternalServerKey::Cpu(cpu_key) => {
+                let (inner_1, inner_2, inner_3) = cpu_key
+                    .string_key()
+                    .split_once(&self.inner.on_cpu(), (&*pat.inner.on_cpu()).into());
+                (
+                    Self::new(inner_1, cpu_key.tag.clone()),
+                    Self::new(inner_2, cpu_key.tag.clone()),
+                    FheBool::new(inner_3, cpu_key.tag.clone()),
+                )
+            }
+            #[cfg(feature = "gpu")]
+            InternalServerKey::Cuda(_) => {
+                panic!("gpu does not support strings split_once");
+            }
+        })
+    }
+
+    /// Splits the encrypted string into two substrings at the last occurrence of the pattern
+    /// and returns a tuple of the two substrings along with a boolean
+    /// indicating if the split occurred.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use tfhe::prelude::*;
+    /// use tfhe::{
+    ///     generate_keys, set_server_key, ClearString, ConfigBuilder, FheAsciiString,
+    ///     FheStringIsEmpty, FheStringLen,
+    /// };
+    ///
+    /// let (client_key, server_key) = generate_keys(ConfigBuilder::default());
+    /// set_server_key(server_key);
+    ///
+    /// let string = FheAsciiString::try_encrypt("tfhe is an fhe scheme", &client_key).unwrap();
+    /// let separator = FheAsciiString::try_encrypt(" ", &client_key).unwrap();
+    ///
+    /// let (lhs, rhs, split_occurred) = string.rsplit_once(&separator);
+    /// assert!(split_occurred.decrypt(&client_key));
+    /// let lhs_decrypted = lhs.decrypt(&client_key);
+    /// assert_eq!(&lhs_decrypted, "tfhe is an fhe");
+    /// let rhs_decrypted = rhs.decrypt(&client_key);
+    /// assert_eq!(&rhs_decrypted, "scheme");
+    /// ```
+    fn rsplit_once(&self, pat: &Self) -> (Self, Self, FheBool) {
+        with_internal_keys(|keys| match keys {
+            InternalServerKey::Cpu(cpu_key) => {
+                let (inner_1, inner_2, inner_3) = cpu_key
+                    .string_key()
+                    .rsplit_once(&self.inner.on_cpu(), (&*pat.inner.on_cpu()).into());
+                (
+                    Self::new(inner_1, cpu_key.tag.clone()),
+                    Self::new(inner_2, cpu_key.tag.clone()),
+                    FheBool::new(inner_3, cpu_key.tag.clone()),
+                )
+            }
+            #[cfg(feature = "gpu")]
+            InternalServerKey::Cuda(_) => {
+                panic!("gpu does not support strings rsplit_once");
+            }
+        })
+    }
+}
+
+impl FheStringSplitOnce<&ClearString> for FheAsciiString {
+    /// Splits the encrypted string into two substrings at the first occurrence of the pattern
+    /// and returns a tuple of the two substrings along with a boolean
+    /// indicating if the split occurred.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use tfhe::prelude::*;
+    /// use tfhe::{
+    ///     generate_keys, set_server_key, ClearString, ConfigBuilder, FheAsciiString,
+    ///     FheStringIsEmpty, FheStringLen,
+    /// };
+    ///
+    /// let (client_key, server_key) = generate_keys(ConfigBuilder::default());
+    /// set_server_key(server_key);
+    ///
+    /// let string = FheAsciiString::try_encrypt("tfhe is an fhe scheme", &client_key).unwrap();
+    /// let separator = ClearString::new(" ".to_string());
+    ///
+    /// // separator is found
+    /// let (lhs, rhs, split_occurred) = string.split_once(&separator);
+    /// assert!(split_occurred.decrypt(&client_key));
+    /// let lhs_decrypted = lhs.decrypt(&client_key);
+    /// assert_eq!(&lhs_decrypted, "tfhe");
+    /// let rhs_decrypted = rhs.decrypt(&client_key);
+    /// assert_eq!(&rhs_decrypted, "is an fhe scheme");
+    /// ```
+    fn split_once(&self, pat: &ClearString) -> (Self, Self, FheBool) {
+        with_internal_keys(|keys| match keys {
+            InternalServerKey::Cpu(cpu_key) => {
+                let (inner_1, inner_2, inner_3) = cpu_key
+                    .string_key()
+                    .split_once(&self.inner.on_cpu(), pat.into());
+                (
+                    Self::new(inner_1, cpu_key.tag.clone()),
+                    Self::new(inner_2, cpu_key.tag.clone()),
+                    FheBool::new(inner_3, cpu_key.tag.clone()),
+                )
+            }
+            #[cfg(feature = "gpu")]
+            InternalServerKey::Cuda(_) => {
+                panic!("gpu does not support strings split_once");
+            }
+        })
+    }
+
+    /// Splits the encrypted string into two substrings at the last occurrence of the pattern
+    /// and returns a tuple of the two substrings along with a boolean
+    /// indicating if the split occurred.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use tfhe::prelude::*;
+    /// use tfhe::{
+    ///     generate_keys, set_server_key, ClearString, ConfigBuilder, FheAsciiString,
+    ///     FheStringIsEmpty, FheStringLen,
+    /// };
+    ///
+    /// let (client_key, server_key) = generate_keys(ConfigBuilder::default());
+    /// set_server_key(server_key);
+    ///
+    /// let string = FheAsciiString::try_encrypt("tfhe is an fhe scheme", &client_key).unwrap();
+    /// let separator = ClearString::new(" ".to_string());
+    ///
+    /// let (lhs, rhs, split_occurred) = string.rsplit_once(&separator);
+    /// assert!(split_occurred.decrypt(&client_key));
+    /// let lhs_decrypted = lhs.decrypt(&client_key);
+    /// assert_eq!(&lhs_decrypted, "tfhe is an fhe");
+    /// let rhs_decrypted = rhs.decrypt(&client_key);
+    /// assert_eq!(&rhs_decrypted, "scheme");
+    /// ```
+    fn rsplit_once(&self, pat: &ClearString) -> (Self, Self, FheBool) {
+        with_internal_keys(|keys| match keys {
+            InternalServerKey::Cpu(cpu_key) => {
+                let (inner_1, inner_2, inner_3) = cpu_key
+                    .string_key()
+                    .rsplit_once(&self.inner.on_cpu(), pat.into());
+                (
+                    Self::new(inner_1, cpu_key.tag.clone()),
+                    Self::new(inner_2, cpu_key.tag.clone()),
+                    FheBool::new(inner_3, cpu_key.tag.clone()),
+                )
+            }
+            #[cfg(feature = "gpu")]
+            InternalServerKey::Cuda(_) => {
+                panic!("gpu does not support strings rsplit_once");
+            }
+        })
+    }
+}

tfhe/src/high_level_api/strings/tests/cpu.rs

@@ -35,3 +35,9 @@ fn test_string_strip() {
     let cks = setup_default_cpu();
     super::test_string_strip(&cks);
 }
+
+#[test]
+fn test_string_split_once() {
+    let cks = setup_default_cpu();
+    super::test_string_split_once(&cks);
+}

tfhe/src/high_level_api/strings/tests/mod.rs

@@ -181,3 +181,54 @@ fn test_string_strip(client_key: &ClientKey) {
     let dec = stripped.decrypt(client_key);
     assert_eq!(dec, "The lazy cat");
 }
+
+fn test_string_split_once(client_key: &ClientKey) {
+    let string = FheAsciiString::try_encrypt("Thespacelazyspacecat", client_key).unwrap();
+    let separator = FheAsciiString::try_encrypt("space", client_key).unwrap();
+
+    let (lhs, rhs, split_occurred) = string.split_once(&separator);
+    assert!(split_occurred.decrypt(client_key));
+    let lhs_decrypted = lhs.decrypt(client_key);
+    assert_eq!(&lhs_decrypted, "The");
+    let rhs_decrypted = rhs.decrypt(client_key);
+    assert_eq!(&rhs_decrypted, "lazyspacecat");
+
+    let (lhs, rhs, split_occurred) = string.rsplit_once(&separator);
+    assert!(split_occurred.decrypt(client_key));
+    let lhs_decrypted = lhs.decrypt(client_key);
+    assert_eq!(&lhs_decrypted, "Thespacelazy");
+    let rhs_decrypted = rhs.decrypt(client_key);
+    assert_eq!(&rhs_decrypted, "cat");
+
+    let separator = FheAsciiString::try_encrypt(" ", client_key).unwrap();
+    let (lhs, rhs, split_occurred) = string.split_once(&separator);
+    assert!(!split_occurred.decrypt(client_key));
+    let lhs_decrypted = lhs.decrypt(client_key);
+    assert_eq!(&lhs_decrypted, "");
+    let rhs_decrypted = rhs.decrypt(client_key);
+    assert_eq!(&rhs_decrypted, "hespacelazyspacecat");
+
+    let string = FheAsciiString::try_encrypt("The_lazy_cat", client_key).unwrap();
+    let separator = ClearString::new("_".into());
+    let (lhs, rhs, split_occurred) = string.split_once(&separator);
+    assert!(split_occurred.decrypt(client_key));
+    let lhs_decrypted = lhs.decrypt(client_key);
+    assert_eq!(&lhs_decrypted, "The");
+    let rhs_decrypted = rhs.decrypt(client_key);
+    assert_eq!(&rhs_decrypted, "lazy_cat");
+
+    let (lhs, rhs, split_occurred) = string.rsplit_once(&separator);
+    assert!(split_occurred.decrypt(client_key));
+    let lhs_decrypted = lhs.decrypt(client_key);
+    assert_eq!(&lhs_decrypted, "The_lazy");
+    let rhs_decrypted = rhs.decrypt(client_key);
+    assert_eq!(&rhs_decrypted, "cat");
+
+    let separator = ClearString::new(" ".into());
+    let (lhs, rhs, split_occurred) = string.split_once(&separator);
+    assert!(!split_occurred.decrypt(client_key));
+    let lhs_decrypted = lhs.decrypt(client_key);
+    assert_eq!(&lhs_decrypted, "");
+    let rhs_decrypted = rhs.decrypt(client_key);
+    assert_eq!(&rhs_decrypted, "he_lazy_cat");
+}

tfhe/src/high_level_api/strings/traits.rs

@@ -43,3 +43,12 @@ where
 {
     fn repeat(&self, count: Count) -> Self;
 }
+
+pub trait FheStringSplitOnce<Pat>
+where
+    Self: Sized,
+{
+    fn split_once(&self, pat: Pat) -> (Self, Self, FheBool);
+
+    fn rsplit_once(&self, pat: Pat) -> (Self, Self, FheBool);
+}