<commit_before>extern crate phf_codegen;

use std::env;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;

fn main() {
    generate_phf_of_svg_attributes();
}

/// Creates a perfect hash function (PHF) to map SVG attribute names to enum values.
fn generate_phf_of_svg_attributes() {
    // (attribute name, Rust enum value)
    //
    #[cfg_attr(rustfmt, rustfmt_skip)]
    let attribute_defs = [
        ( "alternate",          "Alternate" ),
        ( "amplitude",          "Amplitude" ),
        ( "azimuth",            "Azimuth" ),
        ( "baseFrequency",      "BaseFrequency" ),
        ( "baseline-shift",     "BaselineShift" ),
        ( "bias",               "Bias" ),
        ( "class",              "Class" ),
        ( "clip-path",          "ClipPath" ),
        ( "clip-rule",          "ClipRule" ),
        ( "clipPathUnits",      "ClipPathUnits" ),
        ( "color",              "Color" ),
        ( "color-interpolation-filters", "ColorInterpolationFilters" ),
        ( "cx",                 "Cx" ),
        ( "cy",                 "Cy" ),
        ( "d",                  "D" ),
        ( "diffuseConstant",    "DiffuseConstant" ),
        ( "direction",          "Direction" ),
        ( "display",            "Display" ),
        ( "divisor",            "Divisor" ),
        ( "dx",                 "Dx" ),
        ( "dy",                 "Dy" ),
        ( "edgeMode",           "EdgeMode" ),
        ( "elevation",          "Elevation" ),
        ( "enable-background",  "EnableBackground" ),
        ( "encoding",           "Encoding" ),
        ( "exponent",           "Exponent" ),
        ( "fill",               "Fill" ),
        ( "fill-opacity",       "FillOpacity" ),
        ( "fill-rule",          "FillRule" ),
        ( "filter",             "Filter" ),
        ( "filterUnits",        "FilterUnits" ),
        ( "flood-color",        "FloodColor" ),
        ( "flood-opacity",      "FloodOpacity" ),
        ( "font-family",        "FontFamily" ),
        ( "font-size",          "FontSize" ),
        ( "font-stretch",       "FontStretch" ),
        ( "font-style",         "FontStyle" ),
        ( "font-variant",       "FontVariant" ),
        ( "font-weight",        "FontWeight" ),
        ( "fx",                 "Fx" ),
        ( "fy",                 "Fy" ),
        ( "gradientTransform",  "GradientTransform" ),
        ( "gradientUnits",      "GradientUnits" ),
        ( "height",             "Height" ),
        ( "href",               "Href" ),
        ( "id",                 "Id" ),
        ( "in",                 "In" ),
        ( "in2",                "In2" ),
        ( "intercept",          "Intercept" ),
        ( "k1",                 "K1" ),
        ( "k2",                 "K2" ),
        ( "k3",                 "K3" ),
        ( "k4",                 "K4" ),
        ( "kernelMatrix",       "KernelMatrix" ),
        ( "kernelUnitLength",   "KernelUnitLength" ),
        ( "letter-spacing",     "LetterSpacing" ),
        ( "lighting-color",     "LightingColor" ),
        ( "limitingConeAngle",  "LimitingConeAngle" ),
        ( "marker",             "Marker" ),
        ( "marker-end",         "MarkerEnd" ),
        ( "marker-mid",         "MarkerMid" ),
        ( "marker-start",       "MarkerStart" ),
        ( "markerHeight",       "MarkerHeight" ),
        ( "markerUnits",        "MarkerUnits" ),
        ( "markerWidth",        "MarkerWidth" ),
        ( "mask",               "Mask" ),
        ( "maskContentUnits",   "MaskContentUnits" ),
        ( "maskUnits",          "MaskUnits" ),
        ( "mode",               "Mode" ),
        ( "numOctaves",         "NumOctaves" ),
        ( "offset",             "Offset" ),
        ( "opacity",            "Opacity" ),
        ( "operator",           "Operator" ),
        ( "order",              "Order" ),
        ( "orient",             "Orient" ),
        ( "overflow",           "Overflow" ),
        ( "parse",              "Parse" ),
        ( "path",               "Path" ),
        ( "patternContentUnits", "PatternContentUnits" ),
        ( "patternTransform",   "PatternTransform" ),
        ( "patternUnits",       "PatternUnits" ),
        ( "points",             "Points" ),
        ( "pointsAtX",          "PointsAtX" ),
        ( "pointsAtY",          "PointsAtY" ),
        ( "pointsAtZ",          "PointsAtZ" ),
        ( "preserveAlpha",      "PreserveAlpha" ),
        ( "preserveAspectRatio", "PreserveAspectRatio" ),
        ( "primitiveUnits",     "PrimitiveUnits" ),
        ( "r",                  "R" ),
        ( "radius",             "Radius" ),
        ( "refX",               "RefX" ),
        ( "refY",               "RefY" ),
        ( "requiredExtensions", "RequiredExtensions" ),
        ( "requiredFeatures",   "RequiredFeatures" ),
        ( "result",             "Result" ),
        ( "rx",                 "Rx" ),
        ( "ry",                 "Ry" ),
        ( "scale",              "Scale" ),
        ( "seed",               "Seed" ),
        ( "shape-rendering",    "ShapeRendering" ),
        ( "slope",              "Slope" ),
        ( "specularConstant",   "SpecularConstant" ),
        ( "specularExponent",   "SpecularExponent" ),
        ( "spreadMethod",       "SpreadMethod" ),
        ( "stdDeviation",       "StdDeviation" ),
        ( "stitchTiles",        "StitchTiles" ),
        ( "stop-color",         "StopColor" ),
        ( "stop-opacity",       "StopOpacity" ),
        ( "stroke",             "Stroke" ),
        ( "stroke-dasharray",   "StrokeDasharray" ),
        ( "stroke-dashoffset",  "StrokeDashoffset" ),
        ( "stroke-linecap",     "StrokeLinecap" ),
        ( "stroke-linejoin",    "StrokeLinejoin" ),
        ( "stroke-miterlimit",  "StrokeMiterlimit" ),
        ( "stroke-opacity",     "StrokeOpacity" ),
        ( "stroke-width",       "StrokeWidth" ),
        ( "style",              "Style" ),
        ( "surfaceScale",       "SurfaceScale" ),
        ( "systemLanguage",     "SystemLanguage" ),
        ( "tableValues",        "TableValues" ),
        ( "targetX",            "TargetX" ),
        ( "targetY",            "TargetY" ),
        ( "text-anchor",        "TextAnchor" ),
        ( "text-decoration",    "TextDecoration" ),
        ( "text-rendering",     "TextRendering" ),
        ( "transform",          "Transform" ),
        ( "type",               "Type" ),
        ( "unicode-bidi",       "UnicodeBidi" ),
        ( "values",             "Values" ),
        ( "verts",              "Verts" ),
        ( "viewBox",            "ViewBox" ),
        ( "visibility",         "Visibility" ),
        ( "width",              "Width" ),
        ( "writing-mode",       "WritingMode" ),
        ( "x",                  "X" ),
        ( "x1",                 "X1" ),
        ( "y1",                 "Y1" ),
        ( "x2",                 "X2" ),
        ( "y2",                 "Y2" ),
        ( "xChannelSelector",   "XChannelSelector" ),
        ( "xlink:href",         "XlinkHref" ),
        ( "xml:lang",           "XmlLang" ),
        ( "xml:space",          "XmlSpace" ),
        ( "y",                  "Y" ),
        ( "yChannelSelector",   "YChannelSelector" ),
        ( "z",                  "Z" ),
    ];

    let path = Path::new(&env::var("OUT_DIR").unwrap()).join("attributes-codegen.rs");
    let mut file = BufWriter::new(File::create(&path).unwrap());

    writeln!(&mut file, "#[repr(C)]").unwrap();
    writeln!(
        &mut file,
        "#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]"
    )
    .unwrap();
    writeln!(&mut file, "pub enum Attribute {{").unwrap();

    for &(_, valname) in attribute_defs.iter() {
        writeln!(&mut file, "    {},", valname).unwrap();
    }

    writeln!(&mut file, "}}").unwrap();

    writeln!(
        &mut file,
        "static ATTRIBUTES: phf::Map<&'static str, Attribute> = "
    )
    .unwrap();

    let mut map = phf_codegen::Map::new();
    map.phf_path("phf");
    for &(name, valname) in attribute_defs.iter() {
        let valname = ["Attribute::", valname].concat();
        map.entry(name, &valname);
    }

    map.build(&mut file).unwrap();
    writeln!(&mut file, ";").unwrap();

    output_srgb_tables();
}

/// Converts an sRGB color value to a linear sRGB color value (undoes the gamma correction).
///
/// The input and the output are supposed to be in the [0, 1] range.
#[inline]
fn linearize(c: f64) -> f64 {
    if c <= (12.92 * 0.0031308) {
        c / 12.92
    } else {
        ((c + 0.055) / 1.055).powf(2.4)
    }
}

/// Converts a linear sRGB color value to a normal sRGB color value (applies the gamma correction).
///
/// The input and the output are supposed to be in the [0, 1] range.
#[inline]
fn unlinearize(c: f64) -> f64 {
    if c <= 0.0031308 {
        12.92 * c
    } else {
        1.055 * c.powf(1f64 / 2.4) - 0.055
    }
}

fn compute_table<F: Fn(f64) -> f64>(f: F) -> [u8; 256] {
    let mut table = [0; 256];

    for i in 0..=255 {
        let c = i as f64 / 255.0;
        let x = f(c);
        table[i] = (x * 255.0).round() as u8;
    }

    table
}

fn print_table<W: Write>(w: &mut W, name: &str, table: &[u8]) {
    writeln!(w, "const {}: [u8; {}] = [", name, table.len()).unwrap();

    for x in table {
        writeln!(w, "    {},", x).unwrap();
    }

    writeln!(w, "];").unwrap();
}

fn output_srgb_tables() {
    let linearize_table = compute_table(linearize);
    let unlinearize_table = compute_table(unlinearize);

    let path = Path::new(&env::var("OUT_DIR").unwrap()).join("srgb-codegen.rs");
    let mut file = BufWriter::new(File::create(&path).unwrap());

    print_table(&mut file, "LINEARIZE", &linearize_table);
    print_table(&mut file, "UNLINEARIZE", &unlinearize_table);
}
<commit_msg>build.rs: Move the call to generate_srgb_tables() to the correct place<commit_after>extern crate phf_codegen;

use std::env;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::Path;

fn main() {
    generate_phf_of_svg_attributes();
    generate_srgb_tables();
}

