<commit_before>use nom::multispace;
use std::str;

use common::{integer_literal, sql_identifier, Literal};
use column::Column;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ArithmeticOperator {
    Add,
    Subtract,
    Multiply,
    Divide,
}

#[derive(Debug, Clone, PartialEq)]
pub enum ArithmeticBase {
    Column(Column),
    Scalar(Literal),
}

#[derive(Debug, Clone, PartialEq)]
pub struct ArithmeticExpression {
    op: ArithmeticOperator,
    left: ArithmeticBase,
    right: ArithmeticBase,
}

impl ArithmeticExpression {
    pub fn new(op: ArithmeticOperator, left: ArithmeticBase, right: ArithmeticBase) -> Self {
        Self {
            op: op,
            left: left,
            right: right,
        }
    }
}

/// Parse standard math operators.
/// TODO(malte): this doesn't currently observe operator precedence.
named!(pub arithmetic_operator<&[u8], ArithmeticOperator>,
    alt_complete!(
          map!(tag!("+"), |_| ArithmeticOperator::Add)
        | map!(tag!("-"), |_| ArithmeticOperator::Subtract)
        | map!(tag!("*"), |_| ArithmeticOperator::Multiply)
        | map!(tag!("/"), |_| ArithmeticOperator::Divide)
    )
);

/// Base case for nested arithmetic expressions: column name or literal.
named!(pub arithmetic_base<&[u8], ArithmeticBase>,
    alt_complete!(
          map!(integer_literal, |il| ArithmeticBase::Scalar(il))
        | map!(sql_identifier, |si| {
              let c = Column::from(str::from_utf8(si).unwrap());
              ArithmeticBase::Column(c)
          })
    )
);

/// Parse simple arithmetic expressions combining literals, and columns and literals.
/// TODO(malte): this doesn't currently support nested expressions.
named!(pub arithmetic_expression<&[u8], ArithmeticExpression>,
    complete!(chain!(
        left: arithmetic_base ~
        multispace? ~
        op: arithmetic_operator ~
        multispace? ~
        right: arithmetic_base,
        || {
            ArithmeticExpression {
                op: op,
                left: left,
                right: right
            }
        }
    ))
);

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

    #[test]
    fn it_parses_arithmetic_expressions() {
        use super::ArithmeticOperator::*;
        use super::ArithmeticBase::*;

        let lit_ae = ["5 + 42", "5+42", "5 * 42", "5 - 42", "5 / 42"];
        let col_lit_ae = ["foo+5", "foo + 5", "5 + foo"];

        let expected_lit_ae = [
            ArithmeticExpression::new(Add, Scalar(5.into()), Scalar(42.into())),
            ArithmeticExpression::new(Add, Scalar(5.into()), Scalar(42.into())),
            ArithmeticExpression::new(Multiply, Scalar(5.into()), Scalar(42.into())),
            ArithmeticExpression::new(Subtract, Scalar(5.into()), Scalar(42.into())),
            ArithmeticExpression::new(Divide, Scalar(5.into()), Scalar(42.into())),
        ];
        let expected_col_lit_ae = [
            ArithmeticExpression::new(Add, Column("foo".into()), Scalar(5.into())),
            ArithmeticExpression::new(Add, Column("foo".into()), Scalar(5.into())),
            ArithmeticExpression::new(Add, Scalar(5.into()), Column("foo".into())),
        ];

        for (i, e) in lit_ae.iter().enumerate() {
            let res = arithmetic_expression(e.as_bytes());
            assert!(res.is_done());
            assert_eq!(res.unwrap().1, expected_lit_ae[i]);
        }

        for (i, e) in col_lit_ae.iter().enumerate() {
            let res = arithmetic_expression(e.as_bytes());
            assert!(res.is_done());
            assert_eq!(res.unwrap().1, expected_col_lit_ae[i]);
        }

    }
}
<commit_msg>Can do arithmetic on columns defined by functions<commit_after>use nom::multispace;
use std::str;

use common::{column_identifier_no_alias, integer_literal, Literal};
use column::{Column, FunctionExpression};

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ArithmeticOperator {
    Add,
    Subtract,
    Multiply,
    Divide,
}

#[derive(Debug, Clone, PartialEq)]
pub enum ArithmeticBase {
    Column(Column),
    Scalar(Literal),
}

#[derive(Debug, Clone, PartialEq)]
pub struct ArithmeticExpression {
    op: ArithmeticOperator,
    left: ArithmeticBase,
    right: ArithmeticBase,
}

impl ArithmeticExpression {
    pub fn new(op: ArithmeticOperator, left: ArithmeticBase, right: ArithmeticBase) -> Self {
        Self {
            op: op,
            left: left,
            right: right,
        }
    }
}

/// Parse standard math operators.
/// TODO(malte): this doesn't currently observe operator precedence.
named!(pub arithmetic_operator<&[u8], ArithmeticOperator>,
    alt_complete!(
          map!(tag!("+"), |_| ArithmeticOperator::Add)
        | map!(tag!("-"), |_| ArithmeticOperator::Subtract)
        | map!(tag!("*"), |_| ArithmeticOperator::Multiply)
        | map!(tag!("/"), |_| ArithmeticOperator::Divide)
    )
);

/// Base case for nested arithmetic expressions: column name or literal.
named!(pub arithmetic_base<&[u8], ArithmeticBase>,
    alt_complete!(
          map!(integer_literal, |il| ArithmeticBase::Scalar(il))
        | map!(column_identifier_no_alias, |ci| ArithmeticBase::Column(ci))
    )
);

/// Parse simple arithmetic expressions combining literals, and columns and literals.
/// TODO(malte): this doesn't currently support nested expressions.
named!(pub arithmetic_expression<&[u8], ArithmeticExpression>,
    complete!(chain!(
        left: arithmetic_base ~
        multispace? ~
        op: arithmetic_operator ~
        multispace? ~
        right: arithmetic_base,
        || {
            ArithmeticExpression {
                op: op,
                left: left,
                right: right
            }
        }
    ))
);

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

    #[test]
    fn it_parses_arithmetic_expressions() {
        use super::ArithmeticOperator::*;
        use super::ArithmeticBase::Scalar;
        use super::ArithmeticBase::Column as ABColumn;

        let lit_ae = ["5 + 42", "5+42", "5 * 42", "5 - 42", "5 / 42"];
        // N.B. trailing space in "5 + foo " is required because `sql_identifier`'s keyword
        // detection requires a follow-up character (in practice, there always is one because we
        // use semicolon-terminated queries).
        let col_lit_ae = ["foo+5", "foo + 5", "5 + foo ", "MAX(foo)-3333"];

        let expected_lit_ae = [
            ArithmeticExpression::new(Add, Scalar(5.into()), Scalar(42.into())),
            ArithmeticExpression::new(Add, Scalar(5.into()), Scalar(42.into())),
            ArithmeticExpression::new(Multiply, Scalar(5.into()), Scalar(42.into())),
            ArithmeticExpression::new(Subtract, Scalar(5.into()), Scalar(42.into())),
            ArithmeticExpression::new(Divide, Scalar(5.into()), Scalar(42.into())),
        ];
        let expected_col_lit_ae = [
            ArithmeticExpression::new(Add, ABColumn("foo".into()), Scalar(5.into())),
            ArithmeticExpression::new(Add, ABColumn("foo".into()), Scalar(5.into())),
            ArithmeticExpression::new(Add, Scalar(5.into()), ABColumn("foo".into())),
            ArithmeticExpression::new(
                Subtract,
                ABColumn(Column {
                    name: String::from("max(foo)"),
                    alias: None,
                    table: None,
                    function: Some(Box::new(FunctionExpression::Max("foo".into()))),
                }),
                Scalar(3333.into()),
            ),
        ];

        for (i, e) in lit_ae.iter().enumerate() {
            let res = arithmetic_expression(e.as_bytes());
            assert!(res.is_done());
            assert_eq!(res.unwrap().1, expected_lit_ae[i]);
        }

        for (i, e) in col_lit_ae.iter().enumerate() {
            let res = arithmetic_expression(e.as_bytes());
            assert!(res.is_done());
            assert_eq!(res.unwrap().1, expected_col_lit_ae[i]);
        }
    }
}
<|endoftext|>


<commit_before>//! `Attribute`s are optional modifiers to functions, function parameters, and return types.

