From 82d66262264cdbfdd57dec799a229c18c8928ec8 Mon Sep 17 00:00:00 2001
From: Avril <flanchan@cumallover.me>
Date: Wed, 10 Jun 2020 00:05:01 +0100
Subject: [PATCH] fix leak; fix corruption

---
 Cargo.toml  |  2 +-
 src/iter.rs | 21 +++++++++++----------
 src/lib.rs  | 47 +++++++++++++++++++++++++++++++++++++++--------
 src/ptr.rs  |  9 +++++++++
 4 files changed, 60 insertions(+), 19 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 184961a..7864893 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -22,4 +22,4 @@ jemalloc = ["jemalloc-sys"]
 
 [dependencies]
 libc = "0.2"
-jemalloc-sys = { version = "0.3", optional = true }
\ No newline at end of file
+jemalloc-sys = { version = "0.3", optional = true }
diff --git a/src/iter.rs b/src/iter.rs
index b11e93d..5991f5b 100644
--- a/src/iter.rs
+++ b/src/iter.rs
@@ -2,8 +2,6 @@ use super::*;
 
 use std::{
     mem::{
-	replace,
-	MaybeUninit,
 	forget,
     },
 };
@@ -14,19 +12,21 @@ use ptr::{
 pub struct IntoIter<T>
 {
     start: *mut T,
-    current: *mut T,
+    current_offset: usize,
     sz: usize,
 }
 
 impl<T> IntoIter<T>
 {
-    fn current_offset(&self) -> usize
+    fn current(&mut self) -> *mut T
     {
-	(self.current as usize) - (self.start as usize)
+	unsafe {
+	    self.start.offset(self.current_offset as isize)
+	}
     }
     fn free_if_needed(&mut self)
     {
-	if self.start != ptr::null() && self.current_offset() >= self.sz {
+	if self.start != ptr::null() && self.current_offset >= self.sz {
 	    unsafe {
 		alloc::free(self.start as VoidPointer);
 	    }
@@ -40,12 +40,13 @@ impl<T> Iterator for IntoIter<T>
     type Item = T;
     fn next(&mut self) -> Option<Self::Item>
     {
-	let output = if self.current_offset() >= self.sz {
+	let output = if self.current_offset >= self.sz || self.sz == 0 {
 	    None
 	} else {
 	    unsafe {
-		let output = replace(&mut (*self.current), MaybeUninit::zeroed().assume_init());
-		self.current = self.current.offset(1);
+		let output = crate::ptr::take(self.current());//replace(&mut (*self.current), MaybeUninit::zeroed().assume_init());
+		self.current_offset+=1;
+
 		Some(output)
 	    }
 	};
@@ -71,7 +72,7 @@ impl<T> IntoIterator for HeapArray<T>
     {
 	let output = Self::IntoIter {
 	    start: self.ptr,
-	    current: self.ptr,
+	    current_offset: 0,
 	    sz: self.len(),
 	};
 	forget(self);
diff --git a/src/lib.rs b/src/lib.rs
index 5276e06..6b01431 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -23,6 +23,8 @@ mod tests {
 	assert_eq!(heap.as_slice(), ["test one", "test two"]);
     }
 
+    struct Unit;
+    
     #[test]
     fn zero_size() {
 	let heap: HeapArray<u8> = heap![];
@@ -30,6 +32,14 @@ mod tests {
 
 	assert_eq!(&heap.as_slice(), &[]);
 	assert_eq!(&heap_zst.as_slice(), &[(),(),()]);
+
+	let heap = heap![Unit; 32];
+	let mut i=0;
+	for _x in heap.into_iter()
+	{
+	    i+=1;
+	}
+	assert_eq!(i, 32);
     }
 
     #[test]
@@ -43,16 +53,16 @@ mod tests {
 	    assert_eq!(x % 2, 1);
 	}
 
-	let non = heap!["string one".to_owned(), "string two".to_owned()];
+	let non = heap!["string one".to_owned(), "string two!".to_owned(), "string".to_owned()];
 
 	let iter = non.into_iter();
-	assert_eq!(iter.len(), 2);
+	assert_eq!(iter.len(), 3);
 	for x in iter
 	{
-	    assert_eq!(&x[..6], "string");
+	    assert_eq!(&x[..6], "string");	    
 	}
     }
-
+    
     #[test]
     fn vec()
     {
@@ -318,6 +328,26 @@ impl<T> HeapArray<T>
 	op
     }
 
+    pub fn free(self)
+    {
+	if self.ptr != ptr::null() {
+	    unsafe {
+		alloc::free(self.ptr as VoidPointer);
+	    }
+	}
+	std::mem::forget(self);
+    }
+
+    pub fn into_slice(self, slice: &mut [T])
+    {
+	let ptr = &mut slice[0] as *mut T;
+	assert!(slice.len() >= self.len());
+	unsafe{
+	    ptr::memmove(ptr as ptr::VoidPointer, self.ptr as ptr::VoidPointer, self.len_bytes());
+	}
+	self.free();
+    }
+
     pub fn from_boxed_slice(bx: Box<[T]>) -> Self
     {
 	#[cfg(feature="assume_libc")]
@@ -326,7 +356,7 @@ impl<T> HeapArray<T>
 	    Self::from_raw_parts(Box::<[T]>::into_raw(bx) as *mut T, len)
 	}
 	#[cfg(not(feature="assume_libc"))]
-	{ //TODO: wtf?
+	{
 	    let len = bx.len();
 	    let out = Self::from(Vec::from(bx));
 	    assert_eq!(len, out.len());
@@ -445,10 +475,11 @@ impl<T> From<HeapArray<T>> for Vec<T>
     fn from(ha: HeapArray<T>) -> Self
     {
 	let mut output = Vec::with_capacity(ha.len());
-	for item in ha.into_iter()
-	{
-	    output.push(item);
+	unsafe {
+	    ptr::memmove(output.as_mut_ptr() as ptr::VoidPointer, ha.ptr as ptr::VoidPointer, ha.len_bytes());
+	    output.set_len(ha.len());
 	}
+	ha.free();
 	output
     }
 }
diff --git a/src/ptr.rs b/src/ptr.rs
index cfd6ddc..e4124bf 100644
--- a/src/ptr.rs
+++ b/src/ptr.rs
@@ -33,3 +33,12 @@ pub unsafe fn take<T>(ptr: *mut T) -> T
 {
     mem::replace(&mut *ptr, MaybeUninit::zeroed().assume_init())
 }
+
+pub unsafe fn memcpy(dst: VoidPointer, src: ConstVoidPointer, size: usize) -> VoidPointer
+{
+    libc::memcpy(dst as *mut c_void, src as *const c_void, size as size_t) as VoidPointer
+}
+pub unsafe fn memmove(dst: VoidPointer, src: ConstVoidPointer, size: usize) -> VoidPointer
+{
+    libc::memmove(dst as *mut c_void, src as *const c_void, size as size_t) as VoidPointer
+}