<commit_before>use crate::ram::chunk::Chunk;
use crate::ram::entry::EntryContent;
use crate::ram::segs::{Segment, SEGMENT_SIZE};
use itertools::Itertools;
use lightning::linked_map::NodeRef as MapNodeRef;
use rayon::prelude::*;
use std::collections::HashSet;
use std::sync::atomic::Ordering;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use std::sync::Arc;

use libc;

pub struct CombinedCleaner;

#[derive(Clone)]
struct DummyEntry {
    size: usize,
    addr: usize,
    cell_hash: Option<u64>,
    timestamp: u32,
}

struct DummySegment {
    head: usize,
    entries: Vec<DummyEntry>,
}

impl DummySegment {
    fn new() -> DummySegment {
        DummySegment {
            head: 0,
            entries: Vec::new(),
        }
    }
}

// Combine small segments into larger segments
// This cleaner will perform a greedy approach to relocate entries from old segments to fewer
// new segments to reduce number or segments and reclaim spaces from tombstones

// for higher hit rate for fetching cells in segments, we need to put data with close timestamp together
// this optimization is intended for enabling neb to contain data more than it's memory

impl CombinedCleaner {
    pub fn combine_segments(chunk: &Chunk, segments: &Vec<MapNodeRef<Segment>>) -> usize {
        if segments.len() < 2 {
            trace!(
                "too few segments to combine, chunk {}, segments {}",
                chunk.id,
                segments.len()
            );
            return 0;
        }
        debug!("Combining segments");

        let space_to_collect = segments
            .iter()
            .map(|seg| seg.used_spaces() as usize)
            .sum::<usize>();
        let segment_ids_to_combine: HashSet<_> = segments.iter().map(|seg| seg.id).collect();

        // Get all entries in segments to combine and order them by data temperature and size
        let nested_entries = segments
            .iter()
            .flat_map(|seg| chunk.live_entries(seg))
            .filter(|entry| {
                // live entries have done a lot of filtering work already
                // but we still need to remove those tombstones that pointed to segments we are about to combine
                if let EntryContent::Tombstone(ref tombstone) = entry.content {
                    return !segment_ids_to_combine.contains(&tombstone.segment_id);
                }
                return true;
            })
            .map(|entry| {
                let entry_size = entry.meta.entry_size;
                let entry_addr = entry.meta.entry_pos;
                let cell_header = match entry.content {
                    EntryContent::Cell(header) => Some(header),
                    _ => None,
                };
                DummyEntry {
                    size: entry_size,
                    addr: entry_addr,
                    timestamp: cell_header.map(|h| h.timestamp).unwrap_or(0),
                    cell_hash: cell_header.map(|h| h.hash),
                }
            })
            .group_by(|entry| entry.timestamp / 10)
            .into_iter()
            .map(|(t, group)| {
                let mut group: Vec<_> = group.collect();
                group.sort_by_key(|entry| entry.size);
                return (t, group.into_iter());
            })
            .sorted_by_key(|&(t, _)| t)
            .into_iter()
            .map(|(_, group)| group);

        let mut entries: Vec<_> = Iterator::flatten(nested_entries)
            // order by temperature and size from greater to lesser
            .rev()
            // provide additional state for whether entry have been claimed on simulation
            .map(|e| (e, false))
            .collect();

        debug!("Found {} entries to combine", entries.len());

        let mut space_cleaned = 0;
        if entries.len() > 0 {
            // Simulate the combine process to determine the efficiency
            let mut pending_segments = Vec::with_capacity(segments.len());
            let entries_num = entries.len();
            let mut entries_to_claim = entries_num;
            let mut cursor = 0;
            pending_segments.push(DummySegment::new());
            while entries_to_claim > 0 {
                let segment_space_remains = SEGMENT_SIZE - pending_segments.last().unwrap().head;
                let index = cursor % entries_num;
                let entry_pair = entries.get_mut(index).unwrap();
                if entry_pair.1 {
                    // entry claimed
                    cursor += 1;
                    continue;
                }

                let entry = &entry_pair.0;
                let entry_size = entry.size;
                if entry_size > segment_space_remains {
                    if index == entries_num - 1 {
                        // iterated to the last one, which means no more entries can be placed
                        // in the segment, then create a new segment
                        pending_segments.push(DummySegment::new());
                    }
                    cursor += 1;
                    continue;
                }
                let last_segment = pending_segments.last_mut().unwrap();
                last_segment.entries.push(entry.clone());

                // pump dummy segment head pointer
                last_segment.head += entry_size;

                // mark entry claimed
                entry_pair.1 = true;
                entries_to_claim -= 1;

                // move to next entry
                cursor += 1;
            }

            debug!("Checking combine feasibility");
            let pending_segments_len = pending_segments.len();
            let segments_to_combine_len = segments.len();
            let cleaned_total_live_space = AtomicUsize::new(0);
            if pending_segments_len >= segments_to_combine_len {
                warn!(
                    "Trying to combine segments but resulting segments still does not go down {}/{}",
                    pending_segments_len, segments_to_combine_len
                );
            }

            debug!(
                "Updating cell reference, pending segments {}",
                pending_segments.len()
            );
            pending_segments
                .par_iter()
                .map(|dummy_seg| {
                    let new_seg = chunk
                        .allocator
                        .alloc_seg(&chunk.backup_storage, &chunk.wal_storage)
                        .expect("No space left during combine");
                    let new_seg_id = new_seg.id;
                    let mut cell_mapping = Vec::with_capacity(dummy_seg.entries.len());
                    let mut seg_cursor = new_seg.addr;
                    trace!(
                        "Combining segment to new one with id {} with {} cells",
                        new_seg_id,
                        dummy_seg.entries.len()
                    );
                    for entry in &dummy_seg.entries {
                        let entry_addr = entry.addr;
                        unsafe {
                            libc::memcpy(
                                seg_cursor as *mut libc::c_void,
                                entry_addr as *mut libc::c_void,
                                entry.size,
                            );
                        }
                        if let Some(cell_hash) = entry.cell_hash {
                            trace!(
                                "Marked cell relocation hash {}, addr {} to segment {}",
                                cell_hash,
                                entry_addr,
                                new_seg_id
                            );
                            cell_mapping.push((seg_cursor, entry_addr, cell_hash));
                        }
                        seg_cursor += entry.size;
                    }
                    new_seg.append_header.store(seg_cursor, Ordering::Relaxed);
                    cleaned_total_live_space.fetch_add(new_seg.used_spaces() as usize, Relaxed);
                    return (new_seg, cell_mapping);
                })
                .for_each(|(segment, cells)| {
                    trace!("Putting new segment {}, cells {}", segment.id, cells.len());
                    segment.archive().unwrap();
                    let new_seg_id = segment.id as usize;
                    chunk.put_segment(segment);
                    let new_seg = chunk.segs.get(&new_seg_id).unwrap();
                    for (new, old, hash) in cells {
                        trace!("Reset cell {} ptr from {} to {}", hash, old, new);
                        #[cfg(feature = "fast_map")]
                        let index = chunk.index.lock(hash as usize);
                        #[cfg(feature = "slow_map")]
                        let index = chunk.index.get_mut(&hash);

                        if let Some(mut actual_addr) = index {
                            if *actual_addr == old {
                                *actual_addr = new;
                                trace!(
                                    "Cell addr for hash {} set from {} to {} for combine",
                                    hash,
                                    old,
                                    new
                                );
                            } else {
                                trace!(
                                    "cell {} with address {}, have been changed to {} on combine",
                                    hash,
                                    old,
                                    *actual_addr
                                );
                                chunk.mark_dead_entry_with_seg(new, &new_seg);
                            }
                        } else {
                            trace!("cell {} address {} have been removed on combine", hash, old);
                            let _ = chunk.put_tombstone_by_cell_loc(new);
                        }
                    }
                });
            space_cleaned = space_to_collect - cleaned_total_live_space.load(Relaxed);
            debug!(
                "Combined {} segments to {}, total {} bytes",
                segments_to_combine_len,
                pending_segments.len(),
                space_cleaned
            );
        } else {
            debug!("No entries to work on, will remove all selected segments instead");
        }

        debug!("Removing {} old segments", segments.len());
        for old_seg in segments {
            chunk.remove_segment(old_seg.id);
            old_seg.mem_drop(chunk);
        }
        space_cleaned
    }
}
<commit_msg>parallel marking<commit_after>use crate::ram::chunk::Chunk;
use crate::ram::entry::EntryContent;
use crate::ram::segs::{Segment, SEGMENT_SIZE};
use itertools::Itertools;
use lightning::linked_map::NodeRef as MapNodeRef;
use rayon::prelude::*;
use std::collections::HashSet;
use std::sync::atomic::Ordering;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use std::sync::Arc;

use libc;

pub struct CombinedCleaner;

#[derive(Clone)]
struct DummyEntry {
    size: usize,
    addr: usize,
    cell_hash: Option<u64>,
    timestamp: u32,
}

struct DummySegment {
    head: usize,
    entries: Vec<DummyEntry>,
}

impl DummySegment {
    fn new() -> DummySegment {
        DummySegment {
            head: 0,
            entries: Vec::new(),
        }
    }
}

// Combine small segments into larger segments
// This cleaner will perform a greedy approach to relocate entries from old segments to fewer
// new segments to reduce number or segments and reclaim spaces from tombstones

// for higher hit rate for fetching cells in segments, we need to put data with close timestamp together
// this optimization is intended for enabling neb to contain data more than it's memory

