<commit_before>use std::sync::atomic::{AtomicUsize, Ordering};
use std::thread;
use std::time::Duration;

use crossbeam_deque::{Injector, Steal, Stealer, Worker};
use rand::distributions::{Distribution, Uniform};
use rand::thread_rng;
use scoped_threadpool::Pool;

use gc::root::Slot;
use gc::{Address, Region};

pub fn start(rootset: &[Slot], heap: Region, perm: Region, threadpool: &mut Pool) {
    let number_workers = threadpool.thread_count() as usize;
    let mut workers = Vec::with_capacity(number_workers);
    let mut stealers = Vec::with_capacity(number_workers);
    let injector = Injector::new();

    for _ in 0..number_workers {
        let w = Worker::new_lifo();
        let s = w.stealer();
        workers.push(w);
        stealers.push(s);
    }

    for root in rootset {
        let root_ptr = root.get();

        if heap.contains(root_ptr) {
            let root_obj = root_ptr.to_mut_obj();

            if !root_obj.header().is_marked_non_atomic() {
                root_obj.header_mut().mark_non_atomic();
                injector.push(root_ptr);
            }
        } else {
            debug_assert!(root_ptr.is_null() || perm.contains(root_ptr));
        }
    }

    let terminator = Terminator::new(number_workers);

    threadpool.scoped(|scoped| {
        for (task_id, worker) in workers.into_iter().enumerate() {
            let heap_region = heap.clone();
            let perm_region = perm.clone();

            let injector = &injector;
            let stealers = &stealers;
            let terminator = &terminator;

            scoped.execute(move || {
                let mut task = MarkingTask {
                    task_id: task_id,
                    local: Segment::new(),
                    worker: worker,
                    injector: injector,
                    stealers: stealers,
                    terminator: terminator,
                    heap_region: heap_region,
                    perm_region: perm_region,
                    marked: 0,
                };

                task.run();
            });
        }
    });
}

pub struct Terminator {
    const_nworkers: usize,
    nworkers: AtomicUsize,
}

impl Terminator {
    pub fn new(number_workers: usize) -> Terminator {
        Terminator {
            const_nworkers: number_workers,
            nworkers: AtomicUsize::new(number_workers),
        }
    }

    pub fn try_terminate(&self) -> bool {
        if self.const_nworkers == 1 {
            return;
        }

        self.decrease_workers();
        thread::sleep(Duration::from_micros(1));
        self.zero_or_increase_workers()
    }

    fn decrease_workers(&self) -> bool {
        self.nworkers.fetch_sub(1, Ordering::SeqCst) == 1
    }

    fn zero_or_increase_workers(&self) -> bool {
        let mut nworkers = self.nworkers.load(Ordering::Relaxed);

        loop {
            if nworkers == 0 {
                return true;
            }

            let prev_nworkers =
                self.nworkers
                    .compare_and_swap(nworkers, nworkers + 1, Ordering::SeqCst);

            if nworkers == prev_nworkers {
                return false;
            }

            nworkers = prev_nworkers;
        }
    }
}

struct MarkingTask<'a> {
    task_id: usize,
    local: Segment,
    worker: Worker<Address>,
    injector: &'a Injector<Address>,
    stealers: &'a [Stealer<Address>],
    terminator: &'a Terminator,
    heap_region: Region,
    perm_region: Region,
    marked: usize,
}

impl<'a> MarkingTask<'a> {
    fn pop(&mut self) -> Option<Address> {
        self.pop_local()
            .or_else(|| self.pop_worker())
            .or_else(|| self.pop_global())
            .or_else(|| self.steal())
    }

    fn pop_local(&mut self) -> Option<Address> {
        if self.local.is_empty() {
            return None;
        }

        let obj = self.local.pop().expect("should be non-empty");
        Some(obj)
    }

    fn pop_worker(&mut self) -> Option<Address> {
        self.worker.pop()
    }

    fn pop_global(&mut self) -> Option<Address> {
        loop {
            let result = self.injector.steal_batch_and_pop(&mut self.worker);

            match result {
                Steal::Empty => break,
                Steal::Success(value) => return Some(value),
                Steal::Retry => continue,
            }
        }

        None
    }

    fn steal(&self) -> Option<Address> {
        if self.stealers.len() == 1 {
            return None;
        }

        let mut rng = thread_rng();
        let range = Uniform::new(0, self.stealers.len());

        for _ in 0..2 * self.stealers.len() {
            let mut stealer_id = self.task_id;

            while stealer_id == self.task_id {
                stealer_id = range.sample(&mut rng);
            }

            let stealer = &self.stealers[stealer_id];

            loop {
                match stealer.steal_batch_and_pop(&self.worker) {
                    Steal::Empty => break,
                    Steal::Success(address) => return Some(address),
                    Steal::Retry => continue,
                }
            }
        }

        None
    }

    fn run(&mut self) {
        loop {
            let object_addr = if let Some(object_addr) = self.pop() {
                object_addr
            } else if self.terminator.try_terminate() {
                break;
            } else {
                continue;
            };

            let object = object_addr.to_mut_obj();

            object.visit_reference_fields(|field| {
                self.trace(field);
            });
        }
    }

    fn trace(&mut self, slot: Slot) {
        let field_addr = slot.get();

        if self.heap_region.contains(field_addr) {
            let field_obj = field_addr.to_mut_obj();

            if field_obj.header().try_mark_non_atomic() {
                if self.local.has_capacity() {
                    self.local.push(field_addr);
                    self.defensive_push();
                } else {
                    self.worker.push(field_addr);
                }
            }
        } else {
            debug_assert!(field_addr.is_null() || self.perm_region.contains(field_addr));
        }
    }

    fn defensive_push(&mut self) {
        self.marked += 1;

        if self.marked > 256 {
            if self.local.len() > 4 {
                let target_len = self.local.len() / 2;

                while self.local.len() > target_len {
                    let val = self.local.pop().unwrap();
                    self.injector.push(val);
                }
            }

            self.marked = 0;
        }
    }
}

const SEGMENT_SIZE: usize = 64;

struct Segment {
    data: Vec<Address>,
}

impl Segment {
    fn new() -> Segment {
        Segment {
            data: Vec::with_capacity(SEGMENT_SIZE),
        }
    }

    fn empty() -> Segment {
        Segment { data: Vec::new() }
    }

    fn with(addr: Address) -> Segment {
        let mut segment = Segment::new();
        segment.data.push(addr);

        segment
    }

    fn has_capacity(&self) -> bool {
        self.data.len() < SEGMENT_SIZE
    }

    fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    fn push(&mut self, addr: Address) {
        debug_assert!(self.has_capacity());
        self.data.push(addr);
    }

    fn pop(&mut self) -> Option<Address> {
        self.data.pop()
    }

    fn len(&mut self) -> usize {
        self.data.len()
    }
}
<commit_msg>fix build<commit_after>use std::sync::atomic::{AtomicUsize, Ordering};
use std::thread;
use std::time::Duration;

use crossbeam_deque::{Injector, Steal, Stealer, Worker};
use rand::distributions::{Distribution, Uniform};
use rand::thread_rng;
use scoped_threadpool::Pool;

use gc::root::Slot;
use gc::{Address, Region};

pub fn start(rootset: &[Slot], heap: Region, perm: Region, threadpool: &mut Pool) {
    let number_workers = threadpool.thread_count() as usize;
    let mut workers = Vec::with_capacity(number_workers);
    let mut stealers = Vec::with_capacity(number_workers);
    let injector = Injector::new();

    for _ in 0..number_workers {
        let w = Worker::new_lifo();
        let s = w.stealer();
        workers.push(w);
        stealers.push(s);
    }

    for root in rootset {
        let root_ptr = root.get();

        if heap.contains(root_ptr) {
            let root_obj = root_ptr.to_mut_obj();

            if !root_obj.header().is_marked_non_atomic() {
                root_obj.header_mut().mark_non_atomic();
                injector.push(root_ptr);
            }
        } else {
            debug_assert!(root_ptr.is_null() || perm.contains(root_ptr));
        }
    }

    let terminator = Terminator::new(number_workers);

    threadpool.scoped(|scoped| {
        for (task_id, worker) in workers.into_iter().enumerate() {
            let heap_region = heap.clone();
            let perm_region = perm.clone();

            let injector = &injector;
            let stealers = &stealers;
            let terminator = &terminator;

            scoped.execute(move || {
                let mut task = MarkingTask {
                    task_id: task_id,
                    local: Segment::new(),
                    worker: worker,
                    injector: injector,
                    stealers: stealers,
                    terminator: terminator,
                    heap_region: heap_region,
                    perm_region: perm_region,
                    marked: 0,
                };

                task.run();
            });
        }
    });
}

pub struct Terminator {
    const_nworkers: usize,
    nworkers: AtomicUsize,
}

impl Terminator {
    pub fn new(number_workers: usize) -> Terminator {
        Terminator {
            const_nworkers: number_workers,
            nworkers: AtomicUsize::new(number_workers),
        }
    }

    pub fn try_terminate(&self) -> bool {
        if self.const_nworkers == 1 {
            return true;
        }

        self.decrease_workers();
        thread::sleep(Duration::from_micros(1));
        self.zero_or_increase_workers()
    }

    fn decrease_workers(&self) -> bool {
        self.nworkers.fetch_sub(1, Ordering::SeqCst) == 1
    }

    fn zero_or_increase_workers(&self) -> bool {
        let mut nworkers = self.nworkers.load(Ordering::Relaxed);

        loop {
            if nworkers == 0 {
                return true;
            }

            let prev_nworkers =
                self.nworkers
                    .compare_and_swap(nworkers, nworkers + 1, Ordering::SeqCst);

            if nworkers == prev_nworkers {
                return false;
            }

            nworkers = prev_nworkers;
        }
    }
}

