<commit_before>//! Visitor for a run-time value with a given layout: Traverse enums, structs and other compound
//! types until we arrive at the leaves, with custom handling for primitive types.

use rustc::ty::layout::{self, TyLayout, VariantIdx};
use rustc::ty;
use rustc::mir::interpret::{
    EvalResult,
};

use super::{
    Machine, EvalContext, MPlaceTy, OpTy,
};

// A thing that we can project into, and that has a layout.
// This wouldn't have to depend on `Machine` but with the current type inference,
// that's just more convenient to work with (avoids repeating all the `Machine` bounds).
pub trait Value<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>>: Copy
{
    /// Get this value's layout.
    fn layout(&self) -> TyLayout<'tcx>;

    /// Make this into an `OpTy`.
    fn to_op(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, OpTy<'tcx, M::PointerTag>>;

    /// Create this from an `MPlaceTy`.
    fn from_mem_place(MPlaceTy<'tcx, M::PointerTag>) -> Self;

    /// Project to the given enum variant.
    fn project_downcast(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        variant: VariantIdx,
    ) -> EvalResult<'tcx, Self>;

    /// Project to the n-th field.
    fn project_field(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self>;
}

// Operands and memory-places are both values.
// Places in general are not due to `place_field` having to do `force_allocation`.
impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Value<'a, 'mir, 'tcx, M>
    for OpTy<'tcx, M::PointerTag>
{
    #[inline(always)]
    fn layout(&self) -> TyLayout<'tcx> {
        self.layout
    }

    #[inline(always)]
    fn to_op(
        self,
        _ecx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, OpTy<'tcx, M::PointerTag>> {
        Ok(self)
    }

    #[inline(always)]
    fn from_mem_place(mplace: MPlaceTy<'tcx, M::PointerTag>) -> Self {
        mplace.into()
    }

    #[inline(always)]
    fn project_downcast(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        variant: VariantIdx,
    ) -> EvalResult<'tcx, Self> {
        ecx.operand_downcast(self, variant)
    }

    #[inline(always)]
    fn project_field(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self> {
        ecx.operand_field(self, field)
    }
}
impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Value<'a, 'mir, 'tcx, M>
    for MPlaceTy<'tcx, M::PointerTag>
{
    #[inline(always)]
    fn layout(&self) -> TyLayout<'tcx> {
        self.layout
    }

    #[inline(always)]
    fn to_op(
        self,
        _ecx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, OpTy<'tcx, M::PointerTag>> {
        Ok(self.into())
    }

    #[inline(always)]
    fn from_mem_place(mplace: MPlaceTy<'tcx, M::PointerTag>) -> Self {
        mplace
    }

    #[inline(always)]
    fn project_downcast(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        variant: VariantIdx,
    ) -> EvalResult<'tcx, Self> {
        ecx.mplace_downcast(self, variant)
    }

    #[inline(always)]
    fn project_field(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self> {
        ecx.mplace_field(self, field)
    }
}

macro_rules! make_value_visitor {
    ($visitor_trait_name:ident, $($mutability:ident)*) => {
        // How to traverse a value and what to do when we are at the leaves.
        pub trait $visitor_trait_name<'a, 'mir, 'tcx: 'mir+'a, M: Machine<'a, 'mir, 'tcx>>: Sized {
            type V: Value<'a, 'mir, 'tcx, M>;

            /// The visitor must have an `EvalContext` in it.
            fn ecx(&$($mutability)* self)
                -> &$($mutability)* EvalContext<'a, 'mir, 'tcx, M>;

            // Recursive actions, ready to be overloaded.
            /// Visit the given value, dispatching as appropriate to more specialized visitors.
            #[inline(always)]
            fn visit_value(&mut self, v: Self::V) -> EvalResult<'tcx>
            {
                self.walk_value(v)
            }
            /// Visit the given value as a union.  No automatic recursion can happen here.
            #[inline(always)]
            fn visit_union(&mut self, _v: Self::V) -> EvalResult<'tcx>
            {
                Ok(())
            }
            /// Visit this vale as an aggregate, you are even getting an iterator yielding
            /// all the fields (still in an `EvalResult`, you have to do error handling yourself).
            /// Recurses into the fields.
            #[inline(always)]
            fn visit_aggregate(
                &mut self,
                v: Self::V,
                fields: impl Iterator<Item=EvalResult<'tcx, Self::V>>,
            ) -> EvalResult<'tcx> {
                self.walk_aggregate(v, fields)
            }
            /// Called each time we recurse down to a field, passing in old and new value.
            /// This gives the visitor the chance to track the stack of nested fields that
            /// we are descending through.
            #[inline(always)]
            fn visit_field(
                &mut self,
                _old_val: Self::V,
                _field: usize,
                new_val: Self::V,
            ) -> EvalResult<'tcx> {
                self.visit_value(new_val)
            }

            #[inline(always)]
            fn visit_variant(
                &mut self,
                _old_val: Self::V,
                _variant: VariantIdx,
                new_val: Self::V,
            ) -> EvalResult<'tcx> {
                self.visit_value(new_val)
            }

            /// Called whenever we reach a value with uninhabited layout.
            /// Recursing to fields will *always* continue after this!  This is not meant to control
            /// whether and how we descend recursively/ into the scalar's fields if there are any,
            /// it is meant to provide the chance for additional checks when a value of uninhabited
            /// layout is detected.
            #[inline(always)]
            fn visit_uninhabited(&mut self) -> EvalResult<'tcx>
            { Ok(()) }
            /// Called whenever we reach a value with scalar layout.
            /// We do NOT provide a `ScalarMaybeUndef` here to avoid accessing memory if the
            /// visitor is not even interested in scalars.
            /// Recursing to fields will *always* continue after this!  This is not meant to control
            /// whether and how we descend recursively/ into the scalar's fields if there are any,
            /// it is meant to provide the chance for additional checks when a value of scalar
            /// layout is detected.
            #[inline(always)]
            fn visit_scalar(&mut self, _v: Self::V, _layout: &layout::Scalar) -> EvalResult<'tcx>
            { Ok(()) }

            /// Called whenever we reach a value of primitive type.  There can be no recursion
            /// below such a value.  This is the leaf function.
            /// We do *not* provide an `ImmTy` here because some implementations might want
            /// to write to the place this primitive lives in.
            #[inline(always)]
            fn visit_primitive(&mut self, _v: Self::V) -> EvalResult<'tcx>
            { Ok(()) }

            // Default recursors. Not meant to be overloaded.
            fn walk_aggregate(
                &mut self,
                v: Self::V,
                fields: impl Iterator<Item=EvalResult<'tcx, Self::V>>,
            ) -> EvalResult<'tcx> {
                // Now iterate over it.
                for (idx, field_val) in fields.enumerate() {
                    self.visit_field(v, idx, field_val?)?;
                }
                Ok(())
            }
            fn walk_value(&mut self, v: Self::V) -> EvalResult<'tcx>
            {
                trace!("walk_value: type: {}", v.layout().ty);
                // If this is a multi-variant layout, we have find the right one and proceed with
                // that.
                match v.layout().variants {
                    layout::Variants::NicheFilling { .. } |
                    layout::Variants::Tagged { .. } => {
                        let op = v.to_op(self.ecx())?;
                        let idx = self.ecx().read_discriminant(op)?.1;
                        let inner = v.project_downcast(self.ecx(), idx)?;
                        trace!("walk_value: variant layout: {:#?}", inner.layout());
                        // recurse with the inner type
                        return self.visit_variant(v, idx, inner);
                    }
                    layout::Variants::Single { .. } => {}
                }

                // Even for single variants, we might be able to get a more refined type:
                // If it is a trait object, switch to the actual type that was used to create it.
                match v.layout().ty.sty {
                    ty::Dynamic(..) => {
                        // immediate trait objects are not a thing
                        let dest = v.to_op(self.ecx())?.to_mem_place();
                        let inner = self.ecx().unpack_dyn_trait(dest)?.1;
                        trace!("walk_value: dyn object layout: {:#?}", inner.layout);
                        // recurse with the inner type
                        return self.visit_field(v, 0, Value::from_mem_place(inner));
                    },
                    _ => {},
                };

                // If this is a scalar, visit it as such.
                // Things can be aggregates and have scalar layout at the same time, and that
                // is very relevant for `NonNull` and similar structs: We need to visit them
                // at their scalar layout *before* descending into their fields.
                // FIXME: We could avoid some redundant checks here. For newtypes wrapping
                // scalars, we do the same check on every "level" (e.g. first we check
                // MyNewtype and then the scalar in there).
                match v.layout().abi {
                    layout::Abi::Uninhabited => {
                        self.visit_uninhabited()?;
                    }
                    layout::Abi::Scalar(ref layout) => {
                        self.visit_scalar(v, layout)?;
                    }
                    // FIXME: Should we do something for ScalarPair? Vector?
                    _ => {}
                }

                // Check primitive types.  We do this after checking the scalar layout,
                // just to have that done as well.  Primitives can have varying layout,
                // so we check them separately and before aggregate handling.
                // It is CRITICAL that we get this check right, or we might be
                // validating the wrong thing!
                let primitive = match v.layout().fields {
                    // Primitives appear as Union with 0 fields - except for Boxes and fat pointers.
                    layout::FieldPlacement::Union(0) => true,
                    _ => v.layout().ty.builtin_deref(true).is_some(),
                };
                if primitive {
                    return self.visit_primitive(v);
                }

                // Proceed into the fields.
                match v.layout().fields {
                    layout::FieldPlacement::Union(fields) => {
                        // Empty unions are not accepted by rustc. That's great, it means we can
                        // use that as an unambiguous signal for detecting primitives.  Make sure
                        // we did not miss any primitive.
                        debug_assert!(fields > 0);
                        self.visit_union(v)?;
                    },
                    layout::FieldPlacement::Arbitrary { ref offsets, .. } => {
                        // FIXME: We collect in a vec because otherwise there are lifetime errors:
                        // Projecting to a field needs (mutable!) access to `ecx`.
                        let fields: Vec<EvalResult<'tcx, Self::V>> =
                            (0..offsets.len()).map(|i| {
                                v.project_field(self.ecx(), i as u64)
                            })
                            .collect();
                        self.visit_aggregate(v, fields.into_iter())?;
                    },
                    layout::FieldPlacement::Array { .. } => {
                        // Let's get an mplace first.
                        let mplace = if v.layout().is_zst() {
                            // it's a ZST, the memory content cannot matter
                            MPlaceTy::dangling(v.layout(), self.ecx())
                        } else {
                            // non-ZST array/slice/str cannot be immediate
                            v.to_op(self.ecx())?.to_mem_place()
                        };
                        // Now we can go over all the fields.
                        let iter = self.ecx().mplace_array_fields(mplace)?
                            .map(|f| f.and_then(|f| {
                                Ok(Value::from_mem_place(f))
                            }));
                        self.visit_aggregate(v, iter)?;
                    }
                }
                Ok(())
            }
        }
    }
}

