<commit_before>use byteorder::{NativeEndian, ByteOrder};
use common::bitpacker::BitPacker;
use common::compute_num_bits;
use common::Endianness;
use common::{BinarySerializable, FixedSize};
use directory::ReadOnlySource;
use postings::TermInfo;
use std::cmp;
use std::io::{self, Read, Write};
use termdict::TermOrdinal;

const BLOCK_LEN: usize = 256;

#[derive(Debug, Eq, PartialEq, Default)]
struct TermInfoBlockMeta {
    offset: u64,
    ref_term_info: TermInfo,
    doc_freq_nbits: u8,
    postings_offset_nbits: u8,
    positions_idx_nbits: u8,
}

impl BinarySerializable for TermInfoBlockMeta {
    fn serialize<W: Write>(&self, write: &mut W) -> io::Result<()> {
        self.offset.serialize(write)?;
        self.ref_term_info.serialize(write)?;
        write.write_all(&[
            self.doc_freq_nbits,
            self.postings_offset_nbits,
            self.positions_idx_nbits,
        ])?;
        Ok(())
    }

    fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
        let offset = u64::deserialize(reader)?;
        let ref_term_info = TermInfo::deserialize(reader)?;
        let mut buffer = [0u8; 3];
        reader.read_exact(&mut buffer)?;
        Ok(TermInfoBlockMeta {
            offset,
            ref_term_info,
            doc_freq_nbits: buffer[0],
            postings_offset_nbits: buffer[1],
            positions_idx_nbits: buffer[2],
        })
    }
}

impl FixedSize for TermInfoBlockMeta {
    const SIZE_IN_BYTES: usize =
        u64::SIZE_IN_BYTES + TermInfo::SIZE_IN_BYTES + 3 * u8::SIZE_IN_BYTES;
}

impl TermInfoBlockMeta {
    fn num_bits(&self) -> u8 {
        self.doc_freq_nbits + self.postings_offset_nbits + self.positions_idx_nbits
    }

    fn deserialize_term_info(&self, data: &[u8], inner_offset: usize) -> TermInfo {
        let num_bits = self.num_bits() as usize;
        let mut cursor = num_bits * inner_offset;

        let doc_freq = extract_bits(data, cursor, self.doc_freq_nbits) as u32;
        cursor += self.doc_freq_nbits as usize;

        let postings_offset = extract_bits(data, cursor, self.postings_offset_nbits);
        cursor += self.postings_offset_nbits as usize;

        let positions_idx = extract_bits(data, cursor, self.positions_idx_nbits);

        TermInfo {
            doc_freq,
            postings_offset: postings_offset + self.ref_term_info.postings_offset,
            positions_idx: positions_idx + self.ref_term_info.positions_idx,
        }
    }
}

pub struct TermInfoStore {
    num_terms: usize,
    block_meta_source: ReadOnlySource,
    term_info_source: ReadOnlySource,
}

fn extract_bits(data: &[u8], addr_bits: usize, num_bits: u8) -> u64 {
    assert!(num_bits <= 56);
    let addr_byte = addr_bits / 8;
    let bit_shift = (addr_bits % 8) as u64;
    let val_unshifted_unmasked: u64;
    if data.len() >= addr_byte + 8 {
        val_unshifted_unmasked =  NativeEndian::read_u64(&data[addr_byte..][..8]);
    } else {
        let mut buf = [0u8; 8];
        let data_to_copy = &data[addr_byte..];
        let nbytes = data_to_copy.len();
        buf[..nbytes].copy_from_slice(data_to_copy);
        val_unshifted_unmasked =  NativeEndian::read_u64(&buf);
    }
    let val_shifted_unmasked = val_unshifted_unmasked >> bit_shift;
    let mask = (1u64 << u64::from(num_bits)) - 1;
    val_shifted_unmasked & mask

}

impl TermInfoStore {
    pub fn open(data: &ReadOnlySource) -> TermInfoStore {
        let buffer = data.as_slice();
        let len = Endianness::read_u64(&buffer[0..8]) as usize;
        let num_terms = Endianness::read_u64(&buffer[8..16]) as usize;
        let block_meta_source = data.slice(16, 16 + len);
        let term_info_source = data.slice_from(16 + len);
        TermInfoStore {
            num_terms,
            block_meta_source,
            term_info_source,
        }
    }

    pub fn get(&self, term_ord: TermOrdinal) -> TermInfo {
        let block_id = (term_ord as usize) / BLOCK_LEN;
        let buffer = self.block_meta_source.as_slice();
        let mut block_data: &[u8] = &buffer[block_id * TermInfoBlockMeta::SIZE_IN_BYTES..];
        let term_info_block_data = TermInfoBlockMeta::deserialize(&mut block_data)
            .expect("Failed to deserialize terminfoblockmeta");
        let inner_offset = (term_ord as usize) % BLOCK_LEN;
        if inner_offset == 0 {
            term_info_block_data.ref_term_info
        } else {
            let term_info_data = self.term_info_source.as_slice();
            term_info_block_data.deserialize_term_info(
                &term_info_data[term_info_block_data.offset as usize..],
                inner_offset - 1,
            )
        }
    }

    pub fn num_terms(&self) -> usize {
        self.num_terms
    }
}

pub struct TermInfoStoreWriter {
    buffer_block_metas: Vec<u8>,
    buffer_term_infos: Vec<u8>,
    term_infos: Vec<TermInfo>,
    num_terms: u64,
}

fn bitpack_serialize<W: Write>(
    write: &mut W,
    bit_packer: &mut BitPacker,
    term_info_block_meta: &TermInfoBlockMeta,
    term_info: &TermInfo,
) -> io::Result<()> {
    bit_packer.write(
        u64::from(term_info.doc_freq),
        term_info_block_meta.doc_freq_nbits,
        write,
    )?;
    bit_packer.write(
        term_info.postings_offset,
        term_info_block_meta.postings_offset_nbits,
        write,
    )?;
    bit_packer.write(
        term_info.positions_idx,
        term_info_block_meta.positions_idx_nbits,
        write,
    )?;
    Ok(())
}

impl TermInfoStoreWriter {
    pub fn new() -> TermInfoStoreWriter {
        TermInfoStoreWriter {
            buffer_block_metas: Vec::new(),
            buffer_term_infos: Vec::new(),
            term_infos: Vec::with_capacity(BLOCK_LEN),
            num_terms: 0u64,
        }
    }

    fn flush_block(&mut self) -> io::Result<()> {
        if self.term_infos.is_empty() {
            return Ok(());
        }
        let mut bit_packer = BitPacker::new();
        let ref_term_info = self.term_infos[0].clone();
        for term_info in &mut self.term_infos[1..] {
            term_info.postings_offset -= ref_term_info.postings_offset;
            term_info.positions_idx -= ref_term_info.positions_idx;
        }

        let mut max_doc_freq: u32 = 0u32;
        let mut max_postings_offset: u64 = 0u64;
        let mut max_positions_idx: u64 = 0u64;
        for term_info in &self.term_infos[1..] {
            max_doc_freq = cmp::max(max_doc_freq, term_info.doc_freq);
            max_postings_offset = cmp::max(max_postings_offset, term_info.postings_offset);
            max_positions_idx = cmp::max(max_positions_idx, term_info.positions_idx);
        }

        let max_doc_freq_nbits: u8 = compute_num_bits(u64::from(max_doc_freq));
        let max_postings_offset_nbits = compute_num_bits(max_postings_offset);
        let max_positions_idx_nbits = compute_num_bits(max_positions_idx);

        let term_info_block_meta = TermInfoBlockMeta {
            offset: self.buffer_term_infos.len() as u64,
            ref_term_info,
            doc_freq_nbits: max_doc_freq_nbits,
            postings_offset_nbits: max_postings_offset_nbits,
            positions_idx_nbits: max_positions_idx_nbits,
        };

        term_info_block_meta.serialize(&mut self.buffer_block_metas)?;
        for term_info in self.term_infos[1..].iter().cloned() {
            bitpack_serialize(
                &mut self.buffer_term_infos,
                &mut bit_packer,
                &term_info_block_meta,
                &term_info,
            )?;
        }

        // Block need end up at the end of a byte.
        bit_packer.flush(&mut self.buffer_term_infos)?;
        self.term_infos.clear();

        Ok(())
    }

    pub fn write_term_info(&mut self, term_info: &TermInfo) -> io::Result<()> {
        self.num_terms += 1u64;
        self.term_infos.push(term_info.clone());
        if self.term_infos.len() >= BLOCK_LEN {
            self.flush_block()?;
        }
        Ok(())
    }

