<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::collections::HashMap;
use webrender_traits::{FontKey, FontRenderMode, GlyphDimensions};
use webrender_traits::{GlyphKey, GlyphOptions};
use gamma_lut::{GammaLut, Color as ColorLut};

use dwrote;

lazy_static! {
    static ref DEFAULT_FONT_DESCRIPTOR: dwrote::FontDescriptor = dwrote::FontDescriptor {
        family_name: "Arial".to_owned(),
        weight: dwrote::FontWeight::Regular,
        stretch: dwrote::FontStretch::Normal,
        style: dwrote::FontStyle::Normal,
    };
}

pub struct FontContext {
    fonts: HashMap<FontKey, dwrote::FontFace>,
    gamma_lut: GammaLut,
    gdi_gamma_lut: GammaLut,
}

// DirectWrite is safe to use on multiple threads and non-shareable resources are
// all hidden inside their font context.
unsafe impl Send for FontContext {}

pub struct RasterizedGlyph {
    pub width: u32,
    pub height: u32,
    pub bytes: Vec<u8>,
}

fn dwrite_texture_type(render_mode: FontRenderMode) ->
                       dwrote::DWRITE_TEXTURE_TYPE {
    match render_mode {
        FontRenderMode::Mono => dwrote::DWRITE_TEXTURE_ALIASED_1x1 ,
        FontRenderMode::Alpha |
        FontRenderMode::Subpixel => dwrote::DWRITE_TEXTURE_CLEARTYPE_3x1,
    }
}

fn dwrite_measure_mode(render_mode: FontRenderMode, options: Option<GlyphOptions>) ->
                       dwrote::DWRITE_MEASURING_MODE {
    if let Some(GlyphOptions{ force_gdi_rendering: true, .. }) = options {
        return dwrote::DWRITE_MEASURING_MODE_GDI_CLASSIC;
    }

    match render_mode {
        FontRenderMode::Mono => dwrote::DWRITE_MEASURING_MODE_GDI_NATURAL,
        FontRenderMode::Alpha |
        FontRenderMode::Subpixel => dwrote::DWRITE_MEASURING_MODE_NATURAL,
    }
}

fn dwrite_render_mode(font_face: &dwrote::FontFace,
                      render_mode: FontRenderMode,
                      em_size: f32,
                      measure_mode: dwrote::DWRITE_MEASURING_MODE,
                      options: Option<GlyphOptions>) ->
                      dwrote::DWRITE_RENDERING_MODE {
    if let Some(GlyphOptions{ force_gdi_rendering: true, .. }) = options {
        return dwrote::DWRITE_RENDERING_MODE_GDI_CLASSIC;
    }

    let dwrite_render_mode = match render_mode {
        FontRenderMode::Mono => dwrote::DWRITE_RENDERING_MODE_ALIASED,
        FontRenderMode::Alpha |
        FontRenderMode::Subpixel => {
            font_face.get_recommended_rendering_mode_default_params(em_size,
                                                                    1.0,
                                                                    measure_mode)
        },
    };

    if dwrite_render_mode  == dwrote::DWRITE_RENDERING_MODE_OUTLINE {
        // Outline mode is not supported
        return dwrote::DWRITE_RENDERING_MODE_CLEARTYPE_NATURAL_SYMMETRIC;
    }

    dwrite_render_mode
}

fn get_glyph_dimensions_with_analysis(analysis: dwrote::GlyphRunAnalysis,
                                      texture_type: dwrote::DWRITE_TEXTURE_TYPE)
                                      -> Option<GlyphDimensions> {
    let bounds = analysis.get_alpha_texture_bounds(texture_type);

    let width = (bounds.right - bounds.left) as u32;
    let height = (bounds.bottom - bounds.top) as u32;

    // Alpha texture bounds can sometimes return an empty rect
    // Such as for spaces
    if width == 0 || height == 0 {
        return None
    }

    Some(GlyphDimensions {
        left: bounds.left,
        top: -bounds.top,
        width: width,
        height: height,
    })
}

impl FontContext {
    pub fn new() -> FontContext {
        // These are the default values we use in Gecko.
        // We use a gamma value of 2.3 for gdi fonts
        // TODO: Fetch this data from Gecko itself.
        let contrast = 1.0;
        let gamma = 1.8;
        let gdi_gamma = 2.3;
        FontContext {
            fonts: HashMap::new(),
            gamma_lut: GammaLut::new(contrast, gamma, gamma),
            gdi_gamma_lut: GammaLut::new(contrast, gdi_gamma, gdi_gamma),
        }
    }

    pub fn has_font(&self, font_key: &FontKey) -> bool {
        self.fonts.contains_key(font_key)
    }

    pub fn add_raw_font(&mut self, font_key: &FontKey, data: &[u8], index: u32) {
        if self.fonts.contains_key(font_key) {
            return
        }

        if let Some(font_file) = dwrote::FontFile::new_from_data(data) {
            let face = font_file.create_face(index, dwrote::DWRITE_FONT_SIMULATIONS_NONE);
            self.fonts.insert((*font_key).clone(), face);
        } else {
            // XXX add_raw_font needs to have a way to return an error
            debug!("DWrite WR failed to load font from data, using Arial instead");
            self.add_native_font(font_key, DEFAULT_FONT_DESCRIPTOR.clone());
        }
    }

    pub fn add_native_font(&mut self, font_key: &FontKey, font_handle: dwrote::FontDescriptor) {
        if self.fonts.contains_key(font_key) {
            return
        }

        let system_fc = dwrote::FontCollection::system();
        let font = system_fc.get_font_from_descriptor(&font_handle).unwrap();
        let face = font.create_font_face();
        self.fonts.insert((*font_key).clone(), face);
    }

    pub fn delete_font(&mut self, font_key: &FontKey) {
        self.fonts.remove(font_key);
    }

    // Assumes RGB format from dwrite, which is 3 bytes per pixel as dwrite
    // doesn't output an alpha value via GlyphRunAnalysis::CreateAlphaTexture
    #[allow(dead_code)]
    fn print_glyph_data(&self, data: &[u8], width: usize, height: usize) {
        // Rust doesn't have step_by support on stable :(
        for i in 0..height {
            let current_height = i * width * 3;

            for pixel in data[current_height .. current_height + (width * 3)].chunks(3) {
                let r = pixel[0];
                let g = pixel[1];
                let b = pixel[2];
                print!("({}, {}, {}) ", r, g, b, );
            }
            println!("");
        }
    }

    fn create_glyph_analysis(&self, key: &GlyphKey,
                            render_mode: FontRenderMode,
                            options: Option<GlyphOptions>) ->
                            dwrote::GlyphRunAnalysis {
        let face = self.fonts.get(&key.font_key).unwrap();
        let glyph = key.index as u16;
        let advance = 0.0f32;
        let offset = dwrote::GlyphOffset { advanceOffset: 0.0, ascenderOffset: 0.0 };

        let glyph_run = dwrote::DWRITE_GLYPH_RUN {
            fontFace: unsafe { face.as_ptr() },
            fontEmSize: key.size.to_f32_px(), // size in DIPs (1/96", same as CSS pixels)
            glyphCount: 1,
            glyphIndices: &glyph,
            glyphAdvances: &advance,
            glyphOffsets: &offset,
            isSideways: 0,
            bidiLevel: 0,
        };

        let dwrite_measure_mode = dwrite_measure_mode(render_mode, options);
        let dwrite_render_mode = dwrite_render_mode(face,
                                                    render_mode,
                                                    key.size.to_f32_px(),
                                                    dwrite_measure_mode,
                                                    options);

        let (x_offset, y_offset) = key.subpixel_point.to_f64();
        let transform = Some(
                        dwrote::DWRITE_MATRIX { m11: 1.0, m12: 0.0, m21: 0.0, m22: 1.0,
                                                dx: x_offset as f32, dy: y_offset as f32 }
                        );

        dwrote::GlyphRunAnalysis::create(&glyph_run, 1.0, transform,
                                         dwrite_render_mode,
                                         dwrite_measure_mode,
                                         0.0, 0.0)
    }

    // TODO: Pipe GlyphOptions into glyph_dimensions too
    pub fn get_glyph_dimensions(&self,
                                key: &GlyphKey)
                                -> Option<GlyphDimensions> {
        // Probably have to default to something else here.
        let render_mode = FontRenderMode::Subpixel;
        let analysis = self.create_glyph_analysis(key, render_mode, None);

        let texture_type = dwrite_texture_type(render_mode);
        get_glyph_dimensions_with_analysis(analysis, texture_type)
    }

    // DWRITE gives us values in RGB. WR doesn't really touch it after. Note, CG returns in BGR
    // TODO: Decide whether all fonts should return RGB or BGR
    fn convert_to_rgba(&self, pixels: &[u8], render_mode: FontRenderMode) -> Vec<u8> {
        match render_mode {
            FontRenderMode::Mono => {
                let mut rgba_pixels: Vec<u8> = vec![0; pixels.len() * 4];
                for i in 0..pixels.len() {
                    rgba_pixels[i*4+0] = pixels[i];
                    rgba_pixels[i*4+1] = pixels[i];
                    rgba_pixels[i*4+2] = pixels[i];
                    rgba_pixels[i*4+3] = pixels[i];
                }
                rgba_pixels
            }
            FontRenderMode::Alpha => {
                let length = pixels.len() / 3;
                let mut rgba_pixels: Vec<u8> = vec![0; length * 4];
                for i in 0..length {
                    // Only take the G channel, as its closest to D2D
                    let alpha = pixels[i*3 + 1] as u8;
                    rgba_pixels[i*4+0] = alpha;
                    rgba_pixels[i*4+1] = alpha;
                    rgba_pixels[i*4+2] = alpha;
                    rgba_pixels[i*4+3] = alpha;
                }
                rgba_pixels
            }
            FontRenderMode::Subpixel => {
                let length = pixels.len() / 3;
                let mut rgba_pixels: Vec<u8> = vec![0; length * 4];
                for i in 0..length {
                    rgba_pixels[i*4+0] = pixels[i*3+0];
                    rgba_pixels[i*4+1] = pixels[i*3+1];
                    rgba_pixels[i*4+2] = pixels[i*3+2];
                    rgba_pixels[i*4+3] = 0xff;
                }
                rgba_pixels
            }
        }
    }

