<commit_before>use std::{convert, io};
use std::fmt;
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use hash::Sha1;
use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

const PROTOCOL: &'static str = "BitTorrent protocol";
const BLOCK_SIZE: u32 = 16384;

struct Download<'a> {
    metainfo: &'a Metainfo,
    pieces:   Vec<Piece>,
}

impl<'a> Download<'a> {
    fn new(metainfo: &Metainfo) -> Download {
        let file_length = metainfo.info.length;
        let piece_length = metainfo.info.piece_length;
        let num_pieces = metainfo.info.num_pieces;

        // create pieces
        let mut pieces = vec![];
        for i in 0..num_pieces {
            let len = if i < (num_pieces - 1) {
                piece_length
            } else {
                (file_length - (piece_length as u64 * (num_pieces as u64 - 1))) as u32
            };
            pieces.push(Piece::new(i, len, metainfo.info.pieces[i as usize].clone()));
        }

        Download {
            metainfo: metainfo,
            pieces:   pieces,
        }
    }

    fn store(&mut self, piece_index: u32, block_index: u32, data: Vec<u8>) -> Result<(), Error> {
        let piece = &mut self.pieces[piece_index as usize];
        piece.store(block_index, data)
    }

    fn next_block_to_request(&self, peer_has_pieces: &[bool]) -> Option<(u32, u32, u32)> {
        for piece in self.pieces.iter() {
            if peer_has_pieces[piece.index as usize] {
                match piece.next_block_to_request() {
                    Some(block) => {
                        return Some((piece.index, block.index, block.length))
                    },
                    None => {}
                }
            }
        }
        None
    }
}

struct Piece {
    index:  u32,
    length: u32,
    hash:   Sha1,
    blocks: Vec<Block>,
}

impl Piece {
    fn new(index: u32, length: u32, hash: Sha1) -> Piece {
        // create blocks
        let mut blocks = vec![];
        let num_blocks = (length as f64 / BLOCK_SIZE as f64).ceil() as u32;
        for i in 0..num_blocks {
            let len = if i < (num_blocks - 1) {
                BLOCK_SIZE
            } else {
                length - (BLOCK_SIZE * (num_blocks - 1))
            };
            blocks.push(Block::new(i, len));
        }

        Piece {
            index:  index,
            length: length,
            hash:   hash,
            blocks: blocks,
        }
    }

    fn store(&mut self, block_index: u32, data: Vec<u8>) -> Result<(), Error> {
        let block = &mut self.blocks[block_index as usize];
        block.data = Some(data);
        Ok(())
    }

    fn next_block_to_request(&self) -> Option<&Block> {
        for block in self.blocks.iter() {
            if block.data.is_none() {
                return Some(block)
            }
        }
        None
    }
}

struct Block {
    index:  u32,
    length: u32,
    data:   Option<Vec<u8>>,
}

impl Block {
    fn new(index: u32, length: u32) -> Block {
        Block {
            index:  index,
            length: length,
            data:   None,
        }
    }
}

struct PeerConnection<'a> {
    metainfo: &'a Metainfo,
    stream: TcpStream,
    have: Vec<bool>,
    download: Download<'a>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl<'a> PeerConnection<'a> {
    fn connect(peer: &Peer, metainfo: &'a Metainfo) -> Result<PeerConnection<'a>, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.num_pieces;
        let mut conn = PeerConnection {
            metainfo: metainfo,
            stream: stream,
            have: vec![false; num_pieces as usize],
            download: Download::new(metainfo),
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.run());
        Ok(conn)
    }

    fn run(&mut self) -> Result<(), Error> {
        try!(self.send_handshake());
        try!(self.receive_handshake());
        loop {
            let message = try!(self.receive_message());
            println!("Recieved: {:?}", message);
            try!(self.process(message));
        }
        Ok(())
    }

    fn send_handshake(&mut self) -> Result<(), Error> {
        let mut message = vec![];
        message.push(PROTOCOL.len() as u8);
        message.extend(PROTOCOL.bytes());
        message.extend(vec![0; 8].into_iter());
        message.extend(self.metainfo.info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".bytes());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as u32));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn send_message(&mut self, message: Message) -> Result<(), Error> {
        println!("Sending: {:?}", message);
        try!(self.stream.write_all(&message.serialize()));
        Ok(())
    }

    fn receive_message(&mut self) -> Result<Message, Error> {
        let message_size = bytes_to_u32(&try!(self.read_n(4)));
        if message_size > 0 {
            let message = try!(self.read_n(message_size));
            Ok(Message::new(&message[0], &message[1..]))
        } else {
            Ok(Message::KeepAlive)
        }
    }

    fn read_n(&mut self, bytes_to_read: u32) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read as usize => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }

    fn process(&mut self, message: Message) -> Result<(), Error>{
        match message {
            Message::KeepAlive => {},
            Message::Bitfield(bytes) => {
                for have_index in 0..self.have.len() {
                    let bytes_index = have_index / 8;
                    let index_into_byte = have_index % 8;
                    let byte = bytes[bytes_index];
                    let value = (byte & (1 << (7 - index_into_byte))) != 0;
                    self.have[have_index] = value;
                };
                try!(self.send_interested());
            },
            Message::Have(have_index) => {
                self.have[have_index as usize] = true;
                try!(self.send_interested());
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    try!(self.request_next_block());
                }
            }
            Message::Piece(piece_index, offset, data) => {
                let block_index = offset / BLOCK_SIZE;
                try!(self.download.store(piece_index, block_index, data));
                try!(self.request_next_block());
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }

    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            self.am_i_interested = true;
            try!(self.send_message(Message::Interested));
        }
        Ok(())
    }

    fn request_next_block(&mut self) -> Result<(), Error> {
        match self.download.next_block_to_request(&self.have) {
            Some((piece_index, block_index, block_length)) => {
                let offset = block_index * BLOCK_SIZE;
                self.send_message(Message::Request(piece_index, offset, block_length))
            },
            None => {
                println!("Download complete!");
                Ok(())
            }
        }
    }
}

enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(u32),
    Bitfield(Vec<u8>),
    Request(u32, u32, u32),
    Piece(u32, u32, Vec<u8>),
    Cancel,  // TODO add params
    Port,    // TODO add params
}

