<commit_before>//! This modules provides an implementation of the Lempel–Ziv–Welch Compression Algorithm

// Note: This implementation borrows heavily from the work of Julius Pettersson
// See http://www.cplusplus.com/articles/iL18T05o/ for his extensive explanations
// and a C++ implementatation

use std::io;
use std::io::{Read, Write};

use bitstream::{Bits, BitReader, BitWriter};

const MAX_CODESIZE: u8 = 12;
const MAX_ENTRIES: usize = 1 << MAX_CODESIZE as usize;

/// Alias for a LZW code point
type Code = u16;

/// Decoding dictionary.
///
/// It is not generic due to current limitations of Rust
/// Inspired by http://www.cplusplus.com/articles/iL18T05o/
#[derive(Debug)]
struct DecodingDict {
    min_size: u8,
    table: Vec<(Option<Code>, u8)>,
    buffer: Vec<u8>,
}

impl DecodingDict {
    /// Creates a new dict
    fn new(min_size: u8) -> DecodingDict {
        DecodingDict {
            min_size: min_size,
            table: Vec::with_capacity(512),
            buffer: Vec::with_capacity((1 << MAX_CODESIZE as usize) - 1)
        }
    }

    /// Resets the dictionary
    fn reset(&mut self) {
        self.table.clear();
        for i in (0..(1u16 << self.min_size as usize)) {
            self.table.push((None, i as u8));
        }
    }

    /// Inserts a value into the dict
    #[inline(always)]
    fn push(&mut self, key: Option<Code>, value: u8) {
        self.table.push((key, value))
    }

    /// Reconstructs the data for the corresponding code
    fn reconstruct(&mut self, code: Option<Code>) -> io::Result<&[u8]> {
        self.buffer.clear();
        let mut code = code;
        let mut cha;
        // Check the first access more thoroughly since a bad code
        // could occur if the data is malformed
        if let Some(k) = code {
            match self.table.get(k as usize) {
                Some(&(code_, cha_)) => {
                    code = code_;
                    cha = cha_;
                }
                None => return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    &*format!("Invalid code {:X}, expected code <= {:X}", k, self.table.len())
                ))
            }
            self.buffer.push(cha);
        }
        while let Some(k) = code {
            //(code, cha) = self.table[k as usize];
            // Note: This could possibly be replaced with an unchecked array access
            let entry = self.table[k as usize]; code = entry.0; cha = entry.1;
            self.buffer.push(cha);
        }
        self.buffer.reverse();
        Ok(&self.buffer)
    }

    /// Returns the buffer constructed by the last reconstruction
    #[inline(always)]
    fn buffer(&self) -> &[u8] {
        &self.buffer
    }

    /// Number of entries in the dictionary
    #[inline(always)]
    fn next_code(&self) -> u16 {
        self.table.len() as u16
    }
}

macro_rules! define_decoder_struct {
    {$(
        $name:ident, $offset:expr, #[$doc:meta];
    )*} => {

$( // START struct definition

#[$doc]
///
/// The maximum supported code size is 16 bits. The decoder assumes two
/// special code word to be present in the stream:
///
///  * `CLEAR_CODE == 1 << min_code_size`
///  * `END_CODE   == CLEAR_CODE + 1`
///
///Furthermore the decoder expects the stream to start with a `CLEAR_CODE`. This
/// corresponds to the implementation needed for en- and decoding GIF and TIFF files.
#[derive(Debug)]
pub struct $name<R: BitReader> {
    r: R,
    prev: Option<Code>,
    table: DecodingDict,
    buf: [u8; 1],
    code_size: u8,
    min_code_size: u8,
    clear_code: Code,
    end_code: Code,
}

impl<R> $name<R> where R: BitReader {
    /// Creates a new LZW decoder.
    pub fn new(reader: R, min_code_size: u8) -> $name<R> {
        $name {
            r: reader,
            prev: None,
            table: DecodingDict::new(min_code_size),
            buf: [0; 1],
            code_size: min_code_size + 1,
            min_code_size: min_code_size,
            clear_code: 1 << min_code_size,
            end_code: (1 << min_code_size) + 1,
        }
    }

    /// Tries to obtain and decode a code word from `bytes`.
    ///
    /// Returns the number of bytes that have been consumed from `bytes`. An empty
    /// slice does not indicate `EOF`.
    pub fn decode_bytes(&mut self, bytes: &[u8]) -> io::Result<(usize, &[u8])> {
        Ok(match self.r.read_bits(bytes, self.code_size) {
            Bits::Some(consumed, code) => {
                (consumed, if code == self.clear_code {
                    self.table.reset();
                    self.table.push(None, 0); // clear code
                    self.table.push(None, 0); // end code
                    self.code_size = self.min_code_size + 1;
                    self.prev = None;
                    &[]
                } else if code == self.end_code {
                    &[]
                } else {
                    let next_code = self.table.next_code();
                    let prev = self.prev;
                    let result = if prev.is_none() {
                        self.buf = [code as u8];
                        &self.buf[..]
                    } else {
                        let data = if code == next_code {
                            let cha = try!(self.table.reconstruct(prev))[0];
                            self.table.push(prev, cha);
                            try!(self.table.reconstruct(Some(code)))
                        } else if code < next_code {
                            let cha = try!(self.table.reconstruct(Some(code)))[0];
                            self.table.push(prev, cha);
                            self.table.buffer()
                        } else {
                            return Err(io::Error::new(
                                io::ErrorKind::InvalidInput,
                                &*format!("Invalid code {:X}, expected code <= {:X}",
                                          code,
                                          next_code
                                )
                            ))
                        };
                        data
                    };
                    if next_code == (1 << self.code_size as usize) - 1 - $offset
                       && self.code_size < MAX_CODESIZE {
                        self.code_size += 1;
                    }
                    self.prev = Some(code);
                    result
                })
            },
            Bits::None(consumed) => {
                (consumed, &[])
            }
        })
    }
}

)* // END struct definition

    }
}

define_decoder_struct!{
    Decoder, 0, #[doc = "Decoder for a LZW compressed stream (this algorithm is used for GIF files)."];
    DecoderEarlyChange, 1, #[doc = "Decoder for a LZW compressed stream using an “early change” algorithm (used in TIFF files)."];
}

struct Node {
    prefix: Option<Code>,
    c: u8,
    left: Option<Code>,
    right: Option<Code>,
}

impl Node {
    #[inline(always)]
    fn new(c: u8) -> Node {
        Node {
            prefix: None,
            c: c,
            left: None,
            right: None
        }
    }
}

struct EncodingDict {
    table: Vec<Node>,
    min_size: u8,

}

/// Encoding dictionary based on a binary tree
impl EncodingDict {
    fn new(min_size: u8) -> EncodingDict {
        let mut this = EncodingDict {
            table: Vec::with_capacity(MAX_ENTRIES),
            min_size: min_size
        };
        this.reset();
        this
    }

    fn reset(&mut self) {
        self.table.clear();
        for i in 0 .. (1u16 << self.min_size as usize) {
            self.push_node(Node::new(i as u8));
        }
    }

    #[inline(always)]
    fn push_node(&mut self, node: Node) {
        self.table.push(node)
    }

    #[inline(always)]
    fn clear_code(&self) -> Code {
        1u16 << self.min_size
    }

    #[inline(always)]
    fn end_code(&self) -> Code {
        self.clear_code() + 1
    }

    // Searches for a new prefix
    fn search_and_insert(&mut self, i: Option<Code>, c: u8) -> Option<Code> {
        if let Some(i) = i.map(|v| v as usize) {
            let table_size = self.table.len() as Code;
            if let Some(mut j) = self.table[i].prefix {
                loop {
                    let entry = &mut self.table[j as usize];
                    if c < entry.c {
                        if let Some(k) = entry.left {
                            j = k
                        } else {
                            entry.left = Some(table_size);
                            break
                        }
                    } else if c > entry.c {
                        if let Some(k) = entry.right {
                            j = k
                        } else {
                            entry.right = Some(table_size);
                            break
                        }
                    } else {
                        return Some(j)
                    }
                }
            } else {
                self.table[i].prefix = Some(table_size);
            }
            self.table.push(Node::new(c));
            None
        } else {
            Some(self.search_initials(c as Code))
        }
    }

    fn next_code(&self) -> usize {
        self.table.len()
    }

    fn search_initials(&self, i: Code) -> Code {
        self.table[i as usize].c as Code
    }
}

/// Convenience function that reads and compresses all bytes from `R`.
pub fn encode<R, W>(r: R, mut w: W, min_code_size: u8) -> io::Result<()>
where R: Read, W: BitWriter {
    let mut dict = EncodingDict::new(min_code_size);
    dict.push_node(Node::new(0)); // clear code
    dict.push_node(Node::new(0)); // end code
    let mut code_size = min_code_size + 1;
    let mut i = None;
    // gif spec: first clear code
    try!(w.write_bits(dict.clear_code(), code_size));
    let mut r = r.bytes();
    while let Some(Ok(c)) = r.next() {
        let prev = i;
        i = dict.search_and_insert(prev, c);
        if i.is_none() {
            if let Some(code) = prev {
                try!(w.write_bits(code, code_size));
            }
            i = Some(dict.search_initials(c as Code))
        }
        // There is a hit: do not write out code but continue
        let next_code = dict.next_code();
        if next_code > (1 << code_size as usize)
           && code_size < MAX_CODESIZE {
            code_size += 1;
        }
        if next_code > MAX_ENTRIES {
            dict.reset();
            dict.push_node(Node::new(0)); // clear code
            dict.push_node(Node::new(0)); // end code
            try!(w.write_bits(dict.clear_code(), code_size));
            code_size = min_code_size + 1;
        }

    }
    if let Some(code) = i {
        try!(w.write_bits(code, code_size));
    }
    try!(w.write_bits(dict.end_code(), code_size));
    try!(w.flush());
    Ok(())
}

/// LZW encoder using the algorithm of GIF files.
pub struct Encoder<W: BitWriter> {
    w: W,
    dict: EncodingDict,
    min_code_size: u8,
    code_size: u8,
    i: Option<Code>
}

impl<W: BitWriter> Encoder<W> {
    /// Creates a new LZW encoder.
    ///
    /// **Note**: If `min_code_size < 8` then `Self::encode_bytes` might panic when
    /// the supplied data containts values that exceed `1 << min_code_size`.
    pub fn new(mut w: W, min_code_size: u8) -> io::Result<Encoder<W>> {
        let mut dict = EncodingDict::new(min_code_size);
        dict.push_node(Node::new(0)); // clear code
        dict.push_node(Node::new(0)); // end code
        let code_size = min_code_size + 1;
        try!(w.write_bits(dict.clear_code(), code_size));
        Ok(Encoder {
            w: w,
            dict: dict,
            min_code_size: min_code_size,
            code_size: code_size,
            i: None
        })
    }

    /// Compresses `bytes` and writes the result into the writer.
    ///
    /// ## Panics
    ///
    /// This function might panic if any of the input bytes exceeds `1 << min_code_size`.
    /// This cannot happen if `min_code_size >= 8`.
    pub fn encode_bytes(&mut self, bytes: &[u8]) -> io::Result<()> {
        let w = &mut self.w;
        let dict = &mut self.dict;
        let code_size = &mut self.code_size;
        let i = &mut self.i;
        for &c in bytes {
            let prev = *i;
            *i = dict.search_and_insert(prev, c);
            if i.is_none() {
                if let Some(code) = prev {
                    try!(w.write_bits(code, *code_size));
                }
                *i = Some(dict.search_initials(c as Code))
            }
            // There is a hit: do not write out code but continue
            let next_code = dict.next_code();
            if next_code > (1 << *code_size as usize)
               && *code_size < MAX_CODESIZE {
                *code_size += 1;
            }
            if next_code > MAX_ENTRIES {
                dict.reset();
                dict.push_node(Node::new(0)); // clear code
                dict.push_node(Node::new(0)); // end code
                try!(w.write_bits(dict.clear_code(), *code_size));
                *code_size = self.min_code_size + 1;
            }

        }
        Ok(())
    }
}

