<commit_before>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Resource management and resolution for the virtio-video device.

use std::convert::TryInto;
use std::fmt;

use base::{
    self, FromRawDescriptor, IntoRawDescriptor, MemoryMapping, MemoryMappingArena, MmapError,
    SafeDescriptor,
};
use vm_memory::{GuestAddress, GuestMemory, GuestMemoryError};

use thiserror::Error as ThisError;

use crate::virtio::resource_bridge::{self, ResourceBridgeError, ResourceInfo, ResourceRequest};
use crate::virtio::video::format::{FramePlane, PlaneFormat};
use crate::virtio::video::protocol::{virtio_video_mem_entry, virtio_video_object_entry};

/// Defines how resources for a given queue are represented.
#[derive(Clone, Copy, Debug)]
pub enum ResourceType {
    /// Resources are backed by guest memory pages.
    GuestPages,
    /// Resources are backed by virtio objects.
    VirtioObject,
}

impl Default for ResourceType {
    fn default() -> Self {
        ResourceType::VirtioObject
    }
}

#[repr(C)]
#[derive(Clone, Copy)]
/// A guest resource entry which type is not decided yet.
pub union UnresolvedResourceEntry {
    pub object: virtio_video_object_entry,
    pub guest_mem: virtio_video_mem_entry,
}
unsafe impl data_model::DataInit for UnresolvedResourceEntry {}

impl fmt::Debug for UnresolvedResourceEntry {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // Safe because `self.object` and `self.guest_mem` are the same size and both made of
        // integers, making it safe to display them no matter their value.
        write!(
            f,
            "unresolved {:?} or {:?}",
            unsafe { self.object },
            unsafe { self.guest_mem }
        )
    }
}

/// Trait for types that can serve as video buffer backing memory.
pub trait BufferHandle: Sized {
    /// Try to clone this handle. This must only create a new reference to the same backing memory
    /// and not duplicate the buffer itself.
    fn try_clone(&self) -> Result<Self, base::Error>;

    /// Returns a linear mapping of [`offset`..`offset`+`size`] of the memory backing this buffer.
    fn get_mapping(&self, offset: usize, size: usize) -> Result<MemoryMappingArena, MmapError>;
}

/// Linear memory area of a `GuestMemHandle`
#[derive(Clone)]
pub struct GuestMemArea {
    /// Offset within the guest region to the start of the area.
    pub offset: u64,
    /// Length of the area within the memory region.
    pub length: usize,
}

pub struct GuestMemHandle {
    /// Descriptor to the guest memory region containing the buffer.
    pub desc: SafeDescriptor,
    /// Memory areas (i.e. sg list) that make the memory buffer.
    pub mem_areas: Vec<GuestMemArea>,
}

impl BufferHandle for GuestMemHandle {
    fn try_clone(&self) -> Result<Self, base::Error> {
        Ok(Self {
            desc: self.desc.try_clone()?,
            mem_areas: self.mem_areas.clone(),
        })
    }

    fn get_mapping(&self, offset: usize, size: usize) -> Result<MemoryMappingArena, MmapError> {
        let mut arena = MemoryMappingArena::new(size)?;
        let mut mapped_size = 0;
        let mut area_iter = self.mem_areas.iter();
        let mut area_offset = offset;
        while mapped_size < size {
            let area = match area_iter.next() {
                Some(area) => area,
                None => {
                    return Err(MmapError::InvalidRange(
                        offset,
                        size,
                        self.mem_areas.iter().map(|a| a.length).sum(),
                    ));
                }
            };
            if area_offset > area.length {
                area_offset -= area.length;
            } else {
                let mapping_length = std::cmp::min(area.length - area_offset, size - mapped_size);
                arena.add_fd_offset(mapped_size, mapping_length, &self.desc, area.offset)?;
                mapped_size += mapping_length;
                area_offset = 0;
            }
        }
        Ok(arena)
    }
}

pub struct VirtioObjectHandle {
    /// Descriptor for the object.
    pub desc: SafeDescriptor,
    /// Modifier to apply to frame resources.
    pub modifier: u64,
}

impl BufferHandle for VirtioObjectHandle {
    fn try_clone(&self) -> Result<Self, base::Error> {
        Ok(Self {
            desc: self.desc.try_clone()?,
            modifier: self.modifier,
        })
    }

    fn get_mapping(&self, offset: usize, size: usize) -> Result<MemoryMappingArena, MmapError> {
        MemoryMapping::from_fd_offset(&self.desc, size, offset as u64).map(MemoryMappingArena::from)
    }
}

pub enum GuestResourceHandle {
    GuestPages(GuestMemHandle),
    VirtioObject(VirtioObjectHandle),
}

impl BufferHandle for GuestResourceHandle {
    fn try_clone(&self) -> Result<Self, base::Error> {
        Ok(match self {
            Self::GuestPages(handle) => Self::GuestPages(handle.try_clone()?),
            Self::VirtioObject(handle) => Self::VirtioObject(handle.try_clone()?),
        })
    }

    fn get_mapping(&self, offset: usize, size: usize) -> Result<MemoryMappingArena, MmapError> {
        match self {
            GuestResourceHandle::GuestPages(handle) => handle.get_mapping(offset, size),
            GuestResourceHandle::VirtioObject(handle) => handle.get_mapping(offset, size),
        }
    }
}

pub struct GuestResource {
    /// Handle to the backing memory.
    pub handle: GuestResourceHandle,
    /// Layout of color planes, if the resource will receive frames.
    pub planes: Vec<FramePlane>,
}

#[derive(Debug, ThisError)]
pub enum GuestMemResourceCreationError {
    #[error("Provided slice of entries is empty")]
    NoEntriesProvided,
    #[error("cannot get shm region: {0}")]
    CantGetShmRegion(GuestMemoryError),
    #[error("cannot get shm offset: {0}")]
    CantGetShmOffset(GuestMemoryError),
    #[error("error while cloning shm region descriptor: {0}")]
    DescriptorCloneError(base::Error),
}

#[derive(Debug, ThisError)]
pub enum ObjectResourceCreationError {
    #[error("uuid {0:08} is larger than 32 bits")]
    UuidNot32Bits(u128),
    #[error("resource returned by bridge is not a buffer")]
    NotABuffer,
    #[error("resource bridge failure: {0}")]
    ResourceBridgeFailure(ResourceBridgeError),
}

impl GuestResource {
    /// Try to convert an unresolved virtio guest memory entry into a resolved guest memory
    /// resource.
    ///
    /// Convert `mem_entry` into the guest memory resource it represents and resolve it through
    /// `mem`. `planes_format` describes the format of the individual planes for the buffer.
    pub fn from_virtio_guest_mem_entry(
        mem_entries: &[virtio_video_mem_entry],
        mem: &GuestMemory,
        planes_format: &[PlaneFormat],
    ) -> Result<GuestResource, GuestMemResourceCreationError> {
        let region_desc = match mem_entries.first() {
            None => return Err(GuestMemResourceCreationError::NoEntriesProvided),
            Some(entry) => {
                let addr: u64 = entry.addr.into();

                let guest_region = mem
                    .shm_region(GuestAddress(addr))
                    .map_err(GuestMemResourceCreationError::CantGetShmRegion)?;
                let desc = base::clone_descriptor(guest_region)
                    .map_err(GuestMemResourceCreationError::DescriptorCloneError)?;
                // Safe because we are the sole owner of the duplicated descriptor.
                unsafe { SafeDescriptor::from_raw_descriptor(desc) }
            }
        };

        let mem_areas = mem_entries
            .into_iter()
            .map(|entry| {
                let addr: u64 = entry.addr.into();
                let length: u32 = entry.length.into();
                let region_offset = mem
                    .offset_from_base(GuestAddress(addr))
                    .map_err(GuestMemResourceCreationError::CantGetShmOffset)
                    .unwrap();

                GuestMemArea {
                    offset: region_offset,
                    length: length as usize,
                }
            })
            .collect();

        // The plane information can be computed from the currently set format.
        let mut buffer_offset = 0;
        let planes = planes_format
            .iter()
            .map(|p| {
                let plane_offset = buffer_offset;
                buffer_offset += p.plane_size;

                FramePlane {
                    offset: plane_offset as usize,
                    stride: p.stride as usize,
                }
            })
            .collect();

        Ok(GuestResource {
            handle: GuestResourceHandle::GuestPages(GuestMemHandle {
                desc: region_desc,
                mem_areas,
            }),
            planes,
        })
    }

    /// Try to convert an unresolved virtio object entry into a resolved object resource.
    ///
    /// Convert `object` into the object resource it represents and resolve it through `res_bridge`.
    /// Returns an error if the object's UUID is invalid or cannot be resolved to a buffer object
    /// by `res_bridge`.
    pub fn from_virtio_object_entry(
        object: virtio_video_object_entry,
        res_bridge: &base::Tube,
    ) -> Result<GuestResource, ObjectResourceCreationError> {
        // We trust that the caller has chosen the correct object type.
        let uuid = u128::from_be_bytes(object.uuid);

        // TODO(stevensd): `Virtio3DBackend::resource_assign_uuid` is currently implemented to use
        // 32-bits resource_handles as UUIDs. Once it starts using real UUIDs, we need to update
        // this conversion.
        let handle = TryInto::<u32>::try_into(uuid)
            .map_err(|_| ObjectResourceCreationError::UuidNot32Bits(uuid))?;

        let buffer_info = match resource_bridge::get_resource_info(
            res_bridge,
            ResourceRequest::GetBuffer { id: handle },
        ) {
            Ok(ResourceInfo::Buffer(buffer_info)) => buffer_info,
            Ok(_) => return Err(ObjectResourceCreationError::NotABuffer),
            Err(e) => return Err(ObjectResourceCreationError::ResourceBridgeFailure(e)),
        };

        Ok(GuestResource {
            handle: GuestResourceHandle::VirtioObject(VirtioObjectHandle {
                // Safe because `buffer_info.file` is a valid file descriptor and we are stealing
                // it.
                desc: unsafe {
                    SafeDescriptor::from_raw_descriptor(buffer_info.file.into_raw_descriptor())
                },
                modifier: buffer_info.modifier,
            }),
            planes: buffer_info
                .planes
                .iter()
                .take_while(|p| p.offset != 0 || p.stride != 0)
                .map(|p| FramePlane {
                    offset: p.offset as usize,
                    stride: p.stride as usize,
                })
                .collect(),
        })
    }

    #[cfg(feature = "video-encoder")]
    pub fn try_clone(&self) -> Result<Self, base::Error> {
        Ok(Self {
            handle: self.handle.try_clone()?,
            planes: self.planes.clone(),
        })
    }
}

#[cfg(test)]
mod tests {
    use base::{MappedRegion, SafeDescriptor};
    use base::{MemoryMapping, SharedMemory};

    use super::*;

