From 9daf016d6377ed5511b62731923f488b00224568 Mon Sep 17 00:00:00 2001
From: bakaq <bakaq@bakaq.net>
Date: Sun, 2 Mar 2025 18:51:57 -0300
Subject: [PATCH] Fix parsing of \x0\ in partial strings

---
 src/functor_macro.rs | 10 +++---
 src/machine/heap.rs  | 85 +++++++++++++++++++++++++++++++++-----------
 2 files changed, 70 insertions(+), 25 deletions(-)

diff --git a/src/functor_macro.rs b/src/functor_macro.rs
index ee419136..4f9032a7 100644
--- a/src/functor_macro.rs
+++ b/src/functor_macro.rs
@@ -485,7 +485,7 @@ mod tests {
         let mut functor_writer = Heap::functor_writer(functor);
         functor_writer(&mut heap).unwrap();
 
-        assert_eq!(heap.cell_len(), 8);
+        assert_eq!(heap.cell_len(), 10);
 
         assert_eq!(heap[0], atom_as_cell!(atom!("second"), 1));
         assert_eq!(heap[1], pstr_loc_as_cell!(heap_index!(2)));
@@ -493,12 +493,14 @@ mod tests {
             heap.slice_to_str(heap_index!(2), "a stuttered".len()),
             "a stuttered"
         );
-        assert_eq!(heap[4], pstr_loc_as_cell!(heap_index!(5)));
+        assert_eq!(heap[4], list_loc_as_cell!(5));
+        assert_eq!(heap[5], char_as_cell!('\u{0}'));
+        assert_eq!(heap[6], pstr_loc_as_cell!(heap_index!(7)));
         assert_eq!(
-            heap.slice_to_str(heap_index!(5), " string".len()),
+            heap.slice_to_str(heap_index!(7), " string".len()),
             " string"
         );
-        assert_eq!(heap[7], empty_list_as_cell!());
+        assert_eq!(heap[9], empty_list_as_cell!());
     }
 
     #[test]
diff --git a/src/machine/heap.rs b/src/machine/heap.rs
index 11d0fcea..b753809d 100644
--- a/src/machine/heap.rs
+++ b/src/machine/heap.rs
@@ -251,31 +251,73 @@ impl ReservedHeapSection {
     }
 
     pub(crate) fn push_pstr(&mut self, mut src: &str) -> Option<HeapCellValue> {
-        let orig_h = self.cell_len();
-
-        if src.is_empty() {
-            return if orig_h == self.heap_cell_len {
-                // src is empty and always was. nothing allocated
-                // in this case, so nothing to point to in heap.
-                None
-            } else {
-                self.push_cell(heap_loc_as_cell!(orig_h));
-                Some(heap_loc_as_cell!(orig_h))
-            };
-        }
+        let anchor = self.cell_len();
+        let mut ret = None;
 
         loop {
-            let null_char_idx = src.find('\u{0}').unwrap_or(src.len());
-            let cells_written = self.push_pstr_segment(&src[0..null_char_idx]);
-
-            if cells_written == 0 {
-                return None;
-            } else if null_char_idx + 1 < src.len() {
-                let tail_idx = self.cell_len();
-                self.push_cell(pstr_loc_as_cell!(heap_index!(tail_idx + 1)));
+            // Eat the first null chars
+            while let Some('\u{0}') = src.chars().next() {
+                match ret {
+                    Some(_) => {
+                        debug_assert_ne!(anchor, self.cell_len());
+                        self.push_cell(list_loc_as_cell!(self.cell_len() + 1));
+                    }
+                    None => {
+                        debug_assert_eq!(anchor, self.cell_len());
+                        ret = Some(list_loc_as_cell!(self.cell_len()));
+                    }
+                }
+
+                self.push_cell(char_as_cell!('\u{0}'));
+
+                src = &src[1..];
+            }
+
+            if src.is_empty() {
+                return ret;
+            }
+
+            debug_assert!(!src.is_empty());
+
+            if let Some(null_char_idx) = src.find('\u{0}') {
+                debug_assert_ne!(null_char_idx, 0);
+
+                match ret {
+                    Some(_) => {
+                        debug_assert_ne!(anchor, self.cell_len());
+                        self.push_cell(pstr_loc_as_cell!(heap_index!(self.cell_len() + 1)));
+                    }
+                    None => {
+                        debug_assert_eq!(anchor, self.cell_len());
+                        ret = Some(pstr_loc_as_cell!(heap_index!(self.cell_len())));
+                    }
+                }
+
+                self.push_pstr_segment(&src[0..null_char_idx]);
+
+                // Put the \x0\
+                self.push_cell(list_loc_as_cell!(self.cell_len() + 1));
+                self.push_cell(char_as_cell!('\u{0}'));
+
                 src = &src[null_char_idx + 1..];
+                if src.is_empty() {
+                    return ret;
+                }
             } else {
-                return Some(pstr_loc_as_cell!(heap_index!(orig_h)));
+                match ret {
+                    Some(_) => {
+                        debug_assert_ne!(anchor, self.cell_len());
+                        self.push_cell(pstr_loc_as_cell!(heap_index!(self.cell_len() + 1)));
+                    }
+                    None => {
+                        debug_assert_eq!(anchor, self.cell_len());
+                        ret = Some(pstr_loc_as_cell!(heap_index!(self.cell_len())));
+                    }
+                }
+
+                self.push_pstr_segment(&src);
+
+                return ret;
             }
         }
     }
@@ -812,6 +854,7 @@ impl Heap {
 
     pub fn allocate_pstr(&mut self, src: &str) -> Result<Option<PStrWriteInfo>, usize> {
         let size_in_heap = Self::compute_pstr_size(src);
+
         let pstr_loc = heap_index!(self.cell_len());
 
         Ok(if size_in_heap > 0 {
-- 
2.54.0