impl<W: BitWriter> Drop for Encoder<W> {
    #[cfg(feature = "raii_no_panic")]
    fn drop(&mut self) {
        let w = &mut self.w;
        let code_size = &mut self.code_size;

        if let Some(code) = self.i {
            let _ = w.write_bits(code, *code_size);
        }
        let _ = w.write_bits(self.dict.end_code(), *code_size);
    }
    #[cfg(not(feature = "raii_no_panic"))]
    fn drop(&mut self) {
        (|| {
            let w = &mut self.w;
            let code_size = &mut self.code_size;
            if let Some(code) = self.i {
                try!(w.write_bits(code, *code_size));
            }
            w.write_bits(self.dict.end_code(), *code_size)
         })().unwrap()

    }
}

#[cfg(test)]
#[test]
fn round_trip() {
    use {LsbWriter, LsbReader};

    let size = 8;
    let data = b"TOBEORNOTTOBEORTOBEORNOT";
    let mut compressed = vec![];
    {
        let mut enc = Encoder::new(LsbWriter::new(&mut compressed), size).unwrap();
        enc.encode_bytes(data).unwrap();
    }
    println!("{:?}", compressed);
    let mut dec = Decoder::new(LsbReader::new(), size);
    let mut compressed = &compressed[..];
    let mut data2 = vec![];
    while compressed.len() > 0 {
        let (start, bytes) = dec.decode_bytes(&compressed).unwrap();
        compressed = &compressed[start..];
        data2.extend(bytes.iter().map(|&i| i));
    }
    assert_eq!(data2, data)
}<commit_msg>Restore removed flush.<commit_after>//! This modules provides an implementation of the Lempel–Ziv–Welch Compression Algorithm

// Note: This implementation borrows heavily from the work of Julius Pettersson
// See http://www.cplusplus.com/articles/iL18T05o/ for his extensive explanations
// and a C++ implementatation

use std::io;
use std::io::{Read, Write};

use bitstream::{Bits, BitReader, BitWriter};

const MAX_CODESIZE: u8 = 12;
const MAX_ENTRIES: usize = 1 << MAX_CODESIZE as usize;

/// Alias for a LZW code point
type Code = u16;

/// Decoding dictionary.
///
/// It is not generic due to current limitations of Rust
/// Inspired by http://www.cplusplus.com/articles/iL18T05o/
#[derive(Debug)]
struct DecodingDict {
    min_size: u8,
    table: Vec<(Option<Code>, u8)>,
    buffer: Vec<u8>,
}

impl DecodingDict {
    /// Creates a new dict
    fn new(min_size: u8) -> DecodingDict {
        DecodingDict {
            min_size: min_size,
            table: Vec::with_capacity(512),
            buffer: Vec::with_capacity((1 << MAX_CODESIZE as usize) - 1)
        }
    }

    /// Resets the dictionary
    fn reset(&mut self) {
        self.table.clear();
        for i in (0..(1u16 << self.min_size as usize)) {
            self.table.push((None, i as u8));
        }
    }

    /// Inserts a value into the dict
    #[inline(always)]
    fn push(&mut self, key: Option<Code>, value: u8) {
        self.table.push((key, value))
    }

    /// Reconstructs the data for the corresponding code
    fn reconstruct(&mut self, code: Option<Code>) -> io::Result<&[u8]> {
        self.buffer.clear();
        let mut code = code;
        let mut cha;
        // Check the first access more thoroughly since a bad code
        // could occur if the data is malformed
        if let Some(k) = code {
            match self.table.get(k as usize) {
                Some(&(code_, cha_)) => {
                    code = code_;
                    cha = cha_;
                }
                None => return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    &*format!("Invalid code {:X}, expected code <= {:X}", k, self.table.len())
                ))
            }
            self.buffer.push(cha);
        }
        while let Some(k) = code {
            //(code, cha) = self.table[k as usize];
            // Note: This could possibly be replaced with an unchecked array access
            let entry = self.table[k as usize]; code = entry.0; cha = entry.1;
            self.buffer.push(cha);
        }
        self.buffer.reverse();
        Ok(&self.buffer)
    }

    /// Returns the buffer constructed by the last reconstruction
    #[inline(always)]
    fn buffer(&self) -> &[u8] {
        &self.buffer
    }

    /// Number of entries in the dictionary
    #[inline(always)]
    fn next_code(&self) -> u16 {
        self.table.len() as u16
    }
}

macro_rules! define_decoder_struct {
    {$(
        $name:ident, $offset:expr, #[$doc:meta];
    )*} => {

$( // START struct definition

#[$doc]
///
/// The maximum supported code size is 16 bits. The decoder assumes two
/// special code word to be present in the stream:
///
///  * `CLEAR_CODE == 1 << min_code_size`
///  * `END_CODE   == CLEAR_CODE + 1`
///
///Furthermore the decoder expects the stream to start with a `CLEAR_CODE`. This
/// corresponds to the implementation needed for en- and decoding GIF and TIFF files.
#[derive(Debug)]
pub struct $name<R: BitReader> {
    r: R,
    prev: Option<Code>,
    table: DecodingDict,
    buf: [u8; 1],
    code_size: u8,
    min_code_size: u8,
    clear_code: Code,
    end_code: Code,
}

impl<R> $name<R> where R: BitReader {
    /// Creates a new LZW decoder.
    pub fn new(reader: R, min_code_size: u8) -> $name<R> {
        $name {
            r: reader,
            prev: None,
            table: DecodingDict::new(min_code_size),
            buf: [0; 1],
            code_size: min_code_size + 1,
            min_code_size: min_code_size,
            clear_code: 1 << min_code_size,
            end_code: (1 << min_code_size) + 1,
        }
    }

    /// Tries to obtain and decode a code word from `bytes`.
    ///
    /// Returns the number of bytes that have been consumed from `bytes`. An empty
    /// slice does not indicate `EOF`.
    pub fn decode_bytes(&mut self, bytes: &[u8]) -> io::Result<(usize, &[u8])> {
        Ok(match self.r.read_bits(bytes, self.code_size) {
            Bits::Some(consumed, code) => {
                (consumed, if code == self.clear_code {
                    self.table.reset();
                    self.table.push(None, 0); // clear code
                    self.table.push(None, 0); // end code
                    self.code_size = self.min_code_size + 1;
                    self.prev = None;
                    &[]
                } else if code == self.end_code {
                    &[]
                } else {
                    let next_code = self.table.next_code();
                    let prev = self.prev;
                    let result = if prev.is_none() {
                        self.buf = [code as u8];
                        &self.buf[..]
                    } else {
                        let data = if code == next_code {
                            let cha = try!(self.table.reconstruct(prev))[0];
                            self.table.push(prev, cha);
                            try!(self.table.reconstruct(Some(code)))
                        } else if code < next_code {
                            let cha = try!(self.table.reconstruct(Some(code)))[0];
                            self.table.push(prev, cha);
                            self.table.buffer()
                        } else {
                            return Err(io::Error::new(
                                io::ErrorKind::InvalidInput,
                                &*format!("Invalid code {:X}, expected code <= {:X}",
                                          code,
                                          next_code
                                )
                            ))
                        };
                        data
                    };
                    if next_code == (1 << self.code_size as usize) - 1 - $offset
                       && self.code_size < MAX_CODESIZE {
                        self.code_size += 1;
                    }
                    self.prev = Some(code);
                    result
                })
            },
            Bits::None(consumed) => {
                (consumed, &[])
            }
        })
    }
}

)* // END struct definition

    }
}

define_decoder_struct!{
    Decoder, 0, #[doc = "Decoder for a LZW compressed stream (this algorithm is used for GIF files)."];
    DecoderEarlyChange, 1, #[doc = "Decoder for a LZW compressed stream using an “early change” algorithm (used in TIFF files)."];
}

struct Node {
    prefix: Option<Code>,
    c: u8,
    left: Option<Code>,
    right: Option<Code>,
}

impl Node {
    #[inline(always)]
    fn new(c: u8) -> Node {
        Node {
            prefix: None,
            c: c,
            left: None,
            right: None
        }
    }
}

struct EncodingDict {
    table: Vec<Node>,
    min_size: u8,

}

/// Encoding dictionary based on a binary tree
impl EncodingDict {
    fn new(min_size: u8) -> EncodingDict {
        let mut this = EncodingDict {
            table: Vec::with_capacity(MAX_ENTRIES),
            min_size: min_size
        };
        this.reset();
        this
    }

    fn reset(&mut self) {
        self.table.clear();
        for i in 0 .. (1u16 << self.min_size as usize) {
            self.push_node(Node::new(i as u8));
        }
    }

    #[inline(always)]
    fn push_node(&mut self, node: Node) {
        self.table.push(node)
    }

    #[inline(always)]
    fn clear_code(&self) -> Code {
        1u16 << self.min_size
    }

    #[inline(always)]
    fn end_code(&self) -> Code {
        self.clear_code() + 1
    }

    // Searches for a new prefix
    fn search_and_insert(&mut self, i: Option<Code>, c: u8) -> Option<Code> {
        if let Some(i) = i.map(|v| v as usize) {
            let table_size = self.table.len() as Code;
            if let Some(mut j) = self.table[i].prefix {
                loop {
                    let entry = &mut self.table[j as usize];
                    if c < entry.c {
                        if let Some(k) = entry.left {
                            j = k
                        } else {
                            entry.left = Some(table_size);
                            break
                        }
                    } else if c > entry.c {
                        if let Some(k) = entry.right {
                            j = k
                        } else {
                            entry.right = Some(table_size);
                            break
                        }
                    } else {
                        return Some(j)
                    }
                }
            } else {
                self.table[i].prefix = Some(table_size);
            }
            self.table.push(Node::new(c));
            None
        } else {
            Some(self.search_initials(c as Code))
        }
    }

    fn next_code(&self) -> usize {
        self.table.len()
    }

    fn search_initials(&self, i: Code) -> Code {
        self.table[i as usize].c as Code
    }
}

/// Convenience function that reads and compresses all bytes from `R`.
pub fn encode<R, W>(r: R, mut w: W, min_code_size: u8) -> io::Result<()>
where R: Read, W: BitWriter {
    let mut dict = EncodingDict::new(min_code_size);
    dict.push_node(Node::new(0)); // clear code
    dict.push_node(Node::new(0)); // end code
    let mut code_size = min_code_size + 1;
    let mut i = None;
    // gif spec: first clear code
    try!(w.write_bits(dict.clear_code(), code_size));
    let mut r = r.bytes();
    while let Some(Ok(c)) = r.next() {
        let prev = i;
        i = dict.search_and_insert(prev, c);
        if i.is_none() {
            if let Some(code) = prev {
                try!(w.write_bits(code, code_size));
            }
            i = Some(dict.search_initials(c as Code))
        }
        // There is a hit: do not write out code but continue
        let next_code = dict.next_code();
        if next_code > (1 << code_size as usize)
           && code_size < MAX_CODESIZE {
            code_size += 1;
        }
        if next_code > MAX_ENTRIES {
            dict.reset();
            dict.push_node(Node::new(0)); // clear code
            dict.push_node(Node::new(0)); // end code
            try!(w.write_bits(dict.clear_code(), code_size));
            code_size = min_code_size + 1;
        }

    }
    if let Some(code) = i {
        try!(w.write_bits(code, code_size));
    }
    try!(w.write_bits(dict.end_code(), code_size));
    try!(w.flush());
    Ok(())
}

/// LZW encoder using the algorithm of GIF files.
pub struct Encoder<W: BitWriter> {
    w: W,
    dict: EncodingDict,
    min_code_size: u8,
    code_size: u8,
    i: Option<Code>
}

