<commit_before>extern mod std;

use libc::c_char;
use libc::c_int;
use libc::size_t;
use libc::c_ulong;
use libc::c_void;
use ptr::null;
use ptr::addr_of;
use ptr::mut_addr_of;
use str::as_c_str;

struct mpz_struct {
  _mp_alloc: c_int,
  _mp_size: c_int,
  _mp_d: *c_void
}

type mpz_srcptr = *const mpz_struct;
type mpz_ptr = *mut mpz_struct;

extern mod gmp {
  fn __gmpz_init(x: mpz_ptr);
  fn __gmpz_init_set_str(rop: mpz_ptr, str: *c_char, base: c_int) -> c_int;
  fn __gmpz_clear(x: mpz_ptr);
  fn __gmpz_set_str(rop: mpz_ptr, str: *c_char, base: c_int) -> c_int;
  fn __gmpz_get_str(str: *c_char, base: c_int, op: mpz_srcptr) -> *c_char;
  pure fn __gmpz_sizeinbase(op: mpz_srcptr, base: c_int) -> size_t;
  pure fn __gmpz_cmp(op: mpz_srcptr, op2: mpz_srcptr) -> c_int;
  pure fn __gmpz_cmp_ui(op1: mpz_srcptr, op2: c_ulong) -> c_int;
  fn __gmpz_add(rop: mpz_ptr, op1: mpz_srcptr, op2: mpz_srcptr);
  fn __gmpz_sub(rop: mpz_ptr, op1: mpz_srcptr, op2: mpz_srcptr);
  fn __gmpz_mul(rop: mpz_ptr, op1: mpz_srcptr, op2: mpz_srcptr);
  fn __gmpz_neg(rop: mpz_ptr, op: mpz_srcptr);
  fn __gmpz_tdiv_q(r: mpz_ptr, n: mpz_srcptr, d: mpz_srcptr);
  fn __gmpz_mod(r: mpz_ptr, n: mpz_srcptr, d: mpz_srcptr);
}

use gmp::*;

pub struct Mpz {
  priv mpz: mpz_struct,

  drop {
    __gmpz_clear(mut_addr_of(&self.mpz));
  }
}

impl Mpz {
  fn set_str(&self, s: &str, base: int) -> bool {
    let mpz = mut_addr_of(&self.mpz);
    let r = as_c_str(s, { |s| __gmpz_set_str(mpz, s, base as c_int) });
    r == 0
  }

  pure fn size_in_base(&self, base: int) -> uint {
    __gmpz_sizeinbase(addr_of(&self.mpz), base as c_int) as uint
  }
}

impl Mpz: cmp::Eq {
  pure fn eq(other: &Mpz) -> bool {
    __gmpz_cmp(addr_of(&self.mpz), addr_of(&other.mpz)) == 0
  }
  pure fn ne(other: &Mpz) -> bool {
    __gmpz_cmp(addr_of(&self.mpz), addr_of(&other.mpz)) != 0
  }
}

impl Mpz: cmp::Ord {
  pure fn lt(other: &Mpz) -> bool {
    __gmpz_cmp(addr_of(&self.mpz), addr_of(&other.mpz)) < 0
  }
  pure fn le(other: &Mpz) -> bool {
    __gmpz_cmp(addr_of(&self.mpz), addr_of(&other.mpz)) <= 0
  }
  pure fn gt(other: &Mpz) -> bool {
    __gmpz_cmp(addr_of(&self.mpz), addr_of(&other.mpz)) > 0
  }
  pure fn ge(other: &Mpz) -> bool {
    __gmpz_cmp(addr_of(&self.mpz), addr_of(&other.mpz)) >= 0
  }
}

impl Mpz: num::Num {
  pure fn add(other: &Mpz) -> Mpz unsafe {
    let res = init();
    __gmpz_add(mut_addr_of(&res.mpz), addr_of(&self.mpz), addr_of(&other.mpz));
    res
  }
  pure fn sub(other: &Mpz) -> Mpz unsafe {
    let res = init();
    __gmpz_sub(mut_addr_of(&res.mpz), addr_of(&self.mpz), addr_of(&other.mpz));
    res
  }
  pure fn mul(other: &Mpz) -> Mpz unsafe {
    let res = init();
    __gmpz_mul(mut_addr_of(&res.mpz), addr_of(&self.mpz), addr_of(&other.mpz));
    res
  }
  pure fn div(other: &Mpz) -> Mpz unsafe {
    if __gmpz_cmp_ui(addr_of(&self.mpz), 0) == 0 {
      fail ~"divide by zero";
    }

    let res = init();
    __gmpz_tdiv_q(mut_addr_of(&res.mpz), addr_of(&self.mpz), addr_of(&other.mpz));
    res
  }
  pure fn modulo(other: &Mpz) -> Mpz unsafe {
    if __gmpz_cmp_ui(addr_of(&self.mpz), 0) == 0 {
      fail ~"divide by zero";
    }

    let res = init();
    __gmpz_mod(mut_addr_of(&res.mpz), addr_of(&self.mpz), addr_of(&other.mpz));
    res
  }
  pure fn neg() -> Mpz unsafe {
    let res = init();
    __gmpz_neg(mut_addr_of(&res.mpz), addr_of(&self.mpz));
    res
  }
  pure fn to_int() -> int {
    fail ~"not implemented";
  }
  static pure fn from_int(other: int) -> Mpz unsafe {
    let res = init();
    // the gmp functions dealing with longs aren't usable here - long is only
    // guaranteed to be at least 32-bit
    assert(res.set_str(other.to_str(), 10));
    res
  }
}

impl Mpz : from_str::FromStr {
  static fn from_str(s: &str) -> Option<Mpz> {
    init_set_str(s, 10)
  }
}

impl Mpz : to_str::ToStr {
  pure fn to_str() -> ~str unsafe {
    let length = self.size_in_base(10) + 2;
    let dst = vec::to_mut(vec::from_elem(length, '0'));
    let pdst = vec::raw::to_ptr(dst);

    str::raw::from_c_str(__gmpz_get_str(pdst as *c_char, 10, addr_of(&self.mpz)))
  }
}

pub fn init() -> Mpz {
  let mpz = mpz_struct { _mp_alloc: 0, _mp_size: 0, _mp_d: null() };
  __gmpz_init(mut_addr_of(&mpz));
  Mpz { mpz: mpz }
}

pub fn init_set_str(s: &str, base: int) -> Option<Mpz> {
  let mpz = mpz_struct { _mp_alloc: 0, _mp_size: 0, _mp_d: null() };
  let mpz_ptr = mut_addr_of(&mpz);
  let r = as_c_str(s, { |s| __gmpz_init_set_str(mpz_ptr, s, base as c_int) });
  if r == 0 {
    Some(Mpz { mpz: mpz })
  } else {
    __gmpz_clear(mpz_ptr);
    None
  }
}

#[cfg(test)]
mod tests {
  #[test]
  fn size_in_base() {
    let x = init();
    x.set_str("150000", 10);
    assert(x.size_in_base(10) == 6);
  }

  #[test]
  fn eq() {
    let x = option::unwrap(init_set_str("4242142195", 10));
    let y = option::unwrap(init_set_str("4242142195", 10));
    let z = option::unwrap(init_set_str("4242142196", 10));

    assert(x == y);
    assert(x != z);
    assert(y != z);
  }

  #[test]
  fn ord() {
    let x = option::unwrap(init_set_str("40000000000000000000000", 10));
    let y = option::unwrap(init_set_str("45000000000000000000000", 10));
    let z = option::unwrap(init_set_str("50000000000000000000000", 10));

    assert(x < y && x < z && y < z);
    assert(x <= x && x <= y && x <= z && y <= z);
    assert(z > y && z > x && y > x);
    assert(z >= z && z >= y && z >= x && y >= x);
  }

  #[test]
  #[should_fail]
  fn div_zero() {
    let x = init();
    x / x;
  }

  #[test]
  #[should_fail]
  fn modulo_zero() {
    let x = init();
    x % x;
  }

  #[test]
  fn test_div_round() {
    let x = init();
    let y = init();
    let mut z: Mpz;

    x.set_str("2", 10);
    y.set_str("3", 10);
    z = x / y;
    assert(__gmpz_cmp_ui(addr_of(&z.mpz), 2 / 3) == 0);

    x.set_str("2", 10);
    y.set_str("-3", 10);
    z = x / y;
    assert(__gmpz_cmp_ui(addr_of(&z.mpz), 2 / -3) == 0);
  }

  #[test]
  fn to_str() {
    let x = option::unwrap(init_set_str("1234567890", 10));
    assert(x.to_str() == ~"1234567890");
  }

  #[test]
  fn invalid_str() {
    assert(init_set_str("foobar", 10).is_none());
  }
}
<commit_msg>expose from_str at the module level<commit_after>extern mod std;

use libc::c_char;
use libc::c_int;
use libc::size_t;
use libc::c_ulong;
use libc::c_void;
use ptr::null;
use ptr::addr_of;
use ptr::mut_addr_of;
use str::as_c_str;

struct mpz_struct {
  _mp_alloc: c_int,
  _mp_size: c_int,
  _mp_d: *c_void
}

type mpz_srcptr = *const mpz_struct;
type mpz_ptr = *mut mpz_struct;

extern mod gmp {
  fn __gmpz_init(x: mpz_ptr);
  fn __gmpz_init_set_str(rop: mpz_ptr, str: *c_char, base: c_int) -> c_int;
  fn __gmpz_clear(x: mpz_ptr);
  fn __gmpz_set_str(rop: mpz_ptr, str: *c_char, base: c_int) -> c_int;
  fn __gmpz_get_str(str: *c_char, base: c_int, op: mpz_srcptr) -> *c_char;
  pure fn __gmpz_sizeinbase(op: mpz_srcptr, base: c_int) -> size_t;
  pure fn __gmpz_cmp(op: mpz_srcptr, op2: mpz_srcptr) -> c_int;
  pure fn __gmpz_cmp_ui(op1: mpz_srcptr, op2: c_ulong) -> c_int;
  fn __gmpz_add(rop: mpz_ptr, op1: mpz_srcptr, op2: mpz_srcptr);
  fn __gmpz_sub(rop: mpz_ptr, op1: mpz_srcptr, op2: mpz_srcptr);
  fn __gmpz_mul(rop: mpz_ptr, op1: mpz_srcptr, op2: mpz_srcptr);
  fn __gmpz_neg(rop: mpz_ptr, op: mpz_srcptr);
  fn __gmpz_tdiv_q(r: mpz_ptr, n: mpz_srcptr, d: mpz_srcptr);
  fn __gmpz_mod(r: mpz_ptr, n: mpz_srcptr, d: mpz_srcptr);
}

use gmp::*;

pub struct Mpz {
  priv mpz: mpz_struct,

  drop {
    __gmpz_clear(mut_addr_of(&self.mpz));
  }
}

impl Mpz {
  fn set_str(&self, s: &str, base: int) -> bool {
    let mpz = mut_addr_of(&self.mpz);
    let r = as_c_str(s, { |s| __gmpz_set_str(mpz, s, base as c_int) });
    r == 0
  }

  pure fn size_in_base(&self, base: int) -> uint {
    __gmpz_sizeinbase(addr_of(&self.mpz), base as c_int) as uint
  }
}

impl Mpz: cmp::Eq {
  pure fn eq(other: &Mpz) -> bool {
    __gmpz_cmp(addr_of(&self.mpz), addr_of(&other.mpz)) == 0
  }
  pure fn ne(other: &Mpz) -> bool {
    __gmpz_cmp(addr_of(&self.mpz), addr_of(&other.mpz)) != 0
  }
}

impl Mpz: cmp::Ord {
  pure fn lt(other: &Mpz) -> bool {
    __gmpz_cmp(addr_of(&self.mpz), addr_of(&other.mpz)) < 0
  }
  pure fn le(other: &Mpz) -> bool {
    __gmpz_cmp(addr_of(&self.mpz), addr_of(&other.mpz)) <= 0
  }
  pure fn gt(other: &Mpz) -> bool {
    __gmpz_cmp(addr_of(&self.mpz), addr_of(&other.mpz)) > 0
  }
  pure fn ge(other: &Mpz) -> bool {
    __gmpz_cmp(addr_of(&self.mpz), addr_of(&other.mpz)) >= 0
  }
}