make_value_visitor!(ValueVisitor,);
make_value_visitor!(MutValueVisitor,mut);
<commit_msg>Rollup merge of #56100 - RalfJung:visiting-generators, r=oli-obk<commit_after>//! Visitor for a run-time value with a given layout: Traverse enums, structs and other compound
//! types until we arrive at the leaves, with custom handling for primitive types.

use rustc::ty::layout::{self, TyLayout, VariantIdx};
use rustc::ty;
use rustc::mir::interpret::{
    EvalResult,
};

use super::{
    Machine, EvalContext, MPlaceTy, OpTy,
};

// A thing that we can project into, and that has a layout.
// This wouldn't have to depend on `Machine` but with the current type inference,
// that's just more convenient to work with (avoids repeating all the `Machine` bounds).
pub trait Value<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>>: Copy
{
    /// Get this value's layout.
    fn layout(&self) -> TyLayout<'tcx>;

    /// Make this into an `OpTy`.
    fn to_op(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, OpTy<'tcx, M::PointerTag>>;

    /// Create this from an `MPlaceTy`.
    fn from_mem_place(MPlaceTy<'tcx, M::PointerTag>) -> Self;

    /// Project to the given enum variant.
    fn project_downcast(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        variant: VariantIdx,
    ) -> EvalResult<'tcx, Self>;

    /// Project to the n-th field.
    fn project_field(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self>;
}

// Operands and memory-places are both values.
// Places in general are not due to `place_field` having to do `force_allocation`.
impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Value<'a, 'mir, 'tcx, M>
    for OpTy<'tcx, M::PointerTag>
{
    #[inline(always)]
    fn layout(&self) -> TyLayout<'tcx> {
        self.layout
    }

    #[inline(always)]
    fn to_op(
        self,
        _ecx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, OpTy<'tcx, M::PointerTag>> {
        Ok(self)
    }

    #[inline(always)]
    fn from_mem_place(mplace: MPlaceTy<'tcx, M::PointerTag>) -> Self {
        mplace.into()
    }

    #[inline(always)]
    fn project_downcast(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        variant: VariantIdx,
    ) -> EvalResult<'tcx, Self> {
        ecx.operand_downcast(self, variant)
    }

    #[inline(always)]
    fn project_field(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self> {
        ecx.operand_field(self, field)
    }
}
impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Value<'a, 'mir, 'tcx, M>
    for MPlaceTy<'tcx, M::PointerTag>
{
    #[inline(always)]
    fn layout(&self) -> TyLayout<'tcx> {
        self.layout
    }

    #[inline(always)]
    fn to_op(
        self,
        _ecx: &EvalContext<'a, 'mir, 'tcx, M>,
    ) -> EvalResult<'tcx, OpTy<'tcx, M::PointerTag>> {
        Ok(self.into())
    }

    #[inline(always)]
    fn from_mem_place(mplace: MPlaceTy<'tcx, M::PointerTag>) -> Self {
        mplace
    }

    #[inline(always)]
    fn project_downcast(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        variant: VariantIdx,
    ) -> EvalResult<'tcx, Self> {
        ecx.mplace_downcast(self, variant)
    }

    #[inline(always)]
    fn project_field(
        self,
        ecx: &EvalContext<'a, 'mir, 'tcx, M>,
        field: u64,
    ) -> EvalResult<'tcx, Self> {
        ecx.mplace_field(self, field)
    }
}

macro_rules! make_value_visitor {
    ($visitor_trait_name:ident, $($mutability:ident)*) => {
        // How to traverse a value and what to do when we are at the leaves.
        pub trait $visitor_trait_name<'a, 'mir, 'tcx: 'mir+'a, M: Machine<'a, 'mir, 'tcx>>: Sized {
            type V: Value<'a, 'mir, 'tcx, M>;

            /// The visitor must have an `EvalContext` in it.
            fn ecx(&$($mutability)* self)
                -> &$($mutability)* EvalContext<'a, 'mir, 'tcx, M>;

            // Recursive actions, ready to be overloaded.
            /// Visit the given value, dispatching as appropriate to more specialized visitors.
            #[inline(always)]
            fn visit_value(&mut self, v: Self::V) -> EvalResult<'tcx>
            {
                self.walk_value(v)
            }
            /// Visit the given value as a union.  No automatic recursion can happen here.
            #[inline(always)]
            fn visit_union(&mut self, _v: Self::V) -> EvalResult<'tcx>
            {
                Ok(())
            }
            /// Visit this vale as an aggregate, you are even getting an iterator yielding
            /// all the fields (still in an `EvalResult`, you have to do error handling yourself).
            /// Recurses into the fields.
            #[inline(always)]
            fn visit_aggregate(
                &mut self,
                v: Self::V,
                fields: impl Iterator<Item=EvalResult<'tcx, Self::V>>,
            ) -> EvalResult<'tcx> {
                self.walk_aggregate(v, fields)
            }

            /// Called each time we recurse down to a field of a "product-like" aggregate
            /// (structs, tuples, arrays and the like, but not enums), passing in old and new value.
            /// This gives the visitor the chance to track the stack of nested fields that
            /// we are descending through.
            #[inline(always)]
            fn visit_field(
                &mut self,
                _old_val: Self::V,
                _field: usize,
                new_val: Self::V,
            ) -> EvalResult<'tcx> {
                self.visit_value(new_val)
            }

            /// Called for recursing into the field of a generator.  These are not known to be
            /// initialized, so we treat them like unions.
            #[inline(always)]
            fn visit_generator_field(
                &mut self,
                _old_val: Self::V,
                _field: usize,
                new_val: Self::V,
            ) -> EvalResult<'tcx> {
                self.visit_union(new_val)
            }

            /// Called when recursing into an enum variant.
            #[inline(always)]
            fn visit_variant(
                &mut self,
                _old_val: Self::V,
                _variant: VariantIdx,
                new_val: Self::V,
            ) -> EvalResult<'tcx> {
                self.visit_value(new_val)
            }

            /// Called whenever we reach a value with uninhabited layout.
            /// Recursing to fields will *always* continue after this!  This is not meant to control
            /// whether and how we descend recursively/ into the scalar's fields if there are any,
            /// it is meant to provide the chance for additional checks when a value of uninhabited
            /// layout is detected.
            #[inline(always)]
            fn visit_uninhabited(&mut self) -> EvalResult<'tcx>
            { Ok(()) }
            /// Called whenever we reach a value with scalar layout.
            /// We do NOT provide a `ScalarMaybeUndef` here to avoid accessing memory if the
            /// visitor is not even interested in scalars.
            /// Recursing to fields will *always* continue after this!  This is not meant to control
            /// whether and how we descend recursively/ into the scalar's fields if there are any,
            /// it is meant to provide the chance for additional checks when a value of scalar
            /// layout is detected.
            #[inline(always)]
            fn visit_scalar(&mut self, _v: Self::V, _layout: &layout::Scalar) -> EvalResult<'tcx>
            { Ok(()) }

            /// Called whenever we reach a value of primitive type.  There can be no recursion
            /// below such a value.  This is the leaf function.
            /// We do *not* provide an `ImmTy` here because some implementations might want
            /// to write to the place this primitive lives in.
            #[inline(always)]
            fn visit_primitive(&mut self, _v: Self::V) -> EvalResult<'tcx>
            { Ok(()) }

            // Default recursors. Not meant to be overloaded.
            fn walk_aggregate(
                &mut self,
                v: Self::V,
                fields: impl Iterator<Item=EvalResult<'tcx, Self::V>>,
            ) -> EvalResult<'tcx> {
                // Now iterate over it.
                for (idx, field_val) in fields.enumerate() {
                    self.visit_field(v, idx, field_val?)?;
                }
                Ok(())
            }
            fn walk_value(&mut self, v: Self::V) -> EvalResult<'tcx>
            {
                trace!("walk_value: type: {}", v.layout().ty);
                // If this is a multi-variant layout, we have find the right one and proceed with
                // that.
                match v.layout().variants {
                    layout::Variants::NicheFilling { .. } |
                    layout::Variants::Tagged { .. } => {
                        let op = v.to_op(self.ecx())?;
                        let idx = self.ecx().read_discriminant(op)?.1;
                        let inner = v.project_downcast(self.ecx(), idx)?;
                        trace!("walk_value: variant layout: {:#?}", inner.layout());
                        // recurse with the inner type
                        return self.visit_variant(v, idx, inner);
                    }
                    layout::Variants::Single { .. } => {}
                }

                // Even for single variants, we might be able to get a more refined type:
                // If it is a trait object, switch to the actual type that was used to create it.
                match v.layout().ty.sty {
                    ty::Dynamic(..) => {
                        // immediate trait objects are not a thing
                        let dest = v.to_op(self.ecx())?.to_mem_place();
                        let inner = self.ecx().unpack_dyn_trait(dest)?.1;
                        trace!("walk_value: dyn object layout: {:#?}", inner.layout);
                        // recurse with the inner type
                        return self.visit_field(v, 0, Value::from_mem_place(inner));
                    },
                    _ => {},
                };

                // If this is a scalar, visit it as such.
                // Things can be aggregates and have scalar layout at the same time, and that
                // is very relevant for `NonNull` and similar structs: We need to visit them
                // at their scalar layout *before* descending into their fields.
                // FIXME: We could avoid some redundant checks here. For newtypes wrapping
                // scalars, we do the same check on every "level" (e.g. first we check
                // MyNewtype and then the scalar in there).
                match v.layout().abi {
                    layout::Abi::Uninhabited => {
                        self.visit_uninhabited()?;
                    }
                    layout::Abi::Scalar(ref layout) => {
                        self.visit_scalar(v, layout)?;
                    }
                    // FIXME: Should we do something for ScalarPair? Vector?
                    _ => {}
                }

                // Check primitive types.  We do this after checking the scalar layout,
                // just to have that done as well.  Primitives can have varying layout,
                // so we check them separately and before aggregate handling.
                // It is CRITICAL that we get this check right, or we might be
                // validating the wrong thing!
                let primitive = match v.layout().fields {
                    // Primitives appear as Union with 0 fields - except for Boxes and fat pointers.
                    layout::FieldPlacement::Union(0) => true,
                    _ => v.layout().ty.builtin_deref(true).is_some(),
                };
                if primitive {
                    return self.visit_primitive(v);
                }

                // Proceed into the fields.
                match v.layout().fields {
                    layout::FieldPlacement::Union(fields) => {
                        // Empty unions are not accepted by rustc. That's great, it means we can
                        // use that as an unambiguous signal for detecting primitives.  Make sure
                        // we did not miss any primitive.
                        debug_assert!(fields > 0);
                        self.visit_union(v)
                    },
                    layout::FieldPlacement::Arbitrary { ref offsets, .. } => {
                        // Special handling needed for generators: All but the first field
                        // (which is the state) are actually implicitly `MaybeUninit`, i.e.,
                        // they may or may not be initialized, so we cannot visit them.
                        match v.layout().ty.sty {
                            ty::Generator(..) => {
                                let field = v.project_field(self.ecx(), 0)?;
                                self.visit_aggregate(v, std::iter::once(Ok(field)))?;
                                for i in 1..offsets.len() {
                                    let field = v.project_field(self.ecx(), i as u64)?;
                                    self.visit_generator_field(v, i, field)?;
                                }
                                Ok(())
                            }
                            _ => {
                                // FIXME: We collect in a vec because otherwise there are lifetime
                                // errors: Projecting to a field needs access to `ecx`.
                                let fields: Vec<EvalResult<'tcx, Self::V>> =
                                    (0..offsets.len()).map(|i| {
                                        v.project_field(self.ecx(), i as u64)
                                    })
                                    .collect();
                                self.visit_aggregate(v, fields.into_iter())
                            }
                        }
                    },
                    layout::FieldPlacement::Array { .. } => {
                        // Let's get an mplace first.
                        let mplace = if v.layout().is_zst() {
                            // it's a ZST, the memory content cannot matter
                            MPlaceTy::dangling(v.layout(), self.ecx())
                        } else {
                            // non-ZST array/slice/str cannot be immediate
                            v.to_op(self.ecx())?.to_mem_place()
                        };
                        // Now we can go over all the fields.
                        let iter = self.ecx().mplace_array_fields(mplace)?
                            .map(|f| f.and_then(|f| {
                                Ok(Value::from_mem_place(f))
                            }));
                        self.visit_aggregate(v, iter)
                    }
                }
            }
        }
    }
}

