<commit_before>// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.

use kvproto::coprocessor::KeyRange;
use tipb::expression::FieldType;
use tipb::schema::ColumnInfo;

use crate::storage::{Key, Store};

use super::ranges_iter::{PointRangePolicy, RangesIterator};
use crate::coprocessor::codec::batch::LazyBatchColumnVec;
use crate::coprocessor::dag::batch::interface::*;
use crate::coprocessor::dag::expr::EvalContext;
use crate::coprocessor::dag::Scanner;
use crate::coprocessor::Result;

/// Common interfaces for table scan and index scan implementations.
pub trait ScanExecutorImpl: Send {
    /// Gets the schema.
    fn schema(&self) -> &[FieldType];

    /// Gets a mutable reference of the executor context.
    fn mut_context(&mut self) -> &mut EvalContext;

    fn build_scanner<S: Store>(&self, store: &S, desc: bool, range: KeyRange)
        -> Result<Scanner<S>>;

    fn build_column_vec(&self, scan_rows: usize) -> LazyBatchColumnVec;

    /// Accepts a key value pair and fills the column vector.
    ///
    /// The column vector does not need to be regular when there are errors during this process.
    /// However if there is no error, the column vector must be regular.
    fn process_kv_pair(
        &mut self,
        key: &[u8],
        value: &[u8],
        columns: &mut LazyBatchColumnVec,
    ) -> Result<()>;
}

/// A shared executor implementation for both table scan and index scan. Implementation differences
/// between table scan and index scan are further given via `ScanExecutorImpl`.
pub struct ScanExecutor<S: Store, I: ScanExecutorImpl, P: PointRangePolicy> {
    /// The internal scanning implementation.
    imp: I,

    store: S,
    desc: bool,

    /// Number of rows scanned from each range. Notice that this vector does not contain ranges
    /// that have not been scanned.
    scanned_rows_per_range: Vec<usize>,

    /// Iterates ranges.
    ranges: RangesIterator<P>,

    /// Row scanner.
    ///
    /// It is optional because sometimes it is not needed, e.g. when point range is given.
    /// Also, the value may be re-constructed several times if there are multiple key ranges.
    scanner: Option<Scanner<S>>,

    /// A flag indicating whether this executor is ended. When table is drained or there was an
    /// error scanning the table, this flag will be set to `true` and `next_batch` should be never
    /// called again.
    is_ended: bool,
}

impl<S: Store, I: ScanExecutorImpl, P: PointRangePolicy> ScanExecutor<S, I, P> {
    pub fn new(
        imp: I,
        store: S,
        desc: bool,
        mut key_ranges: Vec<KeyRange>,
        point_range_policy: P,
    ) -> Result<Self> {
        crate::coprocessor::codec::table::check_table_ranges(&key_ranges)?;
        if desc {
            key_ranges.reverse();
        }
        Ok(Self {
            imp,
            store,
            desc,
            scanned_rows_per_range: Vec::with_capacity(key_ranges.len()),
            ranges: RangesIterator::new(key_ranges, point_range_policy),
            scanner: None,
            is_ended: false,
        })
    }

    /// Creates or resets the range of inner scanner.
    #[inline]
    fn reset_range(&mut self, range: KeyRange) -> Result<()> {
        self.scanner = Some(self.imp.build_scanner(&self.store, self.desc, range)?);
        Ok(())
    }

    /// Scans next row from the scanner.
    #[inline]
    fn scan_next(&mut self) -> Result<Option<(Vec<u8>, Vec<u8>)>> {
        // TODO: Key and value doesn't have to be owned
        if let Some(scanner) = self.scanner.as_mut() {
            Ok(scanner.next_row()?)
        } else {
            // `self.scanner` should never be `None` when this function is being called.
            unreachable!()
        }
    }

    /// Get one row from the store.
    #[inline]
    fn point_get(&mut self, mut range: KeyRange) -> Result<Option<(Vec<u8>, Vec<u8>)>> {
        let mut statistics = crate::storage::Statistics::default();
        // TODO: Key and value doesn't have to be owned
        let key = range.take_start();
        let value = self.store.get(&Key::from_raw(&key), &mut statistics)?;
        Ok(value.map(move |v| (key, v)))
    }

    /// Fills a column vector and returns whether or not all ranges are drained.
    ///
    /// The columns are ensured to be regular even if there are errors during the process.
    fn fill_column_vec(
        &mut self,
        scan_rows: usize,
        columns: &mut LazyBatchColumnVec,
    ) -> Result<bool> {
        assert!(scan_rows > 0);

        loop {
            let range = self.ranges.next();
            let some_row = match range {
                super::ranges_iter::IterStatus::NewPointRange(r) => {
                    self.scanned_rows_per_range.push(0);
                    self.point_get(r)?
                }
                super::ranges_iter::IterStatus::NewNonPointRange(r) => {
                    self.scanned_rows_per_range.push(0);
                    self.reset_range(r)?;
                    self.scan_next()?
                }
                super::ranges_iter::IterStatus::Continue => self.scan_next()?,
                super::ranges_iter::IterStatus::Drained => {
                    return Ok(true); // drained
                }
            };
            if let Some((key, value)) = some_row {
                // Retrieved one row from point range or non-point range.
                self.scanned_rows_per_range
                    .last_mut()
                    .map_or((), |val| *val += 1);

                if let Err(e) = self.imp.process_kv_pair(&key, &value, columns) {
                    // When there are errors in `process_kv_pair`, columns' length may not be
                    // identical. For example, the filling process may be partially done so that
                    // first several columns have N rows while the rest have N-1 rows. Since we do
                    // not immediately fail when there are errors, these irregular columns may
                    // further cause future executors to panic. So let's truncate these columns to
                    // make they all have N-1 rows in that case.
                    columns.truncate_into_equal_length();
                    return Err(e);
                }

                if columns.rows_len() >= scan_rows {
                    return Ok(false); // not drained
                }
            } else {
                // No more row in the range.
                self.ranges.notify_drained();
            }
        }
    }
}

/// Extracts `FieldType` from `ColumnInfo`.
// TODO: Embed FieldType in ColumnInfo directly in Cop DAG v2 to remove this function.
pub fn field_type_from_column_info(ci: &ColumnInfo) -> FieldType {
    let mut field_type = FieldType::new();
    field_type.set_tp(ci.get_tp());
    field_type.set_flag(ci.get_flag() as u32); // FIXME: This `as u32` is really awful.
    field_type.set_flen(ci.get_columnLen());
    field_type.set_decimal(ci.get_decimal());
    field_type.set_collate(ci.get_collation());
    // Note: Charset is not provided in column info.
    field_type
}

/// Checks whether the given columns info are supported.
pub fn check_columns_info_supported(columns_info: &[ColumnInfo]) -> Result<()> {
    use cop_datatype::EvalType;
    use cop_datatype::FieldTypeAccessor;
    use std::convert::TryFrom;

    for column in columns_info {
        box_try!(EvalType::try_from(column.tp()));
    }
    Ok(())
}

impl<S: Store, I: ScanExecutorImpl, P: PointRangePolicy> BatchExecutor for ScanExecutor<S, I, P> {
    #[inline]
    fn schema(&self) -> &[FieldType] {
        self.imp.schema()
    }

    #[inline]
    fn next_batch(&mut self, scan_rows: usize) -> BatchExecuteResult {
        assert!(!self.is_ended);
        assert!(scan_rows > 0);

        let mut data = self.imp.build_column_vec(scan_rows);
        let is_drained = self.fill_column_vec(scan_rows, &mut data);

        data.assert_columns_equal_length();

        // TODO
        // If `is_drained.is_err()`, it means that there is an error after *successfully* retrieving
        // these rows. After that, if we only consumes some of the rows (TopN / Limit), we should
        // ignore this error.

        match &is_drained {
            // Note: `self.is_ended` is only used for assertion purpose.
            Err(_) | Ok(true) => self.is_ended = true,
            Ok(false) => {}
        };

        BatchExecuteResult {
            data,
            is_drained,
            warnings: self.imp.mut_context().take_warnings(),
        }
    }