#[llvm_versions(3.9..=latest)]
use llvm_sys::core::{
    LLVMGetEnumAttributeKind, LLVMGetEnumAttributeKindForName, LLVMGetEnumAttributeValue, LLVMGetLastEnumAttributeKind,
    LLVMGetStringAttributeKind, LLVMGetStringAttributeValue, LLVMIsEnumAttribute, LLVMIsStringAttribute,
};
#[llvm_versions(12.0..=latest)]
use llvm_sys::core::{LLVMGetTypeAttributeValue, LLVMIsTypeAttribute};
#[llvm_versions(3.9..=latest)]
use llvm_sys::prelude::LLVMAttributeRef;

#[llvm_versions(3.9..=latest)]
use std::ffi::CStr;

use crate::types::AnyTypeEnum;
#[cfg(feature = "internal-getters")]
use crate::LLVMReference;

// SubTypes: Attribute<Enum>, Attribute<String>
/// Functions, function parameters, and return types can have `Attribute`s to indicate
/// how they should be treated by optimizations and code generation.
#[llvm_versions(3.9..=latest)]
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Attribute {
    pub(crate) attribute: LLVMAttributeRef,
}

#[llvm_versions(3.9..=latest)]
impl Attribute {
    pub(crate) unsafe fn new(attribute: LLVMAttributeRef) -> Self {
        debug_assert!(!attribute.is_null());

        Attribute { attribute }
    }

    /// Determines whether or not an `Attribute` is an enum. This method will
    /// likely be removed in the future in favor of `Attribute`s being generically
    /// defined.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    ///
    /// let context = Context::create();
    /// let enum_attribute = context.create_enum_attribute(0, 10);
    ///
    /// assert!(enum_attribute.is_enum());
    /// ```
    pub fn is_enum(self) -> bool {
        unsafe { LLVMIsEnumAttribute(self.attribute) == 1 }
    }

    /// Determines whether or not an `Attribute` is a string. This method will
    /// likely be removed in the future in favor of `Attribute`s being generically
    /// defined.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    ///
    /// let context = Context::create();
    /// let string_attribute = context.create_string_attribute("my_key_123", "my_val");
    ///
    /// assert!(string_attribute.is_string());
    /// ```
    pub fn is_string(self) -> bool {
        unsafe { LLVMIsStringAttribute(self.attribute) == 1 }
    }

    /// Determines whether or not an `Attribute` is a type attribute. This method will
    /// likely be removed in the future in favor of `Attribute`s being generically
    /// defined.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    /// use inkwell::attributes::Attribute;
    ///
    /// let context = Context::create();
    /// let kind_id = Attribute::get_named_enum_kind_id("sret");
    /// let type_attribute = context.create_type_attribute(
    ///     kind_id,
    ///     context.i32_type().into(),
    /// );
    ///
    /// assert!(type_attribute.is_type());
    /// ```
    #[llvm_versions(12.0..=latest)]
    pub fn is_type(self) -> bool {
        unsafe { LLVMIsTypeAttribute(self.attribute) == 1 }
    }

    /// Gets the enum kind id associated with a builtin name.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::attributes::Attribute;
    ///
    /// // This kind id doesn't exist:
    /// assert_eq!(Attribute::get_named_enum_kind_id("foobar"), 0);
    ///
    /// // These are real kind ids:
    /// assert_eq!(Attribute::get_named_enum_kind_id("align"), 1);
    /// assert_eq!(Attribute::get_named_enum_kind_id("builtin"), 5);
    /// ```
    pub fn get_named_enum_kind_id(name: &str) -> u32 {
        unsafe { LLVMGetEnumAttributeKindForName(name.as_ptr() as *const ::libc::c_char, name.len()) }
    }

    /// Gets the kind id associated with an enum `Attribute`.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    ///
    /// let context = Context::create();
    /// let enum_attribute = context.create_enum_attribute(0, 10);
    ///
    /// assert_eq!(enum_attribute.get_enum_kind_id(), 0);
    /// ```
    ///
    /// This function also works for type `Attribute`s.
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    /// use inkwell::attributes::Attribute;
    /// use inkwell::types::AnyType;
    ///
    /// let context = Context::create();
    /// let kind_id = Attribute::get_named_enum_kind_id("sret");
    /// let any_type = context.i32_type().as_any_type_enum();
    /// let type_attribute = context.create_type_attribute(
    ///     kind_id,
    ///     any_type,
    /// );
    ///
    /// assert_eq!(type_attribute.get_enum_kind_id(), kind_id);
    /// ```
    pub fn get_enum_kind_id(self) -> u32 {
        assert!(self.get_enum_kind_id_is_valid()); // FIXME: SubTypes

        unsafe { LLVMGetEnumAttributeKind(self.attribute) }
    }

    #[llvm_versions(3.6..12.0)]
    fn get_enum_kind_id_is_valid(self) -> bool {
        self.is_enum()
    }

    #[llvm_versions(12.0..=latest)]
    fn get_enum_kind_id_is_valid(self) -> bool {
        self.is_enum() || self.is_type()
    }

    /// Gets the last enum kind id associated with builtin names.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::attributes::Attribute;
    ///
    /// assert_eq!(Attribute::get_last_enum_kind_id(), 56);
    /// ```
    pub fn get_last_enum_kind_id() -> u32 {
        unsafe { LLVMGetLastEnumAttributeKind() }
    }

    /// Gets the value associated with an enum `Attribute`.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    ///
    /// let context = Context::create();
    /// let enum_attribute = context.create_enum_attribute(0, 10);
    ///
    /// assert_eq!(enum_attribute.get_enum_value(), 10);
    /// ```
    pub fn get_enum_value(self) -> u64 {
        assert!(self.is_enum()); // FIXME: SubTypes

        unsafe { LLVMGetEnumAttributeValue(self.attribute) }
    }

    /// Gets the string kind id associated with a string attribute.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    ///
    /// let context = Context::create();
    /// let string_attribute = context.create_string_attribute("my_key", "my_val");
    ///
    /// assert_eq!(string_attribute.get_string_kind_id().to_str(), Ok("my_key"));
    /// ```
    pub fn get_string_kind_id(&self) -> &CStr {
        assert!(self.is_string()); // FIXME: SubTypes

        let mut length = 0;
        let cstr_ptr = unsafe { LLVMGetStringAttributeKind(self.attribute, &mut length) };

        unsafe { CStr::from_ptr(cstr_ptr) }
    }

    /// Gets the string value associated with a string attribute.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    ///
    /// let context = Context::create();
    /// let string_attribute = context.create_string_attribute("my_key", "my_val");
    ///
    /// assert_eq!(string_attribute.get_string_value().to_str(), Ok("my_val"));
    /// ```
    pub fn get_string_value(&self) -> &CStr {
        assert!(self.is_string()); // FIXME: SubTypes

        let mut length = 0;
        let cstr_ptr = unsafe { LLVMGetStringAttributeValue(self.attribute, &mut length) };

        unsafe { CStr::from_ptr(cstr_ptr) }
    }

    /// Gets the type associated with a type attribute.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    /// use inkwell::attributes::Attribute;
    /// use inkwell::types::AnyType;
    ///
    /// let context = Context::create();
    /// let kind_id = Attribute::get_named_enum_kind_id("sret");
    /// let any_type = context.i32_type().as_any_type_enum();
    /// let type_attribute = context.create_type_attribute(
    ///     kind_id,
    ///     any_type,
    /// );
    ///
    /// assert!(type_attribute.is_type());
    /// assert_eq!(type_attribute.get_type_value(), any_type);
    /// assert_ne!(type_attribute.get_type_value(), context.i64_type().as_any_type_enum());
    /// ```
    #[llvm_versions(12.0..=latest)]
    pub fn get_type_value(&self) -> AnyTypeEnum {
        assert!(self.is_type()); // FIXME: SubTypes

        unsafe { AnyTypeEnum::new(LLVMGetTypeAttributeValue(self.attribute)) }
    }
}

/// An `AttributeLoc` determines where on a function an attribute is assigned to.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum AttributeLoc {
    /// Assign to the `FunctionValue`'s return type.
    Return,
    /// Assign to one of the `FunctionValue`'s params (0-indexed).
    Param(u32),
    /// Assign to the `FunctionValue` itself.
    Function,
}

impl AttributeLoc {
    pub(crate) fn get_index(self) -> u32 {
        match self {
            AttributeLoc::Return => 0,
            AttributeLoc::Param(index) => {
                assert!(
                    index <= u32::max_value() - 2,
                    "Param index must be <= u32::max_value() - 2"
                );

                index + 1
            },
            AttributeLoc::Function => u32::max_value(),
        }
    }
}

