use crate::common::{

borrow::{BorrowedValue, BorrowedValueMut},

lock::{MapImmutable, PyMutex, PyMutexGuard},

};

use crate::vm::{

AsObject, FromArgs, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,

TryFromBorrowedObject, VirtualMachine, atomic_func,

builtins::{PyBytes, PyBytesRef, PyInt, PyIntRef, PyType, PyTypeRef},

byte::{bytes_from_object, value_from_object},

convert::ToPyException,

function::{ArgBytesLike, FuncArgs, OptionalArg},

protocol::{

BufferDescriptor, BufferMethods, PyBuffer, PyMappingMethods, PySequenceMethods,

},

sliceable::{SaturatedSlice, SequenceIndex, SequenceIndexOp},

types::{AsBuffer, AsMapping, AsSequence, Constructor, Representable},

};

use core::ops::{Deref, DerefMut};

use crossbeam_utils::atomic::AtomicCell;

use num_traits::Signed;

#[cfg(windows)]

use std::io;

use std::io::Write;

#[cfg(unix)]

use rustpython_host_env::crt_fd;

#[cfg(any(unix, windows))]

use rustpython_host_env::mmap as host_mmap;

#[cfg(windows)]

use rustpython_host_env::nt as host_nt;

#[repr(C)]

#[derive(PartialEq, Eq, Debug)]

enum AccessMode {

Default = 0,

Read = 1,

Write = 2,

Copy = 3,

}

impl<'a> TryFromBorrowedObject<'a> for AccessMode {

fn try_from_borrowed_object(vm: &VirtualMachine, obj: &'a PyObject) -> PyResult<Self> {

let i = u32::try_from_borrowed_object(vm, obj)?;

Ok(match i {

0 => Self::Default,

1 => Self::Read,

2 => Self::Write,

3 => Self::Copy,

_ => return Err(vm.new_value_error("Not a valid AccessMode value")),

})

}

}

#[cfg(unix)]

#[pyattr]

use libc::{

MADV_DONTNEED, MADV_NORMAL, MADV_RANDOM, MADV_SEQUENTIAL, MADV_WILLNEED, MAP_ANON,

MAP_ANONYMOUS, MAP_PRIVATE, MAP_SHARED, PROT_EXEC, PROT_READ, PROT_WRITE,

};

#[cfg(target_os = "macos")]

#[pyattr]

use libc::{MADV_FREE_REUSABLE, MADV_FREE_REUSE};

#[cfg(any(

#[pyattr]

use libc::MADV_FREE;

#[cfg(target_os = "linux")]

#[pyattr]

use libc::{

MADV_DODUMP, MADV_DOFORK, MADV_DONTDUMP, MADV_DONTFORK, MADV_HUGEPAGE, MADV_HWPOISON,

MADV_MERGEABLE, MADV_NOHUGEPAGE, MADV_REMOVE, MADV_UNMERGEABLE,

};

#[cfg(any(

#[pyattr]

use libc::MADV_SOFT_OFFLINE;

#[cfg(all(target_os = "linux", target_arch = "x86_64", target_env = "gnu"))]

#[pyattr]

use libc::{MAP_DENYWRITE, MAP_EXECUTABLE, MAP_POPULATE};

// MAP_STACK is available on Linux, OpenBSD, and NetBSD

#[cfg(any(target_os = "linux", target_os = "openbsd", target_os = "netbsd"))]

#[pyattr]

use libc::MAP_STACK;

// FreeBSD-specific MADV constants

#[cfg(target_os = "freebsd")]

#[pyattr]

use libc::{MADV_AUTOSYNC, MADV_CORE, MADV_NOCORE, MADV_NOSYNC, MADV_PROTECT};

#[pyattr]

const ACCESS_DEFAULT: u32 = AccessMode::Default as u32;

#[pyattr]

const ACCESS_READ: u32 = AccessMode::Read as u32;

#[pyattr]

const ACCESS_WRITE: u32 = AccessMode::Write as u32;

#[pyattr]

const ACCESS_COPY: u32 = AccessMode::Copy as u32;

#[pyattr(name = "PAGESIZE", once)]

fn page_size(_vm: &VirtualMachine) -> usize {

rustpython_host_env::os::page_size()

}

