Merge pull request #1059 from czgdp1807/dict09

Supporting tuples as keys in ``dict``

Author: czgdp1807

integration_tests/CMakeLists.txt

@@ -160,6 +160,7 @@ RUN(NAME test_tuple_02       LABELS cpython llvm)
 RUN(NAME test_dict_01        LABELS cpython llvm)
 RUN(NAME test_dict_02        LABELS cpython llvm)
 RUN(NAME test_dict_03        LABELS cpython llvm)
+RUN(NAME test_dict_04        LABELS cpython llvm)
 RUN(NAME modules_01          LABELS cpython llvm)
 RUN(NAME modules_02          LABELS cpython llvm)
 RUN(NAME test_math           LABELS cpython llvm)

integration_tests/test_dict_04.py

@@ -0,0 +1,80 @@
+from ltypes import i32, i64, f64
+from math import pi, sin, cos
+
+def test_dict():
+    terms2poly: dict[tuple[i32, i32], i64] = {}
+    rtheta2coords: dict[tuple[i64, i64], tuple[f64, f64]] = {}
+    i: i32
+    n: i64
+    size: i32 = 7000
+    size1: i32
+    theta: f64
+    r: f64
+    coords: tuple[f64, f64]
+    eps: f64 = 1e-12
+
+    n = 0
+    for i in range(1000, 1000 + size, 7):
+        terms2poly[(i, i*i)] = int(i + i*i)
+
+        theta = float(n) * pi
+        r = float(i)
+        rtheta2coords[(int(i), n)] = (r * sin(theta), r * cos(theta))
+
+        n += int(1)
+
+    size1 = size/7
+    n = 0
+    for i in range(1000, 1000 + size//2, 7):
+        assert terms2poly.pop((i, i*i)) == int(i + i*i)
+
+        theta = float(n) * pi
+        r = float(i)
+        coords = rtheta2coords.pop((int(i), n))
+        assert abs(coords[0] - r * sin(theta)) <= eps
+        assert abs(coords[1] - r * cos(theta)) <= eps
+
+        size1 = size1 - 1
+        assert len(terms2poly) == size1
+        n += int(1)
+
+    n = 0
+    for i in range(1000, 1000 + size//2, 7):
+        terms2poly[(i, i*i)] = int(1 + 2*i + i*i)
+
+        theta = float(n) * pi
+        r = float(i)
+        rtheta2coords[(int(i), n)] = (r * cos(theta), r * sin(theta))
+
+        n += int(1)
+
+    n = 0
+    for i in range(1000, 1000 + size//2, 7):
+        assert terms2poly[(i, i*i)] == (i + 1)*(i + 1)
+
+        theta = float(n) * pi
+        r = float(i)
+        assert abs(rtheta2coords[(int(i), n)][0] - r * cos(theta)) <= eps
+        assert abs(rtheta2coords[(int(i), n)][1] - r * sin(theta)) <= eps
+
+        n += int(1)
+
+    n = 0
+    for i in range(1000, 1000 + size, 7):
+        terms2poly[(i, i*i)] = int(1 + 2*i + i*i)
+
+        theta = float(n) * pi
+        r = float(i)
+        rtheta2coords[(int(i), n)] = (r * cos(theta), r * sin(theta))
+        n += int(1)
+
+    n = 0
+    for i in range(1000, 1000 + size, 7):
+        assert terms2poly[(i, i*i)] == (i + 1)*(i + 1)
+
+        theta = float(n) * pi
+        r = float(i)
+        assert abs(r**2 - rtheta2coords[(int(i), n)][0]**2 - r**2 * sin(theta)**2) <= eps
+        n += int(1)
+
+test_dict()

src/libasr/codegen/asr_to_llvm.cpp

@@ -1369,6 +1369,7 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
         ptr_loads = !LLVM::is_llvm_struct(dict_type->m_key_type);
         this->visit_expr_wrapper(x.m_key, true);
         llvm::Value *key = tmp;
+        ptr_loads = !LLVM::is_llvm_struct(dict_type->m_value_type);
         this->visit_expr_wrapper(x.m_value, true);
         llvm::Value *value = tmp;
         ptr_loads = ptr_loads_copy;

src/libasr/codegen/llvm_utils.cpp

@@ -883,6 +883,12 @@ namespace LFortran {
                           occupancy_ptr);
 
         llvm::Value* linear_prob_happened = builder->CreateICmpNE(key_hash, pos);
+        linear_prob_happened = builder->CreateOr(linear_prob_happened,
+            builder->CreateICmpEQ(
+                LLVM::CreateLoad(*builder, llvm_utils->create_ptr_gep(key_mask, key_hash)),
+                llvm::ConstantInt::get(llvm::Type::getInt8Ty(context), llvm::APInt(8, 2)
+            ))
+        );
         llvm::Value* set_max_2 = builder->CreateSelect(linear_prob_happened,
             llvm::ConstantInt::get(llvm::Type::getInt8Ty(context), llvm::APInt(8, 2)),
             llvm::ConstantInt::get(llvm::Type::getInt8Ty(context), llvm::APInt(8, 1)));
@@ -941,7 +947,7 @@ namespace LFortran {
     }
 
     llvm::Value* LLVMDict::get_key_hash(llvm::Value* capacity, llvm::Value* key,
-                                        ASR::ttype_t* key_asr_type, llvm::Module& /*module*/) {
+                                        ASR::ttype_t* key_asr_type, llvm::Module& module) {
         // Write specialised hash functions for intrinsic types
         // This is to avoid unnecessary calls to C-runtime and do
         // as much as possible in LLVM directly.
@@ -951,11 +957,12 @@ namespace LFortran {
                 // We can update it later to do a better hash function
                 // which produces lesser collisions.
 
-                return builder->CreateZExtOrTrunc(
+                llvm::Value* int_hash = builder->CreateZExtOrTrunc(
                     builder->CreateSRem(key,
                     builder->CreateZExtOrTrunc(capacity, key->getType())),
                     capacity->getType()
                 );
+                return int_hash;
             }
             case ASR::ttypeType::Character: {
                 // Polynomial rolling hash function for strings
@@ -1022,6 +1029,18 @@ namespace LFortran {
                 hash = builder->CreateTrunc(hash, llvm::Type::getInt32Ty(context));
                 return builder->CreateSRem(hash, capacity);
             }
+            case ASR::ttypeType::Tuple: {
+                llvm::Value* tuple_hash = llvm::ConstantInt::get(llvm::Type::getInt32Ty(context), llvm::APInt(32, 0));
+                ASR::Tuple_t* asr_tuple = ASR::down_cast<ASR::Tuple_t>(key_asr_type);
+                for( size_t i = 0; i < asr_tuple->n_type; i++ ) {
+                    llvm::Value* llvm_tuple_i = llvm_utils->tuple_api->read_item(key, i,
+                                                    LLVM::is_llvm_struct(asr_tuple->m_type[i]));
+                    tuple_hash = builder->CreateAdd(tuple_hash, get_key_hash(capacity, llvm_tuple_i,
+                                                                             asr_tuple->m_type[i], module));
+                    tuple_hash = builder->CreateSRem(tuple_hash, capacity);
+                }
+                return tuple_hash;
+            }
             default: {
                 throw LCompilersException("Hashing " + ASRUtils::type_to_str_python(key_asr_type) +
                                           " isn't implemented yet.");

src/lpython/semantics/python_ast_to_asr.cpp

@@ -2198,7 +2198,8 @@ class CommonVisitor : public AST::BaseVisitor<Derived> {
             } else if (ASR::is_a<ASR::Dict_t>(*type)) {
                 throw SemanticError("unhashable type in dict: 'slice'", loc);
             }
-        } else if(AST::is_a<AST::Tuple_t>(*m_slice)) {
+        } else if(AST::is_a<AST::Tuple_t>(*m_slice) &&
+                  !ASR::is_a<ASR::Dict_t>(*type)) {
             bool final_result = true;
             AST::Tuple_t* indices = AST::down_cast<AST::Tuple_t>(m_slice);
             for( size_t i = 0; i < indices->n_elts; i++ ) {