    pub fn rasterize_glyph(&mut self,
                           key: &GlyphKey,
                           render_mode: FontRenderMode,
                           glyph_options: Option<GlyphOptions>)
                           -> Option<RasterizedGlyph> {
        let analysis = self.create_glyph_analysis(key,
                                                  render_mode,
                                                  glyph_options);
        let texture_type = dwrite_texture_type(render_mode);

        let bounds = analysis.get_alpha_texture_bounds(texture_type);
        let width = (bounds.right - bounds.left) as usize;
        let height = (bounds.bottom - bounds.top) as usize;

        // We should not get here since glyph_dimensions would return
        // None for empty glyphs.
        assert!(width > 0 && height > 0);

        let mut pixels = analysis.create_alpha_texture(texture_type, bounds);

        if render_mode != FontRenderMode::Mono {
            let lut_correction = match glyph_options {
                Some(option) => {
                    if option.force_gdi_rendering {
                        &self.gdi_gamma_lut
                    } else {
                        &self.gamma_lut
                    }
                },
                None => &self.gamma_lut
            };

            lut_correction.preblend_rgb(&mut pixels, width, height,
                                        ColorLut::new(key.color.r,
                                                      key.color.g,
                                                      key.color.b,
                                                      key.color.a));
        }

        let rgba_pixels = self.convert_to_rgba(&mut pixels, render_mode);

        Some(RasterizedGlyph {
            width: width as u32,
            height: height as u32,
            bytes: rgba_pixels,
        })
    }
}
<commit_msg>Fix glyph size assertion in windows/font.rs.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use std::collections::HashMap;
use webrender_traits::{FontKey, FontRenderMode, GlyphDimensions};
use webrender_traits::{GlyphKey, GlyphOptions};
use gamma_lut::{GammaLut, Color as ColorLut};

use dwrote;

lazy_static! {
    static ref DEFAULT_FONT_DESCRIPTOR: dwrote::FontDescriptor = dwrote::FontDescriptor {
        family_name: "Arial".to_owned(),
        weight: dwrote::FontWeight::Regular,
        stretch: dwrote::FontStretch::Normal,
        style: dwrote::FontStyle::Normal,
    };
}

pub struct FontContext {
    fonts: HashMap<FontKey, dwrote::FontFace>,
    gamma_lut: GammaLut,
    gdi_gamma_lut: GammaLut,
}

// DirectWrite is safe to use on multiple threads and non-shareable resources are
// all hidden inside their font context.
unsafe impl Send for FontContext {}

pub struct RasterizedGlyph {
    pub width: u32,
    pub height: u32,
    pub bytes: Vec<u8>,
}

fn dwrite_texture_type(render_mode: FontRenderMode) ->
                       dwrote::DWRITE_TEXTURE_TYPE {
    match render_mode {
        FontRenderMode::Mono => dwrote::DWRITE_TEXTURE_ALIASED_1x1 ,
        FontRenderMode::Alpha |
        FontRenderMode::Subpixel => dwrote::DWRITE_TEXTURE_CLEARTYPE_3x1,
    }
}

fn dwrite_measure_mode(render_mode: FontRenderMode, options: Option<GlyphOptions>) ->
                       dwrote::DWRITE_MEASURING_MODE {
    if let Some(GlyphOptions{ force_gdi_rendering: true, .. }) = options {
        return dwrote::DWRITE_MEASURING_MODE_GDI_CLASSIC;
    }

    match render_mode {
        FontRenderMode::Mono => dwrote::DWRITE_MEASURING_MODE_GDI_NATURAL,
        FontRenderMode::Alpha |
        FontRenderMode::Subpixel => dwrote::DWRITE_MEASURING_MODE_NATURAL,
    }
}

fn dwrite_render_mode(font_face: &dwrote::FontFace,
                      render_mode: FontRenderMode,
                      em_size: f32,
                      measure_mode: dwrote::DWRITE_MEASURING_MODE,
                      options: Option<GlyphOptions>) ->
                      dwrote::DWRITE_RENDERING_MODE {
    if let Some(GlyphOptions{ force_gdi_rendering: true, .. }) = options {
        return dwrote::DWRITE_RENDERING_MODE_GDI_CLASSIC;
    }

    let dwrite_render_mode = match render_mode {
        FontRenderMode::Mono => dwrote::DWRITE_RENDERING_MODE_ALIASED,
        FontRenderMode::Alpha |
        FontRenderMode::Subpixel => {
            font_face.get_recommended_rendering_mode_default_params(em_size,
                                                                    1.0,
                                                                    measure_mode)
        },
    };

    if dwrite_render_mode  == dwrote::DWRITE_RENDERING_MODE_OUTLINE {
        // Outline mode is not supported
        return dwrote::DWRITE_RENDERING_MODE_CLEARTYPE_NATURAL_SYMMETRIC;
    }

    dwrite_render_mode
}

fn get_glyph_dimensions_with_analysis(analysis: dwrote::GlyphRunAnalysis,
                                      texture_type: dwrote::DWRITE_TEXTURE_TYPE)
                                      -> Option<GlyphDimensions> {
    let bounds = analysis.get_alpha_texture_bounds(texture_type);

    let width = (bounds.right - bounds.left) as u32;
    let height = (bounds.bottom - bounds.top) as u32;

    // Alpha texture bounds can sometimes return an empty rect
    // Such as for spaces
    if width == 0 || height == 0 {
        return None
    }

    Some(GlyphDimensions {
        left: bounds.left,
        top: -bounds.top,
        width: width,
        height: height,
    })
}

impl FontContext {
    pub fn new() -> FontContext {
        // These are the default values we use in Gecko.
        // We use a gamma value of 2.3 for gdi fonts
        // TODO: Fetch this data from Gecko itself.
        let contrast = 1.0;
        let gamma = 1.8;
        let gdi_gamma = 2.3;
        FontContext {
            fonts: HashMap::new(),
            gamma_lut: GammaLut::new(contrast, gamma, gamma),
            gdi_gamma_lut: GammaLut::new(contrast, gdi_gamma, gdi_gamma),
        }
    }

    pub fn has_font(&self, font_key: &FontKey) -> bool {
        self.fonts.contains_key(font_key)
    }

    pub fn add_raw_font(&mut self, font_key: &FontKey, data: &[u8], index: u32) {
        if self.fonts.contains_key(font_key) {
            return
        }

        if let Some(font_file) = dwrote::FontFile::new_from_data(data) {
            let face = font_file.create_face(index, dwrote::DWRITE_FONT_SIMULATIONS_NONE);
            self.fonts.insert((*font_key).clone(), face);
        } else {
            // XXX add_raw_font needs to have a way to return an error
            debug!("DWrite WR failed to load font from data, using Arial instead");
            self.add_native_font(font_key, DEFAULT_FONT_DESCRIPTOR.clone());
        }
    }

    pub fn add_native_font(&mut self, font_key: &FontKey, font_handle: dwrote::FontDescriptor) {
        if self.fonts.contains_key(font_key) {
            return
        }

        let system_fc = dwrote::FontCollection::system();
        let font = system_fc.get_font_from_descriptor(&font_handle).unwrap();
        let face = font.create_font_face();
        self.fonts.insert((*font_key).clone(), face);
    }

    pub fn delete_font(&mut self, font_key: &FontKey) {
        self.fonts.remove(font_key);
    }

    // Assumes RGB format from dwrite, which is 3 bytes per pixel as dwrite
    // doesn't output an alpha value via GlyphRunAnalysis::CreateAlphaTexture
    #[allow(dead_code)]
    fn print_glyph_data(&self, data: &[u8], width: usize, height: usize) {
        // Rust doesn't have step_by support on stable :(
        for i in 0..height {
            let current_height = i * width * 3;

            for pixel in data[current_height .. current_height + (width * 3)].chunks(3) {
                let r = pixel[0];
                let g = pixel[1];
                let b = pixel[2];
                print!("({}, {}, {}) ", r, g, b, );
            }
            println!("");
        }
    }

    fn create_glyph_analysis(&self, key: &GlyphKey,
                            render_mode: FontRenderMode,
                            options: Option<GlyphOptions>) ->
                            dwrote::GlyphRunAnalysis {
        let face = self.fonts.get(&key.font_key).unwrap();
        let glyph = key.index as u16;
        let advance = 0.0f32;
        let offset = dwrote::GlyphOffset { advanceOffset: 0.0, ascenderOffset: 0.0 };

        let glyph_run = dwrote::DWRITE_GLYPH_RUN {
            fontFace: unsafe { face.as_ptr() },
            fontEmSize: key.size.to_f32_px(), // size in DIPs (1/96", same as CSS pixels)
            glyphCount: 1,
            glyphIndices: &glyph,
            glyphAdvances: &advance,
            glyphOffsets: &offset,
            isSideways: 0,
            bidiLevel: 0,
        };

        let dwrite_measure_mode = dwrite_measure_mode(render_mode, options);
        let dwrite_render_mode = dwrite_render_mode(face,
                                                    render_mode,
                                                    key.size.to_f32_px(),
                                                    dwrite_measure_mode,
                                                    options);

        let (x_offset, y_offset) = key.subpixel_point.to_f64();
        let transform = Some(
                        dwrote::DWRITE_MATRIX { m11: 1.0, m12: 0.0, m21: 0.0, m22: 1.0,
                                                dx: x_offset as f32, dy: y_offset as f32 }
                        );

        dwrote::GlyphRunAnalysis::create(&glyph_run, 1.0, transform,
                                         dwrite_render_mode,
                                         dwrite_measure_mode,
                                         0.0, 0.0)
    }

    // TODO: Pipe GlyphOptions into glyph_dimensions too
    pub fn get_glyph_dimensions(&self,
                                key: &GlyphKey)
                                -> Option<GlyphDimensions> {
        // Probably have to default to something else here.
        let render_mode = FontRenderMode::Subpixel;
        let analysis = self.create_glyph_analysis(key, render_mode, None);

        let texture_type = dwrite_texture_type(render_mode);
        get_glyph_dimensions_with_analysis(analysis, texture_type)
    }

