Skip to content

Commit a649c9a

Browse files
kabra1110kabra1110
committed
add list reserve
1 parent b3add2a commit a649c9a

File tree

7 files changed

+149
-1
lines changed

7 files changed

+149
-1
lines changed

‎integration_tests/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -517,6 +517,7 @@ RUN(NAME test_list_pop2 LABELS cpython llvm NOFAST) # TODO: Remove NOFAST
517517
RUN(NAME test_list_pop3 LABELS cpython llvm)
518518
RUN(NAME test_list_compare LABELS cpython llvm)
519519
RUN(NAME test_list_concat LABELS cpython llvm c NOFAST)
520+
RUN(NAME test_list_reserve LABELS cpython llvm)
520521
RUN(NAME test_tuple_01 LABELS cpython llvm c)
521522
RUN(NAME test_tuple_02 LABELS cpython llvm c NOFAST)
522523
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
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+
l1.reserve(100)
9+
for i in range(50):
10+
l1.append(i)
11+
assert len(l1) == i + 1
12+
13+
l1.reserve(150)
14+
15+
for i in range(50):
16+
l1.pop(0)
17+
assert len(l1) == 49 - i
18+
19+
l2.reserve(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+
l2.reserve(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: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1680,6 +1680,20 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
16801680
tmp = list_api->pop_position(plist, pos, asr_el_type, module.get(), name2memidx);
16811681
}
16821682

1683+
void generate_ListReserve(ASR::expr_t* m_arg, ASR::expr_t* m_ele) {
1684+
ASR::ttype_t* asr_el_type = ASRUtils::get_contained_type(ASRUtils::expr_type(m_arg));
1685+
int64_t ptr_loads_copy = ptr_loads;
1686+
ptr_loads = 0;
1687+
this->visit_expr(*m_arg);
1688+
llvm::Value* plist = tmp;
1689+
1690+
ptr_loads = 2;
1691+
this->visit_expr_wrapper(m_ele, true);
1692+
ptr_loads = ptr_loads_copy;
1693+
llvm::Value* n = tmp;
1694+
list_api->reserve(plist, n, asr_el_type, module.get());
1695+
}
1696+
16831697
void generate_DictElems(ASR::expr_t* m_arg, bool key_or_value) {
16841698
ASR::Dict_t* dict_type = ASR::down_cast<ASR::Dict_t>(
16851699
ASRUtils::expr_type(m_arg));
@@ -1794,6 +1808,10 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
17941808
}
17951809
break;
17961810
}
1811+
case ASRUtils::IntrinsicFunctions::ListReserve: {
1812+
generate_ListReserve(x.m_args[0], x.m_args[1]);
1813+
break;
1814+
}
17971815
case ASRUtils::IntrinsicFunctions::DictKeys: {
17981816
generate_DictElems(x.m_args[0], 0);
17991817
break;

‎src/libasr/codegen/llvm_utils.cpp

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4105,6 +4105,33 @@ namespace LCompilers {
41054105
shift_end_point_by_one(list);
41064106
}
41074107

4108+
void LLVMList::reserve(llvm::Value* list, llvm::Value* n,
4109+
ASR::ttype_t* asr_type, llvm::Module* module) {
4110+
/**
4111+
* C++ equivalent
4112+
*
4113+
* if( n > current_capacity ) {
4114+
* list_data = realloc(list_data, sizeof(el_type) * n);
4115+
* }
4116+
*
4117+
*/
4118+
llvm::Value* capacity = LLVM::CreateLoad(*builder, get_pointer_to_current_capacity(list));
4119+
std::string type_code = ASRUtils::get_type_code(asr_type);
4120+
int type_size = std::get<1>(typecode2listtype[type_code]);
4121+
llvm::Type* el_type = std::get<2>(typecode2listtype[type_code]);
4122+
llvm_utils->create_if_else(builder->CreateICmpSGT(n, capacity), [&]() {
4123+
llvm::Value* arg_size = builder->CreateMul(llvm::ConstantInt::get(context,
4124+
llvm::APInt(32, type_size)), n);
4125+
llvm::Value* copy_data_ptr = get_pointer_to_list_data(list);
4126+
llvm::Value* copy_data = LLVM::CreateLoad(*builder, copy_data_ptr);
4127+
copy_data = LLVM::lfortran_realloc(context, *module, *builder,
4128+
copy_data, arg_size);
4129+
copy_data = builder->CreateBitCast(copy_data, el_type->getPointerTo());
4130+
builder->CreateStore(copy_data, copy_data_ptr);
4131+
builder->CreateStore(n, get_pointer_to_current_capacity(list));
4132+
}, []() {});
4133+
}
4134+
41084135
void LLVMList::reverse(llvm::Value* list, llvm::Module& module) {
41094136

41104137
/* Equivalent in C++:

‎src/libasr/codegen/llvm_utils.h

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

406+
void reserve(llvm::Value* list, llvm::Value* n,
407+
ASR::ttype_t* asr_type, llvm::Module* module);
408+
406409
void remove(llvm::Value* list, llvm::Value* item,
407410
ASR::ttype_t* item_type, llvm::Module& module);
408411

‎src/libasr/pass/intrinsic_function_registry.h

Lines changed: 53 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+
ListReserve,
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(ListReserve)
105107
INTRINSIC_NAME_CASE(DictKeys)
106108
INTRINSIC_NAME_CASE(DictValues)
107109
INTRINSIC_NAME_CASE(SetAdd)
@@ -1262,6 +1264,52 @@ static inline ASR::asr_t* create_ListPop(Allocator& al, const Location& loc,
12621264

12631265
} // namespace ListPop
12641266

1267+
namespace ListReserve {
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 list.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+
"Argument to list.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+
"Argument to list.reserve must be an integer",
1277+
x.base.base.loc, diagnostics);
1278+
ASRUtils::require_impl(x.m_type == nullptr,
1279+
"Return type of list.reserve must be empty",
1280+
x.base.base.loc, diagnostics);
1281+
}
1282+
1283+
static inline ASR::expr_t *eval_list_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_ListReserve(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 list.reserve must have exactly one argument", loc);
1294+
}
1295+
if (!ASR::is_a<ASR::Integer_t>(*ASRUtils::expr_type(args[1]))) {
1296+
err("Argument to list.reserve must be an integer", loc);
1297+
}
1298+
1299+
Vec<ASR::expr_t*> arg_values;
1300+
arg_values.reserve(al, args.size());
1301+
for( size_t i = 0; i < args.size(); i++ ) {
1302+
arg_values.push_back(al, ASRUtils::expr_value(args[i]));
1303+
}
1304+
ASR::expr_t* compile_time_value = eval_list_reserve(al, loc, arg_values);
1305+
return ASR::make_Expr_t(al, loc,
1306+
ASRUtils::EXPR(ASRUtils::make_IntrinsicFunction_t_util(al, loc,
1307+
static_cast<int64_t>(ASRUtils::IntrinsicFunctions::ListReserve),
1308+
args.p, args.size(), 0, nullptr, compile_time_value)));
1309+
}
1310+
1311+
} // namespace ListReserve
1312+
12651313
namespace DictKeys {
12661314

12671315
static inline void verify_args(const ASR::IntrinsicFunction_t& x, diag::Diagnostics& diagnostics) {
@@ -3124,6 +3172,8 @@ namespace IntrinsicFunctionRegistry {
31243172
{nullptr, &DictValues::verify_args}},
31253173
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::ListPop),
31263174
{nullptr, &ListPop::verify_args}},
3175+
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::ListReserve),
3176+
{nullptr, &ListReserve::verify_args}},
31273177
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::SetAdd),
31283178
{nullptr, &SetAdd::verify_args}},
31293179
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::SetRemove),
@@ -3206,6 +3256,8 @@ namespace IntrinsicFunctionRegistry {
32063256
"list.reverse"},
32073257
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::ListPop),
32083258
"list.pop"},
3259+
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::ListReserve),
3260+
"list.reserve"},
32093261
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::DictKeys),
32103262
"dict.keys"},
32113263
{static_cast<int64_t>(ASRUtils::IntrinsicFunctions::DictValues),
@@ -3290,6 +3342,7 @@ namespace IntrinsicFunctionRegistry {
32903342
{"list.index", {&ListIndex::create_ListIndex, &ListIndex::eval_list_index}},
32913343
{"list.reverse", {&ListReverse::create_ListReverse, &ListReverse::eval_list_reverse}},
32923344
{"list.pop", {&ListPop::create_ListPop, &ListPop::eval_list_pop}},
3345+
{"list.reserve", {&ListReserve::create_ListReserve, &ListReserve::eval_list_reserve}},
32933346
{"dict.keys", {&DictKeys::create_DictKeys, &DictKeys::eval_dict_keys}},
32943347
{"dict.values", {&DictValues::create_DictValues, &DictValues::eval_dict_values}},
32953348
{"set.add", {&SetAdd::create_SetAdd, &SetAdd::eval_set_add}},

‎src/lpython/semantics/python_attribute_eval.h

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ struct AttributeHandler {
3030
{"list@pop", &eval_list_pop},
3131
{"list@clear", &eval_list_clear},
3232
{"list@insert", &eval_list_insert},
33+
{"list@reserve", &eval_list_reserve},
3334
{"set@pop", &eval_set_pop},
3435
{"set@add", &eval_set_add},
3536
{"set@remove", &eval_set_remove},
@@ -41,7 +42,8 @@ struct AttributeHandler {
4142

4243
modify_attr_set = {"list@append", "list@remove",
4344
"list@reverse", "list@clear", "list@insert", "list@pop",
44-
"set@pop", "set@add", "set@remove", "dict@pop"};
45+
"list@reserve", "set@pop", "set@add", "set@remove",
46+
"dict@pop"};
4547

4648
symbolic_attribute_map = {
4749
{"diff", &eval_symbolic_diff},
@@ -284,6 +286,20 @@ struct AttributeHandler {
284286
return make_ListClear_t(al, loc, s);
285287
}
286288

289+
static ASR::asr_t* eval_list_reserve(ASR::expr_t *s, Allocator &al, const Location &loc,
290+
Vec<ASR::expr_t*> &args, diag::Diagnostics &/*diag*/) {
291+
Vec<ASR::expr_t*> args_with_list;
292+
args_with_list.reserve(al, args.size() + 1);
293+
args_with_list.push_back(al, s);
294+
for(size_t i = 0; i < args.size(); i++) {
295+
args_with_list.push_back(al, args[i]);
296+
}
297+
ASRUtils::create_intrinsic_function create_function =
298+
ASRUtils::IntrinsicFunctionRegistry::get_create_function("list.reserve");
299+
return create_function(al, loc, args_with_list, [&](const std::string &msg, const Location &loc)
300+
{ throw SemanticError(msg, loc); });
301+
}
302+
287303
static ASR::asr_t* eval_set_pop(ASR::expr_t *s, Allocator &al, const Location &loc,
288304
Vec<ASR::expr_t*> &args, diag::Diagnostics &/*diag*/) {
289305
if (args.size() != 0) {

0 commit comments

Comments
 (0)