    pub fn serialize<W: io::Write>(&mut self, write: &mut W) -> io::Result<()> {
        if !self.term_infos.is_empty() {
            self.flush_block()?;
        }
        let len = self.buffer_block_metas.len() as u64;
        len.serialize(write)?;
        self.num_terms.serialize(write)?;
        write.write_all(&self.buffer_block_metas)?;
        write.write_all(&self.buffer_term_infos)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {

    use super::extract_bits;
    use super::TermInfoBlockMeta;
    use super::{TermInfoStore, TermInfoStoreWriter};
    use common;
    use common::bitpacker::BitPacker;
    use common::compute_num_bits;
    use common::BinarySerializable;
    use directory::ReadOnlySource;
    use postings::TermInfo;

    #[test]
    fn test_term_info_block() {
        common::test::fixed_size_test::<TermInfoBlockMeta>();
    }

    #[test]
    fn test_bitpacked() {
        let mut buffer = Vec::new();
        let mut bitpack = BitPacker::new();
        bitpack.write(321u64, 9, &mut buffer).unwrap();
        assert_eq!(compute_num_bits(321u64), 9);
        bitpack.write(2u64, 2, &mut buffer).unwrap();
        assert_eq!(compute_num_bits(2u64), 2);
        bitpack.write(51, 6, &mut buffer).unwrap();
        assert_eq!(compute_num_bits(51), 6);
        bitpack.close(&mut buffer).unwrap();
        assert_eq!(buffer.len(), 3 + 7);
        assert_eq!(extract_bits(&buffer[..], 0, 9), 321u64);
        assert_eq!(extract_bits(&buffer[..], 9, 2), 2u64);
        assert_eq!(extract_bits(&buffer[..], 11, 6), 51u64);
    }

    #[test]
    fn test_term_info_block_meta_serialization() {
        let term_info_block_meta = TermInfoBlockMeta {
            offset: 2009,
            ref_term_info: TermInfo {
                doc_freq: 512,
                postings_offset: 51,
                positions_idx: 3584,
            },
            doc_freq_nbits: 10,
            postings_offset_nbits: 5,
            positions_idx_nbits: 11,
        };
        let mut buffer: Vec<u8> = Vec::new();
        term_info_block_meta.serialize(&mut buffer).unwrap();
        let mut cursor: &[u8] = &buffer[..];
        let term_info_block_meta_serde = TermInfoBlockMeta::deserialize(&mut cursor).unwrap();
        assert_eq!(term_info_block_meta_serde, term_info_block_meta);
    }

    #[test]
    fn test_pack() {
        let mut store_writer = TermInfoStoreWriter::new();
        let mut term_infos = vec![];
        for i in 0..1000 {
            let term_info = TermInfo {
                doc_freq: i as u32,
                postings_offset: (i / 10) as u64,
                positions_idx: (i * 7) as u64,
            };
            store_writer.write_term_info(&term_info).unwrap();
            term_infos.push(term_info);
        }
        let mut buffer = Vec::new();
        store_writer.serialize(&mut buffer).unwrap();
        let term_info_store = TermInfoStore::open(&ReadOnlySource::from(buffer));
        for i in 0..1000 {
            assert_eq!(term_info_store.get(i as u64), term_infos[i]);
        }
    }

}
<commit_msg>Using LittleEndian explicitely<commit_after>use byteorder::{LittleEndian, ByteOrder};
use common::bitpacker::BitPacker;
use common::compute_num_bits;
use common::Endianness;
use common::{BinarySerializable, FixedSize};
use directory::ReadOnlySource;
use postings::TermInfo;
use std::cmp;
use std::io::{self, Read, Write};
use termdict::TermOrdinal;

const BLOCK_LEN: usize = 256;

#[derive(Debug, Eq, PartialEq, Default)]
struct TermInfoBlockMeta {
    offset: u64,
    ref_term_info: TermInfo,
    doc_freq_nbits: u8,
    postings_offset_nbits: u8,
    positions_idx_nbits: u8,
}

impl BinarySerializable for TermInfoBlockMeta {
    fn serialize<W: Write>(&self, write: &mut W) -> io::Result<()> {
        self.offset.serialize(write)?;
        self.ref_term_info.serialize(write)?;
        write.write_all(&[
            self.doc_freq_nbits,
            self.postings_offset_nbits,
            self.positions_idx_nbits,
        ])?;
        Ok(())
    }

    fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
        let offset = u64::deserialize(reader)?;
        let ref_term_info = TermInfo::deserialize(reader)?;
        let mut buffer = [0u8; 3];
        reader.read_exact(&mut buffer)?;
        Ok(TermInfoBlockMeta {
            offset,
            ref_term_info,
            doc_freq_nbits: buffer[0],
            postings_offset_nbits: buffer[1],
            positions_idx_nbits: buffer[2],
        })
    }
}

impl FixedSize for TermInfoBlockMeta {
    const SIZE_IN_BYTES: usize =
        u64::SIZE_IN_BYTES + TermInfo::SIZE_IN_BYTES + 3 * u8::SIZE_IN_BYTES;
}

impl TermInfoBlockMeta {
    fn num_bits(&self) -> u8 {
        self.doc_freq_nbits + self.postings_offset_nbits + self.positions_idx_nbits
    }

    fn deserialize_term_info(&self, data: &[u8], inner_offset: usize) -> TermInfo {
        let num_bits = self.num_bits() as usize;
        let mut cursor = num_bits * inner_offset;

        let doc_freq = extract_bits(data, cursor, self.doc_freq_nbits) as u32;
        cursor += self.doc_freq_nbits as usize;

        let postings_offset = extract_bits(data, cursor, self.postings_offset_nbits);
        cursor += self.postings_offset_nbits as usize;

        let positions_idx = extract_bits(data, cursor, self.positions_idx_nbits);

        TermInfo {
            doc_freq,
            postings_offset: postings_offset + self.ref_term_info.postings_offset,
            positions_idx: positions_idx + self.ref_term_info.positions_idx,
        }
    }
}

pub struct TermInfoStore {
    num_terms: usize,
    block_meta_source: ReadOnlySource,
    term_info_source: ReadOnlySource,
}

fn extract_bits(data: &[u8], addr_bits: usize, num_bits: u8) -> u64 {
    assert!(num_bits <= 56);
    let addr_byte = addr_bits / 8;
    let bit_shift = (addr_bits % 8) as u64;
    let val_unshifted_unmasked: u64;
    if data.len() >= addr_byte + 8 {
        val_unshifted_unmasked =  LittleEndian::read_u64(&data[addr_byte..][..8]);
    } else {
        let mut buf = [0u8; 8];
        let data_to_copy = &data[addr_byte..];
        let nbytes = data_to_copy.len();
        buf[..nbytes].copy_from_slice(data_to_copy);
        val_unshifted_unmasked =  LittleEndian::read_u64(&buf);
    }
    let val_shifted_unmasked = val_unshifted_unmasked >> bit_shift;
    let mask = (1u64 << u64::from(num_bits)) - 1;
    val_shifted_unmasked & mask

}

impl TermInfoStore {
    pub fn open(data: &ReadOnlySource) -> TermInfoStore {
        let buffer = data.as_slice();
        let len = Endianness::read_u64(&buffer[0..8]) as usize;
        let num_terms = Endianness::read_u64(&buffer[8..16]) as usize;
        let block_meta_source = data.slice(16, 16 + len);
        let term_info_source = data.slice_from(16 + len);
        TermInfoStore {
            num_terms,
            block_meta_source,
            term_info_source,
        }
    }

    pub fn get(&self, term_ord: TermOrdinal) -> TermInfo {
        let block_id = (term_ord as usize) / BLOCK_LEN;
        let buffer = self.block_meta_source.as_slice();
        let mut block_data: &[u8] = &buffer[block_id * TermInfoBlockMeta::SIZE_IN_BYTES..];
        let term_info_block_data = TermInfoBlockMeta::deserialize(&mut block_data)
            .expect("Failed to deserialize terminfoblockmeta");
        let inner_offset = (term_ord as usize) % BLOCK_LEN;
        if inner_offset == 0 {
            term_info_block_data.ref_term_info
        } else {
            let term_info_data = self.term_info_source.as_slice();
            term_info_block_data.deserialize_term_info(
                &term_info_data[term_info_block_data.offset as usize..],
                inner_offset - 1,
            )
        }
    }

    pub fn num_terms(&self) -> usize {
        self.num_terms
    }
}

pub struct TermInfoStoreWriter {
    buffer_block_metas: Vec<u8>,
    buffer_term_infos: Vec<u8>,
    term_infos: Vec<TermInfo>,
    num_terms: u64,
}

fn bitpack_serialize<W: Write>(
    write: &mut W,
    bit_packer: &mut BitPacker,
    term_info_block_meta: &TermInfoBlockMeta,
    term_info: &TermInfo,
) -> io::Result<()> {
    bit_packer.write(
        u64::from(term_info.doc_freq),
        term_info_block_meta.doc_freq_nbits,
        write,
    )?;
    bit_packer.write(
        term_info.postings_offset,
        term_info_block_meta.postings_offset_nbits,
        write,
    )?;
    bit_packer.write(
        term_info.positions_idx,
        term_info_block_meta.positions_idx_nbits,
        write,
    )?;
    Ok(())
}

impl TermInfoStoreWriter {
    pub fn new() -> TermInfoStoreWriter {
        TermInfoStoreWriter {
            buffer_block_metas: Vec::new(),
            buffer_term_infos: Vec::new(),
            term_infos: Vec::with_capacity(BLOCK_LEN),
            num_terms: 0u64,
        }
    }

    fn flush_block(&mut self) -> io::Result<()> {
        if self.term_infos.is_empty() {
            return Ok(());
        }
        let mut bit_packer = BitPacker::new();
        let ref_term_info = self.term_infos[0].clone();
        for term_info in &mut self.term_infos[1..] {
            term_info.postings_offset -= ref_term_info.postings_offset;
            term_info.positions_idx -= ref_term_info.positions_idx;
        }

        let mut max_doc_freq: u32 = 0u32;
        let mut max_postings_offset: u64 = 0u64;
        let mut max_positions_idx: u64 = 0u64;
        for term_info in &self.term_infos[1..] {
            max_doc_freq = cmp::max(max_doc_freq, term_info.doc_freq);
            max_postings_offset = cmp::max(max_postings_offset, term_info.postings_offset);
            max_positions_idx = cmp::max(max_positions_idx, term_info.positions_idx);
        }

        let max_doc_freq_nbits: u8 = compute_num_bits(u64::from(max_doc_freq));
        let max_postings_offset_nbits = compute_num_bits(max_postings_offset);
        let max_positions_idx_nbits = compute_num_bits(max_positions_idx);

        let term_info_block_meta = TermInfoBlockMeta {
            offset: self.buffer_term_infos.len() as u64,
            ref_term_info,
            doc_freq_nbits: max_doc_freq_nbits,
            postings_offset_nbits: max_postings_offset_nbits,
            positions_idx_nbits: max_positions_idx_nbits,
        };

        term_info_block_meta.serialize(&mut self.buffer_block_metas)?;
        for term_info in self.term_infos[1..].iter().cloned() {
            bitpack_serialize(
                &mut self.buffer_term_infos,
                &mut bit_packer,
                &term_info_block_meta,
                &term_info,
            )?;
        }

        // Block need end up at the end of a byte.
        bit_packer.flush(&mut self.buffer_term_infos)?;
        self.term_infos.clear();

        Ok(())
    }

    pub fn write_term_info(&mut self, term_info: &TermInfo) -> io::Result<()> {
        self.num_terms += 1u64;
        self.term_infos.push(term_info.clone());
        if self.term_infos.len() >= BLOCK_LEN {
            self.flush_block()?;
        }
        Ok(())
    }