    // DWRITE gives us values in RGB. WR doesn't really touch it after. Note, CG returns in BGR
    // TODO: Decide whether all fonts should return RGB or BGR
    fn convert_to_rgba(&self, pixels: &[u8], render_mode: FontRenderMode) -> Vec<u8> {
        match render_mode {
            FontRenderMode::Mono => {
                let mut rgba_pixels: Vec<u8> = vec![0; pixels.len() * 4];
                for i in 0..pixels.len() {
                    rgba_pixels[i*4+0] = pixels[i];
                    rgba_pixels[i*4+1] = pixels[i];
                    rgba_pixels[i*4+2] = pixels[i];
                    rgba_pixels[i*4+3] = pixels[i];
                }
                rgba_pixels
            }
            FontRenderMode::Alpha => {
                let length = pixels.len() / 3;
                let mut rgba_pixels: Vec<u8> = vec![0; length * 4];
                for i in 0..length {
                    // Only take the G channel, as its closest to D2D
                    let alpha = pixels[i*3 + 1] as u8;
                    rgba_pixels[i*4+0] = alpha;
                    rgba_pixels[i*4+1] = alpha;
                    rgba_pixels[i*4+2] = alpha;
                    rgba_pixels[i*4+3] = alpha;
                }
                rgba_pixels
            }
            FontRenderMode::Subpixel => {
                let length = pixels.len() / 3;
                let mut rgba_pixels: Vec<u8> = vec![0; length * 4];
                for i in 0..length {
                    rgba_pixels[i*4+0] = pixels[i*3+0];
                    rgba_pixels[i*4+1] = pixels[i*3+1];
                    rgba_pixels[i*4+2] = pixels[i*3+2];
                    rgba_pixels[i*4+3] = 0xff;
                }
                rgba_pixels
            }
        }
    }

    pub fn rasterize_glyph(&mut self,
                           key: &GlyphKey,
                           render_mode: FontRenderMode,
                           glyph_options: Option<GlyphOptions>)
                           -> Option<RasterizedGlyph> {
        let analysis = self.create_glyph_analysis(key,
                                                  render_mode,
                                                  glyph_options);
        let texture_type = dwrite_texture_type(render_mode);

        let bounds = analysis.get_alpha_texture_bounds(texture_type);
        let width = (bounds.right - bounds.left) as usize;
        let height = (bounds.bottom - bounds.top) as usize;

        // Alpha texture bounds can sometimes return an empty rect
        // Such as for spaces
        if width == 0 || height == 0 {
            return None;
        }

        let mut pixels = analysis.create_alpha_texture(texture_type, bounds);

        if render_mode != FontRenderMode::Mono {
            let lut_correction = match glyph_options {
                Some(option) => {
                    if option.force_gdi_rendering {
                        &self.gdi_gamma_lut
                    } else {
                        &self.gamma_lut
                    }
                },
                None => &self.gamma_lut
            };

            lut_correction.preblend_rgb(&mut pixels, width, height,
                                        ColorLut::new(key.color.r,
                                                      key.color.g,
                                                      key.color.b,
                                                      key.color.a));
        }

        let rgba_pixels = self.convert_to_rgba(&mut pixels, render_mode);

        Some(RasterizedGlyph {
            width: width as u32,
            height: height as u32,
            bytes: rgba_pixels,
        })
    }
}
<|endoftext|>


<commit_before>use intern::intern;
use parser;
use normalize::macro_expand::expand_macros;
use normalize::tyinfer::infer_types;
use grammar::parse_tree::NonterminalString;
use grammar::repr::TypeRepr;

fn type_repr(s: &str) -> TypeRepr {
    let type_ref = parser::parse_type_ref(s).unwrap();
    return type_ref.type_repr();
}

fn compare(g1: &str, expected: Vec<(&'static str, &'static str)>) {
    let grammar = parser::parse_grammar(g1).unwrap();
    let grammar = expand_macros(grammar).unwrap();
    let types = infer_types(&grammar).unwrap();

    println!("types table: {:?}", types);

    for (nt_id, nt_type) in expected {
        let id = NonterminalString(intern(nt_id));
        let ty = type_repr(nt_type);
        println!("expected type of {:?} is {:?}", id, ty);
        assert_eq!(types.nonterminal_type(id), &ty);
    }
}

#[test]
fn test_pairs_and_tokens() {
    compare("
grammar;
    extern token { enum Tok { } }
    X = Y Z;
    Y: Foo = \"Hi\";
    Z = \"Ho\";
", vec![
    ("X", "(Foo, Tok)"),
    ("Y", "Foo"),
    ("Z", "Tok")
        ])
}

#[test]
fn test_cycle_direct() {
    let grammar = parser::parse_grammar("
grammar;
    extern token { enum Tok { } }
    X = {
        X Y;
        <Y> => vec![<>];
    };
    Y = \"Hi\";
").unwrap();

    let actual = expand_macros(grammar).unwrap();
    assert!(infer_types(&actual).is_err());
}

#[test]
fn test_cycle_indirect() {
    let grammar = parser::parse_grammar("
grammar;
    extern token { enum Tok { } }
    A = B;
    B = C;
    C = D;
    D = A;
").unwrap();

    let actual = expand_macros(grammar).unwrap();
    assert!(infer_types(&actual).is_err());
}

#[test]
fn test_macro_expansion() {
    compare("
grammar;
    extern token { enum Tok { } }
    Two<X>: (X, X) = X X;
    Ids = Two<\"Id\">;
", vec![
    ("Ids", "(Tok, Tok)"),
    (r#"Two<"Id">"#, "(Tok, Tok)"),
        ])
}

#[test]
fn test_macro_expansion_infer() {
    compare("
grammar;
    extern token { enum Tok { } }
    Two<X> = X X;
    Ids = Two<\"Id\">;
", vec![
    ("Ids", "(Tok, Tok)"),
    (r#"Two<"Id">"#, "(Tok, Tok)"),
        ])
}

#[test]
fn test_type_question() {
    compare("
grammar;
    extern token { enum Tok { } }
    X = Y?;
    Y = \"Hi\";
",vec![
    ("X", "::std::option::Option<Tok>"),
    ("Y", "Tok")
        ])
}

#[test]
fn test_star_plus_question() {
    compare("
grammar;
    extern token { enum Tok { } }
    A = Z*;
    X = \"Hi\"*;
    Y = \"Hi\"+;
    Z = \"Hi\"?;
", vec![
    ("A", "::std::vec::Vec<::std::option::Option<Tok>>"),
    ("X", "::std::vec::Vec<Tok>"),
    ("Y", "::std::vec::Vec<Tok>"),
    ("Z", "::std::option::Option<Tok>")
        ])
}

#[test]
fn test_lookahead() {
    compare(r#"
grammar;
    extern token { type Location = usize; enum Tok { } }
    A = @L;
"#, vec![
    ("A", "::std::option::Option<usize>"),
        ])
}

#[test]
fn test_spanned_macro() {
    compare(r#"
        grammar;
        extern token { type Location = usize; enum Tok { } }
        A = Spanned<"Foo">;
        Spanned<T>: (Option<usize>, Option<usize>) = {
            <@L> T <@R> => (<>);
        };
"#, vec![
    ("A", "(Option<usize>, Option<usize>)"),
        ])
}

#[test]
fn test_action() {
    compare(r#"
grammar;
    extern token { enum Tok { } }

    X = {
        Y;
        <l:X> "+" <r:Y> => l + r;
    };

    Y: i32 = "foo" => 22;
"#,vec![
    ("X", "i32"),
    ("Y", "i32"),
        ])
}

#[test]
fn test_inconsistent_action() {
    let grammar = parser::parse_grammar("
grammar;
    extern token { enum Tok { } }

    X = {
        Y;
        Z;
        <l:X> \"+\" <r:Y> => l + r;
    };

    Y: i32 = \"foo\" => 22;

    Z: u32 = \"bar\" => 22;
").unwrap();

    let actual = expand_macros(grammar).unwrap();
    assert!(infer_types(&actual).is_err());
}
<commit_msg>fix tyinfer test<commit_after>use intern::intern;
use parser;
use normalize::macro_expand::expand_macros;
use normalize::tyinfer::infer_types;
use grammar::parse_tree::NonterminalString;
use grammar::repr::TypeRepr;

fn type_repr(s: &str) -> TypeRepr {
    let type_ref = parser::parse_type_ref(s).unwrap();
    return type_ref.type_repr();
}

fn compare(g1: &str, expected: Vec<(&'static str, &'static str)>) {
    let grammar = parser::parse_grammar(g1).unwrap();
    let grammar = expand_macros(grammar).unwrap();
    let types = infer_types(&grammar).unwrap();

    println!("types table: {:?}", types);

    for (nt_id, nt_type) in expected {
        let id = NonterminalString(intern(nt_id));
        let ty = type_repr(nt_type);
        println!("expected type of {:?} is {:?}", id, ty);
        assert_eq!(types.nonterminal_type(id), &ty);
    }
}

#[test]
fn test_pairs_and_tokens() {
    compare("
grammar;
    extern token { enum Tok { } }
    X = Y Z;
    Y: Foo = \"Hi\";
    Z = \"Ho\";
", vec![
    ("X", "(Foo, Tok)"),
    ("Y", "Foo"),
    ("Z", "Tok")
        ])
}

#[test]
fn test_cycle_direct() {
    let grammar = parser::parse_grammar("
grammar;
    extern token { enum Tok { } }
    X = {
        X Y;
        <Y> => vec![<>];
    };
    Y = \"Hi\";
").unwrap();

    let actual = expand_macros(grammar).unwrap();
    assert!(infer_types(&actual).is_err());
}

#[test]
fn test_cycle_indirect() {
    let grammar = parser::parse_grammar("
grammar;
    extern token { enum Tok { } }
    A = B;
    B = C;
    C = D;
    D = A;
").unwrap();

    let actual = expand_macros(grammar).unwrap();
    assert!(infer_types(&actual).is_err());
}

#[test]
fn test_macro_expansion() {
    compare("
grammar;
    extern token { enum Tok { } }
    Two<X>: (X, X) = X X;
    Ids = Two<\"Id\">;
", vec![
    ("Ids", "(Tok, Tok)"),
    (r#"Two<"Id">"#, "(Tok, Tok)"),
        ])
}

#[test]
fn test_macro_expansion_infer() {
    compare("
grammar;
    extern token { enum Tok { } }
    Two<X> = X X;
    Ids = Two<\"Id\">;
", vec![
    ("Ids", "(Tok, Tok)"),
    (r#"Two<"Id">"#, "(Tok, Tok)"),
        ])
}

#[test]
fn test_type_question() {
    compare("
grammar;
    extern token { enum Tok { } }
    X = Y?;
    Y = \"Hi\";
",vec![
    ("X", "::std::option::Option<Tok>"),
    ("Y", "Tok")
        ])
}

#[test]
fn test_star_plus_question() {
    compare("
grammar;
    extern token { enum Tok { } }
    A = Z*;
    X = \"Hi\"*;
    Y = \"Hi\"+;
    Z = \"Hi\"?;
", vec![
    ("A", "::std::vec::Vec<::std::option::Option<Tok>>"),
    ("X", "::std::vec::Vec<Tok>"),
    ("Y", "::std::vec::Vec<Tok>"),
    ("Z", "::std::option::Option<Tok>")
        ])
}

