<commit_before>//! A wrapper around a Cairo image surface.

use std::io::Write;
use std::fs::File;
use std::default::Default;

use cairo::{Format, ImageSurface};
use glib::translate::ToGlibPtr;
use rlua::{self, LightUserData, Lua, Table,
           UserData, UserDataMethods, Value};
use wlroots::{Area, Origin, Size};
use tempfile;

use common::{class::{self, Class},
             object::{self, Object},
             property::Property};
use wayland_obj::{self, XdgToplevel};
use common::signal::emit_object_signal;

#[derive(Debug)]
pub struct DrawableState {
    temp_file: File,
    wayland_shell: XdgToplevel,
    pub surface: Option<ImageSurface>,
    geo: Area,
    // TODO Use this to determine whether we draw this or not
    refreshed: bool
}

pub type Drawable<'lua> = Object<'lua, DrawableState>;

impl Default for DrawableState {
    fn default() -> Self {
        let wayland_shell = wayland_obj::create_xdg_toplevel(None)
            .expect("Could not construct an xdg toplevel for a drawable");
        // TODO Is a temp file really the best way to do this?
        let temp_file = tempfile::tempfile()
            .expect("Could not make a temp file for the buffer");
        DrawableState { temp_file,
                        wayland_shell,
                        surface: None,
                        geo: Area::default(),
                        refreshed: false }
    }
}

impl<'lua> Drawable<'lua> {
    pub fn new(lua: &Lua) -> rlua::Result<Drawable> {
        let class = class::class_setup(lua, "drawable")?;
        let builder = Drawable::allocate(lua, class)?;
        // TODO Do properly
        let table = lua.create_table()?;
        table.set("geometry", lua.create_function(geometry)?)?;
        table.set("refresh", lua.create_function(refresh)?)?;
        Ok(builder.add_to_meta(table)?.build())
    }

    pub fn get_geometry(&self) -> rlua::Result<Area> {
        let drawable = self.state()?;
        Ok(drawable.geo)
    }

    pub fn get_surface(&self) -> rlua::Result<Value<'lua>> {
        let drawable = self.state()?;
        Ok(match drawable.surface {
            None => Value::Nil,
            Some(ref image) => {
                let stash = image.to_glib_none();
                let ptr = stash.0;
                // NOTE
                // We bump the reference count because now Lua has a reference which
                // it manages via LGI.
                //
                // If there's a bug, worst case scenario there's a memory leak.
                unsafe {
                    let _ = ::cairo_sys::cairo_surface_reference(ptr);
                }
                Value::LightUserData(LightUserData(ptr as _))
            }
        })
    }

    /// Sets the geometry, and allocates a new surface.
    pub fn set_geometry(&mut self, lua: &Lua, geometry: Area) -> rlua::Result<()> {
        use rlua::Error::RuntimeError;
        let obj_clone = self.clone();
        let mut drawable = self.state_mut()?;
        let size_changed = drawable.geo != geometry;
        drawable.geo = geometry;
        if size_changed {
            drawable.refreshed = false;
            drawable.surface = None;
            let size: Size = geometry.size;

            if size.width > 0 && size.height > 0 {
                let temp_file = tempfile::tempfile()
                    .expect("Could not make new temp file");
                temp_file.set_len(size.width as u64 * size.height as u64 * 4)
                    .expect("Could not set file length");
                drawable.surface = Some(ImageSurface::create(Format::ARgb32,
                                                             size.width,
                                                             size.height)
                    .map_err(|err| RuntimeError(format!("Could not allocate {:?}", err)))?);
                drawable.wayland_shell.set_size(size);
                drawable.wayland_shell.set_surface(&temp_file, size)
                    .map_err(|_| RuntimeError(format!("Could not set surface for drawable")))?;
                drawable.temp_file = temp_file;
                emit_object_signal(lua, obj_clone.into(), "property::surface".into(), ())?;
            }
        }
        Ok(())
    }

    /// Signals that the drawable's surface was updated.
    pub fn refresh(&mut self) -> rlua::Result<()> {
        let mut drawable = self.state_mut()?;
        let drawable = &mut *drawable;
        if let Some(data) = drawable.surface.as_mut().map(get_data) {
            drawable.temp_file.write(&*data)
                .expect("Could not write data to buffer");
            drawable.temp_file.flush()
                .expect("Could not flush buffer");
            drawable.refreshed = true;
        }
        Ok(())
    }
}

impl UserData for DrawableState {
    fn add_methods(methods: &mut UserDataMethods<Self>) {
        object::default_add_methods(methods);
    }
}

pub fn init(lua: &Lua) -> rlua::Result<Class<DrawableState>> {
    Class::<DrawableState>::builder(lua, "drawable", None)?
        .method("geometry".into(), lua.create_function(geometry)?)?
        .property(Property::new("surface".into(),
                                None,
                                Some(lua.create_function(get_surface)?),
                                None))?
        .save_class("drawable")?
        .build()
}

fn get_surface<'lua>(_: &'lua Lua, drawable: Drawable<'lua>) -> rlua::Result<Value<'lua>> {
    drawable.get_surface()
}

fn geometry<'lua>(lua: &'lua Lua, drawable: Drawable<'lua>) -> rlua::Result<Table<'lua>> {
    let geometry = drawable.get_geometry()?;
    let Origin { x, y } = geometry.origin;
    let Size { width, height } = geometry.size;
    let table = lua.create_table()?;
    table.set("x", x)?;
    table.set("y", y)?;
    table.set("width", width)?;
    table.set("height", height)?;
    Ok(table)
}

fn refresh<'lua>(_: &'lua Lua, mut drawable: Drawable<'lua>) -> rlua::Result<()> {
    drawable.refresh()
}

/// Get the data associated with the ImageSurface.
fn get_data(surface: &mut ImageSurface) -> &[u8] {
    // NOTE This is safe to do because there's one thread.
    //
    // We know Lua is not modifying it because it's not running.
    use std::slice;
    use cairo_sys;
    unsafe {
        let len = surface.get_stride() as usize * surface.get_height() as usize;
        let surface = surface.to_glib_none().0;
        slice::from_raw_parts(cairo_sys::cairo_image_surface_get_data(surface as _), len)
    }
}
<commit_msg>Remake the top level on drawable change<commit_after>//! A wrapper around a Cairo image surface.

use std::io::Write;
use std::fs::File;
use std::default::Default;

use cairo::{Format, ImageSurface};
use glib::translate::ToGlibPtr;
use rlua::{self, LightUserData, Lua, Table,
           UserData, UserDataMethods, Value};
use wlroots::{Area, Origin, Size};
use tempfile;

use common::{class::{self, Class},
             object::{self, Object},
             property::Property};
use wayland_obj::{self, XdgToplevel};
use common::signal::emit_object_signal;

#[derive(Debug)]
pub struct DrawableState {
    temp_file: File,
    wayland_shell: XdgToplevel,
    pub surface: Option<ImageSurface>,
    geo: Area,
    // TODO Use this to determine whether we draw this or not
    refreshed: bool
}

pub type Drawable<'lua> = Object<'lua, DrawableState>;

impl Default for DrawableState {
    fn default() -> Self {
        let wayland_shell = wayland_obj::create_xdg_toplevel(None)
            .expect("Could not construct an xdg toplevel for a drawable");
        // TODO Is a temp file really the best way to do this?
        let temp_file = tempfile::tempfile()
            .expect("Could not make a temp file for the buffer");
        DrawableState { temp_file,
                        wayland_shell,
                        surface: None,
                        geo: Area::default(),
                        refreshed: false }
    }
}

impl<'lua> Drawable<'lua> {
    pub fn new(lua: &Lua) -> rlua::Result<Drawable> {
        let class = class::class_setup(lua, "drawable")?;
        let builder = Drawable::allocate(lua, class)?;
        // TODO Do properly
        let table = lua.create_table()?;
        table.set("geometry", lua.create_function(geometry)?)?;
        table.set("refresh", lua.create_function(refresh)?)?;
        Ok(builder.add_to_meta(table)?.build())
    }

    pub fn get_geometry(&self) -> rlua::Result<Area> {
        let drawable = self.state()?;
        Ok(drawable.geo)
    }

    pub fn get_surface(&self) -> rlua::Result<Value<'lua>> {
        let drawable = self.state()?;
        Ok(match drawable.surface {
            None => Value::Nil,
            Some(ref image) => {
                let stash = image.to_glib_none();
                let ptr = stash.0;
                // NOTE
                // We bump the reference count because now Lua has a reference which
                // it manages via LGI.
                //
                // If there's a bug, worst case scenario there's a memory leak.
                unsafe {
                    let _ = ::cairo_sys::cairo_surface_reference(ptr);
                }
                Value::LightUserData(LightUserData(ptr as _))
            }
        })
    }

    /// Sets the geometry, and allocates a new surface.
    pub fn set_geometry(&mut self, lua: &Lua, geometry: Area) -> rlua::Result<()> {
        use rlua::Error::RuntimeError;
        let obj_clone = self.clone();
        let mut drawable = self.state_mut()?;
        let size_changed = drawable.geo != geometry;
        drawable.geo = geometry;
        if size_changed {
            drawable.refreshed = false;
            drawable.surface = None;
            // TODO The first one we create is useless.
            // How do we ensure we don't have that overhead? Option<>?
            drawable.wayland_shell = wayland_obj::create_xdg_toplevel(None)
                .expect("Could not construct an xdg toplevel for a drawable");
            let size: Size = geometry.size;

            if size.width > 0 && size.height > 0 {
                let temp_file = tempfile::tempfile()
                    .expect("Could not make new temp file");
                temp_file.set_len(size.width as u64 * size.height as u64 * 4)
                    .expect("Could not set file length");
                drawable.surface = Some(ImageSurface::create(Format::ARgb32,
                                                             size.width,
                                                             size.height)
                    .map_err(|err| RuntimeError(format!("Could not allocate {:?}", err)))?);
                drawable.wayland_shell.set_size(size);
                drawable.wayland_shell.set_surface(&temp_file, size)
                    .map_err(|_| RuntimeError(format!("Could not set surface for drawable")))?;
                drawable.temp_file = temp_file;
                emit_object_signal(lua, obj_clone.into(), "property::surface".into(), ())?;
            }
        }
        Ok(())
    }

    /// Signals that the drawable's surface was updated.
    pub fn refresh(&mut self) -> rlua::Result<()> {
        let mut drawable = self.state_mut()?;
        let drawable = &mut *drawable;
        if let Some(data) = drawable.surface.as_mut().map(get_data) {
            drawable.temp_file.write(&*data)
                .expect("Could not write data to buffer");
            drawable.temp_file.flush()
                .expect("Could not flush buffer");
            drawable.refreshed = true;
        }
        Ok(())
    }
}

impl UserData for DrawableState {
    fn add_methods(methods: &mut UserDataMethods<Self>) {
        object::default_add_methods(methods);
    }
}

pub fn init(lua: &Lua) -> rlua::Result<Class<DrawableState>> {
    Class::<DrawableState>::builder(lua, "drawable", None)?
        .method("geometry".into(), lua.create_function(geometry)?)?
        .property(Property::new("surface".into(),
                                None,
                                Some(lua.create_function(get_surface)?),
                                None))?
        .save_class("drawable")?
        .build()
}