    fn collect_statistics(&mut self, destination: &mut BatchExecuteStatistics) {
        for (index, num) in self.scanned_rows_per_range.iter_mut().enumerate() {
            destination.scanned_rows_per_range[index] += *num;
            *num = 0;
        }
        if let Some(scanner) = &mut self.scanner {
            scanner.collect_statistics_into(&mut destination.cf_stats);
        }
    }
}
<commit_msg>do the has_pk_handle check before type check (#4788)<commit_after>// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.

use kvproto::coprocessor::KeyRange;
use tipb::expression::FieldType;
use tipb::schema::ColumnInfo;

use crate::storage::{Key, Store};

use super::ranges_iter::{PointRangePolicy, RangesIterator};
use crate::coprocessor::codec::batch::LazyBatchColumnVec;
use crate::coprocessor::dag::batch::interface::*;
use crate::coprocessor::dag::expr::EvalContext;
use crate::coprocessor::dag::Scanner;
use crate::coprocessor::Result;

/// Common interfaces for table scan and index scan implementations.
pub trait ScanExecutorImpl: Send {
    /// Gets the schema.
    fn schema(&self) -> &[FieldType];

    /// Gets a mutable reference of the executor context.
    fn mut_context(&mut self) -> &mut EvalContext;

    fn build_scanner<S: Store>(&self, store: &S, desc: bool, range: KeyRange)
        -> Result<Scanner<S>>;

    fn build_column_vec(&self, scan_rows: usize) -> LazyBatchColumnVec;

    /// Accepts a key value pair and fills the column vector.
    ///
    /// The column vector does not need to be regular when there are errors during this process.
    /// However if there is no error, the column vector must be regular.
    fn process_kv_pair(
        &mut self,
        key: &[u8],
        value: &[u8],
        columns: &mut LazyBatchColumnVec,
    ) -> Result<()>;
}

/// A shared executor implementation for both table scan and index scan. Implementation differences
/// between table scan and index scan are further given via `ScanExecutorImpl`.
pub struct ScanExecutor<S: Store, I: ScanExecutorImpl, P: PointRangePolicy> {
    /// The internal scanning implementation.
    imp: I,

    store: S,
    desc: bool,

    /// Number of rows scanned from each range. Notice that this vector does not contain ranges
    /// that have not been scanned.
    scanned_rows_per_range: Vec<usize>,

    /// Iterates ranges.
    ranges: RangesIterator<P>,

    /// Row scanner.
    ///
    /// It is optional because sometimes it is not needed, e.g. when point range is given.
    /// Also, the value may be re-constructed several times if there are multiple key ranges.
    scanner: Option<Scanner<S>>,

    /// A flag indicating whether this executor is ended. When table is drained or there was an
    /// error scanning the table, this flag will be set to `true` and `next_batch` should be never
    /// called again.
    is_ended: bool,
}

impl<S: Store, I: ScanExecutorImpl, P: PointRangePolicy> ScanExecutor<S, I, P> {
    pub fn new(
        imp: I,
        store: S,
        desc: bool,
        mut key_ranges: Vec<KeyRange>,
        point_range_policy: P,
    ) -> Result<Self> {
        crate::coprocessor::codec::table::check_table_ranges(&key_ranges)?;
        if desc {
            key_ranges.reverse();
        }
        Ok(Self {
            imp,
            store,
            desc,
            scanned_rows_per_range: Vec::with_capacity(key_ranges.len()),
            ranges: RangesIterator::new(key_ranges, point_range_policy),
            scanner: None,
            is_ended: false,
        })
    }

    /// Creates or resets the range of inner scanner.
    #[inline]
    fn reset_range(&mut self, range: KeyRange) -> Result<()> {
        self.scanner = Some(self.imp.build_scanner(&self.store, self.desc, range)?);
        Ok(())
    }

    /// Scans next row from the scanner.
    #[inline]
    fn scan_next(&mut self) -> Result<Option<(Vec<u8>, Vec<u8>)>> {
        // TODO: Key and value doesn't have to be owned
        if let Some(scanner) = self.scanner.as_mut() {
            Ok(scanner.next_row()?)
        } else {
            // `self.scanner` should never be `None` when this function is being called.
            unreachable!()
        }
    }

    /// Get one row from the store.
    #[inline]
    fn point_get(&mut self, mut range: KeyRange) -> Result<Option<(Vec<u8>, Vec<u8>)>> {
        let mut statistics = crate::storage::Statistics::default();
        // TODO: Key and value doesn't have to be owned
        let key = range.take_start();
        let value = self.store.get(&Key::from_raw(&key), &mut statistics)?;
        Ok(value.map(move |v| (key, v)))
    }

    /// Fills a column vector and returns whether or not all ranges are drained.
    ///
    /// The columns are ensured to be regular even if there are errors during the process.
    fn fill_column_vec(
        &mut self,
        scan_rows: usize,
        columns: &mut LazyBatchColumnVec,
    ) -> Result<bool> {
        assert!(scan_rows > 0);

        loop {
            let range = self.ranges.next();
            let some_row = match range {
                super::ranges_iter::IterStatus::NewPointRange(r) => {
                    self.scanned_rows_per_range.push(0);
                    self.point_get(r)?
                }
                super::ranges_iter::IterStatus::NewNonPointRange(r) => {
                    self.scanned_rows_per_range.push(0);
                    self.reset_range(r)?;
                    self.scan_next()?
                }
                super::ranges_iter::IterStatus::Continue => self.scan_next()?,
                super::ranges_iter::IterStatus::Drained => {
                    return Ok(true); // drained
                }
            };
            if let Some((key, value)) = some_row {
                // Retrieved one row from point range or non-point range.
                self.scanned_rows_per_range
                    .last_mut()
                    .map_or((), |val| *val += 1);

                if let Err(e) = self.imp.process_kv_pair(&key, &value, columns) {
                    // When there are errors in `process_kv_pair`, columns' length may not be
                    // identical. For example, the filling process may be partially done so that
                    // first several columns have N rows while the rest have N-1 rows. Since we do
                    // not immediately fail when there are errors, these irregular columns may
                    // further cause future executors to panic. So let's truncate these columns to
                    // make they all have N-1 rows in that case.
                    columns.truncate_into_equal_length();
                    return Err(e);
                }

                if columns.rows_len() >= scan_rows {
                    return Ok(false); // not drained
                }
            } else {
                // No more row in the range.
                self.ranges.notify_drained();
            }
        }
    }
}

/// Extracts `FieldType` from `ColumnInfo`.
// TODO: Embed FieldType in ColumnInfo directly in Cop DAG v2 to remove this function.
pub fn field_type_from_column_info(ci: &ColumnInfo) -> FieldType {
    let mut field_type = FieldType::new();
    field_type.set_tp(ci.get_tp());
    field_type.set_flag(ci.get_flag() as u32); // FIXME: This `as u32` is really awful.
    field_type.set_flen(ci.get_columnLen());
    field_type.set_decimal(ci.get_decimal());
    field_type.set_collate(ci.get_collation());
    // Note: Charset is not provided in column info.
    field_type
}

/// Checks whether the given columns info are supported.
pub fn check_columns_info_supported(columns_info: &[ColumnInfo]) -> Result<()> {
    use cop_datatype::EvalType;
    use cop_datatype::FieldTypeAccessor;
    use std::convert::TryFrom;

    for column in columns_info {
        if column.has_pk_handle() {
            box_try!(EvalType::try_from(column.tp()));
        }
    }
    Ok(())
}

impl<S: Store, I: ScanExecutorImpl, P: PointRangePolicy> BatchExecutor for ScanExecutor<S, I, P> {
    #[inline]
    fn schema(&self) -> &[FieldType] {
        self.imp.schema()
    }

    #[inline]
    fn next_batch(&mut self, scan_rows: usize) -> BatchExecuteResult {
        assert!(!self.is_ended);
        assert!(scan_rows > 0);

        let mut data = self.imp.build_column_vec(scan_rows);
        let is_drained = self.fill_column_vec(scan_rows, &mut data);

        data.assert_columns_equal_length();

        // TODO
        // If `is_drained.is_err()`, it means that there is an error after *successfully* retrieving
        // these rows. After that, if we only consumes some of the rows (TopN / Limit), we should
        // ignore this error.

        match &is_drained {
            // Note: `self.is_ended` is only used for assertion purpose.
            Err(_) | Ok(true) => self.is_ended = true,
            Ok(false) => {}
        };

        BatchExecuteResult {
            data,
            is_drained,
            warnings: self.imp.mut_context().take_warnings(),
        }
    }

    fn collect_statistics(&mut self, destination: &mut BatchExecuteStatistics) {
        for (index, num) in self.scanned_rows_per_range.iter_mut().enumerate() {
            destination.scanned_rows_per_range[index] += *num;
            *num = 0;
        }
        if let Some(scanner) = &mut self.scanner {
            scanner.collect_statistics_into(&mut destination.cf_stats);
        }
    }
}
<|endoftext|>


<commit_before>use either::Either;
use ide_db::syntax_helpers::node_ext::walk_ty;
use itertools::Itertools;
use syntax::{
    ast::{self, edit::IndentLevel, AstNode, HasGenericParams, HasName},
    match_ast,
};

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: extract_type_alias
//
// Extracts the selected type as a type alias.
//
// ```
// struct S {
//     field: $0(u8, u8, u8)$0,
// }
// ```
// ->
// ```
// type $0Type = (u8, u8, u8);
//
// struct S {
//     field: Type,
// }
// ```
pub(crate) fn extract_type_alias(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    if ctx.has_empty_selection() {
        return None;
    }

    let ty = ctx.find_node_at_range::<ast::Type>()?;
    let item = ty.syntax().ancestors().find_map(ast::Item::cast)?;
    let assoc_owner = item.syntax().ancestors().nth(2).and_then(|it| {
        match_ast! {
            match it {
                ast::Trait(tr) => Some(Either::Left(tr)),
                ast::Impl(impl_) => Some(Either::Right(impl_)),
                _ => None,
            }
        }
    });
    let node = assoc_owner.as_ref().map_or_else(
        || item.syntax(),
        |impl_| impl_.as_ref().either(AstNode::syntax, AstNode::syntax),
    );
    let insert_pos = node.text_range().start();
    let target = ty.syntax().text_range();

    acc.add(
        AssistId("extract_type_alias", AssistKind::RefactorExtract),
        "Extract type as type alias",
        target,
        |builder| {
            let mut known_generics = match item.generic_param_list() {
                Some(it) => it.generic_params().collect(),
                None => Vec::new(),
            };
            if let Some(it) = assoc_owner.as_ref().and_then(|it| match it {
                Either::Left(it) => it.generic_param_list(),
                Either::Right(it) => it.generic_param_list(),
            }) {
                known_generics.extend(it.generic_params());
            }
            let generics = collect_used_generics(&ty, &known_generics);

            let replacement = if !generics.is_empty() {
                format!(
                    "Type<{}>",
                    generics.iter().format_with(", ", |generic, f| {
                        match generic {
                            ast::GenericParam::ConstParam(cp) => f(&cp.name().unwrap()),
                            ast::GenericParam::LifetimeParam(lp) => f(&lp.lifetime().unwrap()),
                            ast::GenericParam::TypeParam(tp) => f(&tp.name().unwrap()),
                        }
                    })
                )
            } else {
                String::from("Type")
            };
            builder.replace(target, replacement);

            let indent = IndentLevel::from_node(node);
            let generics = if !generics.is_empty() {
                format!("<{}>", generics.iter().format(", "))
            } else {
                String::new()
            };
            match ctx.config.snippet_cap {
                Some(cap) => {
                    builder.insert_snippet(
                        cap,
                        insert_pos,
                        format!("type $0Type{} = {};\n\n{}", generics, ty, indent),
                    );
                }
                None => {
                    builder.insert(
                        insert_pos,
                        format!("type Type{} = {};\n\n{}", generics, ty, indent),
                    );
                }
            }
        },
    )
}

fn collect_used_generics<'gp>(
    ty: &ast::Type,
    known_generics: &'gp [ast::GenericParam],
) -> Vec<&'gp ast::GenericParam> {
    // can't use a closure -> closure here cause lifetime inference fails for that
    fn find_lifetime(text: &str) -> impl Fn(&&ast::GenericParam) -> bool + '_ {
        move |gp: &&ast::GenericParam| match gp {
            ast::GenericParam::LifetimeParam(lp) => {
                lp.lifetime().map_or(false, |lt| lt.text() == text)
            }
            _ => false,
        }
    }

    let mut generics = Vec::new();
    walk_ty(ty, &mut |ty| match ty {
        ast::Type::PathType(ty) => {
            if let Some(path) = ty.path() {
                if let Some(name_ref) = path.as_single_name_ref() {
                    if let Some(param) = known_generics.iter().find(|gp| {
                        match gp {
                            ast::GenericParam::ConstParam(cp) => cp.name(),
                            ast::GenericParam::TypeParam(tp) => tp.name(),
                            _ => None,
                        }
                        .map_or(false, |n| n.text() == name_ref.text())
                    }) {
                        generics.push(param);
                    }
                }
                generics.extend(
                    path.segments()
                        .filter_map(|seg| seg.generic_arg_list())
                        .flat_map(|it| it.generic_args())
                        .filter_map(|it| match it {
                            ast::GenericArg::LifetimeArg(lt) => {
                                let lt = lt.lifetime()?;
                                known_generics.iter().find(find_lifetime(&lt.text()))
                            }
                            _ => None,
                        }),
                );
            }
        }
        ast::Type::ImplTraitType(impl_ty) => {
            if let Some(it) = impl_ty.type_bound_list() {
                generics.extend(
                    it.bounds()
                        .filter_map(|it| it.lifetime())
                        .filter_map(|lt| known_generics.iter().find(find_lifetime(&lt.text()))),
                );
            }
        }
        ast::Type::DynTraitType(dyn_ty) => {
            if let Some(it) = dyn_ty.type_bound_list() {
                generics.extend(
                    it.bounds()
                        .filter_map(|it| it.lifetime())
                        .filter_map(|lt| known_generics.iter().find(find_lifetime(&lt.text()))),
                );
            }
        }
        ast::Type::RefType(ref_) => generics.extend(
            ref_.lifetime().and_then(|lt| known_generics.iter().find(find_lifetime(&lt.text()))),
        ),
        _ => (),
    });
    // stable resort to lifetime, type, const
    generics.sort_by_key(|gp| match gp {
        ast::GenericParam::ConstParam(_) => 2,
        ast::GenericParam::LifetimeParam(_) => 0,
        ast::GenericParam::TypeParam(_) => 1,
    });
    generics
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn test_not_applicable_without_selection() {
        check_assist_not_applicable(
            extract_type_alias,
            r"
struct S {
    field: $0(u8, u8, u8),
}
            ",
        );
    }

    #[test]
    fn test_simple_types() {
        check_assist(
            extract_type_alias,
            r"
struct S {
    field: $0u8$0,
}
            ",
            r#"
type $0Type = u8;

struct S {
    field: Type,
}
            "#,
        );
    }

    #[test]
    fn test_generic_type_arg() {
        check_assist(
            extract_type_alias,
            r"
fn generic<T>() {}

fn f() {
    generic::<$0()$0>();
}
            ",
            r#"
fn generic<T>() {}

type $0Type = ();

fn f() {
    generic::<Type>();
}
            "#,
        );
    }

    #[test]
    fn test_inner_type_arg() {
        check_assist(
            extract_type_alias,
            r"
struct Vec<T> {}
struct S {
    v: Vec<Vec<$0Vec<u8>$0>>,
}
            ",
            r#"
struct Vec<T> {}
type $0Type = Vec<u8>;

struct S {
    v: Vec<Vec<Type>>,
}
            "#,
        );
    }

    #[test]
    fn test_extract_inner_type() {
        check_assist(
            extract_type_alias,
            r"
struct S {
    field: ($0u8$0,),
}
            ",
            r#"
type $0Type = u8;

struct S {
    field: (Type,),
}
            "#,
        );
    }

    #[test]
    fn extract_from_impl_or_trait() {
        // When invoked in an impl/trait, extracted type alias should be placed next to the
        // impl/trait, not inside.
        check_assist(
            extract_type_alias,
            r#"
impl S {
    fn f() -> $0(u8, u8)$0 {}
}
            "#,
            r#"
type $0Type = (u8, u8);

impl S {
    fn f() -> Type {}
}
            "#,
        );
        check_assist(
            extract_type_alias,
            r#"
trait Tr {
    fn f() -> $0(u8, u8)$0 {}
}
            "#,
            r#"
type $0Type = (u8, u8);

trait Tr {
    fn f() -> Type {}
}
            "#,
        );
    }

    #[test]
    fn indentation() {
        check_assist(
            extract_type_alias,
            r#"
mod m {
    fn f() -> $0u8$0 {}
}
            "#,
            r#"
mod m {
    type $0Type = u8;

    fn f() -> Type {}
}
            "#,
        );
    }

    #[test]
    fn generics() {
        check_assist(
            extract_type_alias,
            r#"
struct Struct<const C: usize>;
impl<'outer, Outer, const OUTER: usize> () {
    fn func<'inner, Inner, const INNER: usize>(_: $0&(Struct<INNER>, Struct<OUTER>, Outer, &'inner (), Inner, &'outer ())$0) {}
}
"#,
            r#"
struct Struct<const C: usize>;
type $0Type<'inner, 'outer, Outer, Inner, const INNER: usize, const OUTER: usize> = &(Struct<INNER>, Struct<OUTER>, Outer, &'inner (), Inner, &'outer ());

impl<'outer, Outer, const OUTER: usize> () {
    fn func<'inner, Inner, const INNER: usize>(_: Type<'inner, 'outer, Outer, Inner, INNER, OUTER>) {}
}
"#,
        );
    }
}
<commit_msg>add test<commit_after>use either::Either;
use ide_db::syntax_helpers::node_ext::walk_ty;
use itertools::Itertools;
use syntax::{
    ast::{self, edit::IndentLevel, AstNode, HasGenericParams, HasName},
    match_ast,
};

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: extract_type_alias
//
// Extracts the selected type as a type alias.
//
// ```
// struct S {
//     field: $0(u8, u8, u8)$0,
// }
// ```
// ->
// ```
// type $0Type = (u8, u8, u8);
//
// struct S {
//     field: Type,
// }
// ```
pub(crate) fn extract_type_alias(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
    if ctx.has_empty_selection() {
        return None;
    }

    let ty = ctx.find_node_at_range::<ast::Type>()?;
    let item = ty.syntax().ancestors().find_map(ast::Item::cast)?;
    let assoc_owner = item.syntax().ancestors().nth(2).and_then(|it| {
        match_ast! {
            match it {
                ast::Trait(tr) => Some(Either::Left(tr)),
                ast::Impl(impl_) => Some(Either::Right(impl_)),
                _ => None,
            }
        }
    });
    let node = assoc_owner.as_ref().map_or_else(
        || item.syntax(),
        |impl_| impl_.as_ref().either(AstNode::syntax, AstNode::syntax),
    );
    let insert_pos = node.text_range().start();
    let target = ty.syntax().text_range();

    acc.add(
        AssistId("extract_type_alias", AssistKind::RefactorExtract),
        "Extract type as type alias",
        target,
        |builder| {
            let mut known_generics = match item.generic_param_list() {
                Some(it) => it.generic_params().collect(),
                None => Vec::new(),
            };
            if let Some(it) = assoc_owner.as_ref().and_then(|it| match it {
                Either::Left(it) => it.generic_param_list(),
                Either::Right(it) => it.generic_param_list(),
            }) {
                known_generics.extend(it.generic_params());
            }
            let generics = collect_used_generics(&ty, &known_generics);

            let replacement = if !generics.is_empty() {
                format!(
                    "Type<{}>",
                    generics.iter().format_with(", ", |generic, f| {
                        match generic {
                            ast::GenericParam::ConstParam(cp) => f(&cp.name().unwrap()),
                            ast::GenericParam::LifetimeParam(lp) => f(&lp.lifetime().unwrap()),
                            ast::GenericParam::TypeParam(tp) => f(&tp.name().unwrap()),
                        }
                    })
                )
            } else {
                String::from("Type")
            };
            builder.replace(target, replacement);

            let indent = IndentLevel::from_node(node);
            let generics = if !generics.is_empty() {
                format!("<{}>", generics.iter().format(", "))
            } else {
                String::new()
            };
            match ctx.config.snippet_cap {
                Some(cap) => {
                    builder.insert_snippet(
                        cap,
                        insert_pos,
                        format!("type $0Type{} = {};\n\n{}", generics, ty, indent),
                    );
                }
                None => {
                    builder.insert(
                        insert_pos,
                        format!("type Type{} = {};\n\n{}", generics, ty, indent),
                    );
                }
            }
        },
    )
}

