Skip to main content

mosekcomodel/utils/
mod.rs

1#![allow(unused)]
2#![allow(dead_code)]
3pub mod iter;
4
5use std::{marker::PhantomData, ptr::NonNull};
6
7use itertools::izip;
8
9#[derive(Debug,Clone,Copy)]
10pub struct Strides<const N : usize> {
11    shape   : [usize;N],
12    strides : [usize;N]
13}
14
15impl<const N : usize> Strides<N> {
16    pub fn from_shape(shape : &[usize;N]) -> Strides<N> {
17        let mut strides = [0usize; N]; 
18        strides.iter_mut().zip(shape.iter()).rev().fold(1usize,|c,(t,&s)| { *t = c; c*s });
19        Strides{ strides, shape : *shape }
20    }
21    pub fn to_array(&self) -> [usize;N] { self.strides }
22    pub fn to_linear(&self, index : &[usize;N]) -> usize {
23        index.iter().zip(self.strides.iter()).map(|(a,b)| a*b).sum()
24    }
25    pub fn to_index(&self, i : usize) -> [usize;N] {
26        let mut r = [0usize;N];
27        r.iter_mut().zip(self.strides.iter()).fold(i,|i,(r,&s)| { *r = i/s; i%s} );
28        r
29    }
30
31    // Given coordinates `i`, compute the correponsing linear index. If `i` is not inside the
32    // shape, return None.
33    pub fn from_coords_checked(&self, i : &[usize;N]) -> Option<usize> {
34        if i.iter().zip(self.shape.iter()).all(|v| *v.0 < *v.1) {
35            Some(self.to_linear(i))
36        }
37        else {
38            None
39        }
40    }
41
42    pub fn iter<'a>(&'a self) -> std::slice::Iter<'a,usize> { self.strides.iter() }
43}
44
45
46pub trait ShapeToStridesEx<const N : usize> {
47    fn to_strides(&self) -> Strides<N>;
48}
49
50impl<const N : usize> ShapeToStridesEx<N> for [usize;N] {
51    fn to_strides(&self) -> Strides<N> {
52        Strides::from_shape(self)
53    }
54}
55
56
57
58pub trait Cummulate {
59    fn cummulate(&mut self);
60}
61
62impl<T> Cummulate for [T] where
63    T : Copy+std::ops::AddAssign
64{
65    fn cummulate(& mut self) {
66        if ! self.is_empty() {
67            let v0 = self[0];
68            self[1..].iter_mut().fold(v0,|c,v| { *v += c; *v });
69        }
70    }
71}
72
73impl<T> Cummulate for Vec<T> where
74    T : Copy+std::ops::AddAssign
75{
76    fn cummulate(& mut self) {
77        if ! self.is_empty() {
78            let v0 = self[0];
79            self[1..].iter_mut().fold(v0,|c,v| { *v += c; *v });
80        }
81    }
82}
83
84
85
86
87
88
89
90/// A trait that supplies functionality for appending self to a string.
91pub trait NameAppender {
92    /// Append self to a string
93    fn append_to_string(&self, s : & mut String);
94}
95
96impl<T> NameAppender for [T] where T : NameAppender {
97    fn append_to_string(&self, s : & mut String) {
98        s.push('[');
99        if self.len() > 0 {
100            self[0].append_to_string(s);
101            for i in self[1..].iter() { s.push(','); i.append_to_string(s) }
102        }
103        s.push(']');
104    }
105}
106
107impl NameAppender for usize {
108    fn append_to_string(&self, s : & mut String) {
109        if *self == 0 {
110            s.push('0');
111        }
112        else {
113            let mut buf = [0u8; 20];
114            let n = buf.iter_mut().rev().scan(*self,|v,b| if *v > 0 { let r = (*v%10) as u8; *v = *v/10; *b = r; Some(r) } else { None }).count();
115            for c in &buf[20-n..] {
116                s.push((*c + b'0') as char);
117            }
118
119        }
120    }
121}
122
123
124
125
126
127/// A struct representing a permutation or mutation of indexes.
128#[derive(Clone,Copy)]
129pub struct Permutation<'b> {
130    perm : &'b [usize],
131    max  : usize
132}
133
134/// A struct representing a permutation or mutation of a vector.
135pub struct AppliedPermutation<'a, 'b, T> {
136    data : &'a [T],
137    perm : &'b [usize]
138}
139
140pub struct AppliedPermutationMut<'a, 'b, T> {
141    data : &'a mut [T],
142    perm : &'b [usize]
143}
144
145/// An iterator over a permutation of a vector.
146pub struct AppliedPermutationIterator<'a,'b,T> {
147    perm : &'a [usize],
148    data : &'b [T],
149    index : usize
150}
151
152/// An iterator over a permutation of a vector.
153pub struct AppliedPermutationMutIterator<'a,'b,T> {
154    perm : &'a [usize],
155    ptr : NonNull<T>,
156    _marker : PhantomData<&'b T>,
157    index : usize
158}
159
160impl<'b> Permutation<'b> {
161    /// Create a permutation from a vector of indexes.
162    pub fn from(perm : & 'b[usize]) -> Permutation<'b> {
163        Permutation{
164            perm,
165            max : perm.iter().max().map(|&v| v+1).unwrap_or(0)
166        }
167    }
168    /// Apply the permutation to a vector
169    ///
170    /// # Arguments
171    /// - `data` The array to permute
172    /// # Returns
173    /// If the permutation is valid for the given array (all indexes are within bounds), return an
174    /// applied permutation, otherwise `None`.
175    pub fn apply<'a,T>(&self, data : &'a[T]) -> Option<AppliedPermutation<'a,'b,T>> {
176        if data.len() < self.max { None }
177        else { Some(AppliedPermutation{ data, perm : self.perm }) }
178    }
179    pub fn apply_mut<'a,T>(&self, data : &'a mut[T]) -> Option<AppliedPermutationMut<'a,'b,T>> {
180        if data.len() < self.max { None }
181        else { Some(AppliedPermutationMut{ data, perm : self.perm }) }
182    }
183
184    /// Length of the permutation
185    pub fn len(&self) -> usize { self.perm.len() }
186}
187
188impl<'a,'b,T> std::ops::Index<usize> for AppliedPermutation<'a,'b,T> {
189    type Output = T;
190    fn index(&self, i : usize) -> &T {
191        unsafe {
192            self.data.get_unchecked(self.perm[i])
193        }
194    }
195}
196
197//impl<'a,'b,T> std::ops::Index<usize> for AppliedPermutationMut<'a,'b,T> {
198//    type Output = T;
199//    fn index(&self, i : usize) -> &mut T {
200//        unsafe {
201//            self.data.get_unchecked_mut(self.perm[i])
202//        }
203//    }
204//}
205
206impl<'a,'b,T> Iterator for AppliedPermutationIterator<'a,'b,T> {
207    type Item = &'b T;
208    fn next(&mut self) -> Option<Self::Item> {
209        if self.index < self.perm.len() {
210            let res = unsafe{self.data.get_unchecked(*self.perm.get_unchecked(self.index))};
211            self.index += 1;
212
213            Some(res)
214        }
215        else {
216            None
217        }
218    }
219}
220
221impl<'a,'b,T> Iterator for AppliedPermutationMutIterator<'a,'b,T> {
222    type Item = &'b mut T;
223    fn next(&mut self) -> Option<Self::Item> {
224        self.perm.get(self.index)
225            .and_then(|&index| { 
226                self.index += 1; 
227                Some(unsafe{self.ptr.as_ptr().add(index).as_mut().unwrap()}) 
228            })
229    }
230}
231
232impl<'a,'b,T> AppliedPermutation<'a,'b,T> {
233    /// Return an iterator over the permited elements.
234    pub fn iter(&self) -> AppliedPermutationIterator<'b,'a,T> {
235        AppliedPermutationIterator{
236            perm : self.perm,
237            data : self.data,
238            index : 0
239        }
240    }
241    /// Length of the underlying permutation
242    pub fn len(&self) -> usize { self.perm.len() }
243}
244
245impl<'a,'b,T> AppliedPermutationMut<'a,'b,T> {
246    /// Return an iterator over the permited elements.
247    pub fn iter(self) -> AppliedPermutationMutIterator<'b,'a,T> {
248        AppliedPermutationMutIterator{
249            perm : self.perm,
250            ptr : NonNull::from(self.data).cast(), 
251            _marker : PhantomData,
252            index : 0
253        }
254    }
255    /// Length of the underlying permutation
256    pub fn len(&self) -> usize { self.perm.len() }
257}
258
259
260
261impl<'a> From<&Permutation<'a>> for Permutation<'a> {
262    fn from(value: &Permutation<'a>) -> Self { *value }
263}
264impl<'a> From<&'a[usize]> for Permutation<'a> {
265    fn from(value: &'a[usize]) -> Self { Permutation::from(value) }
266}
267impl<'a> From<&'a Vec<usize>> for Permutation<'a> {
268    fn from(value: &'a Vec<usize>) -> Self {
269        Permutation::from(value.as_slice())
270    }
271}
272
273
274pub trait ApplyPermutationEx<T> {
275    fn try_permute_by<'a,'b,P>(&'b self, perm : P) ->  Option<AppliedPermutationIterator<'a,'b,T>> where P : Into<Permutation<'a>>;
276    fn permute_by<'a,'b,P>(&'b self, perm : P) ->  AppliedPermutationIterator<'a,'b,T> where P : Into<Permutation<'a>> { self.try_permute_by(perm).unwrap() } 
277}
278pub trait ApplyPermutationMutEx<T> {
279    fn try_permute_by_mut<'a,'b,P>(&'b mut self, perm : P) ->  Option<AppliedPermutationMutIterator<'a,'b,T>> where P : Into<Permutation<'a>>;
280    fn permute_by_mut<'a,'b,P>(&'b mut self, perm : P) ->  AppliedPermutationMutIterator<'a,'b,T> where P : Into<Permutation<'a>> { self.try_permute_by_mut(perm).unwrap() }
281}
282
283impl<T> ApplyPermutationEx<T> for Vec<T> {
284    fn try_permute_by<'a,'b,P>(&'b self, perm : P) ->  Option<AppliedPermutationIterator<'a,'b,T>> where P : Into<Permutation<'a>> {
285        perm.into().apply(self.as_ref()).map(|a| a.iter())
286    }
287}
288
289impl<T> ApplyPermutationEx<T> for [T] {
290    fn try_permute_by<'a,'b,P>(&'b self, perm : P) ->  Option<AppliedPermutationIterator<'a,'b,T>> where P : Into<Permutation<'a>> {
291        perm.into().apply(self).map(|a| a.iter())
292    }
293}
294
295impl<T> ApplyPermutationMutEx<T> for Vec<T> {
296    fn try_permute_by_mut<'a,'b,P>(&'b mut self, perm : P) ->  Option<AppliedPermutationMutIterator<'a,'b,T>> where P : Into<Permutation<'a>> {
297        perm.into().apply_mut(self.as_mut()).map(|a| a.iter())
298    }
299}
300
301impl<T> ApplyPermutationMutEx<T> for [T] {
302    fn try_permute_by_mut<'a,'b,P>(&'b mut self, perm : P) ->  Option<AppliedPermutationMutIterator<'a,'b,T>> where P : Into<Permutation<'a>> {
303        perm.into().apply_mut(self).map(|a| a.iter())
304    }
305}
306
307
308
309
310
311
312pub trait SwapEx where Self : Copy {
313    /// Assign a new value to a reference and return the old value.
314    fn swap_out(&mut self, v : Self) -> Self {
315        let tmp = *self;
316        *self = v;
317        tmp
318    }
319}
320
321impl<T> SwapEx for T where T : Copy { }
322
323
324////////////////////////////////////////////////////////////
325
326#[derive(Debug)]
327pub struct IndexHashMap<'a,'b,'c,'d,T : Copy> {
328    data   : & 'a mut[T],
329    index  : & 'b mut[usize],
330    next   : & 'c mut[usize],
331    bucket : & 'd mut[usize],
332    dflt   : T,
333    n      : usize
334}
335
336fn hash(i : usize) -> usize { i }
337
338impl<'a,'b,'c,'d,T : Copy> IndexHashMap<'a,'b,'c,'d,T> {
339    pub fn new(data   : & 'a mut[T],
340               index  : & 'b mut[usize],
341               next   : & 'c mut[usize],
342               bucket : & 'd mut[usize],
343               dflt   : T) -> IndexHashMap<'a,'b,'c,'d,T> {
344        bucket.iter_mut().for_each(|h| *h = usize::MAX);
345        if next.len() != data.len() || next.len() != index.len() {
346            panic!("Mismatching array sizes");
347        }
348
349        IndexHashMap{
350            data,
351            index,
352            next,
353            bucket,
354            dflt,
355            n : 0
356        }
357    }
358
359    pub fn with_data(data   : & 'a mut[T],
360                     index  : & 'b mut[usize],
361                     next   : & 'c mut[usize],
362                     bucket : & 'd mut[usize],
363                     dflt : T) -> IndexHashMap<'a,'b,'c,'d,T> {
364        bucket.iter_mut().for_each(|h| *h = usize::MAX);
365        let n = data.len();
366        let m = bucket.len();
367
368        if next.len() != data.len() || next.len() != index.len() {
369            panic!("Mismatching array sizes");
370        }
371
372        // Assume that index and data contains data to be put in the map
373        for (i,&k,next) in izip!(0..n,index.iter(), next.iter_mut()) {
374            let b = unsafe { &mut *bucket.get_unchecked_mut(hash(k) % m) };
375            *next = *b;
376            *b = i;
377        }
378
379        IndexHashMap{
380            data,
381            index,
382            next,
383            bucket,
384            dflt,
385            n}
386    }
387
388    pub fn at(&self,i : usize) -> Option<&T> {
389        let mut index = unsafe { * self.bucket.get_unchecked(hash(i) % self.bucket.len()) };
390
391        while index < usize::MAX && i != unsafe { * self.index.get_unchecked(index) }  {
392            index = unsafe{ * self.next.get_unchecked(index) };
393        }
394
395        if index < usize::MAX {
396            Some(unsafe { self.data.get_unchecked(index) })
397        }
398        else {
399            None
400        }
401    }
402
403    pub fn at_mut(&mut self, i : usize) -> &mut T {
404        let key = hash(i) % self.bucket.len();
405        let head = unsafe { self.bucket.get_unchecked_mut(key) };
406        let mut index = *head;
407
408        //println!("IndexHashMap, lookup {}\n\thead = {}",i,index);
409        while index < usize::MAX && i != unsafe { * self.index.get_unchecked(index) } {
410            //println!("\tindex = {}",index);
411            index = unsafe{ * self.next.get_unchecked(index) };
412        }
413
414        if index < usize::MAX {
415            unsafe { &mut * self.data.get_unchecked_mut(index) }
416        }
417        else if self.n < self.next.len() {
418            index = self.n; self.n += 1;
419            unsafe { *self.next.get_unchecked_mut(index) = *head; }
420            unsafe { *self.index.get_unchecked_mut(index) = i; }
421            *head = index;
422
423            unsafe { *self.data.get_unchecked_mut(index) = self.dflt; }
424            unsafe { & mut *self.data.get_unchecked_mut(index) }
425        }
426        else {
427            panic!("Hashmap is full");
428        }
429    }
430
431    pub fn len(&self) -> usize { self.n }
432}
433