<commit_before>#![allow(dead_code)]
use std::io::Result as IoResult;
use std::result::Result;
use std::iter::{Iterator, IntoIterator};
use std::collections::HashMap;
use std::collections::hash_map;
use std::vec::Vec;
use std::cmp::{Eq, PartialEq};
use std::hash::{Hash, Hasher};
use std::convert::From;

pub type Pixels = usize;
pub type FileName = String;
pub type FilePath = String;

#[derive(Clone, Eq, PartialEq, Hash)]
pub enum ResolutionUnits {
    PixelsPerInch,
    PixelsPerCentimeter,
}

#[derive(Clone, Eq, PartialEq, Hash)]
pub enum Direction {
    Horizontal,
    Vertical,
}

#[derive(Clone, Eq, PartialEq, Hash)]
pub enum ImageFileFormat {
    TIFF,
    PNG,
    JPEG,
    UNKNOWN,
}

#[derive(Clone, Eq, PartialEq, Hash)]
struct ImageDimensions {
    x_pixels: Pixels,
    y_pixels: Pixels,
}

impl ImageDimensions {
    fn new(x: Pixels, y: Pixels) -> ImageDimensions {
        ImageDimensions {
            x_pixels: x,
            y_pixels: y,
        }
    }
}

#[derive(Clone, Eq, PartialEq, Hash)]
pub struct ImageResolution {
    amount: Pixels,
    units: ResolutionUnits,
}

impl ImageResolution {
    fn new(amount: usize, units: ResolutionUnits) -> ImageResolution {
        ImageResolution {
            amount: amount,
            units: units,
        }
    }
}

#[derive(Clone)]
pub enum PageOps {
    NoOperation,
    Identify(FileName, FilePath),
    Rescale(Pixels, Direction),
    ExpandLeftEdge(Pixels),
    ExpandRightEdge(Pixels),
    ExpandTopEdge(Pixels),
    ExpandBottomEdge(Pixels),
    TrimLeftEdge(Pixels),
    TrimRightEdge(Pixels),
    TrimTopEdge(Pixels),
    TrimBottomEdge(Pixels),
    SetResolution(ImageResolution),
}

pub trait ElementaryPageOperations {
    fn identify(file_name: FileName, path: FilePath) -> Self;
    fn rescale(amount: Pixels, dir: Direction)       -> Self;
    fn expand_left_edge(amount: Pixels)              -> Self;
    fn expand_right_edge(amount: Pixels)             -> Self;
    fn expand_top_edge(amount: Pixels)               -> Self;
    fn expand_bottom_edge(amount: Pixels)            -> Self;
    fn trim_left_edge(amount: Pixels)                -> Self;
    fn trim_right_edge(amount: Pixels)               -> Self;
    fn trim_top_edge(amount: Pixels)                 -> Self;
    fn trim_bottom_edge(amount: Pixels)              -> Self;
    fn set_resolution(res: ImageResolution)          -> Self;
    fn no_operation()                                -> Self;
}

pub trait RunOperation {
    fn run_operation(op: Self) -> OperationResults;
}

trait CompileOperation<OpType, Op> {
    fn compile_operation(op: OpType) -> Op;
    //fn apply_operation(op: Op)   -> ();
}

impl<Op> CompileOperation<PageOps, Op> for Op where Op: ElementaryPageOperations {
    fn compile_operation(op: PageOps) -> Op {
        match op {
            PageOps::Identify(file, path)     => Op::identify(file, path),
            PageOps::Rescale(amount, dir)     => Op::rescale(amount, dir),
            PageOps::ExpandLeftEdge(amount)   => Op::expand_left_edge(amount),
            PageOps::ExpandRightEdge(amount)  => Op::expand_right_edge(amount),
            PageOps::ExpandTopEdge(amount)    => Op::expand_top_edge(amount),
            PageOps::ExpandBottomEdge(amount) => Op::expand_bottom_edge(amount),
            PageOps::TrimLeftEdge(amount)     => Op::trim_left_edge(amount),
            PageOps::TrimRightEdge(amount)    => Op::trim_right_edge(amount),
            PageOps::TrimTopEdge(amount)      => Op::trim_top_edge(amount),
            PageOps::TrimBottomEdge(amount)   => Op::trim_bottom_edge(amount),
            PageOps::SetResolution(res)       => Op::set_resolution(res),
            PageOps::NoOperation              => Op::no_operation(),
        }
    }
}

#[derive(Clone)]
struct CompoundPageOperation<Op> {
    page_name: FileName,
    page_path: FilePath,
    ops: Option<Vec<Op>>,
}

impl<Op> CompoundPageOperation<Op> where Op: Clone {

    fn new(page_name: FileName, page_path: FilePath, ops: &[Op]) -> CompoundPageOperation<Op> {

        if ops.is_empty() {
            return CompoundPageOperation {
                page_name: page_name,
                page_path: page_path,
                ops: None,
            };
        }

        let mut vec = Vec::new();
        for op in ops.iter() {
            vec.push(op.clone());
        }

        CompoundPageOperation {
            page_name: page_name,
            page_path: page_path,
            ops: Some(vec),
        }
    }

    fn make_noop() -> CompoundPageOperation<Op> {
        CompoundPageOperation {
            page_name: String::from(""),
            page_path: String::from(""),
            ops: None,
        }
    }

    fn is_noop(&self) -> bool {
        match self.ops {
            Some(ref vec) => vec.is_empty(),
            None          => true,
        }
    }

/*
    fn run_operation<Op>(&self) -> OperationResults
        where Op: ElementaryPageOperations + RunOperation + CompileOperation<PageOps, Op> {

        match self.ops {
            None      => {
                OperationResults::from(Ok(String::from("No Operation")))
            }
            Some(ref vec) => {
                if self.is_noop() {
                    // Should not happen.
                    unreachable!();
                    //return Ok(String::from("No Operation"));
                }

                let mut result = Ok(String::from(""));
                for op in vec.iter() {
                    let op_results = Op::run_operation(Op::compile_operation(op.clone()));
                    for res in op_results.results {
                        match res {
                            Ok(s) => {
                                continue;
                            }
                            Err(e) => {
                                // Fail fast if there is an error.
                                result = Err(e);
                                break;
                            }
                        }
                    }
                }

                OperationResults::from(result)
            }
        }
    }
*/
}

#[derive(Clone, Eq)]
struct Page {
    file_name: FileName,
    file_extension: ImageFileFormat,
    file_path: FilePath,
    dimensions: ImageDimensions,
    resolution: ImageResolution,
}

impl Page {
    fn new  (
            file_name: FileName,
            file_extension: ImageFileFormat,
            file_path: FilePath,
            dimensions: ImageDimensions,
            resolution: ImageResolution
            ) -> Page
    {
        Page {
            file_name: file_name,
            file_extension: file_extension,
            file_path: file_path,
            dimensions: dimensions,
            resolution: resolution,
        }
    }
}

impl PartialEq for Page {
    fn eq(&self, other: &Page) -> bool {
            self.file_name      == other.file_name
        &&  self.file_extension == other.file_extension
        &&  self.file_path      == other.file_path
        &&  self.dimensions     == other.dimensions
        &&  self.resolution     == other.resolution
    }
}

impl Hash for Page {
    fn hash<H>(&self, state: &mut H) where H: Hasher {

        self.file_name.hash(state);
        self.file_extension.hash(state);
        self.file_path.hash(state);
        self.dimensions.hash(state);
        self.resolution.hash(state);
    }
}

pub type OperationResult = IoResult<String>;

pub struct OperationResults {
    results: Vec<OperationResult>,
}

impl OperationResults {
    pub fn new() -> OperationResults {
        OperationResults {
            results: Vec::new(),
        }
    }

    pub fn push(&mut self, result: OperationResult) {
        self.results.push(result);
    }

    pub fn append(&mut self, other: &mut OperationResults) {
        self.results.append(&mut other.results);
    }

}

impl From<Vec<OperationResult>> for OperationResults {
    fn from(vec: Vec<OperationResult>) -> OperationResults {
        unimplemented!();
    }
}

/// Destructive conversion from a mutable vector of OperationResult
/// to simplify the process of returning results from running operations.
impl<'a> From<&'a mut Vec<OperationResult>> for OperationResults {
    fn from(vec: &mut Vec<OperationResult>) -> OperationResults {
        let mut results = Vec::new();
        results.append(vec);

        OperationResults {
            results: results,
        }
    }
}

/// Generates an OperationResults struct from a single OperationResult.
/// For compatibility between operations that may return multiple results
/// and ones that may return only one result.
impl From<OperationResult> for OperationResults {
    fn from(op_res: OperationResult) -> OperationResults {
        let mut results = Vec::new();
        results.push(op_res);

        OperationResults {
            results: results
        }
    }
}

struct OperationPlan<Op> {
    plan: HashMap<Page, CompoundPageOperation<Op>>,
}

#[derive(Clone, Eq, PartialEq)]
enum OperationPlanError {
    LengthMismatch,
    Aborted,
}

impl<Op> OperationPlan<Op> where Op: Clone {
    fn new() -> OperationPlan<Op> {
        OperationPlan {
            plan: HashMap::new(),
        }
    }

    fn add_operation(&mut self, page: Page, op: CompoundPageOperation<Op>) {
        self.plan.insert(page, op);
    }

    fn build_schedule(pages: &[Page], ops: &[CompoundPageOperation<Op>]) -> Result<Self, OperationPlanError> {
        if pages.len() == ops.len() {

            let mut plan = OperationPlan::new();

            for page_number in 0..pages.len() {
                plan.add_operation(pages[page_number].clone(), ops[page_number].clone());
            }

            Ok(plan)

        } else {
            Err(OperationPlanError::LengthMismatch)
        }

    }

    fn iter(&self) -> OpPlanIter<Op> {
        OpPlanIter {
            inner: self.plan.iter()
        }
    }
/*
    fn run_operation<Op>(&self) -> OperationResults
        where Op: ElementaryPageOperations + RunOperation + CompileOperation<PageOps, Op> {

        let mut results = OperationResults::new();

        for (page, op) in self {

            match op.ops {
                None => {
                    continue;
                }
                Some(ref operations) => {
                    for elem_op in operations {
                        let compiled_op = Op::compile_operation(elem_op.clone());
                        let mut result = Op::run_operation(compiled_op);
                        results.append(&mut result);
                    }
                }
            }

        }

        results
    }
*/
}

/// Iterator implementation for OperationPlan.
struct OpPlanIter<'a, Op: 'a> {
    inner:  hash_map::Iter<'a, Page, CompoundPageOperation<Op>>
}

impl<'a, Op> Iterator for OpPlanIter<'a, Op> {
    type Item = (&'a Page, &'a CompoundPageOperation<Op>);

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }

}

