support str split

Author: Ehsan Totoni

hpat/_str_ext.cpp

@@ -1,12 +1,14 @@
 #include <Python.h>
 #include <string>
 #include <iostream>
+#include <vector>
 
 void* init_string(char*, int64_t);
 void* init_string_const(char* in_str);
 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);
 
 PyMODINIT_FUNC PyInit_hstr_ext(void) {
     PyObject *m;
@@ -26,6 +28,8 @@ PyMODINIT_FUNC PyInit_hstr_ext(void) {
                             PyLong_FromVoidPtr((void*)(&str_concat)));
     PyObject_SetAttrString(m, "str_equal",
                             PyLong_FromVoidPtr((void*)(&str_equal)));
+    PyObject_SetAttrString(m, "str_split",
+                            PyLong_FromVoidPtr((void*)(&str_split)));
     return m;
 }
 
@@ -59,3 +63,27 @@ bool str_equal(std::string* s1, std::string* s2)
     // printf("in str_equal %s %s\n", s1->c_str(), s2->c_str());
     return s1->compare(*s2)==0;
 }
+
+void* str_split(std::string* str, std::string* sep, int64_t *size)
+{
+    // std::cout << *str << " " << *sep << std::endl;
+    std::vector<std::string*> res;
+
+    size_t last = 0;
+    size_t next = 0;
+    while ((next = str->find(*sep, last)) != std::string::npos) {
+        std::string *token = new std::string(str->substr(last, next-last));
+        res.push_back(token);
+        last = next + 1;
+    }
+    std::string *token = new std::string(str->substr(last));
+    res.push_back(token);
+    *size = res.size();
+    // for(int i=0; i<*size; i++)
+    //    std::cout<<*(res[i])<<std::endl;
+    // TODO: avoid extra copy
+    void* out = new void*[*size];
+    memcpy(out, res.data(), (*size)*sizeof(void*));
+    // std::cout<< *(((std::string**)(out))[1])<<std::endl;
+    return out;
+}

hpat/str_ext.py

@@ -1,8 +1,9 @@
+import numba
 from numba.extending import (box, unbox, typeof_impl, register_model, models,
                             NativeValue, lower_builtin)
-from numba.targets.imputils import lower_constant
+from numba.targets.imputils import lower_constant, impl_ret_new_ref
 from numba import types, typing
-from numba.typing.templates import (signature, AbstractTemplate, infer,
+from numba.typing.templates import (signature, AbstractTemplate, infer, infer_getattr,
         ConcreteTemplate, AttributeTemplate, bound_function, infer_global)
 from numba import cgutils
 from llvmlite import ir as lir
@@ -38,12 +39,23 @@ def generic(self, args, kws):
 class StringOpNotEq(StringOpEq):
     key = '!='
 
+@infer_getattr
+class StringAttribute(AttributeTemplate):
+    key = StringType
+
+    @bound_function("str.split")
+    def resolve_split(self, dict, args, kws):
+        assert not kws
+        assert len(args) == 1
+        return signature(types.List(string_type), *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)
 
 @unbox(StringType)
 def unbox_string(typ, obj, c):
@@ -99,3 +111,23 @@ def string_neq_impl(context, builder, sig, args):
                     [lir.IntType(8).as_pointer(), lir.IntType(8).as_pointer()])
     fn = builder.module.get_or_insert_function(fnty, name="str_equal")
     return builder.not_(builder.call(fn, args))
+
+@lower_builtin("str.split", string_type, string_type)
+def string_split_impl(context, builder, sig, args):
+    nitems = cgutils.alloca_once(builder, lir.IntType(64))
+    # input str, sep, size pointer
+    fnty = lir.FunctionType(lir.IntType(8).as_pointer().as_pointer(),
+                [lir.IntType(8).as_pointer(), lir.IntType(8).as_pointer(),
+                lir.IntType(64).as_pointer()])
+    fn = builder.module.get_or_insert_function(fnty, name="str_split")
+    ptr = builder.call(fn, args+[nitems])
+    size = builder.load(nitems)
+    # TODO: use ptr instead of allocating and copying, use NRT_MemInfo_new
+    # TODO: deallocate ptr
+    _list = numba.targets.listobj.ListInstance.allocate(context, builder,
+                                    sig.return_type, size)
+    _list.size = size
+    with cgutils.for_range(builder, size) as loop:
+        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)