#[test]
fn test_lookahead() {
    compare(r#"
grammar;
    extern token { type Location = usize; enum Tok { } }
    A = @L;
"#, vec![
    ("A", "usize"),
        ])
}

#[test]
fn test_spanned_macro() {
    compare(r#"
        grammar;
        extern token { type Location = usize; enum Tok { } }
        A = Spanned<"Foo">;
        Spanned<T> = {
            @L T @R;
        };
"#, vec![
    ("A", "(usize, Tok, usize)"),
        ])
}

#[test]
fn test_action() {
    compare(r#"
grammar;
    extern token { enum Tok { } }

    X = {
        Y;
        <l:X> "+" <r:Y> => l + r;
    };

    Y: i32 = "foo" => 22;
"#,vec![
    ("X", "i32"),
    ("Y", "i32"),
        ])
}

#[test]
fn test_inconsistent_action() {
    let grammar = parser::parse_grammar("
grammar;
    extern token { enum Tok { } }

    X = {
        Y;
        Z;
        <l:X> \"+\" <r:Y> => l + r;
    };

    Y: i32 = \"foo\" => 22;

    Z: u32 = \"bar\" => 22;
").unwrap();

    let actual = expand_macros(grammar).unwrap();
    assert!(infer_types(&actual).is_err());
}
<|endoftext|>


<commit_before>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

extern crate bindgen;

use std::{
    path::{Path, PathBuf},
    str::FromStr,
};

use anyhow::{Context, Result};
use tvm_build::{BuildConfig, CMakeSetting};

/// The necessary information for detecting a TVM installation.
struct TVMInstall {
    source_path: PathBuf,
    build_path: PathBuf,
}

/// Find the TVM install using the provided path.
fn find_using_tvm_path<P: AsRef<Path>>(tvm_path: P) -> Result<TVMInstall> {
    Ok(TVMInstall {
        source_path: tvm_path.as_ref().into(),
        build_path: tvm_path.as_ref().into(),
    })
}

#[allow(unused)]
fn if_unset<K: AsRef<std::ffi::OsStr>, V: AsRef<std::ffi::OsStr>>(k: K, v: V) -> Result<()> {
    match std::env::var(k.as_ref()) {
        Ok(other) if other != "" => {
            println!(
                "cargo:warning=Using existing environment variable setting {:?}={:?}",
                k.as_ref(),
                v.as_ref()
            );
        }
        _ => std::env::set_var(k, v),
    }

    Ok(())
}

/// Find a TVM installation using TVM build by either first installing or detecting.
fn find_using_tvm_build() -> Result<TVMInstall> {
    let mut build_config = BuildConfig::default();
    build_config.repository = Some("https://github.com/apache/tvm".to_string());
    build_config.branch = Some(option_env!("TVM_BRANCH").unwrap_or("main").into());

    if cfg!(feature = "use-cuda") {
        build_config.settings.use_cuda = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-opencl") {
        build_config.settings.use_opencl = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-vulkan") {
        build_config.settings.use_vulkan = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-rocm") {
        build_config.settings.use_rocm = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-metal") {
        build_config.settings.use_metal = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-hexagon-device") {
        build_config.settings.use_hexagon_device = Some(true);
    }
    if cfg!(feature = "use-rpc") {
        build_config.settings.use_rpc = Some(true);
    }
    if cfg!(feature = "use-threads") {
        build_config.settings.use_threads = Some(true);
    }
    if cfg!(feature = "use-llvm") {
        build_config.settings.use_llvm = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-stackvm-runtime") {
        build_config.settings.use_stackvm_runtime = Some(true);
    }
    if cfg!(feature = "use-graph-runtime") {
        build_config.settings.use_graph_runtime = Some(true);
    }
    if cfg!(feature = "use-graph-runtime-debug") {
        build_config.settings.use_graph_runtime_debug = Some(true);
    }
    if cfg!(feature = "use-openmp") {
        build_config.settings.use_openmp = Some(true);
    }
    if cfg!(feature = "use-relay-debug") {
        build_config.settings.use_relay_debug = Some(true);
    }
    if cfg!(feature = "use-rtti") {
        build_config.settings.use_rtti = Some(true);
    }
    if cfg!(feature = "use-mscv-mt") {
        build_config.settings.use_mscv_mt = Some(true);
    }
    if cfg!(feature = "use-micro") {
        build_config.settings.use_micro = Some(true);
    }
    if cfg!(feature = "use-install-dev") {
        build_config.settings.use_install_dev = Some(true);
    }
    if cfg!(feature = "hide_private-symbols") {
        build_config.settings.hide_private_symbols = Some(true);
    }
    if cfg!(feature = "use-fallback-stl-map") {
        build_config.settings.use_fallback_stl_map = Some(true);
    }
    if cfg!(feature = "use-ethosn") {
        build_config.settings.use_ethosn = Some(true);
    }
    if cfg!(feature = "use-index_default-i64") {
        build_config.settings.use_index_default_i64 = Some(true);
    }
    if cfg!(feature = "use-tf-tvmdsoop") {
        build_config.settings.use_tf_tvmdsoop = Some(true);
    }
    if cfg!(feature = "use-byodt-posit") {
        build_config.settings.use_byodt_posit = Some(true);
    }
    if cfg!(feature = "use-mkl") {
        build_config.settings.use_mkl = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-mkldnn") {
        build_config.settings.use_mkldnn = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-dnnl-codegen") {
        build_config.settings.use_dnnl_codegen = Some(true);
    }
    if cfg!(feature = "use-cudnn") {
        build_config.settings.use_cudnn = Some(true);
    }
    if cfg!(feature = "use-cublas") {
        build_config.settings.use_cublas = Some(true);
    }
    if cfg!(feature = "use-thrust") {
        build_config.settings.use_thrust = Some(true);
    }
    if cfg!(feature = "use-miopen") {
        build_config.settings.use_miopen = Some(true);
    }
    if cfg!(feature = "use-rocblas") {
        build_config.settings.use_rocblas = Some(true);
    }
    if cfg!(feature = "use-sort") {
        build_config.settings.use_sort = Some(true);
    }
    if cfg!(feature = "use-nnpack") {
        build_config.settings.use_nnpack = Some(true);
    }
    if cfg!(feature = "use-random") {
        build_config.settings.use_random = Some(true);
    }
    if cfg!(feature = "use-micro-standalone-runtime") {
        build_config.settings.use_micro_standalone_runtime = Some(true);
    }
    if cfg!(feature = "use-cpp-rpc") {
        build_config.settings.use_cpp_rpc = Some(true);
    }
    if cfg!(feature = "use-tflite") {
        build_config.settings.use_tflite = Some(true);
    }
    if cfg!(feature = "use-coreml") {
        build_config.settings.use_coreml = Some(true);
    }
    if cfg!(feature = "use-target-onnx") {
        build_config.settings.use_target_onnx = Some(true);
    }
    if cfg!(feature = "use-arm-compute-lib") {
        build_config.settings.use_arm_compute_lib = Some(true);
    }
    if cfg!(feature = "use-arm-compute-lib-graph-runtime") {
        build_config.settings.use_arm_compute_lib_graph_runtime = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-tensorrt-codegen") {
        build_config.settings.use_tensorrt_codegen = Some(true);
    }
    if cfg!(feature = "use-tensorrt-runtime") {
        build_config.settings.use_tensorrt_runtime = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-vitis-ai") {
        build_config.settings.use_vitis_ai = Some(true);
    }
    if cfg!(any(
        feature = "static-linking",
        feature = "build-static-runtime"
    )) {
        build_config.settings.build_static_runtime = Some(true);
    }

    let build_result = tvm_build::build(build_config)?;
    let source_path = build_result.revision.source_path();
    let build_path = build_result.revision.build_path();
    Ok(TVMInstall {
        source_path,
        build_path,
    })
}

fn main() -> Result<()> {
    let TVMInstall {
        source_path,
        build_path,
    } = match option_env!("TVM_HOME") {
        Some(tvm_path) if tvm_path != "" => find_using_tvm_path(tvm_path),
        _ => find_using_tvm_build(),
    }?;

    // If the TVM_HOME environment variable changed, the LLVM_CONFIG_PATH environment variable
    // changed, the build directory or headers have changed we need to rebuild the Rust bindings.
    println!("cargo:rerun-if-env-changed=TVM_HOME");
    println!("cargo:rerun-if-env-changed=LLVM_CONFIG_PATH");
    println!("cargo:rerun-if-changed={}", build_path.display());
    println!("cargo:rerun-if-changed={}/include", source_path.display());

    let library_name = if cfg!(feature = "runtime-only") {
        "tvm_runtime"
    } else {
        "tvm"
    };

    match &std::env::var("CARGO_CFG_TARGET_ARCH")
        .expect("CARGO_CFG_TARGET_ARCH must be set by CARGO")[..]
    {
        "wasm32" => {}
        _ => {
            if cfg!(feature = "static-linking") {
                println!("cargo:rustc-link-lib=static={}", library_name);
                // TODO(@jroesch): move this to tvm-build as library_path?
                println!(
                    "cargo:rustc-link-search=native={}/build",
                    build_path.display()
                );
            }

            if cfg!(feature = "dynamic-linking") {
                println!("cargo:rustc-link-lib=dylib={}", library_name);
                println!(
                    "cargo:rustc-link-search=native={}/build",
                    build_path.display()
                );
            }
        }
    };

    let runtime_api = source_path.join("include/tvm/runtime/c_runtime_api.h");
    let backend_api = source_path.join("include/tvm/runtime/c_backend_api.h");
    let source_path = source_path.display().to_string();
    let dlpack_include = format!("-I{}/3rdparty/dlpack/include/", source_path);
    let tvm_include = format!("-I{}/include/", source_path);

    let out_file = PathBuf::from(std::env::var("OUT_DIR")?).join("c_runtime_api.rs");

    // @see rust-bindgen#550 for `blacklist_type`
    bindgen::Builder::default()
        .header(runtime_api.display().to_string())
        .header(backend_api.display().to_string())
        .clang_arg(dlpack_include)
        .clang_arg(tvm_include)
        .blacklist_type("max_align_t")
        .layout_tests(false)
        .derive_partialeq(true)
        .derive_eq(true)
        .derive_default(true)
        .generate()
        .map_err(|()| {
            anyhow::anyhow!("bindgen failed to generate the Rust bindings for the C API")
        })?
        .write_to_file(out_file)
        .context("failed to write the generated Rust binding to disk")?;

    Ok(())
}
<commit_msg>Fix an infinite recompilation loop in tvm-sys. (#9450)<commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

extern crate bindgen;

use std::{
    path::{Path, PathBuf},
    str::FromStr,
};

use anyhow::{Context, Result};
use tvm_build::{BuildConfig, CMakeSetting};

/// The necessary information for detecting a TVM installation.
struct TVMInstall {
    source_path: PathBuf,
    build_path: PathBuf,
}

/// Find the TVM install using the provided path.
fn find_using_tvm_path<P: AsRef<Path>>(tvm_path: P) -> Result<TVMInstall> {
    Ok(TVMInstall {
        source_path: tvm_path.as_ref().into(),
        build_path: tvm_path.as_ref().into(),
    })
}

#[allow(unused)]
fn if_unset<K: AsRef<std::ffi::OsStr>, V: AsRef<std::ffi::OsStr>>(k: K, v: V) -> Result<()> {
    match std::env::var(k.as_ref()) {
        Ok(other) if other != "" => {
            println!(
                "cargo:warning=Using existing environment variable setting {:?}={:?}",
                k.as_ref(),
                v.as_ref()
            );
        }
        _ => std::env::set_var(k, v),
    }

    Ok(())
}

/// Find a TVM installation using TVM build by either first installing or detecting.
fn find_using_tvm_build() -> Result<TVMInstall> {
    let mut build_config = BuildConfig::default();
    build_config.repository = Some("https://github.com/apache/tvm".to_string());
    build_config.branch = Some(option_env!("TVM_BRANCH").unwrap_or("main").into());

    if cfg!(feature = "use-cuda") {
        build_config.settings.use_cuda = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-opencl") {
        build_config.settings.use_opencl = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-vulkan") {
        build_config.settings.use_vulkan = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-rocm") {
        build_config.settings.use_rocm = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-metal") {
        build_config.settings.use_metal = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-hexagon-device") {
        build_config.settings.use_hexagon_device = Some(true);
    }
    if cfg!(feature = "use-rpc") {
        build_config.settings.use_rpc = Some(true);
    }
    if cfg!(feature = "use-threads") {
        build_config.settings.use_threads = Some(true);
    }
    if cfg!(feature = "use-llvm") {
        build_config.settings.use_llvm = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-stackvm-runtime") {
        build_config.settings.use_stackvm_runtime = Some(true);
    }
    if cfg!(feature = "use-graph-runtime") {
        build_config.settings.use_graph_runtime = Some(true);
    }
    if cfg!(feature = "use-graph-runtime-debug") {
        build_config.settings.use_graph_runtime_debug = Some(true);
    }
    if cfg!(feature = "use-openmp") {
        build_config.settings.use_openmp = Some(true);
    }
    if cfg!(feature = "use-relay-debug") {
        build_config.settings.use_relay_debug = Some(true);
    }
    if cfg!(feature = "use-rtti") {
        build_config.settings.use_rtti = Some(true);
    }
    if cfg!(feature = "use-mscv-mt") {
        build_config.settings.use_mscv_mt = Some(true);
    }
    if cfg!(feature = "use-micro") {
        build_config.settings.use_micro = Some(true);
    }
    if cfg!(feature = "use-install-dev") {
        build_config.settings.use_install_dev = Some(true);
    }
    if cfg!(feature = "hide_private-symbols") {
        build_config.settings.hide_private_symbols = Some(true);
    }
    if cfg!(feature = "use-fallback-stl-map") {
        build_config.settings.use_fallback_stl_map = Some(true);
    }
    if cfg!(feature = "use-ethosn") {
        build_config.settings.use_ethosn = Some(true);
    }
    if cfg!(feature = "use-index_default-i64") {
        build_config.settings.use_index_default_i64 = Some(true);
    }
    if cfg!(feature = "use-tf-tvmdsoop") {
        build_config.settings.use_tf_tvmdsoop = Some(true);
    }
    if cfg!(feature = "use-byodt-posit") {
        build_config.settings.use_byodt_posit = Some(true);
    }
    if cfg!(feature = "use-mkl") {
        build_config.settings.use_mkl = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-mkldnn") {
        build_config.settings.use_mkldnn = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-dnnl-codegen") {
        build_config.settings.use_dnnl_codegen = Some(true);
    }
    if cfg!(feature = "use-cudnn") {
        build_config.settings.use_cudnn = Some(true);
    }
    if cfg!(feature = "use-cublas") {
        build_config.settings.use_cublas = Some(true);
    }
    if cfg!(feature = "use-thrust") {
        build_config.settings.use_thrust = Some(true);
    }
    if cfg!(feature = "use-miopen") {
        build_config.settings.use_miopen = Some(true);
    }
    if cfg!(feature = "use-rocblas") {
        build_config.settings.use_rocblas = Some(true);
    }
    if cfg!(feature = "use-sort") {
        build_config.settings.use_sort = Some(true);
    }
    if cfg!(feature = "use-nnpack") {
        build_config.settings.use_nnpack = Some(true);
    }
    if cfg!(feature = "use-random") {
        build_config.settings.use_random = Some(true);
    }
    if cfg!(feature = "use-micro-standalone-runtime") {
        build_config.settings.use_micro_standalone_runtime = Some(true);
    }
    if cfg!(feature = "use-cpp-rpc") {
        build_config.settings.use_cpp_rpc = Some(true);
    }
    if cfg!(feature = "use-tflite") {
        build_config.settings.use_tflite = Some(true);
    }
    if cfg!(feature = "use-coreml") {
        build_config.settings.use_coreml = Some(true);
    }
    if cfg!(feature = "use-target-onnx") {
        build_config.settings.use_target_onnx = Some(true);
    }
    if cfg!(feature = "use-arm-compute-lib") {
        build_config.settings.use_arm_compute_lib = Some(true);
    }
    if cfg!(feature = "use-arm-compute-lib-graph-runtime") {
        build_config.settings.use_arm_compute_lib_graph_runtime = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-tensorrt-codegen") {
        build_config.settings.use_tensorrt_codegen = Some(true);
    }
    if cfg!(feature = "use-tensorrt-runtime") {
        build_config.settings.use_tensorrt_runtime = CMakeSetting::from_str("on").ok();
    }
    if cfg!(feature = "use-vitis-ai") {
        build_config.settings.use_vitis_ai = Some(true);
    }
    if cfg!(any(
        feature = "static-linking",
        feature = "build-static-runtime"
    )) {
        build_config.settings.build_static_runtime = Some(true);
    }

    let build_result = tvm_build::build(build_config)?;
    let source_path = build_result.revision.source_path();
    let build_path = build_result.revision.build_path();
    Ok(TVMInstall {
        source_path,
        build_path,
    })
}

fn main() -> Result<()> {
    let TVMInstall {
        source_path,
        build_path,
    } = match option_env!("TVM_HOME") {
        Some(tvm_path) if tvm_path != "" => find_using_tvm_path(tvm_path),
        _ => find_using_tvm_build(),
    }?;

    // If the TVM_HOME environment variable changed, the LLVM_CONFIG_PATH environment variable
    // changed or the source headers have changed we need to rebuild the Rust bindings.
    println!("cargo:rerun-if-env-changed=TVM_HOME");
    println!("cargo:rerun-if-env-changed=LLVM_CONFIG_PATH");
    println!("cargo:rerun-if-changed={}/include", source_path.display());

    let library_name = if cfg!(feature = "runtime-only") {
        "tvm_runtime"
    } else {
        "tvm"
    };

    match &std::env::var("CARGO_CFG_TARGET_ARCH")
        .expect("CARGO_CFG_TARGET_ARCH must be set by CARGO")[..]
    {
        "wasm32" => {}
        _ => {
            if cfg!(feature = "static-linking") {
                println!("cargo:rustc-link-lib=static={}", library_name);
                // TODO(@jroesch): move this to tvm-build as library_path?
                println!(
                    "cargo:rustc-link-search=native={}/build",
                    build_path.display()
                );
            }

            if cfg!(feature = "dynamic-linking") {
                println!("cargo:rustc-link-lib=dylib={}", library_name);
                println!(
                    "cargo:rustc-link-search=native={}/build",
                    build_path.display()
                );
            }
        }
    };

    let runtime_api = source_path.join("include/tvm/runtime/c_runtime_api.h");
    let backend_api = source_path.join("include/tvm/runtime/c_backend_api.h");
    let source_path = source_path.display().to_string();
    let dlpack_include = format!("-I{}/3rdparty/dlpack/include/", source_path);
    let tvm_include = format!("-I{}/include/", source_path);

    let out_file = PathBuf::from(std::env::var("OUT_DIR")?).join("c_runtime_api.rs");

    // @see rust-bindgen#550 for `blacklist_type`
    bindgen::Builder::default()
        .header(runtime_api.display().to_string())
        .header(backend_api.display().to_string())
        .clang_arg(dlpack_include)
        .clang_arg(tvm_include)
        .blacklist_type("max_align_t")
        .layout_tests(false)
        .derive_partialeq(true)
        .derive_eq(true)
        .derive_default(true)
        .generate()
        .map_err(|()| {
            anyhow::anyhow!("bindgen failed to generate the Rust bindings for the C API")
        })?
        .write_to_file(out_file)
        .context("failed to write the generated Rust binding to disk")?;

    Ok(())
}
<|endoftext|>


<commit_before>use toml::Value;

use libimagerror::into::IntoError;

use vcs::git::error::GitHookErrorKind as GHEK;
use vcs::git::error::MapErrInto;
use vcs::git::result::Result;

use vcs::git::action::StoreAction;

pub fn commit_interactive(config: &Value) -> bool {
    warn!("Interactive committing not yet supported, using dummy commit message");
    false
}

pub fn commit_message(config: &Value, action: StoreAction) -> Result<String> {
    if commit_interactive(config) {
        unimplemented!()
    } else {
        Ok(String::from("Dummy commit"))
    }
}

pub fn abort_on_repo_init_err(cfg: Option<&Value>) -> bool {
    get_bool_cfg(cfg, "abort_on_repo_init_failure", true, true)
}

pub fn ensure_branch(cfg: Option<&Value>) -> Result<Option<String>> {
    match cfg {
        Some(cfg) => {
            match cfg.lookup("ensure_branch") {
                Some(&Value::String(ref s)) => Ok(Some(s.clone())),
                Some(_) => {
                    warn!("Configuration error, 'ensure_branch' must be a String.");
                    Err(GHEK::ConfigTypeError.into_error())
                        .map_err_into(GHEK::ConfigTypeError)
                },
                None => {
                    debug!("No key `ensure_branch'");
                    Ok(None)
                },
            }
        },
        None => Ok(None),
    }
}

pub fn do_checkout_ensure_branch(cfg: Option<&Value>) -> bool {
    get_bool_cfg(cfg, "try_checkout_ensure_branch", true, true)
}

fn get_bool_cfg(cfg: Option<&Value>, name: &str, on_fail: bool, on_unavail: bool) -> bool {
    cfg.map(|cfg| {
        match cfg.lookup(name) {
            Some(&Value::Boolean(b)) => b,
            Some(_) => {
                warn!("Configuration error, '{}' must be a Boolean (true|false).", name);
                warn!("Assuming '{}' now.", on_fail);
                on_fail
            },
            None => {
                debug!("No key '{}' - Assuming '{}'", name, on_unavail);
                on_unavail
            },
        }
    })
    .unwrap_or(on_unavail)
}

<commit_msg>Implement commit_message()<commit_after>use toml::Value;

use libimagerror::into::IntoError;

use vcs::git::error::GitHookErrorKind as GHEK;
use vcs::git::error::MapErrInto;
use vcs::git::result::Result;

use vcs::git::action::StoreAction;

pub fn commit_interactive(config: &Value) -> bool {
    match config.lookup("commit.interactive") {
        Some(&Value::Boolean(b)) => b,
        Some(_) => {
            warn!("Configuration error, 'store.hooks.stdhook_git_update.commit.interactive' must be a Boolean.");
            warn!("Defaulting to commit.interactive = false");
            false
        }
        None => {
            warn!("Unavailable configuration for");
            warn!("\t'store.hooks.stdhook_git_update.commit.interactive'");
            warn!("Defaulting to false");
            false
        }
    }
}

fn commit_with_editor(config: &Value) -> bool {
    match config.lookup("commit.interactive_editor") {
        Some(&Value::Boolean(b)) => b,
        Some(_) => {
            warn!("Configuration error, 'store.hooks.stdhook_git_update.commit.interactive_editor' must be a Boolean.");
            warn!("Defaulting to commit.interactive_editor = false");
            false
        }
        None => {
            warn!("Unavailable configuration for");
            warn!("\t'store.hooks.stdhook_git_update.commit.interactive_editor'");
            warn!("Defaulting to false");
            false
        }
    }
}

fn commit_default_msg<'a>(config: &'a Value) -> &'a str {
    match config.lookup("commit.message") {
        Some(&Value::String(ref b)) => b,
        Some(_) => {
            warn!("Configuration error, 'store.hooks.stdhook_git_update.commit.message' must be a String.");
            warn!("Defaulting to commit.message = 'Update'");
            "Update"
        }
        None => {
            warn!("Unavailable configuration for");
            warn!("\t'store.hooks.stdhook_git_update.commit.message'");
            warn!("Defaulting to commit.message = 'Update'");
            "Update"
        }
    }
}

