<commit_before>use sdl2;
use sdl2::AudioSubsystem;
use sdl2::EventPump;
use sdl2::event::Event;
use sdl2::render::Renderer;
use sdl2::pixels::Color;
use sdl2::rect::Rect;
use sdl2::image::LoadTexture;
use sdl2::mixer;
use sdl2::mixer::Sdl2MixerContext;
use sdl2::mixer::Channel;
use sdl2::mixer::Chunk;
use OptionAway;
use input_private;
use input::Input;
use game_object::GameObject;
use component::Message;
use game_object;
use texture::Texture;
use texture;
use PolarCoords;
use Vector;
use transform::Transform;
use camera::Camera;
use nphysics2d::world::World;
use std::mem::replace;
use std::collections::HashMap;
use std::borrow::Borrow;
use std::path::PathBuf;
use std::f32;
use std::time::Instant;
use chrono::Duration;

pub struct App {
    exit: bool,
    renderer: Renderer<'static>,
    _audio: AudioSubsystem,
    _mixer: Sdl2MixerContext,
    event_pump: EventPump,
    game_objects: HashMap<u64, Option<Box<GameObject>>>,
    sound_cache: HashMap<String, Chunk>,
    next_game_object_id: u64,
    pub input: Input,
    pub camera: Camera,
    pub world: World<f32>,
}

impl App {
    ///Constructs a new App using the given name for the window title.
    pub fn new(name: &str) -> App {
        let sdl_context = sdl2::init().expect("Failed to initialize SDL2.");
        let video_subsystem = sdl_context.video().expect("Failed to initialize video subsystem");
        let window = video_subsystem.window(name, 800, 600)
            .position_centered()
            .fullscreen_desktop()
            .opengl()
            .build()
            .unwrap();
        let renderer = window.renderer().build().expect("Failed to initialize renderer");
        let audio = sdl_context.audio().expect("Failed to initialize audio");
        let mixer = mixer::init(mixer::INIT_OGG).expect("Failed to initialize mixer");
        mixer::open_audio(mixer::DEFAULT_FREQUENCY, mixer::DEFAULT_FORMAT, 2, 1024)
            .expect("Failed to open audio");
        mixer::allocate_channels(256);
        mixer::reserve_channels(128);
        App {
            exit: false,
            next_game_object_id: 0,
            game_objects: HashMap::new(),
            input: Input::new(),
            renderer: renderer,
            sound_cache: HashMap::new(),
            _audio: audio,
            _mixer: mixer,
            event_pump: sdl_context.event_pump().expect("Failed to initalize event pump."),
            camera: Camera { transform: Transform::new() },
            world: World::new(),
        }
    }

    /// Called every frame to paint the scene. Do not put game logic here, that goes in update.
    fn render(&mut self) {
        use std::f32;
        let game_objs = &self.game_objects;
        let camera_transform = self.camera.transform;
        // Clear the screen with grey.
        self.renderer.set_draw_color(Color::RGB(240, 240, 240));
        self.renderer.clear();
        // Reset the draw color to white so subsequent draw calls are correct.
        self.renderer.set_draw_color(Color::RGB(255, 255, 255));
        let (screen_width, screen_height) = self.renderer.window().unwrap().size();
        for game_obj in game_objs.values() {
            if let Some(ref game_obj) = *game_obj {
                let mut render_transform = game_obj.transform;
                *render_transform.mut_x() -= *camera_transform.x();
                *render_transform.mut_y() -= *camera_transform.y();
                let mut polar = PolarCoords::from(render_transform.position().clone());
                polar.rotation -= *camera_transform.rotation();
                *render_transform.mut_position() = Vector::from(polar);
                *render_transform.mut_x() += (screen_width / 2) as f32;
                *render_transform.mut_y() += (screen_height / 2) as f32;
                *render_transform.mut_rotation() -= *camera_transform.rotation();
                let (tex_width, tex_height) = texture::size(&game_obj.texture);
                let render_rect =
                    Rect::new(((*render_transform.x()) as i32) - (tex_width as i32 / 2),
                              ((*render_transform.y()) as i32) - (tex_height as i32 / 2),
                              tex_width,
                              tex_height);
                if let &Some(ref texture) = texture::get_raw(&game_obj.texture) {
                    let result = self.renderer.copy_ex(&texture,
                                                       None,
                                                       Some(render_rect),
                                                       (*game_obj.transform.rotation() * 180.0 /
                                                        f32::consts::PI) as
                                                       f64,
                                                       None,
                                                       false,
                                                       false);
                    if let Err(err) = result {
                        println!("Unable to draw texture, Error: {:?}", err);
                    }
                }
            }
        }
        self.renderer.present();
    }

    /// Called every frame to simulate the game. Do not put rendering here, that goes in render.
    fn update(&mut self, delta_time: f32) {
        //Copy game_object to the physics world, step, then copy from physics to game_object
        for mut game_object in self.game_objects.values_mut() {
            if let Some(game_object) = game_object.as_mut() {
                game_object::copy_to_physics(game_object);
            }
        }
        self.world.step(delta_time);
        for mut game_object in self.game_objects.values_mut() {
            if let Some(game_object) = game_object.as_mut() {
                game_object::copy_from_physics(game_object);
            }
        }

        //Send update messages
        let keys = self.game_objects.keys().map(|x| *x).collect::<Vec<u64>>();
        for key in keys {
            let mut game_obj_option = None;
            if let Some(game_obj_ref) = self.game_objects.get_mut(&key) {
                game_obj_option = replace(game_obj_ref, None);
            }
            if let Some(ref mut game_obj) = game_obj_option {
                game_obj.update(self, delta_time);
                *game_obj.transform.mut_rotation() %= 2.0 * f32::consts::PI;
            }
            if let Some(game_obj_ref) = self.game_objects.get_mut(&key) {
                replace(game_obj_ref, game_obj_option);
            }
        }

        // Drop any game objects that need to be dropped.
        let mut dropped_keys = Vec::new();
        for (key, game_obj) in self.game_objects.iter() {
            if let Some(ref game_obj) = *game_obj {
                if game_object::was_dropped(game_obj.borrow()) {
                    dropped_keys.push(*key);
                }
            } else {
                dropped_keys.push(*key);
            }
        }
        for key in dropped_keys {
            let removed = self.game_objects.remove(&key);
            if let Some(Some(mut game_obj)) = removed {
                game_obj.receive_message(self, &Message::OnDestroy);
            }
        }

        input_private::input::shift_frame(&mut self.input);
    }

    /// Begin execution of the game.
    ///
    /// This function won't return until the game has been quit.
    pub fn run(&mut self) {
        mixer::allocate_channels(128);
        let mut last_frame_instant = Instant::now();
        while !self.exit {
            while let Some(e) = self.event_pump.poll_event() {
                if let Event::Quit { .. } = e {
                    self.exit = true;
                }
                input_private::input::process_event(&mut self.input, &e);
            }
            self.update(Duration::from_std(last_frame_instant.elapsed())
                .unwrap()
                .num_microseconds()
                .unwrap() as f32 / 1000000.0);
            last_frame_instant = Instant::now();
            self.render();
        }
    }

    /// Exit the game
    ///
    /// Signals that it's time to quit. Execution will stop at the end of the current frame.
    pub fn quit(&mut self) {
        self.exit = true;
    }

    pub fn camera(&self) -> &Camera {
        &self.camera
    }

    pub fn camera_mut(&mut self) -> &mut Camera {
        &mut self.camera
    }

    /// Plays a sound at a volume between 1.0 and 0.0
    ///
    /// path is a filename relative to assets/sounds (assets\sounds on Windows)
    pub fn play_sound(&mut self, path: &str, volume: f32) -> Result<(), String> {
        if let Some(chunk) = self.sound_cache.get(path) {
            let channel = Channel::all().play(chunk, 0)?;
            channel.set_volume((volume * 128.0) as i32);
            return Ok(());
        }
        let mut file_path = PathBuf::from("assets");
        file_path.push("sounds");
        file_path.push(path);
        let chunk = Chunk::from_file(file_path)?;
        let channel = Channel::all().play(&chunk, 0)?;
        channel.set_volume((volume * 128.0) as i32);
        self.sound_cache.insert(path.to_owned(), chunk);
        Ok(())
    }

    /// Play a sound on a user sound channel.
    ///
    /// There are 128 user sound channels available at indices 0-127. You may use them as you
    /// please.
    pub fn play_sound_on_channel(&mut self, channel: i32, path: &str) -> Result<(), String> {
        if channel >= 0 && channel < 128 {
            mixer::channel(channel).play(fetch_sound_chunk(self, path)?, 0)?;
            return Ok(());
        }
        Err("Channel out of range.".to_string())
    }

    /// Set the volume for a user sound channel. Volume is between 0.0 and 1.0
    ///
    /// There are 128 user sound channels available at indices 0-127. You may use them as you
    /// please.
    pub fn set_channel_volume(&mut self, channel: i32, volume: f32) -> Result<(), String> {
       if channel >= 0 && channel < 128 {
            mixer::channel(channel).set_volume((volume * 128.0) as i32);
            return Ok(());
       }
       Err("Channel out of range.".to_string())
    }

    pub fn get_channel_volume(&self, channel: i32) -> Result<f32, String> {
        if channel >= 0 && channel < 128 {
            return Ok((mixer::channel(channel).get_volume() as f32)/128.0);
        }
        Err("Channel out of range.".to_string())
    }

    pub fn pause_channel(&mut self, channel: i32) -> Result<(), String> {
        if channel >= 0 && channel < 128 {
            mixer::channel(channel).pause();
            return Ok(());
        }
        Err("Channel out of range.".to_string())
    }

    pub fn resume_channel(&mut self, channel: i32) -> Result<(), String> {
        if channel >= 0 && channel < 128 {
            mixer::channel(channel).resume();
            return Ok(());
        }
        Err("Channel out of range.".to_string())
    }

    /// Creates a new GameObject
    ///
    /// GameObjects are typically physical objects in your game world, such as characters or
    /// decorative objects.
    ///
    /// Parameter f will be called on the newly created GameObject, allowing you to initialize it.
    /// Returns the id of the new GameObject, which can be used to retrieve this GameObject later.
    pub fn new_gameobject<F>(&mut self, f: F) -> u64
        where F: FnOnce(&mut App, &mut GameObject)
    {
        let mut game_object = Box::new(game_object::new(self));
        f(self, &mut game_object);
        let id = game_object.id();
        self.game_objects.insert(id, Some(game_object));
        id
    }