    pub fn serialize<W: io::Write>(&mut self, write: &mut W) -> io::Result<()> {
        if !self.term_infos.is_empty() {
            self.flush_block()?;
        }
        let len = self.buffer_block_metas.len() as u64;
        len.serialize(write)?;
        self.num_terms.serialize(write)?;
        write.write_all(&self.buffer_block_metas)?;
        write.write_all(&self.buffer_term_infos)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {

    use super::extract_bits;
    use super::TermInfoBlockMeta;
    use super::{TermInfoStore, TermInfoStoreWriter};
    use common;
    use common::bitpacker::BitPacker;
    use common::compute_num_bits;
    use common::BinarySerializable;
    use directory::ReadOnlySource;
    use postings::TermInfo;

    #[test]
    fn test_term_info_block() {
        common::test::fixed_size_test::<TermInfoBlockMeta>();
    }

    #[test]
    fn test_bitpacked() {
        let mut buffer = Vec::new();
        let mut bitpack = BitPacker::new();
        bitpack.write(321u64, 9, &mut buffer).unwrap();
        assert_eq!(compute_num_bits(321u64), 9);
        bitpack.write(2u64, 2, &mut buffer).unwrap();
        assert_eq!(compute_num_bits(2u64), 2);
        bitpack.write(51, 6, &mut buffer).unwrap();
        assert_eq!(compute_num_bits(51), 6);
        bitpack.close(&mut buffer).unwrap();
        assert_eq!(buffer.len(), 3 + 7);
        assert_eq!(extract_bits(&buffer[..], 0, 9), 321u64);
        assert_eq!(extract_bits(&buffer[..], 9, 2), 2u64);
        assert_eq!(extract_bits(&buffer[..], 11, 6), 51u64);
    }

    #[test]
    fn test_term_info_block_meta_serialization() {
        let term_info_block_meta = TermInfoBlockMeta {
            offset: 2009,
            ref_term_info: TermInfo {
                doc_freq: 512,
                postings_offset: 51,
                positions_idx: 3584,
            },
            doc_freq_nbits: 10,
            postings_offset_nbits: 5,
            positions_idx_nbits: 11,
        };
        let mut buffer: Vec<u8> = Vec::new();
        term_info_block_meta.serialize(&mut buffer).unwrap();
        let mut cursor: &[u8] = &buffer[..];
        let term_info_block_meta_serde = TermInfoBlockMeta::deserialize(&mut cursor).unwrap();
        assert_eq!(term_info_block_meta_serde, term_info_block_meta);
    }

    #[test]
    fn test_pack() {
        let mut store_writer = TermInfoStoreWriter::new();
        let mut term_infos = vec![];
        for i in 0..1000 {
            let term_info = TermInfo {
                doc_freq: i as u32,
                postings_offset: (i / 10) as u64,
                positions_idx: (i * 7) as u64,
            };
            store_writer.write_term_info(&term_info).unwrap();
            term_infos.push(term_info);
        }
        let mut buffer = Vec::new();
        store_writer.serialize(&mut buffer).unwrap();
        let term_info_store = TermInfoStore::open(&ReadOnlySource::from(buffer));
        for i in 0..1000 {
            assert_eq!(term_info_store.get(i as u64), term_infos[i]);
        }
    }

}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::prelude::*;
use std::collections::HashMap;

type Layer = HashMap<(usize, usize), u8>;

fn get_input() -> std::io::Result<String> {
  let mut file = File::open("08.txt")?;
  let mut contents = String::new();
  file.read_to_string(&mut contents)?;

  Ok(contents)
}

fn parse_layers(input: &str, width: usize, height: usize) -> Vec<Layer> {
  let mut layers = vec![];
  let mut digits = input.chars()
    .filter(|c| c.is_ascii_digit())
    .filter_map(|c| c.to_digit(10).map(|d| d as u8));

  loop {
    let mut layer = Layer::new();

    for y in 0..height {
      for x in 0..width {
        if let Some(d) = digits.next() {
          layer.insert((x, y), d);
        } else {
          return layers;
        }
      }
    }

    layers.push(layer);
  }
}

fn merge_layers(layers: &Vec<Layer>, width: usize, height: usize) -> Layer {
  let mut result = Layer::new();

  for y in 0..height {
    for x in 0..width {
      result.insert(
        (x, y),
        layers.iter()
        .filter_map(|layer| layer.get(&(x, y)).cloned())
        .filter(|&pixel| pixel != 2)
        .next()
        .unwrap_or(2)
      );
    }
  }

  result
}

fn render(layer: Layer, width: usize, height: usize) -> String {
  let mut result = String::new();

  for y in 0..height {
    for x in 0..width {
      let output = match layer.get(&(x, y)) {
        Some(&0) => ' ',
        _ => '#'
      };

      result.push(output);
    }

    result.push('\n');
  }

  result
}

fn main() {
  let input = get_input().unwrap();
  let width = 25;
  let height = 6;
  let layers = parse_layers(&input, width, height);

  let (_, ones, twos) = layers.iter()
    .map(|layer| {
      layer.iter()
      .filter_map(|(_, &pixel)| match pixel {
        0 => Some((1, 0, 0)),
        1 => Some((0, 1, 0)),
        2 => Some((0, 0, 1)),
        _ => None
      })
      .fold((0, 0, 0), |mut acc, (a, b, c)| {
        acc.0 += a;
        acc.1 += b;
        acc.2 += c;
        acc
      })
    })
    .min_by_key(|&(zeros, _, _)| zeros)
    .unwrap();

  println!("Part 1: {}", ones * twos);

  let merged_layer = merge_layers(&layers, width, height);

  println!("Part 2:\n{}", render(merged_layer, width, height));
}
<commit_msg>Introduce LayerIter<T><commit_after>use std::fs::File;
use std::io::prelude::*;
use std::collections::HashMap;

type Layer = HashMap<(usize, usize), u8>;

struct LayerIter<T: Iterator<Item = u8>> {
  pixels: T,
  width: usize,
  height: usize
}

impl<T: Iterator<Item = u8>> Iterator for LayerIter<T> {
  type Item = Layer;

  fn next(&mut self) -> Option<Self::Item> {
    let mut layer = Layer::new();

    for y in 0..self.height {
      for x in 0..self.width {
        if let Some(pixel) = self.pixels.next() {
          layer.insert((x, y), pixel);
        } else {
          return None;
        }
      }
    }

    Some(layer)
  }
}

fn get_input() -> std::io::Result<String> {
  let mut file = File::open("08.txt")?;
  let mut contents = String::new();
  file.read_to_string(&mut contents)?;

  Ok(contents)
}

fn parse_layers(input: &str, width: usize, height: usize) -> LayerIter<impl Iterator<Item = u8> + '_> {
  let pixels = input.chars()
    .filter(|c| c.is_ascii_digit())
    .filter_map(|c| c.to_digit(10).map(|d| d as u8));

  LayerIter {
    pixels,
    width,
    height
  }
}

fn merge_layers<T: Iterator<Item = Layer>>(layers: T) -> Layer {
  layers
  .fold(Layer::new(), |mut merged_layer, layer| {
    for (&position, &pixel) in layer.iter() {
      if !merged_layer.contains_key(&position) && pixel != 2 {
        merged_layer.insert(position, pixel);
      }
    }

    merged_layer
  })
}

fn render(layer: Layer, width: usize, height: usize) -> String {
  let mut result = String::new();

  for y in 0..height {
    for x in 0..width {
      let output = match layer.get(&(x, y)) {
        Some(&0) => ' ',
        _ => '#'
      };

      result.push(output);
    }

    result.push('\n');
  }

  result
}

fn main() {
  let input = get_input().unwrap();
  let width = 25;
  let height = 6;

  let (_, ones, twos) = parse_layers(&input, width, height)
    .map(|layer| {
      layer.iter()
      .filter_map(|(_, &pixel)| match pixel {
        0 => Some((1, 0, 0)),
        1 => Some((0, 1, 0)),
        2 => Some((0, 0, 1)),
        _ => None
      })
      .fold((0, 0, 0), |mut acc, (a, b, c)| {
        acc.0 += a;
        acc.1 += b;
        acc.2 += c;
        acc
      })
    })
    .min_by_key(|&(zeros, _, _)| zeros)
    .unwrap();

  println!("Part 1: {}", ones * twos);

  let merged_layer = merge_layers(parse_layers(&input, width, height));

  println!("Part 2:\n{}", render(merged_layer, width, height));
}
<|endoftext|>


<commit_before>// Test case for issue #2843.
//

proto! streamp (
    open:send<T: send> {
        data(T) -> open<T>
    }
)

fn rendezvous() {
    let (c, s) = streamp::init();
    let streams: ~[streamp::client::open<int>] = ~[c];

    error!("%?", streams[0]);
}

fn main(args: ~[~str]) {
    //os::getenv("FOO");
    rendezvous();
}<commit_msg>Uppercase Send bound and make test for #2834 work<commit_after>// Test case for issue #2843.
//

proto! streamp (
    open:send<T: Send> {
        data(T) -> open<T>
    }
)

fn rendezvous() {
    let (c, s) = streamp::init();
    let streams: ~[streamp::client::open<int>] = ~[c];

    error!("%?", streams[0]);
}

fn main(args: ~[~str]) {
    //os::getenv("FOO");
    rendezvous();
}<|endoftext|>


<commit_before>// Copyright 2012-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// ignore-android needs extra network permissions
// ignore-bitrig system ulimit (Too many open files)
// ignore-netbsd system ulimit (Too many open files)
// ignore-openbsd system ulimit (Too many open files)
// ignore-emscripten no threads or sockets support

use std::io::prelude::*;
use std::net::{TcpListener, TcpStream};
use std::process;
use std::sync::mpsc::channel;
use std::time::Duration;
use std::thread::{self, Builder};

fn main() {
    // This test has a chance to time out, try to not let it time out
    thread::spawn(move|| -> () {
        thread::sleep(Duration::from_secs(30));
        process::exit(1);
    });

    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap();
    thread::spawn(move || -> () {
        loop {
            let mut stream = match listener.accept() {
                Ok(stream) => stream.0,
                Err(_) => continue,
            };
            let _ = stream.read(&mut [0]);
            let _ = stream.write(&[2]);
        }
    });

    let (tx, rx) = channel();
    let mut spawned_cnt = 0;
    for _ in 0..1000 {
        let tx = tx.clone();
        let res = Builder::new().stack_size(64 * 1024).spawn(move|| {
            match TcpStream::connect(addr) {
                Ok(mut stream) => {
                    let _ = stream.write(&[1]);
                    let _ = stream.read(&mut [0]);
                },
                Err(..) => {}
            }
            tx.send(()).unwrap();
        });
        if let Ok(_) = res {
            spawned_cnt += 1;
        };
    }

    // Wait for all clients to exit, but don't wait for the server to exit. The
    // server just runs infinitely.
    drop(tx);
    for _ in 0..spawned_cnt {
        rx.recv().unwrap();
    }
    assert_eq!(spawned_cnt, 1000);
    process::exit(0);
}
<commit_msg>tcp-stress-test: Factor out thread count as constant.<commit_after>// Copyright 2012-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// ignore-android needs extra network permissions
// ignore-bitrig system ulimit (Too many open files)
// ignore-netbsd system ulimit (Too many open files)
// ignore-openbsd system ulimit (Too many open files)
// ignore-emscripten no threads or sockets support

use std::io::prelude::*;
use std::net::{TcpListener, TcpStream};
use std::process;
use std::sync::mpsc::channel;
use std::time::Duration;
use std::thread::{self, Builder};

const TARGET_CNT: usize = 1000;

fn main() {
    // This test has a chance to time out, try to not let it time out
    thread::spawn(move|| -> () {
        thread::sleep(Duration::from_secs(30));
        process::exit(1);
    });

    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap();
    thread::spawn(move || -> () {
        loop {
            let mut stream = match listener.accept() {
                Ok(stream) => stream.0,
                Err(_) => continue,
            };
            let _ = stream.read(&mut [0]);
            let _ = stream.write(&[2]);
        }
    });

    let (tx, rx) = channel();

    let mut spawned_cnt = 0;
    for _ in 0..TARGET_CNT {
        let tx = tx.clone();
        let res = Builder::new().stack_size(64 * 1024).spawn(move|| {
            match TcpStream::connect(addr) {
                Ok(mut stream) => {
                    let _ = stream.write(&[1]);
                    let _ = stream.read(&mut [0]);
                },
                Err(..) => {}
            }
            tx.send(()).unwrap();
        });
        if let Ok(_) = res {
            spawned_cnt += 1;
        };
    }

    // Wait for all clients to exit, but don't wait for the server to exit. The
    // server just runs infinitely.
    drop(tx);
    for _ in 0..spawned_cnt {
        rx.recv().unwrap();
    }
    assert_eq!(spawned_cnt, TARGET_CNT);
    process::exit(0);
}
<|endoftext|>


<commit_before>use crate::test_utils::TestRandom;
use crate::{AttestationData, Bitfield, Slot};
use rand::RngCore;
use serde_derive::{Deserialize, Serialize};
use ssz_derive::{Decode, Encode, TreeHash};
use test_random_derive::TestRandom;

/// An attestation that has been included in the state but not yet fully processed.
///
/// Spec v0.4.0
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Encode, Decode, TreeHash, TestRandom)]
pub struct PendingAttestation {
    pub aggregation_bitfield: Bitfield,
    pub data: AttestationData,
    pub custody_bitfield: Bitfield,
    pub inclusion_slot: Slot,
}

