<commit_before>#![cfg_attr(not(feature="std"), allow(unused_imports))]
use core;
use super::encode::{BrotliEncoderCreateInstance, BrotliEncoderDestroyInstance,
                    BrotliEncoderParameter, BrotliEncoderSetParameter, BrotliEncoderOperation,
                    BrotliEncoderStateStruct, BrotliEncoderCompressStream, BrotliEncoderIsFinished};
use super::backward_references::BrotliEncoderParams;
use brotli_decompressor::CustomWrite;
use super::interface;
#[cfg(feature="std")]
pub use brotli_decompressor::{IntoIoWriter, IoWriterWrapper};
use super::combined_alloc::BrotliAlloc;
pub use alloc::{AllocatedStackMemory, Allocator, SliceWrapper, SliceWrapperMut, StackAllocator};
#[cfg(feature="std")]
pub use alloc_stdlib::StandardAlloc;
#[cfg(feature="std")]
use std::io;

#[cfg(feature="std")]
use std::io::{Write, Error, ErrorKind};

#[cfg(feature="std")]
pub struct CompressorWriterCustomAlloc<W: Write,
                                       BufferType : SliceWrapperMut<u8>,
                                       Alloc: BrotliAlloc> (
    CompressorWriterCustomIo<io::Error,
                             IntoIoWriter<W>,
                             BufferType,
                             Alloc>);

#[cfg(feature="std")]
impl<W: Write,
     BufferType : SliceWrapperMut<u8>,
     Alloc: BrotliAlloc>
    CompressorWriterCustomAlloc<W, BufferType, Alloc>
    {

    pub fn new(w: W, buffer : BufferType,
               alloc: Alloc,
               q: u32,
               lgwin: u32) -> Self {
        CompressorWriterCustomAlloc::<W, BufferType, Alloc>(
          CompressorWriterCustomIo::<Error,
                                 IntoIoWriter<W>,
                                 BufferType,
                                 Alloc>::new(
              IntoIoWriter::<W>(w),
              buffer,
              alloc,
              Error::new(ErrorKind::InvalidData,
                         "Invalid Data"),
              q, lgwin))
    }

    pub fn get_ref(&self) -> &W {
      &self.0.get_ref().0
    }
    pub fn into_inner(self) -> W {
      self.0.into_inner().0
    }
}

#[cfg(feature="std")]
impl<W: Write,
     BufferType: SliceWrapperMut<u8>,
     Alloc: BrotliAlloc>
    Write for CompressorWriterCustomAlloc<W, BufferType,
                                         Alloc> {
  	fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
       self.0.write(buf)
    }
    fn flush(&mut self) -> Result<(), Error> {
       self.0.flush()
    }
}

#[cfg(feature="std")]
pub struct CompressorWriter<W: Write>(CompressorWriterCustomAlloc<W,
                                     <StandardAlloc
                                      as Allocator<u8>>::AllocatedMemory,
                                     StandardAlloc>);

#[cfg(feature="std")]
impl<W: Write> CompressorWriter<W> {
  pub fn new(w: W, buffer_size: usize, q: u32, lgwin: u32) -> Self {
    let mut alloc = StandardAlloc::default();
    let buffer = <StandardAlloc as Allocator<u8>>::alloc_cell(&mut alloc, if buffer_size == 0 { 4096} else {buffer_size});
    CompressorWriter::<W>(CompressorWriterCustomAlloc::new(w,
                                                           buffer,
                                                           alloc,
                                                           q,
                                                           lgwin))
  }

  pub fn with_params(w: W, buffer_size: usize, params: &BrotliEncoderParams) -> Self {
    let mut writer = Self::new(w, buffer_size, params.quality as u32, params.lgwin as u32);
    (writer.0).0.state.params = params.clone();
    writer
  }

  pub fn get_ref(&self) -> &W {
    self.0.get_ref()
  }
  pub fn into_inner(self) -> W {
    self.0.into_inner()
  }
}

#[cfg(feature="std")]
impl<W: Write> Write for CompressorWriter<W> {
  fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
    self.0.write(buf)
  }
  fn flush(&mut self) -> Result<(), Error> {
    self.0.flush()
  }
}

pub struct CompressorWriterCustomIo<ErrType,
                                    W: CustomWrite<ErrType>,
                                    BufferType: SliceWrapperMut<u8>,
                                    Alloc: BrotliAlloc>
{
  output_buffer: BufferType,
  total_out: Option<usize>,
  output: Option<W>,
  error_if_invalid_data: Option<ErrType>,
  state: BrotliEncoderStateStruct<Alloc>,
}
pub fn write_all<ErrType, W: CustomWrite<ErrType>>(writer: &mut W, mut buf : &[u8]) -> Result<(), ErrType> {
    while buf.len() != 0 {
          match writer.write(buf) {
                Ok(bytes_written) => buf = &buf[bytes_written..],
                Err(e) => return Err(e),
          }
    }
    Ok(())
}
impl<ErrType,
     W: CustomWrite<ErrType>,
     BufferType : SliceWrapperMut<u8>,
     Alloc: BrotliAlloc>
CompressorWriterCustomIo<ErrType, W, BufferType, Alloc>
{

    pub fn new(w: W, buffer : BufferType,
               alloc : Alloc,
               invalid_data_error_type : ErrType,
               q: u32,
               lgwin: u32) -> Self {
        let mut ret = CompressorWriterCustomIo{
            output_buffer : buffer,
            total_out : Some(0),
            output: Some(w),
            state : BrotliEncoderCreateInstance(alloc),
            error_if_invalid_data : Some(invalid_data_error_type),
        };
        BrotliEncoderSetParameter(&mut ret.state,
                                  BrotliEncoderParameter::BROTLI_PARAM_QUALITY,
                                  q as (u32));
        BrotliEncoderSetParameter(&mut ret.state,
                                  BrotliEncoderParameter::BROTLI_PARAM_LGWIN,
                                  lgwin as (u32));

        ret
    }
    fn flush_or_close(&mut self, op:BrotliEncoderOperation) -> Result<(), ErrType>{
       let mut nop_callback = |_data:&mut interface::PredictionModeContextMap<interface::InputReferenceMut>,
                               _cmds: &mut [interface::StaticCommand],
                               _mb: interface::InputPair, _mfv: &mut Alloc|();

        loop {
            let mut avail_in : usize = 0;
            let mut input_offset : usize = 0;
            let mut avail_out : usize = self.output_buffer.slice_mut().len();
            let mut output_offset : usize = 0;
            let ret = BrotliEncoderCompressStream(
                &mut self.state,
                op,
                &mut avail_in,
                &[],
                &mut input_offset,
                &mut avail_out,
                self.output_buffer.slice_mut(),
                &mut output_offset,
                &mut self.total_out,
                &mut nop_callback);
           if output_offset > 0 {
             match write_all(self.output.as_mut().unwrap(), &self.output_buffer.slice_mut()[..output_offset]) {
               Ok(_) => {},
               Err(e) => return Err(e),
             }
           }
           if ret <= 0 {
              return Err(self.error_if_invalid_data.take().unwrap());
           }
           if let BrotliEncoderOperation::BROTLI_OPERATION_FLUSH = op {
              return Ok(());
           }
           if BrotliEncoderIsFinished(&mut self.state) != 0 {
              return Ok(());
           }
        }
    }

    pub fn get_ref(&self) -> &W {
      self.output.as_ref().unwrap()
    }
    pub fn into_inner(mut self) -> W {
      match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FINISH) {
        Ok(_) => {},
        Err(_) => {},
      }
      core::mem::replace(&mut self.output, None).unwrap()
    }
}

impl<ErrType,
     W: CustomWrite<ErrType>,
     BufferType : SliceWrapperMut<u8>,
     Alloc: BrotliAlloc> Drop for
CompressorWriterCustomIo<ErrType, W, BufferType, Alloc> {
    fn drop(&mut self) {
        if self.output.is_some() {
            match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FINISH) {
                Ok(_) => {},
              Err(_) => {},
            }
        }
        BrotliEncoderDestroyInstance(&mut self.state);
    }
}
impl<ErrType,
     W: CustomWrite<ErrType>,
     BufferType : SliceWrapperMut<u8>,
     Alloc: BrotliAlloc> CustomWrite<ErrType> for
CompressorWriterCustomIo<ErrType, W, BufferType, Alloc> {
	fn write(&mut self, buf: & [u8]) -> Result<usize, ErrType > {
        let mut nop_callback = |_data:&mut interface::PredictionModeContextMap<interface::InputReferenceMut>,
                                _cmds: &mut [interface::StaticCommand],
                                _mb: interface::InputPair, _mfv: &mut Alloc|();
        let mut avail_in = buf.len();
        let mut input_offset : usize = 0;
        while avail_in != 0 {
            let mut output_offset = 0;
            let mut avail_out = self.output_buffer.slice_mut().len();
            let ret = BrotliEncoderCompressStream(
                &mut self.state,
                BrotliEncoderOperation::BROTLI_OPERATION_PROCESS,
                &mut avail_in,
                &buf[..],
                &mut input_offset,
                &mut avail_out,
                self.output_buffer.slice_mut(),
                &mut output_offset,
                &mut self.total_out,
                &mut nop_callback);
           if output_offset > 0 {
             match write_all(self.output.as_mut().unwrap(), &self.output_buffer.slice_mut()[..output_offset]) {
              Ok(_) => {},
              Err(e) => return Err(e),
             }
           }
           if ret <= 0 {
              return Err(self.error_if_invalid_data.take().unwrap());
           }
        }
        Ok(buf.len())
      }
      fn flush(&mut self) -> Result<(), ErrType > {
        match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FLUSH) {
              Ok(_) => {},
              Err(e) => return Err(e),
        }
        self.output.as_mut().unwrap().flush()
      }
}
<commit_msg>add get_mut for `CompressorWriter`<commit_after>#![cfg_attr(not(feature="std"), allow(unused_imports))]
use core;
use super::encode::{BrotliEncoderCreateInstance, BrotliEncoderDestroyInstance,
                    BrotliEncoderParameter, BrotliEncoderSetParameter, BrotliEncoderOperation,
                    BrotliEncoderStateStruct, BrotliEncoderCompressStream, BrotliEncoderIsFinished};
