<commit_before>use crate::decode::lzbuffer;
use crate::decode::rangecoder;
use crate::error;
use byteorder::{LittleEndian, ReadBytesExt};
use std::io;

use crate::decompress::Options;
use crate::decompress::UnpackedSize;

pub struct LZMAParams {
    // most lc significant bits of previous byte are part of the literal context
    lc: u32, // 0..8
    lp: u32, // 0..4
    // context for literal/match is plaintext offset modulo 2^pb
    pb: u32, // 0..4
    dict_size: u32,
    unpacked_size: Option<u64>,
}

impl LZMAParams {
    pub fn read_header<R>(input: &mut R, options: &Options) -> error::Result<LZMAParams>
    where
        R: io::BufRead,
    {
        // Properties
        let props = input.read_u8().or_else(|e| {
            Err(error::Error::LZMAError(format!(
                "LZMA header too short: {}",
                e
            )))
        })?;

        let mut pb = props as u32;
        if pb >= 225 {
            return Err(error::Error::LZMAError(format!(
                "LZMA header invalid properties: {} must be < 225",
                pb
            )));
        }

        let lc: u32 = pb % 9;
        pb /= 9;
        let lp: u32 = pb % 5;
        pb /= 5;

        lzma_info!("Properties {{ lc: {}, lp: {}, pb: {} }}", lc, lp, pb);

        // Dictionary
        let dict_size_provided = input.read_u32::<LittleEndian>().or_else(|e| {
            Err(error::Error::LZMAError(format!(
                "LZMA header too short: {}",
                e
            )))
        })?;
        let dict_size = if dict_size_provided < 0x1000 {
            0x1000
        } else {
            dict_size_provided
        };

        lzma_info!("Dict size: {}", dict_size);

        // Unpacked size
        let unpacked_size: Option<u64> = match options.unpacked_size {
            UnpackedSize::ReadFromHeader => {
                let unpacked_size_provided = input.read_u64::<LittleEndian>().or_else(|e| {
                    Err(error::Error::LZMAError(format!(
                        "LZMA header too short: {}",
                        e
                    )))
                })?;
                let marker_mandatory: bool = unpacked_size_provided == 0xFFFF_FFFF_FFFF_FFFF;
                if marker_mandatory {
                    None
                } else {
                    Some(unpacked_size_provided)
                }
            }
            UnpackedSize::ReadHeaderButUseProvided(x) => {
                input.read_u64::<LittleEndian>()?;
                x
            }
            UnpackedSize::UseProvided(x) => x,
        };

        lzma_info!("Unpacked size: {:?}", unpacked_size);

        let params = LZMAParams {
            lc,
            lp,
            pb,
            dict_size,
            unpacked_size,
        };

        Ok(params)
    }
}

pub struct DecoderState<LZB>
where
    LZB: lzbuffer::LZBuffer,
{
    pub output: LZB,
    // most lc significant bits of previous byte are part of the literal context
    pub lc: u32, // 0..8
    pub lp: u32, // 0..4
    // context for literal/match is plaintext offset modulo 2^pb
    pub pb: u32, // 0..4
    unpacked_size: Option<u64>,
    literal_probs: Vec<Vec<u16>>,
    pos_slot_decoder: Vec<rangecoder::BitTree>,
    align_decoder: rangecoder::BitTree,
    pos_decoders: [u16; 115],
    is_match: [u16; 192], // true = LZ, false = literal
    is_rep: [u16; 12],
    is_rep_g0: [u16; 12],
    is_rep_g1: [u16; 12],
    is_rep_g2: [u16; 12],
    is_rep_0long: [u16; 192],
    state: usize,
    rep: [usize; 4],
    len_decoder: rangecoder::LenDecoder,
    rep_len_decoder: rangecoder::LenDecoder,
}

// Initialize decoder with accumulating buffer
pub fn new_accum<'a, W>(
    output: lzbuffer::LZAccumBuffer<'a, W>,
    lc: u32,
    lp: u32,
    pb: u32,
    unpacked_size: Option<u64>,
) -> DecoderState<lzbuffer::LZAccumBuffer<'a, W>>
where
    W: io::Write,
{
    DecoderState {
        output,
        lc,
        lp,
        pb,
        unpacked_size,
        literal_probs: vec![vec![0x400; 0x300]; 1 << (lc + lp)],
        pos_slot_decoder: vec![rangecoder::BitTree::new(6); 4],
        align_decoder: rangecoder::BitTree::new(4),
        pos_decoders: [0x400; 115],
        is_match: [0x400; 192],
        is_rep: [0x400; 12],
        is_rep_g0: [0x400; 12],
        is_rep_g1: [0x400; 12],
        is_rep_g2: [0x400; 12],
        is_rep_0long: [0x400; 192],
        state: 0,
        rep: [0; 4],
        len_decoder: rangecoder::LenDecoder::new(),
        rep_len_decoder: rangecoder::LenDecoder::new(),
    }
}

// Initialize decoder with circular buffer
pub fn new_circular<'a, W>(
    output: &'a mut W,
    params: LZMAParams,
) -> error::Result<DecoderState<lzbuffer::LZCircularBuffer<'a, W>>>
where
    W: io::Write,
{
    // Decoder
    let decoder = DecoderState {
        output: lzbuffer::LZCircularBuffer::from_stream(output, params.dict_size as usize),
        lc: params.lc,
        lp: params.lp,
        pb: params.pb,
        unpacked_size: params.unpacked_size,
        literal_probs: vec![vec![0x400; 0x300]; 1 << (params.lc + params.lp)],
        pos_slot_decoder: vec![rangecoder::BitTree::new(6); 4],
        align_decoder: rangecoder::BitTree::new(4),
        pos_decoders: [0x400; 115],
        is_match: [0x400; 192],
        is_rep: [0x400; 12],
        is_rep_g0: [0x400; 12],
        is_rep_g1: [0x400; 12],
        is_rep_g2: [0x400; 12],
        is_rep_0long: [0x400; 192],
        state: 0,
        rep: [0; 4],
        len_decoder: rangecoder::LenDecoder::new(),
        rep_len_decoder: rangecoder::LenDecoder::new(),
    };

    Ok(decoder)
}

impl<LZB> DecoderState<LZB>
where
    LZB: lzbuffer::LZBuffer,
{
    pub fn reset_state(&mut self, lc: u32, lp: u32, pb: u32) {
        self.lc = lc;
        self.lp = lp;
        self.pb = pb;
        self.literal_probs = vec![vec![0x400; 0x300]; 1 << (lc + lp)];
        self.pos_slot_decoder = vec![rangecoder::BitTree::new(6); 4];
        self.align_decoder = rangecoder::BitTree::new(4);
        self.pos_decoders = [0x400; 115];
        self.is_match = [0x400; 192];
        self.is_rep = [0x400; 12];
        self.is_rep_g0 = [0x400; 12];
        self.is_rep_g1 = [0x400; 12];
        self.is_rep_g2 = [0x400; 12];
        self.is_rep_0long = [0x400; 192];
        self.state = 0;
        self.rep = [0; 4];
        self.len_decoder = rangecoder::LenDecoder::new();
        self.rep_len_decoder = rangecoder::LenDecoder::new();
    }

    pub fn set_unpacked_size(&mut self, unpacked_size: Option<u64>) {
        self.unpacked_size = unpacked_size;
    }

    pub fn process<'a, R: io::BufRead>(
        &mut self,
        rangecoder: &mut rangecoder::RangeDecoder<'a, R>,
    ) -> error::Result<()> {
        loop {
            if let Some(unpacked_size) = self.unpacked_size {
                if self.output.len() as u64 >= unpacked_size {
                    break;
                }
            } else if rangecoder.is_finished_ok()? {
                break;
            }

            let pos_state = self.output.len() & ((1 << self.pb) - 1);

            // Literal
            if !rangecoder.decode_bit(
                // TODO: assumes pb = 2 ??
                &mut self.is_match[(self.state << 4) + pos_state],
            )? {
                let byte: u8 = self.decode_literal(rangecoder)?;
                lzma_debug!("Literal: {}", byte);
                self.output.append_literal(byte)?;

                self.state = if self.state < 4 {
                    0
                } else if self.state < 10 {
                    self.state - 3
                } else {
                    self.state - 6
                };
                continue;
            }

            // LZ
            let mut len: usize;
            // Distance is repeated from LRU
            if rangecoder.decode_bit(&mut self.is_rep[self.state])? {
                // dist = rep[0]
                if !rangecoder.decode_bit(&mut self.is_rep_g0[self.state])? {
                    // len = 1
                    if !rangecoder
                        .decode_bit(&mut self.is_rep_0long[(self.state << 4) + pos_state])?
                    {
                        // update state (short rep)
                        self.state = if self.state < 7 { 9 } else { 11 };
                        let dist = self.rep[0] + 1;
                        self.output.append_lz(1, dist)?;
                        continue;
                    }
                // dist = rep[i]
                } else {
                    let idx: usize;
                    if !rangecoder.decode_bit(&mut self.is_rep_g1[self.state])? {
                        idx = 1;
                    } else if !rangecoder.decode_bit(&mut self.is_rep_g2[self.state])? {
                        idx = 2;
                    } else {
                        idx = 3;
                    }
                    // Update LRU
                    let dist = self.rep[idx];
                    for i in (0..idx).rev() {
                        self.rep[i + 1] = self.rep[i];
                    }
                    self.rep[0] = dist
                }

                len = self.rep_len_decoder.decode(rangecoder, pos_state)?;
                // update state (rep)
                self.state = if self.state < 7 { 8 } else { 11 };
            // New distance
            } else {
                // Update LRU
                self.rep[3] = self.rep[2];
                self.rep[2] = self.rep[1];
                self.rep[1] = self.rep[0];
                len = self.len_decoder.decode(rangecoder, pos_state)?;

                // update state (match)
                self.state = if self.state < 7 { 7 } else { 10 };
                self.rep[0] = self.decode_distance(rangecoder, len)?;

                if self.rep[0] == 0xFFFF_FFFF {
                    if rangecoder.is_finished_ok()? {
                        break;
                    }
                    return Err(error::Error::LZMAError(String::from(
                        "Found end-of-stream marker but more bytes are available",
                    )));
                }
            }

            len += 2;

            let dist = self.rep[0] + 1;
            self.output.append_lz(len, dist)?;
        }

        if let Some(len) = self.unpacked_size {
            if self.output.len() as u64 != len {
                return Err(error::Error::LZMAError(format!(
                    "Expected unpacked size of {} but decompressed to {}",
                    len,
                    self.output.len()
                )));
            }
        }

        Ok(())
    }

    fn decode_literal<'a, R: io::BufRead>(
        &mut self,
        rangecoder: &mut rangecoder::RangeDecoder<'a, R>,
    ) -> error::Result<u8> {
        let def_prev_byte = 0u8;
        let prev_byte = self.output.last_or(def_prev_byte) as usize;

        let mut result: usize = 1;
        let lit_state =
            ((self.output.len() & ((1 << self.lp) - 1)) << self.lc) + (prev_byte >> (8 - self.lc));
        let probs = &mut self.literal_probs[lit_state];

        if self.state >= 7 {
            let mut match_byte = self.output.last_n(self.rep[0] + 1)? as usize;

            while result < 0x100 {
                let match_bit = (match_byte >> 7) & 1;
                match_byte <<= 1;
                let bit =
                    rangecoder.decode_bit(&mut probs[((1 + match_bit) << 8) + result])? as usize;
                result = (result << 1) ^ bit;
                if match_bit != bit {
                    break;
                }
            }
        }

        while result < 0x100 {
            result = (result << 1) ^ (rangecoder.decode_bit(&mut probs[result])? as usize);
        }

        Ok((result - 0x100) as u8)
    }

    fn decode_distance<'a, R: io::BufRead>(
        &mut self,
        rangecoder: &mut rangecoder::RangeDecoder<'a, R>,
        length: usize,
    ) -> error::Result<usize> {
        let len_state = if length > 3 { 3 } else { length };

        let pos_slot = self.pos_slot_decoder[len_state].parse(rangecoder)? as usize;
        if pos_slot < 4 {
            return Ok(pos_slot);
        }

        let num_direct_bits = (pos_slot >> 1) - 1;
        let mut result = (2 ^ (pos_slot & 1)) << num_direct_bits;

        if pos_slot < 14 {
            result += rangecoder.parse_reverse_bit_tree(
                num_direct_bits,
                &mut self.pos_decoders,
                result - pos_slot,
            )? as usize;
        } else {
            result += (rangecoder.get(num_direct_bits - 4)? as usize) << 4;
            result += self.align_decoder.parse_reverse(rangecoder)? as usize;
        }

        Ok(result)
    }
}
<commit_msg>Fix unpacked_size check<commit_after>use crate::decode::lzbuffer;
use crate::decode::rangecoder;
use crate::error;
use byteorder::{LittleEndian, ReadBytesExt};
use std::io;

