<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Apache Arrow Rust Benchmarks

use std::collections::HashMap;
use std::path::PathBuf;
use std::process;
use std::time::Instant;

use arrow::datatypes::{DataType, Field, Schema};
use arrow::util::pretty;
use datafusion::error::Result;
use datafusion::execution::context::ExecutionContext;

use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "Benchmarks", about = "Apache Arrow Rust Benchmarks.")]
struct Opt {
    /// Activate debug mode to see query results
    #[structopt(short, long)]
    debug: bool,

    /// Number of iterations of each test run
    #[structopt(short = "i", long = "iterations", default_value = "3")]
    iterations: usize,

    /// Batch size when reading CSV or Parquet files
    #[structopt(short = "s", long = "batch-size", default_value = "4096")]
    batch_size: usize,

    /// Path to data files
    #[structopt(parse(from_os_str), required = true, short = "p", long = "path")]
    path: PathBuf,

    /// File format: `csv` or `parquet`
    #[structopt(short = "f", long = "format", default_value = "csv")]
    file_format: String,
}

fn main() -> Result<()> {
    let opt = Opt::from_args();
    println!("Running benchmarks with the following options: {:?}", opt);

    let mut ctx = ExecutionContext::new();

    let path = opt.path.to_str().unwrap();

    match opt.file_format.as_str() {
        "csv" => ctx.register_csv("tripdata", path, &nyctaxi_schema(), true),
        "parquet" => ctx.register_parquet("tripdata", path)?,
        other => {
            println!("Invalid file format '{}'", other);
            process::exit(-1);
        }
    }

    datafusion_sql_benchmarks(&mut ctx, opt.iterations, opt.batch_size, opt.debug)
}

fn datafusion_sql_benchmarks(
    ctx: &mut ExecutionContext,
    iterations: usize,
    batch_size: usize,
    debug: bool,
) -> Result<()> {
    let mut queries = HashMap::new();
    queries.insert("fare_amt_by_passenger", "SELECT passenger_count, MIN(fare_amount), MIN(fare_amount), SUM(fare_amount) FROM tripdata GROUP BY passenger_count");
    for (name, sql) in &queries {
        println!("Executing '{}'", name);
        for i in 0..iterations {
            let start = Instant::now();
            execute_sql(ctx, sql, batch_size, debug)?;
            println!(
                "Query '{}' iteration {} took {} ms",
                name,
                i,
                start.elapsed().as_millis()
            );
        }
    }
    Ok(())
}

fn execute_sql(
    ctx: &mut ExecutionContext,
    sql: &str,
    batch_size: usize,
    debug: bool,
) -> Result<()> {
    let plan = ctx.create_logical_plan(sql)?;
    let plan = ctx.optimize(&plan)?;
    if debug {
        println!("Optimized logical plan:\n{:?}", plan);
    }
    let physical_plan = ctx.create_physical_plan(&plan, batch_size)?;
    let result = ctx.collect(physical_plan.as_ref())?;
    if debug {
        pretty::print_batches(&result)?;
    }
    Ok(())
}

fn nyctaxi_schema() -> Schema {
    Schema::new(vec![
        Field::new("VendorID", DataType::Utf8, true),
        Field::new("tpep_pickup_datetime", DataType::Utf8, true),
        Field::new("tpep_dropoff_datetime", DataType::Utf8, true),
        Field::new("passenger_count", DataType::Int32, true),
        Field::new("trip_distance", DataType::Utf8, true),
        Field::new("RatecodeID", DataType::Utf8, true),
        Field::new("store_and_fwd_flag", DataType::Utf8, true),
        Field::new("PULocationID", DataType::Utf8, true),
        Field::new("DOLocationID", DataType::Utf8, true),
        Field::new("payment_type", DataType::Utf8, true),
        Field::new("fare_amount", DataType::Float64, true),
        Field::new("extra", DataType::Float64, true),
        Field::new("mta_tax", DataType::Float64, true),
        Field::new("tip_amount", DataType::Float64, true),
        Field::new("tolls_amount", DataType::Float64, true),
        Field::new("improvement_surcharge", DataType::Float64, true),
        Field::new("total_amount", DataType::Float64, true),
    ])
}
<commit_msg>ARROW-8959: [Rust] Update benchmark to use new API (fixes broken build)<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Apache Arrow Rust Benchmarks

use std::collections::HashMap;
use std::path::PathBuf;
use std::process;
use std::time::Instant;

use arrow::datatypes::{DataType, Field, Schema};
use arrow::util::pretty;
use datafusion::error::Result;
use datafusion::execution::context::ExecutionContext;

use datafusion::execution::physical_plan::csv::CsvReadOptions;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "Benchmarks", about = "Apache Arrow Rust Benchmarks.")]
struct Opt {
    /// Activate debug mode to see query results
    #[structopt(short, long)]
    debug: bool,

    /// Number of iterations of each test run
    #[structopt(short = "i", long = "iterations", default_value = "3")]
    iterations: usize,

    /// Batch size when reading CSV or Parquet files
    #[structopt(short = "s", long = "batch-size", default_value = "4096")]
    batch_size: usize,

    /// Path to data files
    #[structopt(parse(from_os_str), required = true, short = "p", long = "path")]
    path: PathBuf,

    /// File format: `csv` or `parquet`
    #[structopt(short = "f", long = "format", default_value = "csv")]
    file_format: String,
}

fn main() -> Result<()> {
    let opt = Opt::from_args();
    println!("Running benchmarks with the following options: {:?}", opt);

    let mut ctx = ExecutionContext::new();

    let path = opt.path.to_str().unwrap();

    match opt.file_format.as_str() {
        "csv" => {
            let schema = nyctaxi_schema();
            let options = CsvReadOptions::new().schema(&schema).has_header(true);
            ctx.register_csv("tripdata", path, options)?
        }
        "parquet" => ctx.register_parquet("tripdata", path)?,
        other => {
            println!("Invalid file format '{}'", other);
            process::exit(-1);
        }
    }

    datafusion_sql_benchmarks(&mut ctx, opt.iterations, opt.batch_size, opt.debug)
}

fn datafusion_sql_benchmarks(
    ctx: &mut ExecutionContext,
    iterations: usize,
    batch_size: usize,
    debug: bool,
) -> Result<()> {
    let mut queries = HashMap::new();
    queries.insert("fare_amt_by_passenger", "SELECT passenger_count, MIN(fare_amount), MIN(fare_amount), SUM(fare_amount) FROM tripdata GROUP BY passenger_count");
    for (name, sql) in &queries {
        println!("Executing '{}'", name);
        for i in 0..iterations {
            let start = Instant::now();
            execute_sql(ctx, sql, batch_size, debug)?;
            println!(
                "Query '{}' iteration {} took {} ms",
                name,
                i,
                start.elapsed().as_millis()
            );
        }
    }
    Ok(())
}

fn execute_sql(
    ctx: &mut ExecutionContext,
    sql: &str,
    batch_size: usize,
    debug: bool,
) -> Result<()> {
    let plan = ctx.create_logical_plan(sql)?;
    let plan = ctx.optimize(&plan)?;
    if debug {
        println!("Optimized logical plan:\n{:?}", plan);
    }
    let physical_plan = ctx.create_physical_plan(&plan, batch_size)?;
    let result = ctx.collect(physical_plan.as_ref())?;
    if debug {
        pretty::print_batches(&result)?;
    }
    Ok(())
}

fn nyctaxi_schema() -> Schema {
    Schema::new(vec![
        Field::new("VendorID", DataType::Utf8, true),
        Field::new("tpep_pickup_datetime", DataType::Utf8, true),
        Field::new("tpep_dropoff_datetime", DataType::Utf8, true),
        Field::new("passenger_count", DataType::Int32, true),
        Field::new("trip_distance", DataType::Utf8, true),
        Field::new("RatecodeID", DataType::Utf8, true),
        Field::new("store_and_fwd_flag", DataType::Utf8, true),
        Field::new("PULocationID", DataType::Utf8, true),
        Field::new("DOLocationID", DataType::Utf8, true),
        Field::new("payment_type", DataType::Utf8, true),
        Field::new("fare_amount", DataType::Float64, true),
        Field::new("extra", DataType::Float64, true),
        Field::new("mta_tax", DataType::Float64, true),
        Field::new("tip_amount", DataType::Float64, true),
        Field::new("tolls_amount", DataType::Float64, true),
        Field::new("improvement_surcharge", DataType::Float64, true),
        Field::new("total_amount", DataType::Float64, true),
    ])
}
<|endoftext|>


<commit_before>use fasteval::{ez_eval, Error};
use kernel::{
    async_trait::async_trait,
    eyre::{bail, Result},
    serde::Serialize,
    stencila_schema::{CodeError, Node, Primitive},
    Kernel, KernelStatus, KernelTrait, KernelType, Task, TaskResult,
};
use once_cell::sync::Lazy;
use regex::Regex;
use std::collections::BTreeMap;

/// A kernel that evaluates simple calculator like numerical expressions
///
/// Based on [`fasteval`](https://github.com/likebike/fasteval). See the
/// `fasteval` docs for more on the syntax and functions supported.
#[derive(Debug, Clone, Default, Serialize)]
#[serde(crate = "kernel::serde")]
pub struct CalcKernel {
    symbols: BTreeMap<String, f64>,
}

impl CalcKernel {
    pub fn new() -> Self {
        Self::default()
    }
}

#[async_trait]
impl KernelTrait for CalcKernel {
    fn spec(&self) -> Kernel {
        Kernel::new("calc", KernelType::Builtin, &["calc"], true)
    }

    async fn status(&self) -> Result<KernelStatus> {
        Ok(KernelStatus::Idle)
    }

    async fn get(&mut self, name: &str) -> Result<Node> {
        match self.symbols.get(name) {
            Some(number) => Ok(Node::Number(*number)),
            None => bail!("Symbol `{}` does not exist in this kernel", name),
        }
    }

    async fn set(&mut self, name: &str, value: Node) -> Result<()> {
        let value = match value {
            Node::Null(..) => 0.,
            Node::Boolean(boolean) => match boolean {
                true => 1.,
                false => 0.,
            },
            Node::Integer(integer) => integer as f64,
            Node::Number(number) => number,
            Node::Array(array) => match array.first() {
                Some(Primitive::Null(..)) => 0.,
                Some(Primitive::Boolean(boolean)) => match boolean {
                    true => 1.,
                    false => 0.,
                },
                Some(Primitive::Integer(integer)) => *integer as f64,
                Some(Primitive::Number(number)) => *number,
                None => 0.,
                _ => bail!(
                    "Unable to convert first item of array to a number for use in Calc kernel"
                ),
            },
            _ => bail!(
                "Node is of type that can not be converted to a number for use in Calc kernel"
            ),
        };
        self.symbols.insert(name.to_string(), value);
        Ok(())
    }