fn collect_used_generics<'gp>(
    ty: &ast::Type,
    known_generics: &'gp [ast::GenericParam],
) -> Vec<&'gp ast::GenericParam> {
    // can't use a closure -> closure here cause lifetime inference fails for that
    fn find_lifetime(text: &str) -> impl Fn(&&ast::GenericParam) -> bool + '_ {
        move |gp: &&ast::GenericParam| match gp {
            ast::GenericParam::LifetimeParam(lp) => {
                lp.lifetime().map_or(false, |lt| lt.text() == text)
            }
            _ => false,
        }
    }

    let mut generics = Vec::new();
    walk_ty(ty, &mut |ty| match ty {
        ast::Type::PathType(ty) => {
            if let Some(path) = ty.path() {
                if let Some(name_ref) = path.as_single_name_ref() {
                    if let Some(param) = known_generics.iter().find(|gp| {
                        match gp {
                            ast::GenericParam::ConstParam(cp) => cp.name(),
                            ast::GenericParam::TypeParam(tp) => tp.name(),
                            _ => None,
                        }
                        .map_or(false, |n| n.text() == name_ref.text())
                    }) {
                        generics.push(param);
                    }
                }
                generics.extend(
                    path.segments()
                        .filter_map(|seg| seg.generic_arg_list())
                        .flat_map(|it| it.generic_args())
                        .filter_map(|it| match it {
                            ast::GenericArg::LifetimeArg(lt) => {
                                let lt = lt.lifetime()?;
                                known_generics.iter().find(find_lifetime(&lt.text()))
                            }
                            _ => None,
                        }),
                );
            }
        }
        ast::Type::ImplTraitType(impl_ty) => {
            if let Some(it) = impl_ty.type_bound_list() {
                generics.extend(
                    it.bounds()
                        .filter_map(|it| it.lifetime())
                        .filter_map(|lt| known_generics.iter().find(find_lifetime(&lt.text()))),
                );
            }
        }
        ast::Type::DynTraitType(dyn_ty) => {
            if let Some(it) = dyn_ty.type_bound_list() {
                generics.extend(
                    it.bounds()
                        .filter_map(|it| it.lifetime())
                        .filter_map(|lt| known_generics.iter().find(find_lifetime(&lt.text()))),
                );
            }
        }
        ast::Type::RefType(ref_) => generics.extend(
            ref_.lifetime().and_then(|lt| known_generics.iter().find(find_lifetime(&lt.text()))),
        ),
        _ => (),
    });
    // stable resort to lifetime, type, const
    generics.sort_by_key(|gp| match gp {
        ast::GenericParam::ConstParam(_) => 2,
        ast::GenericParam::LifetimeParam(_) => 0,
        ast::GenericParam::TypeParam(_) => 1,
    });
    generics
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn test_not_applicable_without_selection() {
        check_assist_not_applicable(
            extract_type_alias,
            r"
struct S {
    field: $0(u8, u8, u8),
}
            ",
        );
    }

    #[test]
    fn test_simple_types() {
        check_assist(
            extract_type_alias,
            r"
struct S {
    field: $0u8$0,
}
            ",
            r#"
type $0Type = u8;

struct S {
    field: Type,
}
            "#,
        );
    }

    #[test]
    fn test_generic_type_arg() {
        check_assist(
            extract_type_alias,
            r"
fn generic<T>() {}

fn f() {
    generic::<$0()$0>();
}
            ",
            r#"
fn generic<T>() {}

type $0Type = ();

fn f() {
    generic::<Type>();
}
            "#,
        );
    }

    #[test]
    fn test_inner_type_arg() {
        check_assist(
            extract_type_alias,
            r"
struct Vec<T> {}
struct S {
    v: Vec<Vec<$0Vec<u8>$0>>,
}
            ",
            r#"
struct Vec<T> {}
type $0Type = Vec<u8>;

struct S {
    v: Vec<Vec<Type>>,
}
            "#,
        );
    }

    #[test]
    fn test_extract_inner_type() {
        check_assist(
            extract_type_alias,
            r"
struct S {
    field: ($0u8$0,),
}
            ",
            r#"
type $0Type = u8;

struct S {
    field: (Type,),
}
            "#,
        );
    }

    #[test]
    fn extract_from_impl_or_trait() {
        // When invoked in an impl/trait, extracted type alias should be placed next to the
        // impl/trait, not inside.
        check_assist(
            extract_type_alias,
            r#"
impl S {
    fn f() -> $0(u8, u8)$0 {}
}
            "#,
            r#"
type $0Type = (u8, u8);

impl S {
    fn f() -> Type {}
}
            "#,
        );
        check_assist(
            extract_type_alias,
            r#"
trait Tr {
    fn f() -> $0(u8, u8)$0 {}
}
            "#,
            r#"
type $0Type = (u8, u8);

trait Tr {
    fn f() -> Type {}
}
            "#,
        );
    }

    #[test]
    fn indentation() {
        check_assist(
            extract_type_alias,
            r#"
mod m {
    fn f() -> $0u8$0 {}
}
            "#,
            r#"
mod m {
    type $0Type = u8;

    fn f() -> Type {}
}
            "#,
        );
    }

    #[test]
    fn generics() {
        check_assist(
            extract_type_alias,
            r#"
struct Struct<const C: usize>;
impl<'outer, Outer, const OUTER: usize> () {
    fn func<'inner, Inner, const INNER: usize>(_: $0&(Struct<INNER>, Struct<OUTER>, Outer, &'inner (), Inner, &'outer ())$0) {}
}
"#,
            r#"
struct Struct<const C: usize>;
type $0Type<'inner, 'outer, Outer, Inner, const INNER: usize, const OUTER: usize> = &(Struct<INNER>, Struct<OUTER>, Outer, &'inner (), Inner, &'outer ());

impl<'outer, Outer, const OUTER: usize> () {
    fn func<'inner, Inner, const INNER: usize>(_: Type<'inner, 'outer, Outer, Inner, INNER, OUTER>) {}
}
"#,
        );
    }

    #[test]
    fn issue_11197 () {
        check_assist(extract_type_alias,
            r#"
struct Foo<T, const N: usize>
where
    [T; N]: Sized,
{
    arr: $0[T; N]$0,
}
            "#,
            r#"
type $0Type<T, const N: usize> = [T; N];

struct Foo<T, const N: usize>
where
    [T; N]: Sized,
{
    arr: Type<T, N>,
}
            "#);
    }
}
<|endoftext|>