struct MarkingTask<'a> {
    task_id: usize,
    local: Segment,
    worker: Worker<Address>,
    injector: &'a Injector<Address>,
    stealers: &'a [Stealer<Address>],
    terminator: &'a Terminator,
    heap_region: Region,
    perm_region: Region,
    marked: usize,
}

impl<'a> MarkingTask<'a> {
    fn pop(&mut self) -> Option<Address> {
        self.pop_local()
            .or_else(|| self.pop_worker())
            .or_else(|| self.pop_global())
            .or_else(|| self.steal())
    }

    fn pop_local(&mut self) -> Option<Address> {
        if self.local.is_empty() {
            return None;
        }

        let obj = self.local.pop().expect("should be non-empty");
        Some(obj)
    }

    fn pop_worker(&mut self) -> Option<Address> {
        self.worker.pop()
    }

    fn pop_global(&mut self) -> Option<Address> {
        loop {
            let result = self.injector.steal_batch_and_pop(&mut self.worker);

            match result {
                Steal::Empty => break,
                Steal::Success(value) => return Some(value),
                Steal::Retry => continue,
            }
        }

        None
    }

    fn steal(&self) -> Option<Address> {
        if self.stealers.len() == 1 {
            return None;
        }

        let mut rng = thread_rng();
        let range = Uniform::new(0, self.stealers.len());

        for _ in 0..2 * self.stealers.len() {
            let mut stealer_id = self.task_id;

            while stealer_id == self.task_id {
                stealer_id = range.sample(&mut rng);
            }

            let stealer = &self.stealers[stealer_id];

            loop {
                match stealer.steal_batch_and_pop(&self.worker) {
                    Steal::Empty => break,
                    Steal::Success(address) => return Some(address),
                    Steal::Retry => continue,
                }
            }
        }

        None
    }

    fn run(&mut self) {
        loop {
            let object_addr = if let Some(object_addr) = self.pop() {
                object_addr
            } else if self.terminator.try_terminate() {
                break;
            } else {
                continue;
            };

            let object = object_addr.to_mut_obj();

            object.visit_reference_fields(|field| {
                self.trace(field);
            });
        }
    }

    fn trace(&mut self, slot: Slot) {
        let field_addr = slot.get();

        if self.heap_region.contains(field_addr) {
            let field_obj = field_addr.to_mut_obj();

            if field_obj.header().try_mark_non_atomic() {
                if self.local.has_capacity() {
                    self.local.push(field_addr);
                    self.defensive_push();
                } else {
                    self.worker.push(field_addr);
                }
            }
        } else {
            debug_assert!(field_addr.is_null() || self.perm_region.contains(field_addr));
        }
    }

    fn defensive_push(&mut self) {
        self.marked += 1;

        if self.marked > 256 {
            if self.local.len() > 4 {
                let target_len = self.local.len() / 2;

                while self.local.len() > target_len {
                    let val = self.local.pop().unwrap();
                    self.injector.push(val);
                }
            }

            self.marked = 0;
        }
    }
}

const SEGMENT_SIZE: usize = 64;

struct Segment {
    data: Vec<Address>,
}

impl Segment {
    fn new() -> Segment {
        Segment {
            data: Vec::with_capacity(SEGMENT_SIZE),
        }
    }

    fn empty() -> Segment {
        Segment { data: Vec::new() }
    }

    fn with(addr: Address) -> Segment {
        let mut segment = Segment::new();
        segment.data.push(addr);

        segment
    }

    fn has_capacity(&self) -> bool {
        self.data.len() < SEGMENT_SIZE
    }

    fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    fn push(&mut self, addr: Address) {
        debug_assert!(self.has_capacity());
        self.data.push(addr);
    }

    fn pop(&mut self) -> Option<Address> {
        self.data.pop()
    }

    fn len(&mut self) -> usize {
        self.data.len()
    }
}
<|endoftext|>


<commit_before>use crate::{context::Has, graphics::gpu::bind_group::BindGroupBuilder, GameError, GameResult};

use super::{
    context::GraphicsContext,
    draw::{DrawParam, DrawUniforms, Std140DrawUniforms},
    gpu::arc::{ArcBindGroup, ArcBindGroupLayout, ArcBuffer},
    internal_canvas::InstanceArrayView,
    transform_rect, Canvas, Draw, Drawable, Image, Mesh, Rect, WgpuContext,
};
use crevice::std140::AsStd140;
use std::{
    collections::BTreeMap,
    sync::{
        atomic::{AtomicBool, AtomicU32, Ordering::SeqCst},
        Mutex,
    },
};

/// Array of instances for fast rendering of many meshes.
///
/// Traditionally known as a "batch".
#[derive(Debug)]
pub struct InstanceArray {
    pub(crate) buffer: Mutex<ArcBuffer>,
    pub(crate) indices: Mutex<ArcBuffer>,
    pub(crate) bind_group: Mutex<ArcBindGroup>,
    pub(crate) bind_layout: ArcBindGroupLayout,
    pub(crate) image: Image,
    pub(crate) ordered: bool,
    dirty: AtomicBool,
    capacity: AtomicU32,
    uniforms: Vec<Std140DrawUniforms>,
    params: Vec<DrawParam>,
}

impl InstanceArray {
    /// Creates a new [InstanceArray] capable of storing up to n-`capacity` instances
    /// (this can be changed and is resized automatically when needed).
    ///
    /// If `image` is `None`, a 1x1 white image will be used which can be used to draw solid rectangles.
    ///
    /// This constructor is `unordered` meaning instances will be drawn by their push/index order. Use [`InstanceArray::new_ordered`] to order by z-value.
    pub fn new(
        gfx: &impl Has<GraphicsContext>,
        image: impl Into<Option<Image>>,
        capacity: u32,
    ) -> Self {
        let gfx = gfx.retrieve();
        InstanceArray::new_wgpu(
            &gfx.wgpu,
            gfx.instance_bind_layout.clone(),
            image.into().unwrap_or_else(|| gfx.white_image.clone()),
            capacity,
            false,
        )
    }

    /// See [`InstanceArray::new`] for details.
    ///
    /// This constructor is `ordered` meaning instances will be drawn by their z-value at a slight performance cost. Use [`InstanceArray::new`] to order by index.
    pub fn new_ordered(
        gfx: &impl Has<GraphicsContext>,
        image: impl Into<Option<Image>>,
        capacity: u32,
    ) -> Self {
        let gfx = gfx.retrieve();
        InstanceArray::new_wgpu(
            &gfx.wgpu,
            gfx.instance_bind_layout.clone(),
            image.into().unwrap_or_else(|| gfx.white_image.clone()),
            capacity,
            false,
        )
    }