    async fn exec_sync(&mut self, code: &str) -> Result<Task> {
        static STATEMENTS_REGEX: Lazy<Regex> =
            Lazy::new(|| Regex::new(r"\r?\n|;").expect("Unable to create regex"));
        static ASSIGN_REGEX: Lazy<Regex> = Lazy::new(|| {
            Regex::new(r"\s*([a-zA-Z_][a-zA-Z_0-9]*)\s*=(.*)").expect("Unable to create regex")
        });

        let mut task = Task::start_sync();
        let mut outputs = Vec::new();
        let mut messages = Vec::new();
        for statement in STATEMENTS_REGEX.split(code) {
            let statement = statement.trim();

            // Skip the statement if it is blank or is a comment
            if statement.is_empty() || statement.starts_with('#') {
                continue;
            }

            // Get the name of any assigned variable, and the expression to be evaluated
            let (symbol, expr) = if let Some(captures) = ASSIGN_REGEX.captures(statement) {
                let symbol = captures.get(1).expect("Should always have group 1");
                let expr = captures.get(2).expect("Should always have group 2");
                (Some(symbol.as_str()), expr.as_str())
            } else {
                (None, statement)
            };

            // A fasteval callback function that defines some custom functions
            let mut cb = |name: &str, _args: Vec<f64>| -> Option<f64> {
                match name {
                    "now" => now(),
                    _ => self.symbols.get(name).copied(),
                }
            };

            // Evaluate the expression
            match ez_eval(expr, &mut cb) {
                Ok(num) => {
                    // Either assign the result, or add it to outputs
                    if let Some(symbol) = symbol {
                        self.symbols.insert(symbol.to_string(), num);
                    } else {
                        outputs.push(Node::Number(num))
                    }
                }
                Err(error) => {
                    let error_message = match error {
                        // Custom error strings for common errors
                        Error::EOF | Error::EofWhileParsing(..) => {
                            "Unexpected end of Calc expression".to_string()
                        }
                        Error::Undefined(name) => {
                            format!("Undefined variable or function: {}", name)
                        }
                        Error::WrongArgs(msg) => {
                            format!("Function called with wrong number of arguments: {}", msg)
                        }
                        Error::InvalidValue => "Unexpected value in expression".to_string(),
                        Error::TooLong => "Calc expression was too long".to_string(),
                        Error::TooDeep => "Calc expression was too recursive".to_string(),
                        // Use the debug string for others
                        _ => format!("Could not execute Calc expression: {:?}", error),
                    };
                    messages.push(CodeError {
                        error_message,
                        ..Default::default()
                    });
                }
            }
        }
        task.finished(TaskResult::new(outputs, messages));
        Ok(task)
    }

    async fn exec_fork(&mut self, code: &str) -> Result<Task> {
        let mut fork = self.clone();
        fork.exec_async(code).await
    }
}

// Custom functions

/// The current system time as seconds (to millisecond resolution) since the Unix epoch
fn now() -> Option<f64> {
    use std::time::{SystemTime, UNIX_EPOCH};

    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .ok()
        .map(|duration| duration.as_millis() as f64 / 1000.0)
}

#[cfg(test)]
mod tests {
    use super::*;
    use kernel::{KernelStatus, KernelTrait};
    use test_utils::{assert_json_eq, serde_json::json};

    #[tokio::test]
    async fn status() -> Result<()> {
        let kernel = CalcKernel::new();

        assert_eq!(kernel.status().await?, KernelStatus::Idle);

        Ok(())
    }

    #[tokio::test]
    async fn get_set_exec() -> Result<()> {
        let mut kernel = CalcKernel::new();

        match kernel.get("a").await {
            Ok(..) => bail!("Expected an error"),
            Err(error) => assert!(error.to_string().contains("does not exist")),
        };

        match kernel.set("a", Node::String("A".to_string())).await {
            Ok(..) => bail!("Expected an error"),
            Err(error) => assert!(error
                .to_string()
                .contains("Unable to convert node to a number")),
        };

        kernel.set("a", Node::Number(1.23)).await?;

        let a = kernel.get("a").await?;
        assert!(matches!(a, Node::Number(..)));
        assert_json_eq!(a, json!(1.23));

        let (outputs, errors) = kernel.exec("a * 2").await?;
        assert_json_eq!(outputs, json!([2.46]));
        assert_eq!(errors.len(), 0);

        let (outputs, errors) = kernel.exec("x * 2").await?;
        assert_eq!(outputs.len(), 0);
        assert_json_eq!(
            errors,
            json!([{"type": "CodeError", "errorMessage": "Undefined variable or function: x"}])
        );

        Ok(())
    }
}
<commit_msg>fix(Calc kernel): Handle conversion of strings to numbers<commit_after>use fasteval::{ez_eval, Error};
use kernel::{
    async_trait::async_trait,
    eyre::{bail, Result},
    serde::Serialize,
    stencila_schema::{CodeError, Node, Primitive},
    Kernel, KernelStatus, KernelTrait, KernelType, Task, TaskResult,
};
use once_cell::sync::Lazy;
use regex::Regex;
use std::collections::BTreeMap;

/// A kernel that evaluates simple calculator like numerical expressions
///
/// Based on [`fasteval`](https://github.com/likebike/fasteval). See the
/// `fasteval` docs for more on the syntax and functions supported.
#[derive(Debug, Clone, Default, Serialize)]
#[serde(crate = "kernel::serde")]
pub struct CalcKernel {
    symbols: BTreeMap<String, f64>,
}

impl CalcKernel {
    pub fn new() -> Self {
        Self::default()
    }
}

#[async_trait]
impl KernelTrait for CalcKernel {
    fn spec(&self) -> Kernel {
        Kernel::new("calc", KernelType::Builtin, &["calc"], true)
    }

    async fn status(&self) -> Result<KernelStatus> {
        Ok(KernelStatus::Idle)
    }

    async fn get(&mut self, name: &str) -> Result<Node> {
        match self.symbols.get(name) {
            Some(number) => Ok(Node::Number(*number)),
            None => bail!("Symbol `{}` does not exist in this kernel", name),
        }
    }

    async fn set(&mut self, name: &str, value: Node) -> Result<()> {
        let value = match value {
            Node::Null(..) => 0.,
            Node::Boolean(boolean) => match boolean {
                true => 1.,
                false => 0.,
            },
            Node::Integer(integer) => integer as f64,
            Node::Number(number) => number,
            Node::String(string) => match string.trim().parse() {
                Ok(number) => number,
                Err(..) => bail!("Unable to convert string `{}` to a number", string),
            },
            Node::Array(array) => match array.first() {
                Some(value) => match value {
                    Primitive::Null(..) => 0.,
                    Primitive::Boolean(boolean) => match boolean {
                        true => 1.,
                        false => 0.,
                    },
                    Primitive::Integer(integer) => *integer as f64,
                    Primitive::Number(number) => *number,
                    Primitive::String(string) => match string.trim().parse() {
                        Ok(number) => number,
                        Err(..) => bail!("Unable to convert string `{}` to a number", string),
                    },
                    _ => bail!("Unable to convert first item of array to a number"),
                },
                _ => bail!("Unable to convert empty array to a number"),
            },
            _ => bail!("Node is of type that can not be converted to a number"),
        };
        self.symbols.insert(name.to_string(), value);
        Ok(())
    }

    async fn exec_sync(&mut self, code: &str) -> Result<Task> {
        static STATEMENTS_REGEX: Lazy<Regex> =
            Lazy::new(|| Regex::new(r"\r?\n|;").expect("Unable to create regex"));
        static ASSIGN_REGEX: Lazy<Regex> = Lazy::new(|| {
            Regex::new(r"\s*([a-zA-Z_][a-zA-Z_0-9]*)\s*=(.*)").expect("Unable to create regex")
        });

        let mut task = Task::start_sync();
        let mut outputs = Vec::new();
        let mut messages = Vec::new();
        for statement in STATEMENTS_REGEX.split(code) {
            let statement = statement.trim();

            // Skip the statement if it is blank or is a comment
            if statement.is_empty() || statement.starts_with('#') {
                continue;
            }

            // Get the name of any assigned variable, and the expression to be evaluated
            let (symbol, expr) = if let Some(captures) = ASSIGN_REGEX.captures(statement) {
                let symbol = captures.get(1).expect("Should always have group 1");
                let expr = captures.get(2).expect("Should always have group 2");
                (Some(symbol.as_str()), expr.as_str())
            } else {
                (None, statement)
            };

            // A fasteval callback function that defines some custom functions
            let mut cb = |name: &str, _args: Vec<f64>| -> Option<f64> {
                match name {
                    "now" => now(),
                    _ => self.symbols.get(name).copied(),
                }
            };

            // Evaluate the expression
            match ez_eval(expr, &mut cb) {
                Ok(num) => {
                    // Either assign the result, or add it to outputs
                    if let Some(symbol) = symbol {
                        self.symbols.insert(symbol.to_string(), num);
                    } else {
                        outputs.push(Node::Number(num))
                    }
                }
                Err(error) => {
                    let error_message = match error {
                        // Custom error strings for common errors
                        Error::EOF | Error::EofWhileParsing(..) => {
                            "Unexpected end of Calc expression".to_string()
                        }
                        Error::Undefined(name) => {
                            format!("Undefined variable or function: {}", name)
                        }
                        Error::WrongArgs(msg) => {
                            format!("Function called with wrong number of arguments: {}", msg)
                        }
                        Error::InvalidValue => "Unexpected value in expression".to_string(),
                        Error::TooLong => "Calc expression was too long".to_string(),
                        Error::TooDeep => "Calc expression was too recursive".to_string(),
                        // Use the debug string for others
                        _ => format!("Could not execute Calc expression: {:?}", error),
                    };
                    messages.push(CodeError {
                        error_message,
                        ..Default::default()
                    });
                }
            }
        }
        task.finished(TaskResult::new(outputs, messages));
        Ok(task)
    }

    async fn exec_fork(&mut self, code: &str) -> Result<Task> {
        let mut fork = self.clone();
        fork.exec_async(code).await
    }
}

// Custom functions

/// The current system time as seconds (to millisecond resolution) since the Unix epoch
fn now() -> Option<f64> {
    use std::time::{SystemTime, UNIX_EPOCH};

    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .ok()
        .map(|duration| duration.as_millis() as f64 / 1000.0)
}

#[cfg(test)]
mod tests {
    use super::*;
    use kernel::{KernelStatus, KernelTrait};
    use test_utils::{assert_json_eq, serde_json::json};

    #[tokio::test]
    async fn status() -> Result<()> {
        let kernel = CalcKernel::new();

        assert_eq!(kernel.status().await?, KernelStatus::Idle);

        Ok(())
    }

    #[tokio::test]
    async fn get_set_exec() -> Result<()> {
        let mut kernel = CalcKernel::new();

        match kernel.get("a").await {
            Ok(..) => bail!("Expected an error"),
            Err(error) => assert!(error.to_string().contains("does not exist")),
        };

        match kernel.set("a", Node::String("A".to_string())).await {
            Ok(..) => bail!("Expected an error"),
            Err(error) => assert!(error
                .to_string()
                .contains("Unable to convert string `A` to a number")),
        };

        kernel.set("a", Node::Number(1.23)).await?;

        let a = kernel.get("a").await?;
        assert!(matches!(a, Node::Number(..)));
        assert_json_eq!(a, json!(1.23));

        let (outputs, errors) = kernel.exec("a * 2").await?;
        assert_json_eq!(outputs, json!([2.46]));
        assert_eq!(errors.len(), 0);

        let (outputs, errors) = kernel.exec("x * 2").await?;
        assert_eq!(outputs.len(), 0);
        assert_json_eq!(
            errors,
            json!([{"type": "CodeError", "errorMessage": "Undefined variable or function: x"}])
        );

        Ok(())
    }
}
<|endoftext|>