make_value_visitor!(ValueVisitor,);
make_value_visitor!(MutValueVisitor,mut);
<|endoftext|>


<commit_before>use futures_core::task::{local_waker_from_nonlocal, LocalWaker, Wake};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

/// Number of times the waker was awoken.
///
/// See [`new_count_waker`] for usage.
#[derive(Debug)]
pub struct AwokenCount {
    inner: Arc<WakerInner>,
}

impl AwokenCount {
    /// Get the current count.
    fn get(&self) -> usize {
        self.inner.count.load(Ordering::SeqCst)
    }
}

impl PartialEq<usize> for AwokenCount {
    fn eq(&self, other: &usize) -> bool {
        self.get() == *other
    }
}

#[derive(Debug)]
struct WakerInner {
    count: AtomicUsize,
}

impl Wake for WakerInner {
    fn wake(arc_self: &Arc<Self>) {
        let _ = arc_self.count.fetch_add(1, Ordering::SeqCst);
    }
}

/// Create a new [`LocalWaker`] that counts the number of times it's awoken.
///
/// [`LocalWaker`]: futures_core::task::LocalWaker
///
/// # Examples
///
/// ```
/// #![feature(futures_api)]
/// use futures_test::task::new_count_waker;
///
/// let (lw, count) = new_count_waker();
///
/// assert_eq!(count, 0);
///
/// lw.wake();
/// lw.wake();
///
/// assert_eq!(count, 2);
/// ```
pub fn new_count_waker() -> (LocalWaker, AwokenCount) {
    let inner = Arc::new(WakerInner { count: AtomicUsize::new(0) });
    (local_waker_from_nonlocal(inner.clone()), AwokenCount { inner })
}
<commit_msg>Make AwokenCount.get public<commit_after>use futures_core::task::{local_waker_from_nonlocal, LocalWaker, Wake};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

/// Number of times the waker was awoken.
///
/// See [`new_count_waker`] for usage.
#[derive(Debug)]
pub struct AwokenCount {
    inner: Arc<WakerInner>,
}

impl AwokenCount {
    /// Get the current count.
    pub fn get(&self) -> usize {
        self.inner.count.load(Ordering::SeqCst)
    }
}

impl PartialEq<usize> for AwokenCount {
    fn eq(&self, other: &usize) -> bool {
        self.get() == *other
    }
}

#[derive(Debug)]
struct WakerInner {
    count: AtomicUsize,
}

impl Wake for WakerInner {
    fn wake(arc_self: &Arc<Self>) {
        let _ = arc_self.count.fetch_add(1, Ordering::SeqCst);
    }
}

/// Create a new [`LocalWaker`] that counts the number of times it's awoken.
///
/// [`LocalWaker`]: futures_core::task::LocalWaker
///
/// # Examples
///
/// ```
/// #![feature(futures_api)]
/// use futures_test::task::new_count_waker;
///
/// let (lw, count) = new_count_waker();
///
/// assert_eq!(count, 0);
///
/// lw.wake();
/// lw.wake();
///
/// assert_eq!(count, 2);
/// ```
pub fn new_count_waker() -> (LocalWaker, AwokenCount) {
    let inner = Arc::new(WakerInner { count: AtomicUsize::new(0) });
    (local_waker_from_nonlocal(inner.clone()), AwokenCount { inner })
}
<|endoftext|>


<commit_before>extern crate byteorder;

mod slice;
mod vorbis;

use std::mem;
use std::ops;
use std::borrow::{Borrow, BorrowMut, ToOwned};
use std::io::{Cursor, BufRead};
use std::marker::PhantomData;
use std::borrow::Cow;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use byteorder::Error as ByteOrderError;
use slice::Slice;

const OGG_PAGE_CAPTURE: &'static [u8] = b"OggS";
const POSITION_OFFSET: usize = 6;
const SERIAL_OFFSET: usize = 14;

#[derive(Debug)]
pub enum OggPageCheckError {
    TooShort,
    BadCapture,
    BadVersion,
    BadCrc,
}

#[derive(Clone)]
pub struct OggTrackBuf {
    inner: Vec<u8>,
}

pub struct OggTrack {
    inner: Slice,
}

impl ops::Deref for OggTrackBuf {
    type Target = OggTrack;

    fn deref<'a>(&'a self) -> &'a OggTrack {
        OggTrack::from_u8_slice_unchecked(&self.inner)
    }
}

impl Borrow<OggTrack> for OggTrackBuf {
    fn borrow(&self) -> &OggTrack {
        OggTrack::from_u8_slice_unchecked(&self.inner)
    }
}

impl AsRef<OggTrack> for OggTrackBuf {
    fn as_ref(&self) -> &OggTrack {
        OggTrack::from_u8_slice_unchecked(&self.inner)
    }
}

impl AsMut<OggTrack> for OggTrackBuf {
    fn as_mut(&mut self) -> &mut OggTrack {
        OggTrack::from_u8_slice_unchecked_mut(&mut self.inner)
    }
}

impl BorrowMut<OggTrack> for OggTrackBuf {
    fn borrow_mut(&mut self) -> &mut OggTrack {
        OggTrack::from_u8_slice_unchecked_mut(&mut self.inner)
    }
}