use super::backward_references::BrotliEncoderParams;
use brotli_decompressor::CustomWrite;
use super::interface;
#[cfg(feature="std")]
pub use brotli_decompressor::{IntoIoWriter, IoWriterWrapper};
use super::combined_alloc::BrotliAlloc;
pub use alloc::{AllocatedStackMemory, Allocator, SliceWrapper, SliceWrapperMut, StackAllocator};
#[cfg(feature="std")]
pub use alloc_stdlib::StandardAlloc;
#[cfg(feature="std")]
use std::io;

#[cfg(feature="std")]
use std::io::{Write, Error, ErrorKind};

#[cfg(feature="std")]
pub struct CompressorWriterCustomAlloc<W: Write,
                                       BufferType : SliceWrapperMut<u8>,
                                       Alloc: BrotliAlloc> (
    CompressorWriterCustomIo<io::Error,
                             IntoIoWriter<W>,
                             BufferType,
                             Alloc>);

#[cfg(feature="std")]
impl<W: Write,
     BufferType : SliceWrapperMut<u8>,
     Alloc: BrotliAlloc>
    CompressorWriterCustomAlloc<W, BufferType, Alloc>
    {

    pub fn new(w: W, buffer : BufferType,
               alloc: Alloc,
               q: u32,
               lgwin: u32) -> Self {
        CompressorWriterCustomAlloc::<W, BufferType, Alloc>(
          CompressorWriterCustomIo::<Error,
                                 IntoIoWriter<W>,
                                 BufferType,
                                 Alloc>::new(
              IntoIoWriter::<W>(w),
              buffer,
              alloc,
              Error::new(ErrorKind::InvalidData,
                         "Invalid Data"),
              q, lgwin))
    }

    pub fn get_ref(&self) -> &W {
      &self.0.get_ref().0
    }
    pub fn get_mut(&mut self) -> &mut W {
      &mut self.0.get_mut().0
    }
    pub fn into_inner(self) -> W {
      self.0.into_inner().0
    }
}

#[cfg(feature="std")]
impl<W: Write,
     BufferType: SliceWrapperMut<u8>,
     Alloc: BrotliAlloc>
    Write for CompressorWriterCustomAlloc<W, BufferType,
                                         Alloc> {
  	fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
       self.0.write(buf)
    }
    fn flush(&mut self) -> Result<(), Error> {
       self.0.flush()
    }
}

#[cfg(feature="std")]
pub struct CompressorWriter<W: Write>(CompressorWriterCustomAlloc<W,
                                     <StandardAlloc
                                      as Allocator<u8>>::AllocatedMemory,
                                     StandardAlloc>);

#[cfg(feature="std")]
impl<W: Write> CompressorWriter<W> {
  pub fn new(w: W, buffer_size: usize, q: u32, lgwin: u32) -> Self {
    let mut alloc = StandardAlloc::default();
    let buffer = <StandardAlloc as Allocator<u8>>::alloc_cell(&mut alloc, if buffer_size == 0 { 4096} else {buffer_size});
    CompressorWriter::<W>(CompressorWriterCustomAlloc::new(w,
                                                           buffer,
                                                           alloc,
                                                           q,
                                                           lgwin))
  }

  pub fn with_params(w: W, buffer_size: usize, params: &BrotliEncoderParams) -> Self {
    let mut writer = Self::new(w, buffer_size, params.quality as u32, params.lgwin as u32);
    (writer.0).0.state.params = params.clone();
    writer
  }

  pub fn get_ref(&self) -> &W {
    self.0.get_ref()
  }
  pub fn get_mut(&mut self) -> &mut W {
    self.0.get_mut()
  }
  pub fn into_inner(self) -> W {
    self.0.into_inner()
  }
}

#[cfg(feature="std")]
impl<W: Write> Write for CompressorWriter<W> {
  fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
    self.0.write(buf)
  }
  fn flush(&mut self) -> Result<(), Error> {
    self.0.flush()
  }
}

pub struct CompressorWriterCustomIo<ErrType,
                                    W: CustomWrite<ErrType>,
                                    BufferType: SliceWrapperMut<u8>,
                                    Alloc: BrotliAlloc>
{
  output_buffer: BufferType,
  total_out: Option<usize>,
  output: Option<W>,
  error_if_invalid_data: Option<ErrType>,
  state: BrotliEncoderStateStruct<Alloc>,
}
pub fn write_all<ErrType, W: CustomWrite<ErrType>>(writer: &mut W, mut buf : &[u8]) -> Result<(), ErrType> {
    while buf.len() != 0 {
          match writer.write(buf) {
                Ok(bytes_written) => buf = &buf[bytes_written..],
                Err(e) => return Err(e),
          }
    }
    Ok(())
}
impl<ErrType,
     W: CustomWrite<ErrType>,
     BufferType : SliceWrapperMut<u8>,
     Alloc: BrotliAlloc>
CompressorWriterCustomIo<ErrType, W, BufferType, Alloc>
{

    pub fn new(w: W, buffer : BufferType,
               alloc : Alloc,
               invalid_data_error_type : ErrType,
               q: u32,
               lgwin: u32) -> Self {
        let mut ret = CompressorWriterCustomIo{
            output_buffer : buffer,
            total_out : Some(0),
            output: Some(w),
            state : BrotliEncoderCreateInstance(alloc),
            error_if_invalid_data : Some(invalid_data_error_type),
        };
        BrotliEncoderSetParameter(&mut ret.state,
                                  BrotliEncoderParameter::BROTLI_PARAM_QUALITY,
                                  q as (u32));
        BrotliEncoderSetParameter(&mut ret.state,
                                  BrotliEncoderParameter::BROTLI_PARAM_LGWIN,
                                  lgwin as (u32));

        ret
    }
    fn flush_or_close(&mut self, op:BrotliEncoderOperation) -> Result<(), ErrType>{
       let mut nop_callback = |_data:&mut interface::PredictionModeContextMap<interface::InputReferenceMut>,
                               _cmds: &mut [interface::StaticCommand],
                               _mb: interface::InputPair, _mfv: &mut Alloc|();

        loop {
            let mut avail_in : usize = 0;
            let mut input_offset : usize = 0;
            let mut avail_out : usize = self.output_buffer.slice_mut().len();
            let mut output_offset : usize = 0;
            let ret = BrotliEncoderCompressStream(
                &mut self.state,
                op,
                &mut avail_in,
                &[],
                &mut input_offset,
                &mut avail_out,
                self.output_buffer.slice_mut(),
                &mut output_offset,
                &mut self.total_out,
                &mut nop_callback);
           if output_offset > 0 {
             match write_all(self.output.as_mut().unwrap(), &self.output_buffer.slice_mut()[..output_offset]) {
               Ok(_) => {},
               Err(e) => return Err(e),
             }
           }
           if ret <= 0 {
              return Err(self.error_if_invalid_data.take().unwrap());
           }
           if let BrotliEncoderOperation::BROTLI_OPERATION_FLUSH = op {
              return Ok(());
           }
           if BrotliEncoderIsFinished(&mut self.state) != 0 {
              return Ok(());
           }
        }
    }

    pub fn get_ref(&self) -> &W {
      self.output.as_ref().unwrap()
    }
    pub fn get_mut(&mut self) -> &mut W {
      self.output.as_mut().unwrap()
    }
    pub fn into_inner(mut self) -> W {
      match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FINISH) {
        Ok(_) => {},
        Err(_) => {},
      }
      core::mem::replace(&mut self.output, None).unwrap()
    }
}

impl<ErrType,
     W: CustomWrite<ErrType>,
     BufferType : SliceWrapperMut<u8>,
     Alloc: BrotliAlloc> Drop for
CompressorWriterCustomIo<ErrType, W, BufferType, Alloc> {
    fn drop(&mut self) {
        if self.output.is_some() {
            match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FINISH) {
                Ok(_) => {},
              Err(_) => {},
            }
        }
        BrotliEncoderDestroyInstance(&mut self.state);
    }
}
impl<ErrType,
     W: CustomWrite<ErrType>,
     BufferType : SliceWrapperMut<u8>,
     Alloc: BrotliAlloc> CustomWrite<ErrType> for
CompressorWriterCustomIo<ErrType, W, BufferType, Alloc> {
	fn write(&mut self, buf: & [u8]) -> Result<usize, ErrType > {
        let mut nop_callback = |_data:&mut interface::PredictionModeContextMap<interface::InputReferenceMut>,
                                _cmds: &mut [interface::StaticCommand],
                                _mb: interface::InputPair, _mfv: &mut Alloc|();
        let mut avail_in = buf.len();
        let mut input_offset : usize = 0;
        while avail_in != 0 {
            let mut output_offset = 0;
            let mut avail_out = self.output_buffer.slice_mut().len();
            let ret = BrotliEncoderCompressStream(
                &mut self.state,
                BrotliEncoderOperation::BROTLI_OPERATION_PROCESS,
                &mut avail_in,
                &buf[..],
                &mut input_offset,
                &mut avail_out,
                self.output_buffer.slice_mut(),
                &mut output_offset,
                &mut self.total_out,
                &mut nop_callback);
           if output_offset > 0 {
             match write_all(self.output.as_mut().unwrap(), &self.output_buffer.slice_mut()[..output_offset]) {
              Ok(_) => {},
              Err(e) => return Err(e),
             }
           }
           if ret <= 0 {
              return Err(self.error_if_invalid_data.take().unwrap());
           }
        }
        Ok(buf.len())
      }
      fn flush(&mut self) -> Result<(), ErrType > {
        match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FLUSH) {
              Ok(_) => {},
              Err(e) => return Err(e),
        }
        self.output.as_mut().unwrap().flush()
      }
}
<|endoftext|>


<commit_before>use std::ascii::AsciiExt;
use std::sync::Arc;
use rand::{Rng, XorShiftRng, SeedableRng};
use quickcheck;
use quickcheck::TestResult;
use player::Player;
use zobrist::Zobrist;
use field;
use field::{Pos, Field};