impl Mpz: num::Num {
  pure fn add(other: &Mpz) -> Mpz unsafe {
    let res = init();
    __gmpz_add(mut_addr_of(&res.mpz), addr_of(&self.mpz), addr_of(&other.mpz));
    res
  }
  pure fn sub(other: &Mpz) -> Mpz unsafe {
    let res = init();
    __gmpz_sub(mut_addr_of(&res.mpz), addr_of(&self.mpz), addr_of(&other.mpz));
    res
  }
  pure fn mul(other: &Mpz) -> Mpz unsafe {
    let res = init();
    __gmpz_mul(mut_addr_of(&res.mpz), addr_of(&self.mpz), addr_of(&other.mpz));
    res
  }
  pure fn div(other: &Mpz) -> Mpz unsafe {
    if __gmpz_cmp_ui(addr_of(&self.mpz), 0) == 0 {
      fail ~"divide by zero";
    }

    let res = init();
    __gmpz_tdiv_q(mut_addr_of(&res.mpz), addr_of(&self.mpz), addr_of(&other.mpz));
    res
  }
  pure fn modulo(other: &Mpz) -> Mpz unsafe {
    if __gmpz_cmp_ui(addr_of(&self.mpz), 0) == 0 {
      fail ~"divide by zero";
    }

    let res = init();
    __gmpz_mod(mut_addr_of(&res.mpz), addr_of(&self.mpz), addr_of(&other.mpz));
    res
  }
  pure fn neg() -> Mpz unsafe {
    let res = init();
    __gmpz_neg(mut_addr_of(&res.mpz), addr_of(&self.mpz));
    res
  }
  pure fn to_int() -> int {
    fail ~"not implemented";
  }
  static pure fn from_int(other: int) -> Mpz unsafe {
    let res = init();
    // the gmp functions dealing with longs aren't usable here - long is only
    // guaranteed to be at least 32-bit
    assert(res.set_str(other.to_str(), 10));
    res
  }
}

pub fn from_str(s: &str) -> Option<Mpz> {
  init_set_str(s, 10)
}

impl Mpz : from_str::FromStr {
  static fn from_str(s: &str) -> Option<Mpz> {
    from_str(s)
  }
}

impl Mpz : to_str::ToStr {
  pure fn to_str() -> ~str unsafe {
    let length = self.size_in_base(10) + 2;
    let dst = vec::to_mut(vec::from_elem(length, '0'));
    let pdst = vec::raw::to_ptr(dst);

    str::raw::from_c_str(__gmpz_get_str(pdst as *c_char, 10, addr_of(&self.mpz)))
  }
}

pub fn init() -> Mpz {
  let mpz = mpz_struct { _mp_alloc: 0, _mp_size: 0, _mp_d: null() };
  __gmpz_init(mut_addr_of(&mpz));
  Mpz { mpz: mpz }
}

pub fn init_set_str(s: &str, base: int) -> Option<Mpz> {
  let mpz = mpz_struct { _mp_alloc: 0, _mp_size: 0, _mp_d: null() };
  let mpz_ptr = mut_addr_of(&mpz);
  let r = as_c_str(s, { |s| __gmpz_init_set_str(mpz_ptr, s, base as c_int) });
  if r == 0 {
    Some(Mpz { mpz: mpz })
  } else {
    __gmpz_clear(mpz_ptr);
    None
  }
}

#[cfg(test)]
mod tests {
  #[test]
  fn size_in_base() {
    let x = init();
    x.set_str("150000", 10);
    assert(x.size_in_base(10) == 6);
  }

  #[test]
  fn eq() {
    let x = option::unwrap(from_str("4242142195"));
    let y = option::unwrap(from_str("4242142195"));
    let z = option::unwrap(from_str("4242142196"));

    assert(x == y);
    assert(x != z);
    assert(y != z);
  }

  #[test]
  fn ord() {
    let x = option::unwrap(from_str("40000000000000000000000"));
    let y = option::unwrap(from_str("45000000000000000000000"));
    let z = option::unwrap(from_str("50000000000000000000000"));

    assert(x < y && x < z && y < z);
    assert(x <= x && x <= y && x <= z && y <= z);
    assert(z > y && z > x && y > x);
    assert(z >= z && z >= y && z >= x && y >= x);
  }

  #[test]
  #[should_fail]
  fn div_zero() {
    let x = init();
    x / x;
  }

  #[test]
  #[should_fail]
  fn modulo_zero() {
    let x = init();
    x % x;
  }

  #[test]
  fn test_div_round() {
    let x = init();
    let y = init();
    let mut z: Mpz;

    x.set_str("2", 10);
    y.set_str("3", 10);
    z = x / y;
    assert(__gmpz_cmp_ui(addr_of(&z.mpz), 2 / 3) == 0);

    x.set_str("2", 10);
    y.set_str("-3", 10);
    z = x / y;
    assert(__gmpz_cmp_ui(addr_of(&z.mpz), 2 / -3) == 0);
  }

  #[test]
  fn to_str() {
    let x = option::unwrap(from_str("1234567890"));
    assert(x.to_str() == ~"1234567890");
  }

  #[test]
  fn invalid_str() {
    assert(from_str("foobar").is_none());
  }
}
<|endoftext|>


<commit_before>use std::*;

pub struct UnionFindForest {
    nodes: Vec<UnionFindForestNode>,
}

enum UnionFindForestNode {
    Root(usize),
    Child(usize),
}

impl UnionFindForest {
    pub fn new(size: usize) -> Self {
        UnionFindForest {
            nodes: (0..size)
                .map(|_| UnionFindForestNode::Root(1))
                .collect::<Vec<_>>(),
        }
    }

    pub fn root_node(&mut self, v: usize) -> (usize, usize) {
        match self.nodes[v] {
            UnionFindForestNode::Root(rank) => (v, rank),
            UnionFindForestNode::Child(u) => {
                let (u, rank) = self.root_node(u);
                self.nodes[v] = UnionFindForestNode::Child(u);
                (u, rank)
            }
        }
    }

    pub fn root(&mut self, v: usize) -> usize {
        self.root_node(v).0
    }

    pub fn connects(&mut self, u: usize, v: usize) -> bool {
        self.root(u) == self.root(v)
    }

    pub fn merge(&mut self, u: usize, v: usize) {
        let (u, u_rank) = self.root_node(u);
        let (v, v_rank) = self.root_node(v);

        if u == v {
            return;
        }

        if u_rank < v_rank {
            self.merge(v, u);
            return;
        }

        self.nodes[v] = UnionFindForestNode::Child(u);
        self.nodes[u] = UnionFindForestNode::Root(u_rank + 1);
    }
}
<commit_msg>Rename a private type for brevity<commit_after>use std::*;

pub struct UnionFindForest {
    nodes: Vec<UffNode>,
}

enum UffNode {
    Root(usize),
    Child(usize),
}

impl UnionFindForest {
    pub fn new(size: usize) -> Self {
        UnionFindForest {
            nodes: (0..size)
                .map(|_| UffNode::Root(1))
                .collect::<Vec<_>>(),
        }
    }

    pub fn root_node(&mut self, v: usize) -> (usize, usize) {
        match self.nodes[v] {
            UffNode::Root(rank) => (v, rank),
            UffNode::Child(u) => {
                let (u, rank) = self.root_node(u);
                self.nodes[v] = UffNode::Child(u);
                (u, rank)
            }
        }
    }

    pub fn root(&mut self, v: usize) -> usize {
        self.root_node(v).0
    }

    pub fn connects(&mut self, u: usize, v: usize) -> bool {
        self.root(u) == self.root(v)
    }

    pub fn merge(&mut self, u: usize, v: usize) {
        let (u, u_rank) = self.root_node(u);
        let (v, v_rank) = self.root_node(v);

        if u == v {
            return;
        }

        if u_rank < v_rank {
            self.merge(v, u);
            return;
        }

        self.nodes[v] = UffNode::Child(u);
        self.nodes[u] = UffNode::Root(u_rank + 1);
    }
}
<|endoftext|>


<commit_before>use std::net::IpAddr;
use std::str::FromStr;

use serde::Deserialize;

use super::{MaxMindDBError, Reader};

#[allow(clippy::float_cmp)]
#[test]
fn test_decoder() {
    let _ = env_logger::try_init();

    #[allow(non_snake_case)]
    #[derive(Deserialize, Debug, Eq, PartialEq)]
    struct MapXType {
        arrayX: Vec<u32>,
        utf8_stringX: String,
    }

    #[allow(non_snake_case)]
    #[derive(Deserialize, Debug, Eq, PartialEq)]
    struct MapType {
        mapX: MapXType,
    }

    #[derive(Deserialize, Debug)]
    struct TestType<'a> {
        array: Vec<u32>,
        boolean: bool,
        bytes: &'a [u8],
        double: f64,
        float: f32,
        int32: i32,
        map: MapType,
        uint16: u16,
        uint32: u32,
        uint64: u64,
        uint128: u128,
        utf8_string: String,
    }

    let r = Reader::open_readfile("test-data/test-data/MaxMind-DB-test-decoder.mmdb");
    if let Err(err) = r {
        panic!("error opening mmdb: {:?}", err);
    }
    let r = r.unwrap();
    let ip: IpAddr = FromStr::from_str("1.1.1.0").unwrap();
    let result: TestType = r.lookup(ip).unwrap();

    assert_eq!(result.array, vec![1_u32, 2_u32, 3_u32]);
    assert!(result.boolean);
    assert_eq!(result.bytes, vec![0_u8, 0_u8, 0_u8, 42_u8]);
    assert_eq!(result.double, 42.123_456);
    assert_eq!(result.float, 1.1);
    assert_eq!(result.int32, -268_435_456);

    assert_eq!(
        result.map,
        MapType {
            mapX: MapXType {
                arrayX: vec![7, 8, 9],
                utf8_stringX: "hello".to_string(),
            },
        }
    );

    assert_eq!(result.uint16, 100);
    assert_eq!(result.uint32, 268_435_456);
    assert_eq!(result.uint64, 1_152_921_504_606_846_976);
    assert_eq!(
        result.uint128,
        1_329_227_995_784_915_872_903_807_060_280_344_576
    );

    assert_eq!(
        result.utf8_string,
        "unicode! \u{262f} - \u{266b}".to_string()
    );
}

#[test]
fn test_broken_database() {
    let _ = env_logger::try_init();

    let r = Reader::open_readfile("test-data/test-data/GeoIP2-City-Test-Broken-Double-Format.mmdb")
        .ok()
        .unwrap();
    let ip: IpAddr = FromStr::from_str("2001:220::").unwrap();

    #[derive(Deserialize, Debug)]
    struct TestType {}
    match r.lookup::<TestType>(ip) {
        Err(e) => assert_eq!(
            e,
            MaxMindDBError::InvalidDatabaseError("double of size 2".to_string())
        ),
        Ok(_) => panic!("Error expected"),
    }
}

#[test]
fn test_missing_database() {
    let _ = env_logger::try_init();

    let r = Reader::open_readfile("file-does-not-exist.mmdb");
    match r {
        Ok(_) => panic!("Received Reader when opening non-existent file"),
        Err(e) => assert!(
            e == MaxMindDBError::IoError(
                "The system cannot find the file specified. (os error 2)".to_string()
            ) || e == MaxMindDBError::IoError("No such file or directory (os error 2)".to_string())
        ),
    }
}

#[test]
fn test_non_database() {
    let _ = env_logger::try_init();

    let r = Reader::open_readfile("README.md");
    match r {
        Ok(_) => panic!("Received Reader when opening a non-MMDB file"),
        Err(e) => assert_eq!(
            e,
            MaxMindDBError::InvalidDatabaseError(
                "Could not find MaxMind DB metadata \
                 in file."
                    .to_string(),
            )
        ),
    }
}

#[test]
fn test_reader() {
    let _ = env_logger::try_init();

    let sizes = [24_usize, 28, 32];
    for record_size in &sizes {
        let versions = [4_usize, 6];
        for ip_version in &versions {
            let filename = format!(
                "test-data/test-data/MaxMind-DB-test-ipv{}-{}.mmdb",
                ip_version, record_size
            );
            let reader = Reader::open_readfile(filename).ok().unwrap();

            check_metadata(&reader, *ip_version, *record_size);
            check_ip(&reader, *ip_version);
        }
    }
}

