From b9850da2cff7be3d30f667aceee5c15f31252bfa Mon Sep 17 00:00:00 2001
From: kabra1110 <131529938+kabra1110@users.noreply.github.com>
Date: Sat, 13 May 2023 12:56:11 +0200
Subject: [PATCH] Support ``bool`` typed keys in ``dict`` (#1771)

Co-authored-by: Gagandeep Singh <gdp.1807@gmail.com>
---
 integration_tests/CMakeLists.txt    |  1 +
 integration_tests/test_dict_bool.py | 47 +++++++++++++++++++++++++++++
 src/libasr/codegen/llvm_utils.cpp   |  3 ++
 3 files changed, 51 insertions(+)
 create mode 100644 integration_tests/test_dict_bool.py

diff --git a/integration_tests/CMakeLists.txt b/integration_tests/CMakeLists.txt
index 1d8be7644f..046cb6574b 100644
--- a/integration_tests/CMakeLists.txt
+++ b/integration_tests/CMakeLists.txt
@@ -374,6 +374,7 @@ RUN(NAME test_dict_08        LABELS cpython llvm c)
 RUN(NAME test_dict_09        LABELS cpython llvm c)
 RUN(NAME test_dict_10        LABELS cpython llvm) # TODO: Add support of dict with string in C backend
 RUN(NAME test_dict_11        LABELS cpython llvm c)
+RUN(NAME test_dict_bool      LABELS cpython llvm)
 RUN(NAME test_for_loop       LABELS cpython llvm c)
 RUN(NAME modules_01          LABELS cpython llvm c wasm wasm_x86 wasm_x64)
 RUN(NAME modules_02          LABELS cpython llvm c wasm wasm_x86 wasm_x64)
diff --git a/integration_tests/test_dict_bool.py b/integration_tests/test_dict_bool.py
new file mode 100644
index 0000000000..c1e4e9cca9
--- /dev/null
+++ b/integration_tests/test_dict_bool.py
@@ -0,0 +1,47 @@
+from lpython import i32, f64
+
+def test_dict_bool():
+    d_int: dict[bool, i32] = {}
+    d_float: dict[bool, f64] = {}
+    d_str: dict[bool, str] = {}
+    i: i32
+    j: f64
+    s: str = ""
+    l_str: list[str] = ["a", "b", "c", "d"]
+
+    for i in range(10):
+        d_int[True] = i
+        assert d_int[True] == i
+
+    for i in range(10, 20):
+        d_int[True] = i
+        d_int[False] = i + 1
+        assert d_int[True] == d_int[False] - 1
+        assert d_int[True] == i
+    
+    d_int[True] = 0
+    d_int[False] = d_int[True]
+
+    for i in range(10, 99):
+        d_int[i%2 == 0] = d_int[i%2 == 0] + 1
+    assert d_int[True] == d_int[False] + 1
+    assert d_int[True] == 45
+
+    j = 0.0
+    while j < 1.0:
+        d_float[False] = j + 1.0
+        d_float[True] = d_float[False] * d_float[False]
+        assert d_float[True] == (j + 1.0) * (j + 1.0)
+        assert d_float[False] == j + 1.0
+        j = j + 0.1
+
+    d_str[False] = s
+
+    for i in range(len(l_str)):
+        d_str[True] = d_str[False]
+        s += l_str[i]
+        d_str[False] = s
+        assert d_str[True] + l_str[i] == d_str[False]
+        assert d_str[False] == s
+
+test_dict_bool()
diff --git a/src/libasr/codegen/llvm_utils.cpp b/src/libasr/codegen/llvm_utils.cpp
index 8250d26b38..dd1b2ec973 100644
--- a/src/libasr/codegen/llvm_utils.cpp
+++ b/src/libasr/codegen/llvm_utils.cpp
@@ -1962,6 +1962,9 @@ namespace LCompilers {
                 }
                 return tuple_hash;
             }
+            case ASR::ttypeType::Logical: {
+                return builder->CreateZExt(key, llvm::Type::getInt32Ty(context));
+            }
             default: {
                 throw LCompilersException("Hashing " + ASRUtils::type_to_str_python(key_asr_type) +
                                           " isn't implemented yet.");