#[cfg(not(target_arch = "wasm32"))]

#[pyattr(name = "ALLOCATIONGRANULARITY", once)]

fn granularity(_vm: &VirtualMachine) -> usize {

rustpython_host_env::os::alloc_granularity()

}

#[pyattr(name = "error", once)]

fn error_type(vm: &VirtualMachine) -> PyTypeRef {

vm.ctx.exceptions.os_error.to_owned()

}

#[derive(Debug)]

enum MmapObj {

Mapped(host_mmap::MappedFile),

#[cfg(windows)]

Named(host_mmap::NamedMmap),

}

impl MmapObj {

fn as_slice(&self) -> &[u8] {

match self {

Self::Mapped(mmap) => mmap.as_slice(),

#[cfg(windows)]

Self::Named(named) => named.as_slice(),

}

}

}

#[pyattr]

#[pyclass(name = "mmap")]

#[derive(Debug, PyPayload)]

struct PyMmap {

closed: AtomicCell<bool>,

mmap: PyMutex<Option<MmapObj>>,

#[cfg(unix)]

fd: AtomicCell<i32>,

#[cfg(windows)]

handle: AtomicCell<isize>, // host_mmap::Handle is isize on Windows

offset: i64,

size: AtomicCell<usize>,

pos: AtomicCell<usize>, // relative to offset

exports: AtomicCell<usize>,

access: AccessMode,

}

impl PyMmap {

/// Close the underlying file handle/descriptor if open

fn close_handle(&self) {

#[cfg(unix)]

{

let fd = self.fd.swap(-1);

host_mmap::close_descriptor(fd);

}

#[cfg(windows)]

{

let handle = self.handle.swap(host_mmap::INVALID_HANDLE as isize);

if handle != host_mmap::INVALID_HANDLE as isize {

host_mmap::close_handle(handle as host_mmap::Handle);

}

}

}

}

impl Drop for PyMmap {

fn drop(&mut self) {

self.close_handle();

}

}

#[cfg(unix)]

#[derive(FromArgs)]

struct MmapNewArgs {

#[pyarg(any)]

fileno: i32,

#[pyarg(any)]

length: isize,

#[pyarg(any, default = libc::MAP_SHARED)]

flags: libc::c_int,

#[pyarg(any, default = libc::PROT_WRITE | libc::PROT_READ)]

prot: libc::c_int,

#[pyarg(any, default = AccessMode::Default)]

access: AccessMode,

#[pyarg(any, default = 0)]

offset: i64,

}

#[cfg(windows)]

#[derive(FromArgs)]

struct MmapNewArgs {

#[pyarg(any)]

fileno: i32,

#[pyarg(any)]

length: isize,

#[pyarg(any, default)]

tagname: Option<PyObjectRef>,

#[pyarg(any, default = AccessMode::Default)]

access: AccessMode,

#[pyarg(any, default = 0)]

offset: i64,

}

impl MmapNewArgs {

/// Validate mmap constructor arguments

fn validate_new_args(&self, vm: &VirtualMachine) -> PyResult<usize> {

if self.length < 0 {

return Err(vm.new_overflow_error("memory mapped length must be positive"));

}

if self.offset < 0 {

return Err(vm.new_overflow_error("memory mapped offset must be positive"));

}

Ok(self.length as usize)

}

}

#[derive(FromArgs)]

pub(super) struct FlushOptions {

#[pyarg(positional, default)]

offset: Option<isize>,

#[pyarg(positional, default)]

size: Option<isize>,

}

impl FlushOptions {

fn values(self, len: usize) -> Option<(usize, usize)> {

let offset = match self.offset {

Some(o) if o < 0 => return None,

Some(o) => o as usize,

None => 0,

};

let size = match self.size {

Some(s) if s < 0 => return None,

Some(s) => s as usize,

None => len,

};

if len.checked_sub(offset)? < size {

return None;

}

Some((offset, size))

}

}

#[derive(FromArgs, Clone)]

pub(super) struct FindOptions {

#[pyarg(positional)]

sub: Vec<u8>,

#[pyarg(positional, default)]

start: Option<isize>,

#[pyarg(positional, default)]

end: Option<isize>,

}

#[cfg(all(unix, not(target_os = "redox")))]

#[derive(FromArgs)]