    /// Retrieves a reference to an existing GameObject by the GameObject's id.
    ///
    /// Returns Away if and only if the GameObject exists but is currently loaned out
    /// in another borrow.
    ///
    /// Hint: That borrow is most likely in the game_object parameter of the receive_message
    /// function.
    pub fn game_object_by_id(&self, id: u64) -> OptionAway<&GameObject> {
        OptionAway::from(self.game_objects.get(&id))
    }

    /// Retrieves a mutable reference to an existing GameObject by the GameObject's id.
    ///
    /// Returns Away if and only if the GameObject exists but is currently loaned out
    /// in another borrow.
    ///
    /// Hint: That borrow is most likely in the game_object parameter of the receive_message
    /// function.
    pub fn game_object_by_id_mut(&mut self, id: u64) -> OptionAway<&mut GameObject> {
        OptionAway::from(self.game_objects.get_mut(&id))
    }

    /// Loads a texture and returns it for use.
    ///
    /// texture_name is the file name of the texture relative to assets/textures
    /// (assets\textures on Windows)
    pub fn fetch_texture(&mut self, texture_name: &str) -> Result<Texture, String> {
        let mut texture_path = PathBuf::from("assets");
        texture_path.push("textures");
        texture_path.push(texture_name);
        let sdl_texture = self.renderer.load_texture(texture_path.as_path())?;
        let mut nitro_texture = Texture::new();
        texture::set_raw(&mut nitro_texture, sdl_texture);
        Ok(nitro_texture)
    }
}

pub fn fetch_sound_chunk<'a>(app: &'a mut App, path: &str) -> Result<&'a Chunk, String> {
    if let Some(chunk) = app.sound_cache.get(path) {
        return Ok(chunk);
    }
    let mut file_path = PathBuf::from("assets");
    file_path.push("sounds");
    file_path.push(path);
    let chunk = Chunk::from_file(file_path)?;
    app.sound_cache.insert(path.to_string(), chunk);
    Ok(app.sound_cache.get(path).unwrap())
}

// This function will never return 0.  0 can now be used as a null value.
pub fn next_game_object_id(app: &mut App) -> u64 {
    app.next_game_object_id += 1;
    app.next_game_object_id
}
<commit_msg>Fix Audio API<commit_after>use sdl2;
use sdl2::AudioSubsystem;
use sdl2::EventPump;
use sdl2::event::Event;
use sdl2::render::Renderer;
use sdl2::pixels::Color;
use sdl2::rect::Rect;
use sdl2::image::LoadTexture;
use sdl2::mixer;
use sdl2::mixer::Sdl2MixerContext;
use sdl2::mixer::Channel;
use sdl2::mixer::Chunk;
use OptionAway;
use input_private;
use input::Input;
use game_object::GameObject;
use component::Message;
use game_object;
use texture::Texture;
use texture;
use PolarCoords;
use Vector;
use transform::Transform;
use camera::Camera;
use nphysics2d::world::World;
use std::mem::replace;
use std::collections::HashMap;
use std::borrow::Borrow;
use std::path::PathBuf;
use std::f32;
use std::time::Instant;
use chrono::Duration;

pub struct App {
    exit: bool,
    renderer: Renderer<'static>,
    _audio: AudioSubsystem,
    _mixer: Sdl2MixerContext,
    event_pump: EventPump,
    game_objects: HashMap<u64, Option<Box<GameObject>>>,
    sound_cache: HashMap<String, Chunk>,
    next_game_object_id: u64,
    pub input: Input,
    pub camera: Camera,
    pub world: World<f32>,
}

impl App {
    ///Constructs a new App using the given name for the window title.
    pub fn new(name: &str) -> App {
        let sdl_context = sdl2::init().expect("Failed to initialize SDL2.");
        let video_subsystem = sdl_context.video().expect("Failed to initialize video subsystem");
        let window = video_subsystem.window(name, 800, 600)
            .position_centered()
            .fullscreen_desktop()
            .opengl()
            .build()
            .unwrap();
        let renderer = window.renderer().build().expect("Failed to initialize renderer");
        let audio = sdl_context.audio().expect("Failed to initialize audio");
        let mixer = mixer::init(mixer::INIT_OGG).expect("Failed to initialize mixer");
        mixer::open_audio(mixer::DEFAULT_FREQUENCY, mixer::DEFAULT_FORMAT, 2, 1024)
            .expect("Failed to open audio");
        mixer::allocate_channels(256);
        mixer::reserve_channels(128);
        App {
            exit: false,
            next_game_object_id: 0,
            game_objects: HashMap::new(),
            input: Input::new(),
            renderer: renderer,
            sound_cache: HashMap::new(),
            _audio: audio,
            _mixer: mixer,
            event_pump: sdl_context.event_pump().expect("Failed to initalize event pump."),
            camera: Camera { transform: Transform::new() },
            world: World::new(),
        }
    }

    /// Called every frame to paint the scene. Do not put game logic here, that goes in update.
    fn render(&mut self) {
        use std::f32;
        let game_objs = &self.game_objects;
        let camera_transform = self.camera.transform;
        // Clear the screen with grey.
        self.renderer.set_draw_color(Color::RGB(240, 240, 240));
        self.renderer.clear();
        // Reset the draw color to white so subsequent draw calls are correct.
        self.renderer.set_draw_color(Color::RGB(255, 255, 255));
        let (screen_width, screen_height) = self.renderer.window().unwrap().size();
        for game_obj in game_objs.values() {
            if let Some(ref game_obj) = *game_obj {
                let mut render_transform = game_obj.transform;
                *render_transform.mut_x() -= *camera_transform.x();
                *render_transform.mut_y() -= *camera_transform.y();
                let mut polar = PolarCoords::from(render_transform.position().clone());
                polar.rotation -= *camera_transform.rotation();
                *render_transform.mut_position() = Vector::from(polar);
                *render_transform.mut_x() += (screen_width / 2) as f32;
                *render_transform.mut_y() += (screen_height / 2) as f32;
                *render_transform.mut_rotation() -= *camera_transform.rotation();
                let (tex_width, tex_height) = texture::size(&game_obj.texture);
                let render_rect =
                    Rect::new(((*render_transform.x()) as i32) - (tex_width as i32 / 2),
                              ((*render_transform.y()) as i32) - (tex_height as i32 / 2),
                              tex_width,
                              tex_height);
                if let &Some(ref texture) = texture::get_raw(&game_obj.texture) {
                    let result = self.renderer.copy_ex(&texture,
                                                       None,
                                                       Some(render_rect),
                                                       (*game_obj.transform.rotation() * 180.0 /
                                                        f32::consts::PI) as
                                                       f64,
                                                       None,
                                                       false,
                                                       false);
                    if let Err(err) = result {
                        println!("Unable to draw texture, Error: {:?}", err);
                    }
                }
            }
        }
        self.renderer.present();
    }

    /// Called every frame to simulate the game. Do not put rendering here, that goes in render.
    fn update(&mut self, delta_time: f32) {
        //Copy game_object to the physics world, step, then copy from physics to game_object
        for mut game_object in self.game_objects.values_mut() {
            if let Some(game_object) = game_object.as_mut() {
                game_object::copy_to_physics(game_object);
            }
        }
        self.world.step(delta_time);
        for mut game_object in self.game_objects.values_mut() {
            if let Some(game_object) = game_object.as_mut() {
                game_object::copy_from_physics(game_object);
            }
        }

        //Send update messages
        let keys = self.game_objects.keys().map(|x| *x).collect::<Vec<u64>>();
        for key in keys {
            let mut game_obj_option = None;
            if let Some(game_obj_ref) = self.game_objects.get_mut(&key) {
                game_obj_option = replace(game_obj_ref, None);
            }
            if let Some(ref mut game_obj) = game_obj_option {
                game_obj.update(self, delta_time);
                *game_obj.transform.mut_rotation() %= 2.0 * f32::consts::PI;
            }
            if let Some(game_obj_ref) = self.game_objects.get_mut(&key) {
                replace(game_obj_ref, game_obj_option);
            }
        }

        // Drop any game objects that need to be dropped.
        let mut dropped_keys = Vec::new();
        for (key, game_obj) in self.game_objects.iter() {
            if let Some(ref game_obj) = *game_obj {
                if game_object::was_dropped(game_obj.borrow()) {
                    dropped_keys.push(*key);
                }
            } else {
                dropped_keys.push(*key);
            }
        }
        for key in dropped_keys {
            let removed = self.game_objects.remove(&key);
            if let Some(Some(mut game_obj)) = removed {
                game_obj.receive_message(self, &Message::OnDestroy);
            }
        }

        input_private::input::shift_frame(&mut self.input);
    }

    /// Begin execution of the game.
    ///
    /// This function won't return until the game has been quit.
    pub fn run(&mut self) {
        mixer::allocate_channels(128);
        let mut last_frame_instant = Instant::now();
        while !self.exit {
            while let Some(e) = self.event_pump.poll_event() {
                if let Event::Quit { .. } = e {
                    self.exit = true;
                }
                input_private::input::process_event(&mut self.input, &e);
            }
            self.update(Duration::from_std(last_frame_instant.elapsed())
                .unwrap()
                .num_microseconds()
                .unwrap() as f32 / 1000000.0);
            last_frame_instant = Instant::now();
            self.render();
        }
    }

    /// Exit the game
    ///
    /// Signals that it's time to quit. Execution will stop at the end of the current frame.
    pub fn quit(&mut self) {
        self.exit = true;
    }

    pub fn camera(&self) -> &Camera {
        &self.camera
    }

    pub fn camera_mut(&mut self) -> &mut Camera {
        &mut self.camera
    }

    /// Plays a sound at a volume between 1.0 and 0.0
    ///
    /// path is a filename relative to assets/sounds (assets\sounds on Windows)
    pub fn play_sound(&mut self, path: &str, volume: f32) -> Result<(), String> {
        if let Some(chunk) = self.sound_cache.get(path) {
            let channel = Channel::all().play(chunk, 0)?;
            channel.set_volume((volume * 128.0) as i32);
            return Ok(());
        }
        let mut file_path = PathBuf::from("assets");
        file_path.push("sounds");
        file_path.push(path);
        let chunk = Chunk::from_file(file_path)?;
        let channel = Channel::all().play(&chunk, 0)?;
        channel.set_volume((volume * 128.0) as i32);
        self.sound_cache.insert(path.to_owned(), chunk);
        Ok(())
    }