<commit_before>use crate::wrapper::binary::{alloc, realloc, ErlNifBinary};
use crate::{Decoder, Encoder, Env, Error, NifResult, Term};

use std::borrow::{Borrow, BorrowMut};
use std::io::Write;
use std::mem::MaybeUninit;
use std::ops::{Deref, DerefMut};

// Owned

pub struct OwnedBinary {
    inner: ErlNifBinary,
    release: bool,
}

impl<'a> OwnedBinary {
    pub unsafe fn from_raw(inner: ErlNifBinary) -> OwnedBinary {
        OwnedBinary {
            inner,
            release: true,
        }
    }

    /// Allocates a new OwnedBinary with size `size`.
    ///
    /// Note that the memory is not initially guaranteed to be any particular value.
    /// If an empty buffer is needed, you should manually zero it.
    pub fn new(size: usize) -> Option<OwnedBinary> {
        unsafe { alloc(size) }.map(|binary| OwnedBinary {
            inner: binary,
            release: true,
        })
    }

    /// Copies a given `Binary`.
    pub fn from_unowned(from_bin: &Binary) -> Option<OwnedBinary> {
        let len = from_bin.len();
        if let Some(mut bin) = OwnedBinary::new(len) {
            if let Ok(write_len) = bin.as_mut_slice().write(from_bin.as_slice()) {
                if write_len != len {
                    panic!("Could not copy binary");
                }
                Some(bin)
            } else {
                panic!("Could not copy binary");
            }
        } else {
            None
        }
    }

    /// Attempts to reallocate the buffer with the new size.
    /// Returns false if the buffer cannot be reallocated.
    #[must_use]
    pub fn realloc(&mut self, size: usize) -> bool {
        unsafe { realloc(&mut self.inner, size) }
    }

    /// Attempts to reallocate the buffer with the new size.
    /// If reallocation fails, it will perform a copy instead.

    /// Memory outside the range of the original buffer will
    /// not be initialized. If this needs to be empty, clear it manually.
    pub fn realloc_or_copy(&mut self, size: usize) {
        if !self.realloc(size) {
            let mut new = OwnedBinary::new(size).unwrap();
            if let Ok(num_written) = new.as_mut_slice().write(self.as_slice()) {
                if !(num_written == self.len() || num_written == new.len()) {
                    panic!("Could not copy binary");
                }
                ::std::mem::swap(&mut self.inner, &mut new.inner);
            } else {
                panic!("Could not copy binary");
            }
        }
    }

    pub fn as_slice(&self) -> &'a [u8] {
        unsafe { ::std::slice::from_raw_parts(self.inner.data, self.inner.size) }
    }

    pub fn as_mut_slice(&mut self) -> &'a mut [u8] {
        unsafe { ::std::slice::from_raw_parts_mut(self.inner.data, self.inner.size) }
    }

    /// Releases control of the binary to the VM. After this point
    /// the binary will be immutable.
    pub fn release(self, env: Env) -> Binary {
        Binary::from_owned(self, env)
    }
}

impl Borrow<[u8]> for OwnedBinary {
    fn borrow(&self) -> &[u8] {
        self.as_slice()
    }
}
impl BorrowMut<[u8]> for OwnedBinary {
    fn borrow_mut(&mut self) -> &mut [u8] {
        self.as_mut_slice()
    }
}
impl Deref for OwnedBinary {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        self.as_slice()
    }
}
impl DerefMut for OwnedBinary {
    fn deref_mut(&mut self) -> &mut [u8] {
        self.as_mut_slice()
    }
}

impl Drop for OwnedBinary {
    fn drop(&mut self) {
        if self.release {
            unsafe { rustler_sys::enif_release_binary(&mut self.inner) };
        }
    }
}

unsafe impl Send for OwnedBinary {}

// Borrowed

pub struct Binary<'a> {
    inner: ErlNifBinary,
    term: Term<'a>,
}

impl<'a> Binary<'a> {
    /// Releases a given `OwnedBinary` to the VM.
    /// After this point the binary will be immutable.
    pub fn from_owned(mut bin: OwnedBinary, env: Env<'a>) -> Self {
        bin.release = false;
        let term = unsafe {
            Term::new(
                env,
                rustler_sys::enif_make_binary(env.as_c_arg(), &mut bin.inner),
            )
        };
        Binary {
            inner: bin.inner,
            term,
        }
    }

    pub fn to_owned(&self) -> Option<OwnedBinary> {
        OwnedBinary::from_unowned(self)
    }

    pub fn from_term(term: Term<'a>) -> Result<Self, Error> {
        let mut binary = MaybeUninit::uninit();
        if unsafe {
            rustler_sys::enif_inspect_binary(
                term.get_env().as_c_arg(),
                term.as_c_arg(),
                binary.as_mut_ptr(),
            )
        } == 0
        {
            return Err(Error::BadArg);
        }
        Ok(Binary {
            inner: unsafe { binary.assume_init() },
            term,
        })
    }

    pub fn from_iolist(term: Term<'a>) -> Result<Self, Error> {
        let mut binary = MaybeUninit::uninit();
        if unsafe {
            rustler_sys::enif_inspect_iolist_as_binary(
                term.get_env().as_c_arg(),
                term.as_c_arg(),
                binary.as_mut_ptr(),
            )
        } == 0
        {
            return Err(Error::BadArg);
        }
        Ok(Binary {
            inner: unsafe { binary.assume_init() },
            term,
        })
    }

    pub fn to_term<'b>(&self, env: Env<'b>) -> Term<'b> {
        self.term.in_env(env)
    }

    pub fn as_slice(&self) -> &'a [u8] {
        unsafe { ::std::slice::from_raw_parts(self.inner.data, self.inner.size) }
    }

    /// Returns a new view into the same binary.
    /// This will not copy anything.
    pub fn make_subbinary(&self, offset: usize, length: usize) -> NifResult<Binary<'a>> {
        let min_len = length.checked_add(offset);
        if min_len.ok_or(Error::BadArg)? > self.inner.size {
            return Err(Error::BadArg);
        }

        let raw_term = unsafe {
            rustler_sys::enif_make_sub_binary(
                self.term.get_env().as_c_arg(),
                self.term.as_c_arg(),
                offset,
                length,
            )
        };
        let term = unsafe { Term::new(self.term.get_env(), raw_term) };
        // This should never fail, as we are always passing in a binary term.
        Ok(Binary::from_term(term).ok().unwrap())
    }
}

impl<'a> Borrow<[u8]> for Binary<'a> {
    fn borrow(&self) -> &[u8] {
        self.as_slice()
    }
}
impl<'a> Deref for Binary<'a> {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        self.as_slice()
    }
}

impl<'a> Decoder<'a> for Binary<'a> {
    fn decode(term: Term<'a>) -> Result<Self, Error> {
        Binary::from_term(term)
    }
}
impl<'a> Encoder for Binary<'a> {
    fn encode<'b>(&self, env: Env<'b>) -> Term<'b> {
        self.to_term(env)
    }
}

/// ## Binary terms
impl<'a> Term<'a> {
    pub fn into_binary(self) -> NifResult<Binary<'a>> {
        Binary::from_term(self)
    }
}
<commit_msg>Readd Copy and Clone on Binary<commit_after>use crate::wrapper::binary::{alloc, realloc, ErlNifBinary};
use crate::{Decoder, Encoder, Env, Error, NifResult, Term};

use std::borrow::{Borrow, BorrowMut};
use std::io::Write;
use std::mem::MaybeUninit;
use std::ops::{Deref, DerefMut};

// Owned

pub struct OwnedBinary {
    inner: ErlNifBinary,
    release: bool,
}

impl<'a> OwnedBinary {
    pub unsafe fn from_raw(inner: ErlNifBinary) -> OwnedBinary {
        OwnedBinary {
            inner,
            release: true,
        }
    }

    /// Allocates a new OwnedBinary with size `size`.
    ///
    /// Note that the memory is not initially guaranteed to be any particular value.
    /// If an empty buffer is needed, you should manually zero it.
    pub fn new(size: usize) -> Option<OwnedBinary> {
        unsafe { alloc(size) }.map(|binary| OwnedBinary {
            inner: binary,
            release: true,
        })
    }

    /// Copies a given `Binary`.
    pub fn from_unowned(from_bin: &Binary) -> Option<OwnedBinary> {
        let len = from_bin.len();
        if let Some(mut bin) = OwnedBinary::new(len) {
            if let Ok(write_len) = bin.as_mut_slice().write(from_bin.as_slice()) {
                if write_len != len {
                    panic!("Could not copy binary");
                }
                Some(bin)
            } else {
                panic!("Could not copy binary");
            }
        } else {
            None
        }
    }

    /// Attempts to reallocate the buffer with the new size.
    /// Returns false if the buffer cannot be reallocated.
    #[must_use]
    pub fn realloc(&mut self, size: usize) -> bool {
        unsafe { realloc(&mut self.inner, size) }
    }

    /// Attempts to reallocate the buffer with the new size.
    /// If reallocation fails, it will perform a copy instead.

    /// Memory outside the range of the original buffer will
    /// not be initialized. If this needs to be empty, clear it manually.
    pub fn realloc_or_copy(&mut self, size: usize) {
        if !self.realloc(size) {
            let mut new = OwnedBinary::new(size).unwrap();
            if let Ok(num_written) = new.as_mut_slice().write(self.as_slice()) {
                if !(num_written == self.len() || num_written == new.len()) {
                    panic!("Could not copy binary");
                }
                ::std::mem::swap(&mut self.inner, &mut new.inner);
            } else {
                panic!("Could not copy binary");
            }
        }
    }

    pub fn as_slice(&self) -> &'a [u8] {
        unsafe { ::std::slice::from_raw_parts(self.inner.data, self.inner.size) }
    }

    pub fn as_mut_slice(&mut self) -> &'a mut [u8] {
        unsafe { ::std::slice::from_raw_parts_mut(self.inner.data, self.inner.size) }
    }

    /// Releases control of the binary to the VM. After this point
    /// the binary will be immutable.
    pub fn release(self, env: Env) -> Binary {
        Binary::from_owned(self, env)
    }
}

impl Borrow<[u8]> for OwnedBinary {
    fn borrow(&self) -> &[u8] {
        self.as_slice()
    }
}
impl BorrowMut<[u8]> for OwnedBinary {
    fn borrow_mut(&mut self) -> &mut [u8] {
        self.as_mut_slice()
    }
}
impl Deref for OwnedBinary {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        self.as_slice()
    }
}
impl DerefMut for OwnedBinary {
    fn deref_mut(&mut self) -> &mut [u8] {
        self.as_mut_slice()
    }
}

impl Drop for OwnedBinary {
    fn drop(&mut self) {
        if self.release {
            unsafe { rustler_sys::enif_release_binary(&mut self.inner) };
        }
    }
}

