numo-narray 0.9.1.4 → 0.9.1.5

data/ext/numo/narray/index.c

@@ -1,7 +1,7 @@
 /*
   index.c
-  Numerical Array Extension for Ruby
-    (C) Copyright 1999-2017 by Masahiro TANAKA
+  Ruby/Numo::NArray - Numerical Array class for Ruby
+    Copyright (C) 1999-2019 Masahiro TANAKA
 */
 //#define NARRAY_C
 
@@ -16,24 +16,6 @@
 #define cIndex numo_cInt32
 #endif
 
-// from ruby/enumerator.c
-struct enumerator {
-    VALUE obj;
-    ID    meth;
-    VALUE args;
-    // use only above in this source
-    VALUE fib;
-    VALUE dst;
-    VALUE lookahead;
-    VALUE feedvalue;
-    VALUE stop_exc;
-    VALUE size;
-    // incompatible below depending on ruby version
-    //VALUE procs;                      // ruby 2.4
-    //rb_enumerator_size_func *size_fn; // ruby 2.1-2.4
-    //VALUE (*size_fn)(ANYARGS);        // ruby 2.0
-};
-
 // note: the memory refed by this pointer is not freed and causes memroy leak.
 typedef struct {
     size_t  n; // the number of elements of the dimesnion
@@ -80,6 +62,7 @@ static ID id_dup;
 static ID id_bracket;
 static ID id_shift_left;
 static ID id_mask;
+static ID id_where;
 
 
 static void
@@ -142,25 +125,66 @@ na_parse_array(VALUE ary, int orig_dim, ssize_t size, na_index_arg_t *q)
 static void
 na_parse_narray_index(VALUE a, int orig_dim, ssize_t size, na_index_arg_t *q)
 {
-    VALUE idx;
-    narray_t *na;
-    narray_data_t *nidx;
+    VALUE idx, cls;
+    narray_t *na, *nidx;
     size_t k, n;
-    ssize_t *nidxp;
 
     GetNArray(a,na);
     if (NA_NDIM(na) != 1) {
         rb_raise(rb_eIndexError, "should be 1-d NArray");
     }
-    n = NA_SIZE(na);
-    idx = nary_new(cIndex,1,&n);
-    na_store(idx,a);
-
-    GetNArrayData(idx,nidx);
-    nidxp   = (ssize_t*)nidx->ptr;
-    q->idx  = ALLOC_N(size_t, n);
-    for (k=0; k<n; k++) {
-        q->idx[k] = na_range_check(nidxp[k], size, orig_dim);
+    cls = rb_obj_class(a);
+    if (cls==numo_cBit) {
+        if (NA_SIZE(na) != (size_t)size) {
+            rb_raise(rb_eIndexError, "Bit-NArray size mismatch");
+        }
+        idx = rb_funcall(a,id_where,0);
+        GetNArray(idx,nidx);
+        n = NA_SIZE(nidx);
+        q->idx = ALLOC_N(size_t, n);
+        if (na->type!=NARRAY_DATA_T) {
+            rb_bug("NArray#where returned wrong type of NArray");
+        }
+        if (rb_obj_class(idx)==numo_cInt32) {
+            int32_t *p = (int32_t*)NA_DATA_PTR(nidx);
+            for (k=0; k<n; k++) {
+                q->idx[k] = (size_t)p[k];
+            }
+        } else
+        if (rb_obj_class(idx)==numo_cInt64) {
+            int64_t *p = (int64_t*)NA_DATA_PTR(nidx);
+            for (k=0; k<n; k++) {
+                q->idx[k] = (size_t)p[k];
+            }
+        } else {
+            rb_bug("NArray#where should return Int32 or Int64");
+        }
+        RB_GC_GUARD(idx);
+    } else {
+        n = NA_SIZE(na);
+        q->idx = ALLOC_N(size_t, n);
+        if (cls==numo_cInt32 && na->type==NARRAY_DATA_T) {
+            int32_t *p = (int32_t*)NA_DATA_PTR(na);
+            for (k=0; k<n; k++) {
+                q->idx[k] = na_range_check(p[k], size, orig_dim);
+            }
+        } else
+        if (cls==numo_cInt64 && na->type==NARRAY_DATA_T) {
+            int64_t *p = (int64_t*)NA_DATA_PTR(na);
+            for (k=0; k<n; k++) {
+                q->idx[k] = na_range_check(p[k], size, orig_dim);
+            }
+        } else {
+            ssize_t *p;
+            idx = nary_new(cIndex,1,&n);
+            na_store(idx,a);
+            GetNArray(idx,nidx);
+            p = (ssize_t*)NA_DATA_PTR(nidx);
+            for (k=0; k<n; k++) {
+                q->idx[k] = na_range_check(p[k], size, orig_dim);
+            }
+            RB_GC_GUARD(idx);
+        }
     }
     q->n    = n;
     q->beg  = 0;
@@ -173,10 +197,46 @@ static void
 na_parse_range(VALUE range, ssize_t step, int orig_dim, ssize_t size, na_index_arg_t *q)
 {
     int n;
-    VALUE excl_end;
     ssize_t beg, end, beg_orig, end_orig;
     const char *dot = "..", *edot = "...";
 
+#ifdef HAVE_RB_ARITHMETIC_SEQUENCE_EXTRACT
+    rb_arithmetic_sequence_components_t x;
+    rb_arithmetic_sequence_extract(range, &x);
+    step = NUM2SSIZET(x.step);
+
+    beg = beg_orig = NUM2SSIZET(x.begin);
+    if (beg < 0) {
+        beg += size;
+    }
+    if (T_NIL == TYPE(x.end)) { // endless range
+        end = size - 1;
+        if (RTEST(x.exclude_end)) {
+            dot = edot;
+        }
+        if (beg < 0 || beg >= size) {
+            rb_raise(rb_eRangeError,
+                     "%"SZF"d%s is out of range for size=%"SZF"d",
+                     beg_orig, dot, size);
+        }
+    } else {
+        end = end_orig = NUM2SSIZET(x.end);
+        if (end < 0) {
+            end += size;
+        }
+        if (RTEST(x.exclude_end)) {
+            end--;
+            dot = edot;
+        }
+        if (beg < 0 || beg >= size || end < 0 || end >= size) {
+            rb_raise(rb_eRangeError,
+                     "%"SZF"d%s%"SZF"d is out of range for size=%"SZF"d",
+                     beg_orig, dot, end_orig, size);
+        }
+    }
+#else
+    VALUE excl_end;
+
     beg = beg_orig = NUM2SSIZET(rb_funcall(range,id_beg,0));
     if (beg < 0) {
         beg += size;
@@ -195,17 +255,18 @@ na_parse_range(VALUE range, ssize_t step, int orig_dim, ssize_t size, na_index_a
                  "%"SZF"d%s%"SZF"d is out of range for size=%"SZF"d",
                  beg_orig, dot, end_orig, size);
     }
+#endif
     n = (end-beg)/step+1;
     if (n<0) n=0;
     na_index_set_step(q,orig_dim,n,beg,step);
 
 }
 
-static void
-na_parse_enumerator(VALUE enum_obj, int orig_dim, ssize_t size, na_index_arg_t *q)
+void
+na_parse_enumerator_step(VALUE enum_obj, VALUE *pstep )
 {
     int len;
-    ssize_t step;
+    VALUE step;
     struct enumerator *e;
 
     if (!RB_TYPE_P(enum_obj, T_DATA)) {
@@ -213,26 +274,40 @@ na_parse_enumerator(VALUE enum_obj, int orig_dim, ssize_t size, na_index_arg_t *
     }
     e = (struct enumerator *)DATA_PTR(enum_obj);
 
-    if (rb_obj_is_kind_of(e->obj, rb_cRange)) {
-        if (e->meth == id_each) {
-            na_parse_range(e->obj, 1, orig_dim, size, q);
+    if (!rb_obj_is_kind_of(e->obj, rb_cRange)) {
+        rb_raise(rb_eTypeError,"not Range object");
+    }
+
+    if (e->meth == id_each) {
+        step = INT2NUM(1);
+    }
+    else if (e->meth == id_step) {
+        if (TYPE(e->args) != T_ARRAY) {
+            rb_raise(rb_eArgError,"no argument for step");
         }
-        else if (e->meth == id_step) {
-            if (TYPE(e->args) != T_ARRAY) {
-                rb_raise(rb_eArgError,"no argument for step");
-            }
-            len = RARRAY_LEN(e->args);
-            if (len != 1) {
-                rb_raise(rb_eArgError,"invalid number of step argument (1 for %d)",len);
-            }
-            step = NUM2SSIZET(RARRAY_AREF(e->args,0));
-            na_parse_range(e->obj, step, orig_dim, size, q);
-        } else {
-            rb_raise(rb_eTypeError,"unknown Range method: %s",rb_id2name(e->meth));
+        len = RARRAY_LEN(e->args);
+        if (len != 1) {
+            rb_raise(rb_eArgError,"invalid number of step argument (1 for %d)",len);
         }
+        step = RARRAY_AREF(e->args,0);
     } else {
-        rb_raise(rb_eTypeError,"not Range object");
+        rb_raise(rb_eTypeError,"unknown Range method: %s",rb_id2name(e->meth));
+    }
+    if (pstep) *pstep = step;
+}
+
+static void
+na_parse_enumerator(VALUE enum_obj, int orig_dim, ssize_t size, na_index_arg_t *q)
+{
+    VALUE step;
+    struct enumerator *e;
+
+    if (!RB_TYPE_P(enum_obj, T_DATA)) {
+        rb_raise(rb_eTypeError,"wrong argument type (not T_DATA)");
     }
+    na_parse_enumerator_step(enum_obj, &step);
+    e = (struct enumerator *)DATA_PTR(enum_obj);
+    na_parse_range(e->obj, NUM2SSIZET(step), orig_dim, size, q); // e->obj : Range Object
 }
 
 // Analyze *a* which is *i*-th index object and store the information to q
@@ -289,14 +364,15 @@ na_index_parse_each(volatile VALUE a, ssize_t size, int i, na_index_arg_t *q)
         if (rb_obj_is_kind_of(a, rb_cRange)) {
             na_parse_range(a, 1, i, size, q);
         }
+#ifdef HAVE_RB_ARITHMETIC_SEQUENCE_EXTRACT
+        else if (rb_obj_is_kind_of(a, rb_cArithSeq)) {
+            //na_parse_arith_seq(a, i, size, q);
+            na_parse_range(a, 1, i, size, q);
+        }
+#endif
         else if (rb_obj_is_kind_of(a, rb_cEnumerator)) {
             na_parse_enumerator(a, i, size, q);
         }
-        else if (rb_obj_is_kind_of(a, na_cStep)) {
-            ssize_t beg, step, n;
-            nary_step_array_index(a, size, (size_t*)(&n), &beg, &step);
-            na_index_set_step(q,i,n,beg,step);
-        }
         // NArray index
         else if (NA_IsNArray(a)) {
             na_parse_narray_index(a, i, size, q);
@@ -308,6 +384,102 @@ na_index_parse_each(volatile VALUE a, ssize_t size, int i, na_index_arg_t *q)
 }
 
 
+static void
+na_at_parse_each(volatile VALUE a, ssize_t size, int i, VALUE *idx, ssize_t stride)
+{
+    na_index_arg_t q;
+    size_t n, k;
+    ssize_t *index;
+
+    // NArray index
+    if (NA_IsNArray(a)) {
+        VALUE a2;
+        narray_t *na, *na2;
+        ssize_t *p2;
+        GetNArray(a,na);
+        if (NA_NDIM(na) != 1) {
+            rb_raise(rb_eIndexError, "should be 1-d NArray");
+        }
+        n = NA_SIZE(na);
+        a2 = nary_new(cIndex,1,&n);
+        na_store(a2,a);
+        GetNArray(a2,na2);
+        p2 = (ssize_t*)NA_DATA_PTR(na2);
+        if (*idx == Qnil) {
+            *idx = a2;
+            for (k=0; k<n; k++) {
+                na_range_check(p2[k],size,i);
+            }
+        } else {
+            narray_t *nidx;
+            GetNArray(*idx,nidx);
+            index = (ssize_t*)NA_DATA_PTR(nidx);
+            if (NA_SIZE(nidx) != n) {
+                rb_raise(nary_eShapeError, "index array sizes mismatch");
+            }
+            for (k=0; k<n; k++) {
+                index[k] += na_range_check(p2[k],size,i) * stride;
+            }
+        }
+        RB_GC_GUARD(a2);
+        return;
+    }
+    else if (TYPE(a) == T_ARRAY) {
+        n = RARRAY_LEN(a);
+        if (*idx == Qnil) {
+            *idx = nary_new(cIndex,1,&n);
+            index = (ssize_t*)na_get_pointer_for_write(*idx); // allocate memory
+            for (k=0; k<n; k++) {
+                index[k] = na_range_check(NUM2SSIZET(RARRAY_AREF(a,k)),size,i);
+            }
+        } else {
+            narray_t *nidx;
+            GetNArray(*idx,nidx);
+            index = (ssize_t*)NA_DATA_PTR(nidx);
+            if (NA_SIZE(nidx) != n) {
+                rb_raise(nary_eShapeError, "index array sizes mismatch");
+            }
+            for (k=0; k<n; k++) {
+                index[k] += na_range_check(NUM2SSIZET(RARRAY_AREF(a,k)),size,i) * stride;
+            }
+        }
+        return;
+    }
+    else if (rb_obj_is_kind_of(a, rb_cRange)) {
+        na_parse_range(a, 1, i, size, &q);
+    }
+#ifdef HAVE_RB_ARITHMETIC_SEQUENCE_EXTRACT
+    else if (rb_obj_is_kind_of(a, rb_cArithSeq)) {
+        na_parse_range(a, 1, i, size, &q);
+    }
+#endif
+    else if (rb_obj_is_kind_of(a, rb_cEnumerator)) {
+        na_parse_enumerator(a, i, size, &q);
+    }
+    else {
+        rb_raise(rb_eIndexError, "not allowed type");
+    }
+
+    if (*idx == Qnil) {
+        *idx = nary_new(cIndex,1,&q.n);
+        index = (ssize_t*)na_get_pointer_for_write(*idx); // allocate memory
+        for (k=0; k<q.n; k++) {
+            index[k] = q.beg + q.step*k;
+        }
+    } else {
+        narray_t *nidx;
+        GetNArray(*idx,nidx);
+        index = (ssize_t*)NA_DATA_PTR(nidx);
+        if (NA_SIZE(nidx) != q.n) {
+            rb_raise(nary_eShapeError, "index array sizes mismatch");
+        }
+        for (k=0; k<q.n; k++) {
+            index[k] += (q.beg + q.step*k) * stride;
+        }
+    }
+}
+
+
 static size_t
 na_index_parse_args(VALUE args, narray_t *na, na_index_arg_t *q, int ndim)
 {
@@ -574,7 +746,7 @@ VALUE na_aref_md_protected(VALUE data_value)
 
     na_alloc_shape((narray_t*)na2, ndim_new);
 
-    na2->stridx = ALLOC_N(stridx_t,ndim_new);
+    na2->stridx = ZALLOC_N(stridx_t,ndim_new);
 
     elmsz = nary_element_stride(self);
 
@@ -842,11 +1014,148 @@ static VALUE na_slice(int argc, VALUE *argv, VALUE self)
     return na_aref_main(argc, argv, self, 1, nd);
 }
 
+/*
+  Multi-dimensional element reference.
+  Returns an element at `dim0`, `dim1`, ... are Numeric indices for each dimension, or returns a NArray View as a sliced array if `dim0`, `dim1`, ... includes other than Numeric index, e.g., Range or Array or true.
+  @overload [](dim0,...,dimL)
+  @param [Numeric,Range,Array,Numo::Int32,Numo::Int64,Numo::Bit,TrueClass,FalseClass,Symbol] dim0,...,dimL  multi-dimensional indices.
+  @return [Numeric,Numo::NArray] an element or NArray view.
+  @see #[]=
+  @see #at
+
+  @example
+      a = Numo::DFloat.new(4,5).seq
+      # => Numo::DFloat#shape=[4,5]
+      # [[0, 1, 2, 3, 4],
+      #  [5, 6, 7, 8, 9],
+      #  [10, 11, 12, 13, 14],
+      #  [15, 16, 17, 18, 19]]
+
+      a[1,1]
+      # => 6.0
+
+      a[1..3,1]
+      # => Numo::DFloat#shape=[3]
+      # [6, 11, 16]
+
+      a[1,[1,3,4]]
+      # => Numo::DFloat#shape=[3]
+      # [6, 8, 9]
+
+      a[true,2].fill(99)
+      a
+      # => Numo::DFloat#shape=[4,5]
+      # [[0, 1, 99, 3, 4],
+      #  [5, 6, 99, 8, 9],
+      #  [10, 11, 99, 13, 14],
+      #  [15, 16, 99, 18, 19]]
+ */
+static VALUE na_aref(int argc, VALUE *argv, VALUE self)
+{
+    // implemented in subclasses
+    return rb_f_notimplement(argc,argv,self);
+}
+
+/*
+  Multi-dimensional element assignment.
+  Replace element(s) at `dim0`, `dim1`, ... .
+  Broadcasting mechanism is applied.
+  @overload []=(dim0,...,dimL,val)
+  @param [Numeric,Range,Array,Numo::Int32,Numo::Int64,Numo::Bit,TrueClass,FalseClass,Symbol] dim0,...,dimL  multi-dimensional indices.
+  @param [Numeric,Numo::NArray,Array] val  Value(s) to be set to self.
+  @return [Numeric,Numo::NArray,Array] returns `val` (last argument).
+  @see #[]
+  @example
+      a = Numo::DFloat.new(3,4).seq
+      # => Numo::DFloat#shape=[3,4]
+      # [[0, 1, 2, 3],
+      #  [4, 5, 6, 7],
+      #  [8, 9, 10, 11]]
+
+      a[1,2]=99
+      a
+      # => Numo::DFloat#shape=[3,4]
+      # [[0, 1, 2, 3],
+      #  [4, 5, 99, 7],
+      #  [8, 9, 10, 11]]
+
+      a[1,[0,2]] = [101,102]
+      a
+      # => Numo::DFloat#shape=[3,4]
+      # [[0, 1, 2, 3],
+      #  [101, 5, 102, 7],
+      #  [8, 9, 10, 11]]
+
+      a[1,true]=99
+      a
+      # => Numo::DFloat#shape=[3,4]
+      # [[0, 1, 2, 3],
+      #  [99, 99, 99, 99],
+      #  [8, 9, 10, 11]]
+
+*/
+static VALUE na_aset(int argc, VALUE *argv, VALUE self)
+{
+    // implemented in subclasses
+    return rb_f_notimplement(argc,argv,self);
+}
+
+/*
+  Multi-dimensional array indexing.
+  Similar to numpy's tuple indexing, i.e., `a[[1,2,..],[3,4,..]]`
+  Same as Numo::NArray#[] for one-dimensional NArray.
+  @overload at(dim0,...,dimL)
+  @param [Range,Array,Numo::Int32,Numo::Int64] dim0,...,dimL  multi-dimensional index arrays.
+  @return [Numo::NArray] one-dimensional NArray view.
+  @see #[]
+
+  @example
+      x = Numo::DFloat.new(3,3,3).seq
+      # => Numo::DFloat#shape=[3,3,3]
+      #  [[[0, 1, 2],
+      #    [3, 4, 5],
+      #    [6, 7, 8]],
+      #   [[9, 10, 11],
+      #    [12, 13, 14],
+      #    [15, 16, 17]],
+      #   [[18, 19, 20],
+      #    [21, 22, 23],
+      #    [24, 25, 26]]]
+
+      x.at([0,1,2],[0,1,2],[-1,-2,-3])
+      # => Numo::DFloat(view)#shape=[3]
+      #  [2, 13, 24]
+ */
+static VALUE na_at(int argc, VALUE *argv, VALUE self)
+{
+    int i;
+    size_t n;
+    ssize_t stride=1;
+    narray_t *na;
+    VALUE idx=Qnil;
+
+    na_index_arg_to_internal_order(argc, argv, self);
+
+    GetNArray(self,na);
+    if (NA_NDIM(na) != argc) {
+        rb_raise(rb_eArgError,"the number of argument must be same as dimension");
+    }
+    for (i=argc; i>0; ) {
+        i--;
+        n = NA_SHAPE(na)[i];
+        na_at_parse_each(argv[i], n, i, &idx, stride);
+        stride *= n;
+    }
+    return na_aref_main(1, &idx, self, 1, 1);
+}
 
 void
 Init_nary_index()
 {
     rb_define_method(cNArray, "slice", na_slice, -1);
+    rb_define_method(cNArray, "[]", na_aref, -1);
+    rb_define_method(cNArray, "[]=", na_aset, -1);
+    rb_define_method(cNArray, "at", na_at, -1);
 
     sym_ast        = ID2SYM(rb_intern("*"));
     sym_all        = ID2SYM(rb_intern("all"));
@@ -867,4 +1176,5 @@ Init_nary_index()
     id_bracket     = rb_intern("[]");
     id_shift_left  = rb_intern("<<");
     id_mask        = rb_intern("mask");
+    id_where       = rb_intern("where");
 }