<commit_before>#![cfg(feature = "stream")]
#![allow(dead_code)]

#[macro_use]
extern crate nom;

use nom::{IResult, Needed, be_f32, be_u16, be_u32, be_u64};
//use nom::{Consumer,ConsumerState,Move,Input,Producer,FileProducer,FileProducerState};
//use nom::IResult;
use nom::{Err, ErrorKind};

use std::str;

fn mp4_box(input: &[u8]) -> IResult<&[u8], &[u8]> {
  match be_u32(input) {
    Ok((i, offset)) => {
      let sz: usize = offset as usize;
      if i.len() >= sz - 4 {
        Ok((&i[(sz - 4)..], &i[0..(sz - 4)]))
      } else {
        Err(Err::Incomplete(Needed::Size(offset as usize + 4)))
      }
    }
    Err(e) => Err(e),
  }
}

#[cfg_attr(rustfmt, rustfmt_skip)]
#[derive(PartialEq,Eq,Debug)]
struct FileType<'a> {
  major_brand:         &'a str,
  major_brand_version: &'a [u8],
  compatible_brands:   Vec<&'a str>
}

#[cfg_attr(rustfmt, rustfmt_skip)]
#[allow(non_snake_case)]
#[derive(Debug,Clone)]
pub struct Mvhd32 {
  version_flags: u32, // actually:
  // version: u8,
  // flags: u24       // 3 bytes
  created_date:  u32,
  modified_date: u32,
  scale:         u32,
  duration:      u32,
  speed:         f32,
  volume:        u16, // actually a 2 bytes decimal
  /* 10 bytes reserved */
  scaleA:        f32,
  rotateB:       f32,
  angleU:        f32,
  rotateC:       f32,
  scaleD:        f32,
  angleV:        f32,
  positionX:     f32,
  positionY:     f32,
  scaleW:        f32,
  preview:       u64,
  poster:        u32,
  selection:     u64,
  current_time:  u32,
  track_id:      u32
}

#[cfg_attr(rustfmt, rustfmt_skip)]
#[allow(non_snake_case)]
#[derive(Debug,Clone)]
pub struct Mvhd64 {
  version_flags: u32, // actually:
  // version: u8,
  // flags: u24       // 3 bytes
  created_date:  u64,
  modified_date: u64,
  scale:         u32,
  duration:      u64,
  speed:         f32,
  volume:        u16, // actually a 2 bytes decimal
  /* 10 bytes reserved */
  scaleA:        f32,
  rotateB:       f32,
  angleU:        f32,
  rotateC:       f32,
  scaleD:        f32,
  angleV:        f32,
  positionX:     f32,
  positionY:     f32,
  scaleW:        f32,
  preview:       u64,
  poster:        u32,
  selection:     u64,
  current_time:  u32,
  track_id:      u32
}

#[allow(non_snake_case)]
named!(mvhd32 <&[u8], MvhdBox>,
  do_parse!(
  version_flags: be_u32 >>
  created_date:  be_u32 >>
  modified_date: be_u32 >>
  scale:         be_u32 >>
  duration:      be_u32 >>
  speed:         be_f32 >>
  volume:        be_u16 >> // actually a 2 bytes decimal
              take!(10) >>
  scale_a:       be_f32 >>
  rotate_b:      be_f32 >>
  angle_u:       be_f32 >>
  rotate_c:      be_f32 >>
  scale_d:       be_f32 >>
  angle_v:       be_f32 >>
  position_x:    be_f32 >>
  position_y:    be_f32 >>
  scale_w:       be_f32 >>
  preview:       be_u64 >>
  poster:        be_u32 >>
  selection:     be_u64 >>
  current_time:  be_u32 >>
  track_id:      be_u32 >>
  (
    MvhdBox::M32(Mvhd32 {
      version_flags: version_flags,
      created_date:  created_date,
      modified_date: modified_date,
      scale:         scale,
      duration:      duration,
      speed:         speed,
      volume:        volume,
      scaleA:        scale_a,
      rotateB:       rotate_b,
      angleU:        angle_u,
      rotateC:       rotate_c,
      scaleD:        scale_d,
      angleV:        angle_v,
      positionX:     position_x,
      positionY:     position_y,
      scaleW:        scale_w,
      preview:       preview,
      poster:        poster,
      selection:     selection,
      current_time:  current_time,
      track_id:      track_id
    })
  ))
);

#[allow(non_snake_case)]
named!(mvhd64 <&[u8], MvhdBox>,
  do_parse!(
  version_flags: be_u32 >>
  created_date:  be_u64 >>
  modified_date: be_u64 >>
  scale:         be_u32 >>
  duration:      be_u64 >>
  speed:         be_f32 >>
  volume:        be_u16 >> // actually a 2 bytes decimal
              take!(10) >>
  scale_a:       be_f32 >>
  rotate_b:      be_f32 >>
  angle_u:       be_f32 >>
  rotate_c:      be_f32 >>
  scale_d:       be_f32 >>
  angle_v:       be_f32 >>
  position_x:    be_f32 >>
  position_y:    be_f32 >>
  scale_w:       be_f32 >>
  preview:       be_u64 >>
  poster:        be_u32 >>
  selection:     be_u64 >>
  current_time:  be_u32 >>
  track_id:      be_u32 >>
  (
    MvhdBox::M64(Mvhd64 {
      version_flags: version_flags,
      created_date:  created_date,
      modified_date: modified_date,
      scale:         scale,
      duration:      duration,
      speed:         speed,
      volume:        volume,
      scaleA:        scale_a,
      rotateB:       rotate_b,
      angleU:        angle_u,
      rotateC:       rotate_c,
      scaleD:        scale_d,
      angleV:        angle_v,
      positionX:     position_x,
      positionY:     position_y,
      scaleW:        scale_w,
      preview:       preview,
      poster:        poster,
      selection:     selection,
      current_time:  current_time,
      track_id:      track_id
    })
  ))
);

#[derive(Debug, Clone)]
pub enum MvhdBox {
  M32(Mvhd32),
  M64(Mvhd64),
}

#[derive(Debug, Clone)]
pub enum MoovBox {
  Mdra,
  Dref,
  Cmov,
  Rmra,
  Iods,
  Mvhd(MvhdBox),
  Clip,
  Trak,
  Udta,
}

#[derive(Debug)]
enum MP4BoxType {
  Ftyp,
  Moov,
  Mdat,
  Free,
  Skip,
  Wide,
  Mdra,
  Dref,
  Cmov,
  Rmra,
  Iods,
  Mvhd,
  Clip,
  Trak,
  Udta,
  Unknown,
}

#[derive(Debug)]
struct MP4BoxHeader {
  length: u32,
  tag: MP4BoxType,
}

named!(brand_name<&[u8],&str>, map_res!(take!(4), str::from_utf8));

named!(filetype_parser<&[u8], FileType>,
  do_parse!(
    m: brand_name          >>
    v: take!(4)            >>
    c: many0!(brand_name)  >>
    (FileType{ major_brand: m, major_brand_version:v, compatible_brands: c })
  )
);

fn mvhd_box(input: &[u8]) -> IResult<&[u8], MvhdBox> {
  let res = if input.len() < 100 {
    Err(Err::Incomplete(Needed::Size(100)))
  } else if input.len() == 100 {
    mvhd32(input)
  } else if input.len() == 112 {
    mvhd64(input)
  } else {
    Err(Err::Error(error_position!(input, ErrorKind::Custom(32u32))))
  };
  println!("res: {:?}", res);
  res
}

fn unknown_box_type(input: &[u8]) -> IResult<&[u8], MP4BoxType> {
  Ok((input, MP4BoxType::Unknown))
}

//named!(box_type<&[u8], MP4BoxType>,
fn box_type(input: &[u8]) -> IResult<&[u8], MP4BoxType, u32> {
  alt!(input,
    tag!("ftyp") => { |_| MP4BoxType::Ftyp } |
    tag!("moov") => { |_| MP4BoxType::Moov } |
    tag!("mdat") => { |_| MP4BoxType::Mdat } |
    tag!("free") => { |_| MP4BoxType::Free } |
    tag!("skip") => { |_| MP4BoxType::Skip } |
    tag!("wide") => { |_| MP4BoxType::Wide } |
    unknown_box_type
  )
}

// warning, an alt combinator with 9 branches containing a tag combinator
// can make the compilation very slow. Use functions as sub parsers,
// or split into multiple alt! parsers if it gets slow
named!(moov_type<&[u8], MP4BoxType>,
  alt!(
    tag!("mdra") => { |_| MP4BoxType::Mdra } |
    tag!("dref") => { |_| MP4BoxType::Dref } |
    tag!("cmov") => { |_| MP4BoxType::Cmov } |
    tag!("rmra") => { |_| MP4BoxType::Rmra } |
    tag!("iods") => { |_| MP4BoxType::Iods } |
    tag!("mvhd") => { |_| MP4BoxType::Mvhd } |
    tag!("clip") => { |_| MP4BoxType::Clip } |
    tag!("trak") => { |_| MP4BoxType::Trak } |
    tag!("udta") => { |_| MP4BoxType::Udta }
  )
);

named!(box_header<&[u8],MP4BoxHeader>,
  do_parse!(
    length: be_u32 >>
    tag: box_type  >>
    (MP4BoxHeader{ length: length, tag: tag})
  )
);

named!(moov_header<&[u8],MP4BoxHeader>,
  do_parse!(
    length: be_u32 >>
    tag: moov_type >>
    (MP4BoxHeader{ length: length, tag: tag})
  )
);

