Skip to content

Commit 1180f99

Browse files
kabra1110kabra1110
authored
Add reserve for list (#2238)
1 parent 3927b92 commit 1180f99

File tree

8 files changed

+155
-0
lines changed

8 files changed

+155
-0
lines changed

‎integration_tests/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -518,6 +518,7 @@ RUN(NAME test_list_pop2 LABELS cpython llvm NOFAST) # TODO: Remove NOFAST
518518
RUN(NAME test_list_pop3 LABELS cpython llvm)
519519
RUN(NAME test_list_compare LABELS cpython llvm)
520520
RUN(NAME test_list_concat LABELS cpython llvm c NOFAST)
521+
RUN(NAME test_list_reserve LABELS cpython llvm)
521522
RUN(NAME test_tuple_01 LABELS cpython llvm c)
522523
RUN(NAME test_tuple_02 LABELS cpython llvm c NOFAST)
523524
RUN(NAME test_tuple_03 LABELS cpython llvm c)

‎integration_tests/test_list_reserve.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
from lpython import i32, f64, reserve
2+
3+
def test_list_reserve():
4+
l1: list[i32] = []
5+
l2: list[list[tuple[f64, str, tuple[i32, f64]]]] = []
6+
i: i32
7+
8+
reserve(l1, 100)
9+
for i in range(50):
10+
l1.append(i)
11+
assert len(l1) == i + 1
12+
13+
reserve(l1, 150)
14+
15+
for i in range(50):
16+
l1.pop(0)
17+
assert len(l1) == 49 - i
18+
19+
reserve(l2, 100)
20+
for i in range(50):
21+
l2.append([(f64(i * i), str(i), (i, f64(i + 1))), (f64(i), str(i), (i, f64(i)))])
22+
assert len(l2) == i + 1
23+
24+
reserve(l2, 150)
25+
26+
for i in range(50):
27+
l2.pop(0)
28+
assert len(l2) == 49 - i
29+
30+
test_list_reserve()

‎src/libasr/codegen/asr_to_llvm.cpp

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1687,6 +1687,21 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
16871687
tmp = list_api->pop_position(plist, pos, asr_el_type, module.get(), name2memidx);
16881688
}
16891689