<commit_before>//! API to generate `.rs` files.
//!
//! This API does not require `protoc` command present in `$PATH`.
//!
//! ```
//! extern crate protoc_rust;
//!
//! fn main() {
//!     protobuf_codegen_pure::Codegen::new()
//!         .out_dir("src/protos")
//!         .inputs(&["protos/a.proto", "protos/b.proto"])
//!         .include("protos")
//!         .run()
//!         .expect("Codegen failed.");
//! }
//! ```
//!
//! It is advisable that `protobuf-codegen-pure` build-dependecy version be the same as
//! `protobuf` dependency.
//!
//! The alternative is to use `protoc-rust` crate.

#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]

extern crate protobuf;
extern crate protobuf_codegen;

mod convert;

use std::fs;
use std::path::Path;
use std::path::PathBuf;
use std::process;

mod linked_hash_map;
mod model;
mod parse_dependencies;
mod parser;
mod path;
mod rel_path;

use linked_hash_map::LinkedHashMap;
use protobuf::descriptor::FileDescriptorProto;
use protobuf_codegen::amend_io_error;
pub use protobuf_codegen::Customize;
use protobuf_codegen::ProtoPath;
use protobuf_codegen::ProtoPathBuf;

use crate::parser::ParserErrorWithLocation;

#[cfg(test)]
mod test_against_protobuf_protos;

/// Invoke pure rust codegen. See [crate docs](crate) for example.
// TODO: merge with protoc-rust def
#[derive(Debug, Default)]
pub struct Codegen {
    /// --lang_out= param
    out_dir: PathBuf,
    /// -I args
    includes: Vec<PathBuf>,
    /// List of .proto files to compile
    inputs: Vec<PathBuf>,
    /// Customize code generation
    customize: Customize,
}

impl Codegen {
    /// Fresh new codegen object.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the output directory for codegen.
    pub fn out_dir(&mut self, out_dir: impl AsRef<Path>) -> &mut Self {
        self.out_dir = out_dir.as_ref().to_owned();
        self
    }

    /// Add an include directory.
    pub fn include(&mut self, include: impl AsRef<Path>) -> &mut Self {
        self.includes.push(include.as_ref().to_owned());
        self
    }

    /// Add include directories.
    pub fn includes(&mut self, includes: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for include in includes {
            self.include(include);
        }
        self
    }

    /// Add an input (`.proto` file).
    pub fn input(&mut self, input: impl AsRef<Path>) -> &mut Self {
        self.inputs.push(input.as_ref().to_owned());
        self
    }

    /// Add inputs (`.proto` files).
    pub fn inputs(&mut self, inputs: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for input in inputs {
            self.input(input);
        }
        self
    }

    /// Specify generated code [`Customize`] object.
    pub fn customize(&mut self, customize: Customize) -> &mut Self {
        self.customize = customize;
        self
    }

    /// Like `protoc --rust_out=...` but without requiring `protoc` or `protoc-gen-rust`
    /// commands in `$PATH`.
    pub fn run(&self) -> anyhow::Result<()> {
        let p = parse_and_typecheck(&self.includes, &self.inputs)?;

        protobuf_codegen::gen_and_write(
            &p.file_descriptors,
            &format!("protobuf-codegen-pure={}", env!("CARGO_PKG_VERSION")),
            &p.relative_paths,
            &self.out_dir,
            &self.customize,
        )
    }

    /// Similar to `run`, but prints the message to stderr and exits the process on error.
    pub fn run_from_script(&self) {
        if let Err(e) = self.run() {
            eprintln!("pure rust codegen failed: {}", e);
            process::exit(1);
        }
    }
}

#[derive(Clone)]
pub(crate) struct FileDescriptorPair {
    parsed: model::FileDescriptor,
    descriptor: protobuf::descriptor::FileDescriptorProto,
}

#[derive(Debug, thiserror::Error)]
#[error("error in `{file}`: {error}")]
struct WithFileError {
    file: String,
    #[source]
    error: anyhow::Error,
}

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("protobuf path `{0}` is not found in import path {1}")]
    FileNotFoundInImportPath(String, String),
    #[error("file `{0}` must reside in include path {1}")]
    FileMustResideInImportPath(String, String),
}

struct Run<'a> {
    parsed_files: LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    includes: &'a [PathBuf],
}

impl<'a> Run<'a> {
    fn get_file_and_all_deps_already_parsed(
        &self,
        protobuf_path: &ProtoPath,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        if let Some(_) = result.get(protobuf_path) {
            return;
        }

        let pair = self
            .parsed_files
            .get(protobuf_path)
            .expect("must be already parsed");
        result.insert(protobuf_path.to_proto_path_buf(), pair.clone());

        self.get_all_deps_already_parsed(&pair.parsed, result);
    }

    fn get_all_deps_already_parsed(
        &self,
        parsed: &model::FileDescriptor,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        for import in &parsed.imports {
            self.get_file_and_all_deps_already_parsed(&import.path, result);
        }
    }

    fn add_file(&mut self, protobuf_path: &ProtoPath, fs_path: &Path) -> anyhow::Result<()> {
        if let Some(_) = self.parsed_files.get(protobuf_path) {
            return Ok(());
        }

        let content = fs::read_to_string(fs_path)
            .map_err(|e| amend_io_error(e, format!("failed to read {:?}", fs_path)))?;

        self.add_file_content(protobuf_path, fs_path, &content)
    }

    fn add_file_content(
        &mut self,
        protobuf_path: &ProtoPath,
        fs_path: &Path,
        content: &str,
    ) -> anyhow::Result<()> {
        let parsed = model::FileDescriptor::parse(content).map_err(|e| WithFileError {
            file: format!("{}", fs_path.display()),
            error: e.into(),
        })?;

        for import in &parsed.imports {
            self.add_imported_file(&import.path)?;
        }

        let mut this_file_deps = LinkedHashMap::new();
        self.get_all_deps_already_parsed(&parsed, &mut this_file_deps);

        let this_file_deps: Vec<_> = this_file_deps.into_iter().map(|(_, v)| v).collect();

        let descriptor = convert::file_descriptor(protobuf_path, &parsed, &this_file_deps)
            .map_err(|e| WithFileError {
                file: format!("{}", fs_path.display()),
                error: e.into(),
            })?;

        self.parsed_files.insert(
            protobuf_path.to_proto_path_buf(),
            FileDescriptorPair { parsed, descriptor },
        );

        Ok(())
    }

    fn add_imported_file(&mut self, protobuf_path: &ProtoPath) -> anyhow::Result<()> {
        for include_dir in self.includes {
            let fs_path = include_dir.join(protobuf_path.to_path());
            if fs_path.exists() {
                return self.add_file(protobuf_path, &fs_path);
            }
        }

        let embedded = match protobuf_path.to_str() {
            "rustproto.proto" => Some(RUSTPROTO_PROTO),
            "google/protobuf/any.proto" => Some(ANY_PROTO),
            "google/protobuf/api.proto" => Some(API_PROTO),
            "google/protobuf/descriptor.proto" => Some(DESCRIPTOR_PROTO),
            "google/protobuf/duration.proto" => Some(DURATION_PROTO),
            "google/protobuf/empty.proto" => Some(EMPTY_PROTO),
            "google/protobuf/field_mask.proto" => Some(FIELD_MASK_PROTO),
            "google/protobuf/source_context.proto" => Some(SOURCE_CONTEXT_PROTO),
            "google/protobuf/struct.proto" => Some(STRUCT_PROTO),
            "google/protobuf/timestamp.proto" => Some(TIMESTAMP_PROTO),
            "google/protobuf/type.proto" => Some(TYPE_PROTO),
            "google/protobuf/wrappers.proto" => Some(WRAPPERS_PROTO),
            _ => None,
        };

        match embedded {
            Some(content) => self.add_file_content(protobuf_path, protobuf_path.to_path(), content),
            None => Err(Error::FileNotFoundInImportPath(
                protobuf_path.to_string(),
                format!("{:?}", self.includes),
            )
            .into()),
        }
    }