/*
#[derive(Debug,PartialEq,Eq)]
enum MP4State {
  Main,
  Moov,
  Mvhd(usize)
}

pub struct MP4Consumer {
  state:      MP4State,
  moov_bytes: usize,
  c_state:    ConsumerState<(), (), Move>
}

impl MP4Consumer {
  fn new() -> MP4Consumer {
    MP4Consumer { state: MP4State::Main, moov_bytes: 0, c_state: ConsumerState::Continue(Move::Consume(0)) }
  }

  fn consume_main(&mut self, input: Input<&[u8]>) -> ConsumerState<(), (), Move> {
    //println!("\nparsing box header:\n{}", input.to_hex(8));
    match input {
      Input::Eof(None) => ConsumerState::Done(Move::Consume(0), ()),
      Input::Empty     => ConsumerState::Continue(Move::Consume(0)),
      Input::Element(sl) |  Input::Eof(Some(sl)) => {
        match box_header(sl) {
          Ok((i, header)) => {
            match header.tag {
              MP4BoxType::Ftyp    => {
                println!("-> FTYP");
                match filetype_parser(&i[0..(header.length as usize - 8)]) {
                  Ok((rest, filetype_header)) => {
                    println!("filetype header: {:?}", filetype_header);
                    //return ConsumerState::Await(header.length as usize, header.length as usize - 8);
                    return ConsumerState::Continue(Move::Consume(sl.offset(rest)));
                  }
                  Err(Err::Error(a)) => {
                    println!("ftyp parsing error: {:?}", a);
                    assert!(false);
                    return ConsumerState::Error(());
                  },
                  Err(Err::Incomplete(n)) => {
                    println!("ftyp incomplete -> await: {}", sl.len());
                    return ConsumerState::Continue(Move::Await(n));
                    //return ConsumerState::Await(0, input.len() + 100);
                  }
                }
              },
              MP4BoxType::Moov    => {
                println!("-> MOOV");
                self.state = MP4State::Moov;
                self.moov_bytes = header.length as usize - 8;
                return ConsumerState::Continue(Move::Consume(sl.offset(i)));
              },
              MP4BoxType::Mdat    => println!("-> MDAT"),
              MP4BoxType::Free    => println!("-> FREE"),
              MP4BoxType::Skip    => println!("-> SKIP"),
              MP4BoxType::Wide    => println!("-> WIDE"),
              MP4BoxType::Unknown => {
                println!("-> UNKNOWN");
                println!("bytes:\n{}", (sl).to_hex(8));
                //return ConsumerState::Continue(Move::Consume(sl.offset(i)));
              },
              _                   => { println!("invalid"); return ConsumerState::Error(())}
            }
            return ConsumerState::Continue(Move::Seek(SeekFrom::Current((header.length) as i64)))
          },
          Err(Err::Error(a)) => {
            println!("mp4 parsing error: {:?}", a);
            assert!(false);
            return ConsumerState::Error(());
          },
          Err(Err::Incomplete(i)) => {
            // FIXME: incomplete should send the required size
            println!("mp4 incomplete -> await: {}", sl.len());
            return ConsumerState::Continue(Move::Await(i));
          }
        }
      }
    }
  }

  fn consume_moov(&mut self, input: Input<&[u8]>) -> ConsumerState<(), (), Move> {
    //println!("\nparsing moov box(remaining {} bytes):\n{}", self.moov_bytes, input.to_hex(8));
    match input {
      Input::Eof(None) => return ConsumerState::Error(()),
      Input::Empty => return ConsumerState::Continue(Move::Consume(0)),
      Input::Element(sl) |  Input::Eof(Some(sl)) => {
        if self.moov_bytes == 0 {
          //println!("finished parsing moov atom, continuing with main parser");
          self.state = MP4State::Main;
          return ConsumerState::Continue(Move::Consume(0));
        }
        match moov_header(sl) {
          Ok((i, header)) => {
            match header.tag {
              MP4BoxType::Mvhd    => {
                println!("-> MVHD");
                self.state = MP4State::Mvhd(header.length as usize - 8);
                // TODO: check for overflow here
                self.moov_bytes = self.moov_bytes - (sl.len() - i.len());
                println!("remaining moov_bytes: {}", self.moov_bytes);
                return ConsumerState::Continue(Move::Consume(sl.offset(i)));
              },
              MP4BoxType::Wide    => println!("-> WIDE"),
              MP4BoxType::Mdra    => println!("-> MDRA"),
              MP4BoxType::Dref    => println!("-> DREF"),
              MP4BoxType::Cmov    => println!("-> CMOV"),
              MP4BoxType::Rmra    => println!("-> RMRA"),
              MP4BoxType::Iods    => println!("-> IODS"),
              MP4BoxType::Clip    => println!("-> CLIP"),
              MP4BoxType::Trak    => println!("-> TRAK"),
              MP4BoxType::Udta    => println!("-> UDTA"),
              MP4BoxType::Unknown => println!("-> MOOV UNKNOWN"),
              _                   => { println!("invalid header here: {:?}", header.tag); return ConsumerState::Error(());}
            };
            // TODO: check for overflow here
            self.moov_bytes = self.moov_bytes - header.length as usize;
            println!("remaining moov_bytes: {}", self.moov_bytes);
            return ConsumerState::Continue(Move::Seek(SeekFrom::Current((header.length) as i64)))
          },
          Err(Err::Error(a)) => {
            println!("moov parsing error: {:?}", a);
            println!("data:\n{}", sl.to_hex(8));
            assert!(false);
            return ConsumerState::Error(());
          },
          Err(Err::Incomplete(i)) => {
            println!("moov incomplete -> await: {}", sl.len());
            return ConsumerState::Continue(Move::Await(i));
          }
        }
      }
    };
  }

}

consumer_from_parser!(MvhdConsumer<MvhdBox>, mvhd_box);

impl<'a> Consumer<&'a[u8], (), (), Move> for MP4Consumer {
  fn handle(&mut self, input: Input<&[u8]>) -> &ConsumerState<(), (), Move> {
    match self.state {
      MP4State::Main     => {
        self.c_state = self.consume_main(input);
      },
      MP4State::Moov     => {
        self.c_state = self.consume_moov(input);
      },
      MP4State::Mvhd(sz) => {
        match input {
          Input::Eof(None) => self.c_state = ConsumerState::Error(()),
          Input::Empty     => self.c_state = ConsumerState::Continue(Move::Consume(0)),
          Input::Element(sl) |  Input::Eof(Some(sl)) => {
            let mut c = MvhdConsumer{ state:ConsumerState::Continue(Move::Consume(0)) };
            self.c_state = c.handle(Input::Element(&sl[..sz])).flat_map(|m, _| {
              self.state = MP4State::Moov;
              ConsumerState::Continue(m)
            });
            println!("found mvhd?: {:?}", c.state());
            match self.c_state {
              ConsumerState::Continue(Move::Consume(sz)) => self.moov_bytes = self.moov_bytes - sz,
              ConsumerState::Continue(Move::Seek(SeekFrom::Current(sz))) => self.moov_bytes = self.moov_bytes - (sz as usize),
              _ => ()
            };
            println!("remaining moov_bytes: {}", self.moov_bytes);
          }
        }
      }
    };
    &self.c_state
  }

  fn state(&self) -> &ConsumerState<(), (), Move> {
    &self.c_state
  }
}

#[allow(unused_must_use)]
fn explore_mp4_file(filename: &str) {
  let mut p = FileProducer::new(filename, 400).unwrap();
  let mut c = MP4Consumer{state: MP4State::Main, moov_bytes: 0, c_state: ConsumerState::Continue(Move::Consume(0))};
  //c.run(&mut p);
  while let &ConsumerState::Continue(mv) = p.apply(&mut c) {
    println!("move: {:?}", mv);
  }
  println!("last consumer state: {:?} | last state: {:?}", c.c_state, c.state);

  if let ConsumerState::Done(Move::Consume(0), ()) = c.c_state {
    println!("consumer state ok");
  } else {
    assert!(false, "consumer should have reached Done state");
  }
  assert_eq!(c.state, MP4State::Main);
  assert_eq!(p.state(), FileProducerState::Eof);
  //assert!(false);
}

#[test]
fn small_test() {
  explore_mp4_file("assets/small.mp4");
}

#[test]
fn big_bunny_test() {
  explore_mp4_file("assets/bigbuckbunny.mp4");
}
*/
<commit_msg>fix MP4 test compilation<commit_after>#![allow(dead_code)]

#[macro_use]
extern crate nom;

use nom::{
  IResult, Needed, Err,
  error::ErrorKind,
  number::streaming::{be_u16, be_u32, be_u64, be_f32}
};

use std::str;

fn mp4_box(input: &[u8]) -> IResult<&[u8], &[u8]> {
  match be_u32(input) {
    Ok((i, offset)) => {
      let sz: usize = offset as usize;
      if i.len() >= sz - 4 {
        Ok((&i[(sz - 4)..], &i[0..(sz - 4)]))
      } else {
        Err(Err::Incomplete(Needed::Size(offset as usize + 4)))
      }
    }
    Err(e) => Err(e),
  }
}

#[cfg_attr(rustfmt, rustfmt_skip)]
#[derive(PartialEq,Eq,Debug)]
struct FileType<'a> {
  major_brand:         &'a str,
  major_brand_version: &'a [u8],
  compatible_brands:   Vec<&'a str>
}

#[cfg_attr(rustfmt, rustfmt_skip)]
#[allow(non_snake_case)]
#[derive(Debug,Clone)]
pub struct Mvhd32 {
  version_flags: u32, // actually:
  // version: u8,
  // flags: u24       // 3 bytes
  created_date:  u32,
  modified_date: u32,
  scale:         u32,
  duration:      u32,
  speed:         f32,
  volume:        u16, // actually a 2 bytes decimal
  /* 10 bytes reserved */
  scaleA:        f32,
  rotateB:       f32,
  angleU:        f32,
  rotateC:       f32,
  scaleD:        f32,
  angleV:        f32,
  positionX:     f32,
  positionY:     f32,
  scaleW:        f32,
  preview:       u64,
  poster:        u32,
  selection:     u64,
  current_time:  u32,
  track_id:      u32
}

#[cfg_attr(rustfmt, rustfmt_skip)]
#[allow(non_snake_case)]
#[derive(Debug,Clone)]
pub struct Mvhd64 {
  version_flags: u32, // actually:
  // version: u8,
  // flags: u24       // 3 bytes
  created_date:  u64,
  modified_date: u64,
  scale:         u32,
  duration:      u64,
  speed:         f32,
  volume:        u16, // actually a 2 bytes decimal
  /* 10 bytes reserved */
  scaleA:        f32,
  rotateB:       f32,
  angleU:        f32,
  rotateC:       f32,
  scaleD:        f32,
  angleV:        f32,
  positionX:     f32,
  positionY:     f32,
  scaleW:        f32,
  preview:       u64,
  poster:        u32,
  selection:     u64,
  current_time:  u32,
  track_id:      u32
}