fn construct_field(image: &str) -> Field {
  let lines = image.split('\n').map(|line| line.trim_matches(' ')).filter(|line| !line.is_empty()).collect::<Vec<&str>>();
  let height = lines.len() as u32;
  assert!(height > 0);
  let width = lines.first().unwrap().len() as u32;
  assert!(lines.iter().all(|line| line.len() as u32 == width));
  let mut moves = lines.into_iter().enumerate().flat_map(|(y, line)|
    line.chars().enumerate().filter(|&(_, c)| c.to_ascii_lowercase() != c.to_ascii_uppercase()).map(move |(x, c)| (c, x as u32, y as u32))
  ).collect::<Vec<(char, u32, u32)>>();
  moves.sort_by(|&(c1, _, _), &(c2, _, _)| (c1.to_ascii_lowercase(), c1.is_lowercase()).cmp(&(c2.to_ascii_lowercase(), c2.is_lowercase())));
  let mut rng = XorShiftRng::new_unseeded();
  let zobrist = Arc::new(Zobrist::new(field::length(width, height) * 2, &mut rng));
  let mut field = Field::new(width, height, zobrist);
  for (c, x, y) in moves.into_iter() {
    let player = Player::from_bool(c.is_uppercase());
    let pos = field.to_pos(x, y);
    field.put_point(pos, player);
  }
  field
}

