diff --git a/integration_tests/CMakeLists.txt b/integration_tests/CMakeLists.txt index 39d72d6154..9c1f05791e 100644 --- a/integration_tests/CMakeLists.txt +++ b/integration_tests/CMakeLists.txt @@ -281,6 +281,7 @@ RUN(NAME variable_decl_03 LABELS cpython llvm c) RUN(NAME array_expr_01 LABELS cpython llvm c) RUN(NAME array_expr_02 LABELS cpython llvm c) RUN(NAME array_size_01 LABELS cpython llvm c) +RUN(NAME array_size_02 LABELS cpython llvm c) RUN(NAME array_01 LABELS cpython llvm wasm c) RUN(NAME array_02 LABELS cpython wasm c) RUN(NAME bindc_01 LABELS cpython llvm c) diff --git a/integration_tests/array_size_02.py b/integration_tests/array_size_02.py new file mode 100644 index 0000000000..07bd7af0d7 --- /dev/null +++ b/integration_tests/array_size_02.py @@ -0,0 +1,90 @@ +from lpython import i32, f64, c32, c64, u32 +from numpy import empty, size + +def main0(): + x: i32[4, 5, 2] = empty([4, 5, 2]) + y: f64[24, 100, 2, 5] = empty([24, 100, 2, 5]) + z: i32 + w: i32 + z = 2 + w = 3 + print(size(x)) + print(size(x, 0)) + print(size(x, 1)) + print(size(x, 2)) + print(size(y)) + print(size(y, 0)) + print(size(y, 1)) + print(size(y, z)) + print(size(y, w)) + + assert size(x) == 40 + assert size(x, 0) == 4 + assert size(x, 1) == 5 + assert size(x, 2) == 2 + assert size(y) == 24000 + assert size(y, 0) == 24 + assert size(y, 1) == 100 + assert size(y, z) == 2 + assert size(y, w) == 5 + +def main1(): + a: c32[12] = empty([12]) + b: c64[15, 15, 10] = empty([15, 15, 10]) + c: i32 + d: i32 + c = 1 + d = 2 + print(size(a)) + print(size(a, 0)) + print(size(b)) + print(size(b, 0)) + print(size(b, c)) + print(size(b, d)) + + assert size(a) == 12 + assert size(a, 0) == 12 + assert size(b) == 2250 + assert size(b, 0) == 15 + assert size(b, c) == 15 + assert size(b, d) == 10 + +def main2(): + a: i32[2, 3] = empty([2, 3]) + print(size(a)) + print(size(a, 0)) + print(size(a, 1)) + + assert size(a) == 2*3 + assert size(a, 0) == 2 + assert size(a, 1) == 3 + +def main3(): + a: u32[2, 3, 4] = empty([2, 3, 4]) + b: u64[10, 5] = empty([10, 5]) + c: i32 + d: i32 + c = 1 + d = 2 + print(size(a)) + print(size(a, 0)) + print(size(a, c)) + print(size(a, d)) + + print(size(b)) + print(size(b, 0)) + print(size(b, c)) + + assert size(a) == 2*3*4 + assert size(a, 0) == 2 + assert size(a, c) == 3 + assert size(a, d) == 4 + + assert size(b) == 50 + assert size(b, 0) == 10 + assert size(b, c) == 5 + +main0() +main1() +main2() +main3() diff --git a/src/lpython/semantics/python_ast_to_asr.cpp b/src/lpython/semantics/python_ast_to_asr.cpp index c954ef6405..8134019ef9 100644 --- a/src/lpython/semantics/python_ast_to_asr.cpp +++ b/src/lpython/semantics/python_ast_to_asr.cpp @@ -6749,13 +6749,16 @@ class BodyVisitor : public CommonVisitor { std::to_string(args.size()) + " arguments instead.", x.base.base.loc); } + const Location &loc = x.base.base.loc; ASR::expr_t *var = args[0].m_value; ASR::expr_t *dim = nullptr; + ASR::ttype_t *int_type = ASRUtils::TYPE(ASR::make_Integer_t(al, loc, 4, nullptr, 0)); if (args.size() == 2) { - dim = args[1].m_value; + ASR::expr_t* const_one = ASRUtils::EXPR(make_IntegerConstant_t(al, loc, 1, int_type)); + dim = ASRUtils::EXPR(ASR::make_IntegerBinOp_t(al, loc, + args[1].m_value, ASR::binopType::Add, const_one, int_type, nullptr)); } - ASR::ttype_t *int_type = ASRUtils::TYPE(ASR::make_Integer_t(al, x.base.base.loc, 4, nullptr, 0)); - tmp = ASR::make_ArraySize_t(al, x.base.base.loc, var, dim, int_type, nullptr); + tmp = ASR::make_ArraySize_t(al, loc, var, dim, int_type, nullptr); return; } else if (call_name == "empty") { // TODO: check that the `empty` arguments are compatible