impl CombinedCleaner {
    pub fn combine_segments(chunk: &Chunk, segments: &Vec<MapNodeRef<Segment>>) -> usize {
        if segments.len() < 2 {
            trace!(
                "too few segments to combine, chunk {}, segments {}",
                chunk.id,
                segments.len()
            );
            return 0;
        }
        debug!("Combining segments");

        let space_to_collect = segments
            .iter()
            .map(|seg| seg.used_spaces() as usize)
            .sum::<usize>();
        let segment_ids_to_combine: HashSet<_> = segments.iter().map(|seg| seg.id).collect();

        // Get all entries in segments to combine and order them by data temperature and size
        let nested_entries = segments
            .iter()
            .flat_map(|seg| chunk.live_entries(seg))
            .filter(|entry| {
                // live entries have done a lot of filtering work already
                // but we still need to remove those tombstones that pointed to segments we are about to combine
                if let EntryContent::Tombstone(ref tombstone) = entry.content {
                    return !segment_ids_to_combine.contains(&tombstone.segment_id);
                }
                return true;
            })
            .map(|entry| {
                let entry_size = entry.meta.entry_size;
                let entry_addr = entry.meta.entry_pos;
                let cell_header = match entry.content {
                    EntryContent::Cell(header) => Some(header),
                    _ => None,
                };
                DummyEntry {
                    size: entry_size,
                    addr: entry_addr,
                    timestamp: cell_header.map(|h| h.timestamp).unwrap_or(0),
                    cell_hash: cell_header.map(|h| h.hash),
                }
            })
            .group_by(|entry| entry.timestamp / 10)
            .into_iter()
            .map(|(t, group)| {
                let mut group: Vec<_> = group.collect();
                group.sort_by_key(|entry| entry.size);
                return (t, group.into_iter());
            })
            .sorted_by_key(|&(t, _)| t)
            .into_iter()
            .map(|(_, group)| group);

        let mut entries: Vec<_> = Iterator::flatten(nested_entries)
            // order by temperature and size from greater to lesser
            .rev()
            // provide additional state for whether entry have been claimed on simulation
            .map(|e| (e, false))
            .collect();

        debug!("Found {} entries to combine", entries.len());

        let mut space_cleaned = 0;
        if entries.len() > 0 {
            // Simulate the combine process to determine the efficiency
            let mut pending_segments = Vec::with_capacity(segments.len());
            let entries_num = entries.len();
            let mut entries_to_claim = entries_num;
            let mut cursor = 0;
            pending_segments.push(DummySegment::new());
            while entries_to_claim > 0 {
                let segment_space_remains = SEGMENT_SIZE - pending_segments.last().unwrap().head;
                let index = cursor % entries_num;
                let entry_pair = entries.get_mut(index).unwrap();
                if entry_pair.1 {
                    // entry claimed
                    cursor += 1;
                    continue;
                }

                let entry = &entry_pair.0;
                let entry_size = entry.size;
                if entry_size > segment_space_remains {
                    if index == entries_num - 1 {
                        // iterated to the last one, which means no more entries can be placed
                        // in the segment, then create a new segment
                        pending_segments.push(DummySegment::new());
                    }
                    cursor += 1;
                    continue;
                }
                let last_segment = pending_segments.last_mut().unwrap();
                last_segment.entries.push(entry.clone());

                // pump dummy segment head pointer
                last_segment.head += entry_size;

                // mark entry claimed
                entry_pair.1 = true;
                entries_to_claim -= 1;

                // move to next entry
                cursor += 1;
            }

            debug!("Checking combine feasibility");
            let pending_segments_len = pending_segments.len();
            let segments_to_combine_len = segments.len();
            let cleaned_total_live_space = AtomicUsize::new(0);
            if pending_segments_len >= segments_to_combine_len {
                warn!(
                    "Trying to combine segments but resulting segments still does not go down {}/{}",
                    pending_segments_len, segments_to_combine_len
                );
            }

            debug!(
                "Updating cell reference, pending segments {}",
                pending_segments.len()
            );
            pending_segments
                .par_iter()
                .map(|dummy_seg| {
                    let new_seg = chunk
                        .allocator
                        .alloc_seg(&chunk.backup_storage, &chunk.wal_storage)
                        .expect("No space left during combine");
                    let new_seg_id = new_seg.id;
                    let mut cell_mapping = Vec::with_capacity(dummy_seg.entries.len());
                    let mut seg_cursor = new_seg.addr;
                    trace!(
                        "Combining segment to new one with id {} with {} cells",
                        new_seg_id,
                        dummy_seg.entries.len()
                    );
                    for entry in &dummy_seg.entries {
                        let entry_addr = entry.addr;
                        unsafe {
                            libc::memcpy(
                                seg_cursor as *mut libc::c_void,
                                entry_addr as *mut libc::c_void,
                                entry.size,
                            );
                        }
                        if let Some(cell_hash) = entry.cell_hash {
                            trace!(
                                "Marked cell relocation hash {}, addr {} to segment {}",
                                cell_hash,
                                entry_addr,
                                new_seg_id
                            );
                            cell_mapping.push((seg_cursor, entry_addr, cell_hash));
                        }
                        seg_cursor += entry.size;
                    }
                    new_seg.append_header.store(seg_cursor, Ordering::Relaxed);
                    cleaned_total_live_space.fetch_add(new_seg.used_spaces() as usize, Relaxed);
                    return (new_seg, cell_mapping);
                })
                .for_each(|(segment, cells)| {
                    trace!("Putting new segment {}, cells {}", segment.id, cells.len());
                    segment.archive().unwrap();
                    let new_seg_id = segment.id as usize;
                    chunk.put_segment(segment);
                    let new_seg = chunk.segs.get(&new_seg_id).unwrap();
                    cells.into_par_iter().for_each(|(new, old, hash)|
                    {
                        trace!("Reset cell {} ptr from {} to {}", hash, old, new);
                        #[cfg(feature = "fast_map")]
                        let index = chunk.index.lock(hash as usize);
                        #[cfg(feature = "slow_map")]
                        let index = chunk.index.get_mut(&hash);

                        if let Some(mut actual_addr) = index {
                            if *actual_addr == old {
                                *actual_addr = new;
                                trace!(
                                    "Cell addr for hash {} set from {} to {} for combine",
                                    hash,
                                    old,
                                    new
                                );
                            } else {
                                trace!(
                                    "cell {} with address {}, have been changed to {} on combine",
                                    hash,
                                    old,
                                    *actual_addr
                                );
                                chunk.mark_dead_entry_with_seg(new, &new_seg);
                            }
                        } else {
                            trace!("cell {} address {} have been removed on combine", hash, old);
                            let _ = chunk.put_tombstone_by_cell_loc(new);
                        }
                    });
                });
            space_cleaned = space_to_collect - cleaned_total_live_space.load(Relaxed);
            debug!(
                "Combined {} segments to {}, total {} bytes",
                segments_to_combine_len,
                pending_segments.len(),
                space_cleaned
            );
        } else {
            debug!("No entries to work on, will remove all selected segments instead");
        }

        debug!("Removing {} old segments", segments.len());
        for old_seg in segments {
            chunk.remove_segment(old_seg.id);
            old_seg.mem_drop(chunk);
        }
        space_cleaned
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::proto;
use crate::proto_grpc;
use cgmath::{Decomposed, Matrix4, PerspectiveFov, Point3, Quaternion, Rad, Transform, Vector3};
use collision::Aabb3;
use futures::sync::mpsc;
use futures::{Future, Sink, Stream};
use grpcio::{
    Environment, RpcContext, RpcStatus, RpcStatusCode, Server, ServerBuilder, ServerStreamingSink,
    UnarySink, WriteFlags,
};
use point_viewer::errors::*;
use point_viewer::octree::{NodeId, Octree, OctreeFactory};
use point_viewer::{InternalIterator, Point};
use protobuf::Message;
use std::collections::HashMap;
use std::path::PathBuf;
use std::str::FromStr;
use std::sync::{Arc, RwLock};

struct OctreeServiceData {
    octree: Box<Octree>,
    meta: point_viewer::proto::Meta,
}

#[derive(Clone)]
struct OctreeService {
    location: PathBuf,
    data_cache: Arc<RwLock<HashMap<String, Arc<OctreeServiceData>>>>,
    factory: OctreeFactory,
}

#[derive(Debug)]
enum OctreeQuery {
    Frustum(Matrix4<f32>),
    Box(Aabb3<f64>),
    FullPointcloud,
}

fn send_fail_stream<T>(ctx: &RpcContext, sink: ServerStreamingSink<T>, err_str: String) {
    let f = sink
        .fail(RpcStatus::new(RpcStatusCode::Internal, Some(err_str)))
        .map_err(move |err| eprintln!("Failed to reply: {:?}", err));
    ctx.spawn(f);
}

fn send_fail<T>(ctx: &RpcContext, sink: UnarySink<T>, err_str: String) {
    let f = sink
        .fail(RpcStatus::new(RpcStatusCode::Internal, Some(err_str)))
        .map_err(move |err| eprintln!("Failed to reply: {:?}", err));
    ctx.spawn(f);
}

impl proto_grpc::Octree for OctreeService {
    fn get_meta(
        &self,
        ctx: RpcContext,
        req: proto::GetMetaRequest,
        sink: UnarySink<proto::GetMetaReply>,
    ) {
        let mut resp = proto::GetMetaReply::new();
        let service_data = match self.get_service_data(&req.octree_id) {
            Ok(service_data) => service_data,
            Err(e) => return send_fail(&ctx, sink, e.to_string()),
        };
        resp.set_meta(service_data.meta.clone());
        let f = sink
            .success(resp)
            .map_err(move |e| println!("failed to reply {:?}: {:?}", req, e));
        ctx.spawn(f)
    }

    fn get_node_data(
        &self,
        ctx: RpcContext,
        req: proto::GetNodeDataRequest,
        sink: UnarySink<proto::GetNodeDataReply>,
    ) {
        let service_data = match self.get_service_data(&req.octree_id) {
            Ok(service_data) => service_data,
            Err(e) => return send_fail(&ctx, sink, e.to_string()),
        };
        let node_id = match NodeId::from_str(&req.id) {
            Ok(node_id) => node_id,
            Err(e) => return send_fail(&ctx, sink, e.to_string()),
        };
        let node_data = match service_data.octree.get_node_data(&node_id) {
            Ok(data) => data,
            Err(e) => return send_fail(&ctx, sink, e.to_string()),
        };
        let mut resp = proto::GetNodeDataReply::new();
        resp.mut_node()
            .set_position_encoding(node_data.meta.position_encoding.to_proto());
        resp.mut_node().set_num_points(node_data.meta.num_points);
        resp.set_position(node_data.position);
        resp.set_color(node_data.color);
        let f = sink
            .success(resp)
            .map_err(move |e| println!("failed to reply {:?}: {:?}", req, e));
        ctx.spawn(f)
    }

    fn get_points_in_frustum(
        &self,
        ctx: RpcContext,
        req: proto::GetPointsInFrustumRequest,
        resp: ServerStreamingSink<proto::PointsReply>,
    ) {
        let projection_matrix = Matrix4::from(PerspectiveFov {
            fovy: Rad(req.fovy_rad as f32),
            aspect: (req.aspect as f32),
            near: (req.z_near as f32),
            far: (req.z_far as f32),
        });
        let rotation = {
            let q = req.rotation.unwrap();
            Quaternion::new(q.w, q.x, q.y, q.z)
        };

        let translation = {
            let t = req.translation.unwrap();
            Vector3::new(t.x, t.y, t.z)
        };

        let view_transform = Decomposed {
            scale: 1.0,
            rot: rotation,
            disp: translation,
        };
        let frustum_matrix = projection_matrix.concat(&view_transform.into());
        self.stream_points_back_to_sink(
            OctreeQuery::Frustum(frustum_matrix),
            &req.octree_id,
            &ctx,
            resp,
        )
    }

    fn get_points_in_box(
        &self,
        ctx: RpcContext,
        req: proto::GetPointsInBoxRequest,
        resp: ServerStreamingSink<proto::PointsReply>,
    ) {
        let bounding_box = {
            let bounding_box = req.bounding_box.clone().unwrap();
            let min = bounding_box.min.unwrap();
            let max = bounding_box.max.unwrap();
            Aabb3::new(
                Point3::new(min.x, min.y, min.z),
                Point3::new(max.x, max.y, max.z),
            )
        };
        self.stream_points_back_to_sink(OctreeQuery::Box(bounding_box), &req.octree_id, &ctx, resp)
    }

    fn get_all_points(
        &self,
        ctx: RpcContext,
        req: proto::GetAllPointsRequest,
        resp: ServerStreamingSink<proto::PointsReply>,
    ) {
        self.stream_points_back_to_sink(OctreeQuery::FullPointcloud, &req.octree_id, &ctx, resp)
    }
}

impl OctreeService {
    fn stream_points_back_to_sink(
        &self,
        query: OctreeQuery,
        octree_id: &str,
        ctx: &RpcContext,
        resp: ServerStreamingSink<proto::PointsReply>,
    ) {
        use std::thread;

        let service_data = match self.get_service_data(octree_id) {
            Ok(service_data) => service_data,
            Err(e) => return send_fail_stream(&ctx, resp, e.to_string()),
        };

        // This creates a async-aware (tx, rx) pair that can wake up the event loop when new data
        // is piped through it.

        // We create a channel with a small buffer, which yields better performance than
        // fully blocking without requiring a ton of memory. This has not been carefully benchmarked
        // for best performance though.
        let (tx, rx) = mpsc::channel(4);
        thread::spawn(move || {
            // This is the secret sauce connecting an OS thread to a event-based receiver. Calling
            // wait() on this turns the event aware, i.e. async 'tx' into a blocking 'tx' that will
            // make this thread block when the event loop is not quick enough with piping out data.
            let mut tx = tx.wait();

            let mut reply = proto::PointsReply::new();
            let bytes_per_point = {
                let initial_proto_size = reply.compute_size();
                let mut v = point_viewer::proto::Vector3d::new();
                v.set_x(1.);
                v.set_y(1.);
                v.set_z(1.);
                reply.mut_positions().push(v);

                let mut v = point_viewer::proto::Color::new();
                v.set_red(1.);
                v.set_green(1.);
                v.set_blue(1.);
                v.set_alpha(1.);
                reply.mut_colors().push(v);

                reply.mut_intensities().push(1.);

                let final_proto_size = reply.compute_size();
                reply.mut_positions().clear();
                reply.mut_colors().clear();
                reply.mut_intensities().clear();
                final_proto_size - initial_proto_size
            };

            // Proto message must be below 4 MB.
            let max_message_size = 4 * 1024 * 1024;
            let mut reply_size = 0;

            {
                // Extra scope to make sure that 'func' does not outlive 'reply'.
                let func = |p: &Point| {
                    {
                        let mut v = point_viewer::proto::Vector3d::new();
                        v.set_x(p.position.x);
                        v.set_y(p.position.y);
                        v.set_z(p.position.z);
                        reply.mut_positions().push(v);
                    }

                    {
                        let mut v = point_viewer::proto::Color::new();
                        let clr = p.color.to_f32();
                        v.set_red(clr.red);
                        v.set_green(clr.green);
                        v.set_blue(clr.blue);
                        v.set_alpha(clr.alpha);
                        reply.mut_colors().push(v);
                    }

                    if let Some(i) = p.intensity {
                        reply.mut_intensities().push(i);
                    }

                    reply_size += bytes_per_point;
                    if reply_size > max_message_size - bytes_per_point {
                        tx.send((reply.clone(), WriteFlags::default())).unwrap();
                        reply.mut_positions().clear();
                        reply.mut_colors().clear();
                        reply.mut_intensities().clear();
                        reply_size = 0;
                    }
                };
                match query {
                    OctreeQuery::Box(bounding_box) => service_data
                        .octree
                        .points_in_box(&bounding_box)
                        .for_each(func),
                    OctreeQuery::Frustum(frustum_matrix) => service_data
                        .octree
                        .points_in_frustum(&frustum_matrix)
                        .for_each(func),
                    OctreeQuery::FullPointcloud => service_data.octree.all_points().for_each(func),
                };
            }
            tx.send((reply, WriteFlags::default())).unwrap();
        });

        let rx = rx.map_err(|_| grpcio::Error::RemoteStopped);
        let f = resp
            .send_all(rx)
            .map(|_| {})
            .map_err(|e| println!("failed to reply: {:?}", e));
        ctx.spawn(f)
    }

    fn get_service_data(&self, octree_id: &str) -> Result<Arc<OctreeServiceData>> {
        let read_cache = self.data_cache.read().unwrap();
        match read_cache.get(octree_id) {
            None => {
                drop(read_cache);
                let octree = self
                    .factory
                    .generate_octree(self.location.join(&octree_id).to_string_lossy())?;
                let meta = octree.to_meta_proto();
                let service_data = Arc::new(OctreeServiceData { octree, meta });
                self.data_cache
                    .write()
                    .unwrap()
                    .insert(octree_id.to_string(), Arc::clone(&service_data));
                Ok(service_data)
            }
            Some(service_data) => Ok(Arc::clone(service_data)),
        }
    }
}

pub fn start_grpc_server(
    host: &str,
    port: u16,
    location: impl Into<PathBuf>,
    factory: OctreeFactory,
) -> Server {
    let env = Arc::new(Environment::new(1));
    let data_cache = Arc::new(RwLock::new(HashMap::new()));

    let service = proto_grpc::create_octree(OctreeService {
        location: location.into(),
        data_cache,
        factory,
    });
    ServerBuilder::new(env)
        .register_service(service)
        .bind(host /* ip to bind to */, port)
        .build()
        .unwrap()
}
<commit_msg>Improved rwlock release logic. (#241)<commit_after>// Copyright 2018 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::proto;
use crate::proto_grpc;
use cgmath::{Decomposed, Matrix4, PerspectiveFov, Point3, Quaternion, Rad, Transform, Vector3};
use collision::Aabb3;
use futures::sync::mpsc;
use futures::{Future, Sink, Stream};
use grpcio::{
    Environment, RpcContext, RpcStatus, RpcStatusCode, Server, ServerBuilder, ServerStreamingSink,
    UnarySink, WriteFlags,
};
use point_viewer::errors::*;
use point_viewer::octree::{NodeId, Octree, OctreeFactory};
use point_viewer::{InternalIterator, Point};
use protobuf::Message;
use std::collections::HashMap;
use std::path::PathBuf;
use std::str::FromStr;
use std::sync::{Arc, RwLock};

struct OctreeServiceData {
    octree: Box<Octree>,
    meta: point_viewer::proto::Meta,
}

#[derive(Clone)]
struct OctreeService {
    location: PathBuf,
    data_cache: Arc<RwLock<HashMap<String, Arc<OctreeServiceData>>>>,
    factory: OctreeFactory,
}

#[derive(Debug)]
enum OctreeQuery {
    Frustum(Matrix4<f32>),
    Box(Aabb3<f64>),
    FullPointcloud,
}

fn send_fail_stream<T>(ctx: &RpcContext, sink: ServerStreamingSink<T>, err_str: String) {
    let f = sink
        .fail(RpcStatus::new(RpcStatusCode::Internal, Some(err_str)))
        .map_err(move |err| eprintln!("Failed to reply: {:?}", err));
    ctx.spawn(f);
}

fn send_fail<T>(ctx: &RpcContext, sink: UnarySink<T>, err_str: String) {
    let f = sink
        .fail(RpcStatus::new(RpcStatusCode::Internal, Some(err_str)))
        .map_err(move |err| eprintln!("Failed to reply: {:?}", err));
    ctx.spawn(f);
}

impl proto_grpc::Octree for OctreeService {
    fn get_meta(
        &self,
        ctx: RpcContext,
        req: proto::GetMetaRequest,
        sink: UnarySink<proto::GetMetaReply>,
    ) {
        let mut resp = proto::GetMetaReply::new();
        let service_data = match self.get_service_data(&req.octree_id) {
            Ok(service_data) => service_data,
            Err(e) => return send_fail(&ctx, sink, e.to_string()),
        };
        resp.set_meta(service_data.meta.clone());
        let f = sink
            .success(resp)
            .map_err(move |e| println!("failed to reply {:?}: {:?}", req, e));
        ctx.spawn(f)
    }

    fn get_node_data(
        &self,
        ctx: RpcContext,
        req: proto::GetNodeDataRequest,
        sink: UnarySink<proto::GetNodeDataReply>,
    ) {
        let service_data = match self.get_service_data(&req.octree_id) {
            Ok(service_data) => service_data,
            Err(e) => return send_fail(&ctx, sink, e.to_string()),
        };
        let node_id = match NodeId::from_str(&req.id) {
            Ok(node_id) => node_id,
            Err(e) => return send_fail(&ctx, sink, e.to_string()),
        };
        let node_data = match service_data.octree.get_node_data(&node_id) {
            Ok(data) => data,
            Err(e) => return send_fail(&ctx, sink, e.to_string()),
        };
        let mut resp = proto::GetNodeDataReply::new();
        resp.mut_node()
            .set_position_encoding(node_data.meta.position_encoding.to_proto());
        resp.mut_node().set_num_points(node_data.meta.num_points);
        resp.set_position(node_data.position);
        resp.set_color(node_data.color);
        let f = sink
            .success(resp)
            .map_err(move |e| println!("failed to reply {:?}: {:?}", req, e));
        ctx.spawn(f)
    }

    fn get_points_in_frustum(
        &self,
        ctx: RpcContext,
        req: proto::GetPointsInFrustumRequest,
        resp: ServerStreamingSink<proto::PointsReply>,
    ) {
        let projection_matrix = Matrix4::from(PerspectiveFov {
            fovy: Rad(req.fovy_rad as f32),
            aspect: (req.aspect as f32),
            near: (req.z_near as f32),
            far: (req.z_far as f32),
        });
        let rotation = {
            let q = req.rotation.unwrap();
            Quaternion::new(q.w, q.x, q.y, q.z)
        };

        let translation = {
            let t = req.translation.unwrap();
            Vector3::new(t.x, t.y, t.z)
        };

        let view_transform = Decomposed {
            scale: 1.0,
            rot: rotation,
            disp: translation,
        };
        let frustum_matrix = projection_matrix.concat(&view_transform.into());
        self.stream_points_back_to_sink(
            OctreeQuery::Frustum(frustum_matrix),
            &req.octree_id,
            &ctx,
            resp,
        )
    }

    fn get_points_in_box(
        &self,
        ctx: RpcContext,
        req: proto::GetPointsInBoxRequest,
        resp: ServerStreamingSink<proto::PointsReply>,
    ) {
        let bounding_box = {
            let bounding_box = req.bounding_box.clone().unwrap();
            let min = bounding_box.min.unwrap();
            let max = bounding_box.max.unwrap();
            Aabb3::new(
                Point3::new(min.x, min.y, min.z),
                Point3::new(max.x, max.y, max.z),
            )
        };
        self.stream_points_back_to_sink(OctreeQuery::Box(bounding_box), &req.octree_id, &ctx, resp)
    }

    fn get_all_points(
        &self,
        ctx: RpcContext,
        req: proto::GetAllPointsRequest,
        resp: ServerStreamingSink<proto::PointsReply>,
    ) {
        self.stream_points_back_to_sink(OctreeQuery::FullPointcloud, &req.octree_id, &ctx, resp)
    }
}

impl OctreeService {
    fn stream_points_back_to_sink(
        &self,
        query: OctreeQuery,
        octree_id: &str,
        ctx: &RpcContext,
        resp: ServerStreamingSink<proto::PointsReply>,
    ) {
        use std::thread;

        let service_data = match self.get_service_data(octree_id) {
            Ok(service_data) => service_data,
            Err(e) => return send_fail_stream(&ctx, resp, e.to_string()),
        };

        // This creates a async-aware (tx, rx) pair that can wake up the event loop when new data
        // is piped through it.

        // We create a channel with a small buffer, which yields better performance than
        // fully blocking without requiring a ton of memory. This has not been carefully benchmarked
        // for best performance though.
        let (tx, rx) = mpsc::channel(4);
        thread::spawn(move || {
            // This is the secret sauce connecting an OS thread to a event-based receiver. Calling
            // wait() on this turns the event aware, i.e. async 'tx' into a blocking 'tx' that will
            // make this thread block when the event loop is not quick enough with piping out data.
            let mut tx = tx.wait();

            let mut reply = proto::PointsReply::new();
            let bytes_per_point = {
                let initial_proto_size = reply.compute_size();
                let mut v = point_viewer::proto::Vector3d::new();
                v.set_x(1.);
                v.set_y(1.);
                v.set_z(1.);
                reply.mut_positions().push(v);

                let mut v = point_viewer::proto::Color::new();
                v.set_red(1.);
                v.set_green(1.);
                v.set_blue(1.);
                v.set_alpha(1.);
                reply.mut_colors().push(v);

                reply.mut_intensities().push(1.);

                let final_proto_size = reply.compute_size();
                reply.mut_positions().clear();
                reply.mut_colors().clear();
                reply.mut_intensities().clear();
                final_proto_size - initial_proto_size
            };

            // Proto message must be below 4 MB.
            let max_message_size = 4 * 1024 * 1024;
            let mut reply_size = 0;

            {
                // Extra scope to make sure that 'func' does not outlive 'reply'.
                let func = |p: &Point| {
                    {
                        let mut v = point_viewer::proto::Vector3d::new();
                        v.set_x(p.position.x);
                        v.set_y(p.position.y);
                        v.set_z(p.position.z);
                        reply.mut_positions().push(v);
                    }

                    {
                        let mut v = point_viewer::proto::Color::new();
                        let clr = p.color.to_f32();
                        v.set_red(clr.red);
                        v.set_green(clr.green);
                        v.set_blue(clr.blue);
                        v.set_alpha(clr.alpha);
                        reply.mut_colors().push(v);
                    }

                    if let Some(i) = p.intensity {
                        reply.mut_intensities().push(i);
                    }

                    reply_size += bytes_per_point;
                    if reply_size > max_message_size - bytes_per_point {
                        tx.send((reply.clone(), WriteFlags::default())).unwrap();
                        reply.mut_positions().clear();
                        reply.mut_colors().clear();
                        reply.mut_intensities().clear();
                        reply_size = 0;
                    }
                };
                match query {
                    OctreeQuery::Box(bounding_box) => service_data
                        .octree
                        .points_in_box(&bounding_box)
                        .for_each(func),
                    OctreeQuery::Frustum(frustum_matrix) => service_data
                        .octree
                        .points_in_frustum(&frustum_matrix)
                        .for_each(func),
                    OctreeQuery::FullPointcloud => service_data.octree.all_points().for_each(func),
                };
            }
            tx.send((reply, WriteFlags::default())).unwrap();
        });

        let rx = rx.map_err(|_| grpcio::Error::RemoteStopped);
        let f = resp
            .send_all(rx)
            .map(|_| {})
            .map_err(|e| println!("failed to reply: {:?}", e));
        ctx.spawn(f)
    }

    fn get_service_data(&self, octree_id: &str) -> Result<Arc<OctreeServiceData>> {
        if let Some(service_data) = self.data_cache.read().unwrap().get(octree_id) {
            return Ok(Arc::clone(service_data));
        };
        let octree = self
            .factory
            .generate_octree(self.location.join(&octree_id).to_string_lossy())?;
        let meta = octree.to_meta_proto();
        let service_data = Arc::new(OctreeServiceData { octree, meta });
        self.data_cache
            .write()
            .unwrap()
            .insert(octree_id.to_string(), Arc::clone(&service_data));
        Ok(service_data)
    }
}

pub fn start_grpc_server(
    host: &str,
    port: u16,
    location: impl Into<PathBuf>,
    factory: OctreeFactory,
) -> Server {
    let env = Arc::new(Environment::new(1));
    let data_cache = Arc::new(RwLock::new(HashMap::new()));

    let service = proto_grpc::create_octree(OctreeService {
        location: location.into(),
        data_cache,
        factory,
    });
    ServerBuilder::new(env)
        .register_service(service)
        .bind(host /* ip to bind to */, port)
        .build()
        .unwrap()
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::error;
use std::fmt;
use std::marker::PhantomData;
use std::ptr;
use smallvec::SmallVec;

use command_buffer::cb::UnsafeCommandBuffer;
use command_buffer::pool::CommandPool;
use device::Queue;
use sync::Fence;
use sync::PipelineStages;
use sync::Semaphore;

use check_errors;
use vk;
use Error;
use OomError;
use VulkanObject;
use VulkanPointers;
use SynchronizedVulkanObject;

/// Prototype for a submission that executes command buffers.
// TODO: example here
#[derive(Debug)]
pub struct SubmitCommandBufferBuilder<'a> {
    wait_semaphores: SmallVec<[vk::Semaphore; 16]>,
    dest_stages: SmallVec<[vk::PipelineStageFlags; 8]>,
    signal_semaphores: SmallVec<[vk::Semaphore; 16]>,
    command_buffers: SmallVec<[vk::CommandBuffer; 4]>,
    fence: vk::Fence,
    marker: PhantomData<&'a ()>,
}

impl<'a> SubmitCommandBufferBuilder<'a> {
    /// Builds a new empty `SubmitCommandBufferBuilder`.
    #[inline]
    pub fn new() -> SubmitCommandBufferBuilder<'a> {
        SubmitCommandBufferBuilder {
            wait_semaphores: SmallVec::new(),
            dest_stages: SmallVec::new(),
            signal_semaphores: SmallVec::new(),
            command_buffers: SmallVec::new(),
            fence: 0,
            marker: PhantomData,
        }
    }

    /// Returns true if this builder will signal a fence when submitted.
    ///
    /// # Example
    ///
    /// ```
    /// use vulkano::command_buffer::submit::SubmitCommandBufferBuilder;
    /// use vulkano::sync::Fence;
    /// # let device: std::sync::Arc<vulkano::device::Device> = return;
    ///
    /// unsafe {
    ///     let fence = Fence::new(device.clone()).unwrap();
    ///
    ///     let mut builder = SubmitCommandBufferBuilder::new();
    ///     assert!(!builder.has_fence());
    ///     builder.set_fence_signal(&fence);
    ///     assert!(builder.has_fence());
    /// }
    /// ```
    #[inline]
    pub fn has_fence(&self) -> bool {
        self.fence != 0
    }

    /// Adds an operation that signals a fence after this submission ends.
    ///
    /// # Example
    ///
    /// ```
    /// use std::time::Duration;
    /// use vulkano::command_buffer::submit::SubmitCommandBufferBuilder;
    /// use vulkano::sync::Fence;
    /// # let device: std::sync::Arc<vulkano::device::Device> = return;
    /// # let queue: std::sync::Arc<vulkano::device::Queue> = return;
    ///
    /// unsafe {
    ///     let fence = Fence::new(device.clone()).unwrap();
    ///
    ///     let mut builder = SubmitCommandBufferBuilder::new();
    ///     builder.set_fence_signal(&fence);
    ///
    ///     builder.submit(&queue).unwrap();
    ///
    ///     // We must not destroy the fence before it is signaled.
    ///     fence.wait(Duration::from_secs(5)).unwrap();
    /// }
    /// ```
    ///
    /// # Safety
    ///
    /// - The fence must not be signaled at the time when you call `submit()`.
    ///
    /// - If you use the fence for multiple submissions, only one at a time must be executed by the
    ///   GPU. In other words, you must submit one, wait for the fence to be signaled, then reset
    ///   the fence, and then only submit the second.
    ///
    /// - If you submit this builder, the fence must be kept alive until it is signaled by the GPU.
    ///   Destroying the fence earlier is an undefined behavior.
    ///
    /// - The fence, command buffers, and semaphores must all belong to the same device.
    ///
    #[inline]
    pub unsafe fn set_fence_signal(&mut self, fence: &'a Fence) {
        self.fence = fence.internal_object();
    }

    /// Adds a semaphore to be waited upon before the command buffers are executed.
    ///
    /// Only the given `stages` of the command buffers added afterwards will wait upon
    /// the semaphore. Other stages not included in `stages` can execute before waiting.
    ///
    /// # Safety
    ///
    /// - The stages must be supported by the device.
    ///
    /// - If you submit this builder, the semaphore must be kept alive until you are guaranteed
    ///   that the GPU has at least started executing the command buffers.
    ///
    /// - If you submit this builder, no other queue must be waiting on these semaphores. In other
    ///   words, each semaphore signal can only correspond to one semaphore wait.
    ///
    /// - If you submit this builder, the semaphores must be signaled when the queue execution
    ///   reaches this submission, or there must be one or more submissions in queues that are
    ///   going to signal these semaphores. In other words, you must not block the queue with
    ///   semaphores that can't get signaled.
    ///
    /// - The fence, command buffers, and semaphores must all belong to the same device.
    ///
    #[inline]
    pub unsafe fn add_wait_semaphore(&mut self, semaphore: &'a Semaphore, stages: PipelineStages) {
        debug_assert!(Into::<vk::PipelineStageFlagBits>::into(stages) != 0);
        // TODO: debug assert that the device supports the stages
        self.wait_semaphores.push(semaphore.internal_object());
        self.dest_stages.push(stages.into());
    }

    /// Adds a command buffer that is executed as part of this command.
    ///
    /// The command buffers are submitted in the order in which they are added.
    ///
    /// # Safety
    ///
    /// - If you submit this builder, the command buffer must be kept alive until you are
    ///   guaranteed that the GPU has finished executing it.
    ///
    /// - Any calls to vkCmdSetEvent, vkCmdResetEvent or vkCmdWaitEvents that have been recorded
    ///   into the command buffer must not reference any VkEvent that is referenced by any of
    ///   those commands that is pending execution on another queue.
    ///   TODO: rephrase ^ ?
    ///
    /// - The fence, command buffers, and semaphores must all belong to the same device.
    ///
    #[inline]
    pub unsafe fn add_command_buffer<P>(&mut self, command_buffer: &'a UnsafeCommandBuffer<P>)
        where P: CommandPool
    {
        self.command_buffers.push(command_buffer.internal_object());
    }

    /// Returns the number of semaphores to signal.
    ///
    /// In other words, this is the number of times `add_signal_semaphore` has been called.
    #[inline]
    pub fn num_signal_semaphores(&self) -> usize {
        self.signal_semaphores.len()
    }

    /// Adds a semaphore that is going to be signaled at the end of the submission.
    ///
    /// # Safety
    ///
    /// - If you submit this builder, the semaphore must be kept alive until you are guaranteed
    ///   that the GPU has finished executing this submission.
    ///
    /// - The semaphore must be in the unsignaled state when queue execution reaches this
    ///   submission.
    ///
    /// - The fence, command buffers, and semaphores must all belong to the same device.
    ///
    #[inline]
    pub unsafe fn add_signal_semaphore(&mut self, semaphore: &'a Semaphore) {
        self.signal_semaphores.push(semaphore.internal_object());
    }

    /// Submits the command buffer to the given queue.
    ///
    /// > **Note**: This is an expensive operation, so you may want to merge as many builders as
    /// > possible together and avoid submitting them one by one.
    ///
    pub fn submit(mut self, queue: &Queue) -> Result<(), SubmitCommandBufferError> {
        unsafe {
            let vk = queue.device().pointers();
            let queue = queue.internal_object_guard();

            debug_assert_eq!(self.wait_semaphores.len(), self.dest_stages.len());

            let batch = vk::SubmitInfo {
                sType: vk::STRUCTURE_TYPE_SUBMIT_INFO,
                pNext: ptr::null(),
                waitSemaphoreCount: self.wait_semaphores.len() as u32,
                pWaitSemaphores: self.wait_semaphores.as_ptr(),
                pWaitDstStageMask: self.dest_stages.as_ptr(),
                commandBufferCount: self.command_buffers.len() as u32,
                pCommandBuffers: self.command_buffers.as_ptr(),
                signalSemaphoreCount: self.signal_semaphores.len() as u32,
                pSignalSemaphores: self.signal_semaphores.as_ptr(),
            };

            try!(check_errors(vk.QueueSubmit(*queue, 1, &batch, self.fence)));
            Ok(())
        }
    }

    /// Merges this builder with another builder.
    ///
    /// # Panic
    ///
    /// Panics if both builders have a fence already set.
    // TODO: create multiple batches instead
    pub fn merge(mut self, other: Self) -> Self {
        assert!(self.fence == 0 || other.fence == 0,
               "Can't merge two queue submits that both have a fence");

        self.wait_semaphores.extend(other.wait_semaphores);
        self.dest_stages.extend(other.dest_stages);     // TODO: meh? will be solved if we submit multiple batches
        self.signal_semaphores.extend(other.signal_semaphores);
        self.command_buffers.extend(other.command_buffers);

        if self.fence == 0 {
            self.fence = other.fence;
        }

        self
    }
}

/// Error that can happen when submitting the prototype.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
#[repr(u32)]
pub enum SubmitCommandBufferError {
    /// Not enough memory.
    OomError(OomError),

    /// The connection to the device has been lost.
    DeviceLost,
}

impl error::Error for SubmitCommandBufferError {
    #[inline]
    fn description(&self) -> &str {
        match *self {
            SubmitCommandBufferError::OomError(_) => "not enough memory",
            SubmitCommandBufferError::DeviceLost => "the connection to the device has been lost",
        }
    }

    #[inline]
    fn cause(&self) -> Option<&error::Error> {
        match *self {
            SubmitCommandBufferError::OomError(ref err) => Some(err),
            _ => None
        }
    }
}

impl fmt::Display for SubmitCommandBufferError {
    #[inline]
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", error::Error::description(self))
    }
}

impl From<Error> for SubmitCommandBufferError {
    #[inline]
    fn from(err: Error) -> SubmitCommandBufferError {
        match err {
            err @ Error::OutOfHostMemory => SubmitCommandBufferError::OomError(OomError::from(err)),
            err @ Error::OutOfDeviceMemory => SubmitCommandBufferError::OomError(OomError::from(err)),
            Error::DeviceLost => SubmitCommandBufferError::DeviceLost,
            _ => panic!("unexpected error: {:?}", err)
        }
    }
}

#[cfg(test)]
mod tests {
    use std::thread;
    use std::time::Duration;
    use super::*;
    use sync::Fence;

    #[test]
    fn empty_submit() {
        let (device, queue) = gfx_dev_and_queue!();
        let mut builder = SubmitCommandBufferBuilder::new();
        builder.submit(&queue).unwrap();
    }

    #[test]
    fn signal_fence() {
        unsafe {
            let (device, queue) = gfx_dev_and_queue!();

            let fence = Fence::new(device.clone()).unwrap();
            assert!(!fence.ready().unwrap());

            let mut builder = SubmitCommandBufferBuilder::new();
            builder.set_fence_signal(&fence);

            builder.submit(&queue).unwrap();
            fence.wait(Duration::from_secs(10)).unwrap();
            assert!(fence.ready().unwrap());
        }
    }

    #[test]
    fn has_fence() {
        unsafe {
            let (device, queue) = gfx_dev_and_queue!();

            let fence = Fence::new(device.clone()).unwrap();

            let mut builder = SubmitCommandBufferBuilder::new();
            assert!(!builder.has_fence());
            builder.set_fence_signal(&fence);
            assert!(builder.has_fence());
        }
    }

    #[test]
    #[should_panic(expected = "Can't merge two queue submits that both have a fence")]
    fn merge_both_have_fences() {
        unsafe {
            let (device, _) = gfx_dev_and_queue!();

            let fence1 = Fence::new(device.clone()).unwrap();
            let fence2 = Fence::new(device.clone()).unwrap();

            let mut builder1 = SubmitCommandBufferBuilder::new();
            builder1.set_fence_signal(&fence1);
            let mut builder2 = SubmitCommandBufferBuilder::new();
            builder2.set_fence_signal(&fence2);

            let _ = builder1.merge(builder2);
        }
    }
}
<commit_msg>Small TODO<commit_after>// Copyright (c) 2017 The vulkano developers
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>,
// at your option. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

use std::error;
use std::fmt;
use std::marker::PhantomData;
use std::ptr;
use smallvec::SmallVec;

use command_buffer::cb::UnsafeCommandBuffer;
use command_buffer::pool::CommandPool;
use device::Queue;
use sync::Fence;
use sync::PipelineStages;
use sync::Semaphore;

use check_errors;
use vk;
use Error;
use OomError;
use VulkanObject;
use VulkanPointers;
use SynchronizedVulkanObject;

/// Prototype for a submission that executes command buffers.
// TODO: example here
#[derive(Debug)]
pub struct SubmitCommandBufferBuilder<'a> {
    wait_semaphores: SmallVec<[vk::Semaphore; 16]>,
    dest_stages: SmallVec<[vk::PipelineStageFlags; 8]>,
    signal_semaphores: SmallVec<[vk::Semaphore; 16]>,
    command_buffers: SmallVec<[vk::CommandBuffer; 4]>,
    fence: vk::Fence,
    marker: PhantomData<&'a ()>,
}

impl<'a> SubmitCommandBufferBuilder<'a> {
    /// Builds a new empty `SubmitCommandBufferBuilder`.
    #[inline]
    pub fn new() -> SubmitCommandBufferBuilder<'a> {
        SubmitCommandBufferBuilder {
            wait_semaphores: SmallVec::new(),
            dest_stages: SmallVec::new(),
            signal_semaphores: SmallVec::new(),
            command_buffers: SmallVec::new(),
            fence: 0,
            marker: PhantomData,
        }
    }