    fn new_wgpu(
        wgpu: &WgpuContext,
        bind_layout: ArcBindGroupLayout,
        image: Image,
        capacity: u32,
        ordered: bool,
    ) -> Self {
        assert!(capacity > 0);

        let buffer = ArcBuffer::new(wgpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: DrawUniforms::std140_size_static() as u64 * capacity as u64,
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_DST
                | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        }));

        let indices = ArcBuffer::new(wgpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: if ordered {
                std::mem::size_of::<u32>() as u64 * capacity as u64
            } else {
                4 // min for layout
            },
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_SRC
                | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        }));

        let bind_group = BindGroupBuilder::new()
            .buffer(
                &buffer,
                0,
                wgpu::ShaderStages::VERTEX,
                wgpu::BufferBindingType::Storage { read_only: true },
                false,
                None,
            )
            .buffer(
                &indices,
                0,
                wgpu::ShaderStages::VERTEX,
                wgpu::BufferBindingType::Storage { read_only: true },
                false,
                None,
            );
        let bind_group =
            ArcBindGroup::new(wgpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &bind_layout,
                entries: bind_group.entries(),
            }));

        let uniforms = Vec::with_capacity(capacity as usize);
        let params = Vec::with_capacity(capacity as usize);

        InstanceArray {
            buffer: Mutex::new(buffer),
            indices: Mutex::new(indices),
            bind_group: Mutex::new(bind_group),
            bind_layout,
            image,
            ordered,
            dirty: AtomicBool::new(false),
            capacity: AtomicU32::new(capacity),
            uniforms,
            params,
        }
    }

    /// Resets all the instance data to a set of `DrawParam`.
    pub fn set(&mut self, instances: impl IntoIterator<Item = DrawParam>) {
        self.dirty.store(true, SeqCst);
        self.params.clear();
        self.params.extend(instances);
        self.uniforms.clear();
        self.uniforms.extend(
            self.params
                .iter()
                .map(|x| DrawUniforms::from_param(x, None).as_std140()),
        );
    }

    /// Pushes a new instance onto the end.
    pub fn push(&mut self, instance: DrawParam) {
        self.dirty.store(true, SeqCst);
        self.uniforms
            .push(DrawUniforms::from_param(&instance, None).as_std140());
        self.params.push(instance);
    }

    /// Updates an existing instance at a given index, if it is valid.
    pub fn update(&mut self, index: u32, instance: DrawParam) {
        if let Some((uniform, param)) = self
            .uniforms
            .get_mut(index as usize)
            .and_then(|x| Some((x, self.params.get_mut(index as usize)?)))
        {
            self.dirty.store(true, SeqCst);
            *uniform = DrawUniforms::from_param(&instance, None).as_std140();
            *param = instance;
        }
    }

    /// Clears all instance data.
    pub fn clear(&mut self) {
        // don't need to set dirty here
        self.uniforms.clear();
        self.params.clear();
    }

    /// Returns whether the instance data has been changed without being flushed (i.e., uploaded to the GPU).
    #[inline]
    pub fn is_dirty(&self) -> bool {
        self.dirty.load(SeqCst)
    }

    /// Returns an immutable slice of all the instance data in this [`InstanceArray`].
    #[inline]
    pub fn instances(&self) -> &[DrawParam] {
        &self.params
    }

    #[allow(unsafe_code)]
    pub(crate) fn flush_wgpu(&self, wgpu: &WgpuContext) -> GameResult<()> {
        if !self.dirty.load(SeqCst) {
            return Ok(());
        } else {
            self.dirty.store(false, SeqCst);
        }

        let len = self.uniforms.len() as u32;
        if len > self.capacity.load(SeqCst) {
            let mut resized = InstanceArray::new_wgpu(
                wgpu,
                self.bind_layout.clone(),
                self.image.clone(),
                len,
                self.ordered,
            );
            *self.buffer.lock().map_err(|_| GameError::LockError)? =
                resized.buffer.get_mut().unwrap().clone();
            *self.indices.lock().map_err(|_| GameError::LockError)? =
                resized.indices.get_mut().unwrap().clone();
            *self.bind_group.lock().map_err(|_| GameError::LockError)? =
                resized.bind_group.get_mut().unwrap().clone();
            self.capacity.store(len, SeqCst);
        }

        wgpu.queue.write_buffer(
            &self.buffer.lock().unwrap(),
            0,
            bytemuck::cast_slice(self.uniforms.as_slice()),
        );

        if self.ordered {
            let mut layers = BTreeMap::<_, Vec<_>>::new();
            for (i, param) in self.params.iter().enumerate() {
                layers.entry(param.z).or_default().push(i as u32);
            }
            let indices = layers.into_values().flatten().collect::<Vec<_>>();
            wgpu.queue.write_buffer(
                &self.indices.lock().unwrap(),
                0,
                bytemuck::cast_slice(indices.as_slice()),
            );
        }

        Ok(())
    }

    /// Changes the capacity of this `InstanceArray` while preserving instances.
    ///
    /// If `new_capacity` is less than the `len`, the instances will be truncated.
    pub fn resize(&mut self, gfx: &impl Has<GraphicsContext>, new_capacity: u32) {
        let resized = if self.ordered {
            InstanceArray::new_ordered(gfx, self.image.clone(), new_capacity)
        } else {
            InstanceArray::new(gfx, self.image.clone(), new_capacity)
        };
        self.buffer = resized.buffer;
        self.indices = resized.indices;
        self.capacity.store(new_capacity, SeqCst);
        self.dirty.store(true, SeqCst);
        self.uniforms.truncate(new_capacity as usize);
        self.params.truncate(new_capacity as usize);
        self.uniforms
            .reserve((new_capacity as usize).min(self.uniforms.len()) - self.uniforms.len());
        self.params
            .reserve((new_capacity as usize).min(self.params.len()) - self.params.len());
    }

    /// Returns this `InstanceArray`'s associated `image`.
    #[inline]
    pub fn image(&self) -> Image {
        self.image.clone()
    }

    /// Returns the number of instances this [InstanceArray] is capable of holding.
    /// This number was specified when creating the [InstanceArray], or if the [InstanceArray]
    /// was automatically resized, the greatest length of instances.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.capacity.load(SeqCst) as usize
    }

    /// This is equivalent to `<InstanceArray as Drawable>::dimensions()` (see [`Drawable::dimensions()`]), but with a mesh taken into account.
    ///
    /// Essentially, consider `<InstanceArray as Drawable>::dimensions()` to be the bounds when the [`InstanceArray`] is drawn with `canvas.draw()`,
    /// and consider [`InstanceArray::dimensions_meshed()`] to be the bounds when the [`InstanceArray`] is drawn with `canvas.draw_instanced_mesh()`.
    pub fn dimensions_meshed(&self, gfx: &impl Has<GraphicsContext>, mesh: &Mesh) -> Option<Rect> {
        if self.params.is_empty() {
            return None;
        }
        let dimensions = mesh.dimensions(gfx)?;
        self.params
            .iter()
            .map(|&param| transform_rect(dimensions, param))
            .fold(None, |acc: Option<Rect>, rect| {
                Some(if let Some(acc) = acc {
                    acc.combine_with(rect)
                } else {
                    rect
                })
            })
    }
}

impl Drawable for InstanceArray {
    fn draw(&self, canvas: &mut Canvas, param: impl Into<DrawParam>) {
        self.flush_wgpu(&canvas.wgpu).unwrap();
        canvas.push_draw(
            Draw::MeshInstances {
                mesh: canvas.default_resources().mesh.clone(),
                instances: InstanceArrayView::from_instances(self).unwrap(),
                scale: true,
            },
            param.into(),
        );
    }

    fn dimensions(&self, gfx: &impl Has<GraphicsContext>) -> Option<Rect> {
        let gfx = gfx.retrieve();
        if self.params.is_empty() {
            return None;
        }
        let dimensions = self.image.dimensions(gfx)?;
        self.params
            .iter()
            .map(|&param| transform_rect(dimensions, param))
            .fold(None, |acc: Option<Rect>, rect| {
                Some(if let Some(acc) = acc {
                    acc.combine_with(rect)
                } else {
                    rect
                })
            })
    }
}
<commit_msg>fix InstanceArray::new_ordered typo<commit_after>use crate::{context::Has, graphics::gpu::bind_group::BindGroupBuilder, GameError, GameResult};

use super::{
    context::GraphicsContext,
    draw::{DrawParam, DrawUniforms, Std140DrawUniforms},
    gpu::arc::{ArcBindGroup, ArcBindGroupLayout, ArcBuffer},
    internal_canvas::InstanceArrayView,
    transform_rect, Canvas, Draw, Drawable, Image, Mesh, Rect, WgpuContext,
};
use crevice::std140::AsStd140;
use std::{
    collections::BTreeMap,
    sync::{
        atomic::{AtomicBool, AtomicU32, Ordering::SeqCst},
        Mutex,
    },
};

/// Array of instances for fast rendering of many meshes.
///
/// Traditionally known as a "batch".
#[derive(Debug)]
pub struct InstanceArray {
    pub(crate) buffer: Mutex<ArcBuffer>,
    pub(crate) indices: Mutex<ArcBuffer>,
    pub(crate) bind_group: Mutex<ArcBindGroup>,
    pub(crate) bind_layout: ArcBindGroupLayout,
    pub(crate) image: Image,
    pub(crate) ordered: bool,
    dirty: AtomicBool,
    capacity: AtomicU32,
    uniforms: Vec<Std140DrawUniforms>,
    params: Vec<DrawParam>,
}

impl InstanceArray {
    /// Creates a new [InstanceArray] capable of storing up to n-`capacity` instances
    /// (this can be changed and is resized automatically when needed).
    ///
    /// If `image` is `None`, a 1x1 white image will be used which can be used to draw solid rectangles.
    ///
    /// This constructor is `unordered` meaning instances will be drawn by their push/index order. Use [`InstanceArray::new_ordered`] to order by z-value.
    pub fn new(
        gfx: &impl Has<GraphicsContext>,
        image: impl Into<Option<Image>>,
        capacity: u32,
    ) -> Self {
        let gfx = gfx.retrieve();
        InstanceArray::new_wgpu(
            &gfx.wgpu,
            gfx.instance_bind_layout.clone(),
            image.into().unwrap_or_else(|| gfx.white_image.clone()),
            capacity,
            false,
        )
    }

    /// See [`InstanceArray::new`] for details.
    ///
    /// This constructor is `ordered` meaning instances will be drawn by their z-value at a slight performance cost. Use [`InstanceArray::new`] to order by index.
    pub fn new_ordered(
        gfx: &impl Has<GraphicsContext>,
        image: impl Into<Option<Image>>,
        capacity: u32,
    ) -> Self {
        let gfx = gfx.retrieve();
        InstanceArray::new_wgpu(
            &gfx.wgpu,
            gfx.instance_bind_layout.clone(),
            image.into().unwrap_or_else(|| gfx.white_image.clone()),
            capacity,
            true,
        )
    }