impl Message {
    fn new(id: &u8, body: &[u8]) -> Message {
        match *id {
            0 => Message::Choke,
            1 => Message::Unchoke,
            2 => Message::Interested,
            3 => Message::NotInterested,
            4 => Message::Have(bytes_to_u32(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => {
                let index = bytes_to_u32(&body[0..4]);
                let offset = bytes_to_u32(&body[4..8]);
                let length = bytes_to_u32(&body[8..12]);
                Message::Request(index, offset, length)
            },
            7 => {
                let index = bytes_to_u32(&body[0..4]);
                let offset = bytes_to_u32(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, offset, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }

    fn serialize(self) -> Vec<u8> {
        let mut payload = vec![];
        match self {
            Message::KeepAlive => {},
            Message::Choke => payload.push(0),
            Message::Unchoke => payload.push(1),
            Message::Interested => payload.push(2),
            Message::NotInterested => payload.push(3),
            Message::Have(index) => {
                payload.push(4);
                payload.extend(u32_to_bytes(index).into_iter());
            },
            Message::Bitfield(bytes) => {
                payload.push(5);
                payload.extend(bytes);
            },
            Message::Request(index, offset, amount) => {
                payload.push(6);
                payload.extend(u32_to_bytes(index).into_iter());
                payload.extend(u32_to_bytes(offset).into_iter());
                payload.extend(u32_to_bytes(amount).into_iter());
            },
            Message::Piece(index, offset, data) => {
                payload.push(6);
                payload.extend(u32_to_bytes(index).into_iter());
                payload.extend(u32_to_bytes(offset).into_iter());
                payload.extend(data);
            },
            Message::Cancel => payload.push(8),
            Message::Port => payload.push(9),
        };

        // prepend size
        let mut size = u32_to_bytes(payload.len() as u32);
        size.extend(payload);
        size
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
             Message::KeepAlive => write!(f, "KeepAlive"),
             Message::Choke => write!(f, "Choke"),
             Message::Unchoke => write!(f, "Unchoke"),
             Message::Interested => write!(f, "Interested"),
             Message::NotInterested => write!(f, "NotInterested"),
             Message::Have(ref index) => write!(f, "Have({})", index),
             Message::Bitfield(ref bytes) => write!(f, "Bitfield({:?})", bytes),
             Message::Request(ref index, ref offset, ref length) => write!(f, "Request({}, {}, {})", index, offset, length),
             Message::Piece(ref index, ref offset, ref data) => write!(f, "Piece({}, {}, size={})", index, offset, data.len()),
             Message::Cancel => write!(f, "Cancel"),
             Message::Port => write!(f, "Port"),
        }
    }
}

const BYTE_0: u32 = 256 * 256 * 256;
const BYTE_1: u32 = 256 * 256;
const BYTE_2: u32 = 256;
const BYTE_3: u32 = 1;

fn bytes_to_u32(bytes: &[u8]) -> u32 {
    bytes[0] as u32 * BYTE_0 +
    bytes[1] as u32 * BYTE_1 +
    bytes[2] as u32 * BYTE_2 +
    bytes[3] as u32 * BYTE_3
}

fn u32_to_bytes(integer: u32) -> Vec<u8> {
    let mut rest = integer;
    let first = rest / BYTE_0;
    rest -= first * BYTE_0;
    let second = rest / BYTE_1;
    rest -= second * BYTE_1;
    let third = rest / BYTE_2;
    rest -= third * BYTE_2;
    let fourth = rest;
    vec![first as u8, second as u8, third as u8, fourth as u8]
}

#[derive(Debug)]
pub enum Error {
    IoError(io::Error),
    NotEnoughData,
}

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<commit_msg>Verify hashes of downloaded pieces, delete block data and re-download if verification fails.<commit_after>use std::{convert, io};
use std::fmt;
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use hash::{calculate_sha1, Sha1};
use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

const PROTOCOL: &'static str = "BitTorrent protocol";
const BLOCK_SIZE: u32 = 16384;

struct Download<'a> {
    metainfo: &'a Metainfo,
    pieces:   Vec<Piece>,
}

impl<'a> Download<'a> {
    fn new(metainfo: &Metainfo) -> Download {
        let file_length = metainfo.info.length;
        let piece_length = metainfo.info.piece_length;
        let num_pieces = metainfo.info.num_pieces;

        // create pieces
        let mut pieces = vec![];
        for i in 0..num_pieces {
            let len = if i < (num_pieces - 1) {
                piece_length
            } else {
                (file_length - (piece_length as u64 * (num_pieces as u64 - 1))) as u32
            };
            pieces.push(Piece::new(i, len, metainfo.info.pieces[i as usize].clone()));
        }

        Download {
            metainfo: metainfo,
            pieces:   pieces,
        }
    }

    fn store(&mut self, piece_index: u32, block_index: u32, data: Vec<u8>) -> Result<(), Error> {
        let piece = &mut self.pieces[piece_index as usize];
        piece.store(block_index, data)
        // TODO Detect if file is complete
    }

    fn next_block_to_request(&self, peer_has_pieces: &[bool]) -> Option<(u32, u32, u32)> {
        for piece in self.pieces.iter() {
            if peer_has_pieces[piece.index as usize] {
                match piece.next_block_to_request() {
                    Some(block) => {
                        return Some((piece.index, block.index, block.length))
                    },
                    None => {}
                }
            }
        }
        None
    }
}

struct Piece {
    index:  u32,
    length: u32,
    hash:   Sha1,
    blocks: Vec<Block>,
}

impl Piece {
    fn new(index: u32, length: u32, hash: Sha1) -> Piece {
        // create blocks
        let mut blocks = vec![];
        let num_blocks = (length as f64 / BLOCK_SIZE as f64).ceil() as u32;
        for i in 0..num_blocks {
            let len = if i < (num_blocks - 1) {
                BLOCK_SIZE
            } else {
                length - (BLOCK_SIZE * (num_blocks - 1))
            };
            blocks.push(Block::new(i, len));
        }

        Piece {
            index:  index,
            length: length,
            hash:   hash,
            blocks: blocks,
        }
    }

    fn store(&mut self, block_index: u32, data: Vec<u8>) -> Result<(), Error> {
        {
            let block = &mut self.blocks[block_index as usize];
            block.data = Some(data);
        }

        if self.is_complete() {
            if self.is_correct() {
                println!("Piece {} is complete and correct", self.index);
                // TODO Save to file
            } else {
                println!("Piece is corrupt, deleting downloaded piece data!");
                for block in self.blocks.iter_mut() {
                    block.data = None;
                }
            }
        }
        Ok(())
    }

    fn next_block_to_request(&self) -> Option<&Block> {
        for block in self.blocks.iter() {
            if block.data.is_none() {
                return Some(block)
            }
        }
        None
    }

    fn is_complete(&self) -> bool {
        for block in self.blocks.iter() {
            if block.data.is_none() {
                return false
            }
        }
        true
    }

    fn is_correct(&self) -> bool {
        let mut data = vec![];
        for block in self.blocks.iter() {
            data.extend(block.data.clone().unwrap());
        }

        let hash = calculate_sha1(&data);
        self.hash == hash
    }
}

struct Block {
    index:  u32,
    length: u32,
    data:   Option<Vec<u8>>,
}

impl Block {
    fn new(index: u32, length: u32) -> Block {
        Block {
            index:  index,
            length: length,
            data:   None,
        }
    }
}

struct PeerConnection<'a> {
    metainfo: &'a Metainfo,
    stream: TcpStream,
    have: Vec<bool>,
    download: Download<'a>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl<'a> PeerConnection<'a> {
    fn connect(peer: &Peer, metainfo: &'a Metainfo) -> Result<PeerConnection<'a>, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.num_pieces;
        let mut conn = PeerConnection {
            metainfo: metainfo,
            stream: stream,
            have: vec![false; num_pieces as usize],
            download: Download::new(metainfo),
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.run());
        Ok(conn)
    }

    fn run(&mut self) -> Result<(), Error> {
        try!(self.send_handshake());
        try!(self.receive_handshake());
        loop {
            let message = try!(self.receive_message());
            println!("Recieved: {:?}", message);
            try!(self.process(message));
        }
        Ok(())
    }

    fn send_handshake(&mut self) -> Result<(), Error> {
        let mut message = vec![];
        message.push(PROTOCOL.len() as u8);
        message.extend(PROTOCOL.bytes());
        message.extend(vec![0; 8].into_iter());
        message.extend(self.metainfo.info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".bytes());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as u32));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn send_message(&mut self, message: Message) -> Result<(), Error> {
        println!("Sending: {:?}", message);
        try!(self.stream.write_all(&message.serialize()));
        Ok(())
    }

    fn receive_message(&mut self) -> Result<Message, Error> {
        let message_size = bytes_to_u32(&try!(self.read_n(4)));
        if message_size > 0 {
            let message = try!(self.read_n(message_size));
            Ok(Message::new(&message[0], &message[1..]))
        } else {
            Ok(Message::KeepAlive)
        }
    }

    fn read_n(&mut self, bytes_to_read: u32) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read as usize => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }

    fn process(&mut self, message: Message) -> Result<(), Error>{
        match message {
            Message::KeepAlive => {},
            Message::Bitfield(bytes) => {
                for have_index in 0..self.have.len() {
                    let bytes_index = have_index / 8;
                    let index_into_byte = have_index % 8;
                    let byte = bytes[bytes_index];
                    let value = (byte & (1 << (7 - index_into_byte))) != 0;
                    self.have[have_index] = value;
                };
                try!(self.send_interested());
            },
            Message::Have(have_index) => {
                self.have[have_index as usize] = true;
                try!(self.send_interested());
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    try!(self.request_next_block());
                }
            }
            Message::Piece(piece_index, offset, data) => {
                let block_index = offset / BLOCK_SIZE;
                try!(self.download.store(piece_index, block_index, data));
                try!(self.request_next_block());
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }

    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            self.am_i_interested = true;
            try!(self.send_message(Message::Interested));
        }
        Ok(())
    }

    fn request_next_block(&mut self) -> Result<(), Error> {
        match self.download.next_block_to_request(&self.have) {
            Some((piece_index, block_index, block_length)) => {
                let offset = block_index * BLOCK_SIZE;
                self.send_message(Message::Request(piece_index, offset, block_length))
            },
            None => {
                println!("We've downloaded all the pieces we can from this peer.");
                Ok(())
            }
        }
    }
}

enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(u32),
    Bitfield(Vec<u8>),
    Request(u32, u32, u32),
    Piece(u32, u32, Vec<u8>),
    Cancel,  // TODO Add params
    Port,    // TODO Add params
}