#[allow(non_snake_case)]
named!(mvhd32 <&[u8], MvhdBox>,
  do_parse!(
  version_flags: be_u32 >>
  created_date:  be_u32 >>
  modified_date: be_u32 >>
  scale:         be_u32 >>
  duration:      be_u32 >>
  speed:         be_f32 >>
  volume:        be_u16 >> // actually a 2 bytes decimal
              take!(10) >>
  scale_a:       be_f32 >>
  rotate_b:      be_f32 >>
  angle_u:       be_f32 >>
  rotate_c:      be_f32 >>
  scale_d:       be_f32 >>
  angle_v:       be_f32 >>
  position_x:    be_f32 >>
  position_y:    be_f32 >>
  scale_w:       be_f32 >>
  preview:       be_u64 >>
  poster:        be_u32 >>
  selection:     be_u64 >>
  current_time:  be_u32 >>
  track_id:      be_u32 >>
  (
    MvhdBox::M32(Mvhd32 {
      version_flags: version_flags,
      created_date:  created_date,
      modified_date: modified_date,
      scale:         scale,
      duration:      duration,
      speed:         speed,
      volume:        volume,
      scaleA:        scale_a,
      rotateB:       rotate_b,
      angleU:        angle_u,
      rotateC:       rotate_c,
      scaleD:        scale_d,
      angleV:        angle_v,
      positionX:     position_x,
      positionY:     position_y,
      scaleW:        scale_w,
      preview:       preview,
      poster:        poster,
      selection:     selection,
      current_time:  current_time,
      track_id:      track_id
    })
  ))
);

#[allow(non_snake_case)]
named!(mvhd64 <&[u8], MvhdBox>,
  do_parse!(
  version_flags: be_u32 >>
  created_date:  be_u64 >>
  modified_date: be_u64 >>
  scale:         be_u32 >>
  duration:      be_u64 >>
  speed:         be_f32 >>
  volume:        be_u16 >> // actually a 2 bytes decimal
              take!(10) >>
  scale_a:       be_f32 >>
  rotate_b:      be_f32 >>
  angle_u:       be_f32 >>
  rotate_c:      be_f32 >>
  scale_d:       be_f32 >>
  angle_v:       be_f32 >>
  position_x:    be_f32 >>
  position_y:    be_f32 >>
  scale_w:       be_f32 >>
  preview:       be_u64 >>
  poster:        be_u32 >>
  selection:     be_u64 >>
  current_time:  be_u32 >>
  track_id:      be_u32 >>
  (
    MvhdBox::M64(Mvhd64 {
      version_flags: version_flags,
      created_date:  created_date,
      modified_date: modified_date,
      scale:         scale,
      duration:      duration,
      speed:         speed,
      volume:        volume,
      scaleA:        scale_a,
      rotateB:       rotate_b,
      angleU:        angle_u,
      rotateC:       rotate_c,
      scaleD:        scale_d,
      angleV:        angle_v,
      positionX:     position_x,
      positionY:     position_y,
      scaleW:        scale_w,
      preview:       preview,
      poster:        poster,
      selection:     selection,
      current_time:  current_time,
      track_id:      track_id
    })
  ))
);

#[derive(Debug, Clone)]
pub enum MvhdBox {
  M32(Mvhd32),
  M64(Mvhd64),
}

#[derive(Debug, Clone)]
pub enum MoovBox {
  Mdra,
  Dref,
  Cmov,
  Rmra,
  Iods,
  Mvhd(MvhdBox),
  Clip,
  Trak,
  Udta,
}

#[derive(Debug)]
enum MP4BoxType {
  Ftyp,
  Moov,
  Mdat,
  Free,
  Skip,
  Wide,
  Mdra,
  Dref,
  Cmov,
  Rmra,
  Iods,
  Mvhd,
  Clip,
  Trak,
  Udta,
  Unknown,
}

#[derive(Debug)]
struct MP4BoxHeader {
  length: u32,
  tag: MP4BoxType,
}

named!(brand_name<&[u8],&str>, map_res!(take!(4), str::from_utf8));

named!(filetype_parser<&[u8], FileType>,
  do_parse!(
    m: brand_name          >>
    v: take!(4)            >>
    c: many0!(brand_name)  >>
    (FileType{ major_brand: m, major_brand_version:v, compatible_brands: c })
  )
);

fn mvhd_box(input: &[u8]) -> IResult<&[u8], MvhdBox> {
  let res = if input.len() < 100 {
    Err(Err::Incomplete(Needed::Size(100)))
  } else if input.len() == 100 {
    mvhd32(input)
  } else if input.len() == 112 {
    mvhd64(input)
  } else {
    Err(Err::Error(error_position!(input, ErrorKind::TooLarge)))
  };
  println!("res: {:?}", res);
  res
}

fn unknown_box_type(input: &[u8]) -> IResult<&[u8], MP4BoxType> {
  Ok((input, MP4BoxType::Unknown))
}

//named!(box_type<&[u8], MP4BoxType>,
fn box_type(input: &[u8]) -> IResult<&[u8], MP4BoxType> {
  alt!(input,
    tag!("ftyp") => { |_| MP4BoxType::Ftyp } |
    tag!("moov") => { |_| MP4BoxType::Moov } |
    tag!("mdat") => { |_| MP4BoxType::Mdat } |
    tag!("free") => { |_| MP4BoxType::Free } |
    tag!("skip") => { |_| MP4BoxType::Skip } |
    tag!("wide") => { |_| MP4BoxType::Wide } |
    unknown_box_type
  )
}

// warning, an alt combinator with 9 branches containing a tag combinator
// can make the compilation very slow. Use functions as sub parsers,
// or split into multiple alt! parsers if it gets slow
named!(moov_type<&[u8], MP4BoxType>,
  alt!(
    tag!("mdra") => { |_| MP4BoxType::Mdra } |
    tag!("dref") => { |_| MP4BoxType::Dref } |
    tag!("cmov") => { |_| MP4BoxType::Cmov } |
    tag!("rmra") => { |_| MP4BoxType::Rmra } |
    tag!("iods") => { |_| MP4BoxType::Iods } |
    tag!("mvhd") => { |_| MP4BoxType::Mvhd } |
    tag!("clip") => { |_| MP4BoxType::Clip } |
    tag!("trak") => { |_| MP4BoxType::Trak } |
    tag!("udta") => { |_| MP4BoxType::Udta }
  )
);

named!(box_header<&[u8],MP4BoxHeader>,
  do_parse!(
    length: be_u32 >>
    tag: box_type  >>
    (MP4BoxHeader{ length: length, tag: tag})
  )
);

named!(moov_header<&[u8],MP4BoxHeader>,
  do_parse!(
    length: be_u32 >>
    tag: moov_type >>
    (MP4BoxHeader{ length: length, tag: tag})
  )
);

<|endoftext|>


<commit_before>pub mod cow {
    pub struct Vec3 {
        x: f64,
        y: f64,
        z: f64
    }

    impl Vec3 {
        fn len(&self) -> f64 {
            (self.x * self.x +
             self.y * self.y +
             self.z * self.z).sqrt()
        }

        fn dot(&self, other: &Vec3) -> f64 {
            self.x * other.x +
            self.y * other.y +
            self.z * other.z
        }

        fn cross(&self, other: &Vec3) -> Vec3 {
            Vec3 {
                x: self.y * other.z - self.z * other.y,
                y: self.z * other.x - self.x * other.z,
                z: self.x * other.y - self.y * other.x
            }
        }

        fn unit(&self) -> Vec3 {
            let len = self.len();

            Vec3 {
                x: self.x / len,
                y: self.y / len,
                z: self.z / len
            }
        }
    }

    impl Add<Vec3, Vec3> for Vec3 {
        fn add(&self, other: &Vec3) -> Vec3 {
            Vec3 {
                x: self.x + other.x,
                y: self.y + other.y,
                z: self.z + other.z
            }
        }
    }

    impl Sub<Vec3, Vec3> for Vec3 {
        fn sub(&self, other: &Vec3) -> Vec3 {
            Vec3 {
                x: self.x - other.x,
                y: self.y - other.y,
                z: self.z - other.z
            }
        }
    }

    impl Mul<Vec3, Vec3> for Vec3 {
        fn mul(&self, other: &Vec3) -> Vec3 {
            Vec3 {
                x: self.x * other.x,
                y: self.y * other.y,
                z: self.z * other.z
            }
        }
    }
}
<commit_msg>Remove cow.rs<commit_after><|endoftext|>


<commit_before>use std::{env,
          ffi::OsString,
          fs as stdfs,
          path::PathBuf};

use crate::{common::ui::UI,
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs}};

use crate::{config,
            error::Result};

pub const ARTIFACT_PATH_ENVVAR: &str = "ARTIFACT_PATH";