fn commit_template() -> &'static str {
    "Commit template"
}

pub fn commit_message(repo: &Repository, config: &Value, action: StoreAction) -> Result<String> {
    use libimaginteraction::ask::ask_string;
    use libimagutil::edit::edit_in_tmpfile_with_command;
    use std::process::Command;

    if commit_interactive(config) {
        if commit_with_editor(config) {
            unimplemented!()
        } else {
            unimplemented!()
        }
    } else {
        Ok(String::from(commit_default_msg(config)))
    }
}

pub fn abort_on_repo_init_err(cfg: Option<&Value>) -> bool {
    get_bool_cfg(cfg, "abort_on_repo_init_failure", true, true)
}

pub fn ensure_branch(cfg: Option<&Value>) -> Result<Option<String>> {
    match cfg {
        Some(cfg) => {
            match cfg.lookup("ensure_branch") {
                Some(&Value::String(ref s)) => Ok(Some(s.clone())),
                Some(_) => {
                    warn!("Configuration error, 'ensure_branch' must be a String.");
                    Err(GHEK::ConfigTypeError.into_error())
                        .map_err_into(GHEK::ConfigTypeError)
                },
                None => {
                    debug!("No key `ensure_branch'");
                    Ok(None)
                },
            }
        },
        None => Ok(None),
    }
}