#[cfg(feature = "internal-getters")]
impl LLVMReference<LLVMAttributeRef> for Attribute {
    unsafe fn get_ref(&self) -> LLVMAttributeRef {
        self.attribute
    }
}
<commit_msg>conditional compilation so example does not use function that is not there yet<commit_after>//! `Attribute`s are optional modifiers to functions, function parameters, and return types.

#[llvm_versions(3.9..=latest)]
use llvm_sys::core::{
    LLVMGetEnumAttributeKind, LLVMGetEnumAttributeKindForName, LLVMGetEnumAttributeValue, LLVMGetLastEnumAttributeKind,
    LLVMGetStringAttributeKind, LLVMGetStringAttributeValue, LLVMIsEnumAttribute, LLVMIsStringAttribute,
};
#[llvm_versions(12.0..=latest)]
use llvm_sys::core::{LLVMGetTypeAttributeValue, LLVMIsTypeAttribute};
#[llvm_versions(3.9..=latest)]
use llvm_sys::prelude::LLVMAttributeRef;

#[llvm_versions(3.9..=latest)]
use std::ffi::CStr;

use crate::types::AnyTypeEnum;
#[cfg(feature = "internal-getters")]
use crate::LLVMReference;

// SubTypes: Attribute<Enum>, Attribute<String>
/// Functions, function parameters, and return types can have `Attribute`s to indicate
/// how they should be treated by optimizations and code generation.
#[llvm_versions(3.9..=latest)]
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Attribute {
    pub(crate) attribute: LLVMAttributeRef,
}

#[llvm_versions(3.9..=latest)]
impl Attribute {
    pub(crate) unsafe fn new(attribute: LLVMAttributeRef) -> Self {
        debug_assert!(!attribute.is_null());

        Attribute { attribute }
    }

    /// Determines whether or not an `Attribute` is an enum. This method will
    /// likely be removed in the future in favor of `Attribute`s being generically
    /// defined.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    ///
    /// let context = Context::create();
    /// let enum_attribute = context.create_enum_attribute(0, 10);
    ///
    /// assert!(enum_attribute.is_enum());
    /// ```
    pub fn is_enum(self) -> bool {
        unsafe { LLVMIsEnumAttribute(self.attribute) == 1 }
    }

    /// Determines whether or not an `Attribute` is a string. This method will
    /// likely be removed in the future in favor of `Attribute`s being generically
    /// defined.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    ///
    /// let context = Context::create();
    /// let string_attribute = context.create_string_attribute("my_key_123", "my_val");
    ///
    /// assert!(string_attribute.is_string());
    /// ```
    pub fn is_string(self) -> bool {
        unsafe { LLVMIsStringAttribute(self.attribute) == 1 }
    }

    /// Determines whether or not an `Attribute` is a type attribute. This method will
    /// likely be removed in the future in favor of `Attribute`s being generically
    /// defined.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    /// use inkwell::attributes::Attribute;
    ///
    /// let context = Context::create();
    /// let kind_id = Attribute::get_named_enum_kind_id("sret");
    /// let type_attribute = context.create_type_attribute(
    ///     kind_id,
    ///     context.i32_type().into(),
    /// );
    ///
    /// assert!(type_attribute.is_type());
    /// ```
    #[llvm_versions(12.0..=latest)]
    pub fn is_type(self) -> bool {
        unsafe { LLVMIsTypeAttribute(self.attribute) == 1 }
    }

    /// Gets the enum kind id associated with a builtin name.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::attributes::Attribute;
    ///
    /// // This kind id doesn't exist:
    /// assert_eq!(Attribute::get_named_enum_kind_id("foobar"), 0);
    ///
    /// // These are real kind ids:
    /// assert_eq!(Attribute::get_named_enum_kind_id("align"), 1);
    /// assert_eq!(Attribute::get_named_enum_kind_id("builtin"), 5);
    /// ```
    pub fn get_named_enum_kind_id(name: &str) -> u32 {
        unsafe { LLVMGetEnumAttributeKindForName(name.as_ptr() as *const ::libc::c_char, name.len()) }
    }

    /// Gets the kind id associated with an enum `Attribute`.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    ///
    /// let context = Context::create();
    /// let enum_attribute = context.create_enum_attribute(0, 10);
    ///
    /// assert_eq!(enum_attribute.get_enum_kind_id(), 0);
    /// ```
    #[llvm_versions(3.6..12.0)]
    pub fn get_enum_kind_id(self) -> u32 {
        assert!(self.get_enum_kind_id_is_valid()); // FIXME: SubTypes

        unsafe { LLVMGetEnumAttributeKind(self.attribute) }
    }

    /// Gets the kind id associated with an enum `Attribute`.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    ///
    /// let context = Context::create();
    /// let enum_attribute = context.create_enum_attribute(0, 10);
    ///
    /// assert_eq!(enum_attribute.get_enum_kind_id(), 0);
    /// ```
    ///
    /// This function also works for type `Attribute`s.
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    /// use inkwell::attributes::Attribute;
    /// use inkwell::types::AnyType;
    ///
    /// let context = Context::create();
    /// let kind_id = Attribute::get_named_enum_kind_id("sret");
    /// let any_type = context.i32_type().as_any_type_enum();
    /// let type_attribute = context.create_type_attribute(
    ///     kind_id,
    ///     any_type,
    /// );
    ///
    /// assert_eq!(type_attribute.get_enum_kind_id(), kind_id);
    /// ```
    #[llvm_versions(12.0..=latest)]
    pub fn get_enum_kind_id(self) -> u32 {
        assert!(self.get_enum_kind_id_is_valid()); // FIXME: SubTypes

        unsafe { LLVMGetEnumAttributeKind(self.attribute) }
    }

    #[llvm_versions(3.6..12.0)]
    fn get_enum_kind_id_is_valid(self) -> bool {
        self.is_enum()
    }

    #[llvm_versions(12.0..=latest)]
    fn get_enum_kind_id_is_valid(self) -> bool {
        self.is_enum() || self.is_type()
    }

    /// Gets the last enum kind id associated with builtin names.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::attributes::Attribute;
    ///
    /// assert_eq!(Attribute::get_last_enum_kind_id(), 56);
    /// ```
    pub fn get_last_enum_kind_id() -> u32 {
        unsafe { LLVMGetLastEnumAttributeKind() }
    }

    /// Gets the value associated with an enum `Attribute`.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    ///
    /// let context = Context::create();
    /// let enum_attribute = context.create_enum_attribute(0, 10);
    ///
    /// assert_eq!(enum_attribute.get_enum_value(), 10);
    /// ```
    pub fn get_enum_value(self) -> u64 {
        assert!(self.is_enum()); // FIXME: SubTypes

        unsafe { LLVMGetEnumAttributeValue(self.attribute) }
    }

    /// Gets the string kind id associated with a string attribute.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    ///
    /// let context = Context::create();
    /// let string_attribute = context.create_string_attribute("my_key", "my_val");
    ///
    /// assert_eq!(string_attribute.get_string_kind_id().to_str(), Ok("my_key"));
    /// ```
    pub fn get_string_kind_id(&self) -> &CStr {
        assert!(self.is_string()); // FIXME: SubTypes

        let mut length = 0;
        let cstr_ptr = unsafe { LLVMGetStringAttributeKind(self.attribute, &mut length) };

        unsafe { CStr::from_ptr(cstr_ptr) }
    }

    /// Gets the string value associated with a string attribute.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    ///
    /// let context = Context::create();
    /// let string_attribute = context.create_string_attribute("my_key", "my_val");
    ///
    /// assert_eq!(string_attribute.get_string_value().to_str(), Ok("my_val"));
    /// ```
    pub fn get_string_value(&self) -> &CStr {
        assert!(self.is_string()); // FIXME: SubTypes

        let mut length = 0;
        let cstr_ptr = unsafe { LLVMGetStringAttributeValue(self.attribute, &mut length) };

        unsafe { CStr::from_ptr(cstr_ptr) }
    }

