1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

//! Structs and utility routines related to memory management.

// others
use std::ops::{Add, Index, IndexMut};

// see memory.h

const LOG_BLOCK_SIZE: usize = 2;
const BLOCK_SIZE: usize = 1 << LOG_BLOCK_SIZE;

fn round_up(x: usize) -> usize {
    (x + BLOCK_SIZE - 1) & !(BLOCK_SIZE - 1)
}

#[derive(Debug, Clone, Default)]
pub struct BlockedArray<T> {
    pub data: Vec<T>,
    pub u_res: usize,
    pub v_res: usize,
    pub u_blocks: usize,
    log_block_size: usize,
    // block_size: usize,
}

impl<T> BlockedArray<T>
where
    T: num::Zero + Clone + Add<T, Output = T>,
{
    pub fn new(u_res: usize, v_res: usize) -> BlockedArray<T> {
        let data = vec![num::Zero::zero(); round_up(u_res) * round_up(v_res)];
        BlockedArray {
            u_res,
            v_res,
            u_blocks: round_up(u_res) >> LOG_BLOCK_SIZE,
            log_block_size: LOG_BLOCK_SIZE,
            // block_size: BLOCK_SIZE,
            data,
        }
    }
    pub fn new_from(u_res: usize, v_res: usize, d: &[T]) -> BlockedArray<T> {
        let mut ba = Self::new(u_res, v_res);
        for u in 0..u_res {
            for v in 0..v_res {
                ba[(u, v)] = d[v * u_res + u].clone();
            }
        }
        ba
    }
    pub fn u_size(&self) -> usize {
        self.u_res
    }
    pub fn v_size(&self) -> usize {
        self.v_res
    }
    pub fn block_size(&self) -> usize {
        1 << self.log_block_size
    }
    pub fn block(&self, a: usize) -> usize {
        a >> self.log_block_size
    }
    pub fn offset(&self, a: usize) -> usize {
        a & (self.block_size() - 1)
    }
}

impl<T> Index<(usize, usize)> for BlockedArray<T>
where
    T: num::Zero + std::clone::Clone + Add<T, Output = T>,
{
    type Output = T;
    fn index(&self, i: (usize, usize)) -> &T {
        let (u, v) = i;
        let bu = self.block(u);
        let bv = self.block(v);
        let ou = self.offset(u);
        let ov = self.offset(v);
        let offset = self.block_size() * self.block_size() * (self.u_blocks * bv + bu)
            + self.block_size() * ov
            + ou;
        &self.data[offset]
    }
}

impl<T> IndexMut<(usize, usize)> for BlockedArray<T>
where
    T: num::Zero + std::clone::Clone + Add<T, Output = T>,
{
    fn index_mut(&mut self, i: (usize, usize)) -> &mut T {
        let (u, v) = i;
        let bu = self.block(u);
        let bv = self.block(v);
        let ou = self.offset(u);
        let ov = self.offset(v);
        let offset = self.block_size() * self.block_size() * (self.u_blocks * bv + bu)
            + self.block_size() * ov
            + ou;
        &mut self.data[offset]
    }
}