    fn strip_prefix<'b>(path: &'b Path, prefix: &Path) -> anyhow::Result<ProtoPathBuf> {
        // special handling of `.` to allow successful `strip_prefix("foo.proto", ".")
        if prefix == Path::new(".") && path.is_relative() {
            ProtoPathBuf::from_path(path)
        } else {
            ProtoPathBuf::from_path(path.strip_prefix(prefix)?)
        }
    }

    fn add_fs_file(&mut self, fs_path: &Path) -> anyhow::Result<ProtoPathBuf> {
        let relative_path = self
            .includes
            .iter()
            .filter_map(|include_dir| Self::strip_prefix(fs_path, include_dir).ok())
            .next();

        match relative_path {
            Some(relative_path) => {
                self.add_file(&relative_path, fs_path)?;
                Ok(relative_path)
            }
            None => Err(Error::FileMustResideInImportPath(
                fs_path.display().to_string(),
                format!("{:?}", self.includes),
            )
            .into()),
        }
    }
}

/// Result of parsing `.proto` files.
#[doc(hidden)]
pub struct ParsedAndTypechecked {
    /// One entry for each input `.proto` file.
    pub relative_paths: Vec<ProtoPathBuf>,
    /// All parsed `.proto` files including dependencies of input files.
    pub file_descriptors: Vec<protobuf::descriptor::FileDescriptorProto>,
}

#[doc(hidden)]
pub fn parse_and_typecheck(
    includes: &[PathBuf],
    input: &[PathBuf],
) -> anyhow::Result<ParsedAndTypechecked> {
    let mut run = Run {
        parsed_files: LinkedHashMap::new(),
        includes,
    };

    let mut relative_paths = Vec::new();

    for input in input {
        relative_paths.push(run.add_fs_file(input)?);
    }

    let file_descriptors: Vec<_> = run
        .parsed_files
        .into_iter()
        .map(|(_, v)| v.descriptor)
        .collect();

    Ok(ParsedAndTypechecked {
        relative_paths,
        file_descriptors,
    })
}

/// Parse imports from a `.proto` file.
///
/// The result is [`FileDescriptorProto`] object with only `*dependency` fields filled.
#[doc(hidden)]
pub fn parse_dependencies(content: &str) -> Result<FileDescriptorProto, ParserErrorWithLocation> {
    parse_dependencies::parse_dependencies(content)
}

// TODO: these include don't work when publishing to crates
const RUSTPROTO_PROTO: &str = include_str!("proto/rustproto.proto");
const ANY_PROTO: &str = include_str!("proto/google/protobuf/any.proto");
const API_PROTO: &str = include_str!("proto/google/protobuf/api.proto");
const DESCRIPTOR_PROTO: &str = include_str!("proto/google/protobuf/descriptor.proto");
const DURATION_PROTO: &str = include_str!("proto/google/protobuf/duration.proto");
const EMPTY_PROTO: &str = include_str!("proto/google/protobuf/empty.proto");
const FIELD_MASK_PROTO: &str = include_str!("proto/google/protobuf/field_mask.proto");
const SOURCE_CONTEXT_PROTO: &str = include_str!("proto/google/protobuf/source_context.proto");
const STRUCT_PROTO: &str = include_str!("proto/google/protobuf/struct.proto");
const TIMESTAMP_PROTO: &str = include_str!("proto/google/protobuf/timestamp.proto");
const TYPE_PROTO: &str = include_str!("proto/google/protobuf/type.proto");
const WRAPPERS_PROTO: &str = include_str!("proto/google/protobuf/wrappers.proto");
<commit_msg>Skip already added file<commit_after>//! API to generate `.rs` files.
//!
//! This API does not require `protoc` command present in `$PATH`.
//!
//! ```
//! extern crate protoc_rust;
//!
//! fn main() {
//!     protobuf_codegen_pure::Codegen::new()
//!         .out_dir("src/protos")
//!         .inputs(&["protos/a.proto", "protos/b.proto"])
//!         .include("protos")
//!         .run()
//!         .expect("Codegen failed.");
//! }
//! ```
//!
//! It is advisable that `protobuf-codegen-pure` build-dependecy version be the same as
//! `protobuf` dependency.
//!
//! The alternative is to use `protoc-rust` crate.

#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]

extern crate protobuf;
extern crate protobuf_codegen;

mod convert;

use std::fs;
use std::path::Path;
use std::path::PathBuf;
use std::process;

mod linked_hash_map;
mod model;
mod parse_dependencies;
mod parser;
mod path;
mod rel_path;

use linked_hash_map::LinkedHashMap;
use protobuf::descriptor::FileDescriptorProto;
use protobuf_codegen::amend_io_error;
pub use protobuf_codegen::Customize;
use protobuf_codegen::ProtoPath;
use protobuf_codegen::ProtoPathBuf;

use crate::parser::ParserErrorWithLocation;

#[cfg(test)]
mod test_against_protobuf_protos;

/// Invoke pure rust codegen. See [crate docs](crate) for example.
// TODO: merge with protoc-rust def
#[derive(Debug, Default)]
pub struct Codegen {
    /// --lang_out= param
    out_dir: PathBuf,
    /// -I args
    includes: Vec<PathBuf>,
    /// List of .proto files to compile
    inputs: Vec<PathBuf>,
    /// Customize code generation
    customize: Customize,
}

impl Codegen {
    /// Fresh new codegen object.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the output directory for codegen.
    pub fn out_dir(&mut self, out_dir: impl AsRef<Path>) -> &mut Self {
        self.out_dir = out_dir.as_ref().to_owned();
        self
    }

    /// Add an include directory.
    pub fn include(&mut self, include: impl AsRef<Path>) -> &mut Self {
        self.includes.push(include.as_ref().to_owned());
        self
    }

    /// Add include directories.
    pub fn includes(&mut self, includes: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for include in includes {
            self.include(include);
        }
        self
    }

    /// Add an input (`.proto` file).
    pub fn input(&mut self, input: impl AsRef<Path>) -> &mut Self {
        self.inputs.push(input.as_ref().to_owned());
        self
    }

    /// Add inputs (`.proto` files).
    pub fn inputs(&mut self, inputs: impl IntoIterator<Item = impl AsRef<Path>>) -> &mut Self {
        for input in inputs {
            self.input(input);
        }
        self
    }

    /// Specify generated code [`Customize`] object.
    pub fn customize(&mut self, customize: Customize) -> &mut Self {
        self.customize = customize;
        self
    }

    /// Like `protoc --rust_out=...` but without requiring `protoc` or `protoc-gen-rust`
    /// commands in `$PATH`.
    pub fn run(&self) -> anyhow::Result<()> {
        let p = parse_and_typecheck(&self.includes, &self.inputs)?;

        protobuf_codegen::gen_and_write(
            &p.file_descriptors,
            &format!("protobuf-codegen-pure={}", env!("CARGO_PKG_VERSION")),
            &p.relative_paths,
            &self.out_dir,
            &self.customize,
        )
    }

    /// Similar to `run`, but prints the message to stderr and exits the process on error.
    pub fn run_from_script(&self) {
        if let Err(e) = self.run() {
            eprintln!("pure rust codegen failed: {}", e);
            process::exit(1);
        }
    }
}

#[derive(Clone)]
pub(crate) struct FileDescriptorPair {
    parsed: model::FileDescriptor,
    descriptor: protobuf::descriptor::FileDescriptorProto,
}

#[derive(Debug, thiserror::Error)]
#[error("error in `{file}`: {error}")]
struct WithFileError {
    file: String,
    #[source]
    error: anyhow::Error,
}

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("protobuf path `{0}` is not found in import path {1}")]
    FileNotFoundInImportPath(String, String),
    #[error("file `{0}` must reside in include path {1}")]
    FileMustResideInImportPath(String, String),
}

struct Run<'a> {
    parsed_files: LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    includes: &'a [PathBuf],
}

impl<'a> Run<'a> {
    fn get_file_and_all_deps_already_parsed(
        &self,
        protobuf_path: &ProtoPath,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        if let Some(_) = result.get(protobuf_path) {
            return;
        }

        let pair = self
            .parsed_files
            .get(protobuf_path)
            .expect("must be already parsed");
        result.insert(protobuf_path.to_proto_path_buf(), pair.clone());

        self.get_all_deps_already_parsed(&pair.parsed, result);
    }

    fn get_all_deps_already_parsed(
        &self,
        parsed: &model::FileDescriptor,
        result: &mut LinkedHashMap<ProtoPathBuf, FileDescriptorPair>,
    ) {
        for import in &parsed.imports {
            self.get_file_and_all_deps_already_parsed(&import.path, result);
        }
    }

    fn add_file(&mut self, protobuf_path: &ProtoPath, fs_path: &Path) -> anyhow::Result<()> {
        if let Some(_) = self.parsed_files.get(protobuf_path) {
            return Ok(());
        }

        let content = fs::read_to_string(fs_path)
            .map_err(|e| amend_io_error(e, format!("failed to read {:?}", fs_path)))?;

        self.add_file_content(protobuf_path, fs_path, &content)
    }

    fn add_file_content(
        &mut self,
        protobuf_path: &ProtoPath,
        fs_path: &Path,
        content: &str,
    ) -> anyhow::Result<()> {
        let parsed = model::FileDescriptor::parse(content).map_err(|e| WithFileError {
            file: format!("{}", fs_path.display()),
            error: e.into(),
        })?;

        for import in &parsed.imports {
            self.add_imported_file(&import.path)?;
        }

        let mut this_file_deps = LinkedHashMap::new();
        self.get_all_deps_already_parsed(&parsed, &mut this_file_deps);

        let this_file_deps: Vec<_> = this_file_deps.into_iter().map(|(_, v)| v).collect();

        let descriptor = convert::file_descriptor(protobuf_path, &parsed, &this_file_deps)
            .map_err(|e| WithFileError {
                file: format!("{}", fs_path.display()),
                error: e.into(),
            })?;

        self.parsed_files.insert(
            protobuf_path.to_proto_path_buf(),
            FileDescriptorPair { parsed, descriptor },
        );

        Ok(())
    }