    /// Gets the type associated with a type attribute.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use inkwell::context::Context;
    /// use inkwell::attributes::Attribute;
    /// use inkwell::types::AnyType;
    ///
    /// let context = Context::create();
    /// let kind_id = Attribute::get_named_enum_kind_id("sret");
    /// let any_type = context.i32_type().as_any_type_enum();
    /// let type_attribute = context.create_type_attribute(
    ///     kind_id,
    ///     any_type,
    /// );
    ///
    /// assert!(type_attribute.is_type());
    /// assert_eq!(type_attribute.get_type_value(), any_type);
    /// assert_ne!(type_attribute.get_type_value(), context.i64_type().as_any_type_enum());
    /// ```
    #[llvm_versions(12.0..=latest)]
    pub fn get_type_value(&self) -> AnyTypeEnum {
        assert!(self.is_type()); // FIXME: SubTypes

        unsafe { AnyTypeEnum::new(LLVMGetTypeAttributeValue(self.attribute)) }
    }
}

/// An `AttributeLoc` determines where on a function an attribute is assigned to.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum AttributeLoc {
    /// Assign to the `FunctionValue`'s return type.
    Return,
    /// Assign to one of the `FunctionValue`'s params (0-indexed).
    Param(u32),
    /// Assign to the `FunctionValue` itself.
    Function,
}

impl AttributeLoc {
    pub(crate) fn get_index(self) -> u32 {
        match self {
            AttributeLoc::Return => 0,
            AttributeLoc::Param(index) => {
                assert!(
                    index <= u32::max_value() - 2,
                    "Param index must be <= u32::max_value() - 2"
                );

                index + 1
            },
            AttributeLoc::Function => u32::max_value(),
        }
    }
}

#[cfg(feature = "internal-getters")]
impl LLVMReference<LLVMAttributeRef> for Attribute {
    unsafe fn get_ref(&self) -> LLVMAttributeRef {
        self.attribute
    }
}
<|endoftext|>


<commit_before>use std::error::Error;
use std::io;
use std::str;
use std::io::Write;

use itertools::Itertools;
use rust_htslib::bcf;
use rust_htslib::bcf::record::Numeric;

pub struct Writer {
    inner: io::BufWriter<io::Stdout>,
    field_count: usize
}

impl Writer {
    fn new(inner: io::BufWriter<io::Stdout>) -> Self {
        Writer {
            inner: inner,
            field_count: 0
        }
    }

    fn write_integer(&mut self, value: i32) -> Result<(), Box<Error>> {
        self.write_field(format!("{}", value).as_bytes())
    }

    fn write_float(&mut self, value: f32) -> Result<(), Box<Error>> {
        self.write_field(format!("{}", value).as_bytes())
    }

    fn write_flag(&mut self, value: bool) -> Result<(), Box<Error>> {
        self.write_field(format!("{}", value).as_bytes())
    }

    fn write_field(&mut self, value: &[u8]) -> Result<(), Box<Error>> {
        if self.field_count > 0 {
            try!(self.inner.write(b"\t"));
        }
        try!(self.inner.write(value));
        self.field_count += 1;
        Ok(())
    }

    fn newline(&mut self) -> Result<(), Box<Error>> {
        try!(self.inner.write(b"\n"));
        self.field_count = 0;
        Ok(())
    }
}

const HEADER_COMMON: &'static [u8] = b"VARIANT";

pub fn to_txt(
    info_tags: &[&str],
    format_tags: &[&str],
    show_genotypes: bool
) -> Result<(), Box<Error>> {
    let mut reader = bcf::Reader::from_stdin()?;
    let mut writer = Writer::new(io::BufWriter::new(io::stdout()));

    let common_n = 5 + info_tags.len();
    try!(writer.write_field(HEADER_COMMON));
    for _ in 1..common_n {
        try!(writer.write_field(HEADER_COMMON));
    }
    let show_samples = show_genotypes || !format_tags.is_empty();
    if show_samples {
        for sample in reader.header().samples() {
            try!(writer.write_field(sample));
            for _ in 1..format_tags.len() + show_genotypes as usize {
                try!(writer.write_field(sample));
            }
        }
    }
    try!(writer.newline());
    try!(writer.write_field(b"CHROM"));
    try!(writer.write_field(b"POS"));
    try!(writer.write_field(b"REF"));
    try!(writer.write_field(b"ALT"));
    try!(writer.write_field(b"QUAL"));
    for name in info_tags {
        try!(writer.write_field(name.as_bytes()));
    }
    if show_samples {
        for _ in 0..reader.header().sample_count() {
            if show_genotypes {
                try!(writer.write_field(b"GT"));
            }
            for name in format_tags {
                try!(writer.write_field(name.as_bytes()));
            }
        }
    }
    try!(writer.newline());

    let mut record = bcf::Record::new();
    loop {
        if let Err(e) = reader.read(&mut record) {
            if e.is_eof() {
                break;
            } else {
                return Err(Box::new(e));
            }
        }

        let alleles = record.alleles().into_iter().map(|a| a.to_owned()).collect_vec();
        for (i, allele) in alleles[1..].iter().enumerate() {
            try!(writer.write_field(reader.header().rid2name(record.rid().unwrap())));
            try!(writer.write_integer(record.pos() as i32 + 1));
            try!(writer.write_field(&alleles[0]));
            try!(writer.write_field(allele));
            match record.qual() {
                q if q.is_missing() => try!(writer.write_field(b"")),
                q => try!(writer.write_float(q))
            }

            for name in info_tags {
                let _name = name.as_bytes();
                if let Ok((tag_type, tag_length)) = reader.header().info_type(_name) {
                    let get_idx = || {
                        match tag_length {
                            bcf::header::TagLength::Fixed => {
                                Ok(0)
                            },
                            bcf::header::TagLength::AltAlleles => {
                                Ok(i)
                            },
                            bcf::header::TagLength::Alleles => {
                                Ok(i + 1)
                            },
                            bcf::header::TagLength::Variable => {
                                Ok(i)
                            },
                            _ => Err(Box::new(ParseError::UnsupportedTagLength))
                        }
                    };

                    match tag_type {
                        bcf::header::TagType::Flag => {
                            try!(writer.write_flag(try!(record.info(_name).flag())));
                        },
                        bcf::header::TagType::Integer => {
                            let i = try!(get_idx());
                            if let Some(values) = try!(record.info(_name).integer()) {
                                try!(writer.write_integer(values[i]));
                            } else {
                                try!(writer.write_field(b""));
                            }
                        },
                        bcf::header::TagType::Float => {
                            let i = try!(get_idx());
                            if let Some(values) = try!(record.info(_name).float()) {
                                try!(writer.write_float(values[i]));
                            } else {
                                try!(writer.write_field(b""));
                            }
                        },
                        bcf::header::TagType::String => {
                            let i = try!(get_idx());
                            if let Some(values) = try!(record.info(_name).string()) {
                                try!(writer.write_field(values[i]));
                            } else {
                                try!(writer.write_field(b""));
                            }
                        }
                    }
                } else {
                    // tag undefined, write NA
                    try!(writer.write_field(b""));
                }
            }

            let genotypes = if show_genotypes {
                let genotypes = try!(record.genotypes());
                Some(
                    (0..reader.header().sample_count() as usize).map(|s| {
                        format!("{}", genotypes.get(s))
                    }).collect_vec()
                )
            } else {
                None
            };

            for s in 0..reader.header().sample_count() as usize {
                if let Some(ref genotypes) = genotypes {
                    try!(writer.write_field(genotypes[s].as_bytes()));
                }
                for name in format_tags {
                    let _name = name.as_bytes();
                    if let Ok((tag_type, tag_length)) = reader.header().format_type(_name) {
                        let i = match tag_length {
                            bcf::header::TagLength::Fixed => {
                                0
                            },
                            bcf::header::TagLength::AltAlleles => {
                                i
                            },
                            bcf::header::TagLength::Alleles => {
                                i + 1
                            },
                            _ => return Err(Box::new(ParseError::UnsupportedTagLength))
                        };

                        match tag_type {
                            bcf::header::TagType::Flag => {
                                panic!("there is no flag type for format");
                            },
                            bcf::header::TagType::Integer => {
                                try!(writer.write_field(format!("{}", try!(record.format(_name).integer())[s][i]).as_bytes()));
                            },
                            bcf::header::TagType::Float => {
                                try!(writer.write_field(format!("{}", try!(record.format(_name).float())[s][i]).as_bytes()));
                            },
                            bcf::header::TagType::String => {
                                try!(writer.write_field(try!(record.format(_name).string())[s]));
                            }
                        }
                    } else {
                        // tag undefined, write NA
                        try!(writer.write_field(b""));
                    }
                }
            }
            try!(writer.newline());
        }
    }

    Ok(())
}