    /// Play a sound on a user sound channel.
    ///
    /// There are 128 user sound channels available at indices 0-127. You may use them as you
    /// please.
    pub fn play_sound_on_channel(&mut self, channel: i32, path: &str) -> Result<(), String> {
        if channel >= 0 && channel < 128 {
            if let Some(chunk) = self.sound_cache.get(path) {
                mixer::channel(channel).play(chunk, 0)?;
                return Ok(());
            }
            let mut file_path = PathBuf::from("assets");
            file_path.push("sounds");
            file_path.push(path);
            let chunk = Chunk::from_file(file_path)?;
            Channel::all().play(&chunk, 0)?;
            self.sound_cache.insert(path.to_owned(), chunk);
            return Ok(());
        }
        Err("Channel out of range.".to_string())
    }

    /// Set the volume for a user sound channel. Volume is between 0.0 and 1.0
    ///
    /// There are 128 user sound channels available at indices 0-127. You may use them as you
    /// please.
    pub fn set_channel_volume(&mut self, channel: i32, volume: f32) -> Result<(), String> {
       if channel >= 0 && channel < 128 {
            mixer::channel(channel).set_volume((volume * 128.0) as i32);
            return Ok(());
       }
       Err("Channel out of range.".to_string())
    }

    pub fn get_channel_volume(&self, channel: i32) -> Result<f32, String> {
        if channel >= 0 && channel < 128 {
            return Ok((mixer::channel(channel).get_volume() as f32)/128.0);
        }
        Err("Channel out of range.".to_string())
    }

    pub fn pause_channel(&mut self, channel: i32) -> Result<(), String> {
        if channel >= 0 && channel < 128 {
            mixer::channel(channel).pause();
            return Ok(());
        }
        Err("Channel out of range.".to_string())
    }

    pub fn resume_channel(&mut self, channel: i32) -> Result<(), String> {
        if channel >= 0 && channel < 128 {
            mixer::channel(channel).resume();
            return Ok(());
        }
        Err("Channel out of range.".to_string())
    }

    /// Creates a new GameObject
    ///
    /// GameObjects are typically physical objects in your game world, such as characters or
    /// decorative objects.
    ///
    /// Parameter f will be called on the newly created GameObject, allowing you to initialize it.
    /// Returns the id of the new GameObject, which can be used to retrieve this GameObject later.
    pub fn new_gameobject<F>(&mut self, f: F) -> u64
        where F: FnOnce(&mut App, &mut GameObject)
    {
        let mut game_object = Box::new(game_object::new(self));
        f(self, &mut game_object);
        let id = game_object.id();
        self.game_objects.insert(id, Some(game_object));
        id
    }

    /// Retrieves a reference to an existing GameObject by the GameObject's id.
    ///
    /// Returns Away if and only if the GameObject exists but is currently loaned out
    /// in another borrow.
    ///
    /// Hint: That borrow is most likely in the game_object parameter of the receive_message
    /// function.
    pub fn game_object_by_id(&self, id: u64) -> OptionAway<&GameObject> {
        OptionAway::from(self.game_objects.get(&id))
    }

    /// Retrieves a mutable reference to an existing GameObject by the GameObject's id.
    ///
    /// Returns Away if and only if the GameObject exists but is currently loaned out
    /// in another borrow.
    ///
    /// Hint: That borrow is most likely in the game_object parameter of the receive_message
    /// function.
    pub fn game_object_by_id_mut(&mut self, id: u64) -> OptionAway<&mut GameObject> {
        OptionAway::from(self.game_objects.get_mut(&id))
    }

    /// Loads a texture and returns it for use.
    ///
    /// texture_name is the file name of the texture relative to assets/textures
    /// (assets\textures on Windows)
    pub fn fetch_texture(&mut self, texture_name: &str) -> Result<Texture, String> {
        let mut texture_path = PathBuf::from("assets");
        texture_path.push("textures");
        texture_path.push(texture_name);
        let sdl_texture = self.renderer.load_texture(texture_path.as_path())?;
        let mut nitro_texture = Texture::new();
        texture::set_raw(&mut nitro_texture, sdl_texture);
        Ok(nitro_texture)
    }
}

// This function will never return 0.  0 can now be used as a null value.
pub fn next_game_object_id(app: &mut App) -> u64 {
    app.next_game_object_id += 1;
    app.next_game_object_id
}
<|endoftext|>


<commit_before>extern crate exonum;
extern crate exonum_time;
#[macro_use]
extern crate exonum_testkit;
#[macro_use]
extern crate pretty_assertions;

use std::time::{SystemTime, Duration};

use exonum::helpers::{Height, ValidatorId};
use exonum::crypto::gen_keypair;

use exonum_time::TimeSchema;
use exonum_time::{TimeService, TxTime, Time};
use exonum_testkit::{TestKitBuilder, TestNode};

#[test]
fn test_exonum_time_service() {
    let mut testkit = TestKitBuilder::validator()
        .with_validators(3)
        .with_service(TimeService::new())
        .create();

    let validators = testkit.network().validators().to_vec();

    // Validators time, that is saved in storage, look like this:
    // number | 0    | 1    | 2    |
    // time   | None | None | None |
    //
    // Time, that is saved in storage, is None

    let validators_time_test: Vec<Option<Time>> = vec![None, None, None];

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);
    let validators_time_storage = schema.validators_time();

    for (num, validator) in validators.iter().enumerate() {
        let pub_key = &validator.public_keys().service_key;
        assert_eq!(
            validators_time_test[num],
            validators_time_storage.get(pub_key)
        );
    }
    assert_eq!(schema.time().get(), None);

    // Add first transaction 'tx0' from first validator with time 'time0'.
    // After that validators time look like this:
    // number | 0       | 0    | 0    |
    // time   | 'time0' | None | None |
    //
    // Time, that is saved in storage, does not change.

    let time0 = SystemTime::now();
    let tx0 = {
        let (pub_key, sec_key) = validators[0].service_keypair();
        TxTime::new(time0, pub_key, sec_key)
    };
    testkit.create_block_with_transactions(txvec![tx0.clone()]);

    let validators_time_test: Vec<Option<Time>> = vec![Some(Time::new(time0)), None, None];
    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);
    let validators_time_storage = schema.validators_time();

    for (num, validator) in validators.iter().enumerate() {
        let pub_key = &validator.public_keys().service_key;
        assert_eq!(
            validators_time_test[num],
            validators_time_storage.get(pub_key)
        );
    }
    assert_eq!(schema.time().get(), None);

    // Add second transaction 'tx1' from second validator with time 'time1' = 'time0' + 10 sec.
    // After that validators time look like this:
    // number | 0       | 1       | 2    |
    // time   | 'time0' | 'time1' | None |
    //
    // In sorted order: 'time1' >= 'time0'
    // Time, that is saved in storage, will have the value 'time0'.

    let time1 = time0 + Duration::new(10, 0);
    let tx1 = {
        let (pub_key, sec_key) = validators[1].service_keypair();
        TxTime::new(time1, pub_key, sec_key)
    };
    testkit.create_block_with_transactions(txvec![tx1.clone()]);

    let validators_time_test: Vec<Option<Time>> =
        vec![Some(Time::new(time0)), Some(Time::new(time1)), None];

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);
    let validators_time_storage = schema.validators_time();

    for (num, validator) in validators.iter().enumerate() {
        let pub_key = &validator.public_keys().service_key;
        assert_eq!(
            validators_time_test[num],
            validators_time_storage.get(pub_key)
        );
    }
    assert_eq!(schema.time().get(), Some(Time::new(time0)));

    // Add third transaction 'tx2' from third validator with time 'time2' = 'time1' + 10 sec.
    // After that validators time look like this:
    // number | 0       | 1       | 2       |
    // time   | 'time0' | 'time1' | 'time2' |
    //
    // In sorted order: 'time2' >= 'time1' >= 'time0'
    // Time, that is saved in storage, will have the value 'time1'.

    let time2 = time1 + Duration::new(10, 0);
    let tx2 = {
        let (pub_key, sec_key) = validators[2].service_keypair();
        TxTime::new(time2, pub_key, sec_key)
    };
    testkit.create_block_with_transactions(txvec![tx2.clone()]);

    let validators_time_test: Vec<Option<Time>> = vec![
        Some(Time::new(time0)),
        Some(Time::new(time1)),
        Some(Time::new(time2)),
    ];

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);
    let validators_time_storage = schema.validators_time();

    for (num, validator) in validators.iter().enumerate() {
        let pub_key = &validator.public_keys().service_key;
        assert_eq!(
            validators_time_test[num],
            validators_time_storage.get(pub_key)
        );
    }
    assert_eq!(schema.time().get(), Some(Time::new(time1)));
}

#[test]
fn test_selected_time_less_than_time_in_storage() {
    let mut testkit = TestKitBuilder::validator()
        .with_validators(1)
        .with_service(TimeService::new())
        .create();

    let validators = testkit.network().validators().to_vec();

    let (pub_key_0, _) = validators[0].service_keypair();

    let cfg_change_height = Height(5);
    let new_cfg = {
        let mut cfg = testkit.configuration_change_proposal();
        cfg.set_validators(vec![TestNode::new_validator(ValidatorId(0))]);
        cfg.set_actual_from(cfg_change_height);
        cfg
    };
    testkit.commit_configuration_change(new_cfg);
    testkit.create_blocks_until(cfg_change_height.previous());

    let validators = testkit.network().validators().to_vec();
    let (pub_key_1, sec_key_1) = validators[0].service_keypair();

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);

    assert!(schema.time().get().is_some());
    assert!(schema.validators_time().get(pub_key_0).is_some());
    assert!(schema.validators_time().get(pub_key_1).is_none());
    assert_eq!(schema.time().get(), schema.validators_time().get(pub_key_0));

    if let Some(time_in_storage) = schema.time().get() {
        let time_tx = time_in_storage.time() - Duration::new(10, 0);
        let tx = {
            TxTime::new(time_tx, pub_key_1, sec_key_1)
        };
        testkit.create_block_with_transactions(txvec![tx.clone()]);
    }

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);
    assert!(schema.time().get().is_some());
    assert!(schema.validators_time().get(pub_key_0).is_some());
    assert!(schema.validators_time().get(pub_key_1).is_some());
    assert_eq!(schema.time().get(), schema.validators_time().get(pub_key_0));
}