/// IntoIterator implementation for OperationPlan.
impl<'a, Op> IntoIterator for &'a OperationPlan<Op> where Op: Clone {
    type Item = (&'a Page, &'a CompoundPageOperation<Op>);
    type IntoIter = OpPlanIter<'a, Op>;

    fn into_iter(self) -> OpPlanIter<'a, Op> {
        self.iter()
    }
}

struct OpPlanIntoIter<Op> {
    inner: hash_map::IntoIter<Page, CompoundPageOperation<Op>>,
}

impl<Op> IntoIterator for OperationPlan<Op> {
    type Item = (Page, CompoundPageOperation<Op>);
    type IntoIter = OpPlanIntoIter<Op>;

    fn into_iter(self) -> OpPlanIntoIter<Op> {
        OpPlanIntoIter {
            inner: self.plan.into_iter()
        }
    }
}

impl<Op> Iterator for OpPlanIntoIter<Op> {
    type Item = (Page, CompoundPageOperation<Op>);

    fn next(&mut self) -> Option<(Page, CompoundPageOperation<Op>)> {
        self.inner.next()
    }
}

struct OperationPlanResult {
    results: HashMap<Page, OperationResults>,
}

impl OperationPlanResult {
    fn new() -> OperationPlanResult {
        OperationPlanResult {
            results: HashMap::new(),
        }
    }

    fn insert(&mut self, page: Page, res: OperationResults) {
        self.results.insert(page, res);
    }

    fn iter(&self) -> OpPlanResultIter {
        OpPlanResultIter {
            inner: self.results.iter()
        }
    }
}

struct OpPlanResultIter<'a> {
    inner: hash_map::Iter<'a, Page, OperationResults>,
}

impl<'a> Iterator for OpPlanResultIter<'a> {
    type Item = (&'a Page, &'a OperationResults);

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }
}

/// IntoIterator implementation for OpPlanResultIter.
impl<'a> IntoIterator for &'a OperationPlanResult {
    type Item = (&'a Page, &'a OperationResults);
    type IntoIter = OpPlanResultIter<'a>;

    fn into_iter(self) -> OpPlanResultIter<'a> {
        self.iter()
    }
}

struct OpPlanResultIntoIter {
    inner: hash_map::IntoIter<Page, OperationResults>,
}

impl IntoIterator for OperationPlanResult {
    type Item = (Page, OperationResults);
    type IntoIter = OpPlanResultIntoIter;

    fn into_iter(self) -> OpPlanResultIntoIter {
        OpPlanResultIntoIter {
            inner: self.results.into_iter()
        }
    }
}

impl Iterator for OpPlanResultIntoIter {
    type Item = (Page, OperationResults);

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }
}
<commit_msg>Revise CompoundPageOperation impl.<commit_after>#![allow(dead_code)]
use std::io::Result as IoResult;
use std::result::Result;
use std::iter::{Iterator, IntoIterator};
use std::collections::HashMap;
use std::collections::hash_map;
use std::vec::Vec;
use std::cmp::{Eq, PartialEq};
use std::hash::{Hash, Hasher};
use std::convert::From;

pub type Pixels = usize;
pub type FileName = String;
pub type FilePath = String;

#[derive(Clone, Eq, PartialEq, Hash)]
pub enum ResolutionUnits {
    PixelsPerInch,
    PixelsPerCentimeter,
}

#[derive(Clone, Eq, PartialEq, Hash)]
pub enum Direction {
    Horizontal,
    Vertical,
}

#[derive(Clone, Eq, PartialEq, Hash)]
pub enum ImageFileFormat {
    TIFF,
    PNG,
    JPEG,
    UNKNOWN,
}

#[derive(Clone, Eq, PartialEq, Hash)]
struct ImageDimensions {
    x_pixels: Pixels,
    y_pixels: Pixels,
}

impl ImageDimensions {
    fn new(x: Pixels, y: Pixels) -> ImageDimensions {
        ImageDimensions {
            x_pixels: x,
            y_pixels: y,
        }
    }
}

#[derive(Clone, Eq, PartialEq, Hash)]
pub struct ImageResolution {
    amount: Pixels,
    units: ResolutionUnits,
}

impl ImageResolution {
    fn new(amount: usize, units: ResolutionUnits) -> ImageResolution {
        ImageResolution {
            amount: amount,
            units: units,
        }
    }
}

#[derive(Clone, PartialEq, Eq)]
pub enum PageOps {
    NoOperation,
    Identify(FileName, FilePath),
    Rescale(Pixels, Direction),
    ExpandLeftEdge(Pixels),
    ExpandRightEdge(Pixels),
    ExpandTopEdge(Pixels),
    ExpandBottomEdge(Pixels),
    TrimLeftEdge(Pixels),
    TrimRightEdge(Pixels),
    TrimTopEdge(Pixels),
    TrimBottomEdge(Pixels),
    SetResolution(ImageResolution),
}

pub trait ElementaryPageOperations {
    fn identify(file_name: FileName, path: FilePath) -> Self;
    fn rescale(amount: Pixels, dir: Direction)       -> Self;
    fn expand_left_edge(amount: Pixels)              -> Self;
    fn expand_right_edge(amount: Pixels)             -> Self;
    fn expand_top_edge(amount: Pixels)               -> Self;
    fn expand_bottom_edge(amount: Pixels)            -> Self;
    fn trim_left_edge(amount: Pixels)                -> Self;
    fn trim_right_edge(amount: Pixels)               -> Self;
    fn trim_top_edge(amount: Pixels)                 -> Self;
    fn trim_bottom_edge(amount: Pixels)              -> Self;
    fn set_resolution(res: ImageResolution)          -> Self;
    fn no_operation()                                -> Self;
}

pub trait RunOperation {
    fn run_operation(op: Self) -> OperationResults;
}

trait CompileOperation<OpType, Op> {
    fn compile_operation(op: OpType) -> Op;
}

impl<Op> CompileOperation<PageOps, Op> for Op where Op: ElementaryPageOperations {
    fn compile_operation(op: PageOps) -> Op {
        match op {
            PageOps::Identify(file, path)     => Op::identify(file, path),
            PageOps::Rescale(amount, dir)     => Op::rescale(amount, dir),
            PageOps::ExpandLeftEdge(amount)   => Op::expand_left_edge(amount),
            PageOps::ExpandRightEdge(amount)  => Op::expand_right_edge(amount),
            PageOps::ExpandTopEdge(amount)    => Op::expand_top_edge(amount),
            PageOps::ExpandBottomEdge(amount) => Op::expand_bottom_edge(amount),
            PageOps::TrimLeftEdge(amount)     => Op::trim_left_edge(amount),
            PageOps::TrimRightEdge(amount)    => Op::trim_right_edge(amount),
            PageOps::TrimTopEdge(amount)      => Op::trim_top_edge(amount),
            PageOps::TrimBottomEdge(amount)   => Op::trim_bottom_edge(amount),
            PageOps::SetResolution(res)       => Op::set_resolution(res),
            PageOps::NoOperation              => Op::no_operation(),
        }
    }
}

#[derive(Clone)]
struct CompoundPageOperation<Op> {
    page_name: FileName,
    page_path: FilePath,
    ops: Vec<Op>,
}

impl<Op> CompoundPageOperation<Op> where Op: Clone {

    fn new(page_name: FileName, page_path: FilePath, ops: &[Op]) -> CompoundPageOperation<Op> {
        let mut vec = Vec::new();
        for op in ops.iter() {
            vec.push(op.clone());
        }

        CompoundPageOperation {
            page_name: page_name,
            page_path: page_path,
            ops: vec,
        }
    }

    fn make_no_op(page_name: FileName, page_path: FilePath) -> CompoundPageOperation<Op> {
        CompoundPageOperation {
            page_name: page_name,
            page_path: page_path,
            ops: Vec::new(),
        }
    }

    fn is_no_op(&self) -> bool {
        self.ops.is_empty()
    }

/*  TODO: Move into a RunOperation instance.
    fn run_operation<Op>(&self) -> OperationResults
        where Op: ElementaryPageOperations + RunOperation + CompileOperation<PageOps, Op> {

        match self.ops {
            None      => {
                OperationResults::from(Ok(String::from("No Operation")))
            }
            Some(ref vec) => {
                if self.is_noop() {
                    // Should not happen.
                    unreachable!();
                    //return Ok(String::from("No Operation"));
                }

                let mut result = Ok(String::from(""));
                for op in vec.iter() {
                    let op_results = Op::run_operation(Op::compile_operation(op.clone()));
                    for res in op_results.results {
                        match res {
                            Ok(s) => {
                                continue;
                            }
                            Err(e) => {
                                // Fail fast if there is an error.
                                result = Err(e);
                                break;
                            }
                        }
                    }
                }

                OperationResults::from(result)
            }
        }
    }
*/
}

impl CompoundPageOperation<PageOps> {
    fn is_no_op(&self) -> bool {
        for op in self.ops.iter() {
            if *op == PageOps::NoOperation {
                return true;
            }
        }

        false
    }
}

// TODO: Implement this.
impl<Op, OtherOp> CompileOperation<CompoundPageOperation<Op>, CompoundPageOperation<OtherOp>>
    for CompoundPageOperation<Op> where Op: CompileOperation<Op, OtherOp> {

        fn compile_operation(op: CompoundPageOperation<Op>) -> CompoundPageOperation<OtherOp> {
            unimplemented!();
        }
    }

#[derive(Clone, Eq)]
struct Page {
    file_name: FileName,
    file_extension: ImageFileFormat,
    file_path: FilePath,
    dimensions: ImageDimensions,
    resolution: ImageResolution,
}

impl Page {
    fn new  (
            file_name: FileName,
            file_extension: ImageFileFormat,
            file_path: FilePath,
            dimensions: ImageDimensions,
            resolution: ImageResolution
            ) -> Page
    {
        Page {
            file_name: file_name,
            file_extension: file_extension,
            file_path: file_path,
            dimensions: dimensions,
            resolution: resolution,
        }
    }
}

impl PartialEq for Page {
    fn eq(&self, other: &Page) -> bool {
            self.file_name      == other.file_name
        &&  self.file_extension == other.file_extension
        &&  self.file_path      == other.file_path
        &&  self.dimensions     == other.dimensions
        &&  self.resolution     == other.resolution
    }
}

impl Hash for Page {
    fn hash<H>(&self, state: &mut H) where H: Hasher {

        self.file_name.hash(state);
        self.file_extension.hash(state);
        self.file_path.hash(state);
        self.dimensions.hash(state);
        self.resolution.hash(state);
    }
}

pub type OperationResult = IoResult<String>;

pub struct OperationResults {
    results: Vec<OperationResult>,
}

impl OperationResults {
    pub fn new() -> OperationResults {
        OperationResults {
            results: Vec::new(),
        }
    }

    pub fn push(&mut self, result: OperationResult) {
        self.results.push(result);
    }