impl ToOwned for OggTrack {
    type Owned = OggTrackBuf;

    fn to_owned(&self) -> OggTrackBuf {
        OggTrackBuf { inner: self.inner.to_owned() }
    }
}

impl OggTrackBuf {
    pub fn new(buf: Vec<u8>) -> Result<OggTrackBuf, OggPageCheckError> {
        try!(OggTrack::new(&buf));
        Ok(OggTrackBuf { inner: buf })
    }

    pub fn into_inner(self) -> Vec<u8> {
        self.inner
    }
}

impl OggTrack {
    pub fn new(buf: &[u8]) -> Result<&OggTrack, OggPageCheckError>  {
        let mut offset = 0;
        while offset < buf.len() {
            let page = try!(OggPage::new(&buf[offset..]));
            offset += page.as_u8_slice().len();
        }
        Ok(OggTrack::from_u8_slice_unchecked(buf))
    }

    /// The following function allows unchecked construction of a ogg track
    /// from a u8 slice.  This is private because it does not maintain
    /// the OggPage invariant.
    fn from_u8_slice_unchecked(s: &[u8]) -> &OggTrack {
        unsafe { mem::transmute(s) }
    }

    /// The following (private!) function allows unchecked construction of a
    /// mutable ogg page from a mutable u8 slice.  This is private because it
    /// does not maintain the OggPage invariant.
    fn from_u8_slice_unchecked_mut(s: &mut [u8]) -> &mut OggTrack {
        unsafe { mem::transmute(s) }
    }

    pub fn as_u8_slice(&self) -> &[u8] {
        &self.inner
    }

    pub fn pages(&self) -> TrackPageIter {
        TrackPageIter {
            data: self.as_u8_slice(),
            offset: 0,
        }
    }

    pub fn pages_mut(&self) -> TrackPageIterMut {
        // TODO: impl
        unimplemented!();
    }
}

pub struct TrackPageIter<'a> {
    data: &'a [u8],
    offset: usize,
}

impl<'a> Iterator for TrackPageIter<'a> {
    type Item = &'a OggPage;

    fn next(&mut self) -> Option<&'a OggPage> {
        if self.data.len() == self.offset {
            return None;
        }
        let page = OggPage::new(&self.data[self.offset..]).unwrap();
        self.offset += page.as_u8_slice().len();
        Some(page)
    }
}

// TODO: iterate &mut OggPage
pub struct TrackPageIterMut<'a> {
    data: &'a mut [u8],
    offset: usize,
}

#[derive(Clone)]
pub struct OggPageBuf {
    inner: Vec<u8>,
}

pub struct OggPage {
    inner: Slice,
}

impl AsRef<OggPage> for OggPageBuf {
    fn as_ref(&self) -> &OggPage {
        OggPage::from_u8_slice_unchecked(&self.inner)
    }
}

impl AsMut<OggPage> for OggPageBuf {
    fn as_mut(&mut self) -> &mut OggPage {
        OggPage::from_u8_slice_unchecked_mut(&mut self.inner)
    }
}

impl ops::Deref for OggPageBuf {
    type Target = OggPage;

    fn deref<'a>(&'a self) -> &'a OggPage {
        OggPage::from_u8_slice_unchecked(&self.inner)
    }
}

impl Borrow<OggPage> for OggPageBuf {
    fn borrow(&self) -> &OggPage {
        OggPage::from_u8_slice_unchecked(&self.inner)
    }
}

impl BorrowMut<OggPage> for OggPageBuf {
    fn borrow_mut(&mut self) -> &mut OggPage {
        OggPage::from_u8_slice_unchecked_mut(&mut self.inner)
    }
}

impl ToOwned for OggPage {
    type Owned = OggPageBuf;

    fn to_owned(&self) -> OggPageBuf {
        OggPageBuf { inner: self.inner.to_owned() }
    }
}

impl OggPageBuf {
    pub fn new(mut buf: Vec<u8>) -> Result<OggPageBuf, OggPageCheckError>  {
        let slice_len = try!(OggPage::measure_whole(&buf)).len();
        buf.truncate(slice_len);
        Ok(OggPageBuf { inner: buf })
    }

    pub fn into_cow(self) -> Cow<'static, OggPage> {
        Cow::Owned(self)
    }
}

impl OggPage {
    /// The following function allows unchecked construction of a ogg page
    /// from a u8 slice.  This is private because it does not maintain
    /// the OggPage invariant.
    fn from_u8_slice_unchecked(s: &[u8]) -> &OggPage {
        unsafe { mem::transmute(s) }
    }

    /// The following (private!) function allows unchecked construction of a
    /// mutable ogg page from a mutable u8 slice.  This is private because it
    /// does not maintain the OggPage invariant.
    fn from_u8_slice_unchecked_mut(s: &mut [u8]) -> &mut OggPage {
        unsafe { mem::transmute(s) }
    }

    pub fn as_u8_slice(&self) -> &[u8] {
        unsafe { mem::transmute(self) }
    }

    /// Mutably borrow the underlying storage.  This is private because it
    /// does not maintain the OggPage invariant.
    fn as_u8_slice_mut(&mut self) -> &mut [u8] {
        unsafe { mem::transmute(self) }
    }

    pub fn new(buf: &[u8]) -> Result<&OggPage, OggPageCheckError> {
        let buffer = try!(OggPage::measure_whole(buf));
        Ok(OggPage::from_u8_slice_unchecked(buffer))
    }

    pub fn new_mut(buf: &mut [u8]) -> Result<&mut OggPage, OggPageCheckError> {
        let page_length = {
            let (hbuf, bbuf) = try!(OggPage::measure(buf));
            hbuf.len() + bbuf.len()
        };
        Ok(OggPage::from_u8_slice_unchecked_mut(&mut buf[0..page_length]))
    }

    fn measure(buf: &[u8]) -> Result<(&[u8], &[u8]), OggPageCheckError> {
        impl From<ByteOrderError> for OggPageCheckError {
            fn from(e: ByteOrderError) -> OggPageCheckError {
                match e {
                    ByteOrderError::UnexpectedEOF => OggPageCheckError::TooShort,
                    ByteOrderError::Io(_) => panic!("no I/O allowed"),
                }
            }
        }

        let mut cursor = Cursor::new(buf);

        if buf.len() < 27 {
            return Err(OggPageCheckError::TooShort);
        }
        if &buf[0..4] != OGG_PAGE_CAPTURE {
            return Err(OggPageCheckError::BadCapture);
        }

        cursor.consume(4);  // capture sequence
        if try!(cursor.read_u8()) != 0 {
            return Err(OggPageCheckError::BadVersion);
        }

        // flags(1) + granule(8) + serial(4) + page_seq(4) + csum(4)
        cursor.consume(1 + 8 + 4 + 4 + 4);

        let page_segments = try!(cursor.read_u8());

        let mut body_len = 0;
        for _ in 0..page_segments {
            body_len += try!(cursor.read_u8()) as u64;
        }

        let total_len = (cursor.position() + body_len) as usize;
        if buf.len() < total_len {
            return Err(OggPageCheckError::TooShort);
        }

        let h_end = cursor.position() as usize;
        let b_end = h_end + body_len as usize;
        Ok((
            &buf[0..h_end],
            &buf[h_end..b_end],
        ))
    }

    fn measure_whole(buf: &[u8]) -> Result<&[u8], OggPageCheckError> {
        let page_length = {
            let (h_buf, b_buf) = try!(OggPage::measure(buf));
            h_buf.len() + b_buf.len()
        };
        Ok(&buf[0..page_length])
    }

    fn measure_whole_mut(buf: &mut [u8]) -> Result<&mut [u8], OggPageCheckError> {
        let page_length = {
            let (hbuf, bbuf) = try!(OggPage::measure(buf));
            hbuf.len() + bbuf.len()
        };
        Ok(&mut buf[0..page_length])
    }

    pub fn position(&self) -> u64 {
        let self_buf = self.as_u8_slice();
        let mut cur = Cursor::new(&self_buf[POSITION_OFFSET..POSITION_OFFSET+8]);
        cur.read_u64::<LittleEndian>().unwrap()
    }

    pub fn set_position(&mut self, granule: u64) {
        let mut tx = self.begin();
        tx.set_position(granule);
    }

    pub fn serial(&self) -> u32 {
        let self_buf = self.as_u8_slice();
        let mut cur = Cursor::new(&self_buf[SERIAL_OFFSET..SERIAL_OFFSET+4]);
        cur.read_u32::<LittleEndian>().unwrap()
    }

    pub fn set_serial(&mut self, serial: u32) {
        let mut tx = self.begin();
        tx.set_serial(serial);
    }

    fn recompute_checksum(&mut self) {
        unimplemented!();
    }

    pub fn begin<'a>(&'a mut self) -> ChecksumGuard<'a> {
        ChecksumGuard {
            page: self,
            _marker: PhantomData,
        }
    }

    /// Am iterator of packet slices
    pub fn packets<'a>(&'a self) -> Packets<'a> {
        Packets { page: &self, packet: 0 }
    }

    pub fn into_cow<'a>(&'a self) -> Cow<'a, OggPage> {
        Cow::Borrowed(self)
    }
}

pub struct Packets<'a> {
    page: &'a OggPage,
    packet: usize,
}

impl<'a> Iterator for Packets<'a> {
    type Item = &'a [u8];

    fn next(&mut self) -> Option<&'a [u8]> {
        let packet_slice = self.page;
        unimplemented!();
    }
}