fn get_surface<'lua>(_: &'lua Lua, drawable: Drawable<'lua>) -> rlua::Result<Value<'lua>> {
    drawable.get_surface()
}

fn geometry<'lua>(lua: &'lua Lua, drawable: Drawable<'lua>) -> rlua::Result<Table<'lua>> {
    let geometry = drawable.get_geometry()?;
    let Origin { x, y } = geometry.origin;
    let Size { width, height } = geometry.size;
    let table = lua.create_table()?;
    table.set("x", x)?;
    table.set("y", y)?;
    table.set("width", width)?;
    table.set("height", height)?;
    Ok(table)
}

fn refresh<'lua>(_: &'lua Lua, mut drawable: Drawable<'lua>) -> rlua::Result<()> {
    drawable.refresh()
}

/// Get the data associated with the ImageSurface.
fn get_data(surface: &mut ImageSurface) -> &[u8] {
    // NOTE This is safe to do because there's one thread.
    //
    // We know Lua is not modifying it because it's not running.
    use std::slice;
    use cairo_sys;
    unsafe {
        let len = surface.get_stride() as usize * surface.get_height() as usize;
        let surface = surface.to_glib_none().0;
        slice::from_raw_parts(cairo_sys::cairo_image_surface_get_data(surface as _), len)
    }
}
<|endoftext|>


<commit_before>use std::io::Read;

use crate::common::Lines;
use crate::Solution;

#[derive(Clone, Copy)]
enum Op {
    Addition,
    Multiplication,
}

fn extract_brackets(with_brackets: &str) -> (&str, &str) {
    let mut brackets = 1;

    let mut end = None;

    for (i, c) in with_brackets.chars().enumerate().skip(1) {
        match c {
            '(' => brackets += 1,
            ')' => brackets -= 1,
            _ => {}
        }

        if brackets == 0 {
            end = Some(i);
            break;
        }
    }

    let end = end.expect("Unmatched brackets");
    let next_start = with_brackets.len().min(end + 2);

    (&with_brackets[1..end], &with_brackets[next_start..])
}

fn compute_value(expression: &str) -> u64 {
    let mut remainder = expression;
    let mut value = 0;
    let mut op = None;

    let mut apply_value = |n, op| match op {
        Some(Op::Addition) => value += n,
        Some(Op::Multiplication) => value *= n,
        None => value = n,
    };

    while !remainder.is_empty() {
        let pos = remainder.find(' ');

        let part = if let Some(pos) = pos {
            &remainder[..pos]
        } else {
            remainder
        };

        match part.chars().next().unwrap() {
            '+' => op = Some(Op::Addition),
            '*' => op = Some(Op::Multiplication),
            '(' => {
                let (in_brackets, rest) = extract_brackets(remainder);
                remainder = rest;
                let n = compute_value(in_brackets);

                apply_value(n, op);
                // Skip the remainder update
                continue;
            }
            c if c.is_ascii_digit() => {
                let n = part.parse().unwrap();
                apply_value(n, op);
            }
            _ => panic!("Not a valid expression part {}", part),
        }

        let pos = pos.map(|n| n + 1).unwrap_or(remainder.len());
        remainder = &remainder[pos..];
    }

    value
}

fn compute_value2(expression: &str) -> u64 {
    let mut remainder = expression;
    let mut value = 0;
    let mut result = 1;
    let mut op = None;

    let mut apply_value = |n, op| match op {
        Some(Op::Addition) => value += n,
        Some(Op::Multiplication) => {
            result *= value;
            value = n;
        }
        None => value = n,
    };

    while !remainder.is_empty() {
        let pos = remainder.find(' ');

        let part = if let Some(pos) = pos {
            &remainder[..pos]
        } else {
            remainder
        };

        match part.chars().next().unwrap() {
            '+' => op = Some(Op::Addition),
            '*' => op = Some(Op::Multiplication),
            '(' => {
                let (in_brackets, rest) = extract_brackets(remainder);
                remainder = rest;
                let n = compute_value2(in_brackets);

                apply_value(n, op);
                // Skip the remainder update
                continue;
            }
            c if c.is_ascii_digit() => {
                let n = part.parse().unwrap();

                apply_value(n, op);
            }
            _ => panic!("Not a valid expression part {}", part),
        }

        let pos = pos.map(|n| n + 1).unwrap_or(remainder.len());
        remainder = &remainder[pos..];
    }

    result * value
}

#[derive(Default)]
pub struct Day18;

impl Solution for Day18 {
    fn part1(&mut self, input: &mut dyn Read) -> String {
        Lines::new(input)
            .map(|l| compute_value(&l))
            .sum::<u64>()
            .to_string()
    }

    fn part2(&mut self, input: &mut dyn Read) -> String {
        Lines::new(input)
            .map(|l| compute_value2(&l))
            .sum::<u64>()
            .to_string()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sample_part1() {
        assert_eq!(71, compute_value("1 + 2 * 3 + 4 * 5 + 6"));
        assert_eq!(51, compute_value("1 + (2 * 3) + (4 * (5 + 6))"));
        assert_eq!(26, compute_value("2 * 3 + (4 * 5)"));
        assert_eq!(437, compute_value("5 + (8 * 3 + 9 + 3 * 4 * 3)"));
        assert_eq!(
            12240,
            compute_value("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))")
        );
        assert_eq!(
            13632,
            compute_value("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2")
        );
    }

    #[test]
    fn sample_part2() {
        assert_eq!(231, compute_value2("1 + 2 * 3 + 4 * 5 + 6"));
        assert_eq!(51, compute_value2("1 + (2 * 3) + (4 * (5 + 6))"));
        assert_eq!(46, compute_value2("2 * 3 + (4 * 5)"));
        assert_eq!(1445, compute_value2("5 + (8 * 3 + 9 + 3 * 4 * 3)"));
        assert_eq!(
            669060,
            compute_value2("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))")
        );
        assert_eq!(
            23340,
            compute_value2("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2")
        );
    }
}
<commit_msg>Fix clippy warnings<commit_after>use std::io::Read;

use crate::common::Lines;
use crate::Solution;

#[derive(Clone, Copy)]
enum Op {
    Addition,
    Multiplication,
}

fn extract_brackets(with_brackets: &str) -> (&str, &str) {
    let mut brackets = 1;

    let mut end = None;

    for (i, c) in with_brackets.chars().enumerate().skip(1) {
        match c {
            '(' => brackets += 1,
            ')' => brackets -= 1,
            _ => {}
        }

        if brackets == 0 {
            end = Some(i);
            break;
        }
    }

    let end = end.expect("Unmatched brackets");
    let next_start = with_brackets.len().min(end + 2);

    (&with_brackets[1..end], &with_brackets[next_start..])
}

fn compute_value(expression: &str) -> u64 {
    let mut remainder = expression;
    let mut value = 0;
    let mut op = None;

    let mut apply_value = |n, op| match op {
        Some(Op::Addition) => value += n,
        Some(Op::Multiplication) => value *= n,
        None => value = n,
    };

    while !remainder.is_empty() {
        let pos = remainder.find(' ');

        let part = if let Some(pos) = pos {
            &remainder[..pos]
        } else {
            remainder
        };

        match part.chars().next().unwrap() {
            '+' => op = Some(Op::Addition),
            '*' => op = Some(Op::Multiplication),
            '(' => {
                let (in_brackets, rest) = extract_brackets(remainder);
                remainder = rest;
                let n = compute_value(in_brackets);

                apply_value(n, op);
                // Skip the remainder update
                continue;
            }
            c if c.is_ascii_digit() => {
                let n = part.parse().unwrap();
                apply_value(n, op);
            }
            _ => panic!("Not a valid expression part {}", part),
        }

        let pos = pos.map(|n| n + 1).unwrap_or_else(|| remainder.len());
        remainder = &remainder[pos..];
    }

    value
}

fn compute_value2(expression: &str) -> u64 {
    let mut remainder = expression;
    let mut value = 0;
    let mut result = 1;
    let mut op = None;

    let mut apply_value = |n, op| match op {
        Some(Op::Addition) => value += n,
        Some(Op::Multiplication) => {
            result *= value;
            value = n;
        }
        None => value = n,
    };

    while !remainder.is_empty() {
        let pos = remainder.find(' ');

        let part = if let Some(pos) = pos {
            &remainder[..pos]
        } else {
            remainder
        };

        match part.chars().next().unwrap() {
            '+' => op = Some(Op::Addition),
            '*' => op = Some(Op::Multiplication),
            '(' => {
                let (in_brackets, rest) = extract_brackets(remainder);
                remainder = rest;
                let n = compute_value2(in_brackets);

                apply_value(n, op);
                // Skip the remainder update
                continue;
            }
            c if c.is_ascii_digit() => {
                let n = part.parse().unwrap();

                apply_value(n, op);
            }
            _ => panic!("Not a valid expression part {}", part),
        }

        let pos = pos.map(|n| n + 1).unwrap_or_else(|| remainder.len());
        remainder = &remainder[pos..];
    }

    result * value
}

#[derive(Default)]
pub struct Day18;

impl Solution for Day18 {
    fn part1(&mut self, input: &mut dyn Read) -> String {
        Lines::new(input)
            .map(|l| compute_value(&l))
            .sum::<u64>()
            .to_string()
    }

