miri add write_bytes method to Memory doing bounds-checks and supporting iterators

2 files changed, 27 insertions, 3 deletions

diff --git a/src/librustc/mir/interpret/allocation.rs b/src/librustc/mir/interpret/allocation.rs
index 15e6cb6bcab..3bcde8defdf 100644
--- a/src/librustc/mir/interpret/allocation.rs
+++ b/src/librustc/mir/interpret/allocation.rs
@@ -346,11 +346,16 @@ impl<'tcx, Tag: Copy, Extra: AllocationExtra<Tag>> Allocation<Tag, Extra> {
         &mut self,
         cx: &impl HasDataLayout,
         ptr: Pointer<Tag>,
-        src: &[u8],
+        src: impl IntoIterator<Item=u8, IntoIter: iter::ExactSizeIterator>,
     ) -> InterpResult<'tcx>
     {
+        let mut src = src.into_iter();
         let bytes = self.get_bytes_mut(cx, ptr, Size::from_bytes(src.len() as u64))?;
-        bytes.clone_from_slice(src);
+        // `zip` would stop when the first iterator ends; we want to definitely
+        // cover all of `bytes`.
+        for dest in bytes {
+            *dest = src.next().expect("iterator was shorter than it said it would be");
+        }
         Ok(())
     }
 
diff --git a/src/librustc_mir/interpret/memory.rs b/src/librustc_mir/interpret/memory.rs
index 924474c5317..0b65e9742b6 100644
--- a/src/librustc_mir/interpret/memory.rs
+++ b/src/librustc_mir/interpret/memory.rs
@@ -7,7 +7,7 @@
 //! short-circuiting the empty case!
 
 use std::collections::VecDeque;
-use std::ptr;
+use std::{ptr, iter};
 use std::borrow::Cow;
 
 use rustc::ty::{self, Instance, ParamEnv, query::TyCtxtAt};
@@ -785,6 +785,25 @@ impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> Memory<'mir, 'tcx, M> {
         self.get(ptr.alloc_id)?.read_c_str(self, ptr)
     }
 
+    /// Writes the given stream of bytes into memory.
+    ///
+    /// Performs appropriate bounds checks.
+    pub fn write_bytes(
+        &mut self,
+        ptr: Scalar<M::PointerTag>,
+        src: impl IntoIterator<Item=u8, IntoIter: iter::ExactSizeIterator>,
+    ) -> InterpResult<'tcx>
+    {
+        let src = src.into_iter();
+        let size = Size::from_bytes(src.len() as u64);
+        let ptr = match self.check_ptr_access(ptr, size, Align::from_bytes(1).unwrap())? {
+            Some(ptr) => ptr,
+            None => return Ok(()), // zero-sized access
+        };
+        let tcx = self.tcx.tcx;
+        self.get_mut(ptr.alloc_id)?.write_bytes(&tcx, ptr, src)
+    }
+
     /// Expects the caller to have checked bounds and alignment.
     pub fn copy(
         &mut self,