    pub fn append(&mut self, other: &mut OperationResults) {
        self.results.append(&mut other.results);
    }

}

impl From<Vec<OperationResult>> for OperationResults {
    fn from(vec: Vec<OperationResult>) -> OperationResults {
        unimplemented!();
    }
}

/// Destructive conversion from a mutable vector of OperationResult
/// to simplify the process of returning results from running operations.
impl<'a> From<&'a mut Vec<OperationResult>> for OperationResults {
    fn from(vec: &mut Vec<OperationResult>) -> OperationResults {
        let mut results = Vec::new();
        results.append(vec);

        OperationResults {
            results: results,
        }
    }
}

/// Generates an OperationResults struct from a single OperationResult.
/// For compatibility between operations that may return multiple results
/// and ones that may return only one result.
impl From<OperationResult> for OperationResults {
    fn from(op_res: OperationResult) -> OperationResults {
        let mut results = Vec::new();
        results.push(op_res);

        OperationResults {
            results: results
        }
    }
}

struct OperationPlan<Op> {
    plan: HashMap<Page, CompoundPageOperation<Op>>,
}

#[derive(Clone, Eq, PartialEq)]
enum OperationPlanError {
    LengthMismatch,
    Aborted,
}

impl<Op> OperationPlan<Op> where Op: Clone {
    fn new() -> OperationPlan<Op> {
        OperationPlan {
            plan: HashMap::new(),
        }
    }

    fn add_operation(&mut self, page: Page, op: CompoundPageOperation<Op>) {
        self.plan.insert(page, op);
    }

    fn build_schedule(pages: &[Page], ops: &[CompoundPageOperation<Op>]) -> Result<Self, OperationPlanError> {
        if pages.len() == ops.len() {

            let mut plan = OperationPlan::new();

            for page_number in 0..pages.len() {
                plan.add_operation(pages[page_number].clone(), ops[page_number].clone());
            }

            Ok(plan)

        } else {
            Err(OperationPlanError::LengthMismatch)
        }

    }

    fn iter(&self) -> OpPlanIter<Op> {
        OpPlanIter {
            inner: self.plan.iter()
        }
    }
/*  TODO: Move into a RunOperation instance.
    fn run_operation<Op>(&self) -> OperationResults
        where Op: ElementaryPageOperations + RunOperation + CompileOperation<PageOps, Op> {

        let mut results = OperationResults::new();

        for (page, op) in self {

            match op.ops {
                None => {
                    continue;
                }
                Some(ref operations) => {
                    for elem_op in operations {
                        let compiled_op = Op::compile_operation(elem_op.clone());
                        let mut result = Op::run_operation(compiled_op);
                        results.append(&mut result);
                    }
                }
            }

        }

        results
    }
*/
}

/// Iterator implementation for OperationPlan.
struct OpPlanIter<'a, Op: 'a> {
    inner:  hash_map::Iter<'a, Page, CompoundPageOperation<Op>>
}

impl<'a, Op> Iterator for OpPlanIter<'a, Op> {
    type Item = (&'a Page, &'a CompoundPageOperation<Op>);

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }

}

/// IntoIterator implementation for OperationPlan.
impl<'a, Op> IntoIterator for &'a OperationPlan<Op> where Op: Clone {
    type Item = (&'a Page, &'a CompoundPageOperation<Op>);
    type IntoIter = OpPlanIter<'a, Op>;

    fn into_iter(self) -> OpPlanIter<'a, Op> {
        self.iter()
    }
}

struct OpPlanIntoIter<Op> {
    inner: hash_map::IntoIter<Page, CompoundPageOperation<Op>>,
}

impl<Op> IntoIterator for OperationPlan<Op> {
    type Item = (Page, CompoundPageOperation<Op>);
    type IntoIter = OpPlanIntoIter<Op>;

    fn into_iter(self) -> OpPlanIntoIter<Op> {
        OpPlanIntoIter {
            inner: self.plan.into_iter()
        }
    }
}

impl<Op> Iterator for OpPlanIntoIter<Op> {
    type Item = (Page, CompoundPageOperation<Op>);

    fn next(&mut self) -> Option<(Page, CompoundPageOperation<Op>)> {
        self.inner.next()
    }
}

struct OperationPlanResult {
    results: HashMap<Page, OperationResults>,
}

impl OperationPlanResult {
    fn new() -> OperationPlanResult {
        OperationPlanResult {
            results: HashMap::new(),
        }
    }

    fn insert(&mut self, page: Page, res: OperationResults) {
        self.results.insert(page, res);
    }

    fn iter(&self) -> OpPlanResultIter {
        OpPlanResultIter {
            inner: self.results.iter()
        }
    }
}

struct OpPlanResultIter<'a> {
    inner: hash_map::Iter<'a, Page, OperationResults>,
}

impl<'a> Iterator for OpPlanResultIter<'a> {
    type Item = (&'a Page, &'a OperationResults);

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }
}

/// IntoIterator implementation for OpPlanResultIter.
impl<'a> IntoIterator for &'a OperationPlanResult {
    type Item = (&'a Page, &'a OperationResults);
    type IntoIter = OpPlanResultIter<'a>;

    fn into_iter(self) -> OpPlanResultIter<'a> {
        self.iter()
    }
}

struct OpPlanResultIntoIter {
    inner: hash_map::IntoIter<Page, OperationResults>,
}

impl IntoIterator for OperationPlanResult {
    type Item = (Page, OperationResults);
    type IntoIter = OpPlanResultIntoIter;

    fn into_iter(self) -> OpPlanResultIntoIter {
        OpPlanResultIntoIter {
            inner: self.results.into_iter()
        }
    }
}

impl Iterator for OpPlanResultIntoIter {
    type Item = (Page, OperationResults);

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }
}
<|endoftext|>


<commit_before>use serde::Deserialize;
use std::collections::BTreeMap;
use std::path::Path;

use std::collections::BTreeSet as Set;
use std::fs;

/// A lookup table from TypeName to Entry.
pub type NodeTypeMap = BTreeMap<TypeName, Entry>;

#[derive(Debug)]
pub struct Entry {
    pub dbscheme_name: String,
    pub ql_class_name: String,
    pub kind: EntryKind,
}

#[derive(Debug)]
pub enum EntryKind {
    Union { members: Set<TypeName> },
    Table { name: String, fields: Vec<Field> },
    Token { kind_id: usize },
}

#[derive(Debug, Ord, PartialOrd, Eq, PartialEq)]
pub struct TypeName {
    pub kind: String,
    pub named: bool,
}

#[derive(Debug)]
pub enum FieldTypeInfo {
    /// The field has a single type.
    Single(TypeName),

    /// The field can take one of several types, so we also provide the name of
    /// the database union type that wraps them, and the corresponding QL class
    /// name.
    Multiple {
        types: Set<TypeName>,
        dbscheme_union: String,
        ql_class: String,
    },

    /// The field can be one of several tokens, so the db type will be an `int`
    /// with a `case @foo.kind` for each possiblity.
    ReservedWordInt(BTreeMap<String, (usize, String)>),
}

#[derive(Debug)]
pub struct Field {
    pub parent: TypeName,
    pub type_info: FieldTypeInfo,
    /// The name of the field or None for the anonymous 'children'
    /// entry from node_types.json
    pub name: Option<String>,
    /// The name of the predicate to get this field.
    pub getter_name: String,
    pub storage: Storage,
}

fn name_for_field_or_child(name: &Option<String>) -> String {
    match name {
        Some(name) => name.clone(),
        None => "child".to_owned(),
    }
}

#[derive(Debug)]
pub enum Storage {
    /// the field is stored as a column in the parent table
    Column { name: String },
    /// the field is stored in a link table, and may or may not have an
    /// associated index column
    Table { name: String, has_index: bool },
}

pub fn read_node_types(node_types_path: &Path) -> std::io::Result<NodeTypeMap> {
    let file = fs::File::open(node_types_path)?;
    let node_types = serde_json::from_reader(file)?;
    Ok(convert_nodes(&node_types))
}

pub fn read_node_types_str(node_types_json: &str) -> std::io::Result<NodeTypeMap> {
    let node_types = serde_json::from_str(node_types_json)?;
    Ok(convert_nodes(&node_types))
}

fn convert_type(node_type: &NodeType) -> TypeName {
    TypeName {
        kind: node_type.kind.to_string(),
        named: node_type.named,
    }
}

fn convert_types(node_types: &Vec<NodeType>) -> Set<TypeName> {
    let iter = node_types.iter().map(convert_type).collect();
    std::collections::BTreeSet::from(iter)
}

pub fn convert_nodes(nodes: &Vec<NodeInfo>) -> NodeTypeMap {
    let mut entries = NodeTypeMap::new();
    let mut token_kinds = Set::new();

    // First, find all the token kinds
    for node in nodes {
        if node.subtypes.is_none() {
            if node.fields.as_ref().map_or(0, |x| x.len()) == 0 && node.children.is_none() {
                let type_name = TypeName {
                    kind: node.kind.clone(),
                    named: node.named,
                };
                token_kinds.insert(type_name);
            }
        }
    }

    for node in nodes {
        let flattened_name = &node_type_name(&node.kind, node.named);
        let dbscheme_name = escape_name(&flattened_name);
        let ql_class_name = dbscheme_name_to_class_name(&dbscheme_name);
        if let Some(subtypes) = &node.subtypes {
            // It's a tree-sitter supertype node, for which we create a union
            // type.
            entries.insert(
                TypeName {
                    kind: node.kind.clone(),
                    named: node.named,
                },
                Entry {
                    dbscheme_name,
                    ql_class_name,
                    kind: EntryKind::Union {
                        members: convert_types(&subtypes),
                    },
                },
            );
        } else if node.fields.as_ref().map_or(0, |x| x.len()) == 0 && node.children.is_none() {
            // Token kind, handled above.
        } else {
            // It's a product type, defined by a table.
            let type_name = TypeName {
                kind: node.kind.clone(),
                named: node.named,
            };
            let table_name = escape_name(&(format!("{}_def", &flattened_name)));
            let mut fields = Vec::new();

            // If the type also has fields or children, then we create either
            // auxiliary tables or columns in the defining table for them.
            if let Some(node_fields) = &node.fields {
                for (field_name, field_info) in node_fields {
                    add_field(
                        &type_name,
                        Some(field_name.to_string()),
                        field_info,
                        &mut fields,
                        &token_kinds,
                    );
                }
            }
            if let Some(children) = &node.children {
                // Treat children as if they were a field called 'child'.
                add_field(&type_name, None, children, &mut fields, &token_kinds);
            }
            entries.insert(
                type_name,
                Entry {
                    dbscheme_name,
                    ql_class_name,
                    kind: EntryKind::Table {
                        name: table_name,
                        fields,
                    },
                },
            );
        }
    }
    let mut counter = 0;
    for type_name in token_kinds {
        let entry = if type_name.named {
            counter += 1;
            let unprefixed_name = node_type_name(&type_name.kind, true);
            Entry {
                dbscheme_name: escape_name(&format!("token_{}", &unprefixed_name)),
                ql_class_name: dbscheme_name_to_class_name(&escape_name(&unprefixed_name)),
                kind: EntryKind::Token { kind_id: counter },
            }
        } else {
            Entry {
                dbscheme_name: "reserved_word".to_owned(),
                ql_class_name: "ReservedWord".to_owned(),
                kind: EntryKind::Token { kind_id: 0 },
            }
        };
        entries.insert(type_name, entry);
    }
    entries
}

