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

use tidb_query_codegen::AggrFunction;
use tidb_query_datatype::builder::FieldTypeBuilder;
use tidb_query_datatype::{EvalType, FieldTypeFlag, FieldTypeTp};
use tipb::{Expr, ExprType, FieldType};

use super::summable::Summable;
use super::*;
use tidb_query_common::Result;
use tidb_query_datatype::codec::data_type::*;
use tidb_query_datatype::expr::EvalContext;
use tidb_query_expr::RpnExpression;

/// The parser for AVG aggregate function.
pub struct AggrFnDefinitionParserAvg;

impl super::AggrDefinitionParser for AggrFnDefinitionParserAvg {
    fn check_supported(&self, aggr_def: &Expr) -> Result<()> {
        assert_eq!(aggr_def.get_tp(), ExprType::Avg);
        super::util::check_aggr_exp_supported_one_child(aggr_def)
    }

    #[inline]
    fn parse_rpn(
        &self,
        mut root_expr: Expr,
        mut exp: RpnExpression,
        _ctx: &mut EvalContext,
        src_schema: &[FieldType],
        out_schema: &mut Vec<FieldType>,
        out_exp: &mut Vec<RpnExpression>,
    ) -> Result<Box<dyn AggrFunction>> {
        use std::convert::TryFrom;
        use tidb_query_datatype::FieldTypeAccessor;

        assert_eq!(root_expr.get_tp(), ExprType::Avg);

        let col_sum_ft = root_expr.take_field_type();
        let col_sum_et = box_try!(EvalType::try_from(col_sum_ft.as_accessor().tp()));

        // Rewrite expression to insert CAST() if needed.
        super::util::rewrite_exp_for_sum_avg(src_schema, &mut exp).unwrap();

        let rewritten_eval_type =
            EvalType::try_from(exp.ret_field_type(src_schema).as_accessor().tp()).unwrap();
        if col_sum_et != rewritten_eval_type {
            return Err(other_err!(
                "Unexpected return field type {}",
                col_sum_ft.as_accessor().tp()
            ));
        }

        // AVG outputs two columns.
        out_schema.push(
            FieldTypeBuilder::new()
                .tp(FieldTypeTp::LongLong)
                .flag(FieldTypeFlag::UNSIGNED)
                .build(),
        );
        out_schema.push(col_sum_ft);
        out_exp.push(exp);

        Ok(match rewritten_eval_type {
            EvalType::Decimal => Box::new(AggrFnAvg::<Decimal>::new()),
            EvalType::Real => Box::new(AggrFnAvg::<Real>::new()),
            _ => unreachable!(),
        })
    }
}

/// The AVG aggregate function.
///
/// Note that there are `AVG(Decimal) -> (Int, Decimal)` and `AVG(Double) -> (Int, Double)`.
#[derive(Debug, AggrFunction)]
#[aggr_function(state = AggrFnStateAvg::<T>::new())]
pub struct AggrFnAvg<T>
where
    T: Summable,
    VectorValue: VectorValueExt<T>,
{
    _phantom: std::marker::PhantomData<T>,
}

impl<T> AggrFnAvg<T>
where
    T: Summable,
    VectorValue: VectorValueExt<T>,
{
    pub fn new() -> Self {
        Self {
            _phantom: std::marker::PhantomData,
        }
    }
}

/// The state of the AVG aggregate function.
#[derive(Debug)]
pub struct AggrFnStateAvg<T>
where
    T: Summable,
    VectorValue: VectorValueExt<T>,
{
    sum: T,
    count: usize,
}

impl<T> AggrFnStateAvg<T>
where
    T: Summable,
    VectorValue: VectorValueExt<T>,
{
    pub fn new() -> Self {
        Self {
            sum: T::zero(),
            count: 0,
        }
    }

    #[inline]
    fn update_concrete<'a, TT>(&mut self, ctx: &mut EvalContext, value: Option<TT>) -> Result<()>
    where
        TT: EvaluableRef<'a, EvaluableType = T>,
    {
        match value {
            None => Ok(()),
            Some(value) => {
                self.sum.add_assign(ctx, &value.to_owned_value())?;
                self.count += 1;
                Ok(())
            }
        }
    }
}