    fn part2(&mut self, input: &mut dyn Read) -> String {
        Lines::new(input)
            .map(|l| compute_value2(&l))
            .sum::<u64>()
            .to_string()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sample_part1() {
        assert_eq!(71, compute_value("1 + 2 * 3 + 4 * 5 + 6"));
        assert_eq!(51, compute_value("1 + (2 * 3) + (4 * (5 + 6))"));
        assert_eq!(26, compute_value("2 * 3 + (4 * 5)"));
        assert_eq!(437, compute_value("5 + (8 * 3 + 9 + 3 * 4 * 3)"));
        assert_eq!(
            12240,
            compute_value("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))")
        );
        assert_eq!(
            13632,
            compute_value("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2")
        );
    }

    #[test]
    fn sample_part2() {
        assert_eq!(231, compute_value2("1 + 2 * 3 + 4 * 5 + 6"));
        assert_eq!(51, compute_value2("1 + (2 * 3) + (4 * (5 + 6))"));
        assert_eq!(46, compute_value2("2 * 3 + (4 * 5)"));
        assert_eq!(1445, compute_value2("5 + (8 * 3 + 9 + 3 * 4 * 3)"));
        assert_eq!(
            669060,
            compute_value2("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))")
        );
        assert_eq!(
            23340,
            compute_value2("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2")
        );
    }
}
<|endoftext|>


<commit_before>use std::io::Read;

use itertools::Itertools;

use crate::common::ordered;

fn compute_groups<'a>(it: impl IntoIterator<Item = &'a usize>) -> Vec<(usize, usize)> {
    let mut it = it.into_iter().copied().dedup_with_count();

    let (mut population, mut last_pos) = it.next().unwrap();
    let mut last_cost = 0;

    let mut costs = vec![(last_pos, 0)];

    for (number, pos) in it {
        let (first, last) = ordered(last_pos, pos);
        let new_cost = last_cost + population * (last - first);

        costs.push((pos, new_cost));

        population += number;
        last_pos = pos;
        last_cost = new_cost;
    }

    costs
}

fn read_input(input: &mut dyn Read) -> Vec<usize> {
    let mut buf = String::new();
    input.read_to_string(&mut buf).unwrap();

    let mut crabs: Vec<usize> = buf
        .trim_end()
        .split(',')
        .map(|s| s.parse().unwrap())
        .collect();

    crabs.sort_unstable();

    crabs
}

pub fn part1(input: &mut dyn Read) -> String {
    let crabs = read_input(input);

    let forward_costs = compute_groups(&crabs);
    let backwards_costs = compute_groups(crabs.iter().rev());

    backwards_costs
        .into_iter()
        .rev()
        .zip(forward_costs)
        .map(|((pos_b, cost_b), (pos_f, cost_f))| {
            debug_assert_eq!(pos_f, pos_b);

            cost_f + cost_b
        })
        .min()
        .unwrap()
        .to_string()
}

pub fn sum_until(end: usize) -> usize {
    (end * (1 + end)) / 2
}

fn cost_at(pos: usize, groups: &[(usize, usize)]) -> usize {
    groups
        .iter()
        .map(|&(number, new_pos)| {
            let (first, last) = ordered(pos, new_pos);

            number * sum_until(last - first)
        })
        .sum()
}

fn ternary_search(groups: &[(usize, usize)], min: usize, max: usize) -> usize {
    if max - min <= 6 {
        // ternary search isn't effective in small ranges, just iterate instead.
        (min..=max).map(|pos| cost_at(pos, groups)).min().unwrap()
    } else {
        let mid1 = min + (max - min) / 3;
        let mid2 = max - (max - min) / 3;

        let cost1 = cost_at(mid1, groups);
        let cost2 = cost_at(mid2, groups);

        if cost1 < cost2 {
            ternary_search(groups, min, mid2 - 1)
        } else {
            ternary_search(groups, mid1 + 1, max)
        }
    }
}

pub fn part2(input: &mut dyn Read) -> String {
    let crabs = read_input(input);
    let groups: Vec<_> = crabs.into_iter().dedup_with_count().collect();

    let min = groups.first().unwrap().1;
    let max = groups.last().unwrap().1;

    // Brute force approach, better version later
    ternary_search(&groups, min, max).to_string()
}

#[cfg(test)]
mod tests {
    use crate::test_implementation;

    use super::*;

    const SAMPLE: &[u8] = &*b"16,1,2,0,4,2,7,1,2,14";

    #[test]
    fn sample_part1() {
        test_implementation(part1, SAMPLE, 37);
    }

    #[test]
    fn sample_part2() {
        test_implementation(part2, SAMPLE, 168);
    }

    #[test]
    fn test_maths() {
        assert_eq!(sum_until(1), 1);
        assert_eq!(sum_until(2), 3);
        assert_eq!(sum_until(3), 6);
        assert_eq!(sum_until(4), 10);
    }
}
<commit_msg>Convert recursion to iteration<commit_after>use std::io::Read;

use itertools::Itertools;

use crate::common::ordered;

fn compute_groups<'a>(it: impl IntoIterator<Item = &'a usize>) -> Vec<(usize, usize)> {
    let mut it = it.into_iter().copied().dedup_with_count();

    let (mut population, mut last_pos) = it.next().unwrap();
    let mut last_cost = 0;

    let mut costs = vec![(last_pos, 0)];

    for (number, pos) in it {
        let (first, last) = ordered(last_pos, pos);
        let new_cost = last_cost + population * (last - first);

        costs.push((pos, new_cost));

        population += number;
        last_pos = pos;
        last_cost = new_cost;
    }

    costs
}

fn read_input(input: &mut dyn Read) -> Vec<usize> {
    let mut buf = String::new();
    input.read_to_string(&mut buf).unwrap();

    let mut crabs: Vec<usize> = buf
        .trim_end()
        .split(',')
        .map(|s| s.parse().unwrap())
        .collect();

    crabs.sort_unstable();

    crabs
}

pub fn part1(input: &mut dyn Read) -> String {
    let crabs = read_input(input);

    let forward_costs = compute_groups(&crabs);
    let backwards_costs = compute_groups(crabs.iter().rev());

    backwards_costs
        .into_iter()
        .rev()
        .zip(forward_costs)
        .map(|((_, cost_b), (_, cost_f))| cost_f + cost_b)
        .min()
        .unwrap()
        .to_string()
}

pub fn sum_until(end: usize) -> usize {
    (end * (1 + end)) / 2
}

fn cost_at(pos: usize, groups: &[(usize, usize)]) -> usize {
    groups
        .iter()
        .map(|&(number, new_pos)| {
            let (first, last) = ordered(pos, new_pos);

            number * sum_until(last - first)
        })
        .sum()
}

fn ternary_search(groups: &[(usize, usize)], mut min: usize, mut max: usize) -> usize {
    while max - min > 6 {
        let mid1 = min + (max - min) / 3;
        let mid2 = max - (max - min) / 3;

        let cost1 = cost_at(mid1, groups);
        let cost2 = cost_at(mid2, groups);

        if cost1 < cost2 {
            max = mid2 - 1
        } else {
            min = mid1 + 1
        }
    }

    // Ternary search isn't effective at such small intervals so we iterate the remaining part
    (min..=max).map(|pos| cost_at(pos, groups)).min().unwrap()
}

pub fn part2(input: &mut dyn Read) -> String {
    let crabs = read_input(input);
    let groups: Vec<_> = crabs.into_iter().dedup_with_count().collect();

    let min = groups.first().unwrap().1;
    let max = groups.last().unwrap().1;

    ternary_search(&groups, min, max).to_string()
}

#[cfg(test)]
mod tests {
    use crate::test_implementation;

    use super::*;

    const SAMPLE: &[u8] = &*b"16,1,2,0,4,2,7,1,2,14";

    #[test]
    fn sample_part1() {
        test_implementation(part1, SAMPLE, 37);
    }

    #[test]
    fn sample_part2() {
        test_implementation(part2, SAMPLE, 168);
    }

    #[test]
    fn test_maths() {
        assert_eq!(sum_until(1), 1);
        assert_eq!(sum_until(2), 3);
        assert_eq!(sum_until(3), 6);
        assert_eq!(sum_until(4), 10);
    }
}
<|endoftext|>


<commit_before>use rustc_serialize::base64::{Config, Newline, CharacterSet};
use hyper::header::Headers;
use utils::squeeze_path;
use std::ascii::AsciiExt;

pub mod auth13;
use authentication::auth13::Auth13;

pub static BASE64_AUTH: Config = Config {
    char_set: CharacterSet::Standard,
    newline: Newline::LF,
    pad: true,
    line_length: Some(60),
};

#[derive(Clone)]
pub struct Authentication {
    api_version: Option<String>,
    body: Option<String>,
    keypath: String,
    method: String,
    path: String,
    userid: String,
    version: String,
}

impl Authentication {
    pub fn new<P, K, M, U, V>(path: P, key: K, method: M, userid: U, version: V) -> Authentication
        where P: Into<String>,
              K: Into<String>,
              M: Into<String>,
              U: Into<String>,
              V: Into<String>
              {
                  Authentication {
                      api_version: None,
                      body: None,
                      keypath: key.into(),
                      method: method.into().to_ascii_uppercase(),
                      path: squeeze_path(path.into()),
                      userid: userid.into(),
                      version: version.into(),
                  }

              }

    pub fn api_version<S>(mut self, api_version: S) -> Authentication
        where S: Into<String>
        {
            self.api_version = Some(api_version.into());
            self
        }

    pub fn body<S>(mut self, body: S) -> Authentication
        where S: Into<String>
        {
            self.body = Some(body.into());
            self
        }

    pub fn headers(self) -> Headers {
        Auth13::new(&self.path,
                    &self.keypath,
                    &self.method,
                    &self.userid,
                    &self.api_version.unwrap(),
                    self.body)
            .headers()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::str;

    const PATH: &'static str = "/organizations/clownco";
    const BODY: &'static str = "Spec Body";
    const USER: &'static str = "spec-user";

    const PRIVATE_KEY: &'static str = "fixtures/spec-user.pem";

    #[test]
    fn test_auth_return() {
        let auth = Authentication::new(PATH, PRIVATE_KEY, "GET", USER, "1.3");
        let headers = auth.body(BODY).api_version("1").headers();
        assert!(headers.get_raw("x-ops-authorization-1").is_some())
    }
}
<commit_msg>Unwrap the headers<commit_after>use rustc_serialize::base64::{Config, Newline, CharacterSet};
use hyper::header::Headers;
use utils::squeeze_path;
use std::ascii::AsciiExt;

pub mod auth13;
use authentication::auth13::Auth13;

pub static BASE64_AUTH: Config = Config {
    char_set: CharacterSet::Standard,
    newline: Newline::LF,
    pad: true,
    line_length: Some(60),
};

#[derive(Clone)]
pub struct Authentication {
    api_version: Option<String>,
    body: Option<String>,
    keypath: String,
    method: String,
    path: String,
    userid: String,
    version: String,
}

impl Authentication {
    pub fn new<P, K, M, U, V>(path: P, key: K, method: M, userid: U, version: V) -> Authentication
        where P: Into<String>,
              K: Into<String>,
              M: Into<String>,
              U: Into<String>,
              V: Into<String>
              {
                  Authentication {
                      api_version: None,
                      body: None,
                      keypath: key.into(),
                      method: method.into().to_ascii_uppercase(),
                      path: squeeze_path(path.into()),
                      userid: userid.into(),
                      version: version.into(),
                  }

              }

    pub fn api_version<S>(mut self, api_version: S) -> Authentication
        where S: Into<String>
        {
            self.api_version = Some(api_version.into());
            self
        }

    pub fn body<S>(mut self, body: S) -> Authentication
        where S: Into<String>
        {
            self.body = Some(body.into());
            self
        }

