<commit_before>use crate::background_jobs::Environment;
use crate::schema;
use anyhow::Context;
use cargo_registry_index::{Crate, Repository};
use chrono::Utc;
use diesel::prelude::*;
use std::fs::{self, OpenOptions};
use std::io::ErrorKind;
use std::process::Command;
use swirl::PerformError;

#[swirl::background_job]
pub fn add_crate(env: &Environment, conn: &PgConnection, krate: Crate) -> Result<(), PerformError> {
    use std::io::prelude::*;

    let repo = env.lock_index()?;
    let dst = repo.index_file(&krate.name);

    // Add the crate to its relevant file
    fs::create_dir_all(dst.parent().unwrap())?;
    let mut file = OpenOptions::new().append(true).create(true).open(&dst)?;
    serde_json::to_writer(&mut file, &krate)?;
    file.write_all(b"\n")?;

    let message: String = format!("Updating crate `{}#{}`", krate.name, krate.vers);
    repo.commit_and_push(&message, &dst)?;

    // Queue another background job to update the http-based index as well.
    update_crate_index(krate.name.clone()).enqueue(conn)?;
    Ok(())
}

#[swirl::background_job]
pub fn update_crate_index(env: &Environment, crate_name: String) -> Result<(), PerformError> {
    let repo = env.lock_index()?;
    let dst = repo.index_file(&crate_name);

    let contents = match fs::read_to_string(dst) {
        Ok(contents) => Some(contents),
        Err(e) if e.kind() == ErrorKind::NotFound => None,
        Err(e) => return Err(e.into()),
    };

    env.uploader
        .sync_index(env.http_client(), &crate_name, contents)?;

    if let Some(cloudfront) = env.cloudfront() {
        trace!(?crate_name, "invalidate CloudFront");
        cloudfront.invalidate(
            env.http_client(),
            &Repository::relative_index_file_for_url(&crate_name),
        )?;
    }

    Ok(())
}

/// Yanks or unyanks a crate version. This requires finding the index
/// file, deserlialise the crate from JSON, change the yank boolean to
/// `true` or `false`, write all the lines back out, and commit and
/// push the changes.
#[swirl::background_job]
pub fn sync_yanked(
    env: &Environment,
    conn: &PgConnection,
    krate: String,
    version_num: String,
) -> Result<(), PerformError> {
    trace!(?krate, ?version_num, "Load yanked status from database");

    let yanked: bool = schema::versions::table
        .inner_join(schema::crates::table)
        .filter(schema::crates::name.eq(&krate))
        .filter(schema::versions::num.eq(&version_num))
        .select(schema::versions::yanked)
        .get_result(conn)
        .context("Failed to load yanked status from database")?;

    trace!(?krate, ?version_num, yanked);

    let repo = env.lock_index()?;
    let dst = repo.index_file(&krate);

    let prev = fs::read_to_string(&dst)?;
    let new = prev
        .lines()
        .map(|line| {
            let mut git_crate = serde_json::from_str::<Crate>(line)
                .map_err(|_| format!("couldn't decode: `{line}`"))?;
            if git_crate.name != krate || git_crate.vers != version_num {
                return Ok(line.to_string());
            }
            git_crate.yanked = Some(yanked);
            Ok(serde_json::to_string(&git_crate)?)
        })
        .collect::<Result<Vec<_>, PerformError>>();
    let new = new?.join("\n") + "\n";

    if new != prev {
        fs::write(&dst, new.as_bytes())?;

        let action = if yanked { "Yanking" } else { "Unyanking" };
        let message = format!("{action} crate `{krate}#{version_num}`");

        repo.commit_and_push(&message, &dst)?;
    } else {
        debug!("Skipping `yanked` update because index is up-to-date");
    }

    // Queue another background job to update the http-based index as well.
    update_crate_index(krate.clone()).enqueue(conn)?;

    Ok(())
}

/// Collapse the index into a single commit, archiving the current history in a snapshot branch.
#[swirl::background_job]
pub fn squash_index(env: &Environment) -> Result<(), PerformError> {
    let repo = env.lock_index()?;
    println!("Squashing the index into a single commit.");

    let now = Utc::now().format("%Y-%m-%d");
    let original_head = repo.head_oid()?.to_string();
    let msg = format!("Collapse index into one commit\n\n\
        Previous HEAD was {}, now on the `snapshot-{}` branch\n\n\
        More information about this change can be found [online] and on [this issue].\n\n\
        [online]: https://internals.rust-lang.org/t/cargos-crate-index-upcoming-squash-into-one-commit/8440\n\
        [this issue]: https://github.com/rust-lang/crates-io-cargo-teams/issues/47", original_head, now);

    repo.squash_to_single_commit(&msg)?;

    // Shell out to git because libgit2 does not currently support push leases

    repo.run_command(Command::new("git").args([
        "push",
        // Both updates should succeed or fail together
        "--atomic",
        "origin",
        // Overwrite master, but only if it server matches the expected value
        &format!("--force-with-lease=refs/heads/master:{original_head}"),
        // The new squashed commit is pushed to master
        "HEAD:refs/heads/master",
        // The previous value of HEAD is pushed to a snapshot branch
        &format!("{original_head}:refs/heads/snapshot-{now}"),
    ]))?;

    println!("The index has been successfully squashed.");

    Ok(())
}
<commit_msg>worker::git: Use `tracing` macros instead of `println!()`<commit_after>use crate::background_jobs::Environment;
use crate::schema;
use anyhow::Context;
use cargo_registry_index::{Crate, Repository};
use chrono::Utc;
use diesel::prelude::*;
use std::fs::{self, OpenOptions};
use std::io::ErrorKind;
use std::process::Command;
use swirl::PerformError;

#[swirl::background_job]
pub fn add_crate(env: &Environment, conn: &PgConnection, krate: Crate) -> Result<(), PerformError> {
    use std::io::prelude::*;

    let repo = env.lock_index()?;
    let dst = repo.index_file(&krate.name);

    // Add the crate to its relevant file
    fs::create_dir_all(dst.parent().unwrap())?;
    let mut file = OpenOptions::new().append(true).create(true).open(&dst)?;
    serde_json::to_writer(&mut file, &krate)?;
    file.write_all(b"\n")?;

    let message: String = format!("Updating crate `{}#{}`", krate.name, krate.vers);
    repo.commit_and_push(&message, &dst)?;

    // Queue another background job to update the http-based index as well.
    update_crate_index(krate.name.clone()).enqueue(conn)?;
    Ok(())
}

#[swirl::background_job]
pub fn update_crate_index(env: &Environment, crate_name: String) -> Result<(), PerformError> {
    let repo = env.lock_index()?;
    let dst = repo.index_file(&crate_name);

    let contents = match fs::read_to_string(dst) {
        Ok(contents) => Some(contents),
        Err(e) if e.kind() == ErrorKind::NotFound => None,
        Err(e) => return Err(e.into()),
    };

    env.uploader
        .sync_index(env.http_client(), &crate_name, contents)?;

    if let Some(cloudfront) = env.cloudfront() {
        trace!(?crate_name, "invalidate CloudFront");
        cloudfront.invalidate(
            env.http_client(),
            &Repository::relative_index_file_for_url(&crate_name),
        )?;
    }

    Ok(())
}

/// Yanks or unyanks a crate version. This requires finding the index
/// file, deserlialise the crate from JSON, change the yank boolean to
/// `true` or `false`, write all the lines back out, and commit and
/// push the changes.
#[swirl::background_job]
pub fn sync_yanked(
    env: &Environment,
    conn: &PgConnection,
    krate: String,
    version_num: String,
) -> Result<(), PerformError> {
    trace!(?krate, ?version_num, "Load yanked status from database");

    let yanked: bool = schema::versions::table
        .inner_join(schema::crates::table)
        .filter(schema::crates::name.eq(&krate))
        .filter(schema::versions::num.eq(&version_num))
        .select(schema::versions::yanked)
        .get_result(conn)
        .context("Failed to load yanked status from database")?;

    trace!(?krate, ?version_num, yanked);

    let repo = env.lock_index()?;
    let dst = repo.index_file(&krate);

    let prev = fs::read_to_string(&dst)?;
    let new = prev
        .lines()
        .map(|line| {
            let mut git_crate = serde_json::from_str::<Crate>(line)
                .map_err(|_| format!("couldn't decode: `{line}`"))?;
            if git_crate.name != krate || git_crate.vers != version_num {
                return Ok(line.to_string());
            }
            git_crate.yanked = Some(yanked);
            Ok(serde_json::to_string(&git_crate)?)
        })
        .collect::<Result<Vec<_>, PerformError>>();
    let new = new?.join("\n") + "\n";

    if new != prev {
        fs::write(&dst, new.as_bytes())?;

        let action = if yanked { "Yanking" } else { "Unyanking" };
        let message = format!("{action} crate `{krate}#{version_num}`");

        repo.commit_and_push(&message, &dst)?;
    } else {
        debug!("Skipping `yanked` update because index is up-to-date");
    }

    // Queue another background job to update the http-based index as well.
    update_crate_index(krate.clone()).enqueue(conn)?;

    Ok(())
}