1690+
void generate_Reserve(ASR::expr_t* m_arg, ASR::expr_t* m_ele) {
1691+
// For now, this only handles lists
1692+
ASR::ttype_t* asr_el_type = ASRUtils::get_contained_type(ASRUtils::expr_type(m_arg));
1693+
int64_t ptr_loads_copy = ptr_loads;
1694+
ptr_loads = 0;
1695+
this->visit_expr(*m_arg);
1696+
llvm::Value* plist = tmp;
1697+
1698+
ptr_loads = 2;
1699+
this->visit_expr_wrapper(m_ele, true);
1700+
ptr_loads = ptr_loads_copy;
1701+
llvm::Value* n = tmp;
1702+
list_api->reserve(plist, n, asr_el_type, module.get());
1703+
}
1704+
16901705
void generate_DictElems(ASR::expr_t* m_arg, bool key_or_value) {
16911706
ASR::Dict_t* dict_type = ASR::down_cast<ASR::Dict_t>(
16921707
ASRUtils::expr_type(m_arg));
@@ -1807,6 +1822,10 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
18071822
}
18081823
break;
18091824
}
1825+
case ASRUtils::IntrinsicFunctions::Reserve: {
1826+
generate_Reserve(x.m_args[0], x.m_args[1]);
1827+
break;
1828+
}
18101829
case ASRUtils::IntrinsicFunctions::DictKeys: {
18111830
generate_DictElems(x.m_args[0], 0);
18121831
break;

‎src/libasr/codegen/llvm_utils.cpp

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4128,6 +4128,33 @@ namespace LCompilers {
41284128
shift_end_point_by_one(list);
41294129
}
41304130

4131+
void LLVMList::reserve(llvm::Value* list, llvm::Value* n,
4132+
ASR::ttype_t* asr_type, llvm::Module* module) {
4133+
/**
4134+
* C++ equivalent
4135+
*
4136+
* if( n > current_capacity ) {
4137+
* list_data = realloc(list_data, sizeof(el_type) * n);
4138+
* }
4139+
*
4140+
*/
4141+
llvm::Value* capacity = LLVM::CreateLoad(*builder, get_pointer_to_current_capacity(list));
4142+
std::string type_code = ASRUtils::get_type_code(asr_type);
4143+
int type_size = std::get<1>(typecode2listtype[type_code]);
4144+
llvm::Type* el_type = std::get<2>(typecode2listtype[type_code]);
4145+
llvm_utils->create_if_else(builder->CreateICmpSGT(n, capacity), [&]() {
4146+
llvm::Value* arg_size = builder->CreateMul(llvm::ConstantInt::get(context,
4147+
llvm::APInt(32, type_size)), n);
4148+
llvm::Value* copy_data_ptr = get_pointer_to_list_data(list);
4149+
llvm::Value* copy_data = LLVM::CreateLoad(*builder, copy_data_ptr);
4150+
copy_data = LLVM::lfortran_realloc(context, *module, *builder,
4151+
copy_data, arg_size);
4152+
copy_data = builder->CreateBitCast(copy_data, el_type->getPointerTo());
4153+
builder->CreateStore(copy_data, copy_data_ptr);
4154+
builder->CreateStore(n, get_pointer_to_current_capacity(list));
4155+
}, []() {});
4156+
}
4157+
41314158
void LLVMList::reverse(llvm::Value* list, llvm::Module& module) {
41324159

41334160
/* Equivalent in C++:

‎src/libasr/codegen/llvm_utils.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -407,6 +407,9 @@ namespace LCompilers {
407407
llvm::Module* module,
408408
std::map<std::string, std::map<std::string, int>>& name2memidx);
409409

410+
void reserve(llvm::Value* list, llvm::Value* n,
411+
ASR::ttype_t* asr_type, llvm::Module* module);
412+
410413
void remove(llvm::Value* list, llvm::Value* item,
411414
ASR::ttype_t* item_type, llvm::Module& module);
412415

‎src/libasr/pass/intrinsic_function_registry.h

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ enum class IntrinsicFunctions : int64_t {
4646
Partition,
4747
ListReverse,
4848
ListPop,
49+
Reserve,
4950
DictKeys,
5051
DictValues,
5152
SetAdd,
@@ -102,6 +103,7 @@ inline std::string get_intrinsic_name(int x) {
102103
INTRINSIC_NAME_CASE(Partition)
103104
INTRINSIC_NAME_CASE(ListReverse)
104105
INTRINSIC_NAME_CASE(ListPop)
106+
INTRINSIC_NAME_CASE(Reserve)
105107
INTRINSIC_NAME_CASE(DictKeys)
106108
INTRINSIC_NAME_CASE(DictValues)
107109
INTRINSIC_NAME_CASE(SetAdd)
@@ -1262,6 +1264,55 @@ static inline ASR::asr_t* create_ListPop(Allocator& al, const Location& loc,
12621264

12631265
} // namespace ListPop
12641266

1267+
namespace Reserve {
1268+
1269+
static inline void verify_args(const ASR::IntrinsicFunction_t& x, diag::Diagnostics& diagnostics) {
1270+
ASRUtils::require_impl(x.n_args == 2, "Call to reserve must have exactly one argument",
1271+
x.base.base.loc, diagnostics);
1272+
ASRUtils::require_impl(ASR::is_a<ASR::List_t>(*ASRUtils::expr_type(x.m_args[0])),
1273+
"First argument to reserve must be of list type",
1274+
x.base.base.loc, diagnostics);
1275+
ASRUtils::require_impl(ASR::is_a<ASR::Integer_t>(*ASRUtils::expr_type(x.m_args[1])),
1276+
"Second argument to reserve must be an integer",
1277+
x.base.base.loc, diagnostics);
1278+
ASRUtils::require_impl(x.m_type == nullptr,
1279+
"Return type of reserve must be empty",
1280+
x.base.base.loc, diagnostics);
1281+
}
1282+
1283+
static inline ASR::expr_t *eval_reserve(Allocator &/*al*/,
1284+
const Location &/*loc*/, Vec<ASR::expr_t*>& /*args*/) {
1285+
// TODO: To be implemented for ListConstant expression
1286+
return nullptr;
1287+
}
1288+
1289+
static inline ASR::asr_t* create_Reserve(Allocator& al, const Location& loc,
1290+
Vec<ASR::expr_t*>& args,
1291+
const std::function<void (const std::string &, const Location &)> err) {
1292+
if (args.size() != 2) {
1293+
err("Call to reserve must have exactly two argument", loc);
1294+
}
1295+
if (!ASR::is_a<ASR::List_t>(*ASRUtils::expr_type(args[0]))) {
1296+
err("First argument to reserve must be of list type", loc);
1297+
}
1298+
if (!ASR::is_a<ASR::Integer_t>(*ASRUtils::expr_type(args[1]))) {
1299+
err("Second argument to reserve must be an integer", loc);
1300+
}
1301+
1302+
Vec<ASR::expr_t*> arg_values;
1303+
arg_values.reserve(al, args.size());
1304+
for( size_t i = 0; i < args.size(); i++ ) {
1305+
arg_values.push_back(al, ASRUtils::expr_value(args[i]));
1306+
}
1307+
ASR::expr_t* compile_time_value = eval_reserve(al, loc, arg_values);
1308+
return ASR::make_Expr_t(al, loc,
1309+
ASRUtils::EXPR(ASRUtils::make_IntrinsicFunction_t_util(al, loc,
1310+
static_cast<int64_t>(ASRUtils::IntrinsicFunctions::Reserve),
1311+
args.p, args.size(), 0, nullptr, compile_time_value)));
1312+
}
1313+
1314+
} // namespace Reserve
1315+
12651316
namespace DictKeys {
12661317

12671318
static inline void verify_args(const ASR::IntrinsicFunction_t& x, diag::Diagnostics& diagnostics) {
@@ -3124,6 +3175,8 @@ namespace IntrinsicFunctionRegistry {
31243175
{nullptr, &DictValues::verify_args}},
31253176
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::ListPop),
31263177
{nullptr, &ListPop::verify_args}},
3178+
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::Reserve),
3179+
{nullptr, &Reserve::verify_args}},
31273180
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::SetAdd),
31283181
{nullptr, &SetAdd::verify_args}},
31293182
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::SetRemove),
@@ -3206,6 +3259,8 @@ namespace IntrinsicFunctionRegistry {
32063259
"list.reverse"},
32073260
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::ListPop),
32083261
"list.pop"},
3262+
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::Reserve),
3263+
"reserve"},
32093264
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::DictKeys),
32103265
"dict.keys"},
32113266
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::DictValues),
@@ -3290,6 +3345,7 @@ namespace IntrinsicFunctionRegistry {
32903345
{"list.index", {&ListIndex::create_ListIndex, &ListIndex::eval_list_index}},
32913346
{"list.reverse", {&ListReverse::create_ListReverse, &ListReverse::eval_list_reverse}},
32923347
{"list.pop", {&ListPop::create_ListPop, &ListPop::eval_list_pop}},
3348+
{"reserve", {&Reserve::create_Reserve, &Reserve::eval_reserve}},
32933349
{"dict.keys", {&DictKeys::create_DictKeys, &DictKeys::eval_dict_keys}},
32943350
{"dict.values", {&DictValues::create_DictValues, &DictValues::eval_dict_values}},
32953351
{"set.add", {&SetAdd::create_SetAdd, &SetAdd::eval_set_add}},

‎src/lpython/semantics/python_ast_to_asr.cpp

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6511,6 +6511,9 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
65116511
// Keyword arguments to be handled in make_call_helper
65126512
args.reserve(al, c->n_args);
65136513
visit_expr_list(c->m_args, c->n_args, args);
6514+
// TODO: Avoid overriding of user defined functions with same name as
6515+
// intrinsics like print, quit and reserve. Right now, user defined
6516+
// functions will never be considered.
65146517
if (call_name == "print") {
65156518
ASR::expr_t *fmt = nullptr;
65166519
Vec<ASR::expr_t*> args_expr = ASRUtils::call_arg2expr(al, args);
@@ -6570,6 +6573,17 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
65706573
}
65716574
tmp = ASR::make_Stop_t(al, x.base.base.loc, code);
65726575
return;
6576+
} else if( call_name == "reserve" ) {
6577+
ASRUtils::create_intrinsic_function create_func =
6578+
ASRUtils::IntrinsicFunctionRegistry::get_create_function("reserve");
6579+
Vec<ASR::expr_t*> args_exprs; args_exprs.reserve(al, args.size());
6580+
for( size_t i = 0; i < args.size(); i++ ) {
6581+
args_exprs.push_back(al, args[i].m_value);
6582+
}
6583+
tmp = create_func(al, x.base.base.loc, args_exprs,
6584+
[&](const std::string &msg, const Location &loc) {
6585+
throw SemanticError(msg, loc); });
6586+
return ;
65736587
}
65746588
ASR::symbol_t *s = current_scope->resolve_symbol(call_name);
65756589
if (!s) {

‎src/runtime/lpython/lpython.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -760,6 +760,11 @@ def __lpython(*args, **kwargs):
760760
def bitnot(x, bitsize):
761761
return (~x) % (2 ** bitsize)
762762

763+
def reserve(data_structure, n):
764+
if isinstance(data_structure, list):
765+
data_structure = [None] * n
766+
# no-op
767+
763768
bitnot_u8 = lambda x: bitnot(x, 8)
764769
bitnot_u16 = lambda x: bitnot(x, 16)
765770
bitnot_u32 = lambda x: bitnot(x, 32)

0 commit comments

Comments
 (0)