/// Creates a perfect hash function (PHF) to map SVG attribute names to enum values.
fn generate_phf_of_svg_attributes() {
    // (attribute name, Rust enum value)
    //
    #[cfg_attr(rustfmt, rustfmt_skip)]
    let attribute_defs = [
        ( "alternate",          "Alternate" ),
        ( "amplitude",          "Amplitude" ),
        ( "azimuth",            "Azimuth" ),
        ( "baseFrequency",      "BaseFrequency" ),
        ( "baseline-shift",     "BaselineShift" ),
        ( "bias",               "Bias" ),
        ( "class",              "Class" ),
        ( "clip-path",          "ClipPath" ),
        ( "clip-rule",          "ClipRule" ),
        ( "clipPathUnits",      "ClipPathUnits" ),
        ( "color",              "Color" ),
        ( "color-interpolation-filters", "ColorInterpolationFilters" ),
        ( "cx",                 "Cx" ),
        ( "cy",                 "Cy" ),
        ( "d",                  "D" ),
        ( "diffuseConstant",    "DiffuseConstant" ),
        ( "direction",          "Direction" ),
        ( "display",            "Display" ),
        ( "divisor",            "Divisor" ),
        ( "dx",                 "Dx" ),
        ( "dy",                 "Dy" ),
        ( "edgeMode",           "EdgeMode" ),
        ( "elevation",          "Elevation" ),
        ( "enable-background",  "EnableBackground" ),
        ( "encoding",           "Encoding" ),
        ( "exponent",           "Exponent" ),
        ( "fill",               "Fill" ),
        ( "fill-opacity",       "FillOpacity" ),
        ( "fill-rule",          "FillRule" ),
        ( "filter",             "Filter" ),
        ( "filterUnits",        "FilterUnits" ),
        ( "flood-color",        "FloodColor" ),
        ( "flood-opacity",      "FloodOpacity" ),
        ( "font-family",        "FontFamily" ),
        ( "font-size",          "FontSize" ),
        ( "font-stretch",       "FontStretch" ),
        ( "font-style",         "FontStyle" ),
        ( "font-variant",       "FontVariant" ),
        ( "font-weight",        "FontWeight" ),
        ( "fx",                 "Fx" ),
        ( "fy",                 "Fy" ),
        ( "gradientTransform",  "GradientTransform" ),
        ( "gradientUnits",      "GradientUnits" ),
        ( "height",             "Height" ),
        ( "href",               "Href" ),
        ( "id",                 "Id" ),
        ( "in",                 "In" ),
        ( "in2",                "In2" ),
        ( "intercept",          "Intercept" ),
        ( "k1",                 "K1" ),
        ( "k2",                 "K2" ),
        ( "k3",                 "K3" ),
        ( "k4",                 "K4" ),
        ( "kernelMatrix",       "KernelMatrix" ),
        ( "kernelUnitLength",   "KernelUnitLength" ),
        ( "letter-spacing",     "LetterSpacing" ),
        ( "lighting-color",     "LightingColor" ),
        ( "limitingConeAngle",  "LimitingConeAngle" ),
        ( "marker",             "Marker" ),
        ( "marker-end",         "MarkerEnd" ),
        ( "marker-mid",         "MarkerMid" ),
        ( "marker-start",       "MarkerStart" ),
        ( "markerHeight",       "MarkerHeight" ),
        ( "markerUnits",        "MarkerUnits" ),
        ( "markerWidth",        "MarkerWidth" ),
        ( "mask",               "Mask" ),
        ( "maskContentUnits",   "MaskContentUnits" ),
        ( "maskUnits",          "MaskUnits" ),
        ( "mode",               "Mode" ),
        ( "numOctaves",         "NumOctaves" ),
        ( "offset",             "Offset" ),
        ( "opacity",            "Opacity" ),
        ( "operator",           "Operator" ),
        ( "order",              "Order" ),
        ( "orient",             "Orient" ),
        ( "overflow",           "Overflow" ),
        ( "parse",              "Parse" ),
        ( "path",               "Path" ),
        ( "patternContentUnits", "PatternContentUnits" ),
        ( "patternTransform",   "PatternTransform" ),
        ( "patternUnits",       "PatternUnits" ),
        ( "points",             "Points" ),
        ( "pointsAtX",          "PointsAtX" ),
        ( "pointsAtY",          "PointsAtY" ),
        ( "pointsAtZ",          "PointsAtZ" ),
        ( "preserveAlpha",      "PreserveAlpha" ),
        ( "preserveAspectRatio", "PreserveAspectRatio" ),
        ( "primitiveUnits",     "PrimitiveUnits" ),
        ( "r",                  "R" ),
        ( "radius",             "Radius" ),
        ( "refX",               "RefX" ),
        ( "refY",               "RefY" ),
        ( "requiredExtensions", "RequiredExtensions" ),
        ( "requiredFeatures",   "RequiredFeatures" ),
        ( "result",             "Result" ),
        ( "rx",                 "Rx" ),
        ( "ry",                 "Ry" ),
        ( "scale",              "Scale" ),
        ( "seed",               "Seed" ),
        ( "shape-rendering",    "ShapeRendering" ),
        ( "slope",              "Slope" ),
        ( "specularConstant",   "SpecularConstant" ),
        ( "specularExponent",   "SpecularExponent" ),
        ( "spreadMethod",       "SpreadMethod" ),
        ( "stdDeviation",       "StdDeviation" ),
        ( "stitchTiles",        "StitchTiles" ),
        ( "stop-color",         "StopColor" ),
        ( "stop-opacity",       "StopOpacity" ),
        ( "stroke",             "Stroke" ),
        ( "stroke-dasharray",   "StrokeDasharray" ),
        ( "stroke-dashoffset",  "StrokeDashoffset" ),
        ( "stroke-linecap",     "StrokeLinecap" ),
        ( "stroke-linejoin",    "StrokeLinejoin" ),
        ( "stroke-miterlimit",  "StrokeMiterlimit" ),
        ( "stroke-opacity",     "StrokeOpacity" ),
        ( "stroke-width",       "StrokeWidth" ),
        ( "style",              "Style" ),
        ( "surfaceScale",       "SurfaceScale" ),
        ( "systemLanguage",     "SystemLanguage" ),
        ( "tableValues",        "TableValues" ),
        ( "targetX",            "TargetX" ),
        ( "targetY",            "TargetY" ),
        ( "text-anchor",        "TextAnchor" ),
        ( "text-decoration",    "TextDecoration" ),
        ( "text-rendering",     "TextRendering" ),
        ( "transform",          "Transform" ),
        ( "type",               "Type" ),
        ( "unicode-bidi",       "UnicodeBidi" ),
        ( "values",             "Values" ),
        ( "verts",              "Verts" ),
        ( "viewBox",            "ViewBox" ),
        ( "visibility",         "Visibility" ),
        ( "width",              "Width" ),
        ( "writing-mode",       "WritingMode" ),
        ( "x",                  "X" ),
        ( "x1",                 "X1" ),
        ( "y1",                 "Y1" ),
        ( "x2",                 "X2" ),
        ( "y2",                 "Y2" ),
        ( "xChannelSelector",   "XChannelSelector" ),
        ( "xlink:href",         "XlinkHref" ),
        ( "xml:lang",           "XmlLang" ),
        ( "xml:space",          "XmlSpace" ),
        ( "y",                  "Y" ),
        ( "yChannelSelector",   "YChannelSelector" ),
        ( "z",                  "Z" ),
    ];

    let path = Path::new(&env::var("OUT_DIR").unwrap()).join("attributes-codegen.rs");
    let mut file = BufWriter::new(File::create(&path).unwrap());

    writeln!(&mut file, "#[repr(C)]").unwrap();
    writeln!(
        &mut file,
        "#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]"
    )
    .unwrap();
    writeln!(&mut file, "pub enum Attribute {{").unwrap();

    for &(_, valname) in attribute_defs.iter() {
        writeln!(&mut file, "    {},", valname).unwrap();
    }

    writeln!(&mut file, "}}").unwrap();

    writeln!(
        &mut file,
        "static ATTRIBUTES: phf::Map<&'static str, Attribute> = "
    )
    .unwrap();

    let mut map = phf_codegen::Map::new();
    map.phf_path("phf");
    for &(name, valname) in attribute_defs.iter() {
        let valname = ["Attribute::", valname].concat();
        map.entry(name, &valname);
    }

    map.build(&mut file).unwrap();
    writeln!(&mut file, ";").unwrap();
}

/// Converts an sRGB color value to a linear sRGB color value (undoes the gamma correction).
///
/// The input and the output are supposed to be in the [0, 1] range.
#[inline]
fn linearize(c: f64) -> f64 {
    if c <= (12.92 * 0.0031308) {
        c / 12.92
    } else {
        ((c + 0.055) / 1.055).powf(2.4)
    }
}

/// Converts a linear sRGB color value to a normal sRGB color value (applies the gamma correction).
///
/// The input and the output are supposed to be in the [0, 1] range.
#[inline]
fn unlinearize(c: f64) -> f64 {
    if c <= 0.0031308 {
        12.92 * c
    } else {
        1.055 * c.powf(1f64 / 2.4) - 0.055
    }
}

fn compute_table<F: Fn(f64) -> f64>(f: F) -> [u8; 256] {
    let mut table = [0; 256];

    for i in 0..=255 {
        let c = i as f64 / 255.0;
        let x = f(c);
        table[i] = (x * 255.0).round() as u8;
    }

    table
}

fn print_table<W: Write>(w: &mut W, name: &str, table: &[u8]) {
    writeln!(w, "const {}: [u8; {}] = [", name, table.len()).unwrap();

    for x in table {
        writeln!(w, "    {},", x).unwrap();
    }

    writeln!(w, "];").unwrap();
}

fn generate_srgb_tables() {
    let linearize_table = compute_table(linearize);
    let unlinearize_table = compute_table(unlinearize);

    let path = Path::new(&env::var("OUT_DIR").unwrap()).join("srgb-codegen.rs");
    let mut file = BufWriter::new(File::create(&path).unwrap());

    print_table(&mut file, "LINEARIZE", &linearize_table);
    print_table(&mut file, "UNLINEARIZE", &unlinearize_table);
}
<|endoftext|>


<commit_before>extern crate tokio_core;
extern crate futures;
extern crate tk_http;
extern crate tk_listen;
extern crate wonky_tanks_server;

use std::cell::RefCell;
use std::collections::hash_map::HashMap;
use std::fmt::Write;
use std::rc::Rc;

use tokio_core::reactor::Core;
use tokio_core::net::TcpListener;

use futures::{Stream, Future};
use futures::future::{FutureResult, ok};
use futures::sync::mpsc::{unbounded, UnboundedSender};

use tk_http::{Status};
use tk_http::server::buffered::{Request, BufferedDispatcher};
use tk_http::server::{Encoder, EncoderDone, Config as ServerConfig, Proto, Error as ServerError};
use tk_http::websocket::{Packet, Config as WebsocketConfig, Dispatcher, Frame, Loop, Error as WebsocketError};

use tk_listen::ListenExt;

use wonky_tanks_server::*;

const LOBBY: &'static str = "/";

type SharedData = Rc<RefCell<HashMap<CountedString, ChannelState>>>;

#[derive(Clone, Default)]
pub struct ChannelState {
    capacity: Option<u8>,
    identities: usize,
    participants: Vec<Rc<UnboundedSender<Packet>>>,
}

impl std::fmt::Debug for ChannelState {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        f
            .debug_struct("ChannelState")
            .field("capacity", &self.capacity)
            .field("identities", &self.identities)
            .field(
                "participants",
                &format_args!(
                    "{:x}",
                    IterContentsDebug(||
                        self
                            .participants
                            .iter()
                            .map(|i| (&i) as &UnboundedSender<_> as *const _ as usize)
                    )
                )
            )
            .finish()
    }
}