/// Collapse the index into a single commit, archiving the current history in a snapshot branch.
#[swirl::background_job]
pub fn squash_index(env: &Environment) -> Result<(), PerformError> {
    let repo = env.lock_index()?;
    info!("Squashing the index into a single commit.");

    let now = Utc::now().format("%Y-%m-%d");
    let original_head = repo.head_oid()?.to_string();
    let msg = format!("Collapse index into one commit\n\n\
        Previous HEAD was {}, now on the `snapshot-{}` branch\n\n\
        More information about this change can be found [online] and on [this issue].\n\n\
        [online]: https://internals.rust-lang.org/t/cargos-crate-index-upcoming-squash-into-one-commit/8440\n\
        [this issue]: https://github.com/rust-lang/crates-io-cargo-teams/issues/47", original_head, now);

    repo.squash_to_single_commit(&msg)?;

    // Shell out to git because libgit2 does not currently support push leases

    repo.run_command(Command::new("git").args([
        "push",
        // Both updates should succeed or fail together
        "--atomic",
        "origin",
        // Overwrite master, but only if it server matches the expected value
        &format!("--force-with-lease=refs/heads/master:{original_head}"),
        // The new squashed commit is pushed to master
        "HEAD:refs/heads/master",
        // The previous value of HEAD is pushed to a snapshot branch
        &format!("{original_head}:refs/heads/snapshot-{now}"),
    ]))?;

    info!("The index has been successfully squashed.");

    Ok(())
}
<|endoftext|>


<commit_before>use std::collections::{HashMap, HashSet};
use std::fmt;
use std::time::Duration;

use prost::Message;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tracing::{debug, info, warn};

use crate::xds::istio::workload::Workload;
use crate::xds::service::discovery::v3::aggregated_discovery_service_client::AggregatedDiscoveryServiceClient;
use crate::xds::service::discovery::v3::Resource as ProtoResource;
use crate::xds::service::discovery::v3::*;
use crate::{tls, xds};

use super::Error;

#[derive(Eq, Hash, PartialEq, Debug, Clone)]
pub struct ResourceKey {
    pub name: String,
    pub type_url: String,
}

#[derive(Default)]
pub struct HandlerContext {
    rejects: Vec<RejectedConfig>,
}

impl HandlerContext {
    pub fn new() -> HandlerContext {
        HandlerContext {
            ..Default::default()
        }
    }
}

struct RejectedConfig {
    name: String,
    reason: anyhow::Error,
}

impl fmt::Display for RejectedConfig {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}: {}", self.name, self.reason,)
    }
}

impl HandlerContext {
    pub fn reject(&mut self, name: String, reason: anyhow::Error) {
        self.rejects.push(RejectedConfig { name, reason })
    }
}

pub trait Handler<T: prost::Message>: Send + Sync + 'static {
    fn handle(&self, ctx: &mut HandlerContext, res: Vec<XdsUpdate<T>>);
}

struct NopHandler {}

impl<T: prost::Message> Handler<T> for NopHandler {
    fn handle(&self, _ctx: &mut HandlerContext, _res: Vec<XdsUpdate<T>>) {}
}

pub struct Config {
    workload_handler: Box<dyn Handler<Workload>>,
    initial_watches: Vec<String>,
    on_demand: bool,
}

impl Config {
    pub fn new(config: crate::config::Config) -> Config {
        Config {
            workload_handler: Box::new(NopHandler {}),
            initial_watches: Vec::new(),
            on_demand: config.xds_on_demand,
        }
    }

    pub fn with_workload_handler(mut self, f: impl Handler<Workload>) -> Config {
        self.workload_handler = Box::new(f);
        self
    }

    pub fn watch(mut self, type_url: String) -> Config {
        self.initial_watches.push(type_url);
        self
    }

    pub fn build(self) -> AdsClient {
        let (tx, rx) = mpsc::channel(100);
        AdsClient {
            config: self,
            workloads: HashSet::new(),
            pending: Default::default(),
            demand: rx,
            demand_tx: tx,
        }
    }
}

pub struct AdsClient {
    config: Config,
    /// Stores all known workload resources (by name)
    workloads: HashSet<String>,

    /// pending stores a list of all resources that are pending and XDS push
    pending: HashMap<ResourceKey, oneshot::Sender<()>>,

    demand: mpsc::Receiver<(oneshot::Sender<()>, ResourceKey)>,
    demand_tx: mpsc::Sender<(oneshot::Sender<()>, ResourceKey)>,
}

/// Demanded allows awaiting for an on-demand XDS resource
pub struct Demanded {
    b: oneshot::Receiver<()>,
}

impl Demanded {
    /// recv awaits for the requested resource
    /// Note: the actual resource is not directly returned. Instead, callers are notified that the event
    /// has been handled through the configured resource handler.
    pub async fn recv(self) {
        let _ = self.b.await;
    }
}

/// Demander allows requesting XDS resources on-demand
#[derive(Debug, Clone)]
pub struct Demander {
    demand: mpsc::Sender<(oneshot::Sender<()>, ResourceKey)>,
}

impl Demander {
    /// Demand requests a given workload by name
    pub async fn demand(&self, name: String) -> Demanded {
        let (tx, rx) = oneshot::channel::<()>();
        self.demand
            .send((
                tx,
                ResourceKey {
                    name,
                    type_url: xds::WORKLOAD_TYPE.to_string(),
                },
            ))
            .await
            .unwrap();
        Demanded { b: rx }
    }
}

impl AdsClient {
    /// demander returns a Demander instance which can be used to request resources on-demand
    pub fn demander(&self) -> Demander {
        Demander {
            demand: self.demand_tx.clone(),
        }
    }

    pub async fn run(mut self) -> Result<(), Error> {
        let mut backoff = Duration::from_millis(10);
        let max_backoff = Duration::from_secs(15);
        loop {
            let res = self.run_internal().await;
            match res {
                Err(e @ Error::Connection(_)) => {
                    // For connection errors, we add backoff
                    backoff = std::cmp::min(max_backoff, backoff * 2);
                    warn!("XDS client error: {}, retrying in {:?}", e, backoff);
                    tokio::time::sleep(backoff).await;
                }
                Err(e) => {
                    // For other errors, we connect immediately
                    // TODO: we may need more nuance here; if we fail due to invalid initial request we may overload
                    // But we want to reconnect from MaxConnectionAge immediately.
                    warn!("XDS client error: {}, retrying", e);
                    backoff = Duration::from_millis(10);
                }
                Ok(_) => {
                    warn!("XDS client complete");
                    backoff = Duration::from_millis(10);
                }
            }
        }
    }

    async fn run_internal(&mut self) -> Result<(), Error> {
        let initial_watches = &self.config.initial_watches;
        let workloads = &mut self.workloads;
        match workloads.len() {
            0 => info!("Starting initial ADS client"),
            n => info!("Starting ADS client with {n} workloads"),
        };

        // TODO: copy JWT auth logic from CA client and use TLS here
        let _address = if std::env::var("KUBERNETES_SERVICE_HOST").is_ok() {
            "https://istiod.istio-system:15012"
        } else {
            "https://localhost:15012"
        };
        let address = if std::env::var("KUBERNETES_SERVICE_HOST").is_ok() {
            "http://istiod.istio-system:15010"
        } else {
            "http://localhost:15010"
        };
        let svc = tls::grpc_connector(address).unwrap();
        let mut client = AggregatedDiscoveryServiceClient::new(svc);
        let (discovery_req_tx, mut discovery_req_rx) = mpsc::channel::<DeltaDiscoveryRequest>(100);
        let watches = initial_watches.clone();
        let irv: HashMap<String, String> = workloads
            .iter()
            .map(|n| (n.clone(), "".to_string()))
            .collect();

        let (sub, unsub) = if self.config.on_demand {
            // XDS doesn't have a way to subscribe to zero resources. We workaround this by subscribing and unsubscribing
            // in one event, effectively giving us "subscribe to nothing".
            (vec!["*".to_string()], vec!["*".to_string()])
        } else {
            (vec![], vec![])
        };
        let outbound = async_stream::stream! {
            for request_type in watches {
                let irv = irv.clone();
                let initial = DeltaDiscoveryRequest {
                    type_url: request_type.clone(),
                    node: Some(Node{
                        id: format!("{}~{}~{}.{}~{}.svc.cluster.local", "sidecar", "1.1.1.1", "test", "test", "test"),
                        ..Default::default()
                    }),
                    initial_resource_versions: irv,
                    resource_names_subscribe: sub.clone(),
                    resource_names_unsubscribe: unsub.clone(),
                    ..Default::default()
                };
                info!(type_url=request_type, "sending initial request");
                yield initial;
            }
            while let Some(message) = discovery_req_rx.recv().await {
                info!(type_url=message.type_url, "sending request");
                yield message
            }
            warn!("outbound stream complete");
        };

        info!("Starting stream");
        let mut response_stream = client
            .delta_aggregated_resources(tonic::Request::new(outbound))
            .await
            .map_err(Error::Connection)?
            .into_inner();

        info!("Stream established");

        loop {
            tokio::select! {
                _demand_event = self.demand.recv() => {
                    self.handle_demand_event(None, &discovery_req_tx).await?;
                }
                msg = response_stream.message() =>{
                    self.handle_stream_event(msg?, &discovery_req_tx).await?;
                }

            }
        }
    }