/// Create Reader by explicitly reading the entire file into a buffer.
#[test]
fn test_reader_readfile() {
    let _ = env_logger::try_init();

    let sizes = [24_usize, 28, 32];
    for record_size in &sizes {
        let versions = [4_usize, 6];
        for ip_version in &versions {
            let filename = format!(
                "test-data/test-data/MaxMind-DB-test-ipv{}-{}.mmdb",
                ip_version, record_size
            );
            let reader = Reader::open_readfile(filename).ok().unwrap();

            check_metadata(&reader, *ip_version, *record_size);
            check_ip(&reader, *ip_version);
        }
    }
}

#[test]
#[cfg(feature = "mmap")]
fn test_reader_mmap() {
    let _ = env_logger::try_init();

    let sizes = [24usize, 28, 32];
    for record_size in sizes.iter() {
        let versions = [4usize, 6];
        for ip_version in versions.iter() {
            let filename = format!(
                "test-data/test-data/MaxMind-DB-test-ipv{}-{}.mmdb",
                ip_version, record_size
            );
            let reader = Reader::open_mmap(filename).ok().unwrap();

            check_metadata(&reader, *ip_version, *record_size);
            check_ip(&reader, *ip_version);
        }
    }
}

#[test]
fn test_lookup_city() {
    use super::geoip2::City;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-City-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("89.160.20.112").unwrap();
    let city: City = reader.lookup(ip).unwrap();

    let iso_code = city.country.and_then(|cy| cy.iso_code);

    assert_eq!(iso_code, Some("SE"));
}

#[test]
fn test_lookup_country() {
    use super::geoip2::Country;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-Country-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("89.160.20.112").unwrap();
    let country: Country = reader.lookup(ip).unwrap();
    let country: super::geoip2::model::Country = country.country.unwrap();

    assert_eq!(country.iso_code, Some("SE"));
    assert_eq!(country.is_in_european_union, Some(true));
}

#[test]
fn test_lookup_connection_type() {
    use super::geoip2::ConnectionType;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-Connection-Type-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("96.1.20.112").unwrap();
    let connection_type: ConnectionType = reader.lookup(ip).unwrap();

    assert_eq!(connection_type.connection_type, Some("Cable/DSL"));
}

#[test]
fn test_lookup_annonymous_ip() {
    use super::geoip2::AnonymousIp;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-Anonymous-IP-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("81.2.69.123").unwrap();
    let anonymous_ip: AnonymousIp = reader.lookup(ip).unwrap();

    assert_eq!(anonymous_ip.is_anonymous, Some(true));
    assert_eq!(anonymous_ip.is_public_proxy, Some(true));
    assert_eq!(anonymous_ip.is_anonymous_vpn, Some(true));
    assert_eq!(anonymous_ip.is_hosting_provider, Some(true));
    assert_eq!(anonymous_ip.is_tor_exit_node, Some(true))
}

#[test]
fn test_lookup_density_income() {
    use super::geoip2::DensityIncome;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-DensityIncome-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("5.83.124.123").unwrap();
    let density_income: DensityIncome = reader.lookup(ip).unwrap();

    assert_eq!(density_income.average_income, Some(32323));
    assert_eq!(density_income.population_density, Some(1232))
}

#[test]
fn test_lookup_domain() {
    use super::geoip2::Domain;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-Domain-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("66.92.80.123").unwrap();
    let domain: Domain = reader.lookup(ip).unwrap();

    assert_eq!(domain.domain, Some("speakeasy.net"));
}

#[test]
fn test_lookup_isp() {
    use super::geoip2::Isp;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-ISP-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("12.87.118.123").unwrap();
    let isp: Isp = reader.lookup(ip).unwrap();

    assert_eq!(isp.autonomous_system_number, Some(7018));
    assert_eq!(isp.isp, Some("AT&T Services"));
    assert_eq!(isp.organization, Some("AT&T Worldnet Services"));
}

#[test]
fn test_lookup_asn() {
    use super::geoip2::Asn;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-ISP-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("1.128.0.123").unwrap();
    let asn: Asn = reader.lookup(ip).unwrap();

    assert_eq!(asn.autonomous_system_number, Some(1221));
    assert_eq!(asn.autonomous_system_organization, Some("Telstra Pty Ltd"));
}

fn check_metadata<T: AsRef<[u8]>>(reader: &Reader<T>, ip_version: usize, record_size: usize) {
    let metadata = &reader.metadata;

    assert_eq!(metadata.binary_format_major_version, 2_u16);

    assert_eq!(metadata.binary_format_minor_version, 0_u16);
    assert!(metadata.build_epoch >= 1_397_457_605);
    assert_eq!(metadata.database_type, "Test".to_string());

    assert_eq!(
        *metadata.description[&"en".to_string()],
        "Test Database".to_string()
    );
    assert_eq!(
        *metadata.description[&"zh".to_string()],
        "Test Database Chinese".to_string()
    );

    assert_eq!(metadata.ip_version, ip_version as u16);
    assert_eq!(metadata.languages, vec!["en".to_string(), "zh".to_string()]);

    if ip_version == 4 {
        assert_eq!(metadata.node_count, 164)
    } else {
        assert_eq!(metadata.node_count, 416)
    }

    assert_eq!(metadata.record_size, record_size as u16)
}

fn check_ip<T: AsRef<[u8]>>(reader: &Reader<T>, ip_version: usize) {
    let subnets = match ip_version {
        6 => [
            "::1:ffff:ffff",
            "::2:0:0",
            "::2:0:0",
            "::2:0:0",
            "::2:0:0",
            "::2:0:40",
            "::2:0:40",
            "::2:0:40",
            "::2:0:50",
            "::2:0:50",
            "::2:0:50",
            "::2:0:58",
            "::2:0:58",
        ],
        _ => [
            "1.1.1.1", "1.1.1.2", "1.1.1.2", "1.1.1.4", "1.1.1.4", "1.1.1.4", "1.1.1.4", "1.1.1.8",
            "1.1.1.8", "1.1.1.8", "1.1.1.16", "1.1.1.16", "1.1.1.16",
        ],
    };

    #[derive(Deserialize, Debug)]
    struct IpType {
        ip: String,
    }

    for subnet in &subnets {
        let ip: IpAddr = FromStr::from_str(&subnet).unwrap();
        let value: IpType = reader.lookup(ip).unwrap();

        assert_eq!(value.ip, *subnet);
    }

    let no_record = ["1.1.1.33", "255.254.253.123", "89fa::"];

    for &address in &no_record {
        let ip: IpAddr = FromStr::from_str(address).unwrap();
        match reader.lookup::<IpType>(ip) {
            Ok(v) => panic!("received an unexpected value: {:?}", v),
            Err(e) => assert_eq!(
                e,
                MaxMindDBError::AddressNotFoundError("Address not found in database".to_string())
            ),
        }
    }
}
<commit_msg>Add test for pointers in metadata<commit_after>use std::net::IpAddr;
use std::str::FromStr;

use serde::Deserialize;

use super::{MaxMindDBError, Reader};

#[allow(clippy::float_cmp)]
#[test]
fn test_decoder() {
    let _ = env_logger::try_init();

    #[allow(non_snake_case)]
    #[derive(Deserialize, Debug, Eq, PartialEq)]
    struct MapXType {
        arrayX: Vec<u32>,
        utf8_stringX: String,
    }

    #[allow(non_snake_case)]
    #[derive(Deserialize, Debug, Eq, PartialEq)]
    struct MapType {
        mapX: MapXType,
    }

    #[derive(Deserialize, Debug)]
    struct TestType<'a> {
        array: Vec<u32>,
        boolean: bool,
        bytes: &'a [u8],
        double: f64,
        float: f32,
        int32: i32,
        map: MapType,
        uint16: u16,
        uint32: u32,
        uint64: u64,
        uint128: u128,
        utf8_string: String,
    }

    let r = Reader::open_readfile("test-data/test-data/MaxMind-DB-test-decoder.mmdb");
    if let Err(err) = r {
        panic!("error opening mmdb: {:?}", err);
    }
    let r = r.unwrap();
    let ip: IpAddr = FromStr::from_str("1.1.1.0").unwrap();
    let result: TestType = r.lookup(ip).unwrap();

    assert_eq!(result.array, vec![1_u32, 2_u32, 3_u32]);
    assert!(result.boolean);
    assert_eq!(result.bytes, vec![0_u8, 0_u8, 0_u8, 42_u8]);
    assert_eq!(result.double, 42.123_456);
    assert_eq!(result.float, 1.1);
    assert_eq!(result.int32, -268_435_456);

    assert_eq!(
        result.map,
        MapType {
            mapX: MapXType {
                arrayX: vec![7, 8, 9],
                utf8_stringX: "hello".to_string(),
            },
        }
    );

    assert_eq!(result.uint16, 100);
    assert_eq!(result.uint32, 268_435_456);
    assert_eq!(result.uint64, 1_152_921_504_606_846_976);
    assert_eq!(
        result.uint128,
        1_329_227_995_784_915_872_903_807_060_280_344_576
    );

    assert_eq!(
        result.utf8_string,
        "unicode! \u{262f} - \u{266b}".to_string()
    );
}

#[test]
fn test_pointers_in_metadata() {
    let _ = env_logger::try_init();

    let r = Reader::open_readfile("test-data/test-data/MaxMind-DB-test-metadata-pointers.mmdb");
    if let Err(err) = r {
        panic!("error opening mmdb: {:?}", err);
    }
}

#[test]
fn test_broken_database() {
    let _ = env_logger::try_init();

    let r = Reader::open_readfile("test-data/test-data/GeoIP2-City-Test-Broken-Double-Format.mmdb")
        .ok()
        .unwrap();
    let ip: IpAddr = FromStr::from_str("2001:220::").unwrap();

    #[derive(Deserialize, Debug)]
    struct TestType {}
    match r.lookup::<TestType>(ip) {
        Err(e) => assert_eq!(
            e,
            MaxMindDBError::InvalidDatabaseError("double of size 2".to_string())
        ),
        Ok(_) => panic!("Error expected"),
    }
}

#[test]
fn test_missing_database() {
    let _ = env_logger::try_init();

    let r = Reader::open_readfile("file-does-not-exist.mmdb");
    match r {
        Ok(_) => panic!("Received Reader when opening non-existent file"),
        Err(e) => assert!(
            e == MaxMindDBError::IoError(
                "The system cannot find the file specified. (os error 2)".to_string()
            ) || e == MaxMindDBError::IoError("No such file or directory (os error 2)".to_string())
        ),
    }
}

#[test]
fn test_non_database() {
    let _ = env_logger::try_init();

    let r = Reader::open_readfile("README.md");
    match r {
        Ok(_) => panic!("Received Reader when opening a non-MMDB file"),
        Err(e) => assert_eq!(
            e,
            MaxMindDBError::InvalidDatabaseError(
                "Could not find MaxMind DB metadata \
                 in file."
                    .to_string(),
            )
        ),
    }
}

#[test]
fn test_reader() {
    let _ = env_logger::try_init();

    let sizes = [24_usize, 28, 32];
    for record_size in &sizes {
        let versions = [4_usize, 6];
        for ip_version in &versions {
            let filename = format!(
                "test-data/test-data/MaxMind-DB-test-ipv{}-{}.mmdb",
                ip_version, record_size
            );
            let reader = Reader::open_readfile(filename).ok().unwrap();

            check_metadata(&reader, *ip_version, *record_size);
            check_ip(&reader, *ip_version);
        }
    }
}

/// Create Reader by explicitly reading the entire file into a buffer.
#[test]
fn test_reader_readfile() {
    let _ = env_logger::try_init();

    let sizes = [24_usize, 28, 32];
    for record_size in &sizes {
        let versions = [4_usize, 6];
        for ip_version in &versions {
            let filename = format!(
                "test-data/test-data/MaxMind-DB-test-ipv{}-{}.mmdb",
                ip_version, record_size
            );
            let reader = Reader::open_readfile(filename).ok().unwrap();

            check_metadata(&reader, *ip_version, *record_size);
            check_ip(&reader, *ip_version);
        }
    }
}