fn service<S>(directory: &Directory, token: &mut Option<String>, req: Request, mut e: Encoder<S>) -> FutureResult<EncoderDone<S>, ServerError> {
    if let Some(ws) = req.websocket_handshake() {
        *token = Some(req.path().to_string());
        e.status(Status::SwitchingProtocol);
        e.add_header("Connection", "upgrade").unwrap();
        e.add_header("Upgrade", "websocket").unwrap();
        e.format_header("Sec-Websocket-Accept", &ws.accept).unwrap();
        e.done_headers().unwrap();
        return ok(e.done())
    }

    if let Some((ctype, data)) = directory.get(req.path()) {
        e.status(Status::Ok);
        e.add_length(data.len() as u64).unwrap();
        e.add_header("Content-Type", ctype).unwrap();
        if e.done_headers().unwrap() {
            e.write_body(&data);
        }
        return ok(e.done())
    }

    let not_found = "No data found".as_bytes();
    e.status(Status::NotFound);
    e.add_length(not_found.len() as u64).unwrap();
    e.add_header("Content-Type", "text/plain").unwrap();
    if e.done_headers().unwrap() {
        e.write_body(not_found);
    }
    ok(e.done())
}

struct WebsocketConnection {
    shared_data: SharedData,
    sender: Rc<UnboundedSender<Packet>>,
    token: CountedString,
}

impl WebsocketConnection {
    pub fn new(shared_data: SharedData, token: String, sender: Rc<UnboundedSender<Packet>>) -> Self {
        let token: CountedString  = token.into();
        let mut data: String = Default::default();
        match match {
            let shared_data = &mut *shared_data.borrow_mut();
            let shared_data = shared_data.entry(token.clone());
            #[cfg(debug_assertions)]
            println!("Previous data: {:?}", &shared_data);
            let &mut ChannelState { ref mut participants, capacity, ref mut identities } = shared_data.or_insert_with(Default::default);
            write!(data, "identity: {}", identities).unwrap();
            drop(sender.unbounded_send(Packet::Text(data.clone())));
            *identities += 1;
            participants.push(sender.clone());
            data.clear();
            write!(data, "count: {}", participants.len()).unwrap();
            let (last, participants) = participants.split_last().unwrap();
            for participant in participants {
                drop(participant.unbounded_send(Packet::Text(data.clone())));
            }
            drop(last.unbounded_send(Packet::Text(data.clone())));
            (capacity, participants.len())
        } {
            (Some(capacity), 1) if 1 < capacity => Some(true),
            (None, 1) => Some(true),
            (Some(capacity), count) if count == capacity as usize && count != 1 => Some(false),
            _ => None,
        } {
            Some(toggle) => {
                if let Some(lobby) = shared_data.try_borrow().unwrap().get(LOBBY) {
                    data.clear();
                    write!(data, "{}: {}", if toggle { "+" } else { "-" }, token).unwrap();
                    let (last, participants) = lobby.participants.split_last().unwrap();
                    for participant in participants {
                        drop(participant.unbounded_send(Packet::Text(data.clone())));
                    }
                    drop(last.unbounded_send(Packet::Text(data)));
                }
            },
            None => {},
        }
        #[cfg(debug_assertions)]
        println!("Created connection: {:x}/{}", (&sender) as &UnboundedSender<_> as *const _ as usize, token);
        WebsocketConnection { shared_data, sender, token }
    }

    fn is_lobby(&self) -> bool {
        LOBBY == &*self.token.0
    }
}

impl Drop for WebsocketConnection {
    fn drop(&mut self) {
        let &mut WebsocketConnection {
            ref mut shared_data,
            ref sender,
            ref token
        } = self;
        use std::collections::hash_map::Entry::Occupied;
        if Some(1) != {
            let mut map = shared_data.borrow_mut();
            let mut entry =
                if let Occupied(entry) = map.entry(token.clone()) {
                    entry
                } else {
                    panic!("Bad state; expected {} to be occupied: {:?}", token, shared_data)
                };
            if let (capacity, 0) = {
                let &mut ChannelState { ref mut participants, capacity, identities: _ } = entry.get_mut();
                participants.retain(|target| !Rc::ptr_eq(target, sender));
                if !participants.is_empty() {
                    let data = format!("count: {}", participants.len());
                    let (last, participants) = participants.split_last().unwrap();
                    for participant in participants {
                        drop(participant.unbounded_send(Packet::Text(data.clone())));
                    }
                    drop(last.unbounded_send(Packet::Text(data)));
                }
                (capacity, participants.len())
            } {
                entry.remove_entry();
                capacity
            } else {
                Some(1)
            }
        } {
            if let Some(lobby) = shared_data.try_borrow().unwrap().get(LOBBY) {
                let data = format!("-: {}", token);
                let (last, participants) = lobby.participants.split_last().unwrap();
                for participant in participants {
                    drop(participant.unbounded_send(Packet::Text(data.clone())));
                }
                drop(last.unbounded_send(Packet::Text(data)));
            }
        }
        #[cfg(debug_assertions)]
        println!("Dropped connection: {:x}/{}", (&self.sender) as &UnboundedSender<_> as *const _ as usize, token);
    }
}

impl Dispatcher for WebsocketConnection {
    type Future = FutureResult<(), WebsocketError>;

    fn frame(&mut self, frame: &Frame) -> FutureResult<(), WebsocketError> {
        if self.is_lobby() { return ok(()) }
        let &mut WebsocketConnection { ref shared_data, ref sender, ref token } = self;

        match frame {
            &Frame::Text(t) => {
                // FIXME - parse state changes
            },
            &Frame::Binary(b) => {
                let &ChannelState { ref participants, .. } = &(&*shared_data.borrow())[token];
                if participants.len() == 1 { return ok(()) }
                let mut participants = participants.iter().filter(|&participant| !Rc::ptr_eq(participant, sender));
                let last = participants.next_back().unwrap();
                let to_send = Packet::Binary(Vec::from(b));
                for participant in participants {
                    drop(participant.unbounded_send(to_send.clone()));
                }
                drop(last.unbounded_send(to_send));
            },
            _ => {},
        }

        ok(())
    }
}

fn main() {
    let directory = &Directory::new(std::env::current_dir().unwrap().as_path());
    let shared_data: SharedData = Default::default();
    let mut lp = Core::new().unwrap();
    let h1 = lp.handle();

    let addr = "0.0.0.0:5000".parse().unwrap();
    let listener = TcpListener::bind(&addr, &lp.handle()).unwrap();
    let cfg = ServerConfig::new().done();
    let wcfg = WebsocketConfig::new().done();

    let pipeline = listener.incoming();
    let pipeline =
        pipeline.map(
        move |(socket, addr)| {
            let wcfg = wcfg.clone();
            let h2 = h1.clone();

            let shared_data = shared_data.clone();
            let data = Rc::new(RefCell::new(None));
            let http_data = data.clone();
            let websocket_data = data;

            Proto::new(
                socket,
                &cfg,
                BufferedDispatcher::new_with_websockets(
                    addr,
                    &h1,
                    move |p1, p2| {
                        let mut http_data = http_data.borrow_mut();
                        let http_data: &mut Option<String> = &mut http_data;
                        if let Some(token) = http_data.take() {
                            panic!("http initializer called twice: {}", token);
                        }
                        let ret = service(
                            directory,
                            http_data,
                            p1,
                            p2
                        );
                        ret
                    },
                    move |outp, inp| {
                        let (sender, receiver) = unbounded();
                        let sender = Rc::new(sender);
                        Loop::server(
                            outp,
                            inp,
                            receiver.map_err(move |e| format!("{:?} : stream closed", e)),
                            WebsocketConnection::new(
                                shared_data.clone(),
                                websocket_data
                                    .borrow_mut()
                                    .take()
                                    .expect("http initializer not called yet"),
                                sender
                            ),
                            &wcfg,
                            &h2
                        )
                            .map_err(|e| println!("Websocket err: {:?}", e))
                    }
                ),
                &h1
            )
                .map_err(|e| {println!("Connection error: {:?}", e);})
                .then(|_| Ok(())) // don't fail, please
        }
    );
    let pipeline = pipeline.listen(1000);

    lp.run(pipeline).unwrap();
}<commit_msg>Properly broadcast channel add messages<commit_after>extern crate tokio_core;
extern crate futures;
extern crate tk_http;
extern crate tk_listen;
extern crate wonky_tanks_server;

use std::cell::RefCell;
use std::collections::hash_map::HashMap;
use std::fmt::Write;
use std::rc::Rc;

use tokio_core::reactor::Core;
use tokio_core::net::TcpListener;

use futures::{Stream, Future};
use futures::future::{FutureResult, ok};
use futures::sync::mpsc::{unbounded, UnboundedSender};

use tk_http::{Status};
use tk_http::server::buffered::{Request, BufferedDispatcher};
use tk_http::server::{Encoder, EncoderDone, Config as ServerConfig, Proto, Error as ServerError};
use tk_http::websocket::{Packet, Config as WebsocketConfig, Dispatcher, Frame, Loop, Error as WebsocketError};

use tk_listen::ListenExt;

use wonky_tanks_server::*;

const LOBBY: &'static str = "/";

type SharedData = Rc<RefCell<HashMap<CountedString, ChannelState>>>;

#[derive(Clone, Default)]
pub struct ChannelState {
    capacity: Option<u8>,
    identities: usize,
    participants: Vec<Rc<UnboundedSender<Packet>>>,
}

impl std::fmt::Debug for ChannelState {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        f
            .debug_struct("ChannelState")
            .field("capacity", &self.capacity)
            .field("identities", &self.identities)
            .field(
                "participants",
                &format_args!(
                    "{:x}",
                    IterContentsDebug(||
                        self
                            .participants
                            .iter()
                            .map(|i| (&i) as &UnboundedSender<_> as *const _ as usize)
                    )
                )
            )
            .finish()
    }
}

fn service<S>(directory: &Directory, token: &mut Option<String>, req: Request, mut e: Encoder<S>) -> FutureResult<EncoderDone<S>, ServerError> {
    if let Some(ws) = req.websocket_handshake() {
        *token = Some(req.path().to_string());
        e.status(Status::SwitchingProtocol);
        e.add_header("Connection", "upgrade").unwrap();
        e.add_header("Upgrade", "websocket").unwrap();
        e.format_header("Sec-Websocket-Accept", &ws.accept).unwrap();
        e.done_headers().unwrap();
        return ok(e.done())
    }

    if let Some((ctype, data)) = directory.get(req.path()) {
        e.status(Status::Ok);
        e.add_length(data.len() as u64).unwrap();
        e.add_header("Content-Type", ctype).unwrap();
        if e.done_headers().unwrap() {
            e.write_body(&data);
        }
        return ok(e.done())
    }

    let not_found = "No data found".as_bytes();
    e.status(Status::NotFound);
    e.add_length(not_found.len() as u64).unwrap();
    e.add_header("Content-Type", "text/plain").unwrap();
    if e.done_headers().unwrap() {
        e.write_body(not_found);
    }
    ok(e.done())
}

struct WebsocketConnection {
    shared_data: SharedData,
    sender: Rc<UnboundedSender<Packet>>,
    token: CountedString,
}