#[cfg(test)]
mod tests {
    use super::*;

    ssz_tests!(PendingAttestation);
}
<commit_msg>Fast-forward PendingAttestation to v0.5.0<commit_after>use crate::test_utils::TestRandom;
use crate::{AttestationData, Bitfield, Slot};
use rand::RngCore;
use serde_derive::{Deserialize, Serialize};
use ssz_derive::{Decode, Encode, TreeHash};
use test_random_derive::TestRandom;

/// An attestation that has been included in the state but not yet fully processed.
///
/// Spec v0.5.0
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Encode, Decode, TreeHash, TestRandom)]
pub struct PendingAttestation {
    pub aggregation_bitfield: Bitfield,
    pub data: AttestationData,
    pub custody_bitfield: Bitfield,
    pub inclusion_slot: Slot,
}

#[cfg(test)]
mod tests {
    use super::*;

    ssz_tests!(PendingAttestation);
}
<|endoftext|>


<commit_before>use ide_db::{
    assists::{AssistId, AssistKind},
    famous_defs::FamousDefs,
};
use syntax::{
    ast::{self, make, Expr, HasArgList},
    AstNode,
};

use crate::{AssistContext, Assists};

// Assist: replace_or_with_or_else
//
// Replace `unwrap_or` with `unwrap_or_else` and `ok_or` with `ok_or_else`.
//
// ```
// # //- minicore:option
// fn foo() {
//     let a = Some(1);
//     a.unwra$0p_or(2);
// }
// ```
// ->
// ```
// fn foo() {
//     let a = Some(1);
//     a.unwrap_or_else(|| 2);
// }
// ```
pub(crate) fn replace_or_with_or_else(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    let call: ast::MethodCallExpr = ctx.find_node_at_offset()?;

    is_option_or_result(call.receiver()?, ctx)?;

    let (name, arg_list) = (call.name_ref()?, call.arg_list()?);

    let replace = match &*name.text() {
        "unwrap_or" => "unwrap_or_else".to_string(),
        "ok_or" => "ok_or_else".to_string(),
        _ => return None,
    };

    let arg = match arg_list.args().collect::<Vec<_>>().as_slice() {
        [] => make::arg_list(Vec::new()),
        [first] => {
            let param = (|| {
                if let ast::Expr::CallExpr(call) = first {
                    if call.arg_list()?.args().count() == 0 {
                        Some(call.expr()?.clone())
                    } else {
                        None
                    }
                } else {
                    None
                }
            })()
            .unwrap_or_else(|| make::expr_closure(None, first.clone()));
            make::arg_list(vec![param])
        }
        _ => return None,
    };

    acc.add(
        AssistId("replace_or_with_or_else", AssistKind::RefactorRewrite),
        "Replace unwrap_or or ok_or with lazy version",
        call.syntax().text_range(),
        |builder| {
            builder.replace(name.syntax().text_range(), replace);
            builder.replace_ast(arg_list, arg)
        },
    )
}

// Assist: replace_or_else_with_or
//
// Replace `unwrap_or_else` with `unwrap_or` and `ok_or_else` with `ok_or`.
//
// ```
// # //- minicore:option
// fn foo() {
//     let a = Some(1);
//     a.unwra$0p_or_else(|| 2);
// }
// ```
// ->
// ```
// fn foo() {
//     let a = Some(1);
//     a.unwrap_or(2);
// }
// ```
pub(crate) fn replace_or_else_with_or(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    let call: ast::MethodCallExpr = ctx.find_node_at_offset()?;

    is_option_or_result(call.receiver()?, ctx)?;

    let (name, arg_list) = (call.name_ref()?, call.arg_list()?);

    let replace = match &*name.text() {
        "unwrap_or_else" => "unwrap_or".to_string(),
        "ok_or_else" => "ok_or".to_string(),
        _ => return None,
    };

    let arg = match arg_list.args().collect::<Vec<_>>().as_slice() {
        [] => make::arg_list(Vec::new()),
        [first] => {
            let param = (|| {
                if let ast::Expr::ClosureExpr(closure) = first {
                    if closure.param_list()?.params().count() == 0 {
                        Some(closure.body()?.clone())
                    } else {
                        None
                    }
                } else {
                    None
                }
            })()
            .unwrap_or_else(|| make::expr_call(first.clone(), make::arg_list(Vec::new())));
            make::arg_list(vec![param])
        }
        _ => return None,
    };

    acc.add(
        AssistId("replace_or_else_with_or", AssistKind::RefactorRewrite),
        "Replace unwrap_or_else or ok_or_else with eager version",
        call.syntax().text_range(),
        |builder| {
            builder.replace(name.syntax().text_range(), replace);
            builder.replace_ast(arg_list, arg)
        },
    )
}

fn is_option_or_result(receiver: Expr, ctx: &AssistContext<'_>) -> Option<()> {
    let ty = ctx.sema.type_of_expr(&receiver)?.adjusted().as_adt()?.as_enum()?;
    let option_enum =
        FamousDefs(&ctx.sema, ctx.sema.scope(receiver.syntax())?.krate()).core_option_Option();

    if let Some(option_enum) = option_enum {
        if ty == option_enum {
            return Some(());
        }
    }

    let result_enum =
        FamousDefs(&ctx.sema, ctx.sema.scope(receiver.syntax())?.krate()).core_result_Result();

    if let Some(result_enum) = result_enum {
        if ty == result_enum {
            return Some(());
        }
    }

    None
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn replace_or_with_or_else_simple() {
        check_assist(
            replace_or_with_or_else,
            r#"
//- minicore: option
fn foo() {
    let foo = Some(1);
    return foo.unwrap_$0or(2);
}
"#,
            r#"
fn foo() {
    let foo = Some(1);
    return foo.unwrap_or_else(|| 2);
}
"#,
        )
    }

    #[test]
    fn replace_or_with_or_else_call() {
        check_assist(
            replace_or_with_or_else,
            r#"
//- minicore: option
fn foo() {
    let foo = Some(1);
    return foo.unwrap_$0or(x());
}
"#,
            r#"
fn foo() {
    let foo = Some(1);
    return foo.unwrap_or_else(x);
}
"#,
        )
    }

    #[test]
    fn replace_or_with_or_else_block() {
        check_assist(
            replace_or_with_or_else,
            r#"
//- minicore: option
fn foo() {
    let foo = Some(1);
    return foo.unwrap_$0or({
        let mut x = bar();
        for i in 0..10 {
            x += i;
        }
        x
    });
}
"#,
            r#"
fn foo() {
    let foo = Some(1);
    return foo.unwrap_or_else(|| {
        let mut x = bar();
        for i in 0..10 {
            x += i;
        }
        x
    });
}
"#,
        )
    }

    #[test]
    fn replace_or_else_with_or_simple() {
        check_assist(
            replace_or_else_with_or,
            r#"
//- minicore: option
fn foo() {
    let foo = Some(1);
    return foo.unwrap_$0or_else(|| 2);
}
"#,
            r#"
fn foo() {
    let foo = Some(1);
    return foo.unwrap_or(2);
}
"#,
        )
    }

    #[test]
    fn replace_or_else_with_or_call() {
        check_assist(
            replace_or_else_with_or,
            r#"
//- minicore: option
fn foo() {
    let foo = Some(1);
    return foo.unwrap_$0or_else(x);
}
"#,
            r#"
fn foo() {
    let foo = Some(1);
    return foo.unwrap_or(x());
}
"#,
        )
    }

    #[test]
    fn replace_or_else_with_or_result() {
        check_assist(
            replace_or_else_with_or,
            r#"
//- minicore: result
fn foo() {
    let foo = Ok(1);
    return foo.unwrap_$0or_else(x);
}
"#,
            r#"
fn foo() {
    let foo = Ok(1);
    return foo.unwrap_or(x());
}
"#,
        )
    }

    #[test]
    fn replace_or_else_with_or_not_applicable() {
        check_assist_not_applicable(
            replace_or_else_with_or,
            r#"
fn foo() {
    let foo = Ok(1);
    return foo.unwrap_$0or_else(x);
}
"#,
        )
    }
}
<commit_msg>change as requested<commit_after>use ide_db::{
    assists::{AssistId, AssistKind},
    famous_defs::FamousDefs,
};
use syntax::{
    ast::{self, make, Expr, HasArgList},
    AstNode,
};

use crate::{AssistContext, Assists};

// Assist: replace_or_with_or_else
//
// Replace `unwrap_or` with `unwrap_or_else` and `ok_or` with `ok_or_else`.
//
// ```
// # //- minicore:option
// fn foo() {
//     let a = Some(1);
//     a.unwra$0p_or(2);
// }
// ```
// ->
// ```
// fn foo() {
//     let a = Some(1);
//     a.unwrap_or_else(|| 2);
// }
// ```
pub(crate) fn replace_or_with_or_else(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    let call: ast::MethodCallExpr = ctx.find_node_at_offset()?;

    let kind = is_option_or_result(call.receiver()?, ctx)?;

    let (name, arg_list) = (call.name_ref()?, call.arg_list()?);

    let mut map_or = false;

    let replace = match &*name.text() {
        "unwrap_or" => "unwrap_or_else".to_string(),
        "or" => "or_else".to_string(),
        "ok_or" if kind == Kind::Option => "ok_or_else".to_string(),
        "map_or" => {
            map_or = true;
            "map_or_else".to_string()
        }
        _ => return None,
    };

    let arg = match arg_list.args().collect::<Vec<_>>().as_slice() {
        [] => make::arg_list(Vec::new()),
        [first] => {
            let param = into_closure(first);
            make::arg_list(vec![param])
        }
        [first, second] if map_or => {
            let param = into_closure(first);
            make::arg_list(vec![param, second.clone()])
        }
        _ => return None,
    };

    acc.add(
        AssistId("replace_or_with_or_else", AssistKind::RefactorRewrite),
        "Replace unwrap_or or ok_or with lazy version",
        call.syntax().text_range(),
        |builder| {
            builder.replace(name.syntax().text_range(), replace);
            builder.replace_ast(arg_list, arg)
        },
    )
}