impl Message {
    fn new(id: &u8, body: &[u8]) -> Message {
        match *id {
            0 => Message::Choke,
            1 => Message::Unchoke,
            2 => Message::Interested,
            3 => Message::NotInterested,
            4 => Message::Have(bytes_to_u32(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => {
                let index = bytes_to_u32(&body[0..4]);
                let offset = bytes_to_u32(&body[4..8]);
                let length = bytes_to_u32(&body[8..12]);
                Message::Request(index, offset, length)
            },
            7 => {
                let index = bytes_to_u32(&body[0..4]);
                let offset = bytes_to_u32(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, offset, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }

    fn serialize(self) -> Vec<u8> {
        let mut payload = vec![];
        match self {
            Message::KeepAlive => {},
            Message::Choke => payload.push(0),
            Message::Unchoke => payload.push(1),
            Message::Interested => payload.push(2),
            Message::NotInterested => payload.push(3),
            Message::Have(index) => {
                payload.push(4);
                payload.extend(u32_to_bytes(index).into_iter());
            },
            Message::Bitfield(bytes) => {
                payload.push(5);
                payload.extend(bytes);
            },
            Message::Request(index, offset, amount) => {
                payload.push(6);
                payload.extend(u32_to_bytes(index).into_iter());
                payload.extend(u32_to_bytes(offset).into_iter());
                payload.extend(u32_to_bytes(amount).into_iter());
            },
            Message::Piece(index, offset, data) => {
                payload.push(6);
                payload.extend(u32_to_bytes(index).into_iter());
                payload.extend(u32_to_bytes(offset).into_iter());
                payload.extend(data);
            },
            Message::Cancel => payload.push(8),
            Message::Port => payload.push(9),
        };

        // prepend size
        let mut size = u32_to_bytes(payload.len() as u32);
        size.extend(payload);
        size
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
             Message::KeepAlive => write!(f, "KeepAlive"),
             Message::Choke => write!(f, "Choke"),
             Message::Unchoke => write!(f, "Unchoke"),
             Message::Interested => write!(f, "Interested"),
             Message::NotInterested => write!(f, "NotInterested"),
             Message::Have(ref index) => write!(f, "Have({})", index),
             Message::Bitfield(ref bytes) => write!(f, "Bitfield({:?})", bytes),
             Message::Request(ref index, ref offset, ref length) => write!(f, "Request({}, {}, {})", index, offset, length),
             Message::Piece(ref index, ref offset, ref data) => write!(f, "Piece({}, {}, size={})", index, offset, data.len()),
             Message::Cancel => write!(f, "Cancel"),
             Message::Port => write!(f, "Port"),
        }
    }
}

const BYTE_0: u32 = 256 * 256 * 256;
const BYTE_1: u32 = 256 * 256;
const BYTE_2: u32 = 256;
const BYTE_3: u32 = 1;

fn bytes_to_u32(bytes: &[u8]) -> u32 {
    bytes[0] as u32 * BYTE_0 +
    bytes[1] as u32 * BYTE_1 +
    bytes[2] as u32 * BYTE_2 +
    bytes[3] as u32 * BYTE_3
}

fn u32_to_bytes(integer: u32) -> Vec<u8> {
    let mut rest = integer;
    let first = rest / BYTE_0;
    rest -= first * BYTE_0;
    let second = rest / BYTE_1;
    rest -= second * BYTE_1;
    let third = rest / BYTE_2;
    rest -= third * BYTE_2;
    let fourth = rest;
    vec![first as u8, second as u8, third as u8, fourth as u8]
}

#[derive(Debug)]
pub enum Error {
    IoError(io::Error),
    NotEnoughData,
}

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<|endoftext|>


<commit_before>use rustc::{declare_lint, hir, lint, lint_array, ty};
use syntax::ast;
use crate::utils;

/// **What it does:** Checks for usage of the `offset` pointer method with a `usize` casted to an
/// `isize`.
///
/// **Why is this bad?** If we’re always increasing the pointer address, we can avoid the numeric
/// cast by using the `add` method instead.
///
/// **Known problems:** None
///
/// **Example:**
/// ```rust
/// let vec = vec![b'a', b'b', b'c'];
/// let ptr = vec.as_ptr();
/// let offset = 1_usize;
///
/// unsafe { ptr.offset(offset as isize); }
/// ```
///
/// Could be written:
///
/// ```rust
/// let vec = vec![b'a', b'b', b'c'];
/// let ptr = vec.as_ptr();
/// let offset = 1_usize;
///
/// unsafe { ptr.add(offset); }
/// ```
declare_clippy_lint! {
    pub PTR_OFFSET_WITH_CAST,
    complexity,
    "uneeded pointer offset cast"
}

#[derive(Copy, Clone, Debug)]
pub struct Pass;

impl lint::LintPass for Pass {
    fn get_lints(&self) -> lint::LintArray {
        lint_array!(PTR_OFFSET_WITH_CAST)
    }
}

impl<'a, 'tcx> lint::LateLintPass<'a, 'tcx> for Pass {
    fn check_expr(&mut self, cx: &lint::LateContext<'a, 'tcx>, expr: &'tcx hir::Expr) {
        // Check if the expressions is a ptr.offset method call
        let [receiver_expr, arg_expr] = match expr_as_ptr_offset_call(cx, expr) {
            Some(call_arg) => call_arg,
            None => return,
        };

        // Check if the argument to ptr.offset is a cast from usize
        let cast_lhs_expr = match expr_as_cast_from_usize(cx, arg_expr) {
            Some(cast_lhs_expr) => cast_lhs_expr,
            None => return,
        };

        utils::span_lint_and_sugg(
            cx,
            PTR_OFFSET_WITH_CAST,
            expr.span,
            "use of `offset` with a `usize` casted to an `isize`",
            "try",
            build_suggestion(cx, receiver_expr, cast_lhs_expr),
        );
    }
}

// If the given expression is a cast from a usize, return the lhs of the cast
fn expr_as_cast_from_usize<'a, 'tcx>(
    cx: &lint::LateContext<'a, 'tcx>,
    expr: &'tcx hir::Expr,
) -> Option<&'tcx hir::Expr> {
    if let hir::ExprKind::Cast(ref cast_lhs_expr, _) = expr.node {
        if is_expr_ty_usize(cx, &cast_lhs_expr) {
            return Some(cast_lhs_expr);
        }
    }
    None
}

// If the given expression is a ptr::offset method call, return the receiver and the arg of the
// method call.
fn expr_as_ptr_offset_call<'a, 'tcx>(
    cx: &lint::LateContext<'a, 'tcx>,
    expr: &'tcx hir::Expr,
) -> Option<[&'tcx hir::Expr; 2]> {
    if let hir::ExprKind::MethodCall(ref path_segment, _, ref args) = expr.node {
        if path_segment.ident.name == "offset" && is_expr_ty_raw_ptr(cx, &args[0]) {
            return Some([&args[0], &args[1]]);
        }
    }
    None
}

// Is the type of the expression a usize?
fn is_expr_ty_usize<'a, 'tcx>(
    cx: &lint::LateContext<'a, 'tcx>,
    expr: &hir::Expr,
) -> bool {
    cx.tables.expr_ty(expr).sty == ty::TyKind::Uint(ast::UintTy::Usize)
}

// Is the type of the expression a raw pointer?
fn is_expr_ty_raw_ptr<'a, 'tcx>(
    cx: &lint::LateContext<'a, 'tcx>,
    expr: &hir::Expr,
) -> bool {
    if let ty::RawPtr(..) = cx.tables.expr_ty(expr).sty {
        true
    } else {
        false
    }
}

fn build_suggestion<'a, 'tcx>(
    cx: &lint::LateContext<'a, 'tcx>,
    receiver_expr: &hir::Expr,
    cast_lhs_expr: &hir::Expr,
) -> String {
    match (
        utils::snippet_opt(cx, receiver_expr.span),
        utils::snippet_opt(cx, cast_lhs_expr.span)
    ) {
        (Some(receiver), Some(cast_lhs)) => format!("{}.add({})", receiver, cast_lhs),
        _ => String::new(),
    }
}
<commit_msg>utilize cx.tcx.types.usize<commit_after>use rustc::{declare_lint, hir, lint, lint_array, ty};
use crate::utils;

/// **What it does:** Checks for usage of the `offset` pointer method with a `usize` casted to an
/// `isize`.
///
/// **Why is this bad?** If we’re always increasing the pointer address, we can avoid the numeric
/// cast by using the `add` method instead.
///
/// **Known problems:** None
///
/// **Example:**
/// ```rust
/// let vec = vec![b'a', b'b', b'c'];
/// let ptr = vec.as_ptr();
/// let offset = 1_usize;
///
/// unsafe { ptr.offset(offset as isize); }
/// ```
///
/// Could be written:
///
/// ```rust
/// let vec = vec![b'a', b'b', b'c'];
/// let ptr = vec.as_ptr();
/// let offset = 1_usize;
///
/// unsafe { ptr.add(offset); }
/// ```
declare_clippy_lint! {
    pub PTR_OFFSET_WITH_CAST,
    complexity,
    "uneeded pointer offset cast"
}

#[derive(Copy, Clone, Debug)]
pub struct Pass;

impl lint::LintPass for Pass {
    fn get_lints(&self) -> lint::LintArray {
        lint_array!(PTR_OFFSET_WITH_CAST)
    }
}

impl<'a, 'tcx> lint::LateLintPass<'a, 'tcx> for Pass {
    fn check_expr(&mut self, cx: &lint::LateContext<'a, 'tcx>, expr: &'tcx hir::Expr) {
        // Check if the expressions is a ptr.offset method call
        let [receiver_expr, arg_expr] = match expr_as_ptr_offset_call(cx, expr) {
            Some(call_arg) => call_arg,
            None => return,
        };

        // Check if the argument to ptr.offset is a cast from usize
        let cast_lhs_expr = match expr_as_cast_from_usize(cx, arg_expr) {
            Some(cast_lhs_expr) => cast_lhs_expr,
            None => return,
        };

        utils::span_lint_and_sugg(
            cx,
            PTR_OFFSET_WITH_CAST,
            expr.span,
            "use of `offset` with a `usize` casted to an `isize`",
            "try",
            build_suggestion(cx, receiver_expr, cast_lhs_expr),
        );
    }
}

// If the given expression is a cast from a usize, return the lhs of the cast
fn expr_as_cast_from_usize<'a, 'tcx>(
    cx: &lint::LateContext<'a, 'tcx>,
    expr: &'tcx hir::Expr,
) -> Option<&'tcx hir::Expr> {
    if let hir::ExprKind::Cast(ref cast_lhs_expr, _) = expr.node {
        if is_expr_ty_usize(cx, &cast_lhs_expr) {
            return Some(cast_lhs_expr);
        }
    }
    None
}

// If the given expression is a ptr::offset method call, return the receiver and the arg of the
// method call.
fn expr_as_ptr_offset_call<'a, 'tcx>(
    cx: &lint::LateContext<'a, 'tcx>,
    expr: &'tcx hir::Expr,
) -> Option<[&'tcx hir::Expr; 2]> {
    if let hir::ExprKind::MethodCall(ref path_segment, _, ref args) = expr.node {
        if path_segment.ident.name == "offset" && is_expr_ty_raw_ptr(cx, &args[0]) {
            return Some([&args[0], &args[1]]);
        }
    }
    None
}

// Is the type of the expression a usize?
fn is_expr_ty_usize<'a, 'tcx>(
    cx: &lint::LateContext<'a, 'tcx>,
    expr: &hir::Expr,
) -> bool {
    cx.tables.expr_ty(expr) == cx.tcx.types.usize
}

// Is the type of the expression a raw pointer?
fn is_expr_ty_raw_ptr<'a, 'tcx>(
    cx: &lint::LateContext<'a, 'tcx>,
    expr: &hir::Expr,
) -> bool {
    if let ty::RawPtr(..) = cx.tables.expr_ty(expr).sty {
        true
    } else {
        false
    }
}

fn build_suggestion<'a, 'tcx>(
    cx: &lint::LateContext<'a, 'tcx>,
    receiver_expr: &hir::Expr,
    cast_lhs_expr: &hir::Expr,
) -> String {
    match (
        utils::snippet_opt(cx, receiver_expr.span),
        utils::snippet_opt(cx, cast_lhs_expr.span)
    ) {
        (Some(receiver), Some(cast_lhs)) => format!("{}.add({})", receiver, cast_lhs),
        _ => String::new(),
    }
}
<|endoftext|>


<commit_before>use cgmath::EuclideanSpace;
use gilrs::{Axis, Button, Gilrs};
use std::f64::consts::PI;
use winit::{
    event,
    event_loop::{ControlFlow, EventLoop},
};

fn compute_projection_matrix(width: f32, height: f32) -> cgmath::Matrix4<f32> {
    let aspect = width as f32 / height as f32;
    let f = 1.0 / (45.0f32.to_radians() / aspect).tan();
    let near = 0.1;

    #[cfg_attr(rustfmt, rustfmt_skip)]
    cgmath::Matrix4::new(
        f,           0.0,  0.0,  0.0,
        0.0,        f * aspect,  0.0,  0.0,
        0.0,        0.0,  0.0,  -1.0,
        0.0,        0.0,  near,  0.0)
}

fn make_swapchain(
    device: &wgpu::Device,
    surface: &wgpu::Surface,
    width: u32,
    height: u32,
) -> wgpu::SwapChain {
    device.create_swap_chain(
        &surface,
        &wgpu::SwapChainDescriptor {
            usage: wgpu::TextureUsage::OUTPUT_ATTACHMENT,
            format: wgpu::TextureFormat::Bgra8UnormSrgb,
            width,
            height,
            present_mode: wgpu::PresentMode::Mailbox,
        },
    )
}
fn make_depth_buffer(device: &wgpu::Device, width: u32, height: u32) -> wgpu::TextureView {
    device
        .create_texture(&wgpu::TextureDescriptor {
            size: wgpu::Extent3d { width, height, depth: 1 },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Depth32Float,
            usage: wgpu::TextureUsage::OUTPUT_ATTACHMENT,
            label: None,
        })
        .create_default_view()
}

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

    let mut gilrs = Gilrs::new().unwrap();
    let mut current_gamepad = None;
    for (_id, gamepad) in gilrs.gamepads() {
        current_gamepad = Some(gamepad.id());
    }

    let event_loop = EventLoop::new();

    let instance = wgpu::Instance::new(wgpu::BackendBit::VULKAN);

    let window = winit::window::Window::new(&event_loop).unwrap();
    for monitor in window.available_monitors() {
        if monitor.video_modes().any(|mode| mode.size().width == 1920) {
            window.set_fullscreen(Some(winit::window::Fullscreen::Borderless(monitor)));
            break;
        }
    }
    let mut size = window.inner_size();
    let surface = unsafe { instance.create_surface(&window) };

    let adapter =
        futures::executor::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
            power_preference: wgpu::PowerPreference::Default,
            compatible_surface: Some(&surface),
        }))
        .unwrap();

    let (device, mut queue) = futures::executor::block_on(adapter.request_device(
        &wgpu::DeviceDescriptor {
            features: wgpu::Features::empty(),
            limits: wgpu::Limits::default(),
            shader_validation: true,
        },
        None,
    )).unwrap();

    let mapfile = terra::MapFileBuilder::build().unwrap();

    let mut terrain = terra::Terrain::new(&device, &mut queue, mapfile).unwrap();

    let mut swap_chain = make_swapchain(&device, &surface, size.width, size.height);
    let mut depth_buffer = make_depth_buffer(&device, size.width, size.height);

    let proj = compute_projection_matrix(size.width as f32, size.height as f32);

    let planet_radius = 6371000.0;

    let mut angle = 3.14159f64;

    let mut lat = 0.0f64;
    let mut long = 0.0f64;
    let mut altitude = 100.0f64;

    event_loop.run(move |event, _, control_flow| {
        *control_flow = if cfg!(feature = "metal-auto-capture") {
            ControlFlow::Exit
        } else {
            ControlFlow::Poll
        };
        match event {
            event::Event::WindowEvent { event, .. } => match event {
                event::WindowEvent::CloseRequested => {
                    *control_flow = ControlFlow::Exit;
                }
                event::WindowEvent::KeyboardInput {
                    input:
                        event::KeyboardInput {
                            virtual_keycode: Some(keycode),
                            state: event::ElementState::Pressed,
                            ..
                        },
                    ..
                } => match keycode {
                    event::VirtualKeyCode::Escape => *control_flow = ControlFlow::Exit,
                    event::VirtualKeyCode::Space => altitude += 0.1 * altitude,
                    event::VirtualKeyCode::Semicolon => altitude -= 0.1 * altitude,
                    // event::VirtualKeyCode::Left => eye.x -= 50.0,
                    // event::VirtualKeyCode::Right => eye.x += 50.0,
                    // event::VirtualKeyCode::Up => eye.z -= 50.0,
                    // event::VirtualKeyCode::Down => eye.z += 50.0,
                    _ => {}
                },
                event::WindowEvent::Resized(new_size) => {
                    size = new_size;
                    swap_chain = make_swapchain(&device, &surface, size.width, size.height);
                    depth_buffer = make_depth_buffer(&device, size.width, size.height);
                }
                _ => {}
            },
            event::Event::MainEventsCleared => {
                let frame = &swap_chain
                    .get_next_frame()
                    .unwrap().output.view;

                while let Some(gilrs::Event { id, event: _event, time: _ }) = gilrs.next_event() {
                    current_gamepad = Some(id);
                }
                if let Some(gamepad) = current_gamepad.map(|id| gilrs.gamepad(id)) {
                    lat += angle.cos() * gamepad.value(Axis::LeftStickY) as f64 * 0.01
                        - angle.sin() * gamepad.value(Axis::LeftStickX) as f64 * 0.01;

                    long += -angle.sin() * gamepad.value(Axis::LeftStickY) as f64 * 0.01
                        + angle.cos() * gamepad.value(Axis::LeftStickX) as f64 * 0.01;

                    angle -= gamepad.value(Axis::RightZ) as f64 * 0.01;

                    if gamepad.is_pressed(Button::DPadUp) {
                        altitude += 0.01 * altitude;
                    }
                    if gamepad.is_pressed(Button::DPadDown) {
                        altitude -= 0.01 * altitude;
                    }
                }

                lat = lat.max(-PI).min(PI);
                if long < -PI {
                    long += PI * 2.0;
                }
                if long > PI {
                    long -= PI * 2.0;
                }

                let r = altitude + planet_radius;
                let eye = cgmath::Point3::new(
                    r * lat.cos() * long.cos(),
                    r * lat.cos() * long.sin(),
                    r * lat.sin(),
                );

                let latc = lat + angle.cos() * 0.001;
                let longc = long - angle.sin() * 0.001;

                let center = cgmath::Point3::new(
                    planet_radius * latc.cos() * longc.cos() - eye.x,
                    planet_radius * latc.cos() * longc.sin() - eye.y,
                    planet_radius * latc.sin() - eye.z,
                );
                let up = cgmath::Vector3::new(eye.x as f32, eye.y as f32, eye.z as f32);

                let view = cgmath::Matrix4::look_at(
                    cgmath::Point3::origin(),
                    cgmath::Point3::new(center.x as f32, center.y as f32, center.z as f32),
                    up,
                );

                let view_proj = proj * view;
                let view_proj = mint::ColumnMatrix4 {
                    x: view_proj.x.into(),
                    y: view_proj.y.into(),
                    z: view_proj.z.into(),
                    w: view_proj.w.into(),
                };

                terrain.render(
                    &device,
                    &mut queue,
                    &frame,
                    &depth_buffer,
                    (size.width, size.height),
                    view_proj,
                    eye.into(),
                );
            }
            _ => (),
        }
    });
}
<commit_msg>Preview should start facing north<commit_after>use cgmath::EuclideanSpace;
use gilrs::{Axis, Button, Gilrs};
use std::f64::consts::PI;
use winit::{
    event,
    event_loop::{ControlFlow, EventLoop},
};

