<commit_before>use std::str;
use std::cell::RefCell;
use std::time::Instant;
use std::collections::BTreeMap;
use std::net::SocketAddr;
use mio::net::UdpSocket;
use std::io::{self,Write};
use sozu_command::proxy::{FilteredData,MetricsData};

mod network_drain;
mod local_drain;
mod writer;

use self::network_drain::NetworkDrain;
use self::local_drain::LocalDrain;

thread_local! {
  pub static METRICS: RefCell<Aggregator> = RefCell::new(Aggregator::new(String::from("sozu")));
}

#[derive(Debug,Clone,PartialEq)]
pub enum MetricData {
  Gauge(usize),
  GaugeAdd(i64),
  Count(i64),
  Time(usize),
}

impl MetricData {
  fn is_time(&self) -> bool {
    match self {
      &MetricData::Time(_) => true,
      _ => false,
    }
  }

  fn update(&mut self, key: &'static str, m: MetricData) -> bool {
    match (self, m) {
      (&mut MetricData::Gauge(ref mut v1), MetricData::Gauge(v2)) => {
        let changed = *v1 != v2;
        *v1 = v2;
        changed
      },
      (&mut MetricData::Gauge(ref mut v1), MetricData::GaugeAdd(v2)) => {
        debug_assert!(*v1 as i64 + v2 >= 0, "metric {} underflow: previous value: {}, adding: {}", key, v1, v2);
        let changed = v2 != 0;
        let res = *v1 as i64 + v2;
        *v1 = if res >= 0 {
          res as usize
        } else {
          error!("metric {} underflow: previous value: {}, adding: {}", key, v1, v2);
          0
        };

        changed
      },
      (&mut MetricData::Count(ref mut v1), MetricData::Count(v2)) => {
        let changed = v2 != 0;
        *v1 += v2;
        changed
      },
      (s,m) => panic!("tried to update metric {} of value {:?} with an incompatible metric: {:?}", key, s, m)
    }
  }
}

#[derive(Debug,Clone)]
pub struct StoredMetricData {
  last_sent: Instant,
  updated:   bool,
  data:      MetricData,
}

impl StoredMetricData {
  pub fn new(last_sent: Instant, data: MetricData) -> StoredMetricData {
    StoredMetricData {
      last_sent,
      updated: true,
      data: if let MetricData::GaugeAdd(v) = data {
        if v >= 0 {
          MetricData::Gauge(v as usize)
        } else {
          MetricData::Gauge(0)
        }
      } else { data }
    }
  }

  pub fn update(&mut self, key: &'static str, m: MetricData) {
    let updated = self.data.update(key, m);
    if ! self.updated {
      self.updated = updated;
    }
  }
}

pub fn setup<O: Into<String>>(metrics_host: &SocketAddr, origin: O, use_tagged_metrics: bool, prefix: Option<String>) {
  let metrics_socket = udp_bind();

  debug!("setting up metrics: local address = {:#?}", metrics_socket.local_addr());

  METRICS.with(|metrics| {
    if let Some(p) = prefix {
      (*metrics.borrow_mut()).set_up_prefix(p);
    }
    (*metrics.borrow_mut()).set_up_remote(metrics_socket, metrics_host.clone());
    (*metrics.borrow_mut()).set_up_origin(origin.into());
    (*metrics.borrow_mut()).set_up_tagged_metrics(use_tagged_metrics);
  });
}

pub trait Subscriber {
  fn receive_metric(&mut self, label: &'static str, cluster_id: Option<&str>, backend_id: Option<&str>, metric: MetricData);
}

pub struct Aggregator {
  prefix:  String,
  network: Option<NetworkDrain>,
  local:   LocalDrain,
}

impl Aggregator {
  pub fn new(prefix: String) -> Aggregator {
    Aggregator {
      prefix: prefix.clone(),
      network: None,
      local: LocalDrain::new(prefix),
    }
  }

  pub fn set_up_prefix(&mut self, prefix: String) {
    self.prefix = prefix;
  }

  pub fn set_up_remote(&mut self, socket: UdpSocket, addr: SocketAddr) {
    self.network = Some(NetworkDrain::new(self.prefix.clone(), socket, addr));
  }

  pub fn set_up_origin(&mut self, origin: String) {
    self.network.as_mut().map(|n| n.origin = origin);
  }

  pub fn set_up_tagged_metrics(&mut self, tagged: bool) {
    self.network.as_mut().map(|n| n.use_tagged_metrics = tagged);
  }

  pub fn socket(&self) -> Option<&UdpSocket> {
    self.network.as_ref().map(|n| &n.remote.get_ref().socket)
  }

  pub fn socket_mut(&mut self) -> Option<&mut UdpSocket> {
    self.network.as_mut().map(|n| &mut n.remote.get_mut().socket)
  }

  pub fn count_add(&mut self, key: &'static str, count_value: i64) {
    self.receive_metric(key, None, None, MetricData::Count(count_value));
  }

  pub fn set_gauge(&mut self, key: &'static str, gauge_value: usize) {
    self.receive_metric(key, None, None, MetricData::Gauge(gauge_value));
  }

  pub fn gauge_add(&mut self, key: &'static str, gauge_value: i64) {
    self.receive_metric(key, None, None, MetricData::GaugeAdd(gauge_value));
  }

  pub fn writable(&mut self) {
    if let Some(ref mut net) = self.network.as_mut() {
      net.writable();
    }
  }

  pub fn send_data(&mut self) {
    if let Some(ref mut net) = self.network.as_mut() {
      net.send_data();
    }
  }

  pub fn dump_metrics_data(&mut self) -> MetricsData {
    self.local.dump_metrics_data()
  }

  pub fn dump_process_data(&mut self) -> BTreeMap<String, FilteredData> {
    self.local.dump_process_data()
  }

  pub fn clear_local(&mut self, now: time::OffsetDateTime) {
    self.local.clear(now);
  }
}

impl Subscriber for Aggregator {
  fn receive_metric(&mut self, label: &'static str, cluster_id: Option<&str>, backend_id: Option<&str>, metric: MetricData) {
    if let Some(ref mut net) = self.network.as_mut() {
      net.receive_metric(label, cluster_id, backend_id, metric.clone());
    }
    self.local.receive_metric(label, cluster_id, backend_id, metric.clone());
  }
}

#[derive(Debug,Clone,PartialEq)]
pub struct MetricLine {
  label:      &'static str,
  cluster_id:     Option<String>,
  backend_id: Option<String>,
  /// in milliseconds
  duration:   usize,
}

pub struct MetricSocket {
  pub addr:   SocketAddr,
  pub socket: UdpSocket,
}

impl Write for MetricSocket {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
    self.socket.send_to(buf, self.addr)
  }

  fn flush(&mut self) -> io::Result<()> {
    Ok(())
  }
}

pub fn udp_bind() -> UdpSocket {
  UdpSocket::bind("0.0.0.0:0".parse().unwrap()).expect("could not parse address")
}

#[macro_export]
macro_rules! count (
  ($key:expr, $value: expr) => ({
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).count_add($key, v);
    });
  })
);

#[macro_export]
macro_rules! incr (
  ($key:expr) => (count!($key, 1));
  ($key:expr, $cluster_id:expr, $backend_id:expr) => {
    use $crate::metrics::Subscriber;

    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).receive_metric($key, $cluster_id, $backend_id, $crate::metrics::MetricData::Count(1));
    });
  }
);

#[macro_export]
macro_rules! decr (
  ($key:expr) => (count!($key, -1))
);

#[macro_export]
macro_rules! gauge (
  ($key:expr, $value: expr) => ({
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).set_gauge($key, v);
    });
  })
);

#[macro_export]
macro_rules! gauge_add (
  ($key:expr, $value: expr) => ({
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).gauge_add($key, v);
    });
  });
  ($key:expr, $value:expr, $cluster_id:expr, $backend_id:expr) => {
    use $crate::metrics::Subscriber;
    let v = $value;

    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).receive_metric($key, $cluster_id, $backend_id, $crate::metrics::MetricData::GaugeAdd(v));
    });
  }
);

#[macro_export]
macro_rules! time (
  ($key:expr, $value: expr) => ({
    use $crate::metrics::{MetricData,Subscriber};
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      let ref mut m = *metrics.borrow_mut();

      m.receive_metric($key, None, None, MetricData::Time(v as usize));
    });
  });
  ($key:expr, $cluster_id:expr, $value: expr) => ({
    use $crate::metrics::{MetricData,Subscriber};
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      let ref mut m = *metrics.borrow_mut();
      let cluster: &str = $cluster_id;

      m.receive_metric($key, Some(cluster), None, MetricData::Time(v as usize));
    });
  })
);

#[macro_export]
macro_rules! record_backend_metrics (
  ($cluster_id:expr, $backend_id:expr, $response_time: expr, $backend_connection_time: expr, $bin: expr, $bout: expr) => {
    use $crate::metrics::{MetricData,Subscriber};
    $crate::metrics::METRICS.with(|metrics| {
      let ref mut m = *metrics.borrow_mut();
      let cluster_id: &str = $cluster_id.as_str();
      let backend_id: &str = $backend_id;

      m.receive_metric("bytes_in", Some(cluster_id), Some(backend_id), MetricData::Count($bin as i64));
      m.receive_metric("bytes_out", Some(cluster_id), Some(backend_id), MetricData::Count($bout as i64));
      m.receive_metric("backend_response_time", Some(cluster_id), Some(backend_id), MetricData::Time($response_time as usize));
      if let Some(t) = $backend_connection_time {
        m.receive_metric("backend_connection_time", Some(cluster_id), Some(backend_id), MetricData::Time(t.whole_milliseconds() as usize));
      }
    });
  }
);

<commit_msg>add a request counter per cluster and backend<commit_after>use std::str;
use std::cell::RefCell;
use std::time::Instant;
use std::collections::BTreeMap;
use std::net::SocketAddr;
use mio::net::UdpSocket;
use std::io::{self,Write};
use sozu_command::proxy::{FilteredData,MetricsData};