use crate::decompress::Options;
use crate::decompress::UnpackedSize;

pub struct LZMAParams {
    // most lc significant bits of previous byte are part of the literal context
    lc: u32, // 0..8
    lp: u32, // 0..4
    // context for literal/match is plaintext offset modulo 2^pb
    pb: u32, // 0..4
    dict_size: u32,
    unpacked_size: Option<u64>,
}

impl LZMAParams {
    pub fn read_header<R>(input: &mut R, options: &Options) -> error::Result<LZMAParams>
    where
        R: io::BufRead,
    {
        // Properties
        let props = input.read_u8().or_else(|e| {
            Err(error::Error::LZMAError(format!(
                "LZMA header too short: {}",
                e
            )))
        })?;

        let mut pb = props as u32;
        if pb >= 225 {
            return Err(error::Error::LZMAError(format!(
                "LZMA header invalid properties: {} must be < 225",
                pb
            )));
        }

        let lc: u32 = pb % 9;
        pb /= 9;
        let lp: u32 = pb % 5;
        pb /= 5;

        lzma_info!("Properties {{ lc: {}, lp: {}, pb: {} }}", lc, lp, pb);

        // Dictionary
        let dict_size_provided = input.read_u32::<LittleEndian>().or_else(|e| {
            Err(error::Error::LZMAError(format!(
                "LZMA header too short: {}",
                e
            )))
        })?;
        let dict_size = if dict_size_provided < 0x1000 {
            0x1000
        } else {
            dict_size_provided
        };

        lzma_info!("Dict size: {}", dict_size);

        // Unpacked size
        let unpacked_size: Option<u64> = match options.unpacked_size {
            UnpackedSize::ReadFromHeader => {
                let unpacked_size_provided = input.read_u64::<LittleEndian>().or_else(|e| {
                    Err(error::Error::LZMAError(format!(
                        "LZMA header too short: {}",
                        e
                    )))
                })?;
                let marker_mandatory: bool = unpacked_size_provided == 0xFFFF_FFFF_FFFF_FFFF;
                if marker_mandatory {
                    None
                } else {
                    Some(unpacked_size_provided)
                }
            }
            UnpackedSize::ReadHeaderButUseProvided(x) => {
                input.read_u64::<LittleEndian>()?;
                x
            }
            UnpackedSize::UseProvided(x) => x,
        };

        lzma_info!("Unpacked size: {:?}", unpacked_size);

        let params = LZMAParams {
            lc,
            lp,
            pb,
            dict_size,
            unpacked_size,
        };

        Ok(params)
    }
}

pub struct DecoderState<LZB>
where
    LZB: lzbuffer::LZBuffer,
{
    pub output: LZB,
    // most lc significant bits of previous byte are part of the literal context
    pub lc: u32, // 0..8
    pub lp: u32, // 0..4
    // context for literal/match is plaintext offset modulo 2^pb
    pub pb: u32, // 0..4
    unpacked_size: Option<u64>,
    literal_probs: Vec<Vec<u16>>,
    pos_slot_decoder: Vec<rangecoder::BitTree>,
    align_decoder: rangecoder::BitTree,
    pos_decoders: [u16; 115],
    is_match: [u16; 192], // true = LZ, false = literal
    is_rep: [u16; 12],
    is_rep_g0: [u16; 12],
    is_rep_g1: [u16; 12],
    is_rep_g2: [u16; 12],
    is_rep_0long: [u16; 192],
    state: usize,
    rep: [usize; 4],
    len_decoder: rangecoder::LenDecoder,
    rep_len_decoder: rangecoder::LenDecoder,
}

// Initialize decoder with accumulating buffer
pub fn new_accum<'a, W>(
    output: lzbuffer::LZAccumBuffer<'a, W>,
    lc: u32,
    lp: u32,
    pb: u32,
    unpacked_size: Option<u64>,
) -> DecoderState<lzbuffer::LZAccumBuffer<'a, W>>
where
    W: io::Write,
{
    DecoderState {
        output,
        lc,
        lp,
        pb,
        unpacked_size,
        literal_probs: vec![vec![0x400; 0x300]; 1 << (lc + lp)],
        pos_slot_decoder: vec![rangecoder::BitTree::new(6); 4],
        align_decoder: rangecoder::BitTree::new(4),
        pos_decoders: [0x400; 115],
        is_match: [0x400; 192],
        is_rep: [0x400; 12],
        is_rep_g0: [0x400; 12],
        is_rep_g1: [0x400; 12],
        is_rep_g2: [0x400; 12],
        is_rep_0long: [0x400; 192],
        state: 0,
        rep: [0; 4],
        len_decoder: rangecoder::LenDecoder::new(),
        rep_len_decoder: rangecoder::LenDecoder::new(),
    }
}

// Initialize decoder with circular buffer
pub fn new_circular<'a, W>(
    output: &'a mut W,
    params: LZMAParams,
) -> error::Result<DecoderState<lzbuffer::LZCircularBuffer<'a, W>>>
where
    W: io::Write,
{
    // Decoder
    let decoder = DecoderState {
        output: lzbuffer::LZCircularBuffer::from_stream(output, params.dict_size as usize),
        lc: params.lc,
        lp: params.lp,
        pb: params.pb,
        unpacked_size: params.unpacked_size,
        literal_probs: vec![vec![0x400; 0x300]; 1 << (params.lc + params.lp)],
        pos_slot_decoder: vec![rangecoder::BitTree::new(6); 4],
        align_decoder: rangecoder::BitTree::new(4),
        pos_decoders: [0x400; 115],
        is_match: [0x400; 192],
        is_rep: [0x400; 12],
        is_rep_g0: [0x400; 12],
        is_rep_g1: [0x400; 12],
        is_rep_g2: [0x400; 12],
        is_rep_0long: [0x400; 192],
        state: 0,
        rep: [0; 4],
        len_decoder: rangecoder::LenDecoder::new(),
        rep_len_decoder: rangecoder::LenDecoder::new(),
    };

    Ok(decoder)
}

impl<LZB> DecoderState<LZB>
where
    LZB: lzbuffer::LZBuffer,
{
    pub fn reset_state(&mut self, lc: u32, lp: u32, pb: u32) {
        self.lc = lc;
        self.lp = lp;
        self.pb = pb;
        self.literal_probs = vec![vec![0x400; 0x300]; 1 << (lc + lp)];
        self.pos_slot_decoder = vec![rangecoder::BitTree::new(6); 4];
        self.align_decoder = rangecoder::BitTree::new(4);
        self.pos_decoders = [0x400; 115];
        self.is_match = [0x400; 192];
        self.is_rep = [0x400; 12];
        self.is_rep_g0 = [0x400; 12];
        self.is_rep_g1 = [0x400; 12];
        self.is_rep_g2 = [0x400; 12];
        self.is_rep_0long = [0x400; 192];
        self.state = 0;
        self.rep = [0; 4];
        self.len_decoder = rangecoder::LenDecoder::new();
        self.rep_len_decoder = rangecoder::LenDecoder::new();
    }

    pub fn set_unpacked_size(&mut self, unpacked_size: Option<u64>) {
        self.unpacked_size = unpacked_size;
    }

    pub fn process<'a, R: io::BufRead>(
        &mut self,
        rangecoder: &mut rangecoder::RangeDecoder<'a, R>,
    ) -> error::Result<()> {
        let mut unpacked = 0;
        loop {
            if let Some(unpacked_size) = self.unpacked_size {
                if unpacked as u64 >= unpacked_size {
                    break;
                }
            } else if rangecoder.is_finished_ok()? {
                break;
            }

            let pos_state = self.output.len() & ((1 << self.pb) - 1);

            // Literal
            if !rangecoder.decode_bit(
                // TODO: assumes pb = 2 ??
                &mut self.is_match[(self.state << 4) + pos_state],
            )? {
                let byte: u8 = self.decode_literal(rangecoder)?;
                lzma_debug!("Literal: {}", byte);
                self.output.append_literal(byte)?;
                unpacked += 1;

                self.state = if self.state < 4 {
                    0
                } else if self.state < 10 {
                    self.state - 3
                } else {
                    self.state - 6
                };
                continue;
            }

            // LZ
            let mut len: usize;
            // Distance is repeated from LRU
            if rangecoder.decode_bit(&mut self.is_rep[self.state])? {
                // dist = rep[0]
                if !rangecoder.decode_bit(&mut self.is_rep_g0[self.state])? {
                    // len = 1
                    if !rangecoder
                        .decode_bit(&mut self.is_rep_0long[(self.state << 4) + pos_state])?
                    {
                        // update state (short rep)
                        self.state = if self.state < 7 { 9 } else { 11 };
                        let dist = self.rep[0] + 1;
                        self.output.append_lz(1, dist)?;
                        unpacked += 1;
                        continue;
                    }
                // dist = rep[i]
                } else {
                    let idx: usize;
                    if !rangecoder.decode_bit(&mut self.is_rep_g1[self.state])? {
                        idx = 1;
                    } else if !rangecoder.decode_bit(&mut self.is_rep_g2[self.state])? {
                        idx = 2;
                    } else {
                        idx = 3;
                    }
                    // Update LRU
                    let dist = self.rep[idx];
                    for i in (0..idx).rev() {
                        self.rep[i + 1] = self.rep[i];
                    }
                    self.rep[0] = dist
                }

                len = self.rep_len_decoder.decode(rangecoder, pos_state)?;
                // update state (rep)
                self.state = if self.state < 7 { 8 } else { 11 };
            // New distance
            } else {
                // Update LRU
                self.rep[3] = self.rep[2];
                self.rep[2] = self.rep[1];
                self.rep[1] = self.rep[0];
                len = self.len_decoder.decode(rangecoder, pos_state)?;

                // update state (match)
                self.state = if self.state < 7 { 7 } else { 10 };
                self.rep[0] = self.decode_distance(rangecoder, len)?;

                if self.rep[0] == 0xFFFF_FFFF {
                    if rangecoder.is_finished_ok()? {
                        break;
                    }
                    return Err(error::Error::LZMAError(String::from(
                        "Found end-of-stream marker but more bytes are available",
                    )));
                }
            }

            len += 2;

            let dist = self.rep[0] + 1;
            self.output.append_lz(len, dist)?;
            unpacked += len
        }

        if let Some(len) = self.unpacked_size {
            if unpacked as u64 != len {
                return Err(error::Error::LZMAError(format!(
                    "Expected unpacked size of {} but decompressed to {}",
                    len,
                    unpacked
                )));
            }
        }

        Ok(())
    }

    fn decode_literal<'a, R: io::BufRead>(
        &mut self,
        rangecoder: &mut rangecoder::RangeDecoder<'a, R>,
    ) -> error::Result<u8> {
        let def_prev_byte = 0u8;
        let prev_byte = self.output.last_or(def_prev_byte) as usize;

        let mut result: usize = 1;
        let lit_state =
            ((self.output.len() & ((1 << self.lp) - 1)) << self.lc) + (prev_byte >> (8 - self.lc));
        let probs = &mut self.literal_probs[lit_state];

        if self.state >= 7 {
            let mut match_byte = self.output.last_n(self.rep[0] + 1)? as usize;

            while result < 0x100 {
                let match_bit = (match_byte >> 7) & 1;
                match_byte <<= 1;
                let bit =
                    rangecoder.decode_bit(&mut probs[((1 + match_bit) << 8) + result])? as usize;
                result = (result << 1) ^ bit;
                if match_bit != bit {
                    break;
                }
            }
        }

        while result < 0x100 {
            result = (result << 1) ^ (rangecoder.decode_bit(&mut probs[result])? as usize);
        }

        Ok((result - 0x100) as u8)
    }

    fn decode_distance<'a, R: io::BufRead>(
        &mut self,
        rangecoder: &mut rangecoder::RangeDecoder<'a, R>,
        length: usize,
    ) -> error::Result<usize> {
        let len_state = if length > 3 { 3 } else { length };

        let pos_slot = self.pos_slot_decoder[len_state].parse(rangecoder)? as usize;
        if pos_slot < 4 {
            return Ok(pos_slot);
        }

        let num_direct_bits = (pos_slot >> 1) - 1;
        let mut result = (2 ^ (pos_slot & 1)) << num_direct_bits;

        if pos_slot < 14 {
            result += rangecoder.parse_reverse_bit_tree(
                num_direct_bits,
                &mut self.pos_decoders,
                result - pos_slot,
            )? as usize;
        } else {
            result += (rangecoder.get(num_direct_bits - 4)? as usize) << 4;
            result += self.align_decoder.parse_reverse(rangecoder)? as usize;
        }

        Ok(result)
    }
}
<|endoftext|>