quick_error! {
    #[derive(Debug)]
    pub enum ParseError {
        UnsupportedTagLength {
            description("currently, only R, A, and 1 are supported multiplicities of tags")
        }

    }
}
<commit_msg>Fix index in case of variable info field.<commit_after>use std::error::Error;
use std::io;
use std::str;
use std::io::Write;

use itertools::Itertools;
use rust_htslib::bcf;
use rust_htslib::bcf::record::Numeric;

pub struct Writer {
    inner: io::BufWriter<io::Stdout>,
    field_count: usize
}

impl Writer {
    fn new(inner: io::BufWriter<io::Stdout>) -> Self {
        Writer {
            inner: inner,
            field_count: 0
        }
    }

    fn write_integer(&mut self, value: i32) -> Result<(), Box<Error>> {
        self.write_field(format!("{}", value).as_bytes())
    }

    fn write_float(&mut self, value: f32) -> Result<(), Box<Error>> {
        self.write_field(format!("{}", value).as_bytes())
    }

    fn write_flag(&mut self, value: bool) -> Result<(), Box<Error>> {
        self.write_field(format!("{}", value).as_bytes())
    }

    fn write_field(&mut self, value: &[u8]) -> Result<(), Box<Error>> {
        if self.field_count > 0 {
            try!(self.inner.write(b"\t"));
        }
        try!(self.inner.write(value));
        self.field_count += 1;
        Ok(())
    }

    fn newline(&mut self) -> Result<(), Box<Error>> {
        try!(self.inner.write(b"\n"));
        self.field_count = 0;
        Ok(())
    }
}

const HEADER_COMMON: &'static [u8] = b"VARIANT";

pub fn to_txt(
    info_tags: &[&str],
    format_tags: &[&str],
    show_genotypes: bool
) -> Result<(), Box<Error>> {
    let mut reader = bcf::Reader::from_stdin()?;
    let mut writer = Writer::new(io::BufWriter::new(io::stdout()));

    let common_n = 5 + info_tags.len();
    try!(writer.write_field(HEADER_COMMON));
    for _ in 1..common_n {
        try!(writer.write_field(HEADER_COMMON));
    }
    let show_samples = show_genotypes || !format_tags.is_empty();
    if show_samples {
        for sample in reader.header().samples() {
            try!(writer.write_field(sample));
            for _ in 1..format_tags.len() + show_genotypes as usize {
                try!(writer.write_field(sample));
            }
        }
    }
    try!(writer.newline());
    try!(writer.write_field(b"CHROM"));
    try!(writer.write_field(b"POS"));
    try!(writer.write_field(b"REF"));
    try!(writer.write_field(b"ALT"));
    try!(writer.write_field(b"QUAL"));
    for name in info_tags {
        try!(writer.write_field(name.as_bytes()));
    }
    if show_samples {
        for _ in 0..reader.header().sample_count() {
            if show_genotypes {
                try!(writer.write_field(b"GT"));
            }
            for name in format_tags {
                try!(writer.write_field(name.as_bytes()));
            }
        }
    }
    try!(writer.newline());

    let mut record = bcf::Record::new();
    loop {
        if let Err(e) = reader.read(&mut record) {
            if e.is_eof() {
                break;
            } else {
                return Err(Box::new(e));
            }
        }

        let alleles = record.alleles().into_iter().map(|a| a.to_owned()).collect_vec();
        for (i, allele) in alleles[1..].iter().enumerate() {
            try!(writer.write_field(reader.header().rid2name(record.rid().unwrap())));
            try!(writer.write_integer(record.pos() as i32 + 1));
            try!(writer.write_field(&alleles[0]));
            try!(writer.write_field(allele));
            match record.qual() {
                q if q.is_missing() => try!(writer.write_field(b"")),
                q => try!(writer.write_float(q))
            }

            for name in info_tags {
                let _name = name.as_bytes();
                if let Ok((tag_type, tag_length)) = reader.header().info_type(_name) {
                    let get_idx = || {
                        match tag_length {
                            bcf::header::TagLength::Fixed => {
                                Ok(0)
                            },
                            bcf::header::TagLength::AltAlleles => {
                                Ok(i)
                            },
                            bcf::header::TagLength::Alleles => {
                                Ok(i + 1)
                            },
                            bcf::header::TagLength::Variable => {
                                Ok(0)
                            },
                            _ => Err(Box::new(ParseError::UnsupportedTagLength))
                        }
                    };

                    match tag_type {
                        bcf::header::TagType::Flag => {
                            try!(writer.write_flag(try!(record.info(_name).flag())));
                        },
                        bcf::header::TagType::Integer => {
                            let i = try!(get_idx());
                            if let Some(values) = try!(record.info(_name).integer()) {
                                try!(writer.write_integer(values[i]));
                            } else {
                                try!(writer.write_field(b""));
                            }
                        },
                        bcf::header::TagType::Float => {
                            let i = try!(get_idx());
                            if let Some(values) = try!(record.info(_name).float()) {
                                try!(writer.write_float(values[i]));
                            } else {
                                try!(writer.write_field(b""));
                            }
                        },
                        bcf::header::TagType::String => {
                            let i = try!(get_idx());
                            if let Some(values) = try!(record.info(_name).string()) {
                                try!(writer.write_field(values[i]));
                            } else {
                                try!(writer.write_field(b""));
                            }
                        }
                    }
                } else {
                    // tag undefined, write NA
                    try!(writer.write_field(b""));
                }
            }

            let genotypes = if show_genotypes {
                let genotypes = try!(record.genotypes());
                Some(
                    (0..reader.header().sample_count() as usize).map(|s| {
                        format!("{}", genotypes.get(s))
                    }).collect_vec()
                )
            } else {
                None
            };

            for s in 0..reader.header().sample_count() as usize {
                if let Some(ref genotypes) = genotypes {
                    try!(writer.write_field(genotypes[s].as_bytes()));
                }
                for name in format_tags {
                    let _name = name.as_bytes();
                    if let Ok((tag_type, tag_length)) = reader.header().format_type(_name) {
                        let i = match tag_length {
                            bcf::header::TagLength::Fixed => {
                                0
                            },
                            bcf::header::TagLength::AltAlleles => {
                                i
                            },
                            bcf::header::TagLength::Alleles => {
                                i + 1
                            },
                            _ => return Err(Box::new(ParseError::UnsupportedTagLength))
                        };

                        match tag_type {
                            bcf::header::TagType::Flag => {
                                panic!("there is no flag type for format");
                            },
                            bcf::header::TagType::Integer => {
                                try!(writer.write_field(format!("{}", try!(record.format(_name).integer())[s][i]).as_bytes()));
                            },
                            bcf::header::TagType::Float => {
                                try!(writer.write_field(format!("{}", try!(record.format(_name).float())[s][i]).as_bytes()));
                            },
                            bcf::header::TagType::String => {
                                try!(writer.write_field(try!(record.format(_name).string())[s]));
                            }
                        }
                    } else {
                        // tag undefined, write NA
                        try!(writer.write_field(b""));
                    }
                }
            }
            try!(writer.newline());
        }
    }

    Ok(())
}