mod network_drain;
mod local_drain;
mod writer;

use self::network_drain::NetworkDrain;
use self::local_drain::LocalDrain;

thread_local! {
  pub static METRICS: RefCell<Aggregator> = RefCell::new(Aggregator::new(String::from("sozu")));
}

#[derive(Debug,Clone,PartialEq)]
pub enum MetricData {
  Gauge(usize),
  GaugeAdd(i64),
  Count(i64),
  Time(usize),
}

impl MetricData {
  fn is_time(&self) -> bool {
    match self {
      &MetricData::Time(_) => true,
      _ => false,
    }
  }

  fn update(&mut self, key: &'static str, m: MetricData) -> bool {
    match (self, m) {
      (&mut MetricData::Gauge(ref mut v1), MetricData::Gauge(v2)) => {
        let changed = *v1 != v2;
        *v1 = v2;
        changed
      },
      (&mut MetricData::Gauge(ref mut v1), MetricData::GaugeAdd(v2)) => {
        debug_assert!(*v1 as i64 + v2 >= 0, "metric {} underflow: previous value: {}, adding: {}", key, v1, v2);
        let changed = v2 != 0;
        let res = *v1 as i64 + v2;
        *v1 = if res >= 0 {
          res as usize
        } else {
          error!("metric {} underflow: previous value: {}, adding: {}", key, v1, v2);
          0
        };

        changed
      },
      (&mut MetricData::Count(ref mut v1), MetricData::Count(v2)) => {
        let changed = v2 != 0;
        *v1 += v2;
        changed
      },
      (s,m) => panic!("tried to update metric {} of value {:?} with an incompatible metric: {:?}", key, s, m)
    }
  }
}

#[derive(Debug,Clone)]
pub struct StoredMetricData {
  last_sent: Instant,
  updated:   bool,
  data:      MetricData,
}

impl StoredMetricData {
  pub fn new(last_sent: Instant, data: MetricData) -> StoredMetricData {
    StoredMetricData {
      last_sent,
      updated: true,
      data: if let MetricData::GaugeAdd(v) = data {
        if v >= 0 {
          MetricData::Gauge(v as usize)
        } else {
          MetricData::Gauge(0)
        }
      } else { data }
    }
  }

  pub fn update(&mut self, key: &'static str, m: MetricData) {
    let updated = self.data.update(key, m);
    if ! self.updated {
      self.updated = updated;
    }
  }
}

pub fn setup<O: Into<String>>(metrics_host: &SocketAddr, origin: O, use_tagged_metrics: bool, prefix: Option<String>) {
  let metrics_socket = udp_bind();

  debug!("setting up metrics: local address = {:#?}", metrics_socket.local_addr());

  METRICS.with(|metrics| {
    if let Some(p) = prefix {
      (*metrics.borrow_mut()).set_up_prefix(p);
    }
    (*metrics.borrow_mut()).set_up_remote(metrics_socket, metrics_host.clone());
    (*metrics.borrow_mut()).set_up_origin(origin.into());
    (*metrics.borrow_mut()).set_up_tagged_metrics(use_tagged_metrics);
  });
}

pub trait Subscriber {
  fn receive_metric(&mut self, label: &'static str, cluster_id: Option<&str>, backend_id: Option<&str>, metric: MetricData);
}

pub struct Aggregator {
  prefix:  String,
  network: Option<NetworkDrain>,
  local:   LocalDrain,
}

impl Aggregator {
  pub fn new(prefix: String) -> Aggregator {
    Aggregator {
      prefix: prefix.clone(),
      network: None,
      local: LocalDrain::new(prefix),
    }
  }

  pub fn set_up_prefix(&mut self, prefix: String) {
    self.prefix = prefix;
  }

  pub fn set_up_remote(&mut self, socket: UdpSocket, addr: SocketAddr) {
    self.network = Some(NetworkDrain::new(self.prefix.clone(), socket, addr));
  }

  pub fn set_up_origin(&mut self, origin: String) {
    self.network.as_mut().map(|n| n.origin = origin);
  }

  pub fn set_up_tagged_metrics(&mut self, tagged: bool) {
    self.network.as_mut().map(|n| n.use_tagged_metrics = tagged);
  }

  pub fn socket(&self) -> Option<&UdpSocket> {
    self.network.as_ref().map(|n| &n.remote.get_ref().socket)
  }

  pub fn socket_mut(&mut self) -> Option<&mut UdpSocket> {
    self.network.as_mut().map(|n| &mut n.remote.get_mut().socket)
  }

  pub fn count_add(&mut self, key: &'static str, count_value: i64) {
    self.receive_metric(key, None, None, MetricData::Count(count_value));
  }

  pub fn set_gauge(&mut self, key: &'static str, gauge_value: usize) {
    self.receive_metric(key, None, None, MetricData::Gauge(gauge_value));
  }

  pub fn gauge_add(&mut self, key: &'static str, gauge_value: i64) {
    self.receive_metric(key, None, None, MetricData::GaugeAdd(gauge_value));
  }

  pub fn writable(&mut self) {
    if let Some(ref mut net) = self.network.as_mut() {
      net.writable();
    }
  }

  pub fn send_data(&mut self) {
    if let Some(ref mut net) = self.network.as_mut() {
      net.send_data();
    }
  }

  pub fn dump_metrics_data(&mut self) -> MetricsData {
    self.local.dump_metrics_data()
  }

  pub fn dump_process_data(&mut self) -> BTreeMap<String, FilteredData> {
    self.local.dump_process_data()
  }

  pub fn clear_local(&mut self, now: time::OffsetDateTime) {
    self.local.clear(now);
  }
}

impl Subscriber for Aggregator {
  fn receive_metric(&mut self, label: &'static str, cluster_id: Option<&str>, backend_id: Option<&str>, metric: MetricData) {
    if let Some(ref mut net) = self.network.as_mut() {
      net.receive_metric(label, cluster_id, backend_id, metric.clone());
    }
    self.local.receive_metric(label, cluster_id, backend_id, metric.clone());
  }
}

#[derive(Debug,Clone,PartialEq)]
pub struct MetricLine {
  label:      &'static str,
  cluster_id:     Option<String>,
  backend_id: Option<String>,
  /// in milliseconds
  duration:   usize,
}

pub struct MetricSocket {
  pub addr:   SocketAddr,
  pub socket: UdpSocket,
}

impl Write for MetricSocket {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
    self.socket.send_to(buf, self.addr)
  }

  fn flush(&mut self) -> io::Result<()> {
    Ok(())
  }
}

pub fn udp_bind() -> UdpSocket {
  UdpSocket::bind("0.0.0.0:0".parse().unwrap()).expect("could not parse address")
}

#[macro_export]
macro_rules! count (
  ($key:expr, $value: expr) => ({
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).count_add($key, v);
    });
  })
);

#[macro_export]
macro_rules! incr (
  ($key:expr) => (count!($key, 1));
  ($key:expr, $cluster_id:expr, $backend_id:expr) => {
    use $crate::metrics::Subscriber;

    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).receive_metric($key, $cluster_id, $backend_id, $crate::metrics::MetricData::Count(1));
    });
  }
);

#[macro_export]
macro_rules! decr (
  ($key:expr) => (count!($key, -1))
);

#[macro_export]
macro_rules! gauge (
  ($key:expr, $value: expr) => ({
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).set_gauge($key, v);
    });
  })
);

#[macro_export]
macro_rules! gauge_add (
  ($key:expr, $value: expr) => ({
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).gauge_add($key, v);
    });
  });
  ($key:expr, $value:expr, $cluster_id:expr, $backend_id:expr) => {
    use $crate::metrics::Subscriber;
    let v = $value;

    $crate::metrics::METRICS.with(|metrics| {
      (*metrics.borrow_mut()).receive_metric($key, $cluster_id, $backend_id, $crate::metrics::MetricData::GaugeAdd(v));
    });
  }
);

#[macro_export]
macro_rules! time (
  ($key:expr, $value: expr) => ({
    use $crate::metrics::{MetricData,Subscriber};
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      let ref mut m = *metrics.borrow_mut();

      m.receive_metric($key, None, None, MetricData::Time(v as usize));
    });
  });
  ($key:expr, $cluster_id:expr, $value: expr) => ({
    use $crate::metrics::{MetricData,Subscriber};
    let v = $value;
    $crate::metrics::METRICS.with(|metrics| {
      let ref mut m = *metrics.borrow_mut();
      let cluster: &str = $cluster_id;

      m.receive_metric($key, Some(cluster), None, MetricData::Time(v as usize));
    });
  })
);

#[macro_export]
macro_rules! record_backend_metrics (
  ($cluster_id:expr, $backend_id:expr, $response_time: expr, $backend_connection_time: expr, $bin: expr, $bout: expr) => {
    use $crate::metrics::{MetricData,Subscriber};
    $crate::metrics::METRICS.with(|metrics| {
      let ref mut m = *metrics.borrow_mut();
      let cluster_id: &str = $cluster_id.as_str();
      let backend_id: &str = $backend_id;

      m.receive_metric("bytes_in", Some(cluster_id), Some(backend_id), MetricData::Count($bin as i64));
      m.receive_metric("bytes_out", Some(cluster_id), Some(backend_id), MetricData::Count($bout as i64));
      m.receive_metric("backend_response_time", Some(cluster_id), Some(backend_id), MetricData::Time($response_time as usize));
      if let Some(t) = $backend_connection_time {
        m.receive_metric("backend_connection_time", Some(cluster_id), Some(backend_id), MetricData::Time(t.whole_milliseconds() as usize));
      }

      m.receive_metric("requests", Some(cluster_id), Some(backend_id), MetricData::Count(1));
    });
  }
);

<|endoftext|>