impl WebsocketConnection {
    pub fn new(shared_data: SharedData, token: String, sender: Rc<UnboundedSender<Packet>>) -> Self {
        let token: CountedString  = token.into();
        let mut data: String = Default::default();
        match match {
            let shared_data = &mut *shared_data.borrow_mut();
            let shared_data = shared_data.entry(token.clone());
            #[cfg(debug_assertions)]
            println!("Previous data: {:?}", &shared_data);
            let &mut ChannelState { ref mut participants, capacity, ref mut identities } = shared_data.or_insert_with(Default::default);
            write!(data, "identity: {}", identities).unwrap();
            drop(sender.unbounded_send(Packet::Text(data.clone())));
            *identities += 1;
            participants.push(sender.clone());
            data.clear();
            write!(data, "count: {}", participants.len()).unwrap();
            let (last, participants_sans_last) = participants.split_last().unwrap();
            for participant in participants_sans_last {
                drop(participant.unbounded_send(Packet::Text(data.clone())));
            }
            drop(last.unbounded_send(Packet::Text(data.clone())));
            (capacity, participants.len())
        } {
            (Some(capacity), 1) if 1 < capacity => Some(true),
            (None, 1) => Some(true),
            (Some(capacity), count) if count == capacity as usize && count != 1 => Some(false),
            _ => None,
        } {
            _ if &*token.0 == LOBBY => {
                for (channel, state)in shared_data.try_borrow().unwrap().iter() {
                    if &*channel.0 == LOBBY { continue }
                    if let Some(capacity) = state.capacity {
                        if state.participants.len() >= capacity as usize {
                            continue
                        }
                    }
                    data.clear();
                    write!(data, "{}: {}", "+", channel).unwrap();

                }
            },
            Some(toggle) => {
                if let Some(lobby) = shared_data.try_borrow().unwrap().get(LOBBY) {
                    data.clear();
                    write!(data, "{}: {}", if toggle { "+" } else { "-" }, token).unwrap();
                    let (last, participants) = lobby.participants.split_last().unwrap();
                    for participant in participants {
                        drop(participant.unbounded_send(Packet::Text(data.clone())));
                    }
                    drop(last.unbounded_send(Packet::Text(data)));
                }
            },
            _ => {},
        }
        #[cfg(debug_assertions)]
        println!("Created connection: {:x}/{}", (&sender) as &UnboundedSender<_> as *const _ as usize, token);
        WebsocketConnection { shared_data, sender, token }
    }

    fn is_lobby(&self) -> bool {
        LOBBY == &*self.token.0
    }
}

impl Drop for WebsocketConnection {
    fn drop(&mut self) {
        let &mut WebsocketConnection {
            ref mut shared_data,
            ref sender,
            ref token
        } = self;
        use std::collections::hash_map::Entry::Occupied;
        if Some(1) != {
            let mut map = shared_data.borrow_mut();
            let mut entry =
                if let Occupied(entry) = map.entry(token.clone()) {
                    entry
                } else {
                    panic!("Bad state; expected {} to be occupied: {:?}", token, shared_data)
                };
            if let (capacity, 0) = {
                let &mut ChannelState { ref mut participants, capacity, identities: _ } = entry.get_mut();
                participants.retain(|target| !Rc::ptr_eq(target, sender));
                if !participants.is_empty() {
                    let data = format!("count: {}", participants.len());
                    let (last, participants) = participants.split_last().unwrap();
                    for participant in participants {
                        drop(participant.unbounded_send(Packet::Text(data.clone())));
                    }
                    drop(last.unbounded_send(Packet::Text(data)));
                }
                (capacity, participants.len())
            } {
                entry.remove_entry();
                capacity
            } else {
                Some(1)
            }
        } {
            if let Some(lobby) = shared_data.try_borrow().unwrap().get(LOBBY) {
                let data = format!("-: {}", token);
                let (last, participants) = lobby.participants.split_last().unwrap();
                for participant in participants {
                    drop(participant.unbounded_send(Packet::Text(data.clone())));
                }
                drop(last.unbounded_send(Packet::Text(data)));
            }
        }
        #[cfg(debug_assertions)]
        println!("Dropped connection: {:x}/{}", (&self.sender) as &UnboundedSender<_> as *const _ as usize, token);
    }
}

impl Dispatcher for WebsocketConnection {
    type Future = FutureResult<(), WebsocketError>;

    fn frame(&mut self, frame: &Frame) -> FutureResult<(), WebsocketError> {
        if self.is_lobby() { return ok(()) }
        let &mut WebsocketConnection { ref shared_data, ref sender, ref token } = self;

        match frame {
            &Frame::Text(t) => {
                // FIXME - parse state changes
            },
            &Frame::Binary(b) => {
                let &ChannelState { ref participants, .. } = &(&*shared_data.borrow())[token];
                if participants.len() == 1 { return ok(()) }
                let mut participants = participants.iter().filter(|&participant| !Rc::ptr_eq(participant, sender));
                let last = participants.next_back().unwrap();
                let to_send = Packet::Binary(Vec::from(b));
                for participant in participants {
                    drop(participant.unbounded_send(to_send.clone()));
                }
                drop(last.unbounded_send(to_send));
            },
            _ => {},
        }

        ok(())
    }
}

fn main() {
    let directory = &Directory::new(std::env::current_dir().unwrap().as_path());
    let shared_data: SharedData = Default::default();
    let mut lp = Core::new().unwrap();
    let h1 = lp.handle();

    let addr = "0.0.0.0:5000".parse().unwrap();
    let listener = TcpListener::bind(&addr, &lp.handle()).unwrap();
    let cfg = ServerConfig::new().done();
    let wcfg = WebsocketConfig::new().done();

    let pipeline = listener.incoming();
    let pipeline =
        pipeline.map(
        move |(socket, addr)| {
            let wcfg = wcfg.clone();
            let h2 = h1.clone();

            let shared_data = shared_data.clone();
            let data = Rc::new(RefCell::new(None));
            let http_data = data.clone();
            let websocket_data = data;

            Proto::new(
                socket,
                &cfg,
                BufferedDispatcher::new_with_websockets(
                    addr,
                    &h1,
                    move |p1, p2| {
                        let mut http_data = http_data.borrow_mut();
                        let http_data: &mut Option<String> = &mut http_data;
                        if let Some(token) = http_data.take() {
                            panic!("http initializer called twice: {}", token);
                        }
                        let ret = service(
                            directory,
                            http_data,
                            p1,
                            p2
                        );
                        ret
                    },
                    move |outp, inp| {
                        let (sender, receiver) = unbounded();
                        let sender = Rc::new(sender);
                        Loop::server(
                            outp,
                            inp,
                            receiver.map_err(move |e| format!("{:?} : stream closed", e)),
                            WebsocketConnection::new(
                                shared_data.clone(),
                                websocket_data
                                    .borrow_mut()
                                    .take()
                                    .expect("http initializer not called yet"),
                                sender
                            ),
                            &wcfg,
                            &h2
                        )
                            .map_err(|e| println!("Websocket err: {:?}", e))
                    }
                ),
                &h1
            )
                .map_err(|e| {println!("Connection error: {:?}", e);})
                .then(|_| Ok(())) // don't fail, please
        }
    );
    let pipeline = pipeline.listen(1000);

    lp.run(pipeline).unwrap();
}<|endoftext|>


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

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;

use serde::{Deserialize, Serialize};
use tokio::sync::oneshot::{channel, Sender};
use tokio::time::{self, Instant};

use metrics::Metrics;

use crate::extensions::filter_registry::{CreateFilterArgs, DownstreamContext, DownstreamResponse};
use crate::extensions::{Error, Filter, FilterFactory};

mod metrics;

/// RateLimitFilter applies rate limiting to packets flowing through the proxy
///
/// # Configuration
///
/// ```yaml
/// local:
///   port: 7000 # the port to receive traffic to locally
/// filters:
///   - name: quilkin.extensions.filters.local_rate_limit.v1alpha1.LocalRateLimit
///     config:
///       max_packets: 10
///       period: 500ms
/// client:
///   addresses:
///     - 127.0.0.1:7001
/// ```
///  `config.max_packets` is the maximum number of packets allowed
///  to be forwarded by the rate limiter in a given duration.
///  `config.period` (optional) is the duration during which config.max_packets applies.
///  If none is provided, it defaults to 1 second.
///
/// # Metrics
///
/// `quilkin_filter_LocalRateLimit_packets_dropped`: Total number of packets dropped due to rate limiting
///

/// Config represents a RateLimitFilter's configuration.
#[derive(Serialize, Deserialize, Debug)]
struct Config {
    /// max_packets is the maximum number of packets allowed
    /// to be forwarded by the rate limiter in a given duration.
    max_packets: usize,
    /// period is the duration during which max_packets applies.
    /// If none is provided, it defaults to 1 second.
    #[serde(with = "humantime_serde")]
    period: Option<Duration>,
}

/// Creates instances of RateLimitFilter.
#[derive(Default)]
pub struct RateLimitFilterFactory;

/// A filter that implements rate limiting on packets based on
/// the token-bucket algorithm.
/// Packets that violate the rate limit are dropped.
/// It only applies rate limiting on packets that are destined for the
/// proxy's endpoints. All other packets flow through the filter untouched.
struct RateLimitFilter {
    /// available_tokens is how many tokens are left in the bucket any
    /// any given moment.
    available_tokens: Arc<AtomicUsize>,
    /// metrics reporter for this filter.
    metrics: Metrics,
    /// shutdown_tx signals the spawned token refill future to exit.
    shutdown_tx: Option<Sender<()>>,
}

impl FilterFactory for RateLimitFilterFactory {
    fn name(&self) -> String {
        "quilkin.extensions.filters.local_rate_limit.v1alpha1.LocalRateLimit".into()
    }

    fn create_filter(&self, args: CreateFilterArgs) -> Result<Box<dyn Filter>, Error> {
        let config: Config = args.parse_config()?;

        match config.period {
            Some(period) if period.lt(&Duration::from_millis(100)) => Err(Error::FieldInvalid {
                field: "period".into(),
                reason: "value must be at least 100ms".into(),
            }),
            _ => Ok(Box::new(RateLimitFilter::new(
                config,
                Metrics::new(&args.metrics_registry)?,
            ))),
        }
    }
}

impl RateLimitFilter {
    /// new returns a new RateLimitFilter. It spawns a future in the background
    /// that periodically refills the rate limiter's tokens.
    fn new(config: Config, metrics: Metrics) -> Self {
        let (shutdown_tx, mut shutdown_rx) = channel();

        let tokens = Arc::new(AtomicUsize::new(config.max_packets));

        let max_tokens = config.max_packets;
        let period = config.period.unwrap_or_else(|| Duration::from_secs(1));
        let available_tokens = tokens.clone();
        let _ = tokio::spawn(async move {
            let mut interval = time::interval_at(Instant::now() + period, period);
            loop {
                tokio::select! {
                    _ = interval.tick() => {
                        // Refill tokens.
                        let mut refilled = false;
                        while !refilled {
                            let remaining_tokens = available_tokens.load(Ordering::Relaxed);

                            refilled = available_tokens.compare_and_swap(
                                remaining_tokens,
                                max_tokens,
                                Ordering::Relaxed) == remaining_tokens;
                        }
                    },
                    _ = &mut shutdown_rx => {
                        return;
                    }
                }
            }
        });

        RateLimitFilter {
            available_tokens: tokens,
            metrics,
            shutdown_tx: Some(shutdown_tx),
        }
    }