fn compute_projection_matrix(width: f32, height: f32) -> cgmath::Matrix4<f32> {
    let aspect = width as f32 / height as f32;
    let f = 1.0 / (45.0f32.to_radians() / aspect).tan();
    let near = 0.1;

    #[cfg_attr(rustfmt, rustfmt_skip)]
    cgmath::Matrix4::new(
        f,           0.0,  0.0,  0.0,
        0.0,        f * aspect,  0.0,  0.0,
        0.0,        0.0,  0.0,  -1.0,
        0.0,        0.0,  near,  0.0)
}

fn make_swapchain(
    device: &wgpu::Device,
    surface: &wgpu::Surface,
    width: u32,
    height: u32,
) -> wgpu::SwapChain {
    device.create_swap_chain(
        &surface,
        &wgpu::SwapChainDescriptor {
            usage: wgpu::TextureUsage::OUTPUT_ATTACHMENT,
            format: wgpu::TextureFormat::Bgra8UnormSrgb,
            width,
            height,
            present_mode: wgpu::PresentMode::Mailbox,
        },
    )
}
fn make_depth_buffer(device: &wgpu::Device, width: u32, height: u32) -> wgpu::TextureView {
    device
        .create_texture(&wgpu::TextureDescriptor {
            size: wgpu::Extent3d { width, height, depth: 1 },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Depth32Float,
            usage: wgpu::TextureUsage::OUTPUT_ATTACHMENT,
            label: None,
        })
        .create_default_view()
}

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

    let mut gilrs = Gilrs::new().unwrap();
    let mut current_gamepad = None;
    for (_id, gamepad) in gilrs.gamepads() {
        current_gamepad = Some(gamepad.id());
    }

    let event_loop = EventLoop::new();

    let instance = wgpu::Instance::new(wgpu::BackendBit::VULKAN);

    let window = winit::window::Window::new(&event_loop).unwrap();
    for monitor in window.available_monitors() {
        if monitor.video_modes().any(|mode| mode.size().width == 1920) {
            window.set_fullscreen(Some(winit::window::Fullscreen::Borderless(monitor)));
            break;
        }
    }
    let mut size = window.inner_size();
    let surface = unsafe { instance.create_surface(&window) };

    let adapter =
        futures::executor::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
            power_preference: wgpu::PowerPreference::Default,
            compatible_surface: Some(&surface),
        }))
        .unwrap();

    let (device, mut queue) = futures::executor::block_on(adapter.request_device(
        &wgpu::DeviceDescriptor {
            features: wgpu::Features::empty(),
            limits: wgpu::Limits::default(),
            shader_validation: true,
        },
        None,
    )).unwrap();

    let mapfile = terra::MapFileBuilder::build().unwrap();

    let mut terrain = terra::Terrain::new(&device, &mut queue, mapfile).unwrap();

    let mut swap_chain = make_swapchain(&device, &surface, size.width, size.height);
    let mut depth_buffer = make_depth_buffer(&device, size.width, size.height);

    let proj = compute_projection_matrix(size.width as f32, size.height as f32);

    let planet_radius = 6371000.0;

    let mut angle = 0.0f64;

    let mut lat = 0.0f64;
    let mut long = 0.0f64;
    let mut altitude = 100.0f64;

    event_loop.run(move |event, _, control_flow| {
        *control_flow = if cfg!(feature = "metal-auto-capture") {
            ControlFlow::Exit
        } else {
            ControlFlow::Poll
        };
        match event {
            event::Event::WindowEvent { event, .. } => match event {
                event::WindowEvent::CloseRequested => {
                    *control_flow = ControlFlow::Exit;
                }
                event::WindowEvent::KeyboardInput {
                    input:
                        event::KeyboardInput {
                            virtual_keycode: Some(keycode),
                            state: event::ElementState::Pressed,
                            ..
                        },
                    ..
                } => match keycode {
                    event::VirtualKeyCode::Escape => *control_flow = ControlFlow::Exit,
                    event::VirtualKeyCode::Space => altitude += 0.1 * altitude,
                    event::VirtualKeyCode::Semicolon => altitude -= 0.1 * altitude,
                    // event::VirtualKeyCode::Left => eye.x -= 50.0,
                    // event::VirtualKeyCode::Right => eye.x += 50.0,
                    // event::VirtualKeyCode::Up => eye.z -= 50.0,
                    // event::VirtualKeyCode::Down => eye.z += 50.0,
                    _ => {}
                },
                event::WindowEvent::Resized(new_size) => {
                    size = new_size;
                    swap_chain = make_swapchain(&device, &surface, size.width, size.height);
                    depth_buffer = make_depth_buffer(&device, size.width, size.height);
                }
                _ => {}
            },
            event::Event::MainEventsCleared => {
                let frame = &swap_chain
                    .get_next_frame()
                    .unwrap().output.view;

                while let Some(gilrs::Event { id, event: _event, time: _ }) = gilrs.next_event() {
                    current_gamepad = Some(id);
                }
                if let Some(gamepad) = current_gamepad.map(|id| gilrs.gamepad(id)) {
                    lat += angle.cos() * gamepad.value(Axis::LeftStickY) as f64 * 0.01
                        - angle.sin() * gamepad.value(Axis::LeftStickX) as f64 * 0.01;

                    long += -angle.sin() * gamepad.value(Axis::LeftStickY) as f64 * 0.01
                        + angle.cos() * gamepad.value(Axis::LeftStickX) as f64 * 0.01;

                    angle -= gamepad.value(Axis::RightZ) as f64 * 0.01;

                    if gamepad.is_pressed(Button::DPadUp) {
                        altitude += 0.01 * altitude;
                    }
                    if gamepad.is_pressed(Button::DPadDown) {
                        altitude -= 0.01 * altitude;
                    }
                }

                lat = lat.max(-PI).min(PI);
                if long < -PI {
                    long += PI * 2.0;
                }
                if long > PI {
                    long -= PI * 2.0;
                }

                let r = altitude + planet_radius;
                let eye = cgmath::Point3::new(
                    r * lat.cos() * long.cos(),
                    r * lat.cos() * long.sin(),
                    r * lat.sin(),
                );

                let latc = lat + angle.cos() * 0.001;
                let longc = long - angle.sin() * 0.001;

                let center = cgmath::Point3::new(
                    planet_radius * latc.cos() * longc.cos() - eye.x,
                    planet_radius * latc.cos() * longc.sin() - eye.y,
                    planet_radius * latc.sin() - eye.z,
                );
                let up = cgmath::Vector3::new(eye.x as f32, eye.y as f32, eye.z as f32);

                let view = cgmath::Matrix4::look_at(
                    cgmath::Point3::origin(),
                    cgmath::Point3::new(center.x as f32, center.y as f32, center.z as f32),
                    up,
                );

                let view_proj = proj * view;
                let view_proj = mint::ColumnMatrix4 {
                    x: view_proj.x.into(),
                    y: view_proj.y.into(),
                    z: view_proj.z.into(),
                    w: view_proj.w.into(),
                };

                terrain.render(
                    &device,
                    &mut queue,
                    &frame,
                    &depth_buffer,
                    (size.width, size.height),
                    view_proj,
                    eye.into(),
                );
            }
            _ => (),
        }
    });
}
<|endoftext|>