    async fn handle_stream_event(
        &mut self,
        stream_event: Option<DeltaDiscoveryResponse>,
        send: &mpsc::Sender<DeltaDiscoveryRequest>,
    ) -> Result<(), Error> {
        let Some(response) = stream_event else {
            return Ok(());
        };
        let type_url = response.type_url;
        info!(
            type_url = type_url.clone(),
            size = response.resources.len(),
            "received response"
        );
        let mut updates: Vec<XdsUpdate<Workload>> = Vec::new();
        for res in response.resources.iter().cloned() {
            let resource = Resource::try_from(res).unwrap();
            let _cpy = resource.clone();
            match resource {
                Resource::Workload(inner) => {
                    let name = inner.name.clone();
                    let up = XdsUpdate::Update(inner);
                    updates.push(up);
                    self.workloads.insert(name.clone());
                    let pending = {
                        self.pending.remove(&ResourceKey {
                            name: name.clone(),
                            type_url: type_url.clone(),
                        })
                    };
                    if let Some(send) = pending {
                        debug!("on demand notify {}", name);
                        send.send(()).map_err(|_| Error::OnDemandSend())?;
                    }
                }
            }
        }
        for res in response.removed_resources {
            self.workloads.remove(&res.clone());
            debug!("received delete resource {:#?}", res);
            let pending = {
                self.pending.remove(&ResourceKey {
                    name: res.clone(),
                    type_url: type_url.clone(),
                })
            };
            if let Some(send) = pending {
                debug!("on demand notify {}", res);
                send.send(()).map_err(|_| Error::OnDemandSend())?;
            }
            match type_url.clone().as_str() {
                xds::WORKLOAD_TYPE => {
                    updates.push(XdsUpdate::Remove(res));
                }
                _ => warn!("ignoring unwatched type {}", type_url),
            }
        }

        let mut ctx = HandlerContext::new();
        self.config.workload_handler.handle(&mut ctx, updates);

        let error_detail = match ctx.rejects.len() {
            0 => None,
            _ => Some(
                ctx.rejects
                    .into_iter()
                    .map(|reject| reject.to_string())
                    .collect::<Vec<String>>()
                    .join("; "),
            ),
        };
        info!(
            type_url = type_url.clone(),
            none = response.nonce,
            "sending {}",
            if error_detail.is_some() {
                "NACK"
            } else {
                "ACK"
            }
        );
        send.send(DeltaDiscoveryRequest {
            type_url: type_url.clone(),
            node: Some(Node {
                id: format!(
                    "{}~{}~{}.{}~{}.svc.cluster.local",
                    "sidecar", "1.1.1.1", "test", "test", "test"
                ),
                ..Default::default()
            }),
            response_nonce: response.nonce.clone(),
            error_detail: error_detail.map(|msg| Status {
                message: msg,
                ..Default::default()
            }),
            ..Default::default()
        })
        .await
        .map_err(|e| Error::RequestFailure(Box::new(e)))
    }

    async fn handle_demand_event(
        &mut self,
        demand_event: Option<(oneshot::Sender<()>, ResourceKey)>,
        send: &mpsc::Sender<DeltaDiscoveryRequest>,
    ) -> Result<(), Error> {
        info!("received on demand request {demand_event:?}");
        let Some((tx, demand_event)) = demand_event else {
            return Ok(());
        };
        let ResourceKey { type_url, name } = demand_event.clone();
        self.pending.insert(demand_event, tx);
        send.send(DeltaDiscoveryRequest {
            type_url,
            node: Some(Node {
                id: format!(
                    "{}~{}~{}.{}~{}.svc.cluster.local",
                    "sidecar", "1.1.1.1", "test", "test", "test"
                ),
                ..Default::default()
            }),
            resource_names_subscribe: vec![name],
            ..Default::default()
        })
        .await
        .map_err(|e| Error::RequestFailure(Box::new(e)))?;
        Ok(())
    }
}

#[derive(Clone, Debug)]
pub struct XdsResource<T: prost::Message> {
    pub name: String,
    pub resource: T,
}

#[derive(Debug)]
pub enum XdsUpdate<T: prost::Message> {
    Update(XdsResource<T>),
    Remove(String),
}

impl<T: prost::Message> XdsUpdate<T> {
    pub fn name(&self) -> String {
        match self {
            XdsUpdate::Update(ref r) => r.name.clone(),
            XdsUpdate::Remove(n) => n.to_string(),
        }
    }
}

// TODO: consider using https://crates.io/crates/prost-reflect to dynamically create types allowing any
// compiled protos. For now, we just hardcode
#[derive(Clone, Debug)]
pub enum Resource {
    Workload(XdsResource<Workload>),
}

impl TryFrom<ProtoResource> for Resource {
    type Error = AdsError;

    fn try_from(resource: ProtoResource) -> Result<Self, Self::Error> {
        let res = resource.resource.unwrap();
        Ok(match &*res.type_url {
            xds::WORKLOAD_TYPE => {
                let inner = <Workload>::decode(&*res.value).unwrap();
                Resource::Workload(XdsResource {
                    name: resource.name,
                    resource: inner,
                })
            }
            url => return Err(AdsError::UnknownResourceType(url.into())),
        })
    }
}

#[derive(Clone, Debug)]
pub enum AdsError {
    UnknownResourceType(String),
}
<commit_msg>Use env to generate node id for xds client (#124)<commit_after>use std::collections::{HashMap, HashSet};
use std::fmt;
use std::time::Duration;

use prost::Message;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tracing::{debug, info, warn};

use crate::xds::istio::workload::Workload;
use crate::xds::service::discovery::v3::aggregated_discovery_service_client::AggregatedDiscoveryServiceClient;
use crate::xds::service::discovery::v3::Resource as ProtoResource;
use crate::xds::service::discovery::v3::*;
use crate::{tls, xds};

use super::Error;

#[derive(Eq, Hash, PartialEq, Debug, Clone)]
pub struct ResourceKey {
    pub name: String,
    pub type_url: String,
}

#[derive(Default)]
pub struct HandlerContext {
    rejects: Vec<RejectedConfig>,
}

impl HandlerContext {
    pub fn new() -> HandlerContext {
        HandlerContext {
            ..Default::default()
        }
    }
}

struct RejectedConfig {
    name: String,
    reason: anyhow::Error,
}

impl fmt::Display for RejectedConfig {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}: {}", self.name, self.reason,)
    }
}

impl HandlerContext {
    pub fn reject(&mut self, name: String, reason: anyhow::Error) {
        self.rejects.push(RejectedConfig { name, reason })
    }
}

pub trait Handler<T: prost::Message>: Send + Sync + 'static {
    fn handle(&self, ctx: &mut HandlerContext, res: Vec<XdsUpdate<T>>);
}

struct NopHandler {}

impl<T: prost::Message> Handler<T> for NopHandler {
    fn handle(&self, _ctx: &mut HandlerContext, _res: Vec<XdsUpdate<T>>) {}
}

pub struct Config {
    workload_handler: Box<dyn Handler<Workload>>,
    initial_watches: Vec<String>,
    on_demand: bool,
}

impl Config {
    pub fn new(config: crate::config::Config) -> Config {
        Config {
            workload_handler: Box::new(NopHandler {}),
            initial_watches: Vec::new(),
            on_demand: config.xds_on_demand,
        }
    }

    pub fn with_workload_handler(mut self, f: impl Handler<Workload>) -> Config {
        self.workload_handler = Box::new(f);
        self
    }

    pub fn watch(mut self, type_url: String) -> Config {
        self.initial_watches.push(type_url);
        self
    }

    pub fn build(self) -> AdsClient {
        let (tx, rx) = mpsc::channel(100);
        AdsClient {
            config: self,
            workloads: HashSet::new(),
            pending: Default::default(),
            demand: rx,
            demand_tx: tx,
            node_id: generate_node_id(),
        }
    }
}

fn generate_node_id() -> String {
    // TODO: may add new type, but this requires changing together with pilot-discovery
    let node_type = "sidecar";
    let ip = std::env::var("INSTANCE_IP").unwrap_or_else(|_| "1.1.1.1".to_string());
    let pod_name = std::env::var("POD_NAME").unwrap_or_else(|_| "test".to_string());
    let ns = std::env::var("POD_NAMESPACE").unwrap_or_else(|_| "test".to_string());
    let dns_domain = format!("{}.{}", ns, "svc.cluster.local");
    format!("{}~{}~{}.{}~{}", node_type, ip, pod_name, ns, dns_domain)
}

pub struct AdsClient {
    config: Config,
    /// Stores all known workload resources (by name)
    workloads: HashSet<String>,

    /// pending stores a list of all resources that are pending and XDS push
    pending: HashMap<ResourceKey, oneshot::Sender<()>>,

    demand: mpsc::Receiver<(oneshot::Sender<()>, ResourceKey)>,
    demand_tx: mpsc::Sender<(oneshot::Sender<()>, ResourceKey)>,

    node_id: String,
}

/// Demanded allows awaiting for an on-demand XDS resource
pub struct Demanded {
    b: oneshot::Receiver<()>,
}

impl Demanded {
    /// recv awaits for the requested resource
    /// Note: the actual resource is not directly returned. Instead, callers are notified that the event
    /// has been handled through the configured resource handler.
    pub async fn recv(self) {
        let _ = self.b.await;
    }
}

/// Demander allows requesting XDS resources on-demand
#[derive(Debug, Clone)]
pub struct Demander {
    demand: mpsc::Sender<(oneshot::Sender<()>, ResourceKey)>,
}

impl Demander {
    /// Demand requests a given workload by name
    pub async fn demand(&self, name: String) -> Demanded {
        let (tx, rx) = oneshot::channel::<()>();
        self.demand
            .send((
                tx,
                ResourceKey {
                    name,
                    type_url: xds::WORKLOAD_TYPE.to_string(),
                },
            ))
            .await
            .unwrap();
        Demanded { b: rx }
    }
}

impl AdsClient {
    /// demander returns a Demander instance which can be used to request resources on-demand
    pub fn demander(&self) -> Demander {
        Demander {
            demand: self.demand_tx.clone(),
        }
    }

