Commits

Maciej Fijalkowski  committed d1cbcc3

first approximation of array[array-of-int]

  • Participants
  • Parent commits 54c7d01
  • Branches numpy-indexing-by-arrays-2

Comments (0)

Files changed (3)

File pypy/module/micronumpy/interp_boxes.py

     _attrs_ = ()
 
 class W_IntegerBox(W_NumberBox):
-    pass
+    def convert_to_int(self):
+        from pypy.rpython.lltypesystem import rffi, lltype
+        
+        return rffi.cast(lltype.Signed, self.value)
 
 class W_SignedIntegerBox(W_IntegerBox):
     pass

File pypy/module/micronumpy/interp_iter.py

         if self.step != 0:
             shape.append(self.lgt)
 
+    def get_iter(self):
+        xxx
+
 class IntArrayChunk(BaseChunk):
     def __init__(self, arr):
         self.arr = arr.get_concrete()
     def extend_shape(self, shape):
         shape.extend(self.arr.shape)
 
+    def get_iter(self):
+        return self.arr.create_iter()
+
+    def get_index(self, iter):
+        return self.arr.getitem(iter.offset).convert_to_int()
+
 class BoolArrayChunk(BaseChunk):
     def __init__(self, arr):
         self.arr = arr.get_concrete()
 
+    def extend_shape(self, shape):
+        xxx
+
+    def get_iter(self):
+        xxx
+
 class BaseTransform(object):
     pass
 
         return self._done
 
 # ------ other iterators that are not part of the computation frame ----------
+
+class ChunkIterator(object):
+    def __init__(self, shape, chunks):
+        self.chunks = chunks
+        self.indices = [0] * len(shape)
+        self.shape = shape
+        self.chunk_iters = [chunk.get_iter() for chunk in self.chunks]
+
+    def next(self, shapelen):
+        for i in range(shapelen - 1, -1, -1):
+            if self.indices[i] < self.shape[i] - 1:
+                self.indices[i] += 1
+                self.chunk_iters[i] = self.chunk_iters[i].next()
+                break
+            else:
+                self.indices[i] = 0
+                # XXX reset one dim iter probably
+        return self
+
+    def get_index(self, shapelen):
+        l = []
+        for i in range(shapelen):
+            l.append(self.chunks[i].get_index(self.chunk_iters[i]))
+        return l
     
 class SkipLastAxisIterator(object):
     def __init__(self, arr):

File pypy/module/micronumpy/interp_numarray.py

 from pypy.tool.sourcetools import func_with_new_name
 from pypy.rlib.rstring import StringBuilder
 from pypy.module.micronumpy.interp_iter import ArrayIterator, OneDimIterator,\
-     SkipLastAxisIterator, Chunk, ViewIterator, BoolArrayChunk, IntArrayChunk
+     SkipLastAxisIterator, Chunk, ViewIterator, BoolArrayChunk, IntArrayChunk,\
+     ChunkIterator
 
 numpy_driver = jit.JitDriver(
     greens=['shapelen', 'sig'],
             item = concrete._index_of_single_item(space, w_idx)
             return concrete.getitem(item)
         chunks = self._prepare_slice_args(space, w_idx)
+        for chunk in chunks:
+            if not isinstance(chunk, Chunk):
+                return self.get_concrete().force_slice_getitem(space, chunks)
         return space.wrap(self.create_slice(chunks))
 
     def descr_setitem(self, space, w_idx, w_value):
         view = self.create_slice(chunks).get_concrete()
         view.setslice(space, w_value)
 
-    def force_slice(self, shape, chunks):
-        size = 1
-        for elem in shape:
-            size *= elem
-        res = W_NDimArray(size, shape, self.find_dtype())
-        xxx
-
     @jit.unroll_safe
     def create_slice(self, chunks):
         shape = []
         s = i + 1
         assert s >= 0
         shape += self.shape[s:]
-        for chunk in chunks:
-            if not isinstance(chunk, Chunk):
-                return self.force_slice(shape, chunks)
         if not isinstance(self, ConcreteArray):
             return VirtualSlice(self, chunks, shape)
         r = calculate_slice_strides(self.shape, self.start, self.strides,
         builder.append(']')
 
     @jit.unroll_safe
+    def _index_of_single_item_int(self, space, index):
+        item = self.start
+        for i in range(len(index)):
+            v = index[i]
+            if v < 0:
+                v += self.shape[i]
+            if v < 0 or v >= self.shape[i]:
+                raise operationerrfmt(space.w_IndexError,
+                    "index (%d) out of range (0<=index<%d", i, self.shape[i],
+                )
+            item += v * self.strides[i]
+        return item
+
+    @jit.unroll_safe
     def _index_of_single_item(self, space, w_idx):
         if space.isinstance_w(w_idx, space.w_int):
             idx = space.int_w(w_idx)
             return self.start + idx * self.strides[0]
         index = [space.int_w(w_item)
                  for w_item in space.fixedview(w_idx)]
-        item = self.start
-        for i in range(len(index)):
-            v = index[i]
-            if v < 0:
-                v += self.shape[i]
-            if v < 0 or v >= self.shape[i]:
-                raise operationerrfmt(space.w_IndexError,
-                    "index (%d) out of range (0<=index<%d", i, self.shape[i],
-                )
-            item += v * self.strides[i]
-        return item
+        return self._index_of_single_item_int(space, index)
 
     def setslice(self, space, w_value):
         res_shape = shape_agreement(space, self.shape, w_value.shape)
     def fill(self, space, w_value):
         self.setslice(space, scalar_w(space, self.dtype, w_value))
 
+    def force_slice_getitem(self, space, chunks):
+        shape = []
+        i = -1
+        for i, chunk in enumerate(chunks):
+            chunk.extend_shape(shape)
+        s = i + 1
+        assert s >= 0
+        shape += self.shape[s:]
+        size = 1
+        for elem in shape:
+            size *= elem
+        res = W_NDimArray(size, shape, self.find_dtype())
+        ri = res.create_iter()
+        ci = ChunkIterator(shape, chunks)
+        shapelen = len(shape)
+        while not ri.done():
+            index = ci.get_index(shapelen)
+            v = self.getitem(self._index_of_single_item_int(space, index))
+            res.setitem(ri.offset, v)
+            ri = ri.next(shapelen)
+            ci = ci.next(shapelen)
+        return res
 
 class ViewArray(ConcreteArray):
     def create_sig(self):