pub fn do_checkout_ensure_branch(cfg: Option<&Value>) -> bool {
    get_bool_cfg(cfg, "try_checkout_ensure_branch", true, true)
}

fn get_bool_cfg(cfg: Option<&Value>, name: &str, on_fail: bool, on_unavail: bool) -> bool {
    cfg.map(|cfg| {
        match cfg.lookup(name) {
            Some(&Value::Boolean(b)) => b,
            Some(_) => {
                warn!("Configuration error, '{}' must be a Boolean (true|false).", name);
                warn!("Assuming '{}' now.", on_fail);
                on_fail
            },
            None => {
                debug!("No key '{}' - Assuming '{}'", name, on_unavail);
                on_unavail
            },
        }
    })
    .unwrap_or(on_unavail)
}

<|endoftext|>


<commit_before>use std::path::PathBuf;
use std::path::Path;
use std::fmt::{Debug, Formatter, Error as FmtError};
use std::result::Result as RResult;

use toml::Value;

use libimagerror::into::IntoError;
use libimagerror::trace::trace_error;
use libimagstore::hook::Hook;
use libimagstore::hook::accessor::StoreIdAccessor;
use libimagstore::hook::accessor::{HookDataAccessor, HookDataAccessorProvider};
use libimagstore::hook::error::HookError as HE;
use libimagstore::hook::error::HookErrorKind as HEK;
use libimagstore::hook::position::HookPosition;
use libimagstore::hook::result::HookResult;
use libimagstore::storeid::StoreId;
use libimagutil::debug_result::*;

use vcs::git::error::GitHookError as GHE;
use vcs::git::error::GitHookErrorKind as GHEK;
use vcs::git::error::MapErrInto;
use vcs::git::error::MapIntoHookError;
use vcs::git::result::Result;
use vcs::git::runtime::Runtime as GRuntime;

pub struct UpdateHook {
    storepath: PathBuf,

    runtime: GRuntime,

    position: HookPosition,
}

impl UpdateHook {

    pub fn new(storepath: PathBuf, p: HookPosition) -> UpdateHook {
        UpdateHook {
            runtime: GRuntime::new(&storepath),
            storepath: storepath,
            position: p,
        }
    }

}

impl Debug for UpdateHook {
    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        write!(fmt, "UpdateHook(storepath={:?}, repository={}, pos={:?}, cfg={:?}",
               self.storepath,
               (if self.runtime.has_repository() { "Some(_)" } else { "None" }),
               self.position,
               self.runtime.has_config())
    }
}

impl Hook for UpdateHook {

    fn name(&self) -> &'static str {
        "stdhook_git_update"
    }

    fn set_config(&mut self, config: &Value) {
        if let Err(e) = self.runtime.set_config(config) {
            trace_error(&e);
        }
    }

}

impl HookDataAccessorProvider for UpdateHook {

    fn accessor(&self) -> HookDataAccessor {
        HookDataAccessor::StoreIdAccess(self)
    }
}

impl StoreIdAccessor for UpdateHook {

    /// The implementation of the UpdateHook
    ///
    /// # Scope
    ///
    /// This hook takes the git index and commits it either interactively or with a default message,
    /// if there is no configuration for an interactive commit.
    ///
    fn access(&self, id: &StoreId) -> HookResult<()> {
        use vcs::git::action::StoreAction;
        use vcs::git::config::commit_message;
        use vcs::git::error::MapIntoHookError;
        use vcs::git::util::fetch_index;
        use git2::{Reference as GitReference, Repository, Error as Git2Error};
        use git2::{ADD_DEFAULT, STATUS_WT_NEW, STATUS_WT_MODIFIED, IndexMatchedPath};

        debug!("[GIT UPDATE HOOK]: {:?}", id);

        let action    = StoreAction::Update;
        let cfg       = try!(self.runtime.config_value_or_err(&action));
        let repo      = try!(self.runtime.repository(&action));
        let mut index = try!(fetch_index(repo, &action));

        let path = try!(
            id.clone()
            .into_pathbuf()
            .map_err_into(GHEK::StoreIdHandlingError)
            .map_into_hook_error()
        );

        let tree_id = try!(
            index.write_tree()
                .map_err_into(GHEK::RepositoryIndexWritingError)
                .map_into_hook_error()
        );

        let signature = try!(
            repo.signature()
                .map_err_into(GHEK::MkSignature)
                .map_into_hook_error()
        );

        let head = try!(
            repo.head()
                .map_err_into(GHEK::HeadFetchError)
                .map_into_hook_error()
        );

        let file_status = try!(
            repo
                .status_file(&path)
                .map_err_into(GHEK::RepositoryFileStatusError)
                .map_into_hook_error()
        );

        let cb = &mut |path: &Path, _matched_spec: &[u8]| -> i32 {
            if file_status.contains(STATUS_WT_NEW) || file_status.contains(STATUS_WT_MODIFIED) {
                debug!("[GIT CREATE HOOK]: File is modified/new: {}", path.display());
                0
            } else {
                debug!("[GIT CREATE HOOK]: Ignoring file: {}", path.display());
                1
            }
        };

        try!(
            index.add_all(&[path], ADD_DEFAULT, Some(cb as &mut IndexMatchedPath))
                .map_err_into(GHEK::RepositoryPathAddingError)
                .map_into_hook_error()
        );

        let mut parents = Vec::new();
        {
            let commit = try!(
                repo.find_commit(head.target().unwrap())
                    .map_err_into(GHEK::RepositoryParentFetchingError)
                    .map_into_hook_error()
            );
            parents.push(commit);
        }

        // for converting from Vec<Commit> to Vec<&Commit>
        let parents = parents.iter().collect::<Vec<_>>();

        let tree = try!(
            repo.find_tree(tree_id)
                .map_err_into(GHEK::RepositoryParentFetchingError)
                .map_into_hook_error()
        );

        let message = try!(commit_message(cfg, StoreAction::Update));

        repo.commit(Some("HEAD"), &signature, &signature, &message, &tree, &parents)
            .map_err_into(GHEK::RepositoryCommittingError)
            .map_into_hook_error()
            .map(|_| ())

    }

}