    pub async fn run(mut self) -> Result<(), Error> {
        let mut backoff = Duration::from_millis(10);
        let max_backoff = Duration::from_secs(15);
        loop {
            let res = self.run_internal().await;
            match res {
                Err(e @ Error::Connection(_)) => {
                    // For connection errors, we add backoff
                    backoff = std::cmp::min(max_backoff, backoff * 2);
                    warn!("XDS client error: {}, retrying in {:?}", e, backoff);
                    tokio::time::sleep(backoff).await;
                }
                Err(e) => {
                    // For other errors, we connect immediately
                    // TODO: we may need more nuance here; if we fail due to invalid initial request we may overload
                    // But we want to reconnect from MaxConnectionAge immediately.
                    warn!("XDS client error: {}, retrying", e);
                    backoff = Duration::from_millis(10);
                }
                Ok(_) => {
                    warn!("XDS client complete");
                    backoff = Duration::from_millis(10);
                }
            }
        }
    }

    async fn run_internal(&mut self) -> Result<(), Error> {
        let initial_watches = &self.config.initial_watches;
        let workloads = &mut self.workloads;
        let node_id = self.node_id.clone();
        match workloads.len() {
            0 => info!("Starting initial ADS client"),
            n => info!("Starting ADS client with {n} workloads"),
        };

        // TODO: copy JWT auth logic from CA client and use TLS here
        let _address = if std::env::var("KUBERNETES_SERVICE_HOST").is_ok() {
            "https://istiod.istio-system:15012"
        } else {
            "https://localhost:15012"
        };
        let address = if std::env::var("KUBERNETES_SERVICE_HOST").is_ok() {
            "http://istiod.istio-system:15010"
        } else {
            "http://localhost:15010"
        };
        let svc = tls::grpc_connector(address).unwrap();
        let mut client = AggregatedDiscoveryServiceClient::new(svc);
        let (discovery_req_tx, mut discovery_req_rx) = mpsc::channel::<DeltaDiscoveryRequest>(100);
        let watches = initial_watches.clone();
        let irv: HashMap<String, String> = workloads
            .iter()
            .map(|n| (n.clone(), "".to_string()))
            .collect();

        let (sub, unsub) = if self.config.on_demand {
            // XDS doesn't have a way to subscribe to zero resources. We workaround this by subscribing and unsubscribing
            // in one event, effectively giving us "subscribe to nothing".
            (vec!["*".to_string()], vec!["*".to_string()])
        } else {
            (vec![], vec![])
        };
        let outbound = async_stream::stream! {
            for request_type in watches {
                let irv = irv.clone();
                let initial = DeltaDiscoveryRequest {
                    type_url: request_type.clone(),
                    node: Some(Node{
                        id: node_id.clone(),
                        ..Default::default()
                    }),
                    initial_resource_versions: irv,
                    resource_names_subscribe: sub.clone(),
                    resource_names_unsubscribe: unsub.clone(),
                    ..Default::default()
                };
                info!(type_url=request_type, "sending initial request");
                yield initial;
            }
            while let Some(message) = discovery_req_rx.recv().await {
                info!(type_url=message.type_url, "sending request");
                yield message
            }
            warn!("outbound stream complete");
        };

        info!("Starting stream");
        let mut response_stream = client
            .delta_aggregated_resources(tonic::Request::new(outbound))
            .await
            .map_err(Error::Connection)?
            .into_inner();

        info!("Stream established");

        loop {
            tokio::select! {
                _demand_event = self.demand.recv() => {
                    self.handle_demand_event(None, &discovery_req_tx).await?;
                }
                msg = response_stream.message() =>{
                    self.handle_stream_event(msg?, &discovery_req_tx).await?;
                }

            }
        }
    }

    async fn handle_stream_event(
        &mut self,
        stream_event: Option<DeltaDiscoveryResponse>,
        send: &mpsc::Sender<DeltaDiscoveryRequest>,
    ) -> Result<(), Error> {
        let Some(response) = stream_event else {
            return Ok(());
        };
        let type_url = response.type_url;
        info!(
            type_url = type_url.clone(),
            size = response.resources.len(),
            "received response"
        );
        let mut updates: Vec<XdsUpdate<Workload>> = Vec::new();
        for res in response.resources.iter().cloned() {
            let resource = Resource::try_from(res).unwrap();
            let _cpy = resource.clone();
            match resource {
                Resource::Workload(inner) => {
                    let name = inner.name.clone();
                    let up = XdsUpdate::Update(inner);
                    updates.push(up);
                    self.workloads.insert(name.clone());
                    let pending = {
                        self.pending.remove(&ResourceKey {
                            name: name.clone(),
                            type_url: type_url.clone(),
                        })
                    };
                    if let Some(send) = pending {
                        debug!("on demand notify {}", name);
                        send.send(()).map_err(|_| Error::OnDemandSend())?;
                    }
                }
            }
        }
        for res in response.removed_resources {
            self.workloads.remove(&res.clone());
            debug!("received delete resource {:#?}", res);
            let pending = {
                self.pending.remove(&ResourceKey {
                    name: res.clone(),
                    type_url: type_url.clone(),
                })
            };
            if let Some(send) = pending {
                debug!("on demand notify {}", res);
                send.send(()).map_err(|_| Error::OnDemandSend())?;
            }
            match type_url.clone().as_str() {
                xds::WORKLOAD_TYPE => {
                    updates.push(XdsUpdate::Remove(res));
                }
                _ => warn!("ignoring unwatched type {}", type_url),
            }
        }

        let mut ctx = HandlerContext::new();
        self.config.workload_handler.handle(&mut ctx, updates);

        let error_detail = match ctx.rejects.len() {
            0 => None,
            _ => Some(
                ctx.rejects
                    .into_iter()
                    .map(|reject| reject.to_string())
                    .collect::<Vec<String>>()
                    .join("; "),
            ),
        };
        info!(
            type_url = type_url.clone(),
            none = response.nonce,
            "sending {}",
            if error_detail.is_some() {
                "NACK"
            } else {
                "ACK"
            }
        );
        send.send(DeltaDiscoveryRequest {
            type_url: type_url.clone(),
            node: Some(Node {
                id: self.node_id.clone(),
                ..Default::default()
            }),
            response_nonce: response.nonce.clone(),
            error_detail: error_detail.map(|msg| Status {
                message: msg,
                ..Default::default()
            }),
            ..Default::default()
        })
        .await
        .map_err(|e| Error::RequestFailure(Box::new(e)))
    }

    async fn handle_demand_event(
        &mut self,
        demand_event: Option<(oneshot::Sender<()>, ResourceKey)>,
        send: &mpsc::Sender<DeltaDiscoveryRequest>,
    ) -> Result<(), Error> {
        info!("received on demand request {demand_event:?}");
        let Some((tx, demand_event)) = demand_event else {
            return Ok(());
        };
        let ResourceKey { type_url, name } = demand_event.clone();
        self.pending.insert(demand_event, tx);
        send.send(DeltaDiscoveryRequest {
            type_url,
            node: Some(Node {
                id: self.node_id.clone(),
                ..Default::default()
            }),
            resource_names_subscribe: vec![name],
            ..Default::default()
        })
        .await
        .map_err(|e| Error::RequestFailure(Box::new(e)))?;
        Ok(())
    }
}

#[derive(Clone, Debug)]
pub struct XdsResource<T: prost::Message> {
    pub name: String,
    pub resource: T,
}

#[derive(Debug)]
pub enum XdsUpdate<T: prost::Message> {
    Update(XdsResource<T>),
    Remove(String),
}

impl<T: prost::Message> XdsUpdate<T> {
    pub fn name(&self) -> String {
        match self {
            XdsUpdate::Update(ref r) => r.name.clone(),
            XdsUpdate::Remove(n) => n.to_string(),
        }
    }
}

// TODO: consider using https://crates.io/crates/prost-reflect to dynamically create types allowing any
// compiled protos. For now, we just hardcode
#[derive(Clone, Debug)]
pub enum Resource {
    Workload(XdsResource<Workload>),
}

impl TryFrom<ProtoResource> for Resource {
    type Error = AdsError;

    fn try_from(resource: ProtoResource) -> Result<Self, Self::Error> {
        let res = resource.resource.unwrap();
        Ok(match &*res.type_url {
            xds::WORKLOAD_TYPE => {
                let inner = <Workload>::decode(&*res.value).unwrap();
                Resource::Workload(XdsResource {
                    name: resource.name,
                    resource: inner,
                })
            }
            url => return Err(AdsError::UnknownResourceType(url.into())),
        })
    }
}

#[derive(Clone, Debug)]
pub enum AdsError {
    UnknownResourceType(String),
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2017 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

pub struct Model {
    pub board: Board,
    pub turn: Turn,
    pub white_pieces: u32,
    pub white_hexes: u32,
    pub black_pieces: u32,
    pub black_hexes: u32,
    pub selected_piece: Option<FieldCoord>,
}

impl Model {
    pub fn new() -> Model {
        Model {
            board: Board::new(),
            turn: Turn::White,
            white_pieces: 18,
            white_hexes: 0,
            black_pieces: 18,
            black_hexes: 0,
            selected_piece: None,
        }
    }
}

pub enum Turn {
    White,
    Black,
}

impl Turn {
    pub fn switch_turns(&mut self) {
        *self = match *self {
            Turn::White => Turn::Black,
            Turn::Black => Turn::White,
        }
    }
}

pub struct Board {
    /*
                ____
               /    \
          ____/  +y  \____
         /    \      /    \
        /      \____/  +x  \
        \      /    \      /
         \____/      \____/
         /    \      /    \
        /  -x  \____/      \
        \      /    \      /
         \____/  -y  \____/
              \      /
               \____/

        The hex board uses an axial coordinate system with (0, 0) at the center. The x-axis slopes
        up to the right, and the y-axis goes up and down. The board is stored as a dense 2D array,
        because a ragged array won't work (since each row would have to be a different type).

        See http://www.redblobgames.com/grids/hexagons/#coordinates-axial for more info.
    */
    board: [[Option<Hex>; 5]; 5],
}

// Fields are numbered clockwise from the top. Even indicies are black, odd indicies are white.
type Hex = [Field; 6];

// There's no need to store the color of the piece on this field, since only white pieces can be on
// white fields, and vice versa.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Field {
    Piece,
    Empty,
}