    /// Returns true if this builder will signal a fence when submitted.
    ///
    /// # Example
    ///
    /// ```
    /// use vulkano::command_buffer::submit::SubmitCommandBufferBuilder;
    /// use vulkano::sync::Fence;
    /// # let device: std::sync::Arc<vulkano::device::Device> = return;
    ///
    /// unsafe {
    ///     let fence = Fence::new(device.clone()).unwrap();
    ///
    ///     let mut builder = SubmitCommandBufferBuilder::new();
    ///     assert!(!builder.has_fence());
    ///     builder.set_fence_signal(&fence);
    ///     assert!(builder.has_fence());
    /// }
    /// ```
    #[inline]
    pub fn has_fence(&self) -> bool {
        self.fence != 0
    }

    /// Adds an operation that signals a fence after this submission ends.
    ///
    /// # Example
    ///
    /// ```
    /// use std::time::Duration;
    /// use vulkano::command_buffer::submit::SubmitCommandBufferBuilder;
    /// use vulkano::sync::Fence;
    /// # let device: std::sync::Arc<vulkano::device::Device> = return;
    /// # let queue: std::sync::Arc<vulkano::device::Queue> = return;
    ///
    /// unsafe {
    ///     let fence = Fence::new(device.clone()).unwrap();
    ///
    ///     let mut builder = SubmitCommandBufferBuilder::new();
    ///     builder.set_fence_signal(&fence);
    ///
    ///     builder.submit(&queue).unwrap();
    ///
    ///     // We must not destroy the fence before it is signaled.
    ///     fence.wait(Duration::from_secs(5)).unwrap();
    /// }
    /// ```
    ///
    /// # Safety
    ///
    /// - The fence must not be signaled at the time when you call `submit()`.
    ///
    /// - If you use the fence for multiple submissions, only one at a time must be executed by the
    ///   GPU. In other words, you must submit one, wait for the fence to be signaled, then reset
    ///   the fence, and then only submit the second.
    ///
    /// - If you submit this builder, the fence must be kept alive until it is signaled by the GPU.
    ///   Destroying the fence earlier is an undefined behavior.
    ///
    /// - The fence, command buffers, and semaphores must all belong to the same device.
    ///
    #[inline]
    pub unsafe fn set_fence_signal(&mut self, fence: &'a Fence) {
        self.fence = fence.internal_object();
    }