#[test]
#[cfg(feature = "mmap")]
fn test_reader_mmap() {
    let _ = env_logger::try_init();

    let sizes = [24usize, 28, 32];
    for record_size in sizes.iter() {
        let versions = [4usize, 6];
        for ip_version in versions.iter() {
            let filename = format!(
                "test-data/test-data/MaxMind-DB-test-ipv{}-{}.mmdb",
                ip_version, record_size
            );
            let reader = Reader::open_mmap(filename).ok().unwrap();

            check_metadata(&reader, *ip_version, *record_size);
            check_ip(&reader, *ip_version);
        }
    }
}

#[test]
fn test_lookup_city() {
    use super::geoip2::City;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-City-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("89.160.20.112").unwrap();
    let city: City = reader.lookup(ip).unwrap();

    let iso_code = city.country.and_then(|cy| cy.iso_code);

    assert_eq!(iso_code, Some("SE"));
}

#[test]
fn test_lookup_country() {
    use super::geoip2::Country;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-Country-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("89.160.20.112").unwrap();
    let country: Country = reader.lookup(ip).unwrap();
    let country: super::geoip2::model::Country = country.country.unwrap();

    assert_eq!(country.iso_code, Some("SE"));
    assert_eq!(country.is_in_european_union, Some(true));
}

#[test]
fn test_lookup_connection_type() {
    use super::geoip2::ConnectionType;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-Connection-Type-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("96.1.20.112").unwrap();
    let connection_type: ConnectionType = reader.lookup(ip).unwrap();

    assert_eq!(connection_type.connection_type, Some("Cable/DSL"));
}

#[test]
fn test_lookup_annonymous_ip() {
    use super::geoip2::AnonymousIp;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-Anonymous-IP-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("81.2.69.123").unwrap();
    let anonymous_ip: AnonymousIp = reader.lookup(ip).unwrap();

    assert_eq!(anonymous_ip.is_anonymous, Some(true));
    assert_eq!(anonymous_ip.is_public_proxy, Some(true));
    assert_eq!(anonymous_ip.is_anonymous_vpn, Some(true));
    assert_eq!(anonymous_ip.is_hosting_provider, Some(true));
    assert_eq!(anonymous_ip.is_tor_exit_node, Some(true))
}

#[test]
fn test_lookup_density_income() {
    use super::geoip2::DensityIncome;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-DensityIncome-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("5.83.124.123").unwrap();
    let density_income: DensityIncome = reader.lookup(ip).unwrap();

    assert_eq!(density_income.average_income, Some(32323));
    assert_eq!(density_income.population_density, Some(1232))
}

#[test]
fn test_lookup_domain() {
    use super::geoip2::Domain;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-Domain-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("66.92.80.123").unwrap();
    let domain: Domain = reader.lookup(ip).unwrap();

    assert_eq!(domain.domain, Some("speakeasy.net"));
}

#[test]
fn test_lookup_isp() {
    use super::geoip2::Isp;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-ISP-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("12.87.118.123").unwrap();
    let isp: Isp = reader.lookup(ip).unwrap();

    assert_eq!(isp.autonomous_system_number, Some(7018));
    assert_eq!(isp.isp, Some("AT&T Services"));
    assert_eq!(isp.organization, Some("AT&T Worldnet Services"));
}

#[test]
fn test_lookup_asn() {
    use super::geoip2::Asn;
    let _ = env_logger::try_init();

    let filename = "test-data/test-data/GeoIP2-ISP-Test.mmdb";

    let reader = Reader::open_readfile(filename).unwrap();

    let ip: IpAddr = FromStr::from_str("1.128.0.123").unwrap();
    let asn: Asn = reader.lookup(ip).unwrap();

    assert_eq!(asn.autonomous_system_number, Some(1221));
    assert_eq!(asn.autonomous_system_organization, Some("Telstra Pty Ltd"));
}

fn check_metadata<T: AsRef<[u8]>>(reader: &Reader<T>, ip_version: usize, record_size: usize) {
    let metadata = &reader.metadata;

    assert_eq!(metadata.binary_format_major_version, 2_u16);

    assert_eq!(metadata.binary_format_minor_version, 0_u16);
    assert!(metadata.build_epoch >= 1_397_457_605);
    assert_eq!(metadata.database_type, "Test".to_string());

    assert_eq!(
        *metadata.description[&"en".to_string()],
        "Test Database".to_string()
    );
    assert_eq!(
        *metadata.description[&"zh".to_string()],
        "Test Database Chinese".to_string()
    );

    assert_eq!(metadata.ip_version, ip_version as u16);
    assert_eq!(metadata.languages, vec!["en".to_string(), "zh".to_string()]);

    if ip_version == 4 {
        assert_eq!(metadata.node_count, 164)
    } else {
        assert_eq!(metadata.node_count, 416)
    }

    assert_eq!(metadata.record_size, record_size as u16)
}

fn check_ip<T: AsRef<[u8]>>(reader: &Reader<T>, ip_version: usize) {
    let subnets = match ip_version {
        6 => [
            "::1:ffff:ffff",
            "::2:0:0",
            "::2:0:0",
            "::2:0:0",
            "::2:0:0",
            "::2:0:40",
            "::2:0:40",
            "::2:0:40",
            "::2:0:50",
            "::2:0:50",
            "::2:0:50",
            "::2:0:58",
            "::2:0:58",
        ],
        _ => [
            "1.1.1.1", "1.1.1.2", "1.1.1.2", "1.1.1.4", "1.1.1.4", "1.1.1.4", "1.1.1.4", "1.1.1.8",
            "1.1.1.8", "1.1.1.8", "1.1.1.16", "1.1.1.16", "1.1.1.16",
        ],
    };

    #[derive(Deserialize, Debug)]
    struct IpType {
        ip: String,
    }

    for subnet in &subnets {
        let ip: IpAddr = FromStr::from_str(&subnet).unwrap();
        let value: IpType = reader.lookup(ip).unwrap();

        assert_eq!(value.ip, *subnet);
    }

    let no_record = ["1.1.1.33", "255.254.253.123", "89fa::"];

    for &address in &no_record {
        let ip: IpAddr = FromStr::from_str(address).unwrap();
        match reader.lookup::<IpType>(ip) {
            Ok(v) => panic!("received an unexpected value: {:?}", v),
            Err(e) => assert_eq!(
                e,
                MaxMindDBError::AddressNotFoundError("Address not found in database".to_string())
            ),
        }
    }
}
<|endoftext|>


<commit_before>/*
* Huffman table generation:
* HuffTable::huff_decode()
* HuffTable::initialize()
* and used data structures are originally grabbed from the IJG software,
* and adapted by Hubert Figuiere.
*
* Copyright (C) 1991, 1992, Thomas G. Lane.
* Part of the Independent JPEG Group's software.
* See the file Copyright for more details.
*
* Copyright (c) 1993 Brian C. Smith, The Regents of the University
* of California
* All rights reserved.
*
* Copyright (c) 1994 Kongji Huang and Brian C. Smith.
* Cornell University
* All rights reserved.
*
* Permission to use, copy, modify, and distribute this software and its
* documentation for any purpose, without fee, and without written agreement is
* hereby granted, provided that the above copyright notice and the following
* two paragraphs appear in all copies of this software.
*
* IN NO EVENT SHALL CORNELL UNIVERSITY BE LIABLE TO ANY PARTY FOR
* DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT
* OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN IF CORNELL
* UNIVERSITY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
* CORNELL UNIVERSITY SPECIFICALLY DISCLAIMS ANY WARRANTIES,
* INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
* AND FITNESS FOR A PARTICULAR PURPOSE.  THE SOFTWARE PROVIDED HEREUNDER IS
* ON AN "AS IS" BASIS, AND CORNELL UNIVERSITY HAS NO OBLIGATION TO
* PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
*/

use std::fmt;
use decoders::basics::*;

const HUFF_BITMASK: [u32;32] = [0xffffffff, 0x7fffffff,
                                0x3fffffff, 0x1fffffff,
                                0x0fffffff, 0x07ffffff,
                                0x03ffffff, 0x01ffffff,
                                0x00ffffff, 0x007fffff,
                                0x003fffff, 0x001fffff,
                                0x000fffff, 0x0007ffff,
                                0x0003ffff, 0x0001ffff,
                                0x0000ffff, 0x00007fff,
                                0x00003fff, 0x00001fff,
                                0x00000fff, 0x000007ff,
                                0x000003ff, 0x000001ff,
                                0x000000ff, 0x0000007f,
                                0x0000003f, 0x0000001f,
                                0x0000000f, 0x00000007,
                                0x00000003, 0x00000001];

pub struct HuffTable {
  // These two fields directly represent the contents of a JPEG DHT marker
  pub bits: [u32;17],
  pub huffval: [u32;256],

  // The remaining fields are computed from the above to allow more
  // efficient coding and decoding and thus private
  mincode: [u16;17],
  maxcode: [i32;18],
  valptr: [i16;17],
  numbits: [usize;256],
  bigtable: Vec<i32>,
  precision: usize,
  pub use_bigtable: bool,
  pub dng_compatible: bool,
  pub initialized: bool,
}

impl HuffTable {
  pub fn empty(precision: usize) -> HuffTable {
    HuffTable {
      bits: [0;17],
      huffval: [0;256],
      mincode: [0;17],
      maxcode: [0;18],
      valptr: [0;17],
      numbits: [0;256],
      bigtable: Vec::new(),
      precision: precision,
      use_bigtable: true,
      dng_compatible: true,
      initialized: false,
    }
  }