#[test]
fn simple_surround() {
  let field = construct_field(
    "
    .a.
    cBa
    .a.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(!field.cell(field.to_pos(1, 1)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(0, 1)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(1, 0)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(1, 2)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(2, 1)).is_putting_allowed());
}

#[test]
fn surround_empty_territory() {
  let field = construct_field(
    "
    .a.
    a.a
    .a.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 0);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(field.cell(field.to_pos(1, 1)).is_putting_allowed());
  assert!(field.cell(field.to_pos(1, 1)).is_empty_base());
  assert_eq!(field.cell(field.to_pos(1, 1)).get_empty_base_player(), Some(Player::Red));
  assert!(!field.cell(field.to_pos(0, 1)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(1, 0)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(1, 2)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(2, 1)).is_putting_allowed());
}

#[test]
fn move_priority() {
  let field = construct_field(
    "
    .aB.
    aCaB
    .aB.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 0);
  assert_eq!(field.captured_count(Player::Black), 1);
}

#[test]
fn move_priority_big() {
  let field = construct_field(
    "
    .B..
    BaB.
    aCaB
    .aB.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 0);
  assert_eq!(field.captured_count(Player::Black), 2);
}

#[test]
fn onion_surroundings() {
  let field = construct_field(
    "
    ...c...
    ..cBc..
    .cBaBc.
    ..cBc..
    ...c...
    "
  );
  assert_eq!(field.captured_count(Player::Red), 4);
  assert_eq!(field.captured_count(Player::Black), 0);
}

#[test]
fn apply_control_surrounding_in_same_turn() {
  let field = construct_field(
    "
    .a.
    aBa
    .a.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
}

#[test]
fn double_surround() {
  let field = construct_field(
    "
    .b.b..
    bAzAb.
    .b.b..
    "
  );
  assert_eq!(field.captured_count(Player::Red), 2);
  assert_eq!(field.captured_count(Player::Black), 0);
}

#[test]
fn double_surround_with_empty_part() {
  let field = construct_field(
    "
    .b.b..
    b.zAb.
    .b.b..
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(field.cell(field.to_pos(1, 1)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(3, 1)).is_putting_allowed());
}

#[test]
fn should_not_leave_empty_inside() {
  let field = construct_field(
    "
    .aaaa..
    a....a.
    a.b...a
    .z.bC.a
    a.b...a
    a....a.
    .aaaa..
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(!field.cell(field.to_pos(2, 3)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(2, 4)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(2, 2)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(1, 3)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(3, 3)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(1, 1)).is_putting_allowed());
}

#[test]
fn a_hole_inside_a_surrounding() {
  let field = construct_field(
    "
    ....c....
    ...c.c...
    ..c...c..
    .c..a..c.
    c..a.a..c
    .c..a..c.
    ..c...c..
    ...cBc...
    ....d....
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(!field.cell(field.to_pos(4, 4)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(4, 1)).is_putting_allowed());
}

#[test]
fn a_hole_inside_a_surrounding_after_control_surrounding() {
  let field = construct_field(
    "
    ....b....
    ...b.b...
    ..b...b..
    .b..a..b.
    b..a.a..b
    .b..a..b.
    ..b...b..
    ...bCb...
    ....b....
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(!field.cell(field.to_pos(4, 4)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(4, 1)).is_putting_allowed());
}

#[test]
fn surrounding_does_not_expand() {
  let field = construct_field(
    "
    ....a....
    ...a.a...
    ..a.a.a..
    .a.a.a.a.
    a.a.aBa.a
    .a.a.a.a.
    ..a.a.a..
    ...a.a...
    ....a....
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(field.cell(field.to_pos(6, 3)).is_putting_allowed());
  assert!(field.cell(field.to_pos(4, 3)).is_putting_allowed());
  assert!(field.cell(field.to_pos(4, 5)).is_putting_allowed());
  assert!(field.cell(field.to_pos(6, 5)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(5, 4)).is_putting_allowed());
}

#[test]
fn two_surroundings_with_common_border() {
  let field = construct_field(
    "
    .a..
    aAa.
    .bAa
    ..a.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 2);
  assert_eq!(field.captured_count(Player::Black), 0);
}

#[test]
fn two_surroundings_with_common_dot() {
  let field = construct_field(
    "
    .a.a.
    aAbAa
    .a.a.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 2);
  assert_eq!(field.captured_count(Player::Black), 0);
}

#[test]
fn three_surroundings_with_common_borders() {
  let field = construct_field(
    "
    ..a..
    .aAa.
    ..bAa
    .aAa.
    ..a..
    "
  );
  assert_eq!(field.captured_count(Player::Red), 3);
  assert_eq!(field.captured_count(Player::Black), 0);
}

#[test]
fn undo_check() {
  fn prop(width_seed: u32, height_seed: u32, seed: u64) -> TestResult {
    let width = width_seed % 30;
    let height = height_seed % 30;
    if width < 3 || height < 3 {
      return TestResult::discard();
    }
    let seed_array = [3, seed as u32, 7, (seed >> 32) as u32];
    let mut rng = XorShiftRng::from_seed(seed_array);
    let mut moves = (field::min_pos(width) .. field::max_pos(width, height) + 1).collect::<Vec<Pos>>();
    rng.shuffle(&mut moves);
    let zobrist = Arc::new(Zobrist::new(field::length(width, height) * 2, &mut rng));
    let mut field = Field::new(width, height, zobrist);
    let mut player = Player::Red;
    for pos in moves {
      if field.is_putting_allowed(pos) {
        player = player.next();
        let field_before = field.clone();
        field.put_point(pos, player);
        field.undo();
        if field_before != field {
          return TestResult::failed();
        }
        field.put_point(pos, player);
      }
    }
    TestResult::passed()
  }
  quickcheck::quickcheck(prop as fn(u32, u32, u64) -> TestResult);
}
<commit_msg>Remove duplicated test.<commit_after>use std::ascii::AsciiExt;
use std::sync::Arc;
use rand::{Rng, XorShiftRng, SeedableRng};
use quickcheck;
use quickcheck::TestResult;
use player::Player;
use zobrist::Zobrist;
use field;
use field::{Pos, Field};

fn construct_field(image: &str) -> Field {
  let lines = image.split('\n').map(|line| line.trim_matches(' ')).filter(|line| !line.is_empty()).collect::<Vec<&str>>();
  let height = lines.len() as u32;
  assert!(height > 0);
  let width = lines.first().unwrap().len() as u32;
  assert!(lines.iter().all(|line| line.len() as u32 == width));
  let mut moves = lines.into_iter().enumerate().flat_map(|(y, line)|
    line.chars().enumerate().filter(|&(_, c)| c.to_ascii_lowercase() != c.to_ascii_uppercase()).map(move |(x, c)| (c, x as u32, y as u32))
  ).collect::<Vec<(char, u32, u32)>>();
  moves.sort_by(|&(c1, _, _), &(c2, _, _)| (c1.to_ascii_lowercase(), c1.is_lowercase()).cmp(&(c2.to_ascii_lowercase(), c2.is_lowercase())));
  let mut rng = XorShiftRng::new_unseeded();
  let zobrist = Arc::new(Zobrist::new(field::length(width, height) * 2, &mut rng));
  let mut field = Field::new(width, height, zobrist);
  for (c, x, y) in moves.into_iter() {
    let player = Player::from_bool(c.is_uppercase());
    let pos = field.to_pos(x, y);
    field.put_point(pos, player);
  }
  field
}

#[test]
fn simple_surround() {
  let field = construct_field(
    "
    .a.
    cBa
    .a.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(!field.cell(field.to_pos(1, 1)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(0, 1)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(1, 0)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(1, 2)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(2, 1)).is_putting_allowed());
}

#[test]
fn surround_empty_territory() {
  let field = construct_field(
    "
    .a.
    a.a
    .a.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 0);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(field.cell(field.to_pos(1, 1)).is_putting_allowed());
  assert!(field.cell(field.to_pos(1, 1)).is_empty_base());
  assert_eq!(field.cell(field.to_pos(1, 1)).get_empty_base_player(), Some(Player::Red));
  assert!(!field.cell(field.to_pos(0, 1)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(1, 0)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(1, 2)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(2, 1)).is_putting_allowed());
}

#[test]
fn move_priority() {
  let field = construct_field(
    "
    .aB.
    aCaB
    .aB.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 0);
  assert_eq!(field.captured_count(Player::Black), 1);
}

#[test]
fn move_priority_big() {
  let field = construct_field(
    "
    .B..
    BaB.
    aCaB
    .aB.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 0);
  assert_eq!(field.captured_count(Player::Black), 2);
}

#[test]
fn onion_surroundings() {
  let field = construct_field(
    "
    ...c...
    ..cBc..
    .cBaBc.
    ..cBc..
    ...c...
    "
  );
  assert_eq!(field.captured_count(Player::Red), 4);
  assert_eq!(field.captured_count(Player::Black), 0);
}

#[test]
fn apply_control_surrounding_in_same_turn() {
  let field = construct_field(
    "
    .a.
    aBa
    .a.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
}

#[test]
fn double_surround() {
  let field = construct_field(
    "
    .a.a.
    aAbAa
    .a.a.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 2);
  assert_eq!(field.captured_count(Player::Black), 0);
}

#[test]
fn double_surround_with_empty_part() {
  let field = construct_field(
    "
    .b.b..
    b.zAb.
    .b.b..
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(field.cell(field.to_pos(1, 1)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(3, 1)).is_putting_allowed());
}

#[test]
fn should_not_leave_empty_inside() {
  let field = construct_field(
    "
    .aaaa..
    a....a.
    a.b...a
    .z.bC.a
    a.b...a
    a....a.
    .aaaa..
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(!field.cell(field.to_pos(2, 3)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(2, 4)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(2, 2)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(1, 3)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(3, 3)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(1, 1)).is_putting_allowed());
}

#[test]
fn a_hole_inside_a_surrounding() {
  let field = construct_field(
    "
    ....c....
    ...c.c...
    ..c...c..
    .c..a..c.
    c..a.a..c
    .c..a..c.
    ..c...c..
    ...cBc...
    ....d....
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(!field.cell(field.to_pos(4, 4)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(4, 1)).is_putting_allowed());
}

#[test]
fn a_hole_inside_a_surrounding_after_control_surrounding() {
  let field = construct_field(
    "
    ....b....
    ...b.b...
    ..b...b..
    .b..a..b.
    b..a.a..b
    .b..a..b.
    ..b...b..
    ...bCb...
    ....b....
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(!field.cell(field.to_pos(4, 4)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(4, 1)).is_putting_allowed());
}

#[test]
fn surrounding_does_not_expand() {
  let field = construct_field(
    "
    ....a....
    ...a.a...
    ..a.a.a..
    .a.a.a.a.
    a.a.aBa.a
    .a.a.a.a.
    ..a.a.a..
    ...a.a...
    ....a....
    "
  );
  assert_eq!(field.captured_count(Player::Red), 1);
  assert_eq!(field.captured_count(Player::Black), 0);
  assert!(field.cell(field.to_pos(6, 3)).is_putting_allowed());
  assert!(field.cell(field.to_pos(4, 3)).is_putting_allowed());
  assert!(field.cell(field.to_pos(4, 5)).is_putting_allowed());
  assert!(field.cell(field.to_pos(6, 5)).is_putting_allowed());
  assert!(!field.cell(field.to_pos(5, 4)).is_putting_allowed());
}

#[test]
fn two_surroundings_with_common_border() {
  let field = construct_field(
    "
    .a..
    aAa.
    .bAa
    ..a.
    "
  );
  assert_eq!(field.captured_count(Player::Red), 2);
  assert_eq!(field.captured_count(Player::Black), 0);
}

#[test]
fn three_surroundings_with_common_borders() {
  let field = construct_field(
    "
    ..a..
    .aAa.
    ..bAa
    .aAa.
    ..a..
    "
  );
  assert_eq!(field.captured_count(Player::Red), 3);
  assert_eq!(field.captured_count(Player::Black), 0);
}

#[test]
fn undo_check() {
  fn prop(width_seed: u32, height_seed: u32, seed: u64) -> TestResult {
    let width = width_seed % 30;
    let height = height_seed % 30;
    if width < 3 || height < 3 {
      return TestResult::discard();
    }
    let seed_array = [3, seed as u32, 7, (seed >> 32) as u32];
    let mut rng = XorShiftRng::from_seed(seed_array);
    let mut moves = (field::min_pos(width) .. field::max_pos(width, height) + 1).collect::<Vec<Pos>>();
    rng.shuffle(&mut moves);
    let zobrist = Arc::new(Zobrist::new(field::length(width, height) * 2, &mut rng));
    let mut field = Field::new(width, height, zobrist);
    let mut player = Player::Red;
    for pos in moves {
      if field.is_putting_allowed(pos) {
        player = player.next();
        let field_before = field.clone();
        field.put_point(pos, player);
        field.undo();
        if field_before != field {
          return TestResult::failed();
        }
        field.put_point(pos, player);
      }
    }
    TestResult::passed()
  }
  quickcheck::quickcheck(prop as fn(u32, u32, u64) -> TestResult);
}
<|endoftext|>


<commit_before>pub mod search;

use std::thread;
use std::cmp::min;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver};
use std::time::SystemTime;
use std::num::Wrapping;
use basetypes::*;
use uci::{UciEngine, UciEngineFactory, EngineReply, OptionName, OptionDescription};
use position::Position;
use chess_move::*;
use tt::*;

const VERSION: &'static str = "0.1";
const MAX_DEPTH: u8 = 126;

/// Implements `UciEngine` trait.
pub struct Engine {
    tt: Arc<TranspositionTable>,
    reply_queue: VecDeque<EngineReply>,
    position: Position,
    is_thinking: bool,
    is_pondering: bool,

    // Tells the engine if it will be allowed to ponder. This option
    // is needed because the engine might change its time management
    // algorithm when pondering is allowed.
    pondering_is_allowed: bool,

    // A channel for sending commands to the search thread.
    commands: Sender<search::Command>,

    // A channel for receiving reports from the search thread.
    reports: Receiver<search::Report>,

    // Search status info
    search_id: usize,
    thinking_since: SystemTime,
    silent_since: SystemTime,
    current_depth: u8,
    searched_nodes: NodeCount,
    searched_time: u64, // milliseconds
    nps: u64, // nodes per second
    mangled_pv: bool, // `true` if the primary variation is imperfect
    stop_when: TimeManagement,
}

impl Engine {
    /// Creates a new instance.
    ///
    /// `tt_size_mb` is the preferred size of the transposition
    /// table in Mbytes.
    pub fn new(tt_size_mb: usize) -> Engine {
        let mut tt = TranspositionTable::new();
        tt.resize(tt_size_mb);
        let tt1 = Arc::new(tt);
        let tt2 = tt1.clone();
        let (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        thread::spawn(move || {
            search::run_deepening(tt2, commands_rx, reports_tx);
        });

        Engine {
            tt: tt1,
            reply_queue: VecDeque::new(),
            position: Position::from_fen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 \
                                          1")
                          .ok()
                          .unwrap(),
            is_thinking: false,
            is_pondering: false,
            pondering_is_allowed: false,
            commands: commands_tx,
            reports: reports_rx,
            search_id: 0,
            thinking_since: SystemTime::now(),
            silent_since: SystemTime::now(),
            current_depth: 0,
            searched_nodes: 0,
            searched_time: 0,
            nps: 0,
            mangled_pv: false,
            stop_when: TimeManagement::Infinite,
        }
    }
}

impl UciEngine for Engine {
    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            "Ponder" => {
                self.pondering_is_allowed = value == "true";
            }

            // An invalid option. Notice that we do not support
            // re-sizing of the transposition table once the engine
            // had started.
            _ => (),
        }
    }

    fn new_game(&mut self) {
        // Clearing the transposition table would not change anything,
        // so we do nothing.
    }

    fn position(&mut self, fen: &str, moves: &mut Iterator<Item = &str>) {
        if !self.is_thinking {
            if let Ok(p) = Position::from_history(fen, moves) {
                self.position = p;
                self.tt.new_search();
            }
        }
    }

    #[allow(unused_variables)]
    fn go(&mut self,
          searchmoves: Option<Vec<String>>,
          ponder: bool,
          wtime: Option<u64>,
          btime: Option<u64>,
          winc: Option<u64>,
          binc: Option<u64>,
          movestogo: Option<u64>,
          depth: Option<u64>,
          nodes: Option<u64>,
          mate: Option<u64>,
          movetime: Option<u64>,
          infinite: bool) {
        if !self.is_thinking {
            // Note: We ignore "searchmoves" and "mate" parameters.

            self.is_thinking = true;
            self.is_pondering = ponder;
            self.search_id = (Wrapping(self.search_id) + Wrapping(1)).0;
            self.thinking_since = SystemTime::now();
            self.silent_since = SystemTime::now();
            self.current_depth = 0;
            self.searched_nodes = 0;
            self.searched_time = 0;
            self.mangled_pv = false;

            // Figure out when we should stop thinking.
            //
            // TODO: Implement smarter time management.
            let (time, inc) = if self.position.board().to_move() == WHITE {
                (wtime, winc.unwrap_or(0))
            } else {
                (btime, binc.unwrap_or(0))
            };
            self.stop_when = if infinite {
                TimeManagement::Infinite
            } else if movetime.is_some() {
                TimeManagement::MoveTime(movetime.unwrap())
            } else if nodes.is_some() {
                TimeManagement::Nodes(nodes.unwrap())
            } else if depth.is_some() {
                TimeManagement::Depth(depth.unwrap() as u8)
            } else {
                let time = time.unwrap_or(0);
                let movestogo = movestogo.unwrap_or(40);
                let movetime = (time + inc * movestogo) / movestogo;
                TimeManagement::MoveTimeHint(min(movetime, time / 2))
            };

            // Start deepening search to the maximum possible depth.
            self.commands
                .send(search::Command::Search {
                    search_id: self.search_id,
                    position: self.position.clone(),
                    depth: MAX_DEPTH,
                    lower_bound: -20000,
                    upper_bound: 20000,
                })
                .unwrap();
        }
    }

    fn ponder_hit(&mut self) {
        if self.is_thinking && self.is_pondering {
            self.is_pondering = false;
        }
    }

    fn stop(&mut self) {
        if self.is_thinking {
            self.commands.send(search::Command::Stop).unwrap();
            let mut p = self.position.clone();
            let (best_move, ponder_move) = if let Some(m) = self.do_best_move(&mut p) {
                (m.notation(), self.do_best_move(&mut p).map(|m| m.notation()))
            } else {
                ("0000".to_string(), None)
            };
            self.reply_queue.push_back(EngineReply::BestMove {
                best_move: best_move,
                ponder_move: ponder_move,
            });
            self.is_thinking = false;
        }
    }

    fn is_thinking(&self) -> bool {
        self.is_thinking
    }

    fn get_reply(&mut self) -> Option<EngineReply> {
        // Check if we should stop thinking.
        if self.is_thinking && !self.is_pondering {
            if match self.stop_when {
                TimeManagement::MoveTime(t) => self.searched_time >= t,
                TimeManagement::MoveTimeHint(t) => self.searched_time >= t,
                TimeManagement::Nodes(n) => self.searched_nodes >= n,
                TimeManagement::Depth(d) => self.current_depth > d,
                TimeManagement::Infinite => false,
            } {
                self.stop();
            }
        }

        // Check the reports queue.
        while let Ok(report) = self.reports.try_recv() {
            if self.is_thinking {
                match report {
                    search::Report::Progress { search_id, depth, searched_nodes }
                        if search_id == self.search_id => {
                        // Register search progress.
                        self.register_progress(depth, searched_nodes);
                        if self.silent_since.elapsed().unwrap().as_secs() > 20 {
                            if self.mangled_pv {
                                self.report_pv(depth - 1);
                            } else {
                                self.report_progress();
                            }
                        }
                    }
                    search::Report::Done { search_id, .. } if search_id == self.search_id => {
                        // Terminate the search if not infinite.
                        self.stop_when = match self.stop_when {
                            TimeManagement::Infinite => TimeManagement::Infinite,
                            _ => TimeManagement::MoveTime(0),
                        };
                    }
                    // We may still receive stale reports from already
                    // stopped searches.
                    _ => (),
                }
            }
        }

        self.reply_queue.pop_front()
    }
}

impl Engine {
    // A helper method. It updates the search status info and makes
    // sure that a new PV is sent to the GUI for every newly reached
    // depth.
    fn register_progress(&mut self, depth: u8, searched_nodes: NodeCount) {
        let thinking_duration = self.thinking_since.elapsed().unwrap();
        self.searched_time = 1000 * thinking_duration.as_secs() +
                             (thinking_duration.subsec_nanos() / 1000000) as u64;
        self.searched_nodes = searched_nodes;
        self.nps = 1000 * (self.nps + self.searched_nodes) / (1000 + self.searched_time);
        if self.current_depth < depth {
            self.current_depth = depth;
            self.report_pv(depth - 1);
            if !self.mangled_pv {
                self.report_progress();
            }
        }
    }

    // A helper method. It extracts the primary variation (PV) from
    // the transposition table (TT) and sends it to the GUI.
    fn report_pv(&mut self, depth: u8) {
        if depth == 0 {
            return;
        }
        let mut prev_move = None;
        let mut p = self.position.clone();

        // Extract the PV, the value, and the bound from the TT.
        let mut pv = Vec::new();
        let mut value = -30000;
        let mut bound = BOUND_LOWER;
        while let Some(entry) = self.tt.probe(p.hash()) {
            if pv.len() < depth as usize && entry.bound() != BOUND_NONE {
                if let Some(m) = prev_move {
                    // Extend the PV.
                    pv.push(m);
                }

                // Get the value and the bound type. In half of the
                // cases the value stored in `entry` is from other
                // side's perspective.
                if pv.len() & 1 == 0 {
                    value = entry.value();
                    bound = entry.bound();
                } else {
                    value = -entry.value();
                    bound = match entry.bound() {
                        BOUND_UPPER => BOUND_LOWER,
                        BOUND_LOWER => BOUND_UPPER,
                        x => x,
                    };
                }

                // The values under -19999 and over 19999 carry
                // information in how many moves is the inevitable
                // checkmate. We choose to not show this to the user,
                // because it would complicate unnecessarily the PV
                // extraction procedure.
                if value >= 20000 && bound == BOUND_LOWER {
                    value = 19999;
                    bound = BOUND_EXACT
                }
                if value <= -20000 && bound == BOUND_UPPER {
                    value = -19999;
                    bound = BOUND_EXACT
                }

                // Try to extend the PV.
                if bound == BOUND_EXACT {
                    if let Some(m) = p.try_move_digest(entry.move16()) {
                        if p.do_move(m) && !p.is_repeated() {
                            prev_move = Some(m);
                            continue;
                        }
                    }
                }
            }
            break;
        }

        // Check if the extracted PV is imperfect.
        self.mangled_pv = bound != BOUND_EXACT || pv.len() < depth as usize;

        // Send the newly extracted PV to the GUI.
        let value_suffix = match bound {
            BOUND_EXACT => "",
            BOUND_UPPER => " upperbound",
            BOUND_LOWER => " lowerbound",
            _ => panic!("unexpected bound type"),
        };
        let mut pv_string = String::new();
        for m in &pv {
            pv_string.push_str(&m.notation());
            pv_string.push(' ');
        }
        self.reply_queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", depth)),
            ("score".to_string(), format!("cp {}{}", value, value_suffix)),
            ("time".to_string(), format!("{}", self.searched_time)),
            ("nodes".to_string(), format!("{}", self.searched_nodes)),
            ("nps".to_string(), format!("{}", self.nps)),
            ("pv".to_string(), pv_string),
        ]));
        self.silent_since = SystemTime::now();
    }

    // A helper method. It reports the current depth, the searched
    // node count, and the nodes per second to the GUI.
    fn report_progress(&mut self) {
        self.reply_queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", self.current_depth)),
            ("nodes".to_string(), format!("{}", self.searched_nodes)),
            ("nps".to_string(), format!("{}", self.nps)),
        ]));
        self.silent_since = SystemTime::now();
    }

    // A helper method for `Engine::stop`. It probes the TT for the
    // best move in the position `p`, plays it, and returns it. If the
    // TT gives no move this function will play and return the first
    // legal move. `None` is returned only when there are no legal
    // moves.
    fn do_best_move(&self, p: &mut Position) -> Option<Move> {
        // Try to get best move from the TT.
        let move16 = self.tt.probe(p.hash()).map_or(0, |entry| entry.move16());
        if let Some(m) = p.try_move_digest(move16) {
            if p.do_move(m) {
                return Some(m);
            }
        }

        // Otherwise, pick the first legal move.
        let mut s = MoveStack::new();
        p.generate_moves(&mut s);
        while let Some(m) = s.pop() {
            if p.do_move(m) {
                return Some(m);
            }
        }

        // No legal moves.
        None
    }
}

/// Implements `UciEngineFactory` trait.
pub struct EngineFactory;

impl UciEngineFactory<Engine> for EngineFactory {
    fn name(&self) -> String {
        format!("Alcibiades {}", VERSION)
    }

    fn author(&self) -> String {
        "Evgeni Pandurski".to_string()
    }

    fn options(&self) -> Vec<(OptionName, OptionDescription)> {
        vec![
            // TODO: Calculate a sane limit for the hash size.
            ("Hash".to_string(), OptionDescription::Spin { min: 1, max: 2048, default: 16 }),
            ("Ponder".to_string(), OptionDescription::Check { default: false }),
        ]
    }

    fn create(&self, hash_size_mb: Option<usize>) -> Engine {
        Engine::new(hash_size_mb.unwrap_or(16))
    }
}

enum TimeManagement {
    MoveTime(u64),
    MoveTimeHint(u64),
    Nodes(NodeCount),
    Depth(u8),
    Infinite,
}
<commit_msg>fixed bug in PV extraction<commit_after>pub mod search;

use std::thread;
use std::cmp::min;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver};
use std::time::SystemTime;
use std::num::Wrapping;
use basetypes::*;
use uci::{UciEngine, UciEngineFactory, EngineReply, OptionName, OptionDescription};
use position::Position;
use chess_move::*;
use tt::*;

const VERSION: &'static str = "0.1";
const MAX_DEPTH: u8 = 126;

/// Implements `UciEngine` trait.
pub struct Engine {
    tt: Arc<TranspositionTable>,
    reply_queue: VecDeque<EngineReply>,
    position: Position,
    is_thinking: bool,
    is_pondering: bool,

    // Tells the engine if it will be allowed to ponder. This option
    // is needed because the engine might change its time management
    // algorithm when pondering is allowed.
    pondering_is_allowed: bool,

    // A channel for sending commands to the search thread.
    commands: Sender<search::Command>,

    // A channel for receiving reports from the search thread.
    reports: Receiver<search::Report>,

    // Search status info
    search_id: usize,
    thinking_since: SystemTime,
    silent_since: SystemTime,
    current_depth: u8,
    searched_nodes: NodeCount,
    searched_time: u64, // milliseconds
    nps: u64, // nodes per second
    mangled_pv: bool, // `true` if the primary variation is imperfect
    stop_when: TimeManagement,
}

impl Engine {
    /// Creates a new instance.
    ///
    /// `tt_size_mb` is the preferred size of the transposition
    /// table in Mbytes.
    pub fn new(tt_size_mb: usize) -> Engine {
        let mut tt = TranspositionTable::new();
        tt.resize(tt_size_mb);
        let tt1 = Arc::new(tt);
        let tt2 = tt1.clone();
        let (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        thread::spawn(move || {
            search::run_deepening(tt2, commands_rx, reports_tx);
        });

        Engine {
            tt: tt1,
            reply_queue: VecDeque::new(),
            position: Position::from_fen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 \
                                          1")
                          .ok()
                          .unwrap(),
            is_thinking: false,
            is_pondering: false,
            pondering_is_allowed: false,
            commands: commands_tx,
            reports: reports_rx,
            search_id: 0,
            thinking_since: SystemTime::now(),
            silent_since: SystemTime::now(),
            current_depth: 0,
            searched_nodes: 0,
            searched_time: 0,
            nps: 0,
            mangled_pv: false,
            stop_when: TimeManagement::Infinite,
        }
    }
}

impl UciEngine for Engine {
    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            "Ponder" => {
                self.pondering_is_allowed = value == "true";
            }

            // An invalid option. Notice that we do not support
            // re-sizing of the transposition table once the engine
            // had started.
            _ => (),
        }
    }

    fn new_game(&mut self) {
        // Clearing the transposition table would not change anything,
        // so we do nothing.
    }

    fn position(&mut self, fen: &str, moves: &mut Iterator<Item = &str>) {
        if !self.is_thinking {
            if let Ok(p) = Position::from_history(fen, moves) {
                self.position = p;
                self.tt.new_search();
            }
        }
    }

    #[allow(unused_variables)]
    fn go(&mut self,
          searchmoves: Option<Vec<String>>,
          ponder: bool,
          wtime: Option<u64>,
          btime: Option<u64>,
          winc: Option<u64>,
          binc: Option<u64>,
          movestogo: Option<u64>,
          depth: Option<u64>,
          nodes: Option<u64>,
          mate: Option<u64>,
          movetime: Option<u64>,
          infinite: bool) {
        if !self.is_thinking {
            // Note: We ignore "searchmoves" and "mate" parameters.

            self.is_thinking = true;
            self.is_pondering = ponder;
            self.search_id = (Wrapping(self.search_id) + Wrapping(1)).0;
            self.thinking_since = SystemTime::now();
            self.silent_since = SystemTime::now();
            self.current_depth = 0;
            self.searched_nodes = 0;
            self.searched_time = 0;
            self.mangled_pv = false;

            // Figure out when we should stop thinking.
            //
            // TODO: Implement smarter time management.
            let (time, inc) = if self.position.board().to_move() == WHITE {
                (wtime, winc.unwrap_or(0))
            } else {
                (btime, binc.unwrap_or(0))
            };
            self.stop_when = if infinite {
                TimeManagement::Infinite
            } else if movetime.is_some() {
                TimeManagement::MoveTime(movetime.unwrap())
            } else if nodes.is_some() {
                TimeManagement::Nodes(nodes.unwrap())
            } else if depth.is_some() {
                TimeManagement::Depth(depth.unwrap() as u8)
            } else {
                let time = time.unwrap_or(0);
                let movestogo = movestogo.unwrap_or(40);
                let movetime = (time + inc * movestogo) / movestogo;
                TimeManagement::MoveTimeHint(min(movetime, time / 2))
            };

            // Start deepening search to the maximum possible depth.
            self.commands
                .send(search::Command::Search {
                    search_id: self.search_id,
                    position: self.position.clone(),
                    depth: MAX_DEPTH,
                    lower_bound: -20000,
                    upper_bound: 20000,
                })
                .unwrap();
        }
    }

    fn ponder_hit(&mut self) {
        if self.is_thinking && self.is_pondering {
            self.is_pondering = false;
        }
    }

    fn stop(&mut self) {
        if self.is_thinking {
            self.commands.send(search::Command::Stop).unwrap();
            let mut p = self.position.clone();
            let (best_move, ponder_move) = if let Some(m) = self.do_best_move(&mut p) {
                (m.notation(), self.do_best_move(&mut p).map(|m| m.notation()))
            } else {
                ("0000".to_string(), None)
            };
            self.reply_queue.push_back(EngineReply::BestMove {
                best_move: best_move,
                ponder_move: ponder_move,
            });
            self.is_thinking = false;
        }
    }

    fn is_thinking(&self) -> bool {
        self.is_thinking
    }

    fn get_reply(&mut self) -> Option<EngineReply> {
        // Check if we should stop thinking.
        if self.is_thinking && !self.is_pondering {
            if match self.stop_when {
                TimeManagement::MoveTime(t) => self.searched_time >= t,
                TimeManagement::MoveTimeHint(t) => self.searched_time >= t,
                TimeManagement::Nodes(n) => self.searched_nodes >= n,
                TimeManagement::Depth(d) => self.current_depth > d,
                TimeManagement::Infinite => false,
            } {
                self.stop();
            }
        }

        // Check the reports queue.
        while let Ok(report) = self.reports.try_recv() {
            if self.is_thinking {
                match report {
                    search::Report::Progress { search_id, depth, searched_nodes }
                        if search_id == self.search_id => {
                        // Register search progress.
                        self.register_progress(depth, searched_nodes);
                        if self.silent_since.elapsed().unwrap().as_secs() > 20 {
                            if self.mangled_pv {
                                self.report_pv(depth - 1);
                            } else {
                                self.report_progress();
                            }
                        }
                    }
                    search::Report::Done { search_id, .. } if search_id == self.search_id => {
                        // Terminate the search if not infinite.
                        self.stop_when = match self.stop_when {
                            TimeManagement::Infinite => TimeManagement::Infinite,
                            _ => TimeManagement::MoveTime(0),
                        };
                    }
                    // We may still receive stale reports from already
                    // stopped searches.
                    _ => (),
                }
            }
        }

        self.reply_queue.pop_front()
    }
}

impl Engine {
    // A helper method. It updates the search status info and makes
    // sure that a new PV is sent to the GUI for every newly reached
    // depth.
    fn register_progress(&mut self, depth: u8, searched_nodes: NodeCount) {
        let thinking_duration = self.thinking_since.elapsed().unwrap();
        self.searched_time = 1000 * thinking_duration.as_secs() +
                             (thinking_duration.subsec_nanos() / 1000000) as u64;
        self.searched_nodes = searched_nodes;
        self.nps = 1000 * (self.nps + self.searched_nodes) / (1000 + self.searched_time);
        if self.current_depth < depth {
            self.current_depth = depth;
            self.report_pv(depth - 1);
            if !self.mangled_pv {
                self.report_progress();
            }
        }
    }

    // A helper method. It extracts the primary variation (PV) from
    // the transposition table (TT) and sends it to the GUI.
    fn report_pv(&mut self, depth: u8) {
        if depth == 0 {
            return;
        }
        let mut prev_move = None;
        let mut p = self.position.clone();

        // Extract the PV, the value, and the bound from the TT.
        let mut pv = Vec::new();
        let mut value = -30000;
        let mut bound = BOUND_LOWER;
        while let Some(entry) = self.tt.probe(p.hash()) {
            if pv.len() < depth as usize && entry.bound() != BOUND_NONE {
                if let Some(m) = prev_move {
                    // Extend the PV.
                    pv.push(m);
                }

                // Get the value and the bound type. In half of the
                // cases the value stored in `entry` is from other
                // side's perspective.
                if pv.len() & 1 == 0 {
                    value = entry.value();
                    bound = entry.bound();
                } else {
                    value = -entry.value();
                    bound = match entry.bound() {
                        BOUND_UPPER => BOUND_LOWER,
                        BOUND_LOWER => BOUND_UPPER,
                        x => x,
                    };
                }

                // The values under -19999 and over 19999 carry
                // information in how many moves is the inevitable
                // checkmate. We choose to not show this to the user,
                // because it would complicate unnecessarily the PV
                // extraction procedure.
                if value >= 20000 && bound != BOUND_UPPER {
                    value = 19999;
                    bound = BOUND_EXACT
                }
                if value <= -20000 && bound != BOUND_LOWER {
                    value = -19999;
                    bound = BOUND_EXACT
                }

                // Try to extend the PV.
                if bound == BOUND_EXACT {
                    if let Some(m) = p.try_move_digest(entry.move16()) {
                        if p.do_move(m) && !p.is_repeated() {
                            prev_move = Some(m);
                            continue;
                        }
                    }
                }
            }
            break;
        }

        // Check if the extracted PV is imperfect.
        self.mangled_pv = bound != BOUND_EXACT || pv.len() < depth as usize;

        // Send the newly extracted PV to the GUI.
        let value_suffix = match bound {
            BOUND_EXACT => "",
            BOUND_UPPER => " upperbound",
            BOUND_LOWER => " lowerbound",
            _ => panic!("unexpected bound type"),
        };
        let mut pv_string = String::new();
        for m in &pv {
            pv_string.push_str(&m.notation());
            pv_string.push(' ');
        }
        self.reply_queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", depth)),
            ("score".to_string(), format!("cp {}{}", value, value_suffix)),
            ("time".to_string(), format!("{}", self.searched_time)),
            ("nodes".to_string(), format!("{}", self.searched_nodes)),
            ("nps".to_string(), format!("{}", self.nps)),
            ("pv".to_string(), pv_string),
        ]));
        self.silent_since = SystemTime::now();
    }

    // A helper method. It reports the current depth, the searched
    // node count, and the nodes per second to the GUI.
    fn report_progress(&mut self) {
        self.reply_queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", self.current_depth)),
            ("nodes".to_string(), format!("{}", self.searched_nodes)),
            ("nps".to_string(), format!("{}", self.nps)),
        ]));
        self.silent_since = SystemTime::now();
    }

    // A helper method for `Engine::stop`. It probes the TT for the
    // best move in the position `p`, plays it, and returns it. If the
    // TT gives no move this function will play and return the first
    // legal move. `None` is returned only when there are no legal
    // moves.
    fn do_best_move(&self, p: &mut Position) -> Option<Move> {
        // Try to get best move from the TT.
        let move16 = self.tt.probe(p.hash()).map_or(0, |entry| entry.move16());
        if let Some(m) = p.try_move_digest(move16) {
            if p.do_move(m) {
                return Some(m);
            }
        }

        // Otherwise, pick the first legal move.
        let mut s = MoveStack::new();
        p.generate_moves(&mut s);
        while let Some(m) = s.pop() {
            if p.do_move(m) {
                return Some(m);
            }
        }

        // No legal moves.
        None
    }
}