    /// Creates a sparse guest memory handle using as many pages as there are entries in
    /// `page_order`. The page with index `0` will be the first page, `1` will be the second page,
    /// etc.
    ///
    /// The memory handle is filled with increasing u32s starting from page 0, then page 1, and so
    /// on. Finally the handle is mapped into a linear space and we check that the written integers
    /// appear in the expected order.
    fn check_guest_mem_handle(page_order: &[usize]) {
        const PAGE_SIZE: usize = 0x1000;
        const U32_SIZE: usize = std::mem::size_of::<u32>();
        const ENTRIES_PER_PAGE: usize = PAGE_SIZE as usize / std::mem::size_of::<u32>();

        // Fill a vector of the same size as the handle with u32s of increasing value, following
        // the page layout given as argument.
        let mut data = vec![0u8; PAGE_SIZE * page_order.len()];
        for (page_index, page) in page_order.iter().enumerate() {
            let page_slice =
                &mut data[(page * PAGE_SIZE as usize)..((page + 1) * PAGE_SIZE as usize)];
            for (index, chunk) in page_slice.chunks_exact_mut(4).enumerate() {
                let sized_chunk: &mut [u8; 4] = chunk.try_into().unwrap();
                *sized_chunk = (((page_index * ENTRIES_PER_PAGE) + index) as u32).to_ne_bytes();
            }
        }

        // Copy the initialized vector's content into an anonymous shared memory.
        let mut mem = SharedMemory::anon().unwrap();
        mem.set_size(data.len() as u64).unwrap();
        let mapping = MemoryMapping::from_fd(&mem, mem.size() as usize).unwrap();
        assert_eq!(mapping.write_slice(&data, 0).unwrap(), data.len());

        // Create the `GuestMemHandle` we will try to map and retrieve the data from.
        let mem_handle = GuestResourceHandle::GuestPages(GuestMemHandle {
            desc: unsafe { SafeDescriptor::from_raw_descriptor(base::clone_fd(&mem).unwrap()) },
            mem_areas: page_order
                .iter()
                .map(|&page| GuestMemArea {
                    offset: page as u64 * PAGE_SIZE as u64,
                    length: PAGE_SIZE as usize,
                })
                .collect(),
        });

        // Map the handle into a linear memory area, retrieve its data into a new vector, and check
        // that its u32s appear to increase linearly.
        let mapping = mem_handle.get_mapping(0, mem.size() as usize).unwrap();
        let mut data = vec![0u8; PAGE_SIZE * page_order.len()];
        unsafe { std::ptr::copy_nonoverlapping(mapping.as_ptr(), data.as_mut_ptr(), data.len()) };
        for (index, chunk) in data.chunks_exact(U32_SIZE).enumerate() {
            let sized_chunk: &[u8; 4] = chunk.try_into().unwrap();
            assert_eq!(u32::from_ne_bytes(*sized_chunk), index as u32);
        }
    }

    // Fill a guest memory handle with a single memory page.
    // Then check that the data can be properly mapped and appears in the expected order.
    #[test]
    fn test_single_guest_mem_handle() {
        check_guest_mem_handle(&[0])
    }

    // Fill a guest memory handle with 4 memory pages that are contiguous.
    // Then check that the pages appear in the expected order in the mapping.
    #[test]
    fn test_linear_guest_mem_handle() {
        check_guest_mem_handle(&[0, 1, 2, 3])
    }

    // Fill a guest memory handle with 8 pages mapped in non-linear order.
    // Then check that the pages appear in the expected order in the mapping.
    #[test]
    fn test_sparse_guest_mem_handle() {
        check_guest_mem_handle(&[1, 7, 6, 3, 5, 0, 4, 2])
    }
}
<commit_msg>devices: virtio: video: fix build with base MemoryMapping<commit_after>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Resource management and resolution for the virtio-video device.

use std::convert::TryInto;
use std::fmt;

use base::{
    self, FromRawDescriptor, IntoRawDescriptor, MemoryMappingArena, MemoryMappingBuilder,
    MemoryMappingBuilderUnix, MmapError, SafeDescriptor,
};
use vm_memory::{GuestAddress, GuestMemory, GuestMemoryError};

use thiserror::Error as ThisError;

use crate::virtio::resource_bridge::{self, ResourceBridgeError, ResourceInfo, ResourceRequest};
use crate::virtio::video::format::{FramePlane, PlaneFormat};
use crate::virtio::video::protocol::{virtio_video_mem_entry, virtio_video_object_entry};

/// Defines how resources for a given queue are represented.
#[derive(Clone, Copy, Debug)]
pub enum ResourceType {
    /// Resources are backed by guest memory pages.
    GuestPages,
    /// Resources are backed by virtio objects.
    VirtioObject,
}

impl Default for ResourceType {
    fn default() -> Self {
        ResourceType::VirtioObject
    }
}

#[repr(C)]
#[derive(Clone, Copy)]
/// A guest resource entry which type is not decided yet.
pub union UnresolvedResourceEntry {
    pub object: virtio_video_object_entry,
    pub guest_mem: virtio_video_mem_entry,
}
unsafe impl data_model::DataInit for UnresolvedResourceEntry {}

impl fmt::Debug for UnresolvedResourceEntry {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // Safe because `self.object` and `self.guest_mem` are the same size and both made of
        // integers, making it safe to display them no matter their value.
        write!(
            f,
            "unresolved {:?} or {:?}",
            unsafe { self.object },
            unsafe { self.guest_mem }
        )
    }
}

/// Trait for types that can serve as video buffer backing memory.
pub trait BufferHandle: Sized {
    /// Try to clone this handle. This must only create a new reference to the same backing memory
    /// and not duplicate the buffer itself.
    fn try_clone(&self) -> Result<Self, base::Error>;

    /// Returns a linear mapping of [`offset`..`offset`+`size`] of the memory backing this buffer.
    fn get_mapping(&self, offset: usize, size: usize) -> Result<MemoryMappingArena, MmapError>;
}

/// Linear memory area of a `GuestMemHandle`
#[derive(Clone)]
pub struct GuestMemArea {
    /// Offset within the guest region to the start of the area.
    pub offset: u64,
    /// Length of the area within the memory region.
    pub length: usize,
}

pub struct GuestMemHandle {
    /// Descriptor to the guest memory region containing the buffer.
    pub desc: SafeDescriptor,
    /// Memory areas (i.e. sg list) that make the memory buffer.
    pub mem_areas: Vec<GuestMemArea>,
}

impl BufferHandle for GuestMemHandle {
    fn try_clone(&self) -> Result<Self, base::Error> {
        Ok(Self {
            desc: self.desc.try_clone()?,
            mem_areas: self.mem_areas.clone(),
        })
    }

    fn get_mapping(&self, offset: usize, size: usize) -> Result<MemoryMappingArena, MmapError> {
        let mut arena = MemoryMappingArena::new(size)?;
        let mut mapped_size = 0;
        let mut area_iter = self.mem_areas.iter();
        let mut area_offset = offset;
        while mapped_size < size {
            let area = match area_iter.next() {
                Some(area) => area,
                None => {
                    return Err(MmapError::InvalidRange(
                        offset,
                        size,
                        self.mem_areas.iter().map(|a| a.length).sum(),
                    ));
                }
            };
            if area_offset > area.length {
                area_offset -= area.length;
            } else {
                let mapping_length = std::cmp::min(area.length - area_offset, size - mapped_size);
                arena.add_fd_offset(mapped_size, mapping_length, &self.desc, area.offset)?;
                mapped_size += mapping_length;
                area_offset = 0;
            }
        }
        Ok(arena)
    }
}

pub struct VirtioObjectHandle {
    /// Descriptor for the object.
    pub desc: SafeDescriptor,
    /// Modifier to apply to frame resources.
    pub modifier: u64,
}

impl BufferHandle for VirtioObjectHandle {
    fn try_clone(&self) -> Result<Self, base::Error> {
        Ok(Self {
            desc: self.desc.try_clone()?,
            modifier: self.modifier,
        })
    }

    fn get_mapping(&self, offset: usize, size: usize) -> Result<MemoryMappingArena, MmapError> {
        MemoryMappingBuilder::new(size)
            .from_descriptor(&self.desc)
            .offset(offset as u64)
            .build()
            .map(MemoryMappingArena::from)
    }
}

pub enum GuestResourceHandle {
    GuestPages(GuestMemHandle),
    VirtioObject(VirtioObjectHandle),
}

impl BufferHandle for GuestResourceHandle {
    fn try_clone(&self) -> Result<Self, base::Error> {
        Ok(match self {
            Self::GuestPages(handle) => Self::GuestPages(handle.try_clone()?),
            Self::VirtioObject(handle) => Self::VirtioObject(handle.try_clone()?),
        })
    }

    fn get_mapping(&self, offset: usize, size: usize) -> Result<MemoryMappingArena, MmapError> {
        match self {
            GuestResourceHandle::GuestPages(handle) => handle.get_mapping(offset, size),
            GuestResourceHandle::VirtioObject(handle) => handle.get_mapping(offset, size),
        }
    }
}

pub struct GuestResource {
    /// Handle to the backing memory.
    pub handle: GuestResourceHandle,
    /// Layout of color planes, if the resource will receive frames.
    pub planes: Vec<FramePlane>,
}

#[derive(Debug, ThisError)]
pub enum GuestMemResourceCreationError {
    #[error("Provided slice of entries is empty")]
    NoEntriesProvided,
    #[error("cannot get shm region: {0}")]
    CantGetShmRegion(GuestMemoryError),
    #[error("cannot get shm offset: {0}")]
    CantGetShmOffset(GuestMemoryError),
    #[error("error while cloning shm region descriptor: {0}")]
    DescriptorCloneError(base::Error),
}

#[derive(Debug, ThisError)]
pub enum ObjectResourceCreationError {
    #[error("uuid {0:08} is larger than 32 bits")]
    UuidNot32Bits(u128),
    #[error("resource returned by bridge is not a buffer")]
    NotABuffer,
    #[error("resource bridge failure: {0}")]
    ResourceBridgeFailure(ResourceBridgeError),
}

impl GuestResource {
    /// Try to convert an unresolved virtio guest memory entry into a resolved guest memory
    /// resource.
    ///
    /// Convert `mem_entry` into the guest memory resource it represents and resolve it through
    /// `mem`. `planes_format` describes the format of the individual planes for the buffer.
    pub fn from_virtio_guest_mem_entry(
        mem_entries: &[virtio_video_mem_entry],
        mem: &GuestMemory,
        planes_format: &[PlaneFormat],
    ) -> Result<GuestResource, GuestMemResourceCreationError> {
        let region_desc = match mem_entries.first() {
            None => return Err(GuestMemResourceCreationError::NoEntriesProvided),
            Some(entry) => {
                let addr: u64 = entry.addr.into();

                let guest_region = mem
                    .shm_region(GuestAddress(addr))
                    .map_err(GuestMemResourceCreationError::CantGetShmRegion)?;
                let desc = base::clone_descriptor(guest_region)
                    .map_err(GuestMemResourceCreationError::DescriptorCloneError)?;
                // Safe because we are the sole owner of the duplicated descriptor.
                unsafe { SafeDescriptor::from_raw_descriptor(desc) }
            }
        };

        let mem_areas = mem_entries
            .into_iter()
            .map(|entry| {
                let addr: u64 = entry.addr.into();
                let length: u32 = entry.length.into();
                let region_offset = mem
                    .offset_from_base(GuestAddress(addr))
                    .map_err(GuestMemResourceCreationError::CantGetShmOffset)
                    .unwrap();

                GuestMemArea {
                    offset: region_offset,
                    length: length as usize,
                }
            })
            .collect();

        // The plane information can be computed from the currently set format.
        let mut buffer_offset = 0;
        let planes = planes_format
            .iter()
            .map(|p| {
                let plane_offset = buffer_offset;
                buffer_offset += p.plane_size;

                FramePlane {
                    offset: plane_offset as usize,
                    stride: p.stride as usize,
                }
            })
            .collect();

        Ok(GuestResource {
            handle: GuestResourceHandle::GuestPages(GuestMemHandle {
                desc: region_desc,
                mem_areas,
            }),
            planes,
        })
    }