impl<W: BitWriter> Encoder<W> {
    /// Creates a new LZW encoder.
    ///
    /// **Note**: If `min_code_size < 8` then `Self::encode_bytes` might panic when
    /// the supplied data containts values that exceed `1 << min_code_size`.
    pub fn new(mut w: W, min_code_size: u8) -> io::Result<Encoder<W>> {
        let mut dict = EncodingDict::new(min_code_size);
        dict.push_node(Node::new(0)); // clear code
        dict.push_node(Node::new(0)); // end code
        let code_size = min_code_size + 1;
        try!(w.write_bits(dict.clear_code(), code_size));
        Ok(Encoder {
            w: w,
            dict: dict,
            min_code_size: min_code_size,
            code_size: code_size,
            i: None
        })
    }

    /// Compresses `bytes` and writes the result into the writer.
    ///
    /// ## Panics
    ///
    /// This function might panic if any of the input bytes exceeds `1 << min_code_size`.
    /// This cannot happen if `min_code_size >= 8`.
    pub fn encode_bytes(&mut self, bytes: &[u8]) -> io::Result<()> {
        let w = &mut self.w;
        let dict = &mut self.dict;
        let code_size = &mut self.code_size;
        let i = &mut self.i;
        for &c in bytes {
            let prev = *i;
            *i = dict.search_and_insert(prev, c);
            if i.is_none() {
                if let Some(code) = prev {
                    try!(w.write_bits(code, *code_size));
                }
                *i = Some(dict.search_initials(c as Code))
            }
            // There is a hit: do not write out code but continue
            let next_code = dict.next_code();
            if next_code > (1 << *code_size as usize)
               && *code_size < MAX_CODESIZE {
                *code_size += 1;
            }
            if next_code > MAX_ENTRIES {
                dict.reset();
                dict.push_node(Node::new(0)); // clear code
                dict.push_node(Node::new(0)); // end code
                try!(w.write_bits(dict.clear_code(), *code_size));
                *code_size = self.min_code_size + 1;
            }

        }
        Ok(())
    }
}

impl<W: BitWriter> Drop for Encoder<W> {
    #[cfg(feature = "raii_no_panic")]
    fn drop(&mut self) {
        let w = &mut self.w;
        let code_size = &mut self.code_size;

        if let Some(code) = self.i {
            let _ = w.write_bits(code, *code_size);
        }
        let _ = w.write_bits(self.dict.end_code(), *code_size);
        let _ = w.flush();
    }
    #[cfg(not(feature = "raii_no_panic"))]
    fn drop(&mut self) {
        (|| {
            let w = &mut self.w;
            let code_size = &mut self.code_size;
            if let Some(code) = self.i {
                try!(w.write_bits(code, *code_size));
            }
            try!(w.write_bits(self.dict.end_code(), *code_size));
            w.flush()
         })().unwrap()

    }
}

#[cfg(test)]
#[test]
fn round_trip() {
    use {LsbWriter, LsbReader};

    let size = 8;
    let data = b"TOBEORNOTTOBEORTOBEORNOT";
    let mut compressed = vec![];
    {
        let mut enc = Encoder::new(LsbWriter::new(&mut compressed), size).unwrap();
        enc.encode_bytes(data).unwrap();
    }
    println!("{:?}", compressed);
    let mut dec = Decoder::new(LsbReader::new(), size);
    let mut compressed = &compressed[..];
    let mut data2 = vec![];
    while compressed.len() > 0 {
        let (start, bytes) = dec.decode_bytes(&compressed).unwrap();
        compressed = &compressed[start..];
        data2.extend(bytes.iter().map(|&i| i));
    }
    assert_eq!(data2, data)
}<|endoftext|>


<commit_before>use std::rc::Rc;

use syntax::ext::build::AstBuilder;
use syntax::ext::base::ExtCtxt;
use syntax::ast;
use syntax::ptr::P;
use syntax::symbol::Symbol;
use syntax::codemap::{Span, Spanned};

pub type Ident = Spanned<ast::Ident>;

#[derive(Debug, PartialOrd, PartialEq, Ord, Eq, Hash, Clone, Copy)]
pub enum Size {
    BYTE  = 1,
    WORD  = 2,
    DWORD = 4,
    QWORD = 8,
    PWORD = 10,
    OWORD = 16,
    HWORD = 32
}

impl Size {
    pub fn in_bytes(&self) -> u8 {
        *self as u8
    }

    pub fn as_literal(&self) -> ast::Ident {
        ast::Ident::from_str(match *self {
            Size::BYTE  => "i8",
            Size::WORD  => "i16",
            Size::DWORD => "i32",
            Size::QWORD => "i64",
            Size::PWORD => "i80",
            Size::OWORD => "i128",
            Size::HWORD => "i256"
        })
    }
}

#[derive(Clone, Debug)]
pub enum Stmt {
    Const(u8),
    ExprConst(P<ast::Expr>),

    Var(P<ast::Expr>, Size),
    Extend(P<ast::Expr>),

    DynScale(P<ast::Expr>, P<ast::Expr>),

    Align(P<ast::Expr>),

    GlobalLabel(Ident),
    LocalLabel(Ident),
    DynamicLabel(P<ast::Expr>),

    GlobalJumpTarget(Ident, Size),
    ForwardJumpTarget(Ident, Size),
    BackwardJumpTarget(Ident, Size),
    DynamicJumpTarget(P<ast::Expr>, Size),

    Stmt(ast::Stmt),
}

pub fn serialize(ecx: &mut ExtCtxt, name: P<ast::Expr>, stmts: Vec<Stmt>) -> Vec<ast::Stmt> {
    let mut buffer = Vec::new();

    // construction for `op.push(expr)` is as follows
    // op = Path::from_ident(name)
    // push = Ident::from_str("push")
    // expr = expr_lit(span, LitKind::Byte)
    // expr_method_call(span, op, Vec![expr])

    let mut stmts = stmts.into_iter().peekable();

    while let Some(stmt) = stmts.next() {
        use self::Stmt::*;

        let (method, args) = match stmt {
            Const(byte)            => {
                let mut bytes = vec![byte];
                while let Some(&Const(byte)) = stmts.peek() {
                    bytes.push(byte);
                    stmts.next();
                    if stmts.len() == 32 {
                        break;
                    }
                }

                if bytes.len() == 1 {
                    ("push",   vec![ecx.expr_lit(ecx.call_site(), ast::LitKind::Byte(bytes[0]))])
                } else {
                    ("extend", vec![ecx.expr_lit(ecx.call_site(), ast::LitKind::ByteStr(Rc::new(bytes)))])
                }
            },
            ExprConst(expr)        => ("push",    vec![expr]),

            Var(expr, Size::BYTE)  => ("push_i8",  vec![expr]),
            Var(expr, Size::WORD)  => ("push_i16", vec![expr]),
            Var(expr, Size::DWORD) => ("push_i32", vec![expr]),
            Var(expr, Size::QWORD) => ("push_i64", vec![expr]),
            Var(_, _)           => panic!("immediate serializaiton of this size is not supported yet"),

            Extend(expr)           => ("extend", vec![expr]),

            DynScale(scale, rest)  => {
                let temp = ast::Ident::from_str("temp");
                buffer.push(ecx.stmt_let(ecx.call_site(), false, temp, encoded_size_expr(ecx, &name, scale)));
                ("push", vec![or_mask_shift_expr(ecx, rest, ecx.expr_ident(ecx.call_site(), temp), 3, 6)])
            },

            Align(expr)            => ("align",   vec![expr]),

            GlobalLabel(ident)     => ("global_label", vec![ecx.expr_lit(
                ident.span,
                ast::LitKind::Str(ident.node.name, ast::StrStyle::Cooked)
            )]),
            LocalLabel(ident)      => ("local_label", vec![ecx.expr_lit(
                ident.span,
                ast::LitKind::Str(ident.node.name, ast::StrStyle::Cooked)
            )]),
            DynamicLabel(expr)     => ("dynamic_label", vec![expr]),

            GlobalJumpTarget(ident, size) => ("global_reloc", vec![ecx.expr_lit(
                ident.span,
                ast::LitKind::Str(ident.node.name, ast::StrStyle::Cooked)
            ), ecx.expr_u8(ident.span, size.in_bytes())]),
            ForwardJumpTarget(ident, size) => ("forward_reloc", vec![ecx.expr_lit(
                ident.span,
                ast::LitKind::Str(ident.node.name, ast::StrStyle::Cooked)
            ), ecx.expr_u8(ident.span, size.in_bytes())]),
            BackwardJumpTarget(ident, size) => ("backward_reloc", vec![ecx.expr_lit(
                ident.span,
                ast::LitKind::Str(ident.node.name, ast::StrStyle::Cooked)
            ), ecx.expr_u8(ident.span, size.in_bytes())]),
            DynamicJumpTarget(expr, size) => {
                let span = expr.span;
                ("dynamic_reloc", vec![expr, ecx.expr_u8(span, size.in_bytes())])
            },
            Stmt(stmt) => {
                buffer.push(stmt);
                continue;
            }
        };

        let op = name.clone();
        let method = ast::Ident::from_str(method);
        let expr = ecx.expr_method_call(ecx.call_site(), op, method, args);

        buffer.push(ecx.stmt_semi(expr));
    }

    buffer
}

pub fn add_exprs<T: Iterator<Item=P<ast::Expr>>>(ecx: &ExtCtxt, span: Span, mut exprs: T) -> Option<P<ast::Expr>> {
    exprs.next().map(|mut accum| {
        for next in exprs {
            accum = ecx.expr_binary(span, ast::BinOpKind::Add, accum, next);
        }
        accum
    })
}

pub fn size_of_scale_expr(ecx: &ExtCtxt, ty: ast::Path, value: P<ast::Expr>, size: Size) -> P<ast::Expr> {
    let span = value.span;
    ecx.expr_binary(span,
        ast::BinOpKind::Mul,
        ecx.expr_cast(span,
            size_of_expr(ecx, ty),
            ecx.ty_ident(span, size.as_literal())
        ),
        value
    )
}

pub fn or_mask_shift_expr(ecx: &ExtCtxt, orig: P<ast::Expr>, mut expr: P<ast::Expr>, mask: u64, shift: i8) -> P<ast::Expr> {
    let span = expr.span;
    // take expr and return !((expr & mask) << shift)

    expr = ecx.expr_binary(span, ast::BinOpKind::BitAnd, expr, ecx.expr_lit(
        span, ast::LitKind::Int(mask as u128, ast::LitIntType::Unsuffixed)
    ));

    if shift < 0 {
        expr = ecx.expr_binary(span, ast::BinOpKind::Shr, expr, ecx.expr_lit(
            span, ast::LitKind::Int((-shift) as u128, ast::LitIntType::Unsuffixed)
        ));
    } else if shift > 0 {
        expr = ecx.expr_binary(span, ast::BinOpKind::Shl, expr, ecx.expr_lit(
            span, ast::LitKind::Int(shift as u128, ast::LitIntType::Unsuffixed)
        ));
    }

    ecx.expr_binary(span, ast::BinOpKind::BitOr, orig, expr)
}