const ORIGIN_ENVVAR: &str = "HAB_ORIGIN";
const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    if henv::var(ORIGIN_ENVVAR).is_err() {
        let config = config::load()?;
        if let Some(default_origin) = config.origin {
            debug!("Setting default origin {} via CLI config", &default_origin);
            env::set_var("HAB_ORIGIN", default_origin);
        }
    }

    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        let path = fs::cache_key_path(None::<&str>);
        debug!("Setting {}={}", CACHE_KEY_PATH_ENV_VAR, path.display());
        env::set_var(CACHE_KEY_PATH_ENV_VAR, &path);
    };

    let artifact_path = match henv::var(ARTIFACT_PATH_ENVVAR) {
        Ok(p) => PathBuf::from(p),
        Err(_) => {
            let path = fs::cache_artifact_path(None::<&str>);
            debug!("Setting {}={}", ARTIFACT_PATH_ENVVAR, path.display());
            env::set_var(ARTIFACT_PATH_ENVVAR, &path);
            path
        }
    };
    if !artifact_path.is_dir() {
        debug!("Creating artifact_path at: {}", artifact_path.display());
        stdfs::create_dir_all(&artifact_path)?;
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use crate::{command::studio::docker,
                common::ui::{UIWriter,
                             UI},
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::{am_i_root,
                             find_command},
                        os::process,
                        package::{PackageIdent,
                                  PackageInstall},
                        users::linux as group},
                VERSION};
    use std::{env,
              ffi::OsString,
              path::PathBuf,
              str::FromStr};

    const SUDO_CMD: &str = "sudo";

    pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
        rerun_with_sudo_if_needed(ui, &args)?;
        if is_docker_studio(&args) {
            docker::start_docker_studio(ui, args)
        } else {
            let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
                Ok(command) => PathBuf::from(command),
                Err(_) => {
                    init();
                    let version: Vec<&str> = VERSION.split('/').collect();
                    let ident = PackageIdent::from_str(&format!("{}/{}",
                                                                super::STUDIO_PACKAGE_IDENT,
                                                                version[0]))?;
                    // This is a duplicate of the code in `hab pkg exec` and
                    // should be refactored as part of or after:
                    // https://github.com/habitat-sh/habitat/issues/6633
                    // https://github.com/habitat-sh/habitat/issues/6634
                    let pkg_install = PackageInstall::load(&ident, None)?;
                    let cmd_env = pkg_install.environment_for_command()?;
                    for (key, value) in cmd_env.into_iter() {
                        debug!("Setting: {}='{}'", key, value);
                        env::set_var(key, value);
                    }

                    let mut display_args = super::STUDIO_CMD.to_string();
                    for arg in args {
                        display_args.push(' ');
                        display_args.push_str(arg.to_string_lossy().as_ref());
                    }
                    debug!("Running: {}", display_args);

                    exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
                }
            };

            if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
                process::become_command(cmd, args)?;
                Ok(())
            } else {
                Err(Error::ExecCommandNotFound(command))
            }
        }
    }

    fn is_docker_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "linux")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return true;
            }
        }

        false
    }

    fn has_docker_group() -> bool {
        let current_user = group::get_current_username().unwrap();
        let docker_members = group::get_members_by_groupname("docker");
        docker_members.map_or(false, |d| d.contains(&current_user))
    }

    fn rerun_with_sudo_if_needed(ui: &mut UI, args: &[OsString]) -> Result<()> {
        // If I have root permissions or if I am executing a docker studio
        // and have the appropriate group - early return, we are done.
        if am_i_root() || (is_docker_studio(args) && has_docker_group()) {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into(),];
                args.append(&mut env::args_os().collect());
                process::become_command(sudo_prog, &args)?;
                Ok(())
            }
            None => {
                ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                SUDO_CMD))?;
                ui.warn("Running Habitat Studio requires root or administrator privileges. \
                         Please retry this command as a super user or use a privilege-granting \
                         facility such as sudo.")?;
                ui.br()?;
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use crate::{command::studio::docker,
                common::ui::UI,
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::find_command,
                        os::process,
                        package::PackageIdent},
                VERSION};
    use std::{ffi::OsString,
              path::PathBuf,
              str::FromStr};

    pub fn start(_ui: &mut UI, args: &[OsString]) -> Result<()> {
        if is_windows_studio(&args) {
            start_windows_studio(_ui, args)
        } else {
            docker::start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: &[OsString]) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split('/').collect();
                let ident = PackageIdent::from_str(&format!("{}/{}",
                                                            super::STUDIO_PACKAGE_IDENT,
                                                            version[0]))?;
                exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            process::become_command(cmd, args)?;
        } else {
            return Err(Error::ExecCommandNotFound(command));
        }
        Ok(())
    }

    fn is_windows_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "windows")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return false;
            }
        }

        // -w/--windows is deprecated and should be removed in a post 0.64.0 release
        for arg in args.iter() {
            let str_arg = arg.to_string_lossy().to_lowercase();
            if str_arg == "--windows" || str_arg == "-w" {
                return true;
            }
        }

        true
    }
}
<commit_msg>add an env var... Build env is broken on FC29 so moving to ubuntu vm for testing<commit_after>use std::{env,
          ffi::OsString,
          fs as stdfs,
          path::PathBuf};

use crate::{common::ui::UI,
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs}};

use crate::{config,
            error::Result};

pub const ARTIFACT_PATH_ENVVAR: &str = "ARTIFACT_PATH";

const ORIGIN_ENVVAR: &str = "HAB_ORIGIN";
const AUTH_TOKEN_ENVVAR: &str = "HAB_AUTH_TOKEN"
const STUDIO_CMD: &str = "hab-studio";
const STUDIO_CMD_ENVVAR: &str = "HAB_STUDIO_BINARY";
const STUDIO_PACKAGE_IDENT: &str = "core/hab-studio";

pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
    if henv::var(ORIGIN_ENVVAR).is_err() {
        let config = config::load()?;
        if let Some(default_origin) = config.origin {
            debug!("Setting default origin {} via CLI config", &default_origin);
            env::set_var("HAB_ORIGIN", default_origin);
        }
    }
    if henv::var(AUTH_TOKEN_ENVVAR).is_err() {
        let config = config::load()?;
        if let Some(auth_token) = config.auth_token {
            debug!("Setting default origin {} via CLI config", &auth_token);
            env::set_var("HAB_AUTH_TOKEN", auth_token;
        }
    }

    if henv::var(CACHE_KEY_PATH_ENV_VAR).is_err() {
        let path = fs::cache_key_path(None::<&str>);
        debug!("Setting {}={}", CACHE_KEY_PATH_ENV_VAR, path.display());
        env::set_var(CACHE_KEY_PATH_ENV_VAR, &path);
    };

    let artifact_path = match henv::var(ARTIFACT_PATH_ENVVAR) {
        Ok(p) => PathBuf::from(p),
        Err(_) => {
            let path = fs::cache_artifact_path(None::<&str>);
            debug!("Setting {}={}", ARTIFACT_PATH_ENVVAR, path.display());
            env::set_var(ARTIFACT_PATH_ENVVAR, &path);
            path
        }
    };
    if !artifact_path.is_dir() {
        debug!("Creating artifact_path at: {}", artifact_path.display());
        stdfs::create_dir_all(&artifact_path)?;
    }

    inner::start(ui, args)
}

#[cfg(target_os = "linux")]
mod inner {
    use crate::{command::studio::docker,
                common::ui::{UIWriter,
                             UI},
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::{am_i_root,
                             find_command},
                        os::process,
                        package::{PackageIdent,
                                  PackageInstall},
                        users::linux as group},
                VERSION};
    use std::{env,
              ffi::OsString,
              path::PathBuf,
              str::FromStr};

    const SUDO_CMD: &str = "sudo";

    pub fn start(ui: &mut UI, args: &[OsString]) -> Result<()> {
        rerun_with_sudo_if_needed(ui, &args)?;
        if is_docker_studio(&args) {
            docker::start_docker_studio(ui, args)
        } else {
            let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
                Ok(command) => PathBuf::from(command),
                Err(_) => {
                    init();
                    let version: Vec<&str> = VERSION.split('/').collect();
                    let ident = PackageIdent::from_str(&format!("{}/{}",
                                                                super::STUDIO_PACKAGE_IDENT,
                                                                version[0]))?;
                    // This is a duplicate of the code in `hab pkg exec` and
                    // should be refactored as part of or after:
                    // https://github.com/habitat-sh/habitat/issues/6633
                    // https://github.com/habitat-sh/habitat/issues/6634
                    let pkg_install = PackageInstall::load(&ident, None)?;
                    let cmd_env = pkg_install.environment_for_command()?;
                    for (key, value) in cmd_env.into_iter() {
                        debug!("Setting: {}='{}'", key, value);
                        env::set_var(key, value);
                    }

                    let mut display_args = super::STUDIO_CMD.to_string();
                    for arg in args {
                        display_args.push(' ');
                        display_args.push_str(arg.to_string_lossy().as_ref());
                    }
                    debug!("Running: {}", display_args);

                    exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
                }
            };

            if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
                process::become_command(cmd, args)?;
                Ok(())
            } else {
                Err(Error::ExecCommandNotFound(command))
            }
        }
    }

    fn is_docker_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "linux")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return true;
            }
        }

        false
    }

    fn has_docker_group() -> bool {
        let current_user = group::get_current_username().unwrap();
        let docker_members = group::get_members_by_groupname("docker");
        docker_members.map_or(false, |d| d.contains(&current_user))
    }

    fn rerun_with_sudo_if_needed(ui: &mut UI, args: &[OsString]) -> Result<()> {
        // If I have root permissions or if I am executing a docker studio
        // and have the appropriate group - early return, we are done.
        if am_i_root() || (is_docker_studio(args) && has_docker_group()) {
            return Ok(());
        }

        // Otherwise we will try to re-run this program using `sudo`
        match find_command(SUDO_CMD) {
            Some(sudo_prog) => {
                let mut args: Vec<OsString> = vec!["-p".into(),
                                                   "[sudo hab-studio] password for %u: ".into(),
                                                   "-E".into(),];
                args.append(&mut env::args_os().collect());
                process::become_command(sudo_prog, &args)?;
                Ok(())
            }
            None => {
                ui.warn(format!("Could not find the `{}' command, is it in your PATH?",
                                SUDO_CMD))?;
                ui.warn("Running Habitat Studio requires root or administrator privileges. \
                         Please retry this command as a super user or use a privilege-granting \
                         facility such as sudo.")?;
                ui.br()?;
                Err(Error::RootRequired)
            }
        }
    }
}

#[cfg(not(target_os = "linux"))]
mod inner {
    use crate::{command::studio::docker,
                common::ui::UI,
                error::{Error,
                        Result},
                exec,
                hcore::{crypto::init,
                        env as henv,
                        fs::find_command,
                        os::process,
                        package::PackageIdent},
                VERSION};
    use std::{ffi::OsString,
              path::PathBuf,
              str::FromStr};

    pub fn start(_ui: &mut UI, args: &[OsString]) -> Result<()> {
        if is_windows_studio(&args) {
            start_windows_studio(_ui, args)
        } else {
            docker::start_docker_studio(_ui, args)
        }
    }

    pub fn start_windows_studio(ui: &mut UI, args: &[OsString]) -> Result<()> {
        let command = match henv::var(super::STUDIO_CMD_ENVVAR) {
            Ok(command) => PathBuf::from(command),
            Err(_) => {
                init();
                let version: Vec<&str> = VERSION.split('/').collect();
                let ident = PackageIdent::from_str(&format!("{}/{}",
                                                            super::STUDIO_PACKAGE_IDENT,
                                                            version[0]))?;
                exec::command_from_min_pkg(ui, super::STUDIO_CMD, &ident)?
            }
        };

        if let Some(cmd) = find_command(command.to_string_lossy().as_ref()) {
            process::become_command(cmd, args)?;
        } else {
            return Err(Error::ExecCommandNotFound(command));
        }
        Ok(())
    }

    fn is_windows_studio(args: &[OsString]) -> bool {
        if cfg!(not(target_os = "windows")) {
            return false;
        }

        for arg in args.iter() {
            let str_arg = arg.to_string_lossy();
            if str_arg == "-D" {
                return false;
            }
        }

        // -w/--windows is deprecated and should be removed in a post 0.64.0 release
        for arg in args.iter() {
            let str_arg = arg.to_string_lossy().to_lowercase();
            if str_arg == "--windows" || str_arg == "-w" {
                return true;
            }
        }

        true
    }
}
<|endoftext|>