<commit_before>use cgmath::EuclideanSpace;
use gilrs::{Axis, Button, Gilrs};
use std::{f64::consts::PI, path::PathBuf};
use structopt::StructOpt;
use winit::{
    event,
    event_loop::{ControlFlow, EventLoop},
};

#[derive(Debug, StructOpt)]
struct Opt {
    #[structopt(short, long, default_value = "8FH495PF+29")]
    plus: String,
    #[structopt(short, long, default_value = "0")]
    heading: f64,
    #[structopt(short, long, default_value = "200000")]
    elevation: f64,
    #[structopt(long)]
    generate: Option<PathBuf>,
}

fn compute_projection_matrix(width: f32, height: f32) -> cgmath::Matrix4<f32> {
    let aspect = width / height;
    let f = 1.0 / (45.0f32.to_radians() / aspect).tan();
    let near = 0.1;

    #[cfg_attr(rustfmt, rustfmt_skip)]
    cgmath::Matrix4::new(
        f/aspect,  0.0,  0.0,   0.0,
        0.0,       f,    0.0,   0.0,
        0.0,       0.0,  0.0,  -1.0,
        0.0,       0.0,  near,  0.0)
}

fn make_swapchain(
    device: &wgpu::Device,
    surface: &wgpu::Surface,
    width: u32,
    height: u32,
) -> wgpu::SwapChain {
    device.create_swap_chain(
        &surface,
        &wgpu::SwapChainDescriptor {
            usage: wgpu::TextureUsage::RENDER_ATTACHMENT,
            format: wgpu::TextureFormat::Bgra8UnormSrgb,
            width,
            height,
            present_mode: wgpu::PresentMode::Fifo, // disable vsync by switching to Mailbox,
        },
    )
}
fn make_depth_buffer(device: &wgpu::Device, width: u32, height: u32) -> wgpu::TextureView {
    device
        .create_texture(&wgpu::TextureDescriptor {
            size: wgpu::Extent3d { width, height, depth: 1 },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Depth32Float,
            usage: wgpu::TextureUsage::RENDER_ATTACHMENT,
            label: None,
        })
        .create_view(&Default::default())
}