<commit_msg>Add debug output in UpdateHook::access()<commit_after>use std::path::PathBuf;
use std::path::Path;
use std::fmt::{Debug, Formatter, Error as FmtError};
use std::result::Result as RResult;

use toml::Value;

use libimagerror::into::IntoError;
use libimagerror::trace::trace_error;
use libimagstore::hook::Hook;
use libimagstore::hook::accessor::StoreIdAccessor;
use libimagstore::hook::accessor::{HookDataAccessor, HookDataAccessorProvider};
use libimagstore::hook::error::HookError as HE;
use libimagstore::hook::error::HookErrorKind as HEK;
use libimagstore::hook::position::HookPosition;
use libimagstore::hook::result::HookResult;
use libimagstore::storeid::StoreId;
use libimagutil::debug_result::*;

use vcs::git::error::GitHookError as GHE;
use vcs::git::error::GitHookErrorKind as GHEK;
use vcs::git::error::MapErrInto;
use vcs::git::error::MapIntoHookError;
use vcs::git::result::Result;
use vcs::git::runtime::Runtime as GRuntime;

pub struct UpdateHook {
    storepath: PathBuf,

    runtime: GRuntime,

    position: HookPosition,
}

impl UpdateHook {

    pub fn new(storepath: PathBuf, p: HookPosition) -> UpdateHook {
        UpdateHook {
            runtime: GRuntime::new(&storepath),
            storepath: storepath,
            position: p,
        }
    }

}

impl Debug for UpdateHook {
    fn fmt(&self, fmt: &mut Formatter) -> RResult<(), FmtError> {
        write!(fmt, "UpdateHook(storepath={:?}, repository={}, pos={:?}, cfg={:?}",
               self.storepath,
               (if self.runtime.has_repository() { "Some(_)" } else { "None" }),
               self.position,
               self.runtime.has_config())
    }
}

impl Hook for UpdateHook {

    fn name(&self) -> &'static str {
        "stdhook_git_update"
    }

    fn set_config(&mut self, config: &Value) {
        if let Err(e) = self.runtime.set_config(config) {
            trace_error(&e);
        }
    }

}

impl HookDataAccessorProvider for UpdateHook {

    fn accessor(&self) -> HookDataAccessor {
        HookDataAccessor::StoreIdAccess(self)
    }
}

impl StoreIdAccessor for UpdateHook {

    /// The implementation of the UpdateHook
    ///
    /// # Scope
    ///
    /// This hook takes the git index and commits it either interactively or with a default message,
    /// if there is no configuration for an interactive commit.
    ///
    fn access(&self, id: &StoreId) -> HookResult<()> {
        use vcs::git::action::StoreAction;
        use vcs::git::config::commit_message;
        use vcs::git::error::MapIntoHookError;
        use vcs::git::util::fetch_index;
        use git2::{Reference as GitReference, Repository, Error as Git2Error};
        use git2::{ADD_DEFAULT, STATUS_WT_NEW, STATUS_WT_MODIFIED, IndexMatchedPath};

        debug!("[GIT UPDATE HOOK]: {:?}", id);

        let action    = StoreAction::Update;
        let cfg       = try!(self.runtime.config_value_or_err(&action));
        let repo      = try!(self.runtime.repository(&action));
        let mut index = try!(fetch_index(repo, &action));

        let path = try!(
            id.clone()
            .into_pathbuf()
            .map_err_into(GHEK::StoreIdHandlingError)
            .map_dbg_err_str("Failed to StoreId.into_pathbuf()")
            .map_into_hook_error()
        );

        let signature = try!(
            repo.signature()
                .map_err_into(GHEK::MkSignature)
                .map_dbg_err_str("Failed to fetch signature")
                .map_into_hook_error()
        );

        let head = try!(
            repo.head()
                .map_err_into(GHEK::HeadFetchError)
                .map_dbg_err_str("Failed to fetch HEAD")
                .map_into_hook_error()
        );

        let file_status = try!(
            repo
                .status_file(&path)
                .map_dbg_err_str("Failed to fetch file status")
                .map_dbg_err(|e| format!("\t->  {:?}", e))
                .map_err_into(GHEK::RepositoryFileStatusError)
                .map_into_hook_error()
        );

        let cb = &mut |path: &Path, _matched_spec: &[u8]| -> i32 {
            if file_status.contains(STATUS_WT_NEW) || file_status.contains(STATUS_WT_MODIFIED) {
                debug!("[GIT CREATE HOOK]: File is modified/new: {}", path.display());
                0
            } else {
                debug!("[GIT CREATE HOOK]: Ignoring file: {}", path.display());
                1
            }
        };

        try!(
            index.add_all(&[path], ADD_DEFAULT, Some(cb as &mut IndexMatchedPath))
                .map_err_into(GHEK::RepositoryPathAddingError)
                .map_dbg_err_str("Failed to add to index")
                .map_into_hook_error()
        );

        let mut parents = Vec::new();
        {
            let commit = try!(
                repo.find_commit(head.target().unwrap())
                    .map_err_into(GHEK::RepositoryParentFetchingError)
                    .map_dbg_err_str("Failed to find commit HEAD")
                    .map_into_hook_error()
            );
            parents.push(commit);
        }

        // for converting from Vec<Commit> to Vec<&Commit>
        let parents = parents.iter().collect::<Vec<_>>();

        let tree_id = try!(
            index.write_tree()
                .map_err_into(GHEK::RepositoryIndexWritingError)
                .map_dbg_err_str("Failed to write tree")
                .map_into_hook_error()
        );

        let tree = try!(
            repo.find_tree(tree_id)
                .map_err_into(GHEK::RepositoryParentFetchingError)
                .map_dbg_err_str("Failed to find tree")
                .map_into_hook_error()
        );

        let message = try!(commit_message(cfg, StoreAction::Update)
                .map_dbg_err_str("Failed to get commit message"));

        repo.commit(Some("HEAD"), &signature, &signature, &message, &tree, &parents)
            .map_dbg_err_str("Failed to commit")
            .map_err_into(GHEK::RepositoryCommittingError)
            .map_into_hook_error()
            .map(|_| ())
    }

}

<|endoftext|>


<commit_before>use super::Input;
use crate::flowgger::config::Config;
use crate::flowgger::decoder::Decoder;
use crate::flowgger::encoder::Encoder;
use flate2::read::{GzDecoder, ZlibDecoder};
use std::io::{stderr, Read, Write};
use std::net::UdpSocket;
use std::str;
use std::sync::mpsc::SyncSender;

const DEFAULT_LISTEN: &str = "0.0.0.0:514";
const MAX_UDP_PACKET_SIZE: usize = 65_527;
const MAX_COMPRESSION_RATIO: usize = 5;

pub struct UdpInput {
    listen: String,
}

impl UdpInput {
    pub fn new(config: &Config) -> UdpInput {
        let listen = config
            .lookup("input.listen")
            .map_or(DEFAULT_LISTEN, |x| {
                x.as_str().expect("input.listen must be an ip:port string")
            })
            .to_owned();
        UdpInput { listen }
    }
}

impl Input for UdpInput {
    fn accept(
        &self,
        tx: SyncSender<Vec<u8>>,
        decoder: Box<dyn Decoder + Send>,
        encoder: Box<dyn Encoder + Send>,
    ) {
        let socket = UdpSocket::bind(&self.listen as &str)
            .unwrap_or_else(|_| panic!("Unable to listen to {}", self.listen));
        let tx = tx.clone();
        let (decoder, encoder): (Box<dyn Decoder>, Box<dyn Encoder>) =
            (decoder.clone_boxed(), encoder.clone_boxed());
        let mut buf = [0; MAX_UDP_PACKET_SIZE];
        loop {
            let (length, _src) = match socket.recv_from(&mut buf) {
                Ok(res) => res,
                Err(_) => continue,
            };
            let line = &buf[..length];
            if let Err(e) = handle_record_maybe_compressed(line, &tx, &decoder, &encoder) {
                let _ = writeln!(stderr(), "{}", e);
            }
        }
    }
}

fn handle_record_maybe_compressed(
    line: &[u8],
    tx: &SyncSender<Vec<u8>>,
    decoder: &Box<dyn Decoder>,
    encoder: &Box<dyn Encoder>,
) -> Result<(), &'static str> {
    if line.len() >= 8
        && (line[0] == 0x78 && (line[1] == 0x01 || line[1] == 0x9c || line[1] == 0xda))
    {
        let mut decompressed = Vec::with_capacity(MAX_UDP_PACKET_SIZE * MAX_COMPRESSION_RATIO);
        match ZlibDecoder::new(line).read_to_end(&mut decompressed) {
            Ok(_) => handle_record(&decompressed, tx, decoder, encoder),
            Err(_) => Err("Corrupted compressed (zlib) record"),
        }
    } else if line.len() >= 24 && (line[0] == 0x1f && line[1] == 0x8b && line[2] == 0x08) {
        let mut decompressed = Vec::with_capacity(MAX_UDP_PACKET_SIZE * MAX_COMPRESSION_RATIO);
        match GzDecoder::new(line).read_to_end(&mut decompressed) {
            Ok(_) => handle_record(&decompressed, tx, decoder, encoder),
            Err(_) => Err("Corrupted compressed (gzip) record"),
        }
    } else {
        handle_record(line, tx, decoder, encoder)
    }
}