impl<T> super::ConcreteAggrFunctionState for AggrFnStateAvg<T>
where
    T: Summable,
    VectorValue: VectorValueExt<T>,
{
    type ParameterType = &'static T;

    impl_concrete_state! { Self::ParameterType }

    #[inline]
    fn push_result(&self, _ctx: &mut EvalContext, target: &mut [VectorValue]) -> Result<()> {
        // Note: The result of `AVG()` is returned as `(count, sum)`.
        assert_eq!(target.len(), 2);
        target[0].push_int(Some(self.count as Int));
        if self.count == 0 {
            target[1].push(None);
        } else {
            target[1].push(Some(self.sum.clone()));
        }
        Ok(())
    }
}

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

    use tidb_query_datatype::FieldTypeAccessor;
    use tipb_helper::ExprDefBuilder;

    use crate::parser::AggrDefinitionParser;
    use tidb_query_datatype::codec::batch::{LazyBatchColumn, LazyBatchColumnVec};

    #[test]
    fn test_update() {
        let mut ctx = EvalContext::default();
        let function = AggrFnAvg::<Real>::new();
        let mut state = function.create_state();

        let mut result = [
            VectorValue::with_capacity(0, EvalType::Int),
            VectorValue::with_capacity(0, EvalType::Real),
        ];
        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(result[0].to_int_vec(), &[Some(0)]);
        assert_eq!(result[1].to_real_vec(), &[None]);

        update!(state, &mut ctx, Option::<&Real>::None).unwrap();

        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(result[0].to_int_vec(), &[Some(0), Some(0)]);
        assert_eq!(result[1].to_real_vec(), &[None, None]);

        update!(state, &mut ctx, Real::new(5.0).ok().as_ref()).unwrap();
        update!(state, &mut ctx, Option::<&Real>::None).unwrap();
        update!(state, &mut ctx, Real::new(10.0).ok().as_ref()).unwrap();

        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(result[0].to_int_vec(), &[Some(0), Some(0), Some(2)]);
        assert_eq!(result[1].to_real_vec(), &[None, None, Real::new(15.0).ok()]);

        let x: ChunkedVecSized<Real> = vec![Real::new(0.0).ok(), Real::new(-4.5).ok(), None].into();

        update_vector!(state, &mut ctx, &x, &[0, 1, 2]).unwrap();

        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(
            result[0].to_int_vec(),
            &[Some(0), Some(0), Some(2), Some(4)]
        );
        assert_eq!(
            result[1].to_real_vec(),
            &[None, None, Real::new(15.0).ok(), Real::new(10.5).ok()]
        );
    }

    /// AVG(IntColumn) should produce (Int, Decimal).
    #[test]
    fn test_integration() {
        let expr = ExprDefBuilder::aggr_func(ExprType::Avg, FieldTypeTp::NewDecimal)
            .push_child(ExprDefBuilder::column_ref(0, FieldTypeTp::LongLong))
            .build();
        AggrFnDefinitionParserAvg.check_supported(&expr).unwrap();

        let src_schema = [FieldTypeTp::LongLong.into()];
        let mut columns = LazyBatchColumnVec::from(vec![{
            let mut col = LazyBatchColumn::decoded_with_capacity_and_tp(0, EvalType::Int);
            col.mut_decoded().push_int(Some(100));
            col.mut_decoded().push_int(Some(1));
            col.mut_decoded().push_int(None);
            col.mut_decoded().push_int(Some(42));
            col.mut_decoded().push_int(None);
            col
        }]);

        let mut schema = vec![];
        let mut exp = vec![];

        let mut ctx = EvalContext::default();
        let aggr_fn = AggrFnDefinitionParserAvg
            .parse(expr, &mut ctx, &src_schema, &mut schema, &mut exp)
            .unwrap();
        assert_eq!(schema.len(), 2);
        assert_eq!(schema[0].as_accessor().tp(), FieldTypeTp::LongLong);
        assert_eq!(schema[1].as_accessor().tp(), FieldTypeTp::NewDecimal);

        assert_eq!(exp.len(), 1);

        let mut state = aggr_fn.create_state();
        let mut ctx = EvalContext::default();

        let exp_result = exp[0]
            .eval(&mut ctx, &src_schema, &mut columns, &[4, 1, 2, 3], 4)
            .unwrap();
        let exp_result = exp_result.vector_value().unwrap();
        let slice = exp_result.as_ref().to_decimal_vec();
        let slice: ChunkedVecSized<Decimal> = slice.into();
        update_vector!(state, &mut ctx, &slice, exp_result.logical_rows()).unwrap();

        let mut aggr_result = [
            VectorValue::with_capacity(0, EvalType::Int),
            VectorValue::with_capacity(0, EvalType::Decimal),
        ];
        state.push_result(&mut ctx, &mut aggr_result).unwrap();

        assert_eq!(aggr_result[0].to_int_vec(), &[Some(2)]);
        assert_eq!(
            aggr_result[1].to_decimal_vec(),
            &[Some(Decimal::from(43u64))]
        );
    }

    #[test]
    fn test_illegal_request() {
        let expr = ExprDefBuilder::aggr_func(ExprType::Avg, FieldTypeTp::Double) // Expect NewDecimal but give Real
            .push_child(ExprDefBuilder::column_ref(0, FieldTypeTp::LongLong)) // FIXME: This type can be incorrect as well
            .build();
        AggrFnDefinitionParserAvg.check_supported(&expr).unwrap();

        let src_schema = [FieldTypeTp::LongLong.into()];
        let mut schema = vec![];
        let mut exp = vec![];
        let mut ctx = EvalContext::default();
        AggrFnDefinitionParserAvg
            .parse(expr, &mut ctx, &src_schema, &mut schema, &mut exp)
            .unwrap_err();
    }
}
<commit_msg>copr: Implement ENUM/SET support for aggr AVG  (#9186)<commit_after>// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.

use tidb_query_codegen::AggrFunction;
use tidb_query_common::Result;
use tidb_query_datatype::builder::FieldTypeBuilder;
use tidb_query_datatype::codec::data_type::*;
use tidb_query_datatype::expr::EvalContext;
use tidb_query_datatype::{EvalType, FieldTypeFlag, FieldTypeTp};
use tidb_query_expr::RpnExpression;
use tipb::{Expr, ExprType, FieldType};

use super::summable::Summable;
use super::*;

/// The parser for AVG aggregate function.
pub struct AggrFnDefinitionParserAvg;

impl super::AggrDefinitionParser for AggrFnDefinitionParserAvg {
    fn check_supported(&self, aggr_def: &Expr) -> Result<()> {
        assert_eq!(aggr_def.get_tp(), ExprType::Avg);
        super::util::check_aggr_exp_supported_one_child(aggr_def)
    }

    #[inline]
    fn parse_rpn(
        &self,
        mut root_expr: Expr,
        mut exp: RpnExpression,
        _ctx: &mut EvalContext,
        src_schema: &[FieldType],
        out_schema: &mut Vec<FieldType>,
        out_exp: &mut Vec<RpnExpression>,
    ) -> Result<Box<dyn AggrFunction>> {
        use std::convert::TryFrom;
        use tidb_query_datatype::FieldTypeAccessor;

        assert_eq!(root_expr.get_tp(), ExprType::Avg);

        let col_sum_ft = root_expr.take_field_type();
        let col_sum_et = box_try!(EvalType::try_from(col_sum_ft.as_accessor().tp()));

        // Rewrite expression to insert CAST() if needed.
        super::util::rewrite_exp_for_sum_avg(src_schema, &mut exp).unwrap();

        let rewritten_eval_type =
            EvalType::try_from(exp.ret_field_type(src_schema).as_accessor().tp()).unwrap();
        if col_sum_et != rewritten_eval_type {
            return Err(other_err!(
                "Unexpected return field type {}",
                col_sum_ft.as_accessor().tp()
            ));
        }

        // AVG outputs two columns.
        out_schema.push(
            FieldTypeBuilder::new()
                .tp(FieldTypeTp::LongLong)
                .flag(FieldTypeFlag::UNSIGNED)
                .build(),
        );
        out_schema.push(col_sum_ft);
        out_exp.push(exp);

        Ok(match rewritten_eval_type {
            EvalType::Decimal => Box::new(AggrFnAvg::<Decimal>::new()),
            EvalType::Real => Box::new(AggrFnAvg::<Real>::new()),
            EvalType::Enum => Box::new(AggrFnAvgForEnum::new()),
            EvalType::Set => Box::new(AggrFnAvgForSet::new()),
            _ => unreachable!(),
        })
    }
}

/// The AVG aggregate function.
///
/// Note that there are `AVG(Decimal) -> (Int, Decimal)` and `AVG(Double) -> (Int, Double)`.
#[derive(Debug, AggrFunction)]
#[aggr_function(state = AggrFnStateAvg::<T>::new())]
pub struct AggrFnAvg<T>
where
    T: Summable,
    VectorValue: VectorValueExt<T>,
{
    _phantom: std::marker::PhantomData<T>,
}

impl<T> AggrFnAvg<T>
where
    T: Summable,
    VectorValue: VectorValueExt<T>,
{
    pub fn new() -> Self {
        Self {
            _phantom: std::marker::PhantomData,
        }
    }
}

/// The state of the AVG aggregate function.
#[derive(Debug)]
pub struct AggrFnStateAvg<T>
where
    T: Summable,
    VectorValue: VectorValueExt<T>,
{
    sum: T,
    count: usize,
}

impl<T> AggrFnStateAvg<T>
where
    T: Summable,
    VectorValue: VectorValueExt<T>,
{
    pub fn new() -> Self {
        Self {
            sum: T::zero(),
            count: 0,
        }
    }

    #[inline]
    fn update_concrete<'a, TT>(&mut self, ctx: &mut EvalContext, value: Option<TT>) -> Result<()>
    where
        TT: EvaluableRef<'a, EvaluableType = T>,
    {
        match value {
            None => Ok(()),
            Some(value) => {
                self.sum.add_assign(ctx, &value.to_owned_value())?;
                self.count += 1;
                Ok(())
            }
        }
    }
}

impl<T> super::ConcreteAggrFunctionState for AggrFnStateAvg<T>
where
    T: Summable,
    VectorValue: VectorValueExt<T>,
{
    type ParameterType = &'static T;

    impl_concrete_state! { Self::ParameterType }

    #[inline]
    fn push_result(&self, _ctx: &mut EvalContext, target: &mut [VectorValue]) -> Result<()> {
        // Note: The result of `AVG()` is returned as `(count, sum)`.
        assert_eq!(target.len(), 2);
        target[0].push_int(Some(self.count as Int));
        if self.count == 0 {
            target[1].push(None);
        } else {
            target[1].push(Some(self.sum.clone()));
        }
        Ok(())
    }
}

#[derive(Debug, AggrFunction)]
#[aggr_function(state = AggrFnStateAvgForEnum::new())]
pub struct AggrFnAvgForEnum
where
    VectorValue: VectorValueExt<Enum>,
{
    _phantom: std::marker::PhantomData<Enum>,
}

impl AggrFnAvgForEnum
where
    VectorValue: VectorValueExt<Enum>,
{
    pub fn new() -> Self {
        Self {
            _phantom: std::marker::PhantomData,
        }
    }
}

#[derive(Debug)]
pub struct AggrFnStateAvgForEnum
where
    VectorValue: VectorValueExt<Enum>,
{
    sum: Decimal,
    count: usize,
}

impl AggrFnStateAvgForEnum
where
    VectorValue: VectorValueExt<Enum>,
{
    pub fn new() -> Self {
        Self {
            sum: Decimal::zero(),
            count: 0,
        }
    }

    #[inline]
    fn update_concrete(&mut self, ctx: &mut EvalContext, value: Option<EnumRef>) -> Result<()> {
        match value {
            None => Ok(()),
            Some(value) => {
                self.sum.add_assign(ctx, &Decimal::from(value.value()))?;
                self.count += 1;
                Ok(())
            }
        }
    }
}

impl super::ConcreteAggrFunctionState for AggrFnStateAvgForEnum
where
    VectorValue: VectorValueExt<Enum>,
{
    type ParameterType = EnumRef<'static>;

    impl_concrete_state! { Self::ParameterType }

    #[inline]
    fn push_result(&self, _ctx: &mut EvalContext, target: &mut [VectorValue]) -> Result<()> {
        // Note: The result of `AVG()` is returned as `(count, sum)`.
        assert_eq!(target.len(), 2);
        target[0].push_int(Some(self.count as Int));
        target[1].push(if self.count == 0 {
            None
        } else {
            Some(self.sum)
        });
        Ok(())
    }
}

#[derive(Debug, AggrFunction)]
#[aggr_function(state = AggrFnStateAvgForSet::new())]
pub struct AggrFnAvgForSet
where
    VectorValue: VectorValueExt<Set>,
{
    _phantom: std::marker::PhantomData<Set>,
}

impl AggrFnAvgForSet
where
    VectorValue: VectorValueExt<Set>,
{
    pub fn new() -> Self {
        Self {
            _phantom: std::marker::PhantomData,
        }
    }
}

#[derive(Debug)]
pub struct AggrFnStateAvgForSet
where
    VectorValue: VectorValueExt<Set>,
{
    sum: Decimal,
    count: usize,
}

impl AggrFnStateAvgForSet
where
    VectorValue: VectorValueExt<Set>,
{
    pub fn new() -> Self {
        Self {
            sum: Decimal::zero(),
            count: 0,
        }
    }

    #[inline]
    fn update_concrete(&mut self, ctx: &mut EvalContext, value: Option<SetRef>) -> Result<()> {
        match value {
            None => Ok(()),
            Some(value) => {
                self.sum.add_assign(ctx, &Decimal::from(value.value()))?;
                self.count += 1;
                Ok(())
            }
        }
    }
}

impl super::ConcreteAggrFunctionState for AggrFnStateAvgForSet
where
    VectorValue: VectorValueExt<Set>,
{
    type ParameterType = SetRef<'static>;

    impl_concrete_state! { Self::ParameterType }

    #[inline]
    fn push_result(&self, _ctx: &mut EvalContext, target: &mut [VectorValue]) -> Result<()> {
        // Note: The result of `AVG()` is returned as `(count, sum)`.
        assert_eq!(target.len(), 2);
        target[0].push_int(Some(self.count as Int));
        target[1].push(if self.count == 0 {
            None
        } else {
            Some(self.sum)
        });
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use tidb_query_datatype::codec::batch::{LazyBatchColumn, LazyBatchColumnVec};
    use tidb_query_datatype::FieldTypeAccessor;
    use tikv_util::buffer_vec::BufferVec;
    use tipb_helper::ExprDefBuilder;

    use crate::parser::AggrDefinitionParser;

    use super::super::AggrFunction;
    use super::*;

    #[test]
    fn test_update() {
        let mut ctx = EvalContext::default();
        let function = AggrFnAvg::<Real>::new();
        let mut state = function.create_state();

        let mut result = [
            VectorValue::with_capacity(0, EvalType::Int),
            VectorValue::with_capacity(0, EvalType::Real),
        ];
        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(result[0].to_int_vec(), &[Some(0)]);
        assert_eq!(result[1].to_real_vec(), &[None]);

        update!(state, &mut ctx, Option::<&Real>::None).unwrap();

        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(result[0].to_int_vec(), &[Some(0), Some(0)]);
        assert_eq!(result[1].to_real_vec(), &[None, None]);

        update!(state, &mut ctx, Real::new(5.0).ok().as_ref()).unwrap();
        update!(state, &mut ctx, Option::<&Real>::None).unwrap();
        update!(state, &mut ctx, Real::new(10.0).ok().as_ref()).unwrap();

        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(result[0].to_int_vec(), &[Some(0), Some(0), Some(2)]);
        assert_eq!(result[1].to_real_vec(), &[None, None, Real::new(15.0).ok()]);

        let x: ChunkedVecSized<Real> = vec![Real::new(0.0).ok(), Real::new(-4.5).ok(), None].into();

        update_vector!(state, &mut ctx, &x, &[0, 1, 2]).unwrap();

        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(
            result[0].to_int_vec(),
            &[Some(0), Some(0), Some(2), Some(4)]
        );
        assert_eq!(
            result[1].to_real_vec(),
            &[None, None, Real::new(15.0).ok(), Real::new(10.5).ok()]
        );
    }

    #[test]
    fn test_update_enum() {
        let mut ctx = EvalContext::default();
        let function = AggrFnAvgForEnum::new();
        let mut state = function.create_state();

        // AVG will returns <Int, Decimal>
        let mut result = [
            VectorValue::with_capacity(0, EvalType::Int),
            VectorValue::with_capacity(0, EvalType::Decimal),
        ];

        let mut buf = BufferVec::new();
        buf.push("我好强啊");
        buf.push("我太强啦");
        let buf = Arc::new(buf);

        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(result[0].to_int_vec(), &[Some(0)]);
        assert_eq!(result[1].to_decimal_vec(), &[None]);

        update!(state, &mut ctx, Some(EnumRef::new(&buf, 1))).unwrap();
        update!(state, &mut ctx, Some(EnumRef::new(&buf, 2))).unwrap();
        result[0].clear();
        result[1].clear();
        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(result[0].to_int_vec(), &[Some(2)]);
        assert_eq!(result[1].to_decimal_vec(), &[Some(Decimal::from(3))]);

        update!(state, &mut ctx, Option::<EnumRef>::None).unwrap();
        result[0].clear();
        result[1].clear();
        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(result[0].to_int_vec(), &[Some(2)]);
        assert_eq!(result[1].to_decimal_vec(), &[Some(Decimal::from(3))]);
    }

    #[test]
    fn test_update_set() {
        let mut ctx = EvalContext::default();
        let function = AggrFnAvgForSet::new();
        let mut state = function.create_state();

        // AVG will returns <Int, Decimal>
        let mut result = [
            VectorValue::with_capacity(0, EvalType::Int),
            VectorValue::with_capacity(0, EvalType::Decimal),
        ];

        let mut buf = BufferVec::new();
        buf.push("我好强啊");
        buf.push("我太强啦");
        let buf = Arc::new(buf);

        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(result[0].to_int_vec(), &[Some(0)]);
        assert_eq!(result[1].to_decimal_vec(), &[None]);

        update!(state, &mut ctx, Some(SetRef::new(&buf, 0b01))).unwrap();
        update!(state, &mut ctx, Some(SetRef::new(&buf, 0b10))).unwrap();
        result[0].clear();
        result[1].clear();
        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(result[0].to_int_vec(), &[Some(2)]);
        assert_eq!(result[1].to_decimal_vec(), &[Some(Decimal::from(3))]);

        update!(state, &mut ctx, Option::<SetRef>::None).unwrap();
        result[0].clear();
        result[1].clear();
        state.push_result(&mut ctx, &mut result[..]).unwrap();
        assert_eq!(result[0].to_int_vec(), &[Some(2)]);
        assert_eq!(result[1].to_decimal_vec(), &[Some(Decimal::from(3))]);
    }

    /// AVG(IntColumn) should produce (Int, Decimal).
    #[test]
    fn test_integration() {
        let expr = ExprDefBuilder::aggr_func(ExprType::Avg, FieldTypeTp::NewDecimal)
            .push_child(ExprDefBuilder::column_ref(0, FieldTypeTp::LongLong))
            .build();
        AggrFnDefinitionParserAvg.check_supported(&expr).unwrap();

        let src_schema = [FieldTypeTp::LongLong.into()];
        let mut columns = LazyBatchColumnVec::from(vec![{
            let mut col = LazyBatchColumn::decoded_with_capacity_and_tp(0, EvalType::Int);
            col.mut_decoded().push_int(Some(100));
            col.mut_decoded().push_int(Some(1));
            col.mut_decoded().push_int(None);
            col.mut_decoded().push_int(Some(42));
            col.mut_decoded().push_int(None);
            col
        }]);

        let mut schema = vec![];
        let mut exp = vec![];

        let mut ctx = EvalContext::default();
        let aggr_fn = AggrFnDefinitionParserAvg
            .parse(expr, &mut ctx, &src_schema, &mut schema, &mut exp)
            .unwrap();
        assert_eq!(schema.len(), 2);
        assert_eq!(schema[0].as_accessor().tp(), FieldTypeTp::LongLong);
        assert_eq!(schema[1].as_accessor().tp(), FieldTypeTp::NewDecimal);

        assert_eq!(exp.len(), 1);

        let mut state = aggr_fn.create_state();
        let mut ctx = EvalContext::default();

        let exp_result = exp[0]
            .eval(&mut ctx, &src_schema, &mut columns, &[4, 1, 2, 3], 4)
            .unwrap();
        let exp_result = exp_result.vector_value().unwrap();
        let slice = exp_result.as_ref().to_decimal_vec();
        let slice: ChunkedVecSized<Decimal> = slice.into();
        update_vector!(state, &mut ctx, &slice, exp_result.logical_rows()).unwrap();

        let mut aggr_result = [
            VectorValue::with_capacity(0, EvalType::Int),
            VectorValue::with_capacity(0, EvalType::Decimal),
        ];
        state.push_result(&mut ctx, &mut aggr_result).unwrap();

        assert_eq!(aggr_result[0].to_int_vec(), &[Some(2)]);
        assert_eq!(
            aggr_result[1].to_decimal_vec(),
            &[Some(Decimal::from(43u64))]
        );
    }

    #[test]
    fn test_illegal_request() {
        let expr = ExprDefBuilder::aggr_func(ExprType::Avg, FieldTypeTp::Double) // Expect NewDecimal but give Real
            .push_child(ExprDefBuilder::column_ref(0, FieldTypeTp::LongLong)) // FIXME: This type can be incorrect as well
            .build();
        AggrFnDefinitionParserAvg.check_supported(&expr).unwrap();

        let src_schema = [FieldTypeTp::LongLong.into()];
        let mut schema = vec![];
        let mut exp = vec![];
        let mut ctx = EvalContext::default();
        AggrFnDefinitionParserAvg
            .parse(expr, &mut ctx, &src_schema, &mut schema, &mut exp)
            .unwrap_err();
    }
}
<|endoftext|>


<commit_before>#[derive(Debug, PartialEq)]
pub struct WordProblem {
    command: String,
}

#[derive(Debug, PartialEq)]
struct Token {
    value: String,
}

impl Token {
    fn is_valid(&self) -> bool {
        !self.value.is_empty() && (self.is_operand() || self.is_operator())
    }

    fn is_operand(&self) -> bool {
        self.value.chars().all(|c| c.is_numeric() || c == '-')
    }

    fn is_operator(&self) -> bool {
        self.value == String::from("plus") || self.value == String::from("minus") ||
        self.value == String::from("multiplied") || self.value == String::from("divided")
    }
}

impl WordProblem {
    pub fn new(c: &str) -> Self {
        WordProblem { command: String::from(c) }
    }

    pub fn answer(&self) -> Result<isize, ()> {
        let mut t = self.tokens();
        let mut result: isize = 0;
        let mut opr = "plus".to_string();

        if t.len() <= 1 {
            Err(())
        } else {
            while t.len() > 1 {
                result = self.evaluate(result, opr, self.operand(&t.remove(0)));
                opr = self.operator(&t.remove(0));
            }
            result = self.evaluate(result, opr, self.operand(&t.remove(0)));
            Ok(result)
        }
    }

    fn evaluate(&self, mut r: isize, operator: String, operand: isize) -> isize {
        if operator == String::from("plus") {
            r += operand;
            r
        } else if operator == String::from("minus") {
            r -= operand;
            r
        } else if operator == String::from("multiplied") {
            r *= operand;
            r
        } else if operator == String::from("divided") {
            r /= operand;
            r
        } else {
            r
        }
    }

    fn operand(&self, t: &Token) -> isize {
        let x: isize = t.value.parse().unwrap();
        x
    }

    fn operator(&self, t: &Token) -> String {
        String::from(t.value.clone())
    }

    fn tokens(&self) -> Vec<Token> {
        self.command
            .split(|c: char| c.is_whitespace() || c == '?')
            .map(|w| Token { value: String::from(w) })
            .filter(|t| t.is_valid())
            .collect()
    }
}
<commit_msg>wordy: Immediately return in `operand`<commit_after>#[derive(Debug, PartialEq)]
pub struct WordProblem {
    command: String,
}

#[derive(Debug, PartialEq)]
struct Token {
    value: String,
}

impl Token {
    fn is_valid(&self) -> bool {
        !self.value.is_empty() && (self.is_operand() || self.is_operator())
    }

    fn is_operand(&self) -> bool {
        self.value.chars().all(|c| c.is_numeric() || c == '-')
    }

    fn is_operator(&self) -> bool {
        self.value == String::from("plus") || self.value == String::from("minus") ||
        self.value == String::from("multiplied") || self.value == String::from("divided")
    }
}

impl WordProblem {
    pub fn new(c: &str) -> Self {
        WordProblem { command: String::from(c) }
    }

    pub fn answer(&self) -> Result<isize, ()> {
        let mut t = self.tokens();
        let mut result: isize = 0;
        let mut opr = "plus".to_string();

        if t.len() <= 1 {
            Err(())
        } else {
            while t.len() > 1 {
                result = self.evaluate(result, opr, self.operand(&t.remove(0)));
                opr = self.operator(&t.remove(0));
            }
            result = self.evaluate(result, opr, self.operand(&t.remove(0)));
            Ok(result)
        }
    }

    fn evaluate(&self, mut r: isize, operator: String, operand: isize) -> isize {
        if operator == String::from("plus") {
            r += operand;
            r
        } else if operator == String::from("minus") {
            r -= operand;
            r
        } else if operator == String::from("multiplied") {
            r *= operand;
            r
        } else if operator == String::from("divided") {
            r /= operand;
            r
        } else {
            r
        }
    }

    fn operand(&self, t: &Token) -> isize {
        t.value.parse().unwrap()
    }

    fn operator(&self, t: &Token) -> String {
        String::from(t.value.clone())
    }

    fn tokens(&self) -> Vec<Token> {
        self.command
            .split(|c: char| c.is_whitespace() || c == '?')
            .map(|w| Token { value: String::from(w) })
            .filter(|t| t.is_valid())
            .collect()
    }
}
<|endoftext|>


<commit_before>//! Android specific features

use std::{
    convert::TryInto,
    io::{self, Error, ErrorKind, Read, Write},
    mem,
    net::Shutdown,
    os::unix::io::{AsRawFd, RawFd},
    path::Path,
    pin::Pin,
    ptr,
    slice,
    task::{Context, Poll},
};

use futures::{future, ready};
use mio::net::{SocketAddr, UnixListener as MioUnixListener, UnixStream as MioUnixStream};
use tokio::io::{unix::AsyncFd, AsyncRead, AsyncWrite, ReadBuf};

/// A UnixStream supports transferring FDs between processes
pub struct UnixStream {
    io: AsyncFd<MioUnixStream>,
}

impl UnixStream {
    /// Connects to the socket named by `path`.
    pub async fn connect<P: AsRef<Path>>(path: P) -> io::Result<UnixStream> {
        let uds = MioUnixStream::connect(path)?;
        let io = AsyncFd::new(uds)?;

        let mut ready = future::poll_fn(|cx| io.poll_write_ready(cx)).await?;
        ready.retain_ready();
        Ok(UnixStream { io })
    }

    /// Create from a mio's `UnixStream`
    pub fn from_mio(io: MioUnixStream) -> io::Result<UnixStream> {
        Ok(UnixStream { io: AsyncFd::new(io)? })
    }

    fn poll_send_with_fd(&self, cx: &mut Context, buf: &[u8], fds: &[RawFd]) -> Poll<io::Result<usize>> {
        loop {
            let mut ready = ready!(self.io.poll_write_ready(cx))?;

            let fd = self.io.get_ref().as_raw_fd();
            match send_with_fd(fd, buf, fds) {
                // self.io.poll_write_ready indicates that writable event have been received by tokio,
                // so it is not a common case that sendto returns EAGAIN.
                //
                // Just for double check. If EAGAIN actually returns, clear the readness state.
                Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
                    ready.clear_ready();
                }
                x => return Poll::Ready(x),
            }
        }
    }

    /// Send data with file descriptors
    pub async fn send_with_fd(&mut self, buf: &[u8], fds: &[RawFd]) -> io::Result<usize> {
        future::poll_fn(|cx| self.poll_send_with_fd(cx, buf, fds)).await
    }

    fn poll_recv_with_fd(
        &self,
        cx: &mut Context,
        buf: &mut [u8],
        fds: &mut [RawFd],
    ) -> Poll<io::Result<(usize, usize)>> {
        loop {
            let mut ready = ready!(self.io.poll_write_ready(cx))?;

            let fd = self.io.get_ref().as_raw_fd();
            match recv_with_fd(fd, buf, fds) {
                // self.io.poll_write_ready indicates that writable event have been received by tokio,
                // so it is not a common case that recvto returns EAGAIN.
                //
                // Just for double check. If EAGAIN actually returns, clear the readness state.
                Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
                    ready.clear_ready();
                }
                x => return Poll::Ready(x),
            }
        }
    }

    /// Recv data with file descriptors
    pub async fn recv_with_fd(&mut self, buf: &mut [u8], fds: &mut [RawFd]) -> io::Result<(usize, usize)> {
        future::poll_fn(|cx| self.poll_recv_with_fd(cx, buf, fds)).await
    }

    /// Shuts down the read, write, or both halves of this connection.
    ///
    /// This function will cause all pending and future I/O calls on the
    /// specified portions to immediately return with an appropriate value
    /// (see the documentation of `Shutdown`).
    pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
        self.io.get_ref().shutdown(how)
    }
}