<commit_before>use futures::{Poll, Future, Async, BoxFuture};
use tokio_service::Service;
use store::Store;
use {LogPos, Error};
use std::net::SocketAddr;
use std::io;
use std::mem;
use replica::{HostConfig, ReplicaRequest, ReplicaResponse, LogEntry};
use spki_sexp as sexp;

enum ReplicatorState<S: Store, D>
    where D: DownstreamService
{
    Idle,
    Fetching(S::FetchFut),
    Forwarding(D::Future),
}

pub trait DownstreamService {
    type Future: Future<Item = ReplicaResponse, Error = io::Error>;
}

impl<D> DownstreamService for D
    where D: Service<Request = (SocketAddr, ReplicaRequest),
                     Response = ReplicaResponse,
                     Error = io::Error,
                     Future = BoxFuture<ReplicaResponse, io::Error>>
{
    type Future = D::Future;
}

pub struct Replicator<S: Store, D: DownstreamService> {
    store: S,
    downstream: D,
    last_seen_seq: LogPos,
    downstream_addr: Option<SocketAddr>,
    state: ReplicatorState<S, D>,
}

impl<S: Store, D: DownstreamService> Replicator<S, D> {
    pub fn new(store: S, downstream: D) -> Self {
        Replicator {
            store: store,
            downstream: downstream,
            last_seen_seq: LogPos::zero(),
            downstream_addr: None,
            state: ReplicatorState::Idle,
        }
    }
}

impl<S: Store, D: DownstreamService> Drop for Replicator<S, D> {
    fn drop(&mut self) {
        debug!("Drop Replicator");
    }
}

impl<S: Store, D> Future for Replicator<S, D>
    where S: Store,
          D: Service<Request = (SocketAddr, ReplicaRequest),
                     Response = ReplicaResponse,
                     Error = io::Error,
                     Future = BoxFuture<ReplicaResponse, io::Error>>
{
    type Item = ();
    type Error = Error;
    fn poll(&mut self) -> Poll<(), Error> {
        debug!("Replicator#poll");
        loop {
            match mem::replace(&mut self.state, ReplicatorState::Idle) {
                ReplicatorState::Idle => {
                    debug!("Idle, fetching after: {:?}", self.last_seen_seq);
                    let fetch_f = self.store.fetch_next(self.last_seen_seq);
                    self.state = ReplicatorState::Fetching(fetch_f);
                }
                // Now, what have we learnt from this, children? Surely it's
                // that because we're forwarding LogEntries at least, we
                // should make it easy to get a LogEntry into/out of the
                // store, rather than having to futz around with serialisation
                // rubbish.
                ReplicatorState::Fetching(mut fetch_f) => {
                    debug!("Fetching");
                    match try!(fetch_f.poll()) {
                        Async::NotReady => {
                            self.state = ReplicatorState::Fetching(fetch_f);
                            return Ok(Async::NotReady);
                        }
                        Async::Ready((off, val)) => {
                            debug!("Fetched: {:?}", (&off, &val));
                            if let &LogEntry::Config(ref conf) = &val {
                                // let conf: HostConfig<SocketAddr> = try!(sexp::from_bytes(&val));
                                debug!("logged Config message: {:?}", conf);
                                // FIXME: Well, this is blatantly wrong.
                                self.downstream_addr = Some(conf.head);
                            };

                            if let Some(addr) = self.downstream_addr {
                                debug!("Forward {:?} to: {:?}", off, addr);
                                // XXX: Symmetry with ReplClient?
                                let req = ReplicaRequest::AppendLogEntry {
                                    assumed_offset: self.last_seen_seq,
                                    entry_offset: off,
                                    datum: val,
                                };
                                let f = self.downstream.call((addr, req));
                                self.state = ReplicatorState::Forwarding(f);
                            } else {
                                debug!("No downstream at {:?}", off);
                            }
                            self.last_seen_seq = off;
                        }
                    }
                }
                ReplicatorState::Forwarding(_) => unimplemented!(),
            }
        }
    }
}

#[cfg(test)]
mod test {
    use futures::{self, Future, stream, Poll, BoxFuture, Async};
    use futures::stream::Stream;
    use service::Service;
    use tokio::channel;
    use replica::LogPos;
    use store::{RamStore, Store};
    use replica::{ReplicaRequest, ReplicaResponse, LogEntry, HostConfig};
    use std::net::SocketAddr;
    use super::*;
    use spki_sexp as sexp;
    use tokio::reactor::{Core, Handle};
    use std::sync::Arc;
    use std::time::Duration;
    use std::io;
    use tokio_timer;
    use env_logger;

    #[test]
    #[ignore] // WIP
    fn should_start_replicating_to_downstream_on_committed_config_message() {
        env_logger::init().unwrap_or(());
        let store = RamStore::new();
        let timer = tokio_timer::wheel().tick_duration(Duration::from_millis(1)).build();

        let mut core = Core::new().expect("core::new");
        let (tx, rx) = channel::channel(&core.handle()).expect("channel");

        struct MyMagicalDownstream {
            send: channel::Sender<((SocketAddr, ReplicaRequest),
                                   futures::Complete<ReplicaResponse>)>,
        }

        impl Service for MyMagicalDownstream {
            type Request = (SocketAddr, ReplicaRequest);
            type Response = ReplicaResponse;
            type Error = io::Error;
            type Future = BoxFuture<ReplicaResponse, io::Error>;
            fn poll_ready(&self) -> Async<()> {
                Async::Ready(())
            }
            fn call(&self, req: Self::Request) -> Self::Future {
                let (c, p) = futures::oneshot();
                match self.send.send((req, c)) {
                    Ok(()) => (),
                    Err(e) => return futures::failed(e).boxed(),
                };
                p.map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "broken pipe"))
                    .boxed()
            }
        }
        let downstream = MyMagicalDownstream { send: tx };

        let replica = Replicator::new(store.clone(), downstream);

        core.handle().spawn(replica.map_err(|e| panic!("Replicator failed!: {:?}", e)));
        debug!("Spawned replica task");

        let off = LogPos::zero();

        let head_addr = "1.2.3.4:5".parse().expect("parse");

        let config: HostConfig<::std::net::SocketAddr> = HostConfig {
            head: head_addr,
            tail: "1.2.3.6:7".parse().expect("parse"),
        };

        debug!("append config message");
        core.run(store.append_entry(off, off.next(), LogEntry::Config(config)))
            .expect("append");

        let f = timer.timeout(rx.into_future().map_err(|(e, _)| e),
                              Duration::from_millis(1000));
        debug!("wait for stub message");
        let (data, rx) = core.run(f).expect("receive downstream");
        let ((addr, msg), completer) = data.expect("Some message");
        println!("Some: {:?}", (addr, msg));
        // Expect that the target at `head_addr` receives a set of replication messages.

        assert!(false)
    }
}
<commit_msg>Show that we can think about contemplating the possiblity of maybe considering trying to replicate something.<commit_after>use futures::{Poll, Future, Async, BoxFuture};
use tokio_service::Service;
use store::Store;
use {LogPos, Error};
use std::net::SocketAddr;
use std::io;
use std::mem;
use replica::{HostConfig, ReplicaRequest, ReplicaResponse, LogEntry};
use spki_sexp as sexp;

enum ReplicatorState<S: Store, D>
    where D: DownstreamService
{
    Idle,
    Fetching(S::FetchFut),
    Forwarding(D::Future),
}

pub trait DownstreamService {
    type Future: Future<Item = ReplicaResponse, Error = io::Error>;
}

impl<D> DownstreamService for D
    where D: Service<Request = (SocketAddr, ReplicaRequest),
                     Response = ReplicaResponse,
                     Error = io::Error,
                     Future = BoxFuture<ReplicaResponse, io::Error>>
{
    type Future = D::Future;
}

pub struct Replicator<S: Store, D: DownstreamService> {
    store: S,
    downstream: D,
    last_seen_seq: LogPos,
    downstream_addr: Option<SocketAddr>,
    state: ReplicatorState<S, D>,
}

impl<S: Store, D: DownstreamService> Replicator<S, D> {
    pub fn new(store: S, downstream: D) -> Self {
        Replicator {
            store: store,
            downstream: downstream,
            last_seen_seq: LogPos::zero(),
            downstream_addr: None,
            state: ReplicatorState::Idle,
        }
    }
}

impl<S: Store, D: DownstreamService> Drop for Replicator<S, D> {
    fn drop(&mut self) {
        debug!("Drop Replicator");
    }
}