fn main() {
    // env_logger::init();

    let runtime = tokio::runtime::Runtime::new().unwrap();

    let trace_path: Option<&std::path::Path> = if cfg!(feature = "trace") {
        std::fs::create_dir_all("trace").unwrap();
        Some(std::path::Path::new("trace"))
    } else {
        None
    };

    let event_loop = EventLoop::new();
    let monitor = event_loop
        .available_monitors()
        .find(|monitor| monitor.video_modes().any(|mode| mode.size().width == 1920));
    let window = winit::window::WindowBuilder::new()
        .with_visible(false)
        .with_fullscreen(Some(winit::window::Fullscreen::Borderless(monitor)))
        .build(&event_loop)
        .unwrap();

    let instance = wgpu::Instance::new(wgpu::BackendBit::VULKAN | wgpu::BackendBit::DX12);
    let surface = unsafe { instance.create_surface(&window) };
    let adapter = runtime
        .block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
            power_preference: wgpu::PowerPreference::HighPerformance,
            compatible_surface: Some(&surface),
        }))
        .expect("Unable to create compatible wgpu adapter");
    let (device, mut queue) = runtime
        .block_on(adapter.request_device(
            &wgpu::DeviceDescriptor {
                features: wgpu::Features::TEXTURE_COMPRESSION_BC,
                limits: wgpu::Limits::default(),
                label: None,
            },
            trace_path,
        ))
        .expect("Unable to create compatible wgpu device");

    let mut size = window.inner_size();
    let mut swap_chain = make_swapchain(&device, &surface, size.width, size.height);
    let mut depth_buffer = make_depth_buffer(&device, size.width, size.height);
    let mut proj = compute_projection_matrix(size.width as f32, size.height as f32);

    let mut gilrs = Gilrs::new().unwrap();
    let mut current_gamepad = None;
    for (_id, gamepad) in gilrs.gamepads() {
        current_gamepad = Some(gamepad.id());
    }

    let opt = Opt::from_args();
    let plus_center =
        open_location_code::decode(&opt.plus).expect("Failed to parse plus code").center;

    let planet_radius = 6371000.0;
    let mut angle = opt.heading.to_radians();
    let mut lat = plus_center.y().to_radians();
    let mut long = plus_center.x().to_radians();
    let mut altitude = opt.elevation;

    let mut terrain = terra::Terrain::new(&device, &mut queue).unwrap();

    if let Some(dataset_directory) = opt.generate {
        let pb = indicatif::ProgressBar::new(100);
        pb.set_style(
            indicatif::ProgressStyle::default_bar()
                .template("{msg} {pos}/{len} [{wide_bar}] {percent}% {per_sec} {eta}")
                .progress_chars("=> "),
        );
        let mut last_message = None;
        let mut progress_callback = |l: &str, i: usize, total: usize| {
            if last_message.is_none() || l != last_message.as_ref().unwrap() {
                pb.set_message(l);
                pb.reset_eta();
                last_message = Some(l.to_string());
            }
            pb.set_length(total as u64);
            pb.set_position(i as u64);
        };

        runtime
            .block_on(terrain.generate_heightmaps(
                dataset_directory.join("ETOPO1_Ice_c_geotiff.zip"),
                dataset_directory.join("strm3"),
                &mut progress_callback,
            ))
            .unwrap();
        runtime
            .block_on(
                terrain
                    .generate_albedos(dataset_directory.join("bluemarble"), &mut progress_callback),
            )
            .unwrap();
        runtime.block_on(terrain.generate_roughness(&mut progress_callback)).unwrap();
    }

    {
        let r = altitude + planet_radius;
        let eye = cgmath::Point3::new(
            r * lat.cos() * long.cos(),
            r * lat.cos() * long.sin(),
            r * lat.sin(),
        );
        while terrain.poll_loading_status(&device, &mut queue, eye.into()) {
            std::thread::sleep(std::time::Duration::from_millis(10));
        }
    }

    let mut set_visible = false;
    event_loop.run(move |event, _, control_flow| {
        *control_flow = if cfg!(feature = "metal-auto-capture") {
            ControlFlow::Exit
        } else {
            ControlFlow::Poll
        };
        match event {
            event::Event::WindowEvent { event, .. } => match event {
                event::WindowEvent::CloseRequested => {
                    *control_flow = ControlFlow::Exit;
                }
                event::WindowEvent::KeyboardInput {
                    input:
                        event::KeyboardInput {
                            virtual_keycode: Some(keycode),
                            state: event::ElementState::Pressed,
                            ..
                        },
                    ..
                } => match keycode {
                    event::VirtualKeyCode::Escape => *control_flow = ControlFlow::Exit,
                    event::VirtualKeyCode::Space => altitude += 0.1 * altitude,
                    event::VirtualKeyCode::Semicolon => altitude -= 0.1 * altitude,
                    event::VirtualKeyCode::Left => {
                        lat += -angle.sin() * -(0.0000001 * altitude).min(0.01);
                        long += angle.cos() * -(0.0000001 * altitude).min(0.01);
                    }
                    event::VirtualKeyCode::Right => {
                        lat += -angle.sin() * (0.0000001 * altitude).min(0.01);
                        long += angle.cos() * (0.0000001 * altitude).min(0.01);
                    }
                    event::VirtualKeyCode::Up => {
                        lat += angle.cos() * (0.0000001 * altitude).min(0.01);
                        long += -angle.sin() * (0.0000001 * altitude).min(0.01);
                    }
                    event::VirtualKeyCode::Down => {
                        lat += angle.cos() * -(0.0000001 * altitude).min(0.01);
                        long += -angle.sin() * -(0.0000001 * altitude).min(0.01);
                    }
                    _ => {}
                },
                event::WindowEvent::Resized(new_size) => {
                    size = new_size;
                    swap_chain = make_swapchain(&device, &surface, size.width, size.height);
                    depth_buffer = make_depth_buffer(&device, size.width, size.height);
                    proj = compute_projection_matrix(size.width as f32, size.height as f32);
                }
                _ => {}
            },
            event::Event::MainEventsCleared => {
                let frame = swap_chain.get_current_frame();
                let frame = match frame {
                    Ok(ref f) => &f.output.view,
                    Err(_) => return,
                };

                while let Some(gilrs::Event { id, event: _event, time: _ }) = gilrs.next_event() {
                    current_gamepad = Some(id);
                }
                if let Some(gamepad) = current_gamepad.map(|id| gilrs.gamepad(id)) {
                    lat += angle.cos() * gamepad.value(Axis::LeftStickY) as f64 * 0.01
                        - angle.sin() * gamepad.value(Axis::LeftStickX) as f64 * 0.01;

                    long += -angle.sin() * gamepad.value(Axis::LeftStickY) as f64 * 0.01
                        + angle.cos() * gamepad.value(Axis::LeftStickX) as f64 * 0.01;

                    angle -= gamepad.value(Axis::RightZ) as f64 * 0.01;

                    if gamepad.is_pressed(Button::DPadUp) {
                        altitude += 0.01 * altitude;
                    }
                    if gamepad.is_pressed(Button::DPadDown) {
                        altitude -= 0.01 * altitude;
                    }
                }

                lat = lat.max(-PI).min(PI);
                if long < -PI {
                    long += PI * 2.0;
                }
                if long > PI {
                    long -= PI * 2.0;
                }

                let surface_height = terrain.get_height(lat, long) as f64;
                let r = altitude + planet_radius + surface_height + 2.0;
                let eye = cgmath::Point3::new(
                    r * lat.cos() * long.cos(),
                    r * lat.cos() * long.sin(),
                    r * lat.sin(),
                );

                let dt = (planet_radius / (planet_radius + altitude)).acos() * 0.3;
                let latc = lat + angle.cos() * dt;
                let longc = long - angle.sin() * dt;

                let center = cgmath::Point3::new(
                    planet_radius * latc.cos() * longc.cos() - eye.x,
                    planet_radius * latc.cos() * longc.sin() - eye.y,
                    planet_radius * latc.sin() - eye.z,
                );
                let up = cgmath::Vector3::new(eye.x as f32, eye.y as f32, eye.z as f32);

                let view = cgmath::Matrix4::look_at_rh(
                    cgmath::Point3::origin(),
                    cgmath::Point3::new(center.x as f32, center.y as f32, center.z as f32),
                    up,
                );

                let view_proj = proj * view;
                let view_proj = mint::ColumnMatrix4 {
                    x: view_proj.x.into(),
                    y: view_proj.y.into(),
                    z: view_proj.z.into(),
                    w: view_proj.w.into(),
                };

                terrain.render(
                    &device,
                    &mut queue,
                    &frame,
                    &depth_buffer,
                    (size.width, size.height),
                    view_proj,
                    eye.into(),
                );

                if !set_visible {
                    window.set_visible(true);
                    set_visible = true;
                }
            }
            _ => (),
        }
    });
}
<commit_msg>Create an adapter with any primary wgpu backend<commit_after>use cgmath::EuclideanSpace;
use gilrs::{Axis, Button, Gilrs};
use std::{f64::consts::PI, path::PathBuf};
use structopt::StructOpt;
use winit::{
    event,
    event_loop::{ControlFlow, EventLoop},
};

