carray 3.0.4 → 3.0.6

data/ext/carray_slab.c

@@ -13,7 +13,8 @@
 
   Internals:
     - ca_slab_iter_state_t (see carray_slab.h) is the per-call state struct.
-    - Each entry point stack-allocates a zeroed state, populates it via
+    - Each entry point heap-allocates a zeroed state under a window guard
+      (slab_state_new), populates it via
       slab_state_init (= nofail; only validation + scalar field fills),
       then rb_ensure-runs slab_state_run_body (= acquires T1 substrate +
       scratch + output, runs per-form loop) and slab_state_finish (= frees
@@ -21,9 +22,8 @@
     - All heap acquisition lives in run_body so init-time rb_raise paths
       cannot leak; rb_ensure guarantees finish on both normal and
       exception exit.
-    - Conservative stack scan keeps the state struct's VALUE fields alive
-      across rb_yield; no dmark callback is required because the state is
-      not Ruby-wrapped.
+    - The guard marks the state's VALUE fields, and runs slab_state_finish
+      from a finalizer if the block's Fiber is abandoned mid-walk.
 
   --------------------------------------------------------------------------- */
 
@@ -122,9 +122,9 @@ slab_state_init (ca_slab_iter_state_t *st, VALUE source, int8_t form,
 {
   CArray *src;
 
-  /* Caller stack-allocates with `{0}` which sets VALUE fields to Qfalse
-     (= 0). Body code uses `== Qnil` to distinguish unset slots from
-     populated CArray handles, so we must promote them to Qnil here. */
+  /* The state arrives zeroed, which leaves VALUE fields Qfalse (= 0).
+     Body code uses `== Qnil` to distinguish unset slots from populated
+     CArray handles, so we must promote them to Qnil here. */
   st->self             = Qnil;
   st->slab_view        = Qnil;
   st->carrier          = Qnil;
@@ -346,6 +346,26 @@ ca_slab_build_output_view (int8_t data_type, ca_size_t bytes,
   return TypedData_Wrap_Struct(rb_cCAWrap, &cawrap_data_type, cw);
 }
 
+/* The array the slab views move onto when the walk ends (see
+   slab_state_release_views), made with the carrier and held by it.  It is
+   made here rather than at the end: the end runs in an ensure, after the
+   block, and an allocation failing there would leave the walk open and
+   the views pointing into freed scratch. */
+static void
+ca_slab_prepare_backing (ca_slab_iter_state_t *st)
+{
+  CArray *cw, *backing;
+  volatile VALUE bv;
+  TypedData_Get_Struct(st->carrier, CArray, &carray_data_type, cw);
+  bv = rb_carray_new(cw->data_type, cw->ndim, cw->dim, cw->bytes, NULL);
+  TypedData_Get_Struct(bv, CArray, &carray_data_type, backing);
+  if ( cw->mask ) {
+    ca_create_mask(backing);
+  }
+  rb_ivar_set(st->carrier, rb_intern("__slab_backing__"), bv);
+  RB_GC_GUARD(bv);
+}
+
 /* Common setup of the input slab view for all 4 forms.  Decides ALIAS
    vs SCRATCH from T1's slab_strides[0] vs src->bytes equality, builds
    the slab_view + (when SCRATCH) carrier in place.  The mask travels
@@ -465,6 +485,10 @@ ca_slab_setup_input_slab_view (ca_slab_iter_state_t *st, CArray *src,
     }
   }
 
+  if ( st->carrier != Qnil ) {
+    ca_slab_prepare_backing(st);
+  }
+
   /* The slab is a window onto the source when the fiber is contiguous and
      a copy in the iterator's scratch when it is not, so writing through it
      landed on one axis and vanished on the next.  Nothing in the surface
@@ -502,8 +526,12 @@ ca_slab_advance_slab_view (ca_slab_iter_state_t *st,
     if ( slab_cs->mask != NULL ) {
       CAStride *mask_cs = (CAStride *) slab_cs->mask;
       if ( m_in != NULL ) {
+        /* m_in points into the walk's copy of the source's mask, not into
+           the source's mask itself, so the offset is taken from the data
+           cell: it is what the slab reads once the copy is gone. */
         mask_cs->ptr         = (char *) m_in;
-        mask_cs->base_offset = (ca_size_t) ((char *) m_in - src->mask->ptr);
+        mask_cs->base_offset = (ca_size_t) ((p_in - src->ptr) / src->bytes)
+                               * src->mask->bytes;
       }
       else {
         /* The walk was opened on a source without a mask and the block
@@ -861,7 +889,6 @@ ca_slab_reduce_setup (ca_slab_iter_state_t *st, CArray **out_src,
                       ca_size_t *out_n)
 {
   CArray *src;
-  CArray *out;
   ca_size_t fiber_len;
   ca_size_t out_dim[CA_RANK_MAX];
   int8_t out_ndim;
@@ -904,7 +931,6 @@ ca_slab_reduce_setup (ca_slab_iter_state_t *st, CArray **out_src,
 
   st->output = rb_carray_new(st->out_data_type, out_ndim, out_dim,
                              st->out_bytes, NULL);
-  TypedData_Get_Struct(st->output, CArray, &carray_data_type, out);
 
   /* FIBER_CONTIG is naxes==1 only (T1 substrate contract).  Multi-axis
      omits the flag and uses bare CA_SLAB_AXES; T1 picks an in-src
@@ -1178,8 +1204,11 @@ slab_state_run_body (VALUE arg)
    as it has a ptr.  Before the scratch they point into is freed, the
    input slab and its carrier are moved onto a copy of the last slab,
    held by the carrier, and the internal output view lets go of its
-   buffer.  A slab that is a window onto the source itself points into
-   the source and needs nothing. */
+   buffer.  The copy is made into the array ca_slab_prepare_backing set
+   aside, so nothing here allocates.  A slab that is a window onto the
+   source keeps reading the source's data, but its mask points into the
+   walk's copy of the mask, which finish frees: it lets go of that copy
+   and reads the source's mask from then on. */
 static void
 slab_state_release_views (ca_slab_iter_state_t *st)
 {
@@ -1188,6 +1217,13 @@ slab_state_release_views (ca_slab_iter_state_t *st)
     TypedData_Get_Struct(st->output_slab_view, CArray, &carray_data_type, ov);
     ov->ptr = NULL;
   }
+  if ( st->carrier == Qnil && st->slab_view != Qnil ) {
+    CArray *cs;
+    TypedData_Get_Struct(st->slab_view, CArray, &carray_data_type, cs);
+    if ( cs->mask ) {
+      cs->mask->ptr = NULL;
+    }
+  }
   if ( st->carrier != Qnil && st->slab_view != Qnil ) {
     CArray *cw, *cs, *backing;
     volatile VALUE bv;
@@ -1196,12 +1232,12 @@ slab_state_release_views (ca_slab_iter_state_t *st)
     if ( cw->ptr == NULL ) {
       return;
     }
-    backing = ca_copy(cw);
-    bv = ca_wrap_struct(backing);
-    if ( cw->mask && ! backing->mask ) {
-      ca_create_mask(backing);
+    bv = rb_ivar_get(st->carrier, rb_intern("__slab_backing__"));
+    TypedData_Get_Struct(bv, CArray, &carray_data_type, backing);
+    memcpy(backing->ptr, cw->ptr, (size_t) (cw->elements * cw->bytes));
+    if ( cw->mask && cw->mask->ptr ) {
+      memcpy(backing->mask->ptr, cw->mask->ptr, (size_t) cw->elements);
     }
-    rb_ivar_set(st->carrier, rb_intern("__slab_backing__"), bv);
     if ( cs->ptr ) {
       cs->ptr = backing->ptr;
     }
@@ -1252,6 +1288,72 @@ slab_state_finish (VALUE arg)
 /* Entry points: CArray#map_slab / #reduce_slab / #each_slab           */
 /* ------------------------------------------------------------------- */
 
+/* The walk yields to the block with its windows open.  The state lives on
+   the heap under a window guard, so that if the block's Fiber is abandoned
+   (an Enumerator stopped with #next and dropped) the guard's finalizer
+   still runs slab_state_finish and the source is let go. */
+
+static void
+slab_state_mark (void *p)
+{
+  ca_slab_iter_state_t *st = (ca_slab_iter_state_t *) p;
+  rb_gc_mark(st->self);
+  rb_gc_mark(st->slab_view);
+  rb_gc_mark(st->carrier);
+  rb_gc_mark(st->mask_carrier);
+  rb_gc_mark(st->output);
+  rb_gc_mark(st->output_slab_view);
+  rb_gc_mark(st->init_val);
+}
+
+static void
+slab_state_finish_abandoned (void *p)
+{
+  slab_state_finish((VALUE)(uintptr_t) p);
+}
+
+typedef struct {
+  ca_slab_iter_state_t *st;
+  VALUE guard;
+  VALUE sentinel;
+} slab_run_t;
+
+static VALUE
+slab_run_body (VALUE arg)
+{
+  return slab_state_run_body((VALUE)(uintptr_t) ((slab_run_t *) arg)->st);
+}
+
+static VALUE
+slab_run_close (VALUE arg)
+{
+  slab_run_t *r = (slab_run_t *) arg;
+  ca_window_guard_disarm(r->guard, r->sentinel);
+  return slab_state_finish((VALUE)(uintptr_t) r->st);
+}
+
+/* Allocates the state; slab_state_init fills it, then slab_run walks. */
+static ca_slab_iter_state_t *
+slab_state_new (slab_run_t *r)
+{
+  ca_slab_iter_state_t *st = ZALLOC(ca_slab_iter_state_t);
+  st->self = st->slab_view = st->carrier = st->mask_carrier = Qnil;
+  st->output = st->output_slab_view = st->init_val = Qnil;
+  r->st    = st;
+  r->guard = ca_window_guard_arm(slab_state_finish_abandoned, slab_state_mark,
+                                 st, &r->sentinel);
+  return st;
+}
+
+static VALUE
+slab_run (slab_run_t *r)
+{
+  VALUE result = rb_ensure(slab_run_body, (VALUE) r, slab_run_close, (VALUE) r);
+  RB_GC_GUARD(r->guard);
+  RB_GC_GUARD(r->sentinel);
+  return result;
+}
+
 /* Shared kwarg parse: returns axis (required) + data_type / init via
    out-params.  init_out is set to Qundef if the `:init` key was absent
    (= caller distinguishes per-slab vs per-element reduce forms).      */
@@ -1299,7 +1401,7 @@ slab_extract_kwargs (int argc, VALUE *argv, const char *method_name)
 static VALUE
 rb_ca_map_slab (int argc, VALUE *argv, VALUE self)
 {
-  ca_slab_iter_state_t st = {0};
+  slab_run_t r;
   VALUE kw, axis, data_type, init_unused;
   VALUE result;
 
@@ -1312,10 +1414,9 @@ rb_ca_map_slab (int argc, VALUE *argv, VALUE self)
     rb_raise(rb_eLocalJumpError, "CArray#map_slab: block required");
   }
 
-  slab_state_init(&st, self, CA_SLAB_FORM_MAP, axis, data_type, Qundef);
+  slab_state_init(slab_state_new(&r), self, CA_SLAB_FORM_MAP, axis, data_type, Qundef);
 
-  result = rb_ensure(slab_state_run_body, (VALUE)(uintptr_t)&st,
-                     slab_state_finish,   (VALUE)(uintptr_t)&st);
+  result = slab_run(&r);
   RB_GC_GUARD(self);
   return result;
 }
@@ -1323,7 +1424,7 @@ rb_ca_map_slab (int argc, VALUE *argv, VALUE self)
 static VALUE
 rb_ca_reduce_slab (int argc, VALUE *argv, VALUE self)
 {
-  ca_slab_iter_state_t st = {0};
+  slab_run_t r;
   VALUE kw, axis, data_type, init_val;
   VALUE result;
   int8_t form;
@@ -1341,10 +1442,9 @@ rb_ca_reduce_slab (int argc, VALUE *argv, VALUE self)
   form = (init_val == Qundef) ? CA_SLAB_FORM_REDUCE_SLAB
                               : CA_SLAB_FORM_REDUCE_FIBER;
 
-  slab_state_init(&st, self, form, axis, data_type, init_val);
+  slab_state_init(slab_state_new(&r), self, form, axis, data_type, init_val);
 
-  result = rb_ensure(slab_state_run_body, (VALUE)(uintptr_t)&st,
-                     slab_state_finish,   (VALUE)(uintptr_t)&st);
+  result = slab_run(&r);
   RB_GC_GUARD(self);
   return result;
 }
@@ -1352,7 +1452,7 @@ rb_ca_reduce_slab (int argc, VALUE *argv, VALUE self)
 static VALUE
 rb_ca_each_slab (int argc, VALUE *argv, VALUE self)
 {
-  ca_slab_iter_state_t st = {0};
+  slab_run_t r;
   VALUE kw, axis, data_type_unused, init_unused;
   VALUE result;
 
@@ -1374,10 +1474,9 @@ rb_ca_each_slab (int argc, VALUE *argv, VALUE self)
                           2, args, RB_PASS_KEYWORDS);
   }
 
-  slab_state_init(&st, self, CA_SLAB_FORM_EACH, axis, Qnil, Qundef);
+  slab_state_init(slab_state_new(&r), self, CA_SLAB_FORM_EACH, axis, Qnil, Qundef);
 
-  result = rb_ensure(slab_state_run_body, (VALUE)(uintptr_t)&st,
-                     slab_state_finish,   (VALUE)(uintptr_t)&st);
+  result = slab_run(&r);
   RB_GC_GUARD(self);
   return result;
 }

data/ext/carray_sort.c

@@ -103,16 +103,207 @@ sort_addr_merge (struct sort_addr_key *a, struct sort_addr_key *aux,
                  ca_size_t n, int (*cmp)(const void *, const void *))
 CA_SORT_MERGE_BY_CMP(struct sort_addr_key)
 
+/* -------------------------------------------------------------------
+
+  Multi-key sort by refinement (MSD), for keys of the numeric data types.
+
+  The first key is sorted whole by the typed pair kernel
+  (ca_sort_quick_pair_*); then only the runs of rows that tie on it are
+  sorted by the second key, and so on, until no run is left or the keys
+  run out.  The answer is the one sort_addr_cmp gives:
+
+  - a pair sorts by (value, row), so rows that tie on every key stay in
+    row order (the stable tie-break of sort_addr_cmp);
+  - each run is split into numbers, NaN and masked cells, in the order
+    sort_addr_cmp puts them (numbers, NaN, masked for masked_last;
+    masked, numbers, NaN otherwise).  The NaN and the masked cells of a
+    run tie among themselves (CA_FLOAT_ORDER_CMP / sort_addr_cmp), so
+    each is a run for the next key.  The split is a stable one done in
+    the load, not ca_partition_nan_pair_*, which allocates on every call
+    and is called here once per run;
+  - -0.0 and 0.0 tie, by == as by CA_FLOAT_ORDER_CMP.
+
+  A run is (lo, n): the rows perm[lo..lo+n).  Within a run the rows are
+  in row order, since they tie on every key before it, so the stable
+  split keeps them so.
+
+------------------------------------------------------------------- */
+
+typedef struct {
+  ca_size_t lo;
+  ca_size_t n;
+} msd_run_t;
+
+typedef struct {
+  msd_run_t *runs;
+  ca_size_t  count;
+} msd_runs_t;
+
+static void
+msd_add_run (msd_runs_t *out, ca_size_t lo, ca_size_t n)
+{
+  if ( out && n > 1 ) {
+    out->runs[out->count].lo = lo;
+    out->runs[out->count].n  = n;
+    out->count++;
+  }
+}
+
+/* Sort each run of +in+ by +key+, writing the new order into perm, and
+   add the runs that still tie to +out+ (NULL for the last key). */
+#define DEFINE_MSD_REFINE(TYPE, SUFFIX, IS_FLOAT)                             \
+static void                                                                   \
+msd_refine_##SUFFIX (CArray *key, ca_size_t *perm, void *vbuf,                \
+                     msd_runs_t *in, msd_runs_t *out, int masked_last)        \
+{                                                                             \
+  const CA_U(TYPE) *ptr = (const CA_U(TYPE) *) key->ptr;                      \
+  const boolean8_t *m = ( key->mask ) ? (boolean8_t *) key->mask->ptr : NULL; \
+  ca_pair_##SUFFIX *buf = (ca_pair_##SUFFIX *) vbuf;                          \
+  ca_size_t r;                                                                \
+                                                                              \
+  for (r=0; r<in->count; r++) {                                               \
+    ca_size_t lo = in->runs[r].lo;                                            \
+    ca_size_t n  = in->runs[r].n;                                             \
+    ca_size_t *p = perm + lo;                                                 \
+    ca_size_t n_num = 0, n_nan = 0, n_mask = 0;                               \
+    ca_size_t o_num, o_nan, o_mask;                                           \
+    ca_size_t k, s;                                                           \
+                                                                              \
+    for (k=0; k<n; k++) {                                                     \
+      ca_size_t row = p[k];                                                   \
+      if ( m && m[row] )                 n_mask++;                            \
+      else if ( IS_FLOAT && isnan((double) ptr[row]) ) n_nan++;               \
+      else                               n_num++;                             \
+    }                                                                         \
+    o_mask = masked_last ? n_num + n_nan : 0;                                 \
+    o_num  = masked_last ? 0 : n_mask;                                        \
+    o_nan  = o_num + n_num;                                                   \
+    {                                                                         \
+      ca_size_t i_num = o_num, i_nan = o_nan, i_mask = o_mask;                \
+      for (k=0; k<n; k++) {                                                   \
+        ca_size_t row = p[k];                                                 \
+        TYPE v = ptr[row];                                                    \
+        ca_size_t at;                                                         \
+        if ( m && m[row] )             at = i_mask++;                         \
+        else if ( IS_FLOAT && isnan((double) v) ) at = i_nan++;               \
+        else                           at = i_num++;                          \
+        buf[at].v = v;                                                        \
+        buf[at].i = row;                                                      \
+      }                                                                       \
+    }                                                                         \
+    ca_sort_quick_pair_##SUFFIX(buf + o_num, n_num);                          \
+    for (k=0; k<n; k++) {                                                     \
+      p[k] = buf[k].i;                                                        \
+    }                                                                         \
+    if ( ! out ) continue;                                                    \
+    for (s=0, k=1; k<=n_num; k++) {                                           \
+      if ( k == n_num || buf[o_num + k].v != buf[o_num + s].v ) {             \
+        msd_add_run(out, lo + o_num + s, k - s);                              \
+        s = k;                                                                \
+      }                                                                       \
+    }                                                                         \
+    msd_add_run(out, lo + o_nan,  n_nan);                                     \
+    msd_add_run(out, lo + o_mask, n_mask);                                    \
+  }                                                                           \
+}
+
+DEFINE_MSD_REFINE(int8_t,    i8,  0)
+DEFINE_MSD_REFINE(uint8_t,   u8,  0)
+DEFINE_MSD_REFINE(int16_t,   i16, 0)
+DEFINE_MSD_REFINE(uint16_t,  u16, 0)
+DEFINE_MSD_REFINE(int32_t,   i32, 0)
+DEFINE_MSD_REFINE(uint32_t,  u32, 0)
+DEFINE_MSD_REFINE(int64_t,   i64, 0)
+DEFINE_MSD_REFINE(uint64_t,  u64, 0)
+DEFINE_MSD_REFINE(float32_t, f32, 1)
+DEFINE_MSD_REFINE(double,    f64, 1)
+
+/* Whether every key has a typed pair kernel, so that MSD can sort them. */
+static int
+msd_is_sortable (struct cmp_base *base)
+{
+  int j;
+  for (j=0; j<base->n; j++) {
+    switch ( base->ca[j]->data_type ) {
+    case CA_BOOLEAN: case CA_INT8:  case CA_UINT8:
+    case CA_INT16:   case CA_UINT16:
+    case CA_INT32:   case CA_UINT32:
+    case CA_INT64:   case CA_UINT64:
+    case CA_FLOAT32: case CA_FLOAT64:
+      break;
+    default:
+      return 0;
+    }
+  }
+  return 1;
+}
+
+typedef struct {
+  struct cmp_base *base;
+  ca_size_t       *perm;      /* elements rows, the answer */
+  void            *buf;       /* elements pairs of 16 bytes */
+  msd_run_t       *runs[2];   /* elements / 2 + 1 runs each */
+  ca_size_t        elements;
+} msd_sort_t;
+
+static void
+msd_sort (msd_sort_t *s)
+{
+  struct cmp_base *base = s->base;
+  msd_runs_t a, b, *in = &a, *out = &b;
+  ca_size_t i;
+  int j;
+
+  for (i=0; i<s->elements; i++) {
+    s->perm[i] = i;
+  }
+  a.runs = s->runs[0];  a.count = 0;
+  b.runs = s->runs[1];  b.count = 0;
+  msd_add_run(in, 0, s->elements);
+
+  for (j=0; j<base->n && in->count > 0; j++) {
+    CArray *key = base->ca[j];
+    msd_runs_t *next = ( j < base->n - 1 ) ? out : NULL;
+    msd_runs_t *t;
+    if ( next ) next->count = 0;
+    switch ( key->data_type ) {
+    case CA_BOOLEAN:
+    case CA_INT8:    msd_refine_i8 (key, s->perm, s->buf, in, next, base->masked_last); break;
+    case CA_UINT8:   msd_refine_u8 (key, s->perm, s->buf, in, next, base->masked_last); break;
+    case CA_INT16:   msd_refine_i16(key, s->perm, s->buf, in, next, base->masked_last); break;
+    case CA_UINT16:  msd_refine_u16(key, s->perm, s->buf, in, next, base->masked_last); break;
+    case CA_INT32:   msd_refine_i32(key, s->perm, s->buf, in, next, base->masked_last); break;
+    case CA_UINT32:  msd_refine_u32(key, s->perm, s->buf, in, next, base->masked_last); break;
+    case CA_INT64:   msd_refine_i64(key, s->perm, s->buf, in, next, base->masked_last); break;
+    case CA_UINT64:  msd_refine_u64(key, s->perm, s->buf, in, next, base->masked_last); break;
+    case CA_FLOAT32: msd_refine_f32(key, s->perm, s->buf, in, next, base->masked_last); break;
+    case CA_FLOAT64: msd_refine_f64(key, s->perm, s->buf, in, next, base->masked_last); break;
+    default:
+      rb_raise(rb_eRuntimeError, "%s: [BUG] data type %d has no pair kernel",
+               ca_calling_method_name(), (int) key->data_type);
+    }
+    if ( ! next ) break;
+    t = in; in = out; out = t;
+  }
+}
+
+/* ----------------------------------------------------------------- */
+
 typedef struct {
   struct sort_addr_key *data;
   struct sort_addr_key *aux;     /* NULL: no object key, libc sorts */
   ca_size_t             elements;
+  msd_sort_t           *msd;     /* non-NULL: every key is numeric, sort by MSD */
 } sort_addr_run_t;
 
 static VALUE
 sort_addr_run (VALUE arg)
 {
   sort_addr_run_t *r = (sort_addr_run_t *) arg;
+  if ( r->msd ) {
+    msd_sort(r->msd);
+    return Qnil;
+  }
   if ( r->aux ) {
     sort_addr_merge(r->data, r->aux, r->elements,
                     (int (*)(const void*,const void*)) sort_addr_cmp);
@@ -145,8 +336,9 @@ rb_ca_s_sort_addr_kw (int argc, VALUE *argv, VALUE ropt, VALUE self)
   CArray *co;
   struct cmp_base base_s, *base = &base_s;
   struct sort_addr_key *data;
-  volatile VALUE hlist = 0, hdata = 0, haux = 0, hmodes = 0;
+  volatile VALUE hlist = 0, hdata = 0, haux = 0, hmodes = 0, hruns0 = 0, hruns1 = 0;
   sort_addr_run_t run;
+  msd_sort_t msd;
   char *modes;
   int has_object = 0;
   ca_size_t elements;
@@ -203,26 +395,42 @@ rb_ca_s_sort_addr_kw (int argc, VALUE *argv, VALUE ropt, VALUE self)
   }
   modes[argc] = '\0';
 
-  data = ALLOCV_N(struct sort_addr_key, hdata, elements);
-  for (i=0; i<elements; i++) {
-    data[i].i = i;
-    data[i].base = base;
-  }
-
-  run.data     = data;
-  run.aux      = has_object ? ALLOCV_N(struct sort_addr_key, haux, elements) : NULL;
-  run.elements = elements;
-  ca_attach_window(argc, base->ca, modes, sort_addr_run, (VALUE) &run);
-
   out = rb_ca_template_with_type(argv[0], INT2NUM(CA_SIZE), INT2NUM(0));
   TypedData_Get_Struct(out, CArray, &carray_data_type, co);
   q = (ca_size_t *) co->ptr;
-  
-  for (i=0; i<elements; i++) {
-    *q = data[i].i;
-    q++;
+
+  run.data     = NULL;
+  run.aux      = NULL;
+  run.elements = elements;
+  run.msd      = NULL;
+
+  if ( msd_is_sortable(base) ) {
+    ca_size_t nruns = elements / 2 + 1;
+    msd.base     = base;
+    msd.perm     = q;
+    msd.buf      = ALLOCV(hdata, (size_t) elements * sizeof(ca_pair_f64));
+    msd.runs[0]  = ALLOCV_N(msd_run_t, hruns0, nruns);
+    msd.runs[1]  = ( argc > 1 ) ? ALLOCV_N(msd_run_t, hruns1, nruns) : NULL;
+    msd.elements = elements;
+    run.msd      = &msd;
+    ca_attach_window(argc, base->ca, modes, sort_addr_run, (VALUE) &run);
+  }
+  else {
+    data = ALLOCV_N(struct sort_addr_key, hdata, elements);
+    for (i=0; i<elements; i++) {
+      data[i].i = i;
+      data[i].base = base;
+    }
+    run.data = data;
+    run.aux  = has_object ? ALLOCV_N(struct sort_addr_key, haux, elements) : NULL;
+    ca_attach_window(argc, base->ca, modes, sort_addr_run, (VALUE) &run);
+    for (i=0; i<elements; i++) {
+      q[i] = data[i].i;
+    }
   }
 
+  ALLOCV_END(hruns1);
+  ALLOCV_END(hruns0);
   ALLOCV_END(haux);
   ALLOCV_END(hdata);
   ALLOCV_END(hmodes);

data/ext/carray_sort_kernel.c

@@ -453,16 +453,15 @@ DEFINE_PARTITION_NAN(double,    f64);
    preserved), so argsort on NaN-containing input keeps NaN indices
    in ascending order.
 
-   Implementation: 2-pass stable filter via xmalloc scratch (single
-   xmalloc + xfree per fiber; negligible overhead vs the sort
-   itself). */
+   Implementation: 2-pass stable filter through scratch, n pairs that
+   the caller provides.  The caller takes it before it opens its walks:
+   an allocation here, inside the walk, would raise past their finish. */
 #define DEFINE_PARTITION_NAN_PAIR(TYPE, SUFFIX)                            \
 ca_size_t                                                                  \
-ca_partition_nan_pair_##SUFFIX (ca_pair_##SUFFIX *a, ca_size_t n)          \
+ca_partition_nan_pair_##SUFFIX (ca_pair_##SUFFIX *a, ca_size_t n,          \
+                                ca_pair_##SUFFIX *scratch)                 \
 {                                                                          \
   if ( n == 0 ) return 0;                                                  \
-  ca_pair_##SUFFIX *scratch =                                              \
-    (ca_pair_##SUFFIX *) xmalloc((size_t) n * sizeof(ca_pair_##SUFFIX));   \
   ca_size_t fin = 0, nan_pos = 0;                                          \
   for ( ca_size_t k = 0; k < n; k++ ) {                                    \
     if ( !isnan(a[k].v) ) scratch[fin++] = a[k];                           \
@@ -472,7 +471,6 @@ ca_partition_nan_pair_##SUFFIX (ca_pair_##SUFFIX *a, ca_size_t n)          \
     if ( isnan(a[k].v) ) scratch[finite_count + nan_pos++] = a[k];         \
   }                                                                        \
   memcpy(a, scratch, (size_t) n * sizeof(ca_pair_##SUFFIX));               \
-  xfree(scratch);                                                          \
   return finite_count;                                                     \
 }                                                                          \
 struct ca_partition_nan_pair_##SUFFIX##_eat_semicolon