    fn add_imported_file(&mut self, protobuf_path: &ProtoPath) -> anyhow::Result<()> {
        for include_dir in self.includes {
            if let Some(_) = self.parsed_files.get(protobuf_path) {
                return Ok(());
            }

            let fs_path = include_dir.join(protobuf_path.to_path());
            if fs_path.exists() {
                return self.add_file(protobuf_path, &fs_path);
            }
        }

        let embedded = match protobuf_path.to_str() {
            "rustproto.proto" => Some(RUSTPROTO_PROTO),
            "google/protobuf/any.proto" => Some(ANY_PROTO),
            "google/protobuf/api.proto" => Some(API_PROTO),
            "google/protobuf/descriptor.proto" => Some(DESCRIPTOR_PROTO),
            "google/protobuf/duration.proto" => Some(DURATION_PROTO),
            "google/protobuf/empty.proto" => Some(EMPTY_PROTO),
            "google/protobuf/field_mask.proto" => Some(FIELD_MASK_PROTO),
            "google/protobuf/source_context.proto" => Some(SOURCE_CONTEXT_PROTO),
            "google/protobuf/struct.proto" => Some(STRUCT_PROTO),
            "google/protobuf/timestamp.proto" => Some(TIMESTAMP_PROTO),
            "google/protobuf/type.proto" => Some(TYPE_PROTO),
            "google/protobuf/wrappers.proto" => Some(WRAPPERS_PROTO),
            _ => None,
        };

        match embedded {
            Some(content) => self.add_file_content(protobuf_path, protobuf_path.to_path(), content),
            None => Err(Error::FileNotFoundInImportPath(
                protobuf_path.to_string(),
                format!("{:?}", self.includes),
            )
            .into()),
        }
    }

    fn strip_prefix<'b>(path: &'b Path, prefix: &Path) -> anyhow::Result<ProtoPathBuf> {
        // special handling of `.` to allow successful `strip_prefix("foo.proto", ".")
        if prefix == Path::new(".") && path.is_relative() {
            ProtoPathBuf::from_path(path)
        } else {
            ProtoPathBuf::from_path(path.strip_prefix(prefix)?)
        }
    }

    fn add_fs_file(&mut self, fs_path: &Path) -> anyhow::Result<ProtoPathBuf> {
        let relative_path = self
            .includes
            .iter()
            .filter_map(|include_dir| Self::strip_prefix(fs_path, include_dir).ok())
            .next();

        match relative_path {
            Some(relative_path) => {
                self.add_file(&relative_path, fs_path)?;
                Ok(relative_path)
            }
            None => Err(Error::FileMustResideInImportPath(
                fs_path.display().to_string(),
                format!("{:?}", self.includes),
            )
            .into()),
        }
    }
}

/// Result of parsing `.proto` files.
#[doc(hidden)]
pub struct ParsedAndTypechecked {
    /// One entry for each input `.proto` file.
    pub relative_paths: Vec<ProtoPathBuf>,
    /// All parsed `.proto` files including dependencies of input files.
    pub file_descriptors: Vec<protobuf::descriptor::FileDescriptorProto>,
}

#[doc(hidden)]
pub fn parse_and_typecheck(
    includes: &[PathBuf],
    input: &[PathBuf],
) -> anyhow::Result<ParsedAndTypechecked> {
    let mut run = Run {
        parsed_files: LinkedHashMap::new(),
        includes,
    };

    let mut relative_paths = Vec::new();

    for input in input {
        relative_paths.push(run.add_fs_file(input)?);
    }

    let file_descriptors: Vec<_> = run
        .parsed_files
        .into_iter()
        .map(|(_, v)| v.descriptor)
        .collect();

    Ok(ParsedAndTypechecked {
        relative_paths,
        file_descriptors,
    })
}

/// Parse imports from a `.proto` file.
///
/// The result is [`FileDescriptorProto`] object with only `*dependency` fields filled.
#[doc(hidden)]
pub fn parse_dependencies(content: &str) -> Result<FileDescriptorProto, ParserErrorWithLocation> {
    parse_dependencies::parse_dependencies(content)
}

// TODO: these include don't work when publishing to crates
const RUSTPROTO_PROTO: &str = include_str!("proto/rustproto.proto");
const ANY_PROTO: &str = include_str!("proto/google/protobuf/any.proto");
const API_PROTO: &str = include_str!("proto/google/protobuf/api.proto");
const DESCRIPTOR_PROTO: &str = include_str!("proto/google/protobuf/descriptor.proto");
const DURATION_PROTO: &str = include_str!("proto/google/protobuf/duration.proto");
const EMPTY_PROTO: &str = include_str!("proto/google/protobuf/empty.proto");
const FIELD_MASK_PROTO: &str = include_str!("proto/google/protobuf/field_mask.proto");
const SOURCE_CONTEXT_PROTO: &str = include_str!("proto/google/protobuf/source_context.proto");
const STRUCT_PROTO: &str = include_str!("proto/google/protobuf/struct.proto");
const TIMESTAMP_PROTO: &str = include_str!("proto/google/protobuf/timestamp.proto");
const TYPE_PROTO: &str = include_str!("proto/google/protobuf/type.proto");
const WRAPPERS_PROTO: &str = include_str!("proto/google/protobuf/wrappers.proto");
<|endoftext|>


<commit_before>use std::thread;
use std::cmp::min;
use threadpool::ThreadPool;
use std::collections::HashMap;
use super::actor::{Actor, ActorRef, ActorVecRef};
use router::{Router, RoutingMessage};

// Max amount of messages to process per actor
// during an activation
const THROUGHPUT: usize = 100;

/// Manages a pool of actors,
/// dispatching them to threads when they have queued messages.
pub fn dispatcher<A>(actors: ActorVecRef<A>, n_threads: usize)
    where A: Actor + 'static
{
    thread::spawn(move || {
        // TODO should this just be another futures cpupool?
        let pool = ThreadPool::new(n_threads);
        let router = Router::<A::M>::new();

        // lookup actors by id
        let mut lookup = HashMap::<usize, ActorRef<A>>::new();
        {
            let actors = actors.clone();
            let actors = actors.read().unwrap();
            for actor in actors.iter() {
                let actor_r = actor.read().unwrap();
                lookup.insert(actor_r.id(), actor.clone());
            }
        }

        loop {
            // give the actors the messages from the router
            let mut messages = router.inbox.write().unwrap();
            for msg in messages.drain(..) {
                match msg {
                    RoutingMessage::Message { sender, recipient, message } => {
                        match lookup.get(&recipient) {
                            Some(actor) => {
                                let actor_r = actor.read().unwrap();
                                let mut inbox = actor_r.inbox().write().unwrap();
                                inbox.push(message);
                            }
                            _ => println!("foo"),
                        }
                    }
                    _ => println!("foo"),
                }
            }
            let actors = actors.clone();
            let actors = actors.read().unwrap();
            for actor in actors.iter() {
                let actor_r = actor.read().unwrap();
                let n_queued = {
                    let inbox = actor_r.inbox().read().unwrap();
                    inbox.len()
                };
                if n_queued > 0 {
                    let n_messages = min(n_queued, THROUGHPUT);
                    let mut inbox = actor_r.inbox().write().unwrap();
                    let chunk: Vec<A::M> = inbox.drain(0..n_messages).collect();
                    let actor = actor.clone();
                    pool.execute(move || {
                        for msg in chunk {
                            let actor = actor.write().unwrap();
                            let _ = actor.handle_msg(msg);
                        }
                    });
                }
            }
        }
    });
}
<commit_msg>minor simplification<commit_after>use std::thread;
use std::cmp::min;
use threadpool::ThreadPool;
use std::collections::HashMap;
use super::actor::{Actor, ActorRef, ActorVecRef};
use router::{Router, RoutingMessage};

// Max amount of messages to process per actor
// during an activation
const THROUGHPUT: usize = 100;

/// Manages a pool of actors,
/// dispatching them to threads when they have queued messages.
pub fn dispatcher<A>(actors: ActorVecRef<A>, n_threads: usize)
    where A: Actor + 'static
{
    thread::spawn(move || {
        // TODO should this just be another futures cpupool?
        let pool = ThreadPool::new(n_threads);
        let router = Router::<A::M>::new();

        // lookup actors by id
        let mut lookup = HashMap::<usize, ActorRef<A>>::new();
        {
            let actors = actors.clone();
            let actors = actors.read().unwrap();
            for actor in actors.iter() {
                let actor_r = actor.read().unwrap();
                lookup.insert(actor_r.id(), actor.clone());
            }
        }

        loop {
            // give the actors the messages from the router
            let mut messages = router.inbox.write().unwrap();
            for msg in messages.drain(..) {
                match msg {
                    RoutingMessage::Message { sender, recipient, message } => {
                        match lookup.get(&recipient) {
                            Some(actor) => {
                                let actor_r = actor.read().unwrap();
                                let mut inbox = actor_r.inbox().write().unwrap();
                                inbox.push(message);
                            }
                            _ => (),
                        }
                    }
                    _ => (),
                }
            }
            let actors = actors.clone();
            let actors = actors.read().unwrap();
            for actor in actors.iter() {
                let actor_r = actor.read().unwrap();
                let n_queued = {
                    let inbox = actor_r.inbox().read().unwrap();
                    inbox.len()
                };
                if n_queued > 0 {
                    let n_messages = min(n_queued, THROUGHPUT);
                    let mut inbox = actor_r.inbox().write().unwrap();
                    let chunk: Vec<A::M> = inbox.drain(0..n_messages).collect();
                    let actor = actor.clone();
                    pool.execute(move || {
                        for msg in chunk {
                            let actor = actor.write().unwrap();
                            let _ = actor.handle_msg(msg);
                        }
                    });
                }
            }
        }
    });
}
<|endoftext|>


<commit_before>use crate::common::util::*;
use std::io::{Seek, SeekFrom, Write};

static TFILE1: &'static str = "truncate_test_1";
static TFILE2: &'static str = "truncate_test_2";

#[test]
fn test_increase_file_size() {
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE1);
    ucmd.args(&["-s", "+5K", TFILE1]).succeeds();

    file.seek(SeekFrom::End(0)).unwrap();
    assert!(file.seek(SeekFrom::Current(0)).unwrap() == 5 * 1024);
}

#[test]
fn test_increase_file_size_kb() {
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE1);
    ucmd.args(&["-s", "+5KB", TFILE1]).succeeds();

    file.seek(SeekFrom::End(0)).unwrap();
    assert!(file.seek(SeekFrom::Current(0)).unwrap() == 5 * 1000);
}

#[test]
fn test_reference() {
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    let mut file = at.make_file(TFILE2);

    scene.ucmd().arg("-s").arg("+5KB").arg(TFILE1).run();

    scene
        .ucmd()
        .arg("--reference")
        .arg(TFILE1)
        .arg(TFILE2)
        .run();

    file.seek(SeekFrom::End(0)).unwrap();
    assert!(file.seek(SeekFrom::Current(0)).unwrap() == 5 * 1000);
}

#[test]
fn test_decrease_file_size() {
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size=-4", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    assert!(file.seek(SeekFrom::Current(0)).unwrap() == 6);
}

#[test]
fn test_space_in_size() {
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", " 4", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    assert!(file.seek(SeekFrom::Current(0)).unwrap() == 4);
}