/// Implements `UciEngineFactory` trait.
pub struct EngineFactory;

impl UciEngineFactory<Engine> for EngineFactory {
    fn name(&self) -> String {
        format!("Alcibiades {}", VERSION)
    }

    fn author(&self) -> String {
        "Evgeni Pandurski".to_string()
    }

    fn options(&self) -> Vec<(OptionName, OptionDescription)> {
        vec![
            // TODO: Calculate a sane limit for the hash size.
            ("Hash".to_string(), OptionDescription::Spin { min: 1, max: 2048, default: 16 }),
            ("Ponder".to_string(), OptionDescription::Check { default: false }),
        ]
    }

    fn create(&self, hash_size_mb: Option<usize>) -> Engine {
        Engine::new(hash_size_mb.unwrap_or(16))
    }
}

enum TimeManagement {
    MoveTime(u64),
    MoveTimeHint(u64),
    Nodes(NodeCount),
    Depth(u8),
    Infinite,
}
<|endoftext|>


<commit_before>use crate::{
    platform_mut::DBPlatformMut,
    DataError,
    DatabaseMut,
    DbError,
    Table,
    ToValue,
    Value,
};

use rustorm_dao::{
    FromDao,
    TableName,
    ToColumnNames,
    ToDao,
    ToTableName,
};

