From dd179a850e66f78c86a5c1b632aadaa6acfb966e Mon Sep 17 00:00:00 2001
From: Michal Nazarewicz <mina86@mina86.com>
Date: Mon, 30 Oct 2023 21:25:25 +0100
Subject: [PATCH] Add support for types from ascii ccrate

---
 borsh/Cargo.toml     |  3 ++-
 borsh/src/de/mod.rs  | 22 ++++++++++++++++++++++
 borsh/src/schema.rs  | 25 +++++++++++++++++++++++++
 borsh/src/ser/mod.rs | 26 ++++++++++++++++++++++++++
 4 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/borsh/Cargo.toml b/borsh/Cargo.toml
index f0306c03b..c1981ff99 100644
--- a/borsh/Cargo.toml
+++ b/borsh/Cargo.toml
@@ -27,10 +27,11 @@ required-features = ["std", "unstable__schema"]
 cfg_aliases = "0.1.0"
 
 [dependencies]
+ascii = { version = "1.1", optional = true }
 borsh-derive = { path = "../borsh-derive", version = "~1.1.1", optional = true }
 
 # hashbrown can be used in no-std context.
-# NOTE: There is no reason to restrict use of older versions, but we don't want to get 
+# NOTE: There is no reason to restrict use of older versions, but we don't want to get
 # sudden breaking changes with an open range of versions, so we limit the range by not yet released 0.15.0 version:
 hashbrown = { version = ">=0.11,<0.15.0", optional = true }
 bytes = { version = "1", optional = true }
diff --git a/borsh/src/de/mod.rs b/borsh/src/de/mod.rs
index ed57938e1..b41cb8603 100644
--- a/borsh/src/de/mod.rs
+++ b/borsh/src/de/mod.rs
@@ -370,6 +370,28 @@ impl BorshDeserialize for String {
     }
 }
 
+#[cfg(feature = "ascii")]
+impl BorshDeserialize for ascii::AsciiString {
+    #[inline]
+    fn deserialize_reader<R: Read>(reader: &mut R) -> Result<Self> {
+        let bytes = Vec::<u8>::deserialize_reader(reader)?;
+        ascii::AsciiString::from_ascii(bytes)
+            .map_err(|err| Error::new(ErrorKind::InvalidData, err.to_string()))
+    }
+}
+
+#[cfg(feature = "ascii")]
+#[test]
+fn test_ascii() {
+    let encoded = borsh::to_vec("foo").unwrap();
+
+    let got = ascii::AsciiString::deserialize_reader(&mut &encoded[..]).unwrap();
+    assert_eq!("foo", got.as_str());
+
+    let encoded = borsh::to_vec("żółw").unwrap();
+    ascii::AsciiString::deserialize_reader(&mut &encoded[..]).unwrap_err();
+}
+
 impl<T> BorshDeserialize for Vec<T>
 where
     T: BorshDeserialize,
diff --git a/borsh/src/schema.rs b/borsh/src/schema.rs
index 9657f1a5a..810726ecf 100644
--- a/borsh/src/schema.rs
+++ b/borsh/src/schema.rs
@@ -384,6 +384,19 @@ impl BorshSchema for String {
         str::declaration()
     }
 }
+
+#[cfg(feature = "ascii")]
+impl BorshSchema for ascii::AsciiString {
+    #[inline]
+    fn add_definitions_recursively(definitions: &mut BTreeMap<Declaration, Definition>) {
+        str::add_definitions_recursively(definitions);
+    }
+    #[inline]
+    fn declaration() -> Declaration {
+        str::declaration()
+    }
+}
+
 impl BorshSchema for str {
     #[inline]
     fn add_definitions_recursively(definitions: &mut BTreeMap<Declaration, Definition>) {
@@ -401,6 +414,18 @@ impl BorshSchema for str {
     }
 }
 
+#[cfg(feature = "ascii")]
+impl BorshSchema for ascii::AsciiStr {
+    #[inline]
+    fn add_definitions_recursively(definitions: &mut BTreeMap<Declaration, Definition>) {
+        str::add_definitions_recursively(definitions);
+    }
+    #[inline]
+    fn declaration() -> Declaration {
+        str::declaration()
+    }
+}
+
 impl BorshSchema for core::ops::RangeFull {
     #[inline]
     fn add_definitions_recursively(definitions: &mut BTreeMap<Declaration, Definition>) {
diff --git a/borsh/src/ser/mod.rs b/borsh/src/ser/mod.rs
index 1410edf1a..a7924ba26 100644
--- a/borsh/src/ser/mod.rs
+++ b/borsh/src/ser/mod.rs
@@ -210,6 +210,32 @@ impl BorshSerialize for String {
     }
 }
 
+#[cfg(feature = "ascii")]
+impl BorshSerialize for ascii::AsciiStr {
+    #[inline]
+    fn serialize<W: Write>(&self, writer: &mut W) -> Result<()> {
+        self.as_bytes().serialize(writer)
+    }
+}
+
+#[cfg(feature = "ascii")]
+impl BorshSerialize for ascii::AsciiString {
+    #[inline]
+    fn serialize<W: Write>(&self, writer: &mut W) -> Result<()> {
+        self.as_bytes().serialize(writer)
+    }
+}
+
+#[cfg(feature = "ascii")]
+#[test]
+fn test_ascii() {
+    let val = ascii::AsciiStr::from_ascii("foo").unwrap();
+    let encoded = borsh::to_vec(val).unwrap();
+
+    let got: String = borsh::BorshDeserialize::deserialize_reader(&mut &encoded[..]).unwrap();
+    assert_eq!("foo", got);
+}
+
 /// Helper method that is used to serialize a slice of data (without the length marker).
 #[inline]
 fn serialize_slice<T: BorshSerialize, W: Write>(data: &[T], writer: &mut W) -> Result<()> {