impl AsyncRead for UnixStream {
    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        self.poll_read_priv(cx, buf)
    }
}

impl AsyncWrite for UnixStream {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        self.poll_write_priv(cx, buf)
    }

    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        self.shutdown(Shutdown::Write)?;
        Poll::Ready(Ok(()))
    }
}

impl UnixStream {
    // == Poll IO functions that takes `&self` ==
    //
    // They are not public because (taken from the doc of `PollEvented`):
    //
    // While `PollEvented` is `Sync` (if the underlying I/O type is `Sync`), the
    // caller must ensure that there are at most two tasks that use a
    // `PollEvented` instance concurrently. One for reading and one for writing.
    // While violating this requirement is "safe" from a Rust memory model point
    // of view, it will result in unexpected behavior in the form of lost
    // notifications and tasks hanging.

    pub(crate) fn poll_read_priv(&self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        loop {
            let mut read_guard = ready!(self.io.poll_read_ready(cx))?;

            let b = unsafe { &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) };
            match self.io.get_ref().read(b) {
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    read_guard.clear_ready();
                }
                Ok(n) => {
                    // Safety: We trust `UnixStream::read` to have filled up `n` bytes
                    // in the buffer.
                    unsafe {
                        buf.assume_init(n);
                    }
                    buf.advance(n);
                    return Poll::Ready(Ok(()));
                }
                Err(e) => return Poll::Ready(Err(e)),
            }
        }
    }

    pub(crate) fn poll_write_priv(&self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        loop {
            let mut write_guard = ready!(self.io.poll_write_ready(cx))?;

            match self.io.get_ref().write(buf) {
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    write_guard.clear_ready();
                }
                x => return Poll::Ready(x),
            }
        }
    }
}