    fn new_wgpu(
        wgpu: &WgpuContext,
        bind_layout: ArcBindGroupLayout,
        image: Image,
        capacity: u32,
        ordered: bool,
    ) -> Self {
        assert!(capacity > 0);

        let buffer = ArcBuffer::new(wgpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: DrawUniforms::std140_size_static() as u64 * capacity as u64,
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_DST
                | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        }));

        let indices = ArcBuffer::new(wgpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: if ordered {
                std::mem::size_of::<u32>() as u64 * capacity as u64
            } else {
                4 // min for layout
            },
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_SRC
                | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        }));

        let bind_group = BindGroupBuilder::new()
            .buffer(
                &buffer,
                0,
                wgpu::ShaderStages::VERTEX,
                wgpu::BufferBindingType::Storage { read_only: true },
                false,
                None,
            )
            .buffer(
                &indices,
                0,
                wgpu::ShaderStages::VERTEX,
                wgpu::BufferBindingType::Storage { read_only: true },
                false,
                None,
            );
        let bind_group =
            ArcBindGroup::new(wgpu.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &bind_layout,
                entries: bind_group.entries(),
            }));

        let uniforms = Vec::with_capacity(capacity as usize);
        let params = Vec::with_capacity(capacity as usize);

        InstanceArray {
            buffer: Mutex::new(buffer),
            indices: Mutex::new(indices),
            bind_group: Mutex::new(bind_group),
            bind_layout,
            image,
            ordered,
            dirty: AtomicBool::new(false),
            capacity: AtomicU32::new(capacity),
            uniforms,
            params,
        }
    }

    /// Resets all the instance data to a set of `DrawParam`.
    pub fn set(&mut self, instances: impl IntoIterator<Item = DrawParam>) {
        self.dirty.store(true, SeqCst);
        self.params.clear();
        self.params.extend(instances);
        self.uniforms.clear();
        self.uniforms.extend(
            self.params
                .iter()
                .map(|x| DrawUniforms::from_param(x, None).as_std140()),
        );
    }

    /// Pushes a new instance onto the end.
    pub fn push(&mut self, instance: DrawParam) {
        self.dirty.store(true, SeqCst);
        self.uniforms
            .push(DrawUniforms::from_param(&instance, None).as_std140());
        self.params.push(instance);
    }

    /// Updates an existing instance at a given index, if it is valid.
    pub fn update(&mut self, index: u32, instance: DrawParam) {
        if let Some((uniform, param)) = self
            .uniforms
            .get_mut(index as usize)
            .and_then(|x| Some((x, self.params.get_mut(index as usize)?)))
        {
            self.dirty.store(true, SeqCst);
            *uniform = DrawUniforms::from_param(&instance, None).as_std140();
            *param = instance;
        }
    }

    /// Clears all instance data.
    pub fn clear(&mut self) {
        // don't need to set dirty here
        self.uniforms.clear();
        self.params.clear();
    }

    /// Returns whether the instance data has been changed without being flushed (i.e., uploaded to the GPU).
    #[inline]
    pub fn is_dirty(&self) -> bool {
        self.dirty.load(SeqCst)
    }

    /// Returns an immutable slice of all the instance data in this [`InstanceArray`].
    #[inline]
    pub fn instances(&self) -> &[DrawParam] {
        &self.params
    }

    #[allow(unsafe_code)]
    pub(crate) fn flush_wgpu(&self, wgpu: &WgpuContext) -> GameResult<()> {
        if !self.dirty.load(SeqCst) {
            return Ok(());
        } else {
            self.dirty.store(false, SeqCst);
        }

        let len = self.uniforms.len() as u32;
        if len > self.capacity.load(SeqCst) {
            let mut resized = InstanceArray::new_wgpu(
                wgpu,
                self.bind_layout.clone(),
                self.image.clone(),
                len,
                self.ordered,
            );
            *self.buffer.lock().map_err(|_| GameError::LockError)? =
                resized.buffer.get_mut().unwrap().clone();
            *self.indices.lock().map_err(|_| GameError::LockError)? =
                resized.indices.get_mut().unwrap().clone();
            *self.bind_group.lock().map_err(|_| GameError::LockError)? =
                resized.bind_group.get_mut().unwrap().clone();
            self.capacity.store(len, SeqCst);
        }

        wgpu.queue.write_buffer(
            &self.buffer.lock().unwrap(),
            0,
            bytemuck::cast_slice(self.uniforms.as_slice()),
        );

        if self.ordered {
            let mut layers = BTreeMap::<_, Vec<_>>::new();
            for (i, param) in self.params.iter().enumerate() {
                layers.entry(param.z).or_default().push(i as u32);
            }
            let indices = layers.into_values().flatten().collect::<Vec<_>>();
            wgpu.queue.write_buffer(
                &self.indices.lock().unwrap(),
                0,
                bytemuck::cast_slice(indices.as_slice()),
            );
        }

        Ok(())
    }

    /// Changes the capacity of this `InstanceArray` while preserving instances.
    ///
    /// If `new_capacity` is less than the `len`, the instances will be truncated.
    pub fn resize(&mut self, gfx: &impl Has<GraphicsContext>, new_capacity: u32) {
        let resized = if self.ordered {
            InstanceArray::new_ordered(gfx, self.image.clone(), new_capacity)
        } else {
            InstanceArray::new(gfx, self.image.clone(), new_capacity)
        };
        self.buffer = resized.buffer;
        self.indices = resized.indices;
        self.capacity.store(new_capacity, SeqCst);
        self.dirty.store(true, SeqCst);
        self.uniforms.truncate(new_capacity as usize);
        self.params.truncate(new_capacity as usize);
        self.uniforms
            .reserve((new_capacity as usize).min(self.uniforms.len()) - self.uniforms.len());
        self.params
            .reserve((new_capacity as usize).min(self.params.len()) - self.params.len());
    }

    /// Returns this `InstanceArray`'s associated `image`.
    #[inline]
    pub fn image(&self) -> Image {
        self.image.clone()
    }

    /// Returns the number of instances this [InstanceArray] is capable of holding.
    /// This number was specified when creating the [InstanceArray], or if the [InstanceArray]
    /// was automatically resized, the greatest length of instances.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.capacity.load(SeqCst) as usize
    }

    /// This is equivalent to `<InstanceArray as Drawable>::dimensions()` (see [`Drawable::dimensions()`]), but with a mesh taken into account.
    ///
    /// Essentially, consider `<InstanceArray as Drawable>::dimensions()` to be the bounds when the [`InstanceArray`] is drawn with `canvas.draw()`,
    /// and consider [`InstanceArray::dimensions_meshed()`] to be the bounds when the [`InstanceArray`] is drawn with `canvas.draw_instanced_mesh()`.
    pub fn dimensions_meshed(&self, gfx: &impl Has<GraphicsContext>, mesh: &Mesh) -> Option<Rect> {
        if self.params.is_empty() {
            return None;
        }
        let dimensions = mesh.dimensions(gfx)?;
        self.params
            .iter()
            .map(|&param| transform_rect(dimensions, param))
            .fold(None, |acc: Option<Rect>, rect| {
                Some(if let Some(acc) = acc {
                    acc.combine_with(rect)
                } else {
                    rect
                })
            })
    }
}

impl Drawable for InstanceArray {
    fn draw(&self, canvas: &mut Canvas, param: impl Into<DrawParam>) {
        self.flush_wgpu(&canvas.wgpu).unwrap();
        canvas.push_draw(
            Draw::MeshInstances {
                mesh: canvas.default_resources().mesh.clone(),
                instances: InstanceArrayView::from_instances(self).unwrap(),
                scale: true,
            },
            param.into(),
        );
    }

    fn dimensions(&self, gfx: &impl Has<GraphicsContext>) -> Option<Rect> {
        let gfx = gfx.retrieve();
        if self.params.is_empty() {
            return None;
        }
        let dimensions = self.image.dimensions(gfx)?;
        self.params
            .iter()
            .map(|&param| transform_rect(dimensions, param))
            .fold(None, |acc: Option<Rect>, rect| {
                Some(if let Some(acc) = acc {
                    acc.combine_with(rect)
                } else {
                    rect
                })
            })
    }
}
<|endoftext|>


<commit_before>use std::vec::Vec;
use std::io::Write;

use super::{Result, Error};

const ENC_GROUP_SIZE: usize = 8;
const ENC_MARKER: u8 = b'\xff';
const ENC_PADDING: [u8; ENC_GROUP_SIZE] = [0; ENC_GROUP_SIZE];

// returns the maximum encoded bytes size.
pub fn max_encoded_bytes_size(n: usize) -> usize {
    (n / ENC_GROUP_SIZE + 1) * (ENC_GROUP_SIZE + 1)
}

// Refer: https://github.com/facebook/mysql-5.6/wiki/MyRocks-record-format#memcomparable-format
pub fn encode_bytes(key: &[u8]) -> Vec<u8> {
    let cap = max_encoded_bytes_size(key.len());
    let mut encoded = Vec::<u8>::with_capacity(cap);
    let len = key.len();
    let mut index = 0;
    while index <= len {
        let remain = len - index;
        let mut pad: usize = 0;
        if remain > ENC_GROUP_SIZE {
            encoded.write(&key[index..index + ENC_GROUP_SIZE]).unwrap();
        } else {
            pad = ENC_GROUP_SIZE - remain;
            encoded.write(&key[index..]).unwrap();
            encoded.write(&ENC_PADDING[..pad]).unwrap();
        }
        encoded.push(ENC_MARKER - (pad as u8));
        index += ENC_GROUP_SIZE;
    }
    encoded
}

pub fn decode_bytes(data: &[u8]) -> Result<(Vec<u8>, usize)> {
    let mut key = Vec::<u8>::with_capacity(data.len());
    let mut read: usize = 0;
    for chunk in data.chunks(ENC_GROUP_SIZE + 1) {
        if chunk.len() != ENC_GROUP_SIZE + 1 {
            return Err(Error::KeyLength);
        }
        read += ENC_GROUP_SIZE + 1;

        let (marker, bytes) = chunk.split_last().unwrap();
        let pad_size = (ENC_MARKER - *marker) as usize;
        if pad_size == 0 {
            key.write(bytes).unwrap();
            continue;
        }
        if pad_size > ENC_GROUP_SIZE {
            return Err(Error::KeyPadding);
        }
        let (bytes, padding) = bytes.split_at(ENC_GROUP_SIZE - pad_size);
        key.write(bytes).unwrap();
        if padding.iter().any(|x| *x != 0) {
            return Err(Error::KeyPadding);
        }
        return Ok((key, read));
    }
    Err(Error::KeyLength)
}

#[cfg(test)]
mod tests {
    use super::{encode_bytes, decode_bytes, max_encoded_bytes_size, ENC_GROUP_SIZE};
    use std::cmp::Ordering;

    #[test]
    fn test_enc_dec_bytes() {
        let pairs = vec![(vec![], vec![0, 0, 0, 0, 0, 0, 0, 0, 247]),
                         (vec![1, 2, 3], vec![1, 2, 3, 0, 0, 0, 0, 0, 250]),
                         (vec![0], vec![0, 0, 0, 0, 0, 0, 0, 0, 248]),
                         (vec![1, 2, 3], vec![1, 2, 3, 0, 0, 0, 0, 0, 250]),
                         (vec![1, 2, 3, 0], vec![1, 2, 3, 0, 0, 0, 0, 0, 251]),
                         (vec![1, 2, 3, 4, 5, 6, 7], vec![1, 2, 3, 4, 5, 6, 7, 0, 254]),

                         (vec![0, 0, 0, 0, 0, 0, 0, 0],
                          vec![0, 0, 0, 0, 0, 0, 0, 0, 255, 0, 0, 0, 0, 0, 0, 0, 0, 247]),

                         (vec![1, 2, 3, 4, 5, 6, 7, 8],
                          vec![1, 2, 3, 4, 5, 6, 7, 8, 255, 0, 0, 0, 0, 0, 0, 0, 0, 247]),

                         (vec![1, 2, 3, 4, 5, 6, 7, 8, 9],
                          vec![1, 2, 3, 4, 5, 6, 7, 8, 255, 9, 0, 0, 0, 0, 0, 0, 0, 248])];

        for (x, y) in pairs {
            assert_eq!(encode_bytes(&x), y);
            let (key, size) = decode_bytes(&y).unwrap();
            assert_eq!(key, x);
            assert_eq!(size, y.len());
        }
    }

