From f3a83649364b3d39e2e71510c67168b8719ac797 Mon Sep 17 00:00:00 2001
From: Stephen Chung <schungx@live.com>
Date: Fri, 31 Dec 2021 17:49:38 +0800
Subject: [PATCH] Split blob write into write_utf8 and write_ascii.

---
 src/packages/blob_basic.rs | 52 +++++++++++++++++++++++++++++++-------
 tests/blobs.rs             | 21 +++++++++++++--
 2 files changed, 62 insertions(+), 11 deletions(-)

diff --git a/src/packages/blob_basic.rs b/src/packages/blob_basic.rs
index ecd11b82..06e8ada8 100644
--- a/src/packages/blob_basic.rs
+++ b/src/packages/blob_basic.rs
@@ -645,8 +645,8 @@ pub mod blob_functions {
     pub fn write_be_float(blob: &mut Blob, start: INT, len: INT, value: FLOAT) {
         write_float(blob, start, len, value, false)
     }
-    #[rhai_fn(name = "write")]
-    pub fn write_string(blob: &mut Blob, start: INT, len: INT, string: &str) {
+    #[inline]
+    fn write_string(blob: &mut Blob, start: INT, len: INT, string: &str, ascii_only: bool) {
         if len <= 0 || blob.is_empty() || string.is_empty() {
             return;
         }
@@ -670,18 +670,52 @@ pub mod blob_functions {
 
         let len = usize::min(len, string.len());
 
-        blob[start..][..len].copy_from_slice(string[..len].as_bytes());
+        if ascii_only {
+            string
+                .chars()
+                .filter(char::is_ascii)
+                .take(len)
+                .map(|ch| ch as u8)
+                .enumerate()
+                .for_each(|(i, x)| blob[start + i] = x);
+        } else {
+            blob[start..][..len].copy_from_slice(&string.as_bytes()[..len]);
+        }
     }
-    #[rhai_fn(name = "write")]
-    pub fn write_string_range(blob: &mut Blob, range: ExclusiveRange, string: &str) {
+    #[rhai_fn(name = "write_utf8")]
+    pub fn write_utf8_string(blob: &mut Blob, start: INT, len: INT, string: &str) {
+        write_string(blob, start, len, string, false)
+    }
+    #[rhai_fn(name = "write_utf8")]
+    pub fn write_utf8_string_range(blob: &mut Blob, range: ExclusiveRange, string: &str) {
         let start = INT::max(range.start, 0);
         let end = INT::max(range.end, start);
-        write_string(blob, start, end - start, string)
+        write_string(blob, start, end - start, string, false)
     }
-    #[rhai_fn(name = "write")]
-    pub fn write_string_range_inclusive(blob: &mut Blob, range: InclusiveRange, string: &str) {
+    #[rhai_fn(name = "write_utf8")]
+    pub fn write_utf8_string_range_inclusive(blob: &mut Blob, range: InclusiveRange, string: &str) {
         let start = INT::max(*range.start(), 0);
         let end = INT::max(*range.end(), start);
-        write_string(blob, start, end - start + 1, string)
+        write_string(blob, start, end - start + 1, string, false)
+    }
+    #[rhai_fn(name = "write_ascii")]
+    pub fn write_ascii_string(blob: &mut Blob, start: INT, len: INT, string: &str) {
+        write_string(blob, start, len, string, true)
+    }
+    #[rhai_fn(name = "write_ascii")]
+    pub fn write_ascii_string_range(blob: &mut Blob, range: ExclusiveRange, string: &str) {
+        let start = INT::max(range.start, 0);
+        let end = INT::max(range.end, start);
+        write_string(blob, start, end - start, string, true)
+    }
+    #[rhai_fn(name = "write_ascii")]
+    pub fn write_ascii_string_range_inclusive(
+        blob: &mut Blob,
+        range: InclusiveRange,
+        string: &str,
+    ) {
+        let start = INT::max(*range.start(), 0);
+        let end = INT::max(*range.end(), start);
+        write_string(blob, start, end - start + 1, string, true)
     }
 }
diff --git a/tests/blobs.rs b/tests/blobs.rs
index dcca26d3..f467be2c 100644
--- a/tests/blobs.rs
+++ b/tests/blobs.rs
@@ -170,15 +170,32 @@ fn test_blobs_parse() -> Result<(), Box<EvalAltResult>> {
     );
 
     assert_eq!(
-        engine.eval::<Blob>(r#"let x = blob(16, 0); write(x, 0, 14, "hello, world!"); x"#)?,
+        engine.eval::<Blob>(r#"let x = blob(16, 0); write_ascii(x, 0, 14, "hello, world!"); x"#)?,
         "hello, world!\0\0\0".as_bytes()
     );
 
     assert_eq!(
-        engine.eval::<Blob>(r#"let x = blob(10, 0); write(x, 3, 5, "hello, world!"); x"#)?,
+        engine.eval::<Blob>(r#"let x = blob(10, 0); write_ascii(x, 3..8, "hello, world!"); x"#)?,
         "\0\0\0hello\0\0".as_bytes()
     );
 
+    assert_eq!(
+        engine.eval::<Blob>(
+            r#"let x = blob(10, 0); write_ascii(x, 0..9, "❤ hello, ❤ world! ❤❤❤"); x"#
+        )?,
+        " hello,  \0".as_bytes()
+    );
+
+    assert_eq!(
+        engine.eval::<Blob>(r#"let x = blob(10, 0); write_utf8(x, 3..9, "❤❤❤❤"); x"#)?,
+        "\0\0\0\u{2764}\u{2764}\0".as_bytes()
+    );
+
+    assert_eq!(
+        engine.eval::<Blob>(r#"let x = blob(10, 0); write_utf8(x, 3..7, "❤❤❤❤"); x"#)?,
+        vec![0, 0, 0, 226, 157, 164, 226, 0, 0, 0]
+    );
+
     Ok(())
 }