dict int32 int32 support

Author: Ehsan Totoni

hpat/__init__.py

@@ -3,7 +3,7 @@
 import numba
 from numba import *
 import hpat.dict_ext
-from hpat.dict_ext import DictIntInt, dict_int_int_type
+from hpat.dict_ext import DictIntInt, DictInt32Int32, dict_int_int_type, dict_int32_int32_type
 import hpat.str_ext
 
 def jit(signature_or_function=None, **options):

hpat/_dict_ext.cpp

@@ -14,6 +14,19 @@ int64_t dict_int_int_min(std::unordered_map<int64_t, int64_t>* m);
 int64_t dict_int_int_max(std::unordered_map<int64_t, int64_t>* m);
 bool dict_int_int_not_empty(std::unordered_map<int64_t, int64_t>* m);
 
+// -- int32 versions --
+void* init_dict_int32_int32();
+void dict_int32_int32_setitem(std::unordered_map<int, int>* m, int index, int value);
+void dict_int32_int32_print(std::unordered_map<int, int>* m);
+int dict_int32_int32_get(std::unordered_map<int, int>* m, int index, int default_val);
+int dict_int32_int32_getitem(std::unordered_map<int, int>* m, int index);
+int dict_int32_int32_pop(std::unordered_map<int, int>* m, int index);
+void* dict_int32_int32_keys(std::unordered_map<int, int>* m);
+int dict_int32_int32_min(std::unordered_map<int, int>* m);
+int dict_int32_int32_max(std::unordered_map<int, int>* m);
+bool dict_int32_int32_not_empty(std::unordered_map<int, int>* m);
+
+
 PyMODINIT_FUNC PyInit_hdict_ext(void) {
     PyObject *m;
     static struct PyModuleDef moduledef = {
@@ -42,6 +55,27 @@ PyMODINIT_FUNC PyInit_hdict_ext(void) {
                             PyLong_FromVoidPtr((void*)(&dict_int_int_max)));
     PyObject_SetAttrString(m, "dict_int_int_not_empty",
                             PyLong_FromVoidPtr((void*)(&dict_int_int_not_empty)));
+    // ---- int32 versions ----
+    PyObject_SetAttrString(m, "init_dict_int32_int32",
+                            PyLong_FromVoidPtr((void*)(&init_dict_int32_int32)));
+    PyObject_SetAttrString(m, "dict_int32_int32_setitem",
+                            PyLong_FromVoidPtr((void*)(&dict_int32_int32_setitem)));
+    PyObject_SetAttrString(m, "dict_int32_int32_print",
+                            PyLong_FromVoidPtr((void*)(&dict_int32_int32_print)));
+    PyObject_SetAttrString(m, "dict_int32_int32_get",
+                            PyLong_FromVoidPtr((void*)(&dict_int32_int32_get)));
+    PyObject_SetAttrString(m, "dict_int32_int32_getitem",
+                            PyLong_FromVoidPtr((void*)(&dict_int32_int32_getitem)));
+    PyObject_SetAttrString(m, "dict_int32_int32_pop",
+                            PyLong_FromVoidPtr((void*)(&dict_int32_int32_pop)));
+    PyObject_SetAttrString(m, "dict_int32_int32_keys",
+                            PyLong_FromVoidPtr((void*)(&dict_int32_int32_keys)));
+    PyObject_SetAttrString(m, "dict_int32_int32_min",
+                            PyLong_FromVoidPtr((void*)(&dict_int32_int32_min)));
+    PyObject_SetAttrString(m, "dict_int32_int32_max",
+                            PyLong_FromVoidPtr((void*)(&dict_int32_int32_max)));
+    PyObject_SetAttrString(m, "dict_int32_int32_not_empty",
+                            PyLong_FromVoidPtr((void*)(&dict_int32_int32_not_empty)));
     return m;
 }
 
@@ -117,3 +151,78 @@ bool dict_int_int_not_empty(std::unordered_map<int64_t, int64_t>* m)
 {
     return !m->empty();
 }
+
+
+// --------- int32 versions ------
+void* init_dict_int32_int32()
+{
+    return new std::unordered_map<int, int>();
+}
+
+void dict_int32_int32_setitem(std::unordered_map<int, int>* m, int index, int value)
+{
+    (*m)[index] = value;
+    return;
+}
+
+void dict_int32_int32_print(std::unordered_map<int, int>* m)
+{
+    // TODO: return python string and print in native mode
+    for (auto& x: *m) {
+        std::cout << x.first << ": " << x.second << std::endl;
+    }
+    return;
+}
+
+int dict_int32_int32_get(std::unordered_map<int, int>* m, int index, int default_val)
+{
+    auto val = m->find(index);
+    if (val==m->end())
+        return default_val;
+    return (*val).second;
+}
+
+int dict_int32_int32_getitem(std::unordered_map<int, int>* m, int index)
+{
+    return m->at(index);
+}
+
+int dict_int32_int32_pop(std::unordered_map<int, int>* m, int index)
+{
+    int val = m->at(index);
+    m->erase(index);
+    return val;
+}
+
+void* dict_int32_int32_keys(std::unordered_map<int, int>* m)
+{
+    // TODO: return actual iterator
+    return m;
+}
+
+int dict_int32_int32_min(std::unordered_map<int, int>* m)
+{
+    // TODO: use actual iterator
+    int res = std::numeric_limits<int>::max();
+    for (auto& x: *m) {
+        if (x.first<res)
+            res = x.first;
+    }
+    return res;
+}
+
+int dict_int32_int32_max(std::unordered_map<int, int>* m)
+{
+    // TODO: use actual iterator
+    int res = std::numeric_limits<int>::min();
+    for (auto& x: *m) {
+        if (x.first>res)
+            res = x.first;
+    }
+    return res;
+}
+
+bool dict_int32_int32_not_empty(std::unordered_map<int, int>* m)
+{
+    return !m->empty();
+}

hpat/dict_ext.py

@@ -117,6 +117,7 @@ def __init__(self, key_typ, val_typ):
                             'DictKeyIteratorType{}{}'.format(key_typ, val_typ))
 
 dict_key_iterator_int_int_type = DictKeyIteratorType(types.intp, types.intp)
