Merge pull request lcompilers#2619 from hankluo6/array_size

Simplify array size determination for negative index
sob6262 · Mar 21, 2024 · 85ced05 · 85ced05
2 parents 1e5523f + 2e4057f
commit 85ced05
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 8 deletions.
diff --git a/integration_tests/CMakeLists.txt b/integration_tests/CMakeLists.txt
@@ -435,6 +435,7 @@ RUN(NAME array_expr_06            LABELS cpython llvm c)
 RUN(NAME array_expr_07            LABELS cpython llvm c)
 RUN(NAME array_expr_08            LABELS cpython llvm c)
 RUN(NAME array_expr_09            LABELS cpython llvm c)
+RUN(NAME array_expr_10            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)

diff --git a/integration_tests/array_expr_10.py b/integration_tests/array_expr_10.py
@@ -0,0 +1,17 @@
+from lpython import i32
+from numpy import empty, int32, array
+
+def foo(x: i32[:]):
+    print(x[3], x[4], x[-1], x[-2])
+    assert x[-1] == 5
+    assert x[-2] == 4
+    assert x[-3] == 3
+    assert x[-4] == 2
+    assert x[-5] == 1
+
+def main():
+    x: i32[5] = empty(5, dtype=int32)
+    x = array([1, 2, 3, 4, 5])
+    foo(x)
+
+main()
diff --git a/src/lpython/semantics/python_ast_to_asr.cpp b/src/lpython/semantics/python_ast_to_asr.cpp
@@ -3907,15 +3907,12 @@ class CommonVisitor : public AST::BaseVisitor<Struct> {
                         ASR::ttype_t *int_type = ASRUtils::TYPE(ASR::make_Integer_t(
                                                         al, loc, 4));
                         ASR::expr_t *neg_idx = ASRUtils::expr_value(index);
-                        ASR::expr_t *dim_size;
-                        if (ASRUtils::extract_physical_type(type) != ASR::array_physical_typeType::DescriptorArray)
-                            dim_size = ASR::down_cast<ASR::Array_t>(type)->m_dims[idx].m_length;
-                        else {
-                            ASR::expr_t *idx_expr = ASRUtils::EXPR(ASR::make_IntegerConstant_t(al, loc, idx + 1, int_type));
-                            dim_size = ASRUtils::EXPR(ASRUtils::make_ArraySize_t_util(al, loc, value, idx_expr, int_type, nullptr, false));
-                        }
+                        // null if the dimension is not known at compile time
+                        ASR::expr_t *dim_size = ASR::down_cast<ASR::Array_t>(type)->m_dims[idx].m_length;
+                        ASR::expr_t *idx_expr = ASRUtils::EXPR(ASR::make_IntegerConstant_t(al, loc, idx + 1, int_type));
+                        ASR::expr_t *size_expr = ASRUtils::EXPR(ASRUtils::make_ArraySize_t_util(al, loc, value, idx_expr, int_type, dim_size, false));
                         index = ASRUtils::EXPR(ASR::make_IntegerBinOp_t(al, loc, 
-                            dim_size, ASR::binopType::Add, neg_idx, int_type, nullptr));
+                            size_expr, ASR::binopType::Add, neg_idx, int_type, nullptr));
                     }
                 }
             } else {