<commit_before>use cairo::{self, ImageSurface};
use markup5ever::{expanded_name, local_name, namespace_url, ns, QualName};
use nalgebra::{DMatrix, Dynamic, VecStorage};

use crate::drawing_ctx::DrawingCtx;
use crate::error::*;
use crate::node::{NodeResult, NodeTrait, RsvgNode};
use crate::number_list::{NumberList, NumberListError, NumberListLength};
use crate::parsers;
use crate::property_bag::PropertyBag;
use crate::rect::IRect;
use crate::surface_utils::{
    iterators::{PixelRectangle, Pixels},
    shared_surface::SharedImageSurface,
    EdgeMode, ImageSurfaceDataExt, Pixel,
};
use crate::util::clamp;

use super::context::{FilterContext, FilterOutput, FilterResult};
use super::{FilterEffect, FilterError, PrimitiveWithInput};

/// The `feConvolveMatrix` filter primitive.
pub struct FeConvolveMatrix {
    base: PrimitiveWithInput,
    order: (u32, u32),
    kernel_matrix: Option<DMatrix<f64>>,
    divisor: Option<f64>,
    bias: f64,
    target_x: Option<u32>,
    target_y: Option<u32>,
    edge_mode: EdgeMode,
    kernel_unit_length: Option<(f64, f64)>,
    preserve_alpha: bool,
}

impl Default for FeConvolveMatrix {
    /// Constructs a new `ConvolveMatrix` with empty properties.
    #[inline]
    fn default() -> FeConvolveMatrix {
        FeConvolveMatrix {
            base: PrimitiveWithInput::new::<Self>(),
            order: (3, 3),
            kernel_matrix: None,
            divisor: None,
            bias: 0.0,
            target_x: None,
            target_y: None,
            edge_mode: EdgeMode::Duplicate,
            kernel_unit_length: None,
            preserve_alpha: false,
        }
    }
}

impl NodeTrait for FeConvolveMatrix {
    impl_node_as_filter_effect!();

    fn set_atts(&mut self, parent: Option<&RsvgNode>, pbag: &PropertyBag<'_>) -> NodeResult {
        self.base.set_atts(parent, pbag)?;

        for (attr, value) in pbag.iter() {
            match attr.expanded() {
                expanded_name!(svg "order") => {
                    self.order = parsers::integer_optional_integer(value)
                        .and_then(|(x, y)| {
                            if x > 0 && y > 0 {
                                Ok((x as u32, y as u32))
                            } else {
                                Err(ValueErrorKind::value_error("values must be greater than 0"))
                            }
                        })
                        .attribute(attr)?
                }
                expanded_name!(svg "divisor") => {
                    self.divisor = Some(
                        parsers::number(value)
                            .and_then(|x| {
                                if x != 0.0 {
                                    Ok(x)
                                } else {
                                    Err(ValueErrorKind::value_error("divisor cannot be equal to 0"))
                                }
                            })
                            .attribute(attr)?,
                    )
                }
                expanded_name!(svg "bias") => self.bias = parsers::number(value).attribute(attr)?,
                expanded_name!(svg "edgeMode") => self.edge_mode = EdgeMode::parse(attr, value)?,
                expanded_name!(svg "kernelUnitLength") => {
                    self.kernel_unit_length = Some(
                        parsers::number_optional_number(value)
                            .and_then(|(x, y)| {
                                if x > 0.0 && y > 0.0 {
                                    Ok((x, y))
                                } else {
                                    Err(ValueErrorKind::value_error(
                                        "kernelUnitLength can't be less or equal to zero",
                                    ))
                                }
                            })
                            .attribute(attr)?,
                    )
                }
                expanded_name!(svg "preserveAlpha") => {
                    self.preserve_alpha = match value {
                        "false" => false,
                        "true" => true,
                        _ => {
                            return Err(NodeError::parse_error(attr, "expected false or true"));
                        }
                    }
                }
                _ => (),
            }
        }

        // target_x and target_y depend on order.
        for (attr, value) in pbag.iter() {
            match attr.expanded() {
                expanded_name!(svg "targetX") => {
                    self.target_x = Some(
                        parsers::integer(value)
                            .and_then(|x| {
                                if x >= 0 && x < self.order.0 as i32 {
                                    Ok(x as u32)
                                } else {
                                    Err(ValueErrorKind::value_error(
                                    "targetX must be greater or equal to zero and less than orderX",
                                ))
                                }
                            })
                            .attribute(attr)?,
                    )
                }
                expanded_name!(svg "targetY") => {
                    self.target_y = Some(
                        parsers::integer(value)
                            .and_then(|x| {
                                if x >= 0 && x < self.order.1 as i32 {
                                    Ok(x as u32)
                                } else {
                                    Err(ValueErrorKind::value_error(
                                    "targetY must be greater or equal to zero and less than orderY",
                                ))
                                }
                            })
                            .attribute(attr)?,
                    )
                }
                _ => (),
            }
        }

        // Default values for target_x and target_y.
        if self.target_x.is_none() {
            self.target_x = Some(self.order.0 / 2);
        }
        if self.target_y.is_none() {
            self.target_y = Some(self.order.1 / 2);
        }

        // Finally, parse the kernel matrix.
        for (attr, value) in pbag
            .iter()
            .filter(|(attr, _)| attr.expanded() == expanded_name!(svg "kernelMatrix"))
        {
            self.kernel_matrix = Some({
                let number_of_elements = self.order.0 as usize * self.order.1 as usize;

                // #352: Parse as an unbounded list rather than exact length to prevent aborts due
                //       to huge allocation attempts by underlying Vec::with_capacity().
                let NumberList(v) = NumberList::parse_str(value, NumberListLength::Unbounded)
                    .map_err(|err| match err {
                        NumberListError::IncorrectNumberOfElements => unreachable!(),
                        NumberListError::Parse(ref err) => ValueErrorKind::parse_error(&err),
                    })
                    .attribute(attr.clone())?;

                if v.len() != number_of_elements {
                    return Err(ValueErrorKind::value_error(&format!(
                        "incorrect number of elements: expected {}",
                        number_of_elements
                    )))
                    .attribute(attr);
                }

                DMatrix::from_data(VecStorage::new(
                    Dynamic::new(self.order.1 as usize),
                    Dynamic::new(self.order.0 as usize),
                    v,
                ))
            });
        }

        // kernel_matrix must have been specified.
        if self.kernel_matrix.is_none() {
            return Err(ValueErrorKind::value_error("the value must be set"))
                .attribute(QualName::new(None, ns!(svg), local_name!("kernelMatrix")));
        }

        // Default value for the divisor.
        if self.divisor.is_none() {
            self.divisor = Some(self.kernel_matrix.as_ref().unwrap().iter().sum());

            if self.divisor.unwrap() == 0.0 {
                self.divisor = Some(1.0);
            }
        }

        Ok(())
    }
}

impl FilterEffect for FeConvolveMatrix {
    fn render(
        &self,
        _node: &RsvgNode,
        ctx: &FilterContext,
        draw_ctx: &mut DrawingCtx,
    ) -> Result<FilterResult, FilterError> {
        let input = self.base.get_input(ctx, draw_ctx)?;
        let mut bounds = self
            .base
            .get_bounds(ctx)
            .add_input(&input)
            .into_irect(draw_ctx);
        let original_bounds = bounds;

        let mut input_surface = if self.preserve_alpha {
            // preserve_alpha means we need to premultiply and unpremultiply the values.
            input.surface().unpremultiply(bounds)?
        } else {
            input.surface().clone()
        };

        let scale = self
            .kernel_unit_length
            .map(|(dx, dy)| ctx.paffine().transform_distance(dx, dy));

        if let Some((ox, oy)) = scale {
            // Scale the input surface to match kernel_unit_length.
            let (new_surface, new_bounds) = input_surface.scale(bounds, 1.0 / ox, 1.0 / oy)?;

            input_surface = new_surface;
            bounds = new_bounds;
        }

        let matrix = self.kernel_matrix.as_ref().unwrap();

        let mut output_surface = ImageSurface::create(
            cairo::Format::ARgb32,
            input_surface.width(),
            input_surface.height(),
        )?;

        let output_stride = output_surface.get_stride() as usize;
        {
            let mut output_data = output_surface.get_data().unwrap();

            for (x, y, pixel) in Pixels::new(&input_surface, bounds) {
                // Compute the convolution rectangle bounds.
                let kernel_bounds = IRect::new(
                    x as i32 - self.target_x.unwrap() as i32,
                    y as i32 - self.target_y.unwrap() as i32,
                    x as i32 - self.target_x.unwrap() as i32 + self.order.0 as i32,
                    y as i32 - self.target_y.unwrap() as i32 + self.order.1 as i32,
                );

                // Do the convolution.
                let mut r = 0.0;
                let mut g = 0.0;
                let mut b = 0.0;
                let mut a = 0.0;

                for (x, y, pixel) in
                    PixelRectangle::new(&input_surface, bounds, kernel_bounds, self.edge_mode)
                {
                    let kernel_x = (kernel_bounds.x1 - x - 1) as usize;
                    let kernel_y = (kernel_bounds.y1 - y - 1) as usize;

                    r += f64::from(pixel.r) / 255.0 * matrix[(kernel_y, kernel_x)];
                    g += f64::from(pixel.g) / 255.0 * matrix[(kernel_y, kernel_x)];
                    b += f64::from(pixel.b) / 255.0 * matrix[(kernel_y, kernel_x)];

                    if !self.preserve_alpha {
                        a += f64::from(pixel.a) / 255.0 * matrix[(kernel_y, kernel_x)];
                    }
                }

                // If preserve_alpha is true, set a to the source alpha value.
                if self.preserve_alpha {
                    a = f64::from(pixel.a) / 255.0;
                } else {
                    a = a / self.divisor.unwrap() + self.bias;
                }

                let clamped_a = clamp(a, 0.0, 1.0);

                let compute = |x| {
                    let x = x / self.divisor.unwrap() + self.bias * a;

                    let x = if self.preserve_alpha {
                        // Premultiply the output value.
                        clamp(x, 0.0, 1.0) * clamped_a
                    } else {
                        clamp(x, 0.0, clamped_a)
                    };

                    ((x * 255.0) + 0.5) as u8
                };

                let output_pixel = Pixel {
                    r: compute(r),
                    g: compute(g),
                    b: compute(b),
                    a: ((clamped_a * 255.0) + 0.5) as u8,
                };

                output_data.set_pixel(output_stride, output_pixel, x, y);
            }
        }

        let mut output_surface =
            SharedImageSurface::new(output_surface, input.surface().surface_type())?;

        if let Some((ox, oy)) = scale {
            // Scale the output surface back.
            output_surface = output_surface.scale_to(
                ctx.source_graphic().width(),
                ctx.source_graphic().height(),
                original_bounds,
                ox,
                oy,
            )?;

            bounds = original_bounds;
        }

        Ok(FilterResult {
            name: self.base.result.clone(),
            output: FilterOutput {
                surface: output_surface,
                bounds,
            },
        })
    }