#[test]
fn test_creating_transaction_is_not_validator() {
    let mut testkit = TestKitBuilder::validator()
        .with_validators(1)
        .with_service(TimeService::new())
        .create();

    let (pub_key, sec_key) = gen_keypair();
    let tx = TxTime::new(SystemTime::now(), &pub_key, &sec_key);
    testkit.create_block_with_transactions(txvec![tx.clone()]);

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);
    assert!(schema.time().get().is_none());
    assert!(schema.validators_time().get(&pub_key).is_none());
}

#[test]
fn test_transaction_time_less_than_validator_time_in_storage() {
    let mut testkit = TestKitBuilder::validator()
        .with_validators(1)
        .with_service(TimeService::new())
        .create();

    let validator = &testkit.network().validators().to_vec()[0];
    let (pub_key, sec_key) = validator.service_keypair();

    let time0 = SystemTime::now();
    let tx0 = TxTime::new(time0, pub_key, sec_key);

    testkit.create_block_with_transactions(txvec![tx0.clone()]);

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);

    assert_eq!(schema.time().get(), Some(Time::new(time0)));
    assert_eq!(
        schema.validators_time().get(pub_key),
        Some(Time::new(time0))
    );

    let time1 = time0 - Duration::new(10, 0);
    let tx1 = TxTime::new(time1, pub_key, sec_key);

    testkit.create_block_with_transactions(txvec![tx1.clone()]);

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);

    assert_eq!(schema.time().get(), Some(Time::new(time0)));
    assert_eq!(
        schema.validators_time().get(pub_key),
        Some(Time::new(time0))
    );
}
<commit_msg>Add test for mock provider.<commit_after>extern crate exonum;
extern crate exonum_time;
#[macro_use]
extern crate exonum_testkit;
#[macro_use]
extern crate pretty_assertions;

use std::time::{SystemTime, Duration, UNIX_EPOCH};

use exonum::helpers::{Height, ValidatorId};
use exonum::crypto::gen_keypair;

use exonum_time::{TimeService, TimeSchema, TxTime, Time, TimeProvider};
use exonum_testkit::{TestKitBuilder, TestNode};

#[test]
fn test_exonum_time_service() {
    let mut testkit = TestKitBuilder::validator()
        .with_validators(3)
        .with_service(TimeService::new())
        .create();

    let validators = testkit.network().validators().to_vec();

    // Validators time, that is saved in storage, look like this:
    // number | 0    | 1    | 2    |
    // time   | None | None | None |
    //
    // Time, that is saved in storage, is None

    let validators_time_test: Vec<Option<Time>> = vec![None, None, None];

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);
    let validators_time_storage = schema.validators_time();

    for (num, validator) in validators.iter().enumerate() {
        let pub_key = &validator.public_keys().service_key;
        assert_eq!(
            validators_time_test[num],
            validators_time_storage.get(pub_key)
        );
    }
    assert_eq!(schema.time().get(), None);

    // Add first transaction 'tx0' from first validator with time 'time0'.
    // After that validators time look like this:
    // number | 0       | 0    | 0    |
    // time   | 'time0' | None | None |
    //
    // Time, that is saved in storage, does not change.

    let time0 = SystemTime::now();
    let tx0 = {
        let (pub_key, sec_key) = validators[0].service_keypair();
        TxTime::new(time0, pub_key, sec_key)
    };
    testkit.create_block_with_transactions(txvec![tx0.clone()]);

    let validators_time_test: Vec<Option<Time>> = vec![Some(Time::new(time0)), None, None];
    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);
    let validators_time_storage = schema.validators_time();

    for (num, validator) in validators.iter().enumerate() {
        let pub_key = &validator.public_keys().service_key;
        assert_eq!(
            validators_time_test[num],
            validators_time_storage.get(pub_key)
        );
    }
    assert_eq!(schema.time().get(), None);

    // Add second transaction 'tx1' from second validator with time 'time1' = 'time0' + 10 sec.
    // After that validators time look like this:
    // number | 0       | 1       | 2    |
    // time   | 'time0' | 'time1' | None |
    //
    // In sorted order: 'time1' >= 'time0'
    // Time, that is saved in storage, will have the value 'time0'.

    let time1 = time0 + Duration::new(10, 0);
    let tx1 = {
        let (pub_key, sec_key) = validators[1].service_keypair();
        TxTime::new(time1, pub_key, sec_key)
    };
    testkit.create_block_with_transactions(txvec![tx1.clone()]);

    let validators_time_test: Vec<Option<Time>> =
        vec![Some(Time::new(time0)), Some(Time::new(time1)), None];

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);
    let validators_time_storage = schema.validators_time();

    for (num, validator) in validators.iter().enumerate() {
        let pub_key = &validator.public_keys().service_key;
        assert_eq!(
            validators_time_test[num],
            validators_time_storage.get(pub_key)
        );
    }
    assert_eq!(schema.time().get(), Some(Time::new(time0)));

    // Add third transaction 'tx2' from third validator with time 'time2' = 'time1' + 10 sec.
    // After that validators time look like this:
    // number | 0       | 1       | 2       |
    // time   | 'time0' | 'time1' | 'time2' |
    //
    // In sorted order: 'time2' >= 'time1' >= 'time0'
    // Time, that is saved in storage, will have the value 'time1'.

    let time2 = time1 + Duration::new(10, 0);
    let tx2 = {
        let (pub_key, sec_key) = validators[2].service_keypair();
        TxTime::new(time2, pub_key, sec_key)
    };
    testkit.create_block_with_transactions(txvec![tx2.clone()]);

    let validators_time_test: Vec<Option<Time>> = vec![
        Some(Time::new(time0)),
        Some(Time::new(time1)),
        Some(Time::new(time2)),
    ];

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);
    let validators_time_storage = schema.validators_time();

    for (num, validator) in validators.iter().enumerate() {
        let pub_key = &validator.public_keys().service_key;
        assert_eq!(
            validators_time_test[num],
            validators_time_storage.get(pub_key)
        );
    }
    assert_eq!(schema.time().get(), Some(Time::new(time1)));
}

// A struct that provides the node with the current time.
#[derive(Debug)]
struct MyTimeProvider;
impl TimeProvider for MyTimeProvider {
    fn current_time(&self) -> SystemTime {
        UNIX_EPOCH
    }
}

#[test]
fn test_mock_provider() {
    // Create a simple testkit network.
    let mut testkit = TestKitBuilder::validator()
        .with_service(TimeService::with_provider(
            Box::new(MyTimeProvider) as Box<TimeProvider>,
        ))
        .create();

    // Get the validator public key.
    let validator_public_key = &testkit.network().validators().to_vec()[0]
        .public_keys()
        .service_key;

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);

    // Check that the blockchain does not contain time.
    assert_eq!(schema.time().get(), None);
    // Check that the time for the validator is unknown.
    assert_eq!(schema.validators_time().get(validator_public_key), None);

    // Create two blocks.
    testkit.create_blocks_until(Height(2));

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);

    // Check that the time in the blockchain and for the validator has been updated.
    assert_eq!(schema.time().get(), Some(Time::new(UNIX_EPOCH)));
    assert_eq!(
        schema.validators_time().get(validator_public_key),
        Some(Time::new(UNIX_EPOCH))
    );
}

#[test]
fn test_selected_time_less_than_time_in_storage() {
    let mut testkit = TestKitBuilder::validator()
        .with_validators(1)
        .with_service(TimeService::new())
        .create();

    let validators = testkit.network().validators().to_vec();

    let (pub_key_0, _) = validators[0].service_keypair();

    let cfg_change_height = Height(5);
    let new_cfg = {
        let mut cfg = testkit.configuration_change_proposal();
        cfg.set_validators(vec![TestNode::new_validator(ValidatorId(0))]);
        cfg.set_actual_from(cfg_change_height);
        cfg
    };
    testkit.commit_configuration_change(new_cfg);
    testkit.create_blocks_until(cfg_change_height.previous());

    let validators = testkit.network().validators().to_vec();
    let (pub_key_1, sec_key_1) = validators[0].service_keypair();

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);

    assert!(schema.time().get().is_some());
    assert!(schema.validators_time().get(pub_key_0).is_some());
    assert!(schema.validators_time().get(pub_key_1).is_none());
    assert_eq!(schema.time().get(), schema.validators_time().get(pub_key_0));

    if let Some(time_in_storage) = schema.time().get() {
        let time_tx = time_in_storage.time() - Duration::new(10, 0);
        let tx = {
            TxTime::new(time_tx, pub_key_1, sec_key_1)
        };
        testkit.create_block_with_transactions(txvec![tx.clone()]);
    }

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);
    assert!(schema.time().get().is_some());
    assert!(schema.validators_time().get(pub_key_0).is_some());
    assert!(schema.validators_time().get(pub_key_1).is_some());
    assert_eq!(schema.time().get(), schema.validators_time().get(pub_key_0));
}

#[test]
fn test_creating_transaction_is_not_validator() {
    let mut testkit = TestKitBuilder::validator()
        .with_validators(1)
        .with_service(TimeService::new())
        .create();

    let (pub_key, sec_key) = gen_keypair();
    let tx = TxTime::new(SystemTime::now(), &pub_key, &sec_key);
    testkit.create_block_with_transactions(txvec![tx.clone()]);

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);
    assert!(schema.time().get().is_none());
    assert!(schema.validators_time().get(&pub_key).is_none());
}