#[derive(Debug, StructOpt)]
struct Opt {
    #[structopt(short, long, default_value = "8FH495PF+29")]
    plus: String,
    #[structopt(short, long, default_value = "0")]
    heading: f64,
    #[structopt(short, long, default_value = "200000")]
    elevation: f64,
    #[structopt(long)]
    generate: Option<PathBuf>,
}

fn compute_projection_matrix(width: f32, height: f32) -> cgmath::Matrix4<f32> {
    let aspect = width / height;
    let f = 1.0 / (45.0f32.to_radians() / aspect).tan();
    let near = 0.1;

    #[cfg_attr(rustfmt, rustfmt_skip)]
    cgmath::Matrix4::new(
        f/aspect,  0.0,  0.0,   0.0,
        0.0,       f,    0.0,   0.0,
        0.0,       0.0,  0.0,  -1.0,
        0.0,       0.0,  near,  0.0)
}

fn make_swapchain(
    device: &wgpu::Device,
    surface: &wgpu::Surface,
    width: u32,
    height: u32,
) -> wgpu::SwapChain {
    device.create_swap_chain(
        &surface,
        &wgpu::SwapChainDescriptor {
            usage: wgpu::TextureUsage::RENDER_ATTACHMENT,
            format: wgpu::TextureFormat::Bgra8UnormSrgb,
            width,
            height,
            present_mode: wgpu::PresentMode::Fifo, // disable vsync by switching to Mailbox,
        },
    )
}
fn make_depth_buffer(device: &wgpu::Device, width: u32, height: u32) -> wgpu::TextureView {
    device
        .create_texture(&wgpu::TextureDescriptor {
            size: wgpu::Extent3d { width, height, depth: 1 },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Depth32Float,
            usage: wgpu::TextureUsage::RENDER_ATTACHMENT,
            label: None,
        })
        .create_view(&Default::default())
}

fn main() {
    // env_logger::init();

    let runtime = tokio::runtime::Runtime::new().unwrap();

    let trace_path: Option<&std::path::Path> = if cfg!(feature = "trace") {
        std::fs::create_dir_all("trace").unwrap();
        Some(std::path::Path::new("trace"))
    } else {
        None
    };

    let event_loop = EventLoop::new();
    let monitor = event_loop
        .available_monitors()
        .find(|monitor| monitor.video_modes().any(|mode| mode.size().width == 1920));
    let window = winit::window::WindowBuilder::new()
        .with_visible(false)
        .with_fullscreen(Some(winit::window::Fullscreen::Borderless(monitor)))
        .build(&event_loop)
        .unwrap();

    let instance = wgpu::Instance::new(wgpu::BackendBit::PRIMARY);
    let surface = unsafe { instance.create_surface(&window) };
    let adapter = runtime
        .block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
            power_preference: wgpu::PowerPreference::HighPerformance,
            compatible_surface: Some(&surface),
        }))
        .expect("Unable to create compatible wgpu adapter");
    let (device, mut queue) = runtime
        .block_on(adapter.request_device(
            &wgpu::DeviceDescriptor {
                features: wgpu::Features::TEXTURE_COMPRESSION_BC,
                limits: wgpu::Limits::default(),
                label: None,
            },
            trace_path,
        ))
        .expect("Unable to create compatible wgpu device");

    let mut size = window.inner_size();
    let mut swap_chain = make_swapchain(&device, &surface, size.width, size.height);
    let mut depth_buffer = make_depth_buffer(&device, size.width, size.height);
    let mut proj = compute_projection_matrix(size.width as f32, size.height as f32);

    let mut gilrs = Gilrs::new().unwrap();
    let mut current_gamepad = None;
    for (_id, gamepad) in gilrs.gamepads() {
        current_gamepad = Some(gamepad.id());
    }

    let opt = Opt::from_args();
    let plus_center =
        open_location_code::decode(&opt.plus).expect("Failed to parse plus code").center;

    let planet_radius = 6371000.0;
    let mut angle = opt.heading.to_radians();
    let mut lat = plus_center.y().to_radians();
    let mut long = plus_center.x().to_radians();
    let mut altitude = opt.elevation;

    let mut terrain = terra::Terrain::new(&device, &mut queue).unwrap();

    if let Some(dataset_directory) = opt.generate {
        let pb = indicatif::ProgressBar::new(100);
        pb.set_style(
            indicatif::ProgressStyle::default_bar()
                .template("{msg} {pos}/{len} [{wide_bar}] {percent}% {per_sec} {eta}")
                .progress_chars("=> "),
        );
        let mut last_message = None;
        let mut progress_callback = |l: &str, i: usize, total: usize| {
            if last_message.is_none() || l != last_message.as_ref().unwrap() {
                pb.set_message(l);
                pb.reset_eta();
                last_message = Some(l.to_string());
            }
            pb.set_length(total as u64);
            pb.set_position(i as u64);
        };

        runtime
            .block_on(terrain.generate_heightmaps(
                dataset_directory.join("ETOPO1_Ice_c_geotiff.zip"),
                dataset_directory.join("strm3"),
                &mut progress_callback,
            ))
            .unwrap();
        runtime
            .block_on(
                terrain
                    .generate_albedos(dataset_directory.join("bluemarble"), &mut progress_callback),
            )
            .unwrap();
        runtime.block_on(terrain.generate_roughness(&mut progress_callback)).unwrap();
    }

    {
        let r = altitude + planet_radius;
        let eye = cgmath::Point3::new(
            r * lat.cos() * long.cos(),
            r * lat.cos() * long.sin(),
            r * lat.sin(),
        );
        while terrain.poll_loading_status(&device, &mut queue, eye.into()) {
            std::thread::sleep(std::time::Duration::from_millis(10));
        }
    }

    let mut set_visible = false;
    event_loop.run(move |event, _, control_flow| {
        *control_flow = if cfg!(feature = "metal-auto-capture") {
            ControlFlow::Exit
        } else {
            ControlFlow::Poll
        };
        match event {
            event::Event::WindowEvent { event, .. } => match event {
                event::WindowEvent::CloseRequested => {
                    *control_flow = ControlFlow::Exit;
                }
                event::WindowEvent::KeyboardInput {
                    input:
                        event::KeyboardInput {
                            virtual_keycode: Some(keycode),
                            state: event::ElementState::Pressed,
                            ..
                        },
                    ..
                } => match keycode {
                    event::VirtualKeyCode::Escape => *control_flow = ControlFlow::Exit,
                    event::VirtualKeyCode::Space => altitude += 0.1 * altitude,
                    event::VirtualKeyCode::Semicolon => altitude -= 0.1 * altitude,
                    event::VirtualKeyCode::Left => {
                        lat += -angle.sin() * -(0.0000001 * altitude).min(0.01);
                        long += angle.cos() * -(0.0000001 * altitude).min(0.01);
                    }
                    event::VirtualKeyCode::Right => {
                        lat += -angle.sin() * (0.0000001 * altitude).min(0.01);
                        long += angle.cos() * (0.0000001 * altitude).min(0.01);
                    }
                    event::VirtualKeyCode::Up => {
                        lat += angle.cos() * (0.0000001 * altitude).min(0.01);
                        long += -angle.sin() * (0.0000001 * altitude).min(0.01);
                    }
                    event::VirtualKeyCode::Down => {
                        lat += angle.cos() * -(0.0000001 * altitude).min(0.01);
                        long += -angle.sin() * -(0.0000001 * altitude).min(0.01);
                    }
                    _ => {}
                },
                event::WindowEvent::Resized(new_size) => {
                    size = new_size;
                    swap_chain = make_swapchain(&device, &surface, size.width, size.height);
                    depth_buffer = make_depth_buffer(&device, size.width, size.height);
                    proj = compute_projection_matrix(size.width as f32, size.height as f32);
                }
                _ => {}
            },
            event::Event::MainEventsCleared => {
                let frame = swap_chain.get_current_frame();
                let frame = match frame {
                    Ok(ref f) => &f.output.view,
                    Err(_) => return,
                };

                while let Some(gilrs::Event { id, event: _event, time: _ }) = gilrs.next_event() {
                    current_gamepad = Some(id);
                }
                if let Some(gamepad) = current_gamepad.map(|id| gilrs.gamepad(id)) {
                    lat += angle.cos() * gamepad.value(Axis::LeftStickY) as f64 * 0.01
                        - angle.sin() * gamepad.value(Axis::LeftStickX) as f64 * 0.01;

                    long += -angle.sin() * gamepad.value(Axis::LeftStickY) as f64 * 0.01
                        + angle.cos() * gamepad.value(Axis::LeftStickX) as f64 * 0.01;

                    angle -= gamepad.value(Axis::RightZ) as f64 * 0.01;

                    if gamepad.is_pressed(Button::DPadUp) {
                        altitude += 0.01 * altitude;
                    }
                    if gamepad.is_pressed(Button::DPadDown) {
                        altitude -= 0.01 * altitude;
                    }
                }

                lat = lat.max(-PI).min(PI);
                if long < -PI {
                    long += PI * 2.0;
                }
                if long > PI {
                    long -= PI * 2.0;
                }

                let surface_height = terrain.get_height(lat, long) as f64;
                let r = altitude + planet_radius + surface_height + 2.0;
                let eye = cgmath::Point3::new(
                    r * lat.cos() * long.cos(),
                    r * lat.cos() * long.sin(),
                    r * lat.sin(),
                );

                let dt = (planet_radius / (planet_radius + altitude)).acos() * 0.3;
                let latc = lat + angle.cos() * dt;
                let longc = long - angle.sin() * dt;

                let center = cgmath::Point3::new(
                    planet_radius * latc.cos() * longc.cos() - eye.x,
                    planet_radius * latc.cos() * longc.sin() - eye.y,
                    planet_radius * latc.sin() - eye.z,
                );
                let up = cgmath::Vector3::new(eye.x as f32, eye.y as f32, eye.z as f32);

                let view = cgmath::Matrix4::look_at_rh(
                    cgmath::Point3::origin(),
                    cgmath::Point3::new(center.x as f32, center.y as f32, center.z as f32),
                    up,
                );

                let view_proj = proj * view;
                let view_proj = mint::ColumnMatrix4 {
                    x: view_proj.x.into(),
                    y: view_proj.y.into(),
                    z: view_proj.z.into(),
                    w: view_proj.w.into(),
                };

                terrain.render(
                    &device,
                    &mut queue,
                    &frame,
                    &depth_buffer,
                    (size.width, size.height),
                    view_proj,
                    eye.into(),
                );

                if !set_visible {
                    window.set_visible(true);
                    set_visible = true;
                }
            }
            _ => (),
        }
    });
}
<|endoftext|>