    /// acquire_token is called on behalf of every packet that is eligible
    /// for rate limiting. It returns whether there exists a token in the current
    /// period - determining whether or not the packet should be forwarded or dropped.
    fn acquire_token(&self) -> Option<()> {
        loop {
            let remaining_tokens = self.available_tokens.load(Ordering::Relaxed);

            if remaining_tokens == 0 {
                return None;
            }

            if self.available_tokens.compare_and_swap(
                remaining_tokens,
                remaining_tokens - 1,
                Ordering::Relaxed,
            ) == remaining_tokens
            {
                return Some(());
            }
        }
    }
}

impl Drop for RateLimitFilter {
    fn drop(&mut self) {
        if let Some(shutdown_tx) = self.shutdown_tx.take() {
            shutdown_tx.send(()).ok();
        }
    }
}

impl Filter for RateLimitFilter {
    fn on_downstream_receive(&self, ctx: DownstreamContext) -> Option<DownstreamResponse> {
        self.acquire_token().map(|()| ctx.into()).or_else(|| {
            self.metrics.packets_dropped_total.inc();
            None
        })
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use prometheus::Registry;
    use tokio::time;

    use crate::cluster::Endpoint;
    use crate::config::Endpoints;
    use crate::extensions::filter_registry::DownstreamContext;
    use crate::extensions::filters::local_rate_limit::metrics::Metrics;
    use crate::extensions::filters::local_rate_limit::{Config, RateLimitFilter};
    use crate::extensions::Filter;
    use crate::test_utils::assert_filter_on_upstream_receive_no_change;

    fn rate_limiter(config: Config) -> RateLimitFilter {
        RateLimitFilter::new(config, Metrics::new(&Registry::default()).unwrap())
    }

    #[tokio::test]
    async fn initially_available_tokens() {
        // Test that we always start with the max number of tokens available.
        let r = rate_limiter(Config {
            max_packets: 3,
            period: Some(Duration::from_millis(100)),
        });

        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);
    }

    #[tokio::test]
    async fn token_exhaustion_and_refill() {
        let r = rate_limiter(Config {
            max_packets: 2,
            period: Some(Duration::from_millis(100)),
        });

        // Exhaust tokens
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);

        // Wait for refill
        time::sleep(Duration::from_millis(110)).await;

        // Exhaust tokens again.
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);
    }

    #[tokio::test]
    async fn token_refill_maximum() {
        // Test that we never refill more than the max_tokens specified.

        let r = rate_limiter(Config {
            max_packets: 3,
            period: Some(Duration::from_millis(100)),
        });

        // Use up some of the tokens.
        assert_eq!(r.acquire_token(), Some(()));

        // Wait for refill
        time::sleep(Duration::from_millis(110)).await;

        // Refill should not go over max token limit.
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);
    }

    #[tokio::test]
    async fn filter_with_no_available_tokens() {
        let r = rate_limiter(Config {
            max_packets: 0,
            period: Some(Duration::from_millis(100)),
        });

        // Check that other routes are not affected.
        assert_filter_on_upstream_receive_no_change(&r);

        // Check that we're rate limited.
        assert!(r
            .on_downstream_receive(DownstreamContext::new(
                Endpoints::new(vec![Endpoint::from_address(
                    "127.0.0.1:8080".parse().unwrap(),
                )])
                .unwrap()
                .into(),
                "127.0.0.1:8080".parse().unwrap(),
                vec![9],
            ))
            .is_none(),);
    }

    #[tokio::test]
    async fn filter_with_available_tokens() {
        let r = rate_limiter(Config {
            max_packets: 1,
            period: Some(Duration::from_millis(100)),
        });

        let result = r
            .on_downstream_receive(DownstreamContext::new(
                Endpoints::new(vec![Endpoint::from_address(
                    "127.0.0.1:8080".parse().unwrap(),
                )])
                .unwrap()
                .into(),
                "127.0.0.1:8080".parse().unwrap(),
                vec![9],
            ))
            .unwrap();
        assert_eq!(result.contents, vec![9]);
        // We should be out of tokens now.
        assert_eq!(None, r.acquire_token());

        // Check that other routes are not affected.
        assert_filter_on_upstream_receive_no_change(&r);
    }
}
<commit_msg>Fix rate limiter flaky test (#178)<commit_after>/*
 * Copyright 2020 Google LLC All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;

use serde::{Deserialize, Serialize};
use tokio::sync::oneshot::{channel, Sender};
use tokio::time::{self, Instant};

use metrics::Metrics;

use crate::extensions::filter_registry::{CreateFilterArgs, DownstreamContext, DownstreamResponse};
use crate::extensions::{Error, Filter, FilterFactory};

mod metrics;

/// RateLimitFilter applies rate limiting to packets flowing through the proxy
///
/// # Configuration
///
/// ```yaml
/// local:
///   port: 7000 # the port to receive traffic to locally
/// filters:
///   - name: quilkin.extensions.filters.local_rate_limit.v1alpha1.LocalRateLimit
///     config:
///       max_packets: 10
///       period: 500ms
/// client:
///   addresses:
///     - 127.0.0.1:7001
/// ```
///  `config.max_packets` is the maximum number of packets allowed
///  to be forwarded by the rate limiter in a given duration.
///  `config.period` (optional) is the duration during which config.max_packets applies.
///  If none is provided, it defaults to 1 second.
///
/// # Metrics
///
/// `quilkin_filter_LocalRateLimit_packets_dropped`: Total number of packets dropped due to rate limiting
///

/// Config represents a RateLimitFilter's configuration.
#[derive(Serialize, Deserialize, Debug)]
struct Config {
    /// max_packets is the maximum number of packets allowed
    /// to be forwarded by the rate limiter in a given duration.
    max_packets: usize,
    /// period is the duration during which max_packets applies.
    /// If none is provided, it defaults to 1 second.
    #[serde(with = "humantime_serde")]
    period: Option<Duration>,
}

/// Creates instances of RateLimitFilter.
#[derive(Default)]
pub struct RateLimitFilterFactory;

/// A filter that implements rate limiting on packets based on
/// the token-bucket algorithm.
/// Packets that violate the rate limit are dropped.
/// It only applies rate limiting on packets that are destined for the
/// proxy's endpoints. All other packets flow through the filter untouched.
struct RateLimitFilter {
    /// available_tokens is how many tokens are left in the bucket any
    /// any given moment.
    available_tokens: Arc<AtomicUsize>,
    /// metrics reporter for this filter.
    metrics: Metrics,
    /// shutdown_tx signals the spawned token refill future to exit.
    shutdown_tx: Option<Sender<()>>,
}

impl FilterFactory for RateLimitFilterFactory {
    fn name(&self) -> String {
        "quilkin.extensions.filters.local_rate_limit.v1alpha1.LocalRateLimit".into()
    }

    fn create_filter(&self, args: CreateFilterArgs) -> Result<Box<dyn Filter>, Error> {
        let config: Config = args.parse_config()?;

        match config.period {
            Some(period) if period.lt(&Duration::from_millis(100)) => Err(Error::FieldInvalid {
                field: "period".into(),
                reason: "value must be at least 100ms".into(),
            }),
            _ => Ok(Box::new(RateLimitFilter::new(
                config,
                Metrics::new(&args.metrics_registry)?,
            ))),
        }
    }
}

impl RateLimitFilter {
    /// new returns a new RateLimitFilter. It spawns a future in the background
    /// that periodically refills the rate limiter's tokens.
    fn new(config: Config, metrics: Metrics) -> Self {
        let (shutdown_tx, mut shutdown_rx) = channel();

        let tokens = Arc::new(AtomicUsize::new(config.max_packets));

        let max_tokens = config.max_packets;
        let period = config.period.unwrap_or_else(|| Duration::from_secs(1));
        let available_tokens = tokens.clone();
        let _ = tokio::spawn(async move {
            let mut interval = time::interval_at(Instant::now() + period, period);
            loop {
                tokio::select! {
                    _ = interval.tick() => {
                        // Refill tokens.
                        let mut refilled = false;
                        while !refilled {
                            let remaining_tokens = available_tokens.load(Ordering::Relaxed);

                            refilled = available_tokens.compare_and_swap(
                                remaining_tokens,
                                max_tokens,
                                Ordering::Relaxed) == remaining_tokens;
                        }
                    },
                    _ = &mut shutdown_rx => {
                        return;
                    }
                }
            }
        });

        RateLimitFilter {
            available_tokens: tokens,
            metrics,
            shutdown_tx: Some(shutdown_tx),
        }
    }

    /// acquire_token is called on behalf of every packet that is eligible
    /// for rate limiting. It returns whether there exists a token in the current
    /// period - determining whether or not the packet should be forwarded or dropped.
    fn acquire_token(&self) -> Option<()> {
        loop {
            let remaining_tokens = self.available_tokens.load(Ordering::Relaxed);

            if remaining_tokens == 0 {
                return None;
            }

            if self.available_tokens.compare_and_swap(
                remaining_tokens,
                remaining_tokens - 1,
                Ordering::Relaxed,
            ) == remaining_tokens
            {
                return Some(());
            }
        }
    }
}

impl Drop for RateLimitFilter {
    fn drop(&mut self) {
        if let Some(shutdown_tx) = self.shutdown_tx.take() {
            shutdown_tx.send(()).ok();
        }
    }
}

impl Filter for RateLimitFilter {
    fn on_downstream_receive(&self, ctx: DownstreamContext) -> Option<DownstreamResponse> {
        self.acquire_token().map(|()| ctx.into()).or_else(|| {
            self.metrics.packets_dropped_total.inc();
            None
        })
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use prometheus::Registry;
    use tokio::time;

    use crate::cluster::Endpoint;
    use crate::config::Endpoints;
    use crate::extensions::filter_registry::DownstreamContext;
    use crate::extensions::filters::local_rate_limit::metrics::Metrics;
    use crate::extensions::filters::local_rate_limit::{Config, RateLimitFilter};
    use crate::extensions::Filter;
    use crate::test_utils::assert_filter_on_upstream_receive_no_change;

    fn rate_limiter(config: Config) -> RateLimitFilter {
        RateLimitFilter::new(config, Metrics::new(&Registry::default()).unwrap())
    }

    #[tokio::test]
    async fn initially_available_tokens() {
        // Test that we always start with the max number of tokens available.
        let r = rate_limiter(Config {
            max_packets: 3,
            period: Some(Duration::from_millis(100)),
        });

        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);
    }

    #[tokio::test]
    async fn token_exhaustion_and_refill() {
        let r = rate_limiter(Config {
            max_packets: 2,
            period: Some(Duration::from_millis(100)),
        });

        // Exhaust tokens
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);

        // Wait for refill
        let mut num_iterations = 0;
        let first_token = loop {
            let token = r.acquire_token();
            if token.is_some() {
                break token;
            }
            num_iterations += 1;
            if num_iterations >= 1000 {
                unreachable!("timed-out waiting for token refill");
            }
            time::sleep(Duration::from_millis(10)).await;
        };

        // Exhaust tokens again.
        assert_eq!(first_token, Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);
    }

    #[tokio::test]
    async fn token_refill_maximum() {
        // Test that we never refill more than the max_tokens specified.

        let r = rate_limiter(Config {
            max_packets: 3,
            period: Some(Duration::from_millis(30)),
        });

        // Use up some of the tokens.
        assert_eq!(r.acquire_token(), Some(()));

        // Wait for refill
        time::sleep(Duration::from_millis(110)).await;

        // Refill should not go over max token limit.
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);
    }

    #[tokio::test]
    async fn filter_with_no_available_tokens() {
        let r = rate_limiter(Config {
            max_packets: 0,
            period: Some(Duration::from_millis(100)),
        });

        // Check that other routes are not affected.
        assert_filter_on_upstream_receive_no_change(&r);

        // Check that we're rate limited.
        assert!(r
            .on_downstream_receive(DownstreamContext::new(
                Endpoints::new(vec![Endpoint::from_address(
                    "127.0.0.1:8080".parse().unwrap(),
                )])
                .unwrap()
                .into(),
                "127.0.0.1:8080".parse().unwrap(),
                vec![9],
            ))
            .is_none(),);
    }

    #[tokio::test]
    async fn filter_with_available_tokens() {
        let r = rate_limiter(Config {
            max_packets: 1,
            period: Some(Duration::from_millis(100)),
        });

        let result = r
            .on_downstream_receive(DownstreamContext::new(
                Endpoints::new(vec![Endpoint::from_address(
                    "127.0.0.1:8080".parse().unwrap(),
                )])
                .unwrap()
                .into(),
                "127.0.0.1:8080".parse().unwrap(),
                vec![9],
            ))
            .unwrap();
        assert_eq!(result.contents, vec![9]);
        // We should be out of tokens now.
        assert_eq!(None, r.acquire_token());

        // Check that other routes are not affected.
        assert_filter_on_upstream_receive_no_change(&r);
    }
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;
use std::ops::Sub;
use std::convert::From;