    /// Adds a semaphore to be waited upon before the command buffers are executed.
    ///
    /// Only the given `stages` of the command buffers added afterwards will wait upon
    /// the semaphore. Other stages not included in `stages` can execute before waiting.
    ///
    /// # Safety
    ///
    /// - The stages must be supported by the device.
    ///
    /// - If you submit this builder, the semaphore must be kept alive until you are guaranteed
    ///   that the GPU has at least started executing the command buffers.
    ///
    /// - If you submit this builder, no other queue must be waiting on these semaphores. In other
    ///   words, each semaphore signal can only correspond to one semaphore wait.
    ///
    /// - If you submit this builder, the semaphores must be signaled when the queue execution
    ///   reaches this submission, or there must be one or more submissions in queues that are
    ///   going to signal these semaphores. In other words, you must not block the queue with
    ///   semaphores that can't get signaled.
    ///
    /// - The fence, command buffers, and semaphores must all belong to the same device.
    ///
    #[inline]
    pub unsafe fn add_wait_semaphore(&mut self, semaphore: &'a Semaphore, stages: PipelineStages) {
        debug_assert!(Into::<vk::PipelineStageFlagBits>::into(stages) != 0);
        // TODO: debug assert that the device supports the stages
        self.wait_semaphores.push(semaphore.internal_object());
        self.dest_stages.push(stages.into());
    }

