diff --git a/src/builtins/search.c b/src/builtins/search.c index 216dd487..572c57ee 100644 --- a/src/builtins/search.c +++ b/src/builtins/search.c @@ -168,6 +168,31 @@ B memberOf_c2(B t, B w, B x) { extern B rt_count; B count_c2(B t, B w, B x) { + if (!isArr(w) || RNK(w)==0) thrM("⊒: 𝕨 must have rank at least 1"); + if (RNK(w)==1) { + if (!isArr(x) || IA(x)<=1) return indexOf_c2(m_f64(0), w, x); + usz wia = IA(w); + usz xia = IA(x); + i32* rp; B r = m_i32arrc(&rp, x); + TALLOC(usz, wnext, wia+1); + H_b2i* map = m_b2i(64); + SGetU(x) + SGetU(w) + wnext[wia] = wia; + for (usz i = wia; i--; ) { + bool had; u64 p = mk_b2i(&map, GetU(w,i), &had); + wnext[i] = had ? map->a[p].val : wia; + map->a[p].val = i; + } + for (usz i = 0; i < xia; i++) { + bool had; u64 p = getQ_b2i(map, GetU(x,i), &had); + usz j = wia; + if (had) { j = map->a[p].val; map->a[p].val = wnext[j]; } + rp[i] = j; + } + TFREE(wnext); free_b2i(map); decG(w); decG(x); + return wia<=I8_MAX? taga(cpyI8Arr(r)) : wia<=I16_MAX? taga(cpyI16Arr(r)) : r; + } return c2(rt_count, w, x); } diff --git a/src/utils/hashmapTemplate.h b/src/utils/hashmapTemplate.h index 88e9c0f8..4305863d 100644 --- a/src/utils/hashmapTemplate.h +++ b/src/utils/hashmapTemplate.h @@ -88,6 +88,10 @@ static inline bool N(has) (Map* m, KT k) { u64 p = N(find)(m, k, h1, h2, &has); return has? m->a[p].val : def; } + static inline u64 N(getQ) (Map* m, KT k, bool* had) { + u64 h1 = H1(k); u64 h2 = H2(k, h1); + return N(find)(m, k, h1, h2, had); + } #endif