    #[inline]
    fn is_affected_by_color_interpolation_filters(&self) -> bool {
        true
    }
}

impl EdgeMode {
    fn parse(attr: QualName, s: &str) -> Result<Self, NodeError> {
        match s {
            "duplicate" => Ok(EdgeMode::Duplicate),
            "wrap" => Ok(EdgeMode::Wrap),
            "none" => Ok(EdgeMode::None),
            _ => Err(NodeError::parse_error(attr, "invalid value")),
        }
    }
}
<commit_msg>convolve_matrix.rs: Don't create a NodeError directly<commit_after>use cairo::{self, ImageSurface};
use markup5ever::{expanded_name, local_name, namespace_url, ns, QualName};
use nalgebra::{DMatrix, Dynamic, VecStorage};

use crate::drawing_ctx::DrawingCtx;
use crate::error::*;
use crate::node::{NodeResult, NodeTrait, RsvgNode};
use crate::number_list::{NumberList, NumberListError, NumberListLength};
use crate::parsers;
use crate::property_bag::PropertyBag;
use crate::rect::IRect;
use crate::surface_utils::{
    iterators::{PixelRectangle, Pixels},
    shared_surface::SharedImageSurface,
    EdgeMode, ImageSurfaceDataExt, Pixel,
};
use crate::util::clamp;

use super::context::{FilterContext, FilterOutput, FilterResult};
use super::{FilterEffect, FilterError, PrimitiveWithInput};

/// The `feConvolveMatrix` filter primitive.
pub struct FeConvolveMatrix {
    base: PrimitiveWithInput,
    order: (u32, u32),
    kernel_matrix: Option<DMatrix<f64>>,
    divisor: Option<f64>,
    bias: f64,
    target_x: Option<u32>,
    target_y: Option<u32>,
    edge_mode: EdgeMode,
    kernel_unit_length: Option<(f64, f64)>,
    preserve_alpha: bool,
}

impl Default for FeConvolveMatrix {
    /// Constructs a new `ConvolveMatrix` with empty properties.
    #[inline]
    fn default() -> FeConvolveMatrix {
        FeConvolveMatrix {
            base: PrimitiveWithInput::new::<Self>(),
            order: (3, 3),
            kernel_matrix: None,
            divisor: None,
            bias: 0.0,
            target_x: None,
            target_y: None,
            edge_mode: EdgeMode::Duplicate,
            kernel_unit_length: None,
            preserve_alpha: false,
        }
    }
}

impl NodeTrait for FeConvolveMatrix {
    impl_node_as_filter_effect!();

    fn set_atts(&mut self, parent: Option<&RsvgNode>, pbag: &PropertyBag<'_>) -> NodeResult {
        self.base.set_atts(parent, pbag)?;

        for (attr, value) in pbag.iter() {
            match attr.expanded() {
                expanded_name!(svg "order") => {
                    self.order = parsers::integer_optional_integer(value)
                        .and_then(|(x, y)| {
                            if x > 0 && y > 0 {
                                Ok((x as u32, y as u32))
                            } else {
                                Err(ValueErrorKind::value_error("values must be greater than 0"))
                            }
                        })
                        .attribute(attr)?
                }
                expanded_name!(svg "divisor") => {
                    self.divisor = Some(
                        parsers::number(value)
                            .and_then(|x| {
                                if x != 0.0 {
                                    Ok(x)
                                } else {
                                    Err(ValueErrorKind::value_error("divisor cannot be equal to 0"))
                                }
                            })
                            .attribute(attr)?,
                    )
                }
                expanded_name!(svg "bias") => self.bias = parsers::number(value).attribute(attr)?,
                expanded_name!(svg "edgeMode") => self.edge_mode = EdgeMode::parse(attr, value)?,
                expanded_name!(svg "kernelUnitLength") => {
                    self.kernel_unit_length = Some(
                        parsers::number_optional_number(value)
                            .and_then(|(x, y)| {
                                if x > 0.0 && y > 0.0 {
                                    Ok((x, y))
                                } else {
                                    Err(ValueErrorKind::value_error(
                                        "kernelUnitLength can't be less or equal to zero",
                                    ))
                                }
                            })
                            .attribute(attr)?,
                    )
                }
                expanded_name!(svg "preserveAlpha") => {
                    self.preserve_alpha = match value {
                        "false" => false,
                        "true" => true,
                        _ => {
                            return Err(ValueErrorKind::parse_error("expected false or true"))
                                .attribute(attr);
                        }
                    }
                }
                _ => (),
            }
        }

        // target_x and target_y depend on order.
        for (attr, value) in pbag.iter() {
            match attr.expanded() {
                expanded_name!(svg "targetX") => {
                    self.target_x = Some(
                        parsers::integer(value)
                            .and_then(|x| {
                                if x >= 0 && x < self.order.0 as i32 {
                                    Ok(x as u32)
                                } else {
                                    Err(ValueErrorKind::value_error(
                                    "targetX must be greater or equal to zero and less than orderX",
                                ))
                                }
                            })
                            .attribute(attr)?,
                    )
                }
                expanded_name!(svg "targetY") => {
                    self.target_y = Some(
                        parsers::integer(value)
                            .and_then(|x| {
                                if x >= 0 && x < self.order.1 as i32 {
                                    Ok(x as u32)
                                } else {
                                    Err(ValueErrorKind::value_error(
                                    "targetY must be greater or equal to zero and less than orderY",
                                ))
                                }
                            })
                            .attribute(attr)?,
                    )
                }
                _ => (),
            }
        }

        // Default values for target_x and target_y.
        if self.target_x.is_none() {
            self.target_x = Some(self.order.0 / 2);
        }
        if self.target_y.is_none() {
            self.target_y = Some(self.order.1 / 2);
        }

        // Finally, parse the kernel matrix.
        for (attr, value) in pbag
            .iter()
            .filter(|(attr, _)| attr.expanded() == expanded_name!(svg "kernelMatrix"))
        {
            self.kernel_matrix = Some({
                let number_of_elements = self.order.0 as usize * self.order.1 as usize;

                // #352: Parse as an unbounded list rather than exact length to prevent aborts due
                //       to huge allocation attempts by underlying Vec::with_capacity().
                let NumberList(v) = NumberList::parse_str(value, NumberListLength::Unbounded)
                    .map_err(|err| match err {
                        NumberListError::IncorrectNumberOfElements => unreachable!(),
                        NumberListError::Parse(ref err) => ValueErrorKind::parse_error(&err),
                    })
                    .attribute(attr.clone())?;

                if v.len() != number_of_elements {
                    return Err(ValueErrorKind::value_error(&format!(
                        "incorrect number of elements: expected {}",
                        number_of_elements
                    )))
                    .attribute(attr);
                }

                DMatrix::from_data(VecStorage::new(
                    Dynamic::new(self.order.1 as usize),
                    Dynamic::new(self.order.0 as usize),
                    v,
                ))
            });
        }

        // kernel_matrix must have been specified.
        if self.kernel_matrix.is_none() {
            return Err(ValueErrorKind::value_error("the value must be set"))
                .attribute(QualName::new(None, ns!(svg), local_name!("kernelMatrix")));
        }

        // Default value for the divisor.
        if self.divisor.is_none() {
            self.divisor = Some(self.kernel_matrix.as_ref().unwrap().iter().sum());

            if self.divisor.unwrap() == 0.0 {
                self.divisor = Some(1.0);
            }
        }

        Ok(())
    }
}

impl FilterEffect for FeConvolveMatrix {
    fn render(
        &self,
        _node: &RsvgNode,
        ctx: &FilterContext,
        draw_ctx: &mut DrawingCtx,
    ) -> Result<FilterResult, FilterError> {
        let input = self.base.get_input(ctx, draw_ctx)?;
        let mut bounds = self
            .base
            .get_bounds(ctx)
            .add_input(&input)
            .into_irect(draw_ctx);
        let original_bounds = bounds;

        let mut input_surface = if self.preserve_alpha {
            // preserve_alpha means we need to premultiply and unpremultiply the values.
            input.surface().unpremultiply(bounds)?
        } else {
            input.surface().clone()
        };

        let scale = self
            .kernel_unit_length
            .map(|(dx, dy)| ctx.paffine().transform_distance(dx, dy));

        if let Some((ox, oy)) = scale {
            // Scale the input surface to match kernel_unit_length.
            let (new_surface, new_bounds) = input_surface.scale(bounds, 1.0 / ox, 1.0 / oy)?;

            input_surface = new_surface;
            bounds = new_bounds;
        }

        let matrix = self.kernel_matrix.as_ref().unwrap();

        let mut output_surface = ImageSurface::create(
            cairo::Format::ARgb32,
            input_surface.width(),
            input_surface.height(),
        )?;

        let output_stride = output_surface.get_stride() as usize;
        {
            let mut output_data = output_surface.get_data().unwrap();

            for (x, y, pixel) in Pixels::new(&input_surface, bounds) {
                // Compute the convolution rectangle bounds.
                let kernel_bounds = IRect::new(
                    x as i32 - self.target_x.unwrap() as i32,
                    y as i32 - self.target_y.unwrap() as i32,
                    x as i32 - self.target_x.unwrap() as i32 + self.order.0 as i32,
                    y as i32 - self.target_y.unwrap() as i32 + self.order.1 as i32,
                );

                // Do the convolution.
                let mut r = 0.0;
                let mut g = 0.0;
                let mut b = 0.0;
                let mut a = 0.0;

                for (x, y, pixel) in
                    PixelRectangle::new(&input_surface, bounds, kernel_bounds, self.edge_mode)
                {
                    let kernel_x = (kernel_bounds.x1 - x - 1) as usize;
                    let kernel_y = (kernel_bounds.y1 - y - 1) as usize;

                    r += f64::from(pixel.r) / 255.0 * matrix[(kernel_y, kernel_x)];
                    g += f64::from(pixel.g) / 255.0 * matrix[(kernel_y, kernel_x)];
                    b += f64::from(pixel.b) / 255.0 * matrix[(kernel_y, kernel_x)];

                    if !self.preserve_alpha {
                        a += f64::from(pixel.a) / 255.0 * matrix[(kernel_y, kernel_x)];
                    }
                }

                // If preserve_alpha is true, set a to the source alpha value.
                if self.preserve_alpha {
                    a = f64::from(pixel.a) / 255.0;
                } else {
                    a = a / self.divisor.unwrap() + self.bias;
                }

                let clamped_a = clamp(a, 0.0, 1.0);

                let compute = |x| {
                    let x = x / self.divisor.unwrap() + self.bias * a;

                    let x = if self.preserve_alpha {
                        // Premultiply the output value.
                        clamp(x, 0.0, 1.0) * clamped_a
                    } else {
                        clamp(x, 0.0, clamped_a)
                    };

                    ((x * 255.0) + 0.5) as u8
                };

                let output_pixel = Pixel {
                    r: compute(r),
                    g: compute(g),
                    b: compute(b),
                    a: ((clamped_a * 255.0) + 0.5) as u8,
                };

                output_data.set_pixel(output_stride, output_pixel, x, y);
            }
        }

        let mut output_surface =
            SharedImageSurface::new(output_surface, input.surface().surface_type())?;

        if let Some((ox, oy)) = scale {
            // Scale the output surface back.
            output_surface = output_surface.scale_to(
                ctx.source_graphic().width(),
                ctx.source_graphic().height(),
                original_bounds,
                ox,
                oy,
            )?;

            bounds = original_bounds;
        }

        Ok(FilterResult {
            name: self.base.result.clone(),
            output: FilterOutput {
                surface: output_surface,
                bounds,
            },
        })
    }