pub struct ChecksumGuard<'a> {
    page: &'a mut OggPage,
    _marker: PhantomData<&'a ()>,
}

impl<'a> ChecksumGuard<'a> {
    pub fn set_position(&mut self, granule: u64) {
        let self_buf = self.page.as_u8_slice_mut();
        let mut cur = Cursor::new(&mut self_buf[POSITION_OFFSET..POSITION_OFFSET+8]);
        cur.write_u64::<LittleEndian>(granule).unwrap();
    }

    pub fn set_serial(&mut self, serial: u32) {
        let self_buf = self.page.as_u8_slice_mut();
        let mut cur = Cursor::new(&mut self_buf[SERIAL_OFFSET..SERIAL_OFFSET+4]);
        cur.write_u32::<LittleEndian>(serial).unwrap();
    }
}

impl<'a> Drop for ChecksumGuard<'a> {
    fn drop(&mut self) {
        self.page.recompute_checksum();
    }
}

pub struct Recapture([u8; 4]);

impl Recapture {
    pub fn new() -> Recapture {
        Recapture([0; 4])
    }

    pub fn push_byte(&mut self, byte: u8) {
        let mut buf = [0; 4];
        buf[0] = self.0[1];
        buf[1] = self.0[2];
        buf[2] = self.0[3];
        buf[3] = byte;
        *self = Recapture(buf);
    }

    pub fn is_captured(&self) -> bool {
        &self.0 == OGG_PAGE_CAPTURE
    }
}

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

    #[test]
    fn test_capture() {
        let mut cap = Recapture::new();
        cap.push_byte(b'O');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'g');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'g');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'S');
        assert_eq!(true, cap.is_captured());
        cap.push_byte(b'g');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'g');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'S');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'O');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'g');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'g');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'S');
        assert_eq!(true, cap.is_captured());
    }
}
<commit_msg>ogg: implement OggTrack#pages_mut iterator<commit_after>extern crate byteorder;

mod slice;
mod vorbis;

use std::mem;
use std::ops;
use std::borrow::{Borrow, BorrowMut, ToOwned};
use std::marker;
use std::io::{Cursor, BufRead};
use std::marker::PhantomData;
use std::borrow::Cow;
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use byteorder::Error as ByteOrderError;
use slice::Slice;

const OGG_PAGE_CAPTURE: &'static [u8] = b"OggS";
const POSITION_OFFSET: usize = 6;
const SERIAL_OFFSET: usize = 14;

#[derive(Debug)]
pub enum OggPageCheckError {
    TooShort,
    BadCapture,
    BadVersion,
    BadCrc,
}

#[derive(Clone)]
pub struct OggTrackBuf {
    inner: Vec<u8>,
}

pub struct OggTrack {
    inner: Slice,
}

impl ops::Deref for OggTrackBuf {
    type Target = OggTrack;

    fn deref<'a>(&'a self) -> &'a OggTrack {
        OggTrack::from_u8_slice_unchecked(&self.inner)
    }
}

impl Borrow<OggTrack> for OggTrackBuf {
    fn borrow(&self) -> &OggTrack {
        OggTrack::from_u8_slice_unchecked(&self.inner)
    }
}

impl AsRef<OggTrack> for OggTrackBuf {
    fn as_ref(&self) -> &OggTrack {
        OggTrack::from_u8_slice_unchecked(&self.inner)
    }
}

impl AsMut<OggTrack> for OggTrackBuf {
    fn as_mut(&mut self) -> &mut OggTrack {
        OggTrack::from_u8_slice_unchecked_mut(&mut self.inner)
    }
}

impl BorrowMut<OggTrack> for OggTrackBuf {
    fn borrow_mut(&mut self) -> &mut OggTrack {
        OggTrack::from_u8_slice_unchecked_mut(&mut self.inner)
    }
}

impl ToOwned for OggTrack {
    type Owned = OggTrackBuf;

    fn to_owned(&self) -> OggTrackBuf {
        OggTrackBuf { inner: self.inner.to_owned() }
    }
}

impl OggTrackBuf {
    pub fn new(buf: Vec<u8>) -> Result<OggTrackBuf, OggPageCheckError> {
        try!(OggTrack::new(&buf));
        Ok(OggTrackBuf { inner: buf })
    }

    pub fn into_inner(self) -> Vec<u8> {
        self.inner
    }
}

impl OggTrack {
    pub fn new(buf: &[u8]) -> Result<&OggTrack, OggPageCheckError>  {
        let mut offset = 0;
        while offset < buf.len() {
            let page = try!(OggPage::new(&buf[offset..]));
            offset += page.as_u8_slice().len();
        }
        Ok(OggTrack::from_u8_slice_unchecked(buf))
    }

    /// The following function allows unchecked construction of a ogg track
    /// from a u8 slice.  This is private because it does not maintain
    /// the OggPage invariant.
    fn from_u8_slice_unchecked(s: &[u8]) -> &OggTrack {
        unsafe { mem::transmute(s) }
    }

    /// The following (private!) function allows unchecked construction of a
    /// mutable ogg page from a mutable u8 slice.  This is private because it
    /// does not maintain the OggPage invariant.
    fn from_u8_slice_unchecked_mut(s: &mut [u8]) -> &mut OggTrack {
        unsafe { mem::transmute(s) }
    }

    pub fn as_u8_slice(&self) -> &[u8] {
        &self.inner
    }

    /// Mutably borrow the underlying storage.  This is private because it
    /// does not maintain the OggPage invariant.
    fn as_u8_slice_mut(&mut self) -> &mut [u8] {
        unsafe { mem::transmute(self) }
    }

    pub fn pages(&self) -> TrackPageIter {
        TrackPageIter {
            data: self.as_u8_slice(),
            offset: 0,
        }
    }

    pub fn pages_mut(&mut self) -> TrackPageIterMut {
        let buffer = self.as_u8_slice_mut();
        let ptr = buffer.as_mut_ptr();
        let end = unsafe { ptr.offset(buffer.len() as isize) };
        TrackPageIterMut {
            ptr: ptr,
            end: end,
            _marker: marker::PhantomData,
        }
    }
}

pub struct TrackPageIter<'a> {
    data: &'a [u8],
    offset: usize,
}

impl<'a> Iterator for TrackPageIter<'a> {
    type Item = &'a OggPage;

    fn next(&mut self) -> Option<&'a OggPage> {
        if self.data.len() == self.offset {
            return None;
        }
        let page = OggPage::new(&self.data[self.offset..]).unwrap();
        self.offset += page.as_u8_slice().len();
        Some(page)
    }
}

pub struct TrackPageIterMut<'a> {
    ptr: *mut u8,
    end: *mut u8,
    _marker: marker::PhantomData<&'a mut ()>,
}

impl<'a> Iterator for TrackPageIterMut<'a> {
    type Item = &'a mut OggPage;

    fn next(&mut self) -> Option<&'a mut OggPage> {
        if self.ptr == self.end {
            return None;
        }

        let length = self.end as usize - self.ptr as usize;
        let next_boundary = {
            let remaining_data = unsafe { std::slice::from_raw_parts(self.ptr, length) };
            let page_len = OggPage::new(remaining_data).unwrap().as_u8_slice().len();
            assert!(page_len <= length);
            page_len
        };

        let page_data;
        unsafe {
            page_data = std::slice::from_raw_parts_mut(self.ptr, next_boundary);
            self.ptr = self.ptr.offset(next_boundary as isize);
        }

        Some(OggPage::from_u8_slice_unchecked_mut(page_data))
    }
}

#[derive(Clone)]
pub struct OggPageBuf {
    inner: Vec<u8>,
}

pub struct OggPage {
    inner: Slice,
}

impl AsRef<OggPage> for OggPageBuf {
    fn as_ref(&self) -> &OggPage {
        OggPage::from_u8_slice_unchecked(&self.inner)
    }
}

impl AsMut<OggPage> for OggPageBuf {
    fn as_mut(&mut self) -> &mut OggPage {
        OggPage::from_u8_slice_unchecked_mut(&mut self.inner)
    }
}

impl ops::Deref for OggPageBuf {
    type Target = OggPage;

    fn deref<'a>(&'a self) -> &'a OggPage {
        OggPage::from_u8_slice_unchecked(&self.inner)
    }
}

impl Borrow<OggPage> for OggPageBuf {
    fn borrow(&self) -> &OggPage {
        OggPage::from_u8_slice_unchecked(&self.inner)
    }
}

impl BorrowMut<OggPage> for OggPageBuf {
    fn borrow_mut(&mut self) -> &mut OggPage {
        OggPage::from_u8_slice_unchecked_mut(&mut self.inner)
    }
}

impl ToOwned for OggPage {
    type Owned = OggPageBuf;

    fn to_owned(&self) -> OggPageBuf {
        OggPageBuf { inner: self.inner.to_owned() }
    }
}

impl OggPageBuf {
    pub fn new(mut buf: Vec<u8>) -> Result<OggPageBuf, OggPageCheckError>  {
        let slice_len = try!(OggPage::measure_whole(&buf)).len();
        buf.truncate(slice_len);
        Ok(OggPageBuf { inner: buf })
    }

    pub fn into_cow(self) -> Cow<'static, OggPage> {
        Cow::Owned(self)
    }
}

impl OggPage {
    /// The following function allows unchecked construction of a ogg page
    /// from a u8 slice.  This is private because it does not maintain
    /// the OggPage invariant.
    fn from_u8_slice_unchecked(s: &[u8]) -> &OggPage {
        unsafe { mem::transmute(s) }
    }

