intarray/
lib.rs

1#[macro_use]
2extern crate log;
3
4use rand::Rng;
5use serde::ser::{Serialize, SerializeSeq, Serializer};
6use std::ops::{AddAssign, MulAssign, Range, SubAssign};
7use std::sync::Once;
8use std::{cmp, fmt, mem};
9
10type Element = u64;
11// type Element = u128;
12const ELEMENT_BITS: usize = mem::size_of::<Element>() * 8;
13static INIT_MASK: Once = Once::new();
14static mut MASK_ARRAY: [Element; ELEMENT_BITS / 2 - 1] = [0; ELEMENT_BITS / 2 - 1];
15
16fn init_mask_fn() {
17    unsafe {
18        for i in 1..(ELEMENT_BITS / 2) {
19            let mask0 = ((1 as Element) << i) - 1;
20            let mut mask: Element = mask0;
21            for _j in 0..(ELEMENT_BITS / (i * 2)) {
22                mask = mask.wrapping_shl(i as u32 * 2);
23                mask |= mask0;
24            }
25            MASK_ARRAY[i - 1] = mask;
26        }
27        for (i, j) in MASK_ARRAY.iter().enumerate() {
28            debug!("mask[{}+1]={:b}", i, j);
29        }
30    }
31}
32
33fn get_mask(bits: usize) -> Element {
34    INIT_MASK.call_once(|| init_mask_fn());
35    unsafe {
36        if (bits - 1) < MASK_ARRAY.len() {
37            MASK_ARRAY[bits - 1]
38        } else if bits == ELEMENT_BITS {
39            !0
40        } else {
41            ((1 as Element) << bits) - 1
42        }
43    }
44}
45
46fn get_masks(bits: usize) -> (Element, Element) {
47    let res = get_mask(bits);
48    (res, res.wrapping_shl(bits as u32))
49}
50
51trait ElementTrait {
52    type E;
53
54    // utility
55    fn val_expand(v: Self::E, bits: usize) -> Self::E;
56
57    // a1+a2+a3+...
58    fn sum_bits(self, bits: usize) -> Option<Self::E>;
59
60    // (a1 OP b1), (a2 OP b2), (a3 OP b3),...
61    fn add_bits(self, b: Self::E, bits: usize) -> Option<Self::E>;
62    fn sub_bits(self, b: Self::E, bits: usize) -> Option<Self::E>;
63
64    // (a1 OP b), (a2 OP b), (a3 OP b),...
65    fn addval_bits(self, b: Self::E, bits: usize) -> Option<Self::E>;
66    fn subval_bits(self, b: Self::E, bits: usize) -> Option<Self::E>;
67    fn mulval_bits(self, b: Self::E, bits: usize) -> Option<Self::E>;
68
69    fn ffs(self) -> usize;
70}
71
72impl ElementTrait for Element {
73    type E = Element;
74
75    fn val_expand(v: Self::E, bits: usize) -> Self::E {
76        if bits == 0 {
77            return 0;
78        }
79        let mut v1: Self::E = 0;
80        for _ in 0..(ELEMENT_BITS / bits) {
81            v1 <<= bits;
82            v1 |= v;
83        }
84        v1
85    }
86
87    fn sum_bits(self, bits: usize) -> Option<Self::E> {
88        if bits >= ELEMENT_BITS {
89            debug!("return self: {}, bits={}", self, bits);
90            return Some(self);
91        }
92        if bits == 0 {
93            error!("invalid argument: bits={}", bits);
94            return None;
95        }
96        let mask = get_mask(bits);
97        let res = (self & mask) + ((self >> bits) & mask);
98        res.sum_bits(bits * 2)
99    }
100
101    fn add_bits(self, b: Self::E, bits: usize) -> Option<Self::E> {
102        if bits == 0 {
103            error!("invalid argument: bits={}", bits);
104            return None;
105        }
106        let (mask1, mask2) = get_masks(bits);
107        let r1 = (self & mask1).wrapping_add(b & mask1);
108        let r2 = (self & mask2).wrapping_add(b & mask2);
109        if (r1 & mask2) != 0 || (r2 & mask1) != 0 {
110            None
111        } else {
112            Some(r1 + r2)
113        }
114    }
115
116    fn sub_bits(self, b: Self::E, bits: usize) -> Option<Self::E> {
117        if bits == 0 {
118            error!("invalid argument: bits={}", bits);
119            return None;
120        }
121        let (mask1, mask2) = get_masks(bits);
122        let r1 = (self & mask1).wrapping_sub(b & mask1);
123        let r2 = (self & mask2).wrapping_sub(b & mask2);
124        if (r1 & mask2) != 0 || (r2 & mask1) != 0 {
125            None
126        } else {
127            Some(r1 + r2)
128        }
129    }
130
131    fn mulval_bits(self, b: Self::E, bits: usize) -> Option<Self::E> {
132        if bits == 0 {
133            error!("invalid argument: bits={}", bits);
134            return None;
135        }
136        let (mask1, mask2) = get_masks(bits);
137        let r1 = (self & mask1).wrapping_mul(b);
138        let r2 = (self & mask2).wrapping_mul(b);
139        if (r1 & mask2) != 0 || (r2 & mask1) != 0 {
140            None
141        } else {
142            Some(r1 + r2)
143        }
144    }
145
146    fn addval_bits(self, b: Self::E, bits: usize) -> Option<Self::E> {
147        self.add_bits(Self::val_expand(b, bits), bits)
148    }
149
150    fn subval_bits(self, b: Self::E, bits: usize) -> Option<Self::E> {
151        self.sub_bits(Self::val_expand(b, bits), bits)
152    }
153
154    fn ffs(self) -> usize {
155        ELEMENT_BITS - self.leading_zeros() as usize
156    }
157}
158
159/// 1-64 bits unsigned integer array
160pub struct IntArray {
161    /// bit width
162    pub bits: usize,
163    /// number of element
164    pub length: usize,
165    /// data store
166    data: Vec<Element>,
167}
168
169/// iterator for IntArray
170pub struct IntIter<'a> {
171    /// current position
172    range: Range<usize>,
173    /// refere to
174    a: &'a IntArray,
175}
176
177impl IntArray {
178    fn sizeval(b: usize, len: usize) -> (usize, usize) {
179        let bpd = ELEMENT_BITS / b;
180        return (bpd, (len + bpd - 1) / bpd);
181    }
182    pub fn new(b: usize, len: usize) -> IntArray {
183        let (bpd, cap) = IntArray::sizeval(b, len);
184        debug!("bpd={}, cap={}", bpd, cap);
185        IntArray {
186            bits: b,
187            length: len,
188            data: vec![0; cap as usize],
189        }
190    }
191
192    pub fn new_with_vec(b: usize, vals: Vec<Element>) -> IntArray {
193        let (bpd, cap) = IntArray::sizeval(b, vals.len());
194        debug!("bpd={}, cap={}", bpd, cap);
195        let mut res = IntArray {
196            bits: b,
197            length: vals.len(),
198            data: vec![0; cap],
199        };
200        for (i, v) in vals.iter().enumerate() {
201            res.set(i, *v).unwrap();
202        }
203        res
204    }
205
206    pub fn new_with_iter<'a, I>(b: usize, vals: I) -> IntArray
207    where
208        I: Iterator<Item = Element>,
209    {
210        const UNIT: usize = 1024;
211        let (bpd, cap) = IntArray::sizeval(b, UNIT);
212        let mut cnt = 0usize;
213        debug!("bpd={}, cap={}", bpd, cap);
214        let mut res = IntArray {
215            bits: b,
216            length: UNIT,
217            data: vec![0; cap as usize],
218        };
219        for v in vals {
220            if cnt >= res.length {
221                res.resize(res.length + UNIT);
222            }
223            res.set(cnt, v).unwrap();
224            cnt += 1
225        }
226        // truncate
227        res.resize(cnt);
228        res
229    }
230
231    pub fn subarray<'a>(&'a self, offset: usize, length: usize) -> IntArray {
232        let mut res = IntArray::new(self.bits, length);
233        let (bpd, _) = IntArray::sizeval(self.bits, 0);
234        if offset % bpd == 0 {
235            // fast path
236            debug!("fast path: offset={}, length={}", offset, length);
237            for i in 0..(length / bpd) {
238                res.data[i] = self.data[offset / bpd + i];
239            }
240            // rest
241            if length % bpd != 0 {
242                let i = length / bpd;
243                let mask = (1 << (bpd * (length % bpd))) - 1;
244                res.data[i] = self.data[offset / bpd + i] & (!mask);
245            }
246        } else {
247            // slow path
248            debug!("slow path: offset={}, length={}", offset, length);
249            for i in 0..length {
250                res.set(i, self.get(offset + i).unwrap()).unwrap();
251            }
252        }
253        res
254    }
255
256    pub fn datasize(self) -> usize {
257        mem::size_of::<IntArray>() + (ELEMENT_BITS / 8) * self.data.capacity()
258    }
259
260    pub fn push(&mut self, v: Element) -> Option<usize> {
261        self.resize(self.length + 1);
262        match self.set(self.length - 1, v) {
263            Ok(_) => Some(self.length - 1),
264            Err(_) => None,
265        }
266    }
267
268    pub fn extend<I>(&mut self, vals: I)
269    where
270        I: Iterator<Item = Element>,
271    {
272        for v in vals {
273            self.push(v).unwrap();
274        }
275    }
276
277    pub fn extend_array(&mut self, vals: IntArray) {
278        let (bpd, _) = IntArray::sizeval(self.bits, 0);
279        if vals.bits == self.bits && self.length % bpd == 0 {
280            // fast path
281            debug!("fast path: bits={}, length={}", self.bits, self.length);
282            self.resize(self.length + vals.length);
283            self.data.extend(vals.data);
284            return;
285        }
286        // slow path
287        debug!("slow path: bits={}, length={}", self.bits, self.length);
288        self.extend(vals.iter());
289    }
290
291    pub fn concat(&mut self, vals: IntArray) {
292        self.extend_array(vals);
293    }
294
295    pub fn shape<'a>(self: &'a IntArray, bits: usize) -> IntArray {
296        IntArray::new_with_iter(bits, self.iter())
297    }
298
299    pub fn shape_auto<'a>(self: &'a IntArray) -> IntArray {
300        let mv = self.iter().max().unwrap();
301        let bits = match mv.ffs() {
302            0 => 1,
303            n => n,
304        };
305        IntArray::new_with_iter(bits, self.iter())
306    }
307
308    pub fn pop(&mut self) -> Result<Element, String> {
309        let res = self.get(self.length - 1);
310        self.resize(self.length - 1);
311        res
312    }
313
314    pub fn iter(&self) -> IntIter {
315        IntIter {
316            range: 0..self.length,
317            a: &self,
318        }
319    }
320
321    pub fn capacity(&self) -> usize {
322        let bpd = ELEMENT_BITS / self.bits;
323        self.data.len() * bpd
324    }
325
326    pub fn resize(&mut self, len: usize) {
327        let bpd = ELEMENT_BITS / self.bits;
328        let cap = (len + bpd - 1) / bpd;
329        self.length = len;
330        self.data.resize(cap, 0);
331    }
332
333    pub fn len(&self) -> usize {
334        self.length
335    }
336
337    pub fn max_value(&self) -> Element {
338        if self.bits == ELEMENT_BITS {
339            return Element::max_value();
340        }
341        ((1 as Element) << self.bits) - 1
342    }
343
344    pub fn max(&self) -> Element {
345        self.iter().max().unwrap()
346    }
347
348    pub fn min(&self) -> Element {
349        self.iter().min().unwrap()
350    }
351
352    pub fn average(&self) -> f64 {
353        self.sum().unwrap() as f64 / self.len() as f64
354    }
355
356    fn getoffset(&self, i: usize) -> (usize, usize, usize) {
357        let bpd = ELEMENT_BITS / self.bits as usize;
358        return (bpd, i / bpd, i % bpd);
359    }
360
361    pub fn get(&self, i: usize) -> Result<Element, String> {
362        if self.length <= i {
363            return Err("OB".to_owned());
364        }
365        // debug!("get {}/{}", i, self.capacity());
366        let (_, idx, iv) = self.getoffset(i);
367        let vv = self.data[idx];
368        // debug!("idx={}, iv={}, vv={}", idx, iv, vv);
369        let res = (vv >> (iv * self.bits as usize)) & self.max_value();
370        Ok(res)
371    }
372
373    pub fn set(&mut self, i: usize, v: Element) -> Result<Element, String> {
374        if self.max_value() < v {
375            return Err("TooLarge".to_owned());
376        }
377        if self.length <= i {
378            return Err("OutOfBounds".to_owned());
379        }
380        let (_, idx, iv) = self.getoffset(i);
381        let mask1 = (self.max_value()) << (iv * self.bits);
382        let mask2 = v << (iv * self.bits);
383        let res = self.max_value() & (self.data[idx] >> (iv * self.bits));
384        self.data[idx] = (self.data[idx] & (!mask1)) | mask2;
385        Ok(res)
386    }
387
388    pub fn add(&mut self, i: usize, v: Element) -> Result<Element, String> {
389        match self.get(i) {
390            Ok(n) => self.set(i, n + v),
391            Err(e) => return Err(e),
392        }
393    }
394
395    pub fn sub(&mut self, i: usize, v: Element) -> Result<Element, String> {
396        match self.get(i) {
397            Ok(n) => self.set(i, n - v),
398            Err(e) => return Err(e),
399        }
400    }
401
402    pub fn incr_limit(&mut self, i: usize) -> Option<Element> {
403        match self.get(i) {
404            Ok(n) => {
405                if n != self.max_value() {
406                    self.set(i, n + 1).unwrap();
407                    Some(n)
408                } else {
409                    None
410                }
411            }
412            Err(_e) => None,
413        }
414    }
415
416    pub fn decr_limit(&mut self, i: usize) -> Option<Element> {
417        match self.get(i) {
418            Ok(n) => {
419                if n != 0 && n != self.max_value() {
420                    self.set(i, n - 1).unwrap();
421                    Some(n)
422                } else {
423                    None
424                }
425            }
426            Err(_e) => None,
427        }
428    }
429
430    pub fn incr(&mut self, i: usize) -> Result<Element, String> {
431        self.add(i, 1)
432    }
433
434    pub fn decr(&mut self, i: usize) -> Result<Element, String> {
435        self.sub(i, 1)
436    }
437
438    pub fn sum<'a>(&'a self) -> Option<Element> {
439        let mut res = 0;
440        if self.bits > ELEMENT_BITS / 2 {
441            return self.sum0();
442        }
443        for i in self.data.iter() {
444            res += (*i).sum_bits(self.bits).unwrap();
445            debug!("sum: {} -> {}", *i, res);
446        }
447        Some(res)
448    }
449
450    pub fn sum0<'a>(&'a self) -> Option<Element> {
451        return Some(self.iter().fold(0 as Element, |sum, a| sum + a));
452    }
453
454    pub fn fill_random(&mut self) {
455        let mut rng = rand::thread_rng();
456        if ELEMENT_BITS % self.bits == 0 {
457            for i in 0..(self.data.len() - 1) {
458                self.data[i] = rng.gen();
459            }
460        } else {
461            let mvbits = ELEMENT_BITS - (ELEMENT_BITS % self.bits);
462            let mvval = (1 as Element) << mvbits;
463            for i in 0..(self.data.len() - 1) {
464                self.data[i] = rng.gen_range(0, mvval);
465            }
466        }
467        let bpd = ELEMENT_BITS / self.bits;
468        for i in (self.data.len() - 1) * bpd..self.length {
469            self.set(i, rng.gen_range(0, self.max_value())).unwrap();
470        }
471    }
472}
473
474impl<'a> Iterator for IntIter<'a> {
475    type Item = Element;
476    fn next(&mut self) -> Option<Element> {
477        self.range.next().map(|i| self.a.get(i).unwrap())
478    }
479
480    fn count(self) -> usize {
481        self.a.length
482    }
483}
484
485impl fmt::Display for IntArray {
486    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
487        write!(f, "{}[{}]=", self.bits, self.length).unwrap();
488        write!(
489            f,
490            "{}",
491            self.iter()
492                .map(|x| x.to_string())
493                .collect::<Vec<String>>()
494                .join(",")
495        )
496        .unwrap();
497        Ok(())
498    }
499}
500
501impl Serialize for IntArray {
502    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
503    where
504        S: Serializer,
505    {
506        let mut seq = serializer.serialize_seq(Some(self.length))?;
507        self.iter().for_each(|x| seq.serialize_element(&x).unwrap());
508        seq.end()
509    }
510}
511
512impl AddAssign<Element> for IntArray {
513    fn add_assign(&mut self, v: Element) {
514        let (bpd, _) = IntArray::sizeval(self.bits, 0);
515        debug!("length={}, bits={}, bpd={}", self.length, self.bits, bpd);
516        for i in 0..(self.length / bpd) {
517            debug!("i={}, v={}", i, v);
518            self.data[i] = self.data[i].addval_bits(v, self.bits).unwrap();
519        }
520        if self.length % bpd != 0 {
521            let i = self.length / bpd;
522            let mask = ((1 as Element) << (self.bits * (self.length % bpd))) - 1;
523            debug!("mask={:x} ({} bits)", mask, mask.ffs());
524            self.data[i] = self.data[i].addval_bits(v, self.bits).unwrap() & mask;
525        }
526    }
527}
528
529impl AddAssign<IntArray> for IntArray {
530    fn add_assign(&mut self, v: IntArray) {
531        if self.bits == v.bits && self.length == v.length {
532            // fast path
533            debug!("fast path: bits={}, length={}", self.bits, self.length);
534            let (bpd, _) = IntArray::sizeval(self.bits, 0);
535            for i in 0..(self.length / bpd) {
536                self.data[i] = self.data[i].add_bits(v.data[i], self.bits).unwrap();
537            }
538            if self.length % bpd != 0 {
539                let i = self.length / bpd;
540                let mask = ((1 as Element) << (self.bits * (self.length % bpd))) - 1;
541                self.data[i] = self.data[i].add_bits(v.data[i], self.bits).unwrap() & mask;
542            }
543            return;
544        }
545        // slow path
546        debug!(
547            "slow path: bits={}/{}, length={}/{}",
548            self.bits, v.bits, self.length, v.length
549        );
550        for i in 0..cmp::min(self.length, v.length) {
551            self.set(i, self.get(i).unwrap() + v.get(i).unwrap())
552                .unwrap();
553        }
554    }
555}
556
557impl SubAssign<Element> for IntArray {
558    fn sub_assign(&mut self, v: Element) {
559        let (bpd, _) = IntArray::sizeval(self.bits, 0);
560        for i in 0..(self.length / bpd) {
561            self.data[i] = self.data[i].subval_bits(v, self.bits).unwrap();
562        }
563        if self.length % bpd != 0 {
564            let i = self.length / bpd;
565            let mask = ((1 as Element) << (self.bits * (self.length % bpd))) - 1;
566            self.data[i] = (self.data[i] | (!mask)).subval_bits(v, self.bits).unwrap() & mask;
567        }
568    }
569}
570
571impl SubAssign<IntArray> for IntArray {
572    fn sub_assign(&mut self, v: IntArray) {
573        if self.bits == v.bits && self.length == v.length {
574            // fast path
575            debug!("fast path: bits={}, length={}", self.bits, self.length);
576            let (bpd, _) = IntArray::sizeval(self.bits, 0);
577            for i in 0..(self.length / bpd) {
578                self.data[i] = self.data[i].sub_bits(v.data[i], self.bits).unwrap();
579            }
580            if self.length % bpd != 0 {
581                let i = self.length / bpd;
582                let mask = ((1 as Element) << (self.bits * (self.length % bpd))) - 1;
583                self.data[i] = (self.data[i] | (!mask))
584                    .sub_bits(v.data[i], self.bits)
585                    .unwrap()
586                    & mask;
587            }
588            return;
589        }
590        // slow path
591        debug!(
592            "slow path: bits={}/{}, length={}/{}",
593            self.bits, v.bits, self.length, v.length
594        );
595        for i in 0..cmp::min(self.length, v.length) {
596            self.sub(i, v.get(i).unwrap()).unwrap();
597        }
598    }
599}
600
601impl MulAssign<Element> for IntArray {
602    fn mul_assign(&mut self, v: Element) {
603        let (bpd, _) = IntArray::sizeval(self.bits, 0);
604        for i in 0..(self.length / bpd) {
605            self.data[i] = self.data[i].mulval_bits(v, self.bits).unwrap();
606        }
607        if self.length % bpd != 0 {
608            let i = self.length / bpd;
609            let mask = ((1 as Element) << (self.bits * (self.length % bpd))) - 1;
610            self.data[i] = (self.data[i] & mask).mulval_bits(v, self.bits).unwrap() & mask;
611        }
612    }
613}
614
615impl Clone for IntArray {
616    fn clone(&self) -> IntArray {
617        let mut res = IntArray::new(self.bits, self.length);
618        res.data.clone_from_slice(&self.data);
619        res
620    }
621}
622
623#[cfg(test)]
624mod tests;
intarray/lib.rs

intarray/
lib.rs