    /// Try to convert an unresolved virtio object entry into a resolved object resource.
    ///
    /// Convert `object` into the object resource it represents and resolve it through `res_bridge`.
    /// Returns an error if the object's UUID is invalid or cannot be resolved to a buffer object
    /// by `res_bridge`.
    pub fn from_virtio_object_entry(
        object: virtio_video_object_entry,
        res_bridge: &base::Tube,
    ) -> Result<GuestResource, ObjectResourceCreationError> {
        // We trust that the caller has chosen the correct object type.
        let uuid = u128::from_be_bytes(object.uuid);

        // TODO(stevensd): `Virtio3DBackend::resource_assign_uuid` is currently implemented to use
        // 32-bits resource_handles as UUIDs. Once it starts using real UUIDs, we need to update
        // this conversion.
        let handle = TryInto::<u32>::try_into(uuid)
            .map_err(|_| ObjectResourceCreationError::UuidNot32Bits(uuid))?;

        let buffer_info = match resource_bridge::get_resource_info(
            res_bridge,
            ResourceRequest::GetBuffer { id: handle },
        ) {
            Ok(ResourceInfo::Buffer(buffer_info)) => buffer_info,
            Ok(_) => return Err(ObjectResourceCreationError::NotABuffer),
            Err(e) => return Err(ObjectResourceCreationError::ResourceBridgeFailure(e)),
        };

        Ok(GuestResource {
            handle: GuestResourceHandle::VirtioObject(VirtioObjectHandle {
                // Safe because `buffer_info.file` is a valid file descriptor and we are stealing
                // it.
                desc: unsafe {
                    SafeDescriptor::from_raw_descriptor(buffer_info.file.into_raw_descriptor())
                },
                modifier: buffer_info.modifier,
            }),
            planes: buffer_info
                .planes
                .iter()
                .take_while(|p| p.offset != 0 || p.stride != 0)
                .map(|p| FramePlane {
                    offset: p.offset as usize,
                    stride: p.stride as usize,
                })
                .collect(),
        })
    }

    #[cfg(feature = "video-encoder")]
    pub fn try_clone(&self) -> Result<Self, base::Error> {
        Ok(Self {
            handle: self.handle.try_clone()?,
            planes: self.planes.clone(),
        })
    }
}

#[cfg(test)]
mod tests {
    use base::{MappedRegion, SafeDescriptor};
    use base::{MemoryMapping, SharedMemory};

    use super::*;

    /// Creates a sparse guest memory handle using as many pages as there are entries in
    /// `page_order`. The page with index `0` will be the first page, `1` will be the second page,
    /// etc.
    ///
    /// The memory handle is filled with increasing u32s starting from page 0, then page 1, and so
    /// on. Finally the handle is mapped into a linear space and we check that the written integers
    /// appear in the expected order.
    fn check_guest_mem_handle(page_order: &[usize]) {
        const PAGE_SIZE: usize = 0x1000;
        const U32_SIZE: usize = std::mem::size_of::<u32>();
        const ENTRIES_PER_PAGE: usize = PAGE_SIZE as usize / std::mem::size_of::<u32>();

        // Fill a vector of the same size as the handle with u32s of increasing value, following
        // the page layout given as argument.
        let mut data = vec![0u8; PAGE_SIZE * page_order.len()];
        for (page_index, page) in page_order.iter().enumerate() {
            let page_slice =
                &mut data[(page * PAGE_SIZE as usize)..((page + 1) * PAGE_SIZE as usize)];
            for (index, chunk) in page_slice.chunks_exact_mut(4).enumerate() {
                let sized_chunk: &mut [u8; 4] = chunk.try_into().unwrap();
                *sized_chunk = (((page_index * ENTRIES_PER_PAGE) + index) as u32).to_ne_bytes();
            }
        }

        // Copy the initialized vector's content into an anonymous shared memory.
        let mut mem = SharedMemory::anon().unwrap();
        mem.set_size(data.len() as u64).unwrap();
        let mapping = MemoryMapping::from_fd(&mem, mem.size() as usize).unwrap();
        assert_eq!(mapping.write_slice(&data, 0).unwrap(), data.len());

        // Create the `GuestMemHandle` we will try to map and retrieve the data from.
        let mem_handle = GuestResourceHandle::GuestPages(GuestMemHandle {
            desc: unsafe { SafeDescriptor::from_raw_descriptor(base::clone_fd(&mem).unwrap()) },
            mem_areas: page_order
                .iter()
                .map(|&page| GuestMemArea {
                    offset: page as u64 * PAGE_SIZE as u64,
                    length: PAGE_SIZE as usize,
                })
                .collect(),
        });

        // Map the handle into a linear memory area, retrieve its data into a new vector, and check
        // that its u32s appear to increase linearly.
        let mapping = mem_handle.get_mapping(0, mem.size() as usize).unwrap();
        let mut data = vec![0u8; PAGE_SIZE * page_order.len()];
        unsafe { std::ptr::copy_nonoverlapping(mapping.as_ptr(), data.as_mut_ptr(), data.len()) };
        for (index, chunk) in data.chunks_exact(U32_SIZE).enumerate() {
            let sized_chunk: &[u8; 4] = chunk.try_into().unwrap();
            assert_eq!(u32::from_ne_bytes(*sized_chunk), index as u32);
        }
    }

    // Fill a guest memory handle with a single memory page.
    // Then check that the data can be properly mapped and appears in the expected order.
    #[test]
    fn test_single_guest_mem_handle() {
        check_guest_mem_handle(&[0])
    }

    // Fill a guest memory handle with 4 memory pages that are contiguous.
    // Then check that the pages appear in the expected order in the mapping.
    #[test]
    fn test_linear_guest_mem_handle() {
        check_guest_mem_handle(&[0, 1, 2, 3])
    }

    // Fill a guest memory handle with 8 pages mapped in non-linear order.
    // Then check that the pages appear in the expected order in the mapping.
    #[test]
    fn test_sparse_guest_mem_handle() {
        check_guest_mem_handle(&[1, 7, 6, 3, 5, 0, 4, 2])
    }
}
<|endoftext|>


<commit_before>use std::io;
use std::os::getenv;
use std::path::BytesContainer;
use std::str::from_utf8_lossy;
use get_environ = std::os::env;
use std::io::fs::{mkdir, rmdir_recursive, rename, symlink, unlink, readlink};
use std::io::process::{ExitStatus, Command, Ignored, InheritFd};
use std::io::stdio::{stdout, stderr};

use super::sha256::Sha256;
use super::sha256::Digest;

use argparse::{ArgumentParser, Store, StoreTrue};

use super::env::{Environ, Container};
use super::options::env_options;
use super::linux::ensure_dir;

fn makedirs(path: &Path) -> Result<(),String> {
    if path.exists() {
        return Ok(());
    }
    try!(makedirs(&path.dir_path()));
    return match mkdir(path, io::UserRWX) {
        Ok(()) => Ok(()),
        Err(e) => Err(format!("Can't mkdir: {}", e)),
    };
}

pub fn build_command(environ: &mut Environ, args: Vec<String>)
    -> Result<int, String>
{
    let mut cname = "devel".to_string();
    let mut force = false;
    {
        let mut ap = ArgumentParser::new();
        ap.refer(&mut cname)
            .add_argument("container", box Store::<String>,
                "A name of the container to build")
            .required();
        ap.refer(&mut force)
            .add_option(["--force"], box StoreTrue,
                "Force rebuild of container event if it already exists");
        env_options(environ, &mut ap);
        match ap.parse(args, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => return Ok(122),
        }
    }
    let mut container = try!(environ.get_container(&cname));
    try!(build_container(environ, &mut container, force));
    try!(link_container(environ, &container));

    return Ok(0);
}

pub fn link_container(environ: &Environ, container: &Container)
    -> Result<(), String>
{
    let container_root = container.container_root.as_ref().unwrap();
    let cdir = container_root.filename().unwrap();
    let tmptarget = environ.local_vagga.join(
        (".lnk.".container_into_owned_bytes() + cdir).as_slice());
    if tmptarget.exists() {
        match unlink(&tmptarget) {
            Ok(()) => {}
            Err(x) => return Err(format!("Error removing file: {}", x)),
        }
    }
    let target = environ.local_vagga.join(container.name.as_slice());
    debug!("Linking {} -> {}", container_root.display(), target.display())
    let relative_root = Path::new(".roots").join(cdir);
    match symlink(&relative_root, &tmptarget)
          .and(rename(&tmptarget, &target)) {
        Ok(()) => {}
        Err(x) => return Err(format!("Can't symlink new root: {}", x)),
    }
    if !container.name.eq(&container.fullname) {
        let target = environ.local_vagga.join(container.fullname.as_slice());
        debug!("Linking {} -> {}", container_root.display(), target.display())
        match symlink(&relative_root, &tmptarget)
              .and(rename(&tmptarget, &target)) {
            Ok(()) => {}
            Err(x) => return Err(format!("Can't symlink new root: {}", x)),
        }
    }
    return Ok(());
}

pub fn build_container(environ: &Environ, container: &mut Container,
                       force: bool)
    -> Result <bool, String>
{
    info!("Checking {}", container.name);

    let builder = &container.builder;
    let cache_dir = environ.local_vagga.join_many(
        [".cache", builder.as_slice()]);

    // TODO(tailhook) which path should be here?
    let path = getenv("PATH").unwrap_or(
        "/sbin:/bin:/usr/sbin:/usr/bin:/usr/local/bin:/usr/local/sbin"
        .to_string());

    let mut env: Vec<(&[u8], &[u8])> = Vec::new();
    let mut caller_env: Vec<(String, String)> = Vec::new();
    let mut parameters: Vec<(String, String)> = Vec::new();
    env.push(("HOME".as_bytes(), "/homeless-shelter".as_bytes()));
    env.push(("PATH".as_bytes(), path.as_bytes()));
    env.push(("container_name".as_bytes(), container.name.as_bytes()));
    env.push(("container_fullname".as_bytes(), container.fullname.as_bytes()));
    env.push(("cache_dir".as_bytes(), cache_dir.as_vec()));
    env.push(("project_root".as_bytes(), environ.project_root.as_vec()));
    for (k, v) in get_environ().move_iter() {
        let pk = "CALLER_".to_string() + k;
        let pv = v;
        caller_env.push((pk, pv));
    }
    for (k, v) in container.parameters.iter() {
        parameters.push((builder + "_" + *k, v.clone()));
    }
    for &(ref k, ref v) in caller_env.iter() {
        env.push((k.container_as_bytes(), v.container_as_bytes()));
    }
    for &(ref k, ref v) in parameters.iter() {
        env.push((k.container_as_bytes(), v.container_as_bytes()));
    }

    let version_sh = environ.vagga_path.join_many(
        ["backends", builder.as_slice(), "version.sh"]);
    let mut digest = Sha256::new();
    if version_sh.exists() {
        digest.input_str(builder.as_slice());
        digest.input_str(":");
        match Command::new(&version_sh).env(env.as_slice())
            .cwd(&environ.project_root)
            .stdin(Ignored).stderr(InheritFd(2)).output() {
            Ok(out) => match out.status {
                ExitStatus(0) => {
                    debug!("Version data ```{}```",
                           from_utf8_lossy(out.output.as_slice()));
                    digest.input(out.output.as_slice());
                }
                e => {
                    return Err(format!("Error running {}: {}",
                        version_sh.display(), e));
                }
            },
            Err(e) => {
                return Err(format!("Error running {}: {}",
                    version_sh.display(), e));
            }
        }
    } else {
        warn!("This backend has no versioning. Using primitive one.");
        digest.input_str(builder.as_slice());
        digest.input_str(":");
        for (k, v) in container.parameters.iter() {
            digest.input_str(k.as_slice());
            digest.input_str("=");
            digest.input_str(v.as_slice());
            digest.input_str(";");
        }
    }
    let fullhash = digest.result_str();
    let hash = fullhash.as_slice().slice_to(7);

    let build_sh = environ.vagga_path.join_many(
        ["backends", builder.as_slice(), "build.sh"]);
    if !build_sh.exists() {
        return Err(format!("Builder {} does not exist", builder));
    }

    let cdir = format!("{}.{}", container.fullname, hash);
    let container_root = environ.local_vagga.join_many(
        [".roots", cdir.as_slice()]);

    if container_root.exists() && !force {
        info!("Container {} already built as {}",
            container.name, container_root.display());
        container.container_root = Some(container_root);
        return Ok(false);
    }

    let artifacts_dir = environ.local_vagga.join_many(
        [".artifacts", cdir.as_slice()]);
    let container_tmp = environ.local_vagga.join_many(
        [".roots", (cdir + ".tmp").as_slice()]);

    info!("Building {} by {}", container_root.display(), build_sh.display());

    if container_tmp.exists() {
        match rmdir_recursive(&container_tmp) {
            Ok(()) => {}
            Err(x) => {
                // TODO(tailhook) join container and rm, because may have
                // files of different owners (subuid's)
                return Err(format!("Can't clean temporary root: {}", x));
            }
        }
    }
    try!(makedirs(&container_tmp));

    env.push(("artifacts_dir".as_bytes(), artifacts_dir.as_vec()));
    env.push(("container_hash".as_bytes(), hash.as_bytes()));
    env.push(("container_root".as_bytes(), container_tmp.as_vec()));
    match Command::new(build_sh).env(env.as_slice())
        .cwd(&environ.project_root)
        .stdin(InheritFd(0)).stdout(InheritFd(1)).stderr(InheritFd(2))
        .status() {
        Ok(ExitStatus(0)) => {}
        Ok(x) => return Err(format!("Builder exited with status {}", x)),
        Err(x) => return Err(format!("Can't spawn: {}", x)),
    }

    for dir in ["proc", "sys", "dev", "work", "tmp"].iter() {
        try!(ensure_dir(&container_tmp.join(*dir)));
    }

    let container_old = environ.local_vagga.join_many(
        [".roots", (cdir + ".old").as_slice()]);
    if container_root.exists() {
        if container_old.exists() {
            match rmdir_recursive(&container_old) {
                Ok(()) => {}
                Err(x) => return Err(format!("Can't remove old root: {}", x)),
            }
        }
        match rename(&container_root, &container_old) {
            Ok(()) => {}
            Err(x) => return Err(format!("Can't rename old root: {}", x)),
        }
    }

    match rename(&container_tmp, &container_root) {
        Ok(()) => {}
        Err(x) => return Err(format!("Can't rename root: {}", x)),
    }

    if container_old.exists() {
        match rmdir_recursive(&container_old) {
            Ok(()) => {}
            Err(x) => return Err(format!("Can't remove old root: {}", x)),
        }
    }
    info!("Done building {} as {}",
        container.name, container_root.display());

    container.container_root = Some(container_root);

    return Ok(true);
}