    /// The following (private!) function allows unchecked construction of a
    /// mutable ogg page from a mutable u8 slice.  This is private because it
    /// does not maintain the OggPage invariant.
    fn from_u8_slice_unchecked_mut(s: &mut [u8]) -> &mut OggPage {
        unsafe { mem::transmute(s) }
    }

    pub fn as_u8_slice(&self) -> &[u8] {
        unsafe { mem::transmute(self) }
    }

    /// Mutably borrow the underlying storage.  This is private because it
    /// does not maintain the OggPage invariant.
    fn as_u8_slice_mut(&mut self) -> &mut [u8] {
        unsafe { mem::transmute(self) }
    }

    pub fn new(buf: &[u8]) -> Result<&OggPage, OggPageCheckError> {
        let buffer = try!(OggPage::measure_whole(buf));
        Ok(OggPage::from_u8_slice_unchecked(buffer))
    }

    pub fn new_mut(buf: &mut [u8]) -> Result<&mut OggPage, OggPageCheckError> {
        let page_length = {
            let (hbuf, bbuf) = try!(OggPage::measure(buf));
            hbuf.len() + bbuf.len()
        };
        Ok(OggPage::from_u8_slice_unchecked_mut(&mut buf[0..page_length]))
    }

    fn measure(buf: &[u8]) -> Result<(&[u8], &[u8]), OggPageCheckError> {
        impl From<ByteOrderError> for OggPageCheckError {
            fn from(e: ByteOrderError) -> OggPageCheckError {
                match e {
                    ByteOrderError::UnexpectedEOF => OggPageCheckError::TooShort,
                    ByteOrderError::Io(_) => panic!("no I/O allowed"),
                }
            }
        }

        let mut cursor = Cursor::new(buf);

        if buf.len() < 27 {
            return Err(OggPageCheckError::TooShort);
        }
        if &buf[0..4] != OGG_PAGE_CAPTURE {
            return Err(OggPageCheckError::BadCapture);
        }

        cursor.consume(4);  // capture sequence
        if try!(cursor.read_u8()) != 0 {
            return Err(OggPageCheckError::BadVersion);
        }

        // flags(1) + granule(8) + serial(4) + page_seq(4) + csum(4)
        cursor.consume(1 + 8 + 4 + 4 + 4);

        let page_segments = try!(cursor.read_u8());

        let mut body_len = 0;
        for _ in 0..page_segments {
            body_len += try!(cursor.read_u8()) as u64;
        }

        let total_len = (cursor.position() + body_len) as usize;
        if buf.len() < total_len {
            return Err(OggPageCheckError::TooShort);
        }

        let h_end = cursor.position() as usize;
        let b_end = h_end + body_len as usize;
        Ok((
            &buf[0..h_end],
            &buf[h_end..b_end],
        ))
    }

    fn measure_whole(buf: &[u8]) -> Result<&[u8], OggPageCheckError> {
        let page_length = {
            let (h_buf, b_buf) = try!(OggPage::measure(buf));
            h_buf.len() + b_buf.len()
        };
        Ok(&buf[0..page_length])
    }

    fn measure_whole_mut(buf: &mut [u8]) -> Result<&mut [u8], OggPageCheckError> {
        let page_length = {
            let (hbuf, bbuf) = try!(OggPage::measure(buf));
            hbuf.len() + bbuf.len()
        };
        Ok(&mut buf[0..page_length])
    }

    pub fn position(&self) -> u64 {
        let self_buf = self.as_u8_slice();
        let mut cur = Cursor::new(&self_buf[POSITION_OFFSET..POSITION_OFFSET+8]);
        cur.read_u64::<LittleEndian>().unwrap()
    }

    pub fn set_position(&mut self, granule: u64) {
        let mut tx = self.begin();
        tx.set_position(granule);
    }

    pub fn serial(&self) -> u32 {
        let self_buf = self.as_u8_slice();
        let mut cur = Cursor::new(&self_buf[SERIAL_OFFSET..SERIAL_OFFSET+4]);
        cur.read_u32::<LittleEndian>().unwrap()
    }

    pub fn set_serial(&mut self, serial: u32) {
        let mut tx = self.begin();
        tx.set_serial(serial);
    }

    fn recompute_checksum(&mut self) {
        unimplemented!();
    }

    pub fn begin<'a>(&'a mut self) -> ChecksumGuard<'a> {
        ChecksumGuard {
            page: self,
            _marker: PhantomData,
        }
    }

    /// Am iterator of packet slices
    pub fn packets<'a>(&'a self) -> Packets<'a> {
        Packets { page: &self, packet: 0 }
    }

    pub fn into_cow<'a>(&'a self) -> Cow<'a, OggPage> {
        Cow::Borrowed(self)
    }
}

pub struct Packets<'a> {
    page: &'a OggPage,
    packet: usize,
}

impl<'a> Iterator for Packets<'a> {
    type Item = &'a [u8];

    fn next(&mut self) -> Option<&'a [u8]> {
        let packet_slice = self.page;
        unimplemented!();
    }
}

pub struct ChecksumGuard<'a> {
    page: &'a mut OggPage,
    _marker: PhantomData<&'a ()>,
}

impl<'a> ChecksumGuard<'a> {
    pub fn set_position(&mut self, granule: u64) {
        let self_buf = self.page.as_u8_slice_mut();
        let mut cur = Cursor::new(&mut self_buf[POSITION_OFFSET..POSITION_OFFSET+8]);
        cur.write_u64::<LittleEndian>(granule).unwrap();
    }

    pub fn set_serial(&mut self, serial: u32) {
        let self_buf = self.page.as_u8_slice_mut();
        let mut cur = Cursor::new(&mut self_buf[SERIAL_OFFSET..SERIAL_OFFSET+4]);
        cur.write_u32::<LittleEndian>(serial).unwrap();
    }
}

impl<'a> Drop for ChecksumGuard<'a> {
    fn drop(&mut self) {
        self.page.recompute_checksum();
    }
}

pub struct Recapture([u8; 4]);

impl Recapture {
    pub fn new() -> Recapture {
        Recapture([0; 4])
    }

    pub fn push_byte(&mut self, byte: u8) {
        let mut buf = [0; 4];
        buf[0] = self.0[1];
        buf[1] = self.0[2];
        buf[2] = self.0[3];
        buf[3] = byte;
        *self = Recapture(buf);
    }

    pub fn is_captured(&self) -> bool {
        &self.0 == OGG_PAGE_CAPTURE
    }
}

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

    #[test]
    fn test_capture() {
        let mut cap = Recapture::new();
        cap.push_byte(b'O');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'g');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'g');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'S');
        assert_eq!(true, cap.is_captured());
        cap.push_byte(b'g');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'g');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'S');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'O');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'g');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'g');
        assert_eq!(false, cap.is_captured());
        cap.push_byte(b'S');
        assert_eq!(true, cap.is_captured());
    }
}
<|endoftext|>


<commit_before>//! SCCS-style delta weave stores.

use crate::{
    node,
    store::{
        Store, StoreTags, StoreVersion, StoreWriter, TempCleaner, TempFile, TempLoader, Version,
    },
    Error, Result, SureNode,
};
use std::{
    env,
    fs::{self, File},
    io::{self, BufRead, BufReader, BufWriter, Write},
    path::{Path, PathBuf},
};
use weave::{
    self, DeltaWriter, NamingConvention, NewWeave, NullSink, Parser, PullParser, SimpleNaming,
};

pub struct WeaveStore {
    naming: SimpleNaming,
}

impl WeaveStore {
    pub fn new<P: AsRef<Path>>(path: P, base: &str, compressed: bool) -> WeaveStore {
        WeaveStore {
            naming: SimpleNaming::new(path, base, "dat", compressed),
        }
    }
}

impl Store for WeaveStore {
    fn get_versions(&self) -> Result<Vec<StoreVersion>> {
        let header = Parser::new(&self.naming, NullSink, 1)?.into_header();
        let mut versions: Vec<_> = header
            .deltas
            .iter()
            .map(|v| StoreVersion {
                name: v.name.clone(),
                time: v.time,
                version: Version::Tagged(v.number.to_string()),
            })
            .collect();
        versions.reverse();
        Ok(versions)
    }

    fn load_iter(&self, version: Version) -> Result<Box<dyn Iterator<Item = Result<SureNode>>>> {
        let last = weave::get_last_delta(&self.naming)?;
        let last = match version {
            Version::Latest => last,
            Version::Prior => last - 1,
            Version::Tagged(vers) => vers.parse()?,
        };

        Ok(Box::new(WeaveIter::new(&self.naming, last)?))
    }

    fn make_temp(&self) -> Result<Box<dyn TempFile + '_>> {
        // TODO: Fixup naming to allow uncompressed writes.
        let (path, file) = self.naming.temp_file()?;
        let cpath = path.clone();
        Ok(Box::new(WeaveTemp {
            parent: self,
            path,
            file: BufWriter::new(file),
            cleaner: FileClean(cpath),
        }))
    }

    fn make_new(&self, tags: &StoreTags) -> Result<Box<dyn StoreWriter + '_>> {
        let itags = tags.iter().map(|(k, v)| (k.as_ref(), v.as_ref()));
        match weave::get_last_delta(&self.naming) {
            Ok(base) => {
                let wv = DeltaWriter::new(&self.naming, itags, base)?;
                Ok(Box::new(NewWeaveDelta { weave: wv }))
            }
            Err(_) => {
                // Create a new weave file.
                let wv = NewWeave::new(&self.naming, itags)?;
                Ok(Box::new(NewWeaveWriter { weave: wv }))
            }
        }
    }
}