#[test]
fn test_transaction_time_less_than_validator_time_in_storage() {
    let mut testkit = TestKitBuilder::validator()
        .with_validators(1)
        .with_service(TimeService::new())
        .create();

    let validator = &testkit.network().validators().to_vec()[0];
    let (pub_key, sec_key) = validator.service_keypair();

    let time0 = SystemTime::now();
    let tx0 = TxTime::new(time0, pub_key, sec_key);

    testkit.create_block_with_transactions(txvec![tx0.clone()]);

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);

    assert_eq!(schema.time().get(), Some(Time::new(time0)));
    assert_eq!(
        schema.validators_time().get(pub_key),
        Some(Time::new(time0))
    );

    let time1 = time0 - Duration::new(10, 0);
    let tx1 = TxTime::new(time1, pub_key, sec_key);

    testkit.create_block_with_transactions(txvec![tx1.clone()]);

    let snapshot = testkit.snapshot();
    let schema = TimeSchema::new(snapshot);

    assert_eq!(schema.time().get(), Some(Time::new(time0)));
    assert_eq!(
        schema.validators_time().get(pub_key),
        Some(Time::new(time0))
    );
}
<|endoftext|>


<commit_before>//! Song analysis module.
//!
//! Use features-extraction functions from other modules
//! e.g. tempo features, spectral features, etc to build an
//! Analysis of a Song
#[cfg(feature = "aubio-lib")]
extern crate aubio_lib;

extern crate crossbeam;
extern crate ndarray;
extern crate ndarray_npy;

use std::f32::consts::PI as PI;

use ndarray::{arr1, s, Array, Array2};
use rustfft::num_complex::Complex;
use rustfft::num_traits::Zero;
use rustfft::FFTplanner;

use crate::chroma::ChromaDesc;
use crate::decode::decode_song;
use crate::misc::LoudnessDesc;
use crate::temporal::BPMDesc;
use crate::timbral::{SpectralDesc, ZeroCrossingRateDesc};
use crate::{Analysis, Song};
use crossbeam::thread;

pub fn decode_and_analyze(path: &str) -> Result<Song, String> {
    // TODO error handling here
    let mut song = decode_song(&path)?;

    song.analysis = analyze(&song);
    Ok(song)
}

fn reflect_pad(array: &[f32], pad: usize) -> Vec<f32> {
    let mut prefix = array[1..=pad].iter().rev().copied().collect::<Vec<f32>>();
    let suffix = array[(array.len() - 2) - pad + 1..array.len() - 1]
        .iter()
        .rev()
        .copied()
        .collect::<Vec<f32>>();
    prefix.extend(array);
    prefix.extend(suffix);
    prefix
}

pub fn stft(signal: &[f32], window_length: usize, hop_length: usize) -> Array2<f64> {
    let mut stft = Array2::zeros((
        window_length / 2 + 1,
        (signal.len() as f32 / hop_length as f32).ceil() as usize,
    ));
    let signal = reflect_pad(&signal, window_length / 2);

    // TODO actually have it constant - no need to compute it everytime
    // Periodic, so window_size + 1
    let mut hann_window = Array::zeros(window_length + 1);
    for n in 0..window_length {
        hann_window[[n]] = 0.5 - 0.5 * f32::cos(2. * n as f32 * PI / (window_length as f32));
    }
    hann_window = hann_window.slice_move(s![0..window_length]);
    let mut output_window = Array::from_elem(window_length, Complex::zero());
    let mut planner = FFTplanner::new(false);
    let fft = planner.plan_fft(window_length);

    for (window, mut stft_col) in signal
        .windows(window_length)
        .step_by(hop_length)
        .zip(stft.gencolumns_mut())
    {
        let mut signal = (arr1(&window) * &hann_window).mapv(|x| Complex::new(x as f32, 0.));
        fft.process(
            signal.as_slice_mut().unwrap(),
            output_window.as_slice_mut().unwrap(),
        );
        stft_col.assign(
            &output_window
                .slice(s![..window_length / 2 + 1])
                .mapv(|x| x.norm() as f64),
        );
    }
    stft
}

pub fn analyze(song: &Song) -> Analysis {
    let mut spectral_desc = SpectralDesc::new(song.sample_rate);
    let mut zcr_desc = ZeroCrossingRateDesc::default();
    let mut tempo_desc = BPMDesc::new(song.sample_rate);
    let mut loudness_desc = LoudnessDesc::default();

    thread::scope(|s| {
        let child = s.spawn(|_| {
            let mut chroma_desc = ChromaDesc::new(song.sample_rate, 12);
            chroma_desc.do_(&song.sample_array);
            chroma_desc.get_values()
        });

        // These descriptors can be streamed
        for i in 1..song.sample_array.len() {
            if (i % SpectralDesc::HOP_SIZE) == 0 {
                let beginning = (i / SpectralDesc::HOP_SIZE - 1) * SpectralDesc::HOP_SIZE;
                let end = i;
                spectral_desc.do_(&song.sample_array[beginning..end]);
                zcr_desc.do_(&song.sample_array[beginning..end]);
            }

            if (i % BPMDesc::HOP_SIZE) == 0 {
                let beginning = (i / BPMDesc::HOP_SIZE - 1) * BPMDesc::HOP_SIZE;
                let end = i;
                tempo_desc.do_(&song.sample_array[beginning..end]);
            }

            // Contiguous windows, so WINDOW_SIZE here
            if (i % LoudnessDesc::WINDOW_SIZE) == 0 {
                let beginning = (i / LoudnessDesc::WINDOW_SIZE - 1) * LoudnessDesc::WINDOW_SIZE;
                let end = i;
                loudness_desc.do_(&song.sample_array[beginning..end]);
            }
        }
        // Non-streaming approach for that one
        let (is_major, fifth) = child.join().unwrap();

        Analysis {
            tempo: tempo_desc.get_value(),
            spectral_centroid: spectral_desc.get_centroid(),
            zero_crossing_rate: zcr_desc.get_value(),
            spectral_rolloff: spectral_desc.get_rolloff(),
            spectral_flatness: spectral_desc.get_flatness(),
            loudness: loudness_desc.get_value(),
            is_major,
            fifth,
        }
    })
    .unwrap()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::decode::decode_song;
    use ndarray::{arr1, Array2};
    use ndarray_npy::ReadNpyExt;
    use std::f32::consts::PI;
    use std::fs::File;

    #[test]
    fn test_analyze() {
        let song = decode_song("data/s16_mono_22_5kHz.flac").unwrap();
        let expected_analysis = Analysis {
            tempo: 0.37860596,
            spectral_centroid: -0.75483,
            zero_crossing_rate: -0.85036564,
            spectral_rolloff: -0.6326486,
            spectral_flatness: -0.77610075,
            loudness: 0.27126348,
            is_major: -1.,
            fifth: (f32::cos(5. * PI / 3.), f32::sin(5. * PI / 3.)),
        };
        assert!(expected_analysis.approx_eq(&analyze(&song)));
    }

    #[test]
    fn test_compute_stft() {
        let file = File::open("data/librosa-stft.npy").unwrap();
        let expected_stft = Array2::<f32>::read_npy(file).unwrap().mapv(|x| x as f64);

        let song = decode_song("data/piano.flac").unwrap();

        let stft = stft(&song.sample_array, 2048, 512);

        assert!(!stft.is_empty() && !expected_stft.is_empty());
        for (expected, actual) in expected_stft.iter().zip(stft.iter()) {
            assert!(0.0001 > (expected - actual).abs());
        }
    }

    #[test]
    fn test_foo() {
        let array = arr1(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]);
        for w in array.windows(4).into_iter().step_by(2) {
            println!("{:?}", w);
        }
    }
}
<commit_msg>Run each analysis separately<commit_after>//! Song analysis module.
//!
//! Use features-extraction functions from other modules
//! e.g. tempo features, spectral features, etc to build an
//! Analysis of a Song
#[cfg(feature = "aubio-lib")]
extern crate aubio_lib;

extern crate crossbeam;
extern crate ndarray;
extern crate ndarray_npy;

use std::f32::consts::PI;

use ndarray::{arr1, s, Array, Array2};
use rustfft::num_complex::Complex;
use rustfft::num_traits::Zero;
use rustfft::FFTplanner;

use crate::chroma::ChromaDesc;
use crate::decode::decode_song;
use crate::misc::LoudnessDesc;
use crate::temporal::BPMDesc;
use crate::timbral::{SpectralDesc, ZeroCrossingRateDesc};
use crate::{Analysis, Song};
use crossbeam::thread;

pub fn decode_and_analyze(path: &str) -> Result<Song, String> {
    // TODO error handling here
    let mut song = decode_song(&path)?;

    song.analysis = analyze(&song);
    Ok(song)
}

fn reflect_pad(array: &[f32], pad: usize) -> Vec<f32> {
    let mut prefix = array[1..=pad].iter().rev().copied().collect::<Vec<f32>>();
    let suffix = array[(array.len() - 2) - pad + 1..array.len() - 1]
        .iter()
        .rev()
        .copied()
        .collect::<Vec<f32>>();
    prefix.extend(array);
    prefix.extend(suffix);
    prefix
}

pub fn stft(signal: &[f32], window_length: usize, hop_length: usize) -> Array2<f64> {
    let mut stft = Array2::zeros((
        window_length / 2 + 1,
        (signal.len() as f32 / hop_length as f32).ceil() as usize,
    ));
    let signal = reflect_pad(&signal, window_length / 2);

    // TODO actually have it constant - no need to compute it everytime
    // Periodic, so window_size + 1
    let mut hann_window = Array::zeros(window_length + 1);
    for n in 0..window_length {
        hann_window[[n]] = 0.5 - 0.5 * f32::cos(2. * n as f32 * PI / (window_length as f32));
    }
    hann_window = hann_window.slice_move(s![0..window_length]);
    let mut output_window = Array::from_elem(window_length, Complex::zero());
    let mut planner = FFTplanner::new(false);
    let fft = planner.plan_fft(window_length);

    for (window, mut stft_col) in signal
        .windows(window_length)
        .step_by(hop_length)
        .zip(stft.gencolumns_mut())
    {
        let mut signal = (arr1(&window) * &hann_window).mapv(|x| Complex::new(x as f32, 0.));
        fft.process(
            signal.as_slice_mut().unwrap(),
            output_window.as_slice_mut().unwrap(),
        );
        stft_col.assign(
            &output_window
                .slice(s![..window_length / 2 + 1])
                .mapv(|x| x.norm() as f64),
        );
    }
    stft
}