/// A UnixListener supports transferring FDs between processes
pub struct UnixListener {
    io: AsyncFd<MioUnixListener>,
}

impl UnixListener {
    /// Creates a new `UnixListener` bound to the specified socket.
    pub fn bind<P: AsRef<Path>>(path: P) -> io::Result<UnixListener> {
        Ok(UnixListener {
            io: AsyncFd::new(MioUnixListener::bind(path)?)?,
        })
    }

    /// Accepts a new incoming connection to this listener.
    pub fn poll_accept(&self, cx: &mut Context<'_>) -> Poll<io::Result<(UnixStream, SocketAddr)>> {
        loop {
            let mut read_guard = ready!(self.io.poll_read_ready(cx))?;

            match self.io.get_ref().accept() {
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    read_guard.clear_ready();
                }
                Ok((stream, peer_addr)) => return Poll::Ready(Ok((UnixStream::from_mio(stream)?, peer_addr))),
                Err(err) => return Poll::Ready(Err(err)),
            }
        }
    }

    /// Accepts a new incoming connection to this listener.
    pub async fn accept(&self) -> io::Result<(UnixStream, SocketAddr)> {
        future::poll_fn(|cx| self.poll_accept(cx)).await
    }
}

/// A common implementation of `sendmsg` that sends provided bytes with ancillary file descriptors
/// over either a datagram or stream unix socket.
///
/// Borrowed from: https://github.com/Standard-Cognition/sendfd
fn send_with_fd(socket: RawFd, bs: &[u8], fds: &[RawFd]) -> io::Result<usize> {
    unsafe {
        let mut iov = libc::iovec {
            // NB: this casts *const to *mut, and in doing so we trust the OS to be a good citizen
            // and not mutate our buffer. This is the API we have to live with.
            iov_base: bs.as_ptr() as *const _ as *mut _,
            iov_len: bs.len(),
        };

        // Construct msghdr
        //
        // 1. Allocate memory for msg_control
        let cmsg_fd_len = fds.len() * mem::size_of::<RawFd>();
        let cmsg_buffer_len = libc::CMSG_SPACE(cmsg_fd_len as u32) as usize;
        let mut cmsg_buffer = Vec::with_capacity(cmsg_buffer_len);
        cmsg_buffer.set_len(cmsg_buffer_len);

        let mut msghdr = libc::msghdr {
            msg_name: ptr::null_mut(),
            msg_namelen: 0,
            msg_iov: &mut iov as *mut _,
            msg_iovlen: 1,
            msg_control: cmsg_buffer.as_mut_ptr(),
            msg_controllen: cmsg_buffer_len.try_into().unwrap(),
            ..mem::zeroed()
        };

        // Fill cmsg with the file descriptors we are sending.
        let cmsg_header = libc::CMSG_FIRSTHDR(&mut msghdr as *mut _);
        let mut cmsg: libc::cmsghdr = mem::zeroed();
        cmsg.cmsg_level = libc::SOL_SOCKET;
        cmsg.cmsg_type = libc::SCM_RIGHTS;
        cmsg.cmsg_len = libc::CMSG_LEN(cmsg_fd_len as u32) as _;
        cmsg_header.write(cmsg);

        let cmsg_data = libc::CMSG_DATA(cmsg_header);
        #[allow(clippy::cast_ptr_alignment)] // false positive
        let cmsg_data_slice = slice::from_raw_parts_mut(cmsg_data as *mut RawFd, fds.len());
        cmsg_data_slice.copy_from_slice(fds);

        let count = libc::sendmsg(socket, &msghdr as *const _, 0);
        if count < 0 {
            let err = Error::last_os_error();
            Err(err)
        } else {
            Ok(count as usize)
        }
    }
}

/// A common implementation of `recvmsg` that receives provided bytes and the ancillary file
/// descriptors over either a datagram or stream unix socket.
///
/// Borrowed from: https://github.com/Standard-Cognition/sendfd
fn recv_with_fd(socket: RawFd, bs: &mut [u8], mut fds: &mut [RawFd]) -> io::Result<(usize, usize)> {
    unsafe {
        let mut iov = libc::iovec {
            iov_base: bs.as_mut_ptr() as *mut _,
            iov_len: bs.len(),
        };

        // Construct msghdr
        //
        // 1. Allocate memory for msg_control
        let cmsg_fd_len = fds.len() * mem::size_of::<RawFd>();
        let cmsg_buffer_len = libc::CMSG_SPACE(cmsg_fd_len as u32) as usize;
        let mut cmsg_buffer = Vec::with_capacity(cmsg_buffer_len);
        cmsg_buffer.set_len(cmsg_buffer_len);

        let mut msghdr = libc::msghdr {
            msg_name: ptr::null_mut(),
            msg_namelen: 0,
            msg_iov: &mut iov as *mut _,
            msg_iovlen: 1,
            msg_control: cmsg_buffer.as_mut_ptr(),
            msg_controllen: cmsg_buffer_len.try_into().unwrap(),
            ..mem::zeroed()
        };

        let count = libc::recvmsg(socket, &mut msghdr as *mut _, 0);
        if count < 0 {
            let error = io::Error::last_os_error();
            return Err(error);
        }

        // Walk the ancillary data buffer and copy the raw descriptors from it into the output
        // buffer.
        let mut descriptor_count = 0;
        let mut cmsg_header = libc::CMSG_FIRSTHDR(&mut msghdr as *mut _);
        while !cmsg_header.is_null() {
            if (*cmsg_header).cmsg_level == libc::SOL_SOCKET && (*cmsg_header).cmsg_type == libc::SCM_RIGHTS {
                let data_ptr = libc::CMSG_DATA(cmsg_header);
                let data_offset = data_ptr.offset_from(cmsg_header as *const _);
                debug_assert!(data_offset >= 0);
                let data_byte_count = (*cmsg_header).cmsg_len as usize - data_offset as usize;
                debug_assert!((*cmsg_header).cmsg_len as isize > data_offset);
                debug_assert!(data_byte_count % mem::size_of::<RawFd>() == 0);
                let rawfd_count = (data_byte_count / mem::size_of::<RawFd>()) as isize;
                let fd_ptr = data_ptr as *const RawFd;
                for i in 0..rawfd_count {
                    if let Some((dst, rest)) = { fds }.split_first_mut() {
                        *dst = ptr::read_unaligned(fd_ptr.offset(i));
                        descriptor_count += 1;
                        fds = rest;
                    } else {
                        // This branch is unreachable. We allocate the ancillary data buffer just
                        // large enough to fit exactly the number of `RawFd`s that are in the `fds`
                        // buffer. It is not possible for the OS to return more of them.
                        //
                        // If this branch ended up being reachable for some reason, it would be
                        // necessary for this branch to close the file descriptors to avoid leaking
                        // resources.
                        //
                        // TODO: consider using unreachable_unchecked
                        unreachable!();
                    }
                }
            }
            cmsg_header = libc::CMSG_NXTHDR(&mut msghdr as *mut _, cmsg_header);
        }

        Ok((count as usize, descriptor_count))
    }
}
<commit_msg>msghdr padding are private<commit_after>//! Android specific features

use std::{
    io::{self, Error, ErrorKind, Read, Write},
    mem,
    net::Shutdown,
    os::unix::io::{AsRawFd, RawFd},
    path::Path,
    pin::Pin,
    ptr,
    slice,
    task::{Context, Poll},
};

use futures::{future, ready};
use mio::net::{SocketAddr, UnixListener as MioUnixListener, UnixStream as MioUnixStream};
use tokio::io::{unix::AsyncFd, AsyncRead, AsyncWrite, ReadBuf};

/// A UnixStream supports transferring FDs between processes
pub struct UnixStream {
    io: AsyncFd<MioUnixStream>,
}

impl UnixStream {
    /// Connects to the socket named by `path`.
    pub async fn connect<P: AsRef<Path>>(path: P) -> io::Result<UnixStream> {
        let uds = MioUnixStream::connect(path)?;
        let io = AsyncFd::new(uds)?;

        let mut ready = future::poll_fn(|cx| io.poll_write_ready(cx)).await?;
        ready.retain_ready();
        Ok(UnixStream { io })
    }

    /// Create from a mio's `UnixStream`
    pub fn from_mio(io: MioUnixStream) -> io::Result<UnixStream> {
        Ok(UnixStream { io: AsyncFd::new(io)? })
    }

    fn poll_send_with_fd(&self, cx: &mut Context, buf: &[u8], fds: &[RawFd]) -> Poll<io::Result<usize>> {
        loop {
            let mut ready = ready!(self.io.poll_write_ready(cx))?;

            let fd = self.io.get_ref().as_raw_fd();
            match send_with_fd(fd, buf, fds) {
                // self.io.poll_write_ready indicates that writable event have been received by tokio,
                // so it is not a common case that sendto returns EAGAIN.
                //
                // Just for double check. If EAGAIN actually returns, clear the readness state.
                Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
                    ready.clear_ready();
                }
                x => return Poll::Ready(x),
            }
        }
    }

    /// Send data with file descriptors
    pub async fn send_with_fd(&mut self, buf: &[u8], fds: &[RawFd]) -> io::Result<usize> {
        future::poll_fn(|cx| self.poll_send_with_fd(cx, buf, fds)).await
    }

    fn poll_recv_with_fd(
        &self,
        cx: &mut Context,
        buf: &mut [u8],
        fds: &mut [RawFd],
    ) -> Poll<io::Result<(usize, usize)>> {
        loop {
            let mut ready = ready!(self.io.poll_write_ready(cx))?;

            let fd = self.io.get_ref().as_raw_fd();
            match recv_with_fd(fd, buf, fds) {
                // self.io.poll_write_ready indicates that writable event have been received by tokio,
                // so it is not a common case that recvto returns EAGAIN.
                //
                // Just for double check. If EAGAIN actually returns, clear the readness state.
                Err(ref err) if err.kind() == ErrorKind::WouldBlock => {
                    ready.clear_ready();
                }
                x => return Poll::Ready(x),
            }
        }
    }

    /// Recv data with file descriptors
    pub async fn recv_with_fd(&mut self, buf: &mut [u8], fds: &mut [RawFd]) -> io::Result<(usize, usize)> {
        future::poll_fn(|cx| self.poll_recv_with_fd(cx, buf, fds)).await
    }

    /// Shuts down the read, write, or both halves of this connection.
    ///
    /// This function will cause all pending and future I/O calls on the
    /// specified portions to immediately return with an appropriate value
    /// (see the documentation of `Shutdown`).
    pub fn shutdown(&self, how: Shutdown) -> io::Result<()> {
        self.io.get_ref().shutdown(how)
    }
}

impl AsyncRead for UnixStream {
    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        self.poll_read_priv(cx, buf)
    }
}

impl AsyncWrite for UnixStream {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        self.poll_write_priv(cx, buf)
    }

    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        Poll::Ready(Ok(()))
    }

    fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
        self.shutdown(Shutdown::Write)?;
        Poll::Ready(Ok(()))
    }
}

impl UnixStream {
    // == Poll IO functions that takes `&self` ==
    //
    // They are not public because (taken from the doc of `PollEvented`):
    //
    // While `PollEvented` is `Sync` (if the underlying I/O type is `Sync`), the
    // caller must ensure that there are at most two tasks that use a
    // `PollEvented` instance concurrently. One for reading and one for writing.
    // While violating this requirement is "safe" from a Rust memory model point
    // of view, it will result in unexpected behavior in the form of lost
    // notifications and tasks hanging.