    #[test]
    fn test_dec_bytes_fail() {
        let invalid_bytes = vec![vec![1, 2, 3, 4],
                                 vec![0, 0, 0, 0, 0, 0, 0, 247],
                                 vec![0, 0, 0, 0, 0, 0, 0, 0, 246],
                                 vec![0, 0, 0, 0, 0, 0, 0, 1, 247],
                                 vec![1, 2, 3, 4, 5, 6, 7, 8, 0],
                                 vec![1, 2, 3, 4, 5, 6, 7, 8, 255, 1],
                                 vec![1, 2, 3, 4, 5, 6, 7, 8, 255, 1, 2, 3, 4, 5, 6, 7, 8],
                                 vec![1, 2, 3, 4, 5, 6, 7, 8, 255, 1, 2, 3, 4, 5, 6, 7, 8, 255],
                                 vec![1, 2, 3, 4, 5, 6, 7, 8, 255, 1, 2, 3, 4, 5, 6, 7, 8, 0]];

        for x in invalid_bytes {
            assert!(decode_bytes(&x).is_err());
        }
    }

    #[test]
    fn test_encode_bytes_compare() {
        let pairs: Vec<(&[u8], &[u8], _)> = vec![(b"", b"\x00", Ordering::Less),
                                                 (b"\x00", b"\x00", Ordering::Equal),
                                                 (b"\xFF", b"\x00", Ordering::Greater),
                                                 (b"\xFF", b"\xFF\x00", Ordering::Less),
                                                 (b"a", b"b", Ordering::Less),
                                                 (b"a", b"\x00", Ordering::Greater),
                                                 (b"\x00", b"\x01", Ordering::Less),
                                                 (b"\x00\x01", b"\x00\x00", Ordering::Greater),
                                                 (b"\x00\x00\x00", b"\x00\x00", Ordering::Greater),
                                                 (b"\x00\x00\x00", b"\x00\x00", Ordering::Greater),

                                                 (b"\x00\x00\x00\x00\x00\x00\x00\x00",
                                                  b"\x00\x00\x00\x00\x00\x00\x00\x00\x00",
                                                  Ordering::Less),

                                                 (b"\x01\x02\x03\x00",
                                                  b"\x01\x02\x03",
                                                  Ordering::Greater),
                                                 (b"\x01\x03\x03\x04",
                                                  b"\x01\x03\x03\x05",
                                                  Ordering::Less),

                                                 (b"\x01\x02\x03\x04\x05\x06\x07",
                                                  b"\x01\x02\x03\x04\x05\x06\x07\x08",
                                                  Ordering::Less),

                                                 (b"\x01\x02\x03\x04\x05\x06\x07\x08\x09",
                                                  b"\x01\x02\x03\x04\x05\x06\x07\x08",
                                                  Ordering::Greater),

                                                 (b"\x01\x02\x03\x04\x05\x06\x07\x08\x00",
                                                  b"\x01\x02\x03\x04\x05\x06\x07\x08",
                                                  Ordering::Greater)];

        for (x, y, ord) in pairs {
            assert_eq!(encode_bytes(x).cmp(&encode_bytes(y)), ord);
        }
    }

    #[test]
    fn test_max_encoded_bytes_size() {
        let n = ENC_GROUP_SIZE;
        let tbl: Vec<(usize, usize)> = vec![(0, n + 1), (n /2, n + 1), (n, 2 * (n + 1))];
        for (x, y) in tbl {
            assert_eq!(max_encoded_bytes_size(x), y);
        }
    }

    use test::Bencher;

    #[bench]
    fn bench_encode(b: &mut Bencher) {
        let key = [b'x'; 20];
        b.iter(|| encode_bytes(&key));
    }

    #[bench]
    fn bench_decode(b: &mut Bencher) {
        let key = [b'x'; 20];
        let encoded = encode_bytes(&key);
        b.iter(|| decode_bytes(&encoded));
    }
}
<commit_msg>codec: format code.<commit_after>use std::vec::Vec;
use std::io::Write;

use super::{Result, Error};

const ENC_GROUP_SIZE: usize = 8;
const ENC_MARKER: u8 = b'\xff';
const ENC_PADDING: [u8; ENC_GROUP_SIZE] = [0; ENC_GROUP_SIZE];

// returns the maximum encoded bytes size.
pub fn max_encoded_bytes_size(n: usize) -> usize {
    (n / ENC_GROUP_SIZE + 1) * (ENC_GROUP_SIZE + 1)
}

// Refer: https://github.com/facebook/mysql-5.6/wiki/MyRocks-record-format#memcomparable-format
pub fn encode_bytes(key: &[u8]) -> Vec<u8> {
    let cap = max_encoded_bytes_size(key.len());
    let mut encoded = Vec::<u8>::with_capacity(cap);
    let len = key.len();
    let mut index = 0;
    while index <= len {
        let remain = len - index;
        let mut pad: usize = 0;
        if remain > ENC_GROUP_SIZE {
            encoded.write(&key[index..index + ENC_GROUP_SIZE]).unwrap();
        } else {
            pad = ENC_GROUP_SIZE - remain;
            encoded.write(&key[index..]).unwrap();
            encoded.write(&ENC_PADDING[..pad]).unwrap();
        }
        encoded.push(ENC_MARKER - (pad as u8));
        index += ENC_GROUP_SIZE;
    }
    encoded
}

pub fn decode_bytes(data: &[u8]) -> Result<(Vec<u8>, usize)> {
    let mut key = Vec::<u8>::with_capacity(data.len());
    let mut read: usize = 0;
    for chunk in data.chunks(ENC_GROUP_SIZE + 1) {
        if chunk.len() != ENC_GROUP_SIZE + 1 {
            return Err(Error::KeyLength);
        }
        read += ENC_GROUP_SIZE + 1;

        let (marker, bytes) = chunk.split_last().unwrap();
        let pad_size = (ENC_MARKER - *marker) as usize;
        if pad_size == 0 {
            key.write(bytes).unwrap();
            continue;
        }
        if pad_size > ENC_GROUP_SIZE {
            return Err(Error::KeyPadding);
        }
        let (bytes, padding) = bytes.split_at(ENC_GROUP_SIZE - pad_size);
        key.write(bytes).unwrap();
        if padding.iter().any(|x| *x != 0) {
            return Err(Error::KeyPadding);
        }
        return Ok((key, read));
    }
    Err(Error::KeyLength)
}

#[cfg(test)]
mod tests {
    use super::{encode_bytes, decode_bytes, max_encoded_bytes_size, ENC_GROUP_SIZE};
    use std::cmp::Ordering;

    #[test]
    fn test_enc_dec_bytes() {
        let pairs = vec![(vec![], vec![0, 0, 0, 0, 0, 0, 0, 0, 247]),
                         (vec![1, 2, 3], vec![1, 2, 3, 0, 0, 0, 0, 0, 250]),
                         (vec![0], vec![0, 0, 0, 0, 0, 0, 0, 0, 248]),
                         (vec![1, 2, 3], vec![1, 2, 3, 0, 0, 0, 0, 0, 250]),
                         (vec![1, 2, 3, 0], vec![1, 2, 3, 0, 0, 0, 0, 0, 251]),
                         (vec![1, 2, 3, 4, 5, 6, 7], vec![1, 2, 3, 4, 5, 6, 7, 0, 254]),

                         (vec![0, 0, 0, 0, 0, 0, 0, 0],
                          vec![0, 0, 0, 0, 0, 0, 0, 0, 255, 0, 0, 0, 0, 0, 0, 0, 0, 247]),

                         (vec![1, 2, 3, 4, 5, 6, 7, 8],
                          vec![1, 2, 3, 4, 5, 6, 7, 8, 255, 0, 0, 0, 0, 0, 0, 0, 0, 247]),

                         (vec![1, 2, 3, 4, 5, 6, 7, 8, 9],
                          vec![1, 2, 3, 4, 5, 6, 7, 8, 255, 9, 0, 0, 0, 0, 0, 0, 0, 248])];

        for (x, y) in pairs {
            assert_eq!(encode_bytes(&x), y);
            let (key, size) = decode_bytes(&y).unwrap();
            assert_eq!(key, x);
            assert_eq!(size, y.len());
        }
    }

    #[test]
    fn test_dec_bytes_fail() {
        let invalid_bytes = vec![vec![1, 2, 3, 4],
                                 vec![0, 0, 0, 0, 0, 0, 0, 247],
                                 vec![0, 0, 0, 0, 0, 0, 0, 0, 246],
                                 vec![0, 0, 0, 0, 0, 0, 0, 1, 247],
                                 vec![1, 2, 3, 4, 5, 6, 7, 8, 0],
                                 vec![1, 2, 3, 4, 5, 6, 7, 8, 255, 1],
                                 vec![1, 2, 3, 4, 5, 6, 7, 8, 255, 1, 2, 3, 4, 5, 6, 7, 8],
                                 vec![1, 2, 3, 4, 5, 6, 7, 8, 255, 1, 2, 3, 4, 5, 6, 7, 8, 255],
                                 vec![1, 2, 3, 4, 5, 6, 7, 8, 255, 1, 2, 3, 4, 5, 6, 7, 8, 0]];

        for x in invalid_bytes {
            assert!(decode_bytes(&x).is_err());
        }
    }