    /// Adds a command buffer that is executed as part of this command.
    ///
    /// The command buffers are submitted in the order in which they are added.
    ///
    /// # Safety
    ///
    /// - If you submit this builder, the command buffer must be kept alive until you are
    ///   guaranteed that the GPU has finished executing it.
    ///
    /// - Any calls to vkCmdSetEvent, vkCmdResetEvent or vkCmdWaitEvents that have been recorded
    ///   into the command buffer must not reference any VkEvent that is referenced by any of
    ///   those commands that is pending execution on another queue.
    ///   TODO: rephrase ^ ?
    ///
    /// - The fence, command buffers, and semaphores must all belong to the same device.
    ///
    /// TODO: more here
    ///
    #[inline]
    pub unsafe fn add_command_buffer<P>(&mut self, command_buffer: &'a UnsafeCommandBuffer<P>)
        where P: CommandPool
    {
        self.command_buffers.push(command_buffer.internal_object());
    }

    /// Returns the number of semaphores to signal.
    ///
    /// In other words, this is the number of times `add_signal_semaphore` has been called.
    #[inline]
    pub fn num_signal_semaphores(&self) -> usize {
        self.signal_semaphores.len()
    }

    /// Adds a semaphore that is going to be signaled at the end of the submission.
    ///
    /// # Safety
    ///
    /// - If you submit this builder, the semaphore must be kept alive until you are guaranteed
    ///   that the GPU has finished executing this submission.
    ///
    /// - The semaphore must be in the unsignaled state when queue execution reaches this
    ///   submission.
    ///
    /// - The fence, command buffers, and semaphores must all belong to the same device.
    ///
    #[inline]
    pub unsafe fn add_signal_semaphore(&mut self, semaphore: &'a Semaphore) {
        self.signal_semaphores.push(semaphore.internal_object());
    }

    /// Submits the command buffer to the given queue.
    ///
    /// > **Note**: This is an expensive operation, so you may want to merge as many builders as
    /// > possible together and avoid submitting them one by one.
    ///
    pub fn submit(mut self, queue: &Queue) -> Result<(), SubmitCommandBufferError> {
        unsafe {
            let vk = queue.device().pointers();
            let queue = queue.internal_object_guard();

            debug_assert_eq!(self.wait_semaphores.len(), self.dest_stages.len());

            let batch = vk::SubmitInfo {
                sType: vk::STRUCTURE_TYPE_SUBMIT_INFO,
                pNext: ptr::null(),
                waitSemaphoreCount: self.wait_semaphores.len() as u32,
                pWaitSemaphores: self.wait_semaphores.as_ptr(),
                pWaitDstStageMask: self.dest_stages.as_ptr(),
                commandBufferCount: self.command_buffers.len() as u32,
                pCommandBuffers: self.command_buffers.as_ptr(),
                signalSemaphoreCount: self.signal_semaphores.len() as u32,
                pSignalSemaphores: self.signal_semaphores.as_ptr(),
            };

            try!(check_errors(vk.QueueSubmit(*queue, 1, &batch, self.fence)));
            Ok(())
        }
    }

    /// Merges this builder with another builder.
    ///
    /// # Panic
    ///
    /// Panics if both builders have a fence already set.
    // TODO: create multiple batches instead
    pub fn merge(mut self, other: Self) -> Self {
        assert!(self.fence == 0 || other.fence == 0,
               "Can't merge two queue submits that both have a fence");

        self.wait_semaphores.extend(other.wait_semaphores);
        self.dest_stages.extend(other.dest_stages);     // TODO: meh? will be solved if we submit multiple batches
        self.signal_semaphores.extend(other.signal_semaphores);
        self.command_buffers.extend(other.command_buffers);

        if self.fence == 0 {
            self.fence = other.fence;
        }

        self
    }
}

/// Error that can happen when submitting the prototype.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
#[repr(u32)]
pub enum SubmitCommandBufferError {
    /// Not enough memory.
    OomError(OomError),

