[pypy-commit] pypy numpy-data-buffer: add jits for setslice and getslice, move constants to earlier places.

[timo_jbo]

Author: Timo Paulssen <timonator at perpetuum-immobile.de>
Branch: numpy-data-buffer
Changeset: r47808:0eb532b36b5d
Date: 2011-10-03 23:00 +0200
http://bitbucket.org/pypy/pypy/changeset/0eb532b36b5d/

Log:	add jits for setslice and getslice, move constants to earlier
	places.

diff --git a/lib_pypy/numpy/__init__.py b/lib_pypy/numpy/__init__.py
--- a/lib_pypy/numpy/__init__.py
+++ b/lib_pypy/numpy/__init__.py
@@ -88,6 +88,7 @@
 
 def __from_buffer_or_datastring(buf_or_str, dt, count, offset=0):
     _dtype = dtype(dt)
+    buf_or_str = buffer(buf_or_str)
 
     if count > 0:
         length = count * _dtype.itemsize
diff --git a/pypy/module/_numpy/interp_buffer.py b/pypy/module/_numpy/interp_buffer.py
--- a/pypy/module/_numpy/interp_buffer.py
+++ b/pypy/module/_numpy/interp_buffer.py
@@ -1,21 +1,28 @@
 from pypy.interpreter.buffer import RWBuffer
 from pypy.rpython.lltypesystem import lltype, rffi
 from pypy.rlib.rstring import StringBuilder
+from pypy.rlib import jit
 
 CHAR_TP = lltype.Ptr(lltype.Array(lltype.Char, hints={'nolength': True}))
 
+setslice_driver = jit.JitDriver(greens = ['index', 'end', 'self', 'newstring'],
+                                reds = ['idx'])
+getslice_driver = jit.JitDriver(greens = ['stop', 'step', 'self', 'builder'],
+                                reds = ['idx'])
 class NumpyBuffer(RWBuffer):
     def __init__(self, array):
         self.array = array
+        storage = self.array.get_concrete().get_root_storage()
+        self.char_data = rffi.cast(CHAR_TP, storage)
+        self.length = self.array.get_concrete().find_size() \
+                          * self.array.find_dtype().num_bytes
 
     def getlength(self):
-        return self.array.get_concrete().find_size() * self.array.find_dtype().num_bytes
+        return self.length
 
     def getitem_noboundcheck(self, index):
-        storage = self.array.get_concrete().get_root_storage()
-        char_data = rffi.cast(CHAR_TP, storage)
         index = self.calc_index(index)
-        return char_data[index]
+        return self.char_data[index]
 
     def getitem(self, index):
         if index > self.getlength():
@@ -23,10 +30,8 @@
         return self.getitem_noboundcheck(index)
 
     def setitem_noboundcheck(self, index, value):
-        storage = self.array.get_concrete().get_root_storage()
-        char_ptr = rffi.cast(CHAR_TP, storage)
         index = self.calc_index(index)
-        char_ptr[index] = value
+        self.char_data[index] = value
 
     def setitem(self, index, value):
         if index > self.getlength():
@@ -36,14 +41,23 @@
     def setslice(self, index, newstring):
         if index + len(newstring) > self.getlength():
             raise IndexError("End of slice to set out of bounds (0<=index<=%d)" % self.getlength())
-        for idx in range(0, len(newstring)):
-            self.setitem_noboundcheck(index + idx, newstring[idx])
+        idx = 0
+        end = len(newstring)
+        while idx < end:
+            setslice_driver.jit_merge_point(self=self, newstring=newstring, index=index, end=end,
+                                            idx=idx)
+            self.setitem_noboundcheck(idx + index, newstring[idx])
+            idx += 1
 
     def getslice(self, start, stop, step, size):
         builder = StringBuilder(size)
 
-        for index in range(start, stop, step):
-            builder.append(self.getitem_noboundcheck(index))
+        idx = start
+        while idx < stop:
+            getslice_driver.jit_merge_point(self=self, builder=builder, stop=stop, step=step,
+                                            idx=idx)
+            builder.append(self.getitem_noboundcheck(idx))
+            idx += step
 
         return builder.build()