<commit_before>extern crate image;

use self::image::ImageBuffer;

use std::fs::File;
use std::path::PathBuf;

pub fn prime_filter_to_spiral_png(x_size: usize, from_vec: Vec<bool>, path: String){
    let imgbuf = make_spiral_buffer(x_size, from_vec);
    let mut w_path = PathBuf::from(path);

    w_path.push(format!("ulam_spiral{}", x_size));
    w_path.set_extension("png");

    let ref mut four = File::create(w_path).unwrap();

    let _ = image::ImageLuma8(imgbuf).save(four, image::PNG);
}
fn point_on_spiral_to_array_index_end(spiral_size: usize, x: usize, y: usize)->usize{
    let ss_minus_one = spiral_size-1;
    match (x, y) {
        (x, y) if y == 0 => x,
        (x, y) if x == 0 => ss_minus_one*4 - y,
        (x, y) if y == spiral_size-1 => ss_minus_one*3 - x,
        (x, y) if x == spiral_size-1 => ss_minus_one + y,
        (x, y) => ss_minus_one*4 + point_on_spiral_to_array_index_end(spiral_size-2, x-1, y-1),
    }

}
fn point_on_spiral_to_array_index(spiral_size: usize, x: usize, y: usize)->usize{
    spiral_size*spiral_size - point_on_spiral_to_array_index_end(spiral_size, x, y)
}
fn make_spiral_buffer(x_size: usize, from_vec: Vec<bool>)->
                    ImageBuffer<image::Luma<u8>, Vec<u8>>{
    assert_eq!(from_vec.len(), x_size*x_size+1);
    let img = ImageBuffer::from_fn(x_size as u32, x_size as u32, |x, y| {
        match from_vec[point_on_spiral_to_array_index(x_size, x as usize, y as usize)] {
            true  => image::Luma([0u8]),
            false => image::Luma([255u8]),
        }
    });
    img
}
<commit_msg>Revert "Revert "Fixed recursion stack problem.""<commit_after>extern crate image;

use self::image::ImageBuffer;

use std::fs::File;
use std::path::PathBuf;
use std::cmp;

pub fn prime_filter_to_spiral_png(x_size: usize, from_vec: Vec<bool>, path: String){
    let imgbuf = make_spiral_buffer(x_size, from_vec);
    let mut w_path = PathBuf::from(path);

    w_path.push(format!("ulam_spiral{}", x_size));
    w_path.set_extension("png");

    let ref mut four = File::create(w_path).unwrap();

    let _ = image::ImageLuma8(imgbuf).save(four, image::PNG);
}
fn point_on_spiral_to_array_index_end(spiral_size: usize, x: usize, y: usize)->usize{
    let ss_minus_one = spiral_size-1;
    match (x, y) {
        (x, y) if y == 0 => x,
        (x, y) if x == 0 => ss_minus_one*4 - y,
        (x, y) if y == spiral_size-1 => ss_minus_one*3 - x,
        (x, y) if x == spiral_size-1 => ss_minus_one + y,
        (x, y) => {
            let skip_loop_num = cmp::min(cmp::min(x, y), cmp::min(ss_minus_one-x, ss_minus_one-y));
            (2*spiral_size - skip_loop_num)*skip_loop_num + point_on_spiral_to_array_index_end(
                    spiral_size-2*skip_loop_num, x-skip_loop_num, y-skip_loop_num)
        },
    }

}
fn point_on_spiral_to_array_index(spiral_size: usize, x: usize, y: usize)->usize{
    spiral_size*spiral_size - point_on_spiral_to_array_index_end(spiral_size, x, y)
}
fn make_spiral_buffer(x_size: usize, from_vec: Vec<bool>)->
                    ImageBuffer<image::Luma<u8>, Vec<u8>>{
    assert_eq!(from_vec.len(), x_size*x_size+1);
    let img = ImageBuffer::from_fn(x_size as u32, x_size as u32, |x, y| {
        match from_vec[point_on_spiral_to_array_index(x_size, x as usize, y as usize)] {
            true  => image::Luma([0u8]),
            false => image::Luma([255u8]),
        }
    });
    img
}
<|endoftext|>


<commit_before>use tool::*;

///
/// Represents a tooling environment, which enables tools to be retrieved by name.
///
/// Environments make it possible to use tools without needing to know where the
/// precise implementation is located. Environments also provide a way to perform
/// dependency injection in other tools.
///
pub trait Environment {
    ///
    /// Retrieves a tool using a JSON interface by name
    ///
    fn get_json_tool<'a>(&self, name: &str) -> Result<&'a JsonTool, Error>;
}
<commit_msg>Fix compiler errors for Error<commit_after>use std::error::Error;
use tool::*;

///
/// Represents a tooling environment, which enables tools to be retrieved by name.
///
/// Environments make it possible to use tools without needing to know where the
/// precise implementation is located. Environments also provide a way to perform
/// dependency injection in other tools.
///
pub trait Environment {
    ///
    /// Retrieves a tool using a JSON interface by name
    ///
    fn get_json_tool<'a>(&self, name: &str) -> Result<&'a JsonTool, &Error>;
}
<|endoftext|>


<commit_before>// Lumol, an extensible molecular simulation engine
// Copyright (C) 2015-2016 G. Fraux — BSD license

use energy::{Potential, PairPotential, PairRestriction};

/// The different way to compute non-bonded pair interactions
#[derive(Clone, Copy, Debug)]
enum PairComputation {
    /// Using only a cutoff distance
    Cutoff,
    /// Using a cutoff distance and a shift
    Shifted(f64),
}

/// A non-bonded interaction between two particle. The potential `T` is
/// associated with a [pair restriction][PairRestriction]. For all non-bonded
/// pairs, the potential is computed up to a cutoff distance. An additional
/// shifting of the potential can be used in molecular dynamics, to ensure that
/// the energy is continuous at the cutoff distance.
///
/// [PairRestriction]: enum.PairRestriction.html
#[derive(Clone)]
pub struct PairInteraction {
    /// The potential of this interaction
    potential: Box<PairPotential>,
    /// The associated pair restrictions
    restriction: PairRestriction,
    /// The cutoff distance
    cutoff: f64,
    /// The computation mode
    computation: PairComputation,
}

impl PairInteraction {
    /// Create a new `PairInteraction` with the given cutoff
    pub fn new(potential: Box<PairPotential>, cutoff: f64) -> PairInteraction {
        PairInteraction {
            potential: potential,
            cutoff: cutoff,
            restriction: PairRestriction::None,
            computation: PairComputation::Cutoff,
        }
    }

    /// Create a new `PairInteraction` with the given cutoff, using shifted
    /// computation of the energy.
    pub fn shifted(potential: Box<PairPotential>, cutoff: f64) -> PairInteraction {
        let shift = potential.energy(cutoff);
        PairInteraction {
            potential: potential,
            cutoff: cutoff,
            restriction: PairRestriction::None,
            computation: PairComputation::Shifted(shift),
        }
    }

    /// Get the associated pair restriction
    pub fn restriction(&self) -> PairRestriction {
        self.restriction
    }

    /// Set the restriction associated with this potential
    pub fn set_restriction(&mut self, restriction: PairRestriction) {
        self.restriction = restriction;
    }
}

impl Potential for PairInteraction {
    fn energy(&self, r: f64) -> f64 {
        if r >= self.cutoff {
            0.0
        } else {
            let energy = self.potential.energy(r);
            match self.computation {
                PairComputation::Cutoff => energy,
                PairComputation::Shifted(shift) => energy - shift,
            }
        }
    }

    fn force(&self, r: f64) -> f64 {
        if r >= self.cutoff {
            0.0
        } else {
            self.potential.force(r)
        }
    }
}

impl PairPotential for PairInteraction {}

#[cfg(test)]
mod tests {
    use super::*;
    use energy::{NullPotential, LennardJones, PairRestriction};
    use energy::Potential;

    #[test]
    fn restriction() {
        let mut pairs = PairInteraction::new(Box::new(NullPotential), 10.0);
        assert_eq!(pairs.restriction(), PairRestriction::None);

        pairs.set_restriction(PairRestriction::Exclude13);
        assert_eq!(pairs.restriction(), PairRestriction::Exclude13);
    }

    #[test]
    fn cutoff() {
        let lj = LennardJones{sigma: 1.0, epsilon: 2.0};
        let pairs = PairInteraction::new(Box::new(lj), 4.0);

        assert_eq!(pairs.force(2.5), lj.force(2.5));
        assert_eq!(pairs.energy(2.5), lj.energy(2.5));

        assert_eq!(pairs.force(4.1), 0.0);
        assert_eq!(pairs.energy(4.1), 0.0);
    }

    #[test]
    fn shifted() {
        let lj = LennardJones{sigma: 1.0, epsilon: 2.0};
        let pairs = PairInteraction::shifted(Box::new(lj), 4.0);

        assert_eq!(pairs.force(2.5), lj.force(2.5));
        assert_approx_eq!(pairs.energy(2.5), -0.030681134109158216);

        assert_eq!(pairs.force(4.1), 0.0);
        assert_eq!(pairs.energy(4.1), 0.0);
    }
}
<commit_msg>[doc] Example for the PairInteractions struct<commit_after>// Lumol, an extensible molecular simulation engine
// Copyright (C) 2015-2016 G. Fraux — BSD license