pub fn ensure_container(env: &Environ, container: &mut Container)
    -> Result<(), String>
{
    if env.settings.version_check {
        try!(build_container(env, container, false));
        try!(link_container(env, container));
    } else {
        let lnk = env.local_vagga.join(container.fullname.as_slice());
        container.container_root = match readlink(&lnk) {
            Ok(path) => Some(lnk.dir_path().join(path)),
            Err(e) => return Err(format!("Container {} not found: {}",
                                         container.fullname, e)),
        };
    };
    return Ok(());
}
<commit_msg>Clean `artifacts_dir` before building<commit_after>use std::io;
use std::os::getenv;
use std::path::BytesContainer;
use std::str::from_utf8_lossy;
use get_environ = std::os::env;
use std::io::fs::{mkdir, rmdir_recursive, rename, symlink, unlink, readlink};
use std::io::process::{ExitStatus, Command, Ignored, InheritFd};
use std::io::stdio::{stdout, stderr};

use super::sha256::Sha256;
use super::sha256::Digest;

use argparse::{ArgumentParser, Store, StoreTrue};

use super::env::{Environ, Container};
use super::options::env_options;
use super::linux::ensure_dir;

fn makedirs(path: &Path) -> Result<(),String> {
    if path.exists() {
        return Ok(());
    }
    try!(makedirs(&path.dir_path()));
    return match mkdir(path, io::UserRWX) {
        Ok(()) => Ok(()),
        Err(e) => Err(format!("Can't mkdir: {}", e)),
    };
}

pub fn build_command(environ: &mut Environ, args: Vec<String>)
    -> Result<int, String>
{
    let mut cname = "devel".to_string();
    let mut force = false;
    {
        let mut ap = ArgumentParser::new();
        ap.refer(&mut cname)
            .add_argument("container", box Store::<String>,
                "A name of the container to build")
            .required();
        ap.refer(&mut force)
            .add_option(["--force"], box StoreTrue,
                "Force rebuild of container event if it already exists");
        env_options(environ, &mut ap);
        match ap.parse(args, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => return Ok(122),
        }
    }
    let mut container = try!(environ.get_container(&cname));
    try!(build_container(environ, &mut container, force));
    try!(link_container(environ, &container));

    return Ok(0);
}

pub fn link_container(environ: &Environ, container: &Container)
    -> Result<(), String>
{
    let container_root = container.container_root.as_ref().unwrap();
    let cdir = container_root.filename().unwrap();
    let tmptarget = environ.local_vagga.join(
        (".lnk.".container_into_owned_bytes() + cdir).as_slice());
    if tmptarget.exists() {
        match unlink(&tmptarget) {
            Ok(()) => {}
            Err(x) => return Err(format!("Error removing file: {}", x)),
        }
    }
    let target = environ.local_vagga.join(container.name.as_slice());
    debug!("Linking {} -> {}", container_root.display(), target.display())
    let relative_root = Path::new(".roots").join(cdir);
    match symlink(&relative_root, &tmptarget)
          .and(rename(&tmptarget, &target)) {
        Ok(()) => {}
        Err(x) => return Err(format!("Can't symlink new root: {}", x)),
    }
    if !container.name.eq(&container.fullname) {
        let target = environ.local_vagga.join(container.fullname.as_slice());
        debug!("Linking {} -> {}", container_root.display(), target.display())
        match symlink(&relative_root, &tmptarget)
              .and(rename(&tmptarget, &target)) {
            Ok(()) => {}
            Err(x) => return Err(format!("Can't symlink new root: {}", x)),
        }
    }
    return Ok(());
}

pub fn build_container(environ: &Environ, container: &mut Container,
                       force: bool)
    -> Result <bool, String>
{
    info!("Checking {}", container.name);

    let builder = &container.builder;
    let cache_dir = environ.local_vagga.join_many(
        [".cache", builder.as_slice()]);

    // TODO(tailhook) which path should be here?
    let path = getenv("PATH").unwrap_or(
        "/sbin:/bin:/usr/sbin:/usr/bin:/usr/local/bin:/usr/local/sbin"
        .to_string());

    let mut env: Vec<(&[u8], &[u8])> = Vec::new();
    let mut caller_env: Vec<(String, String)> = Vec::new();
    let mut parameters: Vec<(String, String)> = Vec::new();
    env.push(("HOME".as_bytes(), "/homeless-shelter".as_bytes()));
    env.push(("PATH".as_bytes(), path.as_bytes()));
    env.push(("container_name".as_bytes(), container.name.as_bytes()));
    env.push(("container_fullname".as_bytes(), container.fullname.as_bytes()));
    env.push(("cache_dir".as_bytes(), cache_dir.as_vec()));
    env.push(("project_root".as_bytes(), environ.project_root.as_vec()));
    for (k, v) in get_environ().move_iter() {
        let pk = "CALLER_".to_string() + k;
        let pv = v;
        caller_env.push((pk, pv));
    }
    for (k, v) in container.parameters.iter() {
        parameters.push((builder + "_" + *k, v.clone()));
    }
    for &(ref k, ref v) in caller_env.iter() {
        env.push((k.container_as_bytes(), v.container_as_bytes()));
    }
    for &(ref k, ref v) in parameters.iter() {
        env.push((k.container_as_bytes(), v.container_as_bytes()));
    }

    let version_sh = environ.vagga_path.join_many(
        ["backends", builder.as_slice(), "version.sh"]);
    let mut digest = Sha256::new();
    if version_sh.exists() {
        digest.input_str(builder.as_slice());
        digest.input_str(":");
        match Command::new(&version_sh).env(env.as_slice())
            .cwd(&environ.project_root)
            .stdin(Ignored).stderr(InheritFd(2)).output() {
            Ok(out) => match out.status {
                ExitStatus(0) => {
                    debug!("Version data ```{}```",
                           from_utf8_lossy(out.output.as_slice()));
                    digest.input(out.output.as_slice());
                }
                e => {
                    return Err(format!("Error running {}: {}",
                        version_sh.display(), e));
                }
            },
            Err(e) => {
                return Err(format!("Error running {}: {}",
                    version_sh.display(), e));
            }
        }
    } else {
        warn!("This backend has no versioning. Using primitive one.");
        digest.input_str(builder.as_slice());
        digest.input_str(":");
        for (k, v) in container.parameters.iter() {
            digest.input_str(k.as_slice());
            digest.input_str("=");
            digest.input_str(v.as_slice());
            digest.input_str(";");
        }
    }
    let fullhash = digest.result_str();
    let hash = fullhash.as_slice().slice_to(7);

    let build_sh = environ.vagga_path.join_many(
        ["backends", builder.as_slice(), "build.sh"]);
    if !build_sh.exists() {
        return Err(format!("Builder {} does not exist", builder));
    }

    let cdir = format!("{}.{}", container.fullname, hash);
    let container_root = environ.local_vagga.join_many(
        [".roots", cdir.as_slice()]);

    if container_root.exists() && !force {
        info!("Container {} already built as {}",
            container.name, container_root.display());
        container.container_root = Some(container_root);
        return Ok(false);
    }

    let artifacts_dir = environ.local_vagga.join_many(
        [".artifacts", cdir.as_slice()]);
    let container_tmp = environ.local_vagga.join_many(
        [".roots", (cdir + ".tmp").as_slice()]);

    info!("Building {} by {}", container_root.display(), build_sh.display());

    if container_tmp.exists() {
        // TODO(tailhook) join container and rm, because may have
        // files of different owners (subuid's)
        try!(rmdir_recursive(&container_tmp).map_err(
                |e| format!("Can't clean temporary root: {}", e)));
    }
    if artifacts_dir.exists() {
        // TODO(tailhook) join container and rm, because may have
        // files of different owners (subuid's)
        try!(rmdir_recursive(&artifacts_dir).map_err(
                |e| format!("Can't clean temporary root: {}", e)));
    }
    try!(makedirs(&container_tmp));

    env.push(("artifacts_dir".as_bytes(), artifacts_dir.as_vec()));
    env.push(("container_hash".as_bytes(), hash.as_bytes()));
    env.push(("container_root".as_bytes(), container_tmp.as_vec()));
    match Command::new(build_sh).env(env.as_slice())
        .cwd(&environ.project_root)
        .stdin(InheritFd(0)).stdout(InheritFd(1)).stderr(InheritFd(2))
        .status() {
        Ok(ExitStatus(0)) => {}
        Ok(x) => return Err(format!("Builder exited with status {}", x)),
        Err(x) => return Err(format!("Can't spawn: {}", x)),
    }

    for dir in ["proc", "sys", "dev", "work", "tmp"].iter() {
        try!(ensure_dir(&container_tmp.join(*dir)));
    }

    let container_old = environ.local_vagga.join_many(
        [".roots", (cdir + ".old").as_slice()]);
    if container_root.exists() {
        if container_old.exists() {
            match rmdir_recursive(&container_old) {
                Ok(()) => {}
                Err(x) => return Err(format!("Can't remove old root: {}", x)),
            }
        }
        match rename(&container_root, &container_old) {
            Ok(()) => {}
            Err(x) => return Err(format!("Can't rename old root: {}", x)),
        }
    }

    match rename(&container_tmp, &container_root) {
        Ok(()) => {}
        Err(x) => return Err(format!("Can't rename root: {}", x)),
    }

    if container_old.exists() {
        match rmdir_recursive(&container_old) {
            Ok(()) => {}
            Err(x) => return Err(format!("Can't remove old root: {}", x)),
        }
    }
    info!("Done building {} as {}",
        container.name, container_root.display());

    container.container_root = Some(container_root);

    return Ok(true);
}