    #[inline]
    fn is_affected_by_color_interpolation_filters(&self) -> bool {
        true
    }
}

impl EdgeMode {
    fn parse(attr: QualName, s: &str) -> Result<Self, NodeError> {
        match s {
            "duplicate" => Ok(EdgeMode::Duplicate),
            "wrap" => Ok(EdgeMode::Wrap),
            "none" => Ok(EdgeMode::None),
            _ => Err(ValueErrorKind::parse_error("invalid value")).attribute(attr),
        }
    }
}
<|endoftext|>


<commit_before>use std::io::{Read, Result};

use super::event::{self, Event, Key};

pub struct Events<R> {
    source: R,
    leftover: Option<u8>
}

impl<R: Read> Events<R> {
    pub fn new(source: R) -> Events<R> {
        Events {
            source: source,
            leftover: None
        }
    }
}

impl<R: Read> Iterator for Events<R> {
    type Item = Result<Event>;

    fn next(&mut self) -> Option<Result<Event>> {
        let ref mut source = self.source;

        if let Some(c) = self.leftover {
            // we have a leftover byte, use it
            self.leftover = None;
            return Some(parse_event(c, &mut source.bytes()));
        }

        let mut buf = [0u8; 2];

        let result = match source.read(&mut buf) {
            Ok(0) => return None,
            Ok(1) => {
                match buf[0] {
                    b'\x1b' => Ok(Event::Key(Key::Escape)),
                    b'\t' => Ok(Event::Key(Key::Tab)),
                    item => parse_event(item, &mut source.bytes())
                }
            },
            Ok(2) => {
                let mut option_iter = &mut Some(buf[1]).into_iter();

                let result = {
                    let mut iter = option_iter.
                        map(Ok).
                        chain(source.bytes());

                    parse_event(buf[0], &mut iter)
                };

                // If the option_iter wasn't consumed, keep the byte for later.
                self.leftover = option_iter.next();
                result
            },
            Ok(_) => unreachable!(),
            Err(e) => Err(e)
        };

        Some(result)
    }
}

fn parse_event<I>(item: u8, iter: &mut I) -> Result<Event>
    where I: Iterator<Item = Result<u8>>
{
    let mut buf = vec![item];
    let result = {
        let mut iter = iter.inspect(|byte| {
            if let Ok(byte) = *byte {
                buf.push(byte);
            }
        });
        event::parse_event(item, &mut iter)
    };
    result.or_else(|_| Ok(Event::Unknown(buf)))
}
<commit_msg>Remove unnecessary usage of 'mut'<commit_after>use std::io::{Read, Result};

use super::event::{self, Event, Key};

pub struct Events<R> {
    source: R,
    leftover: Option<u8>
}

impl<R: Read> Events<R> {
    pub fn new(source: R) -> Events<R> {
        Events {
            source: source,
            leftover: None
        }
    }
}

impl<R: Read> Iterator for Events<R> {
    type Item = Result<Event>;

    fn next(&mut self) -> Option<Result<Event>> {
        let ref mut source = self.source;

        if let Some(c) = self.leftover {
            // we have a leftover byte, use it
            self.leftover = None;
            return Some(parse_event(c, &mut source.bytes()));
        }

        let mut buf = [0u8; 2];

        let result = match source.read(&mut buf) {
            Ok(0) => return None,
            Ok(1) => {
                match buf[0] {
                    b'\x1b' => Ok(Event::Key(Key::Escape)),
                    b'\t' => Ok(Event::Key(Key::Tab)),
                    item => parse_event(item, &mut source.bytes())
                }
            },
            Ok(2) => {
                let option_iter = &mut Some(buf[1]).into_iter();

                let result = {
                    let mut iter = option_iter.
                        map(Ok).
                        chain(source.bytes());

                    parse_event(buf[0], &mut iter)
                };

                // If the option_iter wasn't consumed, keep the byte for later.
                self.leftover = option_iter.next();
                result
            },
            Ok(_) => unreachable!(),
            Err(e) => Err(e)
        };

        Some(result)
    }
}

fn parse_event<I>(item: u8, iter: &mut I) -> Result<Event>
    where I: Iterator<Item = Result<u8>>
{
    let mut buf = vec![item];
    let result = {
        let mut iter = iter.inspect(|byte| {
            if let Ok(byte) = *byte {
                buf.push(byte);
            }
        });
        event::parse_event(item, &mut iter)
    };
    result.or_else(|_| Ok(Event::Unknown(buf)))
}
<|endoftext|>


<commit_before>use std::{cmp::min, env, fs::File, io, mem, time::Duration};

use pgn_reader::{BufferedReader, Color, RawHeader, SanPlus, Skip, Visitor};
use rand::{distributions::OpenClosed01, thread_rng, Rng};
use serde::Serialize;
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};

const BATCH_SIZE: usize = 50;

#[derive(Debug, Serialize, Copy, Clone)]
#[serde(rename_all = "camelCase")]
enum Speed {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

impl Speed {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> Speed {
        let total = seconds + 40 * increment;