struct WeaveTemp<'a> {
    parent: &'a WeaveStore,
    path: PathBuf,
    file: BufWriter<File>,
    cleaner: FileClean,
}

impl<'a> TempFile<'a> for WeaveTemp<'a> {
    fn into_loader(self: Box<Self>) -> Result<Box<dyn TempLoader + 'a>> {
        drop(self.file);
        Ok(Box::new(WeaveTempLoader {
            _parent: self.parent,
            path: self.path,
            cleaner: self.cleaner,
        }))
    }

    fn into_cleaner(self: Box<Self>) -> Result<Box<dyn TempCleaner>> {
        Ok(Box::new(self.cleaner))
    }
}

impl<'a> Write for WeaveTemp<'a> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.file.write(buf)
    }
    fn flush(&mut self) -> io::Result<()> {
        self.file.flush()
    }
}

pub struct WeaveTempLoader<'a> {
    _parent: &'a WeaveStore,
    path: PathBuf,
    cleaner: FileClean,
}

impl<'a> TempLoader for WeaveTempLoader<'a> {
    fn new_loader(&self) -> Result<Box<dyn BufRead>> {
        let read = BufReader::new(File::open(&self.path)?);
        Ok(Box::new(read))
    }

    fn path_ref(&self) -> &Path {
        &self.path
    }

    fn into_cleaner(self: Box<Self>) -> Result<Box<dyn TempCleaner>> {
        Ok(Box::new(self.cleaner))
    }
}

pub struct NewWeaveWriter<'a> {
    weave: NewWeave<'a>,
}

impl<'a> StoreWriter<'a> for NewWeaveWriter<'a> {
    fn commit(self: Box<Self>) -> Result<()> {
        self.weave.close()?;
        Ok(())
    }
}

impl<'a> Write for NewWeaveWriter<'a> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.weave.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.weave.flush()
    }
}

pub struct NewWeaveDelta<'a> {
    weave: DeltaWriter<'a>,
}

impl<'a> StoreWriter<'a> for NewWeaveDelta<'a> {
    fn commit(self: Box<Self>) -> Result<()> {
        self.weave.close()?;
        Ok(())
    }
}

impl<'a> Write for NewWeaveDelta<'a> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.weave.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.weave.flush()
    }
}

pub struct WeaveIter {
    pull: Box<dyn Iterator<Item = Result<String>>>,
}

impl WeaveIter {
    fn new(naming: &dyn NamingConvention, delta: usize) -> Result<WeaveIter> {
        let mut pull = PullParser::new(naming, delta)?.filter_map(kept_text);
        fixed(&mut pull, "asure-2.0")?;
        fixed(&mut pull, "-----")?;
        Ok(WeaveIter {
            pull: Box::new(pull),
        })
    }
}

impl Iterator for WeaveIter {
    type Item = Result<SureNode>;

    fn next(&mut self) -> Option<Result<SureNode>> {
        let line = match self.pull.next() {
            Some(Err(e)) => return Some(Err(e)),
            Some(Ok(line)) => line,
            None => return None,
        };
        let line = line.as_bytes();

        match line[0] {
            b'd' => {
                let (dname, datts) = node::decode_entity(&line[1..]);
                Some(Ok(SureNode::Enter {
                    name: dname,
                    atts: datts,
                }))
            }
            b'f' => {
                let (fname, fatts) = node::decode_entity(&line[1..]);
                Some(Ok(SureNode::File {
                    name: fname,
                    atts: fatts,
                }))
            }
            b'-' => Some(Ok(SureNode::Sep)),
            b'u' => Some(Ok(SureNode::Leave)),
            ch => Some(Err(Error::InvalidSurefileChar(ch as char))),
        }
    }
}

// Filter nodes to only include kept text lines.
fn kept_text(node: weave::Result<weave::Entry>) -> Option<Result<String>> {
    match node {
        Err(e) => Some(Err(e.into())),
        Ok(weave::Entry::Plain { text, keep }) if keep => Some(Ok(text)),
        _ => None,
    }
}

/// Try reading a specific line from the given iterator.  Returns Err if
/// the line didn't match, or something went wrong with the read.
fn fixed<I>(pull: &mut I, expect: &str) -> Result<()>
where
    I: Iterator<Item = Result<String>>,
{
    match pull.next() {
        Some(Ok(line)) => {
            if line == expect {
                Ok(())
            } else {
                Err(Error::UnexpectedLine(line.into(), expect.into()))
            }
        }
        Some(Err(e)) => Err(e),
        None => Err(Error::SureFileEof),
    }
}
/*
fn fixed(recv: &Receiver<Option<Result<String>>>, expect: &[u8]) -> Result<()> {
    match recv.recv() {
        Ok(Some(Ok(line))) => {
            if line.as_bytes() == expect {
                Ok(())
            } else {
                Err(format_err!("Unexpect line from channel: {:?} expect {:?}", line, expect))
            }
        }
        Ok(Some(Err(e))) => Err(format_err!("Error reading suredata: {:?}", e)),
        Ok(None) => Err(format_err!("Unexpected eof reading suredata")),
        Err(e) => Err(e.into()),
    }
}
*/

/// Own a PathBuf, and delete this file on drop.  This is in its own type
/// for two reason. 1. I makes it easy to have cleaning in multiple types,
/// passing ownership between them, and 2.  It prevents the need for those
/// types to implement drop, which prevents moves out of the fields.
struct FileClean(PathBuf);

impl Drop for FileClean {
    fn drop(&mut self) {
        if env::var_os("RSURE_KEEP").is_none() {
            let _ = fs::remove_file(&self.0);
        }
    }
}

impl TempCleaner for FileClean {}
<commit_msg>clippy: Remove a redundant `.into()`<commit_after>//! SCCS-style delta weave stores.

use crate::{
    node,
    store::{
        Store, StoreTags, StoreVersion, StoreWriter, TempCleaner, TempFile, TempLoader, Version,
    },
    Error, Result, SureNode,
};
use std::{
    env,
    fs::{self, File},
    io::{self, BufRead, BufReader, BufWriter, Write},
    path::{Path, PathBuf},
};
use weave::{
    self, DeltaWriter, NamingConvention, NewWeave, NullSink, Parser, PullParser, SimpleNaming,
};

pub struct WeaveStore {
    naming: SimpleNaming,
}

impl WeaveStore {
    pub fn new<P: AsRef<Path>>(path: P, base: &str, compressed: bool) -> WeaveStore {
        WeaveStore {
            naming: SimpleNaming::new(path, base, "dat", compressed),
        }
    }
}

impl Store for WeaveStore {
    fn get_versions(&self) -> Result<Vec<StoreVersion>> {
        let header = Parser::new(&self.naming, NullSink, 1)?.into_header();
        let mut versions: Vec<_> = header
            .deltas
            .iter()
            .map(|v| StoreVersion {
                name: v.name.clone(),
                time: v.time,
                version: Version::Tagged(v.number.to_string()),
            })
            .collect();
        versions.reverse();
        Ok(versions)
    }

    fn load_iter(&self, version: Version) -> Result<Box<dyn Iterator<Item = Result<SureNode>>>> {
        let last = weave::get_last_delta(&self.naming)?;
        let last = match version {
            Version::Latest => last,
            Version::Prior => last - 1,
            Version::Tagged(vers) => vers.parse()?,
        };

        Ok(Box::new(WeaveIter::new(&self.naming, last)?))
    }

    fn make_temp(&self) -> Result<Box<dyn TempFile + '_>> {
        // TODO: Fixup naming to allow uncompressed writes.
        let (path, file) = self.naming.temp_file()?;
        let cpath = path.clone();
        Ok(Box::new(WeaveTemp {
            parent: self,
            path,
            file: BufWriter::new(file),
            cleaner: FileClean(cpath),
        }))
    }

    fn make_new(&self, tags: &StoreTags) -> Result<Box<dyn StoreWriter + '_>> {
        let itags = tags.iter().map(|(k, v)| (k.as_ref(), v.as_ref()));
        match weave::get_last_delta(&self.naming) {
            Ok(base) => {
                let wv = DeltaWriter::new(&self.naming, itags, base)?;
                Ok(Box::new(NewWeaveDelta { weave: wv }))
            }
            Err(_) => {
                // Create a new weave file.
                let wv = NewWeave::new(&self.naming, itags)?;
                Ok(Box::new(NewWeaveWriter { weave: wv }))
            }
        }
    }
}

struct WeaveTemp<'a> {
    parent: &'a WeaveStore,
    path: PathBuf,
    file: BufWriter<File>,
    cleaner: FileClean,
}

impl<'a> TempFile<'a> for WeaveTemp<'a> {
    fn into_loader(self: Box<Self>) -> Result<Box<dyn TempLoader + 'a>> {
        drop(self.file);
        Ok(Box::new(WeaveTempLoader {
            _parent: self.parent,
            path: self.path,
            cleaner: self.cleaner,
        }))
    }

    fn into_cleaner(self: Box<Self>) -> Result<Box<dyn TempCleaner>> {
        Ok(Box::new(self.cleaner))
    }
}

impl<'a> Write for WeaveTemp<'a> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.file.write(buf)
    }
    fn flush(&mut self) -> io::Result<()> {
        self.file.flush()
    }
}

pub struct WeaveTempLoader<'a> {
    _parent: &'a WeaveStore,
    path: PathBuf,
    cleaner: FileClean,
}

impl<'a> TempLoader for WeaveTempLoader<'a> {
    fn new_loader(&self) -> Result<Box<dyn BufRead>> {
        let read = BufReader::new(File::open(&self.path)?);
        Ok(Box::new(read))
    }

