Merge pull request #2157 from Shaikh-Ubaid/support_arr_comp

Support array comparison

Author: ubaidsk

integration_tests/CMakeLists.txt

@@ -368,6 +368,7 @@ 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 array_03            LABELS cpython llvm c)
+RUN(NAME array_04            LABELS cpython llvm c)
 RUN(NAME bindc_01            LABELS cpython llvm c)
 RUN(NAME bindc_02            LABELS cpython llvm c)
 RUN(NAME bindc_04            LABELS llvm c NOFAST)

integration_tests/array_04.py

@@ -0,0 +1,19 @@
+from lpython import i32, Const
+from numpy import empty, int32
+
+def main0():
+    n: Const[i32] = 1
+    x: i32[n, n] = empty([n, n], dtype=int32)
+    y: i32[n, n] = empty([n, n], dtype=int32)
+
+    x[0, 0] = -10
+    y[0, 0] = -10
+
+    print(x[0, 0], y[0, 0])
+    assert x == y
+
+    y[0, 0] = 10
+    print(x[0, 0], y[0, 0])
+    assert x != y
+
+main0()

src/libasr/asr_utils.h

@@ -1868,6 +1868,12 @@ static inline int64_t get_fixed_size_of_array(ASR::dimension_t* m_dims, size_t n
     return array_size;
 }
 
+static inline int64_t get_fixed_size_of_array(ASR::ttype_t* type) {
+    ASR::dimension_t* m_dims = nullptr;
+    size_t n_dims = ASRUtils::extract_dimensions_from_ttype(type, m_dims);
+    return ASRUtils::get_fixed_size_of_array(m_dims, n_dims);
+}
+
 inline int extract_n_dims_from_ttype(ASR::ttype_t *x) {
     ASR::dimension_t* m_dims_temp = nullptr;
     return extract_dimensions_from_ttype(x, m_dims_temp);

src/lpython/semantics/python_ast_to_asr.cpp

@@ -6157,6 +6157,37 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
         if( ASR::is_a<ASR::Enum_t>(*dest_type) || ASR::is_a<ASR::Const_t>(*dest_type) ) {
             dest_type = ASRUtils::get_contained_type(dest_type);
         }
+
+        if (ASRUtils::is_array(dest_type)) {
+            ASR::dimension_t* m_dims = nullptr;
+            int n_dims = ASRUtils::extract_dimensions_from_ttype(dest_type, m_dims);
+            int array_size = ASRUtils::get_fixed_size_of_array(m_dims, n_dims);
+            if (array_size == -1) {
+                throw SemanticError("The truth value of an array is ambiguous. Use a.any() or a.all()", x.base.base.loc);
+            } else if (array_size != 1) {
+                throw SemanticError("The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()", x.base.base.loc);
+            } else {
+                Vec<ASR::array_index_t> argsL, argsR;
+                argsL.reserve(al, 1);
+                argsR.reserve(al, 1);
+                for (int i = 0; i < n_dims; i++) {
+                    ASR::array_index_t aiL, aiR;
+                    ASR::ttype_t *int_type = ASRUtils::TYPE(ASR::make_Integer_t(al, x.base.base.loc, 4));
+                    ASR::expr_t* const_zero = ASRUtils::EXPR(ASR::make_IntegerConstant_t(al, x.base.base.loc, 0, int_type));
+                    aiL.m_right = aiR.m_right = const_zero;
+                    aiL.m_left = aiR.m_left = nullptr;
+                    aiL.m_step = aiR.m_step = nullptr;
+                    aiL.loc = left->base.loc;
+                    aiR.loc = right->base.loc;
+                    argsL.push_back(al, aiL);
+                    argsR.push_back(al, aiR);
+                }
+                dest_type = ASRUtils::type_get_past_array(dest_type);
+                left = ASRUtils::EXPR(make_ArrayItem_t(al, left->base.loc, left, argsL.p, argsL.n, dest_type, ASR::arraystorageType::RowMajor, nullptr));
+                right = ASRUtils::EXPR(make_ArrayItem_t(al, right->base.loc, right, argsR.p, argsR.n, dest_type, ASR::arraystorageType::RowMajor, nullptr));
+            }
+        }
+
         if (ASRUtils::is_integer(*dest_type)) {
             if (ASRUtils::expr_value(left) != nullptr && ASRUtils::expr_value(right) != nullptr) {
                 int64_t left_value = -1;

tests/errors/arrays_02.py

@@ -0,0 +1,28 @@
+from lpython import (i8, i32, dataclass)
+from numpy import (empty, int8)
+
+@dataclass
+class Foo:
+    a : i8[4] = empty(4, dtype=int8)
+    dim : i32 = 4
+
+def trinary_majority(x : Foo, y : Foo, z : Foo) -> Foo:
+    foo : Foo = Foo()
+    i : i32
+    for i in range(foo.dim):
+        foo.a[i] = (x.a[i] & y.a[i]) | (y.a[i] & z.a[i]) | (z.a[i] & x.a[i])
+    return foo
+
+
+t1 : Foo = Foo()
+t1.a = empty(4, dtype=int8)
+
+t2 : Foo = Foo()
+t2.a = empty(4, dtype=int8)
+
+t3 : Foo = Foo()
+t3.a = empty(4, dtype=int8)
+
+r1 : Foo = trinary_majority(t1, t2, t3)
+
+assert r1.a == t1.a

tests/reference/asr-arrays_02-da94458.json

@@ -0,0 +1,13 @@
+{
+    "basename": "asr-arrays_02-da94458",
+    "cmd": "lpython --show-asr --no-color {infile} -o {outfile}",
+    "infile": "tests/errors/arrays_02.py",
+    "infile_hash": "05e70a0056dc67dbf3a54ea66965db8746f9de012561ca95cb1fdb43",
+    "outfile": null,
+    "outfile_hash": null,
+    "stdout": null,
+    "stdout_hash": null,
+    "stderr": "asr-arrays_02-da94458.stderr",
+    "stderr_hash": "dc0e5be7cd6de7395421aedf1ce11977206f3e35bb7cba271aed8992",
+    "returncode": 2
+}

tests/reference/asr-arrays_02-da94458.stderr

@@ -0,0 +1,5 @@
+semantic error: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
+  --> tests/errors/arrays_02.py:28:8
+   |
+28 | assert r1.a == t1.a
+   |        ^^^^^^^^^^^^ 

tests/tests.toml

@@ -650,6 +650,10 @@ asr = true
 filename = "errors/arrays_01.py"
 asr = true
 
+[[test]]
+filename = "errors/arrays_02.py"
+asr = true
+
 [[test]]
 filename = "errors/structs_02.py"
 asr = true