make range faster

This commit is contained in:
blueloveTH 2024-05-25 14:35:56 +08:00
parent 52d7133e1b
commit 9a5ff3440a
4 changed files with 35 additions and 8 deletions

View File

@ -6,7 +6,7 @@
namespace pkpy{ namespace pkpy{
struct RangeIter{ struct RangeIter{ // step > 0
Range r; Range r;
i64 current; i64 current;
RangeIter(Range r) : r(r), current(r.start) {} RangeIter(Range r) : r(r), current(r.start) {}
@ -14,6 +14,14 @@ struct RangeIter{
static void _register(VM* vm, PyVar mod, PyVar type); static void _register(VM* vm, PyVar mod, PyVar type);
}; };
struct RangeIterR{ // step < 0
Range r;
i64 current;
RangeIterR(Range r) : r(r), current(r.start) {}
static void _register(VM* vm, PyVar mod, PyVar type);
};
struct ArrayIter{ struct ArrayIter{
PyVar ref; PyVar ref;
PyVar* begin; PyVar* begin;

View File

@ -6,12 +6,19 @@ namespace pkpy{
vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyVar _0){ return _0; }); vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyVar _0){ return _0; });
vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyVar _0) -> unsigned{ vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyVar _0) -> unsigned{
RangeIter& self = PK_OBJ_GET(RangeIter, _0); RangeIter& self = PK_OBJ_GET(RangeIter, _0);
if(self.r.step > 0){
if(self.current >= self.r.stop) return 0; if(self.current >= self.r.stop) return 0;
}else{ vm->s_data.emplace(VM::tp_int, self.current);
if(self.current <= self.r.stop) return 0; self.current += self.r.step;
return 1;
});
} }
vm->s_data.push(VAR(self.current));
void RangeIterR::_register(VM* vm, PyVar mod, PyVar type){
vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyVar _0){ return _0; });
vm->bind__next__(PK_OBJ_GET(Type, type), [](VM* vm, PyVar _0) -> unsigned{
RangeIterR& self = PK_OBJ_GET(RangeIterR, _0);
if(self.current <= self.r.stop) return 0;
vm->s_data.emplace(VM::tp_int, self.current);
self.current += self.r.step; self.current += self.r.step;
return 1; return 1;
}); });

View File

@ -363,7 +363,14 @@ void __init_builtins(VM* _vm) {
return VAR(r); return VAR(r);
}); });
_vm->bind__iter__(VM::tp_range, [](VM* vm, PyVar obj) { return vm->new_user_object<RangeIter>(PK_OBJ_GET(Range, obj)); }); _vm->bind__iter__(VM::tp_range, [](VM* vm, PyVar obj) {
const Range& r = PK_OBJ_GET(Range, obj);
if(r.step > 0){
return vm->new_user_object<RangeIter>(r);
}else{
return vm->new_user_object<RangeIterR>(r);
}
});
// tp_nonetype // tp_nonetype
_vm->bind__repr__(_vm->_tp(_vm->None), [](VM* vm, PyVar _0) -> Str { _vm->bind__repr__(_vm->_tp(_vm->None), [](VM* vm, PyVar _0) -> Str {
@ -1491,6 +1498,7 @@ void __init_builtins(VM* _vm) {
}); });
_vm->register_user_class<RangeIter>(_vm->builtins, "_range_iter"); _vm->register_user_class<RangeIter>(_vm->builtins, "_range_iter");
_vm->register_user_class<RangeIterR>(_vm->builtins, "_range_iter_r");
_vm->register_user_class<ArrayIter>(_vm->builtins, "_array_iter"); _vm->register_user_class<ArrayIter>(_vm->builtins, "_array_iter");
_vm->register_user_class<StringIter>(_vm->builtins, "_string_iter"); _vm->register_user_class<StringIter>(_vm->builtins, "_string_iter");
_vm->register_user_class<Generator>(_vm->builtins, "generator"); _vm->register_user_class<Generator>(_vm->builtins, "generator");

View File

@ -134,3 +134,7 @@ a, b = [1, 2]
assert a == 1 and b == 2 assert a == 1 and b == 2
assert [1, 2].__getitem__(0) == 1 assert [1, 2].__getitem__(0) == 1
assert list(range(1, 5)) == [1, 2, 3, 4]
assert list(range(1, 5, 2)) == [1, 3]
assert list(range(5, 1, -1)) == [5, 4, 3, 2]
assert list(range(5, 1, -2)) == [5, 3]