unsafe impl Send for OwnedBinary {}

// Borrowed
#[derive(Copy, Clone)]
pub struct Binary<'a> {
    inner: ErlNifBinary,
    term: Term<'a>,
}

impl<'a> Binary<'a> {
    /// Releases a given `OwnedBinary` to the VM.
    /// After this point the binary will be immutable.
    pub fn from_owned(mut bin: OwnedBinary, env: Env<'a>) -> Self {
        bin.release = false;
        let term = unsafe {
            Term::new(
                env,
                rustler_sys::enif_make_binary(env.as_c_arg(), &mut bin.inner),
            )
        };
        Binary {
            inner: bin.inner,
            term,
        }
    }

    pub fn to_owned(&self) -> Option<OwnedBinary> {
        OwnedBinary::from_unowned(self)
    }

    pub fn from_term(term: Term<'a>) -> Result<Self, Error> {
        let mut binary = MaybeUninit::uninit();
        if unsafe {
            rustler_sys::enif_inspect_binary(
                term.get_env().as_c_arg(),
                term.as_c_arg(),
                binary.as_mut_ptr(),
            )
        } == 0
        {
            return Err(Error::BadArg);
        }
        Ok(Binary {
            inner: unsafe { binary.assume_init() },
            term,
        })
    }

    pub fn from_iolist(term: Term<'a>) -> Result<Self, Error> {
        let mut binary = MaybeUninit::uninit();
        if unsafe {
            rustler_sys::enif_inspect_iolist_as_binary(
                term.get_env().as_c_arg(),
                term.as_c_arg(),
                binary.as_mut_ptr(),
            )
        } == 0
        {
            return Err(Error::BadArg);
        }
        Ok(Binary {
            inner: unsafe { binary.assume_init() },
            term,
        })
    }

    pub fn to_term<'b>(&self, env: Env<'b>) -> Term<'b> {
        self.term.in_env(env)
    }

    pub fn as_slice(&self) -> &'a [u8] {
        unsafe { ::std::slice::from_raw_parts(self.inner.data, self.inner.size) }
    }

    /// Returns a new view into the same binary.
    /// This will not copy anything.
    pub fn make_subbinary(&self, offset: usize, length: usize) -> NifResult<Binary<'a>> {
        let min_len = length.checked_add(offset);
        if min_len.ok_or(Error::BadArg)? > self.inner.size {
            return Err(Error::BadArg);
        }

        let raw_term = unsafe {
            rustler_sys::enif_make_sub_binary(
                self.term.get_env().as_c_arg(),
                self.term.as_c_arg(),
                offset,
                length,
            )
        };
        let term = unsafe { Term::new(self.term.get_env(), raw_term) };
        // This should never fail, as we are always passing in a binary term.
        Ok(Binary::from_term(term).ok().unwrap())
    }
}

impl<'a> Borrow<[u8]> for Binary<'a> {
    fn borrow(&self) -> &[u8] {
        self.as_slice()
    }
}
impl<'a> Deref for Binary<'a> {
    type Target = [u8];
    fn deref(&self) -> &[u8] {
        self.as_slice()
    }
}

impl<'a> Decoder<'a> for Binary<'a> {
    fn decode(term: Term<'a>) -> Result<Self, Error> {
        Binary::from_term(term)
    }
}
impl<'a> Encoder for Binary<'a> {
    fn encode<'b>(&self, env: Env<'b>) -> Term<'b> {
        self.to_term(env)
    }
}

/// ## Binary terms
impl<'a> Term<'a> {
    pub fn into_binary(self) -> NifResult<Binary<'a>> {
        Binary::from_term(self)
    }
}
<|endoftext|>


<commit_before>use std::str;

use crate::{Error, Result, Row, Rows, Statement};

/// Information about a column of a SQLite query.
#[derive(Debug)]
pub struct Column<'stmt> {
    name: &'stmt str,
    decl_type: Option<&'stmt str>,
}

impl Column<'_> {
    /// Returns the name of the column.
    pub fn name(&self) -> &str {
        self.name
    }

    /// Returns the type of the column (`None` for expression).
    pub fn decl_type(&self) -> Option<&str> {
        self.decl_type
    }
}

impl Statement<'_> {
    /// Get all the column names in the result set of the prepared statement.
    pub fn column_names(&self) -> Vec<&str> {
        let n = self.column_count();
        let mut cols = Vec::with_capacity(n as usize);
        for i in 0..n {
            let s = self.column_name(i).unwrap(); // safe to unwrap as we're using only verified indices
            cols.push(s);
        }
        cols
    }

    /// Return the number of columns in the result set returned by the prepared
    /// statement.
    pub fn column_count(&self) -> usize {
        self.stmt.column_count()
    }

    /// Returns the name of the column by index or None on out of bounds access. Panics when column name is not valid UTF-8
    pub fn column_name(&self, col: usize) -> Option<&str> {
        self.stmt.column_name(col)
           .map(|x| str::from_utf8(x.to_bytes()).expect("Invalid UTF-8 sequence in column name"))
    }

    /// Returns the column index in the result set for a given column name.
    ///
    /// If there is no AS clause then the name of the column is unspecified and
    /// may change from one release of SQLite to the next.
    ///
    /// # Failure
    ///
    /// Will return an `Error::InvalidColumnName` when there is no column with
    /// the specified `name`.
    pub fn column_index(&self, name: &str) -> Result<usize> {
        let bytes = name.as_bytes();
        let n = self.column_count();
        for i in 0..n {
            // Note: `column_name` is only fallible if `i` is out of bounds,
            // which we've already checked.
            if bytes.eq_ignore_ascii_case(self.stmt.column_name(i).unwrap().to_bytes()) {
                return Ok(i);
            }
        }
        Err(Error::InvalidColumnName(String::from(name)))
    }

    /// Returns a slice describing the columns of the result of the query.
    pub fn columns<'stmt>(&'stmt self) -> Vec<Column<'stmt>> {
        let n = self.column_count();
        let mut cols = Vec::with_capacity(n as usize);
        for i in 0..n {
            let name = self.column_name(i).unwrap(); // safe to unwrap as we're using only verified indices
            let slice = self.stmt.column_decltype(i);
            let decl_type = slice.map(|s| str::from_utf8(s.to_bytes()).expect("Invalid UTF-8 sequence in column declaration"));
            cols.push(Column { name, decl_type });
        }
        cols
    }
}

impl<'stmt> Rows<'stmt> {
    /// Get all the column names.
    pub fn column_names(&self) -> Option<Vec<&str>> {
        self.stmt.map(Statement::column_names)
    }

    /// Return the number of columns.
    pub fn column_count(&self) -> Option<usize> {
        self.stmt.map(Statement::column_count)
    }

    /// Return the name of the column.
    pub fn column_name(&self, col: usize) -> Option<&str> {
        self.stmt.and_then(|x| x.column_name(col))
    }

    /// Return the index of the column.
    pub fn column_index(&self, name: &str) -> Option<Result<usize>> {
        self.stmt.map(|x| x.column_index(name))
    }

    /// Returns a slice describing the columns of the Rows.
    pub fn columns(&self) -> Option<Vec<Column<'stmt>>> {
        self.stmt.map(Statement::columns)
    }
}

impl<'stmt> Row<'stmt> {
    /// Get all the column names of the Row.
    pub fn column_names(&self) -> Vec<&str> {
        self.stmt.column_names()
    }

    /// Return the number of columns in the current row.
    pub fn column_count(&self) -> usize {
        self.stmt.column_count()
    }

    /// Return the name of the column.
    pub fn column_name(&self, col: usize) -> Option<&str> {
        self.stmt.column_name(col)
    }

    /// Return the index of the column.
    pub fn column_index(&self, name: &str) -> Result<usize> {
        self.stmt.column_index(name)
    }
    /// Returns a slice describing the columns of the Row.
    pub fn columns(&self) -> Vec<Column<'stmt>> {
        self.stmt.columns()
    }
}

#[cfg(test)]
mod test {
    use super::Column;
    use crate::Connection;

    #[test]
    fn test_columns() {
        let db = Connection::open_in_memory().unwrap();
        let query = db.prepare("SELECT * FROM sqlite_master").unwrap();
        let columns = query.columns();
        let column_names: Vec<&str> = columns.iter().map(Column::name).collect();
        assert_eq!(
            column_names.as_slice(),
            &["type", "name", "tbl_name", "rootpage", "sql"]
        );
        let column_types: Vec<Option<&str>> = columns.iter().map(Column::decl_type).collect();
        assert_eq!(
            &column_types[..3],
            &[Some("text"), Some("text"), Some("text"),]
        );
    }

    #[test]
    fn test_column_name_in_error() {
        use crate::{types::Type, Error};
        let db = Connection::open_in_memory().unwrap();
        db.execute_batch(
            "BEGIN;
             CREATE TABLE foo(x INTEGER, y TEXT);
             INSERT INTO foo VALUES(4, NULL);
             END;",
        )
        .unwrap();
        let mut stmt = db.prepare("SELECT x as renamed, y FROM foo").unwrap();
        let mut rows = stmt.query(crate::NO_PARAMS).unwrap();
        let row = rows.next().unwrap().unwrap();
        match row.get::<_, String>(0).unwrap_err() {
            Error::InvalidColumnType(idx, name, ty) => {
                assert_eq!(idx, 0);
                assert_eq!(name, "renamed");
                assert_eq!(ty, Type::Integer);
            }
            e => {
                panic!("Unexpected error type: {:?}", e);
            }
        }
        match row.get::<_, String>("y").unwrap_err() {
            Error::InvalidColumnType(idx, name, ty) => {
                assert_eq!(idx, 1);
                assert_eq!(name, "y");
                assert_eq!(ty, Type::Null);
            }
            e => {
                panic!("Unexpected error type: {:?}", e);
            }
        }
    }
}
<commit_msg>Use the correct (larger) lifetime of the returned column-related references in Row, Rows and Column<commit_after>use std::str;

use crate::{Error, Result, Row, Rows, Statement};

/// Information about a column of a SQLite query.
#[derive(Debug)]
pub struct Column<'stmt> {
    name: &'stmt str,
    decl_type: Option<&'stmt str>,
}

impl<'stmt> Column<'stmt> {
    /// Returns the name of the column.
    pub fn name(&self) -> &'stmt str {
        self.name
    }

    /// Returns the type of the column (`None` for expression).
    pub fn decl_type(&self) -> Option<&'stmt str> {
        self.decl_type
    }
}

