scrypt_opt/
memory.rs

1use core::ops::{Deref, DerefMut};
2
3#[repr(align(64))]
4#[derive(Default, Debug, PartialEq, Eq, Clone, Copy)]
5/// Align to 64 bytes
6pub struct Align64<T>(pub T);
7
8impl<T> Align64<T> {
9    /// Perform a reference cast to an [`Align64<Block<R>>`] without checking the alignment.
10    ///
11    /// # Safety
12    /// This function is unsafe because it does not check the alignment of the input pointer.
13    #[inline(always)]
14    pub const unsafe fn new_unchecked(input: &T) -> &Self {
15        unsafe { &*(input as *const T as *const Self) }
16    }
17
18    /// Perform a reference cast to an [`Align64<Block<R>>`].
19    ///
20    /// # Panics
21    ///
22    /// Panics if the input pointer is not aligned to 64 bytes.
23    pub fn new(input: &T) -> &Self {
24        let ptr = input as *const T;
25        assert_eq!(
26            ptr.align_offset(64),
27            0,
28            "Input pointer is not aligned to 64 bytes"
29        );
30        unsafe { Self::new_unchecked(input) }
31    }
32
33    /// Perform a mutable reference cast to an [`Align64<Block<R>>`] without checking the alignment.
34    ///
35    /// # Safety
36    /// This function is unsafe because it does not check the alignment of the input pointer.
37    #[inline(always)]
38    pub const unsafe fn new_mut_unchecked(input: &mut T) -> &mut Self {
39        unsafe { &mut *(input as *mut T as *mut Self) }
40    }
41
42    /// Perform a mutable reference cast to an [`Align64<Block<R>>`].
43    ///
44    /// # Panics
45    ///
46    /// Panics if the input pointer is not aligned to 64 bytes.
47    pub fn new_mut(input: &mut T) -> &mut Self {
48        let ptr = input as *mut T;
49        assert_eq!(
50            ptr.align_offset(64),
51            0,
52            "Input pointer is not aligned to 64 bytes"
53        );
54        unsafe { Self::new_mut_unchecked(input) }
55    }
56}
57
58#[repr(align(32))]
59#[derive(Default, Debug, PartialEq, Eq, Clone, Copy)]
60/// Align to 32 bytes
61pub struct Align32<T>(pub T);
62
63impl<T> Align32<T> {
64    /// Perform a reference cast to an [`Align32<Block<R>>`] without checking the alignment.
65    ///
66    /// # Safety
67    /// This function is unsafe because it does not check the alignment of the input pointer.
68    #[inline(always)]
69    pub const unsafe fn new_unchecked(input: &T) -> &Self {
70        unsafe { &*(input as *const T as *const Self) }
71    }
72
73    /// Perform a reference cast to an [`Align32<Block<R>>`].
74    ///
75    /// # Panics
76    ///
77    /// Panics if the input pointer is not aligned to 32 bytes.
78    pub fn new(input: &T) -> &Self {
79        let ptr = input as *const T;
80        assert_eq!(
81            ptr.align_offset(32),
82            0,
83            "Input pointer is not aligned to 32 bytes"
84        );
85        unsafe { Self::new_unchecked(input) }
86    }
87
88    /// Perform a mutable reference cast to an [`Align32<Block<R>>`] without checking the alignment.
89    ///
90    /// # Safety
91    /// This function is unsafe because it does not check the alignment of the input pointer.
92    #[inline(always)]
93    pub const unsafe fn new_mut_unchecked(input: &mut T) -> &mut Self {
94        unsafe { &mut *(input as *mut T as *mut Self) }
95    }
96
97    /// Perform a mutable reference cast to an [`Align32<Block<R>>`].
98    ///
99    /// # Panics
100    ///
101    /// Panics if the input pointer is not aligned to 32 bytes.
102    pub fn new_mut(input: &mut T) -> &mut Self {
103        let ptr = input as *mut T;
104        assert_eq!(
105            ptr.align_offset(32),
106            0,
107            "Input pointer is not aligned to 32 bytes"
108        );
109        unsafe { Self::new_mut_unchecked(input) }
110    }
111}
112
113impl<T> AsRef<T> for Align64<T> {
114    fn as_ref(&self) -> &T {
115        &self.0
116    }
117}
118
119impl<T> AsRef<Align64<T>> for Align64<T> {
120    fn as_ref(&self) -> &Align64<T> {
121        self
122    }
123}
124
125impl<T> AsMut<Align64<T>> for Align64<T> {
126    fn as_mut(&mut self) -> &mut Align64<T> {
127        self
128    }
129}
130
131impl<T> AsMut<T> for Align64<T> {
132    fn as_mut(&mut self) -> &mut T {
133        &mut self.0
134    }
135}
136
137impl<T> Deref for Align64<T> {
138    type Target = T;
139    fn deref(&self) -> &Self::Target {
140        &self.0
141    }
142}
143
144impl<T> DerefMut for Align64<T> {
145    fn deref_mut(&mut self) -> &mut Self::Target {
146        &mut self.0
147    }
148}
149
150impl<T> AsRef<T> for Align32<T> {
151    fn as_ref(&self) -> &T {
152        &self.0
153    }
154}
155
156impl<T> AsRef<Align32<T>> for Align32<T> {
157    fn as_ref(&self) -> &Align32<T> {
158        self
159    }
160}
161
162impl<T> AsMut<Align32<T>> for Align32<T> {
163    fn as_mut(&mut self) -> &mut Align32<T> {
164        self
165    }
166}
167
168impl<T> AsMut<T> for Align32<T> {
169    fn as_mut(&mut self) -> &mut T {
170        &mut self.0
171    }
172}
173
174impl<T> Deref for Align32<T> {
175    type Target = T;
176    fn deref(&self) -> &Self::Target {
177        &self.0
178    }
179}
180
181impl<T> DerefMut for Align32<T> {
182    fn deref_mut(&mut self) -> &mut Self::Target {
183        &mut self.0
184    }
185}
186
187/// A box for a hugepage-backed buffer
188#[cfg(feature = "huge-page")]
189pub struct HugeSlice<T> {
190    ptr: *mut T,
191    len: usize,
192    #[cfg(any(target_os = "android", target_os = "linux"))]
193    capacity: usize,
194    #[cfg(any(target_os = "android", target_os = "linux"))]
195    file: std::fs::File,
196}
197
198#[cfg(feature = "huge-page")]
199unsafe impl<T> Send for HugeSlice<T> where T: Send {}
200
201#[cfg(feature = "huge-page")]
202unsafe impl<T> Sync for HugeSlice<T> where T: Sync {}
203
204#[cfg(feature = "huge-page")]
205impl<T> core::convert::AsMut<T> for HugeSlice<T> {
206    fn as_mut(&mut self) -> &mut T {
207        unsafe { &mut *self.ptr }
208    }
209}
210
211#[cfg(feature = "huge-page")]
212impl<T> HugeSlice<T> {
213    /// Create a new hugepage-backed buffer
214    #[cfg(target_os = "windows")]
215    pub fn new(len: usize) -> Result<Self, std::io::Error> {
216        use windows_sys::Win32::{
217            Security::TOKEN_ADJUST_PRIVILEGES,
218            System::{
219                Memory::{
220                    GetLargePageMinimum, MEM_COMMIT, MEM_LARGE_PAGES, MEM_RESERVE, PAGE_READWRITE,
221                    VirtualAlloc,
222                },
223                Threading::{GetCurrentProcess, OpenProcessToken},
224            },
225        };
226        unsafe {
227            use windows_sys::Win32::Security::{
228                AdjustTokenPrivileges, LUID_AND_ATTRIBUTES, LookupPrivilegeValueA,
229                SE_PRIVILEGE_ENABLED, TOKEN_PRIVILEGES,
230            };
231
232            let large_page_minimum = GetLargePageMinimum();
233            if large_page_minimum == 0 {
234                return Err(std::io::Error::new(
235                    std::io::ErrorKind::Other,
236                    "Large page not supported by the system",
237                ));
238            }
239
240            if core::mem::align_of::<T>() > large_page_minimum {
241                return Err(std::io::Error::new(
242                    std::io::ErrorKind::Other,
243                    format!(
244                        "Alignment of the type is greater than the large page minimum: {} requires {} alignment, large page minimum is {}",
245                        core::any::type_name::<T>(),
246                        core::mem::align_of::<T>(),
247                        large_page_minimum
248                    ),
249                ));
250            }
251
252            let min_alloc_size = core::mem::size_of::<T>()
253                .checked_mul(len)
254                .and_then(|x| x.checked_next_multiple_of(large_page_minimum))
255                .and_then(|x| x.checked_next_multiple_of(8))
256                .ok_or_else(|| {
257                    std::io::Error::new(
258                        std::io::ErrorKind::Other,
259                        format!(
260                            "Allocation size overflow (requested: {}, max: {})",
261                            core::mem::size_of::<T>() as u128 * len as u128,
262                            usize::MAX
263                        ),
264                    )
265                })?;
266
267            if min_alloc_size == 0 {
268                return Err(std::io::Error::new(
269                    std::io::ErrorKind::Other,
270                    "Allocation size must be non-zero",
271                ));
272            }
273
274            let mut token_handle = Default::default();
275            if OpenProcessToken(
276                GetCurrentProcess(),
277                TOKEN_ADJUST_PRIVILEGES,
278                &mut token_handle,
279            ) == 0
280            {
281                return Err(std::io::Error::last_os_error());
282            }
283
284            let mut privs = [TOKEN_PRIVILEGES {
285                PrivilegeCount: 1,
286                Privileges: [LUID_AND_ATTRIBUTES {
287                    Luid: Default::default(),
288                    Attributes: SE_PRIVILEGE_ENABLED,
289                }],
290            }];
291
292            if LookupPrivilegeValueA(
293                core::ptr::null_mut(),
294                c"SeLockMemoryPrivilege".as_ptr().cast(),
295                &mut privs[0].Privileges[0].Luid,
296            ) == 0
297            {
298                return Err(std::io::Error::last_os_error());
299            }
300
301            if AdjustTokenPrivileges(
302                token_handle,
303                0,
304                privs.as_ptr(),
305                0,
306                core::ptr::null_mut(),
307                core::ptr::null_mut(),
308            ) == 0
309            {
310                return Err(std::io::Error::new(
311                    std::io::ErrorKind::Other,
312                    "Failed to adjust token privileges",
313                ));
314            }
315
316            let ptr = VirtualAlloc(
317                core::ptr::null_mut(),
318                min_alloc_size,
319                MEM_COMMIT | MEM_RESERVE | MEM_LARGE_PAGES,
320                PAGE_READWRITE,
321            );
322
323            if ptr == core::ptr::null_mut() {
324                return Err(std::io::Error::last_os_error());
325            }
326
327            Ok(Self {
328                ptr: ptr.cast::<T>(),
329                len,
330            })
331        }
332    }
333
334    /// Create a new hugepage-backed buffer
335    #[cfg(any(target_os = "android", target_os = "linux"))]
336    pub fn new(len: usize) -> Result<Self, std::io::Error> {
337        Self::new_unix(
338            len,
339            #[cfg(all(target_os = "linux", feature = "std"))]
340            None,
341        )
342    }
343
344    /// Create a new hugepage-backed buffer backed by a file
345    #[cfg(any(target_os = "android", target_os = "linux"))]
346    pub fn new_unix(
347        len: usize,
348        #[cfg(target_os = "linux")] file: Option<std::fs::File>,
349    ) -> Result<Self, std::io::Error> {
350        unsafe {
351            let pagesz = libc::sysconf(libc::_SC_PAGESIZE);
352            if pagesz == -1 {
353                return Err(std::io::Error::last_os_error());
354            }
355
356            let page_size = pagesz as usize;
357            if core::mem::align_of::<T>() > page_size {
358                return Err(std::io::Error::new(
359                    std::io::ErrorKind::Other,
360                    format!(
361                        "Alignment of the type is greater than the page size: {} requires {} alignment, page size is {}",
362                        core::any::type_name::<T>(),
363                        core::mem::align_of::<T>(),
364                        page_size
365                    ),
366                ));
367            }
368
369            let alloc_min_len = core::mem::size_of::<T>()
370                .checked_mul(len)
371                .and_then(|x| x.checked_next_multiple_of(page_size))
372                .and_then(|x| x.checked_next_multiple_of(8))
373                .ok_or(std::io::Error::new(
374                    std::io::ErrorKind::Other,
375                    format!(
376                        "Allocation length overflow (requested: {}, max: {})",
377                        core::mem::size_of::<T>() as u128 * len as u128,
378                        usize::MAX
379                    ),
380                ))?;
381
382            if alloc_min_len == 0 {
383                return Err(std::io::Error::new(
384                    std::io::ErrorKind::Other,
385                    "Allocation length must be non-zero",
386                ));
387            }
388
389            #[cfg(target_os = "linux")]
390            if let Some(file) = file {
391                let ptr = libc::mmap64(
392                    core::ptr::null_mut(),
393                    alloc_min_len,
394                    libc::PROT_READ | libc::PROT_WRITE,
395                    libc::MAP_PRIVATE,
396                    std::os::unix::io::AsRawFd::as_raw_fd(&file),
397                    0,
398                );
399                if ptr != libc::MAP_FAILED {
400                    return Ok(HugeSlice {
401                        ptr: ptr.cast::<T>(),
402                        len,
403                        capacity: alloc_min_len,
404                        file,
405                    });
406                }
407
408                return Err(std::io::Error::last_os_error());
409            }
410
411            for (try_page_size, try_flags) in [
412                #[cfg(target_os = "linux")]
413                ((1 << 30), libc::MFD_HUGE_1GB),
414                #[cfg(target_os = "linux")]
415                ((256 << 20), libc::MFD_HUGE_256MB),
416                #[cfg(target_os = "linux")]
417                ((32 << 20), libc::MFD_HUGE_32MB),
418                #[cfg(target_os = "linux")]
419                ((16 << 20), libc::MFD_HUGE_16MB),
420                #[cfg(target_os = "linux")]
421                ((8 << 20), libc::MFD_HUGE_8MB),
422                #[cfg(target_os = "linux")]
423                ((2 << 20), libc::MFD_HUGE_2MB),
424                ((page_size), 0),
425            ]
426            .into_iter()
427            .filter(|(try_page_size, _flags)| *try_page_size >= page_size)
428            // don't grossly over size by capping page size at 2x amount of memory needed
429            .filter(|(try_page_size, flags)| *flags == 0 || alloc_min_len * 2 > *try_page_size)
430            {
431                let fd = libc::memfd_create(
432                    c"scrypt-opt-huge-page-file".as_ptr().cast(),
433                    libc::MFD_CLOEXEC | libc::MFD_HUGETLB | try_flags,
434                );
435
436                if fd == -1 {
437                    continue;
438                }
439
440                let try_file = std::os::unix::io::FromRawFd::from_raw_fd(fd);
441
442                let try_size = alloc_min_len
443                    .checked_next_multiple_of(try_page_size)
444                    .ok_or(std::io::Error::new(
445                        std::io::ErrorKind::Other,
446                        format!(
447                            "Allocation length overflow (requested: {}, max: {})",
448                            alloc_min_len,
449                            usize::MAX
450                        ),
451                    ))?;
452
453                let ptr = libc::mmap64(
454                    core::ptr::null_mut(),
455                    try_size,
456                    libc::PROT_READ | libc::PROT_WRITE,
457                    libc::MAP_PRIVATE | libc::MAP_POPULATE,
458                    fd,
459                    0,
460                );
461                if ptr != libc::MAP_FAILED {
462                    return Ok(HugeSlice {
463                        ptr: ptr.cast::<T>(),
464                        len,
465                        capacity: try_size,
466                        file: try_file,
467                    });
468                }
469            }
470
471            Err(std::io::Error::last_os_error())
472        }
473    }
474
475    #[cfg(all(
476        not(target_os = "windows"),
477        not(target_os = "android"),
478        not(target_os = "linux")
479    ))]
480    /// Create a new hugepage-backed buffer
481    pub fn new(_len: usize) -> Result<Self, std::io::Error> {
482        Err(std::io::Error::new(
483            std::io::ErrorKind::Other,
484            "Huge page not supported on this platform",
485        ))
486    }
487}
488
489#[cfg(feature = "huge-page")]
490impl<T> HugeSlice<core::mem::MaybeUninit<T>> {
491    /// Assume the buffer is initialized
492    pub unsafe fn assume_init(self) -> HugeSlice<T> {
493        let forgotten = core::mem::ManuallyDrop::new(self);
494        HugeSlice {
495            ptr: forgotten.ptr.cast::<T>(),
496            len: forgotten.len,
497            #[cfg(any(target_os = "android", target_os = "linux"))]
498            capacity: forgotten.capacity,
499            #[cfg(any(target_os = "android", target_os = "linux"))]
500            file: unsafe {
501                std::os::unix::io::FromRawFd::from_raw_fd(std::os::unix::io::AsRawFd::as_raw_fd(
502                    &forgotten.file,
503                ))
504            },
505        }
506    }
507}
508
509#[cfg(feature = "huge-page")]
510impl<T> core::convert::AsRef<[T]> for HugeSlice<T> {
511    fn as_ref(&self) -> &[T] {
512        unsafe { core::slice::from_raw_parts(self.ptr, self.len) }
513    }
514}
515
516#[cfg(feature = "huge-page")]
517impl<T> core::convert::AsMut<[T]> for HugeSlice<T> {
518    fn as_mut(&mut self) -> &mut [T] {
519        unsafe { core::slice::from_raw_parts_mut(self.ptr, self.len) }
520    }
521}
522
523#[cfg(feature = "huge-page")]
524impl<T> core::ops::Deref for HugeSlice<T> {
525    type Target = [T];
526    fn deref(&self) -> &Self::Target {
527        unsafe { core::slice::from_raw_parts(self.ptr, self.len) }
528    }
529}
530
531#[cfg(feature = "huge-page")]
532impl<T> core::ops::DerefMut for HugeSlice<T> {
533    fn deref_mut(&mut self) -> &mut Self::Target {
534        unsafe { core::slice::from_raw_parts_mut(self.ptr, self.len) }
535    }
536}
537
538#[cfg(feature = "huge-page")]
539impl<T> Drop for HugeSlice<T> {
540    #[cfg(target_os = "windows")]
541    fn drop(&mut self) {
542        use windows_sys::Win32::System::Memory::{MEM_RELEASE, VirtualFree};
543        unsafe {
544            debug_assert!(
545                VirtualFree(self.ptr as *mut _, 0, MEM_RELEASE) != 0,
546                "Failed to free huge page"
547            );
548        }
549    }
550
551    #[cfg(any(target_os = "android", target_os = "linux"))]
552    fn drop(&mut self) {
553        unsafe {
554            libc::munmap(self.ptr as *mut libc::c_void, self.capacity);
555        }
556    }
557
558    #[cfg(all(
559        not(target_os = "windows"),
560        not(target_os = "android"),
561        not(target_os = "linux")
562    ))]
563    fn drop(&mut self) {
564        // Do nothing
565    }
566}
567
568#[cfg(feature = "alloc")]
569/// A box for a buffer that can be backed by a huge page or a normal box
570pub enum MaybeHugeSlice<T> {
571    #[cfg(feature = "huge-page")]
572    /// A huge page-backed buffer
573    Huge(HugeSlice<T>),
574    /// A normal buffer backed by heap
575    Normal(alloc::boxed::Box<[T]>),
576}
577
578impl<T> core::ops::Deref for MaybeHugeSlice<T> {
579    type Target = [T];
580    fn deref(&self) -> &Self::Target {
581        self.as_ref()
582    }
583}
584
585impl<T> core::ops::DerefMut for MaybeHugeSlice<T> {
586    fn deref_mut(&mut self) -> &mut Self::Target {
587        self.as_mut()
588    }
589}
590
591impl<T> AsRef<[T]> for MaybeHugeSlice<T> {
592    fn as_ref(&self) -> &[T] {
593        match self {
594            #[cfg(feature = "huge-page")]
595            MaybeHugeSlice::Huge(b) => b.as_ref(),
596            MaybeHugeSlice::Normal(b) => b,
597        }
598    }
599}
600
601impl<T> core::convert::AsMut<[T]> for MaybeHugeSlice<T> {
602    fn as_mut(&mut self) -> &mut [T] {
603        match self {
604            #[cfg(feature = "huge-page")]
605            MaybeHugeSlice::Huge(b) => b.as_mut(),
606            MaybeHugeSlice::Normal(b) => b.as_mut(),
607        }
608    }
609}
610
611impl<T> MaybeHugeSlice<T> {
612    /// Check if the buffer is backed by a huge page
613    pub fn is_huge_page(&self) -> bool {
614        match self {
615            #[cfg(feature = "huge-page")]
616            MaybeHugeSlice::Huge(_) => true,
617            MaybeHugeSlice::Normal(_) => false,
618        }
619    }
620
621    /// Create a new huge page-backed buffer
622    #[cfg(feature = "huge-page")]
623    pub fn new_huge_slice_zeroed(len: usize) -> Result<Self, std::io::Error> {
624        let b: HugeSlice<core::mem::MaybeUninit<T>> = HugeSlice::new(len)?;
625        unsafe {
626            core::ptr::write_bytes(b.ptr.cast::<T>(), 0, len);
627            Ok(MaybeHugeSlice::Huge(b.assume_init()))
628        }
629    }
630
631    /// Create a new normal buffer
632    #[cfg(feature = "alloc")]
633    pub fn new_slice_zeroed(len: usize) -> Self {
634        let mut b = alloc::vec::Vec::<T>::with_capacity(len);
635        unsafe {
636            b.set_len(len);
637            MaybeHugeSlice::Normal(b.into())
638        }
639    }
640
641    /// Create a new buffer
642    #[cfg(feature = "alloc")]
643    pub fn new_maybe(len: usize) -> Self {
644        #[cfg(feature = "huge-page")]
645        {
646            match Self::new_huge_slice_zeroed(len) {
647                Ok(huge) => huge,
648                Err(_) => Self::new_slice_zeroed(len),
649            }
650        }
651
652        #[cfg(not(feature = "huge-page"))]
653        Self::new_slice_zeroed(len)
654    }
655
656    /// Create a new huge page-backed buffer backed by a file
657    #[cfg(all(feature = "huge-page", target_os = "linux", feature = "std"))]
658    pub fn new_in(len: usize, file: std::fs::File) -> Result<Self, std::io::Error> {
659        let huge = HugeSlice::new_unix(len, Some(file))?;
660        Ok(Self::Huge(huge))
661    }
662
663    /// Create a new buffer
664    #[cfg(feature = "std")]
665    pub fn new(len: usize) -> (Self, Option<std::io::Error>) {
666        #[cfg(feature = "huge-page")]
667        {
668            match Self::new_huge_slice_zeroed(len) {
669                Ok(huge) => (huge, None),
670                Err(e) => (Self::new_slice_zeroed(len), Some(e.into())),
671            }
672        }
673
674        #[cfg(not(feature = "huge-page"))]
675        (Self::new_slice_zeroed(len), None)
676    }
677}
scrypt_opt/memory.rs

scrypt_opt/
memory.rs