    pub fn headers(self) -> Headers {
        Auth13::new(&self.path,
                    &self.keypath,
                    &self.method,
                    &self.userid,
                    &self.api_version.unwrap(),
                    self.body)
            .headers().unwrap()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::str;

    const PATH: &'static str = "/organizations/clownco";
    const BODY: &'static str = "Spec Body";
    const USER: &'static str = "spec-user";

    const PRIVATE_KEY: &'static str = "fixtures/spec-user.pem";

    #[test]
    fn test_auth_return() {
        let auth = Authentication::new(PATH, PRIVATE_KEY, "GET", USER, "1.3");
        let headers = auth.body(BODY).api_version("1").headers();
        assert!(headers.get_raw("x-ops-authorization-1").is_some())
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::node_identifier::NodeIdentifier;
use crate::util::{is_relay_custom_inline_fragment_directive, PointerAddress};

use fnv::{FnvBuildHasher, FnvHashMap};
use graphql_ir::{
    Condition, FragmentDefinition, InlineFragment, LinkedField, OperationDefinition, Program,
    Selection, Transformed, TransformedValue, Transformer,
};
use im::HashMap;
use std::iter::Iterator;
use std::sync::Arc;
/**
 * A transform that removes redundant fields and fragment spreads. Redundancy is
 * defined in this context as any selection that is guaranteed to already be
 * fetched by an ancestor selection. This can occur in two cases:
 *
 * 1. Simple duplicates at the same level of the document can always be skipped:
 *
 *
 * fragment Foo on FooType {
 *   id
 *   id
 *   ...Bar
 *   ...Bar
 * }
 *
 *
 * Becomes
 *
 *
 * fragment Foo on FooType {
 *   id
 *   ...Bar
 * }
 *
 *
 * 2. Inline fragments and conditions introduce the possibility for duplication
 * at different levels of the tree. Whenever a selection is fetched in a parent,
 * it is redundant to also fetch it in a child:
 *
 *
 * fragment Foo on FooType {
 *   id
 *   ... on OtherType {
 *     id # 1
 *   }
 *   ... on FooType @include(if: $cond) {
 *     id # 2
 *   }
 * }
 *
 *
 * Becomes:
 *
 *
 * fragment Foo on FooType {
 *   id
 * }
 *
 *
 * In this example:
 * - 1 can be skipped because `id` is already fetched by the parent. Even
 *   though the type is different (FooType/OtherType), the inline fragment
 *   cannot match without the outer fragment matching so the outer `id` is
 *   guaranteed to already be fetched.
 * - 2 can be skipped for similar reasons: it doesn't matter if the condition
 *   holds, `id` is already fetched by the parent regardless.
 *
 * This transform also handles more complicated cases in which selections are
 * nested:
 *
 *
 * fragment Foo on FooType {
 *   a {
 *     bb
 *   }
 *   ... on OtherType {
 *     a {
 *       bb # 1
 *       cc
 *     }
 *   }
 *  }
 *
 *
 * Becomes
 *
 *
 * fragment Foo on FooType {
 *   a {
 *     bb
 *   }
 *   ... on OtherType {
 *     a {
 *       cc
 *     }
 *   }
 *  }
 *
 *
 * 1 can be skipped because it is already fetched at the outer level.
 */
pub fn skip_redundant_nodes(program: &Program) -> Program {
    let mut transform = SkipRedundantNodesTransform::new(program);
    transform
        .transform_program(program)
        .replace_or_else(|| program.clone())
}

#[derive(Default, Clone)]
struct SelectionMap(HashMap<NodeIdentifier, Option<SelectionMap>, FnvBuildHasher>);

type Cache = FnvHashMap<PointerAddress, Transformed<Selection>>;

struct SkipRedundantNodesTransform<'s> {
    program: &'s Program,
    cache: Cache,
}

impl<'s> SkipRedundantNodesTransform<'s> {
    fn new(program: &'s Program) -> Self {
        Self {
            program,
            cache: Default::default(),
        }
    }

    fn transform_selection(
        &mut self,
        selection: &Selection,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Selection> {
        // This will optimize a traversal of the same subselections.
        // If it's the same node, and selection_map is empty
        // result of transform_selection has to be the same.
        let is_empty = selection_map.0.is_empty();
        if is_empty {
            let key = PointerAddress::new(selection);
            if let Some(result) = self.cache.get(&key) {
                return result.clone();
            }
        }
        let identifier = NodeIdentifier::from_selection(&self.program.schema, selection);
        let result = match selection {
            Selection::ScalarField(_) => {
                if selection_map.0.contains_key(&identifier) {
                    Transformed::Delete
                } else {
                    selection_map.0.insert(identifier, None);
                    Transformed::Keep
                }
            }
            Selection::FragmentSpread(_) => {
                if selection_map.0.contains_key(&identifier) {
                    Transformed::Delete
                } else {
                    selection_map.0.insert(identifier, None);
                    Transformed::Keep
                }
            }
            Selection::LinkedField(selection) => {
                if let Some(Some(linked_selection_map)) = selection_map.0.get_mut(&identifier) {
                    self.transform_linked_field(selection, linked_selection_map)
                        .map(Selection::LinkedField)
                } else {
                    let mut linked_selection_map = Default::default();
                    let result = self
                        .transform_linked_field(selection, &mut linked_selection_map)
                        .map(Selection::LinkedField);
                    if !linked_selection_map.0.is_empty() {
                        selection_map
                            .0
                            .insert(identifier, Some(linked_selection_map));
                    }
                    result
                }
            }
            Selection::Condition(selection) => {
                if let Some(Some(existing_selection_map)) = selection_map.0.get_mut(&identifier) {
                    self.transform_condition(selection, existing_selection_map)
                        .map(Selection::Condition)
                } else {
                    // Fork the selection map to prevent conditional selections from
                    // affecting the outer "guaranteed" selections.
                    let mut next_selection_map = selection_map.clone();
                    let result = self
                        .transform_condition(selection, &mut next_selection_map)
                        .map(Selection::Condition);
                    selection_map.0.insert(identifier, Some(next_selection_map));
                    result
                }
            }
            Selection::InlineFragment(selection) => {
                if let Some(Some(existing_selection_map)) = selection_map.0.get_mut(&identifier) {
                    self.transform_inline_fragment(selection, existing_selection_map)
                        .map(Selection::InlineFragment)
                } else if selection
                    .directives
                    .iter()
                    .any(is_relay_custom_inline_fragment_directive)
                {
                    let mut linked_selection_map = Default::default();
                    let result = self
                        .transform_inline_fragment(selection, &mut linked_selection_map)
                        .map(Selection::InlineFragment);
                    if !linked_selection_map.0.is_empty() {
                        selection_map
                            .0
                            .insert(identifier, Some(linked_selection_map));
                    }
                    result
                } else {
                    // Fork for inline fragments for the same reason
                    let mut next_selection_map = selection_map.clone();
                    let result = self
                        .transform_inline_fragment(selection, &mut next_selection_map)
                        .map(Selection::InlineFragment);
                    selection_map.0.insert(identifier, Some(next_selection_map));
                    result
                }
            }
        };
        if is_empty {
            let key = PointerAddress::new(selection);
            self.cache.insert(key, result.clone());
        }
        result
    }

    fn transform_linked_field(
        &mut self,
        field: &LinkedField,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Arc<LinkedField>> {
        let selections = self.transform_selections(&field.selections, selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => {
                if selections.is_empty() {
                    Transformed::Delete
                } else {
                    Transformed::Replace(Arc::new(LinkedField {
                        selections,
                        ..field.clone()
                    }))
                }
            }
        }
    }

    fn transform_condition(
        &mut self,
        condition: &Condition,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Arc<Condition>> {
        let selections = self.transform_selections(&condition.selections, selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => {
                if selections.is_empty() {
                    Transformed::Delete
                } else {
                    Transformed::Replace(Arc::new(Condition {
                        selections,
                        ..condition.clone()
                    }))
                }
            }
        }
    }

    fn transform_inline_fragment(
        &mut self,
        fragment: &InlineFragment,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Arc<InlineFragment>> {
        let selections = self.transform_selections(&fragment.selections, selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => {
                if selections.is_empty() {
                    Transformed::Delete
                } else {
                    Transformed::Replace(Arc::new(InlineFragment {
                        selections,
                        ..fragment.clone()
                    }))
                }
            }
        }
    }

    // Mostly a copy from Transformer::transform_list, but does partition and pass down `selection_map`.
    fn transform_selections(
        &mut self,
        selections: &[Selection],
        selection_map: &mut SelectionMap,
    ) -> TransformedValue<Vec<Selection>> {
        if selections.is_empty() {
            return TransformedValue::Keep;
        }
        let mut result: Vec<Selection> = Vec::new();
        let mut has_changes = false;
        let selections = get_partitioned_selections(selections);

        for (index, prev_item) in selections.iter().enumerate() {
            let next_item = self.transform_selection(prev_item, selection_map);
            match next_item {
                Transformed::Keep => {
                    if has_changes {
                        result.push((*prev_item).clone());
                    }
                }
                Transformed::Delete => {
                    if !has_changes {
                        debug_assert!(result.capacity() == 0);
                        // assume most items won't be skipped and allocate space for all items
                        result.reserve(selections.len());
                        result.extend(selections.iter().take(index).map(|&x| x.clone()));
                        has_changes = true;
                    }
                }
                Transformed::Replace(next_item) => {
                    if !has_changes {
                        debug_assert!(result.capacity() == 0);
                        // assume most items won't be skipped and allocate space for all items
                        result.reserve(selections.len());
                        result.extend(selections.iter().take(index).map(|&x| x.clone()));
                        has_changes = true;
                    }
                    result.push(next_item);
                }
            }
        }
        if has_changes {
            TransformedValue::Replace(result)
        } else {
            // TODO: This should return `Transformed::Keep`
            // For JS parity, clone the selections as the JS transform does
            TransformedValue::Replace(selections.iter().map(|&x| x.clone()).collect())
        }
    }
}

impl<'s> Transformer for SkipRedundantNodesTransform<'s> {
    const NAME: &'static str = "SkipRedundantNodesTransform";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_operation(
        &mut self,
        operation: &OperationDefinition,
    ) -> Transformed<OperationDefinition> {
        let mut selection_map = Default::default();
        let selections = self.transform_selections(&operation.selections, &mut selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => Transformed::Replace(OperationDefinition {
                selections,
                ..operation.clone()
            }),
        }
    }

    fn transform_fragment(
        &mut self,
        fragment: &FragmentDefinition,
    ) -> Transformed<FragmentDefinition> {
        let mut selection_map = Default::default();
        let selections = self.transform_selections(&fragment.selections, &mut selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => Transformed::Replace(FragmentDefinition {
                selections,
                ..fragment.clone()
            }),
        }
    }
}

/* Selections are sorted with fields first, "conditionals"
 * (inline fragments & conditions) last. This means that all fields that are
 * guaranteed to be fetched are encountered prior to any duplicates that may be
 * fetched within a conditional.
 */
fn get_partitioned_selections(selections: &[Selection]) -> Vec<&Selection> {
    let (mut left, right): (Vec<_>, Vec<_>) = selections.iter().partition(|sel| match sel {
        Selection::LinkedField(_) | Selection::ScalarField(_) => true,
        _ => false,
    });
    left.extend(right);
    left
}
<commit_msg>fix caching in skip_redundant_nodes<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::node_identifier::NodeIdentifier;
use crate::util::{is_relay_custom_inline_fragment_directive, PointerAddress};

use fnv::{FnvBuildHasher, FnvHashMap};
use graphql_ir::{
    Condition, FragmentDefinition, InlineFragment, LinkedField, OperationDefinition, Program,
    Selection, Transformed, TransformedValue, Transformer,
};
use im::HashMap;
use std::iter::Iterator;
use std::sync::Arc;
/**
 * A transform that removes redundant fields and fragment spreads. Redundancy is
 * defined in this context as any selection that is guaranteed to already be
 * fetched by an ancestor selection. This can occur in two cases:
 *
 * 1. Simple duplicates at the same level of the document can always be skipped:
 *
 *
 * fragment Foo on FooType {
 *   id
 *   id
 *   ...Bar
 *   ...Bar
 * }
 *
 *
 * Becomes
 *
 *
 * fragment Foo on FooType {
 *   id
 *   ...Bar
 * }
 *
 *
 * 2. Inline fragments and conditions introduce the possibility for duplication
 * at different levels of the tree. Whenever a selection is fetched in a parent,
 * it is redundant to also fetch it in a child:
 *
 *
 * fragment Foo on FooType {
 *   id
 *   ... on OtherType {
 *     id # 1
 *   }
 *   ... on FooType @include(if: $cond) {
 *     id # 2
 *   }
 * }
 *
 *
 * Becomes:
 *
 *
 * fragment Foo on FooType {
 *   id
 * }
 *
 *
 * In this example:
 * - 1 can be skipped because `id` is already fetched by the parent. Even
 *   though the type is different (FooType/OtherType), the inline fragment
 *   cannot match without the outer fragment matching so the outer `id` is
 *   guaranteed to already be fetched.
 * - 2 can be skipped for similar reasons: it doesn't matter if the condition
 *   holds, `id` is already fetched by the parent regardless.
 *
 * This transform also handles more complicated cases in which selections are
 * nested:
 *
 *
 * fragment Foo on FooType {
 *   a {
 *     bb
 *   }
 *   ... on OtherType {
 *     a {
 *       bb # 1
 *       cc
 *     }
 *   }
 *  }
 *
 *
 * Becomes
 *
 *
 * fragment Foo on FooType {
 *   a {
 *     bb
 *   }
 *   ... on OtherType {
 *     a {
 *       cc
 *     }
 *   }
 *  }
 *
 *
 * 1 can be skipped because it is already fetched at the outer level.
 */
pub fn skip_redundant_nodes(program: &Program) -> Program {
    let mut transform = SkipRedundantNodesTransform::new(program);
    transform
        .transform_program(program)
        .replace_or_else(|| program.clone())
}

#[derive(Default, Clone)]
struct SelectionMap(HashMap<NodeIdentifier, Option<SelectionMap>, FnvBuildHasher>);

type Cache = FnvHashMap<PointerAddress, (Transformed<Selection>, SelectionMap)>;

struct SkipRedundantNodesTransform<'s> {
    program: &'s Program,
    cache: Cache,
}

impl<'s> SkipRedundantNodesTransform<'s> {
    fn new(program: &'s Program) -> Self {
        Self {
            program,
            cache: Default::default(),
        }
    }

    fn transform_selection(
        &mut self,
        selection: &Selection,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Selection> {
        // This will optimize a traversal of the same subselections.
        // If it's the same node, and selection_map is empty
        // result of transform_selection has to be the same.
        let is_empty = selection_map.0.is_empty();
        if is_empty {
            let key = PointerAddress::new(selection);
            if let Some((cached_result, cached_selection_map)) = self.cache.get(&key).cloned() {
                *selection_map = cached_selection_map;
                return cached_result;
            }
        }
        let identifier = NodeIdentifier::from_selection(&self.program.schema, selection);
        let result = match selection {
            Selection::ScalarField(_) => {
                if selection_map.0.contains_key(&identifier) {
                    Transformed::Delete
                } else {
                    selection_map.0.insert(identifier, None);
                    Transformed::Keep
                }
            }
            Selection::FragmentSpread(_) => {
                if selection_map.0.contains_key(&identifier) {
                    Transformed::Delete
                } else {
                    selection_map.0.insert(identifier, None);
                    Transformed::Keep
                }
            }
            Selection::LinkedField(selection) => {
                if let Some(Some(linked_selection_map)) = selection_map.0.get_mut(&identifier) {
                    self.transform_linked_field(selection, linked_selection_map)
                        .map(Selection::LinkedField)
                } else {
                    let mut linked_selection_map = Default::default();
                    let result = self
                        .transform_linked_field(selection, &mut linked_selection_map)
                        .map(Selection::LinkedField);
                    if !linked_selection_map.0.is_empty() {
                        selection_map
                            .0
                            .insert(identifier, Some(linked_selection_map));
                    }
                    result
                }
            }
            Selection::Condition(selection) => {
                if let Some(Some(existing_selection_map)) = selection_map.0.get_mut(&identifier) {
                    self.transform_condition(selection, existing_selection_map)
                        .map(Selection::Condition)
                } else {
                    // Fork the selection map to prevent conditional selections from
                    // affecting the outer "guaranteed" selections.
                    let mut next_selection_map = selection_map.clone();
                    let result = self
                        .transform_condition(selection, &mut next_selection_map)
                        .map(Selection::Condition);
                    selection_map.0.insert(identifier, Some(next_selection_map));
                    result
                }
            }
            Selection::InlineFragment(selection) => {
                if let Some(Some(existing_selection_map)) = selection_map.0.get_mut(&identifier) {
                    self.transform_inline_fragment(selection, existing_selection_map)
                        .map(Selection::InlineFragment)
                } else if selection
                    .directives
                    .iter()
                    .any(is_relay_custom_inline_fragment_directive)
                {
                    let mut linked_selection_map = Default::default();
                    let result = self
                        .transform_inline_fragment(selection, &mut linked_selection_map)
                        .map(Selection::InlineFragment);
                    if !linked_selection_map.0.is_empty() {
                        selection_map
                            .0
                            .insert(identifier, Some(linked_selection_map));
                    }
                    result
                } else {
                    // Fork for inline fragments for the same reason
                    let mut next_selection_map = selection_map.clone();
                    let result = self
                        .transform_inline_fragment(selection, &mut next_selection_map)
                        .map(Selection::InlineFragment);
                    selection_map.0.insert(identifier, Some(next_selection_map));
                    result
                }
            }
        };
        if is_empty {
            let key = PointerAddress::new(selection);
            self.cache
                .insert(key, (result.clone(), selection_map.clone()));
        }
        result
    }

    fn transform_linked_field(
        &mut self,
        field: &LinkedField,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Arc<LinkedField>> {
        let selections = self.transform_selections(&field.selections, selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => {
                if selections.is_empty() {
                    Transformed::Delete
                } else {
                    Transformed::Replace(Arc::new(LinkedField {
                        selections,
                        ..field.clone()
                    }))
                }
            }
        }
    }

    fn transform_condition(
        &mut self,
        condition: &Condition,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Arc<Condition>> {
        let selections = self.transform_selections(&condition.selections, selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => {
                if selections.is_empty() {
                    Transformed::Delete
                } else {
                    Transformed::Replace(Arc::new(Condition {
                        selections,
                        ..condition.clone()
                    }))
                }
            }
        }
    }

    fn transform_inline_fragment(
        &mut self,
        fragment: &InlineFragment,
        selection_map: &mut SelectionMap,
    ) -> Transformed<Arc<InlineFragment>> {
        let selections = self.transform_selections(&fragment.selections, selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => {
                if selections.is_empty() {
                    Transformed::Delete
                } else {
                    Transformed::Replace(Arc::new(InlineFragment {
                        selections,
                        ..fragment.clone()
                    }))
                }
            }
        }
    }

    // Mostly a copy from Transformer::transform_list, but does partition and pass down `selection_map`.
    fn transform_selections(
        &mut self,
        selections: &[Selection],
        selection_map: &mut SelectionMap,
    ) -> TransformedValue<Vec<Selection>> {
        if selections.is_empty() {
            return TransformedValue::Keep;
        }
        let mut result: Vec<Selection> = Vec::new();
        let mut has_changes = false;
        let selections = get_partitioned_selections(selections);

        for (index, prev_item) in selections.iter().enumerate() {
            let next_item = self.transform_selection(prev_item, selection_map);
            match next_item {
                Transformed::Keep => {
                    if has_changes {
                        result.push((*prev_item).clone());
                    }
                }
                Transformed::Delete => {
                    if !has_changes {
                        debug_assert!(result.capacity() == 0);
                        // assume most items won't be skipped and allocate space for all items
                        result.reserve(selections.len());
                        result.extend(selections.iter().take(index).map(|&x| x.clone()));
                        has_changes = true;
                    }
                }
                Transformed::Replace(next_item) => {
                    if !has_changes {
                        debug_assert!(result.capacity() == 0);
                        // assume most items won't be skipped and allocate space for all items
                        result.reserve(selections.len());
                        result.extend(selections.iter().take(index).map(|&x| x.clone()));
                        has_changes = true;
                    }
                    result.push(next_item);
                }
            }
        }
        if has_changes {
            TransformedValue::Replace(result)
        } else {
            // TODO: This should return `Transformed::Keep`
            // For JS parity, clone the selections as the JS transform does
            TransformedValue::Replace(selections.iter().map(|&x| x.clone()).collect())
        }
    }
}

impl<'s> Transformer for SkipRedundantNodesTransform<'s> {
    const NAME: &'static str = "SkipRedundantNodesTransform";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_operation(
        &mut self,
        operation: &OperationDefinition,
    ) -> Transformed<OperationDefinition> {
        let mut selection_map = Default::default();
        let selections = self.transform_selections(&operation.selections, &mut selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => Transformed::Replace(OperationDefinition {
                selections,
                ..operation.clone()
            }),
        }
    }

    fn transform_fragment(
        &mut self,
        fragment: &FragmentDefinition,
    ) -> Transformed<FragmentDefinition> {
        let mut selection_map = Default::default();
        let selections = self.transform_selections(&fragment.selections, &mut selection_map);
        match selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => Transformed::Replace(FragmentDefinition {
                selections,
                ..fragment.clone()
            }),
        }
    }
}

/* Selections are sorted with fields first, "conditionals"
 * (inline fragments & conditions) last. This means that all fields that are
 * guaranteed to be fetched are encountered prior to any duplicates that may be
 * fetched within a conditional.
 */
fn get_partitioned_selections(selections: &[Selection]) -> Vec<&Selection> {
    let (mut left, right): (Vec<_>, Vec<_>) = selections.iter().partition(|sel| match sel {
        Selection::LinkedField(_) | Selection::ScalarField(_) => true,
        _ => false,
    });
    left.extend(right);
    left
}
<|endoftext|>


<commit_before>//! C API for mp4parse module.
//!
//! Parses ISO Base Media Format aka video/mp4 streams.
//!
//! # Examples
//!
//! ```rust
//! extern crate mp4parse;
//!
//! // Minimal valid mp4 containing no tracks.
//! let data = b"\0\0\0\x0cftypmp42";
//!
//! let context = mp4parse::mp4parse_new();
//! unsafe {
//!     let rv = mp4parse::mp4parse_read(context, data.as_ptr(), data.len());
//!     assert_eq!(0, rv);
//!     mp4parse::mp4parse_free(context);
//! }
//! ```

// 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 https://mozilla.org/MPL/2.0/.

use std;
use std::io::Cursor;

// Symbols we need from our rust api.
use MediaContext;
use TrackType;
use read_mp4;
use Error;
use media_time_to_ms;
use track_time_to_ms;
use SampleEntry;

// These constants *must* match those in include/mp4parse.h.

/// Map Error to int32 return codes.
const MP4PARSE_OK: i32 = 0;
const MP4PARSE_ERROR_BADARG: i32 = 1;
const MP4PARSE_ERROR_INVALID: i32 = 2;
const MP4PARSE_ERROR_UNSUPPORTED: i32 = 3;
const MP4PARSE_ERROR_EOF: i32 = 4;
const MP4PARSE_ASSERT: i32 = 5;
const MP4PARSE_ERROR_IO: i32 = 6;

/// Map TrackType to uint32 constants.
const TRACK_TYPE_H264: u32 = 0;
const TRACK_TYPE_AAC: u32 = 1;

// These structs *must* match those declared in include/mp4parse.h.

#[repr(C)]
pub struct TrackInfo {
    track_type: u32,
    track_id: u32,
    duration: u64,
    media_time: i64, // wants to be u64? understand how elst adjustment works
    // TODO(kinetik): include crypto guff
}

#[repr(C)]
pub struct TrackAudioInfo {
    channels: u16,
    bit_depth: u16,
    sample_rate: u32,
    // TODO(kinetik):
    // int32_t profile;
    // int32_t extended_profile; // check types
    // extra_data
    // codec_specific_config
}

#[repr(C)]
pub struct TrackVideoInfo {
    display_width: u32,
    display_height: u32,
    image_width: u16,
    image_height: u16,
    // TODO(kinetik):
    // extra_data
    // codec_specific_config
}

// C API wrapper functions.

/// Allocate an opaque rust-side parser context.
#[no_mangle]
pub extern "C" fn mp4parse_new() -> *mut MediaContext {
    let context = Box::new(MediaContext::new());
    Box::into_raw(context)
}

/// Free a rust-side parser context.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_free(context: *mut MediaContext) {
    assert!(!context.is_null());
    let _ = Box::from_raw(context);
}

/// Feed a buffer through `read_mp4()` with the given rust-side
/// parser context, returning success or an error code.
///
/// This is safe to call with NULL arguments but will crash
/// if given invalid pointers, as is usual for C.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_read(context: *mut MediaContext, buffer: *const u8, size: usize) -> i32 {
    // Validate arguments from C.
    if context.is_null() || buffer.is_null() || size < 8 {
        return MP4PARSE_ERROR_BADARG;
    }

    let mut context: &mut MediaContext = &mut *context;

    // Wrap the buffer we've been give in a slice.
    let b = std::slice::from_raw_parts(buffer, size);
    let mut c = Cursor::new(b);

    // Parse in a subthread to catch any panics.
    let task = std::thread::spawn(move || read_mp4(&mut c, &mut context));
    // The task's JoinHandle will return an error result if the
    // thread panicked, and will wrap the closure's return'd
    // result in an Ok(..) otherwise, meaning we could see
    // Ok(Err(Error::..)) here. So map thread failures back
    // to an mp4parse::Error before converting to a C return value.
    match task.join().or(Err(Error::AssertCaught)) {
        Ok(_) => MP4PARSE_OK,
        Err(Error::InvalidData) => MP4PARSE_ERROR_INVALID,
        Err(Error::Unsupported) => MP4PARSE_ERROR_UNSUPPORTED,
        Err(Error::UnexpectedEOF) => MP4PARSE_ERROR_EOF,
        Err(Error::AssertCaught) => MP4PARSE_ASSERT,
        Err(Error::Io(_)) => MP4PARSE_ERROR_IO,
    }
}

/// Return the number of tracks parsed by previous `read_mp4()` calls.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_count(context: *const MediaContext) -> u32 {
    // Validate argument from C.
    assert!(!context.is_null());
    let context = &*context;

    // Make sure the track count fits in a u32.
    assert!(context.tracks.len() < u32::max_value() as usize);
    context.tracks.len() as u32
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_info(context: *mut MediaContext, track: u32, info: *mut TrackInfo) -> i32 {
    if context.is_null() || info.is_null() {
        return MP4PARSE_ERROR_BADARG;
    }

    let context: &mut MediaContext = &mut *context;
    let track_index: usize = track as usize;
    let info: &mut TrackInfo = &mut *info;

    if track_index >= context.tracks.len() {
        return MP4PARSE_ERROR_BADARG;
    }

    info.track_type = match context.tracks[track_index].track_type {
        TrackType::Video => TRACK_TYPE_H264,
        TrackType::Audio => TRACK_TYPE_AAC,
        TrackType::Unknown => return MP4PARSE_ERROR_UNSUPPORTED,
    };

    // Maybe context & track should just have a single simple is_valid() instead?
    if context.timescale.is_none() ||
       context.tracks[track_index].timescale.is_none() ||
       context.tracks[track_index].duration.is_none() ||
       context.tracks[track_index].track_id.is_none() {
        return MP4PARSE_ERROR_INVALID;
    }

    std::thread::spawn(move || {
        let track = &context.tracks[track_index];
        let empty_duration = if track.empty_duration.is_some() {
            media_time_to_ms(track.empty_duration.unwrap(), context.timescale.unwrap())
        } else {
            0
        };
        info.media_time = if track.media_time.is_some() {
            track_time_to_ms(track.media_time.unwrap(), track.timescale.unwrap()) as i64 - empty_duration as i64
        } else {
            0
        };
        info.duration = track_time_to_ms(track.duration.unwrap(), track.timescale.unwrap());
        info.track_id = track.track_id.unwrap();
        MP4PARSE_OK
    }).join().unwrap_or(MP4PARSE_ERROR_INVALID)
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_audio_info(context: *mut MediaContext, track: u32, info: *mut TrackAudioInfo) -> i32 {
    if context.is_null() || info.is_null() {
        return MP4PARSE_ERROR_BADARG;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return MP4PARSE_ERROR_BADARG;
    }

    let track = &context.tracks[track as usize];

    match track.track_type {
        TrackType::Audio => {}
        _ => return MP4PARSE_ERROR_INVALID,
    };

    let audio = match track.data {
        Some(ref data) => data,
        None => return MP4PARSE_ERROR_INVALID,
    };

    let audio = match *audio {
        SampleEntry::Audio(ref x) => x,
        _ => return MP4PARSE_ERROR_INVALID,
    };

    (*info).channels = audio.channelcount;
    (*info).bit_depth = audio.samplesize;
    (*info).sample_rate = audio.samplerate >> 16; // 16.16 fixed point

    MP4PARSE_OK
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_video_info(context: *mut MediaContext, track: u32, info: *mut TrackVideoInfo) -> i32 {
    if context.is_null() || info.is_null() {
        return MP4PARSE_ERROR_BADARG;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return MP4PARSE_ERROR_BADARG;
    }

    let track = &context.tracks[track as usize];

    match track.track_type {
        TrackType::Video => {}
        _ => return MP4PARSE_ERROR_INVALID,
    };

    let video = match track.data {
        Some(ref data) => data,
        None => return MP4PARSE_ERROR_INVALID,
    };

    let video = match *video {
        SampleEntry::Video(ref x) => x,
        _ => return MP4PARSE_ERROR_INVALID,
    };

    if let Some(ref tkhd) = track.tkhd {
        (*info).display_width = tkhd.width >> 16; // 16.16 fixed point
        (*info).display_height = tkhd.height >> 16; // 16.16 fixed point
    } else {
        return MP4PARSE_ERROR_INVALID;
    }
    (*info).image_width = video.width;
    (*info).image_width = video.height;

    MP4PARSE_OK
}

#[test]
fn new_context() {
    let context = mp4parse_new();
    assert!(!context.is_null());
    unsafe {
        mp4parse_free(context);
    }
}

#[test]
#[should_panic(expected = "assertion failed")]
fn free_null_context() {
    unsafe {
        mp4parse_free(std::ptr::null_mut());
    }
}

#[test]
fn arg_validation() {
    let null_buffer = std::ptr::null();
    let null_context = std::ptr::null_mut();

    let context = mp4parse_new();
    assert!(!context.is_null());

    let buffer = vec![0u8; 8];

    unsafe {
        assert_eq!(MP4PARSE_ERROR_BADARG,
                   mp4parse_read(null_context, null_buffer, 0));
        assert_eq!(MP4PARSE_ERROR_BADARG,
                   mp4parse_read(context, null_buffer, 0));
    }

    for size in 0..buffer.len() {
        println!("testing buffer length {}", size);
        unsafe {
            assert_eq!(MP4PARSE_ERROR_BADARG,
                       mp4parse_read(context, buffer.as_ptr(), size));
        }
    }

    unsafe {
        mp4parse_free(context);
    }
}
<commit_msg>Unwrap the JoinHandle's result so we can match on its contents.<commit_after>//! C API for mp4parse module.
//!
//! Parses ISO Base Media Format aka video/mp4 streams.
//!
//! # Examples
//!
//! ```rust
//! extern crate mp4parse;
//!
//! // Minimal valid mp4 containing no tracks.
//! let data = b"\0\0\0\x0cftypmp42";
//!
//! let context = mp4parse::mp4parse_new();
//! unsafe {
//!     let rv = mp4parse::mp4parse_read(context, data.as_ptr(), data.len());
//!     assert_eq!(0, rv);
//!     mp4parse::mp4parse_free(context);
//! }
//! ```

// 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 https://mozilla.org/MPL/2.0/.

use std;
use std::io::Cursor;

// Symbols we need from our rust api.
use MediaContext;
use TrackType;
use read_mp4;
use Error;
use media_time_to_ms;
use track_time_to_ms;
use SampleEntry;

// These constants *must* match those in include/mp4parse.h.

/// Map Error to int32 return codes.
const MP4PARSE_OK: i32 = 0;
const MP4PARSE_ERROR_BADARG: i32 = 1;
const MP4PARSE_ERROR_INVALID: i32 = 2;
const MP4PARSE_ERROR_UNSUPPORTED: i32 = 3;
const MP4PARSE_ERROR_EOF: i32 = 4;
const MP4PARSE_ASSERT: i32 = 5;
const MP4PARSE_ERROR_IO: i32 = 6;

/// Map TrackType to uint32 constants.
const TRACK_TYPE_H264: u32 = 0;
const TRACK_TYPE_AAC: u32 = 1;

// These structs *must* match those declared in include/mp4parse.h.

#[repr(C)]
pub struct TrackInfo {
    track_type: u32,
    track_id: u32,
    duration: u64,
    media_time: i64, // wants to be u64? understand how elst adjustment works
    // TODO(kinetik): include crypto guff
}

#[repr(C)]
pub struct TrackAudioInfo {
    channels: u16,
    bit_depth: u16,
    sample_rate: u32,
    // TODO(kinetik):
    // int32_t profile;
    // int32_t extended_profile; // check types
    // extra_data
    // codec_specific_config
}

#[repr(C)]
pub struct TrackVideoInfo {
    display_width: u32,
    display_height: u32,
    image_width: u16,
    image_height: u16,
    // TODO(kinetik):
    // extra_data
    // codec_specific_config
}

// C API wrapper functions.

/// Allocate an opaque rust-side parser context.
#[no_mangle]
pub extern "C" fn mp4parse_new() -> *mut MediaContext {
    let context = Box::new(MediaContext::new());
    Box::into_raw(context)
}

/// Free a rust-side parser context.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_free(context: *mut MediaContext) {
    assert!(!context.is_null());
    let _ = Box::from_raw(context);
}

/// Feed a buffer through `read_mp4()` with the given rust-side
/// parser context, returning success or an error code.
///
/// This is safe to call with NULL arguments but will crash
/// if given invalid pointers, as is usual for C.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_read(context: *mut MediaContext, buffer: *const u8, size: usize) -> i32 {
    // Validate arguments from C.
    if context.is_null() || buffer.is_null() || size < 8 {
        return MP4PARSE_ERROR_BADARG;
    }

    let mut context: &mut MediaContext = &mut *context;

    // Wrap the buffer we've been give in a slice.
    let b = std::slice::from_raw_parts(buffer, size);
    let mut c = Cursor::new(b);

    // Parse in a subthread to catch any panics.
    let task = std::thread::spawn(move || read_mp4(&mut c, &mut context));
    // The task's JoinHandle will return an error result if the
    // thread panicked, and will wrap the closure's return'd
    // result in an Ok(..) otherwise, meaning we could see
    // Ok(Err(Error::..)) here. So map thread failures back
    // to an mp4parse::Error before converting to a C return value.
    match task.join().unwrap_or(Err(Error::AssertCaught)) {
        Ok(_) => MP4PARSE_OK,
        Err(Error::InvalidData) => MP4PARSE_ERROR_INVALID,
        Err(Error::Unsupported) => MP4PARSE_ERROR_UNSUPPORTED,
        Err(Error::UnexpectedEOF) => MP4PARSE_ERROR_EOF,
        Err(Error::AssertCaught) => MP4PARSE_ASSERT,
        Err(Error::Io(_)) => MP4PARSE_ERROR_IO,
    }
}

/// Return the number of tracks parsed by previous `read_mp4()` calls.
#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_count(context: *const MediaContext) -> u32 {
    // Validate argument from C.
    assert!(!context.is_null());
    let context = &*context;

    // Make sure the track count fits in a u32.
    assert!(context.tracks.len() < u32::max_value() as usize);
    context.tracks.len() as u32
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_info(context: *mut MediaContext, track: u32, info: *mut TrackInfo) -> i32 {
    if context.is_null() || info.is_null() {
        return MP4PARSE_ERROR_BADARG;
    }

    let context: &mut MediaContext = &mut *context;
    let track_index: usize = track as usize;
    let info: &mut TrackInfo = &mut *info;

    if track_index >= context.tracks.len() {
        return MP4PARSE_ERROR_BADARG;
    }

    info.track_type = match context.tracks[track_index].track_type {
        TrackType::Video => TRACK_TYPE_H264,
        TrackType::Audio => TRACK_TYPE_AAC,
        TrackType::Unknown => return MP4PARSE_ERROR_UNSUPPORTED,
    };

    // Maybe context & track should just have a single simple is_valid() instead?
    if context.timescale.is_none() ||
       context.tracks[track_index].timescale.is_none() ||
       context.tracks[track_index].duration.is_none() ||
       context.tracks[track_index].track_id.is_none() {
        return MP4PARSE_ERROR_INVALID;
    }

    std::thread::spawn(move || {
        let track = &context.tracks[track_index];
        let empty_duration = if track.empty_duration.is_some() {
            media_time_to_ms(track.empty_duration.unwrap(), context.timescale.unwrap())
        } else {
            0
        };
        info.media_time = if track.media_time.is_some() {
            track_time_to_ms(track.media_time.unwrap(), track.timescale.unwrap()) as i64 - empty_duration as i64
        } else {
            0
        };
        info.duration = track_time_to_ms(track.duration.unwrap(), track.timescale.unwrap());
        info.track_id = track.track_id.unwrap();
        MP4PARSE_OK
    }).join().unwrap_or(MP4PARSE_ERROR_INVALID)
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_audio_info(context: *mut MediaContext, track: u32, info: *mut TrackAudioInfo) -> i32 {
    if context.is_null() || info.is_null() {
        return MP4PARSE_ERROR_BADARG;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return MP4PARSE_ERROR_BADARG;
    }

    let track = &context.tracks[track as usize];

    match track.track_type {
        TrackType::Audio => {}
        _ => return MP4PARSE_ERROR_INVALID,
    };

    let audio = match track.data {
        Some(ref data) => data,
        None => return MP4PARSE_ERROR_INVALID,
    };

    let audio = match *audio {
        SampleEntry::Audio(ref x) => x,
        _ => return MP4PARSE_ERROR_INVALID,
    };

    (*info).channels = audio.channelcount;
    (*info).bit_depth = audio.samplesize;
    (*info).sample_rate = audio.samplerate >> 16; // 16.16 fixed point

    MP4PARSE_OK
}

#[no_mangle]
pub unsafe extern "C" fn mp4parse_get_track_video_info(context: *mut MediaContext, track: u32, info: *mut TrackVideoInfo) -> i32 {
    if context.is_null() || info.is_null() {
        return MP4PARSE_ERROR_BADARG;
    }

    let context: &mut MediaContext = &mut *context;

    if track as usize >= context.tracks.len() {
        return MP4PARSE_ERROR_BADARG;
    }

    let track = &context.tracks[track as usize];

    match track.track_type {
        TrackType::Video => {}
        _ => return MP4PARSE_ERROR_INVALID,
    };

    let video = match track.data {
        Some(ref data) => data,
        None => return MP4PARSE_ERROR_INVALID,
    };

    let video = match *video {
        SampleEntry::Video(ref x) => x,
        _ => return MP4PARSE_ERROR_INVALID,
    };

    if let Some(ref tkhd) = track.tkhd {
        (*info).display_width = tkhd.width >> 16; // 16.16 fixed point
        (*info).display_height = tkhd.height >> 16; // 16.16 fixed point
    } else {
        return MP4PARSE_ERROR_INVALID;
    }
    (*info).image_width = video.width;
    (*info).image_width = video.height;

    MP4PARSE_OK
}

#[test]
fn new_context() {
    let context = mp4parse_new();
    assert!(!context.is_null());
    unsafe {
        mp4parse_free(context);
    }
}

#[test]
#[should_panic(expected = "assertion failed")]
fn free_null_context() {
    unsafe {
        mp4parse_free(std::ptr::null_mut());
    }
}

#[test]
fn arg_validation() {
    let null_buffer = std::ptr::null();
    let null_context = std::ptr::null_mut();

    let context = mp4parse_new();
    assert!(!context.is_null());

    let buffer = vec![0u8; 8];

    unsafe {
        assert_eq!(MP4PARSE_ERROR_BADARG,
                   mp4parse_read(null_context, null_buffer, 0));
        assert_eq!(MP4PARSE_ERROR_BADARG,
                   mp4parse_read(context, null_buffer, 0));
    }

    for size in 0..buffer.len() {
        println!("testing buffer length {}", size);
        unsafe {
            assert_eq!(MP4PARSE_ERROR_BADARG,
                       mp4parse_read(context, buffer.as_ptr(), size));
        }
    }

    unsafe {
        mp4parse_free(context);
    }
}
<|endoftext|>


<commit_before>// only-linux
// compile-flags: -Zmiri-disable-isolation

#![feature(rustc_private)]
extern crate libc;

use std::mem::MaybeUninit;
use std::ptr;
use std::sync::atomic::AtomicI32;
use std::sync::atomic::Ordering;
use std::thread;
use std::time::{Duration, Instant};

fn wake_nobody() {
    let futex = 0;

    // Wake 1 waiter. Expect zero waiters woken up, as nobody is waiting.
    unsafe {
        assert_eq!(libc::syscall(libc::SYS_futex, &futex as *const i32, libc::FUTEX_WAKE, 1), 0);
    }

    // Same, but without omitting the unused arguments.
    unsafe {
        assert_eq!(
            libc::syscall(
                libc::SYS_futex,
                &futex as *const i32,
                libc::FUTEX_WAKE,
                1,
                ptr::null::<libc::timespec>(),
                0usize,
                0,
            ),
            0,
        );
    }
}

fn wake_dangling() {
    let futex = Box::new(0);
    let ptr: *const i32 = &*futex;
    drop(futex);

    // Wake 1 waiter. Expect zero waiters woken up, as nobody is waiting.
    unsafe {
        assert_eq!(libc::syscall(libc::SYS_futex, ptr, libc::FUTEX_WAKE, 1), 0);
    }
}

fn wait_wrong_val() {
    let futex: i32 = 123;

    // Only wait if the futex value is 456.
    unsafe {
        assert_eq!(
            libc::syscall(
                libc::SYS_futex,
                &futex as *const i32,
                libc::FUTEX_WAIT,
                456,
                ptr::null::<libc::timespec>(),
            ),
            -1,
        );
        assert_eq!(*libc::__errno_location(), libc::EAGAIN);
    }
}

fn wait_timeout() {
    let start = Instant::now();

    let futex: i32 = 123;

    // Wait for 200ms, with nobody waking us up early.
    unsafe {
        assert_eq!(
            libc::syscall(
                libc::SYS_futex,
                &futex as *const i32,
                libc::FUTEX_WAIT,
                123,
                &libc::timespec { tv_sec: 0, tv_nsec: 200_000_000 },
            ),
            -1,
        );
        assert_eq!(*libc::__errno_location(), libc::ETIMEDOUT);
    }

    assert!((200..1000).contains(&start.elapsed().as_millis()));
}

fn wait_absolute_timeout() {
    let start = Instant::now();

    // Get the current monotonic timestamp as timespec.
    let mut timeout = unsafe {
        let mut now: MaybeUninit<libc::timespec> = MaybeUninit::uninit();
        assert_eq!(libc::clock_gettime(libc::CLOCK_MONOTONIC, now.as_mut_ptr()), 0);
        now.assume_init()
    };

    // Add 200ms.
    timeout.tv_nsec += 200_000_000;
    if timeout.tv_nsec > 1_000_000_000 {
        timeout.tv_nsec -= 1_000_000_000;
        timeout.tv_sec += 1;
    }

    let futex: i32 = 123;

    // Wait for 200ms from now, with nobody waking us up early.
    unsafe {
        assert_eq!(
            libc::syscall(
                libc::SYS_futex,
                &futex as *const i32,
                libc::FUTEX_WAIT_BITSET,
                123,
                &timeout,
                0usize,
                u32::MAX,
            ),
            -1,
        );
        assert_eq!(*libc::__errno_location(), libc::ETIMEDOUT);
    }

    assert!((200..1000).contains(&start.elapsed().as_millis()));
}

fn wait_wake() {
    let start = Instant::now();

    static FUTEX: i32 = 0;

    let t = thread::spawn(move || {
        thread::sleep(Duration::from_millis(200));
        unsafe {
            assert_eq!(
                libc::syscall(
                    libc::SYS_futex,
                    &FUTEX as *const i32,
                    libc::FUTEX_WAKE,
                    10, // Wake up at most 10 threads.
                ),
                1, // Woken up one thread.
            );
        }
    });

    unsafe {
        assert_eq!(
            libc::syscall(
                libc::SYS_futex,
                &FUTEX as *const i32,
                libc::FUTEX_WAIT,
                0,
                ptr::null::<libc::timespec>(),
            ),
            0,
        );
    }

    assert!((200..1000).contains(&start.elapsed().as_millis()));
    t.join().unwrap();
}

fn wait_wake_bitset() {
    let start = Instant::now();

    static FUTEX: i32 = 0;

    let t = thread::spawn(move || {
        thread::sleep(Duration::from_millis(200));
        unsafe {
            assert_eq!(
                libc::syscall(
                    libc::SYS_futex,
                    &FUTEX as *const i32,
                    libc::FUTEX_WAKE_BITSET,
                    10, // Wake up at most 10 threads.
                    ptr::null::<libc::timespec>(),
                    0usize,
                    0b1001, // bitset
                ),
                0, // Didn't match any thread.
            );
        }
        thread::sleep(Duration::from_millis(200));
        unsafe {
            assert_eq!(
                libc::syscall(
                    libc::SYS_futex,
                    &FUTEX as *const i32,
                    libc::FUTEX_WAKE_BITSET,
                    10, // Wake up at most 10 threads.
                    ptr::null::<libc::timespec>(),
                    0usize,
                    0b0110, // bitset
                ),
                1, // Woken up one thread.
            );
        }
    });

    unsafe {
        assert_eq!(
            libc::syscall(
                libc::SYS_futex,
                &FUTEX as *const i32,
                libc::FUTEX_WAIT_BITSET,
                0,
                ptr::null::<libc::timespec>(),
                0usize,
                0b0100, // bitset
            ),
            0,
        );
    }

    assert!((400..1000).contains(&start.elapsed().as_millis()));
    t.join().unwrap();
}

const FREE: i32 = 0;
const HELD: i32 = 1;
fn concurrent_wait_wake() {
    static FUTEX: AtomicI32 = AtomicI32::new(0);
    for _ in 0..20 {
        // Suppose the main thread is holding a lock implemented using futex...
        FUTEX.store(HELD, Ordering::Relaxed);

        let t = thread::spawn(move || {
            // If this syscall runs first, then we'll be woken up by
            // the main thread's FUTEX_WAKE, and all is fine.
            //
            // If this sycall runs after the main thread's store
            // and FUTEX_WAKE, the syscall must observe that
            // the FUTEX is FREE != HELD and return without waiting
            // or we'll deadlock.
            unsafe {
                libc::syscall(
                    libc::SYS_futex,
                    &FUTEX as *const AtomicI32,
                    libc::FUTEX_WAIT,
                    HELD,
                    ptr::null::<libc::timespec>(),
                );
            }
        });

        FUTEX.store(FREE, Ordering::Relaxed);
        unsafe {
            libc::syscall(libc::SYS_futex, &FUTEX as *const AtomicI32, libc::FUTEX_WAKE, 1);
        }

        t.join().unwrap();
    }
}

fn main() {
    wake_nobody();
    wake_dangling();
    wait_wrong_val();
    wait_timeout();
    wait_absolute_timeout();
    wait_wake();
    wait_wake_bitset();
    concurrent_wait_wake();
}
<commit_msg>test that futexes induce appropriate synchronization<commit_after>// only-linux
// compile-flags: -Zmiri-disable-isolation

#![feature(rustc_private)]
extern crate libc;

use std::mem::MaybeUninit;
use std::ptr;
use std::sync::atomic::AtomicI32;
use std::sync::atomic::Ordering;
use std::thread;
use std::time::{Duration, Instant};

fn wake_nobody() {
    let futex = 0;

    // Wake 1 waiter. Expect zero waiters woken up, as nobody is waiting.
    unsafe {
        assert_eq!(libc::syscall(libc::SYS_futex, &futex as *const i32, libc::FUTEX_WAKE, 1), 0);
    }

    // Same, but without omitting the unused arguments.
    unsafe {
        assert_eq!(
            libc::syscall(
                libc::SYS_futex,
                &futex as *const i32,
                libc::FUTEX_WAKE,
                1,
                ptr::null::<libc::timespec>(),
                0usize,
                0,
            ),
            0,
        );
    }
}

fn wake_dangling() {
    let futex = Box::new(0);
    let ptr: *const i32 = &*futex;
    drop(futex);

    // Wake 1 waiter. Expect zero waiters woken up, as nobody is waiting.
    unsafe {
        assert_eq!(libc::syscall(libc::SYS_futex, ptr, libc::FUTEX_WAKE, 1), 0);
    }
}

fn wait_wrong_val() {
    let futex: i32 = 123;

    // Only wait if the futex value is 456.
    unsafe {
        assert_eq!(
            libc::syscall(
                libc::SYS_futex,
                &futex as *const i32,
                libc::FUTEX_WAIT,
                456,
                ptr::null::<libc::timespec>(),
            ),
            -1,
        );
        assert_eq!(*libc::__errno_location(), libc::EAGAIN);
    }
}

fn wait_timeout() {
    let start = Instant::now();

    let futex: i32 = 123;

    // Wait for 200ms, with nobody waking us up early.
    unsafe {
        assert_eq!(
            libc::syscall(
                libc::SYS_futex,
                &futex as *const i32,
                libc::FUTEX_WAIT,
                123,
                &libc::timespec { tv_sec: 0, tv_nsec: 200_000_000 },
            ),
            -1,
        );
        assert_eq!(*libc::__errno_location(), libc::ETIMEDOUT);
    }

    assert!((200..1000).contains(&start.elapsed().as_millis()));
}

fn wait_absolute_timeout() {
    let start = Instant::now();

    // Get the current monotonic timestamp as timespec.
    let mut timeout = unsafe {
        let mut now: MaybeUninit<libc::timespec> = MaybeUninit::uninit();
        assert_eq!(libc::clock_gettime(libc::CLOCK_MONOTONIC, now.as_mut_ptr()), 0);
        now.assume_init()
    };

    // Add 200ms.
    timeout.tv_nsec += 200_000_000;
    if timeout.tv_nsec > 1_000_000_000 {
        timeout.tv_nsec -= 1_000_000_000;
        timeout.tv_sec += 1;
    }

    let futex: i32 = 123;

    // Wait for 200ms from now, with nobody waking us up early.
    unsafe {
        assert_eq!(
            libc::syscall(
                libc::SYS_futex,
                &futex as *const i32,
                libc::FUTEX_WAIT_BITSET,
                123,
                &timeout,
                0usize,
                u32::MAX,
            ),
            -1,
        );
        assert_eq!(*libc::__errno_location(), libc::ETIMEDOUT);
    }

    assert!((200..1000).contains(&start.elapsed().as_millis()));
}

fn wait_wake() {
    let start = Instant::now();

    static FUTEX: i32 = 0;

    let t = thread::spawn(move || {
        thread::sleep(Duration::from_millis(200));
        unsafe {
            assert_eq!(
                libc::syscall(
                    libc::SYS_futex,
                    &FUTEX as *const i32,
                    libc::FUTEX_WAKE,
                    10, // Wake up at most 10 threads.
                ),
                1, // Woken up one thread.
            );
        }
    });

    unsafe {
        assert_eq!(
            libc::syscall(
                libc::SYS_futex,
                &FUTEX as *const i32,
                libc::FUTEX_WAIT,
                0,
                ptr::null::<libc::timespec>(),
            ),
            0,
        );
    }

    assert!((200..1000).contains(&start.elapsed().as_millis()));
    t.join().unwrap();
}

fn wait_wake_bitset() {
    let start = Instant::now();

    static FUTEX: i32 = 0;

    let t = thread::spawn(move || {
        thread::sleep(Duration::from_millis(200));
        unsafe {
            assert_eq!(
                libc::syscall(
                    libc::SYS_futex,
                    &FUTEX as *const i32,
                    libc::FUTEX_WAKE_BITSET,
                    10, // Wake up at most 10 threads.
                    ptr::null::<libc::timespec>(),
                    0usize,
                    0b1001, // bitset
                ),
                0, // Didn't match any thread.
            );
        }
        thread::sleep(Duration::from_millis(200));
        unsafe {
            assert_eq!(
                libc::syscall(
                    libc::SYS_futex,
                    &FUTEX as *const i32,
                    libc::FUTEX_WAKE_BITSET,
                    10, // Wake up at most 10 threads.
                    ptr::null::<libc::timespec>(),
                    0usize,
                    0b0110, // bitset
                ),
                1, // Woken up one thread.
            );
        }
    });

    unsafe {
        assert_eq!(
            libc::syscall(
                libc::SYS_futex,
                &FUTEX as *const i32,
                libc::FUTEX_WAIT_BITSET,
                0,
                ptr::null::<libc::timespec>(),
                0usize,
                0b0100, // bitset
            ),
            0,
        );
    }

    assert!((400..1000).contains(&start.elapsed().as_millis()));
    t.join().unwrap();
}

fn concurrent_wait_wake() {
    const FREE: i32 = 0;
    const HELD: i32 = 1;

    static FUTEX: AtomicI32 = AtomicI32::new(0);
    static mut DATA: i32 = 0;
    static WOKEN: AtomicI32 = AtomicI32::new(0);

    let rounds = 50;
    for _ in 0..rounds {
        unsafe { DATA = 0 }; // Reset
        // Suppose the main thread is holding a lock implemented using futex...
        FUTEX.store(HELD, Ordering::Relaxed);

        let t = thread::spawn(move || {
            // If this syscall runs first, then we'll be woken up by
            // the main thread's FUTEX_WAKE, and all is fine.
            //
            // If this sycall runs after the main thread's store
            // and FUTEX_WAKE, the syscall must observe that
            // the FUTEX is FREE != HELD and return without waiting
            // or we'll deadlock.
            unsafe {
                let ret = libc::syscall(
                    libc::SYS_futex,
                    &FUTEX as *const AtomicI32,
                    libc::FUTEX_WAIT,
                    HELD,
                    ptr::null::<libc::timespec>(),
                );
                if ret == 0 {
                    // We actually slept. And then woke up again. So we should be ordered-after
                    // what happened-before the FUTEX_WAKE. So this is not a race.
                    assert_eq!(DATA, 1);
                    // Also remember that this happened at least once.
                    WOKEN.fetch_add(1, Ordering::Relaxed);
                }
            }
        });
        // Increase the chance that the other thread actually goes to sleep.
        // (5 yields in a loop seem to make that happen around 40% of the time.)
        for _ in 0..5 {
            thread::yield_now();
        }

        FUTEX.store(FREE, Ordering::Relaxed);
        unsafe {
            DATA = 1;
            libc::syscall(libc::SYS_futex, &FUTEX as *const AtomicI32, libc::FUTEX_WAKE, 1);
        }

        t.join().unwrap();
    }

    // Make sure we got the interesting case (of having woken a thread) at least once, but not *each* time.
    let woken = WOKEN.load(Ordering::Relaxed);
    assert!(woken > 0 && woken < rounds);
    //eprintln!("waking happened {woken} times");
}

fn main() {
    wake_nobody();
    wake_dangling();
    wait_wrong_val();
    wait_timeout();
    wait_absolute_timeout();
    wait_wake();
    wait_wake_bitset();
    concurrent_wait_wake();
}
<|endoftext|>