    #[test]
    fn test_encode_bytes_compare() {
        let pairs: Vec<(&[u8], &[u8], _)> = vec![(b"", b"\x00", Ordering::Less),
                                                 (b"\x00", b"\x00", Ordering::Equal),
                                                 (b"\xFF", b"\x00", Ordering::Greater),
                                                 (b"\xFF", b"\xFF\x00", Ordering::Less),
                                                 (b"a", b"b", Ordering::Less),
                                                 (b"a", b"\x00", Ordering::Greater),
                                                 (b"\x00", b"\x01", Ordering::Less),
                                                 (b"\x00\x01", b"\x00\x00", Ordering::Greater),
                                                 (b"\x00\x00\x00", b"\x00\x00", Ordering::Greater),
                                                 (b"\x00\x00\x00", b"\x00\x00", Ordering::Greater),

                                                 (b"\x00\x00\x00\x00\x00\x00\x00\x00",
                                                  b"\x00\x00\x00\x00\x00\x00\x00\x00\x00",
                                                  Ordering::Less),

                                                 (b"\x01\x02\x03\x00",
                                                  b"\x01\x02\x03",
                                                  Ordering::Greater),
                                                 (b"\x01\x03\x03\x04",
                                                  b"\x01\x03\x03\x05",
                                                  Ordering::Less),

                                                 (b"\x01\x02\x03\x04\x05\x06\x07",
                                                  b"\x01\x02\x03\x04\x05\x06\x07\x08",
                                                  Ordering::Less),

                                                 (b"\x01\x02\x03\x04\x05\x06\x07\x08\x09",
                                                  b"\x01\x02\x03\x04\x05\x06\x07\x08",
                                                  Ordering::Greater),

                                                 (b"\x01\x02\x03\x04\x05\x06\x07\x08\x00",
                                                  b"\x01\x02\x03\x04\x05\x06\x07\x08",
                                                  Ordering::Greater)];

        for (x, y, ord) in pairs {
            assert_eq!(encode_bytes(x).cmp(&encode_bytes(y)), ord);
        }
    }

    #[test]
    fn test_max_encoded_bytes_size() {
        let n = ENC_GROUP_SIZE;
        let tbl: Vec<(usize, usize)> = vec![(0, n + 1), (n / 2, n + 1), (n, 2 * (n + 1))];
        for (x, y) in tbl {
            assert_eq!(max_encoded_bytes_size(x), y);
        }
    }

    use test::Bencher;

    #[bench]
    fn bench_encode(b: &mut Bencher) {
        let key = [b'x'; 20];
        b.iter(|| encode_bytes(&key));
    }

    #[bench]
    fn bench_decode(b: &mut Bencher) {
        let key = [b'x'; 20];
        let encoded = encode_bytes(&key);
        b.iter(|| decode_bytes(&encoded));
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Implementation for the Solaris family
//!
//! Read from `/dev/random`, with chunks of limited size (1040 bytes).
//! `/dev/random` uses the Hash_DRBG with SHA512 algorithm from NIST SP 800-90A.
//! `/dev/urandom` uses the FIPS 186-2 algorithm, which is considered less
//! secure. We choose to read from `/dev/random`.
//!
//! Since Solaris 11.3 the `getrandom` syscall is available. To make sure we can
//! compile on both Solaris and on OpenSolaris derivatives, that do not have the
//! function, we do a direct syscall instead of calling a library function.
//!
//! We have no way to differentiate between Solaris, illumos, SmartOS, etc.
extern crate libc;

use rand_core::{Error, ErrorKind};
use super::random_device;
use super::OsRngImpl;

use std::io;
use std::io::Read;
use std::fs::{File, OpenOptions};
use std::os::unix::fs::OpenOptionsExt;
use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering, AtomicUsize};
use std::cmp;
use std::mem;

#[derive(Clone, Debug)]
pub struct OsRng {
    method: OsRngMethod,
    initialized: bool,
}

#[derive(Clone, Debug)]
enum OsRngMethod {
    GetRandom,
    RandomDevice,
}

impl OsRngImpl for OsRng {
    fn new() -> Result<OsRng, Error> {
        if is_getrandom_available() {
            return Ok(OsRng { method: OsRngMethod::GetRandom,
                              initialized: false });
        }
        let open = |p| OpenOptions::new()
            .read(true)
            .custom_flags(libc::O_NONBLOCK)
            .open(p);
        random_device::open("/dev/random", &open)?;
        Ok(OsRng { method: OsRngMethod::RandomDevice, initialized: false })
    }

    fn fill_chunk(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        match self.method {
            OsRngMethod::GetRandom => getrandom_try_fill(dest, false),
            OsRngMethod::RandomDevice => random_device::read(dest),
        }
    }

    fn test_initialized(&mut self, dest: &mut [u8], blocking: bool)
        -> Result<usize, Error>
    {
        static OS_RNG_INITIALIZED: AtomicBool = ATOMIC_BOOL_INIT;
        if !self.initialized {
            self.initialized = OS_RNG_INITIALIZED.load(Ordering::Relaxed);
        }
        if self.initialized { return Ok(0); }

        let chunk_len = cmp::min(1024, dest.len());
        let dest = &mut dest[..chunk_len];

        match self.method {
            OsRngMethod::GetRandom => getrandom_try_fill(dest, blocking)?,
            OsRngMethod::RandomDevice => {
                if blocking {
                    info!("OsRng: testing random device /dev/random");
                    // We already have a non-blocking handle, but now need a
                    // blocking one. Not much choice except opening it twice
                    let mut file = File::open("/dev/random")
                        .map_err(random_device::map_err)?;
                    file.read(dest).map_err(random_device::map_err)?;
                } else {
                    self.fill_chunk(dest)?;
                }
            }
        };
        OS_RNG_INITIALIZED.store(true, Ordering::Relaxed);
        self.initialized = true;
        Ok(chunk_len)
    }

    fn max_chunk_size(&self) -> usize {
        // This is the largest size that's guaranteed to not block across
        // all the Solarish platforms, though some may allow for larger
        // sizes.
        256
    }

    fn method_str(&self) -> &'static str {
        match self.method {
            OsRngMethod::GetRandom => "getrandom",
            OsRngMethod::RandomDevice => "/dev/random",
        }
    }
}

#[cfg(target_os = "illumos")]
type GetRandomFn = unsafe extern fn(*mut u8, libc::size_t, libc::c_uint)
    -> libc::ssize_t;
#[cfg(target_os = "solaris")]
type GetRandomFn = unsafe extern fn(*mut u8, libc::size_t, libc::c_uint)
    -> libc::c_int;

// Use dlsym to determine if getrandom(2) is present in libc.  On Solarish
// systems, the syscall interface is not stable and can change between
// updates.  Even worse, issuing unsupported syscalls will cause the system
// to generate a SIGSYS signal (usually terminating the program).
// Instead the stable APIs are exposed via libc.  Cache the result of the
// lookup for future calls.  This is loosely modeled after the
// libstd::sys::unix::weak macro which unfortunately is not exported.
fn fetch() -> Option<&'static GetRandomFn> {
    static FPTR: AtomicUsize = AtomicUsize::new(1);

    unsafe {
       if FPTR.load(Ordering::SeqCst) == 1 {
            let name = "getrandom\0";
            let addr = libc::dlsym(libc::RTLD_DEFAULT,
                                   name.as_ptr() as *const _) as usize;
            FPTR.store(addr, Ordering::SeqCst);
       }

       if FPTR.load(Ordering::SeqCst) == 0 {
           return None;
       } else {
           mem::transmute::<&AtomicUsize, Option<&GetRandomFn>>(&FPTR)
       }
   }
}

fn getrandom(buf: &mut [u8], blocking: bool) -> libc::ssize_t {
    const GRND_NONBLOCK: libc::c_uint = 0x0001;
    const GRND_RANDOM: libc::c_uint = 0x0002;

    if let Some(rand) = fetch() {
        unsafe {
            rand(buf.as_mut_ptr(), buf.len(),
                 if blocking { 0 } else { GRND_NONBLOCK } | GRND_RANDOM) as libc::ssize_t
        }
    } else {
        -1
    }
}

fn getrandom_try_fill(dest: &mut [u8], blocking: bool) -> Result<(), Error> {
    let result = getrandom(dest, blocking);
    if result == -1 || result == 0 {
        let err = io::Error::last_os_error();
        let kind = err.kind();
        if kind == io::ErrorKind::WouldBlock {
            return Err(Error::with_cause(
                ErrorKind::NotReady,
                "getrandom not ready",
                err,
            ));
        } else {
            return Err(Error::with_cause(
                ErrorKind::Unavailable,
                "unexpected getrandom error",
                err,
            ));
        }
    } else if result != dest.len() as libc::ssize_t {
        return Err(Error::new(ErrorKind::Unavailable,
                              "unexpected getrandom error"));
    }
    Ok(())
}

fn is_getrandom_available() -> bool {
    let available = match fetch() {
        Some(_) => true,
        None => false,
    };
    info!("OsRng: using {}", if available { "getrandom" } else { "/dev/random" });
    available
}
<commit_msg>Minimize unsafe regions<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Implementation for the Solaris family
//!
//! Read from `/dev/random`, with chunks of limited size (1040 bytes).
//! `/dev/random` uses the Hash_DRBG with SHA512 algorithm from NIST SP 800-90A.
//! `/dev/urandom` uses the FIPS 186-2 algorithm, which is considered less
//! secure. We choose to read from `/dev/random`.
//!
//! Since Solaris 11.3 the `getrandom` syscall is available. To make sure we can
//! compile on both Solaris and on OpenSolaris derivatives, that do not have the
//! function, we do a direct syscall instead of calling a library function.
//!
//! We have no way to differentiate between Solaris, illumos, SmartOS, etc.
extern crate libc;

use rand_core::{Error, ErrorKind};
use super::random_device;
use super::OsRngImpl;

use std::io;
use std::io::Read;
use std::fs::{File, OpenOptions};
use std::os::unix::fs::OpenOptionsExt;
use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering, AtomicUsize};
use std::cmp;
use std::mem;

#[derive(Clone, Debug)]
pub struct OsRng {
    method: OsRngMethod,
    initialized: bool,
}

#[derive(Clone, Debug)]
enum OsRngMethod {
    GetRandom,
    RandomDevice,
}