fn into_closure(param: &Expr) -> Expr {
    (|| {
        if let ast::Expr::CallExpr(call) = param {
            if call.arg_list()?.args().count() == 0 {
                Some(call.expr()?.clone())
            } else {
                None
            }
        } else {
            None
        }
    })()
    .unwrap_or_else(|| make::expr_closure(None, param.clone()))
}

// Assist: replace_or_else_with_or
//
// Replace `unwrap_or_else` with `unwrap_or` and `ok_or_else` with `ok_or`.
//
// ```
// # //- minicore:option
// fn foo() {
//     let a = Some(1);
//     a.unwra$0p_or_else(|| 2);
// }
// ```
// ->
// ```
// fn foo() {
//     let a = Some(1);
//     a.unwrap_or(2);
// }
// ```
pub(crate) fn replace_or_else_with_or(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    let call: ast::MethodCallExpr = ctx.find_node_at_offset()?;

    let kind = is_option_or_result(call.receiver()?, ctx)?;

    let (name, arg_list) = (call.name_ref()?, call.arg_list()?);

    let mut map_or = false;
    let replace = match &*name.text() {
        "unwrap_or_else" => "unwrap_or".to_string(),
        "or_else" => "or".to_string(),
        "ok_or_else" if kind == Kind::Option => "ok_or".to_string(),
        "map_or_else" => {
            map_or = true;
            "map_or".to_string()
        }
        _ => return None,
    };

    let arg = match arg_list.args().collect::<Vec<_>>().as_slice() {
        [] => make::arg_list(Vec::new()),
        [first] => {
            let param = into_call(first);
            make::arg_list(vec![param])
        }
        [first, second] if map_or => {
            let param = into_call(first);
            make::arg_list(vec![param, second.clone()])
        }
        _ => return None,
    };

    acc.add(
        AssistId("replace_or_else_with_or", AssistKind::RefactorRewrite),
        "Replace unwrap_or_else or ok_or_else with eager version",
        call.syntax().text_range(),
        |builder| {
            builder.replace(name.syntax().text_range(), replace);
            builder.replace_ast(arg_list, arg)
        },
    )
}

fn into_call(param: &Expr) -> Expr {
    (|| {
        if let ast::Expr::ClosureExpr(closure) = param {
            if closure.param_list()?.params().count() == 0 {
                Some(closure.body()?.clone())
            } else {
                None
            }
        } else {
            None
        }
    })()
    .unwrap_or_else(|| make::expr_call(param.clone(), make::arg_list(Vec::new())))
}

#[derive(PartialEq, Eq)]
enum Kind {
    Option,
    Result,
}

fn is_option_or_result(receiver: Expr, ctx: &AssistContext<'_>) -> Option<Kind> {
    let ty = ctx.sema.type_of_expr(&receiver)?.adjusted().as_adt()?.as_enum()?;
    let option_enum =
        FamousDefs(&ctx.sema, ctx.sema.scope(receiver.syntax())?.krate()).core_option_Option();

    if let Some(option_enum) = option_enum {
        if ty == option_enum {
            return Some(Kind::Option);
        }
    }

    let result_enum =
        FamousDefs(&ctx.sema, ctx.sema.scope(receiver.syntax())?.krate()).core_result_Result();

    if let Some(result_enum) = result_enum {
        if ty == result_enum {
            return Some(Kind::Result);
        }
    }

    None
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn replace_or_with_or_else_simple() {
        check_assist(
            replace_or_with_or_else,
            r#"
//- minicore: option
fn foo() {
    let foo = Some(1);
    return foo.unwrap_$0or(2);
}
"#,
            r#"
fn foo() {
    let foo = Some(1);
    return foo.unwrap_or_else(|| 2);
}
"#,
        )
    }

    #[test]
    fn replace_or_with_or_else_call() {
        check_assist(
            replace_or_with_or_else,
            r#"
//- minicore: option
fn foo() {
    let foo = Some(1);
    return foo.unwrap_$0or(x());
}
"#,
            r#"
fn foo() {
    let foo = Some(1);
    return foo.unwrap_or_else(x);
}
"#,
        )
    }

    #[test]
    fn replace_or_with_or_else_block() {
        check_assist(
            replace_or_with_or_else,
            r#"
//- minicore: option
fn foo() {
    let foo = Some(1);
    return foo.unwrap_$0or({
        let mut x = bar();
        for i in 0..10 {
            x += i;
        }
        x
    });
}
"#,
            r#"
fn foo() {
    let foo = Some(1);
    return foo.unwrap_or_else(|| {
        let mut x = bar();
        for i in 0..10 {
            x += i;
        }
        x
    });
}
"#,
        )
    }

    #[test]
    fn replace_or_else_with_or_simple() {
        check_assist(
            replace_or_else_with_or,
            r#"
//- minicore: option
fn foo() {
    let foo = Some(1);
    return foo.unwrap_$0or_else(|| 2);
}
"#,
            r#"
fn foo() {
    let foo = Some(1);
    return foo.unwrap_or(2);
}
"#,
        )
    }

    #[test]
    fn replace_or_else_with_or_call() {
        check_assist(
            replace_or_else_with_or,
            r#"
//- minicore: option
fn foo() {
    let foo = Some(1);
    return foo.unwrap_$0or_else(x);
}
"#,
            r#"
fn foo() {
    let foo = Some(1);
    return foo.unwrap_or(x());
}
"#,
        )
    }

    #[test]
    fn replace_or_else_with_or_result() {
        check_assist(
            replace_or_else_with_or,
            r#"
//- minicore: result
fn foo() {
    let foo = Ok(1);
    return foo.unwrap_$0or_else(x);
}
"#,
            r#"
fn foo() {
    let foo = Ok(1);
    return foo.unwrap_or(x());
}
"#,
        )
    }

    #[test]
    fn replace_or_else_with_or_map() {
        check_assist(
            replace_or_else_with_or,
            r#"
//- minicore: result
fn foo() {
    let foo = Ok("foo");
    return foo.map$0_or_else(|| 42, |v| v.len());
}
"#,
            r#"
fn foo() {
    let foo = Ok("foo");
    return foo.map_or(42, |v| v.len());
}
"#,
        )
    }

    #[test]
    fn replace_or_else_with_or_not_applicable() {
        check_assist_not_applicable(
            replace_or_else_with_or,
            r#"
fn foo() {
    let foo = Ok(1);
    return foo.unwrap_$0or_else(x);
}
"#,
        )
    }
}
<|endoftext|>


<commit_before>//! A thin wrapper around `Command` in the standard library which allows us to
//! read the arguments that are built up.

use std::ffi::{OsStr, OsString};
use std::fmt;
use std::io;
use std::mem;
use std::process::{self, Output};

use rustc_span::symbol::Symbol;
use rustc_target::spec::LldFlavor;

#[derive(Clone)]
pub struct Command {
    program: Program,
    args: Vec<OsString>,
    env: Vec<(OsString, OsString)>,
    env_remove: Vec<OsString>,
}

#[derive(Clone)]
enum Program {
    Normal(OsString),
    CmdBatScript(OsString),
    Lld(OsString, LldFlavor),
}

impl Command {
    pub fn new<P: AsRef<OsStr>>(program: P) -> Command {
        Command::_new(Program::Normal(program.as_ref().to_owned()))
    }

    pub fn bat_script<P: AsRef<OsStr>>(program: P) -> Command {
        Command::_new(Program::CmdBatScript(program.as_ref().to_owned()))
    }

    pub fn lld<P: AsRef<OsStr>>(program: P, flavor: LldFlavor) -> Command {
        Command::_new(Program::Lld(program.as_ref().to_owned(), flavor))
    }

    fn _new(program: Program) -> Command {
        Command { program, args: Vec::new(), env: Vec::new(), env_remove: Vec::new() }
    }

    pub fn arg<P: AsRef<OsStr>>(&mut self, arg: P) -> &mut Command {
        self._arg(arg.as_ref());
        self
    }

    pub fn sym_arg(&mut self, arg: Symbol) -> &mut Command {
        self.arg(&*arg.as_str());
        self
    }

    pub fn args<I>(&mut self, args: I) -> &mut Command
    where
        I: IntoIterator<Item: AsRef<OsStr>>,
    {
        for arg in args {
            self._arg(arg.as_ref());
        }
        self
    }

    fn _arg(&mut self, arg: &OsStr) {
        self.args.push(arg.to_owned());
    }

    pub fn env<K, V>(&mut self, key: K, value: V) -> &mut Command
    where
        K: AsRef<OsStr>,
        V: AsRef<OsStr>,
    {
        self._env(key.as_ref(), value.as_ref());
        self
    }

    fn _env(&mut self, key: &OsStr, value: &OsStr) {
        self.env.push((key.to_owned(), value.to_owned()));
    }

    pub fn env_remove<K>(&mut self, key: K) -> &mut Command
    where
        K: AsRef<OsStr>,
    {
        self._env_remove(key.as_ref());
        self
    }

    fn _env_remove(&mut self, key: &OsStr) {
        self.env_remove.push(key.to_owned());
    }

    pub fn output(&mut self) -> io::Result<Output> {
        self.command().output()
    }

    pub fn command(&self) -> process::Command {
        let mut ret = match self.program {
            Program::Normal(ref p) => process::Command::new(p),
            Program::CmdBatScript(ref p) => {
                let mut c = process::Command::new("cmd");
                c.arg("/c").arg(p);
                c
            }
            Program::Lld(ref p, flavor) => {
                let mut c = process::Command::new(p);
                c.arg("-flavor").arg(match flavor {
                    LldFlavor::Wasm => "wasm",
                    LldFlavor::Ld => "gnu",
                    LldFlavor::Link => "link",
                    LldFlavor::Ld64 => "darwin",
                });
                c
            }
        };
        ret.args(&self.args);
        ret.envs(self.env.clone());
        for k in &self.env_remove {
            ret.env_remove(k);
        }
        ret
    }

    // extensions

    pub fn get_args(&self) -> &[OsString] {
        &self.args
    }

    pub fn take_args(&mut self) -> Vec<OsString> {
        mem::take(&mut self.args)
    }