quick_error! {
    #[derive(Debug)]
    pub enum ParseError {
        UnsupportedTagLength {
            description("currently, only R, A, and 1 are supported multiplicities of tags")
        }

    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate mdbook;
#[macro_use]
extern crate clap;
extern crate crossbeam;

// Dependencies for the Watch feature
#[cfg(feature = "watch")]
extern crate notify;
#[cfg(feature = "watch")]
extern crate time;

// Dependencies for the Serve feature
#[cfg(feature = "serve")]
extern crate iron;
#[cfg(feature = "serve")]
extern crate staticfile;
#[cfg(feature = "serve")]
extern crate ws;

use std::env;
use std::error::Error;
use std::io::{self, Write};
use std::path::{Path, PathBuf};

use clap::{App, ArgMatches, SubCommand, AppSettings};

// Uses for the Watch feature
#[cfg(feature = "watch")]
use notify::Watcher;
#[cfg(feature = "watch")]
use std::sync::mpsc::channel;

use mdbook::MDBook;

const NAME: &'static str = "mdbook";

fn main() {
    // Create a list of valid arguments and sub-commands
    let matches = App::new(NAME)
                    .about("Create a book in form of a static website from markdown files")
                    .author("Mathieu David <mathieudavid@mathieudavid.org>")
                    // Get the version from our Cargo.toml using clap's crate_version!() macro
                    .version(&*format!("v{}", crate_version!()))
                    .setting(AppSettings::SubcommandRequired)
                    .after_help("For more information about a specific command, try `mdbook <command> --help`")
                    .subcommand(SubCommand::with_name("init")
                        .about("Create boilerplate structure and files in the directory")
                        // the {n} denotes a newline which will properly aligned in all help messages
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'")
                        .arg_from_usage("--theme 'Copies the default theme into your source folder'")
                        .arg_from_usage("--force 'skip confirmation prompts'"))
                    .subcommand(SubCommand::with_name("build")
                        .about("Build the book from the markdown files")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'"))
                    .subcommand(SubCommand::with_name("watch")
                        .about("Watch the files for changes")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'"))
                    .subcommand(SubCommand::with_name("serve")
                        .about("Serve the book at http://localhost:3000. Rebuild and reload on change.")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'")
                        .arg_from_usage("-p, --port=[port] 'Use another port{n}(Defaults to 3000)'")
                        .arg_from_usage("-w, --websocket-port=[ws-port] 'Use another port for the websocket connection (livereload){n}(Defaults to 3001)'"))
                    .subcommand(SubCommand::with_name("test")
                        .about("Test that code samples compile"))
                    .get_matches();

    // Check which subcomamnd the user ran...
    let res = match matches.subcommand() {
        ("init", Some(sub_matches)) => init(sub_matches),
        ("build", Some(sub_matches)) => build(sub_matches),
        #[cfg(feature = "watch")]
        ("watch", Some(sub_matches)) => watch(sub_matches),
        #[cfg(feature = "serve")]
        ("serve", Some(sub_matches)) => serve(sub_matches),
        ("test", Some(sub_matches)) => test(sub_matches),
        (_, _) => unreachable!(),
    };

    if let Err(e) = res {
        writeln!(&mut io::stderr(), "An error occured:\n{}", e).ok();
    }
}

// Simple function that user comfirmation
fn confirm() -> bool {
    io::stdout().flush().unwrap();
    let mut s = String::new();
    io::stdin().read_line(&mut s).ok();
    match &*s.trim() {
        "Y" | "y" | "yes" | "Yes" => true,
        _ => false,
    }
}

// Init command implementation
fn init(args: &ArgMatches) -> Result<(), Box<Error>> {

    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir);

    // Call the function that does the initialization
    try!(book.init());

    // If flag `--theme` is present, copy theme to src
    if args.is_present("theme") {

        // Skip this if `--force` is present
        if !args.is_present("force") {
            // Print warning
            print!("\nCopying the default theme to {:?}", book.get_src());
            println!("could potentially overwrite files already present in that directory.");
            print!("\nAre you sure you want to continue? (y/n) ");

            // Read answer from user and exit if it's not 'yes'
            if !confirm() {
                println!("\nSkipping...\n");
                println!("All done, no errors...");
                ::std::process::exit(0);
            }
        }

        // Call the function that copies the theme
        try!(book.copy_theme());
        println!("\nTheme copied.");

    }

    // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`
    let is_dest_inside_root = book.get_dest().starts_with(book.get_root());

    if !args.is_present("force") && is_dest_inside_root {
        println!("\nDo you want a .gitignore to be created? (y/n)");

        if confirm() {
            book.create_gitignore();
            println!("\n.gitignore created.");
        }
    }

    println!("\nAll done, no errors...");

    Ok(())
}

// Build command implementation
fn build(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    try!(book.build());

    Ok(())
}

// Watch command implementation
#[cfg(feature = "watch")]
fn watch(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    trigger_on_change(&mut book, |event, book| {
        if let Some(path) = event.path {
            println!("File changed: {:?}\nBuilding book...\n", path);
            match book.build() {
                Err(e) => println!("Error while building: {:?}", e),
                _ => {},
            }
            println!("");
        }
    });

    Ok(())
}

// Watch command implementation
#[cfg(feature = "serve")]
fn serve(args: &ArgMatches) -> Result<(), Box<Error>> {
    const RELOAD_COMMAND: &'static str = "reload";

    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();
    let port = args.value_of("port").unwrap_or("3000");
    let ws_port = args.value_of("ws-port").unwrap_or("3001");

    let address = format!("localhost:{}", port);
    let ws_address = format!("localhost:{}", ws_port);

    book.set_livereload(format!(r#"
        <script type="text/javascript">
            var socket = new WebSocket("ws://localhost:{}");
            socket.onmessage = function (event) {{
                if (event.data === "{}") {{
                    socket.close();
                    location.reload(true); // force reload from server (not from cache)
                }}
            }};

            window.onbeforeunload = function() {{
                socket.close();
            }}
        </script>
    "#, ws_port, RELOAD_COMMAND).to_owned());

    try!(book.build());

    let staticfile = staticfile::Static::new(book.get_dest());
    let iron = iron::Iron::new(staticfile);
    let _iron = iron.http(&*address).unwrap();

    let ws_server = ws::WebSocket::new(|_| {
        |_| {
            Ok(())
        }
    }).unwrap();

    let broadcaster = ws_server.broadcaster();

    std::thread::spawn(move || {
        ws_server.listen(&*ws_address).unwrap();
    });

    trigger_on_change(&mut book, move |event, book| {
        if let Some(path) = event.path {
            println!("File changed: {:?}\nBuilding book...\n", path);
            match book.build() {
                Err(e) => println!("Error while building: {:?}", e),
                _ => broadcaster.send(RELOAD_COMMAND).unwrap(),
            }
            println!("");
        }
    });

    Ok(())
}

fn test(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    try!(book.test());

    Ok(())
}

fn get_book_dir(args: &ArgMatches) -> PathBuf {
    if let Some(dir) = args.value_of("dir") {
        // Check if path is relative from current dir, or absolute...
        let p = Path::new(dir);
        if p.is_relative() {
            env::current_dir().unwrap().join(dir)
        } else {
            p.to_path_buf()
        }
    } else {
        env::current_dir().unwrap()
    }
}

// Calls the closure when a book source file is changed. This is blocking!
fn trigger_on_change<F>(book: &mut MDBook, closure: F) -> ()
    where F: Fn(notify::Event, &mut MDBook) -> ()
{
    // Create a channel to receive the events.
    let (tx, rx) = channel();

    let w: Result<notify::RecommendedWatcher, notify::Error> = notify::Watcher::new(tx);

    match w {
        Ok(mut watcher) => {
            // Add the source directory to the watcher
            if let Err(e) = watcher.watch(book.get_src()) {
                println!("Error while watching {:?}:\n    {:?}", book.get_src(), e);
                ::std::process::exit(0);
            };

            // Add the book.json file to the watcher if it exists, because it's not
            // located in the source directory
            if let Err(_) = watcher.watch(book.get_root().join("book.json")) {
                // do nothing if book.json is not found
            }

            let mut previous_time = time::get_time();

            println!("\nListening for changes...\n");

            loop {
                match rx.recv() {
                    Ok(event) => {
                        // Skip the event if an event has already been issued in the last second
                        let time = time::get_time();
                        if time - previous_time < time::Duration::seconds(1) {
                            continue;
                        } else {
                            previous_time = time;
                        }

                        closure(event, book);
                    },
                    Err(e) => {
                        println!("An error occured: {:?}", e);
                    },
                }
            }
        },
        Err(e) => {
            println!("Error while trying to watch the files:\n\n\t{:?}", e);
            ::std::process::exit(0);
        },
    }
}
<commit_msg>Add address output to mdbook serve<commit_after>#[macro_use]
extern crate mdbook;
#[macro_use]
extern crate clap;
extern crate crossbeam;

// Dependencies for the Watch feature
#[cfg(feature = "watch")]
extern crate notify;
#[cfg(feature = "watch")]
extern crate time;

// Dependencies for the Serve feature
#[cfg(feature = "serve")]
extern crate iron;
#[cfg(feature = "serve")]
extern crate staticfile;
#[cfg(feature = "serve")]
extern crate ws;

use std::env;
use std::error::Error;
use std::io::{self, Write};
use std::path::{Path, PathBuf};

use clap::{App, ArgMatches, SubCommand, AppSettings};

// Uses for the Watch feature
#[cfg(feature = "watch")]
use notify::Watcher;
#[cfg(feature = "watch")]
use std::sync::mpsc::channel;

use mdbook::MDBook;

const NAME: &'static str = "mdbook";

fn main() {
    // Create a list of valid arguments and sub-commands
    let matches = App::new(NAME)
                    .about("Create a book in form of a static website from markdown files")
                    .author("Mathieu David <mathieudavid@mathieudavid.org>")
                    // Get the version from our Cargo.toml using clap's crate_version!() macro
                    .version(&*format!("v{}", crate_version!()))
                    .setting(AppSettings::SubcommandRequired)
                    .after_help("For more information about a specific command, try `mdbook <command> --help`")
                    .subcommand(SubCommand::with_name("init")
                        .about("Create boilerplate structure and files in the directory")
                        // the {n} denotes a newline which will properly aligned in all help messages
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'")
                        .arg_from_usage("--theme 'Copies the default theme into your source folder'")
                        .arg_from_usage("--force 'skip confirmation prompts'"))
                    .subcommand(SubCommand::with_name("build")
                        .about("Build the book from the markdown files")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'"))
                    .subcommand(SubCommand::with_name("watch")
                        .about("Watch the files for changes")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'"))
                    .subcommand(SubCommand::with_name("serve")
                        .about("Serve the book at http://localhost:3000. Rebuild and reload on change.")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when ommitted)'")
                        .arg_from_usage("-p, --port=[port] 'Use another port{n}(Defaults to 3000)'")
                        .arg_from_usage("-w, --websocket-port=[ws-port] 'Use another port for the websocket connection (livereload){n}(Defaults to 3001)'"))
                    .subcommand(SubCommand::with_name("test")
                        .about("Test that code samples compile"))
                    .get_matches();

    // Check which subcomamnd the user ran...
    let res = match matches.subcommand() {
        ("init", Some(sub_matches)) => init(sub_matches),
        ("build", Some(sub_matches)) => build(sub_matches),
        #[cfg(feature = "watch")]
        ("watch", Some(sub_matches)) => watch(sub_matches),
        #[cfg(feature = "serve")]
        ("serve", Some(sub_matches)) => serve(sub_matches),
        ("test", Some(sub_matches)) => test(sub_matches),
        (_, _) => unreachable!(),
    };

    if let Err(e) = res {
        writeln!(&mut io::stderr(), "An error occured:\n{}", e).ok();
    }
}

// Simple function that user comfirmation
fn confirm() -> bool {
    io::stdout().flush().unwrap();
    let mut s = String::new();
    io::stdin().read_line(&mut s).ok();
    match &*s.trim() {
        "Y" | "y" | "yes" | "Yes" => true,
        _ => false,
    }
}

// Init command implementation
fn init(args: &ArgMatches) -> Result<(), Box<Error>> {

    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir);

    // Call the function that does the initialization
    try!(book.init());

    // If flag `--theme` is present, copy theme to src
    if args.is_present("theme") {

        // Skip this if `--force` is present
        if !args.is_present("force") {
            // Print warning
            print!("\nCopying the default theme to {:?}", book.get_src());
            println!("could potentially overwrite files already present in that directory.");
            print!("\nAre you sure you want to continue? (y/n) ");

            // Read answer from user and exit if it's not 'yes'
            if !confirm() {
                println!("\nSkipping...\n");
                println!("All done, no errors...");
                ::std::process::exit(0);
            }
        }

        // Call the function that copies the theme
        try!(book.copy_theme());
        println!("\nTheme copied.");

    }

    // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`
    let is_dest_inside_root = book.get_dest().starts_with(book.get_root());

    if !args.is_present("force") && is_dest_inside_root {
        println!("\nDo you want a .gitignore to be created? (y/n)");

        if confirm() {
            book.create_gitignore();
            println!("\n.gitignore created.");
        }
    }

    println!("\nAll done, no errors...");

    Ok(())
}

// Build command implementation
fn build(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    try!(book.build());

    Ok(())
}

// Watch command implementation
#[cfg(feature = "watch")]
fn watch(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    trigger_on_change(&mut book, |event, book| {
        if let Some(path) = event.path {
            println!("File changed: {:?}\nBuilding book...\n", path);
            match book.build() {
                Err(e) => println!("Error while building: {:?}", e),
                _ => {},
            }
            println!("");
        }
    });

    Ok(())
}

// Watch command implementation
#[cfg(feature = "serve")]
fn serve(args: &ArgMatches) -> Result<(), Box<Error>> {
    const RELOAD_COMMAND: &'static str = "reload";

    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();
    let port = args.value_of("port").unwrap_or("3000");
    let ws_port = args.value_of("ws-port").unwrap_or("3001");

    let address = format!("localhost:{}", port);
    let ws_address = format!("localhost:{}", ws_port);

    book.set_livereload(format!(r#"
        <script type="text/javascript">
            var socket = new WebSocket("ws://localhost:{}");
            socket.onmessage = function (event) {{
                if (event.data === "{}") {{
                    socket.close();
                    location.reload(true); // force reload from server (not from cache)
                }}
            }};

            window.onbeforeunload = function() {{
                socket.close();
            }}
        </script>
    "#, ws_port, RELOAD_COMMAND).to_owned());

    try!(book.build());

    let staticfile = staticfile::Static::new(book.get_dest());
    let iron = iron::Iron::new(staticfile);
    let _iron = iron.http(&*address).unwrap();

    let ws_server = ws::WebSocket::new(|_| {
        |_| {
            Ok(())
        }
    }).unwrap();

    let broadcaster = ws_server.broadcaster();

    std::thread::spawn(move || {
        ws_server.listen(&*ws_address).unwrap();
    });

    println!("\nServing on {}", address);

    trigger_on_change(&mut book, move |event, book| {
        if let Some(path) = event.path {
            println!("File changed: {:?}\nBuilding book...\n", path);
            match book.build() {
                Err(e) => println!("Error while building: {:?}", e),
                _ => broadcaster.send(RELOAD_COMMAND).unwrap(),
            }
            println!("");
        }
    });

    Ok(())
}

fn test(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    try!(book.test());

    Ok(())
}

fn get_book_dir(args: &ArgMatches) -> PathBuf {
    if let Some(dir) = args.value_of("dir") {
        // Check if path is relative from current dir, or absolute...
        let p = Path::new(dir);
        if p.is_relative() {
            env::current_dir().unwrap().join(dir)
        } else {
            p.to_path_buf()
        }
    } else {
        env::current_dir().unwrap()
    }
}

// Calls the closure when a book source file is changed. This is blocking!
fn trigger_on_change<F>(book: &mut MDBook, closure: F) -> ()
    where F: Fn(notify::Event, &mut MDBook) -> ()
{
    // Create a channel to receive the events.
    let (tx, rx) = channel();

    let w: Result<notify::RecommendedWatcher, notify::Error> = notify::Watcher::new(tx);

    match w {
        Ok(mut watcher) => {
            // Add the source directory to the watcher
            if let Err(e) = watcher.watch(book.get_src()) {
                println!("Error while watching {:?}:\n    {:?}", book.get_src(), e);
                ::std::process::exit(0);
            };

            // Add the book.json file to the watcher if it exists, because it's not
            // located in the source directory
            if let Err(_) = watcher.watch(book.get_root().join("book.json")) {
                // do nothing if book.json is not found
            }

            let mut previous_time = time::get_time();

            println!("\nListening for changes...\n");

            loop {
                match rx.recv() {
                    Ok(event) => {
                        // Skip the event if an event has already been issued in the last second
                        let time = time::get_time();
                        if time - previous_time < time::Duration::seconds(1) {
                            continue;
                        } else {
                            previous_time = time;
                        }

                        closure(event, book);
                    },
                    Err(e) => {
                        println!("An error occured: {:?}", e);
                    },
                }
            }
        },
        Err(e) => {
            println!("Error while trying to watch the files:\n\n\t{:?}", e);
            ::std::process::exit(0);
        },
    }
}
<|endoftext|>


<commit_before>#![warn(clippy::all, rust_2018_idioms)]

use cargo_registry::{boot, App, Env};
use std::{
    fs::File,
    sync::{mpsc::channel, Arc, Mutex},
    thread,
    time::Duration,
};

use civet::Server as CivetServer;
use conduit_hyper::Service;
use futures_util::future::FutureExt;
use reqwest::blocking::Client;

const CORE_THREADS: usize = 4;

#[allow(clippy::large_enum_variant)]
enum Server {
    Civet(CivetServer),
    Hyper(tokio::runtime::Runtime, tokio::task::JoinHandle<()>),
}

use Server::*;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize logging
    env_logger::init();

    let config = cargo_registry::Config::default();
    let client = Client::new();

    let app = App::new(config.clone(), Some(client));
    let app = cargo_registry::build_handler(Arc::new(app));

    // On every server restart, ensure the categories available in the database match
    // the information in *src/categories.toml*.
    let categories_toml = include_str!("../boot/categories.toml");
    boot::categories::sync(categories_toml).unwrap();

    let heroku = dotenv::var("HEROKU").is_ok();
    let fastboot = dotenv::var("USE_FASTBOOT").is_ok();

    let port = if heroku {
        8888
    } else {
        dotenv::var("PORT")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(8888)
    };
    let threads = dotenv::var("SERVER_THREADS")
        .map(|s| s.parse().expect("SERVER_THREADS was not a valid number"))
        .unwrap_or_else(|_| {
            if config.env == Env::Development {
                5
            } else {
                50
            }
        });

    let server = if dotenv::var("USE_HYPER").is_ok() {
        use tokio::io::AsyncWriteExt;
        use tokio::signal::unix::{signal, SignalKind};

        println!("Booting with a hyper based server");

        let mut rt = tokio::runtime::Builder::new()
            .threaded_scheduler()
            .enable_all()
            .core_threads(CORE_THREADS)
            .max_threads(threads as usize + CORE_THREADS)
            .build()
            .unwrap();

        let handler = Arc::new(conduit_hyper::BlockingHandler::new(app));
        let make_service =
            hyper::service::make_service_fn(move |socket: &hyper::server::conn::AddrStream| {
                let addr = socket.remote_addr();
                let handler = handler.clone();
                async move { Service::from_blocking(handler, addr) }
            });

        let addr = ([127, 0, 0, 1], port).into();
        let server = rt.block_on(async { hyper::Server::bind(&addr).serve(make_service) });

        let mut sig_int = rt.block_on(async { signal(SignalKind::interrupt()) })?;
        let mut sig_term = rt.block_on(async { signal(SignalKind::terminate()) })?;

        let server = server.with_graceful_shutdown(async move {
            // Wait for either signal
            futures_util::select! {
                _ = sig_int.recv().fuse() => (),
                _ = sig_term.recv().fuse() => (),
            };
            let mut stdout = tokio::io::stdout();
            stdout.write_all(b"Starting graceful shutdown\n").await.ok();
        });

        let server = rt.spawn(async { server.await.unwrap() });
        Hyper(rt, server)
    } else {
        println!("Booting with a civet based server");
        let mut cfg = civet::Config::new();
        cfg.port(port).threads(threads).keep_alive(true);
        Civet(CivetServer::start(cfg, app).unwrap())
    };

    println!("listening on port {}", port);

    // Give tokio a chance to spawn the first worker thread
    thread::sleep(Duration::from_millis(10));

    // Creating this file tells heroku to tell nginx that the application is ready
    // to receive traffic.
    if heroku {
        let path = if fastboot {
            "/tmp/backend-initialized"
        } else {
            "/tmp/app-initialized"
        };
        println!("Writing to {}", path);
        File::create(path).unwrap();
    }

    // Block the main thread until the server has shutdown
    match server {
        Hyper(mut rt, server) => {
            rt.block_on(async { server.await.unwrap() });
        }
        Civet(server) => {
            let (tx, rx) = channel::<()>();
            ctrlc_handler(move || tx.send(()).unwrap_or(()));
            rx.recv().unwrap();
            drop(server);
        }
    }

    println!("Server has gracefully shutdown!");
    Ok(())
}

fn ctrlc_handler<F>(f: F)
where
    F: FnOnce() + Send + 'static,
{
    let call_once = Mutex::new(Some(f));

    ctrlc::set_handler(move || {
        if let Some(f) = call_once.lock().unwrap().take() {
            println!("Starting graceful shutdown");
            f();
        } else {
            println!("Already sent signal to start graceful shutdown");
        }
    })
    .unwrap();
}
<commit_msg>Auto merge of #2621 - jtgeibel:hyper-as-default, r=JohnTitor<commit_after>#![warn(clippy::all, rust_2018_idioms)]

use cargo_registry::{boot, App, Env};
use std::{
    fs::File,
    sync::{mpsc::channel, Arc, Mutex},
    thread,
    time::Duration,
};

use civet::Server as CivetServer;
use conduit_hyper::Service;
use futures_util::future::FutureExt;
use reqwest::blocking::Client;

const CORE_THREADS: usize = 4;

#[allow(clippy::large_enum_variant)]
enum Server {
    Civet(CivetServer),
    Hyper(tokio::runtime::Runtime, tokio::task::JoinHandle<()>),
}

use Server::*;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize logging
    env_logger::init();

    let config = cargo_registry::Config::default();
    let client = Client::new();

    let app = App::new(config.clone(), Some(client));
    let app = cargo_registry::build_handler(Arc::new(app));

    // On every server restart, ensure the categories available in the database match
    // the information in *src/categories.toml*.
    let categories_toml = include_str!("../boot/categories.toml");
    boot::categories::sync(categories_toml).unwrap();

    let heroku = dotenv::var("HEROKU").is_ok();
    let fastboot = dotenv::var("USE_FASTBOOT").is_ok();

    let port = if heroku {
        8888
    } else {
        dotenv::var("PORT")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(8888)
    };
    let threads = dotenv::var("SERVER_THREADS")
        .map(|s| s.parse().expect("SERVER_THREADS was not a valid number"))
        .unwrap_or_else(|_| {
            if config.env == Env::Development {
                5
            } else {
                50
            }
        });

    let server = if dotenv::var("WEB_USE_CIVET").is_err() {
        use tokio::io::AsyncWriteExt;
        use tokio::signal::unix::{signal, SignalKind};

        println!("Booting with a hyper based server");

        let mut rt = tokio::runtime::Builder::new()
            .threaded_scheduler()
            .enable_all()
            .core_threads(CORE_THREADS)
            .max_threads(threads as usize + CORE_THREADS)
            .build()
            .unwrap();

        let handler = Arc::new(conduit_hyper::BlockingHandler::new(app));
        let make_service =
            hyper::service::make_service_fn(move |socket: &hyper::server::conn::AddrStream| {
                let addr = socket.remote_addr();
                let handler = handler.clone();
                async move { Service::from_blocking(handler, addr) }
            });

        let addr = ([127, 0, 0, 1], port).into();
        let server = rt.block_on(async { hyper::Server::bind(&addr).serve(make_service) });

        let mut sig_int = rt.block_on(async { signal(SignalKind::interrupt()) })?;
        let mut sig_term = rt.block_on(async { signal(SignalKind::terminate()) })?;

        let server = server.with_graceful_shutdown(async move {
            // Wait for either signal
            futures_util::select! {
                _ = sig_int.recv().fuse() => (),
                _ = sig_term.recv().fuse() => (),
            };
            let mut stdout = tokio::io::stdout();
            stdout.write_all(b"Starting graceful shutdown\n").await.ok();
        });

        let server = rt.spawn(async { server.await.unwrap() });
        Hyper(rt, server)
    } else {
        println!("Booting with a civet based server");
        let mut cfg = civet::Config::new();
        cfg.port(port).threads(threads).keep_alive(true);
        Civet(CivetServer::start(cfg, app).unwrap())
    };

    println!("listening on port {}", port);

    // Give tokio a chance to spawn the first worker thread
    thread::sleep(Duration::from_millis(10));

    // Creating this file tells heroku to tell nginx that the application is ready
    // to receive traffic.
    if heroku {
        let path = if fastboot {
            "/tmp/backend-initialized"
        } else {
            "/tmp/app-initialized"
        };
        println!("Writing to {}", path);
        File::create(path).unwrap();
    }

    // Block the main thread until the server has shutdown
    match server {
        Hyper(mut rt, server) => {
            rt.block_on(async { server.await.unwrap() });
        }
        Civet(server) => {
            let (tx, rx) = channel::<()>();
            ctrlc_handler(move || tx.send(()).unwrap_or(()));
            rx.recv().unwrap();
            drop(server);
        }
    }

    println!("Server has gracefully shutdown!");
    Ok(())
}

fn ctrlc_handler<F>(f: F)
where
    F: FnOnce() + Send + 'static,
{
    let call_once = Mutex::new(Some(f));

    ctrlc::set_handler(move || {
        if let Some(f) = call_once.lock().unwrap().take() {
            println!("Starting graceful shutdown");
            f();
        } else {
            println!("Already sent signal to start graceful shutdown");
        }
    })
    .unwrap();
}
<|endoftext|>