impl OsRngImpl for OsRng {
    fn new() -> Result<OsRng, Error> {
        if is_getrandom_available() {
            return Ok(OsRng { method: OsRngMethod::GetRandom,
                              initialized: false });
        }
        let open = |p| OpenOptions::new()
            .read(true)
            .custom_flags(libc::O_NONBLOCK)
            .open(p);
        random_device::open("/dev/random", &open)?;
        Ok(OsRng { method: OsRngMethod::RandomDevice, initialized: false })
    }

    fn fill_chunk(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        match self.method {
            OsRngMethod::GetRandom => getrandom_try_fill(dest, false),
            OsRngMethod::RandomDevice => random_device::read(dest),
        }
    }

    fn test_initialized(&mut self, dest: &mut [u8], blocking: bool)
        -> Result<usize, Error>
    {
        static OS_RNG_INITIALIZED: AtomicBool = ATOMIC_BOOL_INIT;
        if !self.initialized {
            self.initialized = OS_RNG_INITIALIZED.load(Ordering::Relaxed);
        }
        if self.initialized { return Ok(0); }

        let chunk_len = cmp::min(1024, dest.len());
        let dest = &mut dest[..chunk_len];

        match self.method {
            OsRngMethod::GetRandom => getrandom_try_fill(dest, blocking)?,
            OsRngMethod::RandomDevice => {
                if blocking {
                    info!("OsRng: testing random device /dev/random");
                    // We already have a non-blocking handle, but now need a
                    // blocking one. Not much choice except opening it twice
                    let mut file = File::open("/dev/random")
                        .map_err(random_device::map_err)?;
                    file.read(dest).map_err(random_device::map_err)?;
                } else {
                    self.fill_chunk(dest)?;
                }
            }
        };
        OS_RNG_INITIALIZED.store(true, Ordering::Relaxed);
        self.initialized = true;
        Ok(chunk_len)
    }

    fn max_chunk_size(&self) -> usize {
        // This is the largest size that's guaranteed to not block across
        // all the Solarish platforms, though some may allow for larger
        // sizes.
        256
    }

    fn method_str(&self) -> &'static str {
        match self.method {
            OsRngMethod::GetRandom => "getrandom",
            OsRngMethod::RandomDevice => "/dev/random",
        }
    }
}

#[cfg(target_os = "illumos")]
type GetRandomFn = unsafe extern fn(*mut u8, libc::size_t, libc::c_uint)
    -> libc::ssize_t;
#[cfg(target_os = "solaris")]
type GetRandomFn = unsafe extern fn(*mut u8, libc::size_t, libc::c_uint)
    -> libc::c_int;

// Use dlsym to determine if getrandom(2) is present in libc.  On Solarish
// systems, the syscall interface is not stable and can change between
// updates.  Even worse, issuing unsupported syscalls will cause the system
// to generate a SIGSYS signal (usually terminating the program).
// Instead the stable APIs are exposed via libc.  Cache the result of the
// lookup for future calls.  This is loosely modeled after the
// libstd::sys::unix::weak macro which unfortunately is not exported.
fn fetch() -> Option<GetRandomFn> {
    static FPTR: AtomicUsize = AtomicUsize::new(1);

    if FPTR.load(Ordering::SeqCst) == 1 {
        let name = "getrandom\0";
        let addr = unsafe {
            libc::dlsym(libc::RTLD_DEFAULT, name.as_ptr() as *const _) as usize
        };
        FPTR.store(addr, Ordering::SeqCst);
    }

    let ptr = FPTR.load(Ordering::SeqCst);
    unsafe {
        mem::transmute::<usize, Option<GetRandomFn>>(ptr)
    }
}

fn getrandom(buf: &mut [u8], blocking: bool) -> libc::ssize_t {
    const GRND_NONBLOCK: libc::c_uint = 0x0001;
    const GRND_RANDOM: libc::c_uint = 0x0002;

    if let Some(rand) = fetch() {
        let flag = if blocking { 0 } else { GRND_NONBLOCK } | GRND_RANDOM;
        unsafe {
            rand(buf.as_mut_ptr(), buf.len(), flag) as libc::ssize_t
        }
    } else {
        -1
    }
}

fn getrandom_try_fill(dest: &mut [u8], blocking: bool) -> Result<(), Error> {
    let result = getrandom(dest, blocking);
    if result == -1 || result == 0 {
        let err = io::Error::last_os_error();
        let kind = err.kind();
        if kind == io::ErrorKind::WouldBlock {
            return Err(Error::with_cause(
                ErrorKind::NotReady,
                "getrandom not ready",
                err,
            ));
        } else {
            return Err(Error::with_cause(
                ErrorKind::Unavailable,
                "unexpected getrandom error",
                err,
            ));
        }
    } else if result != dest.len() as libc::ssize_t {
        return Err(Error::new(ErrorKind::Unavailable,
                              "unexpected getrandom error"));
    }
    Ok(())
}

fn is_getrandom_available() -> bool {
    let available = match fetch() {
        Some(_) => true,
        None => false,
    };
    info!("OsRng: using {}", if available { "getrandom" } else { "/dev/random" });
    available
}
<|endoftext|>


<commit_before>extern crate rand;
extern crate time;
extern crate timely;
extern crate differential_dataflow;

use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::progress::timestamp::RootTimestamp;

use rand::{Rng, SeedableRng, StdRng};

use differential_dataflow::operators::*;
use differential_dataflow::collection::LeastUpperBound;

type Node = u32;
type Edge = (Node, Node);

fn main() {

    let nodes: u32 = std::env::args().nth(1).unwrap().parse().unwrap();
    let edges: u32 = std::env::args().nth(2).unwrap().parse().unwrap();
    let batch: u32 = std::env::args().nth(3).unwrap().parse().unwrap();

    // define a new computational scope, in which to run BFS
    timely::execute_from_args(std::env::args().skip(4), move |computation| {
        let start = time::precise_time_s();

        // define BFS dataflow; return handles to roots and edges inputs
        let (mut graph, probe) = computation.scoped(|scope| {

            let roots = vec![(0,1)].into_iter().to_stream(scope);
            let (edge_input, graph) = scope.new_input();
            let probe = bfs(&graph, &roots).probe().0;

            (edge_input, probe)
        });

        let seed: &[_] = &[1, 2, 3, 4];
        let mut rng1: StdRng = SeedableRng::from_seed(seed);    // rng for edge additions
        let mut rng2: StdRng = SeedableRng::from_seed(seed);    // rng for edge deletions

        println!("performing BFS on {} nodes, {} edges:", nodes, edges);

        if computation.index() == 0 {
            // trickle edges in to dataflow
            for _ in 0..(edges/1000) {
                for _ in 0..1000 {
                    graph.send(((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)), 1));
                }
                computation.step();
            }
        }

        println!("loaded; elapsed: {}s", time::precise_time_s() - start);

        graph.advance_to(1);
        while probe.le(&RootTimestamp::new(0)) { computation.step(); }

        println!("stable; elapsed: {}s", time::precise_time_s() - start);

        let mut changes = Vec::new();
        for wave in 0.. {
            for _ in 0..batch {
                changes.push(((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)), 1));
                changes.push(((rng2.gen_range(0, nodes), rng2.gen_range(0, nodes)),-1));
            }

            let start = time::precise_time_s();
            let round = *graph.epoch();
            if computation.index() == 0 {
                while let Some(change) = changes.pop() {
                    graph.send(change);
                }
            }
            graph.advance_to(round + 1);

            while probe.le(&RootTimestamp::new(round)) { computation.step(); }

            if computation.index() == 0 {
                println!("wave {}: avg {}", wave, (time::precise_time_s() - start) / (batch as f64));
            }
        }
    });
}

// returns pairs (n, s) indicating node n can be reached from a root in s steps.
fn bfs<G: Scope>(edges: &Stream<G, (Edge, i32)>, roots: &Stream<G, (Node, i32)>) -> Stream<G, ((Node, u32), i32)>
where G::Timestamp: LeastUpperBound {

    // initialize roots as reaching themselves at distance 0
    let nodes = roots.map(|(x,w)| ((x, 0), w));
    // let edges = edges.map_in_place(|x| x.0 = ((x.0).1, (x.0).0))
    //                  .concat(&edges);

    // repeatedly update minimal distances each node can be reached from each root
    nodes.iterate(|inner| {

        let edges = inner.scope().enter(&edges);
        let nodes = inner.scope().enter(&nodes);

        inner.join_map(&edges, |_k,l,d| (*d, l+1))
             .concat(&nodes)
             .group(|_, s, t| t.push((*s.peek().unwrap().0, 1)))
     })
}
<commit_msg>batch size of zero means no batches<commit_after>extern crate rand;
extern crate time;
extern crate timely;
extern crate differential_dataflow;

use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::progress::timestamp::RootTimestamp;

use rand::{Rng, SeedableRng, StdRng};

use differential_dataflow::operators::*;
use differential_dataflow::collection::LeastUpperBound;

type Node = u32;
type Edge = (Node, Node);

