Skip to content

Commit e487d07

Browse files
Smit-createSmit-create
authored
Merge pull request #1815 from Smit-create/i-1785
Support allocatable in lpython
2 parents 04bf030 + e15eed4 commit e487d07

File tree

4 files changed

+145
-27
lines changed

4 files changed

+145
-27
lines changed

‎integration_tests/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,6 +284,7 @@ RUN(NAME array_size_01 LABELS cpython llvm c)
284284
RUN(NAME array_size_02 LABELS cpython llvm c)
285285
RUN(NAME array_01 LABELS cpython llvm wasm c)
286286
RUN(NAME array_02 LABELS cpython wasm c)
287+
RUN(NAME array_03 LABELS cpython llvm)
287288
RUN(NAME bindc_01 LABELS cpython llvm c)
288289
RUN(NAME bindc_02 LABELS cpython llvm c)
289290
RUN(NAME bindc_04 LABELS llvm c)

‎integration_tests/array_03.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
from lpython import Allocatable, f64, i32
2+
from numpy import empty, float64, int32
3+
4+
def f():
5+
n: i32 = 5
6+
a: Allocatable[f64[:]] = empty((n,), dtype=float64)
7+
i: i32
8+
for i in range(n):
9+
a[i] = f64(i+1)
10+
b: Allocatable[i32[:]]
11+
n = 10
12+
b = empty((n,), dtype=int32)
13+
for i in range(n):
14+
b[i] = i+1
15+
print(a)
16+
for i in range(n):
17+
assert b[i] == i+1
18+
19+
f()

‎src/lpython/semantics/python_ast_to_asr.cpp