    pub(crate) fn poll_read_priv(&self, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        loop {
            let mut read_guard = ready!(self.io.poll_read_ready(cx))?;

            let b = unsafe { &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) };
            match self.io.get_ref().read(b) {
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    read_guard.clear_ready();
                }
                Ok(n) => {
                    // Safety: We trust `UnixStream::read` to have filled up `n` bytes
                    // in the buffer.
                    unsafe {
                        buf.assume_init(n);
                    }
                    buf.advance(n);
                    return Poll::Ready(Ok(()));
                }
                Err(e) => return Poll::Ready(Err(e)),
            }
        }
    }

    pub(crate) fn poll_write_priv(&self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        loop {
            let mut write_guard = ready!(self.io.poll_write_ready(cx))?;

            match self.io.get_ref().write(buf) {
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    write_guard.clear_ready();
                }
                x => return Poll::Ready(x),
            }
        }
    }
}

/// A UnixListener supports transferring FDs between processes
pub struct UnixListener {
    io: AsyncFd<MioUnixListener>,
}

impl UnixListener {
    /// Creates a new `UnixListener` bound to the specified socket.
    pub fn bind<P: AsRef<Path>>(path: P) -> io::Result<UnixListener> {
        Ok(UnixListener {
            io: AsyncFd::new(MioUnixListener::bind(path)?)?,
        })
    }

    /// Accepts a new incoming connection to this listener.
    pub fn poll_accept(&self, cx: &mut Context<'_>) -> Poll<io::Result<(UnixStream, SocketAddr)>> {
        loop {
            let mut read_guard = ready!(self.io.poll_read_ready(cx))?;

            match self.io.get_ref().accept() {
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                    read_guard.clear_ready();
                }
                Ok((stream, peer_addr)) => return Poll::Ready(Ok((UnixStream::from_mio(stream)?, peer_addr))),
                Err(err) => return Poll::Ready(Err(err)),
            }
        }
    }

    /// Accepts a new incoming connection to this listener.
    pub async fn accept(&self) -> io::Result<(UnixStream, SocketAddr)> {
        future::poll_fn(|cx| self.poll_accept(cx)).await
    }
}

/// A common implementation of `sendmsg` that sends provided bytes with ancillary file descriptors
/// over either a datagram or stream unix socket.
///
/// Borrowed from: https://github.com/Standard-Cognition/sendfd
fn send_with_fd(socket: RawFd, bs: &[u8], fds: &[RawFd]) -> io::Result<usize> {
    unsafe {
        let mut iov = libc::iovec {
            // NB: this casts *const to *mut, and in doing so we trust the OS to be a good citizen
            // and not mutate our buffer. This is the API we have to live with.
            iov_base: bs.as_ptr() as *const _ as *mut _,
            iov_len: bs.len(),
        };

        // Construct msghdr
        //
        // 1. Allocate memory for msg_control
        let cmsg_fd_len = fds.len() * mem::size_of::<RawFd>();
        let cmsg_buffer_len = libc::CMSG_SPACE(cmsg_fd_len as u32) as usize;
        let mut cmsg_buffer = Vec::with_capacity(cmsg_buffer_len);
        cmsg_buffer.set_len(cmsg_buffer_len);

        let mut msghdr: libc::msghdr = mem::zeroed();
        msghdr.msg_iov = &mut iov as *mut _;
        msghdr.msg_iovlen = 1;
        msghdr.msg_control = cmsg_buffer.as_mut_ptr();
        msghdr.msg_controllen = cmsg_buffer_len as _;

        // Fill cmsg with the file descriptors we are sending.
        let cmsg_header = libc::CMSG_FIRSTHDR(&mut msghdr as *mut _);
        let mut cmsg: libc::cmsghdr = mem::zeroed();
        cmsg.cmsg_level = libc::SOL_SOCKET;
        cmsg.cmsg_type = libc::SCM_RIGHTS;
        cmsg.cmsg_len = libc::CMSG_LEN(cmsg_fd_len as u32) as _;
        cmsg_header.write(cmsg);

        let cmsg_data = libc::CMSG_DATA(cmsg_header);
        #[allow(clippy::cast_ptr_alignment)] // false positive
        let cmsg_data_slice = slice::from_raw_parts_mut(cmsg_data as *mut RawFd, fds.len());
        cmsg_data_slice.copy_from_slice(fds);

        let count = libc::sendmsg(socket, &msghdr as *const _, 0);
        if count < 0 {
            let err = Error::last_os_error();
            Err(err)
        } else {
            Ok(count as usize)
        }
    }
}

/// A common implementation of `recvmsg` that receives provided bytes and the ancillary file
/// descriptors over either a datagram or stream unix socket.
///
/// Borrowed from: https://github.com/Standard-Cognition/sendfd
fn recv_with_fd(socket: RawFd, bs: &mut [u8], mut fds: &mut [RawFd]) -> io::Result<(usize, usize)> {
    unsafe {
        let mut iov = libc::iovec {
            iov_base: bs.as_mut_ptr() as *mut _,
            iov_len: bs.len(),
        };

        // Construct msghdr
        //
        // 1. Allocate memory for msg_control
        let cmsg_fd_len = fds.len() * mem::size_of::<RawFd>();
        let cmsg_buffer_len = libc::CMSG_SPACE(cmsg_fd_len as u32) as usize;
        let mut cmsg_buffer = Vec::with_capacity(cmsg_buffer_len);
        cmsg_buffer.set_len(cmsg_buffer_len);

        let mut msghdr: libc::msghdr = mem::zeroed();
        msghdr.msg_iov = &mut iov as *mut _;
        msghdr.msg_iovlen = 1;
        msghdr.msg_control = cmsg_buffer.as_mut_ptr();
        msghdr.msg_controllen = cmsg_buffer_len as _;

        let count = libc::recvmsg(socket, &mut msghdr as *mut _, 0);
        if count < 0 {
            let error = io::Error::last_os_error();
            return Err(error);
        }

        // Walk the ancillary data buffer and copy the raw descriptors from it into the output
        // buffer.
        let mut descriptor_count = 0;
        let mut cmsg_header = libc::CMSG_FIRSTHDR(&mut msghdr as *mut _);
        while !cmsg_header.is_null() {
            if (*cmsg_header).cmsg_level == libc::SOL_SOCKET && (*cmsg_header).cmsg_type == libc::SCM_RIGHTS {
                let data_ptr = libc::CMSG_DATA(cmsg_header);
                let data_offset = data_ptr.offset_from(cmsg_header as *const _);
                debug_assert!(data_offset >= 0);
                let data_byte_count = (*cmsg_header).cmsg_len as usize - data_offset as usize;
                debug_assert!((*cmsg_header).cmsg_len as isize > data_offset);
                debug_assert!(data_byte_count % mem::size_of::<RawFd>() == 0);
                let rawfd_count = (data_byte_count / mem::size_of::<RawFd>()) as isize;
                let fd_ptr = data_ptr as *const RawFd;
                for i in 0..rawfd_count {
                    if let Some((dst, rest)) = { fds }.split_first_mut() {
                        *dst = ptr::read_unaligned(fd_ptr.offset(i));
                        descriptor_count += 1;
                        fds = rest;
                    } else {
                        // This branch is unreachable. We allocate the ancillary data buffer just
                        // large enough to fit exactly the number of `RawFd`s that are in the `fds`
                        // buffer. It is not possible for the OS to return more of them.
                        //
                        // If this branch ended up being reachable for some reason, it would be
                        // necessary for this branch to close the file descriptors to avoid leaking
                        // resources.
                        //
                        // TODO: consider using unreachable_unchecked
                        unreachable!();
                    }
                }
            }
            cmsg_header = libc::CMSG_NXTHDR(&mut msghdr as *mut _, cmsg_header);
        }

        Ok((count as usize, descriptor_count))
    }
}
<|endoftext|>


<commit_before>use common::*;
use super::{Capability};
use super::untyped::{UntypedHalf};
use core::mem::{size_of, align_of};
use core::ops::{Index, IndexMut};
use core::slice::Iter;
use arch;

#[derive(Debug, Clone)]
pub struct CPoolHalf {
    start_paddr: PAddr
}

#[derive(Debug)]
pub struct CPool([Option<Capability>; 32]);

impl Index<usize> for CPool {
    type Output = Option<Capability>;

    fn index<'a>(&'a self, _index: usize) -> &'a Option<Capability> {
        self.0.index(_index)
    }
}

impl IndexMut<usize> for CPool {
    fn index_mut<'a>(&'a mut self, _index: usize) -> &'a mut Option<Capability> {
        self.0.index_mut(_index)
    }
}

impl CPool {
    pub fn insert(&mut self, cap: Capability) {
        for space in self.0.iter_mut() {
            if space.is_none() {
                *space = Some(cap);
                return;
            }
        }
        assert!(false);
    }

    pub fn slice(&self) -> &[Option<Capability>] {
        &self.0
    }

    pub fn slice_mut(&mut self) -> &mut [Option<Capability>] {
        &mut self.0
    }
}

impl CPoolHalf {
    // TODO handle data races
    pub fn with_cpool<Return, F: FnOnce(&CPool) -> Return>(&self, f: F) -> Return {
        unsafe {
            arch::with_object(self.start_paddr, |cpool: &CPool| {
                f(cpool)
            })
        }
    }

    // TODO handle data races
    pub fn with_cpool_mut<Return, F: FnOnce(&mut CPool) -> Return>(&mut self, f: F) -> Return {
        unsafe {
            arch::with_object_mut(self.start_paddr, |cpool: &mut CPool| {
                f(cpool)
            })
        }
    }

    pub fn new(untyped: &mut UntypedHalf) -> CPoolHalf {
        let alignment = align_of::<CPool>();
        let length = size_of::<CPool>();
        let start_paddr = untyped.allocate(length, alignment);

        let mut cap = CPoolHalf {
            start_paddr: start_paddr
        };

        cap.with_cpool_mut(|cpool| {
            *cpool = CPool([None, None, None, None, None, None, None, None,
                            None, None, None, None, None, None, None, None,
                            None, None, None, None, None, None, None, None,
                            None, None, None, None, None, None, None, None]);
        });

        cap
    }
}
<commit_msg>change cpool to length 16 for caddressing<commit_after>use common::*;
use super::{Capability};
use super::untyped::{UntypedHalf};
use core::mem::{size_of, align_of};
use core::ops::{Index, IndexMut};
use core::slice::Iter;
use arch;

#[derive(Debug, Clone)]
pub struct CPoolHalf {
    start_paddr: PAddr
}

#[derive(Debug)]
pub struct CPool([Option<Capability>; 16]);

impl Index<usize> for CPool {
    type Output = Option<Capability>;

    fn index<'a>(&'a self, _index: usize) -> &'a Option<Capability> {
        self.0.index(_index)
    }
}

impl IndexMut<usize> for CPool {
    fn index_mut<'a>(&'a mut self, _index: usize) -> &'a mut Option<Capability> {
        self.0.index_mut(_index)
    }
}

impl CPool {
    pub fn insert(&mut self, cap: Capability) {
        for space in self.0.iter_mut() {
            if space.is_none() {
                *space = Some(cap);
                return;
            }
        }
        assert!(false);
    }

    pub fn slice(&self) -> &[Option<Capability>] {
        &self.0
    }

    pub fn slice_mut(&mut self) -> &mut [Option<Capability>] {
        &mut self.0
    }
}

impl CPoolHalf {
    // TODO handle data races
    pub fn with_cpool<Return, F: FnOnce(&CPool) -> Return>(&self, f: F) -> Return {
        unsafe {
            arch::with_object(self.start_paddr, |cpool: &CPool| {
                f(cpool)
            })
        }
    }

    // TODO handle data races
    pub fn with_cpool_mut<Return, F: FnOnce(&mut CPool) -> Return>(&mut self, f: F) -> Return {
        unsafe {
            arch::with_object_mut(self.start_paddr, |cpool: &mut CPool| {
                f(cpool)
            })
        }
    }