impl Board {
    #[cfg_attr(rustfmt, rustfmt_skip)]
    /// Create a new board with the "Laurentius" starting position.
    pub fn new() -> Board {
        let mut board = Board {
            board: [[None; 5]; 5]
        };

        // (0, 0) is the only empty hex
        board.set_hex(&HexCoord::new(0, 0), Some([Field::Empty; 6]));

        // Conveniently, every other hex has exactly two pieces on it in the starting position.
        let piece_locations = [
            (-2,  2, 0, 4),
            (-2,  1, 0, 3),
            (-2,  0, 3, 5),
            (-1,  2, 1, 4),
            (-1,  1, 0, 4),
            (-1,  0, 3, 5),
            (-1, -1, 2, 5),
            ( 0,  2, 1, 5),
            ( 0,  1, 1, 5),
            ( 0, -1, 2, 4),
            ( 0, -2, 2, 4),
            ( 1,  1, 2, 5),
            ( 1,  0, 0, 2),
            ( 1, -1, 1, 3),
            ( 1, -2, 1, 4),
            ( 2,  0, 0, 2),
            ( 2, -1, 0, 3),
            ( 2, -2, 1, 3),
        ];

        for &(x, y, f1, f2) in &piece_locations {
            let mut hex = [Field::Empty; 6];
            hex[f1] = Field::Piece;
            hex[f2] = Field::Piece;
            board.set_hex(&HexCoord::new(x, y), Some(hex));
        }

        board
    }

    pub fn get_field(&self, coord: &FieldCoord) -> &Field {
        match *self.get_hex(&coord.to_hex()) {
            Some(ref hex) => &hex[coord.f as usize],
            None => panic!("Cannot get field on removed hex {:?}", coord),
        }
    }
    fn set_field(&mut self, coord: &FieldCoord, field: Field) {
        match *self.get_hex(&coord.to_hex()) {
            Some(mut hex) => hex[coord.f as usize] = field,
            None => panic!("Cannot set field on removed hex {:?}", coord),
        }
    }
    /// Return fields that share an edge with the given field. These fields are always the opposite
    /// color of the given field. If all of a piece's edge neighbors are occupied, that piece might
    /// be capturable.
    pub fn get_field_edge_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        let mut neighbors = vec![
            // There are always two edge neighbors on the same hex as the given field
            FieldCoord::new(coord.x, coord.y, (coord.f + 1) % 6),
            FieldCoord::new(coord.x, coord.y, (coord.f + 5) % 6),
        ];

        for (i, neighbor) in coord.to_hex().get_neighbors() {
            if coord.f == i && self.get_hex(&neighbor).is_some() {
                let f = (i + 3) % 6;
                neighbors.push(neighbor.to_field(f));
                break;
            }
        }
        neighbors
    }
    /// Return fields that share a vertex with the given field and have the same color as the given
    /// field. Pieces can move to fields that are vertex neighbors of the field they are on.
    pub fn get_field_vertex_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        // A field's vertex neighbors can be defined as the edge neighbors of its edge neighbors
        let mut neighbors = vec![];
        for neighbor in self.get_field_edge_neighbors(coord) {
            neighbors.append(&mut self.get_field_edge_neighbors(&neighbor));
        }
        // A field is not its own neighbor
        neighbors.into_iter().filter(|n| n != coord).collect()
    }
    pub fn move_piece(&mut self, from: &FieldCoord, to: &FieldCoord) {
        assert_eq!(
            self.get_field(from),
            &Field::Piece,
            "There is no piece at {:?} to move",
            from
        );
        assert!(
            self.get_field_vertex_neighbors(from).contains(to),
            "Cannot move piece at {:?} to non-vertex neighbor {:?}",
            from,
            to
        );
        assert_eq!(
            self.get_field(to),
            &Field::Empty,
            "Cannot move piece at {:?} to occupied neighbor {:?}",
            from,
            to
        );

        self.set_field(from, Field::Empty);
        self.set_field(to, Field::Piece);
    }
    pub fn remove_piece(&mut self, coord: &FieldCoord) {
        assert_eq!(
            self.get_field(coord),
            &Field::Piece,
            "There is no piece at {:?} to remove",
            coord
        );
        self.set_field(coord, Field::Empty);
    }

    fn get_hex(&self, coord: &HexCoord) -> &Option<Hex> {
        &self.board[(coord.x + 2) as usize][(coord.y + 2) as usize]
    }
    fn set_hex(&mut self, coord: &HexCoord, hex: Option<Hex>) {
        self.board[(coord.x + 2) as usize][(coord.y + 2) as usize] = hex;
    }
    pub fn get_hex_neighbors(&self, coord: &HexCoord) -> Vec<HexCoord> {
        coord
            .get_neighbors()
            .iter()
            .filter_map(|&(_, c)| match *self.get_hex(&c) {
                Some(_) => Some(c),
                None => None,
            })
            .collect()
    }
    /// Return fields that share an edge with the given hex and are outside of the given hex. If a
    /// hex is removed from the board, pieces occupying that hex's field neighbors might be
    /// capturable.
    pub fn get_hex_field_neighbors(&self, coord: &HexCoord) -> Vec<FieldCoord> {
        let mut neighbors = vec![];

        for (i, neighbor) in coord.get_neighbors() {
            if self.get_hex(&neighbor).is_some() {
                let f = (i + 3) % 6;
                neighbors.push(neighbor.to_field(f));
            }
        }
        neighbors
    }
    /// A hex is removable (and must be removed) if it is empty and is "attached to the board by 3
    /// or less adjacent sides."
    pub fn is_hex_removable(&self, coord: &HexCoord) -> bool {
        match *self.get_hex(coord) {
            Some(hex) => if hex.iter().any(|&f| f == Field::Piece) {
                return false;
            },
            None => panic!("The hex at {:?} has already been removed", coord),
        }

        let neighbor_idxs: Vec<u32> = coord.get_neighbors().iter().map(|&(i, _)| i).collect();
        let neighbor_idxs_slice = neighbor_idxs.as_slice();

        match neighbor_idxs_slice.len() {
            0 => panic!("A hex at {:?} is empty and has no neighbors", coord),
            1 => true,
            2 => {
                let valid_idx_combos = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [0, 5]];
                valid_idx_combos.iter().any(|c| c == neighbor_idxs_slice)
            }
            3 => {
                let valid_idx_combos = [
                    [0, 1, 2],
                    [1, 2, 3],
                    [2, 3, 4],
                    [3, 4, 5],
                    [0, 4, 5],
                    [0, 1, 5],
                ];
                valid_idx_combos.iter().any(|c| c == neighbor_idxs_slice)
            }
            _ => false,
        }
    }
    pub fn remove_hex(&mut self, coord: &HexCoord) {
        assert!(
            self.is_hex_removable(coord),
            "Cannot remove the hex at {:?} which is {:?} and has {} neighbors",
            coord,
            self.get_hex(coord),
            self.get_hex_neighbors(coord).len(),
        );
        self.set_hex(coord, None);
    }
    /// > extant (adj.): Still in existence; not destroyed, lost, or extinct (The Free Dictionary)
    ///
    /// Return the coordinates of the hexes that have not been removed yet.
    pub fn extant_hexes(&self) -> Vec<HexCoord> {
        let mut coords = vec![];

        for x in -2..3 {
            for y in -2..3 {
                let coord = HexCoord::new(x, y);
                if (x + y).abs() <= 2 && self.get_hex(&coord).is_some() {
                    coords.push(coord);
                }
            }
        }
        coords
    }
}

#[derive(Debug, PartialEq)]
pub struct FieldCoord {
    x: i32,
    y: i32,
    f: u32,
}

#[derive(Clone, Copy, Debug)]
pub struct HexCoord {
    x: i32,
    y: i32,
}

impl FieldCoord {
    pub fn new(x: i32, y: i32, f: u32) -> FieldCoord {
        assert!((x + y).abs() <= 2 && x.abs() <= 2 && y.abs() <= 2 && f < 6);
        FieldCoord { x, y, f }
    }
    pub fn to_hex(&self) -> HexCoord {
        HexCoord::new(self.x, self.y)
    }
}