fn add_field(
    parent_type_name: &TypeName,
    field_name: Option<String>,
    field_info: &FieldInfo,
    fields: &mut Vec<Field>,
    token_kinds: &Set<TypeName>,
) {
    let parent_flattened_name = node_type_name(&parent_type_name.kind, parent_type_name.named);
    let storage = if !field_info.multiple && field_info.required {
        // This field must appear exactly once, so we add it as
        // a column to the main table for the node type.
        Storage::Column {
            name: escape_name(&name_for_field_or_child(&field_name)),
        }
    } else {
        // Put the field in an auxiliary table.
        let has_index = field_info.multiple;
        let field_table_name = escape_name(&format!(
            "{}_{}",
            parent_flattened_name,
            &name_for_field_or_child(&field_name)
        ));
        Storage::Table {
            has_index,
            name: field_table_name,
        }
    };
    let converted_types = convert_types(&field_info.types);
    let type_info = if field_info
        .types
        .iter()
        .all(|t| !t.named && token_kinds.contains(&convert_type(t)))
    {
        // All possible types for this field are reserved words. The db
        // representation will be an `int` with a `case @foo.field = ...` to
        // enumerate the possible values.
        let mut counter = 0;
        let mut field_token_ints: BTreeMap<String, (usize, String)> = BTreeMap::new();
        for t in converted_types {
            let dbscheme_variant_name =
                escape_name(&format!("{}_{}", parent_flattened_name, t.kind));
            field_token_ints.insert(t.kind.to_owned(), (counter, dbscheme_variant_name));
            counter += 1;
        }
        FieldTypeInfo::ReservedWordInt(field_token_ints)
    } else if field_info.types.len() == 1 {
        FieldTypeInfo::Single(converted_types.into_iter().next().unwrap())
    } else {
        // The dbscheme type for this field will be a union. In QL, it'll just be AstNode.
        FieldTypeInfo::Multiple {
            types: converted_types,
            dbscheme_union: format!(
                "{}_{}_type",
                &parent_flattened_name,
                &name_for_field_or_child(&field_name)
            ),
            ql_class: "AstNode".to_owned(),
        }
    };
    let getter_name = format!(
        "get{}",
        dbscheme_name_to_class_name(&escape_name(&name_for_field_or_child(&field_name)))
    );
    fields.push(Field {
        parent: TypeName {
            kind: parent_type_name.kind.to_string(),
            named: parent_type_name.named,
        },
        type_info,
        name: field_name,
        getter_name,
        storage,
    });
}
#[derive(Deserialize)]
pub struct NodeInfo {
    #[serde(rename = "type")]
    pub kind: String,
    pub named: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub fields: Option<BTreeMap<String, FieldInfo>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub children: Option<FieldInfo>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub subtypes: Option<Vec<NodeType>>,
}

#[derive(Deserialize)]
pub struct NodeType {
    #[serde(rename = "type")]
    pub kind: String,
    pub named: bool,
}

#[derive(Deserialize)]
pub struct FieldInfo {
    pub multiple: bool,
    pub required: bool,
    pub types: Vec<NodeType>,
}

/// Given a tree-sitter node type's (kind, named) pair, returns a single string
/// representing the (unescaped) name we'll use to refer to corresponding QL
/// type.
fn node_type_name(kind: &str, named: bool) -> String {
    if named {
        kind.to_string()
    } else {
        format!("{}_unnamed", kind)
    }
}

const RESERVED_KEYWORDS: [&'static str; 14] = [
    "boolean", "case", "date", "float", "int", "key", "of", "order", "ref", "string", "subtype",
    "type", "unique", "varchar",
];

/// Returns a string that's a copy of `name` but suitably escaped to be a valid
/// QL identifier.
fn escape_name(name: &str) -> String {
    let mut result = String::new();

    // If there's a leading underscore, replace it with 'underscore_'.
    if let Some(c) = name.chars().next() {
        if c == '_' {
            result.push_str("underscore");
        }
    }
    for c in name.chars() {
        match c {
            '{' => result.push_str("lbrace"),
            '}' => result.push_str("rbrace"),
            '<' => result.push_str("langle"),
            '>' => result.push_str("rangle"),
            '[' => result.push_str("lbracket"),
            ']' => result.push_str("rbracket"),
            '(' => result.push_str("lparen"),
            ')' => result.push_str("rparen"),
            '|' => result.push_str("pipe"),
            '=' => result.push_str("equal"),
            '~' => result.push_str("tilde"),
            '?' => result.push_str("question"),
            '`' => result.push_str("backtick"),
            '^' => result.push_str("caret"),
            '!' => result.push_str("bang"),
            '#' => result.push_str("hash"),
            '%' => result.push_str("percent"),
            '&' => result.push_str("ampersand"),
            '.' => result.push_str("dot"),
            ',' => result.push_str("comma"),
            '/' => result.push_str("slash"),
            ':' => result.push_str("colon"),
            ';' => result.push_str("semicolon"),
            '"' => result.push_str("dquote"),
            '*' => result.push_str("star"),
            '+' => result.push_str("plus"),
            '-' => result.push_str("minus"),
            '@' => result.push_str("at"),
            _ if c.is_uppercase() => {
                result.push('_');
                result.push_str(&c.to_lowercase().to_string())
            }
            _ => result.push(c),
        }
    }

    for &keyword in &RESERVED_KEYWORDS {
        if result == keyword {
            result.push_str("__");
            break;
        }
    }

    result
}

/// Given a valid dbscheme name (i.e. in snake case), produces the equivalent QL
/// name (i.e. in CamelCase). For example, "foo_bar_baz" becomes "FooBarBaz".
fn dbscheme_name_to_class_name(dbscheme_name: &str) -> String {
    fn to_title_case(word: &str) -> String {
        let mut first = true;
        let mut result = String::new();
        for c in word.chars() {
            if first {
                first = false;
                result.push(c.to_ascii_uppercase());
            } else {
                result.push(c);
            }
        }
        result
    }
    dbscheme_name
        .split('_')
        .map(|word| to_title_case(word))
        .collect::<Vec<String>>()
        .join("")
}
<commit_msg>Escape field names correctly<commit_after>use serde::Deserialize;
use std::collections::BTreeMap;
use std::path::Path;

use std::collections::BTreeSet as Set;
use std::fs;

/// A lookup table from TypeName to Entry.
pub type NodeTypeMap = BTreeMap<TypeName, Entry>;

#[derive(Debug)]
pub struct Entry {
    pub dbscheme_name: String,
    pub ql_class_name: String,
    pub kind: EntryKind,
}

#[derive(Debug)]
pub enum EntryKind {
    Union { members: Set<TypeName> },
    Table { name: String, fields: Vec<Field> },
    Token { kind_id: usize },
}

#[derive(Debug, Ord, PartialOrd, Eq, PartialEq)]
pub struct TypeName {
    pub kind: String,
    pub named: bool,
}

#[derive(Debug)]
pub enum FieldTypeInfo {
    /// The field has a single type.
    Single(TypeName),

    /// The field can take one of several types, so we also provide the name of
    /// the database union type that wraps them, and the corresponding QL class
    /// name.
    Multiple {
        types: Set<TypeName>,
        dbscheme_union: String,
        ql_class: String,
    },

    /// The field can be one of several tokens, so the db type will be an `int`
    /// with a `case @foo.kind` for each possiblity.
    ReservedWordInt(BTreeMap<String, (usize, String)>),
}

#[derive(Debug)]
pub struct Field {
    pub parent: TypeName,
    pub type_info: FieldTypeInfo,
    /// The name of the field or None for the anonymous 'children'
    /// entry from node_types.json
    pub name: Option<String>,
    /// The name of the predicate to get this field.
    pub getter_name: String,
    pub storage: Storage,
}

fn name_for_field_or_child(name: &Option<String>) -> String {
    match name {
        Some(name) => name.clone(),
        None => "child".to_owned(),
    }
}

#[derive(Debug)]
pub enum Storage {
    /// the field is stored as a column in the parent table
    Column { name: String },
    /// the field is stored in a link table, and may or may not have an
    /// associated index column
    Table { name: String, has_index: bool },
}

pub fn read_node_types(node_types_path: &Path) -> std::io::Result<NodeTypeMap> {
    let file = fs::File::open(node_types_path)?;
    let node_types = serde_json::from_reader(file)?;
    Ok(convert_nodes(&node_types))
}

pub fn read_node_types_str(node_types_json: &str) -> std::io::Result<NodeTypeMap> {
    let node_types = serde_json::from_str(node_types_json)?;
    Ok(convert_nodes(&node_types))
}

fn convert_type(node_type: &NodeType) -> TypeName {
    TypeName {
        kind: node_type.kind.to_string(),
        named: node_type.named,
    }
}

fn convert_types(node_types: &Vec<NodeType>) -> Set<TypeName> {
    let iter = node_types.iter().map(convert_type).collect();
    std::collections::BTreeSet::from(iter)
}