    pub fn new(untyped: &mut UntypedHalf) -> CPoolHalf {
        let alignment = align_of::<CPool>();
        let length = size_of::<CPool>();
        let start_paddr = untyped.allocate(length, alignment);

        let mut cap = CPoolHalf {
            start_paddr: start_paddr
        };

        cap.with_cpool_mut(|cpool| {
            *cpool = CPool([None, None, None, None, None, None, None, None,
                            None, None, None, None, None, None, None, None]);
        });

        cap
    }
}
<|endoftext|>


<commit_before>//! Defines a series of convenience modifiers for editing Responses
//!
//! Modifiers can be used to edit Responses through the owning
//! method `set` or the mutating `set_mut`, both of which are
//! defined in the `Set` trait.
//!
//! Instead of having a combinatorial explosion of Response methods
//! and constructors, this provides a series of modifiers that
//! can be used through the `Set` trait.
//!
//! For instance, instead of `Response::redirect` constructing a
//! redirect response, we provide a `Redirect` modifier, so you
//! can just do:
//!
//! ```
//! # use iron::prelude::*;
//! # use iron::status;
//! # use iron::modifiers::Redirect;
//! # use iron::Url;
//!
//! let url = Url::parse("http://doc.rust-lang.org").unwrap();
//! Response::with((status::Found, Redirect(url)));
//! ```
//!
//! This is more extensible as it allows you to combine
//! arbitrary modifiers without having a massive number of
//! Response constructors.
//!
//! For more information about the modifier system, see
//! [rust-modifier](https://github.com/reem/rust-modifier).

use std::fs::File;
use std::path::{Path, PathBuf};

use modifier::Modifier;

use hyper::mime::Mime;

use {status, headers, Request, Response, Set, Url};

use mime_types;
use response::WriteBody;

lazy_static! {
    static ref MIME_TYPES: mime_types::Types = mime_types::Types::new().unwrap();
}

impl Modifier<Response> for Mime {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.headers.set(headers::ContentType(self))
    }
}

impl Modifier<Response> for Box<WriteBody + Send> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.body = Some(self);
    }
}

impl Modifier<Response> for String {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.into_bytes().modify(res);
    }
}

impl Modifier<Response> for Vec<u8> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.headers.set(headers::ContentLength(self.len() as u64));
        res.body = Some(Box::new(self));
    }
}

impl<'a> Modifier<Response> for &'a str {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.to_owned().modify(res);
    }
}

impl<'a> Modifier<Response> for &'a [u8] {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.to_vec().modify(res);
    }
}

impl Modifier<Response> for File {
    fn modify(self, res: &mut Response) {
        // Set the content type based on the file extension if a path is available.
        if let Ok(metadata) = self.metadata() {
            res.headers.set(headers::ContentLength(metadata.len()));
        }

        res.body = Some(Box::new(self));
    }
}

impl<'a> Modifier<Response> for &'a Path {
    /// Set the body to the contents of the File at this path.
    ///
    /// ## Panics
    ///
    /// Panics if there is no file at the passed-in Path.
    fn modify(self, res: &mut Response) {
        File::open(self)
            .expect(&format!("No such file: {}", self.display()))
            .modify(res);

        let mime_str = MIME_TYPES.mime_for_path(self);
        let _ = mime_str.parse().map(|mime: Mime| res.set_mut(mime));
    }
}

impl Modifier<Response> for PathBuf {
    /// Set the body to the contents of the File at this path.
    ///
    /// ## Panics
    ///
    /// Panics if there is no file at the passed-in Path.
    fn modify(self, res: &mut Response) {
        File::open(&self)
            .expect(&format!("No such file: {}", self.display()))
            .modify(res);
    }
}

impl Modifier<Response> for status::Status {
    fn modify(self, res: &mut Response) {
        res.status = Some(self);
    }
}

/// A modifier for changing headers on requests and responses.
pub struct Header<H: headers::Header + headers::HeaderFormat>(pub H);

impl<H> Modifier<Response> for Header<H>
where H: headers::Header + headers::HeaderFormat {
    fn modify(self, res: &mut Response) {
        res.headers.set(self.0);
    }
}

impl<'a, 'b, H> Modifier<Request<'a, 'b>> for Header<H>
where H: headers::Header + headers::HeaderFormat {
    fn modify(self, res: &mut Request) {
        res.headers.set(self.0);
    }
}

/// A modifier for creating redirect responses.
pub struct Redirect(pub Url);

impl Modifier<Response> for Redirect {
    fn modify(self, res: &mut Response) {
        let Redirect(url) = self;
        res.headers.set(headers::Location(url.to_string()));
    }
}
<commit_msg>add a new modifier RedirectRaw .<commit_after>//! Defines a series of convenience modifiers for editing Responses
//!
//! Modifiers can be used to edit Responses through the owning
//! method `set` or the mutating `set_mut`, both of which are
//! defined in the `Set` trait.
//!
//! Instead of having a combinatorial explosion of Response methods
//! and constructors, this provides a series of modifiers that
//! can be used through the `Set` trait.
//!
//! For instance, instead of `Response::redirect` constructing a
//! redirect response, we provide a `Redirect` modifier, so you
//! can just do:
//!
//! ```
//! # use iron::prelude::*;
//! # use iron::status;
//! # use iron::modifiers::Redirect;
//! # use iron::Url;
//!
//! let url = Url::parse("http://doc.rust-lang.org").unwrap();
//! Response::with((status::Found, Redirect(url)));
//! ```
//!
//! This is more extensible as it allows you to combine
//! arbitrary modifiers without having a massive number of
//! Response constructors.
//!
//! For more information about the modifier system, see
//! [rust-modifier](https://github.com/reem/rust-modifier).

use std::fs::File;
use std::path::{Path, PathBuf};

use modifier::Modifier;

use hyper::mime::Mime;

use {status, headers, Request, Response, Set, Url};

use mime_types;
use response::WriteBody;

lazy_static! {
    static ref MIME_TYPES: mime_types::Types = mime_types::Types::new().unwrap();
}

impl Modifier<Response> for Mime {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.headers.set(headers::ContentType(self))
    }
}

impl Modifier<Response> for Box<WriteBody + Send> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.body = Some(self);
    }
}

impl Modifier<Response> for String {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.into_bytes().modify(res);
    }
}

impl Modifier<Response> for Vec<u8> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.headers.set(headers::ContentLength(self.len() as u64));
        res.body = Some(Box::new(self));
    }
}

impl<'a> Modifier<Response> for &'a str {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.to_owned().modify(res);
    }
}

impl<'a> Modifier<Response> for &'a [u8] {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.to_vec().modify(res);
    }
}

impl Modifier<Response> for File {
    fn modify(self, res: &mut Response) {
        // Set the content type based on the file extension if a path is available.
        if let Ok(metadata) = self.metadata() {
            res.headers.set(headers::ContentLength(metadata.len()));
        }

        res.body = Some(Box::new(self));
    }
}

impl<'a> Modifier<Response> for &'a Path {
    /// Set the body to the contents of the File at this path.
    ///
    /// ## Panics
    ///
    /// Panics if there is no file at the passed-in Path.
    fn modify(self, res: &mut Response) {
        File::open(self)
            .expect(&format!("No such file: {}", self.display()))
            .modify(res);

        let mime_str = MIME_TYPES.mime_for_path(self);
        let _ = mime_str.parse().map(|mime: Mime| res.set_mut(mime));
    }
}

impl Modifier<Response> for PathBuf {
    /// Set the body to the contents of the File at this path.
    ///
    /// ## Panics
    ///
    /// Panics if there is no file at the passed-in Path.
    fn modify(self, res: &mut Response) {
        File::open(&self)
            .expect(&format!("No such file: {}", self.display()))
            .modify(res);
    }
}

impl Modifier<Response> for status::Status {
    fn modify(self, res: &mut Response) {
        res.status = Some(self);
    }
}

/// A modifier for changing headers on requests and responses.
pub struct Header<H: headers::Header + headers::HeaderFormat>(pub H);

impl<H> Modifier<Response> for Header<H>
where H: headers::Header + headers::HeaderFormat {
    fn modify(self, res: &mut Response) {
        res.headers.set(self.0);
    }
}

impl<'a, 'b, H> Modifier<Request<'a, 'b>> for Header<H>
where H: headers::Header + headers::HeaderFormat {
    fn modify(self, res: &mut Request) {
        res.headers.set(self.0);
    }
}

/// A modifier for creating redirect responses.
pub struct Redirect(pub Url);

impl Modifier<Response> for Redirect {
    fn modify(self, res: &mut Response) {
        let Redirect(url) = self;
        res.headers.set(headers::Location(url.to_string()));
    }
}

/// A modifier for creating redirect responses.
pub struct RedirectRaw(String);

