Skip to content

Commit 9f8968f

Browse files
Smit-createSmit-create
committed
Support allocatable in lpython
1 parent 04bf030 commit 9f8968f

File tree

1 file changed

+52
-26
lines changed

1 file changed

+52
-26
lines changed

‎src/lpython/semantics/python_ast_to_asr.cpp

Lines changed: 52 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(
@@ -2568,14 +2580,18 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
25682580
void visit_AnnAssignUtil(const AST::AnnAssign_t& x, std::string& var_name,
25692581
bool wrap_derived_type_in_pointer=false,
25702582
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);
2583+
bool is_allocatable = false;
2584+
ASR::ttype_t *type = ast_expr_to_asr_type(x.base.base.loc, *x.m_annotation, is_allocatable);
25722585
ASR::ttype_t* ann_assign_target_type_copy = ann_assign_target_type;
25732586
ann_assign_target_type = type;
25742587
if( ASR::is_a<ASR::Struct_t>(*type) &&
25752588
wrap_derived_type_in_pointer ) {
25762589
type = ASRUtils::TYPE(ASR::make_Pointer_t(al, type->base.loc, type));
25772590
}
2578-
2591+
ASR::storage_typeType storage_type = ASR::storage_typeType::Default;
2592+
if (is_allocatable) {
2593+
storage_type = ASR::storage_typeType::Allocatable;
2594+
}
25792595
bool is_c_p_pointer_call_copy = is_c_p_pointer_call;
25802596
ASR::expr_t *value = nullptr;
25812597
if( !init_expr ) {
@@ -2591,7 +2607,7 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
25912607
}
25922608
if( is_c_p_pointer_call ) {
25932609
create_add_variable_to_scope(var_name, nullptr, type,
2594-
x.base.base.loc, abi);
2610+
x.base.base.loc, abi, storage_type);
25952611
AST::Call_t* c_p_pointer_call = AST::down_cast<AST::Call_t>(x.m_value);
25962612
AST::expr_t* cptr = c_p_pointer_call->m_args[0];
25972613
AST::expr_t* pptr = assign_ast_target;
@@ -2627,7 +2643,7 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
26272643

26282644
if( !is_c_p_pointer_call ) {
26292645
create_add_variable_to_scope(var_name, init_expr, type,
2630-
x.base.base.loc, abi);
2646+
x.base.base.loc, abi, storage_type);
26312647
}
26322648

26332649
if( !is_c_p_pointer_call ) {
@@ -3786,7 +3802,7 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
37863802
if (parent_scope->get_scope().find(sym_name) != parent_scope->get_scope().end()) {
37873803
throw SemanticError("Function " + std::string(x.m_name) + " is already defined", x.base.base.loc);
37883804
}
3789-
3805+
bool is_allocatable = false;
37903806
for (size_t i=0; i<x.m_args.n_args; i++) {
37913807
char *arg=x.m_args.m_args[i].m_arg;
37923808
Location loc = x.m_args.m_args[i].loc;
@@ -3795,7 +3811,8 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
37953811
}
37963812
ASR::intentType s_intent = ASRUtils::intent_unspecified;
37973813
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);
3814+
is_allocatable = false;
3815+
ASR::ttype_t *arg_type = ast_expr_to_asr_type(x.base.base.loc, *arg_annotation_type, is_allocatable);
37993816
// Set the function as generic if an argument is typed with a type parameter
38003817
if (ASRUtils::is_generic(*arg_type)) {
38013818
ASR::ttype_t* arg_type_type = ASRUtils::get_type_parameter(arg_type);
@@ -3833,6 +3850,9 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
38333850
if( ASR::is_a<ASR::Const_t>(*arg_type) ) {
38343851
storage_type = ASR::storage_typeType::Parameter;
38353852
}
3853+
if (is_allocatable) {
3854+
storage_type = ASR::storage_typeType::Allocatable;
3855+
}
38363856
ASR::accessType s_access = ASR::accessType::Public;
38373857
ASR::presenceType s_presence = ASR::presenceType::Required;
38383858
bool value_attr = false;
@@ -3869,8 +3889,12 @@ class SymbolTableVisitor : public CommonVisitor<SymbolTableVisitor> {
38693889
if (x.m_returns && !AST::is_a<AST::ConstantNone_t>(*x.m_returns)) {
38703890
if (AST::is_a<AST::Name_t>(*x.m_returns) || AST::is_a<AST::Subscript_t>(*x.m_returns)) {
38713891
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);
3892+
is_allocatable = false;
3893+
ASR::ttype_t *type = ast_expr_to_asr_type(x.m_returns->base.loc, *x.m_returns, is_allocatable);
38733894
ASR::storage_typeType storage_type = ASR::storage_typeType::Default;
3895+
if (is_allocatable) {
3896+
storage_type = ASR::storage_typeType::Allocatable;
3897+
}
38743898
ASR::ttype_t* return_type_ = type;
38753899
if( ASR::is_a<ASR::Const_t>(*type) ) {
38763900
return_type_ = ASR::down_cast<ASR::Const_t>(type)->m_type;
@@ -6167,7 +6191,8 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
61676191
ASR::expr_t* cptr = ASRUtils::EXPR(tmp);
61686192
visit_expr(*x.m_args[1]);
61696193
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]);
6194+
bool is_allocatable = false;
6195+
ASR::ttype_t* asr_alloc_type = ast_expr_to_asr_type(x.m_args[1]->base.loc, *x.m_args[1], is_allocatable);
61716196
ASR::ttype_t* target_type = ASRUtils::type_get_past_pointer(ASRUtils::expr_type(pptr));
61726197
if( !ASRUtils::types_equal(target_type, asr_alloc_type, true) ) {
61736198
diag.add(diag::Diagnostic(
@@ -7015,8 +7040,9 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
70157040
std::to_string(x.n_args + x.n_keywords) + " instead.",
70167041
x.base.base.loc);
70177042
}
7018-
7019-
ASR::ttype_t* arg_type = ast_expr_to_asr_type(x.base.base.loc, *x.m_args[0], false);
7043+
bool is_allocatable = false;
7044+
ASR::ttype_t* arg_type = ast_expr_to_asr_type(x.base.base.loc, *x.m_args[0],
7045+
is_allocatable, false);
70207046
ASR::expr_t* arg = nullptr;
70217047
if( !arg_type ) {
70227048
visit_expr(*x.m_args[0]);

0 commit comments

Comments
 (0)