        if total < 30 {
            Speed::UltraBullet
        } else if total < 180 {
            Speed::Bullet
        } else if total < 480 {
            Speed::Blitz
        } else if total < 1500 {
            Speed::Rapid
        } else if total < 21_600 {
            Speed::Classical
        } else {
            Speed::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<Speed, ()> {
        if bytes == b"-" {
            return Ok(Speed::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        let increment = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        Ok(Speed::from_seconds_and_increment(seconds, increment))
    }
}

struct Importer {
    filename: String,
    client: reqwest::blocking::Client,

    current: Game,
    skip: bool,
    batch: Vec<Game>,
}

#[serde_as]
#[derive(Default, Serialize, Debug)]
struct Game {
    variant: Option<String>,
    speed: Option<Speed>,
    fen: Option<String>,
    id: Option<String>,
    date: Option<String>,
    white: Player,
    black: Player,
    #[serde_as(as = "Option<DisplayFromStr>")]
    winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, SanPlus>")]
    moves: Vec<SanPlus>,
}

#[derive(Default, Serialize, Debug)]
struct Player {
    name: Option<String>,
    rating: Option<u16>,
}

impl Importer {
    fn new(filename: String) -> Importer {
        Importer {
            filename,
            client: reqwest::blocking::Client::builder()
                .timeout(Duration::from_secs(60))
                .build()
                .expect("client"),
            current: Game::default(),
            skip: false,
            batch: Vec::with_capacity(BATCH_SIZE),
        }
    }

    pub fn send(&mut self) {
        // println!("{}", serde_json::to_string(&self.batch).expect("serialize"));

        let res = self
            .client
            .put("http://127.0.0.1:9001/import/lichess")
            .json(&self.batch)
            .send()
            .expect("send batch");

        println!(
            "{}: {}: {} - {}",
            self.filename,
            self.batch
                .last()
                .and_then(|g| g.date.as_ref())
                .unwrap_or(&String::new()),
            res.status(),
            res.text().expect("decode response")
        );

        self.batch.clear();
    }
}

impl Visitor for Importer {
    type Result = ();

    fn begin_game(&mut self) {
        self.skip = false;
        self.current = Game::default();
    }

    fn header(&mut self, key: &[u8], value: RawHeader<'_>) {
        if key == b"White" {
            self.current.white.name = Some(value.decode_utf8().expect("White").into_owned());
        } else if key == b"Black" {
            self.current.black.name = Some(value.decode_utf8().expect("Black").into_owned());
        } else if key == b"WhiteElo" {
            if value.as_bytes() != b"?" {
                self.current.white.rating = Some(btoi::btoi(value.as_bytes()).expect("WhiteElo"));
            }
        } else if key == b"BlackElo" {
            if value.as_bytes() != b"?" {
                self.current.black.rating = Some(btoi::btoi(value.as_bytes()).expect("BlackElo"));
            }
        } else if key == b"TimeControl" {
            self.current.speed = Some(Speed::from_bytes(value.as_bytes()).expect("TimeControl"));
        } else if key == b"Variant" {
            self.current.variant = Some(value.decode_utf8().expect("Variant").into_owned());
        } else if key == b"Date" || key == b"UTCDate" {
            self.current.date = Some(value.decode_utf8().expect("Date").into_owned());
        } else if key == b"WhiteTitle" || key == b"BlackTitle" {
            if value.as_bytes() == b"BOT" {
                self.skip = true;
            }
        } else if key == b"Site" {
            self.current.id = Some(
                String::from_utf8(
                    value
                        .as_bytes()
                        .rsplitn(2, |ch| *ch == b'/')
                        .next()
                        .expect("Site")
                        .to_owned(),
                )
                .expect("Site"),
            );
        } else if key == b"Result" {
            match value.as_bytes() {
                b"1-0" => self.current.winner = Some(Color::White),
                b"0-1" => self.current.winner = Some(Color::Black),
                b"1/2-1/2" => self.current.winner = None,
                _ => self.skip = true,
            }
        } else if key == b"FEN" {
            self.current.fen = Some(value.decode_utf8().expect("FEN").into_owned());
        }
    }

    fn end_headers(&mut self) -> Skip {
        let rating =
            (self.current.white.rating.unwrap_or(0) + self.current.black.rating.unwrap_or(0)) / 2;

        let standard = self
            .current
            .variant
            .as_ref()
            .map_or(false, |name| name != "Standard");

        let probability = if standard {
            match self.current.speed.unwrap_or(Speed::Correspondence) {
                Speed::Correspondence | Speed::Classical => 1.0,

                _ if rating >= 2500 => 1.0,

                Speed::Rapid if rating >= 2200 => 1.0,
                Speed::Rapid if rating >= 2000 => 0.30,
                Speed::Rapid if rating >= 1800 => 0.18,
                Speed::Rapid if rating >= 1600 => 0.16,

                Speed::Blitz if rating >= 2200 => 0.16,
                Speed::Blitz if rating >= 2000 => 0.07,
                Speed::Blitz if rating >= 1800 => 0.05,
                Speed::Blitz if rating >= 1600 => 0.05,

                Speed::Bullet if rating >= 2200 => 0.16,
                Speed::Bullet if rating >= 2000 => 0.07,
                Speed::Bullet if rating >= 1800 => 0.05,
                Speed::Bullet if rating >= 1600 => 0.05,

                Speed::UltraBullet => 1.0,

                _ => 0.02,
            }
        } else {
            // variant games
            if rating >= 1600 {
                1.0
            } else {
                0.5
            }
        };

        let rnd = thread_rng().sample(OpenClosed01);
        let accept = min(
            self.current.white.rating.unwrap_or(0),
            self.current.black.rating.unwrap_or(0),
        ) >= 1500
            && probability >= rnd
            && !self.skip;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: SanPlus) {
        self.current.moves.push(san);
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self) {
        if !self.skip {
            self.batch
                .push(mem::replace(&mut self.current, Default::default()));

            if self.batch.len() >= BATCH_SIZE {
                self.send();
            }
        }
    }
}

fn main() -> Result<(), io::Error> {
    for arg in env::args().skip(1) {
        let file = File::open(&arg)?;

        let uncompressed: Box<dyn io::Read> = if arg.ends_with(".bz2") {
            Box::new(bzip2::read::MultiBzDecoder::new(file))
        } else {
            Box::new(file)
        };

        let mut reader = BufferedReader::new(uncompressed);

        let mut importer = Importer::new(arg);
        reader.read_all(&mut importer)?;
    }

    Ok(())
}
<commit_msg>do not import games with rating <= 1500 into lichess db<commit_after>use std::{cmp::min, env, fs::File, io, mem, time::Duration};

use pgn_reader::{BufferedReader, Color, RawHeader, SanPlus, Skip, Visitor};
use rand::{distributions::OpenClosed01, thread_rng, Rng};
use serde::Serialize;
use serde_with::{serde_as, DisplayFromStr, SpaceSeparator, StringWithSeparator};

const BATCH_SIZE: usize = 50;

#[derive(Debug, Serialize, Copy, Clone)]
#[serde(rename_all = "camelCase")]
enum Speed {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

impl Speed {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> Speed {
        let total = seconds + 40 * increment;

        if total < 30 {
            Speed::UltraBullet
        } else if total < 180 {
            Speed::Bullet
        } else if total < 480 {
            Speed::Blitz
        } else if total < 1500 {
            Speed::Rapid
        } else if total < 21_600 {
            Speed::Classical
        } else {
            Speed::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<Speed, ()> {
        if bytes == b"-" {
            return Ok(Speed::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        let increment = btoi::btou(parts.next().ok_or(())?).map_err(|_| ())?;
        Ok(Speed::from_seconds_and_increment(seconds, increment))
    }
}

struct Importer {
    filename: String,
    client: reqwest::blocking::Client,

    current: Game,
    skip: bool,
    batch: Vec<Game>,
}

#[serde_as]
#[derive(Default, Serialize, Debug)]
struct Game {
    variant: Option<String>,
    speed: Option<Speed>,
    fen: Option<String>,
    id: Option<String>,
    date: Option<String>,
    white: Player,
    black: Player,
    #[serde_as(as = "Option<DisplayFromStr>")]
    winner: Option<Color>,
    #[serde_as(as = "StringWithSeparator<SpaceSeparator, SanPlus>")]
    moves: Vec<SanPlus>,
}

#[derive(Default, Serialize, Debug)]
struct Player {
    name: Option<String>,
    rating: Option<u16>,
}

impl Importer {
    fn new(filename: String) -> Importer {
        Importer {
            filename,
            client: reqwest::blocking::Client::builder()
                .timeout(Duration::from_secs(60))
                .build()
                .expect("client"),
            current: Game::default(),
            skip: false,
            batch: Vec::with_capacity(BATCH_SIZE),
        }
    }

    pub fn send(&mut self) {
        // println!("{}", serde_json::to_string(&self.batch).expect("serialize"));

        let res = self
            .client
            .put("http://127.0.0.1:9001/import/lichess")
            .json(&self.batch)
            .send()
            .expect("send batch");

        println!(
            "{}: {}: {} - {}",
            self.filename,
            self.batch
                .last()
                .and_then(|g| g.date.as_ref())
                .unwrap_or(&String::new()),
            res.status(),
            res.text().expect("decode response")
        );

        self.batch.clear();
    }
}

impl Visitor for Importer {
    type Result = ();

    fn begin_game(&mut self) {
        self.skip = false;
        self.current = Game::default();
    }

    fn header(&mut self, key: &[u8], value: RawHeader<'_>) {
        if key == b"White" {
            self.current.white.name = Some(value.decode_utf8().expect("White").into_owned());
        } else if key == b"Black" {
            self.current.black.name = Some(value.decode_utf8().expect("Black").into_owned());
        } else if key == b"WhiteElo" {
            if value.as_bytes() != b"?" {
                self.current.white.rating = Some(btoi::btoi(value.as_bytes()).expect("WhiteElo"));
            }
        } else if key == b"BlackElo" {
            if value.as_bytes() != b"?" {
                self.current.black.rating = Some(btoi::btoi(value.as_bytes()).expect("BlackElo"));
            }
        } else if key == b"TimeControl" {
            self.current.speed = Some(Speed::from_bytes(value.as_bytes()).expect("TimeControl"));
        } else if key == b"Variant" {
            self.current.variant = Some(value.decode_utf8().expect("Variant").into_owned());
        } else if key == b"Date" || key == b"UTCDate" {
            self.current.date = Some(value.decode_utf8().expect("Date").into_owned());
        } else if key == b"WhiteTitle" || key == b"BlackTitle" {
            if value.as_bytes() == b"BOT" {
                self.skip = true;
            }
        } else if key == b"Site" {
            self.current.id = Some(
                String::from_utf8(
                    value
                        .as_bytes()
                        .rsplitn(2, |ch| *ch == b'/')
                        .next()
                        .expect("Site")
                        .to_owned(),
                )
                .expect("Site"),
            );
        } else if key == b"Result" {
            match value.as_bytes() {
                b"1-0" => self.current.winner = Some(Color::White),
                b"0-1" => self.current.winner = Some(Color::Black),
                b"1/2-1/2" => self.current.winner = None,
                _ => self.skip = true,
            }
        } else if key == b"FEN" {
            self.current.fen = Some(value.decode_utf8().expect("FEN").into_owned());
        }
    }

    fn end_headers(&mut self) -> Skip {
        let rating =
            (self.current.white.rating.unwrap_or(0) + self.current.black.rating.unwrap_or(0)) / 2;

        let standard = self
            .current
            .variant
            .as_ref()
            .map_or(false, |name| name != "Standard");

        let probability = if standard {
            match self.current.speed.unwrap_or(Speed::Correspondence) {
                Speed::Correspondence | Speed::Classical => 1.0,

                _ if rating >= 2500 => 1.0,

                Speed::Rapid if rating >= 2200 => 1.0,
                Speed::Rapid if rating >= 2000 => 0.30,
                Speed::Rapid if rating >= 1800 => 0.18,
                Speed::Rapid if rating >= 1600 => 0.16,

                Speed::Blitz if rating >= 2200 => 0.16,
                Speed::Blitz if rating >= 2000 => 0.07,
                Speed::Blitz if rating >= 1800 => 0.05,
                Speed::Blitz if rating >= 1600 => 0.05,

                Speed::Bullet if rating >= 2200 => 0.16,
                Speed::Bullet if rating >= 2000 => 0.07,
                Speed::Bullet if rating >= 1800 => 0.05,
                Speed::Bullet if rating >= 1600 => 0.05,

                Speed::UltraBullet => 1.0,

                _ => 0.02,
            }
        } else {
            // variant games
            if rating >= 1600 {
                1.0
            } else {
                0.5
            }
        };

        let accept = min(
            self.current.white.rating.unwrap_or(0),
            self.current.black.rating.unwrap_or(0),
        ) >= 1501
            && probability >= thread_rng().sample(OpenClosed01)
            && !self.skip;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: SanPlus) {
        self.current.moves.push(san);
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self) {
        if !self.skip {
            self.batch
                .push(mem::replace(&mut self.current, Default::default()));

            if self.batch.len() >= BATCH_SIZE {
                self.send();
            }
        }
    }
}

fn main() -> Result<(), io::Error> {
    for arg in env::args().skip(1) {
        let file = File::open(&arg)?;

        let uncompressed: Box<dyn io::Read> = if arg.ends_with(".bz2") {
            Box::new(bzip2::read::MultiBzDecoder::new(file))
        } else {
            Box::new(file)
        };

        let mut reader = BufferedReader::new(uncompressed);

        let mut importer = Importer::new(arg);
        reader.read_all(&mut importer)?;
    }

    Ok(())
}
<|endoftext|>