<commit_before>// Copyright (c) 2017-2018, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

extern crate clap;
extern crate rav1e;
extern crate y4m;
#[macro_use]
extern crate scan_fmt;

mod common;
mod decoder;
use common::*;
use rav1e::*;

use std::slice;
use std::io;
use std::io::Write;
use std::io::Read;
use std::sync::Arc;
use decoder::Decoder;
use decoder::VideoDetails;

fn read_frame_batch<D: Decoder>(ctx: &mut Context, decoder: &mut D, video_info: VideoDetails) {
  loop {
    if ctx.needs_more_lookahead() {
      match decoder.read_frame(&video_info) {
        Ok(frame) => {
          match video_info.bit_depth {
            8 | 10 | 12 => {}
            _ => panic!("unknown input bit depth!")
          }

          let _ = ctx.send_frame(Some(Arc::new(frame)));
          continue;
        }
        _ => {
          let frames_to_be_coded = ctx.get_frame_count();
          // This is a hack, instead when EOF is reached simply "close" the encoder to input (flag)
          ctx.set_limit(frames_to_be_coded);
          ctx.flush();
        }
      }
    } else if !ctx.needs_more_frames(ctx.get_frame_count()) {
      ctx.flush();
    }
    break;
  }
}

// Encode and write a frame.
// Returns frame information in a `Result`.
fn process_frame(
  ctx: &mut Context, output_file: &mut dyn Write,
  y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
  mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>
) -> Result<Vec<FrameSummary>, ()> {
  let y4m_details = y4m_dec.get_video_details();
  let mut frame_summaries = Vec::new();
  read_frame_batch(ctx, y4m_dec, y4m_details);
  let pkt_wrapped = ctx.receive_packet();
  if let Ok(pkt) = pkt_wrapped {
    write_ivf_frame(output_file, pkt.number as u64, pkt.data.as_ref());
    if let Some(y4m_enc_uw) = y4m_enc.as_mut() {
      if let Some(ref rec) = pkt.rec {
        let pitch_y = if y4m_details.bit_depth > 8 { y4m_details.width * 2 } else { y4m_details.width };
        let chroma_sampling_period = y4m_details.chroma_sampling.sampling_period();
        let (pitch_uv, height_uv) = (
          pitch_y / chroma_sampling_period.0,
          y4m_details.height / chroma_sampling_period.1
        );

        let (mut rec_y, mut rec_u, mut rec_v) = (
          vec![128u8; pitch_y * y4m_details.height],
          vec![128u8; pitch_uv * height_uv],
          vec![128u8; pitch_uv * height_uv]
        );

        let (stride_y, stride_u, stride_v) = (
          rec.planes[0].cfg.stride,
          rec.planes[1].cfg.stride,
          rec.planes[2].cfg.stride
        );

        for (line, line_out) in rec.planes[0]
          .data_origin()
          .chunks(stride_y)
          .zip(rec_y.chunks_mut(pitch_y))
        {
          if y4m_details.bit_depth > 8 {
            unsafe {
              line_out.copy_from_slice(slice::from_raw_parts::<u8>(
                line.as_ptr() as (*const u8),
                pitch_y
              ));
            }
          } else {
            line_out.copy_from_slice(
              &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_y]
            );
          }
        }
        for (line, line_out) in rec.planes[1]
          .data_origin()
          .chunks(stride_u)
          .zip(rec_u.chunks_mut(pitch_uv))
        {
          if y4m_details.bit_depth > 8 {
            unsafe {
              line_out.copy_from_slice(slice::from_raw_parts::<u8>(
                line.as_ptr() as (*const u8),
                pitch_uv
              ));
            }
          } else {
            line_out.copy_from_slice(
              &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
            );
          }
        }
        for (line, line_out) in rec.planes[2]
          .data_origin()
          .chunks(stride_v)
          .zip(rec_v.chunks_mut(pitch_uv))
        {
          if y4m_details.bit_depth > 8 {
            unsafe {
              line_out.copy_from_slice(slice::from_raw_parts::<u8>(
                line.as_ptr() as (*const u8),
                pitch_uv
              ));
            }
          } else {
            line_out.copy_from_slice(
              &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
            );
          }
        }

        let rec_frame = y4m::Frame::new([&rec_y, &rec_u, &rec_v], None);
        y4m_enc_uw.write_frame(&rec_frame).unwrap();
      }
    }
    frame_summaries.push(pkt.into());
  }
  Ok(frame_summaries)
}

fn main() {
  let mut cli = parse_cli();
  let mut y4m_dec = y4m::decode(&mut cli.io.input).expect("input is not a y4m file");
  let video_info = y4m_dec.get_video_details();
  let mut y4m_enc = match cli.io.rec.as_mut() {
    Some(rec) => Some(
      y4m::encode(
        video_info.width,
        video_info.height,
        y4m::Ratio::new(video_info.time_base.den as usize, video_info.time_base.num as usize)
      ).with_colorspace(y4m_dec.get_colorspace())
        .write_header(rec)
        .unwrap()
    ),
    None => None
  };

  cli.enc.width = video_info.width;
  cli.enc.height = video_info.height;
  cli.enc.bit_depth = video_info.bit_depth;
  cli.enc.chroma_sampling = video_info.chroma_sampling;
  cli.enc.chroma_sample_position = video_info.chroma_sample_position;
  cli.enc.time_base = video_info.time_base;
  let cfg = Config {
    enc: cli.enc
  };

  let mut ctx = cfg.new_context();

  let stderr = io::stderr();
  let mut err = stderr.lock();

  let _ = writeln!(
    err,
    "{}x{} @ {}/{} fps",
    video_info.width,
    video_info.height,
    video_info.time_base.den,
    video_info.time_base.num
  );

  write_ivf_header(
    &mut cli.io.output,
    video_info.width,
    video_info.height,
    video_info.time_base.den as usize,
    video_info.time_base.num as usize
  );

  let mut progress = ProgressInfo::new(
    Rational { num: video_info.time_base.den, den: video_info.time_base.num },
    if cli.limit == 0 { None } else { Some(cli.limit) },
      cfg.enc.show_psnr
  );

  ctx.set_limit(cli.limit as u64);

  loop {
    match process_frame(&mut ctx, &mut cli.io.output, &mut y4m_dec, y4m_enc.as_mut()) {
      Ok(frame_info) => {
        for frame in frame_info {
          progress.add_frame(frame);
          let _ = if cli.verbose {
            writeln!(err, "{} - {}", frame, progress)
          } else {
            write!(err, "\r{}                    ", progress)
          };
        }
      },
      Err(_) => break,
    };

    if !ctx.needs_more_frames(progress.frames_encoded() as u64) {
      break;
    }

    cli.io.output.flush().unwrap();
  }

  let _ = write!(err, "\n{}\n", progress.print_stats());
}
<commit_msg>Require both reconstructed y4m writer and frame.<commit_after>// Copyright (c) 2017-2018, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

extern crate clap;
extern crate rav1e;
extern crate y4m;
#[macro_use]
extern crate scan_fmt;

mod common;
mod decoder;
use common::*;
use rav1e::*;

use std::slice;
use std::io;
use std::io::Write;
use std::io::Read;
use std::sync::Arc;
use decoder::Decoder;
use decoder::VideoDetails;

fn read_frame_batch<D: Decoder>(ctx: &mut Context, decoder: &mut D, video_info: VideoDetails) {
  loop {
    if ctx.needs_more_lookahead() {
      match decoder.read_frame(&video_info) {
        Ok(frame) => {
          match video_info.bit_depth {
            8 | 10 | 12 => {}
            _ => panic!("unknown input bit depth!")
          }

          let _ = ctx.send_frame(Some(Arc::new(frame)));
          continue;
        }
        _ => {
          let frames_to_be_coded = ctx.get_frame_count();
          // This is a hack, instead when EOF is reached simply "close" the encoder to input (flag)
          ctx.set_limit(frames_to_be_coded);
          ctx.flush();
        }
      }
    } else if !ctx.needs_more_frames(ctx.get_frame_count()) {
      ctx.flush();
    }
    break;
  }
}

// Encode and write a frame.
// Returns frame information in a `Result`.
fn process_frame(
  ctx: &mut Context, output_file: &mut dyn Write,
  y4m_dec: &mut y4m::Decoder<'_, Box<dyn Read>>,
  mut y4m_enc: Option<&mut y4m::Encoder<'_, Box<dyn Write>>>
) -> Result<Vec<FrameSummary>, ()> {
  let y4m_details = y4m_dec.get_video_details();
  let mut frame_summaries = Vec::new();
  read_frame_batch(ctx, y4m_dec, y4m_details);
  let pkt_wrapped = ctx.receive_packet();
  if let Ok(pkt) = pkt_wrapped {
    write_ivf_frame(output_file, pkt.number as u64, pkt.data.as_ref());
    if let (Some(ref mut y4m_enc_uw), Some(ref rec)) = (y4m_enc.as_mut(), &pkt.rec) {
      let pitch_y = if y4m_details.bit_depth > 8 { y4m_details.width * 2 } else { y4m_details.width };
      let chroma_sampling_period = y4m_details.chroma_sampling.sampling_period();
      let (pitch_uv, height_uv) = (
        pitch_y / chroma_sampling_period.0,
        y4m_details.height / chroma_sampling_period.1
      );

      let (mut rec_y, mut rec_u, mut rec_v) = (
        vec![128u8; pitch_y * y4m_details.height],
        vec![128u8; pitch_uv * height_uv],
        vec![128u8; pitch_uv * height_uv]
      );

      let (stride_y, stride_u, stride_v) = (
        rec.planes[0].cfg.stride,
        rec.planes[1].cfg.stride,
        rec.planes[2].cfg.stride
      );

      for (line, line_out) in rec.planes[0]
        .data_origin()
        .chunks(stride_y)
        .zip(rec_y.chunks_mut(pitch_y))
      {
        if y4m_details.bit_depth > 8 {
          unsafe {
            line_out.copy_from_slice(slice::from_raw_parts::<u8>(
              line.as_ptr() as (*const u8),
              pitch_y
            ));
          }
        } else {
          line_out.copy_from_slice(
            &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_y]
          );
        }
      }
      for (line, line_out) in rec.planes[1]
        .data_origin()
        .chunks(stride_u)
        .zip(rec_u.chunks_mut(pitch_uv))
      {
        if y4m_details.bit_depth > 8 {
          unsafe {
            line_out.copy_from_slice(slice::from_raw_parts::<u8>(
              line.as_ptr() as (*const u8),
              pitch_uv
            ));
          }
        } else {
          line_out.copy_from_slice(
            &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
          );
        }
      }
      for (line, line_out) in rec.planes[2]
        .data_origin()
        .chunks(stride_v)
        .zip(rec_v.chunks_mut(pitch_uv))
      {
        if y4m_details.bit_depth > 8 {
          unsafe {
            line_out.copy_from_slice(slice::from_raw_parts::<u8>(
              line.as_ptr() as (*const u8),
              pitch_uv
            ));
          }
        } else {
          line_out.copy_from_slice(
            &line.iter().map(|&v| v as u8).collect::<Vec<u8>>()[..pitch_uv]
          );
        }
      }

      let rec_frame = y4m::Frame::new([&rec_y, &rec_u, &rec_v], None);
      y4m_enc_uw.write_frame(&rec_frame).unwrap();
    }
    frame_summaries.push(pkt.into());
  }
  Ok(frame_summaries)
}