pub fn analyze(song: &Song) -> Analysis {
    thread::scope(|s| {
        let child_chroma = s.spawn(|_| {
            let mut chroma_desc = ChromaDesc::new(song.sample_rate, 12);
            chroma_desc.do_(&song.sample_array);
            chroma_desc.get_values()
        });

        // These descriptors can be streamed
        let child_timbral = s.spawn(|_| {
            let mut spectral_desc = SpectralDesc::new(song.sample_rate);
            let mut zcr_desc = ZeroCrossingRateDesc::default();
            let windows = song
                .sample_array
                .windows(SpectralDesc::WINDOW_SIZE)
                .step_by(SpectralDesc::HOP_SIZE);
            for window in windows {
                spectral_desc.do_(&window);
                zcr_desc.do_(&window);
            }
            let centroid = spectral_desc.get_centroid();
            let rolloff = spectral_desc.get_rolloff();
            let flatness = spectral_desc.get_flatness();
            let zcr = zcr_desc.get_value();
            (centroid, rolloff, flatness, zcr)
        });

        let child_tempo = s.spawn(|_| {
            let mut tempo_desc = BPMDesc::new(song.sample_rate);
            let windows = song
                .sample_array
                .windows(BPMDesc::WINDOW_SIZE)
                .step_by(BPMDesc::HOP_SIZE);

            for window in windows {
                tempo_desc.do_(&window);
            }
            tempo_desc.get_value()
        });

        let child_loudness = s.spawn(|_| {
            let mut loudness_desc = LoudnessDesc::default();
            let windows = song
                .sample_array
                .chunks(LoudnessDesc::WINDOW_SIZE);

            for window in windows {
                loudness_desc.do_(&window);
            }
            loudness_desc.get_value()
        });

        // Non-streaming approach for that one
        let (is_major, fifth) = child_chroma.join().unwrap();
        let (centroid, rolloff, flatness, zcr) = child_timbral.join().unwrap();
        let tempo = child_tempo.join().unwrap();
        let loudness = child_loudness.join().unwrap();

        Analysis {
            tempo,
            spectral_centroid: centroid,
            zero_crossing_rate: zcr,
            spectral_rolloff: rolloff,
            spectral_flatness: flatness,
            loudness,
            is_major,
            fifth,
        }
    })
    .unwrap()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::decode::decode_song;
    use ndarray::{arr1, Array2};
    use ndarray_npy::ReadNpyExt;
    use std::f32::consts::PI;
    use std::fs::File;

    #[test]
    fn test_analyze() {
        let song = decode_song("data/s16_mono_22_5kHz.flac").unwrap();
        let expected_analysis = Analysis {
            tempo: 0.37860596,
            spectral_centroid: -0.75483,
            zero_crossing_rate: -0.85036564,
            spectral_rolloff: -0.6326486,
            spectral_flatness: -0.77610075,
            loudness: 0.27126348,
            is_major: -1.,
            fifth: (f32::cos(5. * PI / 3.), f32::sin(5. * PI / 3.)),
        };
        assert!(expected_analysis.approx_eq(&analyze(&song)));
    }

    #[test]
    fn test_compute_stft() {
        let file = File::open("data/librosa-stft.npy").unwrap();
        let expected_stft = Array2::<f32>::read_npy(file).unwrap().mapv(|x| x as f64);

        let song = decode_song("data/piano.flac").unwrap();

        let stft = stft(&song.sample_array, 2048, 512);

        assert!(!stft.is_empty() && !expected_stft.is_empty());
        for (expected, actual) in expected_stft.iter().zip(stft.iter()) {
            assert!(0.0001 > (expected - actual).abs());
        }
    }

    #[test]
    fn test_foo() {
        let array = arr1(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]);
        for w in array.windows(4).into_iter().step_by(2) {
            println!("{:?}", w);
        }
    }
}
<|endoftext|>


<commit_before>use std::{collections::HashMap, fs::{
        create_dir_all,
        File,
    }, io::{BufReader, BufWriter, Write}, iter::FromIterator, path::{
        Path,
        PathBuf,
    }};

use anyhow::Result;
use chrono::{
    DateTime,
    Utc,
};
use itertools::Itertools;

use serde_json::Deserializer;
use sorted_iter::{
    assume::*,
    SortedPairIterator,
};

use rnc_core::grouper::Grouped;

use crate::{
    genes::{
        gene::Gene,
        gene_member::GeneMember,
        region::{
            Region,
            RegionGrouper,
        },
    },
    sequences::normalized::Normalized,
};

pub fn assembly_writer(base: &Path, assembly: &str) -> Result<BufWriter<File>> {
    let assembly_name = assembly.replace(".", "_");
    let mut path = PathBuf::from(base);
    path.push(assembly_name);
    path.set_extension("json");
    Ok(BufWriter::new(File::create(path)?))
}

pub fn write_split_selected(locus_path: &Path, sequence_file: &Path, output: &Path) -> Result<()> {
    create_dir_all(output)?;

    let locus_reader = BufReader::new(File::open(locus_path)?);
    let locus = Deserializer::from_reader(locus_reader).into_iter::<Grouped<Region>>();
    let locus = locus.map(|r| r.unwrap());
    let locus = locus.filter(|l| !l.is_empty());
    let locus = RegionGrouper::new(locus);
    let locus = locus.map(|l| (*l.id(), l)).assume_sorted_by_key();

    let sequence_reader = BufReader::new(File::open(sequence_file)?);
    let sequences = Deserializer::from_reader(sequence_reader).into_iter::<Normalized>();
    let sequences = sequences.map(|r| r.unwrap());
    let sequences = sequences.map(|n| (*n.id(), n)).assume_sorted_by_key();

    let mut writers = HashMap::new();
    let merged = locus.join(sequences);

    for (_id, pair) in merged {
        let (seq_locus, sequence) = pair;
        for (assembly_id, regions) in seq_locus.regions_by_assembly() {
            let mut writer = writers
                .entry(assembly_id.clone())
                .or_insert_with(|| assembly_writer(&output, &assembly_id).unwrap());
            for region in regions {
                let member = GeneMember::new(region.clone(), sequence.clone());
                serde_json::to_writer(&mut writer, &member)?;
                writeln!(&mut writer)?;
            }
            writer.flush()?;
        }
    }

    for (_assembly, writer) in writers.iter_mut() {
        writer.flush()?;
    }

    Ok(())
}

pub fn write_merged_members(member_file: &Path, output: &Path) -> Result<()> {
    let member_reader = BufReader::new(File::open(member_file)?);
    let members = Deserializer::from_reader(member_reader).into_iter::<GeneMember>();
    let members = members.map(|r| r.unwrap());
    let mut members: Vec<GeneMember> = members.collect();
    members.sort_by_key(|m| m.locus_id());
    let grouped = members.into_iter().group_by(|l| l.locus_id());

    let mut writer = BufWriter::new(File::create(&output)?);
    for (_locus_id, locus) in &grouped {
        let gene = Gene::from_iter(locus);
        serde_json::to_writer(&mut writer, &gene)?;
        writeln!(&mut writer)?;
    }
    writer.flush()?;

    Ok(())
}

pub fn write_search_files(gene_file: &Path, xml_output: &Path, count_output: &Path) -> Result<()> {

    let gene_reader = BufReader::new(File::open(gene_file)?);
    let genes = Deserializer::from_reader(gene_reader).into_iter::<Gene>();

    let mut xml_writer = BufWriter::new(File::create(&xml_output)?);
    let mut count = 0;
    let now: DateTime<Utc> = Utc::now();
    let date = now.format("%d/%m/%Y");
    writeln!(&mut xml_writer, "<database><name>RNAcentral</name><description></description><release>1.0</release><release_date>{}</release_date><entries>", &date)?;
    for gene in genes {
        let mut gene = gene?;
        gene.fill_default_name();
        let xml = gene.as_search();
        quick_xml::se::to_writer(&mut xml_writer, &xml)?;
        writeln!(&mut xml_writer)?;
        count += 1;
    }

    writeln!(&mut xml_writer, "</entries>")?;
    writeln!(&mut xml_writer, "<entry_count>{}</entry_count></database>", &count)?;
    xml_writer.flush()?;

    let mut count_writer = BufWriter::new(File::create(&count_output)?);
    write!(&mut count_writer, "{}", &count)?;
    count_writer.flush()?;

    Ok(())
}
<commit_msg>Add a hack to fix \b escaping<commit_after>use std::{collections::HashMap, fs::{
        create_dir_all,
        File,
    }, io::{BufReader, BufWriter, Write}, iter::FromIterator, path::{
        Path,
        PathBuf,
    }};

use anyhow::Result;
use chrono::{
    DateTime,
    Utc,
};
use itertools::Itertools;

use serde_json::Deserializer;
use sorted_iter::{
    assume::*,
    SortedPairIterator,
};

use rnc_core::grouper::Grouped;

use crate::{
    genes::{
        gene::Gene,
        gene_member::GeneMember,
        region::{
            Region,
            RegionGrouper,
        },
    },
    sequences::normalized::Normalized,
};

pub fn assembly_writer(base: &Path, assembly: &str) -> Result<BufWriter<File>> {
    let assembly_name = assembly.replace(".", "_");
    let mut path = PathBuf::from(base);
    path.push(assembly_name);
    path.set_extension("json");
    Ok(BufWriter::new(File::create(path)?))
}

pub fn write_split_selected(locus_path: &Path, sequence_file: &Path, output: &Path) -> Result<()> {
    create_dir_all(output)?;

    let locus_reader = BufReader::new(File::open(locus_path)?);
    let locus = Deserializer::from_reader(locus_reader).into_iter::<Grouped<Region>>();
    let locus = locus.map(|r| r.unwrap());
    let locus = locus.filter(|l| !l.is_empty());
    let locus = RegionGrouper::new(locus);
    let locus = locus.map(|l| (*l.id(), l)).assume_sorted_by_key();

    let sequence_reader = BufReader::new(File::open(sequence_file)?);
    let sequences = Deserializer::from_reader(sequence_reader).into_iter::<Normalized>();
    let sequences = sequences.map(|r| r.unwrap());
    let sequences = sequences.map(|n| (*n.id(), n)).assume_sorted_by_key();

    let mut writers = HashMap::new();
    let merged = locus.join(sequences);

    for (_id, pair) in merged {
        let (seq_locus, sequence) = pair;
        for (assembly_id, regions) in seq_locus.regions_by_assembly() {
            let mut writer = writers
                .entry(assembly_id.clone())
                .or_insert_with(|| assembly_writer(&output, &assembly_id).unwrap());
            for region in regions {
                let member = GeneMember::new(region.clone(), sequence.clone());
                serde_json::to_writer(&mut writer, &member)?;
                writeln!(&mut writer)?;
            }
            writer.flush()?;
        }
    }

    for (_assembly, writer) in writers.iter_mut() {
        writer.flush()?;
    }

    Ok(())
}