pub fn convert_nodes(nodes: &Vec<NodeInfo>) -> NodeTypeMap {
    let mut entries = NodeTypeMap::new();
    let mut token_kinds = Set::new();

    // First, find all the token kinds
    for node in nodes {
        if node.subtypes.is_none() {
            if node.fields.as_ref().map_or(0, |x| x.len()) == 0 && node.children.is_none() {
                let type_name = TypeName {
                    kind: node.kind.clone(),
                    named: node.named,
                };
                token_kinds.insert(type_name);
            }
        }
    }

    for node in nodes {
        let flattened_name = &node_type_name(&node.kind, node.named);
        let dbscheme_name = escape_name(&flattened_name);
        let ql_class_name = dbscheme_name_to_class_name(&dbscheme_name);
        if let Some(subtypes) = &node.subtypes {
            // It's a tree-sitter supertype node, for which we create a union
            // type.
            entries.insert(
                TypeName {
                    kind: node.kind.clone(),
                    named: node.named,
                },
                Entry {
                    dbscheme_name,
                    ql_class_name,
                    kind: EntryKind::Union {
                        members: convert_types(&subtypes),
                    },
                },
            );
        } else if node.fields.as_ref().map_or(0, |x| x.len()) == 0 && node.children.is_none() {
            // Token kind, handled above.
        } else {
            // It's a product type, defined by a table.
            let type_name = TypeName {
                kind: node.kind.clone(),
                named: node.named,
            };
            let table_name = escape_name(&(format!("{}_def", &flattened_name)));
            let mut fields = Vec::new();

            // If the type also has fields or children, then we create either
            // auxiliary tables or columns in the defining table for them.
            if let Some(node_fields) = &node.fields {
                for (field_name, field_info) in node_fields {
                    add_field(
                        &type_name,
                        Some(field_name.to_string()),
                        field_info,
                        &mut fields,
                        &token_kinds,
                    );
                }
            }
            if let Some(children) = &node.children {
                // Treat children as if they were a field called 'child'.
                add_field(&type_name, None, children, &mut fields, &token_kinds);
            }
            entries.insert(
                type_name,
                Entry {
                    dbscheme_name,
                    ql_class_name,
                    kind: EntryKind::Table {
                        name: table_name,
                        fields,
                    },
                },
            );
        }
    }
    let mut counter = 0;
    for type_name in token_kinds {
        let entry = if type_name.named {
            counter += 1;
            let unprefixed_name = node_type_name(&type_name.kind, true);
            Entry {
                dbscheme_name: escape_name(&format!("token_{}", &unprefixed_name)),
                ql_class_name: dbscheme_name_to_class_name(&escape_name(&unprefixed_name)),
                kind: EntryKind::Token { kind_id: counter },
            }
        } else {
            Entry {
                dbscheme_name: "reserved_word".to_owned(),
                ql_class_name: "ReservedWord".to_owned(),
                kind: EntryKind::Token { kind_id: 0 },
            }
        };
        entries.insert(type_name, entry);
    }
    entries
}

fn add_field(
    parent_type_name: &TypeName,
    field_name: Option<String>,
    field_info: &FieldInfo,
    fields: &mut Vec<Field>,
    token_kinds: &Set<TypeName>,
) {
    let parent_flattened_name = node_type_name(&parent_type_name.kind, parent_type_name.named);
    let storage = if !field_info.multiple && field_info.required {
        // This field must appear exactly once, so we add it as
        // a column to the main table for the node type.
        Storage::Column {
            name: escape_name(&name_for_field_or_child(&field_name)),
        }
    } else {
        // Put the field in an auxiliary table.
        let has_index = field_info.multiple;
        let field_table_name = escape_name(&format!(
            "{}_{}",
            parent_flattened_name,
            &name_for_field_or_child(&field_name)
        ));
        Storage::Table {
            has_index,
            name: field_table_name,
        }
    };
    let converted_types = convert_types(&field_info.types);
    let type_info = if field_info
        .types
        .iter()
        .all(|t| !t.named && token_kinds.contains(&convert_type(t)))
    {
        // All possible types for this field are reserved words. The db
        // representation will be an `int` with a `case @foo.field = ...` to
        // enumerate the possible values.
        let mut counter = 0;
        let mut field_token_ints: BTreeMap<String, (usize, String)> = BTreeMap::new();
        for t in converted_types {
            let dbscheme_variant_name =
                escape_name(&format!("{}_{}", parent_flattened_name, t.kind));
            field_token_ints.insert(t.kind.to_owned(), (counter, dbscheme_variant_name));
            counter += 1;
        }
        FieldTypeInfo::ReservedWordInt(field_token_ints)
    } else if field_info.types.len() == 1 {
        FieldTypeInfo::Single(converted_types.into_iter().next().unwrap())
    } else {
        // The dbscheme type for this field will be a union. In QL, it'll just be AstNode.
        FieldTypeInfo::Multiple {
            types: converted_types,
            dbscheme_union: format!(
                "{}_{}_type",
                &parent_flattened_name,
                &name_for_field_or_child(&field_name)
            ),
            ql_class: "AstNode".to_owned(),
        }
    };
    let getter_name = format!(
        "get{}",
        dbscheme_name_to_class_name(&escape_name(&name_for_field_or_child(&field_name)))
    );
    fields.push(Field {
        parent: TypeName {
            kind: parent_type_name.kind.to_string(),
            named: parent_type_name.named,
        },
        type_info,
        name: field_name.map(|name| escape_name(&name)),
        getter_name,
        storage,
    });
}
#[derive(Deserialize)]
pub struct NodeInfo {
    #[serde(rename = "type")]
    pub kind: String,
    pub named: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub fields: Option<BTreeMap<String, FieldInfo>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub children: Option<FieldInfo>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub subtypes: Option<Vec<NodeType>>,
}

#[derive(Deserialize)]
pub struct NodeType {
    #[serde(rename = "type")]
    pub kind: String,
    pub named: bool,
}

#[derive(Deserialize)]
pub struct FieldInfo {
    pub multiple: bool,
    pub required: bool,
    pub types: Vec<NodeType>,
}

/// Given a tree-sitter node type's (kind, named) pair, returns a single string
/// representing the (unescaped) name we'll use to refer to corresponding QL
/// type.
fn node_type_name(kind: &str, named: bool) -> String {
    if named {
        kind.to_string()
    } else {
        format!("{}_unnamed", kind)
    }
}

const RESERVED_KEYWORDS: [&'static str; 14] = [
    "boolean", "case", "date", "float", "int", "key", "of", "order", "ref", "string", "subtype",
    "type", "unique", "varchar",
];

/// Returns a string that's a copy of `name` but suitably escaped to be a valid
/// QL identifier.
fn escape_name(name: &str) -> String {
    let mut result = String::new();

    // If there's a leading underscore, replace it with 'underscore_'.
    if let Some(c) = name.chars().next() {
        if c == '_' {
            result.push_str("underscore");
        }
    }
    for c in name.chars() {
        match c {
            '{' => result.push_str("lbrace"),
            '}' => result.push_str("rbrace"),
            '<' => result.push_str("langle"),
            '>' => result.push_str("rangle"),
            '[' => result.push_str("lbracket"),
            ']' => result.push_str("rbracket"),
            '(' => result.push_str("lparen"),
            ')' => result.push_str("rparen"),
            '|' => result.push_str("pipe"),
            '=' => result.push_str("equal"),
            '~' => result.push_str("tilde"),
            '?' => result.push_str("question"),
            '`' => result.push_str("backtick"),
            '^' => result.push_str("caret"),
            '!' => result.push_str("bang"),
            '#' => result.push_str("hash"),
            '%' => result.push_str("percent"),
            '&' => result.push_str("ampersand"),
            '.' => result.push_str("dot"),
            ',' => result.push_str("comma"),
            '/' => result.push_str("slash"),
            ':' => result.push_str("colon"),
            ';' => result.push_str("semicolon"),
            '"' => result.push_str("dquote"),
            '*' => result.push_str("star"),
            '+' => result.push_str("plus"),
            '-' => result.push_str("minus"),
            '@' => result.push_str("at"),
            _ if c.is_uppercase() => {
                result.push('_');
                result.push_str(&c.to_lowercase().to_string())
            }
            _ => result.push(c),
        }
    }

    for &keyword in &RESERVED_KEYWORDS {
        if result == keyword {
            result.push_str("__");
            break;
        }
    }

    result
}

/// Given a valid dbscheme name (i.e. in snake case), produces the equivalent QL
/// name (i.e. in CamelCase). For example, "foo_bar_baz" becomes "FooBarBaz".
fn dbscheme_name_to_class_name(dbscheme_name: &str) -> String {
    fn to_title_case(word: &str) -> String {
        let mut first = true;
        let mut result = String::new();
        for c in word.chars() {
            if first {
                first = false;
                result.push(c.to_ascii_uppercase());
            } else {
                result.push(c);
            }
        }
        result
    }
    dbscheme_name
        .split('_')
        .map(|word| to_title_case(word))
        .collect::<Vec<String>>()
        .join("")
}
<|endoftext|>


<commit_before>use std::{
    collections::VecDeque,
    fmt,
    io,
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
    time::{Duration, Instant},
};

use crate::{
    config::{Config, ServerConfig},
    context::{Context, SharedContext},
    relay::{
        socks5::Address,
        tcprelay::client::ServerClient as TcpServerClient,
        udprelay::client::ServerClient as UdpServerClient,
    },
};

use log::{debug, info};
use tokio::{
    self,
    io::{AsyncReadExt, AsyncWriteExt},
    sync::{Barrier, Mutex},
    time,
};

const MAX_LATENCY_QUEUE_SIZE: usize = 37;
const DEFAULT_CHECK_INTERVAL_SEC: u64 = 6;
const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 2; // Latency shouldn't greater than 2 secs, that's too long

/// Identifier of a valid server
pub trait ServerData: Send + Sync {
    fn create_server(context: &SharedContext, server_idx: usize, data: &SharedServerStatisticData) -> Self;
}

struct ServerStatisticData {
    /// Median of latency time
    ///
    /// Use median instead of average time,
    /// because probing result may have some really bad cases
    rtt: u64,
    /// Total_Fail / Total_Probe
    fail_rate: f64,
    /// Recently probe data
    latency_queue: VecDeque<Score>,
}

impl ServerStatisticData {
    fn new() -> ServerStatisticData {
        ServerStatisticData {
            rtt: u64::max_value(),
            fail_rate: 100.0,
            latency_queue: VecDeque::new(),
        }
    }

    fn score(&self) -> u64 {
        // Normalize rtt
        let nrtt = self.rtt as f64 / DEFAULT_CHECK_TIMEOUT_SEC as f64;

        // Score = norm_lat + prop_err
        //
        // 1. The lower latency, the better
        // 2. The lower errored count, the better
        ((nrtt + self.fail_rate) * 1000.0) as u64
    }

    fn push_score(&mut self, score: Score) -> u64 {
        self.latency_queue.push_back(score);

        // Only records recently MAX_LATENCY_QUEUE_SIZE probe data
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.recalculate_score()
    }

    fn recalculate_score(&mut self) -> u64 {
        let mut vlat = Vec::with_capacity(self.latency_queue.len());
        let mut cerr = 0;
        for s in &self.latency_queue {
            match *s {
                Score::Errored => cerr += 1,
                Score::Latency(lat) => vlat.push(lat),
            }
        }

        if !vlat.is_empty() {
            vlat.sort();

            // Find median of latency
            let mid = vlat.len() / 2;

            self.rtt = if vlat.len() % 2 == 0 {
                (vlat[mid] + vlat[mid - 1]) / 2
            } else {
                vlat[mid]
            };
        }

        // Error rate
        self.fail_rate = cerr as f64 / self.latency_queue.len() as f64;

        self.score()
    }

    pub fn report_failure(&mut self) -> u64 {
        self.push_score(Score::Errored)
    }
}