impl<S: Store, D> Future for Replicator<S, D>
    where S: Store,
          D: Service<Request = (SocketAddr, ReplicaRequest),
                     Response = ReplicaResponse,
                     Error = io::Error,
                     Future = BoxFuture<ReplicaResponse, io::Error>>
{
    type Item = ();
    type Error = Error;
    fn poll(&mut self) -> Poll<(), Error> {
        debug!("Replicator#poll");
        loop {
            match mem::replace(&mut self.state, ReplicatorState::Idle) {
                ReplicatorState::Idle => {
                    debug!("Idle, fetching after: {:?}", self.last_seen_seq);
                    let fetch_f = self.store.fetch_next(self.last_seen_seq);
                    self.state = ReplicatorState::Fetching(fetch_f);
                }
                // Now, what have we learnt from this, children? Surely it's
                // that because we're forwarding LogEntries at least, we
                // should make it easy to get a LogEntry into/out of the
                // store, rather than having to futz around with serialisation
                // rubbish.
                ReplicatorState::Fetching(mut fetch_f) => {
                    debug!("Fetching");
                    match try!(fetch_f.poll()) {
                        Async::NotReady => {
                            self.state = ReplicatorState::Fetching(fetch_f);
                            return Ok(Async::NotReady);
                        }
                        Async::Ready((off, val)) => {
                            debug!("Fetched: {:?}", (&off, &val));
                            if let &LogEntry::Config(ref conf) = &val {
                                debug!("logged Config message: {:?}", conf);
                                // FIXME: Well, this is blatantly wrong.
                                self.downstream_addr = Some(conf.head);
                            };

                            if let Some(addr) = self.downstream_addr {
                                debug!("Forward {:?} to: {:?}", off, addr);
                                // XXX: Symmetry with ReplClient?
                                let req = ReplicaRequest::AppendLogEntry {
                                    assumed_offset: self.last_seen_seq,
                                    entry_offset: off,
                                    datum: val,
                                };
                                let f = self.downstream.call((addr, req));
                                self.state = ReplicatorState::Forwarding(f);
                            } else {
                                debug!("No downstream at {:?}", off);
                            }
                            self.last_seen_seq = off;
                        }
                    }
                }
                ReplicatorState::Forwarding(mut f) => {
                    match try!(f.poll()) {
                        Async::NotReady => {
                            self.state = ReplicatorState::Forwarding(f);
                            return Ok(Async::NotReady);
                        }
                        Async::Ready(ReplicaResponse::Done(pos)) => {
                            debug!("woo! {:?}", pos);
                        }
                        Async::Ready(ReplicaResponse::BadSequence(pos)) => {
                            panic!("ARGH! {:?}", pos);
                        }
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod test {
    use futures::{self, Future, stream, Poll, BoxFuture, Async};
    use futures::stream::Stream;
    use service::Service;
    use tokio::channel;
    use replica::LogPos;
    use store::{RamStore, Store};
    use replica::{ReplicaRequest, ReplicaResponse, LogEntry, HostConfig};
    use std::net::SocketAddr;
    use super::*;
    use spki_sexp as sexp;
    use tokio::reactor::{Core, Handle};
    use std::sync::Arc;
    use std::time::Duration;
    use std::io;
    use tokio_timer;
    use env_logger;

    struct MyMagicalDownstream {
        send: channel::Sender<((SocketAddr, ReplicaRequest), futures::Complete<ReplicaResponse>)>,
    }

    impl Service for MyMagicalDownstream {
        type Request = (SocketAddr, ReplicaRequest);
        type Response = ReplicaResponse;
        type Error = io::Error;
        type Future = BoxFuture<ReplicaResponse, io::Error>;
        fn poll_ready(&self) -> Async<()> {
            Async::Ready(())
        }
        fn call(&self, req: Self::Request) -> Self::Future {
            let (c, p) = futures::oneshot();
            match self.send.send((req, c)) {
                Ok(()) => (),
                Err(e) => return futures::failed(e).boxed(),
            };
            p.map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "broken pipe"))
                .boxed()
        }
    }

    #[test]
    // #[ignore] // WIP
    fn should_start_replicating_to_downstream_on_committed_config_message() {
        env_logger::init().unwrap_or(());
        let store = RamStore::new();
        let timer = tokio_timer::wheel().tick_duration(Duration::from_millis(1)).build();

        let mut core = Core::new().expect("core::new");
        let (tx, rx) = channel::channel(&core.handle()).expect("channel");

        let downstream = MyMagicalDownstream { send: tx };

        let replica = Replicator::new(store.clone(), downstream);

        core.handle().spawn(replica.map_err(|e| panic!("Replicator failed!: {:?}", e)));
        debug!("Spawned replica task");

        let off = LogPos::zero();

        let head_addr = "1.2.3.4:5".parse().expect("parse");

        let config: HostConfig<::std::net::SocketAddr> = HostConfig {
            head: head_addr,
            tail: "1.2.3.6:7".parse().expect("parse"),
        };

        debug!("append config message");
        core.run(store.append_entry(off, off.next(), LogEntry::Config(config.clone())))
            .expect("append");

        let f = timer.timeout(rx.into_future().map_err(|(e, _)| e),
                              Duration::from_millis(1000));
        debug!("wait for stub message");
        let (data, rx) = core.run(f).expect("receive downstream");
        let ((addr, msg), completer) = data.expect("Some message");
        println!("Some: {:?}", (&addr, &msg));
        // Expect that the target at `head_addr` receives a set of replication messages.
        let (assumed, off, entry) = match msg {
            ReplicaRequest::AppendLogEntry { assumed_offset, entry_offset, datum } => {
                (assumed_offset, entry_offset, datum)
            }
        };
        assert_eq!((addr, entry), (head_addr, LogEntry::Config(config)));
    }
}
<|endoftext|>


<commit_before>use lsp_types::{
    DidChangeTextDocumentParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams,
    DocumentFilter, DocumentHighlight, DocumentHighlightParams, FoldingRange, FoldingRangeParams,
    GotoDefinitionParams, GotoDefinitionResponse, InitializeParams, InitializeResult,
    InitializedParams, Location, MessageType, ReferenceParams, Registration, SemanticTokens,
    SemanticTokensFullOptions, SemanticTokensOptions, SemanticTokensParams,
    SemanticTokensRegistrationOptions, SemanticTokensResult, SemanticTokensServerCapabilities,
    ServerCapabilities, ServerInfo, StaticRegistrationOptions,
    TextDocumentChangeRegistrationOptions, TextDocumentContentChangeEvent,
    TextDocumentRegistrationOptions, Url, WorkDoneProgressOptions,
};
use lspower::{jsonrpc, Client, LanguageServer};
use tokio::sync::Mutex;

use crate::{
    compiler::{ast_to_hir::AstToHir, cst_to_ast::CstToAst, string_to_cst::StringToCst},
    database::PROJECT_DIRECTORY,
    input::{Input, InputDb},
    language_server::hints::HintsDb,
    Database,
};

use self::{
    definition::find_definition,
    folding_range::FoldingRangeDb,
    hints::HintsNotification,
    references::{find_document_highlights, find_references},
    semantic_tokens::SemanticTokenDb,
    utils::{line_start_utf8_byte_offsets_raw, offset_from_lsp_raw},
};

pub mod definition;
pub mod folding_range;
pub mod hints;
pub mod references;
pub mod semantic_tokens;
pub mod utils;

pub struct CandyLanguageServer {
    pub client: Client,
    pub db: Mutex<Database>,
}
impl CandyLanguageServer {
    pub fn from_client(client: Client) -> Self {
        Self {
            client,
            db: Mutex::new(Database::default()),
        }
    }
}

#[lspower::async_trait]
impl LanguageServer for CandyLanguageServer {
    async fn initialize(&self, params: InitializeParams) -> jsonrpc::Result<InitializeResult> {
        log::info!("LSP: initialize");
        self.client
            .log_message(MessageType::INFO, "Initializing!")
            .await;

        let first_workspace_folder = params
            .workspace_folders
            .unwrap()
            .first()
            .unwrap()
            .uri
            .clone();
        *PROJECT_DIRECTORY.lock().unwrap() = match first_workspace_folder.scheme() {
            "file" => Some(first_workspace_folder.to_file_path().unwrap()),
            _ => panic!("Workspace folder must be a file URI."),
        };

        Ok(InitializeResult {
            // We only support dynamic registration for now.
            capabilities: ServerCapabilities::default(),
            server_info: Some(ServerInfo {
                name: "🍭 Candy Language Server".to_owned(),
                version: None,
            }),
        })
    }

    async fn initialized(&self, _: InitializedParams) {
        log::info!("LSP: initialized");
        let candy_files = vec![
            DocumentFilter {
                language: Some("candy".to_owned()),
                scheme: Some("file".to_owned()),
                pattern: None,
            },
            DocumentFilter {
                language: Some("candy".to_owned()),
                scheme: Some("untitled".to_owned()),
                pattern: None,
            },
        ];
        let text_document_registration_options = TextDocumentRegistrationOptions {
            document_selector: Some(candy_files.clone()),
        };
        self.client
            .register_capability(vec![
                Registration {
                    id: "0".to_owned(),
                    method: "textDocument/didOpen".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "1".to_owned(),
                    method: "textDocument/didOpen".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "2".to_owned(),
                    method: "textDocument/didChange".to_owned(),
                    register_options: Some(
                        serde_json::to_value(TextDocumentChangeRegistrationOptions {
                            document_selector: Some(candy_files),
                            sync_kind: 2, // incremental
                        })
                        .unwrap(),
                    ),
                },
                Registration {
                    id: "3".to_owned(),
                    method: "textDocument/definition".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "4".to_owned(),
                    method: "textDocument/references".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "5".to_owned(),
                    method: "textDocument/documentHighlight".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "6".to_owned(),
                    method: "textDocument/foldingRange".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "7".to_owned(),
                    method: "textDocument/semanticTokens".to_owned(),
                    register_options: Some(
                        serde_json::to_value(
                            SemanticTokensServerCapabilities::SemanticTokensRegistrationOptions(
                                SemanticTokensRegistrationOptions {
                                    text_document_registration_options,
                                    semantic_tokens_options: SemanticTokensOptions {
                                        work_done_progress_options: WorkDoneProgressOptions {
                                            work_done_progress: None,
                                        },
                                        legend: semantic_tokens::LEGEND.clone(),
                                        range: Some(false),
                                        full: Some(SemanticTokensFullOptions::Bool(true)),
                                    },
                                    static_registration_options: StaticRegistrationOptions {
                                        id: None,
                                    },
                                },
                            ),
                        )
                        .unwrap(),
                    ),
                },
            ])
            .await
            .expect("Dynamic capability registration failed.");
        self.client
            .log_message(MessageType::INFO, "server initialized!")
            .await;
    }

    async fn shutdown(&self) -> jsonrpc::Result<()> {
        Ok(())
    }

    async fn did_open(&self, params: DidOpenTextDocumentParams) {
        let input = params.text_document.uri.into();
        {
            let mut db = self.db.lock().await;
            db.did_open_input(&input, params.text_document.text);
        }
        self.analyze_files(vec![input]).await;
    }
    async fn did_change(&self, params: DidChangeTextDocumentParams) {
        let input: Input = params.text_document.uri.into();
        let mut open_inputs = Vec::<Input>::new();
        {
            let mut db = self.db.lock().await;
            let text = apply_text_changes(&db, input.clone(), params.content_changes);
            db.did_change_input(&input, text);
            open_inputs.extend(db.open_inputs.keys().cloned());
        }
        self.analyze_files(open_inputs).await;
    }
    async fn did_close(&self, params: DidCloseTextDocumentParams) {
        let input = params.text_document.uri.into();
        let mut db = self.db.lock().await;
        db.did_close_input(&input);
    }

    async fn goto_definition(
        &self,
        params: GotoDefinitionParams,
    ) -> jsonrpc::Result<Option<GotoDefinitionResponse>> {
        let db = self.db.lock().await;
        Ok(find_definition(&db, params))
    }

    async fn references(&self, params: ReferenceParams) -> jsonrpc::Result<Option<Vec<Location>>> {
        let db = self.db.lock().await;
        Ok(find_references(&db, params))
    }
    async fn document_highlight(
        &self,
        params: DocumentHighlightParams,
    ) -> jsonrpc::Result<Option<Vec<DocumentHighlight>>> {
        let db = self.db.lock().await;
        Ok(find_document_highlights(&db, params))
    }

    async fn folding_range(
        &self,
        params: FoldingRangeParams,
    ) -> jsonrpc::Result<Option<Vec<FoldingRange>>> {
        let db = self.db.lock().await;
        let ranges = db.folding_ranges(params.text_document.uri.into());
        Ok(Some(ranges))
    }

    async fn semantic_tokens_full(
        &self,
        params: SemanticTokensParams,
    ) -> jsonrpc::Result<Option<SemanticTokensResult>> {
        let db = self.db.lock().await;
        let tokens = db.semantic_tokens(params.text_document.uri.into());
        Ok(Some(SemanticTokensResult::Tokens(SemanticTokens {
            result_id: None,
            data: tokens,
        })))
    }
}

impl CandyLanguageServer {
    async fn analyze_files(&self, inputs: Vec<Input>) {
        log::debug!("Analyzing file(s) {:?}", inputs);
        let db = self.db.lock().await;

        for input in inputs {
            let (_, cst_errors) = db.cst_raw(input.clone()).unwrap();
            let (_, _, ast_errors) = db.ast_raw(input.clone()).unwrap();
            let (_, _, hir_errors) = db.hir_raw(input.clone()).unwrap();

            let diagnostics = cst_errors
                .into_iter()
                .chain(ast_errors.into_iter())
                .chain(hir_errors.into_iter())
                .map(|it| it.to_diagnostic(&db, input.clone()))
                .collect();
            self.client
                .publish_diagnostics(input.clone().into(), diagnostics, None)
                .await;
            let hints = db.hints(input.clone());
            self.client
                .send_custom_notification::<HintsNotification>(HintsNotification {
                    uri: Url::from(input).to_string(),
                    hints,
                })
                .await;
        }
    }
}

fn apply_text_changes(
    db: &Database,
    input: Input,
    changes: Vec<TextDocumentContentChangeEvent>,
) -> String {
    let mut text = db.get_input(input.clone()).unwrap().as_ref().to_owned();
    for change in changes {
        match change.range {
            Some(range) => {
                let line_start_offsets = line_start_utf8_byte_offsets_raw(&text);
                let start = offset_from_lsp_raw(&text, &line_start_offsets[..], range.start);
                let end = offset_from_lsp_raw(&text, &line_start_offsets[..], range.end);
                text = format!("{}{}{}", &text[..start], &change.text, &text[end..]);
            }
            None => text = change.text,
        }
    }
    text
}
<commit_msg>feat(compiler): improve analyzing files log message<commit_after>use itertools::Itertools;
use lsp_types::{
    DidChangeTextDocumentParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams,
    DocumentFilter, DocumentHighlight, DocumentHighlightParams, FoldingRange, FoldingRangeParams,
    GotoDefinitionParams, GotoDefinitionResponse, InitializeParams, InitializeResult,
    InitializedParams, Location, MessageType, ReferenceParams, Registration, SemanticTokens,
    SemanticTokensFullOptions, SemanticTokensOptions, SemanticTokensParams,
    SemanticTokensRegistrationOptions, SemanticTokensResult, SemanticTokensServerCapabilities,
    ServerCapabilities, ServerInfo, StaticRegistrationOptions,
    TextDocumentChangeRegistrationOptions, TextDocumentContentChangeEvent,
    TextDocumentRegistrationOptions, Url, WorkDoneProgressOptions,
};
use lspower::{jsonrpc, Client, LanguageServer};
use tokio::sync::Mutex;

use crate::{
    compiler::{ast_to_hir::AstToHir, cst_to_ast::CstToAst, string_to_cst::StringToCst},
    database::PROJECT_DIRECTORY,
    input::{Input, InputDb},
    language_server::hints::HintsDb,
    Database,
};

use self::{
    definition::find_definition,
    folding_range::FoldingRangeDb,
    hints::HintsNotification,
    references::{find_document_highlights, find_references},
    semantic_tokens::SemanticTokenDb,
    utils::{line_start_utf8_byte_offsets_raw, offset_from_lsp_raw},
};

pub mod definition;
pub mod folding_range;
pub mod hints;
pub mod references;
pub mod semantic_tokens;
pub mod utils;

pub struct CandyLanguageServer {
    pub client: Client,
    pub db: Mutex<Database>,
}
impl CandyLanguageServer {
    pub fn from_client(client: Client) -> Self {
        Self {
            client,
            db: Mutex::new(Database::default()),
        }
    }
}

#[lspower::async_trait]
impl LanguageServer for CandyLanguageServer {
    async fn initialize(&self, params: InitializeParams) -> jsonrpc::Result<InitializeResult> {
        log::info!("LSP: initialize");
        self.client
            .log_message(MessageType::INFO, "Initializing!")
            .await;

        let first_workspace_folder = params
            .workspace_folders
            .unwrap()
            .first()
            .unwrap()
            .uri
            .clone();
        *PROJECT_DIRECTORY.lock().unwrap() = match first_workspace_folder.scheme() {
            "file" => Some(first_workspace_folder.to_file_path().unwrap()),
            _ => panic!("Workspace folder must be a file URI."),
        };

        Ok(InitializeResult {
            // We only support dynamic registration for now.
            capabilities: ServerCapabilities::default(),
            server_info: Some(ServerInfo {
                name: "🍭 Candy Language Server".to_owned(),
                version: None,
            }),
        })
    }

    async fn initialized(&self, _: InitializedParams) {
        log::info!("LSP: initialized");
        let candy_files = vec![
            DocumentFilter {
                language: Some("candy".to_owned()),
                scheme: Some("file".to_owned()),
                pattern: None,
            },
            DocumentFilter {
                language: Some("candy".to_owned()),
                scheme: Some("untitled".to_owned()),
                pattern: None,
            },
        ];
        let text_document_registration_options = TextDocumentRegistrationOptions {
            document_selector: Some(candy_files.clone()),
        };
        self.client
            .register_capability(vec![
                Registration {
                    id: "0".to_owned(),
                    method: "textDocument/didOpen".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "1".to_owned(),
                    method: "textDocument/didOpen".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "2".to_owned(),
                    method: "textDocument/didChange".to_owned(),
                    register_options: Some(
                        serde_json::to_value(TextDocumentChangeRegistrationOptions {
                            document_selector: Some(candy_files),
                            sync_kind: 2, // incremental
                        })
                        .unwrap(),
                    ),
                },
                Registration {
                    id: "3".to_owned(),
                    method: "textDocument/definition".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "4".to_owned(),
                    method: "textDocument/references".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "5".to_owned(),
                    method: "textDocument/documentHighlight".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "6".to_owned(),
                    method: "textDocument/foldingRange".to_owned(),
                    register_options: Some(
                        serde_json::to_value(text_document_registration_options.clone()).unwrap(),
                    ),
                },
                Registration {
                    id: "7".to_owned(),
                    method: "textDocument/semanticTokens".to_owned(),
                    register_options: Some(
                        serde_json::to_value(
                            SemanticTokensServerCapabilities::SemanticTokensRegistrationOptions(
                                SemanticTokensRegistrationOptions {
                                    text_document_registration_options,
                                    semantic_tokens_options: SemanticTokensOptions {
                                        work_done_progress_options: WorkDoneProgressOptions {
                                            work_done_progress: None,
                                        },
                                        legend: semantic_tokens::LEGEND.clone(),
                                        range: Some(false),
                                        full: Some(SemanticTokensFullOptions::Bool(true)),
                                    },
                                    static_registration_options: StaticRegistrationOptions {
                                        id: None,
                                    },
                                },
                            ),
                        )
                        .unwrap(),
                    ),
                },
            ])
            .await
            .expect("Dynamic capability registration failed.");
        self.client
            .log_message(MessageType::INFO, "server initialized!")
            .await;
    }

    async fn shutdown(&self) -> jsonrpc::Result<()> {
        Ok(())
    }

    async fn did_open(&self, params: DidOpenTextDocumentParams) {
        let input = params.text_document.uri.into();
        {
            let mut db = self.db.lock().await;
            db.did_open_input(&input, params.text_document.text);
        }
        self.analyze_files(vec![input]).await;
    }
    async fn did_change(&self, params: DidChangeTextDocumentParams) {
        let input: Input = params.text_document.uri.into();
        let mut open_inputs = Vec::<Input>::new();
        {
            let mut db = self.db.lock().await;
            let text = apply_text_changes(&db, input.clone(), params.content_changes);
            db.did_change_input(&input, text);
            open_inputs.extend(db.open_inputs.keys().cloned());
        }
        self.analyze_files(open_inputs).await;
    }
    async fn did_close(&self, params: DidCloseTextDocumentParams) {
        let input = params.text_document.uri.into();
        let mut db = self.db.lock().await;
        db.did_close_input(&input);
    }

    async fn goto_definition(
        &self,
        params: GotoDefinitionParams,
    ) -> jsonrpc::Result<Option<GotoDefinitionResponse>> {
        let db = self.db.lock().await;
        Ok(find_definition(&db, params))
    }

    async fn references(&self, params: ReferenceParams) -> jsonrpc::Result<Option<Vec<Location>>> {
        let db = self.db.lock().await;
        Ok(find_references(&db, params))
    }
    async fn document_highlight(
        &self,
        params: DocumentHighlightParams,
    ) -> jsonrpc::Result<Option<Vec<DocumentHighlight>>> {
        let db = self.db.lock().await;
        Ok(find_document_highlights(&db, params))
    }

    async fn folding_range(
        &self,
        params: FoldingRangeParams,
    ) -> jsonrpc::Result<Option<Vec<FoldingRange>>> {
        let db = self.db.lock().await;
        let ranges = db.folding_ranges(params.text_document.uri.into());
        Ok(Some(ranges))
    }

    async fn semantic_tokens_full(
        &self,
        params: SemanticTokensParams,
    ) -> jsonrpc::Result<Option<SemanticTokensResult>> {
        let db = self.db.lock().await;
        let tokens = db.semantic_tokens(params.text_document.uri.into());
        Ok(Some(SemanticTokensResult::Tokens(SemanticTokens {
            result_id: None,
            data: tokens,
        })))
    }
}

impl CandyLanguageServer {
    async fn analyze_files(&self, inputs: Vec<Input>) {
        log::debug!("Analyzing file(s) {}", inputs.iter().join(", "));
        let db = self.db.lock().await;

        for input in inputs {
            let (_, cst_errors) = db.cst_raw(input.clone()).unwrap();
            let (_, _, ast_errors) = db.ast_raw(input.clone()).unwrap();
            let (_, _, hir_errors) = db.hir_raw(input.clone()).unwrap();

            let diagnostics = cst_errors
                .into_iter()
                .chain(ast_errors.into_iter())
                .chain(hir_errors.into_iter())
                .map(|it| it.to_diagnostic(&db, input.clone()))
                .collect();
            self.client
                .publish_diagnostics(input.clone().into(), diagnostics, None)
                .await;
            let hints = db.hints(input.clone());
            self.client
                .send_custom_notification::<HintsNotification>(HintsNotification {
                    uri: Url::from(input).to_string(),
                    hints,
                })
                .await;
        }
    }
}

fn apply_text_changes(
    db: &Database,
    input: Input,
    changes: Vec<TextDocumentContentChangeEvent>,
) -> String {
    let mut text = db.get_input(input.clone()).unwrap().as_ref().to_owned();
    for change in changes {
        match change.range {
            Some(range) => {
                let line_start_offsets = line_start_utf8_byte_offsets_raw(&text);
                let start = offset_from_lsp_raw(&text, &line_start_offsets[..], range.start);
                let end = offset_from_lsp_raw(&text, &line_start_offsets[..], range.end);
                text = format!("{}{}{}", &text[..start], &change.text, &text[end..]);
            }
            None => text = change.text,
        }
    }
    text
}
<|endoftext|>


<commit_before>//! Implements single-threaded game tree search.

use basetypes::*;
use chess_move::*;
use tt::*;
use position::Position;

/// Represents a terminated search condition.
pub struct TerminatedSearch;

/// Represents a game tree search.
pub struct Search<'a> {
    tt: &'a TranspositionTable,
    position: Position,
    moves: &'a mut MoveStack,
    moves_starting_ply: usize,
    state_stack: Vec<NodeState>,
    reported_nodes: NodeCount,
    unreported_nodes: NodeCount,
    report_function: &'a mut FnMut(NodeCount) -> bool,
}

impl<'a> Search<'a> {
    /// Creates a new instance.
    ///
    /// `report_function` should be a function that registers the
    /// search progress. It will be called with the number of searched
    /// positions from the beginning of the search to this moment. The
    /// function should return `true` if the search should be
    /// terminated, otherwise it should return `false`.
    pub fn new(root: Position,
               tt: &'a TranspositionTable,
               move_stack: &'a mut MoveStack,
               report_function: &'a mut FnMut(NodeCount) -> bool)
               -> Search<'a> {
        let moves_starting_ply = move_stack.ply();
        Search {
            tt: tt,
            position: root,
            moves: move_stack,
            moves_starting_ply: moves_starting_ply,
            state_stack: Vec::with_capacity(32),
            reported_nodes: 0,
            unreported_nodes: 0,
            report_function: report_function,
        }
    }

    /// Performs a principal variation search (PVS) and returns a
    /// result.
    ///
    /// The alpha-beta algorithm is a significant enhancement to the
    /// minimax search algorithm that eliminates the need to search
    /// large portions of the game tree applying a branch-and-bound
    /// technique. Remarkably, it does this without any potential of
    /// overlooking a better move. If one already has found a quite
    /// good move and search for alternatives, one refutation is
    /// enough to avoid it. No need to look for even stronger
    /// refutations. The algorithm maintains two values, alpha and
    /// beta. They represent the minimum score that the maximizing
    /// player is assured of and the maximum score that the minimizing
    /// player is assured of respectively. PVS is an enhancement to
    /// the alpha-beta search, based on zero (null) window searches of
    /// none PV-nodes, to prove a move is worse or not than an already
    /// safe score from the principal variation. (The principal
    /// variation (PV) is a sequence of moves that the program
    /// considers best and therefore expects to be played.)
    ///
    /// **Important note**: This method may leave un-restored move
    /// lists in the move stack. Call `reset` if you want the move
    /// stack to be restored to the state it had when the search
    /// instance was created.
    pub fn run(&mut self,
               mut alpha: Value, // lower bound
               beta: Value, // upper bound
               depth: u8)
               -> Result<Value, TerminatedSearch> {
        assert!(alpha < beta);

        let entry = self.node_begin();

        // Check if the TT entry gives the result.
        if entry.depth() >= depth {
            let value = entry.value();
            let bound = entry.bound();
            if value >= beta && bound & BOUND_LOWER != 0 {
                self.node_end();
                return Ok(beta);
            }
            if value <= alpha && bound & BOUND_UPPER != 0 {
                self.node_end();
                return Ok(alpha);
            }
            if bound == BOUND_EXACT {
                self.node_end();
                return Ok(value);
            }
        }

        // Initial guests for the final result.
        let mut bound = BOUND_UPPER;
        let mut best_move = Move::invalid();

        if depth == 0 {
            // On leaf nodes, do quiescence search.
            let (value, nodes) = self.position
                                     .evaluate_quiescence(alpha, beta, Some(entry.eval_value()));
            try!(self.report_progress(nodes));

            // See how good this position is.
            if value >= beta {
                alpha = beta;
                bound = BOUND_LOWER;
            } else if value > alpha {
                alpha = value;
                bound = BOUND_EXACT;
            }

        } else {
            // On non-leaf nodes, try moves.
            let mut no_moves_yet = true;
            while let Some(m) = self.do_move() {
                try!(self.report_progress(1));

                // Make a recursive call.
                let value = if no_moves_yet {
                    // The first move we analyze with a fully open window
                    // (alpha, beta). If this happens to be a good move,
                    // it will probably raise `alpha`.
                    no_moves_yet = false;
                    -try!(self.run(-beta, -alpha, depth - 1))
                } else {
                    // For the next moves we first try to prove that they
                    // are not better than our current best move. For this
                    // purpose we analyze them with a null window (alpha,
                    // alpha + 1). This is faster than a full window
                    // search. Only if we are certain that the move is
                    // better than our current best move, we do a
                    // full-window search.
                    match -try!(self.run(-alpha - 1, -alpha, depth - 1)) {
                        x if x <= alpha => x,
                        _ => -try!(self.run(-beta, -alpha, depth - 1)),
                    }
                };
                self.undo_move();

                // See how good this move was.
                if value >= beta {
                    // This move is so good, that the opponent will
                    // probably not allow this line of play to
                    // happen. Therefore we should not lose any more time
                    // on this position.
                    alpha = beta;
                    bound = BOUND_LOWER;
                    best_move = m;
                    break;
                }
                if value > alpha {
                    // We found a new best move.
                    alpha = value;
                    bound = BOUND_EXACT;
                    best_move = m;
                }
            }

            // Check if we are in a final position (no legal moves).
            if no_moves_yet {
                let value = self.position.evaluate_final();
                if value >= beta {
                    alpha = beta;
                    bound = BOUND_LOWER;
                }
                if value > alpha {
                    alpha = value;
                    bound = BOUND_EXACT;
                }
            }
        }

        self.store(alpha, bound, depth, best_move);
        self.node_end();
        Ok(alpha)
    }

    /// Returns the number of searched positions.
    #[inline(always)]
    pub fn node_count(&self) -> NodeCount {
        self.reported_nodes + self.unreported_nodes
    }

    /// Resets the instance to the state it had when it was created.
    #[inline]
    pub fn reset(&mut self) {
        while self.moves.ply() > self.moves_starting_ply {
            self.moves.restore();
        }
        self.state_stack.clear();
        self.reported_nodes = 0;
        self.unreported_nodes = 0;
    }

    // Declares that we are starting to process a new node.
    //
    // Each recursive call to `run` begins with a call to
    // `node_begin`. The returned value is a TT entry telling
    // everything we know about the current position.
    #[inline]
    fn node_begin(&mut self) -> EntryData {
        // Consult the transposition table.
        let entry = if let Some(e) = self.tt.probe(self.position.hash()) {
            e
        } else {
            EntryData::new(0, BOUND_NONE, 0, 0, self.position.evaluate_static())
        };
        self.state_stack.push(NodeState {
            phase: NodePhase::Pristine,
            entry: entry,
            checkers: self.position.board().checkers(),
        });
        entry
    }

    // Declares that we are done processing the current node.
    //
    // Each recursive call to `run` ends with a call to `node_end`.
    #[inline]
    fn node_end(&mut self) {
        if let NodePhase::Pristine = self.state_stack.last().unwrap().phase {
            // For pristine nodes we have not saved the move list
            // yet, so we should not restore it.
        } else {
            self.moves.restore();
        }
        self.state_stack.pop();
    }

    // Plays the next legal move in the current position and returns
    // it.
    //
    // Each call to `do_move` for the same position will play and
    // return a different move. When all legal moves has been played,
    // `None` will be returned. `do_move` will do whatever it can to
    // play the best moves first, and the worst last. It will also try
    // to be efficient, for example it will generate the list of all
    // pseudo-legal moves at the last possible moment.
    #[inline]
    fn do_move(&mut self) -> Option<Move> {
        let state = self.state_stack.last_mut().unwrap();

        // Try the hash move first.
        if let NodePhase::Pristine = state.phase {
            // We save the move list at the last possible moment,
            // because most of the nodes are leafs.
            self.moves.save();

            state.phase = NodePhase::TriedHashMove;
            if state.entry.move16() != 0 {
                if let Some(mut m) = self.position.try_move_digest(state.entry.move16()) {
                    if self.position.do_move(m) {
                        m.set_score(MAX_MOVE_SCORE);
                        return Some(m);
                    }
                }
            }
        }

        // After the hash move, we generate all pseudo-legal
        // moves. But we should not forget to remove the already tried
        // hash move from the list.
        if let NodePhase::TriedHashMove = state.phase {
            // TODO: `generate_moves` needs `_checkers` and `_pinned`
            // to do its work. At this time we might already have
            // generated them during the trying of the best move --
            // use them if possible.
            self.position.generate_moves(self.moves);
            if state.entry.move16() != 0 {
                self.moves.remove_move(state.entry.move16());
            }
            state.phase = NodePhase::GeneratedMoves;
        }

        // Spit out the generated moves.
        while let Some(mut m) = self.moves.remove_best_move() {

            // First, the good captures.
            if let NodePhase::GeneratedMoves = state.phase {
                if m.score() == MAX_MOVE_SCORE {
                    if self.position.evaluate_move(m) >= 0 {
                        if self.position.do_move(m) {
                            return Some(m);
                        }
                        continue;
                    }
                    // This is a bad capture -- push it back to the
                    // move stack.
                    m.set_score(MAX_MOVE_SCORE - 1);
                    self.moves.push(m);
                    continue;
                }
                state.phase = NodePhase::TriedGoodCaptures;
            }

            // Second, the bad captures.
            if let NodePhase::TriedGoodCaptures = state.phase {
                if m.score() == MAX_MOVE_SCORE - 1 {
                    if self.position.do_move(m) {
                        return Some(m);
                    }
                    continue;
                }
                state.phase = NodePhase::TriedBadCaptures;
            }

            // Before trying the quiet moves, we should assign proper
            // move scores to them.
            if let NodePhase::TriedBadCaptures = state.phase {
                // TODO: Assign the moves scores here using the killer
                // move heuristics and the history heuristics.

                // We use the score field (2 bits) to properly order
                // quiet movies. Moves which destination square is
                // more advanced into enemy's territory are tried
                // first. The logic is that those moves are riskier,
                // so if such a move loses material this will be
                // detected early and the search tree will be pruned,
                // but if the move does not lose material, chances are
                // that it is a very good move.
                const SCORE_LOOKUP: [[u32; 8]; 2] = [// white
                                                     [0, 1, 2, 3, 4, 5, 6, 7],
                                                     // black
                                                     [7, 6, 5, 4, 3, 2, 1, 0]];
                for m in self.moves.iter_mut() {
                    let rank = rank(m.dest_square());
                    m.set_score(unsafe {
                        *SCORE_LOOKUP.get_unchecked(self.position.board().to_move())
                                     .get_unchecked(rank)
                    });
                }

                state.phase = NodePhase::SortedQuietMoves;
            }

            // Last, quiet moves.
            if self.position.do_move(m) {
                if state.checkers != 0 {
                    // When in check, we set a high move score to all
                    // moves to avoid search depth reductions.
                    m.set_score(MAX_MOVE_SCORE - 1);
                }
                return Some(m);
            }
        }
        None
    }

    // Takes the last played move back.
    #[inline]
    fn undo_move(&mut self) {
        self.position.undo_move();
    }

    // Stores updated node information in the transposition table.
    #[inline]
    fn store(&mut self, value: Value, bound: BoundType, depth: u8, best_move: Move) {
        let entry = &self.state_stack.last().unwrap().entry;
        let move16 = match best_move.digest() {
            0 => entry.move16(),
            x => x,
        };
        self.tt.store(self.position.hash(),
                      EntryData::new(value, bound, depth, move16, entry.eval_value()));
    }

    // Reports search progress.
    //
    // From time to time, we should report how many nodes had been
    // searched since the beginning of the search. This also gives an
    // opportunity for the search to be terminated.
    #[inline]
    fn report_progress(&mut self, new_nodes: NodeCount) -> Result<(), TerminatedSearch> {
        self.unreported_nodes += new_nodes;
        if self.unreported_nodes >= super::NODE_COUNT_REPORT_INTERVAL {
            self.reported_nodes += self.unreported_nodes;
            self.unreported_nodes = 0;
            if (*self.report_function)(self.reported_nodes) {
                return Err(TerminatedSearch);
            }
        }
        Ok(())
    }
}

enum NodePhase {
    Pristine,
    TriedHashMove,
    GeneratedMoves,
    TriedGoodCaptures,
    TriedBadCaptures,
    SortedQuietMoves,
}

struct NodeState {
    phase: NodePhase,
    entry: EntryData,
    checkers: u64,
}

#[cfg(test)]
mod tests {
    use super::Search;
    use chess_move::*;
    use tt::*;
    use position::Position;

    #[test]
    fn test_search() {
        let p = Position::from_fen("8/8/8/8/3q3k/7n/6PP/2Q2R1K b - - 0 1").ok().unwrap();
        let tt = TranspositionTable::new();
        let mut moves = MoveStack::new();
        let mut report = |_| false;
        let mut search = Search::new(p, &tt, &mut moves, &mut report);
        let value = search.run(-30000, 30000, 2)
                          .ok()
                          .unwrap();
        assert!(value < -300);
        search.reset();
        let value = search.run(-30000, 30000, 4)
                          .ok()
                          .unwrap();
        assert!(value >= 20000);
    }
}
<commit_msg>fixed a performance regression with b.checkers()<commit_after>//! Implements single-threaded game tree search.

use basetypes::*;
use bitsets::*;
use chess_move::*;
use tt::*;
use position::Position;

/// Represents a terminated search condition.
pub struct TerminatedSearch;

/// Represents a game tree search.
pub struct Search<'a> {
    tt: &'a TranspositionTable,
    position: Position,
    moves: &'a mut MoveStack,
    moves_starting_ply: usize,
    state_stack: Vec<NodeState>,
    reported_nodes: NodeCount,
    unreported_nodes: NodeCount,
    report_function: &'a mut FnMut(NodeCount) -> bool,
}

impl<'a> Search<'a> {
    /// Creates a new instance.
    ///
    /// `report_function` should be a function that registers the
    /// search progress. It will be called with the number of searched
    /// positions from the beginning of the search to this moment. The
    /// function should return `true` if the search should be
    /// terminated, otherwise it should return `false`.
    pub fn new(root: Position,
               tt: &'a TranspositionTable,
               move_stack: &'a mut MoveStack,
               report_function: &'a mut FnMut(NodeCount) -> bool)
               -> Search<'a> {
        let moves_starting_ply = move_stack.ply();
        Search {
            tt: tt,
            position: root,
            moves: move_stack,
            moves_starting_ply: moves_starting_ply,
            state_stack: Vec::with_capacity(32),
            reported_nodes: 0,
            unreported_nodes: 0,
            report_function: report_function,
        }
    }

    /// Performs a principal variation search (PVS) and returns a
    /// result.
    ///
    /// The alpha-beta algorithm is a significant enhancement to the
    /// minimax search algorithm that eliminates the need to search
    /// large portions of the game tree applying a branch-and-bound
    /// technique. Remarkably, it does this without any potential of
    /// overlooking a better move. If one already has found a quite
    /// good move and search for alternatives, one refutation is
    /// enough to avoid it. No need to look for even stronger
    /// refutations. The algorithm maintains two values, alpha and
    /// beta. They represent the minimum score that the maximizing
    /// player is assured of and the maximum score that the minimizing
    /// player is assured of respectively. PVS is an enhancement to
    /// the alpha-beta search, based on zero (null) window searches of
    /// none PV-nodes, to prove a move is worse or not than an already
    /// safe score from the principal variation. (The principal
    /// variation (PV) is a sequence of moves that the program
    /// considers best and therefore expects to be played.)
    ///
    /// **Important note**: This method may leave un-restored move
    /// lists in the move stack. Call `reset` if you want the move
    /// stack to be restored to the state it had when the search
    /// instance was created.
    pub fn run(&mut self,
               mut alpha: Value, // lower bound
               beta: Value, // upper bound
               depth: u8)
               -> Result<Value, TerminatedSearch> {
        assert!(alpha < beta);

        let entry = self.node_begin();

        // Check if the TT entry gives the result.
        if entry.depth() >= depth {
            let value = entry.value();
            let bound = entry.bound();
            if value >= beta && bound & BOUND_LOWER != 0 {
                self.node_end();
                return Ok(beta);
            }
            if value <= alpha && bound & BOUND_UPPER != 0 {
                self.node_end();
                return Ok(alpha);
            }
            if bound == BOUND_EXACT {
                self.node_end();
                return Ok(value);
            }
        }

        // Initial guests for the final result.
        let mut bound = BOUND_UPPER;
        let mut best_move = Move::invalid();

        if depth == 0 {
            // On leaf nodes, do quiescence search.
            let (value, nodes) = self.position
                                     .evaluate_quiescence(alpha, beta, Some(entry.eval_value()));
            try!(self.report_progress(nodes));

            // See how good this position is.
            if value >= beta {
                alpha = beta;
                bound = BOUND_LOWER;
            } else if value > alpha {
                alpha = value;
                bound = BOUND_EXACT;
            }

        } else {
            // On non-leaf nodes, try moves.
            let mut no_moves_yet = true;
            while let Some(m) = self.do_move() {
                try!(self.report_progress(1));

                // Make a recursive call.
                let value = if no_moves_yet {
                    // The first move we analyze with a fully open window
                    // (alpha, beta). If this happens to be a good move,
                    // it will probably raise `alpha`.
                    no_moves_yet = false;
                    -try!(self.run(-beta, -alpha, depth - 1))
                } else {
                    // For the next moves we first try to prove that they
                    // are not better than our current best move. For this
                    // purpose we analyze them with a null window (alpha,
                    // alpha + 1). This is faster than a full window
                    // search. Only if we are certain that the move is
                    // better than our current best move, we do a
                    // full-window search.
                    match -try!(self.run(-alpha - 1, -alpha, depth - 1)) {
                        x if x <= alpha => x,
                        _ => -try!(self.run(-beta, -alpha, depth - 1)),
                    }
                };
                self.undo_move();

                // See how good this move was.
                if value >= beta {
                    // This move is so good, that the opponent will
                    // probably not allow this line of play to
                    // happen. Therefore we should not lose any more time
                    // on this position.
                    alpha = beta;
                    bound = BOUND_LOWER;
                    best_move = m;
                    break;
                }
                if value > alpha {
                    // We found a new best move.
                    alpha = value;
                    bound = BOUND_EXACT;
                    best_move = m;
                }
            }

            // Check if we are in a final position (no legal moves).
            if no_moves_yet {
                let value = self.position.evaluate_final();
                if value >= beta {
                    alpha = beta;
                    bound = BOUND_LOWER;
                }
                if value > alpha {
                    alpha = value;
                    bound = BOUND_EXACT;
                }
            }
        }

        self.store(alpha, bound, depth, best_move);
        self.node_end();
        Ok(alpha)
    }

    /// Returns the number of searched positions.
    #[inline(always)]
    pub fn node_count(&self) -> NodeCount {
        self.reported_nodes + self.unreported_nodes
    }

    /// Resets the instance to the state it had when it was created.
    #[inline]
    pub fn reset(&mut self) {
        while self.moves.ply() > self.moves_starting_ply {
            self.moves.restore();
        }
        self.state_stack.clear();
        self.reported_nodes = 0;
        self.unreported_nodes = 0;
    }

    // Declares that we are starting to process a new node.
    //
    // Each recursive call to `run` begins with a call to
    // `node_begin`. The returned value is a TT entry telling
    // everything we know about the current position.
    #[inline]
    fn node_begin(&mut self) -> EntryData {
        // Consult the transposition table.
        let entry = if let Some(e) = self.tt.probe(self.position.hash()) {
            e
        } else {
            EntryData::new(0, BOUND_NONE, 0, 0, self.position.evaluate_static())
        };
        self.state_stack.push(NodeState {
            phase: NodePhase::Pristine,
            entry: entry,
            checkers: BB_UNIVERSAL_SET,
        });
        entry
    }

    // Declares that we are done processing the current node.
    //
    // Each recursive call to `run` ends with a call to `node_end`.
    #[inline]
    fn node_end(&mut self) {
        if let NodePhase::Pristine = self.state_stack.last().unwrap().phase {
            // For pristine nodes we have not saved the move list
            // yet, so we should not restore it.
        } else {
            self.moves.restore();
        }
        self.state_stack.pop();
    }

    // Plays the next legal move in the current position and returns
    // it.
    //
    // Each call to `do_move` for the same position will play and
    // return a different move. When all legal moves has been played,
    // `None` will be returned. `do_move` will do whatever it can to
    // play the best moves first, and the worst last. It will also try
    // to be efficient, for example it will generate the list of all
    // pseudo-legal moves at the last possible moment.
    #[inline]
    fn do_move(&mut self) -> Option<Move> {
        let state = self.state_stack.last_mut().unwrap();

        // Try the hash move first.
        if let NodePhase::Pristine = state.phase {
            // We save the move list at the last possible moment,
            // because most of the nodes are leafs.
            self.moves.save();

            // We save the checkers bitboard also, because we will
            // need this information later many times, and we do not
            // want to recalculate it needlessly.
            //
            // TODO: Verify if this is the best place to calculate
            // checkers, especially in relation to trying null moves.
            state.checkers = self.position.board().checkers();

            state.phase = NodePhase::TriedHashMove;
            if state.entry.move16() != 0 {
                if let Some(mut m) = self.position.try_move_digest(state.entry.move16()) {
                    if self.position.do_move(m) {
                        m.set_score(MAX_MOVE_SCORE);
                        return Some(m);
                    }
                }
            }
        }

        // After the hash move, we generate all pseudo-legal
        // moves. But we should not forget to remove the already tried
        // hash move from the list.
        if let NodePhase::TriedHashMove = state.phase {
            // TODO: `generate_moves` needs `_checkers` and `_pinned`
            // to do its work. At this time we already have generated
            // them, so we should use them.
            self.position.board()._checkers.set(state.checkers);

            self.position.generate_moves(self.moves);
            if state.entry.move16() != 0 {
                self.moves.remove_move(state.entry.move16());
            }
            state.phase = NodePhase::GeneratedMoves;
        }

        // Spit out the generated moves.
        while let Some(mut m) = self.moves.remove_best_move() {

            // First, the good captures.
            if let NodePhase::GeneratedMoves = state.phase {
                if m.score() == MAX_MOVE_SCORE {
                    if self.position.evaluate_move(m) >= 0 {
                        if self.position.do_move(m) {
                            return Some(m);
                        }
                        continue;
                    }
                    // This is a bad capture -- push it back to the
                    // move stack.
                    m.set_score(MAX_MOVE_SCORE - 1);
                    self.moves.push(m);
                    continue;
                }
                state.phase = NodePhase::TriedGoodCaptures;
            }

            // Second, the bad captures.
            if let NodePhase::TriedGoodCaptures = state.phase {
                if m.score() == MAX_MOVE_SCORE - 1 {
                    if self.position.do_move(m) {
                        return Some(m);
                    }
                    continue;
                }
                state.phase = NodePhase::TriedBadCaptures;
            }

            // Before trying the quiet moves, we should assign proper
            // move scores to them.
            if let NodePhase::TriedBadCaptures = state.phase {
                // TODO: Assign the moves scores here using the killer
                // move heuristics and the history heuristics.

                // We use the score field (2 bits) to properly order
                // quiet movies. Moves which destination square is
                // more advanced into enemy's territory are tried
                // first. The logic is that those moves are riskier,
                // so if such a move loses material this will be
                // detected early and the search tree will be pruned,
                // but if the move does not lose material, chances are
                // that it is a very good move.
                const SCORE_LOOKUP: [[u32; 8]; 2] = [// white
                                                     [0, 1, 2, 3, 4, 5, 6, 7],
                                                     // black
                                                     [7, 6, 5, 4, 3, 2, 1, 0]];
                for m in self.moves.iter_mut() {
                    let rank = rank(m.dest_square());
                    m.set_score(unsafe {
                        *SCORE_LOOKUP.get_unchecked(self.position.board().to_move())
                                     .get_unchecked(rank)
                    });
                }

                state.phase = NodePhase::SortedQuietMoves;
            }

            // Last, quiet moves.
            if self.position.do_move(m) {
                if state.checkers != 0 {
                    // When in check, we set a high move score to all
                    // moves to avoid search depth reductions.
                    m.set_score(MAX_MOVE_SCORE - 1);
                }
                return Some(m);
            }
        }
        None
    }

    // Takes the last played move back.
    #[inline]
    fn undo_move(&mut self) {
        self.position.undo_move();
    }

    // Stores updated node information in the transposition table.
    #[inline]
    fn store(&mut self, value: Value, bound: BoundType, depth: u8, best_move: Move) {
        let entry = &self.state_stack.last().unwrap().entry;
        let move16 = match best_move.digest() {
            0 => entry.move16(),
            x => x,
        };
        self.tt.store(self.position.hash(),
                      EntryData::new(value, bound, depth, move16, entry.eval_value()));
    }

    // Reports search progress.
    //
    // From time to time, we should report how many nodes had been
    // searched since the beginning of the search. This also gives an
    // opportunity for the search to be terminated.
    #[inline]
    fn report_progress(&mut self, new_nodes: NodeCount) -> Result<(), TerminatedSearch> {
        self.unreported_nodes += new_nodes;
        if self.unreported_nodes >= super::NODE_COUNT_REPORT_INTERVAL {
            self.reported_nodes += self.unreported_nodes;
            self.unreported_nodes = 0;
            if (*self.report_function)(self.reported_nodes) {
                return Err(TerminatedSearch);
            }
        }
        Ok(())
    }
}

enum NodePhase {
    Pristine,
    TriedHashMove,
    GeneratedMoves,
    TriedGoodCaptures,
    TriedBadCaptures,
    SortedQuietMoves,
}

struct NodeState {
    phase: NodePhase,
    entry: EntryData,
    checkers: u64,
}

#[cfg(test)]
mod tests {
    use super::Search;
    use chess_move::*;
    use tt::*;
    use position::Position;

    #[test]
    fn test_search() {
        let p = Position::from_fen("8/8/8/8/3q3k/7n/6PP/2Q2R1K b - - 0 1").ok().unwrap();
        let tt = TranspositionTable::new();
        let mut moves = MoveStack::new();
        let mut report = |_| false;
        let mut search = Search::new(p, &tt, &mut moves, &mut report);
        let value = search.run(-30000, 30000, 2)
                          .ok()
                          .unwrap();
        assert!(value < -300);
        search.reset();
        let value = search.run(-30000, 30000, 4)
                          .ok()
                          .unwrap();
        assert!(value >= 20000);
    }
}
<|endoftext|>