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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
// Copyright 2017 Peter Williams
// Licensed under the MIT License.

/*!

Reading MIRIAD mask-format files, such as UV data flags.

 */

use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use std::io;

#[derive(Debug)]
pub struct MaskDecoder<R: io::Read> {
    stream: R,
    current_val: u32,
    bits_left_in_current: usize,
}

impl<R: io::Read> MaskDecoder<R> {
    pub fn new(stream: R) -> Self {
        MaskDecoder {
            stream: stream,
            current_val: 0,
            bits_left_in_current: 0,
        }
    }

    pub fn expand(&mut self, dest: &mut [bool]) -> Result<(), io::Error> {
        let mut ofs = 0;
        let mut cur = self.current_val;
        let mut n_bits = dest.len();

        while n_bits > 0 {
            if self.bits_left_in_current > 0 {
                let mut toread = ::std::cmp::min(self.bits_left_in_current, n_bits);
                let mut i = 31 - self.bits_left_in_current;

                n_bits -= toread;
                self.bits_left_in_current -= toread;

                while toread > 0 {
                    dest[ofs] = if cur & (1 << i) != 0 { true } else { false };

                    ofs += 1;
                    i += 1;
                    toread -= 1;
                }
            }

            if n_bits == 0 {
                return Ok(());
            }

            cur = self.stream.read_u32::<BigEndian>()?;
            self.current_val = cur;
            self.bits_left_in_current = 31;
        }

        return Ok(());
    }
}

#[derive(Debug)]
pub struct MaskEncoder<W: io::Write> {
    stream: W,
    current_val: u32,
    bits_left_in_current: usize,
    closed: bool,
}

impl<W: io::Write> MaskEncoder<W> {
    pub fn new(stream: W) -> Self {
        MaskEncoder {
            stream: stream,
            current_val: 0,
            bits_left_in_current: 0,
            closed: false,
        }
    }

    pub fn append_mask(&mut self, data: &[bool]) -> Result<(), io::Error> {
        if self.closed {
            panic!("cannot append to mask after closing it");
        }

        let mut ofs = 0;
        let mut n_bits = data.len();
        let mut cur = self.current_val;
        let mut bits_left = self.bits_left_in_current;

        while n_bits > 0 {
            // There should always be at least one bit left up here.

            let mut towrite = ::std::cmp::min(bits_left, n_bits);
            let mut i = 31 - bits_left;

            n_bits -= towrite;
            bits_left -= towrite;

            while towrite > 0 {
                if data[ofs] {
                    cur |= 1 << i;
                }

                ofs += 1;
                i += 1;
                towrite -= 1;
            }

            // We stopped either because this u32 is full, or because there
            // are no more data. (Well, possibly both at once.)

            if bits_left == 0 {
                self.stream.write_u32::<BigEndian>(cur)?;
                cur = 0;
                bits_left = 31;
            }
        }

        self.current_val = cur;
        self.bits_left_in_current = bits_left;
        return Ok(());
    }

    pub fn close(&mut self) -> Result<(), io::Error> {
        if self.closed {
            return Ok(());
        }

        if self.bits_left_in_current != 31 {
            self.stream.write_u32::<BigEndian>(self.current_val)?;
        }

        self.stream.flush()?;
        self.closed = true;
        Ok(())
    }
}

impl<W: io::Write> Drop for MaskEncoder<W> {
    fn drop(&mut self) {
        if !self.closed {
            let _r = self.close();
        }
    }
}