impl Statement<'_> {
    /// Get all the column names in the result set of the prepared statement.
    pub fn column_names(&self) -> Vec<&str> {
        let n = self.column_count();
        let mut cols = Vec::with_capacity(n as usize);
        for i in 0..n {
            let s = self.column_name(i).unwrap(); // safe to unwrap as we're using only verified indices
            cols.push(s);
        }
        cols
    }

    /// Return the number of columns in the result set returned by the prepared
    /// statement.
    pub fn column_count(&self) -> usize {
        self.stmt.column_count()
    }

    /// Returns the name of the column by index or None on out of bounds access. Panics when column name is not valid UTF-8
    pub fn column_name(&self, col: usize) -> Option<&str> {
        self.stmt.column_name(col)
           .map(|x| str::from_utf8(x.to_bytes()).expect("Invalid UTF-8 sequence in column name"))
    }

    /// Returns the column index in the result set for a given column name.
    ///
    /// If there is no AS clause then the name of the column is unspecified and
    /// may change from one release of SQLite to the next.
    ///
    /// # Failure
    ///
    /// Will return an `Error::InvalidColumnName` when there is no column with
    /// the specified `name`.
    pub fn column_index(&self, name: &str) -> Result<usize> {
        let bytes = name.as_bytes();
        let n = self.column_count();
        for i in 0..n {
            // Note: `column_name` is only fallible if `i` is out of bounds,
            // which we've already checked.
            if bytes.eq_ignore_ascii_case(self.stmt.column_name(i).unwrap().to_bytes()) {
                return Ok(i);
            }
        }
        Err(Error::InvalidColumnName(String::from(name)))
    }

    /// Returns a slice describing the columns of the result of the query.
    pub fn columns(&self) -> Vec<Column> {
        let n = self.column_count();
        let mut cols = Vec::with_capacity(n as usize);
        for i in 0..n {
            let name = self.column_name(i).unwrap(); // safe to unwrap as we're using only verified indices
            let slice = self.stmt.column_decltype(i);
            let decl_type = slice.map(|s| str::from_utf8(s.to_bytes()).expect("Invalid UTF-8 sequence in column declaration"));
            cols.push(Column { name, decl_type });
        }
        cols
    }
}

impl<'stmt> Rows<'stmt> {
    /// Get all the column names.
    pub fn column_names(&self) -> Option<Vec<&'stmt str>> {
        self.stmt.map(Statement::column_names)
    }

    /// Return the number of columns.
    pub fn column_count(&self) -> Option<usize> {
        self.stmt.map(Statement::column_count)
    }

    /// Return the name of the column.
    pub fn column_name(&self, col: usize) -> Option<&'stmt str> {
        self.stmt.and_then(|x| x.column_name(col))
    }

    /// Return the index of the column.
    pub fn column_index(&self, name: &str) -> Option<Result<usize>> {
        self.stmt.map(|x| x.column_index(name))
    }

    /// Returns a slice describing the columns of the Rows.
    pub fn columns(&self) -> Option<Vec<Column<'stmt>>> {
        self.stmt.map(Statement::columns)
    }
}

impl<'stmt> Row<'stmt> {
    /// Get all the column names of the Row.
    pub fn column_names(&self) -> Vec<&'stmt str> {
        self.stmt.column_names()
    }

    /// Return the number of columns in the current row.
    pub fn column_count(&self) -> usize {
        self.stmt.column_count()
    }

    /// Return the name of the column.
    pub fn column_name(&self, col: usize) -> Option<&'stmt str> {
        self.stmt.column_name(col)
    }

    /// Return the index of the column.
    pub fn column_index(&self, name: &str) -> Result<usize> {
        self.stmt.column_index(name)
    }
    /// Returns a slice describing the columns of the Row.
    pub fn columns(&self) -> Vec<Column<'stmt>> {
        self.stmt.columns()
    }
}

#[cfg(test)]
mod test {
    use super::Column;
    use crate::Connection;

    #[test]
    fn test_columns() {
        let db = Connection::open_in_memory().unwrap();
        let query = db.prepare("SELECT * FROM sqlite_master").unwrap();
        let columns = query.columns();
        let column_names: Vec<&str> = columns.iter().map(Column::name).collect();
        assert_eq!(
            column_names.as_slice(),
            &["type", "name", "tbl_name", "rootpage", "sql"]
        );
        let column_types: Vec<Option<&str>> = columns.iter().map(Column::decl_type).collect();
        assert_eq!(
            &column_types[..3],
            &[Some("text"), Some("text"), Some("text"),]
        );
    }

    #[test]
    fn test_column_name_in_error() {
        use crate::{types::Type, Error};
        let db = Connection::open_in_memory().unwrap();
        db.execute_batch(
            "BEGIN;
             CREATE TABLE foo(x INTEGER, y TEXT);
             INSERT INTO foo VALUES(4, NULL);
             END;",
        )
        .unwrap();
        let mut stmt = db.prepare("SELECT x as renamed, y FROM foo").unwrap();
        let mut rows = stmt.query(crate::NO_PARAMS).unwrap();
        let row = rows.next().unwrap().unwrap();
        match row.get::<_, String>(0).unwrap_err() {
            Error::InvalidColumnType(idx, name, ty) => {
                assert_eq!(idx, 0);
                assert_eq!(name, "renamed");
                assert_eq!(ty, Type::Integer);
            }
            e => {
                panic!("Unexpected error type: {:?}", e);
            }
        }
        match row.get::<_, String>("y").unwrap_err() {
            Error::InvalidColumnType(idx, name, ty) => {
                assert_eq!(idx, 1);
                assert_eq!(name, "y");
                assert_eq!(ty, Type::Null);
            }
            e => {
                panic!("Unexpected error type: {:?}", e);
            }
        }
    }
}
<|endoftext|>


<commit_before>//! A simple first program in rust, actually I plan to do something even
//! useful. Trying to find all my audiobooks on many machines, identify them,
//! find duplicates (later trying to solve the problem, also including their permissions, different names, but same albums, etc),
//! get all stats about it).
//#![feature(alloc_system)]
//extern crate alloc_system;   // strip down size of binary executable
extern crate adfblib;

extern crate clap;
extern crate hostname;
extern crate rayon;
extern crate uuid;

use std::path::Path; // path, clear
use std::sync::{Arc, Mutex};
use std::sync::mpsc;
use std::thread;

use rayon::prelude::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
use uuid::Uuid;

use adfblib::data;
use adfblib::data::Collection;
use adfblib::ctrl::{Alive, Ctrl, ReceiveDialog, Status, SystemMsg};
use adfblib::net::Net;

static INPUT_FOLDERS: &str = "folders";
static APP_TITLE: &str = "The audiobook finder";
static ARG_NET: &str = "net";
static ARG_TUI: &str = "tui";

fn main() {
    let parse_args = clap::App::new(APP_TITLE)
        .version("0.1")
        .author("S. K. <skroemeke@gmail.com>")
        .about("A audiobook finder")
        .arg(
            clap::Arg::with_name("config")
                .short("c")
                .long("config")
                .value_name("FILE")
                .help("Sets a custom config file")
                .takes_value(true),
        )
        .arg(
            clap::Arg::with_name(ARG_TUI)
                .short("t")
                .long("TUI")
                .help("Starts the TUI")
                .takes_value(false),
        )
        .arg(
            clap::Arg::with_name(ARG_NET)
                .short("n")
                .long("net")
                .help("Starts without net")
                .takes_value(false),
        )
        .arg(
            clap::Arg::with_name(INPUT_FOLDERS)
                .help("Sets the input folder(s) to use")
                .multiple(true)
                .required(false),
        )
        .get_matches();

    // tricky thing, but I really like that
    let all_pathes = if let Some(correct_input) = parse_args.values_of(INPUT_FOLDERS) {
        correct_input.collect()
    } else {
        vec!["."]
    };

    let hostname = hostname::get_hostname().unwrap_or("undefined".to_string());
    let max_threads = rayon::current_num_threads();

    // check if tui and net search is needed
    let has_arg = |x: &str| parse_args.is_present(x);

    // has to be mutable because in case of error this might be changed
    let mut has_tui = has_arg(ARG_TUI);
    let has_net = has_arg(ARG_NET);

    // prepare the message system
    let (tx, rx) = mpsc::channel::<SystemMsg>();

    let tx_sys_mut = Mutex::new(tx.clone());
    let tx_net_mut = Mutex::new(tx.clone());
    let tx_net_alive_mut = Mutex::new(tx.clone());

    // copy to vec<&str>
    let tui_pathes = all_pathes.iter().map(|s| s.to_string()).collect();

    // get an unique id for this client
    let client_id = Uuid::new_v4();

    // the signal to tell further processes, that tui creation worked
    // and the messages actually go somewhere, otherwise it will assume: no tui
    let (send_tui_worked, receiver_tui_worked) = mpsc::channel::<bool>();

    // start the tui thread
    let tui_runner = thread::spawn(move || {
        if has_tui {
            // start animation .... timer and so on
            if let Ok(starter) = tx_net_alive_mut.lock() {
                starter
                    .send(SystemMsg::StartAnimation(Alive::HostSearch, Status::ON))
                    .unwrap();
            }
            let controller = Ctrl::new(client_id.to_string(), &tui_pathes, rx, tx.clone(), has_net);
            match controller {
                Ok(mut controller) => {
                    // signal ok
                    send_tui_worked.send(true).unwrap();
                    drop(send_tui_worked);

                    // do finally the necessary
                    controller.run();
                }
                Err(error_text) => {
                    println!("{:?}", error_text);
                    // no tui could be created
                    send_tui_worked.send(false).unwrap();
                    drop(send_tui_worked);
                }
            }
        }
    });

    // blocks that one signal (but that should be very short time)
    if let Ok(tui_check_receiver) = receiver_tui_worked.recv() {
        has_tui = tui_check_receiver;
    } else {
        println!("Something bad has happenend!!!");
        has_tui = false;
    }
    // make it immutable from now on
    let has_tui = has_tui;

    // start the net runner thread
    let tx_net_mut_arc = Arc::new(tx_net_mut);
    let net_runner = thread::spawn(move || {
        if has_net {
            // need to simplify and clarify this here ......
            // but this lock unwrap is safe
            let mut netfinder = Net::new(
                &client_id.to_string(),
                has_tui,
                tx_net_mut_arc.lock().unwrap().clone(),
            );
            netfinder.lookup();
        }
    });

    // initialize the data collection for all
    let init_collection = Collection::new(hostname, &client_id, max_threads);
    let collection_protected = Arc::new(Mutex::new(init_collection));

    // start the search threads, each path its own
    all_pathes.par_iter().enumerate().for_each(|(index, elem)| {
        if !has_tui {
            println!("[{:?}] looking into path {:?}", index, elem);
        } else {
            // start animation .... timer and so on
            if let Ok(starter) = tx_sys_mut.lock() {
                starter
                    .send(SystemMsg::StartAnimation(
                        Alive::BusyPath(index),
                        Status::ON,
                    ))
                    .unwrap();
            }
        }
        let live_here = collection_protected.clone();
        let locked_collection = live_here.lock();
        if let Ok(mut pure_collection) = locked_collection {
            match pure_collection.visit_dirs(Path::new(elem), &data::Collection::visit_files) {
            Ok(local_stats) => {
                if has_tui {
                    // stop animation
                    if let Ok(stopper) = tx_sys_mut.lock() {
                        stopper
                            .send(SystemMsg::StartAnimation(
                                Alive::BusyPath(index),
                                Status::OFF,
                            ))
                            .unwrap();
                    }
                } else {
                    let text = format!(
                        "\n\
                         analyzed: {an:>width$}, faulty: {fa:>width$}\n\
                         searched: {se:>width$}, other: {ot:>width$}",
                        an = local_stats.analyzed,
                        fa = local_stats.faulty,
                        se = local_stats.searched,
                        ot = local_stats.other,
                        width = 3
                    );
                    println!("[{:?}] done {}", index, text);
                }
            }
            Err(_e) => {
                let text = format!("An error has occurred in search path [{}]!!", index);
                if has_tui {
                    let debug_message_id = ReceiveDialog::Debug;
                    let text = text.to_string();
                    let debug_text = tx_sys_mut
                        .lock();
                    if let Ok(debug_text) = debug_text {
                        debug_text
                        .send(SystemMsg::Update(debug_message_id, text))
                        .unwrap();
                    }
                } else {
                    println!("{:?}", text);
                }
            }
        }
        }
    });

    if !has_tui {
        if let Ok(result_collection) = collection_protected.lock() {
            result_collection.print_stats();
        }
        let _ = net_runner.join();
    } else {
        let _ = tui_runner.join();
        // if tui, net runner shall stop when tui stops
        drop(net_runner);
    }

    println!("Finished!");
}
<commit_msg>Fixing bug when tui is not used at all.<commit_after>//! A simple first program in rust, actually I plan to do something even
//! useful. Trying to find all my audiobooks on many machines, identify them,
//! find duplicates (later trying to solve the problem, also including their permissions, different names, but same albums, etc),
//! get all stats about it).
//#![feature(alloc_system)]
//extern crate alloc_system;   // strip down size of binary executable
extern crate adfblib;

