From 36b99d3505ec771a211b097c6b50a7d92a665a9d Mon Sep 17 00:00:00 2001 From: dzaima Date: Sat, 22 Apr 2023 17:03:09 +0300 Subject: [PATCH] =?UTF-8?q?handle=20more=20cases=20of=20=E2=8A=8F=E2=8E=89?= =?UTF-8?q?n=20&=20=E2=8A=91=E2=8E=89n?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/builtins/cells.c | 83 ++++++++++++++++++++++---------------------- 1 file changed, 41 insertions(+), 42 deletions(-) diff --git a/src/builtins/cells.c b/src/builtins/cells.c index b9e6a7a8..1d48445c 100644 --- a/src/builtins/cells.c +++ b/src/builtins/cells.c @@ -110,54 +110,54 @@ NOINLINE B toKCells(B x, ur k) { // fast special-case implementations -static NOINLINE B select_cells(usz n, B x, ur xr) { +static NOINLINE B select_cells(usz n, B x, usz cam, usz k, bool leaf) { // n {leaf? <∘⊑; ⊏}⎉¯k x; TODO probably can share some parts with takedrop_highrank and/or call ⊏? + ur xr = RNK(x); + assert(xr>1 && k1? toCells(x) : k==0? m_atomUnit(x) : toKCells(x, k); case n_select: if (IA(x)==0) goto noSpecial; - if (cr==0 || k!=1) goto base; // TODO handle more ranks - selectCells:; - return select_cells(0, x, xr); + if (cr==0) goto base; + return select_cells(0, x, cam, k, false); case n_pick: if (IA(x)==0) goto noSpecial; - if (k!=1 || cr!=1 || !TI(x,arrD1)) goto base; // TODO handle more ranks - goto selectCells; + if (cr==0 || !TI(x,arrD1)) goto base; + return select_cells(0, x, cam, k, true); case n_couple: { Arr* r = cpyWithShape(x); xsh=PSH(r); if (xr==UR_MAX) thrF("≍%c: Result rank too large (%i≡=𝕩)", chr, xr); @@ -333,7 +332,7 @@ B for_cells_c1(B f, u32 xr, u32 cr, u32 k, B x, u32 chr) { // F⎉cr x, with arr if (rtid==n_const) { f=fd->f; goto const_f; } if ((rtid==n_fold || rtid==n_insert) && TI(x,elType)!=el_B && k==1 && xr==2 && isPervasiveDyExt(fd->f)) { usz *sh = SH(x); usz m = sh[1]; - if (m == 1) return select_cells(0, x, 2); + if (m == 1) return select_cells(0, x, cam, k, false); if (m <= 64 && m < sh[0]) return fold_rows(fd, x); } } @@ -435,8 +434,8 @@ B cell_c2(Md1D* d, B w, B x) { B f = d->f; if (cam==0) return cell2_empty(f, w, x, wr, xr); if (isFun(f)) { u8 rtid = v(f)->flags-1; - if (rtid==n_select && isF64(w) && xr>1) return select_cells(WRAP(o2i64(w), SH(x)[1], thrF("⊏: Indexing out-of-bounds (𝕨≡%R, %s≡≠𝕩)", w, cam)), x, xr); - if (rtid==n_pick && TI(x,arrD1) && xr>1 && isF64(w)) return select_cells(WRAP(o2i64(w), SH(x)[1], thrF("⊑: Indexing out-of-bounds (𝕨≡%R, %s≡≠𝕩)", w, cam)), x, xr); + if (rtid==n_select && 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 && 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) && xr==2) { 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);