use energy::{Potential, PairPotential, PairRestriction};

/// The different way to compute non-bonded pair interactions
#[derive(Clone, Copy, Debug)]
enum PairComputation {
    /// Using only a cutoff distance
    Cutoff,
    /// Using a cutoff distance and a shift
    Shifted(f64),
}

/// A non-bonded interaction between two particle.
///
/// This is a thin wrapper around a [`Box<PairPotential>`][PairPotential]
/// associated with a [pair restriction][PairRestriction]. It ensure that the
/// potential is computed up to a cutoff distance. An additional shifting of the
/// potential can be used in molecular dynamics, to ensure that the energy is
/// continuous at the cutoff distance.
///
/// [PairPotential]: trait.PairPotential.html
/// [PairRestriction]: enum.PairRestriction.html
#[derive(Clone)]
pub struct PairInteraction {
    /// The potential of this interaction
    potential: Box<PairPotential>,
    /// The associated pair restrictions
    restriction: PairRestriction,
    /// The cutoff distance
    cutoff: f64,
    /// The computation mode
    computation: PairComputation,
}

impl PairInteraction {
    /// Create a new `PairInteraction` for the given `potential` and using the
    /// given `cutoff` distance.
    ///
    /// # Examples
    ///
    /// ```
    /// # use lumol::energy::{Potential, PairInteraction};
    /// use lumol::energy::Harmonic;
    ///
    /// let potential = Box::new(Harmonic{x0: 0.5, k: 4.2});
    /// let interaction = PairInteraction::new(potential, 2.0);
    ///
    /// assert_eq!(interaction.energy(1.0), 0.525);
    /// // energy at and after the cutoff is zero
    /// assert_eq!(interaction.energy(2.0), 0.0);
    /// ```
    pub fn new(potential: Box<PairPotential>, cutoff: f64) -> PairInteraction {
        PairInteraction {
            potential: potential,
            cutoff: cutoff,
            restriction: PairRestriction::None,
            computation: PairComputation::Cutoff,
        }
    }

    /// Create a new `PairInteraction` with the given `cutoff`, using shifted
    /// computation of the energy.
    ///
    /// # Examples
    ///
    /// ```
    /// # use lumol::energy::{Potential, PairInteraction};
    /// use lumol::energy::Harmonic;
    ///
    /// let potential = Box::new(Harmonic{x0: 0.5, k: 4.2});
    /// let interaction = PairInteraction::shifted(potential, 2.0);
    ///
    /// assert_eq!(interaction.energy(1.0), -4.2);
    ///
    /// // energy goes smoothly to zero at the cutoff
    /// assert!(interaction.energy(1.999).abs() < 0.01);
    /// // energy after the cutoff is zero
    /// assert_eq!(interaction.energy(2.0), 0.0);
    /// ```
    pub fn shifted(potential: Box<PairPotential>, cutoff: f64) -> PairInteraction {
        let shift = potential.energy(cutoff);
        PairInteraction {
            potential: potential,
            cutoff: cutoff,
            restriction: PairRestriction::None,
            computation: PairComputation::Shifted(shift),
        }
    }

    /// Get the associated pair restriction. The default is to have no pair
    /// restriction.
    ///
    /// # Examples
    ///
    /// ```
    /// # use lumol::energy::PairInteraction;
    /// use lumol::energy::{NullPotential, PairRestriction};
    /// let interaction = PairInteraction::new(Box::new(NullPotential), 2.0);
    ///
    /// assert_eq!(interaction.restriction(), PairRestriction::None);
    /// ```
    pub fn restriction(&self) -> PairRestriction {
        self.restriction
    }

    /// Set the pair restriction associated with this interaction.
    ///
    /// # Examples
    ///
    /// ```
    /// # use lumol::energy::PairInteraction;
    /// use lumol::energy::{NullPotential, PairRestriction};
    /// let mut interaction = PairInteraction::new(Box::new(NullPotential), 2.0);
    ///
    /// assert_eq!(interaction.restriction(), PairRestriction::None);
    /// interaction.set_restriction(PairRestriction::Exclude12);
    /// assert_eq!(interaction.restriction(), PairRestriction::Exclude12);
    /// ```
    pub fn set_restriction(&mut self, restriction: PairRestriction) {
        self.restriction = restriction;
    }
}

impl Potential for PairInteraction {
    fn energy(&self, r: f64) -> f64 {
        if r >= self.cutoff {
            0.0
        } else {
            let energy = self.potential.energy(r);
            match self.computation {
                PairComputation::Cutoff => energy,
                PairComputation::Shifted(shift) => energy - shift,
            }
        }
    }

    fn force(&self, r: f64) -> f64 {
        if r >= self.cutoff {
            0.0
        } else {
            self.potential.force(r)
        }
    }
}

impl PairPotential for PairInteraction {}

#[cfg(test)]
mod tests {
    use super::*;
    use energy::{NullPotential, LennardJones, PairRestriction};
    use energy::Potential;

    #[test]
    fn restriction() {
        let mut pairs = PairInteraction::new(Box::new(NullPotential), 10.0);
        assert_eq!(pairs.restriction(), PairRestriction::None);

        pairs.set_restriction(PairRestriction::Exclude13);
        assert_eq!(pairs.restriction(), PairRestriction::Exclude13);
    }

    #[test]
    fn cutoff() {
        let lj = LennardJones{sigma: 1.0, epsilon: 2.0};
        let pairs = PairInteraction::new(Box::new(lj), 4.0);

        assert_eq!(pairs.force(2.5), lj.force(2.5));
        assert_eq!(pairs.energy(2.5), lj.energy(2.5));

        assert_eq!(pairs.force(4.1), 0.0);
        assert_eq!(pairs.energy(4.1), 0.0);
    }

    #[test]
    fn shifted() {
        let lj = LennardJones{sigma: 1.0, epsilon: 2.0};
        let pairs = PairInteraction::shifted(Box::new(lj), 4.0);

        assert_eq!(pairs.force(2.5), lj.force(2.5));
        assert_approx_eq!(pairs.energy(2.5), -0.030681134109158216);

        assert_eq!(pairs.force(4.1), 0.0);
        assert_eq!(pairs.energy(4.1), 0.0);
    }
}
<|endoftext|>


<commit_before>/// Maximum number of channels supported in the FLAC format.
pub const MAX_CHANNELS: usize = 8;

/// Audio frame that contains one sample for each channel.
pub struct Frame {
  pub header: Header,
  pub footer: Footer,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ChannelAssignment {
  Independent = 0,
  LeftSide    = 1,
  RightSide   = 2,
  MiddleSide  = 3,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NumberType {
  Frame(u32),
  Sample(u64),
}

/// Information regarding the current audio frame.
#[derive(Debug, PartialEq, Eq)]
pub struct Header {
  pub block_size: u32,
  pub sample_rate: u32,
  pub channels: u8,
  pub channel_assignment: ChannelAssignment,
  pub bits_per_sample: usize,
  pub number: NumberType,
  pub crc: u8,
}

/// End of the audio frame.
///
/// Contains a CRC-16 value that represents the CRC-16 of everything inside
/// the frame before the footer.
#[derive(Debug, PartialEq, Eq)]
pub struct Footer(pub u16);
<commit_msg>Add documentation for ChannelAssignment<commit_after>/// Maximum number of channels supported in the FLAC format.
pub const MAX_CHANNELS: usize = 8;

/// Audio frame that contains one sample for each channel.
pub struct Frame {
  pub header: Header,
  pub footer: Footer,
}

/// Channel assignment order.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ChannelAssignment {
  Independent = 0,
  LeftSide    = 1,
  RightSide   = 2,
  MiddleSide  = 3,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NumberType {
  Frame(u32),
  Sample(u64),
}

/// Information regarding the current audio frame.
#[derive(Debug, PartialEq, Eq)]
pub struct Header {
  pub block_size: u32,
  pub sample_rate: u32,
  pub channels: u8,
  pub channel_assignment: ChannelAssignment,
  pub bits_per_sample: usize,
  pub number: NumberType,
  pub crc: u8,
}

/// End of the audio frame.
///
/// Contains a CRC-16 value that represents the CRC-16 of everything inside
/// the frame before the footer.
#[derive(Debug, PartialEq, Eq)]
pub struct Footer(pub u16);
<|endoftext|>


<commit_before>use crate::{
    backend::renderer::{Frame, ImportAll, Renderer, Texture},
    desktop::{space::*, utils::*},
    utils::{Logical, Point, Rectangle},
    wayland::output::Output,
};
use std::{
    any::{Any, TypeId},
    hash::{Hash, Hasher},
};
use wayland_server::protocol::wl_surface::WlSurface;

/// Indicates default values for some zindexs inside smithay
#[derive(Debug)]
#[repr(u8)]
pub enum RenderZindex {
    /// WlrLayer::Background default zindex
    Background = 10,
    /// WlrLayer::Bottom default zindex
    Bottom = 20,
    /// Not used yet?
    Shell = 30,
    /// Default zindex for Windows
    Top = 40,
    /// Default Layer for RenderElements
    Overlay = 50,
    /// Default Layer for PopUps?
    PopUp = 60,
}

/// Elements rendered by [`Space::render_output`] in addition to windows, layers and popups.
pub type DynamicRenderElements<R> =
    Box<dyn RenderElement<R, <R as Renderer>::Frame, <R as Renderer>::Error, <R as Renderer>::TextureId>>;

pub(super) type SpaceElem<R> =
    dyn SpaceElement<R, <R as Renderer>::Frame, <R as Renderer>::Error, <R as Renderer>::TextureId>;

/// Trait for custom elements to be rendered during [`Space::render_output`].
pub trait RenderElement<R, F, E, T>
where
    R: Renderer<Error = E, TextureId = T, Frame = F> + ImportAll,
    F: Frame<Error = E, TextureId = T>,
    E: std::error::Error,
    T: Texture + 'static,
    Self: Any + 'static,
{
    /// Returns an id unique to this element for the type of Self.
    fn id(&self) -> usize;
    #[doc(hidden)]
    fn type_of(&self) -> TypeId {
        std::any::Any::type_id(self)
    }
    /// Returns the bounding box of this element including its position in the space.
    fn geometry(&self) -> Rectangle<i32, Logical>;
    /// Returns the damage of the element since it's last update.
    /// It should be relative to the elements coordinates.
    ///
    /// If you receive `Some(_)` for `for_values` you may cache that you
    /// send the damage for this `Space` and `Output` combination once
    /// and return an empty vector for subsequent calls until the contents
    /// of this element actually change again for optimization reasons.
    ///
    /// Returning `vec![Rectangle::from_loc_and_size((0, 0), (i32::MAX, i32::MAX))]` is always
    /// correct, but very inefficient.
    fn accumulated_damage(
        &self,
        for_values: Option<SpaceOutputTuple<'_, '_>>,
    ) -> Vec<Rectangle<i32, Logical>>;
    /// Draws the element using the provided `Frame` and `Renderer`.
    ///
    /// - `scale` provides the current fractional scale value to render as
    /// - `damage` provides the regions you need to re-draw and *may* not
    ///   be equivalent to the damage returned by `accumulated_damage`.
    ///   Redrawing other parts of the element is not valid and may cause rendering artifacts.
    fn draw(
        &self,
        renderer: &mut R,
        frame: &mut F,
        scale: f64,
        damage: &[Rectangle<i32, Logical>],
        log: &slog::Logger,
    ) -> Result<(), R::Error>;

    /// Returns z_index of RenderElement, reverf too [`RenderZindex`] for default values
    fn z_index(&self) -> u8 {
        RenderZindex::Overlay as u8
    }
}

pub(crate) trait SpaceElement<R, F, E, T>
where
    R: Renderer<Error = E, TextureId = T, Frame = F> + ImportAll,
    F: Frame<Error = E, TextureId = T>,
    E: std::error::Error,
    T: Texture,
{
    fn id(&self) -> usize;
    fn type_of(&self) -> TypeId;
    fn location(&self, space_id: usize) -> Point<i32, Logical> {
        self.geometry(space_id).loc
    }
    fn geometry(&self, space_id: usize) -> Rectangle<i32, Logical>;
    fn accumulated_damage(&self, for_values: Option<(&Space, &Output)>) -> Vec<Rectangle<i32, Logical>>;
    #[allow(clippy::too_many_arguments)]
    fn draw(
        &self,
        space_id: usize,
        renderer: &mut R,
        frame: &mut F,
        scale: f64,
        location: Point<i32, Logical>,
        damage: &[Rectangle<i32, Logical>],
        log: &slog::Logger,
    ) -> Result<(), R::Error>;
    fn z_index(&self) -> u8; //{
                             //    0
                             //}
}

impl<R, F, E, T> SpaceElement<R, F, E, T> for Box<dyn RenderElement<R, F, E, T>>
where
    R: Renderer<Error = E, TextureId = T, Frame = F> + ImportAll + 'static,
    F: Frame<Error = E, TextureId = T> + 'static,
    E: std::error::Error + 'static,
    T: Texture + 'static,
{
    fn id(&self) -> usize {
        (&**self as &dyn RenderElement<R, F, E, T>).id()
    }
    fn type_of(&self) -> TypeId {
        (&**self as &dyn RenderElement<R, F, E, T>).type_of()
    }
    fn geometry(&self, _space_id: usize) -> Rectangle<i32, Logical> {
        (&**self as &dyn RenderElement<R, F, E, T>).geometry()
    }
    fn accumulated_damage(&self, for_values: Option<(&Space, &Output)>) -> Vec<Rectangle<i32, Logical>> {
        (&**self as &dyn RenderElement<R, F, E, T>).accumulated_damage(for_values.map(SpaceOutputTuple::from))
    }
    fn draw(
        &self,
        _space_id: usize,
        renderer: &mut R,
        frame: &mut F,
        scale: f64,
        _location: Point<i32, Logical>,
        damage: &[Rectangle<i32, Logical>],
        log: &slog::Logger,
    ) -> Result<(), R::Error> {
        (&**self as &dyn RenderElement<R, F, E, T>).draw(renderer, frame, scale, damage, log)
    }

    fn z_index(&self) -> u8 {
        RenderElement::z_index(self.as_ref())
    }
}

/// Generic helper for drawing [`WlSurface`]s and their subsurfaces
/// as custom elements via [`RenderElement`].
///
/// For example useful for cursor or drag-and-drop surfaces.
///
/// Note: This element will render nothing, if you are not using
/// [`crate::backend::renderer::utils::on_commit_buffer_handler`]
/// to let smithay handle buffer management.
#[derive(Debug)]
pub struct SurfaceTree {
    /// Surface to be drawn
    pub surface: WlSurface,
    /// Position to draw add
    pub position: Point<i32, Logical>,
}

impl<R, F, E, T> RenderElement<R, F, E, T> for SurfaceTree
where
    R: Renderer<Error = E, TextureId = T, Frame = F> + ImportAll,
    F: Frame<Error = E, TextureId = T>,
    E: std::error::Error,
    T: Texture + 'static,
{
    fn id(&self) -> usize {
        self.surface.as_ref().id() as usize
    }

    fn geometry(&self) -> Rectangle<i32, Logical> {
        let mut bbox = bbox_from_surface_tree(&self.surface, (0, 0));
        bbox.loc += self.position;
        bbox
    }

    fn accumulated_damage(
        &self,
        for_values: Option<SpaceOutputTuple<'_, '_>>,
    ) -> Vec<Rectangle<i32, Logical>> {
        damage_from_surface_tree(&self.surface, (0, 0), for_values.map(|x| (x.0, x.1)))
    }

    fn draw(
        &self,
        renderer: &mut R,
        frame: &mut F,
        scale: f64,
        damage: &[Rectangle<i32, Logical>],
        log: &slog::Logger,
    ) -> Result<(), R::Error> {
        crate::backend::renderer::utils::draw_surface_tree(
            renderer,
            frame,
            &self.surface,
            scale,
            self.position,
            damage,
            log,
        )
    }
}

/// Newtype for (&Space, &Output) to provide a `Hash` implementation for damage tracking
#[derive(Debug, PartialEq)]
pub struct SpaceOutputTuple<'a, 'b>(pub &'a Space, pub &'b Output);

impl<'a, 'b> Hash for SpaceOutputTuple<'a, 'b> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.id.hash(state);
        (std::sync::Arc::as_ptr(&self.1.inner) as *const () as usize).hash(state);
    }
}