pub fn offset_of_expr(ecx: &ExtCtxt, path: ast::Path, attr: ast::Ident, size: Size) -> P<ast::Expr> {
    // generate a P<Expr> that resolves into the offset of an attribute to a type.
    // this is somewhat ridiculously complex because we can't expand macros here

    let span = path.span;

    let structpat = ecx.pat_struct(span, path.clone(), vec![
        Spanned {span: span, node: ast::FieldPat {
            ident: attr,
            pat: ecx.pat_wild(span),
            is_shorthand: false
        }},
    ]).map(|mut pat| {
        if let ast::PatKind::Struct(_, _, ref mut dotdot) = pat.node {
            *dotdot = true;
        }
        pat
    });

    // there's no default constructor function for let pattern;
    let validation_stmt = ast::Stmt {
        id: ast::DUMMY_NODE_ID,
        span: span,
        node: ast::StmtKind::Local(P(ast::Local {
            pat: structpat,
            ty: None,
            init: None,
            id: ast::DUMMY_NODE_ID,
            span: span,
            attrs: ast::ThinVec::new()
        }))
    };

    let temp     = ast::Ident::from_str("temp");
    let rv       = ast::Ident::from_str("rv");
    let usize_id = ast::Ident::from_str("usize");
    let uninitialized = ["std", "mem", "uninitialized"].iter().cloned().map(ast::Ident::from_str).collect();
    let forget        = ["std", "mem", "forget"       ].iter().cloned().map(ast::Ident::from_str).collect();

    // unsafe {
    let block = ecx.block(span, vec![
        // let path { attr: _, ..};
        validation_stmt,
        // let temp: path = ::std::mem::uninitialized();
        ecx.stmt_let_typed(span, false, temp, ecx.ty_path(path),
            ecx.expr_call_global(span, uninitialized, Vec::new())
        ),
        // let rv = &temp.attr as *const _ as usize - &temp as *const _ as usize;
        ecx.stmt_let(span,
            false,
            rv,
            ecx.expr_binary(span, ast::BinOpKind::Sub,
                ecx.expr_cast(span,
                    ecx.expr_cast(span,
                        ecx.expr_addr_of(span,
                            ecx.expr_field_access(span,
                                ecx.expr_ident(span, temp),
                                attr
                            )
                        ), ecx.ty_ptr(span, ecx.ty_infer(span), ast::Mutability::Immutable)
                    ), ecx.ty_ident(span, usize_id)
                ),
                ecx.expr_cast(span,
                    ecx.expr_cast(span,
                        ecx.expr_addr_of(span, ecx.expr_ident(span, temp)),
                        ecx.ty_ptr(span, ecx.ty_infer(span), ast::Mutability::Immutable)
                    ), ecx.ty_ident(span, usize_id)
                )
            )
        ),
        // ::std::mem::forget(temp);
        ecx.stmt_semi(ecx.expr_call_global(span, forget, vec![ecx.expr_ident(span, temp)])),
        // rv as i32
        ecx.stmt_expr(ecx.expr_cast(span, ecx.expr_ident(span, rv), ecx.ty_ident(span, size.as_literal())))
    ]).map(|mut b| {
        b.rules = ast::BlockCheckMode::Unsafe(ast::UnsafeSource::CompilerGenerated);
        b
    });

    ecx.expr_block(block)
}

pub fn size_of_expr(ecx: &ExtCtxt, path: ast::Path) -> P<ast::Expr> {
    // generate a P<Expr> that returns the size of type at path
    let span = path.span;

    let ty = ecx.ty_path(path);
    let idents = ["std", "mem", "size_of"].iter().cloned().map(ast::Ident::from_str).collect();
    let size_of = ecx.path_all(span, true, idents, Vec::new(), vec![ty], Vec::new());
    ecx.expr_call(span, ecx.expr_path(size_of), Vec::new())
}

pub fn encoded_size_expr(ecx: &ExtCtxt, name: &P<ast::Expr>, size: P<ast::Expr>) -> P<ast::Expr> {
    let span = size.span;

    ecx.expr_match(span, size, vec![
        ecx.arm(span, vec![ecx.pat_lit(span, ecx.expr_usize(span, 8))], ecx.expr_u8(span, 3)),
        ecx.arm(span, vec![ecx.pat_lit(span, ecx.expr_usize(span, 4))], ecx.expr_u8(span, 2)),
        ecx.arm(span, vec![ecx.pat_lit(span, ecx.expr_usize(span, 2))], ecx.expr_u8(span, 1)),
        ecx.arm(span, vec![ecx.pat_lit(span, ecx.expr_usize(span, 1))], ecx.expr_u8(span, 0)),
        ecx.arm(span, vec![ecx.pat_wild(span)], ecx.expr_method_call(span,
            name.clone(),
            ast::Ident::from_str("runtime_error"),
            vec![ecx.expr_str(span,
                Symbol::intern("Type size not representable as scale")
            )]
        ))
    ])
}
<commit_msg>Fix an error in code generation where long sequences of operands without variable substitution would result in compile-time errors<commit_after>use std::rc::Rc;

use syntax::ext::build::AstBuilder;
use syntax::ext::base::ExtCtxt;
use syntax::ast;
use syntax::ptr::P;
use syntax::symbol::Symbol;
use syntax::codemap::{Span, Spanned};

pub type Ident = Spanned<ast::Ident>;

#[derive(Debug, PartialOrd, PartialEq, Ord, Eq, Hash, Clone, Copy)]
pub enum Size {
    BYTE  = 1,
    WORD  = 2,
    DWORD = 4,
    QWORD = 8,
    PWORD = 10,
    OWORD = 16,
    HWORD = 32
}

impl Size {
    pub fn in_bytes(&self) -> u8 {
        *self as u8
    }

    pub fn as_literal(&self) -> ast::Ident {
        ast::Ident::from_str(match *self {
            Size::BYTE  => "i8",
            Size::WORD  => "i16",
            Size::DWORD => "i32",
            Size::QWORD => "i64",
            Size::PWORD => "i80",
            Size::OWORD => "i128",
            Size::HWORD => "i256"
        })
    }
}

#[derive(Clone, Debug)]
pub enum Stmt {
    Const(u8),
    ExprConst(P<ast::Expr>),

    Var(P<ast::Expr>, Size),
    Extend(P<ast::Expr>),

    DynScale(P<ast::Expr>, P<ast::Expr>),

    Align(P<ast::Expr>),

    GlobalLabel(Ident),
    LocalLabel(Ident),
    DynamicLabel(P<ast::Expr>),

    GlobalJumpTarget(Ident, Size),
    ForwardJumpTarget(Ident, Size),
    BackwardJumpTarget(Ident, Size),
    DynamicJumpTarget(P<ast::Expr>, Size),

    Stmt(ast::Stmt),
}

pub fn serialize(ecx: &mut ExtCtxt, name: P<ast::Expr>, stmts: Vec<Stmt>) -> Vec<ast::Stmt> {
    let mut buffer = Vec::new();

    // construction for `op.push(expr)` is as follows
    // op = Path::from_ident(name)
    // push = Ident::from_str("push")
    // expr = expr_lit(span, LitKind::Byte)
    // expr_method_call(span, op, Vec![expr])

    let mut stmts = stmts.into_iter().peekable();

    while let Some(stmt) = stmts.next() {
        use self::Stmt::*;

        let (method, args) = match stmt {
            Const(byte)            => {
                let mut bytes = vec![byte];
                while let Some(&Const(byte)) = stmts.peek() {
                    bytes.push(byte);
                    stmts.next();
                    if bytes.len() == 32 {
                        break;
                    }
                }

                if bytes.len() == 1 {
                    ("push",   vec![ecx.expr_lit(ecx.call_site(), ast::LitKind::Byte(bytes[0]))])
                } else {
                    ("extend", vec![ecx.expr_lit(ecx.call_site(), ast::LitKind::ByteStr(Rc::new(bytes)))])
                }
            },
            ExprConst(expr)        => ("push",    vec![expr]),

            Var(expr, Size::BYTE)  => ("push_i8",  vec![expr]),
            Var(expr, Size::WORD)  => ("push_i16", vec![expr]),
            Var(expr, Size::DWORD) => ("push_i32", vec![expr]),
            Var(expr, Size::QWORD) => ("push_i64", vec![expr]),
            Var(_, _)           => panic!("immediate serializaiton of this size is not supported yet"),

            Extend(expr)           => ("extend", vec![expr]),

            DynScale(scale, rest)  => {
                let temp = ast::Ident::from_str("temp");
                buffer.push(ecx.stmt_let(ecx.call_site(), false, temp, encoded_size_expr(ecx, &name, scale)));
                ("push", vec![or_mask_shift_expr(ecx, rest, ecx.expr_ident(ecx.call_site(), temp), 3, 6)])
            },

            Align(expr)            => ("align",   vec![expr]),

            GlobalLabel(ident)     => ("global_label", vec![ecx.expr_lit(
                ident.span,
                ast::LitKind::Str(ident.node.name, ast::StrStyle::Cooked)
            )]),
            LocalLabel(ident)      => ("local_label", vec![ecx.expr_lit(
                ident.span,
                ast::LitKind::Str(ident.node.name, ast::StrStyle::Cooked)
            )]),
            DynamicLabel(expr)     => ("dynamic_label", vec![expr]),

            GlobalJumpTarget(ident, size) => ("global_reloc", vec![ecx.expr_lit(
                ident.span,
                ast::LitKind::Str(ident.node.name, ast::StrStyle::Cooked)
            ), ecx.expr_u8(ident.span, size.in_bytes())]),
            ForwardJumpTarget(ident, size) => ("forward_reloc", vec![ecx.expr_lit(
                ident.span,
                ast::LitKind::Str(ident.node.name, ast::StrStyle::Cooked)
            ), ecx.expr_u8(ident.span, size.in_bytes())]),
            BackwardJumpTarget(ident, size) => ("backward_reloc", vec![ecx.expr_lit(
                ident.span,
                ast::LitKind::Str(ident.node.name, ast::StrStyle::Cooked)
            ), ecx.expr_u8(ident.span, size.in_bytes())]),
            DynamicJumpTarget(expr, size) => {
                let span = expr.span;
                ("dynamic_reloc", vec![expr, ecx.expr_u8(span, size.in_bytes())])
            },
            Stmt(stmt) => {
                buffer.push(stmt);
                continue;
            }
        };

        let op = name.clone();
        let method = ast::Ident::from_str(method);
        let expr = ecx.expr_method_call(ecx.call_site(), op, method, args);

        buffer.push(ecx.stmt_semi(expr));
    }

    buffer
}

pub fn add_exprs<T: Iterator<Item=P<ast::Expr>>>(ecx: &ExtCtxt, span: Span, mut exprs: T) -> Option<P<ast::Expr>> {
    exprs.next().map(|mut accum| {
        for next in exprs {
            accum = ecx.expr_binary(span, ast::BinOpKind::Add, accum, next);
        }
        accum
    })
}

pub fn size_of_scale_expr(ecx: &ExtCtxt, ty: ast::Path, value: P<ast::Expr>, size: Size) -> P<ast::Expr> {
    let span = value.span;
    ecx.expr_binary(span,
        ast::BinOpKind::Mul,
        ecx.expr_cast(span,
            size_of_expr(ecx, ty),
            ecx.ty_ident(span, size.as_literal())
        ),
        value
    )
}

pub fn or_mask_shift_expr(ecx: &ExtCtxt, orig: P<ast::Expr>, mut expr: P<ast::Expr>, mask: u64, shift: i8) -> P<ast::Expr> {
    let span = expr.span;
    // take expr and return !((expr & mask) << shift)

    expr = ecx.expr_binary(span, ast::BinOpKind::BitAnd, expr, ecx.expr_lit(
        span, ast::LitKind::Int(mask as u128, ast::LitIntType::Unsuffixed)
    ));

    if shift < 0 {
        expr = ecx.expr_binary(span, ast::BinOpKind::Shr, expr, ecx.expr_lit(
            span, ast::LitKind::Int((-shift) as u128, ast::LitIntType::Unsuffixed)
        ));
    } else if shift > 0 {
        expr = ecx.expr_binary(span, ast::BinOpKind::Shl, expr, ecx.expr_lit(
            span, ast::LitKind::Int(shift as u128, ast::LitIntType::Unsuffixed)
        ));
    }

    ecx.expr_binary(span, ast::BinOpKind::BitOr, orig, expr)
}