    /// Returns a `true` if we're pretty sure that this'll blow OS spawn limits,
    /// or `false` if we should attempt to spawn and see what the OS says.
    pub fn very_likely_to_exceed_some_spawn_limit(&self) -> bool {
        // We mostly only care about Windows in this method, on Unix the limits
        // can be gargantuan anyway so we're pretty unlikely to hit them
        if cfg!(unix) {
            return false;
        }

        // Right now LLD doesn't support the `@` syntax of passing an argument
        // through files, so regardless of the platform we try to go to the OS
        // on this one.
        if let Program::Lld(..) = self.program {
            return false;
        }

        // Ok so on Windows to spawn a process is 32,768 characters in its
        // command line [1]. Unfortunately we don't actually have access to that
        // as it's calculated just before spawning. Instead we perform a
        // poor-man's guess as to how long our command line will be. We're
        // assuming here that we don't have to escape every character...
        //
        // Turns out though that `cmd.exe` has even smaller limits, 8192
        // characters [2]. Linkers can often be batch scripts (for example
        // Emscripten, Gecko's current build system) which means that we're
        // running through batch scripts. These linkers often just forward
        // arguments elsewhere (and maybe tack on more), so if we blow 8192
        // bytes we'll typically cause them to blow as well.
        //
        // Basically as a result just perform an inflated estimate of what our
        // command line will look like and test if it's > 8192 (we actually
        // test against 6k to artificially inflate our estimate). If all else
        // fails we'll fall back to the normal unix logic of testing the OS
        // error code if we fail to spawn and automatically re-spawning the
        // linker with smaller arguments.
        //
        // [1]: https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-createprocessa
        // [2]: https://devblogs.microsoft.com/oldnewthing/?p=41553

        let estimated_command_line_len = self.args.iter().map(|a| a.len()).sum::<usize>();
        estimated_command_line_len > 1024 * 6
    }
}

impl fmt::Debug for Command {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.command().fmt(f)
    }
}
<commit_msg>Auto merge of #77543 - Mark-Simulacrum:rsp-quoting, r=eddyb<commit_after>//! A thin wrapper around `Command` in the standard library which allows us to
//! read the arguments that are built up.

use std::ffi::{OsStr, OsString};
use std::fmt;
use std::io;
use std::mem;
use std::process::{self, Output};

use rustc_span::symbol::Symbol;
use rustc_target::spec::LldFlavor;

#[derive(Clone)]
pub struct Command {
    program: Program,
    args: Vec<OsString>,
    env: Vec<(OsString, OsString)>,
    env_remove: Vec<OsString>,
}

#[derive(Clone)]
enum Program {
    Normal(OsString),
    CmdBatScript(OsString),
    Lld(OsString, LldFlavor),
}

impl Command {
    pub fn new<P: AsRef<OsStr>>(program: P) -> Command {
        Command::_new(Program::Normal(program.as_ref().to_owned()))
    }

    pub fn bat_script<P: AsRef<OsStr>>(program: P) -> Command {
        Command::_new(Program::CmdBatScript(program.as_ref().to_owned()))
    }

    pub fn lld<P: AsRef<OsStr>>(program: P, flavor: LldFlavor) -> Command {
        Command::_new(Program::Lld(program.as_ref().to_owned(), flavor))
    }

    fn _new(program: Program) -> Command {
        Command { program, args: Vec::new(), env: Vec::new(), env_remove: Vec::new() }
    }

    pub fn arg<P: AsRef<OsStr>>(&mut self, arg: P) -> &mut Command {
        self._arg(arg.as_ref());
        self
    }

    pub fn sym_arg(&mut self, arg: Symbol) -> &mut Command {
        self.arg(&*arg.as_str());
        self
    }

    pub fn args<I>(&mut self, args: I) -> &mut Command
    where
        I: IntoIterator<Item: AsRef<OsStr>>,
    {
        for arg in args {
            self._arg(arg.as_ref());
        }
        self
    }

    fn _arg(&mut self, arg: &OsStr) {
        self.args.push(arg.to_owned());
    }

    pub fn env<K, V>(&mut self, key: K, value: V) -> &mut Command
    where
        K: AsRef<OsStr>,
        V: AsRef<OsStr>,
    {
        self._env(key.as_ref(), value.as_ref());
        self
    }

    fn _env(&mut self, key: &OsStr, value: &OsStr) {
        self.env.push((key.to_owned(), value.to_owned()));
    }

    pub fn env_remove<K>(&mut self, key: K) -> &mut Command
    where
        K: AsRef<OsStr>,
    {
        self._env_remove(key.as_ref());
        self
    }

    fn _env_remove(&mut self, key: &OsStr) {
        self.env_remove.push(key.to_owned());
    }

    pub fn output(&mut self) -> io::Result<Output> {
        self.command().output()
    }

    pub fn command(&self) -> process::Command {
        let mut ret = match self.program {
            Program::Normal(ref p) => process::Command::new(p),
            Program::CmdBatScript(ref p) => {
                let mut c = process::Command::new("cmd");
                c.arg("/c").arg(p);
                c
            }
            Program::Lld(ref p, flavor) => {
                let mut c = process::Command::new(p);
                c.arg("-flavor").arg(match flavor {
                    LldFlavor::Wasm => "wasm",
                    LldFlavor::Ld => "gnu",
                    LldFlavor::Link => "link",
                    LldFlavor::Ld64 => "darwin",
                });
                if let LldFlavor::Wasm = flavor {
                    // LLVM expects host-specific formatting for @file
                    // arguments, but we always generate posix formatted files
                    // at this time. Indicate as such.
                    c.arg("--rsp-quoting=posix");
                }
                c
            }
        };
        ret.args(&self.args);
        ret.envs(self.env.clone());
        for k in &self.env_remove {
            ret.env_remove(k);
        }
        ret
    }

    // extensions

    pub fn get_args(&self) -> &[OsString] {
        &self.args
    }

    pub fn take_args(&mut self) -> Vec<OsString> {
        mem::take(&mut self.args)
    }

    /// Returns a `true` if we're pretty sure that this'll blow OS spawn limits,
    /// or `false` if we should attempt to spawn and see what the OS says.
    pub fn very_likely_to_exceed_some_spawn_limit(&self) -> bool {
        // We mostly only care about Windows in this method, on Unix the limits
        // can be gargantuan anyway so we're pretty unlikely to hit them
        if cfg!(unix) {
            return false;
        }

        // Right now LLD doesn't support the `@` syntax of passing an argument
        // through files, so regardless of the platform we try to go to the OS
        // on this one.
        if let Program::Lld(..) = self.program {
            return false;
        }

        // Ok so on Windows to spawn a process is 32,768 characters in its
        // command line [1]. Unfortunately we don't actually have access to that
        // as it's calculated just before spawning. Instead we perform a
        // poor-man's guess as to how long our command line will be. We're
        // assuming here that we don't have to escape every character...
        //
        // Turns out though that `cmd.exe` has even smaller limits, 8192
        // characters [2]. Linkers can often be batch scripts (for example
        // Emscripten, Gecko's current build system) which means that we're
        // running through batch scripts. These linkers often just forward
        // arguments elsewhere (and maybe tack on more), so if we blow 8192
        // bytes we'll typically cause them to blow as well.
        //
        // Basically as a result just perform an inflated estimate of what our
        // command line will look like and test if it's > 8192 (we actually
        // test against 6k to artificially inflate our estimate). If all else
        // fails we'll fall back to the normal unix logic of testing the OS
        // error code if we fail to spawn and automatically re-spawning the
        // linker with smaller arguments.
        //
        // [1]: https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-createprocessa
        // [2]: https://devblogs.microsoft.com/oldnewthing/?p=41553

        let estimated_command_line_len = self.args.iter().map(|a| a.len()).sum::<usize>();
        estimated_command_line_len > 1024 * 6
    }
}

impl fmt::Debug for Command {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.command().fmt(f)
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::io::prelude::*;

struct Triplets {
    column1: Vec<u32>,
    column2: Vec<u32>,
    column3: Vec<u32>,
    triangles: u32,
}

impl Triplets {
    fn new() -> Triplets {
        Triplets {
            column1: Vec::with_capacity(3),
            column2: Vec::with_capacity(3),
            column3: Vec::with_capacity(3),
            triangles: 0,
        }
    }

    fn push_triplet(&mut self, a: u32, b: u32, c: u32) {
        let valid_triangle =
            |ref v: &Vec<u32>| v[0] + v[1] > v[2] && v[0] + v[2] > v[1] && v[1] + v[2] > v[0];
        for &mut (ref mut column, n) in &mut [(&mut self.column1, a),
                                              (&mut self.column2, b),
                                              (&mut self.column3, c)] {
            column.push(n);
            if column.len() == 3 {
                if valid_triangle(&column) {
                    self.triangles += 1;
                }
                column.clear();
            }
        }
    }
}

fn main() {
    let stdin = io::stdin();
    let stdin = stdin.lock().lines();
    let valid_triangle = |a, b, c| a + b > c && a + c > b && b + c > a;
    let mut part2 = Triplets::new();
    let mut count = 0;
    for line in stdin {
        let line = line.unwrap();
        let mut sides = line.split_whitespace();
        let a = sides.next().unwrap().parse::<u32>().unwrap();
        let b = sides.next().unwrap().parse::<u32>().unwrap();
        let c = sides.next().unwrap().parse::<u32>().unwrap();
        if valid_triangle(a, b, c) {
            count += 1;
        }
        part2.push_triplet(a, b, c);
    }
    println!("part 1: {}", count);
    println!("part 2: {}", part2.triangles);
}
<commit_msg>avoid needless borrows<commit_after>use std::io;
use std::io::prelude::*;

struct Triplets {
    column1: Vec<u32>,
    column2: Vec<u32>,
    column3: Vec<u32>,
    triangles: u32,
}

impl Triplets {
    fn new() -> Triplets {
        Triplets {
            column1: Vec::with_capacity(3),
            column2: Vec::with_capacity(3),
            column3: Vec::with_capacity(3),
            triangles: 0,
        }
    }

