string getitem char

Author: Ehsan Totoni

hpat/_str_ext.cpp

@@ -9,6 +9,7 @@ const char* get_c_str(std::string* s);
 void* str_concat(std::string* s1, std::string* s2);
 bool str_equal(std::string* s1, std::string* s2);
 void* str_split(std::string* str, std::string* sep, int64_t *size);
+void* str_substr_int(std::string* str, int64_t index);
 
 PyMODINIT_FUNC PyInit_hstr_ext(void) {
     PyObject *m;
@@ -30,6 +31,8 @@ PyMODINIT_FUNC PyInit_hstr_ext(void) {
                             PyLong_FromVoidPtr((void*)(&str_equal)));
     PyObject_SetAttrString(m, "str_split",
                             PyLong_FromVoidPtr((void*)(&str_split)));
+    PyObject_SetAttrString(m, "str_substr_int",
+                            PyLong_FromVoidPtr((void*)(&str_substr_int)));
     return m;
 }
 
@@ -87,3 +90,8 @@ void* str_split(std::string* str, std::string* sep, int64_t *size)
     // std::cout<< *(((std::string**)(out))[1])<<std::endl;
     return out;
 }
+
+void* str_substr_int(std::string* str, int64_t index)
+{
+    return new std::string(*str, index, 1);
+}

hpat/str_ext.py

@@ -49,13 +49,24 @@ def resolve_split(self, dict, args, kws):
         assert len(args) == 1
         return signature(types.List(string_type), *args)
 
+@infer
+class GetItemString(AbstractTemplate):
+    key = "getitem"
+
+    def generic(self, args, kws):
+        assert not kws
+        if (len(args) == 2 and isinstance(args[0], StringType)
+                and isinstance(args[1], types.Integer)):
+            return signature(args[0], *args)
+
 import hstr_ext
 ll.add_symbol('init_string', hstr_ext.init_string)
 ll.add_symbol('init_string_const', hstr_ext.init_string_const)
 ll.add_symbol('get_c_str', hstr_ext.get_c_str)
 ll.add_symbol('str_concat', hstr_ext.str_concat)
 ll.add_symbol('str_equal', hstr_ext.str_equal)
 ll.add_symbol('str_split', hstr_ext.str_split)
+ll.add_symbol('str_substr_int', hstr_ext.str_substr_int)
 
 @unbox(StringType)
 def unbox_string(typ, obj, c):
@@ -131,3 +142,12 @@ def string_split_impl(context, builder, sig, args):
         value = builder.load(cgutils.gep_inbounds(builder, ptr, loop.index))
         _list.setitem(loop.index, value)
     return impl_ret_new_ref(context, builder, sig.return_type, _list.value)
+
+@lower_builtin('getitem', StringType, types.Integer)
+def getitem_string(context, builder, sig, args):
+    fnty = lir.FunctionType(lir.IntType(8).as_pointer(),
+                    [lir.IntType(8).as_pointer(), lir.IntType(64)])
+    fn = builder.module.get_or_insert_function(fnty, name="str_substr_int")
+    # TODO: handle reference counting
+    #return impl_ret_new_ref(builder.call(fn, args))
+    return (builder.call(fn, args))