/// Shared handle for mutating server's statistic data
#[derive(Clone)]
pub struct SharedServerStatisticData(Arc<Mutex<ServerStatisticData>>);

impl SharedServerStatisticData {
    fn new() -> SharedServerStatisticData {
        SharedServerStatisticData(Arc::new(Mutex::new(ServerStatisticData::new())))
    }

    pub async fn report_failure(&self) -> u64 {
        let mut data = self.0.lock().await;
        data.report_failure()
    }

    async fn push_score(&self, score: Score) -> u64 {
        let mut data = self.0.lock().await;
        data.push_score(score)
    }

    pub async fn score(&self) -> u64 {
        let data = self.0.lock().await;
        data.score()
    }
}

/// Server Statistic scores
pub struct ServerStatistic<S: ServerData> {
    server: S,
    context: SharedContext,
    server_idx: usize,
    data: SharedServerStatisticData,
}

pub type SharedServerStatistic<S> = Arc<ServerStatistic<S>>;

impl<S: ServerData> ServerStatistic<S> {
    fn new(context: SharedContext, server_idx: usize) -> ServerStatistic<S> {
        let data = SharedServerStatisticData::new();

        ServerStatistic {
            server: S::create_server(&context, server_idx, &data),
            context,
            server_idx,
            data,
        }
    }

    fn new_shared(context: SharedContext, server_idx: usize) -> SharedServerStatistic<S> {
        Arc::new(ServerStatistic::new(context, server_idx))
    }

    pub fn server_config(&self) -> &ServerConfig {
        self.context.server_config(self.server_idx)
    }

    pub fn server(&self) -> &S {
        &self.server
    }

    pub fn context(&self) -> &Context {
        &*self.context
    }

    pub fn config(&self) -> &Config {
        self.context.config()
    }

    async fn push_score(&self, score: Score) -> u64 {
        self.data.push_score(score).await
    }

    pub async fn score(&self) -> u64 {
        self.data.score().await
    }

    pub async fn report_failure(&self) -> u64 {
        self.data.report_failure().await
    }
}

#[derive(Debug, Copy, Clone)]
enum Score {
    Latency(u64),
    Errored,
}

#[derive(Debug, Clone, Copy)]
pub enum ServerType {
    Tcp,
    Udp,
}

impl fmt::Display for ServerType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ServerType::Tcp => f.write_str("TCP"),
            ServerType::Udp => f.write_str("UDP"),
        }
    }
}

struct BestServer<S: ServerData> {
    servers: Vec<SharedServerStatistic<S>>,
    best_idx: AtomicUsize,
}

type SharedBestServer<S> = Arc<BestServer<S>>;

impl<S: ServerData> BestServer<S> {
    fn new(servers: Vec<SharedServerStatistic<S>>) -> BestServer<S> {
        BestServer {
            servers,
            best_idx: AtomicUsize::new(0),
        }
    }

    fn new_shared(servers: Vec<SharedServerStatistic<S>>) -> SharedBestServer<S> {
        Arc::new(BestServer::new(servers))
    }

    fn pick_server(&self) -> SharedServerStatistic<S> {
        let idx = self.best_idx.load(Ordering::Relaxed);
        self.servers[idx].clone()
    }

    async fn recalculate_best_server(&self) -> Option<(usize, usize)> {
        let current_best_idx = self.best_idx.load(Ordering::Relaxed);

        let mut best_idx = 0;
        let mut best_score = u64::max_value();

        for (idx, svr) in self.servers.iter().enumerate() {
            let score = svr.score().await;
            if score < best_score {
                best_idx = idx;
                best_score = score;
            }
        }

        if best_idx != current_best_idx {
            self.best_idx.store(best_idx, Ordering::Relaxed);

            Some((current_best_idx, best_idx))
        } else {
            None
        }
    }
}

/// Load balancer based on pinging latencies of all servers
#[derive(Clone)]
pub struct PingBalancer<S: ServerData> {
    best: SharedBestServer<S>,
}

impl<S: ServerData + 'static> PingBalancer<S> {
    /// Create a PingBalancer
    pub async fn new(context: SharedContext, server_type: ServerType) -> PingBalancer<S> {
        let server_count = context.config().server.len();
        let mut servers = Vec::with_capacity(server_count);

        // Check only required if servers count > 1, otherwise, always use the first one
        let check_required = server_count > 1;
        // Barrier count = current + probing tasks
        let check_barrier = Arc::new(Barrier::new(1 + server_count));

        for idx in 0..server_count {
            let stat = ServerStatistic::<S>::new_shared(context.clone(), idx);

            if check_required {
                let stat = stat.clone();
                let context = context.clone();
                let check_barrier = check_barrier.clone();

                // Start a background task for probing
                tokio::spawn(async move {
                    // Check once for initializing data
                    PingBalancer::<S>::check_update_score(&stat, server_type).await;
                    check_barrier.wait().await;

                    while context.server_running() {
                        PingBalancer::<S>::check_update_score(&stat, server_type).await;
                        time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
                    }

                    debug!(
                        "Probing task for remote {} server {} exited",
                        server_type,
                        stat.server_config().addr()
                    );
                });
            }

            servers.push(stat);
        }

        let best = BestServer::new_shared(servers);

        if check_required {
            // Wait all tasks start (run at least one round)
            check_barrier.wait().await;

            // Reinitialize a Barrier for waiting choosing task
            let check_barrier = Arc::new(Barrier::new(2));
            let best = best.clone();

            {
                let context = context.clone();
                let check_barrier = check_barrier.clone();

                tokio::spawn(async move {
                    // Check once for initializing data
                    best.recalculate_best_server().await;
                    check_barrier.wait().await;

                    while context.server_running() {
                        if let Some((old_idx, new_idx)) = best.recalculate_best_server().await {
                            info!(
                                "Switched {} server from {} to {}",
                                server_type,
                                context.server_config(old_idx).addr(),
                                context.server_config(new_idx).addr()
                            );
                        }

                        time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
                    }
                });
            }

            // Wait for choosing task to check at least once
            check_barrier.wait().await;
        }

        PingBalancer { best }
    }

    async fn check_update_score(stat: &ServerStatistic<S>, server_type: ServerType) {
        let score = match PingBalancer::<S>::check_delay(stat, server_type).await {
            Ok(d) => stat.push_score(Score::Latency(d)).await,
            Err(..) => stat.push_score(Score::Errored).await, // Penalty
        };

        debug!(
            "Updated remote {} server {} (score: {})",
            server_type,
            stat.server_config().addr(),
            score
        );
    }

    async fn check_request_tcp(stat: &ServerStatistic<S>) -> io::Result<()> {
        static GET_BODY: &[u8] =
            b"GET /generate_204 HTTP/1.1\r\nHost: dl.google.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";

        let addr = Address::DomainNameAddress("dl.google.com".to_owned(), 80);

        let TcpServerClient { mut stream } =
            TcpServerClient::connect(stat.context(), &addr, stat.server_config()).await?;
        stream.write_all(GET_BODY).await?;
        stream.flush().await?;
        let mut buf = [0u8; 1];
        stream.read_exact(&mut buf).await?;

        Ok(())
    }

    async fn check_request_udp(stat: &ServerStatistic<S>) -> io::Result<()> {
        static DNS_QUERY: &[u8] = b"\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\x05\x62\x61\x69\x64\x75\x03\x63\x6f\x6d\x00\x00\x01\x00\x01";

        let addr = Address::SocketAddress(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53));

        let mut client = UdpServerClient::new(stat.server_config()).await?;
        client.send_to(stat.context(), &addr, DNS_QUERY).await?;
        let _ = client.recv_from(stat.context()).await?;

        Ok(())
    }

    async fn check_request(stat: &ServerStatistic<S>, server_type: ServerType) -> io::Result<()> {
        match server_type {
            ServerType::Tcp => PingBalancer::<S>::check_request_tcp(stat).await,
            ServerType::Udp => PingBalancer::<S>::check_request_udp(stat).await,
        }
    }

    async fn check_delay(stat: &ServerStatistic<S>, server_type: ServerType) -> io::Result<u64> {
        let start = Instant::now();

        // Send HTTP GET and read the first byte
        let timeout = Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC);
        let res = time::timeout(timeout, PingBalancer::<S>::check_request(stat, server_type)).await;

        let elapsed = Instant::now() - start;
        let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
        match res {
            Ok(Ok(..)) => {
                // Got the result ... record its time
                debug!(
                    "Checked remote {} server {} latency with {} ms",
                    server_type,
                    stat.server_config().addr(),
                    elapsed
                );
                Ok(elapsed)
            }
            Ok(Err(err)) => {
                debug!(
                    "Failed to check {} server {}, error: {}",
                    server_type,
                    stat.server_config().addr(),
                    err
                );

                // NOTE: connection / handshake error, server is down
                Err(err)
            }
            Err(..) => {
                // Timeout
                debug!(
                    "Checked remote {} server {} latency timeout, elapsed {} ms",
                    server_type,
                    stat.server_config().addr(),
                    elapsed
                );

                // NOTE: timeout is still available, but server is too slow
                Ok(elapsed)
            }
        }
    }
}

impl<S: ServerData> PingBalancer<S> {
    /// Pick the best server with current known statistic data
    ///
    /// Return a `Arc` shared server statistic reference
    pub fn pick_server(&self) -> SharedServerStatistic<S> {
        self.best.pick_server()
    }
}

/// A default struct for default ping balancer
pub struct EmptyServerData;

impl ServerData for EmptyServerData {
    fn create_server(_: &SharedContext, _: usize, _: &SharedServerStatisticData) -> EmptyServerData {
        EmptyServerData
    }
}

/// A PingBalancer without customized ServerData
pub type PlainPingBalancer = PingBalancer<EmptyServerData>;

/// Shared PlainServerStatistic
pub type SharedPlainServerStatistic = SharedServerStatistic<EmptyServerData>;
<commit_msg>RTT is millisecond<commit_after>use std::{
    collections::VecDeque,
    fmt,
    io,
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
    time::{Duration, Instant},
};

use crate::{
    config::{Config, ServerConfig},
    context::{Context, SharedContext},
    relay::{
        socks5::Address,
        tcprelay::client::ServerClient as TcpServerClient,
        udprelay::client::ServerClient as UdpServerClient,
    },
};

use log::{debug, info};
use tokio::{
    self,
    io::{AsyncReadExt, AsyncWriteExt},
    sync::{Barrier, Mutex},
    time,
};

const MAX_LATENCY_QUEUE_SIZE: usize = 37;
const DEFAULT_CHECK_INTERVAL_SEC: u64 = 6;
const DEFAULT_CHECK_TIMEOUT_SEC: u64 = 2; // Latency shouldn't greater than 2 secs, that's too long

/// Identifier of a valid server
pub trait ServerData: Send + Sync {
    fn create_server(context: &SharedContext, server_idx: usize, data: &SharedServerStatisticData) -> Self;
}