impl<'a, 'b> SpaceOutputTuple<'a, 'b> {
    /// Returns an owned version that produces and equivalent hash
    pub fn owned_hash(&self) -> SpaceOutputHash {
        SpaceOutputHash(
            self.0.id,
            std::sync::Arc::as_ptr(&self.1.inner) as *const () as usize,
        )
    }
}

impl<'a, 'b> From<(&'a Space, &'b Output)> for SpaceOutputTuple<'a, 'b> {
    fn from((space, output): (&'a Space, &'b Output)) -> SpaceOutputTuple<'a, 'b> {
        SpaceOutputTuple(space, output)
    }
}

/// Type to use as an owned hashable value equal to [`SpaceOutputTuple`]
#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
pub struct SpaceOutputHash(usize, usize);
<commit_msg>add extra derives to RenderZindex<commit_after>use crate::{
    backend::renderer::{Frame, ImportAll, Renderer, Texture},
    desktop::{space::*, utils::*},
    utils::{Logical, Point, Rectangle},
    wayland::output::Output,
};
use std::{
    any::{Any, TypeId},
    hash::{Hash, Hasher},
};
use wayland_server::protocol::wl_surface::WlSurface;

/// Indicates default values for some zindexs inside smithay
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(u8)]
pub enum RenderZindex {
    /// WlrLayer::Background default zindex
    Background = 10,
    /// WlrLayer::Bottom default zindex
    Bottom = 20,
    /// Default zindex for Windows
    Shell = 30,
    /// Default zindex for Windows PopUps
    PopUpsShell = 40,
    /// WlrLayer::Top default zindex
    Top = 50,
    /// Default Layer for RenderElements
    Overlay = 80,
    /// Default Layer for Overlay PopUp
    PopUpsOverlay = 100,
}

/// Elements rendered by [`Space::render_output`] in addition to windows, layers and popups.
pub type DynamicRenderElements<R> =
    Box<dyn RenderElement<R, <R as Renderer>::Frame, <R as Renderer>::Error, <R as Renderer>::TextureId>>;

pub(super) type SpaceElem<R> =
    dyn SpaceElement<R, <R as Renderer>::Frame, <R as Renderer>::Error, <R as Renderer>::TextureId>;

/// Trait for custom elements to be rendered during [`Space::render_output`].
pub trait RenderElement<R, F, E, T>
where
    R: Renderer<Error = E, TextureId = T, Frame = F> + ImportAll,
    F: Frame<Error = E, TextureId = T>,
    E: std::error::Error,
    T: Texture + 'static,
    Self: Any + 'static,
{
    /// Returns an id unique to this element for the type of Self.
    fn id(&self) -> usize;
    #[doc(hidden)]
    fn type_of(&self) -> TypeId {
        std::any::Any::type_id(self)
    }
    /// Returns the bounding box of this element including its position in the space.
    fn geometry(&self) -> Rectangle<i32, Logical>;
    /// Returns the damage of the element since it's last update.
    /// It should be relative to the elements coordinates.
    ///
    /// If you receive `Some(_)` for `for_values` you may cache that you
    /// send the damage for this `Space` and `Output` combination once
    /// and return an empty vector for subsequent calls until the contents
    /// of this element actually change again for optimization reasons.
    ///
    /// Returning `vec![Rectangle::from_loc_and_size((0, 0), (i32::MAX, i32::MAX))]` is always
    /// correct, but very inefficient.
    fn accumulated_damage(
        &self,
        for_values: Option<SpaceOutputTuple<'_, '_>>,
    ) -> Vec<Rectangle<i32, Logical>>;
    /// Draws the element using the provided `Frame` and `Renderer`.
    ///
    /// - `scale` provides the current fractional scale value to render as
    /// - `damage` provides the regions you need to re-draw and *may* not
    ///   be equivalent to the damage returned by `accumulated_damage`.
    ///   Redrawing other parts of the element is not valid and may cause rendering artifacts.
    fn draw(
        &self,
        renderer: &mut R,
        frame: &mut F,
        scale: f64,
        damage: &[Rectangle<i32, Logical>],
        log: &slog::Logger,
    ) -> Result<(), R::Error>;

    /// Returns z_index of RenderElement, reverf too [`RenderZindex`] for default values
    fn z_index(&self) -> u8 {
        RenderZindex::Overlay as u8
    }
}

pub(crate) trait SpaceElement<R, F, E, T>
where
    R: Renderer<Error = E, TextureId = T, Frame = F> + ImportAll,
    F: Frame<Error = E, TextureId = T>,
    E: std::error::Error,
    T: Texture,
{
    fn id(&self) -> usize;
    fn type_of(&self) -> TypeId;
    fn location(&self, space_id: usize) -> Point<i32, Logical> {
        self.geometry(space_id).loc
    }
    fn geometry(&self, space_id: usize) -> Rectangle<i32, Logical>;
    fn accumulated_damage(&self, for_values: Option<(&Space, &Output)>) -> Vec<Rectangle<i32, Logical>>;
    #[allow(clippy::too_many_arguments)]
    fn draw(
        &self,
        space_id: usize,
        renderer: &mut R,
        frame: &mut F,
        scale: f64,
        location: Point<i32, Logical>,
        damage: &[Rectangle<i32, Logical>],
        log: &slog::Logger,
    ) -> Result<(), R::Error>;
    fn z_index(&self) -> u8; //{
                             //    0
                             //}
}

impl<R, F, E, T> SpaceElement<R, F, E, T> for Box<dyn RenderElement<R, F, E, T>>
where
    R: Renderer<Error = E, TextureId = T, Frame = F> + ImportAll + 'static,
    F: Frame<Error = E, TextureId = T> + 'static,
    E: std::error::Error + 'static,
    T: Texture + 'static,
{
    fn id(&self) -> usize {
        (&**self as &dyn RenderElement<R, F, E, T>).id()
    }
    fn type_of(&self) -> TypeId {
        (&**self as &dyn RenderElement<R, F, E, T>).type_of()
    }
    fn geometry(&self, _space_id: usize) -> Rectangle<i32, Logical> {
        (&**self as &dyn RenderElement<R, F, E, T>).geometry()
    }
    fn accumulated_damage(&self, for_values: Option<(&Space, &Output)>) -> Vec<Rectangle<i32, Logical>> {
        (&**self as &dyn RenderElement<R, F, E, T>).accumulated_damage(for_values.map(SpaceOutputTuple::from))
    }
    fn draw(
        &self,
        _space_id: usize,
        renderer: &mut R,
        frame: &mut F,
        scale: f64,
        _location: Point<i32, Logical>,
        damage: &[Rectangle<i32, Logical>],
        log: &slog::Logger,
    ) -> Result<(), R::Error> {
        (&**self as &dyn RenderElement<R, F, E, T>).draw(renderer, frame, scale, damage, log)
    }

    fn z_index(&self) -> u8 {
        RenderElement::z_index(self.as_ref())
    }
}