extern crate clap;
extern crate hostname;
extern crate rayon;
extern crate uuid;

use std::path::Path; // path, clear
use std::sync::{Arc, Mutex};
use std::sync::mpsc;
use std::thread;

use rayon::prelude::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
use uuid::Uuid;

use adfblib::data;
use adfblib::data::Collection;
use adfblib::ctrl::{Alive, Ctrl, ReceiveDialog, Status, SystemMsg};
use adfblib::net::Net;

static INPUT_FOLDERS: &str = "folders";
static APP_TITLE: &str = "The audiobook finder";
static ARG_NET: &str = "net";
static ARG_TUI: &str = "tui";

fn main() {
    let parse_args = clap::App::new(APP_TITLE)
        .version("0.1")
        .author("S. K. <skroemeke@gmail.com>")
        .about("A audiobook finder")
        .arg(
            clap::Arg::with_name("config")
                .short("c")
                .long("config")
                .value_name("FILE")
                .help("Sets a custom config file")
                .takes_value(true),
        )
        .arg(
            clap::Arg::with_name(ARG_TUI)
                .short("t")
                .long("TUI")
                .help("Starts the TUI")
                .takes_value(false),
        )
        .arg(
            clap::Arg::with_name(ARG_NET)
                .short("n")
                .long("net")
                .help("Starts without net")
                .takes_value(false),
        )
        .arg(
            clap::Arg::with_name(INPUT_FOLDERS)
                .help("Sets the input folder(s) to use")
                .multiple(true)
                .required(false),
        )
        .get_matches();

    // tricky thing, but I really like that
    let all_pathes = if let Some(correct_input) = parse_args.values_of(INPUT_FOLDERS) {
        correct_input.collect()
    } else {
        vec!["."]
    };

    let hostname = hostname::get_hostname().unwrap_or("undefined".to_string());
    let max_threads = rayon::current_num_threads();

    // check if tui and net search is needed
    let has_arg = |x: &str| parse_args.is_present(x);

    // has to be mutable because in case of error this might be changed
    let mut has_tui = has_arg(ARG_TUI);
    let has_net = has_arg(ARG_NET);

    // prepare the message system
    let (tx, rx) = mpsc::channel::<SystemMsg>();

    let tx_sys_mut = Mutex::new(tx.clone());
    let tx_net_mut = Mutex::new(tx.clone());
    let tx_net_alive_mut = Mutex::new(tx.clone());

    // copy to vec<&str>
    let tui_pathes = all_pathes.iter().map(|s| s.to_string()).collect();

    // get an unique id for this client
    let client_id = Uuid::new_v4();

    // the signal to tell further processes, that tui creation worked
    // and the messages actually go somewhere, otherwise it will assume: no tui
    let (send_tui_worked, receiver_tui_worked) = mpsc::channel::<bool>();

    // start the tui thread
    let tui_runner = thread::spawn(move || {
        if has_tui {
            // start animation .... timer and so on
            if let Ok(starter) = tx_net_alive_mut.lock() {
                starter
                    .send(SystemMsg::StartAnimation(Alive::HostSearch, Status::ON))
                    .unwrap();
            }
            let controller = Ctrl::new(client_id.to_string(), &tui_pathes, rx, tx.clone(), has_net);
            match controller {
                Ok(mut controller) => {
                    // signal ok
                    send_tui_worked.send(true).unwrap();
                    drop(send_tui_worked);

                    // do finally the necessary
                    controller.run();
                }
                Err(error_text) => {
                    println!("{:?}", error_text);
                    // no tui could be created
                    send_tui_worked.send(false).unwrap();
                    drop(send_tui_worked);
                }
            }
        }
    });

    // blocks that one signal (but that should be very short time)
    if has_tui {
        if let Ok(tui_check_receiver) = receiver_tui_worked.recv() {
            has_tui = tui_check_receiver;
        } else {
            println!("Something bad has happenend!!!");
            has_tui = false;
        }
    }
    // make it immutable from now on
    let has_tui = has_tui;

    // start the net runner thread
    let tx_net_mut_arc = Arc::new(tx_net_mut);
    let net_runner = thread::spawn(move || {
        if has_net {
            // need to simplify and clarify this here ......
            // but this lock unwrap is safe
            let mut netfinder = Net::new(
                &client_id.to_string(),
                has_tui,
                tx_net_mut_arc.lock().unwrap().clone(),
            );
            netfinder.lookup();
        }
    });

    // initialize the data collection for all
    let init_collection = Collection::new(hostname, &client_id, max_threads);
    let collection_protected = Arc::new(Mutex::new(init_collection));

    // start the search threads, each path its own
    all_pathes.par_iter().enumerate().for_each(|(index, elem)| {
        if !has_tui {
            println!("[{:?}] looking into path {:?}", index, elem);
        } else {
            // start animation .... timer and so on
            if let Ok(starter) = tx_sys_mut.lock() {
                starter
                    .send(SystemMsg::StartAnimation(
                        Alive::BusyPath(index),
                        Status::ON,
                    ))
                    .unwrap();
            }
        }
        let live_here = collection_protected.clone();
        let locked_collection = live_here.lock();
        if let Ok(mut pure_collection) = locked_collection {
            match pure_collection.visit_dirs(Path::new(elem), &data::Collection::visit_files) {
            Ok(local_stats) => {
                if has_tui {
                    // stop animation
                    if let Ok(stopper) = tx_sys_mut.lock() {
                        stopper
                            .send(SystemMsg::StartAnimation(
                                Alive::BusyPath(index),
                                Status::OFF,
                            ))
                            .unwrap();
                    }
                } else {
                    let text = format!(
                        "\n\
                         analyzed: {an:>width$}, faulty: {fa:>width$}\n\
                         searched: {se:>width$}, other: {ot:>width$}",
                        an = local_stats.analyzed,
                        fa = local_stats.faulty,
                        se = local_stats.searched,
                        ot = local_stats.other,
                        width = 3
                    );
                    println!("[{:?}] done {}", index, text);
                }
            }
            Err(_e) => {
                let text = format!("An error has occurred in search path [{}]!!", index);
                if has_tui {
                    let debug_message_id = ReceiveDialog::Debug;
                    let text = text.to_string();
                    let debug_text = tx_sys_mut
                        .lock();
                    if let Ok(debug_text) = debug_text {
                        debug_text
                        .send(SystemMsg::Update(debug_message_id, text))
                        .unwrap();
                    }
                } else {
                    println!("{:?}", text);
                }
            }
        }
        }
    });

    if !has_tui {
        if let Ok(result_collection) = collection_protected.lock() {
            result_collection.print_stats();
        }
        let _ = net_runner.join();
    } else {
        let _ = tui_runner.join();
        // if tui, net runner shall stop when tui stops
        drop(net_runner);
    }

    println!("Finished!");
}
<|endoftext|>


<commit_before>extern crate osmio;
extern crate image;
extern crate gif;
extern crate clap;

use clap::{Arg, App, AppSettings};

use std::fs;
use osmio::OSMReader;
use osmio::pbf::PBFReader;
use std::io::{Read, Write, BufReader, BufRead, BufWriter};
use std::collections::HashMap;

use gif::SetParameter;

// (frame_no, (pixel_index, num_changes))
type Frames = Vec<(u32, Vec<(u32, u16)>)>;

struct ColourRamp {
    empty_colour: (u8, u8, u8),
    steps: Vec<(u32, (u8, u8, u8))>,
}

impl ColourRamp {
    fn new_from_filename(filename: &str) -> Self {
        let mut contents = String::new();
        let mut file = fs::File::open(filename).unwrap();
        file.read_to_string(&mut contents).expect("Couldn't read colour ramp file");
        Self::new_from_text(&contents)

    }
    fn new_from_text(source: &str) -> Self {
        let lines: Vec<_> = source.lines().collect();
        let empty_vec = lines[0].split(",").filter_map(|x| x.parse::<u8>().ok()).take(3).collect::<Vec<_>>();
        let empty = (empty_vec[0], empty_vec[1], empty_vec[2]);

        let mut steps = Vec::new();
        for line in lines.iter().skip(1) {
            let line = line.split(",").filter_map(|x| x.parse::<u32>().ok()).take(4).collect::<Vec<_>>();
            let age = line[0];
            let colour = (line[1] as u8, line[2] as u8, line[3] as u8);
            steps.push((age, colour));
        }

        if steps.len() > 254 {
            panic!("Too many steps");
        }

        ColourRamp{ empty_colour: empty, steps: steps }
    }

    fn palette(&self) -> Vec<u8> {
        let mut results = Vec::with_capacity((self.steps.len()+1)*3);
        results.push(self.empty_colour.0);
        results.push(self.empty_colour.1);
        results.push(self.empty_colour.2);

        for &(_, (r, g, b)) in self.steps.iter() {
            results.push(r);
            results.push(g);
            results.push(b);
        }

        results
    }

    fn index_for_magnitude(&self, magnitude: Option<u32>) -> u8 {
        match magnitude {
            None => 0,
            Some(magnitude) => {
                if magnitude > 255 {
                    1
                } else {
                    (255 - magnitude) as u8
                }
            },
            //Some(age) => match self.steps.iter().position(|&(x, _)| x == age) {
            //    None => 0,
            //    Some(i) => (i+1) as u8,
            //}
        }
    }
}