fn main() {
  let mut cli = parse_cli();
  let mut y4m_dec = y4m::decode(&mut cli.io.input).expect("input is not a y4m file");
  let video_info = y4m_dec.get_video_details();
  let mut y4m_enc = match cli.io.rec.as_mut() {
    Some(rec) => Some(
      y4m::encode(
        video_info.width,
        video_info.height,
        y4m::Ratio::new(video_info.time_base.den as usize, video_info.time_base.num as usize)
      ).with_colorspace(y4m_dec.get_colorspace())
        .write_header(rec)
        .unwrap()
    ),
    None => None
  };

  cli.enc.width = video_info.width;
  cli.enc.height = video_info.height;
  cli.enc.bit_depth = video_info.bit_depth;
  cli.enc.chroma_sampling = video_info.chroma_sampling;
  cli.enc.chroma_sample_position = video_info.chroma_sample_position;
  cli.enc.time_base = video_info.time_base;
  let cfg = Config {
    enc: cli.enc
  };

  let mut ctx = cfg.new_context();

  let stderr = io::stderr();
  let mut err = stderr.lock();

  let _ = writeln!(
    err,
    "{}x{} @ {}/{} fps",
    video_info.width,
    video_info.height,
    video_info.time_base.den,
    video_info.time_base.num
  );

  write_ivf_header(
    &mut cli.io.output,
    video_info.width,
    video_info.height,
    video_info.time_base.den as usize,
    video_info.time_base.num as usize
  );

  let mut progress = ProgressInfo::new(
    Rational { num: video_info.time_base.den, den: video_info.time_base.num },
    if cli.limit == 0 { None } else { Some(cli.limit) },
      cfg.enc.show_psnr
  );

  ctx.set_limit(cli.limit as u64);

  loop {
    match process_frame(&mut ctx, &mut cli.io.output, &mut y4m_dec, y4m_enc.as_mut()) {
      Ok(frame_info) => {
        for frame in frame_info {
          progress.add_frame(frame);
          let _ = if cli.verbose {
            writeln!(err, "{} - {}", frame, progress)
          } else {
            write!(err, "\r{}                    ", progress)
          };
        }
      },
      Err(_) => break,
    };

    if !ctx.needs_more_frames(progress.frames_encoded() as u64) {
      break;
    }

    cli.io.output.flush().unwrap();
  }

  let _ = write!(err, "\n{}\n", progress.print_stats());
}
<|endoftext|>


<commit_before>
//! main entry for `gerrit-rust`

extern crate chrono;
extern crate clap;
extern crate env_logger;
extern crate gerritlib;
extern crate git2;
#[macro_use]
extern crate log;
extern crate rustc_serialize;
extern crate toml_config;

pub mod changes;
pub mod config;
pub mod topic;

use clap::App;
use gerritlib::error::GGRError;
use std::error::Error;

fn init_log() {
    let format = |record: &log::LogRecord| {
        format!("[{:5.5}] [{}] [{}] - {}", record.level(),
                chrono::Local::now().to_rfc3339(),
                record.location().module_path(),
                record.args())
    };

    let mut builder = env_logger::LogBuilder::new();
    builder.format(format).filter(None, log::LogLevelFilter::Info);

    if let Ok(ref rl) = std::env::var("RUST_LOG") {
        builder.parse(rl);
    }

    let _ = builder.init();
}

fn main() {
    init_log();

    let mut app = App::new("gerrit-rust")
        .author("Silvio Fricke <silvio.fricke@gmail.com>")
        .about("some gerrit tools")
        .subcommand(topic::menu())
        .subcommand(changes::menu())
        .subcommand(config::menu())
        ;

    let matches = app.clone().get_matches();

    let configfile = config::ConfigFile::discover(".", ".ggr.conf").expect("ConfigFile problem");
    let config = config::Config::from_configfile(configfile);
    if ! config.is_valid() {
        panic!(GGRError::General("problem with configfile".to_string()));
    }

    let out = match matches.subcommand() {
        ("topic", Some(x)) => { topic::manage(x, config) },
        ("changes", Some(x)) => { changes::manage(x, config) },
        ("config", Some(x)) => { config::manage(x) },
        _ => { let _ = app.print_help(); Ok(()) },
    };

    match out {
        Ok(_) => {},
        Err(e) => {
            println!("Error: {}", e.description())
        },
    }
}
<commit_msg>print a regular error message if config file is not found, this fixes #1<commit_after>
//! main entry for `gerrit-rust`

extern crate chrono;
extern crate clap;
extern crate env_logger;
extern crate gerritlib;
extern crate git2;
#[macro_use]
extern crate log;
extern crate rustc_serialize;
extern crate toml_config;

pub mod changes;
pub mod config;
pub mod topic;

use clap::App;
use gerritlib::error::GGRError;
use std::error::Error;
use std::process::exit;

fn init_log() {
    let format = |record: &log::LogRecord| {
        format!("[{:5.5}] [{}] [{}] - {}", record.level(),
                chrono::Local::now().to_rfc3339(),
                record.location().module_path(),
                record.args())
    };

    let mut builder = env_logger::LogBuilder::new();
    builder.format(format).filter(None, log::LogLevelFilter::Info);

    if let Ok(ref rl) = std::env::var("RUST_LOG") {
        builder.parse(rl);
    }

    let _ = builder.init();
}

fn main() {
    init_log();

    let mut app = App::new("gerrit-rust")
        .author("Silvio Fricke <silvio.fricke@gmail.com>")
        .about("some gerrit tools")
        .subcommand(topic::menu())
        .subcommand(changes::menu())
        .subcommand(config::menu())
        ;

    let matches = app.clone().get_matches();

    let configfile = match config::ConfigFile::discover(".", ".ggr.conf") {
        Ok(c) => c,
        Err(x) => {
            println!("Problem with loading of config file:");
            println!("{}", x.to_string());
            exit(-1);
        },
    };
    let config = config::Config::from_configfile(configfile);
    if ! config.is_valid() {
        panic!(GGRError::General("problem with configfile".to_string()));
    }

    let out = match matches.subcommand() {
        ("topic", Some(x)) => { topic::manage(x, config) },
        ("changes", Some(x)) => { changes::manage(x, config) },
        ("config", Some(x)) => { config::manage(x) },
        _ => { let _ = app.print_help(); Ok(()) },
    };

    match out {
        Ok(_) => {},
        Err(e) => {
            println!("Error: {}", e.description())
        },
    }
}
<|endoftext|>


<commit_before>/// Character level parsers

use internal::{IResult,Needed};

/// matches one of the provided characters
#[macro_export]
macro_rules! one_of (
  ($i:expr, $inp: expr) => (
    {
      if $i.is_empty() {
        $crate::IResult::Incomplete::<_, _>($crate::Needed::Size(1))
      } else {
        #[inline(always)]
        fn as_bytes<T: $crate::AsBytes>(b: &T) -> &[u8] {
          b.as_bytes()
        }

        let expected = $inp;
        let bytes = as_bytes(&expected);
        one_of_bytes!($i, bytes)
      }
    }
  );
);

#[doc(hidden)]
#[macro_export]
macro_rules! one_of_bytes (
  ($i:expr, $bytes: expr) => (
    {
      if $i.is_empty() {
        $crate::IResult::Incomplete::<_, _>($crate::Needed::Size(1))
      } else {
        let mut found = false;

        for &i in $bytes {
          if i == $i[0] {
            found = true;
            break;
          }
        }

        if found {
          $crate::IResult::Done(&$i[1..], $i[0] as char)
        } else {
          $crate::IResult::Error(error_position!($crate::ErrorKind::OneOf, $i))
        }
      }
    }
  );
);

/// matches anything but the provided characters
#[macro_export]
macro_rules! none_of (
  ($i:expr, $inp: expr) => (
    {
      if $i.is_empty() {
        $crate::IResult::Incomplete::<_, _>($crate::Needed::Size(1))
      } else {
        #[inline(always)]
        fn as_bytes<T: $crate::AsBytes>(b: &T) -> &[u8] {
          b.as_bytes()
        }

        let expected = $inp;
        let bytes = as_bytes(&expected);
        none_of_bytes!($i, bytes)
      }
    }
  );
);

#[doc(hidden)]
#[macro_export]
macro_rules! none_of_bytes (
  ($i:expr, $bytes: expr) => (
    {
      if $i.is_empty() {
        $crate::IResult::Incomplete::<_, _>($crate::Needed::Size(1))
      } else {
        let mut found = false;

        for &i in $bytes {
          if i == $i[0] {
            found = true;
            break;
          }
        }

        if !found {
          $crate::IResult::Done(&$i[1..], $i[0] as char)
        } else {
          $crate::IResult::Error(error_position!($crate::ErrorKind::NoneOf, $i))
        }
      }
    }
  );
);

/// matches one character: `char!(char) => &[u8] -> IResult<&[u8], char>
#[macro_export]
macro_rules! char (
  ($i:expr, $c: expr) => (
    {
      if $i.is_empty() {
        let res: $crate::IResult<&[u8], char> = $crate::IResult::Incomplete($crate::Needed::Size(1));
        res
      } else {
        if $i[0] == $c as u8 {
          $crate::IResult::Done(&$i[1..], $i[0] as char)
        } else {
          $crate::IResult::Error(error_position!($crate::ErrorKind::Char, $i))
        }
      }
    }
  );
);