struct ServerStatisticData {
    /// Median of latency time
    ///
    /// Use median instead of average time,
    /// because probing result may have some really bad cases
    rtt: u64,
    /// Total_Fail / Total_Probe
    fail_rate: f64,
    /// Recently probe data
    latency_queue: VecDeque<Score>,
}

impl ServerStatisticData {
    fn new() -> ServerStatisticData {
        ServerStatisticData {
            rtt: u64::max_value(),
            fail_rate: 100.0,
            latency_queue: VecDeque::new(),
        }
    }

    fn score(&self) -> u64 {
        // Normalize rtt
        let nrtt = self.rtt as f64 / (DEFAULT_CHECK_TIMEOUT_SEC * 1000) as f64;

        // Score = norm_lat + prop_err
        //
        // 1. The lower latency, the better
        // 2. The lower errored count, the better
        ((nrtt + self.fail_rate) * 1000.0) as u64
    }

    fn push_score(&mut self, score: Score) -> u64 {
        self.latency_queue.push_back(score);

        // Only records recently MAX_LATENCY_QUEUE_SIZE probe data
        if self.latency_queue.len() > MAX_LATENCY_QUEUE_SIZE {
            self.latency_queue.pop_front();
        }

        self.recalculate_score()
    }

    fn recalculate_score(&mut self) -> u64 {
        let mut vlat = Vec::with_capacity(self.latency_queue.len());
        let mut cerr = 0;
        for s in &self.latency_queue {
            match *s {
                Score::Errored => cerr += 1,
                Score::Latency(lat) => vlat.push(lat),
            }
        }

        if !vlat.is_empty() {
            vlat.sort();

            // Find median of latency
            let mid = vlat.len() / 2;

            self.rtt = if vlat.len() % 2 == 0 {
                (vlat[mid] + vlat[mid - 1]) / 2
            } else {
                vlat[mid]
            };
        }

        // Error rate
        self.fail_rate = cerr as f64 / self.latency_queue.len() as f64;

        self.score()
    }

    pub fn report_failure(&mut self) -> u64 {
        self.push_score(Score::Errored)
    }
}

/// Shared handle for mutating server's statistic data
#[derive(Clone)]
pub struct SharedServerStatisticData(Arc<Mutex<ServerStatisticData>>);

impl SharedServerStatisticData {
    fn new() -> SharedServerStatisticData {
        SharedServerStatisticData(Arc::new(Mutex::new(ServerStatisticData::new())))
    }

    pub async fn report_failure(&self) -> u64 {
        let mut data = self.0.lock().await;
        data.report_failure()
    }

    async fn push_score(&self, score: Score) -> u64 {
        let mut data = self.0.lock().await;
        data.push_score(score)
    }

    pub async fn score(&self) -> u64 {
        let data = self.0.lock().await;
        data.score()
    }
}

/// Server Statistic scores
pub struct ServerStatistic<S: ServerData> {
    server: S,
    context: SharedContext,
    server_idx: usize,
    data: SharedServerStatisticData,
}

pub type SharedServerStatistic<S> = Arc<ServerStatistic<S>>;

impl<S: ServerData> ServerStatistic<S> {
    fn new(context: SharedContext, server_idx: usize) -> ServerStatistic<S> {
        let data = SharedServerStatisticData::new();

        ServerStatistic {
            server: S::create_server(&context, server_idx, &data),
            context,
            server_idx,
            data,
        }
    }

    fn new_shared(context: SharedContext, server_idx: usize) -> SharedServerStatistic<S> {
        Arc::new(ServerStatistic::new(context, server_idx))
    }

    pub fn server_config(&self) -> &ServerConfig {
        self.context.server_config(self.server_idx)
    }

    pub fn server(&self) -> &S {
        &self.server
    }

    pub fn context(&self) -> &Context {
        &*self.context
    }

    pub fn config(&self) -> &Config {
        self.context.config()
    }

    async fn push_score(&self, score: Score) -> u64 {
        self.data.push_score(score).await
    }

    pub async fn score(&self) -> u64 {
        self.data.score().await
    }

    pub async fn report_failure(&self) -> u64 {
        self.data.report_failure().await
    }
}

#[derive(Debug, Copy, Clone)]
enum Score {
    Latency(u64),
    Errored,
}

#[derive(Debug, Clone, Copy)]
pub enum ServerType {
    Tcp,
    Udp,
}

impl fmt::Display for ServerType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            ServerType::Tcp => f.write_str("TCP"),
            ServerType::Udp => f.write_str("UDP"),
        }
    }
}

struct BestServer<S: ServerData> {
    servers: Vec<SharedServerStatistic<S>>,
    best_idx: AtomicUsize,
}

type SharedBestServer<S> = Arc<BestServer<S>>;

impl<S: ServerData> BestServer<S> {
    fn new(servers: Vec<SharedServerStatistic<S>>) -> BestServer<S> {
        BestServer {
            servers,
            best_idx: AtomicUsize::new(0),
        }
    }

    fn new_shared(servers: Vec<SharedServerStatistic<S>>) -> SharedBestServer<S> {
        Arc::new(BestServer::new(servers))
    }

    fn pick_server(&self) -> SharedServerStatistic<S> {
        let idx = self.best_idx.load(Ordering::Relaxed);
        self.servers[idx].clone()
    }

    async fn recalculate_best_server(&self) -> Option<(usize, usize)> {
        let current_best_idx = self.best_idx.load(Ordering::Relaxed);

        let mut best_idx = 0;
        let mut best_score = u64::max_value();

        for (idx, svr) in self.servers.iter().enumerate() {
            let score = svr.score().await;
            if score < best_score {
                best_idx = idx;
                best_score = score;
            }
        }

        if best_idx != current_best_idx {
            self.best_idx.store(best_idx, Ordering::Relaxed);

            Some((current_best_idx, best_idx))
        } else {
            None
        }
    }
}

/// Load balancer based on pinging latencies of all servers
#[derive(Clone)]
pub struct PingBalancer<S: ServerData> {
    best: SharedBestServer<S>,
}

impl<S: ServerData + 'static> PingBalancer<S> {
    /// Create a PingBalancer
    pub async fn new(context: SharedContext, server_type: ServerType) -> PingBalancer<S> {
        let server_count = context.config().server.len();
        let mut servers = Vec::with_capacity(server_count);

        // Check only required if servers count > 1, otherwise, always use the first one
        let check_required = server_count > 1;
        // Barrier count = current + probing tasks
        let check_barrier = Arc::new(Barrier::new(1 + server_count));

        for idx in 0..server_count {
            let stat = ServerStatistic::<S>::new_shared(context.clone(), idx);

            if check_required {
                let stat = stat.clone();
                let context = context.clone();
                let check_barrier = check_barrier.clone();

                // Start a background task for probing
                tokio::spawn(async move {
                    // Check once for initializing data
                    PingBalancer::<S>::check_update_score(&stat, server_type).await;
                    check_barrier.wait().await;

                    while context.server_running() {
                        PingBalancer::<S>::check_update_score(&stat, server_type).await;
                        time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
                    }

                    debug!(
                        "Probing task for remote {} server {} exited",
                        server_type,
                        stat.server_config().addr()
                    );
                });
            }

            servers.push(stat);
        }

        let best = BestServer::new_shared(servers);

        if check_required {
            // Wait all tasks start (run at least one round)
            check_barrier.wait().await;

            // Reinitialize a Barrier for waiting choosing task
            let check_barrier = Arc::new(Barrier::new(2));
            let best = best.clone();

            {
                let context = context.clone();
                let check_barrier = check_barrier.clone();

                tokio::spawn(async move {
                    // Check once for initializing data
                    best.recalculate_best_server().await;
                    check_barrier.wait().await;

                    while context.server_running() {
                        if let Some((old_idx, new_idx)) = best.recalculate_best_server().await {
                            info!(
                                "Switched {} server from {} to {}",
                                server_type,
                                context.server_config(old_idx).addr(),
                                context.server_config(new_idx).addr()
                            );
                        }

                        time::delay_for(Duration::from_secs(DEFAULT_CHECK_INTERVAL_SEC)).await;
                    }
                });
            }

            // Wait for choosing task to check at least once
            check_barrier.wait().await;
        }

        PingBalancer { best }
    }

    async fn check_update_score(stat: &ServerStatistic<S>, server_type: ServerType) {
        let score = match PingBalancer::<S>::check_delay(stat, server_type).await {
            Ok(d) => stat.push_score(Score::Latency(d)).await,
            Err(..) => stat.push_score(Score::Errored).await, // Penalty
        };

        debug!(
            "Updated remote {} server {} (score: {})",
            server_type,
            stat.server_config().addr(),
            score
        );
    }

    async fn check_request_tcp(stat: &ServerStatistic<S>) -> io::Result<()> {
        static GET_BODY: &[u8] =
            b"GET /generate_204 HTTP/1.1\r\nHost: dl.google.com\r\nConnection: close\r\nAccept: */*\r\n\r\n";

        let addr = Address::DomainNameAddress("dl.google.com".to_owned(), 80);

        let TcpServerClient { mut stream } =
            TcpServerClient::connect(stat.context(), &addr, stat.server_config()).await?;
        stream.write_all(GET_BODY).await?;
        stream.flush().await?;
        let mut buf = [0u8; 1];
        stream.read_exact(&mut buf).await?;

        Ok(())
    }

    async fn check_request_udp(stat: &ServerStatistic<S>) -> io::Result<()> {
        static DNS_QUERY: &[u8] = b"\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\x05\x62\x61\x69\x64\x75\x03\x63\x6f\x6d\x00\x00\x01\x00\x01";

        let addr = Address::SocketAddress(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53));

        let mut client = UdpServerClient::new(stat.server_config()).await?;
        client.send_to(stat.context(), &addr, DNS_QUERY).await?;
        let _ = client.recv_from(stat.context()).await?;

        Ok(())
    }

    async fn check_request(stat: &ServerStatistic<S>, server_type: ServerType) -> io::Result<()> {
        match server_type {
            ServerType::Tcp => PingBalancer::<S>::check_request_tcp(stat).await,
            ServerType::Udp => PingBalancer::<S>::check_request_udp(stat).await,
        }
    }

    async fn check_delay(stat: &ServerStatistic<S>, server_type: ServerType) -> io::Result<u64> {
        let start = Instant::now();

        // Send HTTP GET and read the first byte
        let timeout = Duration::from_secs(DEFAULT_CHECK_TIMEOUT_SEC);
        let res = time::timeout(timeout, PingBalancer::<S>::check_request(stat, server_type)).await;

        let elapsed = Instant::now() - start;
        let elapsed = elapsed.as_secs() * 1000 + u64::from(elapsed.subsec_millis()); // Converted to ms
        match res {
            Ok(Ok(..)) => {
                // Got the result ... record its time
                debug!(
                    "Checked remote {} server {} latency with {} ms",
                    server_type,
                    stat.server_config().addr(),
                    elapsed
                );
                Ok(elapsed)
            }
            Ok(Err(err)) => {
                debug!(
                    "Failed to check {} server {}, error: {}",
                    server_type,
                    stat.server_config().addr(),
                    err
                );

                // NOTE: connection / handshake error, server is down
                Err(err)
            }
            Err(..) => {
                // Timeout
                debug!(
                    "Checked remote {} server {} latency timeout, elapsed {} ms",
                    server_type,
                    stat.server_config().addr(),
                    elapsed
                );

                // NOTE: timeout is still available, but server is too slow
                Ok(elapsed)
            }
        }
    }
}