pub fn ensure_container(env: &Environ, container: &mut Container)
    -> Result<(), String>
{
    if env.settings.version_check {
        try!(build_container(env, container, false));
        try!(link_container(env, container));
    } else {
        let lnk = env.local_vagga.join(container.fullname.as_slice());
        container.container_root = match readlink(&lnk) {
            Ok(path) => Some(lnk.dir_path().join(path)),
            Err(e) => return Err(format!("Container {} not found: {}",
                                         container.fullname, e)),
        };
    };
    return Ok(());
}
<|endoftext|>


<commit_before>// Copyright (C) 2016 Élisabeth HENRY.
//
// This file is part of Crowbook.
//
// Crowbook is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published
// by the Free Software Foundation, either version 2.1 of the License, or
// (at your option) any later version.
//
// Caribon is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
//
// You should have received ba copy of the GNU Lesser General Public License
// along with Crowbook.  If not, see <http://www.gnu.org/licenses/>.

use error::{Error,Result, Source};
use html::HtmlRenderer;
use book::Book;
use token::Token;
use templates::{html};
use resource_handler::ResourceHandler;
use renderer::Renderer;

use mustache;

use std::io::{Read,Write};
use std::fs;
use std::fs::File;
use std::path::Path;
use std::path::PathBuf;
use std::borrow::Cow;
use std::convert::{AsRef, AsMut};

/// Multiple files HTML renderer
///
/// Renders HTML in a given directory.
pub struct HtmlDirRenderer<'a> {
    html: HtmlRenderer<'a>,
}