named!(#[doc="Matches a newline character '\\n'"], pub newline<char>, char!('\n'));

named!(#[doc="Matches a tab character '\\t'"], pub tab<char>, char!('\t'));

pub fn anychar(input:&[u8]) -> IResult<&[u8], char> {
  if input.is_empty() {
    IResult::Incomplete(Needed::Size(1))
  } else {
    IResult::Done(&input[1..], input[0] as char)
  }
}

#[cfg(test)]
mod tests {
  use internal::IResult::*;
  use util::ErrorKind;

  #[test]
  fn one_of() {
    named!(f<char>, one_of!("ab"));

    let a = &b"abcd"[..];
    assert_eq!(f(a), Done(&b"bcd"[..], 'a'));

    let b = &b"cde"[..];
    assert_eq!(f(b), Error(error_position!(ErrorKind::OneOf, b)));
  }

  #[test]
  fn none_of() {
    named!(f<char>, none_of!("ab"));

    let a = &b"abcd"[..];
    assert_eq!(f(a), Error(error_position!(ErrorKind::NoneOf, a)));

    let b = &b"cde"[..];
    assert_eq!(f(b), Done(&b"de"[..], 'c'));
  }

  #[test]
  fn char() {
    named!(f<char>, char!('c'));

    let a = &b"abcd"[..];
    assert_eq!(f(a), Error(error_position!(ErrorKind::Char, a)));

    let b = &b"cde"[..];
    assert_eq!(f(b), Done(&b"de"[..], 'c'));
  }

}
<commit_msg>convert the character combinators to the input traits<commit_after>/// Character level parsers

use internal::{IResult,Needed};

/// matches one of the provided characters
#[macro_export]
macro_rules! one_of (
  ($i:expr, $inp: expr) => (
    {
      use $crate::Slice;
      use $crate::AsChar;
      use $crate::FindToken;
      use $crate::InputIter;

      match ($i).iter_elements().next().map(|c| {
        c.find_token($inp)
      }) {
        None        => $crate::IResult::Incomplete::<_, _>($crate::Needed::Size(1)),
        Some(false) => $crate::IResult::Error(error_position!($crate::ErrorKind::OneOf, $i)),
        //the unwrap should be safe here
        Some(true)  => $crate::IResult::Done($i.slice(1..), $i.iter_elements().next().unwrap().as_char())
      }
    }
  );
);

/// matches anything but the provided characters
#[macro_export]
macro_rules! none_of (
  ($i:expr, $inp: expr) => (
    {
      use $crate::Slice;
      use $crate::AsChar;
      use $crate::FindToken;
      use $crate::InputIter;

      match ($i).iter_elements().next().map(|c| {
        !c.find_token($inp)
      }) {
        None        => $crate::IResult::Incomplete::<_, _>($crate::Needed::Size(1)),
        Some(false) => $crate::IResult::Error(error_position!($crate::ErrorKind::NoneOf, $i)),
        //the unwrap should be safe here
        Some(true)  => $crate::IResult::Done($i.slice(1..), $i.iter_elements().next().unwrap().as_char())
      }
    }
  );
);

/// matches one character: `char!(char) => &[u8] -> IResult<&[u8], char>
#[macro_export]
macro_rules! char (
  ($i:expr, $c: expr) => (
    {
      use $crate::Slice;
      use $crate::AsChar;
      use $crate::InputIter;

      match ($i).iter_elements().next().map(|c| {
        c.as_char() == $c
      }) {
        None        => $crate::IResult::Incomplete::<_, _>($crate::Needed::Size(1)),
        Some(false) => $crate::IResult::Error(error_position!($crate::ErrorKind::Char, $i)),
        //the unwrap should be safe here
        Some(true)  => $crate::IResult::Done($i.slice(1..), $i.iter_elements().next().unwrap().as_char())
      }
    }
  );
);

named!(#[doc="Matches a newline character '\\n'"], pub newline<char>, char!('\n'));

named!(#[doc="Matches a tab character '\\t'"], pub tab<char>, char!('\t'));

pub fn anychar(input:&[u8]) -> IResult<&[u8], char> {
  if input.is_empty() {
    IResult::Incomplete(Needed::Size(1))
  } else {
    IResult::Done(&input[1..], input[0] as char)
  }
}

#[cfg(test)]
mod tests {
  use internal::IResult::*;
  use util::ErrorKind;

  #[test]
  fn one_of() {
    named!(f<char>, one_of!("ab"));

    let a = &b"abcd"[..];
    assert_eq!(f(a), Done(&b"bcd"[..], 'a'));

    let b = &b"cde"[..];
    assert_eq!(f(b), Error(error_position!(ErrorKind::OneOf, b)));
  }

  #[test]
  fn none_of() {
    named!(f<char>, none_of!("ab"));

    let a = &b"abcd"[..];
    assert_eq!(f(a), Error(error_position!(ErrorKind::NoneOf, a)));

    let b = &b"cde"[..];
    assert_eq!(f(b), Done(&b"de"[..], 'c'));
  }

  #[test]
  fn char() {
    named!(f<char>, char!('c'));

    let a = &b"abcd"[..];
    assert_eq!(f(a), Error(error_position!(ErrorKind::Char, a)));

    let b = &b"cde"[..];
    assert_eq!(f(b), Done(&b"de"[..], 'c'));
  }

}
<|endoftext|>


<commit_before>//! ROM image loading code

use std::cmp;
use std::str;

/// The (decoded) SNES header
pub struct RomHeader {
    /// ASCII title, filled with spaces to 21 Bytes
    #[allow(dead_code)] // FIXME Use this or drop this
    title: [u8; 21],
    rom_size: u32,
    ram_size: u32,
    checksum: u16,
    rom_type: RomType,
}

#[derive(Debug, PartialEq, Eq)]
enum RomType {
    LoRom,
    HiRom,
}

impl RomHeader {
    /// Loads the ROM header from the given byte slice (must be exactly 64 bytes large).
    /// `rom_type` is the expected type of the ROM header, based on its location. It will be
    /// used for the actual type if the type inside the header mismatches. This is apparently
    /// required for some broken ROMs.
    fn load(bytes: &[u8], rom_type: RomType) -> Result<RomHeader, ()> {
        // The header size must be correct (the ROM loader won't pass a wrong size)
        assert_eq!(bytes.len(), 64);

        debug!("raw rom header: {:?}", bytes);

        // Byte 28/29 is the checksum's complement followed by the checksum itself in Byte 30/31.
        // The checksum is the sum of all bytes in the ROM (truncated to 16 bits). It can be
        // validated by the loader.
        // 16 bit, little-endian.
        let check_inv = (bytes[29] as u16) << 8 | bytes[28] as u16;
        let checksum = (bytes[31] as u16) << 8 | bytes[30] as u16;
        if check_inv != !checksum {
            debug!("checksum invalid: stored complement is ${:04X}, stored checksum is ${:04X}",
                check_inv, checksum);
            return Err(());
        }

        let mut title = [0; 21];
        let mut warned = false;
        for (i, c) in bytes[0..21].iter().enumerate() {
            match *c {
                0x20 ... 0x7E => {
                    title[i] = *c;
                }
                _ => {
                    if !warned {
                        warn!("title contains non-ascii byte: ${:02X}", *c);
                        warned = true;
                    }
                    title[i] = b' ';
                }
            }
        }

        info!("loading '{}'", str::from_utf8(&title).unwrap().trim_right());

        // 21) ROM makeup byte
        // `---smmmm`
        // * `s`: Speed (0: SlowROM, 1: FastROM)
        // * `m`: Map mode
        //  * `0000`: LoROM
        //  * `0001`: HiROM
        //  * `0010`: LoROM + S-DD1
        //  * `0011`: LoROM + SA-1
        //  * `0101`: ExHiROM
        //  * `1010`: HiROM + SPC7110
        if bytes[21] & 0x10 == 0x10 {
            warn!("FastROM not yet implemented!");
        }

        let header_rom_type = match bytes[21] & 0x0f {
            0 => RomType::LoRom,
            1 => RomType::HiRom,
            t => {
                error!("unknown / unimplemented ROM type {}", t);
                return Err(())
            }
        };

        if header_rom_type == rom_type {
            info!("type: {:?}", rom_type);
        } else {
            warn!("expected rom type {:?}, got {:?} (assuming {:?})", rom_type, header_rom_type,
                rom_type);
        }

        // bytes[22] is the chipset info. For now, we don't care about that.
        debug!("chipset: 0x{:02X}", bytes[22]);

        let rom_size = 0x400 << bytes[23] as u32;
        let ram_size = 0x400 << bytes[24] as u32;
        debug!("ROM/RAM size values: {:02X} {:02X}", bytes[23], bytes[24]);
        info!("{} KB of ROM, {} KB of cartridge RAM", rom_size / 1024, ram_size / 1024);

        // bytes[25-26] is a vendor code (doesn't matter)
        debug!("vendor code: 0x{:02X}{:02X}", bytes[25], bytes[26]);
        // 27 = version (also doesn't matter for us)
        debug!("version: 0x{:02X}", bytes[27]);

        Ok(RomHeader {
            title: title,
            rom_size: rom_size,
            ram_size: ram_size,
            checksum: checksum,
            rom_type: rom_type,
        })
    }
}

/// A ROM image
pub struct Rom {
    header: RomHeader,
    ram: Vec<u8>,
    rom: Vec<u8>,
}

impl_save_state!(Rom { ram } ignore { header, rom });

impl Rom {
    /// Loads a ROM from raw data.
    pub fn from_bytes(mut bytes: &[u8]) -> Result<Rom, ()> {
        // FIXME Return a proper error!

        debug!("raw size: {} bytes (${:X})", bytes.len(), bytes.len());

        // ROMs may begin with a 512 Bytes SMC header. It needs to go.
        match bytes.len() % 1024 {
            512 => {
                debug!("stripping SMC header");
                bytes = &bytes[512..];
            }
            0 => {},
            n => panic!("len() % 1024 == {} (expected 512 or 0)", n),
        }

        // Read the SNES header. It is located in a really stupid location which depends on whether
        // the ROM is HiROM or LoROM, so we must check both locations and pick the right one.
        // LoROM header is at 0x7FFF - 63, HiROM is at 0xFFFF - 63 (it is 64 Bytes large)

        // Test LoROM first (it's more common)
        let header = try!(RomHeader::load(&bytes[0x7FFF - 63..0x7FFF + 1], RomType::LoRom)
            .or_else(|_| RomHeader::load(&bytes[0xFFFF - 63..0xFFFF + 1], RomType::HiRom)));

        if bytes.len() != header.rom_size as usize {
            warn!("raw ROM is {} KB, but header specifies {} KB",
                bytes.len() / 1024, header.rom_size / 1024);
        }

        // Create the right amount of RAM...
        let ram = vec![0; header.ram_size as usize];
        // ...and copy the ROM
        let rom = bytes.iter().cloned().cycle()
            .take(cmp::max(header.rom_size as usize, bytes.len())).collect();

        // Calculate the ROM's checksum
        let mut checksum: u16 = 0;
        for &byte in &rom {
            checksum = checksum.wrapping_add(byte as u16);
        }

        debug!("computed checksum: ${:04X}", checksum);

        if header.checksum != checksum {
            warn!("incorrect checksum: computed ${:04X}, expected ${:04X}",
                checksum, header.checksum);
        }

        Ok(Rom {
            header: header,
            ram: ram,
            rom: rom,
        })
    }

    pub fn get_title(&self) -> Option<&str> {
        str::from_utf8(&self.header.title).ok().map(|s| s.trim_right())
    }

    fn resolve_lorom(&mut self, bank: u8, addr: u16) -> &mut u8 {
        match addr {
            0x0000 ... 0x7fff => {
                // Cartridge RAM mapped to the low 32 KB
                // (there's other stuff here, but that's handled much earlier than we are called)
                match bank {
                    0x70 ... 0x7d => {
                        let a = (bank as u32 - 0x70) * 0x8000 + addr as u32;
                        self.ram.get_mut(a as usize).unwrap_or_else(|| out_of_ram_bounds(bank, addr, a))
                    }
                    0xfe ... 0xff => {
                        // last 64k of RAM
                        let start = self.ram.len() as u32 - 64 * 1024;
                        let a = start + (bank - 0xfe) as u32 * 0x8000 + addr as u32;
                        self.ram.get_mut(a as usize).unwrap_or_else(|| out_of_ram_bounds(bank, addr, a))
                    }
                    _ => {
                        // 0x40 ... 0x6f | 0x7e ... 0xfd
                        panic!("attempted to access unmapped address: ${:02X}:{:04X}", bank, addr);
                    }
                }
            },
            0x8000 ... 0xffff => match bank {
                // LoROM is mapped to the higher 8 pages
                0xfe => {
                    let a = 0x3f0000 + addr as u32 - 0x8000;
                    self.rom.get_mut(a as usize).unwrap_or_else(|| out_of_rom_bounds(bank, addr, a))
                }
                0xff => {
                    let a = 0x3f8000 + addr as u32 - 0x8000;
                    self.rom.get_mut(a as usize).unwrap_or_else(|| out_of_rom_bounds(bank, addr, a))
                }
                0x80 ... 0xfd | 0x00 ... 0x7d => {
                    // `& !0x80` because 0x80-0xFD mirrors 0x00-0x7D
                    let a = (bank as u32 & !0x80) * 0x8000 + addr as u32 - 0x8000;
                    self.rom.get_mut(a as usize).unwrap_or_else(|| out_of_rom_bounds(bank, addr, a))
                }
                _ => panic!("attempted to access unmapped address: ${:02X}:{:04X}", bank, addr)
            },
            _ => unreachable!()
        }
    }

    fn resolve_hirom(&mut self, bank: u8, addr: u16) -> &mut u8 {
        let addr = addr as usize;
        match bank {
            0x00 ... 0x3f | 0x80 ... 0xbf if addr >= 0x8000 => {
                &mut self.rom[(bank as usize & 0x3f) << 16 | addr]
            }
            0x20 ... 0x3f | 0xa0 ... 0xbf if addr >= 0x6000 && addr <= 0x7fff => {
                let a = ((bank as usize & 0x3f) - 0x20) * 0x2000 + (addr - 0x6000);
                &mut self.ram[a]
            }
            0x40 ... 0x7d | 0xc0 ... 0xfd => {
                &mut self.rom[((bank as usize & 0x7f) - 0x40) << 16 | addr]
            }
            0x7e ... 0x7f => unreachable!(),    // WRAM banks
            0xfe ... 0xff => {
                &mut self.rom[(bank as usize - 0xfe + 0x3e) << 16 | addr]
            }
            _ => panic!("attempted to access unmapped address: ${:02}:{:04X}", bank, addr),
        }
    }

    fn resolve_addr(&mut self, bank: u8, addr: u16) -> &mut u8 {
        match self.header.rom_type {
            RomType::LoRom => self.resolve_lorom(bank, addr),
            RomType::HiRom => self.resolve_hirom(bank, addr),
        }
    }
}

impl Rom {
    pub fn load(&mut self, bank: u8, addr: u16) -> u8 {
        *self.resolve_addr(bank, addr)
    }

    pub fn store(&mut self, bank: u8, addr: u16, value: u8) {
        if addr >= 0x8000 {
            warn!("writing ${:02X} to ROM address ${:02X}:{:04X}", value, bank, addr);
        }
        *self.resolve_addr(bank, addr) = value;
    }
}

fn out_of_ram_bounds(bank: u8, addr: u16, abs: u32) -> ! {
    panic!("RAM access out of bounds at {:02X}:{:04X} -> {:04X}",
        bank, addr, abs)
}

fn out_of_rom_bounds(bank: u8, addr: u16, abs: u32) -> ! {
    panic!("ROM access out of bounds at {:02X}:{:04X} -> {:06X}",
        bank, addr, abs)
}
<commit_msg>Use score-based ROM type detection<commit_after>//! ROM image loading code

use std::cmp;
use std::str;

/// The (decoded) SNES header
pub struct RomHeader {
    /// ASCII title, filled with spaces to 21 Bytes
    #[allow(dead_code)] // FIXME Use this or drop this
    title: [u8; 21],
    rom_size: u32,
    ram_size: u32,
    checksum: u16,
    rom_type: RomType,
}

#[derive(Debug, PartialEq, Eq)]
enum RomType {
    LoRom,
    HiRom,
}

impl RomHeader {
    fn dump(&self) {
        info!("'{}'", str::from_utf8(&self.title).unwrap_or("").trim_right());
        info!("{} KB ROM / {} KB Cartridge RAM", self.rom_size / 1024, self.ram_size / 1024);
    }

    /// Loads the ROM header from the given byte slice (must be exactly 64 bytes large).
    /// `rom_type` is the expected type of the ROM header, based on its location.
    ///
    /// Returns the decoded `RomHeader` and a scoring value. The higher the score, the more likely
    /// the header matches the `RomType`.
    fn load(bytes: &[u8], rom_type: RomType) -> (RomHeader, i16) {
        // The header size must be correct (the ROM loader won't pass a wrong size)
        assert_eq!(bytes.len(), 64);

        // Score value. Decremented whenever something isn't right. Subject to tweaking.
        let mut score = 0;

        debug!("loading {:?} header", rom_type);
        debug!("raw rom header: {:?}", bytes);

        // Byte 28/29 is the checksum's complement followed by the checksum itself in Byte 30/31.
        // The checksum is the sum of all bytes in the ROM (truncated to 16 bits). It can be
        // validated by the loader.
        // 16 bit, little-endian.
        let check_inv = (bytes[29] as u16) << 8 | bytes[28] as u16;
        let rom_checksum = (bytes[31] as u16) << 8 | bytes[30] as u16;
        if check_inv != !rom_checksum {
            debug!("checksum invalid: stored complement is ${:04X}, stored checksum is ${:04X}",
                check_inv, rom_checksum);
            score -= 4;
        }

        let mut title = [0; 21];
        let mut warned = false;
        for (i, c) in bytes[0..21].iter().enumerate() {
            match *c {
                0x20 ... 0x7E => {
                    title[i] = *c;
                }
                _ => {
                    if !warned {
                        debug!("title contains non-ascii byte: ${:02X}", *c);
                        warned = true;
                        score -= 2;
                    }
                    title[i] = b' ';
                }
            }
        }

        debug!("loading '{}'", str::from_utf8(&title).unwrap().trim_right());

        // 21) ROM makeup byte
        // `---smmmm`
        // * `s`: Speed (0: SlowROM, 1: FastROM)
        // * `m`: Map mode
        //  * `0000`: LoROM
        //  * `0001`: HiROM
        //  * `0010`: LoROM + S-DD1
        //  * `0011`: LoROM + SA-1
        //  * `0101`: ExHiROM
        //  * `1010`: HiROM + SPC7110
        if bytes[21] & 0x10 == 0x10 {
            debug!("FastROM not yet implemented!");
        }

        let header_rom_type = match bytes[21] & 0x0f {
            0 => RomType::LoRom,
            1 => RomType::HiRom,
            t => {
                debug!("unknown / unimplemented ROM type {}", t);
                score -= 10;    // until we actually implement this (FIXME Dirty hack)
                RomType::LoRom
            }
        };

        if header_rom_type == rom_type {
            debug!("type: {:?}", rom_type);
        } else {
            debug!("expected rom type {:?}, got {:?} (assuming {:?})", rom_type, header_rom_type,
                rom_type);
            score -= 3;
        }

        // bytes[22] is the chipset info. For now, we don't care about that.
        debug!("chipset: 0x{:02X}", bytes[22]);

        let rom_size = 0x400 << bytes[23] as u32;
        let ram_size = 0x400 << bytes[24] as u32;
        debug!("ROM/RAM size values: {:02X} {:02X}", bytes[23], bytes[24]);
        debug!("{} KB of ROM, {} KB of cartridge RAM", rom_size / 1024, ram_size / 1024);

        // bytes[25-26] is a vendor code (doesn't matter)
        debug!("vendor code: 0x{:02X}{:02X}", bytes[25], bytes[26]);
        // 27 = version (also doesn't matter for us)
        debug!("version: 0x{:02X}", bytes[27]);

        (RomHeader {
            title: title,
            rom_size: rom_size,
            ram_size: ram_size,
            checksum: rom_checksum,
            rom_type: rom_type,
        }, score)
    }
}

/// A ROM image
pub struct Rom {
    header: RomHeader,
    ram: Vec<u8>,
    rom: Vec<u8>,
}

impl_save_state!(Rom { ram } ignore { header, rom });

impl Rom {
    /// Loads a ROM from raw data.
    pub fn from_bytes(mut bytes: &[u8]) -> Result<Rom, ()> {
        // FIXME Return a proper error!

        debug!("raw size: {} bytes (${:X})", bytes.len(), bytes.len());

        // ROMs may begin with a 512 Bytes SMC header. It needs to go.
        match bytes.len() % 1024 {
            512 => {
                info!("stripping SMC header");
                bytes = &bytes[512..];
            }
            0 => {},
            n => panic!("len() % 1024 == {} (expected 512 or 0)", n),
        }

        // Read the SNES header. It is located in a really stupid location which depends on whether
        // the ROM is HiROM or LoROM, so we must check both locations and pick the right one.
        // LoROM header is at 0x7FFF - 63, HiROM is at 0xFFFF - 63 (it is 64 Bytes large)

        let (lo_header, lo_score) = RomHeader::load(&bytes[0x7FFF - 63..0x7FFF + 1],
                                                    RomType::LoRom);
        let (hi_header, hi_score) = RomHeader::load(&bytes[0xFFFF - 63..0xFFFF + 1],
                                                    RomType::HiRom);

        info!("LoROM/HiROM scores: {}, {}", lo_score, hi_score);
        let header = if lo_score > hi_score {
            lo_header
        } else {
            hi_header
        };

        header.dump();

        if bytes.len() != header.rom_size as usize {
            warn!("raw ROM is {} KB, but header specifies {} KB",
                bytes.len() / 1024, header.rom_size / 1024);
        }

        // Create the right amount of RAM...
        let ram = vec![0; header.ram_size as usize];
        // ...and copy the ROM
        let rom = bytes.iter().cloned().cycle()
            .take(cmp::max(header.rom_size as usize, bytes.len())).collect();

        // Calculate the ROM's checksum
        let mut checksum: u16 = 0;
        for &byte in &rom {
            checksum = checksum.wrapping_add(byte as u16);
        }

        debug!("computed checksum: ${:04X}", checksum);

        if header.checksum != checksum {
            warn!("incorrect checksum: computed ${:04X}, expected ${:04X}",
                checksum, header.checksum);
        }

        Ok(Rom {
            header: header,
            ram: ram,
            rom: rom,
        })
    }

    pub fn get_title(&self) -> Option<&str> {
        str::from_utf8(&self.header.title).ok().map(|s| s.trim_right())
    }

    fn resolve_lorom(&mut self, bank: u8, addr: u16) -> &mut u8 {
        match addr {
            0x0000 ... 0x7fff => {
                // Cartridge RAM mapped to the low 32 KB
                // (there's other stuff here, but that's handled much earlier than we are called)
                match bank {
                    0x70 ... 0x7d => {
                        let a = (bank as u32 - 0x70) * 0x8000 + addr as u32;
                        self.ram.get_mut(a as usize).unwrap_or_else(|| out_of_ram_bounds(bank, addr, a))
                    }
                    0xfe ... 0xff => {
                        // last 64k of RAM
                        let start = self.ram.len() as u32 - 64 * 1024;
                        let a = start + (bank - 0xfe) as u32 * 0x8000 + addr as u32;
                        self.ram.get_mut(a as usize).unwrap_or_else(|| out_of_ram_bounds(bank, addr, a))
                    }
                    _ => {
                        // 0x40 ... 0x6f | 0x7e ... 0xfd
                        panic!("attempted to access unmapped address: ${:02X}:{:04X}", bank, addr);
                    }
                }
            },
            0x8000 ... 0xffff => match bank {
                // LoROM is mapped to the higher 8 pages
                0xfe => {
                    let a = 0x3f0000 + addr as u32 - 0x8000;
                    self.rom.get_mut(a as usize).unwrap_or_else(|| out_of_rom_bounds(bank, addr, a))
                }
                0xff => {
                    let a = 0x3f8000 + addr as u32 - 0x8000;
                    self.rom.get_mut(a as usize).unwrap_or_else(|| out_of_rom_bounds(bank, addr, a))
                }
                0x80 ... 0xfd | 0x00 ... 0x7d => {
                    // `& !0x80` because 0x80-0xFD mirrors 0x00-0x7D
                    let a = (bank as u32 & !0x80) * 0x8000 + addr as u32 - 0x8000;
                    self.rom.get_mut(a as usize).unwrap_or_else(|| out_of_rom_bounds(bank, addr, a))
                }
                _ => panic!("attempted to access unmapped address: ${:02X}:{:04X}", bank, addr)
            },
            _ => unreachable!()
        }
    }

    fn resolve_hirom(&mut self, bank: u8, addr: u16) -> &mut u8 {
        let addr = addr as usize;
        match bank {
            0x00 ... 0x3f | 0x80 ... 0xbf if addr >= 0x8000 => {
                &mut self.rom[(bank as usize & 0x3f) << 16 | addr]
            }
            0x20 ... 0x3f | 0xa0 ... 0xbf if addr >= 0x6000 && addr <= 0x7fff => {
                let a = ((bank as usize & 0x3f) - 0x20) * 0x2000 + (addr - 0x6000);
                &mut self.ram[a]
            }
            0x40 ... 0x7d | 0xc0 ... 0xfd => {
                &mut self.rom[((bank as usize & 0x7f) - 0x40) << 16 | addr]
            }
            0x7e ... 0x7f => unreachable!(),    // WRAM banks
            0xfe ... 0xff => {
                &mut self.rom[(bank as usize - 0xfe + 0x3e) << 16 | addr]
            }
            _ => panic!("attempted to access unmapped address: ${:02}:{:04X}", bank, addr),
        }
    }

    fn resolve_addr(&mut self, bank: u8, addr: u16) -> &mut u8 {
        match self.header.rom_type {
            RomType::LoRom => self.resolve_lorom(bank, addr),
            RomType::HiRom => self.resolve_hirom(bank, addr),
        }
    }
}

impl Rom {
    pub fn load(&mut self, bank: u8, addr: u16) -> u8 {
        *self.resolve_addr(bank, addr)
    }

    pub fn store(&mut self, bank: u8, addr: u16, value: u8) {
        if addr >= 0x8000 {
            warn!("writing ${:02X} to ROM address ${:02X}:{:04X}", value, bank, addr);
        }
        *self.resolve_addr(bank, addr) = value;
    }
}

fn out_of_ram_bounds(bank: u8, addr: u16, abs: u32) -> ! {
    panic!("RAM access out of bounds at {:02X}:{:04X} -> {:04X}",
        bank, addr, abs)
}

fn out_of_rom_bounds(bank: u8, addr: u16, abs: u32) -> ! {
    panic!("ROM access out of bounds at {:02X}:{:04X} -> {:06X}",
        bank, addr, abs)
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod function_wrapper_cpp;
pub(crate) mod type_to_cpp;

use crate::{types::QualifiedName, CppFilePair};
use itertools::Itertools;
use std::collections::HashSet;
use syn::Type;
use type_to_cpp::type_to_cpp;

use super::{
    analysis::fun::{
        function_wrapper::{FunctionWrapper, FunctionWrapperPayload},
        FnAnalysis,
    },
    api::Api,
    ConvertError,
};

/// Instructions for new C++ which we need to generate.
#[derive(Clone)]
pub(crate) enum AdditionalNeed {
    MakeStringConstructor,
    FunctionWrapper(Box<FunctionWrapper>),
    CTypeTypedef(QualifiedName),
    ConcreteTemplatedTypeTypedef(QualifiedName, Box<Type>),
}

#[derive(Ord, PartialOrd, Eq, PartialEq, Clone, Hash)]
struct Header {
    name: &'static str,
    system: bool,
}

impl Header {
    fn system(name: &'static str) -> Self {
        Header { name, system: true }
    }

    fn user(name: &'static str) -> Self {
        Header {
            name,
            system: false,
        }
    }

    fn include_stmt(&self) -> String {
        if self.system {
            format!("#include <{}>", self.name)
        } else {
            format!("#include \"{}\"", self.name)
        }
    }
}

struct AdditionalFunction {
    type_definition: String, // are output before main declarations
    declaration: String,
    definition: String,
    headers: Vec<Header>,
}

/// Generates additional C++ glue functions needed by autocxx.
/// In some ways it would be preferable to be able to pass snippets
/// of C++ through to `cxx` for inclusion in the C++ file which it
/// generates, and perhaps we'll explore that in future. But for now,
/// autocxx generates its own _additional_ C++ files which therefore
/// need to be built and included in linking procedures.
pub(crate) struct CppCodeGenerator {
    additional_functions: Vec<AdditionalFunction>,
    inclusions: String,
}

impl CppCodeGenerator {
    pub(crate) fn generate_cpp_code(
        inclusions: String,
        apis: &[Api<FnAnalysis>],
    ) -> Result<Option<CppFilePair>, ConvertError> {
        let mut gen = CppCodeGenerator::new(inclusions);
        gen.add_needs(apis.iter().filter_map(|api| api.additional_cpp()))?;
        Ok(gen.generate())
    }

    fn new(inclusions: String) -> Self {
        CppCodeGenerator {
            additional_functions: Vec::new(),
            inclusions,
        }
    }

    fn add_needs(
        &mut self,
        additions: impl Iterator<Item = AdditionalNeed>,
    ) -> Result<(), ConvertError> {
        for need in additions {
            match need {
                AdditionalNeed::MakeStringConstructor => self.generate_string_constructor(),
                AdditionalNeed::FunctionWrapper(by_value_wrapper) => {
                    self.generate_by_value_wrapper(&by_value_wrapper)?
                }
                AdditionalNeed::CTypeTypedef(tn) => self.generate_ctype_typedef(&tn),
                AdditionalNeed::ConcreteTemplatedTypeTypedef(tn, def) => {
                    self.generate_typedef(&tn, type_to_cpp(&def)?)
                }
            }
        }
        Ok(())
    }

    fn generate(&self) -> Option<CppFilePair> {
        if self.additional_functions.is_empty() {
            None
        } else {
            let headers: HashSet<Header> = self
                .additional_functions
                .iter()
                .map(|x| x.headers.iter().cloned())
                .flatten()
                .collect();
            let headers = headers.iter().map(|x| x.include_stmt()).join("\n");
            let type_definitions = self.concat_additional_items(|x| &x.type_definition);
            let declarations = self.concat_additional_items(|x| &x.declaration);
            let declarations = format!(
                "#pragma once\n\n{}\n{}\n{}\n{}",
                headers, self.inclusions, type_definitions, declarations
            );
            let definitions = self.concat_additional_items(|x| &x.definition);
            let definitions = format!("#include \"autocxxgen.h\"\n{}", definitions);
            log::info!("Additional C++ decls:\n{}", declarations);
            log::info!("Additional C++ defs:\n{}", definitions);
            Some(CppFilePair {
                header: declarations.into_bytes(),
                implementation: definitions.into_bytes(),
                header_name: "autocxxgen.h".into(),
            })
        }
    }

    fn concat_additional_items<F>(&self, field_access: F) -> String
    where
        F: FnMut(&AdditionalFunction) -> &str,
    {
        let mut s = self
            .additional_functions
            .iter()
            .map(field_access)
            .collect::<Vec<&str>>()
            .join("\n");
        s.push('\n');
        s
    }

    fn generate_string_constructor(&mut self) {
        let declaration = "std::unique_ptr<std::string> make_string(::rust::Str str)";
        let definition = format!(
            "{} {{ return std::make_unique<std::string>(std::string(str)); }}",
            declaration
        );
        let declaration = format!("{};", declaration);
        self.additional_functions.push(AdditionalFunction {
            type_definition: "".into(),
            declaration,
            definition,
            headers: vec![
                Header::system("memory"),
                Header::system("string"),
                Header::user("cxx.h"),
            ],
        })
    }

    fn generate_by_value_wrapper(&mut self, details: &FunctionWrapper) -> Result<(), ConvertError> {
        // Even if the original function call is in a namespace,
        // we generate this wrapper in the global namespace.
        // We could easily do this the other way round, and when
        // cxx::bridge comes to support nested namespace mods then
        // we wil wish to do that to avoid name conflicts. However,
        // at the moment this is simpler because it avoids us having
        // to generate namespace blocks in the generated C++.
        let is_a_method = details.is_a_method;
        let name = &details.wrapper_function_name;
        let get_arg_name = |counter: usize| -> String {
            if is_a_method && counter == 0 {
                // For method calls that we generate, the first
                // argument name needs to be such that we recognize
                // it as a method in the second invocation of
                // bridge_converter after it's flowed again through
                // bindgen.
                // TODO this may not be the case any longer. We
                // may be able to remove this.
                "autocxx_gen_this".to_string()
            } else {
                format!("arg{}", counter)
            }
        };
        let args: Result<Vec<_>, _> = details
            .argument_conversion
            .iter()
            .enumerate()
            .map(|(counter, ty)| {
                Ok(format!(
                    "{} {}",
                    ty.unconverted_type()?,
                    get_arg_name(counter)
                ))
            })
            .collect();
        let args = args?.join(", ");
        let ret_type = details
            .return_conversion
            .as_ref()
            .map_or(Ok("void".to_string()), |x| x.converted_type())?;
        let declaration = format!("{} {}({})", ret_type, name, args);
        let arg_list: Result<Vec<_>, _> = details
            .argument_conversion
            .iter()
            .enumerate()
            .map(|(counter, conv)| conv.cpp_conversion(&get_arg_name(counter)))
            .collect();
        let mut arg_list = arg_list?.into_iter();
        let receiver = if is_a_method { arg_list.next() } else { None };
        let arg_list = arg_list.join(", ");
        let mut underlying_function_call = match &details.payload {
            FunctionWrapperPayload::Constructor => arg_list,
            FunctionWrapperPayload::FunctionCall(ns, id) => match receiver {
                Some(receiver) => format!("{}.{}({})", receiver, id.to_string(), arg_list),
                None => {
                    let underlying_function_call = ns
                        .into_iter()
                        .cloned()
                        .chain(std::iter::once(id.to_string()))
                        .join("::");
                    format!("{}({})", underlying_function_call, arg_list)
                }
            },
            FunctionWrapperPayload::StaticMethodCall(ns, ty_id, fn_id) => {
                let underlying_function_call = ns
                    .into_iter()
                    .cloned()
                    .chain([ty_id.to_string(), fn_id.to_string()].iter().cloned())
                    .join("::");
                format!("{}({})", underlying_function_call, arg_list)
            }
        };
        if let Some(ret) = &details.return_conversion {
            underlying_function_call =
                format!("return {}", ret.cpp_conversion(&underlying_function_call)?);
        };
        let definition = format!("{} {{ {}; }}", declaration, underlying_function_call,);
        let declaration = format!("{};", declaration);
        self.additional_functions.push(AdditionalFunction {
            type_definition: "".into(),
            declaration,
            definition,
            headers: vec![Header::system("memory")],
        });
        Ok(())
    }

    fn generate_ctype_typedef(&mut self, tn: &QualifiedName) {
        let cpp_name = tn.to_cpp_name();
        self.generate_typedef(tn, cpp_name)
    }

    fn generate_typedef(&mut self, tn: &QualifiedName, definition: String) {
        let our_name = tn.get_final_item();
        self.additional_functions.push(AdditionalFunction {
            type_definition: format!("typedef {} {};", definition, our_name),
            declaration: "".into(),
            definition: "".into(),
            headers: Vec::new(),
        })
    }
}
<commit_msg>Doing a longstanding TODO for include guards.<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod function_wrapper_cpp;
pub(crate) mod type_to_cpp;

use crate::{types::QualifiedName, CppFilePair};
use itertools::Itertools;
use std::collections::HashSet;
use syn::Type;
use type_to_cpp::type_to_cpp;

use super::{
    analysis::fun::{
        function_wrapper::{FunctionWrapper, FunctionWrapperPayload},
        FnAnalysis,
    },
    api::Api,
    ConvertError,
};

/// Instructions for new C++ which we need to generate.
#[derive(Clone)]
pub(crate) enum AdditionalNeed {
    MakeStringConstructor,
    FunctionWrapper(Box<FunctionWrapper>),
    CTypeTypedef(QualifiedName),
    ConcreteTemplatedTypeTypedef(QualifiedName, Box<Type>),
}

#[derive(Ord, PartialOrd, Eq, PartialEq, Clone, Hash)]
struct Header {
    name: &'static str,
    system: bool,
}

impl Header {
    fn system(name: &'static str) -> Self {
        Header { name, system: true }
    }

    fn user(name: &'static str) -> Self {
        Header {
            name,
            system: false,
        }
    }

    fn include_stmt(&self) -> String {
        if self.system {
            format!("#include <{}>", self.name)
        } else {
            format!("#include \"{}\"", self.name)
        }
    }
}

struct AdditionalFunction {
    type_definition: String, // are output before main declarations
    declaration: String,
    definition: String,
    headers: Vec<Header>,
}

/// Generates additional C++ glue functions needed by autocxx.
/// In some ways it would be preferable to be able to pass snippets
/// of C++ through to `cxx` for inclusion in the C++ file which it
/// generates, and perhaps we'll explore that in future. But for now,
/// autocxx generates its own _additional_ C++ files which therefore
/// need to be built and included in linking procedures.
pub(crate) struct CppCodeGenerator {
    additional_functions: Vec<AdditionalFunction>,
    inclusions: String,
}

impl CppCodeGenerator {
    pub(crate) fn generate_cpp_code(
        inclusions: String,
        apis: &[Api<FnAnalysis>],
    ) -> Result<Option<CppFilePair>, ConvertError> {
        let mut gen = CppCodeGenerator::new(inclusions);
        gen.add_needs(apis.iter().filter_map(|api| api.additional_cpp()))?;
        Ok(gen.generate())
    }

    fn new(inclusions: String) -> Self {
        CppCodeGenerator {
            additional_functions: Vec::new(),
            inclusions,
        }
    }

    fn add_needs(
        &mut self,
        additions: impl Iterator<Item = AdditionalNeed>,
    ) -> Result<(), ConvertError> {
        for need in additions {
            match need {
                AdditionalNeed::MakeStringConstructor => self.generate_string_constructor(),
                AdditionalNeed::FunctionWrapper(by_value_wrapper) => {
                    self.generate_by_value_wrapper(&by_value_wrapper)?
                }
                AdditionalNeed::CTypeTypedef(tn) => self.generate_ctype_typedef(&tn),
                AdditionalNeed::ConcreteTemplatedTypeTypedef(tn, def) => {
                    self.generate_typedef(&tn, type_to_cpp(&def)?)
                }
            }
        }
        Ok(())
    }

    fn generate(&self) -> Option<CppFilePair> {
        if self.additional_functions.is_empty() {
            None
        } else {
            let headers: HashSet<Header> = self
                .additional_functions
                .iter()
                .map(|x| x.headers.iter().cloned())
                .flatten()
                .collect();
            let headers = headers.iter().map(|x| x.include_stmt()).join("\n");
            let type_definitions = self.concat_additional_items(|x| &x.type_definition);
            let declarations = self.concat_additional_items(|x| &x.declaration);
            let declarations = format!(
                "#ifndef __AUTOCXXGEN_H__\n#define __AUTOCXXGEN_H__\n\n{}\n{}\n{}\n{}#endif // __AUTOCXXGEN_H__\n",
                headers, self.inclusions, type_definitions, declarations
            );
            let definitions = self.concat_additional_items(|x| &x.definition);
            let definitions = format!("#include \"autocxxgen.h\"\n{}", definitions);
            log::info!("Additional C++ decls:\n{}", declarations);
            log::info!("Additional C++ defs:\n{}", definitions);
            Some(CppFilePair {
                header: declarations.into_bytes(),
                implementation: definitions.into_bytes(),
                header_name: "autocxxgen.h".into(),
            })
        }
    }

    fn concat_additional_items<F>(&self, field_access: F) -> String
    where
        F: FnMut(&AdditionalFunction) -> &str,
    {
        let mut s = self
            .additional_functions
            .iter()
            .map(field_access)
            .collect::<Vec<&str>>()
            .join("\n");
        s.push('\n');
        s
    }

    fn generate_string_constructor(&mut self) {
        let declaration = "std::unique_ptr<std::string> make_string(::rust::Str str)";
        let definition = format!(
            "{} {{ return std::make_unique<std::string>(std::string(str)); }}",
            declaration
        );
        let declaration = format!("{};", declaration);
        self.additional_functions.push(AdditionalFunction {
            type_definition: "".into(),
            declaration,
            definition,
            headers: vec![
                Header::system("memory"),
                Header::system("string"),
                Header::user("cxx.h"),
            ],
        })
    }

    fn generate_by_value_wrapper(&mut self, details: &FunctionWrapper) -> Result<(), ConvertError> {
        // Even if the original function call is in a namespace,
        // we generate this wrapper in the global namespace.
        // We could easily do this the other way round, and when
        // cxx::bridge comes to support nested namespace mods then
        // we wil wish to do that to avoid name conflicts. However,
        // at the moment this is simpler because it avoids us having
        // to generate namespace blocks in the generated C++.
        let is_a_method = details.is_a_method;
        let name = &details.wrapper_function_name;
        let get_arg_name = |counter: usize| -> String {
            if is_a_method && counter == 0 {
                // For method calls that we generate, the first
                // argument name needs to be such that we recognize
                // it as a method in the second invocation of
                // bridge_converter after it's flowed again through
                // bindgen.
                // TODO this may not be the case any longer. We
                // may be able to remove this.
                "autocxx_gen_this".to_string()
            } else {
                format!("arg{}", counter)
            }
        };
        let args: Result<Vec<_>, _> = details
            .argument_conversion
            .iter()
            .enumerate()
            .map(|(counter, ty)| {
                Ok(format!(
                    "{} {}",
                    ty.unconverted_type()?,
                    get_arg_name(counter)
                ))
            })
            .collect();
        let args = args?.join(", ");
        let ret_type = details
            .return_conversion
            .as_ref()
            .map_or(Ok("void".to_string()), |x| x.converted_type())?;
        let declaration = format!("{} {}({})", ret_type, name, args);
        let arg_list: Result<Vec<_>, _> = details
            .argument_conversion
            .iter()
            .enumerate()
            .map(|(counter, conv)| conv.cpp_conversion(&get_arg_name(counter)))
            .collect();
        let mut arg_list = arg_list?.into_iter();
        let receiver = if is_a_method { arg_list.next() } else { None };
        let arg_list = arg_list.join(", ");
        let mut underlying_function_call = match &details.payload {
            FunctionWrapperPayload::Constructor => arg_list,
            FunctionWrapperPayload::FunctionCall(ns, id) => match receiver {
                Some(receiver) => format!("{}.{}({})", receiver, id.to_string(), arg_list),
                None => {
                    let underlying_function_call = ns
                        .into_iter()
                        .cloned()
                        .chain(std::iter::once(id.to_string()))
                        .join("::");
                    format!("{}({})", underlying_function_call, arg_list)
                }
            },
            FunctionWrapperPayload::StaticMethodCall(ns, ty_id, fn_id) => {
                let underlying_function_call = ns
                    .into_iter()
                    .cloned()
                    .chain([ty_id.to_string(), fn_id.to_string()].iter().cloned())
                    .join("::");
                format!("{}({})", underlying_function_call, arg_list)
            }
        };
        if let Some(ret) = &details.return_conversion {
            underlying_function_call =
                format!("return {}", ret.cpp_conversion(&underlying_function_call)?);
        };
        let definition = format!("{} {{ {}; }}", declaration, underlying_function_call,);
        let declaration = format!("{};", declaration);
        self.additional_functions.push(AdditionalFunction {
            type_definition: "".into(),
            declaration,
            definition,
            headers: vec![Header::system("memory")],
        });
        Ok(())
    }

    fn generate_ctype_typedef(&mut self, tn: &QualifiedName) {
        let cpp_name = tn.to_cpp_name();
        self.generate_typedef(tn, cpp_name)
    }

    fn generate_typedef(&mut self, tn: &QualifiedName, definition: String) {
        let our_name = tn.get_final_item();
        self.additional_functions.push(AdditionalFunction {
            type_definition: format!("typedef {} {};", definition, our_name),
            declaration: "".into(),
            definition: "".into(),
            headers: Vec::new(),
        })
    }
}
<|endoftext|>


<commit_before>use std::ops::BitOr;
use super::{chtype, COLOR_PAIR};
use attributes::{Attribute, Attributes};

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct ColorPair(pub u8);

impl From<ColorPair> for chtype {
    fn from(color_pair: ColorPair) -> chtype {
        COLOR_PAIR(chtype::from(color_pair.0))
    }
}

/// Implement the | operator for setting a color pair on an Attributes object
///
/// # Example
///
/// ```
/// use pancurses::{Attribute, Attributes};
/// use pancurses::colorpair::ColorPair;
///
/// let mut attributes = Attributes::new();
/// assert!(attributes.color_pair().0 == 0);
/// attributes = attributes | ColorPair(1);
/// assert!(attributes.color_pair().0 == 1);
/// ```
impl BitOr<ColorPair> for Attributes {
    type Output = Attributes;

    fn bitor(mut self, rhs: ColorPair) -> Attributes {
        self.set_color_pair(rhs);
        self
    }
}

/// Implement the | operator for combining a `ColorPair` and an `Attribute` to produce `Attributes`
///
/// # Example
///
/// ```
/// use pancurses::Attribute;
/// use pancurses::colorpair::ColorPair;
///
/// let attributes = ColorPair(5) | Attribute::Blink;
/// assert!(attributes.color_pair().0 == 5);
/// assert!(!attributes.is_bold());
/// assert!(attributes.is_blink());
/// ```
impl BitOr<Attribute> for ColorPair {
    type Output = Attributes;

    fn bitor(self, rhs: Attribute) -> Attributes {
        Attributes::new() | self | rhs
    }
}

/// Implement the | operator for combining an Attribute and a  ColorPair to produce Attributes
///
/// # Example
///
/// ```
/// use pancurses::Attribute;
/// use pancurses::colorpair::ColorPair;
///
/// let attributes = Attribute::Blink | ColorPair(2);
/// assert!(attributes.color_pair().0 == 2);
/// assert!(!attributes.is_bold());
/// assert!(attributes.is_blink());
/// ```
impl BitOr<ColorPair> for Attribute {
    type Output = Attributes;

    fn bitor(self, rhs: ColorPair) -> Attributes {
        Attributes::new() | self | rhs
    }
}
<commit_msg>More ticks in docs<commit_after>use std::ops::BitOr;
use super::{chtype, COLOR_PAIR};
use attributes::{Attribute, Attributes};

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct ColorPair(pub u8);

impl From<ColorPair> for chtype {
    fn from(color_pair: ColorPair) -> chtype {
        COLOR_PAIR(chtype::from(color_pair.0))
    }
}

/// Implement the | operator for setting a color pair on an `Attributes` object
///
/// # Example
///
/// ```
/// use pancurses::{Attribute, Attributes};
/// use pancurses::colorpair::ColorPair;
///
/// let mut attributes = Attributes::new();
/// assert!(attributes.color_pair().0 == 0);
/// attributes = attributes | ColorPair(1);
/// assert!(attributes.color_pair().0 == 1);
/// ```
impl BitOr<ColorPair> for Attributes {
    type Output = Attributes;

    fn bitor(mut self, rhs: ColorPair) -> Attributes {
        self.set_color_pair(rhs);
        self
    }
}

/// Implement the | operator for combining a `ColorPair` and an `Attribute` to produce `Attributes`
///
/// # Example
///
/// ```
/// use pancurses::Attribute;
/// use pancurses::colorpair::ColorPair;
///
/// let attributes = ColorPair(5) | Attribute::Blink;
/// assert!(attributes.color_pair().0 == 5);
/// assert!(!attributes.is_bold());
/// assert!(attributes.is_blink());
/// ```
impl BitOr<Attribute> for ColorPair {
    type Output = Attributes;

    fn bitor(self, rhs: Attribute) -> Attributes {
        Attributes::new() | self | rhs
    }
}

/// Implement the | operator for combining an `Attribute` and a `ColorPair` to produce `Attributes`
///
/// # Example
///
/// ```
/// use pancurses::Attribute;
/// use pancurses::colorpair::ColorPair;
///
/// let attributes = Attribute::Blink | ColorPair(2);
/// assert!(attributes.color_pair().0 == 2);
/// assert!(!attributes.is_bold());
/// assert!(attributes.is_blink());
/// ```
impl BitOr<ColorPair> for Attribute {
    type Output = Attributes;

    fn bitor(self, rhs: ColorPair) -> Attributes {
        Attributes::new() | self | rhs
    }
}
<|endoftext|>


<commit_before>use anyhow::*;
use diesel;
use diesel::prelude::*;
#[cfg(feature = "profile-index")]
use flame;
use log::{error, info};
use rayon::prelude::*;
use regex::Regex;
use std::fs;
use std::path::Path;
use std::time;
use std::sync::mpsc::*;

use crate::config::MiscSettings;
use crate::db::{directories, misc_settings, songs, DB};
use crate::index::metadata;
use crate::vfs::VFSSource;

const INDEX_BUILDING_INSERT_BUFFER_SIZE: usize = 1000; // Insertions in each transaction
const INDEX_BUILDING_CLEAN_BUFFER_SIZE: usize = 500; // Insertions in each transaction

pub fn update(db: &DB) -> Result<()> {
	let start = time::Instant::now();
	info!("Beginning library index update");
	clean(db)?;
	populate(db)?;
	info!(
		"Library index update took {} seconds",
		start.elapsed().as_millis() as f32 / 1000.0
	);
	#[cfg(feature = "profile-index")]
	flame::dump_html(&mut fs::File::create("index-flame-graph.html").unwrap()).unwrap();
	Ok(())
}

#[derive(Debug, Insertable)]
#[table_name = "songs"]
struct NewSong {
	path: String,
	parent: String,
	track_number: Option<i32>,
	disc_number: Option<i32>,
	title: Option<String>,
	artist: Option<String>,
	album_artist: Option<String>,
	year: Option<i32>,
	album: Option<String>,
	artwork: Option<String>,
	duration: Option<i32>,
}

#[derive(Debug, Insertable)]
#[table_name = "directories"]
struct NewDirectory {
	path: String,
	parent: Option<String>,
	artist: Option<String>,
	year: Option<i32>,
	album: Option<String>,
	artwork: Option<String>,
	date_added: i32,
}

struct IndexUpdater {
	directory_sender: Sender<NewDirectory>,
	song_sender: Sender<NewSong>,
	album_art_pattern: Regex,
}

impl IndexUpdater {
	#[cfg_attr(feature = "profile-index", flame)]
	fn new(album_art_pattern: Regex, directory_sender: Sender<NewDirectory>, song_sender: Sender<NewSong>) -> Result<IndexUpdater> {
		Ok(IndexUpdater {
			directory_sender,
			song_sender,
			album_art_pattern,
		})
	}

	#[cfg_attr(feature = "profile-index", flame)]
	fn push_song(&mut self, song: NewSong) -> Result<()> {
		self.song_sender.send(song).map_err(Error::new)
	}

	#[cfg_attr(feature = "profile-index", flame)]
	fn push_directory(&mut self, directory: NewDirectory) -> Result<()> {
		self.directory_sender.send(directory).map_err(Error::new)
	}

	fn get_artwork(&self, dir: &Path) -> Result<Option<String>> {
		for file in fs::read_dir(dir)? {
			let file = file?;
			if let Some(name_string) = file.file_name().to_str() {
				if self.album_art_pattern.is_match(name_string) {
					return Ok(file.path().to_str().map(|p| p.to_owned()));
				}
			}
		}
		Ok(None)
	}

	#[cfg_attr(feature = "profile-index", flame)]
	fn populate_directory(&mut self, parent: Option<&Path>, path: &Path) -> Result<()> {
		// Find artwork
		let artwork = {
			#[cfg(feature = "profile-index")]
			let _guard = flame::start_guard("artwork");
			self.get_artwork(path).unwrap_or(None)
		};

		// Extract path and parent path
		let parent_string = parent.and_then(|p| p.to_str()).map(|s| s.to_owned());
		let path_string = path.to_str().ok_or(anyhow!("Invalid directory path"))?;

		// Find date added
		let metadata = {
			#[cfg(feature = "profile-index")]
			let _guard = flame::start_guard("metadata");
			fs::metadata(path_string)?
		};
		let created = {
			#[cfg(feature = "profile-index")]
			let _guard = flame::start_guard("created_date");
			metadata
			.created()
			.or_else(|_| metadata.modified())?
			.duration_since(time::UNIX_EPOCH)?
			.as_secs() as i32
		};

		let mut directory_album = None;
		let mut directory_year = None;
		let mut directory_artist = None;
		let mut inconsistent_directory_album = false;
		let mut inconsistent_directory_year = false;
		let mut inconsistent_directory_artist = false;

		// Sub directories
		let mut sub_directories = Vec::new();

		// Insert content
		for file in fs::read_dir(path)? {
			#[cfg(feature = "profile-index")]
			let _guard = flame::start_guard("directory-entry");
			let file_path = match file {
				Ok(ref f) => f.path(),
				_ => {
					error!("File read error within {}", path_string);
					break;
				}
			};

			if file_path.is_dir() {
				sub_directories.push(file_path.to_path_buf());
				continue;
			}

			if let Some(file_path_string) = file_path.to_str() {
				if let Some(tags) = metadata::read(file_path.as_path()) {

					#[cfg(feature = "profile-index")]
					let _guard = flame::start_guard("process_song");

					if tags.year.is_some() {
						inconsistent_directory_year |=
							directory_year.is_some() && directory_year != tags.year;
						directory_year = tags.year;
					}

					if tags.album.is_some() {
						inconsistent_directory_album |=
							directory_album.is_some() && directory_album != tags.album;
						directory_album = tags.album.as_ref().cloned();
					}

					if tags.album_artist.is_some() {
						inconsistent_directory_artist |=
							directory_artist.is_some() && directory_artist != tags.album_artist;
						directory_artist = tags.album_artist.as_ref().cloned();
					} else if tags.artist.is_some() {
						inconsistent_directory_artist |=
							directory_artist.is_some() && directory_artist != tags.artist;
						directory_artist = tags.artist.as_ref().cloned();
					}

					let song = NewSong {
						path: file_path_string.to_owned(),
						parent: path_string.to_owned(),
						disc_number: tags.disc_number.map(|n| n as i32),
						track_number: tags.track_number.map(|n| n as i32),
						title: tags.title,
						duration: tags.duration.map(|n| n as i32),
						artist: tags.artist,
						album_artist: tags.album_artist,
						album: tags.album,
						year: tags.year,
						artwork: artwork.as_ref().cloned(),
					};

					self.push_song(song)?;
				}
			}
		}

		// Insert directory
		let directory = {
			#[cfg(feature = "profile-index")]
			let _guard = flame::start_guard("create_directory");

			if inconsistent_directory_year {
				directory_year = None;
			}
			if inconsistent_directory_album {
				directory_album = None;
			}
			if inconsistent_directory_artist {
				directory_artist = None;
			}

			NewDirectory {
				path: path_string.to_owned(),
				parent: parent_string,
				artwork,
				album: directory_album,
				artist: directory_artist,
				year: directory_year,
				date_added: created,
			}
		};

		self.push_directory(directory)?;

		// Populate subdirectories
		for sub_directory in sub_directories {
			self.populate_directory(Some(path), &sub_directory)?;
		}

		Ok(())
	}
}

#[cfg_attr(feature = "profile-index", flame)]
pub fn clean(db: &DB) -> Result<()> {
	let vfs = db.get_vfs()?;

	{
		let all_songs: Vec<String>;
		{
			let connection = db.connect()?;
			all_songs = songs::table.select(songs::path).load(&connection)?;
		}

		let missing_songs = all_songs
			.par_iter()
			.filter(|ref song_path| {
				let path = Path::new(&song_path);
				!path.exists() || vfs.real_to_virtual(path).is_err()
			})
			.collect::<Vec<_>>();

		{
			let connection = db.connect()?;
			for chunk in missing_songs[..].chunks(INDEX_BUILDING_CLEAN_BUFFER_SIZE) {
				diesel::delete(songs::table.filter(songs::path.eq_any(chunk)))
					.execute(&connection)?;
			}
		}
	}

	{
		let all_directories: Vec<String>;
		{
			let connection = db.connect()?;
			all_directories = directories::table
				.select(directories::path)
				.load(&connection)?;
		}

		let missing_directories = all_directories
			.par_iter()
			.filter(|ref directory_path| {
				let path = Path::new(&directory_path);
				!path.exists() || vfs.real_to_virtual(path).is_err()
			})
			.collect::<Vec<_>>();

		{
			let connection = db.connect()?;
			for chunk in missing_directories[..].chunks(INDEX_BUILDING_CLEAN_BUFFER_SIZE) {
				diesel::delete(directories::table.filter(directories::path.eq_any(chunk)))
					.execute(&connection)?;
			}
		}
	}

	Ok(())
}

#[cfg_attr(feature = "profile-index", flame)]
pub fn populate(db: &DB) -> Result<()> {
	let vfs = db.get_vfs()?;
	let mount_points = vfs.get_mount_points();

	let album_art_pattern  = {
		let connection = db.connect()?;
		let settings: MiscSettings = misc_settings::table.get_result(&connection)?;
		Regex::new(&settings.index_album_art_pattern)?
	};

	let (directory_sender, directory_receiver) = channel();
	let (song_sender, song_receiver) = channel();

	let songs_db = db.clone();
	let directories_db = db.clone();

	let directories_thread = std::thread::spawn(move || {
		insert_directories(directory_receiver, directories_db);
	});

	let songs_thread = std::thread::spawn(move || {
		insert_songs(song_receiver, songs_db);
	});

	{
		let mut updater = IndexUpdater::new(album_art_pattern, directory_sender, song_sender)?;
		for target in mount_points.values() {
			updater.populate_directory(None, target.as_path())?;
		}
	}

	match directories_thread.join() {
		Err(e) => error!("Error while waiting for directory insertions to complete: {:?}", e),
		_ => (),
	}

	match songs_thread.join() {
		Err(e) => error!("Error while waiting for song insertions to complete: {:?}", e),
		_ => (),
	}

	Ok(())
}

fn insert_songs(receiver: Receiver<NewSong>, db: DB) {
	let mut new_entries = Vec::new();
	new_entries.reserve_exact(INDEX_BUILDING_INSERT_BUFFER_SIZE);

	loop {
		match receiver.recv() {
			Ok(s) => {
				new_entries.push(s);
				if new_entries.len() >= INDEX_BUILDING_INSERT_BUFFER_SIZE {
					if db.connect()
					.and_then(|connection|{
						diesel::insert_into(songs::table)
							.values(&new_entries)
							.execute(&*connection) // TODO https://github.com/diesel-rs/diesel/issues/1822
							.map_err(Error::new)
					})
					.is_err() {
						error!("Could not insert new songs in database");
					}
					new_entries.clear();
				}
			},
			Err(_) => break,
		}
	}

	if new_entries.len() > 0 {
		if db.connect()
		.and_then(|connection|{
			diesel::insert_into(songs::table)
				.values(&new_entries)
				.execute(&*connection) // TODO https://github.com/diesel-rs/diesel/issues/1822
				.map_err(Error::new)
		})
		.is_err() {
			error!("Could not insert new songs in database");
		}
	}
}

fn insert_directories(receiver: Receiver<NewDirectory>, db: DB) {
	let mut new_entries = Vec::new();
	new_entries.reserve_exact(INDEX_BUILDING_INSERT_BUFFER_SIZE);

	loop {
		match receiver.recv() {
			Ok(s) => {
				new_entries.push(s);
				if new_entries.len() >= INDEX_BUILDING_INSERT_BUFFER_SIZE {
					if db.connect()
					.and_then(|connection|{
						diesel::insert_into(directories::table)
							.values(&new_entries)
							.execute(&*connection) // TODO https://github.com/diesel-rs/diesel/issues/1822
							.map_err(Error::new)
					})
					.is_err() {
						error!("Could not insert new directories in database");
					}
					new_entries.clear();
				}
			},
			Err(_) => break,
		}
	}

	if new_entries.len() > 0 {
		if db.connect()
		.and_then(|connection|{
			diesel::insert_into(directories::table)
				.values(&new_entries)
				.execute(&*connection) // TODO https://github.com/diesel-rs/diesel/issues/1822
				.map_err(Error::new)
		})
		.is_err() {
			error!("Could not insert new directories in database");
		}
	}
}
<commit_msg>Re-factored duplicated code<commit_after>use anyhow::*;
use diesel;
use diesel::prelude::*;
#[cfg(feature = "profile-index")]
use flame;
use log::{error, info};
use rayon::prelude::*;
use regex::Regex;
use std::fs;
use std::path::Path;
use std::time;
use std::sync::mpsc::*;

use crate::config::MiscSettings;
use crate::db::{directories, misc_settings, songs, DB};
use crate::index::metadata;
use crate::vfs::VFSSource;

const INDEX_BUILDING_INSERT_BUFFER_SIZE: usize = 1000; // Insertions in each transaction
const INDEX_BUILDING_CLEAN_BUFFER_SIZE: usize = 500; // Insertions in each transaction

pub fn update(db: &DB) -> Result<()> {
	let start = time::Instant::now();
	info!("Beginning library index update");
	clean(db)?;
	populate(db)?;
	info!(
		"Library index update took {} seconds",
		start.elapsed().as_millis() as f32 / 1000.0
	);
	#[cfg(feature = "profile-index")]
	flame::dump_html(&mut fs::File::create("index-flame-graph.html").unwrap()).unwrap();
	Ok(())
}

#[derive(Debug, Insertable)]
#[table_name = "songs"]
struct NewSong {
	path: String,
	parent: String,
	track_number: Option<i32>,
	disc_number: Option<i32>,
	title: Option<String>,
	artist: Option<String>,
	album_artist: Option<String>,
	year: Option<i32>,
	album: Option<String>,
	artwork: Option<String>,
	duration: Option<i32>,
}

#[derive(Debug, Insertable)]
#[table_name = "directories"]
struct NewDirectory {
	path: String,
	parent: Option<String>,
	artist: Option<String>,
	year: Option<i32>,
	album: Option<String>,
	artwork: Option<String>,
	date_added: i32,
}

struct IndexUpdater {
	directory_sender: Sender<NewDirectory>,
	song_sender: Sender<NewSong>,
	album_art_pattern: Regex,
}

impl IndexUpdater {
	#[cfg_attr(feature = "profile-index", flame)]
	fn new(album_art_pattern: Regex, directory_sender: Sender<NewDirectory>, song_sender: Sender<NewSong>) -> Result<IndexUpdater> {
		Ok(IndexUpdater {
			directory_sender,
			song_sender,
			album_art_pattern,
		})
	}

	#[cfg_attr(feature = "profile-index", flame)]
	fn push_song(&mut self, song: NewSong) -> Result<()> {
		self.song_sender.send(song).map_err(Error::new)
	}

	#[cfg_attr(feature = "profile-index", flame)]
	fn push_directory(&mut self, directory: NewDirectory) -> Result<()> {
		self.directory_sender.send(directory).map_err(Error::new)
	}

	fn get_artwork(&self, dir: &Path) -> Result<Option<String>> {
		for file in fs::read_dir(dir)? {
			let file = file?;
			if let Some(name_string) = file.file_name().to_str() {
				if self.album_art_pattern.is_match(name_string) {
					return Ok(file.path().to_str().map(|p| p.to_owned()));
				}
			}
		}
		Ok(None)
	}

	#[cfg_attr(feature = "profile-index", flame)]
	fn populate_directory(&mut self, parent: Option<&Path>, path: &Path) -> Result<()> {
		// Find artwork
		let artwork = {
			#[cfg(feature = "profile-index")]
			let _guard = flame::start_guard("artwork");
			self.get_artwork(path).unwrap_or(None)
		};

		// Extract path and parent path
		let parent_string = parent.and_then(|p| p.to_str()).map(|s| s.to_owned());
		let path_string = path.to_str().ok_or(anyhow!("Invalid directory path"))?;

		// Find date added
		let metadata = {
			#[cfg(feature = "profile-index")]
			let _guard = flame::start_guard("metadata");
			fs::metadata(path_string)?
		};
		let created = {
			#[cfg(feature = "profile-index")]
			let _guard = flame::start_guard("created_date");
			metadata
			.created()
			.or_else(|_| metadata.modified())?
			.duration_since(time::UNIX_EPOCH)?
			.as_secs() as i32
		};

		let mut directory_album = None;
		let mut directory_year = None;
		let mut directory_artist = None;
		let mut inconsistent_directory_album = false;
		let mut inconsistent_directory_year = false;
		let mut inconsistent_directory_artist = false;

		// Sub directories
		let mut sub_directories = Vec::new();

		// Insert content
		for file in fs::read_dir(path)? {
			#[cfg(feature = "profile-index")]
			let _guard = flame::start_guard("directory-entry");
			let file_path = match file {
				Ok(ref f) => f.path(),
				_ => {
					error!("File read error within {}", path_string);
					break;
				}
			};

			if file_path.is_dir() {
				sub_directories.push(file_path.to_path_buf());
				continue;
			}

			if let Some(file_path_string) = file_path.to_str() {
				if let Some(tags) = metadata::read(file_path.as_path()) {

					#[cfg(feature = "profile-index")]
					let _guard = flame::start_guard("process_song");

					if tags.year.is_some() {
						inconsistent_directory_year |=
							directory_year.is_some() && directory_year != tags.year;
						directory_year = tags.year;
					}

					if tags.album.is_some() {
						inconsistent_directory_album |=
							directory_album.is_some() && directory_album != tags.album;
						directory_album = tags.album.as_ref().cloned();
					}

					if tags.album_artist.is_some() {
						inconsistent_directory_artist |=
							directory_artist.is_some() && directory_artist != tags.album_artist;
						directory_artist = tags.album_artist.as_ref().cloned();
					} else if tags.artist.is_some() {
						inconsistent_directory_artist |=
							directory_artist.is_some() && directory_artist != tags.artist;
						directory_artist = tags.artist.as_ref().cloned();
					}

					let song = NewSong {
						path: file_path_string.to_owned(),
						parent: path_string.to_owned(),
						disc_number: tags.disc_number.map(|n| n as i32),
						track_number: tags.track_number.map(|n| n as i32),
						title: tags.title,
						duration: tags.duration.map(|n| n as i32),
						artist: tags.artist,
						album_artist: tags.album_artist,
						album: tags.album,
						year: tags.year,
						artwork: artwork.as_ref().cloned(),
					};

					self.push_song(song)?;
				}
			}
		}

		// Insert directory
		let directory = {
			#[cfg(feature = "profile-index")]
			let _guard = flame::start_guard("create_directory");

			if inconsistent_directory_year {
				directory_year = None;
			}
			if inconsistent_directory_album {
				directory_album = None;
			}
			if inconsistent_directory_artist {
				directory_artist = None;
			}

			NewDirectory {
				path: path_string.to_owned(),
				parent: parent_string,
				artwork,
				album: directory_album,
				artist: directory_artist,
				year: directory_year,
				date_added: created,
			}
		};

		self.push_directory(directory)?;

		// Populate subdirectories
		for sub_directory in sub_directories {
			self.populate_directory(Some(path), &sub_directory)?;
		}

		Ok(())
	}
}

#[cfg_attr(feature = "profile-index", flame)]
pub fn clean(db: &DB) -> Result<()> {
	let vfs = db.get_vfs()?;

	{
		let all_songs: Vec<String>;
		{
			let connection = db.connect()?;
			all_songs = songs::table.select(songs::path).load(&connection)?;
		}

		let missing_songs = all_songs
			.par_iter()
			.filter(|ref song_path| {
				let path = Path::new(&song_path);
				!path.exists() || vfs.real_to_virtual(path).is_err()
			})
			.collect::<Vec<_>>();

		{
			let connection = db.connect()?;
			for chunk in missing_songs[..].chunks(INDEX_BUILDING_CLEAN_BUFFER_SIZE) {
				diesel::delete(songs::table.filter(songs::path.eq_any(chunk)))
					.execute(&connection)?;
			}
		}
	}

	{
		let all_directories: Vec<String>;
		{
			let connection = db.connect()?;
			all_directories = directories::table
				.select(directories::path)
				.load(&connection)?;
		}

		let missing_directories = all_directories
			.par_iter()
			.filter(|ref directory_path| {
				let path = Path::new(&directory_path);
				!path.exists() || vfs.real_to_virtual(path).is_err()
			})
			.collect::<Vec<_>>();

		{
			let connection = db.connect()?;
			for chunk in missing_directories[..].chunks(INDEX_BUILDING_CLEAN_BUFFER_SIZE) {
				diesel::delete(directories::table.filter(directories::path.eq_any(chunk)))
					.execute(&connection)?;
			}
		}
	}

	Ok(())
}

#[cfg_attr(feature = "profile-index", flame)]
pub fn populate(db: &DB) -> Result<()> {
	let vfs = db.get_vfs()?;
	let mount_points = vfs.get_mount_points();

	let album_art_pattern  = {
		let connection = db.connect()?;
		let settings: MiscSettings = misc_settings::table.get_result(&connection)?;
		Regex::new(&settings.index_album_art_pattern)?
	};

	let (directory_sender, directory_receiver) = channel();
	let (song_sender, song_receiver) = channel();

	let songs_db = db.clone();
	let directories_db = db.clone();

	let directories_thread = std::thread::spawn(move || {
		insert_directories(directory_receiver, directories_db);
	});

	let songs_thread = std::thread::spawn(move || {
		insert_songs(song_receiver, songs_db);
	});

	{
		let mut updater = IndexUpdater::new(album_art_pattern, directory_sender, song_sender)?;
		for target in mount_points.values() {
			updater.populate_directory(None, target.as_path())?;
		}
	}

	match directories_thread.join() {
		Err(e) => error!("Error while waiting for directory insertions to complete: {:?}", e),
		_ => (),
	}

	match songs_thread.join() {
		Err(e) => error!("Error while waiting for song insertions to complete: {:?}", e),
		_ => (),
	}

	Ok(())
}

fn flush_directories(db: &DB, entries: &Vec<NewDirectory>) {
	if db.connect()
	.and_then(|connection|{
		diesel::insert_into(directories::table)
			.values(entries)
			.execute(&*connection) // TODO https://github.com/diesel-rs/diesel/issues/1822
			.map_err(Error::new)
	})
	.is_err() {
		error!("Could not insert new directories in database");
	}
}

fn flush_songs(db: &DB, entries: &Vec<NewSong>) {
	if db.connect()
	.and_then(|connection|{
		diesel::insert_into(songs::table)
			.values(entries)
			.execute(&*connection) // TODO https://github.com/diesel-rs/diesel/issues/1822
			.map_err(Error::new)
	})
	.is_err() {
		error!("Could not insert new songs in database");
	}
}

fn insert_directories(receiver: Receiver<NewDirectory>, db: DB) {
	let mut new_entries = Vec::new();
	new_entries.reserve_exact(INDEX_BUILDING_INSERT_BUFFER_SIZE);

	loop {
		match receiver.recv() {
			Ok(s) => {
				new_entries.push(s);
				if new_entries.len() >= INDEX_BUILDING_INSERT_BUFFER_SIZE {
					flush_directories(&db, &new_entries);
					new_entries.clear();
				}
			},
			Err(_) => break,
		}
	}

	if new_entries.len() > 0 {
		flush_directories(&db, &new_entries);
	}
}

fn insert_songs(receiver: Receiver<NewSong>, db: DB) {
	let mut new_entries = Vec::new();
	new_entries.reserve_exact(INDEX_BUILDING_INSERT_BUFFER_SIZE);

	loop {
		match receiver.recv() {
			Ok(s) => {
				new_entries.push(s);
				if new_entries.len() >= INDEX_BUILDING_INSERT_BUFFER_SIZE {
					flush_songs(&db, &new_entries);
					new_entries.clear();
				}
			},
			Err(_) => break,
		}
	}

	if new_entries.len() > 0 {
		flush_songs(&db, &new_entries);
	}
}
<|endoftext|>