pub(super) struct AdviseOptions {

#[pyarg(positional)]

option: libc::c_int,

#[pyarg(positional, default)]

start: Option<PyIntRef>,

#[pyarg(positional, default)]

length: Option<PyIntRef>,

}

#[cfg(all(unix, not(target_os = "redox")))]

impl AdviseOptions {

fn values(self, len: usize, vm: &VirtualMachine) -> PyResult<(libc::c_int, usize, usize)> {

let start = self

.start

.map(|s| {

s.try_to_primitive::<usize>(vm)

.ok()

.filter(|s| *s < len)

.ok_or_else(|| vm.new_value_error("madvise start out of bounds"))

})

.transpose()?

.unwrap_or(0);

let length = self

.length

.map(|s| {

s.try_to_primitive::<usize>(vm)

.map_err(|_| vm.new_value_error("madvise length invalid"))

})

.transpose()?

.unwrap_or(len);

if isize::MAX as usize - start < length {

return Err(vm.new_overflow_error("madvise length too large"));

}

let length = if start + length > len {

len - start

} else {

length

};

Ok((self.option, start, length))

}

}

impl Constructor for PyMmap {

type Args = MmapNewArgs;

#[cfg(unix)]

fn py_new(_cls: &Py<PyType>, args: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {

use libc::{MAP_PRIVATE, MAP_SHARED, PROT_READ, PROT_WRITE};

let mut map_size = args.validate_new_args(vm)?;

let MmapNewArgs {

fileno: fd,

flags,

prot,

access,

offset,

..

} = args;

if (access != AccessMode::Default)

&& ((flags != MAP_SHARED) || (prot != (PROT_WRITE | PROT_READ)))

{

return Err(vm.new_value_error("mmap can't specify both access and flags, prot."));

}

// TODO: memmap2 doesn't support mapping with prot and flags right now

let (_flags, _prot, access) = match access {

AccessMode::Read => (MAP_SHARED, PROT_READ, access),

AccessMode::Write => (MAP_SHARED, PROT_READ | PROT_WRITE, access),

AccessMode::Copy => (MAP_PRIVATE, PROT_READ | PROT_WRITE, access),

AccessMode::Default => {

let access = if (prot & PROT_READ) != 0 && (prot & PROT_WRITE) != 0 {

access

} else if (prot & PROT_WRITE) != 0 {

AccessMode::Write

} else {

AccessMode::Read

};

(flags, prot, access)

}

};

let fd = unsafe { crt_fd::Borrowed::try_borrow_raw(fd) };

// macOS: Issue #11277: fsync(2) is not enough on OS X - a special, OS X specific

// fcntl(2) is necessary to force DISKSYNC and get around mmap(2) bug

#[cfg(target_os = "macos")]

if let Ok(fd) = fd {

host_mmap::prepare_file_mapping(fd);

}

if let Ok(fd) = fd {

let file_len = host_mmap::file_len(fd).map_err(|err| err.to_pyexception(vm))?;

if map_size == 0 {

if file_len == 0 {

return Err(vm.new_value_error("cannot mmap an empty file"));

}

if offset > file_len {

return Err(vm.new_value_error("mmap offset is greater than file size"));

}

map_size = (file_len - offset)

.try_into()

.map_err(|_| vm.new_value_error("mmap length is too large"))?;

} else if offset > file_len || file_len - offset < map_size as i64 {

return Err(vm.new_value_error("mmap length is greater than file size"));

}

}

let (fd, mmap) = || -> std::io::Result<_> {

if let Ok(fd) = fd {

let (new_fd, mmap) = host_mmap::map_file(

fd,

offset,

map_size,

match access {

AccessMode::Default => host_mmap::AccessMode::Default,

AccessMode::Read => host_mmap::AccessMode::Read,

AccessMode::Write => host_mmap::AccessMode::Write,

AccessMode::Copy => host_mmap::AccessMode::Copy,

},

)?;

Ok((Some(new_fd), mmap))

} else {

let mmap = host_mmap::map_anon(map_size)?;

Ok((None, mmap))

}

}()

.map_err(|e| e.to_pyexception(vm))?;

Ok(Self {

closed: AtomicCell::new(false),

mmap: PyMutex::new(Some(MmapObj::Mapped(mmap))),

fd: AtomicCell::new(fd.map_or(-1, |fd| fd.into_raw())),

offset,

size: AtomicCell::new(map_size),

pos: AtomicCell::new(0),

exports: AtomicCell::new(0),

access,

})

}

#[cfg(windows)]

fn py_new(_cls: &Py<PyType>, args: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {

let mut map_size = args.validate_new_args(vm)?;

let MmapNewArgs {

fileno,

tagname,

access,

offset,

..

} = args;

// Parse tagname: None or a string

let tag_str: Option<String> = match tagname {

Some(ref obj) if !vm.is_none(obj) => {

let s = obj

.try_to_value::<String>(vm)

.map_err(|_| vm.new_type_error("tagname must be a string or None"))?;

if s.contains('\0') {

return Err(vm.new_value_error("tagname must not contain null characters"));

}

Some(s)

}

_ => None,

};

// Get file handle from fileno. fileno -1 means anonymous mapping.

let fh: Option<host_mmap::Handle> = if fileno != -1 {

// Convert CRT file descriptor to a Windows file mapping handle.

// Use suppress_iph! to avoid crashes when the fd is invalid.

// This is critical because socket fds wrapped via _open_osfhandle

// may cause crashes in _get_osfhandle on Windows.

// See Python bug https://bugs.python.org/issue30114

let handle = host_nt::handle_from_fd(fileno);

// Check for invalid handle value (-1 on Windows)

if host_mmap::is_invalid_handle_value(handle as isize) {

return Err(vm.new_os_error(format!("Invalid file descriptor: {fileno}")));

}

Some(handle as host_mmap::Handle)

} else {

None

};

// Get file size if we have a file handle and map_size is 0

let mut duplicated_handle: host_mmap::Handle = host_mmap::INVALID_HANDLE;

if let Some(fh) = fh {

// Duplicate handle so Python code can close the original

duplicated_handle =

host_mmap::duplicate_handle(fh).map_err(|e| e.to_pyexception(vm))?;

// Get file size

let file_len = match host_mmap::get_file_len(fh) {

Ok(len) => len,

Err(err) => {

host_mmap::close_handle(duplicated_handle);

return Err(err.to_pyexception(vm));

}

};

if map_size == 0 {

if file_len == 0 {

host_mmap::close_handle(duplicated_handle);

return Err(vm.new_value_error("cannot mmap an empty file"));

}

if offset >= file_len {

host_mmap::close_handle(duplicated_handle);

return Err(vm.new_value_error("mmap offset is greater than file size"));

}

if file_len - offset > isize::MAX as i64 {

host_mmap::close_handle(duplicated_handle);

return Err(vm.new_value_error("mmap length is too large"));

}

map_size = (file_len - offset) as usize;

} else {

// If map_size > file_len, extend the file (Windows behavior)

let required_size = offset.checked_add(map_size as i64).ok_or_else(|| {

host_mmap::close_handle(duplicated_handle);

vm.new_overflow_error("mmap size would cause file size overflow")

})?;

if required_size > file_len

&& let Err(err) = host_mmap::extend_file(duplicated_handle, required_size)

{

host_mmap::close_handle(duplicated_handle);

return Err(err.to_pyexception(vm));

}

}

}

// When tagname is provided, use raw Win32 APIs for named shared memory

if let Some(ref tag) = tag_str {

let fh = if let Some(fh) = fh {

// Close the duplicated handle - we'll use the original

// file handle for CreateFileMappingW

if duplicated_handle != host_mmap::INVALID_HANDLE {

host_mmap::close_handle(duplicated_handle);

}

fh

} else {

host_mmap::INVALID_HANDLE

};

let named = host_mmap::create_named_mapping(

fh,

tag,

match access {

AccessMode::Default => host_mmap::AccessMode::Default,

AccessMode::Read => host_mmap::AccessMode::Read,

AccessMode::Write => host_mmap::AccessMode::Write,

AccessMode::Copy => host_mmap::AccessMode::Copy,

},

offset,

map_size,

)

.map_err(|err| {

if err.raw_os_error() == Some(libc::EOVERFLOW) {

vm.new_overflow_error("mmap offset plus size would overflow")

} else {

err.to_pyexception(vm)

}

})?;

return Ok(Self {

closed: AtomicCell::new(false),

mmap: PyMutex::new(Some(MmapObj::Named(named))),

handle: AtomicCell::new(host_mmap::INVALID_HANDLE as isize),

offset,

size: AtomicCell::new(map_size),

pos: AtomicCell::new(0),

exports: AtomicCell::new(0),

access,

});

}

let (handle, mmap) = if duplicated_handle != host_mmap::INVALID_HANDLE {

let mmap = Self::create_mmap_windows(duplicated_handle, offset, map_size, &access)

.map_err(|e| e.to_pyexception(vm))?;

(duplicated_handle as isize, mmap)

} else {

// Anonymous mapping

let mmap = host_mmap::map_anon(map_size).map_err(|e| e.to_pyexception(vm))?;

(host_mmap::INVALID_HANDLE as isize, MmapObj::Mapped(mmap))

};

Ok(Self {

closed: AtomicCell::new(false),

mmap: PyMutex::new(Some(mmap)),

handle: AtomicCell::new(handle),

offset,

size: AtomicCell::new(map_size),

pos: AtomicCell::new(0),

exports: AtomicCell::new(0),

access,

})

}

}

static BUFFER_METHODS: BufferMethods = BufferMethods {

obj_bytes: |buffer| buffer.obj_as::<PyMmap>().as_bytes(),

obj_bytes_mut: |buffer| buffer.obj_as::<PyMmap>().as_bytes_mut(),

release: |buffer| {

buffer.obj_as::<PyMmap>().exports.fetch_sub(1);

},

retain: |buffer| {

buffer.obj_as::<PyMmap>().exports.fetch_add(1);

},

};

impl AsBuffer for PyMmap {

fn as_buffer(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {

let readonly = matches!(zelf.access, AccessMode::Read);

let buf = PyBuffer::new(

zelf.to_owned().into(),

BufferDescriptor::simple(zelf.__len__(), readonly),

&BUFFER_METHODS,

);

Ok(buf)

}

}

impl AsMapping for PyMmap {

fn as_mapping() -> &'static PyMappingMethods {

static AS_MAPPING: PyMappingMethods = PyMappingMethods {

length: atomic_func!(

|mapping, _vm| Ok(PyMmap::mapping_downcast(mapping).__len__())

),

subscript: atomic_func!(|mapping, needle, vm| {

PyMmap::mapping_downcast(mapping).getitem_inner(needle, vm)

}),

ass_subscript: atomic_func!(|mapping, needle, value, vm| {

let zelf = PyMmap::mapping_downcast(mapping);

if let Some(value) = value {

PyMmap::setitem_inner(zelf, needle, value, vm)

} else {

Err(vm

.new_type_error("mmap object doesn't support item deletion".to_owned()))

}

}),

};

&AS_MAPPING

}

}

impl AsSequence for PyMmap {

fn as_sequence() -> &'static PySequenceMethods {

use rustpython_common::lock::LazyLock;

static AS_SEQUENCE: LazyLock<PySequenceMethods> = LazyLock::new(|| PySequenceMethods {

length: atomic_func!(|seq, _vm| Ok(PyMmap::sequence_downcast(seq).__len__())),

item: atomic_func!(|seq, i, vm| {

let zelf = PyMmap::sequence_downcast(seq);

zelf.getitem_by_index(i, vm)

}),

ass_item: atomic_func!(|seq, i, value, vm| {

let zelf = PyMmap::sequence_downcast(seq);

if let Some(value) = value {

PyMmap::setitem_by_index(zelf, i, value, vm)

} else {

Err(vm

.new_type_error("mmap object doesn't support item deletion".to_owned()))

}

}),

..PySequenceMethods::NOT_IMPLEMENTED

});

&AS_SEQUENCE

}

}

#[pyclass(

impl PyMmap {

fn as_bytes_mut(&self) -> BorrowedValueMut<'_, [u8]> {

PyMutexGuard::map(self.mmap.lock(), |m| {

match m.as_mut().expect("mmap closed or invalid") {

MmapObj::Mapped(mmap) => mmap.as_mut_slice(),

#[cfg(windows)]

MmapObj::Named(named) => named.as_mut_slice(),

}

})

.into()

}

fn as_bytes(&self) -> BorrowedValue<'_, [u8]> {

PyMutexGuard::map_immutable(self.mmap.lock(), |m| {

m.as_ref().expect("mmap closed or invalid").as_slice()

})

.into()

}