pub trait MultiObjective {
    fn num_objectives(&self) -> usize;

    fn cmp_objective(&self, other: &Self, objective: usize) -> Ordering;

    /// Calculates the distance of objective between self and other
    fn dist_objective(&self, other: &Self, objective: usize) -> f32;
}

#[derive(Debug, Clone)]
pub struct MultiObjective2<T>
    where T: Sized + PartialOrd + Copy + Clone
{
    pub objectives: [T; 2],
}

impl<T: Sized + PartialOrd + Copy + Clone> From<(T, T)> for MultiObjective2<T> {
    #[inline]
    fn from(t: (T, T)) -> MultiObjective2<T> {
        MultiObjective2 { objectives: [t.0, t.1] }
    }
}

impl<T, R> MultiObjective for MultiObjective2<T>
    where T: Copy + PartialOrd + Sub<Output = R>,
          R: Into<f32>
{
    #[inline]
    fn num_objectives(&self) -> usize {
        2
    }
    #[inline]
    fn cmp_objective(&self, other: &Self, objective: usize) -> Ordering {
        self.objectives[objective].partial_cmp(&other.objectives[objective]).unwrap()
    }
    #[inline]
    fn dist_objective(&self, other: &Self, objective: usize) -> f32 {
        (self.objectives[objective] - other.objectives[objective]).into()
    }
}

#[derive(Debug, Clone, Default)]
pub struct MultiObjective3<T>
    where T: Sized + PartialOrd + Copy + Clone + Default
{
    pub objectives: [T; 3],
}

impl<T, I> From<I> for MultiObjective3<T>
    where T: Sized + PartialOrd + Copy + Clone + Default,
          I: Iterator<Item = T>
{
    #[inline]
    fn from(iter: I) -> MultiObjective3<T> {
        let mut mo = Self::default();
        for (i, elm) in iter.into_iter().enumerate() {
            assert!(i < 3);
            mo.objectives[i] = elm;
        }
        mo
    }
}

impl<T, R> MultiObjective for MultiObjective3<T>
    where T: Copy + PartialOrd + Sub<Output = R> + Default,
          R: Into<f32>
{
    #[inline]
    fn num_objectives(&self) -> usize {
        3
    }
    #[inline]
    fn cmp_objective(&self, other: &Self, objective: usize) -> Ordering {
        self.objectives[objective].partial_cmp(&other.objectives[objective]).unwrap()
    }
    #[inline]
    fn dist_objective(&self, other: &Self, objective: usize) -> f32 {
        (self.objectives[objective] - other.objectives[objective]).into()
    }
}

#[derive(Debug, Clone)]
pub struct MultiObjective4<T>
    where T: Sized + PartialOrd + Copy + Clone
{
    pub objectives: [T; 4],
}

impl<T: Sized + PartialOrd + Copy + Clone> From<(T, T, T, T)> for MultiObjective4<T> {
    #[inline]
    fn from(t: (T, T, T, T)) -> MultiObjective4<T> {
        MultiObjective4 { objectives: [t.0, t.1, t.2, t.3] }
    }
}

impl<T, R> MultiObjective for MultiObjective4<T>
    where T: Copy + PartialOrd + Sub<Output = R>,
          R: Into<f32>
{
    #[inline]
    fn num_objectives(&self) -> usize {
        4
    }
    #[inline]
    fn cmp_objective(&self, other: &Self, objective: usize) -> Ordering {
        self.objectives[objective].partial_cmp(&other.objectives[objective]).unwrap()
    }
    #[inline]
    fn dist_objective(&self, other: &Self, objective: usize) -> f32 {
        (self.objectives[objective] - other.objectives[objective]).into()
    }
}
<commit_msg>Add MultiObjectiveVec<commit_after>use std::cmp::Ordering;
use std::ops::Sub;
use std::convert::From;

pub trait MultiObjective {
    fn num_objectives(&self) -> usize;

    fn cmp_objective(&self, other: &Self, objective: usize) -> Ordering;

    /// Calculates the distance of objective between self and other
    fn dist_objective(&self, other: &Self, objective: usize) -> f32;
}

#[derive(Debug, Clone)]
pub struct MultiObjective2<T>
    where T: Sized + PartialOrd + Copy + Clone
{
    pub objectives: [T; 2],
}

impl<T: Sized + PartialOrd + Copy + Clone> From<(T, T)> for MultiObjective2<T> {
    #[inline]
    fn from(t: (T, T)) -> MultiObjective2<T> {
        MultiObjective2 { objectives: [t.0, t.1] }
    }
}

impl<T, R> MultiObjective for MultiObjective2<T>
    where T: Copy + PartialOrd + Sub<Output = R>,
          R: Into<f32>
{
    #[inline]
    fn num_objectives(&self) -> usize {
        2
    }
    #[inline]
    fn cmp_objective(&self, other: &Self, objective: usize) -> Ordering {
        self.objectives[objective].partial_cmp(&other.objectives[objective]).unwrap()
    }
    #[inline]
    fn dist_objective(&self, other: &Self, objective: usize) -> f32 {
        (self.objectives[objective] - other.objectives[objective]).into()
    }
}

#[derive(Debug, Clone, Default)]
pub struct MultiObjective3<T>
    where T: Sized + PartialOrd + Copy + Clone + Default
{
    pub objectives: [T; 3],
}

impl<T, I> From<I> for MultiObjective3<T>
    where T: Sized + PartialOrd + Copy + Clone + Default,
          I: Iterator<Item = T>
{
    #[inline]
    fn from(iter: I) -> MultiObjective3<T> {
        let mut mo = Self::default();
        for (i, elm) in iter.into_iter().enumerate() {
            assert!(i < 3);
            mo.objectives[i] = elm;
        }
        mo
    }
}

impl<T, R> MultiObjective for MultiObjective3<T>
    where T: Copy + PartialOrd + Sub<Output = R> + Default,
          R: Into<f32>
{
    #[inline]
    fn num_objectives(&self) -> usize {
        3
    }
    #[inline]
    fn cmp_objective(&self, other: &Self, objective: usize) -> Ordering {
        self.objectives[objective].partial_cmp(&other.objectives[objective]).unwrap()
    }
    #[inline]
    fn dist_objective(&self, other: &Self, objective: usize) -> f32 {
        (self.objectives[objective] - other.objectives[objective]).into()
    }
}

#[derive(Debug, Clone)]
pub struct MultiObjective4<T>
    where T: Sized + PartialOrd + Copy + Clone
{
    pub objectives: [T; 4],
}

impl<T: Sized + PartialOrd + Copy + Clone> From<(T, T, T, T)> for MultiObjective4<T> {
    #[inline]
    fn from(t: (T, T, T, T)) -> MultiObjective4<T> {
        MultiObjective4 { objectives: [t.0, t.1, t.2, t.3] }
    }
}

impl<T, R> MultiObjective for MultiObjective4<T>
    where T: Copy + PartialOrd + Sub<Output = R>,
          R: Into<f32>
{
    #[inline]
    fn num_objectives(&self) -> usize {
        4
    }
    #[inline]
    fn cmp_objective(&self, other: &Self, objective: usize) -> Ordering {
        self.objectives[objective].partial_cmp(&other.objectives[objective]).unwrap()
    }
    #[inline]
    fn dist_objective(&self, other: &Self, objective: usize) -> f32 {
        (self.objectives[objective] - other.objectives[objective]).into()
    }
}

#[derive(Debug, Clone)]
pub struct MultiObjectiveVec<T>
    where T: Sized + PartialOrd + Copy + Clone
{
    objectives: Vec<T>,
}

impl<T: Sized + PartialOrd + Copy + Clone> From<Vec<T>> for MultiObjectiveVec<T> {
    #[inline]
    fn from(t: Vec<T>) -> MultiObjectiveVec<T> {
        MultiObjectiveVec { objectives: t }
    }
}

impl<T: Sized + PartialOrd + Copy + Clone> AsRef<[T]> for MultiObjectiveVec<T> {
    fn as_ref(&self) -> &[T] {
        &self.objectives
    }
}

impl<T, R> MultiObjective for MultiObjectiveVec<T>
    where T: Copy + PartialOrd + Sub<Output = R>,
          R: Into<f32>
{
    #[inline]
    fn num_objectives(&self) -> usize {
        self.objectives.len()
    }
    #[inline]
    fn cmp_objective(&self, other: &Self, objective: usize) -> Ordering {
        self.objectives[objective].partial_cmp(&other.objectives[objective]).unwrap()
    }
    #[inline]
    fn dist_objective(&self, other: &Self, objective: usize) -> f32 {
        (self.objectives[objective] - other.objectives[objective]).into()
    }
}
<|endoftext|>


<commit_before>//! Syntax of the language.
//!
//! For now, most of the language is an EDSL describing an augmented GLSL with a few keywords.
//!
//! This module also exports some useful functions to cope with the syntax types.

pub use glsl::syntax::*;
use glsl::writer;
use std::error::Error;
use std::fmt::{self, Write};
use std::iter::once;

/// A module.
///
/// A module has a list of imports and a list of GLSL extern declaration.
#[derive(Clone, Debug, PartialEq)]
pub struct Module {
  pub imports: Vec<ImportList>,
  pub glsl: Vec<ExternalDeclaration>
}

/// A non-empty import list.
///
/// It consists of a module path, like `Foo.Bar.Zoo`, and a list of symbols to load from that path,
/// as in `rick, marty, doggo`.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ImportList {
  pub module: ModulePath,
  pub list: Vec<ModuleSymbol>
}

/// A module path is a non-empty list of module(s), representing a hierarchy.
#[derive(Clone, Debug, Hash, Eq, PartialEq)]
pub struct ModulePath {
  pub path: Vec<ModuleName>
}

/// Name of a module.
pub type ModuleName = String;

/// A symbol, living in a module.
pub type ModuleSymbol = String;