/// Generic helper for drawing [`WlSurface`]s and their subsurfaces
/// as custom elements via [`RenderElement`].
///
/// For example useful for cursor or drag-and-drop surfaces.
///
/// Note: This element will render nothing, if you are not using
/// [`crate::backend::renderer::utils::on_commit_buffer_handler`]
/// to let smithay handle buffer management.
#[derive(Debug)]
pub struct SurfaceTree {
    /// Surface to be drawn
    pub surface: WlSurface,
    /// Position to draw add
    pub position: Point<i32, Logical>,
}

impl<R, F, E, T> RenderElement<R, F, E, T> for SurfaceTree
where
    R: Renderer<Error = E, TextureId = T, Frame = F> + ImportAll,
    F: Frame<Error = E, TextureId = T>,
    E: std::error::Error,
    T: Texture + 'static,
{
    fn id(&self) -> usize {
        self.surface.as_ref().id() as usize
    }

    fn geometry(&self) -> Rectangle<i32, Logical> {
        let mut bbox = bbox_from_surface_tree(&self.surface, (0, 0));
        bbox.loc += self.position;
        bbox
    }

    fn accumulated_damage(
        &self,
        for_values: Option<SpaceOutputTuple<'_, '_>>,
    ) -> Vec<Rectangle<i32, Logical>> {
        damage_from_surface_tree(&self.surface, (0, 0), for_values.map(|x| (x.0, x.1)))
    }

    fn draw(
        &self,
        renderer: &mut R,
        frame: &mut F,
        scale: f64,
        damage: &[Rectangle<i32, Logical>],
        log: &slog::Logger,
    ) -> Result<(), R::Error> {
        crate::backend::renderer::utils::draw_surface_tree(
            renderer,
            frame,
            &self.surface,
            scale,
            self.position,
            damage,
            log,
        )
    }
}

/// Newtype for (&Space, &Output) to provide a `Hash` implementation for damage tracking
#[derive(Debug, PartialEq)]
pub struct SpaceOutputTuple<'a, 'b>(pub &'a Space, pub &'b Output);

impl<'a, 'b> Hash for SpaceOutputTuple<'a, 'b> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.id.hash(state);
        (std::sync::Arc::as_ptr(&self.1.inner) as *const () as usize).hash(state);
    }
}

impl<'a, 'b> SpaceOutputTuple<'a, 'b> {
    /// Returns an owned version that produces and equivalent hash
    pub fn owned_hash(&self) -> SpaceOutputHash {
        SpaceOutputHash(
            self.0.id,
            std::sync::Arc::as_ptr(&self.1.inner) as *const () as usize,
        )
    }
}

impl<'a, 'b> From<(&'a Space, &'b Output)> for SpaceOutputTuple<'a, 'b> {
    fn from((space, output): (&'a Space, &'b Output)) -> SpaceOutputTuple<'a, 'b> {
        SpaceOutputTuple(space, output)
    }
}

/// Type to use as an owned hashable value equal to [`SpaceOutputTuple`]
#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
pub struct SpaceOutputHash(usize, usize);
<|endoftext|>


<commit_before>use std::fmt;
use std::from_str::FromStr;

/// HTTP methods, as defined in RFC 2616, §5.1.1.
///
/// Method names are case-sensitive.
#[deriving(Eq,Clone)]
pub enum Method {
    Options,
    Get,
    Head,
    Post,
    Put,
    Delete,
    Trace,
    Connect,
    Patch,  // RFC 5789
    ExtensionMethod(StrBuf),
}

impl FromStr for Method {
    /**
     * Get a *known* `Method` from an *ASCII* string, regardless of case.
     *
     * If you want to support unregistered methods, use `from_str_or_new` instead.
     *
     * (If the string isn't ASCII, this will at present fail: TODO fix that.)
     */
    fn from_str(method: &str) -> Option<Method> {
        if !method.is_ascii() {
            return None;
        }
        match method {
            "OPTIONS" => Some(Options),
            "GET"     => Some(Get),
            "HEAD"    => Some(Head),
            "POST"    => Some(Post),
            "PUT"     => Some(Put),
            "DELETE"  => Some(Delete),
            "TRACE"   => Some(Trace),
            "CONNECT" => Some(Connect),
            "PATCH"   => Some(Patch),
            _         => None
        }
    }
}

impl fmt::Show for Method {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.buf.write(match *self {
            Options                => "OPTIONS".as_bytes(),
            Get                    => "GET".as_bytes(),
            Head                   => "HEAD".as_bytes(),
            Post                   => "POST".as_bytes(),
            Put                    => "PUT".as_bytes(),
            Delete                 => "DELETE".as_bytes(),
            Trace                  => "TRACE".as_bytes(),
            Connect                => "CONNECT".as_bytes(),
            Patch                  => "PATCH".as_bytes(),
            ExtensionMethod(ref s) => s.as_bytes(),
        })
    }
}

impl Method {
    /**
     * Get a `Method` from an *ASCII* string.
     *
     * (If the string isn't ASCII, this will at present fail.)
     */
    pub fn from_str_or_new(method: &str) -> Option<Method> {
        assert!(method.is_ascii());
        Some(match method {
            "OPTIONS" => Options,
            "GET"     => Get,
            "HEAD"    => Head,
            "POST"    => Post,
            "PUT"     => Put,
            "DELETE"  => Delete,
            "TRACE"   => Trace,
            "CONNECT" => Connect,
            "PATCH"   => Patch,
            _         => ExtensionMethod(StrBuf::from_str(method)),
        })
    }
}
<commit_msg>Make Method case insensitive, add spec compliance (#83)<commit_after>use std::fmt;
use std::from_str::FromStr;
use std::ascii::StrAsciiExt;

/// HTTP methods, as defined in RFC 2616, §5.1.1.
///
/// Method names are case-sensitive.
#[deriving(Eq,Clone)]
pub enum Method {
    Options,
    Get,
    Head,
    Post,
    Put,
    Delete,
    Trace,
    Connect,
    Patch,  // RFC 5789
    ExtensionMethod(StrBuf),
}

impl FromStr for Method {
    /**
     * Get a *known* `Method` from an *ASCII* string, regardless of case.
     *
     * If you want to support unregistered methods, use `from_str_or_new` instead.
     *
     * (If the string isn't ASCII, this will at present fail: TODO fix that.)
     */
    fn from_str(method: &str) -> Option<Method> {
        match Method::from_str_or_new(method) {
            Some(ExtensionMethod(_)) => None,
            method_or_none => method_or_none
        }
    }
}

impl fmt::Show for Method {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.buf.write(match *self {
            Options                => "OPTIONS".as_bytes(),
            Get                    => "GET".as_bytes(),
            Head                   => "HEAD".as_bytes(),
            Post                   => "POST".as_bytes(),
            Put                    => "PUT".as_bytes(),
            Delete                 => "DELETE".as_bytes(),
            Trace                  => "TRACE".as_bytes(),
            Connect                => "CONNECT".as_bytes(),
            Patch                  => "PATCH".as_bytes(),
            ExtensionMethod(ref s) => s.as_bytes(),
        })
    }
}

impl Method {
    /**
     * Get a `Method` from an *ASCII* string.
     *
     * (If the string isn't ASCII, this will at present fail.)
     */
    pub fn from_str_or_new(method: &str) -> Option<Method> {
        if !method.is_ascii() {
            return None;
        }
        match method.to_ascii_upper().as_slice() {
            "OPTIONS" => Some(Options),
            "GET"     => Some(Get),
            "HEAD"    => Some(Head),
            "POST"    => Some(Post),
            "PUT"     => Some(Put),
            "DELETE"  => Some(Delete),
            "TRACE"   => Some(Trace),
            "CONNECT" => Some(Connect),
            "PATCH"   => Some(Patch),
            _         => {
                        let is_token = method.as_bytes().iter().all(|&x| {
                            // http://tools.ietf.org/html/rfc2616#section-2.2
                            match x {
                                0..31 | 127 => false, // CTLs
                                40 | 41 | 60 | 62 | 64 |
                                44 | 59 | 58 | 92 | 34 |
                                47 | 91 | 93 | 63 | 61 |
                                123 | 125 | 32  => false, // separators
                                _ => true
                            }
                        });
                        if is_token {
                            Some(ExtensionMethod(StrBuf::from_str(method)) )
                        } else {
                            None
                        }
                    }
        }
    }
}
<|endoftext|>


<commit_before>use std::cmp;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

pub struct HyperLogLog {
    registers: Vec<u64>,
    b: usize,
}

impl HyperLogLog {
    pub fn new(b: usize) -> HyperLogLog {
        assert!(b >= 4);
        assert!(b <= 16);

        let m = (1 as usize) << b;
        let registers = (0..m).map(|_| 0).collect();
        HyperLogLog {
            registers: registers,
            b: b,
        }
    }

    pub fn add<T>(&mut self, obj: &T)
    where
        T: Hash,
    {
        let mut hasher = DefaultHasher::new();
        obj.hash(&mut hasher);
        let h: u64 = hasher.finish();

        let w = h >> self.b;
        let j = h - (w << self.b);
        let p = w.leading_zeros() + 1 - (self.b as u32);

        let m_old = self.registers[j as usize];
        self.registers[j as usize] = cmp::max(m_old, p as u64);
    }

    pub fn count(&self) -> usize {
        let m = self.registers.len() as f64;

        let z = 1f64 /
            self.registers
                .iter()
                .map(|&x| 2f64.powi(-(x as i32)))
                .sum::<f64>();

        let am = if m >= 128. {
            0.7213 / (1. + 1.079 / m)
        } else if m >= 64. {
            0.709
        } else if m >= 32. {
            0.697
        } else {
            0.673
        };

        let e = am * m * m * z;

        let e_star = if e <= 5. / 2. * m {
            // small range correction
            let v = self.registers.iter().filter(|&&x| x == 0).count();
            if v != 0 {
                m * (m / (v as f64)).ln()
            } else {
                e
            }
        } else if e <= 1. / 30. * 2f64.powi(32) {
            // intermediate range => no correction
            e
        } else {
            // large range correction
            -2f64.powi(32) * (1. - e / 2f64.powi(32)).ln()
        };

        e_star as usize
    }
}

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

    #[test]
    fn empty() {
        let hll = HyperLogLog::new(8);
        assert_eq!(hll.count(), 0);
    }

    #[test]
    fn add_b4_n1k() {
        let mut hll = HyperLogLog::new(4);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 571);
    }

    #[test]
    fn add_b8_n1k() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 966);
    }

    #[test]
    fn add_b12_n1k() {
        let mut hll = HyperLogLog::new(12);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 984);
    }

    #[test]
    fn add_b16_n1k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 998);
    }

    #[test]
    fn add_b8_n10k() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10196);
    }

    #[test]
    fn add_b12_n10k() {
        let mut hll = HyperLogLog::new(12);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10303);
    }

    #[test]
    fn add_b16_n10k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10055);
    }

    #[test]
    fn add_b16_n100k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..100000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 100551);
    }

    #[test]
    fn add_b16_n1m() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..1000000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 1000226);
    }
}
<commit_msg>hyperloglog: docs<commit_after>use std::cmp;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

/// A simple implementation of a [HyperLogLog](https://en.wikipedia.org/wiki/HyperLogLog)
pub struct HyperLogLog {
    registers: Vec<u64>,
    b: usize,
}