    fn path_ref(&self) -> &Path {
        &self.path
    }

    fn into_cleaner(self: Box<Self>) -> Result<Box<dyn TempCleaner>> {
        Ok(Box::new(self.cleaner))
    }
}

pub struct NewWeaveWriter<'a> {
    weave: NewWeave<'a>,
}

impl<'a> StoreWriter<'a> for NewWeaveWriter<'a> {
    fn commit(self: Box<Self>) -> Result<()> {
        self.weave.close()?;
        Ok(())
    }
}

impl<'a> Write for NewWeaveWriter<'a> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.weave.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.weave.flush()
    }
}

pub struct NewWeaveDelta<'a> {
    weave: DeltaWriter<'a>,
}

impl<'a> StoreWriter<'a> for NewWeaveDelta<'a> {
    fn commit(self: Box<Self>) -> Result<()> {
        self.weave.close()?;
        Ok(())
    }
}

impl<'a> Write for NewWeaveDelta<'a> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.weave.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.weave.flush()
    }
}

pub struct WeaveIter {
    pull: Box<dyn Iterator<Item = Result<String>>>,
}

impl WeaveIter {
    fn new(naming: &dyn NamingConvention, delta: usize) -> Result<WeaveIter> {
        let mut pull = PullParser::new(naming, delta)?.filter_map(kept_text);
        fixed(&mut pull, "asure-2.0")?;
        fixed(&mut pull, "-----")?;
        Ok(WeaveIter {
            pull: Box::new(pull),
        })
    }
}

impl Iterator for WeaveIter {
    type Item = Result<SureNode>;

    fn next(&mut self) -> Option<Result<SureNode>> {
        let line = match self.pull.next() {
            Some(Err(e)) => return Some(Err(e)),
            Some(Ok(line)) => line,
            None => return None,
        };
        let line = line.as_bytes();

        match line[0] {
            b'd' => {
                let (dname, datts) = node::decode_entity(&line[1..]);
                Some(Ok(SureNode::Enter {
                    name: dname,
                    atts: datts,
                }))
            }
            b'f' => {
                let (fname, fatts) = node::decode_entity(&line[1..]);
                Some(Ok(SureNode::File {
                    name: fname,
                    atts: fatts,
                }))
            }
            b'-' => Some(Ok(SureNode::Sep)),
            b'u' => Some(Ok(SureNode::Leave)),
            ch => Some(Err(Error::InvalidSurefileChar(ch as char))),
        }
    }
}

// Filter nodes to only include kept text lines.
fn kept_text(node: weave::Result<weave::Entry>) -> Option<Result<String>> {
    match node {
        Err(e) => Some(Err(e.into())),
        Ok(weave::Entry::Plain { text, keep }) if keep => Some(Ok(text)),
        _ => None,
    }
}

/// Try reading a specific line from the given iterator.  Returns Err if
/// the line didn't match, or something went wrong with the read.
fn fixed<I>(pull: &mut I, expect: &str) -> Result<()>
where
    I: Iterator<Item = Result<String>>,
{
    match pull.next() {
        Some(Ok(line)) => {
            if line == expect {
                Ok(())
            } else {
                Err(Error::UnexpectedLine(line, expect.into()))
            }
        }
        Some(Err(e)) => Err(e),
        None => Err(Error::SureFileEof),
    }
}
/*
fn fixed(recv: &Receiver<Option<Result<String>>>, expect: &[u8]) -> Result<()> {
    match recv.recv() {
        Ok(Some(Ok(line))) => {
            if line.as_bytes() == expect {
                Ok(())
            } else {
                Err(format_err!("Unexpect line from channel: {:?} expect {:?}", line, expect))
            }
        }
        Ok(Some(Err(e))) => Err(format_err!("Error reading suredata: {:?}", e)),
        Ok(None) => Err(format_err!("Unexpected eof reading suredata")),
        Err(e) => Err(e.into()),
    }
}
*/

/// Own a PathBuf, and delete this file on drop.  This is in its own type
/// for two reason. 1. I makes it easy to have cleaning in multiple types,
/// passing ownership between them, and 2.  It prevents the need for those
/// types to implement drop, which prevents moves out of the fields.
struct FileClean(PathBuf);

impl Drop for FileClean {
    fn drop(&mut self) {
        if env::var_os("RSURE_KEEP").is_none() {
            let _ = fs::remove_file(&self.0);
        }
    }
}

impl TempCleaner for FileClean {}
<|endoftext|>


<commit_before>extern crate fitsio;
extern crate ndarray;

use std::ops::{Deref, Range};
use fitsio::fitsfile::ReadImage;
use fitsio::types::HduInfo;
use fitsio::{FitsFile, FitsHdu};
use fitsio::errors::Result;

// Newtype so we can implement ReadImage
#[derive(Debug, Default)]
pub struct NdArray(ndarray::Array2<u32>);

impl Deref for NdArray {
    type Target = ndarray::Array2<u32>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl ReadImage for NdArray {
    fn read_section(
        _fits_file: &mut FitsFile,
        _hdu: &FitsHdu,
        _range: Range<usize>,
    ) -> Result<Self> {
        unimplemented!()
    }

    fn read_rows(
        _fits_file: &mut FitsFile,
        _hdu: &FitsHdu,
        _start_row: usize,
        _num_rows: usize,
    ) -> Result<Self> {
        unimplemented!()
    }

    fn read_row(_fits_file: &mut FitsFile, _hdu: &FitsHdu, _row: usize) -> Result<Self> {
        unimplemented!()
    }

    fn read_region(
        _fits_file: &mut FitsFile,
        _hdu: &FitsHdu,
        _ranges: &[&Range<usize>],
    ) -> Result<Self> {
        unimplemented!()
    }

    fn read_image(fits_file: &mut FitsFile, hdu: &FitsHdu) -> Result<Self> {
        let data: Vec<u32> = ReadImage::read_image(fits_file, hdu)?;
        let arr = ndarray::Array::from_vec(data);
        match hdu.info {
            HduInfo::ImageInfo { ref shape, .. } => {
                Ok(NdArray(arr.into_shape((shape[0], shape[1])).unwrap()))
            }
            _ => unreachable!(),
        }
    }
}

fn run() -> std::result::Result<(), Box<std::error::Error>> {
    let mut f = FitsFile::open("../testdata/full_example.fits")?;
    let hdu = f.primary_hdu()?;

    /* Check read_image */
    let data: NdArray = hdu.read_image(&mut f)?;
    assert_eq!(data.dim(), (100, 100));
    assert_eq!(data[[20, 5]], 152);

    /* Check read_row */

    Ok(())
}

fn main() {
    run().unwrap();
}
<commit_msg>Initially implement read_rows<commit_after>extern crate fitsio;
extern crate ndarray;

use std::ops::{Deref, Range};
use fitsio::fitsfile::ReadImage;
use fitsio::types::HduInfo;
use fitsio::{FitsFile, FitsHdu};
use fitsio::errors::Result;

// Newtype so we can implement ReadImage
#[derive(Debug, Default)]
pub struct NdArray(ndarray::Array2<u32>);

impl Deref for NdArray {
    type Target = ndarray::Array2<u32>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl ReadImage for NdArray {
    fn read_section(
        _fits_file: &mut FitsFile,
        _hdu: &FitsHdu,
        _range: Range<usize>,
    ) -> Result<Self> {
        unimplemented!()
    }

    fn read_rows(
        fits_file: &mut FitsFile,
        hdu: &FitsHdu,
        start_row: usize,
        num_rows: usize,
    ) -> Result<Self> {
        if num_rows == 0 || num_rows == 1 {
            unimplemented!("not implemented for ndarray::Array2")
        }
        let data: Vec<u32> = ReadImage::read_rows(fits_file, hdu, start_row, num_rows)?;
        let arr = ndarray::Array::from_vec(data);
        let row_length = arr.len() / num_rows;
        Ok(NdArray(arr.into_shape((num_rows, row_length)).unwrap()))
    }

    fn read_row(_fits_file: &mut FitsFile, _hdu: &FitsHdu, _row: usize) -> Result<Self> {
        unimplemented!()
    }

    fn read_region(
        _fits_file: &mut FitsFile,
        _hdu: &FitsHdu,
        _ranges: &[&Range<usize>],
    ) -> Result<Self> {
        unimplemented!()
    }

    fn read_image(fits_file: &mut FitsFile, hdu: &FitsHdu) -> Result<Self> {
        let data: Vec<u32> = ReadImage::read_image(fits_file, hdu)?;
        let arr = ndarray::Array::from_vec(data);
        match hdu.info {
            HduInfo::ImageInfo { ref shape, .. } => {
                Ok(NdArray(arr.into_shape((shape[0], shape[1])).unwrap()))
            }
            _ => unreachable!(),
        }
    }
}

fn run() -> std::result::Result<(), Box<std::error::Error>> {
    let mut f = FitsFile::open("../testdata/full_example.fits")?;
    let hdu = f.primary_hdu()?;

    {
        /* Check read_image */
        let data: NdArray = hdu.read_image(&mut f)?;
        assert_eq!(data.dim(), (100, 100));
        assert_eq!(data[[20, 5]], 152);
    }

    {
        /* Check read_rows */
        let data: NdArray = hdu.read_rows(&mut f, 0, 2)?;
        assert_eq!(data.dim(), (2, 100));
        assert_eq!(data[[1, 52]], 184);
    }

    Ok(())
}

fn main() {
    run().unwrap();
}
<|endoftext|>