pub type ExpectedNumberOfArgs = usize;
pub type FoundNumberOfArgs = usize;

/// GLSL conversion error.
///
/// Such an errors can happen when a module is ill-formed.
#[derive(Clone, Debug, PartialEq)]
pub enum GLSLConversionError {
  /// No vertex shader was found. A vertex shader is required in order to build a shading pipeline.
  NoVertexShader,
  /// No fragment shader was found. A fragment shader is required in order to build a shading pipeline.
  NoFragmentShader,
  /// The output must not have a qualifier.
  OutputHasMainQualifier,
  /// The returned value must not be a struct.
  ReturnTypeMustBeAStruct(TypeSpecifier),
  /// The first field has the wrong type.
  WrongOutputFirstField(StructFieldSpecifier),
  /// The field of a type used as output cannot be a struct.
  ///
  /// This variant also gives the index of the field.
  OutputFieldCannotBeStruct(usize, StructSpecifier),
  /// The field of a type used as output cannot be a type name.
  ///
  /// This variant also gives the index of the field.
  OutputFieldCannotBeTypeName(usize, TypeName),
  /// The field of a type used as output cannot have several identifiers (only one).
  ///
  /// This variant also gives the index of the field.
  OutputFieldCannotHaveSeveralIdentifiers(usize, StructFieldSpecifier),
  /// The input type is unknown.
  UnknownInputType(TypeName),
  /// Wrong number of arguments.
  WrongNumberOfArgs(ExpectedNumberOfArgs, FoundNumberOfArgs),
  /// The type is not a required type name.
  NotTypeName,
  /// The geometry input is wrong.
  WrongGeometryInput,
  /// The geometry input’s dimension is wrong.
  WrongGeometryInputDim(usize),
  /// The geometry output layout is wrong.
  WrongGeometryOutputLayout(Option<TypeQualifier>)
}

impl fmt::Display for GLSLConversionError {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    f.write_str(self.description())
  }
}

impl Error for GLSLConversionError {
  fn description(&self) -> &str {
    match *self {
      GLSLConversionError::NoVertexShader => "no vertex shader",
      GLSLConversionError::NoFragmentShader => "no fragment shader",
      GLSLConversionError::OutputHasMainQualifier => "output has main qualifier(s)",
      GLSLConversionError::ReturnTypeMustBeAStruct(_) => "return type is not a struct",
      GLSLConversionError::WrongOutputFirstField(_) => "first field’s type is forbidden as output",
      GLSLConversionError::OutputFieldCannotBeStruct(..) => "output field cannot be a struct",
      GLSLConversionError::OutputFieldCannotBeTypeName(..) => "output field cannot be a type name",
      GLSLConversionError::OutputFieldCannotHaveSeveralIdentifiers(..) => "output field cannot have several identifiers",
      GLSLConversionError::UnknownInputType(_) => "unknown input type",
      GLSLConversionError::WrongNumberOfArgs(..) => "wrong number of arguments",
      GLSLConversionError::NotTypeName => "not a type name",
      GLSLConversionError::WrongGeometryInput => "wrong geometry input",
      GLSLConversionError::WrongGeometryInputDim(_) => "wrong geometry input’s dimension",
      GLSLConversionError::WrongGeometryOutputLayout(_) => "wrong geometry output layout"
    }
  }
}

/// Sink single declarations as external declarations.
pub fn sink_single_as_ext_decls<'a, F, I>(sink: &mut F, s: I)
                                      where I: IntoIterator<Item = &'a SingleDeclaration>,
                                            F: Write {
  for sd in s {
    let ed = single_to_external_declaration(sd.clone());
    writer::glsl::show_external_declaration(sink, &ed);
  }
}

/// Turn a `SingleDeclaration` into an `ExternalDeclaration`.
pub fn single_to_external_declaration(sd: SingleDeclaration) -> ExternalDeclaration {
  ExternalDeclaration::Declaration(
    Declaration::InitDeclaratorList(
      InitDeclaratorList {
        head: sd,
        tail: Vec::new()
      }
    )
  )
}

/// Replace an output declaration by its input declaration dual.
///
/// Useful when an input interface must match an output one.
pub fn input_from_output(output: SingleDeclaration, has_array: bool) -> SingleDeclaration {
  let
    qualifier = output.ty.qualifier.map(|q| {
      TypeQualifier {
        qualifiers: q.qualifiers.into_iter().map(|qs| {
          match qs {
            TypeQualifierSpec::Storage(StorageQualifier::Out) =>
              TypeQualifierSpec::Storage(StorageQualifier::In),
            _ => qs
          }
        }).collect()
      }
    });
  let ty =
    TypeSpecifier {
      array_specifier: if has_array { Some(ArraySpecifier::Unsized) } else { None },
      .. output.ty.ty
    };

  SingleDeclaration {
    ty: FullySpecifiedType { qualifier, ty },
    .. output
  }
}

/// Replace outputs by inputs.
pub fn inputs_from_outputs(outputs: &[SingleDeclaration], has_array: bool) -> Vec<SingleDeclaration> {
  outputs.into_iter().map(|sd| input_from_output(sd.clone(), has_array)).collect()
}

/// Map a StructFieldSpecifier to an ExternalDeclaration.
///
/// Typically suitable for generating an output from a struct field.
pub fn field_to_single_decl(field: &StructFieldSpecifier, prefix: &str) -> SingleDeclaration {
  let base_qualifier = TypeQualifierSpec::Storage(StorageQualifier::Out);
  let qualifier = match field.qualifier {
    Some(ref qual) =>
      TypeQualifier {
        qualifiers: qual.clone().qualifiers.into_iter().chain(once(base_qualifier)).collect()
      },
    None => TypeQualifier {
      qualifiers: vec![base_qualifier]
    }
  };
  let fsty = FullySpecifiedType {
    qualifier: Some(qualifier),
    ty: field.ty.clone()
  };

  SingleDeclaration {
    ty: fsty,
    name: Some(prefix.to_owned() + &field.identifiers[0].0),
    array_specifier: field.identifiers[0].1.clone(),
    initializer: None
  }
}

/// Map a struct’s fields to a Vec<ExternalDeclaration>.
///
/// Typically suitable for generating outputs from a struct fields.
pub fn fields_to_single_decls(fields: &[StructFieldSpecifier], prefix: &str)
                              -> Result<Vec<SingleDeclaration>, GLSLConversionError> {
  let mut outputs = Vec::new();

  for (i, field) in fields.into_iter().enumerate() {
    match field.ty.ty {
      TypeSpecifierNonArray::Struct(ref s) => {
        return Err(GLSLConversionError::OutputFieldCannotBeStruct(i + 1, s.clone()));
      }
      TypeSpecifierNonArray::TypeName(ref t) => {
        return Err(GLSLConversionError::OutputFieldCannotBeTypeName(i + 1, t.clone()));
      }
      _ => ()
    }

    if field.identifiers.len() > 1 {
      return Err(GLSLConversionError::OutputFieldCannotHaveSeveralIdentifiers(i + 1, field.clone()));
    }

    outputs.push(field_to_single_decl(&field, prefix));
  }

  Ok(outputs)
}

/// Filter out a function definition by removing its unused arguments.
pub fn remove_unused_args_fn(f: &FunctionDefinition) -> FunctionDefinition {
  let f = f.clone();

  FunctionDefinition {
    prototype: FunctionPrototype {
      parameters: f.prototype.parameters.into_iter().filter(is_fn_arg_named).collect(),
      .. f.prototype
    },
    .. f
  }
}

pub fn is_fn_arg_named(arg: &FunctionParameterDeclaration) -> bool {
  if let FunctionParameterDeclaration::Named(..) = *arg {
    true
  } else {
    false
  }
}

/// Extract the type name of a function argument. If the argument’s type is not a typename,
/// nothing is returned.
/// Get the fully specified type of a function’s argument.
pub fn fn_arg_as_fully_spec_ty(arg: &FunctionParameterDeclaration) -> FullySpecifiedType {
  match *arg {
    FunctionParameterDeclaration::Named(ref qualifier, FunctionParameterDeclarator {
      ref ty,
      ..
    }) => FullySpecifiedType {
      qualifier: qualifier.clone(),
      ty: ty.clone()
    },
    FunctionParameterDeclaration::Unnamed(ref qualifier, ref ty) => {
      FullySpecifiedType {
        qualifier: qualifier.clone(),
        ty: ty.clone()
      }
    }
  }
}

/// Extract the type name of a fully specified type. If the type is not a typename, nothing is
/// returned.
pub fn get_ty_name_from_fully_spec_ty(fst: &FullySpecifiedType) -> Result<TypeName, GLSLConversionError> {
  if let TypeSpecifierNonArray::TypeName(ref n) = fst.ty.ty {
    Ok(n.clone())
  } else {
    Err(GLSLConversionError::NotTypeName)
  }
}

/// Get the type name of the argument of a unary function. If the argument is not unary, fail
/// with the approriate error.
pub fn get_fn1_input_ty_name(f: &FunctionDefinition) -> Result<TypeName, GLSLConversionError> {
  let slice = f.prototype.parameters.as_slice();
  match slice {
    &[ref arg] => {
      let fst = fn_arg_as_fully_spec_ty(arg);
      get_ty_name_from_fully_spec_ty(&fst)
    }
    _ => Err(GLSLConversionError::WrongNumberOfArgs(1, slice.len()))
  }
}

/// Get the return type of a function by looking up its definition in the provided slice.
pub fn get_fn_ret_ty(f: &FunctionDefinition, structs: &[StructSpecifier]) -> Result<StructSpecifier, GLSLConversionError> {
  struct_from_ty_spec(&f.prototype.ty.ty, structs)
}

/// Get the struct definition associated with a type specifier.
pub fn struct_from_ty_spec(
  ty_spec: &TypeSpecifier,
  structs: &[StructSpecifier]
) -> Result<StructSpecifier, GLSLConversionError> {
  if let TypeSpecifierNonArray::TypeName(ref name) = ty_spec.ty {
    if let Some(ref ty) = structs.iter().find(|ref s| s.name.as_ref() == Some(name)) {
      Ok((*ty).clone())
    } else {
      Err(GLSLConversionError::ReturnTypeMustBeAStruct(ty_spec.clone()))
    }
  } else {
    Err(GLSLConversionError::ReturnTypeMustBeAStruct(ty_spec.clone()))
  }
}

/// Drop the first field of a struct.
pub fn drop_1st_field(s: &StructSpecifier) -> StructSpecifier {
  StructSpecifier {
    name: s.name.clone(),
    fields: s.fields.iter().skip(1).cloned().collect(),
  }
}
<commit_msg>Add more doc. #37<commit_after>//! Syntax of the language.
//!
//! For now, most of the language is an EDSL describing an augmented GLSL with a few keywords.
//!
//! This module also exports some useful functions to cope with the syntax types.

pub use glsl::syntax::*;
use glsl::writer;
use std::error::Error;
use std::fmt::{self, Write};
use std::iter::once;

/// A module.
///
/// A module has a list of imports and a list of GLSL extern declaration.
#[derive(Clone, Debug, PartialEq)]
pub struct Module {
  /// List of imports for this module.
  pub imports: Vec<ImportList>,
  /// The GLSL body of the module.
  pub glsl: Vec<ExternalDeclaration>
}