fn handle_record(
    line: &[u8],
    tx: &SyncSender<Vec<u8>>,
    decoder: &Box<dyn Decoder>,
    encoder: &Box<dyn Encoder>,
) -> Result<(), &'static str> {
    let line = match str::from_utf8(line) {
        Err(_) => return Err("Invalid UTF-8 input"),
        Ok(line) => line,
    };
    let decoded = decoder.decode(line)?;
    let reencoded = encoder.encode(decoded)?;
    tx.send(reencoded).unwrap();
    Ok(())
}
<commit_msg>UDP input module operational improvements (#17)<commit_after>use super::Input;
use crate::flowgger::config::Config;
use crate::flowgger::decoder::Decoder;
use crate::flowgger::encoder::Encoder;
use flate2::read::{GzDecoder, ZlibDecoder};
use std::io::{stderr, Read, Write};
use std::net::SocketAddr;
use std::net::UdpSocket;
use std::str;
use std::sync::mpsc::SyncSender;

const DEFAULT_LISTEN: &str = "0.0.0.0:514";
const MAX_UDP_PACKET_SIZE: usize = 65_527;
const MAX_COMPRESSION_RATIO: usize = 5;

/// UDP input structure for flowgger
/// It will receive messages from the network, decode them and reencoded them as configured
/// in the [`Config`][] object it takes as input
///
/// [`Config`]: ../config/struct.Config.html
pub struct UdpInput {
    listen: SocketAddr,
}

impl UdpInput {
    /// Attemps to create a new UdpInput instance by parsing the a Config object in the toml format
    /// the only field needed for this to work in input.listen, if input.listen is missing it will
    /// bind itself to a default ip:port address `0.0.0.0:514`
    ///
    /// # Parameters
    /// `config`: Configuration object in toml format
    ///
    /// # Panic
    /// `input.listen must be an ip:port string`:  input.listen is not parsable as a string
    /// `Unable to parse ip:port string from input.listen` input.listen is not a valid ip:port
    pub fn new(config: &Config) -> UdpInput {
        let listen = config
            .lookup("input.listen")
            .map_or(DEFAULT_LISTEN, |x| {
                x.as_str().expect("input.listen must be an ip:port string")
            })
            .to_owned();
        let bind_address: SocketAddr = listen
            .parse()
            .expect("unable to parse ip:port string from input.listen");
        UdpInput {
            listen: bind_address,
        }
    }
}

impl Input for UdpInput {
    /// Bind a [`UdpSocket`][] to the configured listen address and starts a loop for accepting
    /// incoming upd packets
    ///
    /// [`UdpSocket`]: https://doc.rust-lang.org/std/net/struct.UdpSocket.html
    ///
    /// # Parameters
    /// `tx`: Sender channel
    /// `decoder`: Box containing a dynamically allocated Decoder
    /// `encoder`: Box containing a dynamically allocated Encoder
    ///
    /// # Panics
    /// `Unable to listen to <socket>`: Socket is already open by another program or current
    /// permissions are insufficent to open the specified socket
    fn accept(
        &self,
        tx: SyncSender<Vec<u8>>,
        decoder: Box<dyn Decoder + Send>,
        encoder: Box<dyn Encoder + Send>,
    ) {
        let socket = UdpSocket::bind(&self.listen)
            .unwrap_or_else(|_| panic!("Unable to listen to {}", self.listen));
        let tx = tx.clone();
        let (decoder, encoder): (Box<dyn Decoder>, Box<dyn Encoder>) =
            (decoder.clone_boxed(), encoder.clone_boxed());
        let mut buf = [0; MAX_UDP_PACKET_SIZE];
        loop {
            let (length, _src) = match socket.recv_from(&mut buf) {
                Ok(res) => res,
                Err(_) => continue,
            };
            let line = &buf[..length];
            if let Err(e) = handle_record_maybe_compressed(line, &tx, &decoder, &encoder) {
                let _ = writeln!(stderr(), "{}", e);
            }
        }
    }
}

/// Handle a line that could be compressed in the Zlib or Gz format, uncompress it if compressed
/// with a known algoritm and passed it to handle_record to decoded it from the input format to the
/// output one and send it over for being sent in output
///
/// # Errors
/// `Corrupted compressed (gzip/zlib) record`: The record has been identified as a compressed record in a known format
/// but could not be handled
/// `Invalid UTF-8 input`: Bubble up from handle_record, the record is not in a valid utf-8 format, it could be a non
/// supported compression format
fn handle_record_maybe_compressed(
    line: &[u8],
    tx: &SyncSender<Vec<u8>>,
    decoder: &Box<dyn Decoder>,
    encoder: &Box<dyn Encoder>,
) -> Result<(), &'static str> {
    if line.len() >= 8
        && (line[0] == 0x78 && (line[1] == 0x01 || line[1] == 0x9c || line[1] == 0xda))
    {
        let mut decompressed = Vec::with_capacity(MAX_UDP_PACKET_SIZE * MAX_COMPRESSION_RATIO);
        match ZlibDecoder::new(line).read_to_end(&mut decompressed) {
            Ok(_) => handle_record(&decompressed, tx, decoder, encoder),
            Err(_) => Err("Corrupted compressed (zlib) record"),
        }
    } else if line.len() >= 24 && (line[0] == 0x1f && line[1] == 0x8b && line[2] == 0x08) {
        let mut decompressed = Vec::with_capacity(MAX_UDP_PACKET_SIZE * MAX_COMPRESSION_RATIO);
        match GzDecoder::new(line).read_to_end(&mut decompressed) {
            Ok(_) => handle_record(&decompressed, tx, decoder, encoder),
            Err(_) => Err("Corrupted compressed (gzip) record"),
        }
    } else {
        handle_record(line, tx, decoder, encoder)
    }
}

/// Decode a byte line in a valid utf-8 format, encodes it and sends it over throught a channel
///
/// # Errors
/// `Invalid UTF-8 input`: The record is not in a valid utf-8 format, it could be a non supported compression format
fn handle_record(
    line: &[u8],
    tx: &SyncSender<Vec<u8>>,
    decoder: &Box<dyn Decoder>,
    encoder: &Box<dyn Encoder>,
) -> Result<(), &'static str> {
    let line = match str::from_utf8(line) {
        Err(_) => return Err("Invalid UTF-8 input"),
        Ok(line) => line,
    };
    let decoded = decoder.decode(line)?;
    let reencoded = encoder.encode(decoded)?;
    tx.send(reencoded).unwrap();
    Ok(())
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::flowgger::config::Config;
    use crate::flowgger::get_decoder_rfc3164;
    use crate::flowgger::get_encoder_rfc3164;
    use flate2::write::{GzEncoder, ZlibEncoder};
    use flate2::Compression;
    use std::net;
    use std::sync::mpsc::{sync_channel, Receiver};

    const DEFAULT_QUEUE_SIZE: usize = 10_000_000;

    #[test]
    fn test_udp_input_constructor() {
        let listen_ip = "127.0.0.1:5000";
        let config =
            Config::from_string(format!("[input]\nlisten = \"{}\"", listen_ip).as_str()).unwrap();
        let input = UdpInput::new(&config);
        let listen_addr: SocketAddr = listen_ip.parse().unwrap();
        assert_eq!(input.listen, listen_addr);
    }

    #[test]
    #[should_panic(expected = "unable to parse ip:port string from input.listen")]
    fn test_udp_input_constructor_bad_input() {
        let config = Config::from_string("[input]\nlisten = \"wrongaddress\"").unwrap();
        UdpInput::new(&config);
    }

    #[test]
    fn test_udp_input_default_constructor() {
        let config = Config::from_string("").unwrap();
        let input = UdpInput::new(&config);
        let default_addr: SocketAddr = DEFAULT_LISTEN.parse().unwrap();
        assert_eq!(input.listen, default_addr);
    }

    fn handle_record_set_up() -> (
        &'static str,
        SyncSender<Vec<u8>>,
        Receiver<Vec<u8>>,
        Box<dyn Decoder>,
        Box<dyn Encoder>,
    ) {
        let line = "Aug  6 11:15:24 testhostname appname 69 42 [origin@123 software=\"te\\st sc\"ript\" swVersion=\"0.0.1\"] test message";
        let (tx, rx): (SyncSender<Vec<u8>>, Receiver<Vec<u8>>) = sync_channel(DEFAULT_QUEUE_SIZE);
        let config = Config::from_string("").unwrap();
        let encoder = get_encoder_rfc3164(&config);
        let decoder = get_decoder_rfc3164(&config);
        let (decoder, encoder): (Box<dyn Decoder>, Box<dyn Encoder>) =
            (decoder.clone_boxed(), encoder.clone_boxed());
        (line, tx, rx, decoder, encoder)
    }

    #[test]
    fn test_udp_input_handle_record_uncompressed() {
        let (line, tx, rx, decoder, encoder) = handle_record_set_up();
        handle_record_maybe_compressed(line.as_bytes(), &tx, &decoder, &encoder).unwrap();
        let transmitted = rx.recv().unwrap();
        assert_eq!(str::from_utf8(&transmitted).unwrap(), line);
    }

    #[test]
    fn test_handle_record_compressed_zlib() {
        let (line, tx, rx, decoder, encoder) = handle_record_set_up();
        let mut compressor = ZlibEncoder::new(Vec::new(), Compression::default());
        match compressor.write_all(line.as_bytes()) {
            Ok(e) => e,
            Err(e) => panic!("Compressing line {}, raised Error {:?}", line, e),
        }
        let compressed_line = compressor.finish().unwrap();
        handle_record_maybe_compressed(&compressed_line, &tx, &decoder, &encoder).unwrap();
        let transmitted = rx.recv().unwrap();
        assert_eq!(str::from_utf8(&transmitted).unwrap(), line);
    }

    #[test]
    fn test_handle_record_compressed_gz() {
        let (line, tx, rx, decoder, encoder) = handle_record_set_up();
        let mut compressor = GzEncoder::new(Vec::new(), Compression::default());
        match compressor.write_all(line.as_bytes()) {
            Ok(e) => e,
            Err(e) => panic!("Compressing line {}, raised Error {:?}", line, e),
        }
        let compressed_line = compressor.finish().unwrap();
        handle_record_maybe_compressed(&compressed_line, &tx, &decoder, &encoder).unwrap();
        let transmitted = rx.recv().unwrap();
        assert_eq!(str::from_utf8(&transmitted).unwrap(), line);
    }

    #[test]
    #[should_panic(expected = "Invalid UTF-8 input")]
    fn test_handle_record_bad_record() {
        let (line, tx, _rx, decoder, encoder) = handle_record_set_up();
        let mut compressor = GzEncoder::new(Vec::new(), Compression::default());
        match compressor.write_all(line.as_bytes()) {
            Ok(e) => e,
            Err(e) => panic!("Compressing line {}, raised Error {:?}", line, e),
        }
        let mut compressed_line = compressor.finish().unwrap();
        compressed_line.truncate(5);
        handle_record_maybe_compressed(&compressed_line, &tx, &decoder, &encoder).unwrap();
    }
}
<|endoftext|>