fn read_pbf(filename: &str, height: u32, sec_per_frame: u32, bbox: &[f32; 4]) -> Frames {
    let file = BufReader::new(fs::File::open(&filename).unwrap());
    let mut node_reader = PBFReader::new(file);
    let node_reader = node_reader.nodes();

    let mut results: HashMap<u32, HashMap<u32, u16>> = HashMap::new();

    let left = bbox[0]; let bottom = bbox[1]; let right = bbox[2]; let top = bbox[3];
    let bbox_width = right - left;
    let bbox_height = top - bottom;
    let width = ((bbox_width / bbox_height) * (height as f32)) as u32;

    // 1st April 2005, midnight GMT. We presume no OSM editing before then.
    let osm_epoch = 1109635200;

    // FIXME use a btreemap?

    let mut first_frame_no = std::u32::MAX;
    let mut last_frame_no = 0;

    let mut num_nodes: u64 = 0;
    for node in node_reader {
        if let (Some(lat), Some(lon)) = (node.lat, node.lon) {
            // FIXME should be able to do non-equals but it fails for point at south pole
            if lat >= top || lat <= bottom || lon >= right || lon <= left {
                continue;
            }

            let timestamp = node.timestamp.to_epoch_number() as u64;
            if timestamp < osm_epoch {
                panic!("timestamp before epoch. Change code. {}", timestamp);
            }
            let offset = (timestamp - osm_epoch) as u32;
            // TODO double check that this rounds down
            let frame_no = offset / sec_per_frame;

            if frame_no < first_frame_no {
                first_frame_no = frame_no;
            }
            if frame_no > last_frame_no {
                last_frame_no = frame_no;
            }

            let pixel_idx = latlon_to_pixel_index(lat, lon, width, height, &bbox);
            let curr_val = results.entry(frame_no).or_insert(HashMap::new()).entry(pixel_idx).or_insert(0);
            if *curr_val < std::u16::MAX {
                *curr_val += 1;
            }

            num_nodes += 1;
            if num_nodes % 50_000_000 == 0 {
                println!("Done {} million points", num_nodes/1_000_000);
            }
        }
    }
    println!("Done {} points", num_nodes);
    let num_frames = last_frame_no - first_frame_no + 1;
    println!("There are {} frames, which is {} sec", num_frames, num_frames as f32/30.);

    let mut sorted_results = Vec::with_capacity((last_frame_no-first_frame_no+1) as usize);

    for frame_no in 0..num_frames {
        match results.remove(&(frame_no+first_frame_no)) {
            None => { sorted_results.push((frame_no, Vec::with_capacity(0))) },
            Some(pixels) => {
                let pixels = pixels.into_iter().collect();
                sorted_results.push((frame_no, pixels));
            }
        }
    }

    sorted_results
}

fn write_frames(frames: Frames, filename: &str, height: u32, sec_per_frame: u32, bbox: &[f32; 4]) {
    let left = bbox[0]; let bottom = bbox[1]; let right = bbox[2]; let top = bbox[3];
    let bbox_width = right - left;
    let bbox_height = top - bottom;
    let width = ((bbox_width / bbox_height) * (height as f32)) as u32;

    let mut file = BufWriter::new(fs::File::create(&filename).unwrap());

    writeln!(file, "metadata version {}", env!("CARGO_PKG_VERSION")).expect("Couldn't write metadata");
    writeln!(file, "metadata height {}", height).expect("Couldn't write metadata");
    writeln!(file, "metadata width {}", width).expect("Couldn't write metadata");
    writeln!(file, "metadata sec_per_frame {}", sec_per_frame).expect("Couldn't write metadata");
    writeln!(file, "metadata left {}", left).expect("Couldn't write metadata");
    writeln!(file, "metadata bottom {}", bottom).expect("Couldn't write metadata");
    writeln!(file, "metadata right {}", right).expect("Couldn't write metadata");
    writeln!(file, "metadata top {}", top).expect("Couldn't write metadata");
    writeln!(file, "").expect("Couldn't write metadata");

    for (frame_no, pixels) in frames.into_iter() {
        write!(file, "{}", frame_no).unwrap();
        for p in pixels {
            write!(file, ",{},{}", p.0, p.1).unwrap();
        }
        write!(file, "\n").unwrap()
    }
}

fn read_frames(filename: &str) -> Frames {
    let file = BufReader::new(fs::File::open(&filename).unwrap());
    let mut results = Frames::new();

    //let metadata: Vec<(String, String)> = file.lines().filter_map(|x| x.ok()).take_while(|x| x.len() > 0).map(|x| { let words: Vec<_> = x.split(" ").skip(1).take(2).collect(); (words[0].clone(), words[1].clone())}).collect();

    for line in file.lines().filter_map(|x| x.ok()).skip_while(|x| x.starts_with("metadata ") || x.len() == 0) {
        let frame_no = line.split(",").nth(0).unwrap().parse().unwrap();
        let pixels: Vec<(u32, u16)> = line.split(",").skip(1).collect::<Vec<_>>().chunks(2).map(|pair| (pair[0].parse().unwrap(), pair[1].parse().unwrap())).collect();
        results.push((frame_no, pixels))
    }

    results
}

fn latlon_to_pixel_index(lat: f32, lon: f32, width: u32, height: u32, bbox: &[f32; 4]) -> u32 {
    let left = bbox[0]; let bottom = bbox[1]; let right = bbox[2]; let top = bbox[3];
    let bbox_width = right - left;
    let bbox_height = top - bottom;

    let lat0 = top - lat;
    let lon0 = lon - left;

    let x = ((lon0/bbox_width)*(width as f32)) as u32;
    let y = ((lat0/bbox_height)*(height as f32)) as u32;

    let i = y * width + x;

    assert!(i < width*height, "{} L{}, lat = {} lon = {} width = {} height = {} bbox = {:?} x = {} y = {} i = {}", file!(), line!(), lat, lon, width, height, bbox, x, y, i);

    i
}

fn create_gif(frames: Frames, output_image_filename: &str, height: u32, bbox: &[f32; 4], colour_ramp: &ColourRamp) {
    let mut output_file = fs::File::create(output_image_filename).expect("Can't create image");

    let left = bbox[0]; let bottom = bbox[1]; let right = bbox[2]; let top = bbox[3];
    let bbox_width = right - left;
    let bbox_height = top - bottom;
    let width = ((bbox_width / bbox_height) * (height as f32)) as u32;

    // FIXME change width/height to u16?
    let mut encoder = gif::Encoder::new(&mut output_file, width as u16, height as u16, &colour_ramp.palette()).expect("Couldn't create encoder");
    encoder.set(gif::Repeat::Infinite).expect("Couldn't get inf repeat");

    // TODO have a global palette

    let mut image = vec![None; (width*height) as usize];

    for (frame_no, pixels) in frames.into_iter() {

        for i in 0..image.len() {
            if image[i].is_some() && image[i].unwrap() > 0 {
                image[i] = image[i].map(|x| (x as f32*0.9) as u32);
            }
        }

        for (i, magnitude) in pixels {
            // FIXME sometimes the value is invalid
            //assert!(i < width*height, "{} L{}, width = {} height = {} i = {}", file!(), line!(), width, height, i);
            if i < width*height {
                let new_value: u32 = image[i as usize].unwrap_or(0) + magnitude as u32;
                image[i as usize] = Some(new_value);
            }
        }

        let mut pixels = Vec::with_capacity(image.len() * 4);
        for p in image.iter().cloned() {
            let index = colour_ramp.index_for_magnitude(p);
            pixels.push(index);
        }

        let mut frame = gif::Frame::from_indexed_pixels(width as u16, height as u16, pixels.as_mut_slice(), None);
        // 30 fps, and delay is in units of 10ms.
        frame.delay = 100 / 30;

        encoder.write_frame(&frame).expect("Couldn't write frame");

        if frame_no % 30 == 0 {
            println!("Wrote frame {}", frame_no);
        }

    }

}

fn main() {
    let matches = App::new("osm-history-animation")
                          .arg(Arg::with_name("input").long("input").short("i")
                               .takes_value(true))
                          .arg(Arg::with_name("output") .long("output").short("o")
                               .takes_value(true))
                          .arg(Arg::with_name("height") .long("height").short("h")
                               .takes_value(true))
                          .arg(Arg::with_name("spf") .long("sec-per-frame").short("s")
                               .takes_value(true))
                          .arg(Arg::with_name("colour_ramp") .long("colour-ramp")
                               .takes_value(true))
                          .arg(Arg::with_name("save-intermediate").long("save-intermediate"))
                          .arg(Arg::with_name("load-intermediate").long("load-intermediate"))
                          .arg(Arg::with_name("bbox").long("bbox").takes_value(true).short("b"))
                          .setting(AppSettings::AllowLeadingHyphen)
                          .get_matches();

    let input_filename = matches.value_of("input").unwrap();
    let output_filename = matches.value_of("output").unwrap();
    let height: u32 = matches.value_of("height").unwrap().parse().unwrap();
    let sec_per_frame: u32 = matches.value_of("spf").unwrap().parse().unwrap();

    let bbox = match matches.value_of("bbox") {
        None => [-180., -90., 180., 90.],
        Some(text) => {
            let coords: Vec<f32> = text.split(",").map(|x| x.parse().unwrap() ).collect();
            [coords[0], coords[1], coords[2], coords[3]]
        }
    };

    let frames = if matches.is_present("load-intermediate") {
        println!("Reading frames from {}", input_filename);
        read_frames(&input_filename)
    } else {
        println!("Reading PBF file {}", input_filename);
        read_pbf(&input_filename, height, sec_per_frame, &bbox)
    };

    if matches.is_present("save-intermediate") {
        println!("Saving frame details to {}", output_filename);
        write_frames(frames, &output_filename, height, sec_per_frame, &bbox);
    } else {
        let colour_ramp = ColourRamp::new_from_filename(matches.value_of("colour_ramp").unwrap());
        println!("Creating image {}", output_filename);
        create_gif(frames, &output_filename, height, &bbox, &colour_ramp);
    }
    println!("\nFinished");
}
<commit_msg>tweak fading<commit_after>extern crate osmio;
extern crate image;
extern crate gif;
extern crate clap;

use clap::{Arg, App, AppSettings};

use std::fs;
use osmio::OSMReader;
use osmio::pbf::PBFReader;
use std::io::{Read, Write, BufReader, BufRead, BufWriter};
use std::collections::HashMap;

use gif::SetParameter;

// (frame_no, (pixel_index, num_changes))
type Frames = Vec<(u32, Vec<(u32, u16)>)>;

struct ColourRamp {
    empty_colour: (u8, u8, u8),
    steps: Vec<(u32, (u8, u8, u8))>,
}

impl ColourRamp {
    fn new_from_filename(filename: &str) -> Self {
        let mut contents = String::new();
        let mut file = fs::File::open(filename).unwrap();
        file.read_to_string(&mut contents).expect("Couldn't read colour ramp file");
        Self::new_from_text(&contents)

    }
    fn new_from_text(source: &str) -> Self {
        let lines: Vec<_> = source.lines().collect();
        let empty_vec = lines[0].split(",").filter_map(|x| x.parse::<u8>().ok()).take(3).collect::<Vec<_>>();
        let empty = (empty_vec[0], empty_vec[1], empty_vec[2]);

        let mut steps = Vec::new();
        for line in lines.iter().skip(1) {
            let line = line.split(",").filter_map(|x| x.parse::<u32>().ok()).take(4).collect::<Vec<_>>();
            let age = line[0];
            let colour = (line[1] as u8, line[2] as u8, line[3] as u8);
            steps.push((age, colour));
        }

        if steps.len() > 254 {
            panic!("Too many steps");
        }

        ColourRamp{ empty_colour: empty, steps: steps }
    }

