X Tutup
Skip to content

Commit 8fd4640

Browse files
committed
Fixed some issues, all test_mmap tests pass
1 parent 18e3402 commit 8fd4640

File tree

1 file changed

+144
-51
lines changed

1 file changed

+144
-51
lines changed

stdlib/src/mmap.rs

Lines changed: 144 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -6,21 +6,19 @@ mod mmap {
66
use crate::common::{
77
borrow::{BorrowedValue, BorrowedValueMut},
88
lock::{MapImmutable, PyMutex, PyMutexGuard},
9+
rc::PyRc,
910
};
1011
use crate::vm::{
11-
builtins::PyTypeRef,
12-
builtins::{PyBytes, PyBytesRef, PyInt, PyIntRef},
12+
builtins::{PyBytes, PyBytesRef, PyInt, PyIntRef, PyTypeRef},
1313
bytesinner::{bytes_from_object, value_from_object},
14-
convert::ToPyException,
15-
function::{ArgBytesLike, OptionalArg},
14+
function::{ArgBytesLike, FuncArgs, OptionalArg},
1615
protocol::{
1716
BufferDescriptor, BufferMethods, PyBuffer, PyMappingMethods, PySequenceMethods,
1817
},
1918
sliceable::{SaturatedSlice, SequenceIndex},
20-
types::Constructor,
21-
types::{AsBuffer, AsMapping, AsSequence},
22-
FromArgs, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromBorrowedObject,
23-
VirtualMachine,
19+
types::{AsBuffer, AsMapping, AsSequence, Constructor},
20+
AsObject, FromArgs, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
21+
TryFromBorrowedObject, VirtualMachine,
2422
};
2523
use crossbeam_utils::atomic::AtomicCell;
2624
use memmap2::{Advice, Mmap, MmapMut, MmapOptions};
@@ -29,8 +27,7 @@ mod mmap {
2927
use std::io::Write;
3028
use std::ops::{Deref, DerefMut};
3129
#[cfg(all(unix, not(target_os = "redox")))]
32-
use std::os::unix::io::{FromRawFd, RawFd, IntoRawFd};
33-
use std::sync::Arc;
30+
use std::os::unix::io::{FromRawFd, IntoRawFd, RawFd};
3431

3532
fn advice_try_from_i32(vm: &VirtualMachine, i: i32) -> PyResult<Advice> {
3633
Ok(match i {
@@ -117,10 +114,25 @@ mod mmap {
117114
#[pyattr]
118115
const ACCESS_COPY: u32 = AccessMode::Copy as u32;
119116

117+
#[pyattr(name = "PAGESIZE", once)]
118+
fn page_size(_vm: &VirtualMachine) -> usize {
119+
page_size::get()
120+
}
121+
122+
#[pyattr(name = "error", once)]
123+
fn error_type(vm: &VirtualMachine) -> PyTypeRef {
124+
vm.ctx.exceptions.os_error.to_owned()
125+
}
126+
127+
#[pyattr(name = "ALLOCATIONGRANULARITY", once)]
128+
fn granularity(_vm: &VirtualMachine) -> usize {
129+
page_size::get_granularity()
130+
}
131+
120132
#[derive(Debug)]
121133
enum MmapObj {
122134
Write(MmapMut),
123-
Read(Arc<Mmap>),
135+
Read(PyRc<Mmap>),
124136
None,
125137
}
126138

@@ -137,11 +149,10 @@ mod mmap {
137149
closed: AtomicCell<bool>,
138150
mmap: PyMutex<MmapObj>,
139151
_fd: RawFd,
140-
_offset: isize,
152+
offset: isize,
141153
size: AtomicCell<isize>,
142154
pos: AtomicCell<isize>, // relative to offset
143155
exports: AtomicCell<usize>,
144-
// PyObject *weakreflist;
145156
access: AccessMode,
146157
}
147158

@@ -208,7 +219,7 @@ mod mmap {
208219
fn py_new(
209220
cls: PyTypeRef,
210221
MmapNewArgs {
211-
fileno:mut fd,
222+
fileno: mut fd,
212223
length,
213224
flags,
214225
prot,
@@ -288,23 +299,32 @@ mod mmap {
288299
}
289300

290301
let mut mmap_opt = MmapOptions::new();
291-
let mmap_opt = mmap_opt.offset(offset.try_into().unwrap()).len(map_size.try_into().unwrap());
302+
let mmap_opt = mmap_opt
303+
.offset(offset.try_into().unwrap())
304+
.len(map_size.try_into().unwrap());
292305

293306
let (fd, mmap) = if fd == -1 {
294-
(fd, MmapObj::Write(
295-
mmap_opt.map_anon().map_err(|e| vm.new_os_error(e.to_string()))?,
296-
))
307+
(
308+
fd,
309+
MmapObj::Write(
310+
mmap_opt
311+
.map_anon()
312+
.map_err(|e| vm.new_os_error(e.to_string()))?,
313+
),
314+
)
297315
} else {
298-
let new_fd = nix::unistd::dup(fd).map_err(|e| e.to_pyexception(vm))?;
316+
let new_fd = nix::unistd::dup(fd).map_err(|e| vm.new_os_error(e.to_string()))?;
299317
let mmap = match access {
300318
AccessMode::Default | AccessMode::Write => MmapObj::Write(
301-
unsafe { mmap_opt.map_mut(fd) }.map_err(|e| vm.new_os_error(e.to_string()))?,
319+
unsafe { mmap_opt.map_mut(fd) }
320+
.map_err(|e| vm.new_os_error(e.to_string()))?,
302321
),
303-
AccessMode::Read => MmapObj::Read(Arc::new(
322+
AccessMode::Read => MmapObj::Read(PyRc::new(
304323
unsafe { mmap_opt.map(fd) }.map_err(|e| vm.new_os_error(e.to_string()))?,
305324
)),
306325
AccessMode::Copy => MmapObj::Write(
307-
unsafe { mmap_opt.map_copy(fd) }.map_err(|e| vm.new_os_error(e.to_string()))?,
326+
unsafe { mmap_opt.map_copy(fd) }
327+
.map_err(|e| vm.new_os_error(e.to_string()))?,
308328
),
309329
};
310330
(new_fd, mmap)
@@ -314,7 +334,7 @@ mod mmap {
314334
closed: AtomicCell::new(false),
315335
mmap: PyMutex::new(mmap),
316336
_fd: fd,
317-
_offset: offset,
337+
offset: offset,
318338
size: AtomicCell::new(map_size.into()),
319339
pos: AtomicCell::new(0),
320340
exports: AtomicCell::new(0),
@@ -331,7 +351,6 @@ mod mmap {
331351
release: |buffer| {
332352
buffer.obj_as::<PyMmap>().exports.fetch_sub(1);
333353
},
334-
335354
retain: |buffer| {
336355
buffer.obj_as::<PyMmap>().exports.fetch_add(1);
337356
},
@@ -385,15 +404,15 @@ mod mmap {
385404
};
386405
}
387406

388-
#[pyimpl(with(Constructor, AsMapping, AsSequence, AsBuffer))]
407+
#[pyimpl(with(Constructor, AsMapping, AsSequence, AsBuffer), flags(BASETYPE))]
389408
impl PyMmap {
390409
fn as_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
391410
PyMutexGuard::map(self.mmap.lock(), |m| match m {
392411
MmapObj::Read(_) => panic!("mmap can't modify a readonly memory map."),
393412
MmapObj::Write(mmap) => &mut mmap[..],
394413
_ => panic!("mmap closed or invalid"),
395414
})
396-
.into()
415+
.into()
397416
}
398417

399418
fn as_bytes(&self) -> BorrowedValue<[u8]> {
@@ -402,7 +421,7 @@ mod mmap {
402421
MmapObj::Write(mmap) => &mmap[..],
403422
_ => panic!("mmap closed or invalid"),
404423
})
405-
.into()
424+
.into()
406425
}
407426

408427
#[pymethod(magic)]
@@ -469,6 +488,32 @@ mod mmap {
469488
self.closed.load()
470489
}
471490

491+
#[pymethod(magic)]
492+
fn repr(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyResult<String> {
493+
let mmap = zelf.mmap.lock();
494+
495+
if mmap.is_none() {
496+
return Ok("<mmap.mmap closed=True>".to_owned());
497+
}
498+
499+
let access_str = match zelf.access {
500+
AccessMode::Default => "ACCESS_DEFAULT",
501+
AccessMode::Read => "ACCESS_READ",
502+
AccessMode::Write => "ACCESS_WRITE",
503+
AccessMode::Copy => "ACCESS_COPY",
504+
};
505+
506+
let repr = format!(
507+
"<mmap.mmap closed=False, access={}, length={}, pos={}, offset={}>",
508+
access_str,
509+
zelf.len(),
510+
zelf.inner_pos(),
511+
zelf.offset
512+
);
513+
514+
Ok(repr)
515+
}
516+
472517
#[pymethod]
473518
fn close(&self, vm: &VirtualMachine) -> PyResult<()> {
474519
if self.closed() {
@@ -484,7 +529,7 @@ mod mmap {
484529
Ok(())
485530
}
486531

487-
fn get_find_start_end(&self, options: FindOptions) -> (usize, usize) {
532+
fn get_find_range(&self, options: FindOptions) -> (usize, usize) {
488533
let pos = self.inner_pos();
489534
let size = self.inner_size();
490535
let mut start = options.start.map_or(pos, saturate_to_isize);
@@ -512,35 +557,55 @@ mod mmap {
512557
}
513558

514559
#[pymethod]
515-
fn find(&self, options: FindOptions, vm: &VirtualMachine) -> PyResult<Option<usize>> {
516-
let (start, end) = self.get_find_start_end(options.clone());
560+
fn find(&self, options: FindOptions, vm: &VirtualMachine) -> PyResult<isize> {
561+
let (start, end) = self.get_find_range(options.clone());
517562

518563
let sub = &options.sub;
519-
match self.check_valid(vm)?.deref() {
520-
MmapObj::Read(mmap) => Ok(mmap[start..end]
521-
.windows(sub.len())
522-
.position(|window| window == sub)),
523-
MmapObj::Write(mmap) => Ok(mmap[start..end]
524-
.windows(sub.len())
525-
.position(|window| window == sub)),
564+
565+
if sub.is_empty() {
566+
return Ok(0);
567+
}
568+
let pos = match self.check_valid(vm)?.deref() {
569+
MmapObj::Read(mmap) => mmap[start..end]
570+
.windows(sub.len())
571+
.position(|window| window == sub),
572+
MmapObj::Write(mmap) => mmap[start..end]
573+
.windows(sub.len())
574+
.position(|window| window == sub),
526575
_ => unreachable!("alreay check"),
576+
};
577+
578+
if let Some(i) = pos {
579+
return Ok((start + i) as isize);
527580
}
581+
582+
Ok(-1)
528583
}
529584

530585
#[pymethod]
531-
fn rfind(&self, options: FindOptions, vm: &VirtualMachine) -> PyResult<Option<usize>> {
532-
let (start, end) = self.get_find_start_end(options.clone());
586+
fn rfind(&self, options: FindOptions, vm: &VirtualMachine) -> PyResult<isize> {
587+
let (start, end) = self.get_find_range(options.clone());
533588

534589
let sub = &options.sub;
535-
match self.check_valid(vm)?.deref() {
536-
MmapObj::Read(mmap) => Ok(mmap[start..end]
537-
.windows(sub.len())
538-
.rposition(|window| window == sub)),
539-
MmapObj::Write(mmap) => Ok(mmap[start..end]
540-
.windows(sub.len())
541-
.rposition(|window| window == sub)),
590+
if sub.is_empty() {
591+
return Ok(0);
592+
}
593+
594+
let pos = match self.check_valid(vm)?.deref() {
595+
MmapObj::Read(mmap) => mmap[start..end]
596+
.windows(sub.len())
597+
.rposition(|window| window == sub),
598+
MmapObj::Write(mmap) => mmap[start..end]
599+
.windows(sub.len())
600+
.rposition(|window| window == sub),
542601
_ => unreachable!("alreay check"),
602+
};
603+
604+
if let Some(i) = pos {
605+
return Ok((start + i) as isize);
543606
}
607+
608+
Ok(-1)
544609
}
545610

546611
#[pymethod]
@@ -583,7 +648,7 @@ mod mmap {
583648
}
584649

585650
if isize::MAX - start < length {
586-
return Err(vm.new_value_error("madvise length too large".to_owned()));
651+
return Err(vm.new_overflow_error("madvise length too large".to_owned()));
587652
}
588653

589654
if start + length > self.inner_size() {
@@ -651,8 +716,21 @@ mod mmap {
651716
}
652717

653718
#[pymethod]
654-
fn read(&self, n: OptionalArg<PyIntRef>, vm: &VirtualMachine) -> PyResult<PyBytesRef> {
655-
let mut num_bytes = n.map_or(isize::MAX, saturate_to_isize);
719+
fn read(&self, n: OptionalArg<PyObjectRef>, vm: &VirtualMachine) -> PyResult<PyBytesRef> {
720+
let mut num_bytes = n.map_or(Ok(isize::MAX), |obj| {
721+
if vm.is_none(&obj) {
722+
Ok(isize::MAX)
723+
} else {
724+
let n = obj.downcast::<PyInt>().map_err(|obj| {
725+
vm.new_type_error(format!(
726+
"read argument must be int or None, not {}",
727+
obj.class().name()
728+
))
729+
})?;
730+
731+
Ok(saturate_to_isize(n))
732+
}
733+
})?;
656734
let mmap = self.check_valid(vm)?;
657735
let pos = self.inner_pos();
658736

@@ -741,6 +819,7 @@ mod mmap {
741819
//TODO: supports resize
742820
#[pymethod]
743821
fn resize(&self, _newsize: PyIntRef, vm: &VirtualMachine) -> PyResult<()> {
822+
self.check_resizeable(vm)?;
744823
Err(vm.new_system_error("mmap: resizing not available--no mremap()".to_owned()))
745824
}
746825

@@ -990,19 +1069,22 @@ mod mmap {
9901069
_ => unreachable!("alreay check"),
9911070
}
9921071
} else {
1072+
let mut bi = 0; // bytes index
9931073
if step.is_negative() {
9941074
for i in range.rev().step_by(step.unsigned_abs()) {
9951075
match mmap.deref_mut() {
996-
MmapObj::Write(mmap) => mmap[i] = bytes[i],
1076+
MmapObj::Write(mmap) => mmap[i] = bytes[bi],
9971077
_ => unreachable!("alreay check"),
9981078
}
1079+
bi = bi + 1;
9991080
}
10001081
} else {
10011082
for i in range.step_by(step.unsigned_abs()) {
10021083
match mmap.deref_mut() {
1003-
MmapObj::Write(mmap) => mmap[i] = bytes[i],
1084+
MmapObj::Write(mmap) => mmap[i] = bytes[bi],
10041085
_ => unreachable!("alreay check"),
10051086
}
1087+
bi = bi + 1;
10061088
}
10071089
}
10081090
}
@@ -1019,5 +1101,16 @@ mod mmap {
10191101
) -> PyResult<()> {
10201102
Self::_setitem(zelf, &needle, value, vm)
10211103
}
1104+
1105+
#[pymethod(magic)]
1106+
fn enter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
1107+
let _m = zelf.check_valid(vm)?;
1108+
Ok(zelf.to_owned())
1109+
}
1110+
1111+
#[pymethod(magic)]
1112+
fn exit(zelf: PyRef<Self>, _args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> {
1113+
zelf.close(vm)
1114+
}
10221115
}
10231116
}

0 commit comments

Comments
 (0)
X Tutup