+dict_key_iterator_int32_int32_type = DictKeyIteratorType(types.int32, types.int32)
 
 register_model(DictKeyIteratorType)(models.OpaqueModel)
 
@@ -139,6 +140,19 @@ def generic(self, args, kws):
 ll.add_symbol('dict_int_int_max', hdict_ext.dict_int_int_max)
 ll.add_symbol('dict_int_int_not_empty', hdict_ext.dict_int_int_not_empty)
 
+# int32 versions
+ll.add_symbol('init_dict_int32_int32', hdict_ext.init_dict_int32_int32)
+ll.add_symbol('dict_int32_int32_setitem', hdict_ext.dict_int32_int32_setitem)
+ll.add_symbol('dict_int32_int32_print', hdict_ext.dict_int32_int32_print)
+ll.add_symbol('dict_int32_int32_get', hdict_ext.dict_int32_int32_get)
+ll.add_symbol('dict_int32_int32_getitem', hdict_ext.dict_int32_int32_getitem)
+ll.add_symbol('dict_int32_int32_pop', hdict_ext.dict_int32_int32_pop)
+ll.add_symbol('dict_int32_int32_keys', hdict_ext.dict_int32_int32_keys)
+ll.add_symbol('dict_int32_int32_min', hdict_ext.dict_int32_int32_min)
+ll.add_symbol('dict_int32_int32_max', hdict_ext.dict_int32_int32_max)
+ll.add_symbol('dict_int32_int32_not_empty', hdict_ext.dict_int32_int32_not_empty)
+
+
 @lower_builtin(DictIntInt)
 def impl_dict_int_int(context, builder, sig, args):
     fnty = lir.FunctionType(lir.IntType(8).as_pointer(), [])
@@ -151,7 +165,7 @@ def setitem_dict(context, builder, sig, args):
     fn = builder.module.get_or_insert_function(fnty, name="dict_int_int_setitem")
     return builder.call(fn, args)
 
-@lower_builtin("print_item", DictType)
+@lower_builtin("print_item", dict_int_int_type)
 def print_dict(context, builder, sig, args):
     # pyapi = context.get_python_api(builder)
     # strobj = pyapi.unserialize(pyapi.serialize_object("hello!"))
@@ -180,26 +194,93 @@ def lower_dict_pop(context, builder, sig, args):
     fn = builder.module.get_or_insert_function(fnty, name="dict_int_int_pop")
     return builder.call(fn, args)
 
-@lower_builtin("dict.keys", DictType)
+@lower_builtin("dict.keys", dict_int_int_type)
 def lower_dict_keys(context, builder, sig, args):
     fnty = lir.FunctionType(lir.IntType(8).as_pointer(), [lir.IntType(8).as_pointer()])
     fn = builder.module.get_or_insert_function(fnty, name="dict_int_int_keys")
     return builder.call(fn, args)
 
-@lower_builtin(min, DictKeyIteratorType)
+@lower_builtin(min, dict_key_iterator_int_int_type)
 def lower_dict_min(context, builder, sig, args):
     fnty = lir.FunctionType(lir.IntType(64), [lir.IntType(8).as_pointer()])
     fn = builder.module.get_or_insert_function(fnty, name="dict_int_int_min")
     return builder.call(fn, args)
 