impl HyperLogLog {
    /// Creates a new, empty HyperLogLog.
    ///
    /// - `b` number of bits used for register selection, number of registers within the
    ///   HyperLogLog will be `2^b`. `b` must be in `[4, 16]`
    ///
    /// Panics when `b` is out of bounds.
    pub fn new(b: usize) -> HyperLogLog {
        assert!(b >= 4);
        assert!(b <= 16);

        let m = (1 as usize) << b;
        let registers = (0..m).map(|_| 0).collect();
        HyperLogLog {
            registers: registers,
            b: b,
        }
    }

    /// Adds an element to the HyperLogLog.
    pub fn add<T>(&mut self, obj: &T)
    where
        T: Hash,
    {
        let mut hasher = DefaultHasher::new();
        obj.hash(&mut hasher);
        let h: u64 = hasher.finish();

        let w = h >> self.b;
        let j = h - (w << self.b);
        let p = w.leading_zeros() + 1 - (self.b as u32);

        let m_old = self.registers[j as usize];
        self.registers[j as usize] = cmp::max(m_old, p as u64);
    }

    /// Guess the number of unique elements seen by the HyperLogLog.
    pub fn count(&self) -> usize {
        let m = self.registers.len() as f64;

        let z = 1f64 /
            self.registers
                .iter()
                .map(|&x| 2f64.powi(-(x as i32)))
                .sum::<f64>();

        let am = if m >= 128. {
            0.7213 / (1. + 1.079 / m)
        } else if m >= 64. {
            0.709
        } else if m >= 32. {
            0.697
        } else {
            0.673
        };

        let e = am * m * m * z;

        let e_star = if e <= 5. / 2. * m {
            // small range correction
            let v = self.registers.iter().filter(|&&x| x == 0).count();
            if v != 0 {
                m * (m / (v as f64)).ln()
            } else {
                e
            }
        } else if e <= 1. / 30. * 2f64.powi(32) {
            // intermediate range => no correction
            e
        } else {
            // large range correction
            -2f64.powi(32) * (1. - e / 2f64.powi(32)).ln()
        };

        e_star as usize
    }
}

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

    #[test]
    fn empty() {
        let hll = HyperLogLog::new(8);
        assert_eq!(hll.count(), 0);
    }

    #[test]
    fn add_b4_n1k() {
        let mut hll = HyperLogLog::new(4);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 571);
    }

    #[test]
    fn add_b8_n1k() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 966);
    }

    #[test]
    fn add_b12_n1k() {
        let mut hll = HyperLogLog::new(12);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 984);
    }

    #[test]
    fn add_b16_n1k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..1000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 998);
    }

    #[test]
    fn add_b8_n10k() {
        let mut hll = HyperLogLog::new(8);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10196);
    }

    #[test]
    fn add_b12_n10k() {
        let mut hll = HyperLogLog::new(12);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10303);
    }

    #[test]
    fn add_b16_n10k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..10000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 10055);
    }

    #[test]
    fn add_b16_n100k() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..100000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 100551);
    }

    #[test]
    fn add_b16_n1m() {
        let mut hll = HyperLogLog::new(16);
        for i in 0..1000000 {
            hll.add(&i);
        }
        assert_eq!(hll.count(), 1000226);
    }
}
<|endoftext|>


<commit_before>use std::{
    collections::{
        hash_map::{Entry, RandomState},
        HashMap,
    },
    hash::{BuildHasher, Hash, Hasher},
    sync::Arc,
    time::SystemTime,
};

use axum::http::StatusCode;
use clap::Clap;
use futures_util::StreamExt;
use rustc_hash::FxHashMap;
use shakmaty::{
    uci::Uci, variant::VariantPosition, zobrist::Zobrist, ByColor, CastlingMode, Outcome, Position,
};
use tokio::{
    sync::{
        mpsc::{self, error::TrySendError},
        watch, RwLock,
    },
    task::JoinHandle,
};

use crate::{
    db::Database,
    model::{
        GameInfo, GameInfoPlayer, Mode, Month, PersonalEntry, PersonalKeyBuilder, PersonalStatus,
        UserId,
    },
};

mod lila;

use lila::{Game, Lila};

#[derive(Clap, Clone)]
pub struct IndexerOpt {
    #[clap(long = "lila", default_value = "https://lichess.org")]
    lila: String,
    #[clap(long = "bearer")]
    bearer: Option<String>,
    #[clap(long = "indexers", default_value = "16")]
    indexers: usize,
}

#[derive(Clone)]
pub struct IndexerStub {
    db: Arc<Database>,
    indexing: Arc<RwLock<HashMap<UserId, watch::Sender<()>>>>,
    random_state: RandomState,
    txs: Vec<mpsc::Sender<IndexerMessage>>,
}

impl IndexerStub {
    pub fn spawn(db: Arc<Database>, opt: IndexerOpt) -> (IndexerStub, Vec<JoinHandle<()>>) {
        let random_state = RandomState::new();
        let indexing = Arc::new(RwLock::new(HashMap::new()));
        let mut txs = Vec::with_capacity(opt.indexers);
        let mut join_handles = Vec::with_capacity(opt.indexers);
        for idx in 0..opt.indexers {
            let (tx, rx) = mpsc::channel(500);
            txs.push(tx);
            join_handles.push(tokio::spawn(
                IndexerActor {
                    idx,
                    rx,
                    indexing: Arc::clone(&indexing),
                    db: Arc::clone(&db),
                    lila: Lila::new(opt.clone()),
                }
                .run(),
            ));
        }
        (
            IndexerStub {
                db,
                random_state,
                indexing,
                txs,
            },
            join_handles,
        )
    }

    pub async fn index_player(&self, player: &UserId) -> Option<watch::Receiver<()>> {
        // Optimization: First try subscribing to an existing indexing run,
        // without acquiring a write lock.
        {
            let guard = self.indexing.read().await;
            if let Some(sender) = guard.get(player) {
                return Some(sender.subscribe());
            }
        }

        // Check player indexing status.
        let mut status = self
            .db
            .queryable()
            .get_player_status(player)
            .expect("get player status")
            .unwrap_or_default();

        let since_created_at = match status
            .maybe_revisit_ongoing()
            .or_else(|| status.maybe_index())
        {
            Some(since) => since,
            None => return None, // Do not reindex so soon!
        };

        // Queue indexing request.
        let responsible_indexer = {
            let mut hasher = self.random_state.build_hasher();
            player.hash(&mut hasher);
            hasher.finish() as usize % self.txs.len()
        };

        let mut guard = self.indexing.write().await;
        let entry = match guard.entry(player.to_owned()) {
            Entry::Occupied(entry) => return Some(entry.get().subscribe()),
            Entry::Vacant(entry) => entry,
        };

        match self.txs[responsible_indexer].try_send(IndexerMessage::IndexPlayer {
            player: player.to_owned(),
            status,
            since_created_at,
        }) {
            Ok(_) => {
                let (sender, receiver) = watch::channel(());
                entry.insert(sender);
                Some(receiver)
            }
            Err(TrySendError::Full(_)) => {
                log::error!(
                    "indexer {}: not queuing {} because indexer queue is full",
                    responsible_indexer,
                    player.as_str()
                );
                None
            }
            Err(TrySendError::Closed(_)) => panic!("indexer {} died", responsible_indexer),
        }
    }
}

struct IndexerActor {
    idx: usize,
    indexing: Arc<RwLock<HashMap<UserId, watch::Sender<()>>>>,
    rx: mpsc::Receiver<IndexerMessage>,
    db: Arc<Database>,
    lila: Lila,
}

impl IndexerActor {
    async fn run(mut self) {
        while let Some(msg) = self.rx.recv().await {
            match msg {
                IndexerMessage::IndexPlayer {
                    player,
                    status,
                    since_created_at,
                } => {
                    self.index_player(&player, status, since_created_at).await;

                    let mut guard = self.indexing.write().await;
                    guard.remove(&player);
                }
            }
        }
    }

    async fn index_player(
        &self,
        player: &UserId,
        mut status: PersonalStatus,
        since_created_at: u64,
    ) {
        log::info!(
            "indexer {} starting {} (created_at >= {})",
            self.idx,
            player.as_str(),
            since_created_at
        );
        let mut games = match self.lila.user_games(player, since_created_at).await {
            Ok(games) => games,
            Err(err) if err.status() == Some(StatusCode::NOT_FOUND) => {
                log::warn!("indexer did not find player {}", player.as_str());
                return;
            }
            Err(err) => {
                log::error!("indexer {}: request failed: {}", self.idx, err);
                return;
            }
        };

        let hash = ByColor::new_with(|color| PersonalKeyBuilder::with_user_pov(player, color));
        let mut num_games = 0;
        while let Some(game) = games.next().await {
            let game = match game {
                Ok(game) => game,
                Err(err) => {
                    log::error!("{}", err);
                    continue;
                }
            };

            self.index_game(player, &hash, game, &mut status);

            num_games += 1;
            if num_games % 1024 == 0 {
                log::info!(
                    "indexer {}: indexed {} games for {}",
                    self.idx,
                    num_games,
                    player.as_str()
                );
            }
        }

        status.indexed_at = SystemTime::now();
        self.db
            .queryable()
            .put_player_status(player, status)
            .expect("put player status");
        log::info!(
            "indexer {}: finished indexing {} games for {}",
            self.idx,
            num_games,
            player.as_str()
        );
    }

    fn index_game(
        &self,
        player: &UserId,
        hash: &ByColor<PersonalKeyBuilder>,
        game: Game,
        status: &mut PersonalStatus,
    ) {
        status.latest_created_at = game.created_at;

        if game.status.is_ongoing() {
            if status.revisit_ongoing_created_at.is_none() {
                log::debug!("will revisit ongoing game {} eventually", game.id);
                status.revisit_ongoing_created_at = Some(game.created_at);
            }
            return;
        }

        if game.status.is_unindexable() {
            log::debug!("not indexing {} with status {:?}", game.id, game.status);
            return;
        }

        if game.players.any(|p| p.user.is_none()) {
            return;
        }

        let color = match game
            .players
            .find(|p| p.user.as_ref().map_or(false, |user| user.name == *player))
        {
            Some(color) => color,
            None => {
                log::error!("{} did not play in {}", player.as_str(), game.id);
                return;
            }
        };

        // Skip game if already indexed from this side. This cannot race with
        // writes, because all writes for the same player are sequenced by
        // this actor. So making a transaction is not required.
        let queryable = self.db.queryable();
        if queryable
            .get_game_info(game.id)
            .expect("get game info")
            .map_or(false, |info| info.indexed.into_color(color))
        {
            log::debug!("{}/{} already indexed", game.id, color);
            return;
        }

        // Prepare basic information and setup initial position.
        let month = Month::from_time_saturating(game.last_move_at);
        let outcome = Outcome::from_winner(game.winner);
        let variant = game.variant.into();
        let pos = match game.initial_fen {
            Some(fen) => VariantPosition::from_setup(variant, &fen, CastlingMode::Chess960),
            None => Ok(VariantPosition::new(variant)),
        };

        let mut pos: Zobrist<_, u128> = match pos {
            Ok(pos) => Zobrist::new(pos),
            Err(err) => {
                log::warn!("not indexing {}: {}", game.id, err);
                return;
            }
        };

        // Build an intermediate table to remove loops (due to repetitions).
        let mut table: FxHashMap<u128, Uci> =
            FxHashMap::with_capacity_and_hasher(game.moves.len(), Default::default());

        for (ply, san) in game.moves.into_iter().enumerate() {
            let m = match san.to_move(&pos) {
                Ok(m) => m,
                Err(err) => {
                    log::warn!("cutting off {} at ply {}: {}: {}", game.id, ply, err, san);
                    break;
                }
            };

            let uci = m.to_uci(CastlingMode::Chess960);
            table.insert(pos.zobrist_hash(), uci);

            pos.play_unchecked(&m);
        }

        // Write to database. All writes regarding this game are batched and
        // atomically committed, so the database will always be in a consistent
        // state.
        let mut batch = queryable.batch();

        batch.merge_game_info(
            game.id,
            GameInfo {
                winner: outcome.winner(),
                speed: game.speed,
                rated: game.rated,
                month,
                players: game.players.map(|p| GameInfoPlayer {
                    name: p.user.map(|p| p.name.to_string()),
                    rating: p.rating,
                }),
                indexed: ByColor::new_with(|c| color == c),
            },
        );

        for (zobrist, uci) in table {
            batch.merge_personal(
                hash.by_color(color)
                    .with_zobrist(variant, zobrist)
                    .with_month(month),
                PersonalEntry::new_single(
                    uci.clone(),
                    game.speed,
                    Mode::from_rated(game.rated),
                    game.id,
                    outcome,
                ),
            );
        }

        batch.write().expect("atomically commit game and moves");
    }
}