pub struct EntityManagerMut(pub DBPlatformMut);

impl EntityManagerMut {
    pub fn set_session_user(&mut self, username: &str) -> Result<(), DbError> {
        let sql = format!("SET SESSION ROLE '{}'", username);
        self.0.execute_sql_with_return(&sql, &[])?;
        Ok(())
    }

    // pub fn get_role(&self, username: &str) -> Result<Option<Role>, DbError> {
    //     let result = self.0.get_roles(&self, username);
    //     match result {
    //         Ok(mut result) => {
    //             match result.len() {
    //                 0 => Ok(None),
    //                 1 => Ok(Some(result.remove(0))),
    //                 _ => {
    //                     Err(DbError::DataError(
    //                         DataError::MoreThan1RecordReturned,
    //                     ))
    //                 }
    //             }
    //         }
    //         Err(e) => Err(e),
    //     }
    // }

    pub fn db(&self) -> &dyn DatabaseMut { &*self.0 }

    /// get all the records of this table
    pub fn get_all<T>(&mut self) -> Result<Vec<T>, DbError>
    where
        T: ToTableName + ToColumnNames + FromDao,
    {
        let table = T::to_table_name();
        let columns = T::to_column_names();
        let enumerated_columns = columns
            .iter()
            .map(|c| c.name.to_owned())
            .collect::<Vec<_>>()
            .join(", ");
        let sql = format!(
            "SELECT {} FROM {}",
            enumerated_columns,
            table.complete_name()
        );
        let rows = self.0.execute_sql_with_return(&sql, &[])?;
        let mut entities = vec![];
        for dao in rows.iter() {
            let entity = T::from_dao(&dao);
            entities.push(entity)
        }
        Ok(entities)
    }