    fn push_triplet(&mut self, a: u32, b: u32, c: u32) {
        let valid_triangle =
            |v: &Vec<u32>| v[0] + v[1] > v[2] && v[0] + v[2] > v[1] && v[1] + v[2] > v[0];
        for &mut (ref mut column, n) in &mut [(&mut self.column1, a),
                                              (&mut self.column2, b),
                                              (&mut self.column3, c)] {
            column.push(n);
            if column.len() == 3 {
                if valid_triangle(column) {
                    self.triangles += 1;
                }
                column.clear();
            }
        }
    }
}

fn main() {
    let stdin = io::stdin();
    let stdin = stdin.lock().lines();
    let valid_triangle = |a, b, c| a + b > c && a + c > b && b + c > a;
    let mut part2 = Triplets::new();
    let mut count = 0;
    for line in stdin {
        let line = line.unwrap();
        let mut sides = line.split_whitespace();
        let a = sides.next().unwrap().parse::<u32>().unwrap();
        let b = sides.next().unwrap().parse::<u32>().unwrap();
        let c = sides.next().unwrap().parse::<u32>().unwrap();
        if valid_triangle(a, b, c) {
            count += 1;
        }
        part2.push_triplet(a, b, c);
    }
    println!("part 1: {}", count);
    println!("part 2: {}", part2.triangles);
}
<|endoftext|>


<commit_before>use std::process::{Command, Stdio};
use std::io::Cursor;
use std::sync::Arc;
use std::time::Duration;
use std::os::unix::process::ExitStatusExt;

use tokio::prelude::*;
use tokio::prelude::future::Either;
use tokio::timer::timeout;
use tokio_process::CommandExt;
use tokio::{io::{flush, read_exact, write_all}, net::unix::UnixStream};
use bytes::{BytesMut, Buf, BufMut};

use crate::{ExecBackend, UnixSocketBackend};

fn strsig(sig: i32) -> &'static str {
    match sig {
        1 => "Hangup",
        2 => "Interrupt",
        3 => "Quit",
        4 => "Illegal instruction",
        5 => "Trace/breakpoint trap",
        6 => "Aborted",
        7 => "Bus error",
        8 => "Floating point exception",
        9 => "Killed",
        10 => "User defined signal 1",
        11 => "Segmentation fault",
        12 => "User defined signal 2",
        13 => "Broken pipe",
        14 => "Alarm clock",
        15 => "Terminated",
        16 => "Stack fault",
        17 => "Child exited",
        18 => "Continued",
        19 => "Stopped (signal)",
        20 => "Stopped",
        21 => "Stopped (tty input)",
        22 => "Stopped (tty output)",
        23 => "Urgent I/O condition",
        24 => "CPU time limit exceeded",
        25 => "File size limit exceeded",
        26 => "Virtual timer expired",
        27 => "Profiling timer expired",
        28 => "Window changed",
        29 => "I/O possible",
        30 => "Power failure",
        31 => "Bad system call",
        _ => "Unknown signal"
    }
}

fn strsigabbrev(sig: i32) -> &'static str {
    match sig {
        1 => "SIGHUP",
        2 => "SIGINT",
        3 => "SIGQUIT",
        4 => "SIGILL",
        5 => "SIGTRAP",
        6 => "SIGABRT",
        7 => "SIGBUS",
        8 => "SIGFPE",
        9 => "SIGKILL",
        10 => "SIGUSR1",
        11 => "SIGSEGV",
        12 => "SIGUSR2",
        13 => "SIGPIPE",
        14 => "SIGALRM",
        15 => "SIGTERM",
        16 => "SIGSTKFLT",
        17 => "SIGCHLD",
        18 => "SIGCONT",
        19 => "SIGSTOP",
        20 => "SIGTSTP",
        21 => "SIGTTIN",
        22 => "SIGTTOU",
        23 => "SIGURG",
        24 => "SIGXCPU",
        25 => "SIGXFSZ",
        26 => "SIGVTALRM",
        27 => "SIGPROF",
        28 => "SIGWINCH",
        29 => "SIGPOLL",
        30 => "SIGPWR",
        31 => "SIGSYS",
        _ => "(unknown)"
    }
}

pub fn exec<'a, T>(
    lang: Arc<ExecBackend>,
    timeout: Option<usize>,
    code: T) -> impl Future<Item = String, Error = String> + 'a
        where T: AsRef<[u8]> + 'a {
    let timeout_arg = timeout
        .map(|t| format!("{}{}", lang.timeout_prefix.as_ref().map(String::as_str).unwrap_or(""), t));
    let timeout_arg_ref = timeout_arg.as_ref().map(String::as_str);
    if let Some(path) = lang.cmdline.iter().nth(0) {
        let mut cmd = Command::new(path);
        cmd.args(lang.cmdline.iter()
            .skip(1)
            .filter_map(|a| if a == "{TIMEOUT}" {
                timeout_arg_ref
            } else {
                Some(a.as_ref())
            }))
            .stdin(Stdio::piped()).stdout(Stdio::piped()).stderr(Stdio::piped());
        debug!("spawning {:?}", cmd);
        Either::A(cmd.spawn_async()
            .map_err(|e| format!("failed to exec: {}", e))
            .into_future()
            .and_then(|mut child| {
                child.stdin().take().ok_or_else(|| "stdin missing".to_owned()).into_future()
                    .and_then(|stdin|
                        write_all(stdin, code)
                            .map_err(|e| format!("failed to write to stdin: {}", e))
                    )
                    .and_then(move |_| {
                        let stdout = match child.stdout().take() {
                            // FIXME configurable max
                            Some(io) => Either::A(helper::read_up_to(io, vec![0u8; 1024])
                                .map(|(_, mut v, s)| {
                                    v.truncate(s);
                                    v
                                })),
                            None => Either::B(Ok(Vec::new()).into_future()),
                        };
                        let stderr = match child.stderr().take() {
                            // FIXME configurable max
                            Some(io) => Either::A(helper::read_up_to(io, vec![0u8; 1024])
                                .map(|(_, mut v, s)| {
                                    v.truncate(s);
                                    v
                                })),
                            None => Either::B(Ok(Vec::new()).into_future()),
                        };
                        child.join3(stdout, stderr)
                    }.map_err(|e| format!("failed to wait for process: {}", e)))
            })
            .map_err(|e| format!("unknown error in exec: {}", e))
            .map(|(status, stdout, stderr)| {
                let mut r = format!("{}{}",
                    String::from_utf8_lossy(&stderr),
                    String::from_utf8_lossy(&stdout),
                );
                if !status.success() {
                    if !r.ends_with('\n') {
                        r.push_str("\n");
                    }
                    if let Some(code) = status.code() {
                        r.push_str(&format!("exited with status {}\n", code));
                    } else if let Some(code) = status.signal() {
                        r.push_str(&format!("signalled with {} ({})\n", strsig(code), strsigabbrev(code)));
                    } else {
                        r.push_str("exited with unknown failure\n");
                    }
                }
                r
            }))
    } else {
        Either::B(Err("empty cmdline".to_owned()).into_future())
    }
}

macro_rules! persistent {
    ($lang:expr, $connfut:expr, $timeout:expr, $buf:expr) => ({
        let buf = $buf;
        let fut = $connfut
            .map_err(|e| format!("error connecting: {}", e))
            .and_then(move |s| write_all(s, buf)
                .map_err(|e| format!("error writing: {}", e)))
            .and_then(|(s, _)| flush(s)
                .map_err(|e| format!("error flushing: {}", e)))
            .and_then(|s| read_exact(s, [0u8; 4])
                .map_err(|e| format!("error reading result length: {}", e)));
        if let Some(timeout) = $timeout {
            Either::A(fut.timeout(Duration::from_secs(timeout as u64)))
        } else {
            Either::B(fut.map_err(|e| timeout::Error::inner(e)))
        }.then(move |r| match r {
            Ok((s, lenb)) => Either::A(read_exact(s, {
                // FIXME configurable max
                let outlen = Cursor::new(lenb).get_u32_le().min(1024) as usize;
                let mut buf = BytesMut::with_capacity(outlen);
                buf.resize(outlen, 0);
                buf
            }).map_err(|e| format!("error reading result: {}", e))
                .map(|(_, ref outb)| String::from_utf8_lossy(outb).into_owned())),
            Err(e) => Either::B(if e.is_elapsed() {
                do_persistent_timeout(&$lang.timeout_cmdline);
                Ok("time limit exceeded".to_owned()).into_future()
            } else {
                Err(format!("error from timeout: {}", e)).into_future()
            })
        })
    });
}

pub fn unix<'a, T, U>(
    lang: Arc<UnixSocketBackend>,
    timeout: Option<usize>,
    context: Option<U>,
    code: T) -> impl Future<Item = String, Error = String> + 'a
        where
            T: AsRef<[u8]>,
            U: AsRef<[u8]> {
    persistent!(lang,
        UnixStream::connect(&lang.socket_addr),
        timeout,
        make_persistent_input(timeout, context, code))
}

fn do_persistent_timeout(cmdline: &Option<Vec<String>>) {
    if let Some(cmdline) = cmdline.as_ref() {
        if let Some(path) = cmdline.iter().nth(0) {
            debug!("timeout kill: launching {:?}", cmdline);
            tokio::spawn(Command::new(path)
                .args(cmdline.iter().skip(1))
                .spawn_async()
                .map_err(|e| error!("failed to exec for timeout kill: {}", e))
                .into_future()
                .and_then(|c| c
                    .map_err(|e| error!("failed to exec for timeout kill: {}", e)))
                    .map(|_| ()));
        }
    }
}

fn make_persistent_input<T, U>(timeout: Option<usize>, context: Option<T>, code: U) -> BytesMut
    where
        T: AsRef<[u8]>,
        U: AsRef<[u8]> {
    let timeout = timeout.unwrap_or(0usize) as u32;
    let contextb = context.as_ref().map(|x| x.as_ref()).unwrap_or(&super::EMPTY_U8);
    let codeb = code.as_ref();
    let contextblen = contextb.len() as u32;
    let codeblen = codeb.len() as u32;

    let mut buf = BytesMut::with_capacity(12usize + contextblen as usize + codeblen as usize);
    buf.put_u32_le(timeout*1000);
    buf.put_u32_le(contextblen);
    buf.put_u32_le(codeblen);
    buf.put(&contextb[..contextblen as usize]);
    buf.put(&codeb[..codeblen as usize]);
    buf
}

mod helper {
    // https://docs.rs/tokio-io/0.1.10/src/tokio_io/io/read_exact.rs.html

    use std::io;
    use std::mem;

    use futures::{Poll, Future};

    use tokio::io::AsyncRead;

    #[derive(Debug)]
    pub struct ReadUpTo<A, T> {
        state: State<A, T>,
    }

    #[derive(Debug)]
    enum State<A, T> {
        Reading {
            a: A,
            buf: T,
            pos: usize,
        },
        Empty,
    }

    pub fn read_up_to<A, T>(a: A, buf: T) -> ReadUpTo<A, T>
        where A: AsyncRead,
            T: AsMut<[u8]>,
    {
        ReadUpTo {
            state: State::Reading {
                a,
                buf,
                pos: 0,
            },
        }
    }

    impl<A, T> Future for ReadUpTo<A, T>
        where A: AsyncRead,
            T: AsMut<[u8]>,
    {
        type Item = (A, T, usize);
        type Error = io::Error;

        fn poll(&mut self) -> Poll<(A, T, usize), io::Error> {
            match self.state {
                State::Reading { ref mut a, ref mut buf, ref mut pos } => {
                    let buf = buf.as_mut();
                    while *pos < buf.len() {
                        let n = try_ready!(a.poll_read(&mut buf[*pos..]));
                        *pos += n;
                        if n == 0 {
                            break;
                        }
                    }
                }
                State::Empty => panic!("poll a ReadUpTo after it's done"),
            }

            match mem::replace(&mut self.state, State::Empty) {
                State::Reading { a, buf, pos } => Ok((a, buf, pos).into()),
                State::Empty => panic!(),
            }
        }
    }
}
<commit_msg>evalbotlib: Kill the daemon if any error occurs, not just timeout<commit_after>use std::process::{Command, Stdio};
use std::io::Cursor;
use std::sync::Arc;
use std::time::Duration;
use std::os::unix::process::ExitStatusExt;