Lines changed: 121 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1029,7 +1029,8 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
10291029
} else if (var_annotation == "pointer") {
10301030
LCOMPILERS_ASSERT(n_args == 1);
10311031
AST::expr_t* underlying_type = m_args[0];
1032-
type = ast_expr_to_asr_type(underlying_type->base.loc, *underlying_type);
1032+
bool is_allocatable = false;
1033+
type = ast_expr_to_asr_type(underlying_type->base.loc, *underlying_type, is_allocatable);
10331034
type = ASRUtils::TYPE(ASR::make_Pointer_t(al, loc, type));
10341035
} else {
10351036
ASR::symbol_t *s = current_scope->resolve_symbol(var_annotation);
@@ -1609,7 +1610,7 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
16091610
// i32, i64, f32, f64
16101611
// f64[256], i32[:]
16111612
ASR::ttype_t * ast_expr_to_asr_type(const Location &loc, const AST::expr_t &annotation,
1612-
bool raise_error=true) {
1613+
bool &is_allocatable, bool raise_error=true) {
16131614
Vec<ASR::dimension_t> dims;
16141615
dims.reserve(al, 4);
16151616
AST::expr_t** m_args = nullptr; size_t n_args = 0;
@@ -1632,11 +1633,11 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
16321633
Vec<ASR::ttype_t*> types;
16331634
types.reserve(al, 4);
16341635
if (AST::is_a<AST::Name_t>(*s->m_slice)) {
1635-
types.push_back(al, ast_expr_to_asr_type(loc, *s->m_slice));
1636+
types.push_back(al, ast_expr_to_asr_type(loc, *s->m_slice, is_allocatable));
16361637
} else if (AST::is_a<AST::Tuple_t>(*s->m_slice)) {
16371638
AST::Tuple_t *t = AST::down_cast<AST::Tuple_t>(s->m_slice);
16381639
for (size_t i=0; i<t->n_elts; i++) {
1639-
types.push_back(al, ast_expr_to_asr_type(loc, *t->m_elts[i]));
1640+
types.push_back(al, ast_expr_to_asr_type(loc, *t->m_elts[i], is_allocatable));
16401641
}
16411642
} else {
16421643
throw SemanticError("Only Name or Tuple in Subscript supported for now in `tuple` annotation",
@@ -1656,14 +1657,15 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
16561657
if (arg_list->n_elts > 0) {
16571658
arg_types.reserve(al, arg_list->n_elts);
16581659
for (size_t i=0; i<arg_list->n_elts; i++) {
1659-
arg_types.push_back(al, ast_expr_to_asr_type(loc, *arg_list->m_elts[i]));
1660+
arg_types.push_back(al, ast_expr_to_asr_type(loc, *arg_list->m_elts[i],
1661+
is_allocatable));
16601662
}
16611663
} else {
16621664
arg_types.reserve(al, 1);
16631665
}
16641666
ASR::ttype_t* ret_type = nullptr;
16651667
if (t->n_elts == 2) {
1666-
ret_type = ast_expr_to_asr_type(loc, *t->m_elts[1]);
1668+
ret_type = ast_expr_to_asr_type(loc, *t->m_elts[1], is_allocatable);
16671669
}
16681670
ASR::ttype_t *type = ASRUtils::TYPE(ASR::make_FunctionType_t(al, loc, arg_types.p,
16691671
arg_types.size(), ret_type, ASR::abiType::Source,
@@ -1672,7 +1674,7 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
16721674
return type;
16731675
} else if (var_annotation == "set") {
16741676
if (AST::is_a<AST::Name_t>(*s->m_slice)) {
1675-
ASR::ttype_t *type = ast_expr_to_asr_type(loc, *s->m_slice);
1677+
ASR::ttype_t *type = ast_expr_to_asr_type(loc, *s->m_slice, is_allocatable);
16761678
return ASRUtils::TYPE(ASR::make_Set_t(al, loc, type));
16771679
} else {
16781680
throw SemanticError("Only Name in Subscript supported for now in `set`"
@@ -1681,32 +1683,42 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
16811683
} else if (var_annotation == "list") {
16821684
ASR::ttype_t *type = nullptr;
16831685
if (AST::is_a<AST::Name_t>(*s->m_slice) || AST::is_a<AST::Subscript_t>(*s->m_slice)) {
1684-
type = ast_expr_to_asr_type(loc, *s->m_slice);
1686+
type = ast_expr_to_asr_type(loc, *s->m_slice, is_allocatable);
16851687
return ASRUtils::TYPE(ASR::make_List_t(al, loc, type));
16861688
} else {
16871689
throw SemanticError("Only Name or Subscript inside Subscript supported for now in `list`"
16881690
" annotation", loc);
16891691
}
1692+
} else if (var_annotation == "Allocatable") {
1693+
ASR::ttype_t *type = nullptr;
1694+
if (AST::is_a<AST::Name_t>(*s->m_slice) || AST::is_a<AST::Subscript_t>(*s->m_slice)) {
1695+
type = ast_expr_to_asr_type(loc, *s->m_slice, is_allocatable);
1696+
is_allocatable = true;
1697+
return type;
1698+
} else {
1699+
throw SemanticError("Only Name or Subscript inside Subscript supported for now in `list`"
1700+
" annotation", loc);
1701+
}
16901702
} else if (var_annotation == "dict") {
16911703
if (AST::is_a<AST::Tuple_t>(*s->m_slice)) {
16921704
AST::Tuple_t *t = AST::down_cast<AST::Tuple_t>(s->m_slice);
16931705
if (t->n_elts != 2) {
16941706
throw SemanticError("`dict` annotation must have 2 elements: types"
16951707
" of both keys and values", loc);
16961708
}
1697-
ASR::ttype_t *key_type = ast_expr_to_asr_type(loc, *t->m_elts[0]);
1698-
ASR::ttype_t *value_type = ast_expr_to_asr_type(loc, *t->m_elts[1]);
1709+
ASR::ttype_t *key_type = ast_expr_to_asr_type(loc, *t->m_elts[0], is_allocatable);
1710+
ASR::ttype_t *value_type = ast_expr_to_asr_type(loc, *t->m_elts[1], is_allocatable);
16991711
raise_error_when_dict_key_is_float_or_complex(key_type, loc);
17001712
return ASRUtils::TYPE(ASR::make_Dict_t(al, loc, key_type, value_type));
17011713
} else {
17021714
throw SemanticError("`dict` annotation must have 2 elements: types of"
17031715
" both keys and values", loc);
17041716
}
17051717
} else if (var_annotation == "Pointer") {
1706-
ASR::ttype_t *type = ast_expr_to_asr_type(loc, *s->m_slice);
1718+
ASR::ttype_t *type = ast_expr_to_asr_type(loc, *s->m_slice, is_allocatable);
17071719
return ASRUtils::TYPE(ASR::make_Pointer_t(al, loc, type));
17081720
} else if (var_annotation == "Const") {
1709-
ASR::ttype_t *type = ast_expr_to_asr_type(loc, *s->m_slice);
1721+
ASR::ttype_t *type = ast_expr_to_asr_type(loc, *s->m_slice, is_allocatable);
17101722
return ASRUtils::TYPE(ASR::make_Const_t(al, loc, type));
17111723
} else {
17121724
if (AST::is_a<AST::Slice_t>(*s->m_slice)) {
@@ -2446,7 +2458,8 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
24462458
}
24472459

24482460
void create_add_variable_to_scope(std::string& var_name, ASR::expr_t* init_expr,
2449-
ASR::ttype_t* type, const Location& loc, ASR::abiType abi) {
2461+
ASR::ttype_t* type, const Location& loc, ASR::abiType abi,
2462+
ASR::storage_typeType storage_type=ASR::storage_typeType::Default) {
24502463

24512464
ASR::expr_t* value = nullptr;
24522465
if( init_expr ) {
@@ -2461,8 +2474,6 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
24612474
}
24622475

24632476
ASR::intentType s_intent = ASRUtils::intent_local;
2464-
ASR::storage_typeType storage_type =
2465-
ASR::storage_typeType::Default;
24662477
if( ASR::is_a<ASR::Const_t>(*type) ) {
24672478
storage_type = ASR::storage_typeType::Parameter;
24682479
}
@@ -2548,7 +2559,8 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
25482559
ASR::expr_t* cptr = ASRUtils::EXPR(tmp);
25492560
this->visit_expr(*ast_pptr);
25502561
ASR::expr_t* pptr = ASRUtils::EXPR(tmp);
2551-
ASR::ttype_t* asr_alloc_type = ast_expr_to_asr_type(ast_type_expr->base.loc, *ast_type_expr);
2562+
bool is_allocatable = false;
2563+
ASR::ttype_t* asr_alloc_type = ast_expr_to_asr_type(ast_type_expr->base.loc, *ast_type_expr, is_allocatable);
25522564
ASR::ttype_t* target_type = ASRUtils::type_get_past_pointer(ASRUtils::expr_type(pptr));
25532565
if( !ASRUtils::types_equal(target_type, asr_alloc_type, true) ) {
25542566
diag.add(diag::Diagnostic(
@@ -2565,17 +2577,70 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
25652577
return ASR::make_CPtrToPointer_t(al, loc, cptr, pptr, target_shape, lower_bounds);
25662578
}
25672579

2580+
ASR::asr_t* check_to_allocate_array(AST::expr_t *value, std::string var_name,
2581+
const Location &loc) {
2582+
if (AST::is_a<AST::Call_t>(*value)) {
2583+
AST::Call_t *ct = AST::down_cast<AST::Call_t>(value);
2584+
if (AST::is_a<AST::Name_t>(*ct->m_func)) {
2585+
std::string call_name = AST::down_cast<AST::Name_t>(ct->m_func)->m_id;
2586+
if (call_name == "empty") {
2587+
LCOMPILERS_ASSERT(ct->n_args > 0);
2588+
if (AST::is_a<AST::Tuple_t>(*ct->m_args[0])) {
2589+
AST::Tuple_t *tt = AST::down_cast<AST::Tuple_t>(ct->m_args[0]);
2590+
Vec<ASR::alloc_arg_t> alloc_args_vec;
2591+
alloc_args_vec.reserve(al, 1);
2592+
ASR::alloc_arg_t new_arg;
2593+
new_arg.loc = loc;
2594+
ASR::ttype_t *int32_type = ASRUtils::TYPE(ASR::make_Integer_t(al, loc,
2595+
4, nullptr, 0));
2596+
ASR::expr_t* const_0 = ASRUtils::EXPR(ASR::make_IntegerConstant_t(al,
2597+
loc, 0, int32_type));
2598+
Vec<ASR::dimension_t> dims_vec;
2599+
dims_vec.reserve(al, tt->n_elts);
2600+
for (size_t i=0; i<tt->n_elts; i++) {
2601+
ASR::dimension_t new_dim;
2602+
new_dim.loc = loc;
2603+
this->visit_expr(*tt->m_elts[0]);
2604+
new_dim.m_start = const_0;
2605+
new_dim.m_length = ASRUtils::EXPR(tmp);
2606+
dims_vec.push_back(al, new_dim);
2607+
}
2608+
new_arg.m_dims = dims_vec.p;
2609+
new_arg.n_dims = dims_vec.size();
2610+
ASR::symbol_t *v_sym = current_scope->resolve_symbol(var_name);
2611+
ASR::expr_t* v_expr = ASRUtils::EXPR(ASR::make_Var_t(al,
2612+
loc, v_sym));
2613+
new_arg.m_a = v_expr;
2614+
alloc_args_vec.push_back(al, new_arg);
2615+
tmp = ASR::make_Allocate_t(al, loc,
2616+
alloc_args_vec.p, alloc_args_vec.size(),
2617+
nullptr, nullptr, nullptr);
2618+
return tmp;
2619+
} else {
2620+
throw SemanticError("Only tuple argument is accepted as dimensions "
2621+
"for allocating using empty()", ct->base.base.loc);
2622+
}
2623+
}
2624+
}
2625+
}
2626+
return nullptr;
2627+
}
2628+
25682629
void visit_AnnAssignUtil(const AST::AnnAssign_t& x, std::string& var_name,
25692630
bool wrap_derived_type_in_pointer=false,
25702631
ASR::expr_t* init_expr=nullptr, ASR::abiType abi=ASR::abiType::Source) {
2571-
ASR::ttype_t *type = ast_expr_to_asr_type(x.base.base.loc, *x.m_annotation);
2632+
bool is_allocatable = false;
2633+
ASR::ttype_t *type = ast_expr_to_asr_type(x.base.base.loc, *x.m_annotation, is_allocatable);
25722634
ASR::ttype_t* ann_assign_target_type_copy = ann_assign_target_type;
25732635
ann_assign_target_type = type;
25742636
if( ASR::is_a<ASR::Struct_t>(*type) &&
25752637
wrap_derived_type_in_pointer ) {
25762638
type = ASRUtils::TYPE(ASR::make_Pointer_t(al, type->base.loc, type));
25772639
}
2578-
2640+
ASR::storage_typeType storage_type = ASR::storage_typeType::Default;
2641+
if (is_allocatable) {
2642+
storage_type = ASR::storage_typeType::Allocatable;
2643+
}
25792644
bool is_c_p_pointer_call_copy = is_c_p_pointer_call;
25802645
ASR::expr_t *value = nullptr;
25812646
if( !init_expr ) {
@@ -2591,7 +2656,7 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
25912656
}
25922657
if( is_c_p_pointer_call ) {
25932658
create_add_variable_to_scope(var_name, nullptr, type,
2594-
x.base.base.loc, abi);
2659+
x.base.base.loc, abi, storage_type);
25952660
AST::Call_t* c_p_pointer_call = AST::down_cast<AST::Call_t>(x.m_value);
25962661
AST::expr_t* cptr = c_p_pointer_call->m_args[0];
25972662
AST::expr_t* pptr = assign_ast_target;
@@ -2627,7 +2692,14 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
26272692

26282693
if( !is_c_p_pointer_call ) {
26292694
create_add_variable_to_scope(var_name, init_expr, type,
2630-
x.base.base.loc, abi);
2695+
x.base.base.loc, abi, storage_type);
2696+
}
2697+
2698+
if (is_allocatable && x.m_value && AST::is_a<AST::Call_t>(*x.m_value)) {
2699+
tmp = check_to_allocate_array(x.m_value, var_name, x.base.base.loc);
2700+
if( current_body && tmp) {
2701+
current_body->push_back(al, ASRUtils::STMT(tmp));
2702+
}
26312703
}
26322704

26332705
if( !is_c_p_pointer_call ) {
@@ -3786,7 +3858,7 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
37863858
if (parent_scope->get_scope().find(sym_name) != parent_scope->get_scope().end()) {
37873859
throw SemanticError("Function " + std::string(x.m_name) + " is already defined", x.base.base.loc);
37883860
}
3789-
3861+
bool is_allocatable = false;
37903862
for (size_t i=0; i<x.m_args.n_args; i++) {
37913863
char *arg=x.m_args.m_args[i].m_arg;
37923864
Location loc = x.m_args.m_args[i].loc;
@@ -3795,7 +3867,8 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
37953867
}
37963868
ASR::intentType s_intent = ASRUtils::intent_unspecified;
37973869
AST::expr_t* arg_annotation_type = get_var_intent_and_annotation(x.m_args.m_args[i].m_annotation, s_intent);
3798-
ASR::ttype_t *arg_type = ast_expr_to_asr_type(x.base.base.loc, *arg_annotation_type);
3870+
is_allocatable = false;
3871+
ASR::ttype_t *arg_type = ast_expr_to_asr_type(x.base.base.loc, *arg_annotation_type, is_allocatable);
37993872
// Set the function as generic if an argument is typed with a type parameter
38003873
if (ASRUtils::is_generic(*arg_type)) {
38013874
ASR::ttype_t* arg_type_type = ASRUtils::get_type_parameter(arg_type);
@@ -3833,6 +3906,9 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
38333906
if( ASR::is_a<ASR::Const_t>(*arg_type) ) {
38343907
storage_type = ASR::storage_typeType::Parameter;
38353908
}
3909+
if (is_allocatable) {
3910+
storage_type = ASR::storage_typeType::Allocatable;
3911+
}
38363912
ASR::accessType s_access = ASR::accessType::Public;
38373913
ASR::presenceType s_presence = ASR::presenceType::Required;
38383914
bool value_attr = false;
@@ -3869,8 +3945,12 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
38693945
if (x.m_returns && !AST::is_a<AST::ConstantNone_t>(*x.m_returns)) {
38703946
if (AST::is_a<AST::Name_t>(*x.m_returns) || AST::is_a<AST::Subscript_t>(*x.m_returns)) {
38713947
std::string return_var_name = "_lpython_return_variable";
3872-
ASR::ttype_t *type = ast_expr_to_asr_type(x.m_returns->base.loc, *x.m_returns);
3948+
is_allocatable = false;
3949+
ASR::ttype_t *type = ast_expr_to_asr_type(x.m_returns->base.loc, *x.m_returns, is_allocatable);
38733950
ASR::storage_typeType storage_type = ASR::storage_typeType::Default;
3951+
if (is_allocatable) {
3952+
storage_type = ASR::storage_typeType::Allocatable;
3953+
}
38743954
ASR::ttype_t* return_type_ = type;
38753955
if( ASR::is_a<ASR::Const_t>(*type) ) {
38763956
return_type_ = ASR::down_cast<ASR::Const_t>(type)->m_type;
@@ -4706,6 +4786,19 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
47064786
tmp_value = ASRUtils::EXPR(ASR::make_ListConstant_t(al, x.base.base.loc, list_ele.p,
47074787
list_ele.size(), target_type));
47084788
}
4789+
if (tmp_value == nullptr && ASR::is_a<ASR::Var_t>(*target)) {
4790+
ASR::Var_t *var_tar = ASR::down_cast<ASR::Var_t>(target);
4791+
if (ASR::is_a<ASR::Variable_t>(*var_tar->m_v)) {
4792+
if (ASR::down_cast<ASR::Variable_t>(var_tar->m_v)->m_storage == ASR::storage_typeType::Allocatable) {
4793+
ASR::asr_t *st = check_to_allocate_array(x.m_value, ASRUtils::symbol_name(var_tar->m_v),
4794+
x.base.base.loc);
4795+
if (st) {
4796+
tmp_vec.push_back(st);
4797+
continue;
4798+
}
4799+
}
4800+
}
4801+
}
47094802
if (!tmp_value) continue;
47104803
ASR::ttype_t *value_type = ASRUtils::expr_type(tmp_value);
47114804
if( ASR::is_a<ASR::Pointer_t>(*target_type) &&
@@ -6167,7 +6260,8 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
61676260
ASR::expr_t* cptr = ASRUtils::EXPR(tmp);
61686261
visit_expr(*x.m_args[1]);
61696262
ASR::expr_t* pptr = ASRUtils::EXPR(tmp);
6170-
ASR::ttype_t* asr_alloc_type = ast_expr_to_asr_type(x.m_args[1]->base.loc, *x.m_args[1]);
6263+
bool is_allocatable = false;
6264+
ASR::ttype_t* asr_alloc_type = ast_expr_to_asr_type(x.m_args[1]->base.loc, *x.m_args[1], is_allocatable);
61716265
ASR::ttype_t* target_type = ASRUtils::type_get_past_pointer(ASRUtils::expr_type(pptr));
61726266
if( !ASRUtils::types_equal(target_type, asr_alloc_type, true) ) {
61736267
diag.add(diag::Diagnostic(
@@ -7015,8 +7109,9 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
70157109
std::to_string(x.n_args + x.n_keywords) + " instead.",
70167110
x.base.base.loc);
70177111
}
7018-
7019-
ASR::ttype_t* arg_type = ast_expr_to_asr_type(x.base.base.loc, *x.m_args[0], false);
7112+
bool is_allocatable = false;
7113+
ASR::ttype_t* arg_type = ast_expr_to_asr_type(x.base.base.loc, *x.m_args[0],
7114+
is_allocatable, false);
70207115
ASR::expr_t* arg = nullptr;
70217116
if( !arg_type ) {
70227117
visit_expr(*x.m_args[0]);

0 commit comments

Comments
 (0)