diff --git a/src/builtins/cells.c b/src/builtins/cells.c index 13a47fb6..4b2ee713 100644 --- a/src/builtins/cells.c +++ b/src/builtins/cells.c @@ -656,38 +656,52 @@ NOINLINE B for_cells_SA(B f, B w, B x, ur xcr, ur xr, u32 chr) { usz* xsh=SH(x); usz cam=shProd(xsh,0,xk); if (cam==0) return rank2_empty(f, w, 0, x, xk, chr); if (isFun(f)) { - if (IA(x)==0) goto generic; u8 rtid = v(f)->flags-1; - if (rtid==n_rtack) { dec(w); return x; } - if (rtid==n_ltack) return const_cells(x, xk, xsh, w, chr); - if (rtid==n_select && xk==1 && isF64(w) && xr==2) return select_cells(WRAP(o2i64(w), SH(x)[1], thrF("āŠ: Indexing out-of-bounds (š•Øā‰”%R, %sā‰”ā‰ š•©)", w, cam)), x, cam, 1, false); - if (rtid==n_pick && xk==1 && TI(x,arrD1) && xr==2 && isF64(w)) return select_cells(WRAP(o2i64(w), SH(x)[1], thrF("āŠ‘: Indexing out-of-bounds (š•Øā‰”%R, %sā‰”ā‰ š•©)", w, cam)), x, cam, 1, true); - if ((rtid==n_shifta || rtid==n_shiftb) && xk==1 && xr==2 && isAtm(w)) { - if (isArr(w)) { B w0=w; w = IGet(w,0); decG(w0); } - return shift_cells(w, x, SH(x)[0], SH(x)[1], el_or(TI(x,elType), selfElType(w)), rtid); - } - if (rtid==n_take || rtid==n_drop) { - if (xcr==0 || !isF64(w)) goto generic; - bool take = rtid==n_take; - f64* ap; - B a = m_f64arrv(&ap, xk+1); - if (!take) { FILL_TO(ap, el_f64, 0, m_f64(0), xk); } - else { usz* xsh=SH(x); PLAINLOOP for (usz i=0; i(2048*8)>>arrTypeBitsLog(TY(w)) && IA(w)!=IA(x))) goto generic; - return c2(f, C2(shape, C1(fne, incG(x)), w), x); + switch(rtid) { + case n_rtack: dec(w); return x; + case n_ltack: return const_cells(x, xk, xsh, w, chr); + case n_select: if (isF64(w) && xcr>=1) { + usz l = xsh[xk]; + return select_cells(WRAP(o2i64(w), l, thrF("āŠ: Indexing out-of-bounds (š•Øā‰”%R, %sā‰”ā‰ š•©)", w, l)), x, cam, xk, false); + } break; + case n_pick: if (isF64(w) && xcr==1 && TI(x,arrD1)) { + usz l = xsh[xk]; + return select_cells(WRAP(o2i64(w), l, thrF("āŠ‘: Indexing out-of-bounds (š•Øā‰”%R, %sā‰”ā‰ š•©)", w, l)), x, cam, xk, true); + } break; + case n_shifta: case n_shiftb: if (isAtm(w)) { + if (IA(x)==0) return x; + if (xcr!=1) { + if (xcr==0) break; + if (!(xsh[xk]==1 || shProd(xsh, xk+1, xr)==1)) break; + } + if (isArr(w)) { B w0=w; w = IGet(w,0); decG(w0); } + return shift_cells(w, x, cam, xsh[xk], el_or(TI(x,elType), selfElType(w)), rtid); + } break; + case n_take: case n_drop: { + if (xcr==0 || !isF64(w)) break; + bool take = rtid==n_take; + f64* ap; + B a = m_f64arrv(&ap, xk+1); + if (!take) { FILL_TO(ap, el_f64, 0, m_f64(0), xk); } + else { PLAINLOOP for (usz i=0; i(2048*8)>>arrTypeBitsLog(TY(w)) && IA(w)!=IA(x))) break; + return c2(f, C2(shape, C1(fne, incG(x)), w), x); + } } } else if (!isMd(f)) { dec(w); return const_cells(x, xk, xsh, inc(f), chr); } - generic:; S_KSLICES(x, xsh, xk, cam, 1) incBy(w, cam-1); M_APD_SH(r, xk, xsh); FC2 fc2 = c2fn(f); for (usz i=0,xp=0; i