pub fn write_merged_members(member_file: &Path, output: &Path) -> Result<()> {
    let member_reader = BufReader::new(File::open(member_file)?);
    let members = Deserializer::from_reader(member_reader).into_iter::<GeneMember>();
    let members = members.map(|r| r.unwrap());
    let mut members: Vec<GeneMember> = members.collect();
    members.sort_by_key(|m| m.locus_id());
    let grouped = members.into_iter().group_by(|l| l.locus_id());

    let mut writer = BufWriter::new(File::create(&output)?);
    for (_locus_id, locus) in &grouped {
        let gene = Gene::from_iter(locus);
        serde_json::to_writer(&mut writer, &gene)?;
        writeln!(&mut writer)?;
    }
    writer.flush()?;

    Ok(())
}

pub fn write_search_files(gene_file: &Path, xml_output: &Path, count_output: &Path) -> Result<()> {

    let gene_reader = BufReader::new(File::open(gene_file)?);
    let genes = Deserializer::from_reader(gene_reader).into_iter::<Gene>();

    let mut xml_writer = BufWriter::new(File::create(&xml_output)?);
    let mut count = 0;
    let now: DateTime<Utc> = Utc::now();
    let date = now.format("%d/%m/%Y");
    writeln!(&mut xml_writer, "<database><name>RNAcentral</name><description></description><release>1.0</release><release_date>{}</release_date><entries>", &date)?;
    for gene in genes {
        let mut gene = gene?;
        gene.fill_default_name();
        let xml = gene.as_search();
        let element = quick_xml::se::to_string(&xml)?;
        let element = element.replace(r"", "\\b");
        writeln!(&mut xml_writer, "{}", element)?;
        count += 1;
    }

    writeln!(&mut xml_writer, "</entries>")?;
    writeln!(&mut xml_writer, "<entry_count>{}</entry_count></database>", &count)?;
    xml_writer.flush()?;

    let mut count_writer = BufWriter::new(File::create(&count_output)?);
    write!(&mut count_writer, "{}", &count)?;
    count_writer.flush()?;

    Ok(())
}
<|endoftext|>


<commit_before>mod binding_files;
mod build_tables;
mod char_tree;
mod dedup;
mod grammars;
mod nfa;
mod node_types;
pub mod parse_grammar;
mod prepare_grammar;
mod render;
mod rules;
mod tables;

use self::build_tables::build_tables;
use self::grammars::{InlinedProductionMap, LexicalGrammar, SyntaxGrammar};
use self::parse_grammar::parse_grammar;
use self::prepare_grammar::prepare_grammar;
use self::render::render_c_code;
use self::rules::AliasMap;
use anyhow::{anyhow, Context, Result};
use lazy_static::lazy_static;
use regex::{Regex, RegexBuilder};
use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

lazy_static! {
    static ref JSON_COMMENT_REGEX: Regex = RegexBuilder::new("^\\s*//.*")
        .multi_line(true)
        .build()
        .unwrap();
}

struct GeneratedParser {
    c_code: String,
    node_types_json: String,
}

pub fn generate_parser_in_directory(
    repo_path: &PathBuf,
    grammar_path: Option<&str>,
    next_abi: bool,
    generate_bindings: bool,
    report_symbol_name: Option<&str>,
) -> Result<()> {
    let src_path = repo_path.join("src");
    let header_path = src_path.join("tree_sitter");

    // Ensure that the output directories exist.
    fs::create_dir_all(&src_path)?;
    fs::create_dir_all(&header_path)?;

    // Read the grammar.json.
    let grammar_json;
    match grammar_path {
        Some(path) => {
            grammar_json = load_grammar_file(path.as_ref())?;
        }
        None => {
            let grammar_js_path = grammar_path.map_or(repo_path.join("grammar.js"), |s| s.into());
            grammar_json = load_grammar_file(&grammar_js_path)?;
            fs::write(&src_path.join("grammar.json"), &grammar_json)?;
        }
    }

    // Parse and preprocess the grammar.
    let input_grammar = parse_grammar(&grammar_json)?;
    let (syntax_grammar, lexical_grammar, inlines, simple_aliases) =
        prepare_grammar(&input_grammar)?;
    let language_name = input_grammar.name;

    // Generate the parser and related files.
    let GeneratedParser {
        c_code,
        node_types_json,
    } = generate_parser_for_grammar_with_opts(
        &language_name,
        syntax_grammar,
        lexical_grammar,
        inlines,
        simple_aliases,
        next_abi,
        report_symbol_name,
    )?;

    write_file(&src_path.join("parser.c"), c_code)?;
    write_file(&src_path.join("node-types.json"), node_types_json)?;

    if next_abi {
        write_file(&header_path.join("parser.h"), tree_sitter::PARSER_HEADER)?;
    }

    if generate_bindings {
        binding_files::generate_binding_files(&repo_path, &language_name)?;
    }

    Ok(())
}

pub fn generate_parser_for_grammar(grammar_json: &str) -> Result<(String, String)> {
    let grammar_json = JSON_COMMENT_REGEX.replace_all(grammar_json, "\n");
    let input_grammar = parse_grammar(&grammar_json)?;
    let (syntax_grammar, lexical_grammar, inlines, simple_aliases) =
        prepare_grammar(&input_grammar)?;
    let parser = generate_parser_for_grammar_with_opts(
        &input_grammar.name,
        syntax_grammar,
        lexical_grammar,
        inlines,
        simple_aliases,
        true,
        None,
    )?;
    Ok((input_grammar.name, parser.c_code))
}

fn generate_parser_for_grammar_with_opts(
    name: &String,
    syntax_grammar: SyntaxGrammar,
    lexical_grammar: LexicalGrammar,
    inlines: InlinedProductionMap,
    simple_aliases: AliasMap,
    next_abi: bool,
    report_symbol_name: Option<&str>,
) -> Result<GeneratedParser> {
    let variable_info =
        node_types::get_variable_info(&syntax_grammar, &lexical_grammar, &simple_aliases)?;
    let node_types_json = node_types::generate_node_types_json(
        &syntax_grammar,
        &lexical_grammar,
        &simple_aliases,
        &variable_info,
    );
    let (parse_table, main_lex_table, keyword_lex_table, keyword_capture_token) = build_tables(
        &syntax_grammar,
        &lexical_grammar,
        &simple_aliases,
        &variable_info,
        &inlines,
        report_symbol_name,
    )?;
    let c_code = render_c_code(
        name,
        parse_table,
        main_lex_table,
        keyword_lex_table,
        keyword_capture_token,
        syntax_grammar,
        lexical_grammar,
        simple_aliases,
        next_abi,
    );
    Ok(GeneratedParser {
        c_code,
        node_types_json: serde_json::to_string_pretty(&node_types_json).unwrap(),
    })
}

fn load_grammar_file(grammar_path: &Path) -> Result<String> {
    match grammar_path.extension().and_then(|e| e.to_str()) {
        Some("js") => Ok(load_js_grammar_file(grammar_path)?),
        Some("json") => Ok(fs::read_to_string(grammar_path)?),
        _ => Err(anyhow!(
            "Unknown grammar file extension: {:?}",
            grammar_path
        )),
    }
}

fn load_js_grammar_file(grammar_path: &Path) -> Result<String> {
    let mut node_process = Command::new("node")
        .env("TREE_SITTER_GRAMMAR_PATH", grammar_path)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .spawn()
        .expect("Failed to run `node`");

    let mut node_stdin = node_process
        .stdin
        .take()
        .expect("Failed to open stdin for node");
    let javascript_code = include_bytes!("./dsl.js");
    node_stdin
        .write(javascript_code)
        .expect("Failed to write to node's stdin");
    drop(node_stdin);
    let output = node_process
        .wait_with_output()
        .expect("Failed to read output from node");
    match output.status.code() {
        None => panic!("Node process was killed"),
        Some(0) => {}
        Some(code) => return Err(anyhow!("Node process exited with status {}", code)),
    }

    let mut result = String::from_utf8(output.stdout).expect("Got invalid UTF8 from node");
    result.push('\n');
    Ok(result)
}

fn write_file(path: &Path, body: impl AsRef<[u8]>) -> Result<()> {
    fs::write(path, body)
        .with_context(|| format!("Failed to write {:?}", path.file_name().unwrap()))
}
<commit_msg>fix(cli): Remove tree-sitter grammar ./... call limitation<commit_after>mod binding_files;
mod build_tables;
mod char_tree;
mod dedup;
mod grammars;
mod nfa;
mod node_types;
pub mod parse_grammar;
mod prepare_grammar;
mod render;
mod rules;
mod tables;

use self::build_tables::build_tables;
use self::grammars::{InlinedProductionMap, LexicalGrammar, SyntaxGrammar};
use self::parse_grammar::parse_grammar;
use self::prepare_grammar::prepare_grammar;
use self::render::render_c_code;
use self::rules::AliasMap;
use anyhow::{anyhow, Context, Result};
use lazy_static::lazy_static;
use regex::{Regex, RegexBuilder};
use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

lazy_static! {
    static ref JSON_COMMENT_REGEX: Regex = RegexBuilder::new("^\\s*//.*")
        .multi_line(true)
        .build()
        .unwrap();
}

struct GeneratedParser {
    c_code: String,
    node_types_json: String,
}