    /// The connection to the device has been lost.
    DeviceLost,
}

impl error::Error for SubmitCommandBufferError {
    #[inline]
    fn description(&self) -> &str {
        match *self {
            SubmitCommandBufferError::OomError(_) => "not enough memory",
            SubmitCommandBufferError::DeviceLost => "the connection to the device has been lost",
        }
    }

    #[inline]
    fn cause(&self) -> Option<&error::Error> {
        match *self {
            SubmitCommandBufferError::OomError(ref err) => Some(err),
            _ => None
        }
    }
}

impl fmt::Display for SubmitCommandBufferError {
    #[inline]
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{}", error::Error::description(self))
    }
}

impl From<Error> for SubmitCommandBufferError {
    #[inline]
    fn from(err: Error) -> SubmitCommandBufferError {
        match err {
            err @ Error::OutOfHostMemory => SubmitCommandBufferError::OomError(OomError::from(err)),
            err @ Error::OutOfDeviceMemory => SubmitCommandBufferError::OomError(OomError::from(err)),
            Error::DeviceLost => SubmitCommandBufferError::DeviceLost,
            _ => panic!("unexpected error: {:?}", err)
        }
    }
}

#[cfg(test)]
mod tests {
    use std::thread;
    use std::time::Duration;
    use super::*;
    use sync::Fence;

    #[test]
    fn empty_submit() {
        let (device, queue) = gfx_dev_and_queue!();
        let mut builder = SubmitCommandBufferBuilder::new();
        builder.submit(&queue).unwrap();
    }

    #[test]
    fn signal_fence() {
        unsafe {
            let (device, queue) = gfx_dev_and_queue!();

            let fence = Fence::new(device.clone()).unwrap();
            assert!(!fence.ready().unwrap());

            let mut builder = SubmitCommandBufferBuilder::new();
            builder.set_fence_signal(&fence);

            builder.submit(&queue).unwrap();
            fence.wait(Duration::from_secs(10)).unwrap();
            assert!(fence.ready().unwrap());
        }
    }

    #[test]
    fn has_fence() {
        unsafe {
            let (device, queue) = gfx_dev_and_queue!();

            let fence = Fence::new(device.clone()).unwrap();

            let mut builder = SubmitCommandBufferBuilder::new();
            assert!(!builder.has_fence());
            builder.set_fence_signal(&fence);
            assert!(builder.has_fence());
        }
    }

    #[test]
    #[should_panic(expected = "Can't merge two queue submits that both have a fence")]
    fn merge_both_have_fences() {
        unsafe {
            let (device, _) = gfx_dev_and_queue!();

            let fence1 = Fence::new(device.clone()).unwrap();
            let fence2 = Fence::new(device.clone()).unwrap();

            let mut builder1 = SubmitCommandBufferBuilder::new();
            builder1.set_fence_signal(&fence1);
            let mut builder2 = SubmitCommandBufferBuilder::new();
            builder2.set_fence_signal(&fence2);

            let _ = builder1.merge(builder2);
        }
    }
}
<|endoftext|>


<commit_before>extern crate bytes;
#[macro_use]
extern crate combine;

extern crate futures;
extern crate partial_io;
extern crate tokio_io;

use std::rc::Rc;
use std::cell::Cell;
use std::str;
use std::io::Cursor;

use bytes::BytesMut;

use tokio_io::codec::Decoder;

use combine::primitives::{ParseError, PointerOffset, StreamOnce};
use combine::range::{range, recognize, take};
use combine::{skip_many, Parser};
use combine::easy::{self, Error as CombineError, Errors};
use combine::byte::digit;

fn combine_decode<'a, P, R>(
    mut parser: P,
    src: &'a [u8],
    partial_state: &mut P::PartialState,
) -> Result<Option<(R, usize)>, Errors<usize, u8, String>>
where
    P: Parser<Input = easy::Stream<&'a [u8]>, Output = R>,
{
    let mut input = easy::Stream(&src[..]);
    match parser.parse_with_state(&mut input, partial_state) {
        Ok(message) => Ok(Some((message, src.len() - input.0.len()))),
        Err(err) => {
            return if err.errors
                .iter()
                .any(|err| *err == CombineError::end_of_input())
            {
                Ok(None)
            } else {
                Err(err.map_range(|r| {
                    str::from_utf8(r)
                        .ok()
                        .map_or_else(|| format!("{:?}", r), |s| s.to_string())
                }).map_position(|p| p.translate_position(&src[..])))
            }
        }
    }
}

macro_rules! decode {
    ($parser: expr, $src: expr, $state: expr) => {
        {
            let (output, removed_len) = {
                match combine_decode($parser, &$src[..], $state)? {
                    None => return Ok(None),
                    Some(x) => x,
                }
            };
            $src.split_to(removed_len);
            output
        }
    };
}

pub struct LanguageServerDecoder {
    state: decode_parser::PartialState,
    content_length_parses: Rc<Cell<i32>>,
}

impl LanguageServerDecoder {
    pub fn new() -> LanguageServerDecoder {
        LanguageServerDecoder {
            state: Default::default(),
            content_length_parses: Rc::new(Cell::new(0)),
        }
    }
}

parser! {
    fn decode_parser['a](content_length_parses: Rc<Cell<i32>>)(easy::Stream<&'a [u8]>) -> Vec<u8>
    where [ <easy::Stream<&'a [u8]> as StreamOnce>::Error:
                ParseError<u8, &'a [u8], PointerOffset, StreamError = easy::Error<u8, &'a [u8]>>
        ]
    {
        (
            skip_many(range(&b"\r\n"[..])),
            range(&b"Content-Length: "[..]).map(|_| content_length_parses.set(content_length_parses.get() + 1)),
            recognize(digit()),
            range(&b"\r\n\r\n"[..]),
        ).map(|t| t.2)
            .and_then(|digits: &[u8]| unsafe {
                str::from_utf8_unchecked(digits).parse::<usize>()
            })
            .then(|message_length| take(message_length).map(|bytes: &[u8]| bytes.to_owned()))
    }
}