impl<'a> HtmlDirRenderer<'a> {
    /// Creates a new HtmlDirRenderer
    pub fn new(book: &'a Book) -> HtmlDirRenderer<'a> {
        let mut html = HtmlRenderer::new(book);
        html.handler.set_images_mapping(true);
        html.handler.set_base64(false);
        HtmlDirRenderer {
            html: html,
        }
    }

    /// Render a book
    pub fn render_book(&mut self) -> Result<()> {
        // Add internal files to resource handler
        for (i, filename) in self.html.book.filenames.iter().enumerate() {
            self.html.handler.add_link(filename.clone(), filenamer(i));
        }

        // Create the directory
        let dest_path = try!(self.html.book.options.get_path("output.html_dir"));
        match fs::metadata(&dest_path) {
            Ok(metadata) => if metadata.is_file() {
                return Err(Error::Render(format!("{} already exists and is not a directory", &dest_path)));
            } else if metadata.is_dir() {
                self.html.book.logger.warning(format!("{} already exists, deleting it", &dest_path));
                try!(fs::remove_dir_all(&dest_path)
                     .map_err(|e| Error::Render(format!("error deleting directory {}: {}", &dest_path, e))));
            },
            Err(_) => (),
        }
        try!(fs::DirBuilder::new()
             .recursive(true)
             .create(&dest_path)
             .map_err(|e| Error::Render(format!("could not create HTML directory {}:{}", &dest_path, e))));

        // Write CSS
        try!(self.write_css());
        // Write print.css
        try!(self.write_file("print.css",
                             &self.html.book.get_template("html.print_css").unwrap().as_bytes()));
        // Write index.html and chapter_xxx.html
        try!(self.write_html());
        // Write menu.svg
        try!(self.write_file("menu.svg", html::MENU_SVG));

        // Write highlight files if they are needed
        if self.html.book.options.get_bool("html.highlight_code") == Ok(true) {
            try!(self.write_file("highlight.js", self.html.book.get_template("html.highlight.js").unwrap().as_bytes()));
            try!(self.write_file("highlight.css", self.html.book.get_template("html.highlight.css").unwrap().as_bytes()));
        }

        // Write all images (including cover)
        let images_path = PathBuf::from(&self.html.book.options.get_path("resources.base_path.images").unwrap());
        for (source, dest) in self.html.handler.images_mapping() {
            let mut f = try!(File::open(images_path.join(source)).map_err(|_| Error::FileNotFound(self.html.book.source.clone(),
                                                                                "image or cover".to_owned(),
                                                                                source.to_owned())));
            let mut content = vec!();
            try!(f.read_to_end(&mut content).map_err(|e| Error::Render(format!("error while reading image file {}: {}", source, e))));
            try!(self.write_file(dest, &content));
        }

        // Write additional files
        if let Ok(list) = self.html.book.options.get_paths_list("resources.files") {
            let files_path = self.html.book.options.get_path("resources.base_path.files").unwrap();
            let data_path = Path::new(self.html.book.options.get_relative_path("resources.out_path").unwrap());
            let list = try!(ResourceHandler::get_files(list, &files_path));
            for path in list {
                let abs_path = Path::new(&files_path).join(&path);
                let mut f = try!(File::open(&abs_path)
                                 .map_err(|_| Error::FileNotFound(self.html.book.source.clone(),
                                                                  "additional resource from resources.files".to_owned(),
                                                                  abs_path.to_string_lossy().into_owned())));
                let mut content = vec!();
                try!(f.read_to_end(&mut content).map_err(|e| Error::Render(format!("error while reading resource file: {}", e))));
                try!(self.write_file(data_path.join(&path).to_str().unwrap(), &content));
            }
        }

        Ok(())
    }

    // Render each chapter and write them, and index.html too
    fn write_html(&mut self) -> Result<()> {
        let mut chapters = vec!();
        let mut titles = vec!();
        for (i, &(n, ref v)) in self.html.book.chapters.iter().enumerate() {
            self.html.chapter_config(i, n, filenamer(i));
            let mut title = String::new();
            for token in v {
                match *token {
                    Token::Header(1, ref vec) => {
                        if self.html.current_hide || self.html.current_numbering == 0 {
                            title = try!(self.html.render_vec(vec));
                        } else {
                            title = try!(self.html.book.get_header(self.html.current_chapter[0] + 1,
                                                              &try!(self.html.render_vec(vec))));
                        }
                        break;
                    },
                    _ => {
                        continue;
                    }
                }
            }
            titles.push(title);

            let chapter = HtmlRenderer::render_html(self, v);
            chapters.push(chapter);
        }
        self.html.source = Source::empty();
        let toc = self.html.toc.render();

        for (i, content) in chapters.into_iter().enumerate() {
            let prev_chapter = if i > 0 {
                format!("

    « {}

",
                        filenamer(i-1),
                        titles[i-1])
            } else {
                String::new()
            };

            let next_chapter = if i < titles.len() - 1 {
                format!("

    {} »

",
                        filenamer(i+1),
                        titles[i+1])
            } else {
                String::new()
            };

            // Render each HTML document
            let mut mapbuilder = self.html.book.get_mapbuilder("none")
                .insert_str("content", try!(content))
                .insert_str("chapter_title", format!("{} – {}",
                                             self.html.book.options.get_str("title").unwrap(),
                                             titles[i]))
                .insert_str("toc", toc.clone())
                .insert_str("prev_chapter", prev_chapter)
                .insert_str("next_chapter", next_chapter)
                .insert_str("footer", self.html.get_footer())
                .insert_str("top", self.html.get_top())
                .insert_str("script", self.html.book.get_template("html_dir.script").unwrap())
                .insert_bool(self.html.book.options.get_str("lang").unwrap(), true);

            if self.html.book.options.get_bool("html.highlight_code").unwrap() == true {
                mapbuilder = mapbuilder.insert_bool("highlight_code", true);
            }
            let data = mapbuilder.build();
            let template = mustache::compile_str(try!(self.html.book.get_template("html_dir.chapter.html")).as_ref());
            let mut res = vec!();
            template.render_data(&mut res, &data);
            try!(self.write_file(&filenamer(i), &res));
        }

        let mut content = if let Ok(cover) = self.html.book.options.get_path("cover") {
            // checks first that cover exists
            if fs::metadata(&cover).is_err() {
                return Err(Error::FileNotFound(self.html.book.source.clone(),
                                               "cover".to_owned(),
                                               cover));

            }
            format!("

",
                    self.html.book.options.get_str("title").unwrap(),
                    try!(self.html.handler.map_image(&self.html.book.source,
                                                     Cow::Owned(cover))).as_ref())
        } else {
            String::new()
        };
        if titles.len() > 1 {
            content.push_str(&format!("

    {} »

",
                        filenamer(0),
                        titles[0]));
        }
        // Render index.html and write it too
        let mut mapbuilder = self.html.book.get_mapbuilder("none")
            .insert_str("content", content)
            .insert_str("top", self.html.get_top())
            .insert_str("footer", self.html.get_footer())
            .insert_str("toc", toc.clone())
            .insert_str("script", self.html.book.get_template("html_dir.script").unwrap())
            .insert_bool(self.html.book.options.get_str("lang").unwrap(), true);
        if self.html.book.options.get_bool("html.highlight_code").unwrap() == true {
            mapbuilder = mapbuilder.insert_bool("highlight_code", true);
        }
        let data = mapbuilder.build();
        let template = mustache::compile_str(try!(self.html.book.get_template("html_dir.index.html")).as_ref());
        let mut res = vec!();
        template.render_data(&mut res, &data);
        try!(self.write_file("index.html", &res));

        Ok(())
    }

    // Render the CSS file and write it
    fn write_css(&self) -> Result<()> {
        // Render the CSS
        let template_css = mustache::compile_str(try!(self.html.book.get_template("html_dir.css")).as_ref());
        let data = self.html.book.get_mapbuilder("none")
            .insert_bool(self.html.book.options.get_str("lang").unwrap(), true)
            .build();
        let mut res:Vec<u8> = vec!();
        template_css.render_data(&mut res, &data);
        let css = String::from_utf8_lossy(&res);

        // Write it
        self.write_file("stylesheet.css", css.as_bytes())
    }

    // Write content to a file
    fn write_file(&self, file: &str, content: &[u8]) -> Result<()> {
        let dest_path = PathBuf::from(&self.html.book.options.get_path("output.html_dir").unwrap());
        if dest_path.starts_with("..") {
            panic!("html dir is asked to create a file outside of its directory, no way!");
        }
        let dest_file = dest_path.join(file);
        let dest_dir = dest_file.parent().unwrap();
        if !fs::metadata(dest_dir).is_ok() { // dir does not exist, create it
            try!(fs::DirBuilder::new()
                 .recursive(true)
                 .create(&dest_dir)
                 .map_err(|e| Error::Render(format!("could not create directory in {}:{}", dest_dir.display(), e))));
        }
        let mut f = try!(File::create(&dest_file)
                         .map_err(|e| Error::Render(format!("could not create file {}:{}", dest_file.display(), e))));
        f.write_all(content)
            .map_err(|e| Error::Render(format!("could not write to file {}:{}", dest_file.display(), e)))
    }
}

/// Generate a file name given an int
fn filenamer(i: usize) -> String {
    format!("chapter_{:03}.html", i)
}

derive_html!{HtmlDirRenderer<'a>, HtmlRenderer::static_render_token}
<commit_msg>Remove unused variable<commit_after>// Copyright (C) 2016 Élisabeth HENRY.
//
// This file is part of Crowbook.
//
// Crowbook is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published
// by the Free Software Foundation, either version 2.1 of the License, or
// (at your option) any later version.
//
// Caribon is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
//
// You should have received ba copy of the GNU Lesser General Public License
// along with Crowbook.  If not, see <http://www.gnu.org/licenses/>.

use error::{Error,Result, Source};
use html::HtmlRenderer;
use book::Book;
use token::Token;
use templates::{html};
use resource_handler::ResourceHandler;
use renderer::Renderer;

use mustache;

use std::io::{Read,Write};
use std::fs;
use std::fs::File;
use std::path::Path;
use std::path::PathBuf;
use std::borrow::Cow;
use std::convert::{AsRef, AsMut};

/// Multiple files HTML renderer
///
/// Renders HTML in a given directory.
pub struct HtmlDirRenderer<'a> {
    html: HtmlRenderer<'a>,
}

impl<'a> HtmlDirRenderer<'a> {
    /// Creates a new HtmlDirRenderer
    pub fn new(book: &'a Book) -> HtmlDirRenderer<'a> {
        let mut html = HtmlRenderer::new(book);
        html.handler.set_images_mapping(true);
        html.handler.set_base64(false);
        HtmlDirRenderer {
            html: html,
        }
    }

    /// Render a book
    pub fn render_book(&mut self) -> Result<()> {
        // Add internal files to resource handler
        for (i, filename) in self.html.book.filenames.iter().enumerate() {
            self.html.handler.add_link(filename.clone(), filenamer(i));
        }

        // Create the directory
        let dest_path = try!(self.html.book.options.get_path("output.html_dir"));
        match fs::metadata(&dest_path) {
            Ok(metadata) => if metadata.is_file() {
                return Err(Error::Render(format!("{} already exists and is not a directory", &dest_path)));
            } else if metadata.is_dir() {
                self.html.book.logger.warning(format!("{} already exists, deleting it", &dest_path));
                try!(fs::remove_dir_all(&dest_path)
                     .map_err(|e| Error::Render(format!("error deleting directory {}: {}", &dest_path, e))));
            },
            Err(_) => (),
        }
        try!(fs::DirBuilder::new()
             .recursive(true)
             .create(&dest_path)
             .map_err(|e| Error::Render(format!("could not create HTML directory {}:{}", &dest_path, e))));

        // Write CSS
        try!(self.write_css());
        // Write print.css
        try!(self.write_file("print.css",
                             &self.html.book.get_template("html.print_css").unwrap().as_bytes()));
        // Write index.html and chapter_xxx.html
        try!(self.write_html());
        // Write menu.svg
        try!(self.write_file("menu.svg", html::MENU_SVG));

        // Write highlight files if they are needed
        if self.html.book.options.get_bool("html.highlight_code") == Ok(true) {
            try!(self.write_file("highlight.js", self.html.book.get_template("html.highlight.js").unwrap().as_bytes()));
            try!(self.write_file("highlight.css", self.html.book.get_template("html.highlight.css").unwrap().as_bytes()));
        }

        // Write all images (including cover)
        for (source, dest) in self.html.handler.images_mapping() {
            let mut f = try!(File::open(source).map_err(|_| Error::FileNotFound(self.html.book.source.clone(),
                                                                                "image or cover".to_owned(),
                                                                                source.to_owned())));
            let mut content = vec!();
            try!(f.read_to_end(&mut content).map_err(|e| Error::Render(format!("error while reading image file {}: {}", source, e))));
            try!(self.write_file(dest, &content));
        }

        // Write additional files
        if let Ok(list) = self.html.book.options.get_paths_list("resources.files") {
            let files_path = self.html.book.options.get_path("resources.base_path.files").unwrap();
            let data_path = Path::new(self.html.book.options.get_relative_path("resources.out_path").unwrap());
            let list = try!(ResourceHandler::get_files(list, &files_path));
            for path in list {
                let abs_path = Path::new(&files_path).join(&path);
                let mut f = try!(File::open(&abs_path)
                                 .map_err(|_| Error::FileNotFound(self.html.book.source.clone(),
                                                                  "additional resource from resources.files".to_owned(),
                                                                  abs_path.to_string_lossy().into_owned())));
                let mut content = vec!();
                try!(f.read_to_end(&mut content).map_err(|e| Error::Render(format!("error while reading resource file: {}", e))));
                try!(self.write_file(data_path.join(&path).to_str().unwrap(), &content));
            }
        }

        Ok(())
    }

    // Render each chapter and write them, and index.html too
    fn write_html(&mut self) -> Result<()> {
        let mut chapters = vec!();
        let mut titles = vec!();
        for (i, &(n, ref v)) in self.html.book.chapters.iter().enumerate() {
            self.html.chapter_config(i, n, filenamer(i));
            let mut title = String::new();
            for token in v {
                match *token {
                    Token::Header(1, ref vec) => {
                        if self.html.current_hide || self.html.current_numbering == 0 {
                            title = try!(self.html.render_vec(vec));
                        } else {
                            title = try!(self.html.book.get_header(self.html.current_chapter[0] + 1,
                                                              &try!(self.html.render_vec(vec))));
                        }
                        break;
                    },
                    _ => {
                        continue;
                    }
                }
            }
            titles.push(title);

            let chapter = HtmlRenderer::render_html(self, v);
            chapters.push(chapter);
        }
        self.html.source = Source::empty();
        let toc = self.html.toc.render();

        for (i, content) in chapters.into_iter().enumerate() {
            let prev_chapter = if i > 0 {
                format!("

    « {}

",
                        filenamer(i-1),
                        titles[i-1])
            } else {
                String::new()
            };

            let next_chapter = if i < titles.len() - 1 {
                format!("

    {} »

",
                        filenamer(i+1),
                        titles[i+1])
            } else {
                String::new()
            };

            // Render each HTML document
            let mut mapbuilder = self.html.book.get_mapbuilder("none")
                .insert_str("content", try!(content))
                .insert_str("chapter_title", format!("{} – {}",
                                             self.html.book.options.get_str("title").unwrap(),
                                             titles[i]))
                .insert_str("toc", toc.clone())
                .insert_str("prev_chapter", prev_chapter)
                .insert_str("next_chapter", next_chapter)
                .insert_str("footer", self.html.get_footer())
                .insert_str("top", self.html.get_top())
                .insert_str("script", self.html.book.get_template("html_dir.script").unwrap())
                .insert_bool(self.html.book.options.get_str("lang").unwrap(), true);

            if self.html.book.options.get_bool("html.highlight_code").unwrap() == true {
                mapbuilder = mapbuilder.insert_bool("highlight_code", true);
            }
            let data = mapbuilder.build();
            let template = mustache::compile_str(try!(self.html.book.get_template("html_dir.chapter.html")).as_ref());
            let mut res = vec!();
            template.render_data(&mut res, &data);
            try!(self.write_file(&filenamer(i), &res));
        }

        let mut content = if let Ok(cover) = self.html.book.options.get_path("cover") {
            // checks first that cover exists
            if fs::metadata(&cover).is_err() {
                return Err(Error::FileNotFound(self.html.book.source.clone(),
                                               "cover".to_owned(),
                                               cover));

            }
            format!("

",
                    self.html.book.options.get_str("title").unwrap(),
                    try!(self.html.handler.map_image(&self.html.book.source,
                                                     Cow::Owned(cover))).as_ref())
        } else {
            String::new()
        };
        if titles.len() > 1 {
            content.push_str(&format!("

    {} »

",
                        filenamer(0),
                        titles[0]));
        }
        // Render index.html and write it too
        let mut mapbuilder = self.html.book.get_mapbuilder("none")
            .insert_str("content", content)
            .insert_str("top", self.html.get_top())
            .insert_str("footer", self.html.get_footer())
            .insert_str("toc", toc.clone())
            .insert_str("script", self.html.book.get_template("html_dir.script").unwrap())
            .insert_bool(self.html.book.options.get_str("lang").unwrap(), true);
        if self.html.book.options.get_bool("html.highlight_code").unwrap() == true {
            mapbuilder = mapbuilder.insert_bool("highlight_code", true);
        }
        let data = mapbuilder.build();
        let template = mustache::compile_str(try!(self.html.book.get_template("html_dir.index.html")).as_ref());
        let mut res = vec!();
        template.render_data(&mut res, &data);
        try!(self.write_file("index.html", &res));

        Ok(())
    }

    // Render the CSS file and write it
    fn write_css(&self) -> Result<()> {
        // Render the CSS
        let template_css = mustache::compile_str(try!(self.html.book.get_template("html_dir.css")).as_ref());
        let data = self.html.book.get_mapbuilder("none")
            .insert_bool(self.html.book.options.get_str("lang").unwrap(), true)
            .build();
        let mut res:Vec<u8> = vec!();
        template_css.render_data(&mut res, &data);
        let css = String::from_utf8_lossy(&res);

        // Write it
        self.write_file("stylesheet.css", css.as_bytes())
    }

    // Write content to a file
    fn write_file(&self, file: &str, content: &[u8]) -> Result<()> {
        let dest_path = PathBuf::from(&self.html.book.options.get_path("output.html_dir").unwrap());
        if dest_path.starts_with("..") {
            panic!("html dir is asked to create a file outside of its directory, no way!");
        }
        let dest_file = dest_path.join(file);
        let dest_dir = dest_file.parent().unwrap();
        if !fs::metadata(dest_dir).is_ok() { // dir does not exist, create it
            try!(fs::DirBuilder::new()
                 .recursive(true)
                 .create(&dest_dir)
                 .map_err(|e| Error::Render(format!("could not create directory in {}:{}", dest_dir.display(), e))));
        }
        let mut f = try!(File::create(&dest_file)
                         .map_err(|e| Error::Render(format!("could not create file {}:{}", dest_file.display(), e))));
        f.write_all(content)
            .map_err(|e| Error::Render(format!("could not write to file {}:{}", dest_file.display(), e)))
    }
}

/// Generate a file name given an int
fn filenamer(i: usize) -> String {
    format!("chapter_{:03}.html", i)
}

derive_html!{HtmlDirRenderer<'a>, HtmlRenderer::static_render_token}
<|endoftext|>


<commit_before>#[doc = "Random number generation"];

export rng, seed, seeded_rng, weighted, extensions;
export xorshift, seeded_xorshift;

enum rctx {}

#[abi = "cdecl"]
native mod rustrt {
    fn rand_seed() -> [u8]/~;
    fn rand_new() -> *rctx;
    fn rand_new_seeded(seed: [u8]/~) -> *rctx;
    fn rand_next(c: *rctx) -> u32;
    fn rand_free(c: *rctx);
}

#[doc = "A random number generator"]
iface rng {
    #[doc = "Return the next random integer"]
    fn next() -> u32;
}

#[doc = "A value with a particular weight compared to other values"]
type weighted<T> = { weight: uint, item: T };

#[doc = "Extension methods for random number generators"]
impl extensions for rng {

    #[doc = "Return a random int"]
    fn gen_int() -> int {
        self.gen_i64() as int
    }

    #[doc = "Return an int randomly chosen from the range [start, end), \
             failing if start >= end"]
    fn gen_int_range(start: int, end: int) -> int {
        assert start < end;
        start + int::abs(self.gen_int() % (end - start))
    }

    #[doc = "Return a random i8"]
    fn gen_i8() -> i8 {
        self.next() as i8
    }

    #[doc = "Return a random i16"]
    fn gen_i16() -> i16 {
        self.next() as i16
    }

    #[doc = "Return a random i32"]
    fn gen_i32() -> i32 {
        self.next() as i32
    }

    #[doc = "Return a random i64"]
    fn gen_i64() -> i64 {
        (self.next() as i64 << 32) | self.next() as i64
    }

    #[doc = "Return a random uint"]
    fn gen_uint() -> uint {
        self.gen_u64() as uint
    }

    #[doc = "Return a uint randomly chosen from the range [start, end), \
             failing if start >= end"]
    fn gen_uint_range(start: uint, end: uint) -> uint {
        assert start < end;
        start + (self.gen_uint() % (end - start))
    }

    #[doc = "Return a random u8"]
    fn gen_u8() -> u8 {
        self.next() as u8
    }

    #[doc = "Return a random u16"]
    fn gen_u16() -> u16 {
        self.next() as u16
    }

    #[doc = "Return a random u32"]
    fn gen_u32() -> u32 {
        self.next()
    }

    #[doc = "Return a random u64"]
    fn gen_u64() -> u64 {
        (self.next() as u64 << 32) | self.next() as u64
    }

    #[doc = "Return a random float"]
    fn gen_float() -> float {
        self.gen_f64() as float
    }

    #[doc = "Return a random f32"]
    fn gen_f32() -> f32 {
        self.gen_f64() as f32
    }

    #[doc = "Return a random f64"]
    fn gen_f64() -> f64 {
        let u1 = self.next() as f64;
        let u2 = self.next() as f64;
        let u3 = self.next() as f64;
        let scale = u32::max_value as f64;
        ret ((u1 / scale + u2) / scale + u3) / scale;
    }

    #[doc = "Return a random char"]
    fn gen_char() -> char {
        self.next() as char
    }

    #[doc = "Return a char randomly chosen from chars, failing if chars is \
             empty"]
    fn gen_char_from(chars: str) -> char {
        assert !chars.is_empty();
        self.choose(str::chars(chars))
    }

    #[doc = "Return a random bool"]
    fn gen_bool() -> bool {
        self.next() & 1u32 == 1u32
    }

    #[doc = "Return a bool with a 1 in n chance of true"]
    fn gen_weighted_bool(n: uint) -> bool {
        if n == 0u {
            true
        } else {
            self.gen_uint_range(1u, n + 1u) == 1u
        }
    }

    #[doc = "Return a random string of the specified length composed of A-Z, \
             a-z, 0-9"]
    fn gen_str(len: uint) -> str {
        let charset = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +
                      "abcdefghijklmnopqrstuvwxyz" +
                      "0123456789";
        let mut s = "";
        let mut i = 0u;
        while (i < len) {
            s = s + str::from_char(self.gen_char_from(charset));
            i += 1u;
        }
        s
    }

    #[doc = "Return a random byte string of the specified length"]
    fn gen_bytes(len: uint) -> [u8]/~ {
        vec::from_fn(len) {|_i|
            self.gen_u8()
        }
    }

    #[doc = "Choose an item randomly, failing if values is empty"]
    fn choose<T:copy>(values: [T]/~) -> T {
        self.choose_option(values).get()
    }

    #[doc = "Choose some(item) randomly, returning none if values is empty"]
    fn choose_option<T:copy>(values: [T]/~) -> option<T> {
        if values.is_empty() {
            none
        } else {
            some(values[self.gen_uint_range(0u, values.len())])
        }
    }

    #[doc = "Choose an item respecting the relative weights, failing if \
             the sum of the weights is 0"]
    fn choose_weighted<T: copy>(v : [weighted<T>]/~) -> T {
        self.choose_weighted_option(v).get()
    }

    #[doc = "Choose some(item) respecting the relative weights, returning \
             none if the sum of the weights is 0"]
    fn choose_weighted_option<T:copy>(v: [weighted<T>]/~) -> option<T> {
        let mut total = 0u;
        for v.each {|item|
            total += item.weight;
        }
        if total == 0u {
            ret none;
        }
        let chosen = self.gen_uint_range(0u, total);
        let mut so_far = 0u;
        for v.each {|item|
            so_far += item.weight;
            if so_far > chosen {
                ret some(item.item);
            }
        }
        unreachable();
    }

    #[doc = "Return a vec containing copies of the items, in order, where \
             the weight of the item determines how many copies there are"]
    fn weighted_vec<T:copy>(v: [weighted<T>]/~) -> [T]/~ {
        let mut r = []/~;
        for v.each {|item|
            for uint::range(0u, item.weight) {|_i|
                vec::push(r, item.item);
            }
        }
        r
    }

    #[doc = "Shuffle a vec"]
    fn shuffle<T:copy>(values: [T]/~) -> [T]/~ {
        let mut m = vec::to_mut(values);
        self.shuffle_mut(m);
        ret vec::from_mut(m);
    }

    #[doc = "Shuffle a mutable vec in place"]
    fn shuffle_mut<T>(&&values: [mut T]/~) {
        let mut i = values.len();
        while i >= 2u {
            // invariant: elements with index >= i have been locked in place.
            i -= 1u;
            // lock element i in place.
            vec::swap(values, i, self.gen_uint_range(0u, i + 1u));
        }
    }

}

class rand_res {
    let c: *rctx;
    new(c: *rctx) { self.c = c; }
    drop { rustrt::rand_free(self.c); }
}

impl of rng for @rand_res {
    fn next() -> u32 { ret rustrt::rand_next((*self).c); }
}

#[doc = "Create a new random seed for seeded_rng"]
fn seed() -> [u8]/~ {
    rustrt::rand_seed()
}

#[doc = "Create a random number generator with a system specified seed"]
fn rng() -> rng {
    @rand_res(rustrt::rand_new()) as rng
}

#[doc = "Create a random number generator using the specified seed. A \
         generator constructed with a given seed will generate the same \
         sequence of values as all other generators constructed with the \
         same seed. The seed may be any length."]
fn seeded_rng(seed: [u8]/~) -> rng {
    @rand_res(rustrt::rand_new_seeded(seed)) as rng
}

type xorshift_state = {
    mut x: u32,
    mut y: u32,
    mut z: u32,
    mut w: u32
};

impl of rng for xorshift_state {
    fn next() -> u32 {
        let x = self.x;
        let mut t = x ^ (x << 11);
        self.x = self.y;
        self.y = self.z;
        self.z = self.w;
        let w = self.w;
        self.w = w ^ (w >> 19) ^ (t ^ (t >> 8));
        self.w
    }
}

fn xorshift() -> rng {
    // constants taken from http://en.wikipedia.org/wiki/Xorshift
    seeded_xorshift(123456789u32, 362436069u32, 521288629u32, 88675123u32)
}

fn seeded_xorshift(x: u32, y: u32, z: u32, w: u32) -> rng {
    {mut x: x, mut y: y, mut z: z, mut w: w} as rng
}

#[cfg(test)]
mod tests {

    #[test]
    fn rng_seeded() {
        let seed = rand::seed();
        let ra = rand::seeded_rng(seed);
        let rb = rand::seeded_rng(seed);
        assert ra.gen_str(100u) == rb.gen_str(100u);
    }

    #[test]
    fn rng_seeded_custom_seed() {
        // much shorter than generated seeds which are 1024 bytes
        let seed = [2u8, 32u8, 4u8, 32u8, 51u8]/~;
        let ra = rand::seeded_rng(seed);
        let rb = rand::seeded_rng(seed);
        assert ra.gen_str(100u) == rb.gen_str(100u);
    }

    #[test]
    fn rng_seeded_custom_seed2() {
        let seed = [2u8, 32u8, 4u8, 32u8, 51u8]/~;
        let ra = rand::seeded_rng(seed);
        // Regression test that isaac is actually using the above vector
        let r = ra.next();
        #error("%?", r);
        assert r == 890007737u32 // on x86_64
            || r == 2935188040u32; // on x86
    }

    #[test]
    fn gen_int_range() {
        let r = rand::rng();
        let a = r.gen_int_range(-3, 42);
        assert a >= -3 && a < 42;
        assert r.gen_int_range(0, 1) == 0;
        assert r.gen_int_range(-12, -11) == -12;
    }

    #[test]
    #[should_fail]
    #[ignore(cfg(windows))]
    fn gen_int_from_fail() {
        rand::rng().gen_int_range(5, -2);
    }

    #[test]
    fn gen_uint_range() {
        let r = rand::rng();
        let a = r.gen_uint_range(3u, 42u);
        assert a >= 3u && a < 42u;
        assert r.gen_uint_range(0u, 1u) == 0u;
        assert r.gen_uint_range(12u, 13u) == 12u;
    }

    #[test]
    #[should_fail]
    #[ignore(cfg(windows))]
    fn gen_uint_range_fail() {
        rand::rng().gen_uint_range(5u, 2u);
    }

    #[test]
    fn gen_float() {
        let r = rand::rng();
        let a = r.gen_float();
        let b = r.gen_float();
        log(debug, (a, b));
    }

    #[test]
    fn gen_weighted_bool() {
        let r = rand::rng();
        assert r.gen_weighted_bool(0u) == true;
        assert r.gen_weighted_bool(1u) == true;
    }

    #[test]
    fn gen_str() {
        let r = rand::rng();
        log(debug, r.gen_str(10u));
        log(debug, r.gen_str(10u));
        log(debug, r.gen_str(10u));
        assert r.gen_str(0u).len() == 0u;
        assert r.gen_str(10u).len() == 10u;
        assert r.gen_str(16u).len() == 16u;
    }

    #[test]
    fn gen_bytes() {
        let r = rand::rng();
        assert r.gen_bytes(0u).len() == 0u;
        assert r.gen_bytes(10u).len() == 10u;
        assert r.gen_bytes(16u).len() == 16u;
    }

    #[test]
    fn choose() {
        let r = rand::rng();
        assert r.choose([1, 1, 1]/~) == 1;
    }

    #[test]
    fn choose_option() {
        let r = rand::rng();
        assert r.choose_option([]/~) == none::<int>;
        assert r.choose_option([1, 1, 1]/~) == some(1);
    }

    #[test]
    fn choose_weighted() {
        let r = rand::rng();
        assert r.choose_weighted([{weight: 1u, item: 42}]/~) == 42;
        assert r.choose_weighted([
            {weight: 0u, item: 42},
            {weight: 1u, item: 43}
        ]/~) == 43;
    }

    #[test]
    fn choose_weighted_option() {
        let r = rand::rng();
        assert r.choose_weighted_option([{weight: 1u, item: 42}]/~) ==
               some(42);
        assert r.choose_weighted_option([
            {weight: 0u, item: 42},
            {weight: 1u, item: 43}
        ]/~) == some(43);
        assert r.choose_weighted_option([]/~) == none::<int>;
    }

    #[test]
    fn weighted_vec() {
        let r = rand::rng();
        let empty: [int]/~ = []/~;
        assert r.weighted_vec([]/~) == empty;
        assert r.weighted_vec([
            {weight: 0u, item: 3u},
            {weight: 1u, item: 2u},
            {weight: 2u, item: 1u}
        ]/~) == [2u, 1u, 1u]/~;
    }

    #[test]
    fn shuffle() {
        let r = rand::rng();
        let empty: [int]/~ = []/~;
        assert r.shuffle([]/~) == empty;
        assert r.shuffle([1, 1, 1]/~) == [1, 1, 1]/~;
    }
}

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// End:
<commit_msg>Switch the compiler over to using ~[] notation instead of []/~. Closes #2759.<commit_after>#[doc = "Random number generation"];

export rng, seed, seeded_rng, weighted, extensions;
export xorshift, seeded_xorshift;

enum rctx {}

#[abi = "cdecl"]
native mod rustrt {
    fn rand_seed() -> ~[u8];
    fn rand_new() -> *rctx;
    fn rand_new_seeded(seed: ~[u8]) -> *rctx;
    fn rand_next(c: *rctx) -> u32;
    fn rand_free(c: *rctx);
}

#[doc = "A random number generator"]
iface rng {
    #[doc = "Return the next random integer"]
    fn next() -> u32;
}

#[doc = "A value with a particular weight compared to other values"]
type weighted<T> = { weight: uint, item: T };

#[doc = "Extension methods for random number generators"]
impl extensions for rng {

    #[doc = "Return a random int"]
    fn gen_int() -> int {
        self.gen_i64() as int
    }

    #[doc = "Return an int randomly chosen from the range [start, end), \
             failing if start >= end"]
    fn gen_int_range(start: int, end: int) -> int {
        assert start < end;
        start + int::abs(self.gen_int() % (end - start))
    }

    #[doc = "Return a random i8"]
    fn gen_i8() -> i8 {
        self.next() as i8
    }

    #[doc = "Return a random i16"]
    fn gen_i16() -> i16 {
        self.next() as i16
    }

    #[doc = "Return a random i32"]
    fn gen_i32() -> i32 {
        self.next() as i32
    }

    #[doc = "Return a random i64"]
    fn gen_i64() -> i64 {
        (self.next() as i64 << 32) | self.next() as i64
    }

    #[doc = "Return a random uint"]
    fn gen_uint() -> uint {
        self.gen_u64() as uint
    }

    #[doc = "Return a uint randomly chosen from the range [start, end), \
             failing if start >= end"]
    fn gen_uint_range(start: uint, end: uint) -> uint {
        assert start < end;
        start + (self.gen_uint() % (end - start))
    }

    #[doc = "Return a random u8"]
    fn gen_u8() -> u8 {
        self.next() as u8
    }

    #[doc = "Return a random u16"]
    fn gen_u16() -> u16 {
        self.next() as u16
    }

    #[doc = "Return a random u32"]
    fn gen_u32() -> u32 {
        self.next()
    }

    #[doc = "Return a random u64"]
    fn gen_u64() -> u64 {
        (self.next() as u64 << 32) | self.next() as u64
    }

    #[doc = "Return a random float"]
    fn gen_float() -> float {
        self.gen_f64() as float
    }

    #[doc = "Return a random f32"]
    fn gen_f32() -> f32 {
        self.gen_f64() as f32
    }

    #[doc = "Return a random f64"]
    fn gen_f64() -> f64 {
        let u1 = self.next() as f64;
        let u2 = self.next() as f64;
        let u3 = self.next() as f64;
        let scale = u32::max_value as f64;
        ret ((u1 / scale + u2) / scale + u3) / scale;
    }

    #[doc = "Return a random char"]
    fn gen_char() -> char {
        self.next() as char
    }

    #[doc = "Return a char randomly chosen from chars, failing if chars is \
             empty"]
    fn gen_char_from(chars: str) -> char {
        assert !chars.is_empty();
        self.choose(str::chars(chars))
    }

    #[doc = "Return a random bool"]
    fn gen_bool() -> bool {
        self.next() & 1u32 == 1u32
    }

    #[doc = "Return a bool with a 1 in n chance of true"]
    fn gen_weighted_bool(n: uint) -> bool {
        if n == 0u {
            true
        } else {
            self.gen_uint_range(1u, n + 1u) == 1u
        }
    }

    #[doc = "Return a random string of the specified length composed of A-Z, \
             a-z, 0-9"]
    fn gen_str(len: uint) -> str {
        let charset = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +
                      "abcdefghijklmnopqrstuvwxyz" +
                      "0123456789";
        let mut s = "";
        let mut i = 0u;
        while (i < len) {
            s = s + str::from_char(self.gen_char_from(charset));
            i += 1u;
        }
        s
    }

    #[doc = "Return a random byte string of the specified length"]
    fn gen_bytes(len: uint) -> ~[u8] {
        vec::from_fn(len) {|_i|
            self.gen_u8()
        }
    }

    #[doc = "Choose an item randomly, failing if values is empty"]
    fn choose<T:copy>(values: ~[T]) -> T {
        self.choose_option(values).get()
    }

    #[doc = "Choose some(item) randomly, returning none if values is empty"]
    fn choose_option<T:copy>(values: ~[T]) -> option<T> {
        if values.is_empty() {
            none
        } else {
            some(values[self.gen_uint_range(0u, values.len())])
        }
    }

    #[doc = "Choose an item respecting the relative weights, failing if \
             the sum of the weights is 0"]
    fn choose_weighted<T: copy>(v : ~[weighted<T>]) -> T {
        self.choose_weighted_option(v).get()
    }

    #[doc = "Choose some(item) respecting the relative weights, returning \
             none if the sum of the weights is 0"]
    fn choose_weighted_option<T:copy>(v: ~[weighted<T>]) -> option<T> {
        let mut total = 0u;
        for v.each {|item|
            total += item.weight;
        }
        if total == 0u {
            ret none;
        }
        let chosen = self.gen_uint_range(0u, total);
        let mut so_far = 0u;
        for v.each {|item|
            so_far += item.weight;
            if so_far > chosen {
                ret some(item.item);
            }
        }
        unreachable();
    }

    #[doc = "Return a vec containing copies of the items, in order, where \
             the weight of the item determines how many copies there are"]
    fn weighted_vec<T:copy>(v: ~[weighted<T>]) -> ~[T] {
        let mut r = ~[];
        for v.each {|item|
            for uint::range(0u, item.weight) {|_i|
                vec::push(r, item.item);
            }
        }
        r
    }

    #[doc = "Shuffle a vec"]
    fn shuffle<T:copy>(values: ~[T]) -> ~[T] {
        let mut m = vec::to_mut(values);
        self.shuffle_mut(m);
        ret vec::from_mut(m);
    }

    #[doc = "Shuffle a mutable vec in place"]
    fn shuffle_mut<T>(&&values: ~[mut T]) {
        let mut i = values.len();
        while i >= 2u {
            // invariant: elements with index >= i have been locked in place.
            i -= 1u;
            // lock element i in place.
            vec::swap(values, i, self.gen_uint_range(0u, i + 1u));
        }
    }

}

class rand_res {
    let c: *rctx;
    new(c: *rctx) { self.c = c; }
    drop { rustrt::rand_free(self.c); }
}

impl of rng for @rand_res {
    fn next() -> u32 { ret rustrt::rand_next((*self).c); }
}

#[doc = "Create a new random seed for seeded_rng"]
fn seed() -> ~[u8] {
    rustrt::rand_seed()
}

#[doc = "Create a random number generator with a system specified seed"]
fn rng() -> rng {
    @rand_res(rustrt::rand_new()) as rng
}

#[doc = "Create a random number generator using the specified seed. A \
         generator constructed with a given seed will generate the same \
         sequence of values as all other generators constructed with the \
         same seed. The seed may be any length."]
fn seeded_rng(seed: ~[u8]) -> rng {
    @rand_res(rustrt::rand_new_seeded(seed)) as rng
}

type xorshift_state = {
    mut x: u32,
    mut y: u32,
    mut z: u32,
    mut w: u32
};

impl of rng for xorshift_state {
    fn next() -> u32 {
        let x = self.x;
        let mut t = x ^ (x << 11);
        self.x = self.y;
        self.y = self.z;
        self.z = self.w;
        let w = self.w;
        self.w = w ^ (w >> 19) ^ (t ^ (t >> 8));
        self.w
    }
}

fn xorshift() -> rng {
    // constants taken from http://en.wikipedia.org/wiki/Xorshift
    seeded_xorshift(123456789u32, 362436069u32, 521288629u32, 88675123u32)
}

fn seeded_xorshift(x: u32, y: u32, z: u32, w: u32) -> rng {
    {mut x: x, mut y: y, mut z: z, mut w: w} as rng
}

#[cfg(test)]
mod tests {

    #[test]
    fn rng_seeded() {
        let seed = rand::seed();
        let ra = rand::seeded_rng(seed);
        let rb = rand::seeded_rng(seed);
        assert ra.gen_str(100u) == rb.gen_str(100u);
    }

    #[test]
    fn rng_seeded_custom_seed() {
        // much shorter than generated seeds which are 1024 bytes
        let seed = ~[2u8, 32u8, 4u8, 32u8, 51u8];
        let ra = rand::seeded_rng(seed);
        let rb = rand::seeded_rng(seed);
        assert ra.gen_str(100u) == rb.gen_str(100u);
    }

    #[test]
    fn rng_seeded_custom_seed2() {
        let seed = ~[2u8, 32u8, 4u8, 32u8, 51u8];
        let ra = rand::seeded_rng(seed);
        // Regression test that isaac is actually using the above vector
        let r = ra.next();
        #error("%?", r);
        assert r == 890007737u32 // on x86_64
            || r == 2935188040u32; // on x86
    }

    #[test]
    fn gen_int_range() {
        let r = rand::rng();
        let a = r.gen_int_range(-3, 42);
        assert a >= -3 && a < 42;
        assert r.gen_int_range(0, 1) == 0;
        assert r.gen_int_range(-12, -11) == -12;
    }

    #[test]
    #[should_fail]
    #[ignore(cfg(windows))]
    fn gen_int_from_fail() {
        rand::rng().gen_int_range(5, -2);
    }

    #[test]
    fn gen_uint_range() {
        let r = rand::rng();
        let a = r.gen_uint_range(3u, 42u);
        assert a >= 3u && a < 42u;
        assert r.gen_uint_range(0u, 1u) == 0u;
        assert r.gen_uint_range(12u, 13u) == 12u;
    }

    #[test]
    #[should_fail]
    #[ignore(cfg(windows))]
    fn gen_uint_range_fail() {
        rand::rng().gen_uint_range(5u, 2u);
    }

    #[test]
    fn gen_float() {
        let r = rand::rng();
        let a = r.gen_float();
        let b = r.gen_float();
        log(debug, (a, b));
    }

    #[test]
    fn gen_weighted_bool() {
        let r = rand::rng();
        assert r.gen_weighted_bool(0u) == true;
        assert r.gen_weighted_bool(1u) == true;
    }

    #[test]
    fn gen_str() {
        let r = rand::rng();
        log(debug, r.gen_str(10u));
        log(debug, r.gen_str(10u));
        log(debug, r.gen_str(10u));
        assert r.gen_str(0u).len() == 0u;
        assert r.gen_str(10u).len() == 10u;
        assert r.gen_str(16u).len() == 16u;
    }

    #[test]
    fn gen_bytes() {
        let r = rand::rng();
        assert r.gen_bytes(0u).len() == 0u;
        assert r.gen_bytes(10u).len() == 10u;
        assert r.gen_bytes(16u).len() == 16u;
    }

    #[test]
    fn choose() {
        let r = rand::rng();
        assert r.choose(~[1, 1, 1]) == 1;
    }

    #[test]
    fn choose_option() {
        let r = rand::rng();
        assert r.choose_option(~[]) == none::<int>;
        assert r.choose_option(~[1, 1, 1]) == some(1);
    }

    #[test]
    fn choose_weighted() {
        let r = rand::rng();
        assert r.choose_weighted(~[{weight: 1u, item: 42}]) == 42;
        assert r.choose_weighted(~[
            {weight: 0u, item: 42},
            {weight: 1u, item: 43}
        ]) == 43;
    }

    #[test]
    fn choose_weighted_option() {
        let r = rand::rng();
        assert r.choose_weighted_option(~[{weight: 1u, item: 42}]) ==
               some(42);
        assert r.choose_weighted_option(~[
            {weight: 0u, item: 42},
            {weight: 1u, item: 43}
        ]) == some(43);
        assert r.choose_weighted_option(~[]) == none::<int>;
    }

    #[test]
    fn weighted_vec() {
        let r = rand::rng();
        let empty: ~[int] = ~[];
        assert r.weighted_vec(~[]) == empty;
        assert r.weighted_vec(~[
            {weight: 0u, item: 3u},
            {weight: 1u, item: 2u},
            {weight: 2u, item: 1u}
        ]) == ~[2u, 1u, 1u];
    }

    #[test]
    fn shuffle() {
        let r = rand::rng();
        let empty: ~[int] = ~[];
        assert r.shuffle(~[]) == empty;
        assert r.shuffle(~[1, 1, 1]) == ~[1, 1, 1];
    }
}

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// End:
<|endoftext|>