pub fn offset_of_expr(ecx: &ExtCtxt, path: ast::Path, attr: ast::Ident, size: Size) -> P<ast::Expr> {
    // generate a P<Expr> that resolves into the offset of an attribute to a type.
    // this is somewhat ridiculously complex because we can't expand macros here

    let span = path.span;

    let structpat = ecx.pat_struct(span, path.clone(), vec![
        Spanned {span: span, node: ast::FieldPat {
            ident: attr,
            pat: ecx.pat_wild(span),
            is_shorthand: false
        }},
    ]).map(|mut pat| {
        if let ast::PatKind::Struct(_, _, ref mut dotdot) = pat.node {
            *dotdot = true;
        }
        pat
    });

    // there's no default constructor function for let pattern;
    let validation_stmt = ast::Stmt {
        id: ast::DUMMY_NODE_ID,
        span: span,
        node: ast::StmtKind::Local(P(ast::Local {
            pat: structpat,
            ty: None,
            init: None,
            id: ast::DUMMY_NODE_ID,
            span: span,
            attrs: ast::ThinVec::new()
        }))
    };

    let temp     = ast::Ident::from_str("temp");
    let rv       = ast::Ident::from_str("rv");
    let usize_id = ast::Ident::from_str("usize");
    let uninitialized = ["std", "mem", "uninitialized"].iter().cloned().map(ast::Ident::from_str).collect();
    let forget        = ["std", "mem", "forget"       ].iter().cloned().map(ast::Ident::from_str).collect();

    // unsafe {
    let block = ecx.block(span, vec![
        // let path { attr: _, ..};
        validation_stmt,
        // let temp: path = ::std::mem::uninitialized();
        ecx.stmt_let_typed(span, false, temp, ecx.ty_path(path),
            ecx.expr_call_global(span, uninitialized, Vec::new())
        ),
        // let rv = &temp.attr as *const _ as usize - &temp as *const _ as usize;
        ecx.stmt_let(span,
            false,
            rv,
            ecx.expr_binary(span, ast::BinOpKind::Sub,
                ecx.expr_cast(span,
                    ecx.expr_cast(span,
                        ecx.expr_addr_of(span,
                            ecx.expr_field_access(span,
                                ecx.expr_ident(span, temp),
                                attr
                            )
                        ), ecx.ty_ptr(span, ecx.ty_infer(span), ast::Mutability::Immutable)
                    ), ecx.ty_ident(span, usize_id)
                ),
                ecx.expr_cast(span,
                    ecx.expr_cast(span,
                        ecx.expr_addr_of(span, ecx.expr_ident(span, temp)),
                        ecx.ty_ptr(span, ecx.ty_infer(span), ast::Mutability::Immutable)
                    ), ecx.ty_ident(span, usize_id)
                )
            )
        ),
        // ::std::mem::forget(temp);
        ecx.stmt_semi(ecx.expr_call_global(span, forget, vec![ecx.expr_ident(span, temp)])),
        // rv as i32
        ecx.stmt_expr(ecx.expr_cast(span, ecx.expr_ident(span, rv), ecx.ty_ident(span, size.as_literal())))
    ]).map(|mut b| {
        b.rules = ast::BlockCheckMode::Unsafe(ast::UnsafeSource::CompilerGenerated);
        b
    });

    ecx.expr_block(block)
}

pub fn size_of_expr(ecx: &ExtCtxt, path: ast::Path) -> P<ast::Expr> {
    // generate a P<Expr> that returns the size of type at path
    let span = path.span;

    let ty = ecx.ty_path(path);
    let idents = ["std", "mem", "size_of"].iter().cloned().map(ast::Ident::from_str).collect();
    let size_of = ecx.path_all(span, true, idents, Vec::new(), vec![ty], Vec::new());
    ecx.expr_call(span, ecx.expr_path(size_of), Vec::new())
}

pub fn encoded_size_expr(ecx: &ExtCtxt, name: &P<ast::Expr>, size: P<ast::Expr>) -> P<ast::Expr> {
    let span = size.span;

    ecx.expr_match(span, size, vec![
        ecx.arm(span, vec![ecx.pat_lit(span, ecx.expr_usize(span, 8))], ecx.expr_u8(span, 3)),
        ecx.arm(span, vec![ecx.pat_lit(span, ecx.expr_usize(span, 4))], ecx.expr_u8(span, 2)),
        ecx.arm(span, vec![ecx.pat_lit(span, ecx.expr_usize(span, 2))], ecx.expr_u8(span, 1)),
        ecx.arm(span, vec![ecx.pat_lit(span, ecx.expr_usize(span, 1))], ecx.expr_u8(span, 0)),
        ecx.arm(span, vec![ecx.pat_wild(span)], ecx.expr_method_call(span,
            name.clone(),
            ast::Ident::from_str("runtime_error"),
            vec![ecx.expr_str(span,
                Symbol::intern("Type size not representable as scale")
            )]
        ))
    ])
}
<|endoftext|>


<commit_before>use std::sync::Arc;

use buffer::AbstractBuffer;
use descriptor_set::AbstractDescriptorSet;
use descriptor_set::AbstractDescriptorSetLayout;
use image::AbstractImageView;
use image::Layout as ImageLayout;
use sampler::Sampler;

use vk;

/// Types that describe the layout of a pipeline (descriptor sets and push constants).
pub unsafe trait PipelineLayoutDesc {
    /// Represents a collection of `DescriptorSet` structs. A parameter of this type must be
    /// passed when you add a draw command to a command buffer that uses this layout.
    type DescriptorSets;

    /// Represents a collection of `DescriptorSetLayout` structs. A parameter of this type must
    /// be passed when creating a `PipelineLayout` struct.
    type DescriptorSetLayouts;

    /// Not yet implemented. Useless for now.
    type PushConstants;

    /// Turns the `DescriptorSets` associated type into something vulkano can understand.
    fn decode_descriptor_sets(&self, Self::DescriptorSets) -> Vec<Arc<AbstractDescriptorSet>>;  // TODO: vec is slow

    /// Turns the `DescriptorSetLayouts` associated type into something vulkano can understand.
    fn decode_descriptor_set_layouts(&self, Self::DescriptorSetLayouts)
                                     -> Vec<Arc<AbstractDescriptorSetLayout>>;  // TODO: vec is slow

    // FIXME: implement this correctly
    fn is_compatible_with<P>(&self, _: &P) -> bool where P: PipelineLayoutDesc { true }
}

/// Types that describe a single descriptor set.
pub unsafe trait SetLayout {
    /// Returns the list of descriptors contained in this set.
    fn descriptors(&self) -> Vec<DescriptorDesc>;       // TODO: better perfs

    // FIXME: implement this correctly
    fn is_compatible_with<S>(&self, _: &S) -> bool where S: SetLayout { true }
}

/// Extension for the `SetLayout` trait.
pub unsafe trait SetLayoutWrite<Data>: SetLayout {
    /// Turns the data into something vulkano can understand.
    fn decode(&self, Data) -> Vec<DescriptorWrite>;        // TODO: better perfs
}

/// Extension for the `SetLayout` trait.
pub unsafe trait SetLayoutInit<Data>: SetLayout {
    /// Turns the data into something vulkano can understand.
    fn decode(&self, Data) -> Vec<DescriptorWrite>;        // TODO: better perfs
}

// FIXME: shoud allow multiple array binds at once
pub struct DescriptorWrite {
    pub binding: u32,
    pub array_element: u32,
    pub content: DescriptorBind,
}

// FIXME: incomplete
#[derive(Clone)]        // TODO: Debug
pub enum DescriptorBind {
    StorageImage(Arc<AbstractImageView>, ImageLayout),
    Sampler(Arc<Sampler>),
    SampledImage(Arc<AbstractImageView>, ImageLayout),
    CombinedImageSampler(Arc<Sampler>, Arc<AbstractImageView>, ImageLayout),
    //UniformTexelBuffer(Arc<AbstractBuffer>),      // FIXME: requires buffer views
    //StorageTexelBuffer(Arc<AbstractBuffer>),      // FIXME: requires buffer views
    UniformBuffer { buffer: Arc<AbstractBuffer>, offset: usize, size: usize },
    StorageBuffer { buffer: Arc<AbstractBuffer>, offset: usize, size: usize },
    DynamicUniformBuffer { buffer: Arc<AbstractBuffer>, offset: usize, size: usize },
    DynamicStorageBuffer { buffer: Arc<AbstractBuffer>, offset: usize, size: usize },
    InputAttachment(Arc<AbstractImageView>, ImageLayout),
}

impl DescriptorBind {
    /// Returns the type corresponding to this bind.
    #[inline]
    pub fn ty(&self) -> DescriptorType {
        match *self {
            DescriptorBind::Sampler(_) => DescriptorType::Sampler,
            DescriptorBind::CombinedImageSampler(_, _, _) => DescriptorType::CombinedImageSampler,
            DescriptorBind::SampledImage(_, _) => DescriptorType::SampledImage,
            DescriptorBind::StorageImage(_, _) => DescriptorType::StorageImage,
            //DescriptorBind::UniformTexelBuffer(_) => DescriptorType::UniformTexelBuffer,
            //DescriptorBind::StorageTexelBuffer(_) => DescriptorType::StorageTexelBuffer,
            DescriptorBind::UniformBuffer { .. } => DescriptorType::UniformBuffer,
            DescriptorBind::StorageBuffer { .. } => DescriptorType::StorageBuffer,
            DescriptorBind::DynamicUniformBuffer { .. } => DescriptorType::UniformBufferDynamic,
            DescriptorBind::DynamicStorageBuffer { .. } => DescriptorType::StorageBufferDynamic,
            DescriptorBind::InputAttachment(_, _) => DescriptorType::InputAttachment,
        }
    }
}

/// Describes a single descriptor.
#[derive(Debug, Copy, Clone)]
pub struct DescriptorDesc {
    /// Offset of the binding within the descriptor.
    pub binding: u32,

    /// What kind of resource can later be bind to this descriptor.
    pub ty: DescriptorType,

    /// How many array elements this descriptor is made of.
    pub array_count: u32,

    /// Which shader stages are going to access this descriptor.
    pub stages: ShaderStages,
}

/// Describes what kind of resource may later be bound to a descriptor.
// FIXME: add immutable sampler when relevant
#[derive(Debug, Copy, Clone)]
#[repr(u32)]
pub enum DescriptorType {
    Sampler = vk::DESCRIPTOR_TYPE_SAMPLER,
    CombinedImageSampler = vk::DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
    SampledImage = vk::DESCRIPTOR_TYPE_SAMPLED_IMAGE,
    StorageImage = vk::DESCRIPTOR_TYPE_STORAGE_IMAGE,
    UniformTexelBuffer = vk::DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER,
    StorageTexelBuffer = vk::DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER,
    UniformBuffer = vk::DESCRIPTOR_TYPE_UNIFORM_BUFFER,
    StorageBuffer = vk::DESCRIPTOR_TYPE_STORAGE_BUFFER,
    UniformBufferDynamic = vk::DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC,
    StorageBufferDynamic = vk::DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC,
    InputAttachment = vk::DESCRIPTOR_TYPE_INPUT_ATTACHMENT,
}

impl DescriptorType {
    /// Turns the `DescriptorType` into the corresponding Vulkan constant.
    // this function exists because when immutable samplers are added, it will no longer be possible to do `as u32`
    // TODO: hacky
    #[inline]
    #[doc(hidden)]
    pub fn vk_enum(&self) -> u32 {
        *self as u32
    }
}

/// Describes which shader stages have access to a descriptor.
#[derive(Debug, Copy, Clone)]
pub struct ShaderStages {
    /// `True` means that the descriptor will be used by the vertex shader.
    pub vertex: bool,
    /// `True` means that the descriptor will be used by the tessellation control shader.
    pub tessellation_control: bool,
    /// `True` means that the descriptor will be used by the tessellation evaluation shader.
    pub tessellation_evaluation: bool,
    /// `True` means that the descriptor will be used by the geometry shader.
    pub geometry: bool,
    /// `True` means that the descriptor will be used by the fragment shader.
    pub fragment: bool,
    /// `True` means that the descriptor will be used by the compute shader.
    pub compute: bool,
}

impl ShaderStages {
    /// Creates a `ShaderStages` struct will all graphics stages set to `true`.
    #[inline]
    pub fn all_graphics() -> ShaderStages {
        ShaderStages {
            vertex: true,
            tessellation_control: true,
            tessellation_evaluation: true,
            geometry: true,
            fragment: true,
            compute: false,
        }
    }

    /// Creates a `ShaderStages` struct will the compute stage set to `true`.
    #[inline]
    pub fn compute() -> ShaderStages {
        ShaderStages {
            vertex: false,
            tessellation_control: false,
            tessellation_evaluation: false,
            geometry: false,
            fragment: false,
            compute: true,
        }
    }
}