impl Decoder for LanguageServerDecoder {
    type Item = String;
    type Error = Box<::std::error::Error + Send + Sync>;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        let message = decode!(
            decode_parser(self.content_length_parses.clone()),
            src,
            &mut self.state
        );
        Ok(Some(String::from_utf8(message)?))
    }
}

fn main() {
    use futures::{Future, Stream};
    use tokio_io::codec::FramedRead;
    use partial_io::{PartialAsyncRead, PartialOp};

    let input = "Content-Length: 6\r\n\
                 \r\n\
                 123456\r\n\
                 Content-Length: 4\r\n\
                 \r\n\
                 true";

    let seq = vec![
        PartialOp::Limited(2),
        PartialOp::Limited(1),
        PartialOp::Limited(20),
    ];
    let ref mut reader = Cursor::new(input.as_bytes());
    let partial_reader = PartialAsyncRead::new(reader, seq);

    let decoder = LanguageServerDecoder::new();
    let content_length_parses = decoder.content_length_parses.clone();

    let result = FramedRead::new(partial_reader, decoder).collect().wait();

    assert!(result.as_ref().is_ok(), "{}", result.unwrap_err());
    assert_eq!(
        result.unwrap(),
        vec!["123456".to_string(), "true".to_string()]
    );

    assert_eq!(content_length_parses.get(), 2);
}
<commit_msg>Parse multiple digits in the tokio example<commit_after>extern crate bytes;
#[macro_use]
extern crate combine;

extern crate futures;
extern crate partial_io;
extern crate tokio_io;

use std::rc::Rc;
use std::cell::Cell;
use std::str;
use std::io::Cursor;

use bytes::BytesMut;

use tokio_io::codec::Decoder;

use combine::primitives::{ParseError, PointerOffset, StreamOnce};
use combine::range::{range, recognize, take};
use combine::{skip_many, Parser, skip_many1};
use combine::easy::{self, Error as CombineError, Errors};
use combine::byte::digit;

fn combine_decode<'a, P, R>(
    mut parser: P,
    src: &'a [u8],
    partial_state: &mut P::PartialState,
) -> Result<Option<(R, usize)>, Errors<usize, u8, String>>
where
    P: Parser<Input = easy::Stream<&'a [u8]>, Output = R>,
{
    let mut input = easy::Stream(&src[..]);
    match parser.parse_with_state(&mut input, partial_state) {
        Ok(message) => Ok(Some((message, src.len() - input.0.len()))),
        Err(err) => {
            return if err.errors
                .iter()
                .any(|err| *err == CombineError::end_of_input())
            {
                Ok(None)
            } else {
                Err(err.map_range(|r| {
                    str::from_utf8(r)
                        .ok()
                        .map_or_else(|| format!("{:?}", r), |s| s.to_string())
                }).map_position(|p| p.translate_position(&src[..])))
            }
        }
    }
}

macro_rules! decode {
    ($parser: expr, $src: expr, $state: expr) => {
        {
            let (output, removed_len) = {
                match combine_decode($parser, &$src[..], $state)? {
                    None => return Ok(None),
                    Some(x) => x,
                }
            };
            $src.split_to(removed_len);
            output
        }
    };
}

pub struct LanguageServerDecoder {
    state: decode_parser::PartialState,
    content_length_parses: Rc<Cell<i32>>,
}

impl LanguageServerDecoder {
    pub fn new() -> LanguageServerDecoder {
        LanguageServerDecoder {
            state: Default::default(),
            content_length_parses: Rc::new(Cell::new(0)),
        }
    }
}

parser! {
    fn decode_parser['a](content_length_parses: Rc<Cell<i32>>)(easy::Stream<&'a [u8]>) -> Vec<u8>
    where [ <easy::Stream<&'a [u8]> as StreamOnce>::Error:
                ParseError<u8, &'a [u8], PointerOffset, StreamError = easy::Error<u8, &'a [u8]>>
        ]
    {
        (
            skip_many(range(&b"\r\n"[..])),
            range(&b"Content-Length: "[..]).map(|_| content_length_parses.set(content_length_parses.get() + 1)),
            recognize(skip_many1(digit())),
            range(&b"\r\n\r\n"[..]),
        ).map(|t| t.2)
            .and_then(|digits: &[u8]| unsafe {
                str::from_utf8_unchecked(digits).parse::<usize>()
            })
            .then(|message_length| take(message_length).map(|bytes: &[u8]| bytes.to_owned()))
    }
}

impl Decoder for LanguageServerDecoder {
    type Item = String;
    type Error = Box<::std::error::Error + Send + Sync>;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        let message = decode!(
            decode_parser(self.content_length_parses.clone()),
            src,
            &mut self.state
        );
        Ok(Some(String::from_utf8(message)?))
    }
}

fn main() {
    use futures::{Future, Stream};
    use tokio_io::codec::FramedRead;
    use partial_io::{PartialAsyncRead, PartialOp};

    let input = "Content-Length: 6\r\n\
                 \r\n\
                 123456\r\n\
                 Content-Length: 4\r\n\
                 \r\n\
                 true";

    let seq = vec![
        PartialOp::Limited(2),
        PartialOp::Limited(1),
        PartialOp::Limited(20),
    ];
    let ref mut reader = Cursor::new(input.as_bytes());
    let partial_reader = PartialAsyncRead::new(reader, seq);

    let decoder = LanguageServerDecoder::new();
    let content_length_parses = decoder.content_length_parses.clone();

    let result = FramedRead::new(partial_reader, decoder).collect().wait();

    assert!(result.as_ref().is_ok(), "{}", result.unwrap_err());
    assert_eq!(
        result.unwrap(),
        vec!["123456".to_string(), "true".to_string()]
    );

    assert_eq!(content_length_parses.get(), 2);
}
<|endoftext|>


<commit_before>use std::fs::File;
use std::io::prelude::*;
use std::env;
use std::process;

fn main() {
    let args: Vec<_> = env::args().collect();
    if args.len() < 2 {
        println!("No file was passed in to be parsed");
        process::exit(1);
    }

    let mut f = File::open(&args[1]).unwrap();
    let mut instructions = String::new();
    f.read_to_string(&mut instructions).unwrap();
    let final_floor = process_instructions(instructions);

    println!("Santa ends up on floor {}", final_floor);
}

fn process_instructions(ins: String) -> i32{
    let mut floor = 0;

    for c in ins.chars() {
        match c {
            '(' => { floor += 1; },
            ')' => { floor -= 1; },
            _ => {}
        }
    }
    return return floor;
}
<commit_msg>Fixed some styling issues with returns<commit_after>use std::fs::File;
use std::io::prelude::*;
use std::env;
use std::process;

fn main() {
    let args: Vec<_> = env::args().collect();
    if args.len() < 2 {
        println!("No file was passed in to be parsed");
        process::exit(1);
    }

    let mut f = File::open(&args[1]).unwrap();
    let mut instructions = String::new();
    f.read_to_string(&mut instructions).unwrap();
    let final_floor = process_instructions(instructions);

    println!("Santa ends up on floor {}", final_floor);
}

fn process_instructions(ins: String) -> i32{
    let mut floor = 0;

    for c in ins.chars() {
        match c {
            '(' => { floor += 1; },
            ')' => { floor -= 1; },
            _ => {}
        }
    }
    floor
}
<|endoftext|>


<commit_before>use std::convert::TryInto;

use crate::ImageResult;
use super::lossless::DecoderError;
use super::lossless::BitReader;

/// Rudimentary utility for reading Canonical Huffman Codes.
/// Based off https://github.com/webmproject/libwebp/blob/7f8472a610b61ec780ef0a8873cd954ac512a505/src/utils/huffman.c
///

const MAX_ALLOWED_CODE_LENGTH: usize = 15;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum HuffmanTreeNode {
    Branch(usize), //offset in vector to children
    Leaf(u16), //symbol stored in leaf
    Empty,
}

/// Huffman tree
#[derive(Clone, Debug, Default)]
pub(crate) struct HuffmanTree {
    tree: Vec<HuffmanTreeNode>,
    max_nodes: usize,
    num_nodes: usize,
}

impl HuffmanTree {

    fn is_full(&self) -> bool {
        self.num_nodes == self.max_nodes
    }

    /// Turns a node from empty into a branch and assigns its children
    fn assign_children(&mut self, node_index: usize) -> usize {
        let offset_index = self.num_nodes - node_index;
        self.tree[node_index] = HuffmanTreeNode::Branch(offset_index);
        self.num_nodes += 2;

        offset_index
    }

    /// Init a huffman tree
    fn init(num_leaves: usize) -> ImageResult<HuffmanTree> {
        if num_leaves == 0 {
            return Err(DecoderError::HuffmanError.into());
        }

        let max_nodes = 2 * num_leaves - 1;
        let tree = vec![HuffmanTreeNode::Empty; max_nodes];
        let num_nodes = 1;

        let tree = HuffmanTree {
            tree,
            max_nodes,
            num_nodes,
        };

        Ok(tree)
    }

    /// Converts code lengths to codes
    fn code_lengths_to_codes(code_lengths: &[u16]) -> ImageResult<Vec<Option<u16>>> {

        let max_code_length = *code_lengths.iter().reduce(|a, b| {
            if a >= b {a} else {b}
        }).unwrap();

        if max_code_length > MAX_ALLOWED_CODE_LENGTH.try_into().unwrap() {
            return Err(DecoderError::HuffmanError.into());
        }

        let mut code_length_hist = vec![0; MAX_ALLOWED_CODE_LENGTH + 1];

        for &length in code_lengths.iter() {
            code_length_hist[usize::from(length)] += 1;
        }

        code_length_hist[0] = 0;

        let mut curr_code = 0;
        let mut next_codes = vec![None; MAX_ALLOWED_CODE_LENGTH + 1];

        for code_len in 1..=usize::from(max_code_length) {
            curr_code = (curr_code + code_length_hist[code_len - 1]) << 1;
            next_codes[code_len] = Some(curr_code);
        }

        let mut huff_codes = vec![None; code_lengths.len()];

        for (symbol, &length) in code_lengths.iter().enumerate() {
            let length = usize::from(length);
            if length > 0 {
                huff_codes[symbol] = next_codes[length];
                if let Some(value) = next_codes[length].as_mut() {
                    *value += 1;
                }
            } else {
                huff_codes[symbol] = None;
            }
        }

        Ok(huff_codes)
    }