    /// get the table from database based on this column name
    pub fn get_table(&mut self, table_name: &TableName) -> Result<Table, DbError> {
        self.0.get_table(table_name)
    }

    // /// get all the user table and views from the database
    // pub fn get_all_tables(&self) -> Result<Vec<Table>, DbError> {
    //     info!("EXPENSIVE DB OPERATION: get_all_tables");
    //     self.0.get_all_tables(self)
    // }

    /// Get the total count of records
    pub fn get_total_records(&mut self, table_name: &TableName) -> Result<usize, DbError> {
        #[derive(crate::FromDao)]
        struct Count {
            count: i64,
        }
        let sql = format!(
            "SELECT COUNT(*) AS count FROM {}",
            table_name.complete_name()
        );
        let count: Result<Count, DbError> = self.execute_sql_with_one_return(&sql, &[]);
        count.map(|c| c.count as usize)
    }

    // pub fn get_users(&self) -> Result<Vec<User>, DbError> {
    //     self.0.get_users(self)
    // }

    // pub fn get_database_name(&self) -> Result<Option<DatabaseName>, DbError> {
    //     self.0.get_database_name(self)
    // }

    // /// get all table and views grouped per schema
    // pub fn get_grouped_tables(&self) -> Result<Vec<SchemaContent>, DbError> {
    //     self.0.get_grouped_tables(self)
    // }

    pub fn insert<T, R>(&mut self, _entities: &[&T]) -> Result<Vec<R>, DbError>
    where
        T: ToTableName + ToColumnNames + ToDao,
        R: FromDao + ToColumnNames,
    {
        match self.0 {
            // #[cfg(feature = "with-sqlite")]
            // DBPlatform::Sqlite(_) => self.insert_simple(_entities),
            // #[cfg(feature = "with-postgres")]
            // DBPlatform::Postgres(_) => {
            //     self.insert_bulk_with_returning_support(_entities)
            // }
            #[cfg(feature = "with-mysql")]
            DBPlatformMut::Mysql(_) => self.insert_simple(_entities),
        }
    }

    /// called when the platform used is postgresql
    pub fn insert_bulk_with_returning_support<T, R>(
        &mut self,
        entities: &[&T],
    ) -> Result<Vec<R>, DbError>
    where
        T: ToTableName + ToColumnNames + ToDao,
        R: FromDao + ToColumnNames,
    {
        let columns = T::to_column_names();
        let mut sql = self.build_insert_clause(entities);
        let return_columns = R::to_column_names();
        sql += &self.build_returning_clause(return_columns);

        let mut values: Vec<Value> = Vec::with_capacity(entities.len() * columns.len());
        for entity in entities {
            let dao = entity.to_dao();
            for col in columns.iter() {
                let value = dao.get_value(&col.name);
                match value {
                    Some(value) => values.push(value.clone()),
                    None => values.push(Value::Nil),
                }
            }
        }
        let bvalues: Vec<&Value> = values.iter().collect();
        let rows = self.0.execute_sql_with_return(&sql, &bvalues)?;
        let mut retrieved_entities = vec![];
        for dao in rows.iter() {
            let retrieved = R::from_dao(&dao);
            retrieved_entities.push(retrieved);
        }
        Ok(retrieved_entities)
    }

    /// called multiple times when using database platform that doesn;t support multiple value
    /// insert such as sqlite
    pub fn single_insert<T>(&mut self, entity: &T) -> Result<(), DbError>
    where
        T: ToTableName + ToColumnNames + ToDao,
    {
        let columns = T::to_column_names();
        let sql = self.build_insert_clause(&[entity]);
        let dao = entity.to_dao();
        let mut values: Vec<Value> = Vec::with_capacity(columns.len());
        for col in columns.iter() {
            let value = dao.get_value(&col.name);
            match value {
                Some(value) => values.push(value.clone()),
                None => values.push(Value::Nil),
            }
        }
        let bvalues: Vec<&Value> = values.iter().collect();
        self.0.execute_sql_with_return(&sql, &bvalues)?;
        Ok(())
    }

    /// this is soly for use with sqlite since sqlite doesn't support bulk insert
    pub fn insert_simple<T, R>(&mut self, entities: &[&T]) -> Result<Vec<R>, DbError>
    where
        T: ToTableName + ToColumnNames + ToDao,
        R: FromDao + ToColumnNames,
    {
        let return_columns = R::to_column_names();
        let return_column_names = return_columns
            .iter()
            .map(|rc| rc.name.to_owned())
            .collect::<Vec<_>>()
            .join(", ");

        let table = T::to_table_name();
        //TODO: move this specific query to sqlite
        let last_insert_sql = format!(
            "\
             SELECT {} \
             FROM {} \
             WHERE ROWID = (\
             SELECT LAST_INSERT_ROWID() FROM {})",
            return_column_names,
            table.complete_name(),
            table.complete_name()
        );
        let mut retrieved_entities = vec![];
        println!("sql: {}", last_insert_sql);
        for entity in entities {
            self.single_insert(*entity)?;
            let retrieved = self.execute_sql_with_return(&last_insert_sql, &[])?;
            retrieved_entities.extend(retrieved);
        }
        Ok(retrieved_entities)
    }

    /// build the returning clause
    fn build_returning_clause(&self, return_columns: Vec<rustorm_dao::ColumnName>) -> String {
        format!(
            "\nRETURNING \n{}",
            return_columns
                .iter()
                .map(|rc| rc.name.to_owned())
                .collect::<Vec<_>>()
                .join(", ")
        )
    }

    /// build an insert clause
    fn build_insert_clause<T>(&self, entities: &[&T]) -> String
    where
        T: ToTableName + ToColumnNames + ToDao,
    {
        let table = T::to_table_name();
        let columns = T::to_column_names();
        let columns_len = columns.len();
        let mut sql = String::new();
        sql += &format!("INSERT INTO {} ", table.complete_name());
        sql += &format!(
            "({})\n",
            columns
                .iter()
                .map(|c| c.name.to_owned())
                .collect::<Vec<_>>()
                .join(", ")
        );
        sql += "VALUES ";
        sql += &entities
            .iter()
            .enumerate()
            .map(|(y, _)| {
                format!(
                    "\n\t({})",
                    columns
                        .iter()
                        .enumerate()
                        .map(|(x, _)| {
                            #[allow(unreachable_patterns)]
                            match self.0 {
                                // #[cfg(feature = "with-sqlite")]
                                // DBPlatform::Sqlite(_) => format!("${}", _y * columns_len + _x + 1),
                                // #[cfg(feature = "with-postgres")]
                                // DBPlatform::Postgres(_) => format!("${}", _y * columns_len + _x + 1),
                                #[cfg(feature = "with-mysql")]
                                DBPlatformMut::Mysql(_) => "?".to_string(),
                                _ => format!("${}", y * columns_len + x + 1),
                            }
                        })
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            })
            .collect::<Vec<_>>()
            .join(", ");
        sql
    }

    #[allow(clippy::redundant_closure)]
    pub fn execute_sql_with_return<'a, R>(
        &mut self,
        sql: &str,
        params: &[&'a dyn ToValue],
    ) -> Result<Vec<R>, DbError>
    where
        R: FromDao,
    {
        let values: Vec<Value> = params.iter().map(|p| p.to_value()).collect();
        let bvalues: Vec<&Value> = values.iter().collect();
        let rows = self.0.execute_sql_with_return(sql, &bvalues)?;
        Ok(rows.iter().map(|dao| R::from_dao(&dao)).collect::<Vec<R>>())
    }

    pub fn execute_sql_with_one_return<'a, R>(
        &mut self,
        sql: &str,
        params: &[&'a dyn ToValue],
    ) -> Result<R, DbError>
    where
        R: FromDao,
    {
        let result: Result<Vec<R>, DbError> = self.execute_sql_with_return(sql, &params);
        match result {
            Ok(mut result) => {
                match result.len() {
                    0 => Err(DbError::DataError(DataError::ZeroRecordReturned)),
                    1 => Ok(result.remove(0)),
                    _ => Err(DbError::DataError(DataError::MoreThan1RecordReturned)),
                }
            }
            Err(e) => Err(e),
        }
    }

    pub fn execute_sql_with_maybe_one_return<'a, R>(
        &mut self,
        sql: &str,
        params: &[&'a dyn ToValue],
    ) -> Result<Option<R>, DbError>
    where
        R: FromDao,
    {
        let result: Result<Vec<R>, DbError> = self.execute_sql_with_return(sql, &params);
        match result {
            Ok(mut result) => {
                match result.len() {
                    0 => Ok(None),
                    1 => Ok(Some(result.remove(0))),
                    _ => Err(DbError::DataError(DataError::MoreThan1RecordReturned)),
                }
            }
            Err(e) => Err(e),
        }
    }
}
<commit_msg>Make the db() function return a mutable reference to DatabaseMut<commit_after>use crate::{
    platform_mut::DBPlatformMut,
    DataError,
    DatabaseMut,
    DbError,
    Table,
    ToValue,
    Value,
};

use rustorm_dao::{
    FromDao,
    TableName,
    ToColumnNames,
    ToDao,
    ToTableName,
};

pub struct EntityManagerMut(pub DBPlatformMut);

impl EntityManagerMut {
    pub fn set_session_user(&mut self, username: &str) -> Result<(), DbError> {
        let sql = format!("SET SESSION ROLE '{}'", username);
        self.0.execute_sql_with_return(&sql, &[])?;
        Ok(())
    }