impl<S: ServerData> PingBalancer<S> {
    /// Pick the best server with current known statistic data
    ///
    /// Return a `Arc` shared server statistic reference
    pub fn pick_server(&self) -> SharedServerStatistic<S> {
        self.best.pick_server()
    }
}

/// A default struct for default ping balancer
pub struct EmptyServerData;

impl ServerData for EmptyServerData {
    fn create_server(_: &SharedContext, _: usize, _: &SharedServerStatisticData) -> EmptyServerData {
        EmptyServerData
    }
}

/// A PingBalancer without customized ServerData
pub type PlainPingBalancer = PingBalancer<EmptyServerData>;

/// Shared PlainServerStatistic
pub type SharedPlainServerStatistic = SharedServerStatistic<EmptyServerData>;
<|endoftext|>


<commit_before>use hyper::Client;
use hyper::header::{Headers, Accept, qitem};
use hyper::mime::{Mime, TopLevel, SubLevel};
use hyper::status::StatusCode;
use std::io::Read;
use serde_json;

use Stream;
use types::Result;
use error::HesError;
use error::UserErrorKind;
use connection::ConnectionInfo;
use api::ESResolveLinkTos;

pub struct Reader<'a> {
    connection_info: &'a ConnectionInfo,
}

impl<'a> Reader<'a> {
    pub fn new(connection_info: &'a ConnectionInfo) -> Reader {
        Reader { connection_info: connection_info }
    }

    pub fn read_stream_events_forward(&self,
                                      stream_name: &str,
                                      start: u32,
                                      count: u32,
                                      resolve_link_tos: bool)
                                      -> Result<Stream> {
        let http_client = Client::default();

        let mut headers = Headers::new();
        headers.set(
            Accept(vec![
            qitem(Mime(TopLevel::Application,
                       SubLevel::Ext("vnd.eventstore.atom+json".to_owned()), vec![]))]));
        headers.set(ESResolveLinkTos(resolve_link_tos));

        let mut response = try!(http_client.get(&self.url(stream_name, start, count))
            .headers(headers)
            .send());
        match response.status {
            StatusCode::Ok => {
                let mut body = String::new();
                try!(response.read_to_string(&mut body));
                let stream: Stream = try!(serde_json::from_str(&body));
                Ok(stream)
            },
            StatusCode::NotFound => {
                Err(HesError::UserError(UserErrorKind::StreamNotFound))
            },
            StatusCode::Gone => {
                Err(HesError::UserError(UserErrorKind::StreamDeleted))
            },
            _ => {
                Err(HesError::UserError(UserErrorKind::UnexpectedResponse(response)))
            }
        }
    }

    fn url(&self, stream_name: &str, start: u32, count: u32) -> String {
        format!("http://{}:{}/streams/{}/{}/forward/{}?embed=body",
                self.connection_info.host,
                self.connection_info.port,
                stream_name,
                start,
                count)
    }
}
<commit_msg>Refactor Reader<commit_after>use hyper::Client;
use hyper::client::Response as HyperResponse;
use hyper::header::{Headers, Accept, qitem};
use hyper::mime::{Mime, TopLevel, SubLevel};
use hyper::status::StatusCode;
use std::io::Read;
use serde_json;

use Stream;
use types::Result;
use error::HesError;
use error::UserErrorKind;
use connection::ConnectionInfo;
use api::ESResolveLinkTos;

pub struct Reader<'a> {
    connection_info: &'a ConnectionInfo,
}

impl<'a> Reader<'a> {
    pub fn new(connection_info: &'a ConnectionInfo) -> Reader {
        Reader { connection_info: connection_info }
    }

    pub fn read_stream_events_forward(&self,
                                      stream_name: &str,
                                      start: u32,
                                      count: u32,
                                      resolve_link_tos: bool)
                                      -> Result<Stream> {
        let http_client = Client::default();

        let mut headers = Headers::new();
        headers.set(
            Accept(vec![
            qitem(Mime(TopLevel::Application,
                       SubLevel::Ext("vnd.eventstore.atom+json".to_owned()), vec![]))]));
        headers.set(ESResolveLinkTos(resolve_link_tos));

        let response = try!(http_client.get(&self.url(stream_name, start, count))
            .headers(headers)
            .send());

        to_hes_result(response)
    }

    fn url(&self, stream_name: &str, start: u32, count: u32) -> String {
        format!("http://{}:{}/streams/{}/{}/forward/{}?embed=body",
                self.connection_info.host,
                self.connection_info.port,
                stream_name,
                start,
                count)
    }
}

fn to_hes_result(mut response: HyperResponse) -> Result<Stream> {
    match response.status {
        StatusCode::Ok => {
            let mut body = String::new();
            try!(response.read_to_string(&mut body));
            let stream: Stream = try!(serde_json::from_str(&body));
            Ok(stream)
        },
        StatusCode::NotFound => {
            Err(HesError::UserError(UserErrorKind::StreamNotFound))
        },
        StatusCode::Gone => {
            Err(HesError::UserError(UserErrorKind::StreamDeleted))
        },
        _ => {
            Err(HesError::UserError(UserErrorKind::UnexpectedResponse(response)))
        }
    }
}
<|endoftext|>


<commit_before>//! Procedural macro implementing `#[derive(EnumMap)]`
//!
//! This is supposed to used with `enum-map` crate, which provides the
//! actual usage documentation.

#![recursion_limit = "128"]

extern crate proc_macro;
extern crate syn;
#[macro_use]
extern crate quote;

use std::iter;

use proc_macro::TokenStream;
use syn::{Body, Ident, Variant, VariantData};
use quote::Tokens;

fn generate_enum_code(name: &Ident, variants: &[Variant]) -> Tokens {
    let mut enum_count = 0usize;
    for &Variant {
            ref data,
            ref discriminant,
            ..
        } in variants {
        if data != &VariantData::Unit {
            panic!("#[derive(EnumMap)] requires C style style enum");
        }
        if discriminant.is_some() {
            panic!("#[derive(EnumMap)] doesn't currently support discriminants");
        }
        enum_count += 1;
    }

    let variant_a = variants.iter().map(|variant| &variant.ident);
    let variant_b = variants.iter().map(|variant| &variant.ident);
    let repeat_name_a = iter::repeat(name).take(variants.len());
    let repeat_name_b = repeat_name_a.clone();
    let counter = 0..variants.len();

    let to_usize = if variants.len() == 0 {
        quote! { unreachable!() }
    } else {
        quote! { self as usize }
    };

    quote! {
        impl<V> ::enum_map::Internal<V> for #name {
            type Array = [V; #enum_count];
            fn slice(array: &Self::Array) -> &[V] {
                array
            }
            fn slice_mut(array: &mut Self::Array) -> &mut [V] {
                array
            }
            fn from_usize(value: usize) -> Self {
                match value {
                    #(
                        #counter => #repeat_name_a::#variant_a,
                    )*
                    _ => unreachable!()
                }
            }
            fn to_usize(self) -> usize {
                #to_usize
            }
            fn from_function<F: FnMut(Self) -> V>(mut _f: F) -> Self::Array {
                [#(
                    _f(#repeat_name_b::#variant_b),
                )*]
            }
        }
    }
}

#[proc_macro_derive(EnumMap)]
pub fn derive_enum_map(input: TokenStream) -> TokenStream {
    let input = syn::parse_macro_input(&input.to_string()).unwrap();
    match input.body {
            Body::Enum(ref variants) => generate_enum_code(&input.ident, variants),
            _ => panic!("#[derive(EnumMap)] is only defined for enums"),
        }
        .parse()
        .unwrap()
}
<commit_msg>Reformat enum-map-derive<commit_after>//! Procedural macro implementing `#[derive(EnumMap)]`
//!
//! This is supposed to used with `enum-map` crate, which provides the
//! actual usage documentation.

#![recursion_limit = "128"]

extern crate proc_macro;
extern crate syn;
#[macro_use]
extern crate quote;

use std::iter;

use proc_macro::TokenStream;
use syn::{Body, Ident, Variant, VariantData};
use quote::Tokens;

fn generate_enum_code(name: &Ident, variants: &[Variant]) -> Tokens {
    let mut enum_count = 0usize;
    for &Variant {
        ref data,
        ref discriminant,
        ..
    } in variants
    {
        if data != &VariantData::Unit {
            panic!("#[derive(EnumMap)] requires C style style enum");
        }
        if discriminant.is_some() {
            panic!("#[derive(EnumMap)] doesn't currently support discriminants");
        }
        enum_count += 1;
    }

    let variant_a = variants.iter().map(|variant| &variant.ident);
    let variant_b = variants.iter().map(|variant| &variant.ident);
    let repeat_name_a = iter::repeat(name).take(variants.len());
    let repeat_name_b = repeat_name_a.clone();
    let counter = 0..variants.len();

    let to_usize = if variants.len() == 0 {
        quote! { unreachable!() }
    } else {
        quote! { self as usize }
    };

    quote! {
        impl<V> ::enum_map::Internal<V> for #name {
            type Array = [V; #enum_count];
            fn slice(array: &Self::Array) -> &[V] {
                array
            }
            fn slice_mut(array: &mut Self::Array) -> &mut [V] {
                array
            }
            fn from_usize(value: usize) -> Self {
                match value {
                    #(
                        #counter => #repeat_name_a::#variant_a,
                    )*
                    _ => unreachable!()
                }
            }
            fn to_usize(self) -> usize {
                #to_usize
            }
            fn from_function<F: FnMut(Self) -> V>(mut _f: F) -> Self::Array {
                [#(
                    _f(#repeat_name_b::#variant_b),
                )*]
            }
        }
    }
}

#[proc_macro_derive(EnumMap)]
pub fn derive_enum_map(input: TokenStream) -> TokenStream {
    let input = syn::parse_macro_input(&input.to_string()).unwrap();
    match input.body {
        Body::Enum(ref variants) => generate_enum_code(&input.ident, variants),
        _ => panic!("#[derive(EnumMap)] is only defined for enums"),
    }.parse()
        .unwrap()
}
<|endoftext|>