#[test]
fn test_failed() {
    new_ucmd!().fails();
}

#[test]
fn test_failed_2() {
    let (_at, mut ucmd) = at_and_ucmd!();
    ucmd.args(&[TFILE1]).fails();
}

#[test]
fn test_failed_incorrect_arg() {
    let (_at, mut ucmd) = at_and_ucmd!();
    ucmd.args(&["-s", "+5A", TFILE1]).fails();
}

#[test]
fn test_at_most_shrinks() {
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", "<4", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    assert!(file.seek(SeekFrom::Current(0)).unwrap() == 4);
}

#[test]
fn test_at_most_no_change() {
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", "<40", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    assert!(file.seek(SeekFrom::Current(0)).unwrap() == 10);
}

#[test]
fn test_at_least_grows() {
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", ">15", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    assert!(file.seek(SeekFrom::Current(0)).unwrap() == 15);
}

#[test]
fn test_at_least_no_change() {
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", ">4", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    assert!(file.seek(SeekFrom::Current(0)).unwrap() == 10);
}

#[test]
fn test_round_down() {
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", "/4", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    assert!(file.seek(SeekFrom::Current(0)).unwrap() == 8);
}

#[test]
fn test_round_up() {
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", "*4", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    assert!(file.seek(SeekFrom::Current(0)).unwrap() == 12);
}
<commit_msg>improve assert error messages<commit_after>use crate::common::util::*;
use std::io::{Seek, SeekFrom, Write};

static TFILE1: &'static str = "truncate_test_1";
static TFILE2: &'static str = "truncate_test_2";

#[test]
fn test_increase_file_size() {
    let expected = 5 * 1024;
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE1);
    ucmd.args(&["-s", "+5K", TFILE1]).succeeds();

    file.seek(SeekFrom::End(0)).unwrap();
    let actual = file.seek(SeekFrom::Current(0)).unwrap();
    assert!(expected == actual, "expected '{}' got '{}'", expected, actual);
}

#[test]
fn test_increase_file_size_kb() {
    let expected = 5 * 1000;
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE1);
    ucmd.args(&["-s", "+5KB", TFILE1]).succeeds();

    file.seek(SeekFrom::End(0)).unwrap();
    let actual = file.seek(SeekFrom::Current(0)).unwrap();
    assert!(expected == actual, "expected '{}' got '{}'", expected, actual);
}

#[test]
fn test_reference() {
    let expected = 5 * 1000;
    let scene = TestScenario::new(util_name!());
    let at = &scene.fixtures;
    let mut file = at.make_file(TFILE2);

    scene.ucmd().arg("-s").arg("+5KB").arg(TFILE1).run();

    scene
        .ucmd()
        .arg("--reference")
        .arg(TFILE1)
        .arg(TFILE2)
        .run();

    file.seek(SeekFrom::End(0)).unwrap();
    let actual = file.seek(SeekFrom::Current(0)).unwrap();
    assert!(expected == actual, "expected '{}' got '{}'", expected, actual);
}

#[test]
fn test_decrease_file_size() {
    let expected = 6;
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size=-4", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    let actual = file.seek(SeekFrom::Current(0)).unwrap();
    assert!(expected == actual, "expected '{}' got '{}'", expected, actual);
}

#[test]
fn test_space_in_size() {
    let expected = 4;
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", " 4", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    let actual = file.seek(SeekFrom::Current(0)).unwrap();
    assert!(expected == actual, "expected '{}' got '{}'", expected, actual);
}

#[test]
fn test_failed() {
    new_ucmd!().fails();
}

#[test]
fn test_failed_2() {
    let (_at, mut ucmd) = at_and_ucmd!();
    ucmd.args(&[TFILE1]).fails();
}

#[test]
fn test_failed_incorrect_arg() {
    let (_at, mut ucmd) = at_and_ucmd!();
    ucmd.args(&["-s", "+5A", TFILE1]).fails();
}

#[test]
fn test_at_most_shrinks() {
    let expected = 4;
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", "<4", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    let actual = file.seek(SeekFrom::Current(0)).unwrap();
    assert!(expected == actual, "expected '{}' got '{}'", expected, actual);
}

#[test]
fn test_at_most_no_change() {
    let expected = 10;
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", "<40", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    let actual = file.seek(SeekFrom::Current(0)).unwrap();
    assert!(expected == actual, "expected '{}' got '{}'", expected, actual);
}

#[test]
fn test_at_least_grows() {
    let expected = 15;
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", ">15", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    let actual = file.seek(SeekFrom::Current(0)).unwrap();
    assert!(expected == actual, "expected '{}' got '{}'", expected, actual);
}

#[test]
fn test_at_least_no_change() {
    let expected = 10;
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", ">4", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    let actual = file.seek(SeekFrom::Current(0)).unwrap();
    assert!(expected == actual, "expected '{}' got '{}'", expected, actual);
}

#[test]
fn test_round_down() {
    let expected = 8;
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", "/4", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    let actual = file.seek(SeekFrom::Current(0)).unwrap();
    assert!(expected == actual, "expected '{}' got '{}'", expected, actual);
}

#[test]
fn test_round_up() {
    let expected = 12;
    let (at, mut ucmd) = at_and_ucmd!();
    let mut file = at.make_file(TFILE2);
    file.write_all(b"1234567890").unwrap();
    ucmd.args(&["--size", "*4", TFILE2]).succeeds();
    file.seek(SeekFrom::End(0)).unwrap();
    let actual = file.seek(SeekFrom::Current(0)).unwrap();
    assert!(expected == actual, "expected '{}' got '{}'", expected, actual);
}
<|endoftext|>


<commit_before>//@ignore-windows: Concurrency on Windows is not supported yet.
//@compile-flags: -Zmiri-ignore-leaks -Zmiri-preemption-rate=0

// Tests showing weak memory behaviours are exhibited. All tests
// return true when the desired behaviour is seen.
// This is scheduler and pseudo-RNG dependent, so each test is
// run multiple times until one try returns true.
// Spurious failure is possible, if you are really unlucky with
// the RNG and always read the latest value from the store buffer.

use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::*;
use std::thread::spawn;

#[derive(Copy, Clone)]
struct EvilSend<T>(pub T);

unsafe impl<T> Send for EvilSend<T> {}
unsafe impl<T> Sync for EvilSend<T> {}

// We can't create static items because we need to run each test
// multiple times
fn static_atomic(val: usize) -> &'static AtomicUsize {
    let ret = Box::leak(Box::new(AtomicUsize::new(val)));
    ret
}

// Spins until it reads the given value
fn reads_value(loc: &AtomicUsize, val: usize) -> usize {
    while loc.load(Relaxed) != val {
        std::hint::spin_loop();
    }
    val
}

fn relaxed() -> bool {
    let x = static_atomic(0);
    let j1 = spawn(move || {
        x.store(1, Relaxed);
        x.store(2, Relaxed);
    });

    let j2 = spawn(move || x.load(Relaxed));

    j1.join().unwrap();
    let r2 = j2.join().unwrap();

    r2 == 1
}

// https://www.doc.ic.ac.uk/~afd/homepages/papers/pdfs/2017/POPL.pdf Figure 8
fn seq_cst() -> bool {
    let x = static_atomic(0);

    let j1 = spawn(move || {
        x.store(1, Relaxed);
    });

    let j2 = spawn(move || {
        x.store(2, SeqCst);
        x.store(3, SeqCst);
    });

    let j3 = spawn(move || x.load(SeqCst));

    j1.join().unwrap();
    j2.join().unwrap();
    let r3 = j3.join().unwrap();

    r3 == 1
}

fn initialization_write() -> bool {
    let x = static_atomic(11);
    assert_eq!(x.load(Relaxed), 11); // work around https://github.com/rust-lang/miri/issues/2164

    let wait = static_atomic(0);

    let j1 = spawn(move || {
        x.store(22, Relaxed);
        // Relaxed is intentional. We want to test if the thread 2 reads the initialisation write
        // after a relaxed write
        wait.store(1, Relaxed);
    });

    let j2 = spawn(move || {
        reads_value(wait, 1);
        x.load(Relaxed)
    });

    j1.join().unwrap();
    let r2 = j2.join().unwrap();

    r2 == 11
}

/// Asserts that the function returns true at least once in 100 runs
fn assert_once(f: fn() -> bool) {
    assert!(std::iter::repeat_with(|| f()).take(100).any(|x| x));
}

pub fn main() {
    assert_once(relaxed);
    assert_once(seq_cst);
    assert_once(initialization_write);
}
<commit_msg>test that we can see this weak behavior<commit_after>//@ignore-windows: Concurrency on Windows is not supported yet.
//@compile-flags: -Zmiri-ignore-leaks -Zmiri-preemption-rate=0

// Tests showing weak memory behaviours are exhibited. All tests
// return true when the desired behaviour is seen.
// This is scheduler and pseudo-RNG dependent, so each test is
// run multiple times until one try returns true.
// Spurious failure is possible, if you are really unlucky with
// the RNG and always read the latest value from the store buffer.

use std::sync::atomic::Ordering::*;
use std::sync::atomic::{fence, AtomicUsize};
use std::thread::spawn;

#[derive(Copy, Clone)]
struct EvilSend<T>(pub T);

unsafe impl<T> Send for EvilSend<T> {}
unsafe impl<T> Sync for EvilSend<T> {}

// We can't create static items because we need to run each test
// multiple times
fn static_atomic(val: usize) -> &'static AtomicUsize {
    let ret = Box::leak(Box::new(AtomicUsize::new(val)));
    ret
}

// Spins until it reads the given value
fn reads_value(loc: &AtomicUsize, val: usize) -> usize {
    while loc.load(Relaxed) != val {
        std::hint::spin_loop();
    }
    val
}

fn relaxed() -> bool {
    let x = static_atomic(0);
    let j1 = spawn(move || {
        x.store(1, Relaxed);
        x.store(2, Relaxed);
    });

    let j2 = spawn(move || x.load(Relaxed));

    j1.join().unwrap();
    let r2 = j2.join().unwrap();

    r2 == 1
}

// https://www.doc.ic.ac.uk/~afd/homepages/papers/pdfs/2017/POPL.pdf Figure 8
fn seq_cst() -> bool {
    let x = static_atomic(0);

    let j1 = spawn(move || {
        x.store(1, Relaxed);
    });

    let j2 = spawn(move || {
        x.store(2, SeqCst);
        x.store(3, SeqCst);
    });

    let j3 = spawn(move || x.load(SeqCst));

    j1.join().unwrap();
    j2.join().unwrap();
    let r3 = j3.join().unwrap();

    r3 == 1
}