  pub fn initialize(&mut self, use_bigtable: bool) -> Result<(), String> {
    // Figure C.1: make table of Huffman code length for each symbol
    // Note that this is in code-length order.
    let mut p = 0;
    let mut huffsize: [u8;257] = [0;257];
    for l in 1..17 {
      for _ in 1..((self.bits[l] as usize)+1) {
        huffsize[p] = l as u8;
        p += 1;
        if p > 256 {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
      }
    }
    huffsize[p] = 0;
    let lastp = p;

    // Figure C.2: generate the codes themselves
    // Note that this is in code-length order.
    let mut code: u16 = 0;
    let mut si: u32 = huffsize[0] as u32;
    let mut huffcode: [u16;257] = [0;257];
    p = 0;
    while huffsize[p] > 0 {
      while (huffsize[p] as u32) == si {
        huffcode[p] = code;
        p += 1;
        code += 1;
      }
      code <<= 1;
      si += 1;
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    //Figure F.15: generate decoding tables
    self.mincode[0] = 0;
    self.maxcode[0] = 0;
    p = 0;
    for l in 1..17 {
      if self.bits[l] > 0 {
        self.valptr[l] = p as i16;
        self.mincode[l] = huffcode[p];
        p += self.bits[l] as usize;
        self.maxcode[l] = huffcode[p - 1] as i32;
      } else {
        self.valptr[l] = 0xff;   // This check must be present to avoid crash on junk
        self.maxcode[l] = -1;
      }
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    // We put in this value to ensure HuffDecode terminates
    self.maxcode[17] = 0xFFFFF;

    // Build the numbits, value lookup tables.
    // These table allow us to gather 8 bits from the bits stream,
    // and immediately lookup the size and value of the huffman codes.
    // If size is zero, it means that more than 8 bits are in the huffman
    // code (this happens about 3-4% of the time).
    for p in 0..lastp {
      let size = huffsize[p];
      if size <= 8 {
        let value: i32 = self.huffval[p] as i32;
        let code = huffcode[p];
        let ll: i32 = (code << (8 - size)) as i32;
        let ul: i32 = if size < 8 {
          ll | (HUFF_BITMASK[(24+size) as usize]) as i32
        } else {
          ll
        };
        if ul > 256 || ll > ul {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
        for i in ll..(ul+1) {
          self.numbits[i as usize] = (size as usize) | ((value as usize) << 4);
        }
      }
    }

    if use_bigtable {
      try!(self.initialize_bigtable());
    }
    self.initialized = true;
    self.use_bigtable = use_bigtable;
    Ok(())
  }

  fn initialize_bigtable(&mut self) -> Result<(), String> {
    let bits: usize = 14; // HuffDecode functions must be changed, if this is modified.
    let size: usize = 1 << bits;

    self.bigtable = vec![0 as i32;size];
    let mut rv: i32;
    for i in 0..size {
      let input = (i << 2) as u16;
      let mut code: i32 = (input >> 8) as i32;
      let val: u32 = self.numbits[code as usize] as u32;
      let mut l: u32 = val & 15;
      if l > 0 {
        rv = (val >> 4) as i32;
      } else {
        l = 8;
        while code > self.maxcode[l as usize] && l < 16 {
          let temp: i32 = (input >> (15 - l) & 1) as i32;
          code = (code << 1) | temp;
          l += 1;
        }

        //With garbage input we may reach the sentinel value l = 17.
        if l > self.precision as u32 || self.valptr[l as usize] == 0xff {
          self.bigtable[i] = 0xff;
          continue
        } else {
          rv = self.huffval[
            self.valptr[l as usize] as usize +
            (code - (self.mincode[l as usize] as i32)) as usize
          ] as i32;
        }
      }

      if rv == 16 {
        self.bigtable[i] = if self.dng_compatible {
          (-(32768 << 8)) | (16 + l as i32)
        } else {
          (-(32768 << 8)) | l as i32
        };
        continue
      }

      if rv + l as i32 > bits as i32 {
        self.bigtable[i] = 0xff;
        continue
      }

      if rv != 0 {
        let mut x = (input >> (16 - (l as i32) - rv) & ((1 << rv) - 1)) as i32;
        if (x & (1 << (rv - 1))) == 0 {
          x -= (1 << rv) - 1;
        }
        self.bigtable[i] = (x << 8) | ((l as i32) + rv);
      } else {
        self.bigtable[i] = l as i32;
      }
    }
    self.use_bigtable = true;
    Ok(())
  }

  // Taken from Figure F.16: extract next coded symbol from input stream
  pub fn huff_decode(&self, pump: &mut BitPump) -> Result<i32,String> {
    if !self.initialized {
      return Err("ljpeg: Trying to read from uninitialized HuffTable".to_string());
    }

    //First attempt to do complete decode, by using the first 14 bits
    let mut code: i32 = pump.peek_bits(14) as i32;
    if self.use_bigtable {
      let val: i32 = self.bigtable[code as usize];
      if val & 0xff != 0xff {
        pump.consume_bits((val & 0xff) as u32);
        return Ok(val >> 8)
      }
    }

    // If the huffman code is less than 8 bits, we can use the fast
    // table lookup to get its value.  It's more than 8 bits about
    // 3-4% of the time.
    let rv: i32;
    code >>= 6;
    let val = self.numbits[code as usize] as i32;
    let mut l = val & 15;
    if l != 0 {
      pump.consume_bits(l as u32);
      rv = val >> 4;
    } else {
      pump.consume_bits(8);
      l = 8;
      while code > self.maxcode[l as usize] {
        let temp = pump.get_bits(1) as i32;
        code = (code << 1) | temp;
        l += 1;
      }

      // With garbage input we may reach the sentinel value l = 17.
      if l > self.precision as i32 || self.valptr[l as usize] == 0xff {
        return Err(format!("ljpeg: bad Huffman code: {}", l).to_string())
      } else {
        rv = self.huffval[
          self.valptr[l as usize] as usize +
          (code - (self.mincode[l as usize] as i32)) as usize
        ] as i32;
      }
    }

    if rv == 16 {
      if self.dng_compatible {
        pump.consume_bits(16);
      }
      return Ok(-32768);
    }

    // Section F.2.2.1: decode the difference and
    // Figure F.12: extend sign bit
    if rv != 0 {
      let mut x: i32 = pump.get_bits(rv as u32) as i32;
      if (x & (1 << (rv - 1))) == 0 {
        x -= (1 << rv) - 1;
      }
      return Ok(x)
    }

    return Ok(0)
  }

  pub fn huff_len(&self, pump: &mut BitPump) -> Result<i32,String> {
    let mut code: i32 = pump.peek_bits(8) as i32;
    let val = self.numbits[code as usize] as i32;
    let mut l = val & 15;
    //println!("code {} val {} l {}", code, val, l);
    if l != 0 {
      pump.consume_bits(l as u32);
      return Ok(val >> 4)
    }
    pump.consume_bits(8);
    l = 8;
    while code > self.maxcode[l as usize] {
      let temp = pump.get_bits(1) as i32;
      code = (code << 1) | temp;
      l += 1;
    }

    // With garbage input we may reach the sentinel value l = 17.
    if l > self.precision as i32 || self.valptr[l as usize] == 0xff {
      return Err(format!("ljpeg: bad Huffman code: {}", l).to_string())
    } else {
      return Ok(self.huffval[
        self.valptr[l as usize] as usize +
        (code - (self.mincode[l as usize] as i32)) as usize
      ] as i32);
    }
  }
}

impl fmt::Debug for HuffTable {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    if self.initialized {
      write!(f, "HuffTable {{ bits: {:?} huffval: {:?} }}", self.bits, &self.huffval[..])
    } else {
      write!(f, "HuffTable {{ uninitialized }}")
    }
  }
}

<commit_msg>ljpeg: reuse length code<commit_after>/*
* Huffman table generation:
* HuffTable::huff_decode()
* HuffTable::initialize()
* and used data structures are originally grabbed from the IJG software,
* and adapted by Hubert Figuiere.
*
* Copyright (C) 1991, 1992, Thomas G. Lane.
* Part of the Independent JPEG Group's software.
* See the file Copyright for more details.
*
* Copyright (c) 1993 Brian C. Smith, The Regents of the University
* of California
* All rights reserved.
*
* Copyright (c) 1994 Kongji Huang and Brian C. Smith.
* Cornell University
* All rights reserved.
*
* Permission to use, copy, modify, and distribute this software and its
* documentation for any purpose, without fee, and without written agreement is
* hereby granted, provided that the above copyright notice and the following
* two paragraphs appear in all copies of this software.
*
* IN NO EVENT SHALL CORNELL UNIVERSITY BE LIABLE TO ANY PARTY FOR
* DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT
* OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION, EVEN IF CORNELL
* UNIVERSITY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
* CORNELL UNIVERSITY SPECIFICALLY DISCLAIMS ANY WARRANTIES,
* INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
* AND FITNESS FOR A PARTICULAR PURPOSE.  THE SOFTWARE PROVIDED HEREUNDER IS
* ON AN "AS IS" BASIS, AND CORNELL UNIVERSITY HAS NO OBLIGATION TO
* PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
*/

use std::fmt;
use decoders::basics::*;

const HUFF_BITMASK: [u32;32] = [0xffffffff, 0x7fffffff,
                                0x3fffffff, 0x1fffffff,
                                0x0fffffff, 0x07ffffff,
                                0x03ffffff, 0x01ffffff,
                                0x00ffffff, 0x007fffff,
                                0x003fffff, 0x001fffff,
                                0x000fffff, 0x0007ffff,
                                0x0003ffff, 0x0001ffff,
                                0x0000ffff, 0x00007fff,
                                0x00003fff, 0x00001fff,
                                0x00000fff, 0x000007ff,
                                0x000003ff, 0x000001ff,
                                0x000000ff, 0x0000007f,
                                0x0000003f, 0x0000001f,
                                0x0000000f, 0x00000007,
                                0x00000003, 0x00000001];

pub struct HuffTable {
  // These two fields directly represent the contents of a JPEG DHT marker
  pub bits: [u32;17],
  pub huffval: [u32;256],

  // The remaining fields are computed from the above to allow more
  // efficient coding and decoding and thus private
  mincode: [u16;17],
  maxcode: [i32;18],
  valptr: [i16;17],
  numbits: [usize;256],
  bigtable: Vec<i32>,
  precision: usize,
  pub use_bigtable: bool,
  pub dng_compatible: bool,
  pub initialized: bool,
}

impl HuffTable {
  pub fn empty(precision: usize) -> HuffTable {
    HuffTable {
      bits: [0;17],
      huffval: [0;256],
      mincode: [0;17],
      maxcode: [0;18],
      valptr: [0;17],
      numbits: [0;256],
      bigtable: Vec::new(),
      precision: precision,
      use_bigtable: true,
      dng_compatible: true,
      initialized: false,
    }
  }

  pub fn initialize(&mut self, use_bigtable: bool) -> Result<(), String> {
    // Figure C.1: make table of Huffman code length for each symbol
    // Note that this is in code-length order.
    let mut p = 0;
    let mut huffsize: [u8;257] = [0;257];
    for l in 1..17 {
      for _ in 1..((self.bits[l] as usize)+1) {
        huffsize[p] = l as u8;
        p += 1;
        if p > 256 {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
      }
    }
    huffsize[p] = 0;
    let lastp = p;

    // Figure C.2: generate the codes themselves
    // Note that this is in code-length order.
    let mut code: u16 = 0;
    let mut si: u32 = huffsize[0] as u32;
    let mut huffcode: [u16;257] = [0;257];
    p = 0;
    while huffsize[p] > 0 {
      while (huffsize[p] as u32) == si {
        huffcode[p] = code;
        p += 1;
        code += 1;
      }
      code <<= 1;
      si += 1;
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    //Figure F.15: generate decoding tables
    self.mincode[0] = 0;
    self.maxcode[0] = 0;
    p = 0;
    for l in 1..17 {
      if self.bits[l] > 0 {
        self.valptr[l] = p as i16;
        self.mincode[l] = huffcode[p];
        p += self.bits[l] as usize;
        self.maxcode[l] = huffcode[p - 1] as i32;
      } else {
        self.valptr[l] = 0xff;   // This check must be present to avoid crash on junk
        self.maxcode[l] = -1;
      }
      if p > 256 {
        return Err("ljpeg: Code length too long. Corrupt data.".to_string())
      }
    }

    // We put in this value to ensure HuffDecode terminates
    self.maxcode[17] = 0xFFFFF;

    // Build the numbits, value lookup tables.
    // These table allow us to gather 8 bits from the bits stream,
    // and immediately lookup the size and value of the huffman codes.
    // If size is zero, it means that more than 8 bits are in the huffman
    // code (this happens about 3-4% of the time).
    for p in 0..lastp {
      let size = huffsize[p];
      if size <= 8 {
        let value: i32 = self.huffval[p] as i32;
        let code = huffcode[p];
        let ll: i32 = (code << (8 - size)) as i32;
        let ul: i32 = if size < 8 {
          ll | (HUFF_BITMASK[(24+size) as usize]) as i32
        } else {
          ll
        };
        if ul > 256 || ll > ul {
          return Err("ljpeg: Code length too long. Corrupt data.".to_string())
        }
        for i in ll..(ul+1) {
          self.numbits[i as usize] = (size as usize) | ((value as usize) << 4);
        }
      }
    }

    if use_bigtable {
      try!(self.initialize_bigtable());
    }
    self.initialized = true;
    self.use_bigtable = use_bigtable;
    Ok(())
  }

  fn initialize_bigtable(&mut self) -> Result<(), String> {
    let bits: usize = 14; // HuffDecode functions must be changed, if this is modified.
    let size: usize = 1 << bits;

    self.bigtable = vec![0 as i32;size];
    let mut rv: i32;
    for i in 0..size {
      let input = (i << 2) as u16;
      let mut code: i32 = (input >> 8) as i32;
      let val: u32 = self.numbits[code as usize] as u32;
      let mut l: u32 = val & 15;
      if l > 0 {
        rv = (val >> 4) as i32;
      } else {
        l = 8;
        while code > self.maxcode[l as usize] && l < 16 {
          let temp: i32 = (input >> (15 - l) & 1) as i32;
          code = (code << 1) | temp;
          l += 1;
        }

        //With garbage input we may reach the sentinel value l = 17.
        if l > self.precision as u32 || self.valptr[l as usize] == 0xff {
          self.bigtable[i] = 0xff;
          continue
        } else {
          rv = self.huffval[
            self.valptr[l as usize] as usize +
            (code - (self.mincode[l as usize] as i32)) as usize
          ] as i32;
        }
      }

      if rv == 16 {
        self.bigtable[i] = if self.dng_compatible {
          (-(32768 << 8)) | (16 + l as i32)
        } else {
          (-(32768 << 8)) | l as i32
        };
        continue
      }

      if rv + l as i32 > bits as i32 {
        self.bigtable[i] = 0xff;
        continue
      }

      if rv != 0 {
        let mut x = (input >> (16 - (l as i32) - rv) & ((1 << rv) - 1)) as i32;
        if (x & (1 << (rv - 1))) == 0 {
          x -= (1 << rv) - 1;
        }
        self.bigtable[i] = (x << 8) | ((l as i32) + rv);
      } else {
        self.bigtable[i] = l as i32;
      }
    }
    self.use_bigtable = true;
    Ok(())
  }

  // Taken from Figure F.16: extract next coded symbol from input stream
  pub fn huff_decode(&self, pump: &mut BitPump) -> Result<i32,String> {
    if !self.initialized {
      return Err("ljpeg: Trying to read from uninitialized HuffTable".to_string());
    }

    //First attempt to do complete decode, by using the first 14 bits
    let code: i32 = pump.peek_bits(14) as i32;
    if self.use_bigtable {
      let val: i32 = self.bigtable[code as usize];
      if val & 0xff != 0xff {
        pump.consume_bits((val & 0xff) as u32);
        return Ok(val >> 8)
      }
    }

    let rv = try!(self.huff_len(pump));
    if rv == 16 {
      if self.dng_compatible {
        pump.consume_bits(16);
      }
      return Ok(-32768);
    }

    // Section F.2.2.1: decode the difference and
    // Figure F.12: extend sign bit
    if rv != 0 {
      let mut x: i32 = pump.get_bits(rv as u32) as i32;
      if (x & (1 << (rv - 1))) == 0 {
        x -= (1 << rv) - 1;
      }
      return Ok(x)
    }

    return Ok(0)
  }

  pub fn huff_len(&self, pump: &mut BitPump) -> Result<i32,String> {
    let mut code: i32 = pump.peek_bits(8) as i32;
    let val = self.numbits[code as usize] as i32;
    let mut l = val & 15;
    //println!("code {} val {} l {}", code, val, l);
    if l != 0 {
      pump.consume_bits(l as u32);
      return Ok(val >> 4)
    }
    pump.consume_bits(8);
    l = 8;
    while code > self.maxcode[l as usize] {
      let temp = pump.get_bits(1) as i32;
      code = (code << 1) | temp;
      l += 1;
    }

    // With garbage input we may reach the sentinel value l = 17.
    if l > self.precision as i32 || self.valptr[l as usize] == 0xff {
      return Err(format!("ljpeg: bad Huffman code: {}", l).to_string())
    } else {
      return Ok(self.huffval[
        self.valptr[l as usize] as usize +
        (code - (self.mincode[l as usize] as i32)) as usize
      ] as i32);
    }
  }
}

impl fmt::Debug for HuffTable {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    if self.initialized {
      write!(f, "HuffTable {{ bits: {:?} huffval: {:?} }}", self.bits, &self.huffval[..])
    } else {
      write!(f, "HuffTable {{ uninitialized }}")
    }
  }
}

<|endoftext|>


<commit_before>use backend::common::{self, Arguments};
use gtk;
use gdk::enums::key;
use gtk::prelude::*;
use gtk::{
    Builder, Button, CheckButton, Entry, FileChooserDialog, ListBox, ListBoxRow, ListStore,
    SpinButton, TreeView, TreeViewColumn, Type, Window, WindowType
};
use std::fs;
use std::path::PathBuf;

// TODO:
// - Implement GTK3 notifications for when errors occur, and when renaming has completed successfully
// - Add a tab for checking the rename log.
// - Add support for a configuration file.
// - Reduce source code redundancy with macros and functions.

pub fn launch() {
    gtk::init().unwrap_or_else(|_| panic!("tv-renamer: failed to initialize GTK."));

    // Open the Glade GTK UI and import key GTK objects from the UI.
    let builder = Builder::new_from_string(include_str!("gtk_interface.glade"));
    let window: Window                  = builder.get_object("main_window").unwrap();
    let preview_button: Button          = builder.get_object("preview_button").unwrap();
    let rename_button: Button           = builder.get_object("rename_button").unwrap();
    let input_list: ListBox             = builder.get_object("input_list").unwrap();
    let series_name_entry: Entry        = builder.get_object("series_name_entry").unwrap();
    let series_directory_entry: Entry   = builder.get_object("series_directory_entry").unwrap();
    let series_directory_button: Button = builder.get_object("series_directory_button").unwrap();
    let episode_spin_button: SpinButton = builder.get_object("episode_spin_button").unwrap();
    let season_spin_button: SpinButton  = builder.get_object("season_spin_button").unwrap();
    let preview_tree: TreeView          = builder.get_object("preview_tree").unwrap();
    let info_bar: gtk::InfoBar          = builder.get_object("info_bar").unwrap();
    let info_button: Button             = builder.get_object("info_close").unwrap();
    let notification_label: gtk::Label  = builder.get_object("notification_label").unwrap();

    // Create rows for the input_list
    let automatic_row   = ListBoxRow::new();
    let log_changes_row = ListBoxRow::new();
    let series_name_row = ListBoxRow::new();
    let tvdb_row        = ListBoxRow::new();
    automatic_row.set_selectable(false);
    log_changes_row.set_selectable(false);
    series_name_row.set_selectable(false);
    tvdb_row.set_selectable(false);

    // Create check boxes for the rows
    let automatic_check   = CheckButton::new_with_label("Automatic");
    let no_name_check     = CheckButton::new_with_label("No Series Name");
    let tvdb_check        = CheckButton::new_with_label("TVDB Titles");
    let log_changes_check = CheckButton::new_with_label("Log Changes");

    // Add the check boxes to the rows
    automatic_row.add(&automatic_check);
    series_name_row.add(&no_name_check);
    tvdb_row.add(&tvdb_check);
    log_changes_row.add(&log_changes_check);

    // Add the rows to the list box
    input_list.insert(&automatic_row, -1);
    input_list.insert(&series_name_row, -1);
    input_list.insert(&tvdb_row, -1);
    input_list.insert(&log_changes_row, -1);

    // TreeView's List Store
    // Link these up to the preview_tree and then start renaming
    let preview_list = ListStore::new(&[
        Type::String, // Before
        Type::String, // After
    ]);

    // Create and append the Before column to the preview tree
    let before_column = TreeViewColumn::new();
    let renderer      = gtk::CellRendererText::new();
    before_column.set_title("Before");
    before_column.set_resizable(true);
    before_column.pack_start(&renderer, true);
    before_column.add_attribute(&renderer, "text", 0);
    preview_tree.append_column(&before_column);

    // Create and append the After column to the preview tree
    let after_column = TreeViewColumn::new();
    let renderer     = gtk::CellRendererText::new();
    after_column.set_title("After");
    after_column.set_resizable(true);
    after_column.pack_start(&renderer, true);
    after_column.add_attribute(&renderer, "text", 1);
    preview_tree.append_column(&after_column);

    // Connect the preview_list to the preview tree
    preview_tree.set_model(Some(&preview_list));
    preview_tree.set_headers_visible(true);

    // A simple macro that is shared among all widgets that trigger the action to update the preview.
    macro_rules! gtk_preview {
        ($widget:ident) => {{
            let $widget             = $widget.clone();
            let auto                = automatic_check.clone();
            let no_name             = no_name_check.clone();
            let tvdb                = tvdb_check.clone();
            let log_changes         = log_changes_check.clone();
            let season_spin_button  = season_spin_button.clone();
            let episode_spin_button = episode_spin_button.clone();
            let series_entry        = series_name_entry.clone();
            let directory_entry     = series_directory_entry.clone();
            let preview_list        = preview_list.clone();
            let info_bar            = info_bar.clone();
            let notification_label  = notification_label.clone();
            $widget.connect_clicked(move |_| {
                if let Some(directory) = directory_entry.get_text() {
                    let mut program = &mut Arguments {
                        automatic:     auto.get_active(),
                        dry_run:       true,
                        log_changes:   log_changes.get_active(),
                        no_name:       no_name.get_active(),
                        tvdb:          tvdb.get_active(),
                        verbose:       false,
                        directory:     directory,
                        series_name:   series_entry.get_text().unwrap_or_default(),
                        season_number: season_spin_button.get_value_as_int() as usize,
                        episode_count: episode_spin_button.get_value_as_int() as usize,
                        pad_length:    2,
                    };

                    if !program.directory.is_empty() { program.update_preview(&preview_list, &info_bar, &notification_label); }
                }
            });
        }}
    }

    // All of the widgets that implement the update preview action
    gtk_preview!(no_name_check);
    gtk_preview!(automatic_check);
    gtk_preview!(tvdb_check);
    gtk_preview!(preview_button);

    { // Hide the Info Bar when the Info Bar is closed
        let info_bar = info_bar.clone();
        info_button.connect_clicked(move |_| {
            info_bar.hide();
        });
    }

    { // NOTE: Programs the Choose Directory button with a File Chooser Dialog.
        let auto                = automatic_check.clone();
        let no_name             = no_name_check.clone();
        let tvdb                = tvdb_check.clone();
        let log_changes         = log_changes_check.clone();
        let season_spin_button  = season_spin_button.clone();
        let episode_spin_button = episode_spin_button.clone();
        let series_entry        = series_name_entry.clone();
        let directory_entry     = series_directory_entry.clone();
        let preview_list        = preview_list.clone();
        let info_bar            = info_bar.clone();
        let notification_label  = notification_label.clone();
        series_directory_button.connect_clicked(move |_| {
            // Open file chooser dialog to modify series_directory_entry.
            let dialog = FileChooserDialog::new (
                Some("Choose Directory"),
                Some(&Window::new(WindowType::Popup)),
                gtk::FileChooserAction::SelectFolder,
            );
            dialog.add_button("Cancel", gtk::ResponseType::Cancel as i32);
            dialog.add_button("OK", gtk::ResponseType::Ok as i32);

            if dialog.run() == gtk::ResponseType::Ok as i32 {
                if let Some(path) = dialog.get_filename() {
                    if let Some(text) = path.to_str() {
                        directory_entry.set_text(text);
                    }
                }
            }
            dialog.destroy();

            if let Some(directory) = directory_entry.get_text() {
                let mut program = &mut Arguments {
                    automatic:     auto.get_active(),
                    dry_run:       false,
                    log_changes:   log_changes.get_active(),
                    no_name:       no_name.get_active(),
                    tvdb:          tvdb.get_active(),
                    verbose:       false,
                    directory:     directory,
                    series_name:   series_entry.get_text().unwrap_or_default(),
                    season_number: season_spin_button.get_value_as_int() as usize,
                    episode_count: episode_spin_button.get_value_as_int() as usize,
                    pad_length:    2,
                };

                if !program.directory.is_empty() { program.update_preview(&preview_list, &info_bar, &notification_label); }
            }
        });
    }

    { // NOTE: Controls what happens when rename button is pressed
        let button              = rename_button.clone();
        let auto                = automatic_check.clone();
        let no_name             = no_name_check.clone();
        let tvdb                = tvdb_check.clone();
        let log_changes         = log_changes_check.clone();
        let season_spin_button  = season_spin_button.clone();
        let episode_spin_button = episode_spin_button.clone();
        let series_entry        = series_name_entry.clone();
        let directory_entry     = series_directory_entry.clone();
        let preview_list        = preview_list.clone();
        let info_bar            = info_bar.clone();
        let notification_label  = notification_label.clone();
        button.connect_clicked(move |_| {
            if let Some(directory) = directory_entry.get_text() {
                let mut program = &mut Arguments {
                    automatic:     auto.get_active(),
                    dry_run:       false,
                    log_changes:   log_changes.get_active(),
                    no_name:       no_name.get_active(),
                    tvdb:          tvdb.get_active(),
                    verbose:       false,
                    directory:     directory,
                    series_name:   series_entry.get_text().unwrap_or_default(),
                    season_number: season_spin_button.get_value_as_int() as usize,
                    episode_count: episode_spin_button.get_value_as_int() as usize,
                    pad_length:    2,
                };

                if !program.directory.is_empty() {
                    program.rename_series(&preview_list, &info_bar, &notification_label);
                }
            }
        });
    }

    window.show_all();
    info_bar.hide();

    // Quit the program when the program has been exited
    window.connect_delete_event(|_, _| {
        gtk::main_quit();
        Inhibit(false)
    });

    // Define custom actions on keypress
    window.connect_key_press_event(move |_, key| {
        if let key::Escape = key.get_keyval() { gtk::main_quit() }
        gtk::Inhibit(false)
    });

    gtk::main();

}

impl Arguments {
    fn update_preview(&mut self, preview_list: &ListStore, info_bar: &gtk::InfoBar, notification_label: &gtk::Label) {
        preview_list.clear();
        if self.automatic {
            let series = PathBuf::from(&self.directory);
            self.series_name = series.components().last().unwrap().as_os_str().to_str().unwrap().to_string();
            match common::get_seasons(&self.directory) {
                Ok(seasons) => {
                    for season in seasons {
                        match common::derive_season_number(&season) {
                            Some(number) => self.season_number = number,
                            None         => continue
                        }
                        if let Some(error) = self.rename_episodes(season.as_os_str().to_str().unwrap(), preview_list) {
                            info_bar.set_message_type(gtk::MessageType::Error);
                            notification_label.set_text(&error);
                            info_bar.show();
                        }
                    }
                },
                Err(err) => {
                    info_bar.set_message_type(gtk::MessageType::Error);
                    notification_label.set_text(err);
                    info_bar.show();
                }
            }
        } else if let Some(error) = self.rename_episodes(&self.directory, preview_list) {
            info_bar.set_message_type(gtk::MessageType::Error);
            notification_label.set_text(&error);
            info_bar.show();
        }
    }

    fn rename_series(&mut self, preview_list: &ListStore, info_bar: &gtk::InfoBar, notification_label: &gtk::Label) {
        preview_list.clear();
        if self.automatic {
            let series = PathBuf::from(&self.directory);
            self.series_name = series.components().last().unwrap().as_os_str().to_str().unwrap().to_string();
            match common::get_seasons(&self.directory) {
                Ok(seasons) => {
                    for season in seasons {
                        match common::derive_season_number(&season) {
                            Some(number) => self.season_number = number,
                            None         => continue
                        }
                        if let Some(error) = self.rename_episodes(season.as_os_str().to_str().unwrap(), preview_list) {
                            info_bar.set_message_type(gtk::MessageType::Error);
                            notification_label.set_text(&error);
                        } else {
                            info_bar.set_message_type(gtk::MessageType::Info);
                            notification_label.set_text("Rename Success");
                        }
                    }
                },
                Err(err) => {
                    info_bar.set_message_type(gtk::MessageType::Error);
                    notification_label.set_text(err);
                }
            }
        } else if let Some(error) = self.rename_episodes(&self.directory, preview_list) {
            info_bar.set_message_type(gtk::MessageType::Error);
            notification_label.set_text(&error);
        } else {
            info_bar.set_message_type(gtk::MessageType::Info);
            notification_label.set_text("Rename Success");
        }
        info_bar.show();
    }

    fn rename_episodes(&self, directory: &str, preview_list: &ListStore) -> Option<String> {
        match common::get_episodes(directory) {
            Ok(episodes) => {
                match self.get_targets(directory, &episodes, self.episode_count) {
                    Ok(targets) => {
                        if self.log_changes { common::log_append_time(); }
                        let mut error_occurred = false;
                        for (source, target) in episodes.iter().zip(targets) {
                            if !self.dry_run {
                                if fs::rename(&source, &target).is_err() { error_occurred = true; };
                                if self.log_changes { common::log_append_change(source.as_path(), target.as_path()); }
                            }

                            // Update the preview
                            let source = source.components().last().unwrap().as_os_str().to_str().unwrap().to_string();
                            let target = target.components().last().unwrap().as_os_str().to_str().unwrap().to_string();
                            preview_list.insert_with_values(None, &[0, 1], &[&source, &target]);
                        }
                        if error_occurred { Some(String::from("Rename Failed")) } else { None }
                    },
                    Err(err) => Some(err)
                }
            },
            Err(err) => Some(String::from(err))
        }
    }
}
<commit_msg>Drag-and-Drop Directory Entry Support<commit_after>use backend::common::{self, Arguments};
use gtk;
use gdk::enums::key;
use gtk::prelude::*;
use gtk::{
    Builder, Button, CheckButton, Entry, FileChooserDialog, ListBox, ListBoxRow, ListStore,
    SpinButton, TreeView, TreeViewColumn, Type, Window, WindowType
};
use std::fs;
use std::path::PathBuf;

// TODO:
// - Add a tab for checking the rename log.
// - Add support for a configuration file.

/// Allow drag-and-drop support in the directory entry text field by fixing the URI generated by dropped files.
#[inline]
fn parse_directory(directory: &str) -> String {
    let mut output = String::from(directory);
    if output.starts_with("file://") {
        output = output.replace("file://", "");
        output = output.replace("%20", " ");
        output = output.replace("\n", "");
        output = output.replace("\r", "");
    }
    output
}

pub fn launch() {
    gtk::init().unwrap_or_else(|_| panic!("tv-renamer: failed to initialize GTK."));

    // Open the Glade GTK UI and import key GTK objects from the UI.
    let builder = Builder::new_from_string(include_str!("gtk_interface.glade"));
    let window: Window                  = builder.get_object("main_window").unwrap();
    let preview_button: Button          = builder.get_object("preview_button").unwrap();
    let rename_button: Button           = builder.get_object("rename_button").unwrap();
    let input_list: ListBox             = builder.get_object("input_list").unwrap();
    let series_name_entry: Entry        = builder.get_object("series_name_entry").unwrap();
    let series_directory_entry: Entry   = builder.get_object("series_directory_entry").unwrap();
    let series_directory_button: Button = builder.get_object("series_directory_button").unwrap();
    let episode_spin_button: SpinButton = builder.get_object("episode_spin_button").unwrap();
    let season_spin_button: SpinButton  = builder.get_object("season_spin_button").unwrap();
    let preview_tree: TreeView          = builder.get_object("preview_tree").unwrap();
    let info_bar: gtk::InfoBar          = builder.get_object("info_bar").unwrap();
    let info_button: Button             = builder.get_object("info_close").unwrap();
    let notification_label: gtk::Label  = builder.get_object("notification_label").unwrap();

    // Create rows for the input_list
    let automatic_row   = ListBoxRow::new();
    let log_changes_row = ListBoxRow::new();
    let series_name_row = ListBoxRow::new();
    let tvdb_row        = ListBoxRow::new();
    automatic_row.set_selectable(false);
    log_changes_row.set_selectable(false);
    series_name_row.set_selectable(false);
    tvdb_row.set_selectable(false);

    // Create check boxes for the rows
    let automatic_check   = CheckButton::new_with_label("Automatic");
    let no_name_check     = CheckButton::new_with_label("No Series Name");
    let tvdb_check        = CheckButton::new_with_label("TVDB Titles");
    let log_changes_check = CheckButton::new_with_label("Log Changes");

    // Add the check boxes to the rows
    automatic_row.add(&automatic_check);
    series_name_row.add(&no_name_check);
    tvdb_row.add(&tvdb_check);
    log_changes_row.add(&log_changes_check);

    // Add the rows to the list box
    input_list.insert(&automatic_row, -1);
    input_list.insert(&series_name_row, -1);
    input_list.insert(&tvdb_row, -1);
    input_list.insert(&log_changes_row, -1);

    // TreeView's List Store
    // Link these up to the preview_tree and then start renaming
    let preview_list = ListStore::new(&[
        Type::String, // Before
        Type::String, // After
    ]);

    // A simple macro for adding a column to the preview tree.
    macro_rules! add_column {
        ($preview_tree:ident, $title:expr, $id:expr) => {{
            let column   = TreeViewColumn::new();
            let renderer = gtk::CellRendererText::new();
            column.set_title($title);
            column.set_resizable(true);
            column.pack_start(&renderer, true);
            column.add_attribute(&renderer, "text", $id);
            preview_tree.append_column(&column);
        }}
    }

    // Create and append the Before column to the preview tree
    add_column!(preview_tree, "Before", 0);
    add_column!(preview_tree, "After", 1);

    // Connect the preview_list to the preview tree
    preview_tree.set_model(Some(&preview_list));
    preview_tree.set_headers_visible(true);

    // A simple macro that is shared among all widgets that trigger the action to update the preview.
    macro_rules! gtk_preview {
        ($widget:ident) => {{
            let $widget             = $widget.clone();
            let auto                = automatic_check.clone();
            let no_name             = no_name_check.clone();
            let tvdb                = tvdb_check.clone();
            let log_changes         = log_changes_check.clone();
            let season_spin_button  = season_spin_button.clone();
            let episode_spin_button = episode_spin_button.clone();
            let series_entry        = series_name_entry.clone();
            let directory_entry     = series_directory_entry.clone();
            let preview_list        = preview_list.clone();
            let info_bar            = info_bar.clone();
            let notification_label  = notification_label.clone();
            $widget.connect_clicked(move |_| {
                if let Some(directory) = directory_entry.get_text() {
                    let mut program = &mut Arguments {
                        automatic:     auto.get_active(),
                        dry_run:       true,
                        log_changes:   log_changes.get_active(),
                        no_name:       no_name.get_active(),
                        tvdb:          tvdb.get_active(),
                        verbose:       false,
                        directory:     parse_directory(&directory),
                        series_name:   series_entry.get_text().unwrap_or_default(),
                        season_number: season_spin_button.get_value_as_int() as usize,
                        episode_count: episode_spin_button.get_value_as_int() as usize,
                        pad_length:    2,
                    };

                    if !program.directory.is_empty() {
                        program.update_preview(&preview_list, &info_bar, &notification_label);
                    }
                }
            });
        }}
    }

    // All of the widgets that implement the update preview action
    gtk_preview!(no_name_check);
    gtk_preview!(automatic_check);
    gtk_preview!(tvdb_check);
    gtk_preview!(preview_button);

    { // Hide the Info Bar when the Info Bar is closed
        let info_bar = info_bar.clone();
        info_button.connect_clicked(move |_| {
            info_bar.hide();
        });
    }

    { // NOTE: Programs the Choose Directory button with a File Chooser Dialog.
        let auto                = automatic_check.clone();
        let no_name             = no_name_check.clone();
        let tvdb                = tvdb_check.clone();
        let log_changes         = log_changes_check.clone();
        let season_spin_button  = season_spin_button.clone();
        let episode_spin_button = episode_spin_button.clone();
        let series_entry        = series_name_entry.clone();
        let directory_entry     = series_directory_entry.clone();
        let preview_list        = preview_list.clone();
        let info_bar            = info_bar.clone();
        let notification_label  = notification_label.clone();
        series_directory_button.connect_clicked(move |_| {
            // Open file chooser dialog to modify series_directory_entry.
            let dialog = FileChooserDialog::new (
                Some("Choose Directory"),
                Some(&Window::new(WindowType::Popup)),
                gtk::FileChooserAction::SelectFolder,
            );
            dialog.add_button("Cancel", gtk::ResponseType::Cancel as i32);
            dialog.add_button("OK", gtk::ResponseType::Ok as i32);

            if dialog.run() == gtk::ResponseType::Ok as i32 {
                if let Some(path) = dialog.get_filename() {
                    if let Some(text) = path.to_str() {
                        directory_entry.set_text(text);
                    }
                }
            }
            dialog.destroy();

            if let Some(directory) = directory_entry.get_text() {
                let mut program = &mut Arguments {
                    automatic:     auto.get_active(),
                    dry_run:       true,
                    log_changes:   log_changes.get_active(),
                    no_name:       no_name.get_active(),
                    tvdb:          tvdb.get_active(),
                    verbose:       false,
                    directory:     parse_directory(&directory),
                    series_name:   series_entry.get_text().unwrap_or_default(),
                    season_number: season_spin_button.get_value_as_int() as usize,
                    episode_count: episode_spin_button.get_value_as_int() as usize,
                    pad_length:    2,
                };

                if !program.directory.is_empty() { program.update_preview(&preview_list, &info_bar, &notification_label); }
            }
        });
    }

    { // NOTE: Controls what happens when rename button is pressed
        let button              = rename_button.clone();
        let auto                = automatic_check.clone();
        let no_name             = no_name_check.clone();
        let tvdb                = tvdb_check.clone();
        let log_changes         = log_changes_check.clone();
        let season_spin_button  = season_spin_button.clone();
        let episode_spin_button = episode_spin_button.clone();
        let series_entry        = series_name_entry.clone();
        let directory_entry     = series_directory_entry.clone();
        let preview_list        = preview_list.clone();
        let info_bar            = info_bar.clone();
        let notification_label  = notification_label.clone();
        button.connect_clicked(move |_| {
            if let Some(directory) = directory_entry.get_text() {
                let mut program = &mut Arguments {
                    automatic:     auto.get_active(),
                    dry_run:       false,
                    log_changes:   log_changes.get_active(),
                    no_name:       no_name.get_active(),
                    tvdb:          tvdb.get_active(),
                    verbose:       false,
                    directory:     parse_directory(&directory),
                    series_name:   series_entry.get_text().unwrap_or_default(),
                    season_number: season_spin_button.get_value_as_int() as usize,
                    episode_count: episode_spin_button.get_value_as_int() as usize,
                    pad_length:    2,
                };

                if !program.directory.is_empty() {
                    program.rename_series(&preview_list, &info_bar, &notification_label);
                }
            }
        });
    }

    window.show_all();
    info_bar.hide();

    // Quit the program when the program has been exited
    window.connect_delete_event(|_, _| {
        gtk::main_quit();
        Inhibit(false)
    });

    // Define custom actions on keypress
    window.connect_key_press_event(move |_, key| {
        if let key::Escape = key.get_keyval() { gtk::main_quit() }
        gtk::Inhibit(false)
    });

    gtk::main();

}

impl Arguments {
    fn update_preview(&mut self, preview_list: &ListStore, info_bar: &gtk::InfoBar, notification_label: &gtk::Label) {
        preview_list.clear();
        if self.automatic {
            let series = PathBuf::from(&self.directory);
            self.series_name = series.components().last().unwrap().as_os_str().to_str().unwrap().to_string();
            match common::get_seasons(&self.directory) {
                Ok(seasons) => {
                    for season in seasons {
                        match common::derive_season_number(&season) {
                            Some(number) => self.season_number = number,
                            None         => continue
                        }
                        if let Some(error) = self.rename_episodes(season.as_os_str().to_str().unwrap(), preview_list) {
                            info_bar.set_message_type(gtk::MessageType::Error);
                            notification_label.set_text(&error);
                            info_bar.show();
                        }
                    }
                },
                Err(err) => {
                    info_bar.set_message_type(gtk::MessageType::Error);
                    notification_label.set_text(err);
                    info_bar.show();
                }
            }
        } else if let Some(error) = self.rename_episodes(&self.directory, preview_list) {
            info_bar.set_message_type(gtk::MessageType::Error);
            notification_label.set_text(&error);
            info_bar.show();
        }
    }

    fn rename_series(&mut self, preview_list: &ListStore, info_bar: &gtk::InfoBar, notification_label: &gtk::Label) {
        preview_list.clear();
        if self.automatic {
            let series = PathBuf::from(&self.directory);
            self.series_name = series.components().last().unwrap().as_os_str().to_str().unwrap().to_string();
            match common::get_seasons(&self.directory) {
                Ok(seasons) => {
                    for season in seasons {
                        match common::derive_season_number(&season) {
                            Some(number) => self.season_number = number,
                            None         => continue
                        }
                        if let Some(error) = self.rename_episodes(season.as_os_str().to_str().unwrap(), preview_list) {
                            info_bar.set_message_type(gtk::MessageType::Error);
                            notification_label.set_text(&error);
                        } else {
                            info_bar.set_message_type(gtk::MessageType::Info);
                            notification_label.set_text("Rename Success");
                        }
                    }
                },
                Err(err) => {
                    info_bar.set_message_type(gtk::MessageType::Error);
                    notification_label.set_text(err);
                }
            }
        } else if let Some(error) = self.rename_episodes(&self.directory, preview_list) {
            info_bar.set_message_type(gtk::MessageType::Error);
            notification_label.set_text(&error);
        } else {
            info_bar.set_message_type(gtk::MessageType::Info);
            notification_label.set_text("Rename Success");
        }
        info_bar.show();
    }

    fn rename_episodes(&self, directory: &str, preview_list: &ListStore) -> Option<String> {
        match common::get_episodes(directory) {
            Ok(episodes) => {
                match self.get_targets(directory, &episodes, self.episode_count) {
                    Ok(targets) => {
                        if self.log_changes { common::log_append_time(); }
                        let mut error_occurred = false;
                        for (source, target) in episodes.iter().zip(targets) {
                            if !self.dry_run {
                                if fs::rename(&source, &target).is_err() { error_occurred = true; };
                                if self.log_changes { common::log_append_change(source.as_path(), target.as_path()); }
                            }

                            // Update the preview
                            let source = source.components().last().unwrap().as_os_str().to_str().unwrap().to_string();
                            let target = target.components().last().unwrap().as_os_str().to_str().unwrap().to_string();
                            preview_list.insert_with_values(None, &[0, 1], &[&source, &target]);
                        }
                        if error_occurred { Some(String::from("Rename Failed")) } else { None }
                    },
                    Err(err) => Some(err)
                }
            },
            Err(err) => Some(String::from(err))
        }
    }
}
<|endoftext|>


<commit_before>use crate::postings::compression::COMPRESSION_BLOCK_SIZE;

unsafe fn binary_search_step(ptr: *const u32, target: u32, half_size: isize) -> *const u32 {
    let mid = ptr.offset(half_size);
    if *mid < target {
        mid.offset(1)
    } else {
        ptr
    }
}

/// Search the first index containing an element greater or equal to
/// the target.
///
/// The results should be equivalent to
/// ```compile_fail
/// block[..]
//       .iter()
//       .take_while(|&&val| val < target)
//       .count()
/// ```
///
/// the `start` argument is just used to hint that the response is
/// greater than beyond `start`. the implementation may or may not use
/// it for optimization.
///
/// # Assumption
///
/// - The block is sorted. Some elements may appear several times. This is the case at the
/// end of the last block for instance.
/// - The target is assumed smaller or equal to the last element of the block.
pub fn branchless_binary_search(arr: &[u32; COMPRESSION_BLOCK_SIZE], target: u32) -> usize {
    let start_ptr: *const u32 = &arr[0] as *const u32;
    unsafe {
        let mut ptr = start_ptr;
        ptr = binary_search_step(ptr, target, 63);
        ptr = binary_search_step(ptr, target, 31);
        ptr = binary_search_step(ptr, target, 15);
        ptr = binary_search_step(ptr, target, 7);
        ptr = binary_search_step(ptr, target, 3);
        ptr = binary_search_step(ptr, target, 1);
        let extra = if *ptr < target { 1 } else { 0 };
        (ptr.offset_from(start_ptr) as usize) + extra
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;
    use super::branchless_binary_search;
    use crate::docset::TERMINATED;
    use crate::postings::compression::COMPRESSION_BLOCK_SIZE;
    use proptest::prelude::*;

    fn search_in_block_trivial_but_slow(block: &[u32], target: u32) -> usize {
        block.iter().take_while(|&&val| val < target).count()
    }

    fn util_test_search_in_block(block: &[u32], target: u32) {
        let cursor = search_in_block_trivial_but_slow(block, target);
        assert!(cursor < COMPRESSION_BLOCK_SIZE);
        assert!(block[cursor] >= target);
        if cursor > 0 {
            assert!(block[cursor - 1] < target);
        }
        assert_eq!(block.len(), COMPRESSION_BLOCK_SIZE);
        let mut output_buffer = [TERMINATED; COMPRESSION_BLOCK_SIZE];
        output_buffer[..block.len()].copy_from_slice(block);
        assert_eq!(branchless_binary_search(&output_buffer, target), cursor);
    }

    fn util_test_search_in_block_all(block: &[u32]) {
        let mut targets = HashSet::new();
        targets.insert(0);
        for &val in block {
            if val > 0 {
                targets.insert(val - 1);
            }
            targets.insert(val);
        }
        for target in targets {
            util_test_search_in_block(block, target);
        }
    }

    #[test]
    fn test_search_in_branchless_binary_search() {
        let v: Vec<u32> = (0..COMPRESSION_BLOCK_SIZE).map(|i| i as u32 * 2).collect();
        util_test_search_in_block_all(&v[..]);
    }

    fn monotonous_block() -> impl Strategy<Value = Vec<u32>> {
        prop::collection::vec(0u32..5u32, COMPRESSION_BLOCK_SIZE).prop_map(|mut deltas| {
            let mut el = 0;
            for i in 0..COMPRESSION_BLOCK_SIZE {
                el += deltas[i];
                deltas[i] = el;
            }
            deltas
        })
    }

    proptest! {
        #[test]
        fn test_proptest_branchless_binary_search(block in monotonous_block()) {
            util_test_search_in_block_all(&block[..]);
        }
    }
}
<commit_msg>Cargo fmt<commit_after>use crate::postings::compression::COMPRESSION_BLOCK_SIZE;

unsafe fn binary_search_step(ptr: *const u32, target: u32, half_size: isize) -> *const u32 {
    let mid = ptr.offset(half_size);
    if *mid < target {
        mid.offset(1)
    } else {
        ptr
    }
}

/// Search the first index containing an element greater or equal to
/// the target.
///
/// The results should be equivalent to
/// ```compile_fail
/// block[..]
//       .iter()
//       .take_while(|&&val| val < target)
//       .count()
/// ```
///
/// the `start` argument is just used to hint that the response is
/// greater than beyond `start`. the implementation may or may not use
/// it for optimization.
///
/// # Assumption
///
/// - The block is sorted. Some elements may appear several times. This is the case at the
/// end of the last block for instance.
/// - The target is assumed smaller or equal to the last element of the block.
pub fn branchless_binary_search(arr: &[u32; COMPRESSION_BLOCK_SIZE], target: u32) -> usize {
    let start_ptr: *const u32 = &arr[0] as *const u32;
    unsafe {
        let mut ptr = start_ptr;
        ptr = binary_search_step(ptr, target, 63);
        ptr = binary_search_step(ptr, target, 31);
        ptr = binary_search_step(ptr, target, 15);
        ptr = binary_search_step(ptr, target, 7);
        ptr = binary_search_step(ptr, target, 3);
        ptr = binary_search_step(ptr, target, 1);
        let extra = if *ptr < target { 1 } else { 0 };
        (ptr.offset_from(start_ptr) as usize) + extra
    }
}

#[cfg(test)]
mod tests {
    use super::branchless_binary_search;
    use crate::docset::TERMINATED;
    use crate::postings::compression::COMPRESSION_BLOCK_SIZE;
    use proptest::prelude::*;
    use std::collections::HashSet;

    fn search_in_block_trivial_but_slow(block: &[u32], target: u32) -> usize {
        block.iter().take_while(|&&val| val < target).count()
    }

    fn util_test_search_in_block(block: &[u32], target: u32) {
        let cursor = search_in_block_trivial_but_slow(block, target);
        assert!(cursor < COMPRESSION_BLOCK_SIZE);
        assert!(block[cursor] >= target);
        if cursor > 0 {
            assert!(block[cursor - 1] < target);
        }
        assert_eq!(block.len(), COMPRESSION_BLOCK_SIZE);
        let mut output_buffer = [TERMINATED; COMPRESSION_BLOCK_SIZE];
        output_buffer[..block.len()].copy_from_slice(block);
        assert_eq!(branchless_binary_search(&output_buffer, target), cursor);
    }

    fn util_test_search_in_block_all(block: &[u32]) {
        let mut targets = HashSet::new();
        targets.insert(0);
        for &val in block {
            if val > 0 {
                targets.insert(val - 1);
            }
            targets.insert(val);
        }
        for target in targets {
            util_test_search_in_block(block, target);
        }
    }

    #[test]
    fn test_search_in_branchless_binary_search() {
        let v: Vec<u32> = (0..COMPRESSION_BLOCK_SIZE).map(|i| i as u32 * 2).collect();
        util_test_search_in_block_all(&v[..]);
    }

    fn monotonous_block() -> impl Strategy<Value = Vec<u32>> {
        prop::collection::vec(0u32..5u32, COMPRESSION_BLOCK_SIZE).prop_map(|mut deltas| {
            let mut el = 0;
            for i in 0..COMPRESSION_BLOCK_SIZE {
                el += deltas[i];
                deltas[i] = el;
            }
            deltas
        })
    }

    proptest! {
        #[test]
        fn test_proptest_branchless_binary_search(block in monotonous_block()) {
            util_test_search_in_block_all(&block[..]);
        }
    }
}
<|endoftext|>