-@lower_builtin(max, DictKeyIteratorType)
+@lower_builtin(max, dict_key_iterator_int_int_type)
 def lower_dict_max(context, builder, sig, args):
     fnty = lir.FunctionType(lir.IntType(64), [lir.IntType(8).as_pointer()])
     fn = builder.module.get_or_insert_function(fnty, name="dict_int_int_max")
     return builder.call(fn, args)
 
-@lower_cast(DictType, types.boolean)
+@lower_cast(dict_int_int_type, types.boolean)
 def dict_empty(context, builder, fromty, toty, val):
     fnty = lir.FunctionType(lir.IntType(1), [lir.IntType(8).as_pointer()])
     fn = builder.module.get_or_insert_function(fnty, name="dict_int_int_not_empty")
     return builder.call(fn, (val,))
+
+
+# ------ int32 versions ------
+@lower_builtin(DictInt32Int32)
+def impl_dict_int_int(context, builder, sig, args):
+    fnty = lir.FunctionType(lir.IntType(8).as_pointer(), [])
+    fn = builder.module.get_or_insert_function(fnty, name="init_dict_int32_int32")
+    return builder.call(fn, [])
+
+@lower_builtin('setitem', DictType, types.int32, types.int32)
+def setitem_dict(context, builder, sig, args):
+    fnty = lir.FunctionType(lir.VoidType(), [lir.IntType(8).as_pointer(), lir.IntType(32), lir.IntType(32)])
+    fn = builder.module.get_or_insert_function(fnty, name="dict_int32_int32_setitem")
+    return builder.call(fn, args)
+
+@lower_builtin("print_item", dict_int32_int32_type)
+def print_dict(context, builder, sig, args):
+    # pyapi = context.get_python_api(builder)
+    # strobj = pyapi.unserialize(pyapi.serialize_object("hello!"))
+    # pyapi.print_object(strobj)
+    # pyapi.decref(strobj)
+    # return context.get_dummy_value()
+    fnty = lir.FunctionType(lir.VoidType(), [lir.IntType(8).as_pointer()])
+    fn = builder.module.get_or_insert_function(fnty, name="dict_int32_int32_print")
+    return builder.call(fn, args)
+
+@lower_builtin("dict.get", DictType, types.int32, types.int32)
+def lower_dict_get(context, builder, sig, args):
+    fnty = lir.FunctionType(lir.IntType(32), [lir.IntType(8).as_pointer(), lir.IntType(32), lir.IntType(32)])
+    fn = builder.module.get_or_insert_function(fnty, name="dict_int32_int32_get")
+    return builder.call(fn, args)
+
+@lower_builtin("getitem", DictType, types.int32)
+def lower_dict_getitem(context, builder, sig, args):
+    fnty = lir.FunctionType(lir.IntType(32), [lir.IntType(8).as_pointer(), lir.IntType(32)])
+    fn = builder.module.get_or_insert_function(fnty, name="dict_int32_int32_getitem")
+    return builder.call(fn, args)
+
+@lower_builtin("dict.pop", DictType, types.int32)
+def lower_dict_pop(context, builder, sig, args):
+    fnty = lir.FunctionType(lir.IntType(32), [lir.IntType(8).as_pointer(), lir.IntType(32)])
+    fn = builder.module.get_or_insert_function(fnty, name="dict_int32_int32_pop")
+    return builder.call(fn, args)
+
+@lower_builtin("dict.keys", dict_int32_int32_type)
+def lower_dict_keys(context, builder, sig, args):
+    fnty = lir.FunctionType(lir.IntType(8).as_pointer(), [lir.IntType(8).as_pointer()])
+    fn = builder.module.get_or_insert_function(fnty, name="dict_int32_int32_keys")
+    return builder.call(fn, args)
+
+@lower_builtin(min, dict_key_iterator_int32_int32_type)
+def lower_dict_min(context, builder, sig, args):
+    fnty = lir.FunctionType(lir.IntType(32), [lir.IntType(8).as_pointer()])
+    fn = builder.module.get_or_insert_function(fnty, name="dict_int32_int32_min")
+    return builder.call(fn, args)
+
+@lower_builtin(max, dict_key_iterator_int32_int32_type)
+def lower_dict_max(context, builder, sig, args):
+    fnty = lir.FunctionType(lir.IntType(32), [lir.IntType(8).as_pointer()])
+    fn = builder.module.get_or_insert_function(fnty, name="dict_int32_int32_max")
+    return builder.call(fn, args)
+
+@lower_cast(dict_int32_int32_type, types.boolean)
+def dict_empty(context, builder, fromty, toty, val):
+    fnty = lir.FunctionType(lir.IntType(1), [lir.IntType(8).as_pointer()])
+    fn = builder.module.get_or_insert_function(fnty, name="dict_int32_int32_not_empty")
+    return builder.call(fn, (val,))