fn __len__(&self) -> usize {

self.size.load()

}

#[inline]

fn pos(&self) -> usize {

self.pos.load()

}

#[inline]

fn advance_pos(&self, step: usize) {

self.pos.store(self.pos() + step);

}

#[inline]

fn try_writable<R>(

&self,

vm: &VirtualMachine,

f: impl FnOnce(&mut [u8]) -> R,

) -> PyResult<R> {

if matches!(self.access, AccessMode::Read) {

return Err(vm.new_type_error("mmap can't modify a readonly memory map."));

}

match self.check_valid(vm)?.deref_mut().as_mut().unwrap() {

MmapObj::Mapped(mmap) => Ok(f(mmap.as_mut_slice())),

#[cfg(windows)]

MmapObj::Named(named) => Ok(f(named.as_mut_slice())),

}

}

fn check_valid(&self, vm: &VirtualMachine) -> PyResult<PyMutexGuard<'_, Option<MmapObj>>> {

let m = self.mmap.lock();

if m.is_none() {

return Err(vm.new_value_error("mmap closed or invalid"));

}

Ok(m)

}

/// TODO: impl resize

#[allow(dead_code)]

fn check_resizeable(&self, vm: &VirtualMachine) -> PyResult<()> {

if self.exports.load() > 0 {

return Err(vm.new_buffer_error("mmap can't resize with extant buffers exported."));

}

if self.access == AccessMode::Write || self.access == AccessMode::Default {

return Ok(());

}

Err(vm.new_type_error("mmap can't resize a readonly or copy-on-write memory map."))

}