impl HexCoord {
    pub fn new(x: i32, y: i32) -> HexCoord {
        assert!((x + y).abs() <= 2 && x.abs() <= 2 && y.abs() <= 2);
        HexCoord { x, y }
    }
    pub fn to_field(&self, f: u32) -> FieldCoord {
        FieldCoord::new(self.x, self.y, f)
    }
    // We return a Vec of tuples so that get_hex_field_neighbors and is_hex_removable know which
    // neighbors are on which side of the hex. They need to know this for different reasons:
    //   * get_hex_field_neighbors: the index of each neighboring field depends on which hex
    //                              neighbor that field neighbor is on
    //   * is_hex_removable: a hex is removable only if it is "attached to the board by 3 or less
    //                       adjacent sides"
    fn get_neighbors(&self) -> Vec<(u32, HexCoord)> {
        let mut neighbors = vec![];

        if self.y < 2 && (self.x + self.y) != 2 {
            neighbors.push((0, HexCoord::new(self.x, self.y + 1)));
        }
        if (self.x + self.y) != 2 && self.x < 2 {
            neighbors.push((1, HexCoord::new(self.x + 1, self.y)));
        }
        if self.x < 2 && self.y > -2 {
            neighbors.push((2, HexCoord::new(self.x + 1, self.y - 1)));
        }
        if self.y > -2 && (self.x + self.y) > -2 {
            neighbors.push((3, HexCoord::new(self.x, self.y - 1)));
        }
        if (self.x + self.y) > -2 && self.x > -2 {
            neighbors.push((4, HexCoord::new(self.x - 1, self.y)));
        }
        if self.x > -2 && self.y < 2 {
            neighbors.push((5, HexCoord::new(self.x - 1, self.y + 1)));
        }

        neighbors
    }
}
<commit_msg>Revise Model comments and error messages<commit_after>/*
 * Copyright (C) 2017 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

pub struct Model {
    pub board: Board,
    pub turn: Turn,
    pub white_pieces: u32,
    pub white_hexes: u32,
    pub black_pieces: u32,
    pub black_hexes: u32,
    pub selected_piece: Option<FieldCoord>,
}

impl Model {
    pub fn new() -> Model {
        Model {
            board: Board::new(),
            turn: Turn::White,
            white_pieces: 18,
            white_hexes: 0,
            black_pieces: 18,
            black_hexes: 0,
            selected_piece: None,
        }
    }
}

pub enum Turn {
    White,
    Black,
}

impl Turn {
    pub fn switch_turns(&mut self) {
        *self = match *self {
            Turn::White => Turn::Black,
            Turn::Black => Turn::White,
        }
    }
}

pub struct Board {
    /*
                ____
               /    \
          ____/  +y  \____
         /    \      /    \
        /      \____/  +x  \
        \      /    \      /
         \____/      \____/
         /    \      /    \
        /  -x  \____/      \
        \      /    \      /
         \____/  -y  \____/
              \      /
               \____/

        The hex board uses an axial coordinate system with (0, 0) at the center. The x-axis slopes
        up to the right, and the y-axis goes up and down. The board is stored as a dense 2D array,
        because a ragged array won't work (since each row would have to be a different type).

        See http://www.redblobgames.com/grids/hexagons/#coordinates-axial for more info.
    */
    board: [[Option<Hex>; 5]; 5],
}

// Fields are numbered clockwise from the top. Even indicies are black, odd indicies are white.
type Hex = [Field; 6];

// There's no need to store the color of the piece on this field, since only white pieces can be on
// white fields, and vice versa.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Field {
    Piece,
    Empty,
}

impl Board {
    #[cfg_attr(rustfmt, rustfmt_skip)]
    /// Create a new board with the "Laurentius" starting position.
    pub fn new() -> Board {
        let mut board = Board {
            board: [[None; 5]; 5]
        };

        // (0, 0) is the only empty hex.
        board.set_hex(&HexCoord::new(0, 0), Some([Field::Empty; 6]));

        // All other hexes have exactly two pieces on them in the starting position.
        let piece_locations = [
            (-2,  2, 0, 4),
            (-2,  1, 0, 3),
            (-2,  0, 3, 5),
            (-1,  2, 1, 4),
            (-1,  1, 0, 4),
            (-1,  0, 3, 5),
            (-1, -1, 2, 5),
            ( 0,  2, 1, 5),
            ( 0,  1, 1, 5),
            ( 0, -1, 2, 4),
            ( 0, -2, 2, 4),
            ( 1,  1, 2, 5),
            ( 1,  0, 0, 2),
            ( 1, -1, 1, 3),
            ( 1, -2, 1, 4),
            ( 2,  0, 0, 2),
            ( 2, -1, 0, 3),
            ( 2, -2, 1, 3),
        ];

        for &(x, y, f1, f2) in &piece_locations {
            let mut hex = [Field::Empty; 6];
            hex[f1] = Field::Piece;
            hex[f2] = Field::Piece;
            board.set_hex(&HexCoord::new(x, y), Some(hex));
        }

        board
    }