#[doc(hidden)]
impl Into<vk::ShaderStageFlags> for ShaderStages {
    #[inline]
    fn into(self) -> vk::ShaderStageFlags {
        let mut result = 0;
        if self.vertex { result |= vk::SHADER_STAGE_VERTEX_BIT; }
        if self.tessellation_control { result |= vk::SHADER_STAGE_TESSELLATION_CONTROL_BIT; }
        if self.tessellation_evaluation { result |= vk::SHADER_STAGE_TESSELLATION_EVALUATION_BIT; }
        if self.geometry { result |= vk::SHADER_STAGE_GEOMETRY_BIT; }
        if self.fragment { result |= vk::SHADER_STAGE_FRAGMENT_BIT; }
        if self.compute { result |= vk::SHADER_STAGE_COMPUTE_BIT; }
        result
    }
}
<commit_msg>Minor draft for SetLayoutPossibleSuperset trait<commit_after>use std::sync::Arc;

use buffer::AbstractBuffer;
use descriptor_set::AbstractDescriptorSet;
use descriptor_set::AbstractDescriptorSetLayout;
use image::AbstractImageView;
use image::Layout as ImageLayout;
use sampler::Sampler;

use vk;

/// Types that describe the layout of a pipeline (descriptor sets and push constants).
pub unsafe trait PipelineLayoutDesc {
    /// Represents a collection of `DescriptorSet` structs. A parameter of this type must be
    /// passed when you add a draw command to a command buffer that uses this layout.
    type DescriptorSets;

    /// Represents a collection of `DescriptorSetLayout` structs. A parameter of this type must
    /// be passed when creating a `PipelineLayout` struct.
    type DescriptorSetLayouts;

    /// Not yet implemented. Useless for now.
    type PushConstants;

    /// Turns the `DescriptorSets` associated type into something vulkano can understand.
    fn decode_descriptor_sets(&self, Self::DescriptorSets) -> Vec<Arc<AbstractDescriptorSet>>;  // TODO: vec is slow

    /// Turns the `DescriptorSetLayouts` associated type into something vulkano can understand.
    fn decode_descriptor_set_layouts(&self, Self::DescriptorSetLayouts)
                                     -> Vec<Arc<AbstractDescriptorSetLayout>>;  // TODO: vec is slow

    // FIXME: implement this correctly
    fn is_compatible_with<P>(&self, _: &P) -> bool where P: PipelineLayoutDesc { true }
}

/// Types that describe a single descriptor set.
pub unsafe trait SetLayout {
    /// Returns the list of descriptors contained in this set.
    fn descriptors(&self) -> Vec<DescriptorDesc>;       // TODO: better perfs
}

/// Extension for the `SetLayout` trait.
pub unsafe trait SetLayoutWrite<Data>: SetLayout {
    /// Turns the data into something vulkano can understand.
    fn decode(&self, Data) -> Vec<DescriptorWrite>;        // TODO: better perfs
}

/// Extension for the `SetLayout` trait.
pub unsafe trait SetLayoutInit<Data>: SetLayout {
    /// Turns the data into something vulkano can understand.
    fn decode(&self, Data) -> Vec<DescriptorWrite>;        // TODO: better perfs
}

/// Extension for `SetLayout`.
pub unsafe trait SetLayoutPossibleSuperset<Other>: SetLayout where Other: SetLayout {
    /// Returns true if `self` is a superset of `Other`.
    fn is_superset_of(&self, &Other) -> bool;
}

// FIXME: shoud allow multiple array binds at once
pub struct DescriptorWrite {
    pub binding: u32,
    pub array_element: u32,
    pub content: DescriptorBind,
}

// FIXME: incomplete
#[derive(Clone)]        // TODO: Debug
pub enum DescriptorBind {
    StorageImage(Arc<AbstractImageView>, ImageLayout),
    Sampler(Arc<Sampler>),
    SampledImage(Arc<AbstractImageView>, ImageLayout),
    CombinedImageSampler(Arc<Sampler>, Arc<AbstractImageView>, ImageLayout),
    //UniformTexelBuffer(Arc<AbstractBuffer>),      // FIXME: requires buffer views
    //StorageTexelBuffer(Arc<AbstractBuffer>),      // FIXME: requires buffer views
    UniformBuffer { buffer: Arc<AbstractBuffer>, offset: usize, size: usize },
    StorageBuffer { buffer: Arc<AbstractBuffer>, offset: usize, size: usize },
    DynamicUniformBuffer { buffer: Arc<AbstractBuffer>, offset: usize, size: usize },
    DynamicStorageBuffer { buffer: Arc<AbstractBuffer>, offset: usize, size: usize },
    InputAttachment(Arc<AbstractImageView>, ImageLayout),
}

impl DescriptorBind {
    /// Returns the type corresponding to this bind.
    #[inline]
    pub fn ty(&self) -> DescriptorType {
        match *self {
            DescriptorBind::Sampler(_) => DescriptorType::Sampler,
            DescriptorBind::CombinedImageSampler(_, _, _) => DescriptorType::CombinedImageSampler,
            DescriptorBind::SampledImage(_, _) => DescriptorType::SampledImage,
            DescriptorBind::StorageImage(_, _) => DescriptorType::StorageImage,
            //DescriptorBind::UniformTexelBuffer(_) => DescriptorType::UniformTexelBuffer,
            //DescriptorBind::StorageTexelBuffer(_) => DescriptorType::StorageTexelBuffer,
            DescriptorBind::UniformBuffer { .. } => DescriptorType::UniformBuffer,
            DescriptorBind::StorageBuffer { .. } => DescriptorType::StorageBuffer,
            DescriptorBind::DynamicUniformBuffer { .. } => DescriptorType::UniformBufferDynamic,
            DescriptorBind::DynamicStorageBuffer { .. } => DescriptorType::StorageBufferDynamic,
            DescriptorBind::InputAttachment(_, _) => DescriptorType::InputAttachment,
        }
    }
}

/// Describes a single descriptor.
#[derive(Debug, Copy, Clone)]
pub struct DescriptorDesc {
    /// Offset of the binding within the descriptor.
    pub binding: u32,

    /// What kind of resource can later be bind to this descriptor.
    pub ty: DescriptorType,

    /// How many array elements this descriptor is made of.
    pub array_count: u32,

    /// Which shader stages are going to access this descriptor.
    pub stages: ShaderStages,
}

/// Describes what kind of resource may later be bound to a descriptor.
// FIXME: add immutable sampler when relevant
#[derive(Debug, Copy, Clone)]
#[repr(u32)]
pub enum DescriptorType {
    Sampler = vk::DESCRIPTOR_TYPE_SAMPLER,
    CombinedImageSampler = vk::DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
    SampledImage = vk::DESCRIPTOR_TYPE_SAMPLED_IMAGE,
    StorageImage = vk::DESCRIPTOR_TYPE_STORAGE_IMAGE,
    UniformTexelBuffer = vk::DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER,
    StorageTexelBuffer = vk::DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER,
    UniformBuffer = vk::DESCRIPTOR_TYPE_UNIFORM_BUFFER,
    StorageBuffer = vk::DESCRIPTOR_TYPE_STORAGE_BUFFER,
    UniformBufferDynamic = vk::DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC,
    StorageBufferDynamic = vk::DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC,
    InputAttachment = vk::DESCRIPTOR_TYPE_INPUT_ATTACHMENT,
}

impl DescriptorType {
    /// Turns the `DescriptorType` into the corresponding Vulkan constant.
    // this function exists because when immutable samplers are added, it will no longer be possible to do `as u32`
    // TODO: hacky
    #[inline]
    #[doc(hidden)]
    pub fn vk_enum(&self) -> u32 {
        *self as u32
    }
}

/// Describes which shader stages have access to a descriptor.
#[derive(Debug, Copy, Clone)]
pub struct ShaderStages {
    /// `True` means that the descriptor will be used by the vertex shader.
    pub vertex: bool,
    /// `True` means that the descriptor will be used by the tessellation control shader.
    pub tessellation_control: bool,
    /// `True` means that the descriptor will be used by the tessellation evaluation shader.
    pub tessellation_evaluation: bool,
    /// `True` means that the descriptor will be used by the geometry shader.
    pub geometry: bool,
    /// `True` means that the descriptor will be used by the fragment shader.
    pub fragment: bool,
    /// `True` means that the descriptor will be used by the compute shader.
    pub compute: bool,
}

impl ShaderStages {
    /// Creates a `ShaderStages` struct will all graphics stages set to `true`.
    #[inline]
    pub fn all_graphics() -> ShaderStages {
        ShaderStages {
            vertex: true,
            tessellation_control: true,
            tessellation_evaluation: true,
            geometry: true,
            fragment: true,
            compute: false,
        }
    }

    /// Creates a `ShaderStages` struct will the compute stage set to `true`.
    #[inline]
    pub fn compute() -> ShaderStages {
        ShaderStages {
            vertex: false,
            tessellation_control: false,
            tessellation_evaluation: false,
            geometry: false,
            fragment: false,
            compute: true,
        }
    }
}

#[doc(hidden)]
impl Into<vk::ShaderStageFlags> for ShaderStages {
    #[inline]
    fn into(self) -> vk::ShaderStageFlags {
        let mut result = 0;
        if self.vertex { result |= vk::SHADER_STAGE_VERTEX_BIT; }
        if self.tessellation_control { result |= vk::SHADER_STAGE_TESSELLATION_CONTROL_BIT; }
        if self.tessellation_evaluation { result |= vk::SHADER_STAGE_TESSELLATION_EVALUATION_BIT; }
        if self.geometry { result |= vk::SHADER_STAGE_GEOMETRY_BIT; }
        if self.fragment { result |= vk::SHADER_STAGE_FRAGMENT_BIT; }
        if self.compute { result |= vk::SHADER_STAGE_COMPUTE_BIT; }
        result
    }
}
<|endoftext|>


<commit_before>extern crate mdbook;
#[macro_use]
extern crate clap;
extern crate log;
extern crate env_logger;
extern crate open;

// Dependencies for the Watch feature
#[cfg(feature = "watch")]
extern crate notify;
#[cfg(feature = "watch")]
extern crate time;
#[cfg(feature = "watch")]
extern crate crossbeam;

// Dependencies for the Serve feature
#[cfg(feature = "serve")]
extern crate iron;
#[cfg(feature = "serve")]
extern crate staticfile;
#[cfg(feature = "serve")]
extern crate ws;

#[cfg(feature = "serve")]
use iron::{Iron, AfterMiddleware, IronResult, IronError, Request, Response, status, Set, Chain};

use std::env;
use std::error::Error;
use std::ffi::OsStr;
use std::io::{self, Write};
use std::path::{Path, PathBuf};

use clap::{App, ArgMatches, SubCommand, AppSettings};

// Uses for the Watch feature
#[cfg(feature = "watch")]
use notify::Watcher;
#[cfg(feature = "watch")]
use std::time::Duration;
#[cfg(feature = "watch")]
use std::sync::mpsc::channel;

use mdbook::MDBook;

const NAME: &'static str = "mdbook";