    /// Adds a symbol to a huffman tree
    fn add_symbol(&mut self, symbol: u16, code: u16, code_length: u16) -> ImageResult<()> {
        let mut node_index = 0;
        let code = usize::from(code);

        let mut length = code_length;

        while length > 0 {
            length -= 1;
            if node_index >= self.max_nodes {
                return Err(DecoderError::HuffmanError.into());
            }

            let node = self.tree[node_index];

            let offset = match node {
                HuffmanTreeNode::Empty => {
                    if self.is_full() {
                        return Err(DecoderError::HuffmanError.into());
                    }
                    self.assign_children(node_index)
                }
                HuffmanTreeNode::Leaf(_) => return Err(DecoderError::HuffmanError.into()),
                HuffmanTreeNode::Branch(offset) => offset,
            };

            node_index += offset + ((code >> length) & 1);
        }

        match self.tree[node_index] {
            HuffmanTreeNode::Empty => self.tree[node_index] = HuffmanTreeNode::Leaf(symbol),
            HuffmanTreeNode::Leaf(_) => return Err(DecoderError::HuffmanError.into()),
            HuffmanTreeNode::Branch(_offset) => return Err(DecoderError::HuffmanError.into()),
        }

        Ok(())
    }

    /// Builds a tree implicitly, just from code lengths
    pub(crate) fn build_implicit(code_lengths: Vec<u16>) -> ImageResult<HuffmanTree> {
        let mut num_symbols = 0;
        let mut root_symbol = 0;

        for symbol in 0..code_lengths.len() {
            if code_lengths[symbol] > 0 {
                num_symbols += 1;
                root_symbol = symbol.try_into().unwrap();
            }
        }

        let mut tree = HuffmanTree::init(num_symbols)?;

        if num_symbols == 1 {
            tree.add_symbol(root_symbol, 0, 0)?;
        } else {

            let codes = HuffmanTree::code_lengths_to_codes(&code_lengths)?;

            for (symbol, &length) in code_lengths.iter().enumerate() {
                if length > 0 && codes[symbol].is_some() {
                    tree.add_symbol(symbol.try_into().unwrap(), codes[symbol].unwrap(), length)?;
                }
            }
        }

        Ok(tree)
    }

    /// Builds a tree explicitly from lengths, codes and symbols
    pub(crate) fn build_explicit(code_lengths: Vec<u16>, codes: Vec<u16>, symbols: Vec<u16>) -> ImageResult<HuffmanTree> {
        let mut tree = HuffmanTree::init(symbols.len())?;

        for i in 0..symbols.len() {
            tree.add_symbol(symbols[i], codes[i], code_lengths[i])?;
        }

        Ok(tree)
    }

    /// Reads a symbol using the bitstream
    pub(crate) fn read_symbol(&self, bit_reader: &mut BitReader) -> ImageResult<u16> {
        let mut index = 0;
        let mut node = self.tree[index];

        while let HuffmanTreeNode::Branch(children_offset) = node {
            index += children_offset + bit_reader.read_bits::<usize>(1)?;
            node = self.tree[index];
        }

        let symbol = match node {
            HuffmanTreeNode::Branch(_) => unreachable!(),
            HuffmanTreeNode::Empty => return Err(DecoderError::HuffmanError.into()),
            HuffmanTreeNode::Leaf(symbol) => symbol,
        };

        Ok(symbol)
    }
}

<commit_msg>Update src/codecs/webp/huffman.rs<commit_after>use std::convert::TryInto;

use crate::ImageResult;
use super::lossless::DecoderError;
use super::lossless::BitReader;

/// Rudimentary utility for reading Canonical Huffman Codes.
/// Based off https://github.com/webmproject/libwebp/blob/7f8472a610b61ec780ef0a8873cd954ac512a505/src/utils/huffman.c
///

const MAX_ALLOWED_CODE_LENGTH: usize = 15;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum HuffmanTreeNode {
    Branch(usize), //offset in vector to children
    Leaf(u16), //symbol stored in leaf
    Empty,
}

/// Huffman tree
#[derive(Clone, Debug, Default)]
pub(crate) struct HuffmanTree {
    tree: Vec<HuffmanTreeNode>,
    max_nodes: usize,
    num_nodes: usize,
}

impl HuffmanTree {

    fn is_full(&self) -> bool {
        self.num_nodes == self.max_nodes
    }

    /// Turns a node from empty into a branch and assigns its children
    fn assign_children(&mut self, node_index: usize) -> usize {
        let offset_index = self.num_nodes - node_index;
        self.tree[node_index] = HuffmanTreeNode::Branch(offset_index);
        self.num_nodes += 2;

        offset_index
    }

    /// Init a huffman tree
    fn init(num_leaves: usize) -> ImageResult<HuffmanTree> {
        if num_leaves == 0 {
            return Err(DecoderError::HuffmanError.into());
        }

        let max_nodes = 2 * num_leaves - 1;
        let tree = vec![HuffmanTreeNode::Empty; max_nodes];
        let num_nodes = 1;

        let tree = HuffmanTree {
            tree,
            max_nodes,
            num_nodes,
        };

        Ok(tree)
    }

    /// Converts code lengths to codes
    fn code_lengths_to_codes(code_lengths: &[u16]) -> ImageResult<Vec<Option<u16>>> {

        let max_code_length = *code_lengths.iter().reduce(|a, b| {
            if a >= b {a} else {b}
        }).unwrap();

        if max_code_length > MAX_ALLOWED_CODE_LENGTH.try_into().unwrap() {
            return Err(DecoderError::HuffmanError.into());
        }

        let mut code_length_hist = vec![0; MAX_ALLOWED_CODE_LENGTH + 1];

        for &length in code_lengths.iter() {
            code_length_hist[usize::from(length)] += 1;
        }

        code_length_hist[0] = 0;

        let mut curr_code = 0;
        let mut next_codes = vec![None; MAX_ALLOWED_CODE_LENGTH + 1];

        for code_len in 1..=usize::from(max_code_length) {
            curr_code = (curr_code + code_length_hist[code_len - 1]) << 1;
            next_codes[code_len] = Some(curr_code);
        }

        let mut huff_codes = vec![None; code_lengths.len()];

        for (symbol, &length) in code_lengths.iter().enumerate() {
            let length = usize::from(length);
            if length > 0 {
                huff_codes[symbol] = next_codes[length];
                if let Some(value) = next_codes[length].as_mut() {
                    *value += 1;
                }
            } else {
                huff_codes[symbol] = None;
            }
        }

        Ok(huff_codes)
    }

    /// Adds a symbol to a huffman tree
    fn add_symbol(&mut self, symbol: u16, code: u16, code_length: u16) -> ImageResult<()> {
        let mut node_index = 0;
        let code = usize::from(code);

        for length in (0..code_length).rev() {
            if node_index >= self.max_nodes {
                return Err(DecoderError::HuffmanError.into());
            }

            let node = self.tree[node_index];

            let offset = match node {
                HuffmanTreeNode::Empty => {
                    if self.is_full() {
                        return Err(DecoderError::HuffmanError.into());
                    }
                    self.assign_children(node_index)
                }
                HuffmanTreeNode::Leaf(_) => return Err(DecoderError::HuffmanError.into()),
                HuffmanTreeNode::Branch(offset) => offset,
            };

            node_index += offset + ((code >> length) & 1);
        }

        match self.tree[node_index] {
            HuffmanTreeNode::Empty => self.tree[node_index] = HuffmanTreeNode::Leaf(symbol),
            HuffmanTreeNode::Leaf(_) => return Err(DecoderError::HuffmanError.into()),
            HuffmanTreeNode::Branch(_offset) => return Err(DecoderError::HuffmanError.into()),
        }

        Ok(())
    }

    /// Builds a tree implicitly, just from code lengths
    pub(crate) fn build_implicit(code_lengths: Vec<u16>) -> ImageResult<HuffmanTree> {
        let mut num_symbols = 0;
        let mut root_symbol = 0;

        for symbol in 0..code_lengths.len() {
            if code_lengths[symbol] > 0 {
                num_symbols += 1;
                root_symbol = symbol.try_into().unwrap();
            }
        }

        let mut tree = HuffmanTree::init(num_symbols)?;

        if num_symbols == 1 {
            tree.add_symbol(root_symbol, 0, 0)?;
        } else {

            let codes = HuffmanTree::code_lengths_to_codes(&code_lengths)?;

            for (symbol, &length) in code_lengths.iter().enumerate() {
                if length > 0 && codes[symbol].is_some() {
                    tree.add_symbol(symbol.try_into().unwrap(), codes[symbol].unwrap(), length)?;
                }
            }
        }

        Ok(tree)
    }

    /// Builds a tree explicitly from lengths, codes and symbols
    pub(crate) fn build_explicit(code_lengths: Vec<u16>, codes: Vec<u16>, symbols: Vec<u16>) -> ImageResult<HuffmanTree> {
        let mut tree = HuffmanTree::init(symbols.len())?;

        for i in 0..symbols.len() {
            tree.add_symbol(symbols[i], codes[i], code_lengths[i])?;
        }

        Ok(tree)
    }

    /// Reads a symbol using the bitstream
    pub(crate) fn read_symbol(&self, bit_reader: &mut BitReader) -> ImageResult<u16> {
        let mut index = 0;
        let mut node = self.tree[index];

        while let HuffmanTreeNode::Branch(children_offset) = node {
            index += children_offset + bit_reader.read_bits::<usize>(1)?;
            node = self.tree[index];
        }

        let symbol = match node {
            HuffmanTreeNode::Branch(_) => unreachable!(),
            HuffmanTreeNode::Empty => return Err(DecoderError::HuffmanError.into()),
            HuffmanTreeNode::Leaf(symbol) => symbol,
        };

        Ok(symbol)
    }
}

<|endoftext|>


<commit_before>use super::distance_proxy::DistanceProxy;
use super::super::common::Sweep;

pub struct TOIInput {
    pub proxy_a: DistanceProxy,
    pub proxy_b: DistanceProxy,
    pub sweep_a: Sweep,
    pub sweep_b: Sweep,
    pub t_max: f32
}

impl TOIInput {
    pub fn new() -> TOIInput {
        TOIInput {
            proxy_a: DistanceProxy::new(),
            proxy_b: DistanceProxy::new(),
            sweep_a: Sweep::new(),
            sweep_b: Sweep::new(),
            t_max: 0.0
        }
    }
}
<commit_msg>Adding comment for t_max of TOIInput<commit_after>use super::distance_proxy::DistanceProxy;
use super::super::common::Sweep;

pub struct TOIInput {
    pub proxy_a: DistanceProxy,
    pub proxy_b: DistanceProxy,
    pub sweep_a: Sweep,
    pub sweep_b: Sweep,
    // defines sweep interval [0, tMax]
    pub t_max: f32
}

impl TOIInput {
    pub fn new() -> TOIInput {
        TOIInput {
            proxy_a: DistanceProxy::new(),
            proxy_b: DistanceProxy::new(),
            sweep_a: Sweep::new(),
            sweep_b: Sweep::new(),
            t_max: 0.0
        }
    }
}
<|endoftext|>