enum IndexerMessage {
    IndexPlayer {
        player: UserId,
        status: PersonalStatus,
        since_created_at: u64,
    },
}
<commit_msg>give context in all log statements<commit_after>use std::{
    collections::{
        hash_map::{Entry, RandomState},
        HashMap,
    },
    hash::{BuildHasher, Hash, Hasher},
    sync::Arc,
    time::SystemTime,
};

use axum::http::StatusCode;
use clap::Clap;
use futures_util::StreamExt;
use rustc_hash::FxHashMap;
use shakmaty::{
    uci::Uci, variant::VariantPosition, zobrist::Zobrist, ByColor, CastlingMode, Outcome, Position,
};
use tokio::{
    sync::{
        mpsc::{self, error::TrySendError},
        watch, RwLock,
    },
    task::JoinHandle,
};

use crate::{
    db::Database,
    model::{
        GameInfo, GameInfoPlayer, Mode, Month, PersonalEntry, PersonalKeyBuilder, PersonalStatus,
        UserId,
    },
};

mod lila;

use lila::{Game, Lila};

#[derive(Clap, Clone)]
pub struct IndexerOpt {
    #[clap(long = "lila", default_value = "https://lichess.org")]
    lila: String,
    #[clap(long = "bearer")]
    bearer: Option<String>,
    #[clap(long = "indexers", default_value = "16")]
    indexers: usize,
}

#[derive(Clone)]
pub struct IndexerStub {
    db: Arc<Database>,
    indexing: Arc<RwLock<HashMap<UserId, watch::Sender<()>>>>,
    random_state: RandomState,
    txs: Vec<mpsc::Sender<IndexerMessage>>,
}

impl IndexerStub {
    pub fn spawn(db: Arc<Database>, opt: IndexerOpt) -> (IndexerStub, Vec<JoinHandle<()>>) {
        let random_state = RandomState::new();
        let indexing = Arc::new(RwLock::new(HashMap::new()));
        let mut txs = Vec::with_capacity(opt.indexers);
        let mut join_handles = Vec::with_capacity(opt.indexers);
        for idx in 0..opt.indexers {
            let (tx, rx) = mpsc::channel(500);
            txs.push(tx);
            join_handles.push(tokio::spawn(
                IndexerActor {
                    idx,
                    rx,
                    indexing: Arc::clone(&indexing),
                    db: Arc::clone(&db),
                    lila: Lila::new(opt.clone()),
                }
                .run(),
            ));
        }
        (
            IndexerStub {
                db,
                random_state,
                indexing,
                txs,
            },
            join_handles,
        )
    }

    pub async fn index_player(&self, player: &UserId) -> Option<watch::Receiver<()>> {
        // Optimization: First try subscribing to an existing indexing run,
        // without acquiring a write lock.
        {
            let guard = self.indexing.read().await;
            if let Some(sender) = guard.get(player) {
                return Some(sender.subscribe());
            }
        }

        // Check player indexing status.
        let mut status = self
            .db
            .queryable()
            .get_player_status(player)
            .expect("get player status")
            .unwrap_or_default();

        let since_created_at = match status
            .maybe_revisit_ongoing()
            .or_else(|| status.maybe_index())
        {
            Some(since) => since,
            None => return None, // Do not reindex so soon!
        };

        // Queue indexing request.
        let responsible_indexer = {
            let mut hasher = self.random_state.build_hasher();
            player.hash(&mut hasher);
            hasher.finish() as usize % self.txs.len()
        };

        let mut guard = self.indexing.write().await;
        let entry = match guard.entry(player.to_owned()) {
            Entry::Occupied(entry) => return Some(entry.get().subscribe()),
            Entry::Vacant(entry) => entry,
        };

        match self.txs[responsible_indexer].try_send(IndexerMessage::IndexPlayer {
            player: player.to_owned(),
            status,
            since_created_at,
        }) {
            Ok(_) => {
                let (sender, receiver) = watch::channel(());
                entry.insert(sender);
                Some(receiver)
            }
            Err(TrySendError::Full(_)) => {
                log::error!(
                    "indexer {}: not queuing {} because indexer queue is full",
                    responsible_indexer,
                    player.as_str()
                );
                None
            }
            Err(TrySendError::Closed(_)) => panic!("indexer {} died", responsible_indexer),
        }
    }
}

struct IndexerActor {
    idx: usize,
    indexing: Arc<RwLock<HashMap<UserId, watch::Sender<()>>>>,
    rx: mpsc::Receiver<IndexerMessage>,
    db: Arc<Database>,
    lila: Lila,
}

impl IndexerActor {
    async fn run(mut self) {
        while let Some(msg) = self.rx.recv().await {
            match msg {
                IndexerMessage::IndexPlayer {
                    player,
                    status,
                    since_created_at,
                } => {
                    self.index_player(&player, status, since_created_at).await;

                    let mut guard = self.indexing.write().await;
                    guard.remove(&player);
                }
            }
        }
    }

    async fn index_player(
        &self,
        player: &UserId,
        mut status: PersonalStatus,
        since_created_at: u64,
    ) {
        log::info!(
            "indexer {}: starting {} (created_at >= {})",
            self.idx,
            player.as_str(),
            since_created_at
        );
        let mut games = match self.lila.user_games(player, since_created_at).await {
            Ok(games) => games,
            Err(err) if err.status() == Some(StatusCode::NOT_FOUND) => {
                log::warn!(
                    "indexer {}: did not find player {}",
                    self.idx,
                    player.as_str()
                );
                return;
            }
            Err(err) => {
                log::error!("indexer {}: request failed: {}", self.idx, err);
                return;
            }
        };

        let hash = ByColor::new_with(|color| PersonalKeyBuilder::with_user_pov(player, color));
        let mut num_games = 0;
        while let Some(game) = games.next().await {
            let game = match game {
                Ok(game) => game,
                Err(err) => {
                    log::error!("indexer {}: {}", self.idx, err);
                    continue;
                }
            };

            self.index_game(player, &hash, game, &mut status);

            num_games += 1;
            if num_games % 1024 == 0 {
                log::info!(
                    "indexer {}: indexed {} games for {}",
                    self.idx,
                    num_games,
                    player.as_str()
                );
            }
        }

        status.indexed_at = SystemTime::now();
        self.db
            .queryable()
            .put_player_status(player, status)
            .expect("put player status");
        log::info!(
            "indexer {}: finished indexing {} games for {}",
            self.idx,
            num_games,
            player.as_str()
        );
    }

    fn index_game(
        &self,
        player: &UserId,
        hash: &ByColor<PersonalKeyBuilder>,
        game: Game,
        status: &mut PersonalStatus,
    ) {
        status.latest_created_at = game.created_at;

        if game.status.is_ongoing() {
            if status.revisit_ongoing_created_at.is_none() {
                log::debug!(
                    "indexer {}: will revisit ongoing game {} eventually",
                    self.idx,
                    game.id
                );
                status.revisit_ongoing_created_at = Some(game.created_at);
            }
            return;
        }

        if game.status.is_unindexable() {
            log::debug!(
                "indexer {}: not indexing {} with status {:?}",
                self.idx,
                game.id,
                game.status
            );
            return;
        }

        if game.players.any(|p| p.user.is_none()) {
            return;
        }

        let color = match game
            .players
            .find(|p| p.user.as_ref().map_or(false, |user| user.name == *player))
        {
            Some(color) => color,
            None => {
                log::error!(
                    "indexer {}: {} did not play in {}",
                    self.idx,
                    player.as_str(),
                    game.id
                );
                return;
            }
        };

        // Skip game if already indexed from this side. This cannot race with
        // writes, because all writes for the same player are sequenced by
        // this actor. So making a transaction is not required.
        let queryable = self.db.queryable();
        if queryable
            .get_game_info(game.id)
            .expect("get game info")
            .map_or(false, |info| info.indexed.into_color(color))
        {
            log::debug!(
                "indexer {}: {}/{} already indexed",
                self.idx,
                game.id,
                color
            );
            return;
        }

        // Prepare basic information and setup initial position.
        let month = Month::from_time_saturating(game.last_move_at);
        let outcome = Outcome::from_winner(game.winner);
        let variant = game.variant.into();
        let pos = match game.initial_fen {
            Some(fen) => VariantPosition::from_setup(variant, &fen, CastlingMode::Chess960),
            None => Ok(VariantPosition::new(variant)),
        };

        let mut pos: Zobrist<_, u128> = match pos {
            Ok(pos) => Zobrist::new(pos),
            Err(err) => {
                log::warn!("indexer {}: not indexing {}: {}", self.idx, game.id, err);
                return;
            }
        };

        // Build an intermediate table to remove loops (due to repetitions).
        let mut table: FxHashMap<u128, Uci> =
            FxHashMap::with_capacity_and_hasher(game.moves.len(), Default::default());

        for (ply, san) in game.moves.into_iter().enumerate() {
            let m = match san.to_move(&pos) {
                Ok(m) => m,
                Err(err) => {
                    log::warn!(
                        "indexer {}: cutting off {} at ply {}: {}: {}",
                        self.idx,
                        game.id,
                        ply,
                        err,
                        san
                    );
                    break;
                }
            };

            let uci = m.to_uci(CastlingMode::Chess960);
            table.insert(pos.zobrist_hash(), uci);

            pos.play_unchecked(&m);
        }

        // Write to database. All writes regarding this game are batched and
        // atomically committed, so the database will always be in a consistent
        // state.
        let mut batch = queryable.batch();

        batch.merge_game_info(
            game.id,
            GameInfo {
                winner: outcome.winner(),
                speed: game.speed,
                rated: game.rated,
                month,
                players: game.players.map(|p| GameInfoPlayer {
                    name: p.user.map(|p| p.name.to_string()),
                    rating: p.rating,
                }),
                indexed: ByColor::new_with(|c| color == c),
            },
        );

        for (zobrist, uci) in table {
            batch.merge_personal(
                hash.by_color(color)
                    .with_zobrist(variant, zobrist)
                    .with_month(month),
                PersonalEntry::new_single(
                    uci.clone(),
                    game.speed,
                    Mode::from_rated(game.rated),
                    game.id,
                    outcome,
                ),
            );
        }

        batch.write().expect("atomically commit game and moves");
    }
}

enum IndexerMessage {
    IndexPlayer {
        player: UserId,
        status: PersonalStatus,
        since_created_at: u64,
    },
}
<|endoftext|>