fn initialization_write(add_fence: bool) -> bool {
    let x = static_atomic(11);
    assert_eq!(x.load(Relaxed), 11); // work around https://github.com/rust-lang/miri/issues/2164

    let wait = static_atomic(0);

    let j1 = spawn(move || {
        x.store(22, Relaxed);
        // Relaxed is intentional. We want to test if the thread 2 reads the initialisation write
        // after a relaxed write
        wait.store(1, Relaxed);
    });

    let j2 = spawn(move || {
        reads_value(wait, 1);
        if add_fence {
            fence(AcqRel);
        }
        x.load(Relaxed)
    });

    j1.join().unwrap();
    let r2 = j2.join().unwrap();

    r2 == 11
}

fn faa_replaced_by_load() -> bool {
    // Example from https://github.com/llvm/llvm-project/issues/56450#issuecomment-1183695905
    #[no_mangle]
    pub fn rdmw(storing: &AtomicUsize, sync: &AtomicUsize, loading: &AtomicUsize) -> usize {
        storing.store(1, Relaxed);
        fence(Release);
        // sync.fetch_add(0, Relaxed);
        sync.load(Relaxed);
        fence(Acquire);
        loading.load(Relaxed)
    }

    let x = static_atomic(0);
    assert_eq!(x.load(Relaxed), 0); // work around https://github.com/rust-lang/miri/issues/2164
    let y = static_atomic(0);
    assert_eq!(y.load(Relaxed), 0); // work around https://github.com/rust-lang/miri/issues/2164
    let z = static_atomic(0);
    assert_eq!(z.load(Relaxed), 0); // work around https://github.com/rust-lang/miri/issues/2164

    // Since each thread is so short, we need to make sure that they truely run at the same time
    // Otherwise t1 will finish before t2 even starts
    let go = static_atomic(0);

    let t1 = spawn(move || {
        while go.load(Relaxed) == 0 {}
        rdmw(y, x, z)
    });

    let t2 = spawn(move || {
        while go.load(Relaxed) == 0 {}
        rdmw(z, x, y)
    });

    go.store(1, Relaxed);

    let a = t1.join().unwrap();
    let b = t2.join().unwrap();
    (a, b) == (0, 0)
}

/// Asserts that the function returns true at least once in 100 runs
fn assert_once(f: fn() -> bool) {
    assert!(std::iter::repeat_with(|| f()).take(100).any(|x| x));
}

pub fn main() {
    assert_once(relaxed);
    assert_once(seq_cst);
    assert_once(|| initialization_write(false));
    assert_once(|| initialization_write(true));
    assert_once(faa_replaced_by_load);
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate gcc;
extern crate build_helper;

use std::env;
use std::fs;
use std::path::PathBuf;
use std::process::Command;

use build_helper::run;

fn main() {
    println!("cargo:rustc-cfg=cargobuild");

    let target = env::var("TARGET").unwrap();
    let host = env::var("HOST").unwrap();
    if !target.contains("apple") && !target.contains("msvc") && !target.contains("emscripten"){
        build_libbacktrace(&host, &target);
    }

    if target.contains("unknown-linux") {
        if target.contains("musl") && (target.contains("x86_64") || target.contains("i686")) {
            println!("cargo:rustc-link-lib=static=unwind");
        } else {
            println!("cargo:rustc-link-lib=dl");
            println!("cargo:rustc-link-lib=rt");
            println!("cargo:rustc-link-lib=pthread");
            println!("cargo:rustc-link-lib=gcc_s");
        }
    } else if target.contains("android") {
        println!("cargo:rustc-link-lib=dl");
        println!("cargo:rustc-link-lib=log");
        println!("cargo:rustc-link-lib=gcc");
    } else if target.contains("freebsd") {
        println!("cargo:rustc-link-lib=execinfo");
        println!("cargo:rustc-link-lib=pthread");
        println!("cargo:rustc-link-lib=gcc_s");
    } else if target.contains("dragonfly") || target.contains("bitrig") ||
              target.contains("netbsd") || target.contains("openbsd") {
        println!("cargo:rustc-link-lib=pthread");

        if target.contains("rumprun") {
            println!("cargo:rustc-link-lib=unwind");
        } else if target.contains("netbsd") || target.contains("openbsd") {
            println!("cargo:rustc-link-lib=gcc");
        } else if target.contains("bitrig") {
            println!("cargo:rustc-link-lib=c++abi");
        } else if target.contains("dragonfly") {
            println!("cargo:rustc-link-lib=gcc_pic");
        }
    } else if target.contains("apple-darwin") {
        println!("cargo:rustc-link-lib=System");
    } else if target.contains("apple-ios") {
        println!("cargo:rustc-link-lib=System");
        println!("cargo:rustc-link-lib=objc");
        println!("cargo:rustc-link-lib=framework=Security");
        println!("cargo:rustc-link-lib=framework=Foundation");
    } else if target.contains("windows") {
        if target.contains("windows-gnu") {
            println!("cargo:rustc-link-lib=gcc_eh");
        }
        println!("cargo:rustc-link-lib=advapi32");
        println!("cargo:rustc-link-lib=ws2_32");
        println!("cargo:rustc-link-lib=userenv");
        println!("cargo:rustc-link-lib=shell32");
    }
}

fn build_libbacktrace(host: &str, target: &str) {
    let src_dir = env::current_dir().unwrap().join("../libbacktrace");
    let build_dir = PathBuf::from(env::var_os("OUT_DIR").unwrap());

    println!("cargo:rustc-link-lib=static=backtrace");
    println!("cargo:rustc-link-search=native={}/.libs", build_dir.display());

    if fs::metadata(&build_dir.join(".libs/libbacktrace.a")).is_ok() {
        return
    }

    let compiler = gcc::Config::new().get_compiler();
    let ar = build_helper::cc2ar(compiler.path(), target);
    let cflags = compiler.args().iter().map(|s| s.to_str().unwrap())
                         .collect::<Vec<_>>().join(" ");
    run(Command::new("sh")
                .current_dir(&build_dir)
                .arg(src_dir.join("configure").to_str().unwrap()
                            .replace("C:\\", "/c/")
                            .replace("\\", "/"))
                .arg("--with-pic")
                .arg("--disable-multilib")
                .arg("--disable-shared")
                .arg("--disable-host-shared")
                .arg(format!("--host={}", build_helper::gnu_target(target)))
                .arg(format!("--build={}", build_helper::gnu_target(host)))
                .env("CC", compiler.path())
                .env("AR", &ar)
                .env("RANLIB", format!("{} s", ar.display()))
                .env("CFLAGS", cflags));
    run(Command::new("make")
                .current_dir(&build_dir)
                .arg(format!("INCDIR={}", src_dir.display()))
                .arg("-j").arg(env::var("NUM_JOBS").unwrap()));
}
<commit_msg>Auto merge of #32132 - arcnmx:cargobuild-std-target, r=alexcrichton<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

extern crate gcc;
extern crate build_helper;

use std::env;
use std::fs;
use std::path::PathBuf;
use std::process::Command;

use build_helper::run;

fn main() {
    println!("cargo:rustc-cfg=cargobuild");

    let target = env::var("TARGET").unwrap();
    let host = env::var("HOST").unwrap();
    if !target.contains("apple") && !target.contains("msvc") && !target.contains("emscripten"){
        build_libbacktrace(&host, &target);
    }

    if target.contains("linux") {
        if target.contains("musl") && (target.contains("x86_64") || target.contains("i686")) {
            println!("cargo:rustc-link-lib=static=unwind");
        } else if target.contains("android") {
            println!("cargo:rustc-link-lib=dl");
            println!("cargo:rustc-link-lib=log");
            println!("cargo:rustc-link-lib=gcc");
        } else {
            println!("cargo:rustc-link-lib=dl");
            println!("cargo:rustc-link-lib=rt");
            println!("cargo:rustc-link-lib=pthread");
            println!("cargo:rustc-link-lib=gcc_s");
        }
    } else if target.contains("freebsd") {
        println!("cargo:rustc-link-lib=execinfo");
        println!("cargo:rustc-link-lib=pthread");
        println!("cargo:rustc-link-lib=gcc_s");
    } else if target.contains("dragonfly") || target.contains("bitrig") ||
              target.contains("netbsd") || target.contains("openbsd") {
        println!("cargo:rustc-link-lib=pthread");

        if target.contains("rumprun") {
            println!("cargo:rustc-link-lib=unwind");
        } else if target.contains("netbsd") || target.contains("openbsd") {
            println!("cargo:rustc-link-lib=gcc");
        } else if target.contains("bitrig") {
            println!("cargo:rustc-link-lib=c++abi");
        } else if target.contains("dragonfly") {
            println!("cargo:rustc-link-lib=gcc_pic");
        }
    } else if target.contains("apple-darwin") {
        println!("cargo:rustc-link-lib=System");
    } else if target.contains("apple-ios") {
        println!("cargo:rustc-link-lib=System");
        println!("cargo:rustc-link-lib=objc");
        println!("cargo:rustc-link-lib=framework=Security");
        println!("cargo:rustc-link-lib=framework=Foundation");
    } else if target.contains("windows") {
        if target.contains("windows-gnu") {
            println!("cargo:rustc-link-lib=gcc_eh");
        }
        println!("cargo:rustc-link-lib=advapi32");
        println!("cargo:rustc-link-lib=ws2_32");
        println!("cargo:rustc-link-lib=userenv");
        println!("cargo:rustc-link-lib=shell32");
    }
}

fn build_libbacktrace(host: &str, target: &str) {
    let src_dir = env::current_dir().unwrap().join("../libbacktrace");
    let build_dir = PathBuf::from(env::var_os("OUT_DIR").unwrap());

    println!("cargo:rustc-link-lib=static=backtrace");
    println!("cargo:rustc-link-search=native={}/.libs", build_dir.display());

    if fs::metadata(&build_dir.join(".libs/libbacktrace.a")).is_ok() {
        return
    }

    let compiler = gcc::Config::new().get_compiler();
    let ar = build_helper::cc2ar(compiler.path(), target);
    let cflags = compiler.args().iter().map(|s| s.to_str().unwrap())
                         .collect::<Vec<_>>().join(" ");
    run(Command::new("sh")
                .current_dir(&build_dir)
                .arg(src_dir.join("configure").to_str().unwrap()
                            .replace("C:\\", "/c/")
                            .replace("\\", "/"))
                .arg("--with-pic")
                .arg("--disable-multilib")
                .arg("--disable-shared")
                .arg("--disable-host-shared")
                .arg(format!("--host={}", build_helper::gnu_target(target)))
                .arg(format!("--build={}", build_helper::gnu_target(host)))
                .env("CC", compiler.path())
                .env("AR", &ar)
                .env("RANLIB", format!("{} s", ar.display()))
                .env("CFLAGS", cflags));
    run(Command::new("make")
                .current_dir(&build_dir)
                .arg(format!("INCDIR={}", src_dir.display()))
                .arg("-j").arg(env::var("NUM_JOBS").unwrap()));
}
<|endoftext|>