impl Modifier<Response> for RedirectRaw {
    fn modify(self, res: &mut Response) {
        let RedirectRaw(path) = self;
        res.headers.set(headers::Location(path));
    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2014 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.

pub use self::VarValue::*;

use std::marker;

use middle::ty::{expected_found, IntVarValue};
use middle::ty::{self, Ty};
use std::fmt::Debug;
use std::marker::PhantomData;
use syntax::ast;
use util::snapshot_vec as sv;

/// This trait is implemented by any type that can serve as a type
/// variable. We call such variables *unification keys*. For example,
/// this trait is implemented by `IntVid`, which represents integral
/// variables.
///
/// Each key type has an associated value type `V`. For example, for
/// `IntVid`, this is `Option<IntVarValue>`, representing some
/// (possibly not yet known) sort of integer.
///
/// Implementations of this trait are at the end of this file.
pub trait UnifyKey : Clone + Debug + PartialEq {
    type Value : UnifyValue;

    fn index(&self) -> u32;

    fn from_index(u: u32) -> Self;

    fn tag(k: Option<Self>) -> &'static str;
}

/// Trait for valid types that a type variable can be set to. Note that
/// this is typically not the end type that the value will take on, but
/// rather an `Option` wrapper (where `None` represents a variable
/// whose value is not yet set).
///
/// Implementations of this trait are at the end of this file.
pub trait UnifyValue : Clone + PartialEq + Debug {
}

/// Value of a unification key. We implement Tarjan's union-find
/// algorithm: when two keys are unified, one of them is converted
/// into a "redirect" pointing at the other. These redirects form a
/// DAG: the roots of the DAG (nodes that are not redirected) are each
/// associated with a value of type `V` and a rank. The rank is used
/// to keep the DAG relatively balanced, which helps keep the running
/// time of the algorithm under control. For more information, see
/// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
#[derive(PartialEq,Clone,Debug)]
pub enum VarValue<K:UnifyKey> {
    Redirect(K),
    Root(K::Value, usize),
}

/// Table of unification keys and their values.
pub struct UnificationTable<K:UnifyKey> {
    /// Indicates the current value of each key.
    values: sv::SnapshotVec<Delegate<K>>,
}

/// At any time, users may snapshot a unification table.  The changes
/// made during the snapshot may either be *committed* or *rolled back*.
pub struct Snapshot<K:UnifyKey> {
    // Link snapshot to the key type `K` of the table.
    marker: marker::PhantomData<K>,
    snapshot: sv::Snapshot,
}

/// Internal type used to represent the result of a `get()` operation.
/// Conveys the current root and value of the key.
pub struct Node<K:UnifyKey> {
    pub key: K,
    pub value: K::Value,
    pub rank: usize,
}

#[derive(Copy)]
pub struct Delegate<K>(PhantomData<K>);

// We can't use V:LatticeValue, much as I would like to,
// because frequently the pattern is that V=Option<U> for some
// other type parameter U, and we have no way to say
// Option<U>:LatticeValue.

impl<K:UnifyKey> UnificationTable<K> {
    pub fn new() -> UnificationTable<K> {
        UnificationTable {
            values: sv::SnapshotVec::new(Delegate(PhantomData)),
        }
    }

    /// Starts a new snapshot. Each snapshot must be either
    /// rolled back or committed in a "LIFO" (stack) order.
    pub fn snapshot(&mut self) -> Snapshot<K> {
        Snapshot { marker: marker::PhantomData::<K>,
                   snapshot: self.values.start_snapshot() }
    }

    /// Reverses all changes since the last snapshot. Also
    /// removes any keys that have been created since then.
    pub fn rollback_to(&mut self, snapshot: Snapshot<K>) {
        debug!("{}: rollback_to()", UnifyKey::tag(None::<K>));
        self.values.rollback_to(snapshot.snapshot);
    }

    /// Commits all changes since the last snapshot. Of course, they
    /// can still be undone if there is a snapshot further out.
    pub fn commit(&mut self, snapshot: Snapshot<K>) {
        debug!("{}: commit()", UnifyKey::tag(None::<K>));
        self.values.commit(snapshot.snapshot);
    }

    pub fn new_key(&mut self, value: K::Value) -> K {
        let index = self.values.push(Root(value, 0));
        let k = UnifyKey::from_index(index as u32);
        debug!("{}: created new key: {:?}",
               UnifyKey::tag(None::<K>),
               k);
        k
    }

    /// Find the root node for `vid`. This uses the standard
    /// union-find algorithm with path compression:
    /// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
    ///
    /// NB. This is a building-block operation and you would probably
    /// prefer to call `probe` below.
    fn get(&mut self, vid: K) -> Node<K> {
        let index = vid.index() as usize;
        let value = (*self.values.get(index)).clone();
        match value {
            Redirect(redirect) => {
                let node: Node<K> = self.get(redirect.clone());
                if node.key != redirect {
                    // Path compression
                    self.values.set(index, Redirect(node.key.clone()));
                }
                node
            }
            Root(value, rank) => {
                Node { key: vid, value: value, rank: rank }
            }
        }
    }

    fn is_root(&self, key: &K) -> bool {
        let index = key.index() as usize;
        match *self.values.get(index) {
            Redirect(..) => false,
            Root(..) => true,
        }
    }

    /// Sets the value for `vid` to `new_value`. `vid` MUST be a root
    /// node! This is an internal operation used to impl other things.
    fn set(&mut self, key: K, new_value: VarValue<K>) {
        assert!(self.is_root(&key));

        debug!("Updating variable {:?} to {:?}",
               key, new_value);

        let index = key.index() as usize;
        self.values.set(index, new_value);
    }

    /// Either redirects `node_a` to `node_b` or vice versa, depending
    /// on the relative rank. The value associated with the new root
    /// will be `new_value`.
    ///
    /// NB: This is the "union" operation of "union-find". It is
    /// really more of a building block. If the values associated with
    /// your key are non-trivial, you would probably prefer to call
    /// `unify_var_var` below.
    fn unify(&mut self, node_a: &Node<K>, node_b: &Node<K>, new_value: K::Value) {
        debug!("unify(node_a(id={:?}, rank={:?}), node_b(id={:?}, rank={:?}))",
               node_a.key,
               node_a.rank,
               node_b.key,
               node_b.rank);

        let (new_root, new_rank) = if node_a.rank > node_b.rank {
            // a has greater rank, so a should become b's parent,
            // i.e., b should redirect to a.
            self.set(node_b.key.clone(), Redirect(node_a.key.clone()));
            (node_a.key.clone(), node_a.rank)
        } else if node_a.rank < node_b.rank {
            // b has greater rank, so a should redirect to b.
            self.set(node_a.key.clone(), Redirect(node_b.key.clone()));
            (node_b.key.clone(), node_b.rank)
        } else {
            // If equal, redirect one to the other and increment the
            // other's rank.
            assert_eq!(node_a.rank, node_b.rank);
            self.set(node_b.key.clone(), Redirect(node_a.key.clone()));
            (node_a.key.clone(), node_a.rank + 1)
        };

        self.set(new_root, Root(new_value, new_rank));
    }
}

impl<K:UnifyKey> sv::SnapshotVecDelegate for Delegate<K> {
    type Value = VarValue<K>;
    type Undo = ();

    fn reverse(&mut self, _: &mut Vec<VarValue<K>>, _: ()) {
        panic!("Nothing to reverse");
    }
}

///////////////////////////////////////////////////////////////////////////
// Code to handle keys which carry a value, like ints,
// floats---anything that doesn't have a subtyping relationship we
// need to worry about.

impl<'tcx,K,V> UnificationTable<K>
    where K: UnifyKey<Value=Option<V>>,
          V: Clone+PartialEq,
          Option<V>: UnifyValue,
{
    pub fn unify_var_var(&mut self,
                         a_id: K,
                         b_id: K)
                         -> Result<(),(V,V)>
    {
        let node_a = self.get(a_id);
        let node_b = self.get(b_id);
        let a_id = node_a.key.clone();
        let b_id = node_b.key.clone();

        if a_id == b_id { return Ok(()); }

        let combined = {
            match (&node_a.value, &node_b.value) {
                (&None, &None) => {
                    None
                }
                (&Some(ref v), &None) | (&None, &Some(ref v)) => {
                    Some(v.clone())
                }
                (&Some(ref v1), &Some(ref v2)) => {
                    if *v1 != *v2 {
                        return Err((v1.clone(), v2.clone()));
                    }
                    Some(v1.clone())
                }
            }
        };

        Ok(self.unify(&node_a, &node_b, combined))
    }

    /// Sets the value of the key `a_id` to `b`. Because simple keys do not have any subtyping
    /// relationships, if `a_id` already has a value, it must be the same as `b`.
    pub fn unify_var_value(&mut self,
                           a_id: K,
                           b: V)
                           -> Result<(),(V,V)>
    {
        let node_a = self.get(a_id);
        let a_id = node_a.key.clone();

        match node_a.value {
            None => {
                self.set(a_id, Root(Some(b), node_a.rank));
                Ok(())
            }

            Some(ref a_t) => {
                if *a_t == b {
                    Ok(())
                } else {
                    Err((a_t.clone(), b))
                }
            }
        }
    }

    pub fn has_value(&mut self, id: K) -> bool {
        self.get(id).value.is_some()
    }

    pub fn probe(&mut self, a_id: K) -> Option<V> {
        self.get(a_id).value.clone()
    }
}

///////////////////////////////////////////////////////////////////////////

// Integral type keys

pub trait ToType<'tcx> {
    fn to_type(&self, tcx: &ty::ctxt<'tcx>) -> Ty<'tcx>;
}

impl UnifyKey for ty::IntVid {
    type Value = Option<IntVarValue>;
    fn index(&self) -> u32 { self.index }
    fn from_index(i: u32) -> ty::IntVid { ty::IntVid { index: i } }
    fn tag(_: Option<ty::IntVid>) -> &'static str { "IntVid" }
}

impl<'tcx> ToType<'tcx> for IntVarValue {
    fn to_type(&self, tcx: &ty::ctxt<'tcx>) -> Ty<'tcx> {
        match *self {
            ty::IntType(i) => ty::mk_mach_int(tcx, i),
            ty::UintType(i) => ty::mk_mach_uint(tcx, i),
        }
    }
}

impl UnifyValue for Option<IntVarValue> { }

// Floating point type keys

impl UnifyKey for ty::FloatVid {
    type Value = Option<ast::FloatTy>;
    fn index(&self) -> u32 { self.index }
    fn from_index(i: u32) -> ty::FloatVid { ty::FloatVid { index: i } }
    fn tag(_: Option<ty::FloatVid>) -> &'static str { "FloatVid" }
}

impl UnifyValue for Option<ast::FloatTy> {
}

impl<'tcx> ToType<'tcx> for ast::FloatTy {
    fn to_type(&self, tcx: &ty::ctxt<'tcx>) -> Ty<'tcx> {
        ty::mk_mach_float(tcx, *self)
    }
}
<commit_msg>Remove unused import<commit_after>// Copyright 2012-2014 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.

pub use self::VarValue::*;

use std::marker;

use middle::ty::{IntVarValue};
use middle::ty::{self, Ty};
use std::fmt::Debug;
use std::marker::PhantomData;
use syntax::ast;
use util::snapshot_vec as sv;

/// This trait is implemented by any type that can serve as a type
/// variable. We call such variables *unification keys*. For example,
/// this trait is implemented by `IntVid`, which represents integral
/// variables.
///
/// Each key type has an associated value type `V`. For example, for
/// `IntVid`, this is `Option<IntVarValue>`, representing some
/// (possibly not yet known) sort of integer.
///
/// Implementations of this trait are at the end of this file.
pub trait UnifyKey : Clone + Debug + PartialEq {
    type Value : UnifyValue;

    fn index(&self) -> u32;

    fn from_index(u: u32) -> Self;

    fn tag(k: Option<Self>) -> &'static str;
}

/// Trait for valid types that a type variable can be set to. Note that
/// this is typically not the end type that the value will take on, but
/// rather an `Option` wrapper (where `None` represents a variable
/// whose value is not yet set).
///
/// Implementations of this trait are at the end of this file.
pub trait UnifyValue : Clone + PartialEq + Debug {
}

/// Value of a unification key. We implement Tarjan's union-find
/// algorithm: when two keys are unified, one of them is converted
/// into a "redirect" pointing at the other. These redirects form a
/// DAG: the roots of the DAG (nodes that are not redirected) are each
/// associated with a value of type `V` and a rank. The rank is used
/// to keep the DAG relatively balanced, which helps keep the running
/// time of the algorithm under control. For more information, see
/// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
#[derive(PartialEq,Clone,Debug)]
pub enum VarValue<K:UnifyKey> {
    Redirect(K),
    Root(K::Value, usize),
}

/// Table of unification keys and their values.
pub struct UnificationTable<K:UnifyKey> {
    /// Indicates the current value of each key.
    values: sv::SnapshotVec<Delegate<K>>,
}

/// At any time, users may snapshot a unification table.  The changes
/// made during the snapshot may either be *committed* or *rolled back*.
pub struct Snapshot<K:UnifyKey> {
    // Link snapshot to the key type `K` of the table.
    marker: marker::PhantomData<K>,
    snapshot: sv::Snapshot,
}

/// Internal type used to represent the result of a `get()` operation.
/// Conveys the current root and value of the key.
pub struct Node<K:UnifyKey> {
    pub key: K,
    pub value: K::Value,
    pub rank: usize,
}

#[derive(Copy)]
pub struct Delegate<K>(PhantomData<K>);

// We can't use V:LatticeValue, much as I would like to,
// because frequently the pattern is that V=Option<U> for some
// other type parameter U, and we have no way to say
// Option<U>:LatticeValue.

impl<K:UnifyKey> UnificationTable<K> {
    pub fn new() -> UnificationTable<K> {
        UnificationTable {
            values: sv::SnapshotVec::new(Delegate(PhantomData)),
        }
    }

    /// Starts a new snapshot. Each snapshot must be either
    /// rolled back or committed in a "LIFO" (stack) order.
    pub fn snapshot(&mut self) -> Snapshot<K> {
        Snapshot { marker: marker::PhantomData::<K>,
                   snapshot: self.values.start_snapshot() }
    }

    /// Reverses all changes since the last snapshot. Also
    /// removes any keys that have been created since then.
    pub fn rollback_to(&mut self, snapshot: Snapshot<K>) {
        debug!("{}: rollback_to()", UnifyKey::tag(None::<K>));
        self.values.rollback_to(snapshot.snapshot);
    }

    /// Commits all changes since the last snapshot. Of course, they
    /// can still be undone if there is a snapshot further out.
    pub fn commit(&mut self, snapshot: Snapshot<K>) {
        debug!("{}: commit()", UnifyKey::tag(None::<K>));
        self.values.commit(snapshot.snapshot);
    }

    pub fn new_key(&mut self, value: K::Value) -> K {
        let index = self.values.push(Root(value, 0));
        let k = UnifyKey::from_index(index as u32);
        debug!("{}: created new key: {:?}",
               UnifyKey::tag(None::<K>),
               k);
        k
    }

    /// Find the root node for `vid`. This uses the standard
    /// union-find algorithm with path compression:
    /// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>.
    ///
    /// NB. This is a building-block operation and you would probably
    /// prefer to call `probe` below.
    fn get(&mut self, vid: K) -> Node<K> {
        let index = vid.index() as usize;
        let value = (*self.values.get(index)).clone();
        match value {
            Redirect(redirect) => {
                let node: Node<K> = self.get(redirect.clone());
                if node.key != redirect {
                    // Path compression
                    self.values.set(index, Redirect(node.key.clone()));
                }
                node
            }
            Root(value, rank) => {
                Node { key: vid, value: value, rank: rank }
            }
        }
    }

    fn is_root(&self, key: &K) -> bool {
        let index = key.index() as usize;
        match *self.values.get(index) {
            Redirect(..) => false,
            Root(..) => true,
        }
    }

    /// Sets the value for `vid` to `new_value`. `vid` MUST be a root
    /// node! This is an internal operation used to impl other things.
    fn set(&mut self, key: K, new_value: VarValue<K>) {
        assert!(self.is_root(&key));

        debug!("Updating variable {:?} to {:?}",
               key, new_value);

        let index = key.index() as usize;
        self.values.set(index, new_value);
    }

    /// Either redirects `node_a` to `node_b` or vice versa, depending
    /// on the relative rank. The value associated with the new root
    /// will be `new_value`.
    ///
    /// NB: This is the "union" operation of "union-find". It is
    /// really more of a building block. If the values associated with
    /// your key are non-trivial, you would probably prefer to call
    /// `unify_var_var` below.
    fn unify(&mut self, node_a: &Node<K>, node_b: &Node<K>, new_value: K::Value) {
        debug!("unify(node_a(id={:?}, rank={:?}), node_b(id={:?}, rank={:?}))",
               node_a.key,
               node_a.rank,
               node_b.key,
               node_b.rank);

        let (new_root, new_rank) = if node_a.rank > node_b.rank {
            // a has greater rank, so a should become b's parent,
            // i.e., b should redirect to a.
            self.set(node_b.key.clone(), Redirect(node_a.key.clone()));
            (node_a.key.clone(), node_a.rank)
        } else if node_a.rank < node_b.rank {
            // b has greater rank, so a should redirect to b.
            self.set(node_a.key.clone(), Redirect(node_b.key.clone()));
            (node_b.key.clone(), node_b.rank)
        } else {
            // If equal, redirect one to the other and increment the
            // other's rank.
            assert_eq!(node_a.rank, node_b.rank);
            self.set(node_b.key.clone(), Redirect(node_a.key.clone()));
            (node_a.key.clone(), node_a.rank + 1)
        };

        self.set(new_root, Root(new_value, new_rank));
    }
}

impl<K:UnifyKey> sv::SnapshotVecDelegate for Delegate<K> {
    type Value = VarValue<K>;
    type Undo = ();

    fn reverse(&mut self, _: &mut Vec<VarValue<K>>, _: ()) {
        panic!("Nothing to reverse");
    }
}

///////////////////////////////////////////////////////////////////////////
// Code to handle keys which carry a value, like ints,
// floats---anything that doesn't have a subtyping relationship we
// need to worry about.

impl<'tcx,K,V> UnificationTable<K>
    where K: UnifyKey<Value=Option<V>>,
          V: Clone+PartialEq,
          Option<V>: UnifyValue,
{
    pub fn unify_var_var(&mut self,
                         a_id: K,
                         b_id: K)
                         -> Result<(),(V,V)>
    {
        let node_a = self.get(a_id);
        let node_b = self.get(b_id);
        let a_id = node_a.key.clone();
        let b_id = node_b.key.clone();

        if a_id == b_id { return Ok(()); }

        let combined = {
            match (&node_a.value, &node_b.value) {
                (&None, &None) => {
                    None
                }
                (&Some(ref v), &None) | (&None, &Some(ref v)) => {
                    Some(v.clone())
                }
                (&Some(ref v1), &Some(ref v2)) => {
                    if *v1 != *v2 {
                        return Err((v1.clone(), v2.clone()));
                    }
                    Some(v1.clone())
                }
            }
        };

        Ok(self.unify(&node_a, &node_b, combined))
    }

    /// Sets the value of the key `a_id` to `b`. Because simple keys do not have any subtyping
    /// relationships, if `a_id` already has a value, it must be the same as `b`.
    pub fn unify_var_value(&mut self,
                           a_id: K,
                           b: V)
                           -> Result<(),(V,V)>
    {
        let node_a = self.get(a_id);
        let a_id = node_a.key.clone();

        match node_a.value {
            None => {
                self.set(a_id, Root(Some(b), node_a.rank));
                Ok(())
            }

            Some(ref a_t) => {
                if *a_t == b {
                    Ok(())
                } else {
                    Err((a_t.clone(), b))
                }
            }
        }
    }

    pub fn has_value(&mut self, id: K) -> bool {
        self.get(id).value.is_some()
    }

    pub fn probe(&mut self, a_id: K) -> Option<V> {
        self.get(a_id).value.clone()
    }
}

///////////////////////////////////////////////////////////////////////////

// Integral type keys

pub trait ToType<'tcx> {
    fn to_type(&self, tcx: &ty::ctxt<'tcx>) -> Ty<'tcx>;
}

impl UnifyKey for ty::IntVid {
    type Value = Option<IntVarValue>;
    fn index(&self) -> u32 { self.index }
    fn from_index(i: u32) -> ty::IntVid { ty::IntVid { index: i } }
    fn tag(_: Option<ty::IntVid>) -> &'static str { "IntVid" }
}

impl<'tcx> ToType<'tcx> for IntVarValue {
    fn to_type(&self, tcx: &ty::ctxt<'tcx>) -> Ty<'tcx> {
        match *self {
            ty::IntType(i) => ty::mk_mach_int(tcx, i),
            ty::UintType(i) => ty::mk_mach_uint(tcx, i),
        }
    }
}

impl UnifyValue for Option<IntVarValue> { }

// Floating point type keys

impl UnifyKey for ty::FloatVid {
    type Value = Option<ast::FloatTy>;
    fn index(&self) -> u32 { self.index }
    fn from_index(i: u32) -> ty::FloatVid { ty::FloatVid { index: i } }
    fn tag(_: Option<ty::FloatVid>) -> &'static str { "FloatVid" }
}

impl UnifyValue for Option<ast::FloatTy> {
}

impl<'tcx> ToType<'tcx> for ast::FloatTy {
    fn to_type(&self, tcx: &ty::ctxt<'tcx>) -> Ty<'tcx> {
        ty::mk_mach_float(tcx, *self)
    }
}
<|endoftext|>


<commit_before>// Copyright 2014-2016 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.

// FIXME: This needs an audit for correctness and completeness.

use abi::{FnType, ArgType, LayoutExt, Reg, RegKind, Uniform};
use context::CodegenCx;

fn is_homogeneous_aggregate<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>, arg: &mut ArgType<'tcx>)
                                     -> Option<Uniform> {
    arg.layout.homogeneous_aggregate(cx).and_then(|unit| {
        // Ensure we have at most eight uniquely addressable members.
        if arg.layout.size > unit.size.checked_mul(8, cx).unwrap() {
            return None;
        }

        let valid_unit = match unit.kind {
            RegKind::Integer => false,
            RegKind::Float => true,
            RegKind::Vector => arg.layout.size.bits() == 256
        };

        if valid_unit {
            Some(Uniform {
                unit,
                total: arg.layout.size
            })
        } else {
            None
        }
    })
}

fn classify_ret_ty<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>, ret: &mut ArgType<'tcx>) {
    if !ret.layout.is_aggregate() {
        ret.extend_integer_width_to(64);
        return;
    }

    if let Some(uniform) = is_homogeneous_aggregate(cx, ret) {
        ret.cast_to(uniform);
        return;
    }
    let size = ret.layout.size;
    let bits = size.bits();
    if bits <= 256 {
        let unit = if bits <= 8 {
            Reg::i8()
        } else if bits <= 16 {
            Reg::i16()
        } else if bits <= 32 {
            Reg::i32()
        } else {
            Reg::i64()
        };

        ret.cast_to(Uniform {
            unit,
            total: size
        });
        return;
    }

    // don't return aggregates in registers
    ret.make_indirect();
}

fn classify_arg_ty<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>, arg: &mut ArgType<'tcx>) {
    if !arg.layout.is_aggregate() {
        arg.extend_integer_width_to(64);
        return;
    }

    if let Some(uniform) = is_homogeneous_aggregate(cx, arg) {
        arg.cast_to(uniform);
        return;
    }

    let total = arg.layout.size;
    if total.bits() > 128 {
        arg.make_indirect(cx);
        return;
    }

    arg.cast_to(Uniform {
        unit: Reg::i64(),
        total
    });
}

pub fn compute_abi_info<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>, fty: &mut FnType<'tcx>) {
    if !fty.ret.is_ignore() {
        classify_ret_ty(cx, &mut fty.ret);
    }

    for arg in &mut fty.args {
        if arg.is_ignore() { continue; }
        classify_arg_ty(cx, arg);
    }
}
<commit_msg>Removed uneeded change.<commit_after>// Copyright 2014-2016 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.

// FIXME: This needs an audit for correctness and completeness.

use abi::{FnType, ArgType, LayoutExt, Reg, RegKind, Uniform};
use context::CodegenCx;

fn is_homogeneous_aggregate<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>, arg: &mut ArgType<'tcx>)
                                     -> Option<Uniform> {
    arg.layout.homogeneous_aggregate(cx).and_then(|unit| {
        // Ensure we have at most eight uniquely addressable members.
        if arg.layout.size > unit.size.checked_mul(8, cx).unwrap() {
            return None;
        }

        let valid_unit = match unit.kind {
            RegKind::Integer => false,
            RegKind::Float => true,
            RegKind::Vector => arg.layout.size.bits() == 128
        };

        if valid_unit {
            Some(Uniform {
                unit,
                total: arg.layout.size
            })
        } else {
            None
        }
    })
}

fn classify_ret_ty<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>, ret: &mut ArgType<'tcx>) {
    if !ret.layout.is_aggregate() {
        ret.extend_integer_width_to(64);
        return;
    }

    if let Some(uniform) = is_homogeneous_aggregate(cx, ret) {
        ret.cast_to(uniform);
        return;
    }
    let size = ret.layout.size;
    let bits = size.bits();
    if bits <= 256 {
        let unit = if bits <= 8 {
            Reg::i8()
        } else if bits <= 16 {
            Reg::i16()
        } else if bits <= 32 {
            Reg::i32()
        } else {
            Reg::i64()
        };

        ret.cast_to(Uniform {
            unit,
            total: size
        });
        return;
    }

    // don't return aggregates in registers
    ret.make_indirect();
}

fn classify_arg_ty<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>, arg: &mut ArgType<'tcx>) {
    if !arg.layout.is_aggregate() {
        arg.extend_integer_width_to(64);
        return;
    }

    if let Some(uniform) = is_homogeneous_aggregate(cx, arg) {
        arg.cast_to(uniform);
        return;
    }

    let total = arg.layout.size;
    if total.bits() > 128 {
        arg.make_indirect(cx);
        return;
    }

    arg.cast_to(Uniform {
        unit: Reg::i64(),
        total
    });
}

pub fn compute_abi_info<'a, 'tcx>(cx: &CodegenCx<'a, 'tcx>, fty: &mut FnType<'tcx>) {
    if !fty.ret.is_ignore() {
        classify_ret_ty(cx, &mut fty.ret);
    }

    for arg in &mut fty.args {
        if arg.is_ignore() { continue; }
        classify_arg_ty(cx, arg);
    }
}
<|endoftext|>