diff --git a/src/builtins/sfns.c b/src/builtins/sfns.c index e0f93ec1..49d9fa97 100644 --- a/src/builtins/sfns.c +++ b/src/builtins/sfns.c @@ -488,12 +488,13 @@ static NOINLINE B takedrop_highrank(bool take, B w, B x) { wn[i] = cw<0; usz xshc = i=xshc? 0 : xshc-cwa; + u64 c = take? cwa : cwa>=xshc? 0 : xshc-cwa; if (c!=xshc) cellStart = i; anyFill|= c>xshc; rsh->a[i] = c; if (mulOn(ria, c)) thrOOM(); } + CHECK_IA(ria, 8); if (cellStart<=0) { if (xr==rr) { @@ -609,14 +610,15 @@ static NOINLINE B takedrop_highrank(bool take, B w, B x) { if (!isNum(w)) return takedrop_highrank(TAKE, w, x); \ Arr* a; \ i64 wv = o2i64(w); \ + i64 n = wv; \ ur xr = RNK(x); \ usz csz=1; usz* xsh; \ if (xr>1) { \ csz = arr_csz(x); \ xsh = SH(x); \ ptr_inc(shObjS(xsh)); \ + if (mulOn(n, csz)) thrOOM(); \ } else xr=1; \ - i64 n=wv; if (mulOn(n, csz)) thrOOM(); #define TAKEDROP_SHAPE(SH0) \ if (xr>1) { \ @@ -632,10 +634,12 @@ B take_c2(B t, B w, B x) { TAKEDROP_INIT(1); if (n>=0) { + CHECK_IA(n, 8); a = take_impl(n, x); if (xr==1) return taga(arr_shVec(a)); } else { n = -n; + CHECK_IA(n, 8); usz xia = IA(x); if (n>xia) { B xf = getFillE(x); diff --git a/src/h.h b/src/h.h index 1d4e2b72..d61e6f46 100644 --- a/src/h.h +++ b/src/h.h @@ -164,11 +164,11 @@ typedef double f64; #if USZ_64 typedef u64 usz; #define USZ_MAX ((u64)(1ULL<<48)) - #define CHECK_IA(IA,W) if(IA>USZ_MAX) thrOOM() + #define CHECK_IA(IA,W) if((IA) > USZ_MAX) thrOOM() #else typedef u32 usz; #define USZ_MAX ((u32)((1LL<<32)-1)) - #define CHECK_IA(IA,W) if (IA > ((1LL<<31)/W - 1000)) thrOOM() + #define CHECK_IA(IA,W) if ((IA) > ((1LL<<31)/W - 1000)) thrOOM() #endif #if UNSAFE_SIZES #undef CHECK_IA