/// A non-empty import list.
///
/// It consists of a module path, like `Foo.Bar.Zoo`, and a list of symbols to load from that path,
/// as in `rick, marty, doggo`.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ImportList {
  /// The module path to load symbols from.
  pub module: ModulePath,
  /// List of symbols to import.
  pub list: Vec<ModuleSymbol>
}

/// A module path is a non-empty list of module(s), representing a hierarchy.
#[derive(Clone, Debug, Hash, Eq, PartialEq)]
pub struct ModulePath {
  /// Hierarchical path of modules leading to the module we want to import (included in the path).
  pub path: Vec<ModuleName>
}

/// Name of a module.
pub type ModuleName = String;

/// A symbol, living in a module.
pub type ModuleSymbol = String;

pub type ExpectedNumberOfArgs = usize;
pub type FoundNumberOfArgs = usize;

/// GLSL conversion error.
///
/// Such an errors can happen when a module is ill-formed.
#[derive(Clone, Debug, PartialEq)]
pub enum GLSLConversionError {
  /// No vertex shader was found. A vertex shader is required in order to build a shading pipeline.
  NoVertexShader,
  /// No fragment shader was found. A fragment shader is required in order to build a shading pipeline.
  NoFragmentShader,
  /// The output must not have a qualifier.
  OutputHasMainQualifier,
  /// The returned value must not be a struct.
  ReturnTypeMustBeAStruct(TypeSpecifier),
  /// The first field has the wrong type.
  WrongOutputFirstField(StructFieldSpecifier),
  /// The field of a type used as output cannot be a struct.
  ///
  /// This variant also gives the index of the field.
  OutputFieldCannotBeStruct(usize, StructSpecifier),
  /// The field of a type used as output cannot be a type name.
  ///
  /// This variant also gives the index of the field.
  OutputFieldCannotBeTypeName(usize, TypeName),
  /// The field of a type used as output cannot have several identifiers (only one).
  ///
  /// This variant also gives the index of the field.
  OutputFieldCannotHaveSeveralIdentifiers(usize, StructFieldSpecifier),
  /// The input type is unknown.
  UnknownInputType(TypeName),
  /// Wrong number of arguments.
  WrongNumberOfArgs(ExpectedNumberOfArgs, FoundNumberOfArgs),
  /// The type is not a required type name.
  NotTypeName,
  /// The geometry input is wrong.
  WrongGeometryInput,
  /// The geometry input’s dimension is wrong.
  WrongGeometryInputDim(usize),
  /// The geometry output layout is wrong.
  WrongGeometryOutputLayout(Option<TypeQualifier>)
}

impl fmt::Display for GLSLConversionError {
  fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
    f.write_str(self.description())
  }
}

impl Error for GLSLConversionError {
  fn description(&self) -> &str {
    match *self {
      GLSLConversionError::NoVertexShader => "no vertex shader",
      GLSLConversionError::NoFragmentShader => "no fragment shader",
      GLSLConversionError::OutputHasMainQualifier => "output has main qualifier(s)",
      GLSLConversionError::ReturnTypeMustBeAStruct(_) => "return type is not a struct",
      GLSLConversionError::WrongOutputFirstField(_) => "first field’s type is forbidden as output",
      GLSLConversionError::OutputFieldCannotBeStruct(..) => "output field cannot be a struct",
      GLSLConversionError::OutputFieldCannotBeTypeName(..) => "output field cannot be a type name",
      GLSLConversionError::OutputFieldCannotHaveSeveralIdentifiers(..) => "output field cannot have several identifiers",
      GLSLConversionError::UnknownInputType(_) => "unknown input type",
      GLSLConversionError::WrongNumberOfArgs(..) => "wrong number of arguments",
      GLSLConversionError::NotTypeName => "not a type name",
      GLSLConversionError::WrongGeometryInput => "wrong geometry input",
      GLSLConversionError::WrongGeometryInputDim(_) => "wrong geometry input’s dimension",
      GLSLConversionError::WrongGeometryOutputLayout(_) => "wrong geometry output layout"
    }
  }
}

/// Sink single declarations as external declarations.
pub fn sink_single_as_ext_decls<'a, F, I>(sink: &mut F, s: I)
                                      where I: IntoIterator<Item = &'a SingleDeclaration>,
                                            F: Write {
  for sd in s {
    let ed = single_to_external_declaration(sd.clone());
    writer::glsl::show_external_declaration(sink, &ed);
  }
}

/// Turn a `SingleDeclaration` into an `ExternalDeclaration`.
pub fn single_to_external_declaration(sd: SingleDeclaration) -> ExternalDeclaration {
  ExternalDeclaration::Declaration(
    Declaration::InitDeclaratorList(
      InitDeclaratorList {
        head: sd,
        tail: Vec::new()
      }
    )
  )
}

/// Replace an output declaration by its input declaration dual.
///
/// Useful when an input interface must match an output one.
pub fn input_from_output(output: SingleDeclaration, has_array: bool) -> SingleDeclaration {
  let
    qualifier = output.ty.qualifier.map(|q| {
      TypeQualifier {
        qualifiers: q.qualifiers.into_iter().map(|qs| {
          match qs {
            TypeQualifierSpec::Storage(StorageQualifier::Out) =>
              TypeQualifierSpec::Storage(StorageQualifier::In),
            _ => qs
          }
        }).collect()
      }
    });
  let ty =
    TypeSpecifier {
      array_specifier: if has_array { Some(ArraySpecifier::Unsized) } else { None },
      .. output.ty.ty
    };

  SingleDeclaration {
    ty: FullySpecifiedType { qualifier, ty },
    .. output
  }
}

/// Replace outputs by inputs.
pub fn inputs_from_outputs(outputs: &[SingleDeclaration], has_array: bool) -> Vec<SingleDeclaration> {
  outputs.into_iter().map(|sd| input_from_output(sd.clone(), has_array)).collect()
}

/// Map a StructFieldSpecifier to an ExternalDeclaration.
///
/// Typically suitable for generating an output from a struct field.
pub fn field_to_single_decl(field: &StructFieldSpecifier, prefix: &str) -> SingleDeclaration {
  let base_qualifier = TypeQualifierSpec::Storage(StorageQualifier::Out);
  let qualifier = match field.qualifier {
    Some(ref qual) =>
      TypeQualifier {
        qualifiers: qual.clone().qualifiers.into_iter().chain(once(base_qualifier)).collect()
      },
    None => TypeQualifier {
      qualifiers: vec![base_qualifier]
    }
  };
  let fsty = FullySpecifiedType {
    qualifier: Some(qualifier),
    ty: field.ty.clone()
  };

  SingleDeclaration {
    ty: fsty,
    name: Some(prefix.to_owned() + &field.identifiers[0].0),
    array_specifier: field.identifiers[0].1.clone(),
    initializer: None
  }
}

/// Map a struct’s fields to a Vec<ExternalDeclaration>.
///
/// Typically suitable for generating outputs from a struct fields.
pub fn fields_to_single_decls(fields: &[StructFieldSpecifier], prefix: &str)
                              -> Result<Vec<SingleDeclaration>, GLSLConversionError> {
  let mut outputs = Vec::new();

  for (i, field) in fields.into_iter().enumerate() {
    match field.ty.ty {
      TypeSpecifierNonArray::Struct(ref s) => {
        return Err(GLSLConversionError::OutputFieldCannotBeStruct(i + 1, s.clone()));
      }
      TypeSpecifierNonArray::TypeName(ref t) => {
        return Err(GLSLConversionError::OutputFieldCannotBeTypeName(i + 1, t.clone()));
      }
      _ => ()
    }

    if field.identifiers.len() > 1 {
      return Err(GLSLConversionError::OutputFieldCannotHaveSeveralIdentifiers(i + 1, field.clone()));
    }

    outputs.push(field_to_single_decl(&field, prefix));
  }

  Ok(outputs)
}

/// Filter out a function definition by removing its unused arguments.
pub fn remove_unused_args_fn(f: &FunctionDefinition) -> FunctionDefinition {
  let f = f.clone();

  FunctionDefinition {
    prototype: FunctionPrototype {
      parameters: f.prototype.parameters.into_iter().filter(is_fn_arg_named).collect(),
      .. f.prototype
    },
    .. f
  }
}

pub fn is_fn_arg_named(arg: &FunctionParameterDeclaration) -> bool {
  if let FunctionParameterDeclaration::Named(..) = *arg {
    true
  } else {
    false
  }
}

/// Extract the type name of a function argument. If the argument’s type is not a typename,
/// nothing is returned.
/// Get the fully specified type of a function’s argument.
pub fn fn_arg_as_fully_spec_ty(arg: &FunctionParameterDeclaration) -> FullySpecifiedType {
  match *arg {
    FunctionParameterDeclaration::Named(ref qualifier, FunctionParameterDeclarator {
      ref ty,
      ..
    }) => FullySpecifiedType {
      qualifier: qualifier.clone(),
      ty: ty.clone()
    },
    FunctionParameterDeclaration::Unnamed(ref qualifier, ref ty) => {
      FullySpecifiedType {
        qualifier: qualifier.clone(),
        ty: ty.clone()
      }
    }
  }
}

/// Extract the type name of a fully specified type. If the type is not a typename, nothing is
/// returned.
pub fn get_ty_name_from_fully_spec_ty(fst: &FullySpecifiedType) -> Result<TypeName, GLSLConversionError> {
  if let TypeSpecifierNonArray::TypeName(ref n) = fst.ty.ty {
    Ok(n.clone())
  } else {
    Err(GLSLConversionError::NotTypeName)
  }
}

/// Get the type name of the argument of a unary function. If the argument is not unary, fail
/// with the approriate error.
pub fn get_fn1_input_ty_name(f: &FunctionDefinition) -> Result<TypeName, GLSLConversionError> {
  let slice = f.prototype.parameters.as_slice();
  match slice {
    &[ref arg] => {
      let fst = fn_arg_as_fully_spec_ty(arg);
      get_ty_name_from_fully_spec_ty(&fst)
    }
    _ => Err(GLSLConversionError::WrongNumberOfArgs(1, slice.len()))
  }
}

/// Get the return type of a function by looking up its definition in the provided slice.
pub fn get_fn_ret_ty(f: &FunctionDefinition, structs: &[StructSpecifier]) -> Result<StructSpecifier, GLSLConversionError> {
  struct_from_ty_spec(&f.prototype.ty.ty, structs)
}

/// Get the struct definition associated with a type specifier.
pub fn struct_from_ty_spec(
  ty_spec: &TypeSpecifier,
  structs: &[StructSpecifier]
) -> Result<StructSpecifier, GLSLConversionError> {
  if let TypeSpecifierNonArray::TypeName(ref name) = ty_spec.ty {
    if let Some(ref ty) = structs.iter().find(|ref s| s.name.as_ref() == Some(name)) {
      Ok((*ty).clone())
    } else {
      Err(GLSLConversionError::ReturnTypeMustBeAStruct(ty_spec.clone()))
    }
  } else {
    Err(GLSLConversionError::ReturnTypeMustBeAStruct(ty_spec.clone()))
  }
}

/// Drop the first field of a struct.
pub fn drop_1st_field(s: &StructSpecifier) -> StructSpecifier {
  StructSpecifier {
    name: s.name.clone(),
    fields: s.fields.iter().skip(1).cloned().collect(),
  }
}
<|endoftext|>