]> git.lizzy.rs Git - rust.git/blob - compiler/rustc_middle/src/mir/interpret/allocation.rs
BPF: review fixes
[rust.git] / compiler / rustc_middle / src / mir / interpret / allocation.rs
1 //! The virtual memory representation of the MIR interpreter.
2
3 use std::borrow::Cow;
4 use std::convert::TryFrom;
5 use std::iter;
6 use std::ops::{Deref, DerefMut, Range};
7 use std::ptr;
8
9 use rustc_ast::Mutability;
10 use rustc_data_structures::sorted_map::SortedMap;
11 use rustc_target::abi::{Align, HasDataLayout, Size};
12
13 use super::{
14     read_target_uint, write_target_uint, AllocId, InterpError, Pointer, Scalar, ScalarMaybeUninit,
15     UndefinedBehaviorInfo, UninitBytesAccess, UnsupportedOpInfo,
16 };
17
18 /// This type represents an Allocation in the Miri/CTFE core engine.
19 ///
20 /// Its public API is rather low-level, working directly with allocation offsets and a custom error
21 /// type to account for the lack of an AllocId on this level. The Miri/CTFE core engine `memory`
22 /// module provides higher-level access.
23 #[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, TyEncodable, TyDecodable)]
24 #[derive(HashStable)]
25 pub struct Allocation<Tag = (), Extra = ()> {
26     /// The actual bytes of the allocation.
27     /// Note that the bytes of a pointer represent the offset of the pointer.
28     bytes: Vec<u8>,
29     /// Maps from byte addresses to extra data for each pointer.
30     /// Only the first byte of a pointer is inserted into the map; i.e.,
31     /// every entry in this map applies to `pointer_size` consecutive bytes starting
32     /// at the given offset.
33     relocations: Relocations<Tag>,
34     /// Denotes which part of this allocation is initialized.
35     init_mask: InitMask,
36     /// The alignment of the allocation to detect unaligned reads.
37     /// (`Align` guarantees that this is a power of two.)
38     pub align: Align,
39     /// `true` if the allocation is mutable.
40     /// Also used by codegen to determine if a static should be put into mutable memory,
41     /// which happens for `static mut` and `static` with interior mutability.
42     pub mutability: Mutability,
43     /// Extra state for the machine.
44     pub extra: Extra,
45 }
46
47 /// We have our own error type that does not know about the `AllocId`; that information
48 /// is added when converting to `InterpError`.
49 #[derive(Debug)]
50 pub enum AllocError {
51     /// Encountered a pointer where we needed raw bytes.
52     ReadPointerAsBytes,
53     /// Using uninitialized data where it is not allowed.
54     InvalidUninitBytes(Option<UninitBytesAccess>),
55 }
56 pub type AllocResult<T = ()> = Result<T, AllocError>;
57
58 impl AllocError {
59     pub fn to_interp_error<'tcx>(self, alloc_id: AllocId) -> InterpError<'tcx> {
60         match self {
61             AllocError::ReadPointerAsBytes => {
62                 InterpError::Unsupported(UnsupportedOpInfo::ReadPointerAsBytes)
63             }
64             AllocError::InvalidUninitBytes(info) => InterpError::UndefinedBehavior(
65                 UndefinedBehaviorInfo::InvalidUninitBytes(info.map(|b| (alloc_id, b))),
66             ),
67         }
68     }
69 }
70
71 /// The information that makes up a memory access: offset and size.
72 #[derive(Copy, Clone, Debug)]
73 pub struct AllocRange {
74     pub start: Size,
75     pub size: Size,
76 }
77
78 /// Free-starting constructor for less syntactic overhead.
79 #[inline(always)]
80 pub fn alloc_range(start: Size, size: Size) -> AllocRange {
81     AllocRange { start, size }
82 }
83
84 impl AllocRange {
85     #[inline(always)]
86     pub fn end(self) -> Size {
87         self.start + self.size // This does overflow checking.
88     }
89
90     /// Returns the `subrange` within this range; panics if it is not a subrange.
91     #[inline]
92     pub fn subrange(self, subrange: AllocRange) -> AllocRange {
93         let sub_start = self.start + subrange.start;
94         let range = alloc_range(sub_start, subrange.size);
95         assert!(range.end() <= self.end(), "access outside the bounds for given AllocRange");
96         range
97     }
98 }
99
100 // The constructors are all without extra; the extra gets added by a machine hook later.
101 impl<Tag> Allocation<Tag> {
102     /// Creates a read-only allocation initialized by the given bytes
103     pub fn from_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>, align: Align) -> Self {
104         let bytes = slice.into().into_owned();
105         let size = Size::from_bytes(bytes.len());
106         Self {
107             bytes,
108             relocations: Relocations::new(),
109             init_mask: InitMask::new(size, true),
110             align,
111             mutability: Mutability::Not,
112             extra: (),
113         }
114     }
115
116     pub fn from_byte_aligned_bytes<'a>(slice: impl Into<Cow<'a, [u8]>>) -> Self {
117         Allocation::from_bytes(slice, Align::ONE)
118     }
119
120     pub fn uninit(size: Size, align: Align) -> Self {
121         Allocation {
122             bytes: vec![0; size.bytes_usize()],
123             relocations: Relocations::new(),
124             init_mask: InitMask::new(size, false),
125             align,
126             mutability: Mutability::Mut,
127             extra: (),
128         }
129     }
130 }
131
132 impl Allocation<()> {
133     /// Add Tag and Extra fields
134     pub fn with_tags_and_extra<T, E>(
135         self,
136         mut tagger: impl FnMut(AllocId) -> T,
137         extra: E,
138     ) -> Allocation<T, E> {
139         Allocation {
140             bytes: self.bytes,
141             relocations: Relocations::from_presorted(
142                 self.relocations
143                     .iter()
144                     // The allocations in the relocations (pointers stored *inside* this allocation)
145                     // all get the base pointer tag.
146                     .map(|&(offset, ((), alloc))| {
147                         let tag = tagger(alloc);
148                         (offset, (tag, alloc))
149                     })
150                     .collect(),
151             ),
152             init_mask: self.init_mask,
153             align: self.align,
154             mutability: self.mutability,
155             extra,
156         }
157     }
158 }
159
160 /// Raw accessors. Provide access to otherwise private bytes.
161 impl<Tag, Extra> Allocation<Tag, Extra> {
162     pub fn len(&self) -> usize {
163         self.bytes.len()
164     }
165
166     pub fn size(&self) -> Size {
167         Size::from_bytes(self.len())
168     }
169
170     /// Looks at a slice which may describe uninitialized bytes or describe a relocation. This differs
171     /// from `get_bytes_with_uninit_and_ptr` in that it does no relocation checks (even on the
172     /// edges) at all.
173     /// This must not be used for reads affecting the interpreter execution.
174     pub fn inspect_with_uninit_and_ptr_outside_interpreter(&self, range: Range<usize>) -> &[u8] {
175         &self.bytes[range]
176     }
177
178     /// Returns the mask indicating which bytes are initialized.
179     pub fn init_mask(&self) -> &InitMask {
180         &self.init_mask
181     }
182
183     /// Returns the relocation list.
184     pub fn relocations(&self) -> &Relocations<Tag> {
185         &self.relocations
186     }
187 }
188
189 /// Byte accessors.
190 impl<Tag: Copy, Extra> Allocation<Tag, Extra> {
191     /// The last argument controls whether we error out when there are uninitialized
192     /// or pointer bytes. You should never call this, call `get_bytes` or
193     /// `get_bytes_with_uninit_and_ptr` instead,
194     ///
195     /// This function also guarantees that the resulting pointer will remain stable
196     /// even when new allocations are pushed to the `HashMap`. `copy_repeatedly` relies
197     /// on that.
198     ///
199     /// It is the caller's responsibility to check bounds and alignment beforehand.
200     fn get_bytes_internal(
201         &self,
202         cx: &impl HasDataLayout,
203         range: AllocRange,
204         check_init_and_ptr: bool,
205     ) -> AllocResult<&[u8]> {
206         if check_init_and_ptr {
207             self.check_init(range)?;
208             self.check_relocations(cx, range)?;
209         } else {
210             // We still don't want relocations on the *edges*.
211             self.check_relocation_edges(cx, range)?;
212         }
213
214         Ok(&self.bytes[range.start.bytes_usize()..range.end().bytes_usize()])
215     }
216
217     /// Checks that these bytes are initialized and not pointer bytes, and then return them
218     /// as a slice.
219     ///
220     /// It is the caller's responsibility to check bounds and alignment beforehand.
221     /// Most likely, you want to use the `PlaceTy` and `OperandTy`-based methods
222     /// on `InterpCx` instead.
223     #[inline]
224     pub fn get_bytes(&self, cx: &impl HasDataLayout, range: AllocRange) -> AllocResult<&[u8]> {
225         self.get_bytes_internal(cx, range, true)
226     }
227
228     /// It is the caller's responsibility to handle uninitialized and pointer bytes.
229     /// However, this still checks that there are no relocations on the *edges*.
230     ///
231     /// It is the caller's responsibility to check bounds and alignment beforehand.
232     #[inline]
233     pub fn get_bytes_with_uninit_and_ptr(
234         &self,
235         cx: &impl HasDataLayout,
236         range: AllocRange,
237     ) -> AllocResult<&[u8]> {
238         self.get_bytes_internal(cx, range, false)
239     }
240
241     /// Just calling this already marks everything as defined and removes relocations,
242     /// so be sure to actually put data there!
243     ///
244     /// It is the caller's responsibility to check bounds and alignment beforehand.
245     /// Most likely, you want to use the `PlaceTy` and `OperandTy`-based methods
246     /// on `InterpCx` instead.
247     pub fn get_bytes_mut(&mut self, cx: &impl HasDataLayout, range: AllocRange) -> &mut [u8] {
248         self.mark_init(range, true);
249         self.clear_relocations(cx, range);
250
251         &mut self.bytes[range.start.bytes_usize()..range.end().bytes_usize()]
252     }
253
254     /// A raw pointer variant of `get_bytes_mut` that avoids invalidating existing aliases into this memory.
255     pub fn get_bytes_mut_ptr(&mut self, cx: &impl HasDataLayout, range: AllocRange) -> *mut [u8] {
256         self.mark_init(range, true);
257         self.clear_relocations(cx, range);
258
259         assert!(range.end().bytes_usize() <= self.bytes.len()); // need to do our own bounds-check
260         let begin_ptr = self.bytes.as_mut_ptr().wrapping_add(range.start.bytes_usize());
261         let len = range.end().bytes_usize() - range.start.bytes_usize();
262         ptr::slice_from_raw_parts_mut(begin_ptr, len)
263     }
264 }
265
266 /// Reading and writing.
267 impl<Tag: Copy, Extra> Allocation<Tag, Extra> {
268     /// Validates that `ptr.offset` and `ptr.offset + size` do not point to the middle of a
269     /// relocation. If `allow_uninit_and_ptr` is `false`, also enforces that the memory in the
270     /// given range contains neither relocations nor uninitialized bytes.
271     pub fn check_bytes(
272         &self,
273         cx: &impl HasDataLayout,
274         range: AllocRange,
275         allow_uninit_and_ptr: bool,
276     ) -> AllocResult {
277         // Check bounds and relocations on the edges.
278         self.get_bytes_with_uninit_and_ptr(cx, range)?;
279         // Check uninit and ptr.
280         if !allow_uninit_and_ptr {
281             self.check_init(range)?;
282             self.check_relocations(cx, range)?;
283         }
284         Ok(())
285     }
286
287     /// Reads a *non-ZST* scalar.
288     ///
289     /// ZSTs can't be read because in order to obtain a `Pointer`, we need to check
290     /// for ZSTness anyway due to integer pointers being valid for ZSTs.
291     ///
292     /// It is the caller's responsibility to check bounds and alignment beforehand.
293     /// Most likely, you want to call `InterpCx::read_scalar` instead of this method.
294     pub fn read_scalar(
295         &self,
296         cx: &impl HasDataLayout,
297         range: AllocRange,
298     ) -> AllocResult<ScalarMaybeUninit<Tag>> {
299         // `get_bytes_unchecked` tests relocation edges.
300         let bytes = self.get_bytes_with_uninit_and_ptr(cx, range)?;
301         // Uninit check happens *after* we established that the alignment is correct.
302         // We must not return `Ok()` for unaligned pointers!
303         if self.is_init(range).is_err() {
304             // This inflates uninitialized bytes to the entire scalar, even if only a few
305             // bytes are uninitialized.
306             return Ok(ScalarMaybeUninit::Uninit);
307         }
308         // Now we do the actual reading.
309         let bits = read_target_uint(cx.data_layout().endian, bytes).unwrap();
310         // See if we got a pointer.
311         if range.size != cx.data_layout().pointer_size {
312             // Not a pointer.
313             // *Now*, we better make sure that the inside is free of relocations too.
314             self.check_relocations(cx, range)?;
315         } else {
316             // Maybe a pointer.
317             if let Some(&(tag, alloc_id)) = self.relocations.get(&range.start) {
318                 let ptr = Pointer::new_with_tag(alloc_id, Size::from_bytes(bits), tag);
319                 return Ok(ScalarMaybeUninit::Scalar(ptr.into()));
320             }
321         }
322         // We don't. Just return the bits.
323         Ok(ScalarMaybeUninit::Scalar(Scalar::from_uint(bits, range.size)))
324     }
325
326     /// Writes a *non-ZST* scalar.
327     ///
328     /// ZSTs can't be read because in order to obtain a `Pointer`, we need to check
329     /// for ZSTness anyway due to integer pointers being valid for ZSTs.
330     ///
331     /// It is the caller's responsibility to check bounds and alignment beforehand.
332     /// Most likely, you want to call `InterpCx::write_scalar` instead of this method.
333     pub fn write_scalar(
334         &mut self,
335         cx: &impl HasDataLayout,
336         range: AllocRange,
337         val: ScalarMaybeUninit<Tag>,
338     ) -> AllocResult {
339         let val = match val {
340             ScalarMaybeUninit::Scalar(scalar) => scalar,
341             ScalarMaybeUninit::Uninit => {
342                 self.mark_init(range, false);
343                 return Ok(());
344             }
345         };
346
347         let bytes = match val.to_bits_or_ptr(range.size, cx) {
348             Err(val) => u128::from(val.offset.bytes()),
349             Ok(data) => data,
350         };
351
352         let endian = cx.data_layout().endian;
353         let dst = self.get_bytes_mut(cx, range);
354         write_target_uint(endian, dst, bytes).unwrap();
355
356         // See if we have to also write a relocation.
357         if let Scalar::Ptr(val) = val {
358             self.relocations.insert(range.start, (val.tag, val.alloc_id));
359         }
360
361         Ok(())
362     }
363 }
364
365 /// Relocations.
366 impl<Tag: Copy, Extra> Allocation<Tag, Extra> {
367     /// Returns all relocations overlapping with the given pointer-offset pair.
368     pub fn get_relocations(
369         &self,
370         cx: &impl HasDataLayout,
371         range: AllocRange,
372     ) -> &[(Size, (Tag, AllocId))] {
373         // We have to go back `pointer_size - 1` bytes, as that one would still overlap with
374         // the beginning of this range.
375         let start = range.start.bytes().saturating_sub(cx.data_layout().pointer_size.bytes() - 1);
376         self.relocations.range(Size::from_bytes(start)..range.end())
377     }
378
379     /// Checks that there are no relocations overlapping with the given range.
380     #[inline(always)]
381     fn check_relocations(&self, cx: &impl HasDataLayout, range: AllocRange) -> AllocResult {
382         if self.get_relocations(cx, range).is_empty() {
383             Ok(())
384         } else {
385             Err(AllocError::ReadPointerAsBytes)
386         }
387     }
388
389     /// Removes all relocations inside the given range.
390     /// If there are relocations overlapping with the edges, they
391     /// are removed as well *and* the bytes they cover are marked as
392     /// uninitialized. This is a somewhat odd "spooky action at a distance",
393     /// but it allows strictly more code to run than if we would just error
394     /// immediately in that case.
395     fn clear_relocations(&mut self, cx: &impl HasDataLayout, range: AllocRange) {
396         // Find the start and end of the given range and its outermost relocations.
397         let (first, last) = {
398             // Find all relocations overlapping the given range.
399             let relocations = self.get_relocations(cx, range);
400             if relocations.is_empty() {
401                 return;
402             }
403
404             (
405                 relocations.first().unwrap().0,
406                 relocations.last().unwrap().0 + cx.data_layout().pointer_size,
407             )
408         };
409         let start = range.start;
410         let end = range.end();
411
412         // Mark parts of the outermost relocations as uninitialized if they partially fall outside the
413         // given range.
414         if first < start {
415             self.init_mask.set_range(first, start, false);
416         }
417         if last > end {
418             self.init_mask.set_range(end, last, false);
419         }
420
421         // Forget all the relocations.
422         self.relocations.remove_range(first..last);
423     }
424
425     /// Errors if there are relocations overlapping with the edges of the
426     /// given memory range.
427     #[inline]
428     fn check_relocation_edges(&self, cx: &impl HasDataLayout, range: AllocRange) -> AllocResult {
429         self.check_relocations(cx, alloc_range(range.start, Size::ZERO))?;
430         self.check_relocations(cx, alloc_range(range.end(), Size::ZERO))?;
431         Ok(())
432     }
433 }
434
435 /// Uninitialized bytes.
436 impl<Tag: Copy, Extra> Allocation<Tag, Extra> {
437     /// Checks whether the given range  is entirely initialized.
438     ///
439     /// Returns `Ok(())` if it's initialized. Otherwise returns the range of byte
440     /// indexes of the first contiguous uninitialized access.
441     fn is_init(&self, range: AllocRange) -> Result<(), Range<Size>> {
442         self.init_mask.is_range_initialized(range.start, range.end()) // `Size` addition
443     }
444
445     /// Checks that a range of bytes is initialized. If not, returns the `InvalidUninitBytes`
446     /// error which will report the first range of bytes which is uninitialized.
447     fn check_init(&self, range: AllocRange) -> AllocResult {
448         self.is_init(range).or_else(|idx_range| {
449             Err(AllocError::InvalidUninitBytes(Some(UninitBytesAccess {
450                 access_offset: range.start,
451                 access_size: range.size,
452                 uninit_offset: idx_range.start,
453                 uninit_size: idx_range.end - idx_range.start, // `Size` subtraction
454             })))
455         })
456     }
457
458     pub fn mark_init(&mut self, range: AllocRange, is_init: bool) {
459         if range.size.bytes() == 0 {
460             return;
461         }
462         self.init_mask.set_range(range.start, range.end(), is_init);
463     }
464 }
465
466 /// Run-length encoding of the uninit mask.
467 /// Used to copy parts of a mask multiple times to another allocation.
468 pub struct InitMaskCompressed {
469     /// Whether the first range is initialized.
470     initial: bool,
471     /// The lengths of ranges that are run-length encoded.
472     /// The initialization state of the ranges alternate starting with `initial`.
473     ranges: smallvec::SmallVec<[u64; 1]>,
474 }
475
476 impl InitMaskCompressed {
477     pub fn no_bytes_init(&self) -> bool {
478         // The `ranges` are run-length encoded and of alternating initialization state.
479         // So if `ranges.len() > 1` then the second block is an initialized range.
480         !self.initial && self.ranges.len() == 1
481     }
482 }
483
484 /// Transferring the initialization mask to other allocations.
485 impl<Tag, Extra> Allocation<Tag, Extra> {
486     /// Creates a run-length encoding of the initialization mask.
487     pub fn compress_uninit_range(&self, src: Pointer<Tag>, size: Size) -> InitMaskCompressed {
488         // Since we are copying `size` bytes from `src` to `dest + i * size` (`for i in 0..repeat`),
489         // a naive initialization mask copying algorithm would repeatedly have to read the initialization mask from
490         // the source and write it to the destination. Even if we optimized the memory accesses,
491         // we'd be doing all of this `repeat` times.
492         // Therefore we precompute a compressed version of the initialization mask of the source value and
493         // then write it back `repeat` times without computing any more information from the source.
494
495         // A precomputed cache for ranges of initialized / uninitialized bits
496         // 0000010010001110 will become
497         // `[5, 1, 2, 1, 3, 3, 1]`,
498         // where each element toggles the state.
499
500         let mut ranges = smallvec::SmallVec::<[u64; 1]>::new();
501         let initial = self.init_mask.get(src.offset);
502         let mut cur_len = 1;
503         let mut cur = initial;
504
505         for i in 1..size.bytes() {
506             // FIXME: optimize to bitshift the current uninitialized block's bits and read the top bit.
507             if self.init_mask.get(src.offset + Size::from_bytes(i)) == cur {
508                 cur_len += 1;
509             } else {
510                 ranges.push(cur_len);
511                 cur_len = 1;
512                 cur = !cur;
513             }
514         }
515
516         ranges.push(cur_len);
517
518         InitMaskCompressed { ranges, initial }
519     }
520
521     /// Applies multiple instances of the run-length encoding to the initialization mask.
522     pub fn mark_compressed_init_range(
523         &mut self,
524         defined: &InitMaskCompressed,
525         dest: Pointer<Tag>,
526         size: Size,
527         repeat: u64,
528     ) {
529         // An optimization where we can just overwrite an entire range of initialization
530         // bits if they are going to be uniformly `1` or `0`.
531         if defined.ranges.len() <= 1 {
532             self.init_mask.set_range_inbounds(
533                 dest.offset,
534                 dest.offset + size * repeat, // `Size` operations
535                 defined.initial,
536             );
537             return;
538         }
539
540         for mut j in 0..repeat {
541             j *= size.bytes();
542             j += dest.offset.bytes();
543             let mut cur = defined.initial;
544             for range in &defined.ranges {
545                 let old_j = j;
546                 j += range;
547                 self.init_mask.set_range_inbounds(
548                     Size::from_bytes(old_j),
549                     Size::from_bytes(j),
550                     cur,
551                 );
552                 cur = !cur;
553             }
554         }
555     }
556 }
557
558 /// Relocations.
559 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, TyEncodable, TyDecodable)]
560 pub struct Relocations<Tag = (), Id = AllocId>(SortedMap<Size, (Tag, Id)>);
561
562 impl<Tag, Id> Relocations<Tag, Id> {
563     pub fn new() -> Self {
564         Relocations(SortedMap::new())
565     }
566
567     // The caller must guarantee that the given relocations are already sorted
568     // by address and contain no duplicates.
569     pub fn from_presorted(r: Vec<(Size, (Tag, Id))>) -> Self {
570         Relocations(SortedMap::from_presorted_elements(r))
571     }
572 }
573
574 impl<Tag> Deref for Relocations<Tag> {
575     type Target = SortedMap<Size, (Tag, AllocId)>;
576
577     fn deref(&self) -> &Self::Target {
578         &self.0
579     }
580 }
581
582 impl<Tag> DerefMut for Relocations<Tag> {
583     fn deref_mut(&mut self) -> &mut Self::Target {
584         &mut self.0
585     }
586 }
587
588 /// A partial, owned list of relocations to transfer into another allocation.
589 pub struct AllocationRelocations<Tag> {
590     relative_relocations: Vec<(Size, (Tag, AllocId))>,
591 }
592
593 impl<Tag: Copy, Extra> Allocation<Tag, Extra> {
594     pub fn prepare_relocation_copy(
595         &self,
596         cx: &impl HasDataLayout,
597         src: AllocRange,
598         dest: Size,
599         count: u64,
600     ) -> AllocationRelocations<Tag> {
601         let relocations = self.get_relocations(cx, src);
602         if relocations.is_empty() {
603             return AllocationRelocations { relative_relocations: Vec::new() };
604         }
605
606         let size = src.size;
607         let mut new_relocations = Vec::with_capacity(relocations.len() * (count as usize));
608
609         for i in 0..count {
610             new_relocations.extend(relocations.iter().map(|&(offset, reloc)| {
611                 // compute offset for current repetition
612                 let dest_offset = dest + size * i; // `Size` operations
613                 (
614                     // shift offsets from source allocation to destination allocation
615                     (offset + dest_offset) - src.start, // `Size` operations
616                     reloc,
617                 )
618             }));
619         }
620
621         AllocationRelocations { relative_relocations: new_relocations }
622     }
623
624     /// Applies a relocation copy.
625     /// The affected range, as defined in the parameters to `prepare_relocation_copy` is expected
626     /// to be clear of relocations.
627     pub fn mark_relocation_range(&mut self, relocations: AllocationRelocations<Tag>) {
628         self.relocations.insert_presorted(relocations.relative_relocations);
629     }
630 }
631
632 ////////////////////////////////////////////////////////////////////////////////
633 // Uninitialized byte tracking
634 ////////////////////////////////////////////////////////////////////////////////
635
636 type Block = u64;
637
638 /// A bitmask where each bit refers to the byte with the same index. If the bit is `true`, the byte
639 /// is initialized. If it is `false` the byte is uninitialized.
640 #[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, TyEncodable, TyDecodable)]
641 #[derive(HashStable)]
642 pub struct InitMask {
643     blocks: Vec<Block>,
644     len: Size,
645 }
646
647 impl InitMask {
648     pub const BLOCK_SIZE: u64 = 64;
649
650     pub fn new(size: Size, state: bool) -> Self {
651         let mut m = InitMask { blocks: vec![], len: Size::ZERO };
652         m.grow(size, state);
653         m
654     }
655
656     /// Checks whether the range `start..end` (end-exclusive) is entirely initialized.
657     ///
658     /// Returns `Ok(())` if it's initialized. Otherwise returns a range of byte
659     /// indexes for the first contiguous span of the uninitialized access.
660     #[inline]
661     pub fn is_range_initialized(&self, start: Size, end: Size) -> Result<(), Range<Size>> {
662         if end > self.len {
663             return Err(self.len..end);
664         }
665
666         // FIXME(oli-obk): optimize this for allocations larger than a block.
667         let idx = (start.bytes()..end.bytes()).map(Size::from_bytes).find(|&i| !self.get(i));
668
669         match idx {
670             Some(idx) => {
671                 let uninit_end = (idx.bytes()..end.bytes())
672                     .map(Size::from_bytes)
673                     .find(|&i| self.get(i))
674                     .unwrap_or(end);
675                 Err(idx..uninit_end)
676             }
677             None => Ok(()),
678         }
679     }
680
681     pub fn set_range(&mut self, start: Size, end: Size, new_state: bool) {
682         let len = self.len;
683         if end > len {
684             self.grow(end - len, new_state);
685         }
686         self.set_range_inbounds(start, end, new_state);
687     }
688
689     pub fn set_range_inbounds(&mut self, start: Size, end: Size, new_state: bool) {
690         let (blocka, bita) = bit_index(start);
691         let (blockb, bitb) = bit_index(end);
692         if blocka == blockb {
693             // First set all bits except the first `bita`,
694             // then unset the last `64 - bitb` bits.
695             let range = if bitb == 0 {
696                 u64::MAX << bita
697             } else {
698                 (u64::MAX << bita) & (u64::MAX >> (64 - bitb))
699             };
700             if new_state {
701                 self.blocks[blocka] |= range;
702             } else {
703                 self.blocks[blocka] &= !range;
704             }
705             return;
706         }
707         // across block boundaries
708         if new_state {
709             // Set `bita..64` to `1`.
710             self.blocks[blocka] |= u64::MAX << bita;
711             // Set `0..bitb` to `1`.
712             if bitb != 0 {
713                 self.blocks[blockb] |= u64::MAX >> (64 - bitb);
714             }
715             // Fill in all the other blocks (much faster than one bit at a time).
716             for block in (blocka + 1)..blockb {
717                 self.blocks[block] = u64::MAX;
718             }
719         } else {
720             // Set `bita..64` to `0`.
721             self.blocks[blocka] &= !(u64::MAX << bita);
722             // Set `0..bitb` to `0`.
723             if bitb != 0 {
724                 self.blocks[blockb] &= !(u64::MAX >> (64 - bitb));
725             }
726             // Fill in all the other blocks (much faster than one bit at a time).
727             for block in (blocka + 1)..blockb {
728                 self.blocks[block] = 0;
729             }
730         }
731     }
732
733     #[inline]
734     pub fn get(&self, i: Size) -> bool {
735         let (block, bit) = bit_index(i);
736         (self.blocks[block] & (1 << bit)) != 0
737     }
738
739     #[inline]
740     pub fn set(&mut self, i: Size, new_state: bool) {
741         let (block, bit) = bit_index(i);
742         self.set_bit(block, bit, new_state);
743     }
744
745     #[inline]
746     fn set_bit(&mut self, block: usize, bit: usize, new_state: bool) {
747         if new_state {
748             self.blocks[block] |= 1 << bit;
749         } else {
750             self.blocks[block] &= !(1 << bit);
751         }
752     }
753
754     pub fn grow(&mut self, amount: Size, new_state: bool) {
755         if amount.bytes() == 0 {
756             return;
757         }
758         let unused_trailing_bits =
759             u64::try_from(self.blocks.len()).unwrap() * Self::BLOCK_SIZE - self.len.bytes();
760         if amount.bytes() > unused_trailing_bits {
761             let additional_blocks = amount.bytes() / Self::BLOCK_SIZE + 1;
762             self.blocks.extend(
763                 // FIXME(oli-obk): optimize this by repeating `new_state as Block`.
764                 iter::repeat(0).take(usize::try_from(additional_blocks).unwrap()),
765             );
766         }
767         let start = self.len;
768         self.len += amount;
769         self.set_range_inbounds(start, start + amount, new_state); // `Size` operation
770     }
771 }
772
773 #[inline]
774 fn bit_index(bits: Size) -> (usize, usize) {
775     let bits = bits.bytes();
776     let a = bits / InitMask::BLOCK_SIZE;
777     let b = bits % InitMask::BLOCK_SIZE;
778     (usize::try_from(a).unwrap(), usize::try_from(b).unwrap())
779 }