Fix allocate_pstr randomly refusing to properly allocate memory

This one was a toughie: it turns out that using `ptr::align_of()`` was
a bad idea, since the buffer in `Heap` itself is not aligned to
`Heap::heap_cell_alignment()`, so `ptr::align_of()` would sometimes
return lower values than expected.

That made for an heisenbug: if the alignment of the heap happened to be 4,
then the bug wouldn't trigger.

Author: adri326

src/heap_iter.rs

@@ -2979,6 +2979,19 @@ mod tests {
 
         wam.machine_st.heap.clear();
 
+    }
+
+    #[test]
+    fn heap_stackless_post_order_iter_pstr() {
+        let mut wam = MockWAM::new();
+
+        let f_atom = atom!("f");
+        let a_atom = atom!("a");
+        let b_atom = atom!("b");
+
+        // clear the heap of resource error data etc
+        wam.machine_st.heap.clear();
+
         // first a 'dangling' partial string, later modified to be a
         // two-part complete string, then a three-part cyclic string
         // involving an uncompacted list of chars.
@@ -3004,9 +3017,13 @@ mod tests {
         }
 
         wam.machine_st.heap[2] = heap_loc_as_cell!(2);
+        assert_eq!(wam.machine_st.heap.cell_len(), 3);
+
         wam.machine_st.allocate_pstr("def").unwrap();
+        assert_eq!(wam.machine_st.heap.cell_len(), 4);
 
         wam.machine_st.heap.push_cell(pstr_loc_as_cell!(0)).unwrap();
+        assert_eq!(wam.machine_st.heap.cell_len(), 5);
 
         {
             let mut iter = stackless_post_order_iter(&mut wam.machine_st.heap, 4);

src/machine/heap.rs

@@ -353,15 +353,7 @@ impl<'a> ReservedHeapSection<'a> {
 
             let zero_region_idx = heap_index!(self.heap_cell_len) + str_byte_len;
 
-            let align_offset = self.heap_ptr
-                .add(zero_region_idx)
-                .align_offset(ALIGN_CELL);
-
-            let align_offset = if align_offset == 0 {
-                ALIGN_CELL
-            } else {
-                align_offset
-            };
+            let align_offset = pstr_sentinel_length(zero_region_idx);
 
             ptr::write_bytes(
                 self.heap_ptr.add(zero_region_idx),
@@ -481,6 +473,22 @@ impl<'a> Index<usize> for ReservedHeapSection<'a> {
     }
 }
 
+/// Computes the number of bytes required to pad a string of length `chunk_len`
+/// with zeroes, such that `chunk_len + pstr_sentinel_length(chunk_len)` is a
+/// multiple of `Heap::heap_cell_alignement()`.
+fn pstr_sentinel_length(chunk_len: usize) -> usize {
+    const ALIGN: usize = Heap::heap_cell_alignment();
+
+    let res = chunk_len.next_multiple_of(ALIGN) - chunk_len;
+
+    // No bytes available in last chunk
+    if res == 0 {
+        ALIGN
+    } else {
+        res
+    }
+}
+
 #[must_use]
 #[derive(Debug)]
 pub struct HeapWriter<'a> {
@@ -635,7 +643,7 @@ impl Heap {
                 if self.free_space() >= len {
                     section = ReservedHeapSection {
                         heap_ptr: self.inner.ptr,
-                        heap_cell_len: cell_index!(self.inner.byte_len),
+                        heap_cell_len: self.cell_len(),
                         pstr_vec: &mut self.pstr_vec,
                     };
                     break;
@@ -810,6 +818,11 @@ impl Heap {
                 }
             }
 
+            // SAFETY:
+            // - Postcondition: from `self.grow()`, `self.inner.byte_len + size_of::<HeapCellValue>()`
+            //   is strictly less than `self.inner.byte_cap`.
+            // - Asserted: `self.cell_len() * size_of::<HeapCellvalue>() <= self.inner.byte_cap`.
+            // - Invariant: from `InnerHeap`, `self.inner.byte_cap < isize::MAX`.
             let cell_ptr = (self.inner.ptr as *mut HeapCellValue).add(self.cell_len());
             cell_ptr.write(cell);
             self.pstr_vec.push(false);
@@ -967,17 +980,7 @@ impl Heap {
 
         const ALIGN_CELL: usize = Heap::heap_cell_alignment();
 
-        let align_offset = unsafe {
-            self.inner.ptr
-                .add(self.inner.byte_len + s_len)
-                .align_offset(ALIGN_CELL)
-        };
-
-        let align_offset = if align_offset == 0 {
-            ALIGN_CELL
-        } else {
-            align_offset
-        };
+        let align_offset = pstr_sentinel_length(s_len);
 
         let copy_size = s_len + align_offset;
 
@@ -1040,8 +1043,9 @@ impl Heap {
         Ok(())
     }
 
-    // assumes the string will be allocated on a ALIGN_CELL-byte boundary
-    pub(crate) const fn compute_pstr_size(src: &str) -> usize {
+    /// Returns the number of bytes needed to store `src` as a `PStr`.
+    /// Assumes the string will be allocated on a ALIGN_CELL-byte boundary.
+    pub(crate) fn compute_pstr_size(src: &str) -> usize {
         const ALIGN_CELL: usize = Heap::heap_cell_alignment();
 
         if src.is_empty() {
@@ -1062,7 +1066,7 @@ impl Heap {
                 null_idx += 1;
             }
 
-            byte_size += (null_idx & !(ALIGN_CELL - 1)) + ALIGN_CELL;
+            byte_size += null_idx.next_multiple_of(ALIGN_CELL);
 
             if (null_idx + 1) % ALIGN_CELL == 0 {
                 byte_size += 2 * mem::size_of::<HeapCellValue>();