use tokio::prelude::*;
use tokio::prelude::future::Either;
use tokio::timer::timeout;
use tokio_process::CommandExt;
use tokio::{io::{flush, read_exact, write_all}, net::unix::UnixStream};
use bytes::{BytesMut, Buf, BufMut};

use crate::{ExecBackend, UnixSocketBackend};

fn strsig(sig: i32) -> &'static str {
    match sig {
        1 => "Hangup",
        2 => "Interrupt",
        3 => "Quit",
        4 => "Illegal instruction",
        5 => "Trace/breakpoint trap",
        6 => "Aborted",
        7 => "Bus error",
        8 => "Floating point exception",
        9 => "Killed",
        10 => "User defined signal 1",
        11 => "Segmentation fault",
        12 => "User defined signal 2",
        13 => "Broken pipe",
        14 => "Alarm clock",
        15 => "Terminated",
        16 => "Stack fault",
        17 => "Child exited",
        18 => "Continued",
        19 => "Stopped (signal)",
        20 => "Stopped",
        21 => "Stopped (tty input)",
        22 => "Stopped (tty output)",
        23 => "Urgent I/O condition",
        24 => "CPU time limit exceeded",
        25 => "File size limit exceeded",
        26 => "Virtual timer expired",
        27 => "Profiling timer expired",
        28 => "Window changed",
        29 => "I/O possible",
        30 => "Power failure",
        31 => "Bad system call",
        _ => "Unknown signal"
    }
}

fn strsigabbrev(sig: i32) -> &'static str {
    match sig {
        1 => "SIGHUP",
        2 => "SIGINT",
        3 => "SIGQUIT",
        4 => "SIGILL",
        5 => "SIGTRAP",
        6 => "SIGABRT",
        7 => "SIGBUS",
        8 => "SIGFPE",
        9 => "SIGKILL",
        10 => "SIGUSR1",
        11 => "SIGSEGV",
        12 => "SIGUSR2",
        13 => "SIGPIPE",
        14 => "SIGALRM",
        15 => "SIGTERM",
        16 => "SIGSTKFLT",
        17 => "SIGCHLD",
        18 => "SIGCONT",
        19 => "SIGSTOP",
        20 => "SIGTSTP",
        21 => "SIGTTIN",
        22 => "SIGTTOU",
        23 => "SIGURG",
        24 => "SIGXCPU",
        25 => "SIGXFSZ",
        26 => "SIGVTALRM",
        27 => "SIGPROF",
        28 => "SIGWINCH",
        29 => "SIGPOLL",
        30 => "SIGPWR",
        31 => "SIGSYS",
        _ => "(unknown)"
    }
}

pub fn exec<'a, T>(
    lang: Arc<ExecBackend>,
    timeout: Option<usize>,
    code: T) -> impl Future<Item = String, Error = String> + 'a
        where T: AsRef<[u8]> + 'a {
    let timeout_arg = timeout
        .map(|t| format!("{}{}", lang.timeout_prefix.as_ref().map(String::as_str).unwrap_or(""), t));
    let timeout_arg_ref = timeout_arg.as_ref().map(String::as_str);
    if let Some(path) = lang.cmdline.iter().nth(0) {
        let mut cmd = Command::new(path);
        cmd.args(lang.cmdline.iter()
            .skip(1)
            .filter_map(|a| if a == "{TIMEOUT}" {
                timeout_arg_ref
            } else {
                Some(a.as_ref())
            }))
            .stdin(Stdio::piped()).stdout(Stdio::piped()).stderr(Stdio::piped());
        debug!("spawning {:?}", cmd);
        Either::A(cmd.spawn_async()
            .map_err(|e| format!("failed to exec: {}", e))
            .into_future()
            .and_then(|mut child| {
                child.stdin().take().ok_or_else(|| "stdin missing".to_owned()).into_future()
                    .and_then(|stdin|
                        write_all(stdin, code)
                            .map_err(|e| format!("failed to write to stdin: {}", e))
                    )
                    .and_then(move |_| {
                        let stdout = match child.stdout().take() {
                            // FIXME configurable max
                            Some(io) => Either::A(helper::read_up_to(io, vec![0u8; 1024])
                                .map(|(_, mut v, s)| {
                                    v.truncate(s);
                                    v
                                })),
                            None => Either::B(Ok(Vec::new()).into_future()),
                        };
                        let stderr = match child.stderr().take() {
                            // FIXME configurable max
                            Some(io) => Either::A(helper::read_up_to(io, vec![0u8; 1024])
                                .map(|(_, mut v, s)| {
                                    v.truncate(s);
                                    v
                                })),
                            None => Either::B(Ok(Vec::new()).into_future()),
                        };
                        child.join3(stdout, stderr)
                    }.map_err(|e| format!("failed to wait for process: {}", e)))
            })
            .map_err(|e| format!("unknown error in exec: {}", e))
            .map(|(status, stdout, stderr)| {
                let mut r = format!("{}{}",
                    String::from_utf8_lossy(&stderr),
                    String::from_utf8_lossy(&stdout),
                );
                if !status.success() {
                    if !r.ends_with('\n') {
                        r.push_str("\n");
                    }
                    if let Some(code) = status.code() {
                        r.push_str(&format!("exited with status {}\n", code));
                    } else if let Some(code) = status.signal() {
                        r.push_str(&format!("signalled with {} ({})\n", strsig(code), strsigabbrev(code)));
                    } else {
                        r.push_str("exited with unknown failure\n");
                    }
                }
                r
            }))
    } else {
        Either::B(Err("empty cmdline".to_owned()).into_future())
    }
}

macro_rules! persistent {
    ($lang:expr, $connfut:expr, $timeout:expr, $buf:expr) => ({
        let buf = $buf;
        let fut = $connfut
            .map_err(|e| format!("error connecting: {}", e))
            .and_then(move |s| write_all(s, buf)
                .map_err(|e| format!("error writing: {}", e)))
            .and_then(|(s, _)| flush(s)
                .map_err(|e| format!("error flushing: {}", e)))
            .and_then(|s| read_exact(s, [0u8; 4])
                .map_err(|e| format!("error reading result length: {}", e)));
        if let Some(timeout) = $timeout {
            Either::A(fut.timeout(Duration::from_secs(timeout as u64)))
        } else {
            Either::B(fut.map_err(|e| timeout::Error::inner(e)))
        }.then(|r| match r {
            Ok((s, lenb)) => Either::A(read_exact(s, {
                // FIXME configurable max
                let outlen = Cursor::new(lenb).get_u32_le().min(1024) as usize;
                let mut buf = BytesMut::with_capacity(outlen);
                buf.resize(outlen, 0);
                buf
            }).map_err(|e| format!("error reading result: {}", e))
                .map(|(_, ref outb)| String::from_utf8_lossy(outb).into_owned())),
            Err(e) => Either::B(if e.is_elapsed() {
                Err("time limit exceeded".to_owned()).into_future()
            } else {
                Err(format!("error: {}", e)).into_future()
            })
        }).map_err(move |e| {
            do_persistent_timeout(&$lang.timeout_cmdline);
            e
        })
    });
}

pub fn unix<'a, T, U>(
    lang: Arc<UnixSocketBackend>,
    timeout: Option<usize>,
    context: Option<U>,
    code: T) -> impl Future<Item = String, Error = String> + 'a
        where
            T: AsRef<[u8]>,
            U: AsRef<[u8]> {
    persistent!(lang,
        UnixStream::connect(&lang.socket_addr),
        timeout,
        make_persistent_input(timeout, context, code))
}

fn do_persistent_timeout(cmdline: &Option<Vec<String>>) {
    if let Some(cmdline) = cmdline.as_ref() {
        if let Some(path) = cmdline.iter().nth(0) {
            debug!("timeout kill: launching {:?}", cmdline);
            tokio::spawn(Command::new(path)
                .args(cmdline.iter().skip(1))
                .spawn_async()
                .map_err(|e| error!("failed to exec for timeout kill: {}", e))
                .into_future()
                .and_then(|c| c
                    .map_err(|e| error!("failed to exec for timeout kill: {}", e)))
                    .map(|_| ()));
        }
    }
}

fn make_persistent_input<T, U>(timeout: Option<usize>, context: Option<T>, code: U) -> BytesMut
    where
        T: AsRef<[u8]>,
        U: AsRef<[u8]> {
    let timeout = timeout.unwrap_or(0usize) as u32;
    let contextb = context.as_ref().map(|x| x.as_ref()).unwrap_or(&super::EMPTY_U8);
    let codeb = code.as_ref();
    let contextblen = contextb.len() as u32;
    let codeblen = codeb.len() as u32;

    let mut buf = BytesMut::with_capacity(12usize + contextblen as usize + codeblen as usize);
    buf.put_u32_le(timeout*1000);
    buf.put_u32_le(contextblen);
    buf.put_u32_le(codeblen);
    buf.put(&contextb[..contextblen as usize]);
    buf.put(&codeb[..codeblen as usize]);
    buf
}

mod helper {
    // https://docs.rs/tokio-io/0.1.10/src/tokio_io/io/read_exact.rs.html

    use std::io;
    use std::mem;

    use futures::{Poll, Future};

    use tokio::io::AsyncRead;

    #[derive(Debug)]
    pub struct ReadUpTo<A, T> {
        state: State<A, T>,
    }

    #[derive(Debug)]
    enum State<A, T> {
        Reading {
            a: A,
            buf: T,
            pos: usize,
        },
        Empty,
    }

    pub fn read_up_to<A, T>(a: A, buf: T) -> ReadUpTo<A, T>
        where A: AsyncRead,
            T: AsMut<[u8]>,
    {
        ReadUpTo {
            state: State::Reading {
                a,
                buf,
                pos: 0,
            },
        }
    }

    impl<A, T> Future for ReadUpTo<A, T>
        where A: AsyncRead,
            T: AsMut<[u8]>,
    {
        type Item = (A, T, usize);
        type Error = io::Error;

        fn poll(&mut self) -> Poll<(A, T, usize), io::Error> {
            match self.state {
                State::Reading { ref mut a, ref mut buf, ref mut pos } => {
                    let buf = buf.as_mut();
                    while *pos < buf.len() {
                        let n = try_ready!(a.poll_read(&mut buf[*pos..]));
                        *pos += n;
                        if n == 0 {
                            break;
                        }
                    }
                }
                State::Empty => panic!("poll a ReadUpTo after it's done"),
            }

            match mem::replace(&mut self.state, State::Empty) {
                State::Reading { a, buf, pos } => Ok((a, buf, pos).into()),
                State::Empty => panic!(),
            }
        }
    }
}
<|endoftext|>