pub fn generate_parser_in_directory(
    repo_path: &PathBuf,
    grammar_path: Option<&str>,
    next_abi: bool,
    generate_bindings: bool,
    report_symbol_name: Option<&str>,
) -> Result<()> {
    let src_path = repo_path.join("src");
    let header_path = src_path.join("tree_sitter");

    // Ensure that the output directories exist.
    fs::create_dir_all(&src_path)?;
    fs::create_dir_all(&header_path)?;

    // Read the grammar.json.
    let grammar_json;
    match grammar_path {
        Some(path) => {
            grammar_json = load_grammar_file(path.as_ref())?;
        }
        None => {
            let grammar_js_path = grammar_path.map_or(repo_path.join("grammar.js"), |s| s.into());
            grammar_json = load_grammar_file(&grammar_js_path)?;
            fs::write(&src_path.join("grammar.json"), &grammar_json)?;
        }
    }

    // Parse and preprocess the grammar.
    let input_grammar = parse_grammar(&grammar_json)?;
    let (syntax_grammar, lexical_grammar, inlines, simple_aliases) =
        prepare_grammar(&input_grammar)?;
    let language_name = input_grammar.name;

    // Generate the parser and related files.
    let GeneratedParser {
        c_code,
        node_types_json,
    } = generate_parser_for_grammar_with_opts(
        &language_name,
        syntax_grammar,
        lexical_grammar,
        inlines,
        simple_aliases,
        next_abi,
        report_symbol_name,
    )?;

    write_file(&src_path.join("parser.c"), c_code)?;
    write_file(&src_path.join("node-types.json"), node_types_json)?;

    if next_abi {
        write_file(&header_path.join("parser.h"), tree_sitter::PARSER_HEADER)?;
    }

    if generate_bindings {
        binding_files::generate_binding_files(&repo_path, &language_name)?;
    }

    Ok(())
}

pub fn generate_parser_for_grammar(grammar_json: &str) -> Result<(String, String)> {
    let grammar_json = JSON_COMMENT_REGEX.replace_all(grammar_json, "\n");
    let input_grammar = parse_grammar(&grammar_json)?;
    let (syntax_grammar, lexical_grammar, inlines, simple_aliases) =
        prepare_grammar(&input_grammar)?;
    let parser = generate_parser_for_grammar_with_opts(
        &input_grammar.name,
        syntax_grammar,
        lexical_grammar,
        inlines,
        simple_aliases,
        true,
        None,
    )?;
    Ok((input_grammar.name, parser.c_code))
}

fn generate_parser_for_grammar_with_opts(
    name: &String,
    syntax_grammar: SyntaxGrammar,
    lexical_grammar: LexicalGrammar,
    inlines: InlinedProductionMap,
    simple_aliases: AliasMap,
    next_abi: bool,
    report_symbol_name: Option<&str>,
) -> Result<GeneratedParser> {
    let variable_info =
        node_types::get_variable_info(&syntax_grammar, &lexical_grammar, &simple_aliases)?;
    let node_types_json = node_types::generate_node_types_json(
        &syntax_grammar,
        &lexical_grammar,
        &simple_aliases,
        &variable_info,
    );
    let (parse_table, main_lex_table, keyword_lex_table, keyword_capture_token) = build_tables(
        &syntax_grammar,
        &lexical_grammar,
        &simple_aliases,
        &variable_info,
        &inlines,
        report_symbol_name,
    )?;
    let c_code = render_c_code(
        name,
        parse_table,
        main_lex_table,
        keyword_lex_table,
        keyword_capture_token,
        syntax_grammar,
        lexical_grammar,
        simple_aliases,
        next_abi,
    );
    Ok(GeneratedParser {
        c_code,
        node_types_json: serde_json::to_string_pretty(&node_types_json).unwrap(),
    })
}

fn load_grammar_file(grammar_path: &Path) -> Result<String> {
    match grammar_path.extension().and_then(|e| e.to_str()) {
        Some("js") => Ok(load_js_grammar_file(grammar_path)?),
        Some("json") => Ok(fs::read_to_string(grammar_path)?),
        _ => Err(anyhow!(
            "Unknown grammar file extension: {:?}",
            grammar_path
        )),
    }
}

fn load_js_grammar_file(grammar_path: &Path) -> Result<String> {
    let grammar_path = fs::canonicalize(grammar_path)?;
    let mut node_process = Command::new("node")
        .env("TREE_SITTER_GRAMMAR_PATH", grammar_path)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .spawn()
        .expect("Failed to run `node`");

    let mut node_stdin = node_process
        .stdin
        .take()
        .expect("Failed to open stdin for node");
    let javascript_code = include_bytes!("./dsl.js");
    node_stdin
        .write(javascript_code)
        .expect("Failed to write to node's stdin");
    drop(node_stdin);
    let output = node_process
        .wait_with_output()
        .expect("Failed to read output from node");
    match output.status.code() {
        None => panic!("Node process was killed"),
        Some(0) => {}
        Some(code) => return Err(anyhow!("Node process exited with status {}", code)),
    }

    let mut result = String::from_utf8(output.stdout).expect("Got invalid UTF8 from node");
    result.push('\n');
    Ok(result)
}

fn write_file(path: &Path, body: impl AsRef<[u8]>) -> Result<()> {
    fs::write(path, body)
        .with_context(|| format!("Failed to write {:?}", path.file_name().unwrap()))
}
<|endoftext|>


<commit_before>//! `bless` updates the reference files in the repo with changed output files
//! from the last test run.

use std::env;
use std::ffi::OsStr;
use std::fs;
use std::lazy::SyncLazy;
use std::path::PathBuf;
use walkdir::WalkDir;

use crate::clippy_project_root;

// NOTE: this is duplicated with tests/cargo/mod.rs What to do?
pub static CARGO_TARGET_DIR: SyncLazy<PathBuf> = SyncLazy::new(|| match env::var_os("CARGO_TARGET_DIR") {
    Some(v) => v.into(),
    None => env::current_dir().unwrap().join("target"),
});

pub fn bless() {
    let test_dirs = [
        clippy_project_root().join("tests").join("ui"),
        clippy_project_root().join("tests").join("ui-toml"),
        clippy_project_root().join("tests").join("ui-cargo"),
    ];
    for test_dir in &test_dirs {
        WalkDir::new(test_dir)
            .into_iter()
            .filter_map(Result::ok)
            .filter(|f| f.path().extension() == Some(OsStr::new("rs")))
            .for_each(|f| {
                update_reference_file(f.path().with_extension("stdout"));
                update_reference_file(f.path().with_extension("stderr"));
                update_reference_file(f.path().with_extension("fixed"));
            });
    }
}

fn update_reference_file(reference_file_path: PathBuf) {
    let test_output_path = build_dir().join(PathBuf::from(reference_file_path.file_name().unwrap()));
    let relative_reference_file_path = reference_file_path.strip_prefix(clippy_project_root()).unwrap();

    // If compiletest did not write any changes during the test run,
    // we don't have to update anything
    if !test_output_path.exists() {
        return;
    }

    let test_output_file = fs::read(&test_output_path).expect("Unable to read test output file");
    let reference_file = fs::read(&reference_file_path).unwrap_or_default();

    if test_output_file != reference_file {
        // If a test run caused an output file to change, update the reference file
        println!("updating {}", &relative_reference_file_path.display());
        fs::copy(test_output_path, &reference_file_path).expect("Could not update reference file");

        // We need to re-read the file now because it was potentially updated from copying
        let reference_file = fs::read(&reference_file_path).unwrap_or_default();

        if reference_file.is_empty() {
            // If we copied over an empty output file, we remove the now empty reference file
            println!("removing {}", &relative_reference_file_path.display());
            fs::remove_file(reference_file_path).expect("Could not remove reference file");
        }
    }
}

fn build_dir() -> PathBuf {
    let profile = env::var("PROFILE").unwrap_or_else(|_| "debug".to_string());
    let mut path = PathBuf::new();
    path.push(CARGO_TARGET_DIR.clone());
    path.push(profile);
    path.push("test_build_base");
    path
}
<commit_msg>Fix blessing of test output in subdirectories<commit_after>//! `bless` updates the reference files in the repo with changed output files
//! from the last test run.

use std::env;
use std::ffi::OsStr;
use std::fs;
use std::lazy::SyncLazy;
use std::path::PathBuf;
use walkdir::WalkDir;

use crate::clippy_project_root;

// NOTE: this is duplicated with tests/cargo/mod.rs What to do?
pub static CARGO_TARGET_DIR: SyncLazy<PathBuf> = SyncLazy::new(|| match env::var_os("CARGO_TARGET_DIR") {
    Some(v) => v.into(),
    None => env::current_dir().unwrap().join("target"),
});

pub fn bless() {
    let test_dirs = [
        clippy_project_root().join("tests").join("ui"),
        clippy_project_root().join("tests").join("ui-toml"),
        clippy_project_root().join("tests").join("ui-cargo"),
    ];
    for test_dir in &test_dirs {
        WalkDir::new(test_dir)
            .into_iter()
            .filter_map(Result::ok)
            .filter(|f| f.path().extension() == Some(OsStr::new("rs")))
            .for_each(|f| {
                let test_name = f.path().strip_prefix(test_dir).unwrap();

                update_reference_file(f.path().with_extension("stdout"), test_name.with_extension("stdout"));
                update_reference_file(f.path().with_extension("stderr"), test_name.with_extension("stderr"));
                update_reference_file(f.path().with_extension("fixed"), test_name.with_extension("fixed"));
            });
    }
}

fn update_reference_file(reference_file_path: PathBuf, test_name: PathBuf) {
    let test_output_path = build_dir().join(test_name);
    let relative_reference_file_path = reference_file_path.strip_prefix(clippy_project_root()).unwrap();

    // If compiletest did not write any changes during the test run,
    // we don't have to update anything
    if !test_output_path.exists() {
        return;
    }

    let test_output_file = fs::read(&test_output_path).expect("Unable to read test output file");
    let reference_file = fs::read(&reference_file_path).unwrap_or_default();

    if test_output_file != reference_file {
        // If a test run caused an output file to change, update the reference file
        println!("updating {}", &relative_reference_file_path.display());
        fs::copy(test_output_path, &reference_file_path).expect("Could not update reference file");

        // We need to re-read the file now because it was potentially updated from copying
        let reference_file = fs::read(&reference_file_path).unwrap_or_default();

        if reference_file.is_empty() {
            // If we copied over an empty output file, we remove the now empty reference file
            println!("removing {}", &relative_reference_file_path.display());
            fs::remove_file(reference_file_path).expect("Could not remove reference file");
        }
    }
}

fn build_dir() -> PathBuf {
    let profile = env::var("PROFILE").unwrap_or_else(|_| "debug".to_string());
    let mut path = PathBuf::new();
    path.push(CARGO_TARGET_DIR.clone());
    path.push(profile);
    path.push("test_build_base");
    path
}
<|endoftext|>