    // pub fn get_role(&self, username: &str) -> Result<Option<Role>, DbError> {
    //     let result = self.0.get_roles(&self, username);
    //     match result {
    //         Ok(mut result) => {
    //             match result.len() {
    //                 0 => Ok(None),
    //                 1 => Ok(Some(result.remove(0))),
    //                 _ => {
    //                     Err(DbError::DataError(
    //                         DataError::MoreThan1RecordReturned,
    //                     ))
    //                 }
    //             }
    //         }
    //         Err(e) => Err(e),
    //     }
    // }

    pub fn db(&mut self) -> &mut dyn DatabaseMut { &mut *self.0 }

    /// get all the records of this table
    pub fn get_all<T>(&mut self) -> Result<Vec<T>, DbError>
    where
        T: ToTableName + ToColumnNames + FromDao,
    {
        let table = T::to_table_name();
        let columns = T::to_column_names();
        let enumerated_columns = columns
            .iter()
            .map(|c| c.name.to_owned())
            .collect::<Vec<_>>()
            .join(", ");
        let sql = format!(
            "SELECT {} FROM {}",
            enumerated_columns,
            table.complete_name()
        );
        let rows = self.0.execute_sql_with_return(&sql, &[])?;
        let mut entities = vec![];
        for dao in rows.iter() {
            let entity = T::from_dao(&dao);
            entities.push(entity)
        }
        Ok(entities)
    }

    /// get the table from database based on this column name
    pub fn get_table(&mut self, table_name: &TableName) -> Result<Table, DbError> {
        self.0.get_table(table_name)
    }

    // /// get all the user table and views from the database
    // pub fn get_all_tables(&self) -> Result<Vec<Table>, DbError> {
    //     info!("EXPENSIVE DB OPERATION: get_all_tables");
    //     self.0.get_all_tables(self)
    // }

    /// Get the total count of records
    pub fn get_total_records(&mut self, table_name: &TableName) -> Result<usize, DbError> {
        #[derive(crate::FromDao)]
        struct Count {
            count: i64,
        }
        let sql = format!(
            "SELECT COUNT(*) AS count FROM {}",
            table_name.complete_name()
        );
        let count: Result<Count, DbError> = self.execute_sql_with_one_return(&sql, &[]);
        count.map(|c| c.count as usize)
    }

    // pub fn get_users(&self) -> Result<Vec<User>, DbError> {
    //     self.0.get_users(self)
    // }

    // pub fn get_database_name(&self) -> Result<Option<DatabaseName>, DbError> {
    //     self.0.get_database_name(self)
    // }

    // /// get all table and views grouped per schema
    // pub fn get_grouped_tables(&self) -> Result<Vec<SchemaContent>, DbError> {
    //     self.0.get_grouped_tables(self)
    // }

    pub fn insert<T, R>(&mut self, _entities: &[&T]) -> Result<Vec<R>, DbError>
    where
        T: ToTableName + ToColumnNames + ToDao,
        R: FromDao + ToColumnNames,
    {
        match self.0 {
            // #[cfg(feature = "with-sqlite")]
            // DBPlatform::Sqlite(_) => self.insert_simple(_entities),
            // #[cfg(feature = "with-postgres")]
            // DBPlatform::Postgres(_) => {
            //     self.insert_bulk_with_returning_support(_entities)
            // }
            #[cfg(feature = "with-mysql")]
            DBPlatformMut::Mysql(_) => self.insert_simple(_entities),
        }
    }

    /// called when the platform used is postgresql
    pub fn insert_bulk_with_returning_support<T, R>(
        &mut self,
        entities: &[&T],
    ) -> Result<Vec<R>, DbError>
    where
        T: ToTableName + ToColumnNames + ToDao,
        R: FromDao + ToColumnNames,
    {
        let columns = T::to_column_names();
        let mut sql = self.build_insert_clause(entities);
        let return_columns = R::to_column_names();
        sql += &self.build_returning_clause(return_columns);

        let mut values: Vec<Value> = Vec::with_capacity(entities.len() * columns.len());
        for entity in entities {
            let dao = entity.to_dao();
            for col in columns.iter() {
                let value = dao.get_value(&col.name);
                match value {
                    Some(value) => values.push(value.clone()),
                    None => values.push(Value::Nil),
                }
            }
        }
        let bvalues: Vec<&Value> = values.iter().collect();
        let rows = self.0.execute_sql_with_return(&sql, &bvalues)?;
        let mut retrieved_entities = vec![];
        for dao in rows.iter() {
            let retrieved = R::from_dao(&dao);
            retrieved_entities.push(retrieved);
        }
        Ok(retrieved_entities)
    }

    /// called multiple times when using database platform that doesn;t support multiple value
    /// insert such as sqlite
    pub fn single_insert<T>(&mut self, entity: &T) -> Result<(), DbError>
    where
        T: ToTableName + ToColumnNames + ToDao,
    {
        let columns = T::to_column_names();
        let sql = self.build_insert_clause(&[entity]);
        let dao = entity.to_dao();
        let mut values: Vec<Value> = Vec::with_capacity(columns.len());
        for col in columns.iter() {
            let value = dao.get_value(&col.name);
            match value {
                Some(value) => values.push(value.clone()),
                None => values.push(Value::Nil),
            }
        }
        let bvalues: Vec<&Value> = values.iter().collect();
        self.0.execute_sql_with_return(&sql, &bvalues)?;
        Ok(())
    }

    /// this is soly for use with sqlite since sqlite doesn't support bulk insert
    pub fn insert_simple<T, R>(&mut self, entities: &[&T]) -> Result<Vec<R>, DbError>
    where
        T: ToTableName + ToColumnNames + ToDao,
        R: FromDao + ToColumnNames,
    {
        let return_columns = R::to_column_names();
        let return_column_names = return_columns
            .iter()
            .map(|rc| rc.name.to_owned())
            .collect::<Vec<_>>()
            .join(", ");

        let table = T::to_table_name();
        //TODO: move this specific query to sqlite
        let last_insert_sql = format!(
            "\
             SELECT {} \
             FROM {} \
             WHERE ROWID = (\
             SELECT LAST_INSERT_ROWID() FROM {})",
            return_column_names,
            table.complete_name(),
            table.complete_name()
        );
        let mut retrieved_entities = vec![];
        println!("sql: {}", last_insert_sql);
        for entity in entities {
            self.single_insert(*entity)?;
            let retrieved = self.execute_sql_with_return(&last_insert_sql, &[])?;
            retrieved_entities.extend(retrieved);
        }
        Ok(retrieved_entities)
    }

    /// build the returning clause
    fn build_returning_clause(&self, return_columns: Vec<rustorm_dao::ColumnName>) -> String {
        format!(
            "\nRETURNING \n{}",
            return_columns
                .iter()
                .map(|rc| rc.name.to_owned())
                .collect::<Vec<_>>()
                .join(", ")
        )
    }

    /// build an insert clause
    fn build_insert_clause<T>(&self, entities: &[&T]) -> String
    where
        T: ToTableName + ToColumnNames + ToDao,
    {
        let table = T::to_table_name();
        let columns = T::to_column_names();
        let columns_len = columns.len();
        let mut sql = String::new();
        sql += &format!("INSERT INTO {} ", table.complete_name());
        sql += &format!(
            "({})\n",
            columns
                .iter()
                .map(|c| c.name.to_owned())
                .collect::<Vec<_>>()
                .join(", ")
        );
        sql += "VALUES ";
        sql += &entities
            .iter()
            .enumerate()
            .map(|(y, _)| {
                format!(
                    "\n\t({})",
                    columns
                        .iter()
                        .enumerate()
                        .map(|(x, _)| {
                            #[allow(unreachable_patterns)]
                            match self.0 {
                                // #[cfg(feature = "with-sqlite")]
                                // DBPlatform::Sqlite(_) => format!("${}", _y * columns_len + _x + 1),
                                // #[cfg(feature = "with-postgres")]
                                // DBPlatform::Postgres(_) => format!("${}", _y * columns_len + _x + 1),
                                #[cfg(feature = "with-mysql")]
                                DBPlatformMut::Mysql(_) => "?".to_string(),
                                _ => format!("${}", y * columns_len + x + 1),
                            }
                        })
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            })
            .collect::<Vec<_>>()
            .join(", ");
        sql
    }

    #[allow(clippy::redundant_closure)]
    pub fn execute_sql_with_return<'a, R>(
        &mut self,
        sql: &str,
        params: &[&'a dyn ToValue],
    ) -> Result<Vec<R>, DbError>
    where
        R: FromDao,
    {
        let values: Vec<Value> = params.iter().map(|p| p.to_value()).collect();
        let bvalues: Vec<&Value> = values.iter().collect();
        let rows = self.0.execute_sql_with_return(sql, &bvalues)?;
        Ok(rows.iter().map(|dao| R::from_dao(&dao)).collect::<Vec<R>>())
    }

    pub fn execute_sql_with_one_return<'a, R>(
        &mut self,
        sql: &str,
        params: &[&'a dyn ToValue],
    ) -> Result<R, DbError>
    where
        R: FromDao,
    {
        let result: Result<Vec<R>, DbError> = self.execute_sql_with_return(sql, &params);
        match result {
            Ok(mut result) => {
                match result.len() {
                    0 => Err(DbError::DataError(DataError::ZeroRecordReturned)),
                    1 => Ok(result.remove(0)),
                    _ => Err(DbError::DataError(DataError::MoreThan1RecordReturned)),
                }
            }
            Err(e) => Err(e),
        }
    }

    pub fn execute_sql_with_maybe_one_return<'a, R>(
        &mut self,
        sql: &str,
        params: &[&'a dyn ToValue],
    ) -> Result<Option<R>, DbError>
    where
        R: FromDao,
    {
        let result: Result<Vec<R>, DbError> = self.execute_sql_with_return(sql, &params);
        match result {
            Ok(mut result) => {
                match result.len() {
                    0 => Ok(None),
                    1 => Ok(Some(result.remove(0))),
                    _ => Err(DbError::DataError(DataError::MoreThan1RecordReturned)),
                }
            }
            Err(e) => Err(e),
        }
    }
}
<|endoftext|>