fn main() {

    let nodes: u32 = std::env::args().nth(1).unwrap().parse().unwrap();
    let edges: u32 = std::env::args().nth(2).unwrap().parse().unwrap();
    let batch: u32 = std::env::args().nth(3).unwrap().parse().unwrap();

    // define a new computational scope, in which to run BFS
    timely::execute_from_args(std::env::args().skip(4), move |computation| {
        let start = time::precise_time_s();

        // define BFS dataflow; return handles to roots and edges inputs
        let (mut graph, probe) = computation.scoped(|scope| {

            let roots = vec![(0,1)].into_iter().to_stream(scope);
            let (edge_input, graph) = scope.new_input();
            let probe = bfs(&graph, &roots).probe().0;

            (edge_input, probe)
        });

        let seed: &[_] = &[1, 2, 3, 4];
        let mut rng1: StdRng = SeedableRng::from_seed(seed);    // rng for edge additions
        let mut rng2: StdRng = SeedableRng::from_seed(seed);    // rng for edge deletions

        println!("performing BFS on {} nodes, {} edges:", nodes, edges);

        if computation.index() == 0 {
            // trickle edges in to dataflow
            for _ in 0..(edges/1000) {
                for _ in 0..1000 {
                    graph.send(((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)), 1));
                }
                computation.step();
            }
        }

        println!("loaded; elapsed: {}s", time::precise_time_s() - start);

        graph.advance_to(1);
        while probe.le(&RootTimestamp::new(0)) { computation.step(); }

        println!("stable; elapsed: {}s", time::precise_time_s() - start);

        if batch > 0 {
            let mut changes = Vec::new();
            for wave in 0.. {
                for _ in 0..batch {
                    changes.push(((rng1.gen_range(0, nodes), rng1.gen_range(0, nodes)), 1));
                    changes.push(((rng2.gen_range(0, nodes), rng2.gen_range(0, nodes)),-1));
                }

                let start = time::precise_time_s();
                let round = *graph.epoch();
                if computation.index() == 0 {
                    while let Some(change) = changes.pop() {
                        graph.send(change);
                    }
                }
                graph.advance_to(round + 1);

                while probe.le(&RootTimestamp::new(round)) { computation.step(); }

                if computation.index() == 0 {
                    println!("wave {}: avg {}", wave, (time::precise_time_s() - start) / (batch as f64));
                }
            }
        }
    });
}

// returns pairs (n, s) indicating node n can be reached from a root in s steps.
fn bfs<G: Scope>(edges: &Stream<G, (Edge, i32)>, roots: &Stream<G, (Node, i32)>) -> Stream<G, ((Node, u32), i32)>
where G::Timestamp: LeastUpperBound {

    // initialize roots as reaching themselves at distance 0
    let nodes = roots.map(|(x,w)| ((x, 0), w));
    // let edges = edges.map_in_place(|x| x.0 = ((x.0).1, (x.0).0))
    //                  .concat(&edges);

    // repeatedly update minimal distances each node can be reached from each root
    nodes.iterate(|inner| {

        let edges = inner.scope().enter(&edges);
        let nodes = inner.scope().enter(&nodes);

        inner.join_map(&edges, |_k,l,d| (*d, l+1))
             .concat(&nodes)
             .group(|_, s, t| t.push((*s.peek().unwrap().0, 1)))
     })
}
<|endoftext|>


<commit_before>use std::collections::VecDeque;

fn main() {
    let mut dst = VecDeque::new();
    dst.push_front(Box::new(1));
    dst.push_front(Box::new(2));
    dst.pop_back();

    let mut src = VecDeque::new();
    src.push_front(Box::new(2));
    dst.append(&mut src);
    for a in dst {
      assert_eq!(*a, 2);
    }
}
<commit_msg>disable VecDeque test until the fix lands in rustc<commit_after>// FIXME: Validation disabled until https://github.com/rust-lang/rust/pull/56161 lands
// compile-flags: -Zmiri-disable-validation

use std::collections::VecDeque;

fn main() {
    let mut dst = VecDeque::new();
    dst.push_front(Box::new(1));
    dst.push_front(Box::new(2));
    dst.pop_back();

    let mut src = VecDeque::new();
    src.push_front(Box::new(2));
    dst.append(&mut src);
    for a in dst {
      assert_eq!(*a, 2);
    }
}
<|endoftext|>


<commit_before>use hyper::Client;
use super::helper::KinesisHelper;
use utils::screen::Screen;
use utils::screen::Status;

use rusoto::{
    DefaultCredentialsProvider,
    Region
};

enum State {
    Root,
    End,
    StreamList(Vec<String>),
    ShardList(String, Vec<String>),
    RecordList(String),  // iterator_id, records
    Record(String, String), // itrator_id, record
}

pub fn run(credential_provider: DefaultCredentialsProvider, client: Client, region: Region) {

    let kinesis_helper = KinesisHelper::new(client, credential_provider, region);
    let mut screen: Screen = Screen::new();
    let mut state: State = State::Root;

    loop {

        state = match state {

            State::Root => {

                let commands = vec!["l".to_string(), "q".to_string()];
                screen.update_screen("🍓  Kinesis", &commands);

                match screen.select_line() {
                    Status::Error | Status::Quit => State::End,
                    Status::Selected(ref c) if c == "l" => {
                        match kinesis_helper.list_streams() {
                            Ok(streams) => State::StreamList(streams),
                            Err(e) =>  State::Root
                        }
                    }
                    _ => State::Root
                }
            },
            State::StreamList(streams) => {

                screen.update_screen("🍓  Kinesis > Streams", &streams);

                match screen.select_line() {
                    Status::Error | Status::Quit => State::End,
                    Status::Selected(stream_name) => {
                        match kinesis_helper.describe_shards(&stream_name) {
                            Ok(shards) => State::ShardList(stream_name.to_string(), shards),
                            Err(e) => State::Root,
                        }
                    }
                    _ => State::Root
                }
            },
            State::ShardList(stream_name, shards) => {

                screen.update_screen("🍓  Kinesis > Streams > Shards", &shards);

                match screen.select_line() {
                    Status::Error | Status::Quit => State::End,
                    Status::Selected(shard_id) => {
                        match kinesis_helper.get_shard_iterator(&stream_name, &shard_id) {
                            Ok(shard_iterator) => State::RecordList(shard_iterator),
                            Err(e) => State::Root,
                        }
                    }
                    _ => State::Root
                }

            },
            State::RecordList(shard_iterator) => {

                match kinesis_helper.get_records(&shard_iterator) {

                    Ok((results, iterator)) => {

                        screen.update_screen("🍓  Kinesis > Streams > Shards > Records",
                                             &kinesis_helper.decode_records(&results));

                        match screen.select_line() {
                            Status::Error | Status::Quit => State::End,
                            Status::Selected(ref c) if c == "n" => {
                                match iterator {
                                    Some(iterator) => State::RecordList(iterator),
                                    None => State::Root
                                }
                            },
                            Status::Selected(ref record) => State::Record(shard_iterator, record.to_string()),
                            _ => State::Root
                        }
                    },
                    Err(e) => State::Root,
                }
            },
            State::Record(shard_iterator, record) => {

                screen.update_screen("🍓  Kinesis > Streams > Shards > Records > Record", &vec!(String::new()));

                match screen.render(&kinesis_helper.format_record(&record)) {
                    Status::Error | Status::Quit => State::End,
                    Status::Escaped  => State::RecordList(shard_iterator),
                    _ => State::Root
                }
            }
            State::End => {
                break;
            }
        };

    }
}
<commit_msg>Refactoring<commit_after>use hyper::Client;
use super::helper::KinesisHelper;
use utils::screen::Screen;
use utils::screen::Status;

use rusoto::{
    DefaultCredentialsProvider,
    Region
};

enum State {
    Root,
    End,
    StreamList(Vec<String>),
    ShardList(String, Vec<String>),
    RecordList(String),  // iterator_id, records
    Record(String, String), // itrator_id, record
}

pub fn run(credential_provider: DefaultCredentialsProvider, client: Client, region: Region) {

    let kinesis_helper = KinesisHelper::new(client, credential_provider, region);
    let mut screen: Screen = Screen::new();
    let mut state: State = State::Root;

    loop {

        state = match state {

            State::Root => root(screen, kinesis_helper),
            State::StreamList(streams) => {

                screen.update_screen("🍓  Kinesis > Streams", &streams);

                match screen.select_line() {
                    Status::Error | Status::Quit => State::End,
                    Status::Selected(stream_name) => {
                        match kinesis_helper.describe_shards(&stream_name) {
                            Ok(shards) => State::ShardList(stream_name.to_string(), shards),
                            Err(e) => State::Root,
                        }
                    }
                    _ => State::Root
                }
            },
            State::ShardList(stream_name, shards) => {

                screen.update_screen("🍓  Kinesis > Streams > Shards", &shards);

                match screen.select_line() {
                    Status::Error | Status::Quit => State::End,
                    Status::Selected(shard_id) => {
                        match kinesis_helper.get_shard_iterator(&stream_name, &shard_id) {
                            Ok(shard_iterator) => State::RecordList(shard_iterator),
                            Err(e) => State::Root,
                        }
                    }
                    _ => State::Root
                }

            },
            State::RecordList(shard_iterator) => {

                match kinesis_helper.get_records(&shard_iterator) {

                    Ok((results, iterator)) => {

                        screen.update_screen("🍓  Kinesis > Streams > Shards > Records",
                                             &kinesis_helper.decode_records(&results));

                        match screen.select_line() {
                            Status::Error | Status::Quit => State::End,
                            Status::Selected(ref c) if c == "n" => {
                                match iterator {
                                    Some(iterator) => State::RecordList(iterator),
                                    None => State::Root
                                }
                            },
                            Status::Selected(ref record) => State::Record(shard_iterator, record.to_string()),
                            _ => State::Root
                        }
                    },
                    Err(e) => State::Root,
                }
            },
            State::Record(shard_iterator, record) => {

                screen.update_screen("🍓  Kinesis > Streams > Shards > Records > Record", &vec!(String::new()));

                match screen.render(&kinesis_helper.format_record(&record)) {
                    Status::Error | Status::Quit => State::End,
                    Status::Escaped  => State::RecordList(shard_iterator),
                    _ => State::Root
                }
            }
            State::End => {
                break;
            }
        };

    }
}

fn root(screen: &mut Screen, kinesis_helper: &KinesisHelper) -> State {

    let commands = vec!["l".to_string(), "q".to_string()];
    screen.update_screen("🍓  Kinesis", &commands);

    match screen.select_line() {
        Status::Error | Status::Quit => State::End,
        Status::Selected(ref c) if c == "l" => {
            match kinesis_helper.list_streams() {
                Ok(streams) => State::StreamList(streams),
                Err(e) =>  State::Root
            }
        }
        _ => State::Root
    }
}
<|endoftext|>