<commit_before>// Options
// Space before /
// Quote style
// Ordering of attributes

#[cfg(test)]
mod test {
    use std::io::{IoResult,MemWriter};

    use super::super::{Document,Element,Text};
    use super::super::{ElementElementChild,TextElementChild,CommentElementChild,PIElementChild};

    macro_rules! assert_str_eq(
        ($l:expr, $r:expr) => (assert_eq!($l.as_slice(), $r.as_slice()));
    )

    enum Content {
        WElement(Element),
        WElementEnd(String),
        WText(Text),
    }

    fn format_element<W : Writer>(element: Element, todo: &mut Vec<Content>, writer: &mut W) -> IoResult<()> {
        try!(write!(writer, "<{}", element.name()));

        for attr in element.attributes().iter() {
            try!(write!(writer, " {}='{}'", attr.name(), attr.value()));
        }

        let mut children = element.children();
        if children.is_empty() {
            writer.write_str("/>")
        } else {
            try!(writer.write_str(">"));

            todo.push(WElementEnd(element.name()));
            children.reverse();
            let x = children.move_iter().map(|c| match c {
                ElementElementChild(element) => WElement(element),
                TextElementChild(t) => WText(t),
                CommentElementChild(_) => fail!("comment?"),
                PIElementChild(_) => fail!("PI?"),
            }).collect();
            todo.push_all_move(x);

            Ok(())
        }
    }

    fn format_document<W : Writer>(doc: &Document, writer: &mut W) -> IoResult<()> {
        let mut todo = Vec::new();

        try!(writer.write_str("<?xml version='1.0'?>"));
        let e = doc.root().children()[0].element().expect("not an element?");

        todo.push(WElement(e));

        while ! todo.is_empty() {
            match todo.pop().unwrap() {
                WElement(e) => try!(format_element(e, &mut todo, writer)),
                WElementEnd(name) => try!(write!(writer, "</{}>", name)),
                WText(t) => try!(writer.write_str(t.text().as_slice())),
            }
        }

        Ok(())
    }

    fn format_xml(doc: &Document) -> String {
        let mut w = MemWriter::new();
        format_document(doc, &mut w).ok().expect("Not formatted");
        String::from_utf8(w.unwrap()).ok().expect("Not a string")
    }

    #[test]
    fn top_element() {
        let d = Document::new();
        let e = d.new_element("hello".to_string());
        d.root().append_child(e);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello/>");
    }

    #[test]
    fn element_with_attributes() {
        let d = Document::new();
        let e = d.new_element("hello".to_string());
        e.set_attribute("a".to_string(), "b".to_string());
        d.root().append_child(e);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello a='b'/>");
    }

    #[test]
    fn nested_element() {
        let d = Document::new();
        let hello = d.new_element("hello".to_string());
        let world = d.new_element("world".to_string());
        hello.append_child(world);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello><world/></hello>");
    }

    #[test]
    fn nested_text() {
        let d = Document::new();
        let hello = d.new_element("hello".to_string());
        let text = d.new_text("A fine day to you!".to_string());
        hello.append_child(text);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello>A fine day to you!</hello>");
    }
}
<commit_msg>Write comment nodes<commit_after>// Options
// Space before /
// Quote style
// Ordering of attributes

#[cfg(test)]
mod test {
    use std::io::{IoResult,MemWriter};

    use super::super::{Document,Element,Text,Comment};
    use super::super::{ElementElementChild,TextElementChild,CommentElementChild,PIElementChild};

    macro_rules! assert_str_eq(
        ($l:expr, $r:expr) => (assert_eq!($l.as_slice(), $r.as_slice()));
    )

    enum Content {
        WElement(Element),
        WElementEnd(String),
        WText(Text),
        WComment(Comment),
    }

    fn format_element<W : Writer>(element: Element, todo: &mut Vec<Content>, writer: &mut W) -> IoResult<()> {
        try!(write!(writer, "<{}", element.name()));

        for attr in element.attributes().iter() {
            try!(write!(writer, " {}='{}'", attr.name(), attr.value()));
        }

        let mut children = element.children();
        if children.is_empty() {
            writer.write_str("/>")
        } else {
            try!(writer.write_str(">"));

            todo.push(WElementEnd(element.name()));
            children.reverse();
            let x = children.move_iter().map(|c| match c {
                ElementElementChild(element) => WElement(element),
                TextElementChild(t) => WText(t),
                CommentElementChild(c) => WComment(c),
                PIElementChild(_) => fail!("PI?"),
            }).collect();
            todo.push_all_move(x);

            Ok(())
        }
    }

    fn format_document<W : Writer>(doc: &Document, writer: &mut W) -> IoResult<()> {
        let mut todo = Vec::new();

        try!(writer.write_str("<?xml version='1.0'?>"));
        let e = doc.root().children()[0].element().expect("not an element?");

        todo.push(WElement(e));

        while ! todo.is_empty() {
            match todo.pop().unwrap() {
                WElement(e) => try!(format_element(e, &mut todo, writer)),
                WElementEnd(name) => try!(write!(writer, "</{}>", name)),
                WText(t) => try!(writer.write_str(t.text().as_slice())),
                WComment(c) => try!(write!(writer, "<!--{}-->", c.text())),
            }
        }

        Ok(())
    }

    fn format_xml(doc: &Document) -> String {
        let mut w = MemWriter::new();
        format_document(doc, &mut w).ok().expect("Not formatted");
        String::from_utf8(w.unwrap()).ok().expect("Not a string")
    }

    #[test]
    fn top_element() {
        let d = Document::new();
        let e = d.new_element("hello".to_string());
        d.root().append_child(e);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello/>");
    }

    #[test]
    fn element_with_attributes() {
        let d = Document::new();
        let e = d.new_element("hello".to_string());
        e.set_attribute("a".to_string(), "b".to_string());
        d.root().append_child(e);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello a='b'/>");
    }

    #[test]
    fn nested_element() {
        let d = Document::new();
        let hello = d.new_element("hello".to_string());
        let world = d.new_element("world".to_string());
        hello.append_child(world);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello><world/></hello>");
    }

    #[test]
    fn nested_text() {
        let d = Document::new();
        let hello = d.new_element("hello".to_string());
        let text = d.new_text("A fine day to you!".to_string());
        hello.append_child(text);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello>A fine day to you!</hello>");
    }

    #[test]
    fn nested_comment() {
        let d = Document::new();
        let hello = d.new_element("hello".to_string());
        let comment = d.new_comment(" Fill this in ".to_string());
        hello.append_child(comment);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello><!-- Fill this in --></hello>");
    }
}
<|endoftext|>