#[pygetset]

fn closed(&self) -> bool {

self.closed.load()

}

#[pymethod]

fn close(&self, vm: &VirtualMachine) -> PyResult<()> {

if self.closed() {

return Ok(());

}

if self.exports.load() > 0 {

return Err(vm.new_buffer_error("cannot close exported pointers exist."));

}

let mut mmap = self.mmap.lock();

self.closed.store(true);

*mmap = None;

self.close_handle();

Ok(())

}

fn get_find_range(&self, options: FindOptions) -> (usize, usize) {

let size = self.__len__();

let start = options

.start

.map_or_else(|| self.pos(), |start| start.saturated_at(size));

let end = options.end.map_or(size, |end| end.saturated_at(size));

(start, end)

}

#[pymethod]

fn find(&self, options: FindOptions, vm: &VirtualMachine) -> PyResult<PyInt> {

let (start, end) = self.get_find_range(options.clone());

let sub = &options.sub;

// returns start position for empty string

if sub.is_empty() {

return Ok(PyInt::from(start as isize));

}

let mmap = self.check_valid(vm)?;

let buf = &mmap.as_ref().unwrap().as_slice()[start..end];

let pos = buf.windows(sub.len()).position(|window| window == sub);

Ok(pos.map_or_else(|| PyInt::from(-1isize), |i| PyInt::from(start + i)))

}

#[pymethod]

fn rfind(&self, options: FindOptions, vm: &VirtualMachine) -> PyResult<PyInt> {

let (start, end) = self.get_find_range(options.clone());

let sub = &options.sub;

// returns start position for empty string

if sub.is_empty() {

return Ok(PyInt::from(start as isize));

}

let mmap = self.check_valid(vm)?;

let buf = &mmap.as_ref().unwrap().as_slice()[start..end];

let pos = buf.windows(sub.len()).rposition(|window| window == sub);

Ok(pos.map_or_else(|| PyInt::from(-1isize), |i| PyInt::from(start + i)))

}

#[pymethod]

fn flush(&self, options: FlushOptions, vm: &VirtualMachine) -> PyResult<()> {

let (offset, size) = options

.values(self.__len__())

.ok_or_else(|| vm.new_value_error("flush values out of range"))?;

if self.access == AccessMode::Read || self.access == AccessMode::Copy {

return Ok(());

}

match self.check_valid(vm)?.deref().as_ref().unwrap() {

MmapObj::Mapped(mmap) => {