fn main() {
    env_logger::init().unwrap();

    // Create a list of valid arguments and sub-commands
    let matches = App::new(NAME)
                    .about("Create a book in form of a static website from markdown files")
                    .author("Mathieu David <mathieudavid@mathieudavid.org>")
                    // Get the version from our Cargo.toml using clap's crate_version!() macro
                    .version(&*format!("v{}", crate_version!()))
                    .setting(AppSettings::SubcommandRequired)
                    .after_help("For more information about a specific command, try `mdbook <command> --help`\nSource code for mdbook available at: https://github.com/azerupi/mdBook")
                    .subcommand(SubCommand::with_name("init")
                        .about("Create boilerplate structure and files in the directory")
                        // the {n} denotes a newline which will properly aligned in all help messages
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'")
                        .arg_from_usage("--theme 'Copies the default theme into your source folder'")
                        .arg_from_usage("--force 'skip confirmation prompts'"))
                    .subcommand(SubCommand::with_name("build")
                        .about("Build the book from the markdown files")
                        .arg_from_usage("-o, --open 'Open the compiled book in a web browser'")
                        .arg_from_usage("-d, --dest-dir=[dest-dir] 'The output directory for your book{n}(Defaults to ./book when omitted)'")
                        .arg_from_usage("--no-create 'Will not create non-existent files linked from SUMMARY.md'")
                        .arg_from_usage("--curly-quotes 'Convert straight quotes to curly quotes, except for those that occur in code blocks and code spans'")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'"))
                    .subcommand(SubCommand::with_name("watch")
                        .about("Watch the files for changes")
                        .arg_from_usage("-o, --open 'Open the compiled book in a web browser'")
                        .arg_from_usage("-d, --dest-dir=[dest-dir] 'The output directory for your book{n}(Defaults to ./book when omitted)'")
                        .arg_from_usage("--curly-quotes 'Convert straight quotes to curly quotes, except for those that occur in code blocks and code spans'")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'"))
                    .subcommand(SubCommand::with_name("serve")
                        .about("Serve the book at http://localhost:3000. Rebuild and reload on change.")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'")
                        .arg_from_usage("-d, --dest-dir=[dest-dir] 'The output directory for your book{n}(Defaults to ./book when omitted)'")
                        .arg_from_usage("--curly-quotes 'Convert straight quotes to curly quotes, except for those that occur in code blocks and code spans'")
                        .arg_from_usage("-p, --port=[port] 'Use another port{n}(Defaults to 3000)'")
                        .arg_from_usage("-w, --websocket-port=[ws-port] 'Use another port for the websocket connection (livereload){n}(Defaults to 3001)'")
                        .arg_from_usage("-i, --interface=[interface] 'Interface to listen on{n}(Defaults to localhost)'")
                        .arg_from_usage("-a, --address=[address] 'Address that the browser can reach the websocket server from{n}(Defaults to the interface address)'")
                        .arg_from_usage("-o, --open 'Open the book server in a web browser'"))
                    .subcommand(SubCommand::with_name("test")
                        .about("Test that code samples compile"))
                    .get_matches();

    // Check which subcomamnd the user ran...
    let res = match matches.subcommand() {
        ("init", Some(sub_matches)) => init(sub_matches),
        ("build", Some(sub_matches)) => build(sub_matches),
        #[cfg(feature = "watch")]
        ("watch", Some(sub_matches)) => watch(sub_matches),
        #[cfg(feature = "serve")]
        ("serve", Some(sub_matches)) => serve(sub_matches),
        ("test", Some(sub_matches)) => test(sub_matches),
        (_, _) => unreachable!(),
    };

    if let Err(e) = res {
        writeln!(&mut io::stderr(), "An error occured:\n{}", e).ok();
        ::std::process::exit(101);
    }
}

// Simple function that user comfirmation
fn confirm() -> bool {
    io::stdout().flush().unwrap();
    let mut s = String::new();
    io::stdin().read_line(&mut s).ok();
    match &*s.trim() {
        "Y" | "y" | "yes" | "Yes" => true,
        _ => false,
    }
}

// Init command implementation
fn init(args: &ArgMatches) -> Result<(), Box<Error>> {

    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir);

    // Call the function that does the initialization
    book.init()?;

    // If flag `--theme` is present, copy theme to src
    if args.is_present("theme") {

        // Skip this if `--force` is present
        if !args.is_present("force") {
            // Print warning
            print!("\nCopying the default theme to {:?}", book.get_source());
            println!("could potentially overwrite files already present in that directory.");
            print!("\nAre you sure you want to continue? (y/n) ");

            // Read answer from user and exit if it's not 'yes'
            if !confirm() {
                println!("\nSkipping...\n");
                println!("All done, no errors...");
                ::std::process::exit(0);
            }
        }

        // Call the function that copies the theme
        book.copy_theme()?;
        println!("\nTheme copied.");

    }

    // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`
    let is_dest_inside_root = book.get_destination()
                                  .map(|p| p.starts_with(book.get_root()))
                                  .unwrap_or(false);

    if !args.is_present("force") && is_dest_inside_root {
        println!("\nDo you want a .gitignore to be created? (y/n)");

        if confirm() {
            book.create_gitignore();
            println!("\n.gitignore created.");
        }
    }

    println!("\nAll done, no errors...");

    Ok(())
}

// Build command implementation
fn build(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let book = MDBook::new(&book_dir).read_config()?;

    let mut book = match args.value_of("dest-dir") {
        Some(dest_dir) => book.with_destination(dest_dir),
        None => book,
    };

    if args.is_present("no-create") {
        book.create_missing = false;
    }

    if args.is_present("curly-quotes") {
        book = book.with_curly_quotes(true);
    }

    book.build()?;

    if let Some(d) = book.get_destination() {
        if args.is_present("open") {
            open(d.join("index.html"));
        }
    }

    Ok(())
}

// Watch command implementation
#[cfg(feature = "watch")]
fn watch(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let book = MDBook::new(&book_dir).read_config()?;

    let mut book = match args.value_of("dest-dir") {
        Some(dest_dir) => book.with_destination(dest_dir),
        None => book,
    };

    if args.is_present("curly-quotes") {
        book = book.with_curly_quotes(true);
    }

    if args.is_present("open") {
        book.build()?;
        if let Some(d) = book.get_destination() {
            open(d.join("index.html"));
        }
    }

    trigger_on_change(&mut book, |path, book| {
        println!("File changed: {:?}\nBuilding book...\n", path);
        if let Err(e) = book.build() {
            println!("Error while building: {:?}", e);
        }
        println!("");
    });

    Ok(())
}

#[cfg(feature = "serve")]
struct ErrorRecover;

#[cfg(feature = "serve")]
impl AfterMiddleware for ErrorRecover {
    fn catch(&self, _: &mut Request, err: IronError) -> IronResult<Response> {
        match err.response.status {
            Some(_) => Ok(err.response.set(status::NotFound)),
            _ => Err(err)
        }
    }
}

// Watch command implementation
#[cfg(feature = "serve")]
fn serve(args: &ArgMatches) -> Result<(), Box<Error>> {
    const RELOAD_COMMAND: &'static str = "reload";

    let book_dir = get_book_dir(args);
    let book = MDBook::new(&book_dir).read_config()?;

    let mut book = match args.value_of("dest-dir") {
        Some(dest_dir) => book.with_destination(dest_dir),
        None => book,
    };

    if let None = book.get_destination() {
        println!("The HTML renderer is not set up, impossible to serve the files.");
        std::process::exit(2);
    }

    if args.is_present("curly-quotes") {
        book = book.with_curly_quotes(true);
    }

    let port = args.value_of("port").unwrap_or("3000");
    let ws_port = args.value_of("websocket-port").unwrap_or("3001");
    let interface = args.value_of("interface").unwrap_or("localhost");
    let public_address = args.value_of("address").unwrap_or(interface);
    let open_browser = args.is_present("open");

    let address = format!("{}:{}", interface, port);
    let ws_address = format!("{}:{}", interface, ws_port);

    book.set_livereload(format!(r#"
        <script type="text/javascript">
            var socket = new WebSocket("ws://{}:{}");
            socket.onmessage = function (event) {{
                if (event.data === "{}") {{
                    socket.close();
                    location.reload(true); // force reload from server (not from cache)
                }}
            }};

            window.onbeforeunload = function() {{
                socket.close();
            }}
        </script>
    "#,
                                public_address,
                                ws_port,
                                RELOAD_COMMAND));

    book.build()?;

    let mut chain = Chain::new(staticfile::Static::new(book.get_destination().expect("destination is present, checked before")));
    chain.link_after(ErrorRecover);
    let _iron = Iron::new(chain).http(&*address).unwrap();

    let ws_server = ws::WebSocket::new(|_| |_| Ok(())).unwrap();

    let broadcaster = ws_server.broadcaster();

    std::thread::spawn(move || { ws_server.listen(&*ws_address).unwrap(); });

    let serving_url = format!("http://{}", address);
    println!("\nServing on: {}", serving_url);

    if open_browser {
        open(serving_url);
    }

    trigger_on_change(&mut book, move |path, book| {
        println!("File changed: {:?}\nBuilding book...\n", path);
        match book.build() {
            Err(e) => println!("Error while building: {:?}", e),
            _ => broadcaster.send(RELOAD_COMMAND).unwrap(),
        }
        println!("");
    });

    Ok(())
}

fn test(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config()?;

    book.test()?;

    Ok(())
}

fn get_book_dir(args: &ArgMatches) -> PathBuf {
    if let Some(dir) = args.value_of("dir") {
        // Check if path is relative from current dir, or absolute...
        let p = Path::new(dir);
        if p.is_relative() {
            env::current_dir().unwrap().join(dir)
        } else {
            p.to_path_buf()
        }
    } else {
        env::current_dir().unwrap()
    }
}

fn open<P: AsRef<OsStr>>(path: P) {
    if let Err(e) = open::that(path) {
        println!("Error opening web browser: {}", e);
    }
}

// Calls the closure when a book source file is changed. This is blocking!
#[cfg(feature = "watch")]
fn trigger_on_change<F>(book: &mut MDBook, closure: F) -> ()
    where F: Fn(&Path, &mut MDBook) -> ()
{
    use notify::RecursiveMode::*;
    use notify::DebouncedEvent::*;

    // Create a channel to receive the events.
    let (tx, rx) = channel();

    let mut watcher = match notify::watcher(tx, Duration::from_secs(1)) {
        Ok(w) => w,
        Err(e) => {
            println!("Error while trying to watch the files:\n\n\t{:?}", e);
            ::std::process::exit(0);
        },
    };

    // Add the source directory to the watcher
    if let Err(e) = watcher.watch(book.get_source(), Recursive) {
        println!("Error while watching {:?}:\n    {:?}", book.get_source(), e);
        ::std::process::exit(0);
    };

    // Add the theme directory to the watcher
    if let Some(t) = book.get_theme_path() {
        watcher.watch(t, Recursive).unwrap_or_default();
    }

    // Add the book.{json,toml} file to the watcher if it exists, because it's not
    // located in the source directory
    if watcher
           .watch(book.get_root().join("book.json"), NonRecursive)
           .is_err() {
        // do nothing if book.json is not found
    }
    if watcher
           .watch(book.get_root().join("book.toml"), NonRecursive)
           .is_err() {
        // do nothing if book.toml is not found
    }

    println!("\nListening for changes...\n");

    loop {
        match rx.recv() {
            Ok(event) => {
                match event {
                    NoticeWrite(path) |
                    NoticeRemove(path) |
                    Create(path) |
                    Write(path) |
                    Remove(path) |
                    Rename(_, path) => {
                        closure(&path, book);
                    },
                    _ => {},
                }
            },
            Err(e) => {
                println!("An error occured: {:?}", e);
            },
        }
    }
}
<commit_msg>extract serving code to a separate module<commit_after>extern crate mdbook;
#[macro_use]
extern crate clap;
extern crate log;
extern crate env_logger;
extern crate open;

// Dependencies for the Watch feature
#[cfg(feature = "watch")]
extern crate notify;
#[cfg(feature = "watch")]
extern crate time;
#[cfg(feature = "watch")]
extern crate crossbeam;

use std::env;
use std::error::Error;
use std::ffi::OsStr;
use std::io::{self, Write};
use std::path::{Path, PathBuf};

use clap::{App, ArgMatches, SubCommand, AppSettings};

// Uses for the Watch feature
#[cfg(feature = "watch")]
use notify::Watcher;
#[cfg(feature = "watch")]
use std::time::Duration;
#[cfg(feature = "watch")]
use std::sync::mpsc::channel;

use mdbook::MDBook;

const NAME: &'static str = "mdbook";

fn main() {
    env_logger::init().unwrap();

    // Create a list of valid arguments and sub-commands
    let matches = App::new(NAME)
                    .about("Create a book in form of a static website from markdown files")
                    .author("Mathieu David <mathieudavid@mathieudavid.org>")
                    // Get the version from our Cargo.toml using clap's crate_version!() macro
                    .version(&*format!("v{}", crate_version!()))
                    .setting(AppSettings::SubcommandRequired)
                    .after_help("For more information about a specific command, try `mdbook <command> --help`\nSource code for mdbook available at: https://github.com/azerupi/mdBook")
                    .subcommand(SubCommand::with_name("init")
                        .about("Create boilerplate structure and files in the directory")
                        // the {n} denotes a newline which will properly aligned in all help messages
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'")
                        .arg_from_usage("--theme 'Copies the default theme into your source folder'")
                        .arg_from_usage("--force 'skip confirmation prompts'"))
                    .subcommand(SubCommand::with_name("build")
                        .about("Build the book from the markdown files")
                        .arg_from_usage("-o, --open 'Open the compiled book in a web browser'")
                        .arg_from_usage("-d, --dest-dir=[dest-dir] 'The output directory for your book{n}(Defaults to ./book when omitted)'")
                        .arg_from_usage("--no-create 'Will not create non-existent files linked from SUMMARY.md'")
                        .arg_from_usage("--curly-quotes 'Convert straight quotes to curly quotes, except for those that occur in code blocks and code spans'")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'"))
                    .subcommand(SubCommand::with_name("watch")
                        .about("Watch the files for changes")
                        .arg_from_usage("-o, --open 'Open the compiled book in a web browser'")
                        .arg_from_usage("-d, --dest-dir=[dest-dir] 'The output directory for your book{n}(Defaults to ./book when omitted)'")
                        .arg_from_usage("--curly-quotes 'Convert straight quotes to curly quotes, except for those that occur in code blocks and code spans'")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'"))
                    .subcommand(SubCommand::with_name("serve")
                        .about("Serve the book at http://localhost:3000. Rebuild and reload on change.")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'")
                        .arg_from_usage("-d, --dest-dir=[dest-dir] 'The output directory for your book{n}(Defaults to ./book when omitted)'")
                        .arg_from_usage("--curly-quotes 'Convert straight quotes to curly quotes, except for those that occur in code blocks and code spans'")
                        .arg_from_usage("-p, --port=[port] 'Use another port{n}(Defaults to 3000)'")
                        .arg_from_usage("-w, --websocket-port=[ws-port] 'Use another port for the websocket connection (livereload){n}(Defaults to 3001)'")
                        .arg_from_usage("-i, --interface=[interface] 'Interface to listen on{n}(Defaults to localhost)'")
                        .arg_from_usage("-a, --address=[address] 'Address that the browser can reach the websocket server from{n}(Defaults to the interface address)'")
                        .arg_from_usage("-o, --open 'Open the book server in a web browser'"))
                    .subcommand(SubCommand::with_name("test")
                        .about("Test that code samples compile"))
                    .get_matches();

    // Check which subcomamnd the user ran...
    let res = match matches.subcommand() {
        ("init", Some(sub_matches)) => init(sub_matches),
        ("build", Some(sub_matches)) => build(sub_matches),
        #[cfg(feature = "watch")]
        ("watch", Some(sub_matches)) => watch(sub_matches),
        #[cfg(feature = "serve")]
        ("serve", Some(sub_matches)) => serve::serve(sub_matches),
        ("test", Some(sub_matches)) => test(sub_matches),
        (_, _) => unreachable!(),
    };

    if let Err(e) = res {
        writeln!(&mut io::stderr(), "An error occured:\n{}", e).ok();
        ::std::process::exit(101);
    }
}

// Simple function that user comfirmation
fn confirm() -> bool {
    io::stdout().flush().unwrap();
    let mut s = String::new();
    io::stdin().read_line(&mut s).ok();
    match &*s.trim() {
        "Y" | "y" | "yes" | "Yes" => true,
        _ => false,
    }
}

// Init command implementation
fn init(args: &ArgMatches) -> Result<(), Box<Error>> {

    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir);

    // Call the function that does the initialization
    book.init()?;

    // If flag `--theme` is present, copy theme to src
    if args.is_present("theme") {

        // Skip this if `--force` is present
        if !args.is_present("force") {
            // Print warning
            print!("\nCopying the default theme to {:?}", book.get_source());
            println!("could potentially overwrite files already present in that directory.");
            print!("\nAre you sure you want to continue? (y/n) ");

            // Read answer from user and exit if it's not 'yes'
            if !confirm() {
                println!("\nSkipping...\n");
                println!("All done, no errors...");
                ::std::process::exit(0);
            }
        }

        // Call the function that copies the theme
        book.copy_theme()?;
        println!("\nTheme copied.");

    }

    // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`
    let is_dest_inside_root = book.get_destination()
                                  .map(|p| p.starts_with(book.get_root()))
                                  .unwrap_or(false);

    if !args.is_present("force") && is_dest_inside_root {
        println!("\nDo you want a .gitignore to be created? (y/n)");

        if confirm() {
            book.create_gitignore();
            println!("\n.gitignore created.");
        }
    }

    println!("\nAll done, no errors...");

    Ok(())
}

// Build command implementation
fn build(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let book = MDBook::new(&book_dir).read_config()?;

    let mut book = match args.value_of("dest-dir") {
        Some(dest_dir) => book.with_destination(dest_dir),
        None => book,
    };

    if args.is_present("no-create") {
        book.create_missing = false;
    }

    if args.is_present("curly-quotes") {
        book = book.with_curly_quotes(true);
    }

    book.build()?;

    if let Some(d) = book.get_destination() {
        if args.is_present("open") {
            open(d.join("index.html"));
        }
    }

    Ok(())
}

// Watch command implementation
#[cfg(feature = "watch")]
fn watch(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let book = MDBook::new(&book_dir).read_config()?;

    let mut book = match args.value_of("dest-dir") {
        Some(dest_dir) => book.with_destination(dest_dir),
        None => book,
    };

    if args.is_present("curly-quotes") {
        book = book.with_curly_quotes(true);
    }

    if args.is_present("open") {
        book.build()?;
        if let Some(d) = book.get_destination() {
            open(d.join("index.html"));
        }
    }

    trigger_on_change(&mut book, |path, book| {
        println!("File changed: {:?}\nBuilding book...\n", path);
        if let Err(e) = book.build() {
            println!("Error while building: {:?}", e);
        }
        println!("");
    });

    Ok(())
}

#[cfg(feature = "serve")]
mod serve {
    extern crate iron;
    extern crate staticfile;
    extern crate ws;

    use std;
    use std::path::Path;
    use std::error::Error;
    use self::iron::{Iron, AfterMiddleware, IronResult, IronError, Request, Response, status, Set, Chain};
    use clap::ArgMatches;
    use mdbook::MDBook;

    use {get_book_dir, open, trigger_on_change};

    struct ErrorRecover;

    impl AfterMiddleware for ErrorRecover {
        fn catch(&self, _: &mut Request, err: IronError) -> IronResult<Response> {
            match err.response.status {
                // each error will result in 404 response
                Some(_) => Ok(err.response.set(status::NotFound)),
                _ => Err(err),
            }
        }
    }

    // Watch command implementation
    pub fn serve(args: &ArgMatches) -> Result<(), Box<Error>> {
        const RELOAD_COMMAND: &'static str = "reload";

        let book_dir = get_book_dir(args);
        let book = MDBook::new(&book_dir).read_config()?;

        let mut book = match args.value_of("dest-dir") {
            Some(dest_dir) => book.with_destination(Path::new(dest_dir)),
            None => book,
        };

        if let None = book.get_destination() {
            println!("The HTML renderer is not set up, impossible to serve the files.");
            std::process::exit(2);
        }

        if args.is_present("curly-quotes") {
            book = book.with_curly_quotes(true);
        }

        let port = args.value_of("port").unwrap_or("3000");
        let ws_port = args.value_of("websocket-port").unwrap_or("3001");
        let interface = args.value_of("interface").unwrap_or("localhost");
        let public_address = args.value_of("address").unwrap_or(interface);
        let open_browser = args.is_present("open");

        let address = format!("{}:{}", interface, port);
        let ws_address = format!("{}:{}", interface, ws_port);

        book.set_livereload(format!(r#"
        <script type="text/javascript">
            var socket = new WebSocket("ws://{}:{}");
            socket.onmessage = function (event) {{
                if (event.data === "{}") {{
                    socket.close();
                    location.reload(true); // force reload from server (not from cache)
                }}
            }};

            window.onbeforeunload = function() {{
                socket.close();
            }}
        </script>
    "#,
                                    public_address,
                                    ws_port,
                                    RELOAD_COMMAND));

        book.build()?;

        let mut chain = Chain::new(staticfile::Static::new(book.get_destination()
                                                               .expect("destination is present, checked before")));
        chain.link_after(ErrorRecover);
        let _iron = Iron::new(chain).http(&*address).unwrap();

        let ws_server = ws::WebSocket::new(|_| |_| Ok(())).unwrap();

        let broadcaster = ws_server.broadcaster();

        std::thread::spawn(move || { ws_server.listen(&*ws_address).unwrap(); });

        let serving_url = format!("http://{}", address);
        println!("\nServing on: {}", serving_url);

        if open_browser {
            open(serving_url);
        }

        trigger_on_change(&mut book, move |path, book| {
            println!("File changed: {:?}\nBuilding book...\n", path);
            match book.build() {
                Err(e) => println!("Error while building: {:?}", e),
                _ => broadcaster.send(RELOAD_COMMAND).unwrap(),
            }
            println!("");
        });

        Ok(())
    }
}

fn test(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config()?;

    book.test()?;

    Ok(())
}

fn get_book_dir(args: &ArgMatches) -> PathBuf {
    if let Some(dir) = args.value_of("dir") {
        // Check if path is relative from current dir, or absolute...
        let p = Path::new(dir);
        if p.is_relative() {
            env::current_dir().unwrap().join(dir)
        } else {
            p.to_path_buf()
        }
    } else {
        env::current_dir().unwrap()
    }
}

fn open<P: AsRef<OsStr>>(path: P) {
    if let Err(e) = open::that(path) {
        println!("Error opening web browser: {}", e);
    }
}

// Calls the closure when a book source file is changed. This is blocking!
#[cfg(feature = "watch")]
fn trigger_on_change<F>(book: &mut MDBook, closure: F) -> ()
    where F: Fn(&Path, &mut MDBook) -> ()
{
    use notify::RecursiveMode::*;
    use notify::DebouncedEvent::*;

    // Create a channel to receive the events.
    let (tx, rx) = channel();

    let mut watcher = match notify::watcher(tx, Duration::from_secs(1)) {
        Ok(w) => w,
        Err(e) => {
            println!("Error while trying to watch the files:\n\n\t{:?}", e);
            ::std::process::exit(0);
        },
    };

    // Add the source directory to the watcher
    if let Err(e) = watcher.watch(book.get_source(), Recursive) {
        println!("Error while watching {:?}:\n    {:?}", book.get_source(), e);
        ::std::process::exit(0);
    };

    // Add the theme directory to the watcher
    if let Some(t) = book.get_theme_path() {
        watcher.watch(t, Recursive).unwrap_or_default();
    }

    // Add the book.{json,toml} file to the watcher if it exists, because it's not
    // located in the source directory
    if watcher
           .watch(book.get_root().join("book.json"), NonRecursive)
           .is_err() {
        // do nothing if book.json is not found
    }
    if watcher
           .watch(book.get_root().join("book.toml"), NonRecursive)
           .is_err() {
        // do nothing if book.toml is not found
    }

    println!("\nListening for changes...\n");

    loop {
        match rx.recv() {
            Ok(event) => {
                match event {
                    NoticeWrite(path) |
                    NoticeRemove(path) |
                    Create(path) |
                    Write(path) |
                    Remove(path) |
                    Rename(_, path) => {
                        closure(&path, book);
                    },
                    _ => {},
                }
            },
            Err(e) => {
                println!("An error occured: {:?}", e);
            },
        }
    }
}
<|endoftext|>