    pub fn get_field(&self, coord: &FieldCoord) -> &Field {
        match *self.get_hex(&coord.to_hex()) {
            Some(ref hex) => &hex[coord.f as usize],
            None => panic!(
                "Cannot get field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    fn set_field(&mut self, coord: &FieldCoord, field: Field) {
        match *self.get_hex(&coord.to_hex()) {
            Some(mut hex) => hex[coord.f as usize] = field,
            None => panic!(
                "Cannot set field {} on removed hex at {:?}",
                coord.f,
                coord.to_hex()
            ),
        }
    }
    /// Return fields that share an edge with the given field. These fields are always the opposite
    /// color of the given field. If all of a piece's edge neighbors are occupied, that piece might
    /// be capturable.
    pub fn get_field_edge_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        let mut neighbors = vec![
            // There are always two edge neighbors on the same hex as the given field
            FieldCoord::new(coord.x, coord.y, (coord.f + 1) % 6),
            FieldCoord::new(coord.x, coord.y, (coord.f + 5) % 6),
        ];

        for (i, neighbor) in coord.to_hex().get_neighbors() {
            if coord.f == i && self.get_hex(&neighbor).is_some() {
                let f = (i + 3) % 6;
                neighbors.push(neighbor.to_field(f));
                break;
            }
        }

        neighbors
    }
    /// Return fields that share a vertex with the given field and have the same color as the given
    /// field. Pieces can move to fields that are vertex neighbors of the field they are on.
    pub fn get_field_vertex_neighbors(&self, coord: &FieldCoord) -> Vec<FieldCoord> {
        // A field's vertex neighbors can be defined as the edge neighbors of its edge neighbors
        let mut neighbors = vec![];
        for neighbor in self.get_field_edge_neighbors(coord) {
            neighbors.append(&mut self.get_field_edge_neighbors(&neighbor));
        }
        // A field is not its own neighbor
        neighbors.into_iter().filter(|n| n != coord).collect()
    }
    pub fn move_piece(&mut self, from: &FieldCoord, to: &FieldCoord) {
        assert_eq!(
            self.get_field(from),
            &Field::Piece,
            "There is no piece at {:?} to move",
            from
        );
        assert!(
            self.get_field_vertex_neighbors(from).contains(to),
            "Cannot move piece at {:?} to non-(vertex neighbor) {:?}",
            from,
            to
        );
        assert_eq!(
            self.get_field(to),
            &Field::Empty,
            "Cannot move piece at {:?} to occupied neighbor {:?}",
            from,
            to
        );

        self.set_field(from, Field::Empty);
        self.set_field(to, Field::Piece);
    }
    pub fn remove_piece(&mut self, coord: &FieldCoord) {
        assert_eq!(
            self.get_field(coord),
            &Field::Piece,
            "There is no piece at {:?} to remove",
            coord
        );
        self.set_field(coord, Field::Empty);
    }

    fn get_hex(&self, coord: &HexCoord) -> &Option<Hex> {
        &self.board[(coord.x + 2) as usize][(coord.y + 2) as usize]
    }
    fn set_hex(&mut self, coord: &HexCoord, hex: Option<Hex>) {
        self.board[(coord.x + 2) as usize][(coord.y + 2) as usize] = hex;
    }
    pub fn get_hex_neighbors(&self, coord: &HexCoord) -> Vec<HexCoord> {
        coord
            .get_neighbors()
            .iter()
            .filter_map(|&(_, c)| match *self.get_hex(&c) {
                Some(_) => Some(c),
                None => None,
            })
            .collect()
    }
    /// Return fields that share an edge with the given hex and are outside of the given hex. If a
    /// hex is removed from the board, pieces occupying that hex's field neighbors might be
    /// capturable.
    pub fn get_hex_field_neighbors(&self, coord: &HexCoord) -> Vec<FieldCoord> {
        let mut neighbors = vec![];

        for (i, neighbor) in coord.get_neighbors() {
            if self.get_hex(&neighbor).is_some() {
                let f = (i + 3) % 6;
                neighbors.push(neighbor.to_field(f));
            }
        }
        neighbors
    }
    /// A hex is removable (and must be removed) if it is empty and is "attached to the board by 3
    /// or less adjacent sides."
    pub fn is_hex_removable(&self, coord: &HexCoord) -> bool {
        match *self.get_hex(coord) {
            Some(hex) => if hex.iter().any(|&f| f == Field::Piece) {
                return false;
            },
            None => panic!("The hex at {:?} has already been removed", coord),
        }

        let neighbor_idxs: Vec<u32> = coord.get_neighbors().iter().map(|&(i, _)| i).collect();
        let neighbor_idxs_slice = neighbor_idxs.as_slice();

        match neighbor_idxs_slice.len() {
            0 => panic!("A hex at {:?} is empty and has no neighbors", coord),
            1 => true,
            2 => {
                let valid_idx_combos = [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [0, 5]];
                valid_idx_combos.iter().any(|c| c == neighbor_idxs_slice)
            }
            3 => {
                let valid_idx_combos = [
                    [0, 1, 2],
                    [1, 2, 3],
                    [2, 3, 4],
                    [3, 4, 5],
                    [0, 4, 5],
                    [0, 1, 5],
                ];
                valid_idx_combos.iter().any(|c| c == neighbor_idxs_slice)
            }
            _ => false,
        }
    }
    pub fn remove_hex(&mut self, coord: &HexCoord) {
        assert!(
            self.is_hex_removable(coord),
            "Cannot remove the hex at {:?} which is {:?} and has {} neighbors",
            coord,
            self.get_hex(coord),
            self.get_hex_neighbors(coord).len(),
        );
        self.set_hex(coord, None);
    }
    /// > extant (adj.): Still in existence; not destroyed, lost, or extinct (The Free Dictionary)
    ///
    /// Return the coordinates of the hexes that have not been removed yet.
    pub fn extant_hexes(&self) -> Vec<HexCoord> {
        let mut coords = vec![];

        for x in -2..3 {
            for y in -2..3 {
                let coord = HexCoord::new(x, y);
                if (x + y).abs() <= 2 && self.get_hex(&coord).is_some() {
                    coords.push(coord);
                }
            }
        }
        coords
    }
}

#[derive(Debug, PartialEq)]
pub struct FieldCoord {
    x: i32,
    y: i32,
    f: u32,
}

#[derive(Clone, Copy, Debug)]
pub struct HexCoord {
    x: i32,
    y: i32,
}

impl FieldCoord {
    pub fn new(x: i32, y: i32, f: u32) -> FieldCoord {
        assert!((x + y).abs() <= 2 && x.abs() <= 2 && y.abs() <= 2 && f < 6);
        FieldCoord { x, y, f }
    }
    pub fn to_hex(&self) -> HexCoord {
        HexCoord::new(self.x, self.y)
    }
}

impl HexCoord {
    pub fn new(x: i32, y: i32) -> HexCoord {
        assert!((x + y).abs() <= 2 && x.abs() <= 2 && y.abs() <= 2);
        HexCoord { x, y }
    }
    pub fn to_field(&self, f: u32) -> FieldCoord {
        FieldCoord::new(self.x, self.y, f)
    }
    // We return a Vec of tuples so that get_hex_field_neighbors and is_hex_removable know which
    // neighbors are on which side of the hex. They need to know this for different reasons:
    //   * get_hex_field_neighbors: the index of each neighboring field depends on which hex
    //                              neighbor that field neighbor is on
    //   * is_hex_removable: a hex is removable only if it is "attached to the board by 3 or less
    //                       adjacent sides"
    fn get_neighbors(&self) -> Vec<(u32, HexCoord)> {
        let mut neighbors = vec![];

        if self.y < 2 && (self.x + self.y) != 2 {
            neighbors.push((0, HexCoord::new(self.x, self.y + 1)));
        }
        if (self.x + self.y) != 2 && self.x < 2 {
            neighbors.push((1, HexCoord::new(self.x + 1, self.y)));
        }
        if self.x < 2 && self.y > -2 {
            neighbors.push((2, HexCoord::new(self.x + 1, self.y - 1)));
        }
        if self.y > -2 && (self.x + self.y) > -2 {
            neighbors.push((3, HexCoord::new(self.x, self.y - 1)));
        }
        if (self.x + self.y) > -2 && self.x > -2 {
            neighbors.push((4, HexCoord::new(self.x - 1, self.y)));
        }
        if self.x > -2 && self.y < 2 {
            neighbors.push((5, HexCoord::new(self.x - 1, self.y + 1)));
        }

        neighbors
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Nikita Pekin and the wolfram_alpha_rs contributors
// See the README.md file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![allow(missing_docs)]
#![cfg_attr(feature = "clippy", allow(cyclomatic_complexity))]

//! Struct and enum definitions of values in the Wolfram|Alpha model.
//!
//! For more information, see [Wolfram|Alpha's API
//! documentation](http://products.wolframalpha.com/api/documentation.html).

use serde::{Deserialize, Deserializer, Error as SerdeError};
use serde::de::{MapVisitor, SeqVisitor, Visitor};
use url::Url;

#[cfg(feature = "nightly")]
include!("model.in.rs");

#[cfg(feature = "with-syntex")]
include!(concat!(env!("OUT_DIR"), "/model.rs"));

#[cfg(test)]
mod tests {
    use std::fs::File;
    use std::io::Read;
    use std::path::Path;

    use serde_xml::from_str;

    use super::{
        Definitions,
        DidYouMean,
        DidYouMeans,
        Img,
        Infos,
        LanguageMsg,
        Notes,
        Plaintext,
        Pod,
        QueryResult,
        Spellcheck,
        State,
        Statelist,
        States,
        Subpod,
        Tips,
        Warnings,
    };

    fn read_sample_data_from_path<P>(path: P) -> String
        where P: AsRef<Path>,
    {
        let mut file = File::open(path).unwrap();
        let mut body = String::new();
        file.read_to_string(&mut body).unwrap();
        body
    }

    #[test]
    fn test_query_result_deserializer() {
        from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_1.xml")).unwrap();
        from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_2.xml")).unwrap();
        //from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_3.xml")).unwrap();
        from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_4.xml")).unwrap();
        from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_5.xml")).unwrap();
        from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_6.xml")).unwrap();
        from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_7.xml")).unwrap();
    }

    #[test]
    fn test_notes_deserializer() {
        from_str::<Notes>(&read_sample_data_from_path("tests/sample-data/notes.xml")).unwrap();
    }

    #[test]
    fn test_languagemsg_deserializer() {
        from_str::<LanguageMsg>(&read_sample_data_from_path("tests/sample-data/languagemsg.xml")).unwrap();
    }

    #[test]
    fn test_definitions_deserializer() {
        from_str::<Definitions>(&read_sample_data_from_path("tests/sample-data/definitions.xml")).unwrap();
    }

    #[test]
    fn test_didyoumean_deserializer() {
        from_str::<DidYouMean>(&read_sample_data_from_path("tests/sample-data/didyoumean.xml")).unwrap();
    }

    #[test]
    fn test_didyoumeans_deserializer() {
        from_str::<DidYouMeans>(&read_sample_data_from_path("tests/sample-data/didyoumeans.xml")).unwrap();
    }

    #[test]
    fn test_warning_deserializer() {
        from_str::<Spellcheck>(&read_sample_data_from_path("tests/sample-data/spellcheck.xml")).unwrap();
    }

    #[test]
    fn test_warnings_deserializer() {
        from_str::<Warnings>(&read_sample_data_from_path("tests/sample-data/warnings.xml")).unwrap();
    }

    #[test]
    fn test_tips_deserializer() {
        from_str::<Tips>(&read_sample_data_from_path("tests/sample-data/tips.xml")).unwrap();
    }

    #[test]
    fn test_state_deserializer() {
        from_str::<State>(&read_sample_data_from_path("tests/sample-data/state/state.xml")).unwrap();

        from_str::<Statelist>(&read_sample_data_from_path("tests/sample-data/state/statelist.xml")).unwrap();

        from_str::<States>(&read_sample_data_from_path("tests/sample-data/state/states.xml")).unwrap();
        from_str::<States>(&read_sample_data_from_path("tests/sample-data/state/states-multiple-states.xml")).unwrap();
        from_str::<States>(&read_sample_data_from_path("tests/sample-data/state/states-multiple-statelists.xml")).unwrap();
        from_str::<States>(&read_sample_data_from_path("tests/sample-data/state/states-out-of-order.xml")).unwrap();
    }

    #[test]
    fn test_pod_deserializer() {
        from_str::<Pod>(&read_sample_data_from_path("tests/sample-data/pod.xml")).unwrap();
    }

    #[test]
    fn test_subpod_deserializer() {
        from_str::<Subpod>(&read_sample_data_from_path("tests/sample-data/subpod.xml")).unwrap();
    }

    #[test]
    fn test_img_deserializer() {
        from_str::<Img>(&read_sample_data_from_path("tests/sample-data/img.xml")).unwrap();
    }

    #[test]
    fn test_plaintext_deserializer() {
        from_str::<Plaintext>(&"<plaintext>pi</plaintext>".to_owned()).unwrap();
        from_str::<Option<Plaintext>>(&"<plaintext/>".to_owned()).unwrap();
    }

    #[test]
    fn test_infos_deserializer() {
        from_str::<Infos>(&read_sample_data_from_path("tests/sample-data/infos.xml")).unwrap();
    }
}
<commit_msg>chore: split `States` tests<commit_after>// Copyright (c) 2016 Nikita Pekin and the wolfram_alpha_rs contributors
// See the README.md file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![allow(missing_docs)]
#![cfg_attr(feature = "clippy", allow(cyclomatic_complexity))]

//! Struct and enum definitions of values in the Wolfram|Alpha model.
//!
//! For more information, see [Wolfram|Alpha's API
//! documentation](http://products.wolframalpha.com/api/documentation.html).

use serde::{Deserialize, Deserializer, Error as SerdeError};
use serde::de::{MapVisitor, SeqVisitor, Visitor};
use url::Url;

#[cfg(feature = "nightly")]
include!("model.in.rs");

#[cfg(feature = "with-syntex")]
include!(concat!(env!("OUT_DIR"), "/model.rs"));

#[cfg(test)]
mod tests {
    use std::fs::File;
    use std::io::Read;
    use std::path::Path;

    use serde_xml::from_str;

    use super::{
        Definitions,
        DidYouMean,
        DidYouMeans,
        Img,
        Infos,
        LanguageMsg,
        Notes,
        Plaintext,
        Pod,
        QueryResult,
        Spellcheck,
        State,
        Statelist,
        States,
        Subpod,
        Tips,
        Warnings,
    };

    fn read_sample_data_from_path<P>(path: P) -> String
        where P: AsRef<Path>,
    {
        let mut file = File::open(path).unwrap();
        let mut body = String::new();
        file.read_to_string(&mut body).unwrap();
        body
    }

    #[test]
    fn test_query_result_deserializer() {
        from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_1.xml")).unwrap();
        from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_2.xml")).unwrap();
        //from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_3.xml")).unwrap();
        from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_4.xml")).unwrap();
        from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_5.xml")).unwrap();
        from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_6.xml")).unwrap();
        from_str::<QueryResult>(&read_sample_data_from_path("tests/sample-data/query_result_7.xml")).unwrap();
    }

    #[test]
    fn test_notes_deserializer() {
        from_str::<Notes>(&read_sample_data_from_path("tests/sample-data/notes.xml")).unwrap();
    }

    #[test]
    fn test_languagemsg_deserializer() {
        from_str::<LanguageMsg>(&read_sample_data_from_path("tests/sample-data/languagemsg.xml")).unwrap();
    }

    #[test]
    fn test_definitions_deserializer() {
        from_str::<Definitions>(&read_sample_data_from_path("tests/sample-data/definitions.xml")).unwrap();
    }

    #[test]
    fn test_didyoumean_deserializer() {
        from_str::<DidYouMean>(&read_sample_data_from_path("tests/sample-data/didyoumean.xml")).unwrap();
    }

    #[test]
    fn test_didyoumeans_deserializer() {
        from_str::<DidYouMeans>(&read_sample_data_from_path("tests/sample-data/didyoumeans.xml")).unwrap();
    }

    #[test]
    fn test_warning_deserializer() {
        from_str::<Spellcheck>(&read_sample_data_from_path("tests/sample-data/spellcheck.xml")).unwrap();
    }

    #[test]
    fn test_warnings_deserializer() {
        from_str::<Warnings>(&read_sample_data_from_path("tests/sample-data/warnings.xml")).unwrap();
    }

    #[test]
    fn test_tips_deserializer() {
        from_str::<Tips>(&read_sample_data_from_path("tests/sample-data/tips.xml")).unwrap();
    }

    #[test]
    fn test_state_deserializer() {
        from_str::<State>(&read_sample_data_from_path("tests/sample-data/state/state.xml")).unwrap();
    }

    #[test]
    fn test_statelist_deserializer() {
        from_str::<Statelist>(&read_sample_data_from_path("tests/sample-data/state/statelist.xml")).unwrap();
    }

    #[test]
    fn test_states_deserializer() {
        from_str::<States>(&read_sample_data_from_path("tests/sample-data/state/states.xml")).unwrap();
        from_str::<States>(&read_sample_data_from_path("tests/sample-data/state/states-multiple-states.xml")).unwrap();
        from_str::<States>(&read_sample_data_from_path("tests/sample-data/state/states-multiple-statelists.xml")).unwrap();
        from_str::<States>(&read_sample_data_from_path("tests/sample-data/state/states-out-of-order.xml")).unwrap();
    }

    #[test]
    fn test_pod_deserializer() {
        from_str::<Pod>(&read_sample_data_from_path("tests/sample-data/pod.xml")).unwrap();
    }

    #[test]
    fn test_subpod_deserializer() {
        from_str::<Subpod>(&read_sample_data_from_path("tests/sample-data/subpod.xml")).unwrap();
    }

    #[test]
    fn test_img_deserializer() {
        from_str::<Img>(&read_sample_data_from_path("tests/sample-data/img.xml")).unwrap();
    }

    #[test]
    fn test_plaintext_deserializer() {
        from_str::<Plaintext>(&"<plaintext>pi</plaintext>".to_owned()).unwrap();
        from_str::<Option<Plaintext>>(&"<plaintext/>".to_owned()).unwrap();
    }

    #[test]
    fn test_infos_deserializer() {
        from_str::<Infos>(&read_sample_data_from_path("tests/sample-data/infos.xml")).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::str;

//TODO: Should take json state. Don't check for special values if parsing key
pub fn sanitise(remainder: &[u8], current: &mut String) {
	//Parse to a change of state, sending in the remainder and current
	if remainder.len() == 0 {
		return;
	}

	match remainder[0] {
		//Key
		b'"'|b'\'' => {
			let quote_byte = remainder[0];
			let (rest, key) = take_while(&remainder[1..], |c| c != quote_byte);

			current.push('"');
			current.push_str(key);
			current.push('"');

			// TODO: Make this work for empty strings
			sanitise(&rest[1..], current)
		},
		//Start of item
		b'{'|b'['|b':' => {
			let (rest, c) = take_first(remainder, |_| true);
			current.push(c as char);

			sanitise(&rest[1..], current)
		},
		//Replacements
		b'$' => {
			let (rest, key) = take_while1(&remainder[1..], |c|
				(c as char).is_alphanumeric() ||
				c == b'_' ||
				c == b'.'
			);

			current.push('$');
			current.push_str(key);

			sanitise(&rest[1..], current)
		},
		//Unquoted strings
		b if (b as char).is_alphabetic() => {
			let (rest, key) = take_while1(&remainder, |c|
				(c as char).is_alphabetic() ||
				c == b'_' ||
				c == b'.'
			);

			//Check if the string is a special value; true, false or null
			//For special values, push them as straight unquoted values. Otherwise, quote them
			match key {
				"true"|"false"|"null" =>
					current.push_str(key),
				_ => {
					current.push('"');
					current.push_str(key);
					current.push('"');
				}
			}

			sanitise(rest, current)
		},
		//Trim whitespace
		b' '|b'\r'|b'\n'|b'\t' => {
			sanitise(&remainder[1..], current)
		},
		//Other chars
		_ => {
			let (rest, c) = take_first(remainder, |_| true);
			current.push(c as char);

			sanitise(&rest[1..], current)
		}
	}
}

pub fn shift_while<F>(i: &[u8], f: F) -> &[u8] where F: Fn(u8) -> bool {
	let mut ctr = 0;
	for c in i {
		if f(*c) {
			ctr += 1;
		}
		else {
			break;
		}
	}

	&i[ctr..]
}

pub fn take_while<F>(i: &[u8], f: F) -> (&[u8], &str) where F: Fn(u8) -> bool {
	let mut ctr = 0;

	for c in i {
		if f(*c) {
			ctr += 1;
		}
		else {
			break;
		}
	}

	(&i[ctr..], str::from_utf8(&i[0..ctr]).unwrap())
}

pub fn take_while1<F>(i: &[u8], f: F) -> (&[u8], &str) where F: Fn(u8) -> bool {
	let mut ctr = 0;

	for c in i {
		if f(*c) || ctr == 0 {
			ctr += 1;
		}
		else {
			break;
		}
	}

	(&i[ctr..], str::from_utf8(&i[0..ctr]).unwrap())
}

pub fn take_first<F>(i: &[u8], f: F) -> (&[u8], u8) where F: Fn(u8) -> bool {
	let size = i.len();

	let mut ctr = 0;

	for c in i {
		if f(*c) || ctr == size - 1 {
			break;
		}
		else {
			ctr += 1;
		}
	}

	(&i[ctr..], i[ctr])
}

pub fn shift(i: &[u8], c: usize) -> &[u8] {
	match c {
		c if c >= i.len() => &[],
		_ => &i[c..]
	}
}
<commit_msg>tidy up parser a bit<commit_after>use std::str;

//TODO: Should take json state. Don't check for special values if parsing key
pub fn sanitise(remainder: &[u8], sanitised: &mut String) {
	if remainder.len() == 0 {
		return;
	}

	let current = remainder[0];

	match current {
		//Key
		b'"'|b'\'' => {
			let (rest, key) = take_while(&remainder[1..], |c| c != current);

			sanitised.push('"');
			sanitised.push_str(key);
			sanitised.push('"');

			sanitise(&rest[1..], sanitised)
		},
		//Start of item
		b'{'|b'['|b':' => {
			sanitised.push(current as char);

			sanitise(&remainder[1..], sanitised)
		},
		//Unquoted strings
		b if (b as char).is_alphabetic() => {
			let (rest, key) = take_while1(&remainder, |c|
				(c as char).is_alphabetic() ||
				c == b'_' ||
				c == b'.'
			);

			//Check if the string is a special value; true, false or null
			//For special values, push them as straight unquoted values. Otherwise, quote them
			match key {
				"true"|"false"|"null" =>
					sanitised.push_str(key),
				_ => {
					sanitised.push('"');
					sanitised.push_str(key);
					sanitised.push('"');
				}
			}

			sanitise(rest, sanitised)
		},
		//Trim whitespace
		b' '|b'\r'|b'\n'|b'\t' => {
			sanitise(&remainder[1..], sanitised)
		},
		//Other chars
		_ => {
			sanitised.push(current as char);

			sanitise(&remainder[1..], sanitised)
		}
	}
}

pub fn shift_while<F>(i: &[u8], f: F) -> &[u8] where F: Fn(u8) -> bool {
	let mut ctr = 0;
	for c in i {
		if f(*c) {
			ctr += 1;
		}
		else {
			break;
		}
	}

	&i[ctr..]
}

pub fn take_while<F>(i: &[u8], f: F) -> (&[u8], &str) where F: Fn(u8) -> bool {
	let mut ctr = 0;

	for c in i {
		if f(*c) {
			ctr += 1;
		}
		else {
			break;
		}
	}

	(&i[ctr..], unsafe { str::from_utf8_unchecked(&i[0..ctr]) })
}

pub fn take_while1<F>(i: &[u8], f: F) -> (&[u8], &str) where F: Fn(u8) -> bool {
	let mut ctr = 0;

	for c in i {
		if f(*c) || ctr == 0 {
			ctr += 1;
		}
		else {
			break;
		}
	}

	(&i[ctr..], unsafe { str::from_utf8_unchecked(&i[0..ctr]) })
}
<|endoftext|>