    fn palette(&self) -> Vec<u8> {
        let mut results = Vec::with_capacity((self.steps.len()+1)*3);
        results.push(self.empty_colour.0);
        results.push(self.empty_colour.1);
        results.push(self.empty_colour.2);

        for &(_, (r, g, b)) in self.steps.iter() {
            results.push(r);
            results.push(g);
            results.push(b);
        }

        results
    }

    fn index_for_magnitude(&self, magnitude: Option<u32>) -> u8 {
        match magnitude {
            None => 0,
            Some(magnitude) => {
                if magnitude > 255 {
                    1
                } else {
                    (255 - magnitude) as u8
                }
            },
            //Some(age) => match self.steps.iter().position(|&(x, _)| x == age) {
            //    None => 0,
            //    Some(i) => (i+1) as u8,
            //}
        }
    }
}

fn read_pbf(filename: &str, height: u32, sec_per_frame: u32, bbox: &[f32; 4]) -> Frames {
    let file = BufReader::new(fs::File::open(&filename).unwrap());
    let mut node_reader = PBFReader::new(file);
    let node_reader = node_reader.nodes();

    let mut results: HashMap<u32, HashMap<u32, u16>> = HashMap::new();

    let left = bbox[0]; let bottom = bbox[1]; let right = bbox[2]; let top = bbox[3];
    let bbox_width = right - left;
    let bbox_height = top - bottom;
    let width = ((bbox_width / bbox_height) * (height as f32)) as u32;

    // 1st April 2005, midnight GMT. We presume no OSM editing before then.
    let osm_epoch = 1109635200;

    // FIXME use a btreemap?

    let mut first_frame_no = std::u32::MAX;
    let mut last_frame_no = 0;

    let mut num_nodes: u64 = 0;
    for node in node_reader {
        if let (Some(lat), Some(lon)) = (node.lat, node.lon) {
            // FIXME should be able to do non-equals but it fails for point at south pole
            if lat >= top || lat <= bottom || lon >= right || lon <= left {
                continue;
            }

            let timestamp = node.timestamp.to_epoch_number() as u64;
            if timestamp < osm_epoch {
                panic!("timestamp before epoch. Change code. {}", timestamp);
            }
            let offset = (timestamp - osm_epoch) as u32;
            // TODO double check that this rounds down
            let frame_no = offset / sec_per_frame;

            if frame_no < first_frame_no {
                first_frame_no = frame_no;
            }
            if frame_no > last_frame_no {
                last_frame_no = frame_no;
            }

            let pixel_idx = latlon_to_pixel_index(lat, lon, width, height, &bbox);
            let curr_val = results.entry(frame_no).or_insert(HashMap::new()).entry(pixel_idx).or_insert(0);
            if *curr_val < std::u16::MAX {
                *curr_val += 1;
            }

            num_nodes += 1;
            if num_nodes % 50_000_000 == 0 {
                println!("Done {} million points", num_nodes/1_000_000);
            }
        }
    }
    println!("Done {} points", num_nodes);
    let num_frames = last_frame_no - first_frame_no + 1;
    println!("There are {} frames, which is {} sec", num_frames, num_frames as f32/30.);

    let mut sorted_results = Vec::with_capacity((last_frame_no-first_frame_no+1) as usize);

    for frame_no in 0..num_frames {
        match results.remove(&(frame_no+first_frame_no)) {
            None => { sorted_results.push((frame_no, Vec::with_capacity(0))) },
            Some(pixels) => {
                let pixels = pixels.into_iter().collect();
                sorted_results.push((frame_no, pixels));
            }
        }
    }

    sorted_results
}

fn write_frames(frames: Frames, filename: &str, height: u32, sec_per_frame: u32, bbox: &[f32; 4]) {
    let left = bbox[0]; let bottom = bbox[1]; let right = bbox[2]; let top = bbox[3];
    let bbox_width = right - left;
    let bbox_height = top - bottom;
    let width = ((bbox_width / bbox_height) * (height as f32)) as u32;

    let mut file = BufWriter::new(fs::File::create(&filename).unwrap());

    writeln!(file, "metadata version {}", env!("CARGO_PKG_VERSION")).expect("Couldn't write metadata");
    writeln!(file, "metadata height {}", height).expect("Couldn't write metadata");
    writeln!(file, "metadata width {}", width).expect("Couldn't write metadata");
    writeln!(file, "metadata sec_per_frame {}", sec_per_frame).expect("Couldn't write metadata");
    writeln!(file, "metadata left {}", left).expect("Couldn't write metadata");
    writeln!(file, "metadata bottom {}", bottom).expect("Couldn't write metadata");
    writeln!(file, "metadata right {}", right).expect("Couldn't write metadata");
    writeln!(file, "metadata top {}", top).expect("Couldn't write metadata");
    writeln!(file, "").expect("Couldn't write metadata");

    for (frame_no, pixels) in frames.into_iter() {
        write!(file, "{}", frame_no).unwrap();
        for p in pixels {
            write!(file, ",{},{}", p.0, p.1).unwrap();
        }
        write!(file, "\n").unwrap()
    }
}

fn read_frames(filename: &str) -> Frames {
    let file = BufReader::new(fs::File::open(&filename).unwrap());
    let mut results = Frames::new();

    //let metadata: Vec<(String, String)> = file.lines().filter_map(|x| x.ok()).take_while(|x| x.len() > 0).map(|x| { let words: Vec<_> = x.split(" ").skip(1).take(2).collect(); (words[0].clone(), words[1].clone())}).collect();

    for line in file.lines().filter_map(|x| x.ok()).skip_while(|x| x.starts_with("metadata ") || x.len() == 0) {
        let frame_no = line.split(",").nth(0).unwrap().parse().unwrap();
        let pixels: Vec<(u32, u16)> = line.split(",").skip(1).collect::<Vec<_>>().chunks(2).map(|pair| (pair[0].parse().unwrap(), pair[1].parse().unwrap())).collect();
        results.push((frame_no, pixels))
    }

    results
}

fn latlon_to_pixel_index(lat: f32, lon: f32, width: u32, height: u32, bbox: &[f32; 4]) -> u32 {
    let left = bbox[0]; let bottom = bbox[1]; let right = bbox[2]; let top = bbox[3];
    let bbox_width = right - left;
    let bbox_height = top - bottom;

    let lat0 = top - lat;
    let lon0 = lon - left;

    let x = ((lon0/bbox_width)*(width as f32)) as u32;
    let y = ((lat0/bbox_height)*(height as f32)) as u32;

    let i = y * width + x;

    assert!(i < width*height, "{} L{}, lat = {} lon = {} width = {} height = {} bbox = {:?} x = {} y = {} i = {}", file!(), line!(), lat, lon, width, height, bbox, x, y, i);

    i
}

fn create_gif(frames: Frames, output_image_filename: &str, height: u32, bbox: &[f32; 4], colour_ramp: &ColourRamp) {
    let mut output_file = fs::File::create(output_image_filename).expect("Can't create image");

    let left = bbox[0]; let bottom = bbox[1]; let right = bbox[2]; let top = bbox[3];
    let bbox_width = right - left;
    let bbox_height = top - bottom;
    let width = ((bbox_width / bbox_height) * (height as f32)) as u32;

    // FIXME change width/height to u16?
    let mut encoder = gif::Encoder::new(&mut output_file, width as u16, height as u16, &colour_ramp.palette()).expect("Couldn't create encoder");
    encoder.set(gif::Repeat::Infinite).expect("Couldn't get inf repeat");

    // TODO have a global palette

    let mut image = vec![None; (width*height) as usize];

    for (frame_no, pixels) in frames.into_iter() {

        for i in 0..image.len() {
            if image[i].is_some() && image[i].unwrap() > 0. {
                image[i] = image[i].map(|x| x*0.95);
            }
        }

        for (i, magnitude) in pixels {
            // FIXME sometimes the value is invalid
            //assert!(i < width*height, "{} L{}, width = {} height = {} i = {}", file!(), line!(), width, height, i);
            if i < width*height {
                let new_value= image[i as usize].unwrap_or(0f32) + (magnitude as f32);
                image[i as usize] = Some(new_value);
            }
        }

        let mut pixels = Vec::with_capacity(image.len() * 4);
        for p in image.iter().cloned() {
            let index = colour_ramp.index_for_magnitude(p.map(|x| x.round() as u32));
            pixels.push(index);
        }

        let mut frame = gif::Frame::from_indexed_pixels(width as u16, height as u16, pixels.as_mut_slice(), None);
        // 30 fps, and delay is in units of 10ms.
        frame.delay = 100 / 30;

        encoder.write_frame(&frame).expect("Couldn't write frame");

        if frame_no % 30 == 0 {
            println!("Wrote frame {}", frame_no);
        }

    }

}

fn main() {
    let matches = App::new("osm-history-animation")
                          .arg(Arg::with_name("input").long("input").short("i")
                               .takes_value(true))
                          .arg(Arg::with_name("output") .long("output").short("o")
                               .takes_value(true))
                          .arg(Arg::with_name("height") .long("height").short("h")
                               .takes_value(true))
                          .arg(Arg::with_name("spf") .long("sec-per-frame").short("s")
                               .takes_value(true))
                          .arg(Arg::with_name("colour_ramp") .long("colour-ramp")
                               .takes_value(true))
                          .arg(Arg::with_name("save-intermediate").long("save-intermediate"))
                          .arg(Arg::with_name("load-intermediate").long("load-intermediate"))
                          .arg(Arg::with_name("bbox").long("bbox").takes_value(true).short("b"))
                          .setting(AppSettings::AllowLeadingHyphen)
                          .get_matches();

    let input_filename = matches.value_of("input").unwrap();
    let output_filename = matches.value_of("output").unwrap();
    let height: u32 = matches.value_of("height").unwrap().parse().unwrap();
    let sec_per_frame: u32 = matches.value_of("spf").unwrap().parse().unwrap();

    let bbox = match matches.value_of("bbox") {
        None => [-180., -90., 180., 90.],
        Some(text) => {
            let coords: Vec<f32> = text.split(",").map(|x| x.parse().unwrap() ).collect();
            [coords[0], coords[1], coords[2], coords[3]]
        }
    };

    let frames = if matches.is_present("load-intermediate") {
        println!("Reading frames from {}", input_filename);
        read_frames(&input_filename)
    } else {
        println!("Reading PBF file {}", input_filename);
        read_pbf(&input_filename, height, sec_per_frame, &bbox)
    };

    if matches.is_present("save-intermediate") {
        println!("Saving frame details to {}", output_filename);
        write_frames(frames, &output_filename, height, sec_per_frame, &bbox);
    } else {
        let colour_ramp = ColourRamp::new_from_filename(matches.value_of("colour_ramp").unwrap());
        println!("Creating image {}", output_filename);
        create_gif(frames, &output_filename, height, &bbox, &colour_ramp);
    }
    println!("\nFinished");
}
<|endoftext|>