<commit_before>//
// Copyright 2018 The Project Oak Authors
//
// 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.
//

//! Functionality to help Oak Nodes interact with gRPC.

pub use crate::proto::code::Code;
use crate::{proto, OakError, OakStatus, ReadHandle};
use log::{info, warn};
use proto::grpc_encap::{GrpcRequest, GrpcResponse};
use protobuf::{Message, ProtobufEnum};

mod invocation;
pub use invocation::Invocation;

/// Result type that uses a [`proto::status::Status`] type for error values.
pub type Result<T> = std::result::Result<T, proto::status::Status>;

/// Helper to create a gRPC status object.
pub fn build_status(code: Code, msg: &str) -> proto::status::Status {
    let mut status = proto::status::Status::new();
    status.set_code(code.value());
    status.set_message(msg.to_string());
    status
}

impl crate::io::Encodable for GrpcRequest {
    fn encode(&self) -> std::result::Result<crate::io::Message, crate::OakError> {
        let bytes = self.write_to_bytes()?;
        let handles = Vec::new();
        Ok(crate::io::Message { bytes, handles })
    }
}

impl crate::io::Decodable for GrpcRequest {
    fn decode(message: &crate::io::Message) -> std::result::Result<Self, crate::OakError> {
        let value = protobuf::parse_from_bytes(&message.bytes)?;
        Ok(value)
    }
}

/// Channel-holding object that encapsulates response messages into
/// `GrpcResponse` wrapper messages and writes serialized versions to a send
/// channel.
pub struct ChannelResponseWriter {
    sender: crate::io::Sender<GrpcResponse>,
}

impl crate::io::Encodable for GrpcResponse {
    fn encode(&self) -> std::result::Result<crate::io::Message, crate::OakError> {
        let bytes = self.write_to_bytes()?;
        let handles = Vec::new();
        Ok(crate::io::Message { bytes, handles })
    }
}

impl crate::io::Decodable for GrpcResponse {
    fn decode(message: &crate::io::Message) -> std::result::Result<Self, crate::OakError> {
        let value = protobuf::parse_from_bytes(&message.bytes)?;
        Ok(value)
    }
}

/// Indicate whether a write method should leave the current gRPC method
/// invocation open or close it.
#[derive(PartialEq, Clone, Debug)]
pub enum WriteMode {
    KeepOpen,
    Close,
}

impl ChannelResponseWriter {
    pub fn new(sender: crate::io::Sender<GrpcResponse>) -> Self {
        ChannelResponseWriter { sender }
    }

    /// Retrieve the Oak handle underlying the writer object.
    pub fn handle(self) -> crate::WriteHandle {
        self.sender.handle
    }

    /// Write out a gRPC response and optionally close out the method
    /// invocation.  Any errors from the channel are silently dropped.
    pub fn write<T: protobuf::Message>(
        &mut self,
        rsp: &T,
        mode: WriteMode,
    ) -> std::result::Result<(), OakError> {
        // Put the serialized response into a GrpcResponse message wrapper and
        // serialize it into the channel.
        let mut grpc_rsp = GrpcResponse::new();
        let mut any = protobuf::well_known_types::Any::new();
        rsp.write_to_writer(&mut any.value)?;
        grpc_rsp.set_rsp_msg(any);
        grpc_rsp.set_last(match mode {
            WriteMode::KeepOpen => false,
            WriteMode::Close => true,
        });
        self.sender.send(&grpc_rsp)?;
        if mode == WriteMode::Close {
            self.sender.close()?;
        }
        Ok(())
    }

    /// Write an empty gRPC response and optionally close out the method
    /// invocation. Any errors from the channel are silently dropped.
    pub fn write_empty(&mut self, mode: WriteMode) -> std::result::Result<(), OakError> {
        let mut grpc_rsp = GrpcResponse::new();
        grpc_rsp.set_rsp_msg(protobuf::well_known_types::Any::new());
        grpc_rsp.set_last(match mode {
            WriteMode::KeepOpen => false,
            WriteMode::Close => true,
        });
        self.sender.send(&grpc_rsp)?;
        if mode == WriteMode::Close {
            self.sender.close()?;
        }
        Ok(())
    }

    /// Close out the gRPC method invocation with the given final result.
    pub fn close(&mut self, result: Result<()>) -> std::result::Result<(), OakError> {
        // Build a final GrpcResponse message wrapper and serialize it into the
        // channel.
        let mut grpc_rsp = GrpcResponse::new();
        grpc_rsp.set_last(true);
        if let Err(status) = result {
            grpc_rsp.set_status(status);
        }
        self.sender.send(&grpc_rsp)?;
        self.sender.close()?;
        Ok(())
    }
}

/// Trait for Oak Nodes that act as a gRPC services.
///
/// An `OakNode` instance is normally passed to [`event_loop`], to allow
/// repeated invocation of its `invoke()` method.
pub trait OakNode {
    /// Construct the (single) instance of the node.
    ///
    /// This method may choose to initialize logging by invoking
    /// [`oak_log::init()`].
    ///
    /// [`oak_log::init()`]: ../../oak_log/fn.init.html
    fn new() -> Self
    where
        Self: Sized;

    /// Process a single gRPC method invocation.
    ///
    /// The method name is provided by `method` and the incoming serialized gRPC
    /// request is held in `req`.  Response messages should be written using
    /// `writer.write`, followed by `writer.close`.
    fn invoke(&mut self, method: &str, req: &[u8], writer: ChannelResponseWriter);
}

/// Encapsulate a protocol buffer message in a GrpcRequest wrapper using the
/// given method name.
pub fn encap_request<T: protobuf::Message>(req: T, method_name: &str) -> Option<GrpcRequest> {
    // Put the request in a GrpcRequest wrapper and serialize it.
    let mut grpc_req = GrpcRequest::new();
    grpc_req.set_method_name(method_name.to_string());
    let mut any = protobuf::well_known_types::Any::new();
    if let Err(e) = req.write_to_writer(&mut any.value) {
        warn!("failed to serialize gRPC request: {}", e);
        return None;
    };
    grpc_req.set_req_msg(any);
    grpc_req.set_last(true);
    Some(grpc_req)
}

/// Perform a gRPC event loop for a Node.
///
/// Invoking the given `node`'s [`invoke`] method for each incoming request that
/// arrives on the inbound channel as a serialized [`GrpcRequest`] message,
/// giving the [`invoke`] method the outbound channel for encapsulated responses
/// to be written to.
///
/// [`invoke`]: OakNode::invoke
/// [`GrpcRequest`]: crate::proto::grpc_encap::GrpcRequest
pub fn event_loop<T: OakNode>(
    mut node: T,
    grpc_in_handle: ReadHandle,
) -> std::result::Result<(), crate::OakStatus> {
    info!("start event loop for node with handle {:?}", grpc_in_handle);
    if !grpc_in_handle.handle.is_valid() {
        return Err(OakStatus::ERR_CHANNEL_CLOSED);
    }
    let invocation_receiver = crate::io::Receiver::new(grpc_in_handle);
    loop {
        // Explicitly call `wait` and then `try_receive` here because calling `receive` would hide
        // any `OakStatus::ERR_TERMINATED` errors that may occur in the `wait` phase, since they
        // would be wrapped in a `OakError` value that would then always be unwrapped and panic.
        invocation_receiver.wait()?;
        let invocation: Invocation = invocation_receiver.try_receive().unwrap_or_else(|err| {
            panic!("could not receive gRPC invocation: {}", err);
        });

        // Read a single encapsulated request message from the read half.
        let req: GrpcRequest = invocation.request_receiver.receive().unwrap_or_else(|err| {
            panic!("could not read gRPC request: {:?}", err);
        });
        // Since we are expecting a single message, close the channel immediately.
        // This will change when we implement client streaming (#97).
        invocation.request_receiver.close().unwrap_or_else(|err| {
            panic!("could not close gRPC request channel: {:?}", err);
        });
        if !req.last {
            // TODO(#97): Implement client streaming.
            panic!("Support for streaming requests not yet implemented");
        }
        node.invoke(
            &req.method_name,
            req.get_req_msg().value.as_slice(),
            ChannelResponseWriter::new(invocation.response_sender),
        );
    }
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// request/response pair.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_req_rsp<C, R, Q>(mut node_fn: C, req: &[u8], mut writer: ChannelResponseWriter)
where
    C: FnMut(R) -> Result<Q>,
    R: protobuf::Message,
    Q: protobuf::Message,
{
    let r: R = protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message");
    let result = match node_fn(r) {
        Ok(rsp) => writer.write(&rsp, WriteMode::Close),
        Err(status) => writer.close(Err(status)),
    };
    if let Err(e) = result {
        panic!("Failed to process response: {:?}", e)
    }
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// request and streams responses.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_req_stream<C, R>(mut node_fn: C, req: &[u8], writer: ChannelResponseWriter)
where
    C: FnMut(R, ChannelResponseWriter),
    R: protobuf::Message,
{
    let r: R = protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message");
    node_fn(r, writer)
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// stream of requests to produce a response.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_stream_rsp<C, R, Q>(mut node_fn: C, req: &[u8], mut writer: ChannelResponseWriter)
where
    C: FnMut(Vec<R>) -> Result<Q>,
    R: protobuf::Message,
    Q: protobuf::Message,
{
    // TODO(#97): better client-side streaming
    let rr: Vec<R> =
        vec![protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message")];
    let result = match node_fn(rr) {
        Ok(rsp) => writer.write(&rsp, WriteMode::Close),
        Err(status) => writer.close(Err(status)),
    };
    if let Err(e) = result {
        panic!("Failed to process response: {:?}", e)
    }
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// stream of requests to produce a stream of responses.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_stream_stream<C, R>(mut node_fn: C, req: &[u8], writer: ChannelResponseWriter)
where
    C: FnMut(Vec<R>, ChannelResponseWriter),
    R: protobuf::Message,
{
    // TODO(#97): better client-side streaming
    let rr: Vec<R> =
        vec![protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message")];
    node_fn(rr, writer)
}
<commit_msg>sdk: add oak::grpc::invoke_grpc_method<commit_after>//
// Copyright 2018 The Project Oak Authors
//
// 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.
//

//! Functionality to help Oak Nodes interact with gRPC.

pub use crate::proto::code::Code;
use crate::{proto, OakError, OakStatus, ReadHandle};
use log::{error, info, warn};
use proto::grpc_encap::{GrpcRequest, GrpcResponse};
use protobuf::{Message, ProtobufEnum};

mod invocation;
pub use invocation::Invocation;

/// Result type that uses a [`proto::status::Status`] type for error values.
pub type Result<T> = std::result::Result<T, proto::status::Status>;

/// Helper to create a gRPC status object.
pub fn build_status(code: Code, msg: &str) -> proto::status::Status {
    let mut status = proto::status::Status::new();
    status.set_code(code.value());
    status.set_message(msg.to_string());
    status
}

impl crate::io::Encodable for GrpcRequest {
    fn encode(&self) -> std::result::Result<crate::io::Message, crate::OakError> {
        let bytes = self.write_to_bytes()?;
        let handles = Vec::new();
        Ok(crate::io::Message { bytes, handles })
    }
}

impl crate::io::Decodable for GrpcRequest {
    fn decode(message: &crate::io::Message) -> std::result::Result<Self, crate::OakError> {
        let value = protobuf::parse_from_bytes(&message.bytes)?;
        Ok(value)
    }
}

/// Channel-holding object that encapsulates response messages into
/// `GrpcResponse` wrapper messages and writes serialized versions to a send
/// channel.
pub struct ChannelResponseWriter {
    sender: crate::io::Sender<GrpcResponse>,
}

impl crate::io::Encodable for GrpcResponse {
    fn encode(&self) -> std::result::Result<crate::io::Message, crate::OakError> {
        let bytes = self.write_to_bytes()?;
        let handles = Vec::new();
        Ok(crate::io::Message { bytes, handles })
    }
}

impl crate::io::Decodable for GrpcResponse {
    fn decode(message: &crate::io::Message) -> std::result::Result<Self, crate::OakError> {
        let value = protobuf::parse_from_bytes(&message.bytes)?;
        Ok(value)
    }
}

/// Indicate whether a write method should leave the current gRPC method
/// invocation open or close it.
#[derive(PartialEq, Clone, Debug)]
pub enum WriteMode {
    KeepOpen,
    Close,
}

impl ChannelResponseWriter {
    pub fn new(sender: crate::io::Sender<GrpcResponse>) -> Self {
        ChannelResponseWriter { sender }
    }

    /// Retrieve the Oak handle underlying the writer object.
    pub fn handle(self) -> crate::WriteHandle {
        self.sender.handle
    }

    /// Write out a gRPC response and optionally close out the method
    /// invocation.  Any errors from the channel are silently dropped.
    pub fn write<T: protobuf::Message>(
        &mut self,
        rsp: &T,
        mode: WriteMode,
    ) -> std::result::Result<(), OakError> {
        // Put the serialized response into a GrpcResponse message wrapper and
        // serialize it into the channel.
        let mut grpc_rsp = GrpcResponse::new();
        let mut any = protobuf::well_known_types::Any::new();
        rsp.write_to_writer(&mut any.value)?;
        grpc_rsp.set_rsp_msg(any);
        grpc_rsp.set_last(match mode {
            WriteMode::KeepOpen => false,
            WriteMode::Close => true,
        });
        self.sender.send(&grpc_rsp)?;
        if mode == WriteMode::Close {
            self.sender.close()?;
        }
        Ok(())
    }

    /// Write an empty gRPC response and optionally close out the method
    /// invocation. Any errors from the channel are silently dropped.
    pub fn write_empty(&mut self, mode: WriteMode) -> std::result::Result<(), OakError> {
        let mut grpc_rsp = GrpcResponse::new();
        grpc_rsp.set_rsp_msg(protobuf::well_known_types::Any::new());
        grpc_rsp.set_last(match mode {
            WriteMode::KeepOpen => false,
            WriteMode::Close => true,
        });
        self.sender.send(&grpc_rsp)?;
        if mode == WriteMode::Close {
            self.sender.close()?;
        }
        Ok(())
    }

    /// Close out the gRPC method invocation with the given final result.
    pub fn close(&mut self, result: Result<()>) -> std::result::Result<(), OakError> {
        // Build a final GrpcResponse message wrapper and serialize it into the
        // channel.
        let mut grpc_rsp = GrpcResponse::new();
        grpc_rsp.set_last(true);
        if let Err(status) = result {
            grpc_rsp.set_status(status);
        }
        self.sender.send(&grpc_rsp)?;
        self.sender.close()?;
        Ok(())
    }
}

/// Trait for Oak Nodes that act as a gRPC services.
///
/// An `OakNode` instance is normally passed to [`event_loop`], to allow
/// repeated invocation of its `invoke()` method.
pub trait OakNode {
    /// Construct the (single) instance of the node.
    ///
    /// This method may choose to initialize logging by invoking
    /// [`oak_log::init()`].
    ///
    /// [`oak_log::init()`]: ../../oak_log/fn.init.html
    fn new() -> Self
    where
        Self: Sized;

    /// Process a single gRPC method invocation.
    ///
    /// The method name is provided by `method` and the incoming serialized gRPC
    /// request is held in `req`.  Response messages should be written using
    /// `writer.write`, followed by `writer.close`.
    fn invoke(&mut self, method: &str, req: &[u8], writer: ChannelResponseWriter);
}

/// Encapsulate a protocol buffer message in a GrpcRequest wrapper using the
/// given method name.
pub fn encap_request<T: protobuf::Message>(req: T, method_name: &str) -> Option<GrpcRequest> {
    // Put the request in a GrpcRequest wrapper and serialize it.
    let mut grpc_req = GrpcRequest::new();
    grpc_req.set_method_name(method_name.to_string());
    let mut any = protobuf::well_known_types::Any::new();
    if let Err(e) = req.write_to_writer(&mut any.value) {
        warn!("failed to serialize gRPC request: {}", e);
        return None;
    };
    grpc_req.set_req_msg(any);
    grpc_req.set_last(true);
    Some(grpc_req)
}

/// Helper to inject a (single) gRPC request message via a notification channel,
/// in the same manner as the gRPC pseudo-Node does, and collect a (single)
/// response.
pub fn invoke_grpc_method<R, Q>(
    method_name: &str,
    req: R,
    invocation_channel: &crate::io::Sender<Invocation>,
) -> Result<Q>
where
    R: protobuf::Message,
    Q: protobuf::Message,
{
    // Create a new channel for request message delivery.
    let (req_sender, req_receiver) =
        crate::io::channel_create::<GrpcRequest>().expect("failed to create channel");

    // Put the request in a GrpcRequest wrapper and send it into the request
    // message channel.
    let req = encap_request(req, method_name).expect("failed to serialize GrpcRequest");
    req_sender.send(&req).expect("failed to write to channel");
    req_sender.close().expect("failed to close channel");

    // Create a new channel for responses to arrive on.
    let (rsp_sender, rsp_receiver) =
        crate::io::channel_create::<GrpcResponse>().expect("failed to create channel");

    // Build an Invocation holding the two channels and send it down the
    // specified channel.
    let invocation = Invocation {
        request_receiver: req_receiver.clone(),
        response_sender: rsp_sender.clone(),
    };
    invocation_channel
        .send(&invocation)
        .expect("failed to write invocation to channel");
    req_receiver.close().expect("failed to close channel");
    rsp_sender.close().expect("failed to close channel");

    // Read a single encapsulated response.
    let mut rsp = rsp_receiver.receive().map_err(|status| {
        error!("failed to receive response: {:?}", status);
        build_status(
            Code::INTERNAL,
            &format!("failed to receive gRPC response: {:?}", status),
        )
    })?;
    rsp_receiver.close().expect("failed to close channel");
    if !rsp.last {
        panic!("Expected single final response");
    }

    if rsp.get_status().code != Code::OK.value() {
        return Err(rsp.take_status());
    }
    Ok(protobuf::parse_from_bytes(&rsp.get_rsp_msg().value)
        .expect("Failed to parse response protobuf message"))
}

/// Perform a gRPC event loop for a Node.
///
/// Invoking the given `node`'s [`invoke`] method for each incoming request that
/// arrives on the inbound channel as a serialized [`GrpcRequest`] message,
/// giving the [`invoke`] method the outbound channel for encapsulated responses
/// to be written to.
///
/// [`invoke`]: OakNode::invoke
/// [`GrpcRequest`]: crate::proto::grpc_encap::GrpcRequest
pub fn event_loop<T: OakNode>(
    mut node: T,
    grpc_in_handle: ReadHandle,
) -> std::result::Result<(), crate::OakStatus> {
    info!("start event loop for node with handle {:?}", grpc_in_handle);
    if !grpc_in_handle.handle.is_valid() {
        return Err(OakStatus::ERR_CHANNEL_CLOSED);
    }
    let invocation_receiver = crate::io::Receiver::new(grpc_in_handle);
    loop {
        // Explicitly call `wait` and then `try_receive` here because calling `receive` would hide
        // any `OakStatus::ERR_TERMINATED` errors that may occur in the `wait` phase, since they
        // would be wrapped in a `OakError` value that would then always be unwrapped and panic.
        invocation_receiver.wait()?;
        let invocation: Invocation = invocation_receiver.try_receive().unwrap_or_else(|err| {
            panic!("could not receive gRPC invocation: {}", err);
        });

        // Read a single encapsulated request message from the read half.
        let req: GrpcRequest = invocation.request_receiver.receive().unwrap_or_else(|err| {
            panic!("could not read gRPC request: {:?}", err);
        });
        // Since we are expecting a single message, close the channel immediately.
        // This will change when we implement client streaming (#97).
        invocation.request_receiver.close().unwrap_or_else(|err| {
            panic!("could not close gRPC request channel: {:?}", err);
        });
        if !req.last {
            // TODO(#97): Implement client streaming.
            panic!("Support for streaming requests not yet implemented");
        }
        node.invoke(
            &req.method_name,
            req.get_req_msg().value.as_slice(),
            ChannelResponseWriter::new(invocation.response_sender),
        );
    }
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// request/response pair.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_req_rsp<C, R, Q>(mut node_fn: C, req: &[u8], mut writer: ChannelResponseWriter)
where
    C: FnMut(R) -> Result<Q>,
    R: protobuf::Message,
    Q: protobuf::Message,
{
    let r: R = protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message");
    let result = match node_fn(r) {
        Ok(rsp) => writer.write(&rsp, WriteMode::Close),
        Err(status) => writer.close(Err(status)),
    };
    if let Err(e) = result {
        panic!("Failed to process response: {:?}", e)
    }
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// request and streams responses.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_req_stream<C, R>(mut node_fn: C, req: &[u8], writer: ChannelResponseWriter)
where
    C: FnMut(R, ChannelResponseWriter),
    R: protobuf::Message,
{
    let r: R = protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message");
    node_fn(r, writer)
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// stream of requests to produce a response.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_stream_rsp<C, R, Q>(mut node_fn: C, req: &[u8], mut writer: ChannelResponseWriter)
where
    C: FnMut(Vec<R>) -> Result<Q>,
    R: protobuf::Message,
    Q: protobuf::Message,
{
    // TODO(#97): better client-side streaming
    let rr: Vec<R> =
        vec![protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message")];
    let result = match node_fn(rr) {
        Ok(rsp) => writer.write(&rsp, WriteMode::Close),
        Err(status) => writer.close(Err(status)),
    };
    if let Err(e) = result {
        panic!("Failed to process response: {:?}", e)
    }
}

/// Generic function that handles request deserialization and response
/// serialization (and sending down a channel) for a function that handles a
/// stream of requests to produce a stream of responses.
///
/// Panics if [de-]serialization or channel operations fail.
pub fn handle_stream_stream<C, R>(mut node_fn: C, req: &[u8], writer: ChannelResponseWriter)
where
    C: FnMut(Vec<R>, ChannelResponseWriter),
    R: protobuf::Message,
{
    // TODO(#97): better client-side streaming
    let rr: Vec<R> =
        vec![protobuf::parse_from_bytes(&req).expect("Failed to parse request protobuf message")];
    node_fn(rr, writer)
}
<|endoftext|>


<commit_before>// Copyright (c) 2015, Mikhail Vorotilov
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice, this
//   list of conditions and the following disclaimer.
//
// * Redistributions in binary form must reproduce the above copyright notice,
//   this list of conditions and the following disclaimer in the documentation
//   and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use super::FloatType;
use std::error::Error;
use std::fmt;

/// Pair of the independent variable x and the function value y=F(x)
#[derive(Debug, PartialEq)]
pub struct Sample<F>
where
    F: FloatType,
{
    /// Value of the independent variable (X-axis)
    x: F,
    /// Value of the dependent variable (Y-axis)
    y: F,
}

impl<F> Sample<F>
where
    F: FloatType,
{
    fn is_bracketed_with(&self, other: &Self) -> bool {
        self.y * other.y <= F::zero()
    }
}

/// Interval between two samples, including these samples
#[derive(Debug, PartialEq)]
pub struct Interval<F>
where
    F: FloatType,
{
    /// First sample
    begin: Sample<F>,
    /// Last sample
    end: Sample<F>,
}

impl<F> Interval<F>
where
    F: FloatType,
{
    fn is_bracketed(&self) -> bool {
        self.begin.is_bracketed_with(&self.end)
    }
    fn is_converged(&self, convergency: &mut Convergency<F>) -> bool {
        convergency.is_converged(self.begin.x, self.end.x)
    }
    fn middle(&self) -> F {
        if self.is_bracketed() && self.begin.y != self.end.y {
            self.begin.x - self.begin.y * (self.end.x - self.begin.x) / (self.end.y - self.begin.y)
        } else {
            (self.begin.x + self.end.x) / F::two()
        }
    }
}

/// Possible errors
#[derive(Debug, PartialEq)]
pub enum SearchError {
    /// The algorithm could not converge within the given number of iterations
    NoConvergency,
    /// Initial values do not bracket zero
    NoBracketing,
    /// The algorithm cannot continue from the point where the derivative is zero
    ZeroDerivative,
}

impl fmt::Display for SearchError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            SearchError::NoConvergency => write!(f, "Convergency Error"),
            SearchError::NoBracketing => write!(f, "Bracketing Error"),
            SearchError::ZeroDerivative => write!(f, "Zero Derivative Error"),
        }
    }
}
impl Error for SearchError {
    fn description(&self) -> &str {
        match self {
            SearchError::NoConvergency => "The algorithm could not converge within the given number of iterations",
            SearchError::NoBracketing => "Initial values do not bracket zero",
            SearchError::ZeroDerivative => "The algorithm cannot continue from the point where the derivative is zero",
        }
    }
}

/// The way to check if the algorithm has finished by either finding a root
/// or reaching the iteration limit.
pub trait Convergency<F: FloatType> {
    /// Return true if the given Y value is close enough to the zero
    fn is_root_found(&mut self, y: F) -> bool;
    /// Return true if given x values are close enough to each other
    fn is_converged(&mut self, x1: F, x2: F) -> bool;
    /// Return true if no more iterations desired
    fn is_iteration_limit_reached(&mut self, iter: usize) -> bool;
}

impl<F: FloatType> Convergency<F> for F {
    /// Return true if the given Y value is close enough to the zero
    fn is_root_found(&mut self, y: F) -> bool {
        y.abs() < self.abs()
    }
    /// Return true if given x values are close enough to each other
    fn is_converged(&mut self, x1: F, x2: F) -> bool {
        (x1 - x2).abs() < self.abs()
    }
    /// Return true if no more iterations desired
    fn is_iteration_limit_reached(&mut self, iter: usize) -> bool {
        iter >= 30
    }
}

pub mod brent;
pub mod newton_raphson;
pub mod polynom;
pub mod regula_falsi;
pub mod secant;

pub mod debug_convergency;
pub mod simple_convergency;

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

    #[test]
    fn sample_bracketed() {
        let sample1 = Sample { x: 0f64, y: 0f64 };
        let sample2 = Sample { x: 1f64, y: 1f64 };
        let sample3 = Sample { x: 1f64, y: -1f64 };
        let sample4 = Sample { x: -1f64, y: 0f64 };
        let sample5 = Sample { x: -1f64, y: 1f64 };
        assert_eq!(true, sample1.is_bracketed_with(&sample2));
        assert_eq!(true, sample1.is_bracketed_with(&sample3));
        assert_eq!(true, sample1.is_bracketed_with(&sample4));
        assert_eq!(true, sample1.is_bracketed_with(&sample5));
        assert_eq!(true, sample2.is_bracketed_with(&sample3));
        assert_eq!(true, sample2.is_bracketed_with(&sample4));
        assert_eq!(false, sample2.is_bracketed_with(&sample5));
        assert_eq!(true, sample3.is_bracketed_with(&sample4));
        assert_eq!(true, sample3.is_bracketed_with(&sample5));
        assert_eq!(true, sample4.is_bracketed_with(&sample5));
    }

    #[test]
    fn root_interval_bracketed() {
        let sut1 = Interval {
            begin: Sample { x: 0f64, y: 0f64 },
            end: Sample { x: 0f64, y: 0f64 },
        };
        let sut2 = Interval {
            begin: Sample { x: 0f32, y: 0f32 },
            end: Sample { x: 1f32, y: 0f32 },
        };
        let sut3 = Interval {
            begin: Sample { x: 0f64, y: 0f64 },
            end: Sample { x: 0f64, y: 1f64 },
        };
        let sut4 = Interval {
            begin: Sample { x: -1f64, y: 0f64 },
            end: Sample { x: 0f64, y: 0f64 },
        };
        let sut5 = Interval {
            begin: Sample { x: -1f64, y: 0f64 },
            end: Sample { x: 0f64, y: 1f64 },
        };
        let sut6 = Interval {
            begin: Sample { x: -1f32, y: -1f32 },
            end: Sample { x: 0f32, y: 1f32 },
        };
        let sut7 = Interval {
            begin: Sample { x: 0f64, y: 1f64 },
            end: Sample { x: 1f64, y: -1f64 },
        };
        assert_eq!(true, sut1.is_bracketed());
        assert_eq!(true, sut2.is_bracketed());
        assert_eq!(true, sut3.is_bracketed());
        assert_eq!(true, sut4.is_bracketed());
        assert_eq!(true, sut5.is_bracketed());
        assert_eq!(true, sut6.is_bracketed());
        assert_eq!(true, sut7.is_bracketed());
    }

    #[test]
    fn root_interval_not_bracketed() {
        let sut1 = Interval {
            begin: Sample { x: 0f64, y: 1f64 },
            end: Sample { x: 1f64, y: 1f64 },
        };
        let sut2 = Interval {
            begin: Sample { x: -1f64, y: -1f64 },
            end: Sample { x: 1f64, y: -1f64 },
        };
        assert_eq!(false, sut1.is_bracketed());
        assert_eq!(false, sut2.is_bracketed());
    }

    #[test]
    fn root_interval_middle() {
        let sut1 = Interval {
            begin: Sample { x: 0f64, y: 1f64 },
            end: Sample { x: 2f64, y: -3f64 },
        };
        let sut2 = Interval {
            begin: Sample { x: -1f64, y: 0f64 },
            end: Sample { x: 1f64, y: 0f64 },
        };
        assert_eq!(0.5f64, sut1.middle());
        assert_eq!(0f64, sut2.middle());
    }

}
<commit_msg>Do not let the middle of the interval be too close to the border.<commit_after>// Copyright (c) 2015, Mikhail Vorotilov
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice, this
//   list of conditions and the following disclaimer.
//
// * Redistributions in binary form must reproduce the above copyright notice,
//   this list of conditions and the following disclaimer in the documentation
//   and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use super::FloatType;
use std::error::Error;
use std::fmt;

/// Pair of the independent variable x and the function value y=F(x)
#[derive(Debug, PartialEq)]
pub struct Sample<F>
where
    F: FloatType,
{
    /// Value of the independent variable (X-axis)
    x: F,
    /// Value of the dependent variable (Y-axis)
    y: F,
}

impl<F> Sample<F>
where
    F: FloatType,
{
    fn is_bracketed_with(&self, other: &Self) -> bool {
        self.y * other.y <= F::zero()
    }
}

/// Interval between two samples, including these samples
#[derive(Debug, PartialEq)]
pub struct Interval<F>
where
    F: FloatType,
{
    /// First sample
    begin: Sample<F>,
    /// Last sample
    end: Sample<F>,
}

impl<F> Interval<F>
where
    F: FloatType,
{
    fn is_bracketed(&self) -> bool {
        self.begin.is_bracketed_with(&self.end)
    }
    fn is_converged(&self, convergency: &mut Convergency<F>) -> bool {
        convergency.is_converged(self.begin.x, self.end.x)
    }
    /// Returns a point somewhere in middle of the interval for narrowing this interval down.
    /// Rules are as follows:
    /// * If the interval is bracketed, use the secant to find the middle point.
    /// ** The middle point may not be too close to either range of the interval.
    /// * If the interval is not bracketed (why would one use an unbracketed interval?), bisect it.
    fn middle(&self) -> F {
        if self.is_bracketed() && self.begin.y != self.end.y {
            let mut shift = -self.begin.y * (self.end.x - self.begin.x) / (self.end.y - self.begin.y);
            if shift < (self.end.x - self.begin.x) / F::twenty_seven() {
                shift = (self.end.x - self.begin.x) / F::twenty_seven();
            }
            if shift > (self.end.x - self.begin.x) * (F::twenty_seven() - F::one()) / F::twenty_seven() {
                shift = (self.end.x - self.begin.x) * (F::twenty_seven() - F::one()) / F::twenty_seven();
            }
            self.begin.x + shift
        } else {
            (self.begin.x + self.end.x) / F::two()
        }
    }
}

/// Possible errors
#[derive(Debug, PartialEq)]
pub enum SearchError {
    /// The algorithm could not converge within the given number of iterations
    NoConvergency,
    /// Initial values do not bracket zero
    NoBracketing,
    /// The algorithm cannot continue from the point where the derivative is zero
    ZeroDerivative,
}

impl fmt::Display for SearchError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            SearchError::NoConvergency => write!(f, "Convergency Error"),
            SearchError::NoBracketing => write!(f, "Bracketing Error"),
            SearchError::ZeroDerivative => write!(f, "Zero Derivative Error"),
        }
    }
}
impl Error for SearchError {
    fn description(&self) -> &str {
        match self {
            SearchError::NoConvergency => "The algorithm could not converge within the given number of iterations",
            SearchError::NoBracketing => "Initial values do not bracket zero",
            SearchError::ZeroDerivative => "The algorithm cannot continue from the point where the derivative is zero",
        }
    }
}

/// The way to check if the algorithm has finished by either finding a root
/// or reaching the iteration limit.
pub trait Convergency<F: FloatType> {
    /// Return true if the given Y value is close enough to the zero
    fn is_root_found(&mut self, y: F) -> bool;
    /// Return true if given x values are close enough to each other
    fn is_converged(&mut self, x1: F, x2: F) -> bool;
    /// Return true if no more iterations desired
    fn is_iteration_limit_reached(&mut self, iter: usize) -> bool;
}

impl<F: FloatType> Convergency<F> for F {
    /// Return true if the given Y value is close enough to the zero
    fn is_root_found(&mut self, y: F) -> bool {
        y.abs() < self.abs()
    }
    /// Return true if given x values are close enough to each other
    fn is_converged(&mut self, x1: F, x2: F) -> bool {
        (x1 - x2).abs() < self.abs()
    }
    /// Return true if no more iterations desired
    fn is_iteration_limit_reached(&mut self, iter: usize) -> bool {
        iter >= 30
    }
}

pub mod brent;
pub mod newton_raphson;
pub mod polynom;
pub mod regula_falsi;
pub mod secant;

pub mod debug_convergency;
pub mod simple_convergency;

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

    #[test]
    fn sample_bracketed() {
        let sample1 = Sample { x: 0f64, y: 0f64 };
        let sample2 = Sample { x: 1f64, y: 1f64 };
        let sample3 = Sample { x: 1f64, y: -1f64 };
        let sample4 = Sample { x: -1f64, y: 0f64 };
        let sample5 = Sample { x: -1f64, y: 1f64 };
        assert_eq!(true, sample1.is_bracketed_with(&sample2));
        assert_eq!(true, sample1.is_bracketed_with(&sample3));
        assert_eq!(true, sample1.is_bracketed_with(&sample4));
        assert_eq!(true, sample1.is_bracketed_with(&sample5));
        assert_eq!(true, sample2.is_bracketed_with(&sample3));
        assert_eq!(true, sample2.is_bracketed_with(&sample4));
        assert_eq!(false, sample2.is_bracketed_with(&sample5));
        assert_eq!(true, sample3.is_bracketed_with(&sample4));
        assert_eq!(true, sample3.is_bracketed_with(&sample5));
        assert_eq!(true, sample4.is_bracketed_with(&sample5));
    }

    #[test]
    fn root_interval_bracketed() {
        let sut1 = Interval {
            begin: Sample { x: 0f64, y: 0f64 },
            end: Sample { x: 0f64, y: 0f64 },
        };
        let sut2 = Interval {
            begin: Sample { x: 0f32, y: 0f32 },
            end: Sample { x: 1f32, y: 0f32 },
        };
        let sut3 = Interval {
            begin: Sample { x: 0f64, y: 0f64 },
            end: Sample { x: 0f64, y: 1f64 },
        };
        let sut4 = Interval {
            begin: Sample { x: -1f64, y: 0f64 },
            end: Sample { x: 0f64, y: 0f64 },
        };
        let sut5 = Interval {
            begin: Sample { x: -1f64, y: 0f64 },
            end: Sample { x: 0f64, y: 1f64 },
        };
        let sut6 = Interval {
            begin: Sample { x: -1f32, y: -1f32 },
            end: Sample { x: 0f32, y: 1f32 },
        };
        let sut7 = Interval {
            begin: Sample { x: 0f64, y: 1f64 },
            end: Sample { x: 1f64, y: -1f64 },
        };
        assert_eq!(true, sut1.is_bracketed());
        assert_eq!(true, sut2.is_bracketed());
        assert_eq!(true, sut3.is_bracketed());
        assert_eq!(true, sut4.is_bracketed());
        assert_eq!(true, sut5.is_bracketed());
        assert_eq!(true, sut6.is_bracketed());
        assert_eq!(true, sut7.is_bracketed());
    }

    #[test]
    fn root_interval_not_bracketed() {
        let sut1 = Interval {
            begin: Sample { x: 0f64, y: 1f64 },
            end: Sample { x: 1f64, y: 1f64 },
        };
        let sut2 = Interval {
            begin: Sample { x: -1f64, y: -1f64 },
            end: Sample { x: 1f64, y: -1f64 },
        };
        assert_eq!(false, sut1.is_bracketed());
        assert_eq!(false, sut2.is_bracketed());
    }

    #[test]
    fn root_interval_middle() {
        let sut1 = Interval {
            begin: Sample { x: 0f64, y: 1f64 },
            end: Sample { x: 2f64, y: -3f64 },
        };
        let sut2 = Interval {
            begin: Sample { x: -1f64, y: 0f64 },
            end: Sample { x: 1f64, y: 0f64 },
        };
        assert_eq!(0.5f64, sut1.middle());
        assert_eq!(0f64, sut2.middle());
    }

}
<|endoftext|>


<commit_before>use std::borrow::BorrowMut;
use std::marker::PhantomData;
use std::collections::{HashMap, HashSet};
use context::ContextOf;
use framebuffer::FramebufferBinding;
use program::ProgramAttrib;
use vertex_data::{VertexData, VertexBytes, VertexAttribute};
use index_data::{IndexData, IndexDatum};
use buffer::{Buffer, BufferBinderOf, BufferBinding,
             ArrayBufferBinder, ElementArrayBufferBinder,
             ArrayBufferBinding, ElementArrayBufferBinding};
use types::DrawingMode;

#[derive(Debug)]
pub enum AttribAddError {
    DuplicateAttrib(String)
}

#[derive(Debug)]
pub struct AttribError {
    missing_attribs: Vec<String>,
    unknown_attribs: Vec<String>
}

pub struct AttribBinder {
    attribs: HashMap<String, ProgramAttrib>
}

impl AttribBinder {
    pub fn new() -> Self {
        AttribBinder {
            attribs: HashMap::new()
        }
    }

    pub fn add(&mut self, name: &str, attrib: ProgramAttrib)
        -> Result<(), AttribAddError>
    {
        match self.attribs.insert(name.into(), attrib) {
            None => Ok(()),
            Some(_) => Err(AttribAddError::DuplicateAttrib(name.into()))
        }
    }

    fn for_each<T, F>(&self, mut f: F) -> Result<(), AttribError>
        where T: VertexData, F: FnMut(VertexAttribute, ProgramAttrib)
    {
        // TODO: Avoid heap allocations
        // TODO: Avoid redundant calls to T::visit_attributes
        let mut attribs =
            HashMap::<String, (VertexAttribute, ProgramAttrib)>::new();
        let mut missing = Vec::<String>::new();

        T::visit_attributes(|vertex_attrib| {
            match self.attribs.get(&vertex_attrib.name) {
                Some(program_attrib) => {
                    let pair = (vertex_attrib.clone(), *program_attrib);
                    attribs.insert(vertex_attrib.name, pair);
                },
                None => {
                    missing.push(vertex_attrib.name);
                }
            }
        });

        let unknown: Vec<_> = {
            let expected: HashSet<_> = self.attribs.keys().collect();
            let actual: HashSet<_> = attribs.keys().collect();
            expected.difference(&actual).cloned().cloned().collect()
        };

        if missing.is_empty() && unknown.is_empty() {
            for (_, (vertex_attrib, program_attrib)) in attribs.into_iter() {
                f(vertex_attrib, program_attrib);
            }
            Ok(())
        }
        else {
            Err(AttribError {
                missing_attribs: missing,
                unknown_attribs: unknown
            })
        }
    }

    pub fn enable<V, B, F, P, R, T>(&self, gl: &ContextOf<B, F, P, R, T>)
        -> Result<(), AttribError>
        where V: VertexData
    {
        self.for_each::<V, _>(|_, program_attrib| {
            gl.enable_vertex_attrib_array(program_attrib);
        })
    }

    pub fn bind<V: VertexData>(&self, gl_buffer: &ArrayBufferBinding)
        -> Result<(), AttribError>
    {
        self.for_each::<V, _>(|vertex_attrib, program_attrib| {
            unsafe {
                gl_buffer.vertex_attrib_pointer(
                    program_attrib,
                    vertex_attrib.ty.components,
                    vertex_attrib.ty.data,
                    vertex_attrib.ty.normalize,
                    vertex_attrib.stride,
                    vertex_attrib.offset
                );
            }
        })
    }
}

#[derive(Debug)]
pub enum VertexBindError {
    BindingError(AttribError),
    NoAttributeBindings
}

impl From<AttribError> for VertexBindError {
    fn from(attrib_error: AttribError) -> VertexBindError {
        VertexBindError::BindingError(attrib_error)
    }
}

pub struct VertexBuffer<T: VertexData> {
    attrib_binder: Option<AttribBinder>,
    buffer: Buffer,
    count: usize,
    phantom: PhantomData<*const T>
}

impl<V: VertexData> VertexBuffer<V> {
    pub fn bind_attrib_pointers(&mut self, binder: AttribBinder) {
        self.attrib_binder = Some(binder);
    }

    pub fn bind<BA, BE, F, P, R, T>(&mut self,
                                    gl: ContextOf<BufferBinderOf<BA, BE>,
                                                  F,
                                                  P,
                                                  R,
                                                  T>)
        -> Result<(), VertexBindError>
        where BA: BorrowMut<ArrayBufferBinder>
    {
        match self.attrib_binder {
            Some(ref binder) => {
                let (mut ba_binder, mut gl) = gl.split_array_buffer();
                let mut ba_binder = ba_binder.borrow_mut();
                let gl_buffer = ba_binder.bind(&mut self.buffer);
                try!(binder.enable::<V, _, _, _, _, _>(&mut gl));
                try!(binder.bind::<V>(&gl_buffer));
                Ok(())
            },
            None => { Err(VertexBindError::NoAttributeBindings) }
        }
    }

    pub fn buffer(&self) -> &Buffer {
        &self.buffer
    }

    pub fn buffer_mut(&mut self) -> &mut Buffer {
        &mut self.buffer
    }
}

pub struct VertexBufferBinding<'a, T: VertexData + 'a> {
    gl_buffer: ArrayBufferBinding<'a>,
    count: &'a mut usize,
    _phantom: PhantomData<*const VertexBuffer<T>>
}

impl<'a, T: VertexData + 'a> VertexBufferBinding<'a, T> {
    pub fn buffer_data(&mut self, data: &[T], usage: super::BufferDataUsage)
        where [T]: VertexBytes
    {
        *self.count = data.len();
        self.gl_buffer.buffer_bytes(data.vertex_bytes(), usage);
    }
}

impl<'a> FramebufferBinding<'a> {
    pub fn draw_arrays_range_vbo<V>(&mut self,
                                    gl_vbo: &VertexBufferBinding<V>,
                                    mode: DrawingMode,
                                    start: u32,
                                    length: usize)
        where V: VertexData
    {
        debug_assert!((start as usize) + length <= *gl_vbo.count);

        unsafe {
            self.draw_arrays_range(&gl_vbo.gl_buffer, mode, start, length);
        }
    }

    pub fn draw_arrays_vbo<V>(&mut self,
                              gl_vbo: &VertexBufferBinding<V>,
                              mode: DrawingMode)
        where V: VertexData
    {
        unsafe {
            self.draw_arrays_range(&gl_vbo.gl_buffer,
                                   mode,
                                   0,
                                   *gl_vbo.count);
        }
    }

    pub fn draw_n_elements_buffered_vbo<V, I>(&mut self,
                                              gl_vbo: &VertexBufferBinding<V>,
                                              gl_ibo: &IndexBufferBinding<I>,
                                              mode: DrawingMode,
                                              length: usize)
        where V: VertexData, I: IndexDatum
    {
        debug_assert!(length <= *gl_ibo.count);

        unsafe {
            self.draw_n_elements_buffered(&gl_vbo.gl_buffer,
                                          &gl_ibo.gl_buffer,
                                          mode,
                                          length,
                                          I::index_datum_type());
        }
    }

    pub fn draw_elements_buffered_vbo<V, I>(&mut self,
                                            gl_vbo: &VertexBufferBinding<V>,
                                            gl_ibo: &IndexBufferBinding<I>,
                                            mode: DrawingMode)
        where V: VertexData, I: IndexDatum
    {
        unsafe {
            self.draw_n_elements_buffered(&gl_vbo.gl_buffer,
                                          &gl_ibo.gl_buffer,
                                          mode,
                                          *gl_ibo.count,
                                          I::index_datum_type());
        }
    }

    pub fn draw_n_elements_vbo<V, I>(&mut self,
                                     gl_vbo: &VertexBufferBinding<'a, V>,
                                     mode: DrawingMode,
                                     count: usize,
                                     indices: &[I])
        where V: VertexData, I: IndexDatum, [I]: IndexData
    {
        unsafe {
            self.draw_n_elements(&gl_vbo.gl_buffer, mode, count, indices);
        }
    }

    pub fn draw_elements_vbo<V, I>(&mut self,
                                   gl_vbo: &VertexBufferBinding<'a, V>,
                                   mode: DrawingMode,
                                   indices: &[I])
        where V: VertexData, I: IndexDatum, [I]: IndexData
    {
        unsafe {
            self.draw_elements(&gl_vbo.gl_buffer, mode, indices);
        }
    }
}

impl<B, F, P, R, T> ContextOf<B, F, P, R, T> {
    pub fn new_vertex_buffer<V: VertexData>(&self) -> VertexBuffer<V> {
        VertexBuffer {
            attrib_binder: None,
            buffer: self.gen_buffer(),
            count: 0,
            phantom: PhantomData
        }
    }
}

// TODO: Refactor to work with less restrictive generic params
impl<BA, BE, F, P, R, T> ContextOf<BufferBinderOf<BA, BE>, F, P, R, T> {
    pub fn bind_vertex_buffer<'a, V>(mut self, vbo: &'a mut VertexBuffer<V>)
        -> (
            VertexBufferBinding<V>,
            ContextOf<BufferBinderOf<(), BE>, F, P, R, T>
        )
        where V: VertexData,
              BA: BorrowMut<ArrayBufferBinder>
    {
        {
            let gl = self.borrowed_mut().map_buffers(|b| b.borrowed_mut());

            vbo.bind(gl).unwrap();
        }
        let (mut ba_binder, gl) = self.split_array_buffer();
        let mut ba_binder = ba_binder.borrow_mut();
        let gl_array_buffer = ba_binder.bind(&mut vbo.buffer);
        (
            VertexBufferBinding {
                gl_buffer: gl_array_buffer,
                count: &mut vbo.count,
                _phantom: PhantomData
            },
            gl
        )
    }

    pub fn bind_index_buffer<'a, I>(self, ibo: &'a mut IndexBuffer<I>)
        -> (
            IndexBufferBinding<I>,
            ContextOf<BufferBinderOf<BA, ()>, F, P, R, T>
        )
        where I: IndexDatum + 'a,
              BE: BorrowMut<ElementArrayBufferBinder> + 'a
    {
        let (mut be_binder, gl) = self.split_element_array_buffer();
        let be_binder = be_binder.borrow_mut();
        let gl_be = be_binder.bind(&mut ibo.buffer);
        (
            IndexBufferBinding {
                gl_buffer: gl_be,
                count: &mut ibo.count,
                _phantom: PhantomData
            },
            gl
        )
    }
}

pub struct IndexBuffer<T: IndexDatum> {
    buffer: Buffer,
    count: usize,
    phantom: PhantomData<*const T>
}

impl<T: IndexDatum> IndexBuffer<T> {
    pub fn buffer(&self) -> &Buffer {
        &self.buffer
    }

    pub fn buffer_mut(&mut self) -> &mut Buffer {
        &mut self.buffer
    }
}

pub struct IndexBufferBinding<'a, T: IndexDatum + 'a> {
    gl_buffer: ElementArrayBufferBinding<'a>,
    count: &'a mut usize,
    _phantom: PhantomData<*const IndexBuffer<T>>
}

impl<'a, T: IndexDatum + 'a> IndexBufferBinding<'a, T> {
    pub fn buffer_data(&mut self, data: &[T], usage: super::BufferDataUsage)
        where [T]: IndexData
    {
        *self.count = data.len();
        self.gl_buffer.buffer_bytes(data.index_bytes(), usage);
    }
}

impl<B, F, P, R, T> ContextOf<B, F, P, R, T> {
    pub fn new_index_buffer<I: IndexDatum>(&self) -> IndexBuffer<I> {
        IndexBuffer {
            buffer: self.gen_buffer(),
            count: 0,
            phantom: PhantomData
        }
    }
}

#[macro_export]
macro_rules! attrib_pointers {
    ($gl:expr, $vbo:expr, {
        $($field_name:ident => $field_attrib:expr),*
    }) => {
        {
            let mut binder = $crate::AttribBinder::new();
            $(binder.add(stringify!($field_name), $field_attrib).unwrap());*;
            binder
        }
    }
}

#[macro_export]
macro_rules! bind_attrib_pointers {
    ($gl:expr, $vbo:expr, {
        $($field_name:ident => $field_attrib:expr),*
    }) => {
        {
            let vbo = $vbo;
            let binder = {
                attrib_pointers!($gl, vbo, {
                    $($field_name => $field_attrib),*
                })
            };
            vbo.bind_attrib_pointers(binder);
        }
    }
}

#[macro_export]
macro_rules! bind_vertex_buffer {
    ($gl:expr, $vbo:expr) => {
        {
            use std::borrow::BorrowMut;
            let mut vbo = $vbo;

            let (gl_vbo, _) = $gl.bind_vertex_buffer(vbo.borrow_mut());
            gl_vbo
        }
    }
}

#[macro_export]
macro_rules! bind_index_buffer {
    ($gl:expr, $ibo:expr) => {
        {
            use std::borrow::BorrowMut;
            let mut ibo = $ibo;

            let (gl_ibo, _) = $gl.bind_index_buffer(ibo.borrow_mut());
            gl_ibo
        }
    }
}
<commit_msg>Move `bind_vertex_buffer()` and `bind_index_buffer()` into new traits<commit_after>use std::borrow::BorrowMut;
use std::marker::PhantomData;
use std::collections::{HashMap, HashSet};
use context::ContextOf;
use framebuffer::FramebufferBinding;
use program::ProgramAttrib;
use vertex_data::{VertexData, VertexBytes, VertexAttribute};
use index_data::{IndexData, IndexDatum};
use buffer::{Buffer, BufferBinderOf, BufferBinding,
             ArrayBufferBinder, ArrayBufferBinding,
             ElementArrayBufferBinding, ElementArrayBufferContext};
use types::DrawingMode;

#[derive(Debug)]
pub enum AttribAddError {
    DuplicateAttrib(String)
}

#[derive(Debug)]
pub struct AttribError {
    missing_attribs: Vec<String>,
    unknown_attribs: Vec<String>
}

pub struct AttribBinder {
    attribs: HashMap<String, ProgramAttrib>
}

impl AttribBinder {
    pub fn new() -> Self {
        AttribBinder {
            attribs: HashMap::new()
        }
    }

    pub fn add(&mut self, name: &str, attrib: ProgramAttrib)
        -> Result<(), AttribAddError>
    {
        match self.attribs.insert(name.into(), attrib) {
            None => Ok(()),
            Some(_) => Err(AttribAddError::DuplicateAttrib(name.into()))
        }
    }

    fn for_each<T, F>(&self, mut f: F) -> Result<(), AttribError>
        where T: VertexData, F: FnMut(VertexAttribute, ProgramAttrib)
    {
        // TODO: Avoid heap allocations
        // TODO: Avoid redundant calls to T::visit_attributes
        let mut attribs =
            HashMap::<String, (VertexAttribute, ProgramAttrib)>::new();
        let mut missing = Vec::<String>::new();

        T::visit_attributes(|vertex_attrib| {
            match self.attribs.get(&vertex_attrib.name) {
                Some(program_attrib) => {
                    let pair = (vertex_attrib.clone(), *program_attrib);
                    attribs.insert(vertex_attrib.name, pair);
                },
                None => {
                    missing.push(vertex_attrib.name);
                }
            }
        });

        let unknown: Vec<_> = {
            let expected: HashSet<_> = self.attribs.keys().collect();
            let actual: HashSet<_> = attribs.keys().collect();
            expected.difference(&actual).cloned().cloned().collect()
        };

        if missing.is_empty() && unknown.is_empty() {
            for (_, (vertex_attrib, program_attrib)) in attribs.into_iter() {
                f(vertex_attrib, program_attrib);
            }
            Ok(())
        }
        else {
            Err(AttribError {
                missing_attribs: missing,
                unknown_attribs: unknown
            })
        }
    }

    pub fn enable<V, B, F, P, R, T>(&self, gl: &ContextOf<B, F, P, R, T>)
        -> Result<(), AttribError>
        where V: VertexData
    {
        self.for_each::<V, _>(|_, program_attrib| {
            gl.enable_vertex_attrib_array(program_attrib);
        })
    }

    pub fn bind<V: VertexData>(&self, gl_buffer: &ArrayBufferBinding)
        -> Result<(), AttribError>
    {
        self.for_each::<V, _>(|vertex_attrib, program_attrib| {
            unsafe {
                gl_buffer.vertex_attrib_pointer(
                    program_attrib,
                    vertex_attrib.ty.components,
                    vertex_attrib.ty.data,
                    vertex_attrib.ty.normalize,
                    vertex_attrib.stride,
                    vertex_attrib.offset
                );
            }
        })
    }
}

#[derive(Debug)]
pub enum VertexBindError {
    BindingError(AttribError),
    NoAttributeBindings
}

impl From<AttribError> for VertexBindError {
    fn from(attrib_error: AttribError) -> VertexBindError {
        VertexBindError::BindingError(attrib_error)
    }
}

pub struct VertexBuffer<T: VertexData> {
    attrib_binder: Option<AttribBinder>,
    buffer: Buffer,
    count: usize,
    phantom: PhantomData<*const T>
}

impl<V: VertexData> VertexBuffer<V> {
    pub fn bind_attrib_pointers(&mut self, binder: AttribBinder) {
        self.attrib_binder = Some(binder);
    }

    pub fn bind<BA, BE, F, P, R, T>(&mut self,
                                    gl: ContextOf<BufferBinderOf<BA, BE>,
                                                  F,
                                                  P,
                                                  R,
                                                  T>)
        -> Result<(), VertexBindError>
        where BA: BorrowMut<ArrayBufferBinder>
    {
        match self.attrib_binder {
            Some(ref binder) => {
                let (mut ba_binder, mut gl) = gl.split_array_buffer();
                let mut ba_binder = ba_binder.borrow_mut();
                let gl_buffer = ba_binder.bind(&mut self.buffer);
                try!(binder.enable::<V, _, _, _, _, _>(&mut gl));
                try!(binder.bind::<V>(&gl_buffer));
                Ok(())
            },
            None => { Err(VertexBindError::NoAttributeBindings) }
        }
    }

    pub fn buffer(&self) -> &Buffer {
        &self.buffer
    }

    pub fn buffer_mut(&mut self) -> &mut Buffer {
        &mut self.buffer
    }
}

pub struct VertexBufferBinding<'a, T: VertexData + 'a> {
    gl_buffer: ArrayBufferBinding<'a>,
    count: &'a mut usize,
    _phantom: PhantomData<*const VertexBuffer<T>>
}

impl<'a, T: VertexData + 'a> VertexBufferBinding<'a, T> {
    pub fn buffer_data(&mut self, data: &[T], usage: super::BufferDataUsage)
        where [T]: VertexBytes
    {
        *self.count = data.len();
        self.gl_buffer.buffer_bytes(data.vertex_bytes(), usage);
    }
}

impl<'a> FramebufferBinding<'a> {
    pub fn draw_arrays_range_vbo<V>(&mut self,
                                    gl_vbo: &VertexBufferBinding<V>,
                                    mode: DrawingMode,
                                    start: u32,
                                    length: usize)
        where V: VertexData
    {
        debug_assert!((start as usize) + length <= *gl_vbo.count);

        unsafe {
            self.draw_arrays_range(&gl_vbo.gl_buffer, mode, start, length);
        }
    }

    pub fn draw_arrays_vbo<V>(&mut self,
                              gl_vbo: &VertexBufferBinding<V>,
                              mode: DrawingMode)
        where V: VertexData
    {
        unsafe {
            self.draw_arrays_range(&gl_vbo.gl_buffer,
                                   mode,
                                   0,
                                   *gl_vbo.count);
        }
    }

    pub fn draw_n_elements_buffered_vbo<V, I>(&mut self,
                                              gl_vbo: &VertexBufferBinding<V>,
                                              gl_ibo: &IndexBufferBinding<I>,
                                              mode: DrawingMode,
                                              length: usize)
        where V: VertexData, I: IndexDatum
    {
        debug_assert!(length <= *gl_ibo.count);

        unsafe {
            self.draw_n_elements_buffered(&gl_vbo.gl_buffer,
                                          &gl_ibo.gl_buffer,
                                          mode,
                                          length,
                                          I::index_datum_type());
        }
    }

    pub fn draw_elements_buffered_vbo<V, I>(&mut self,
                                            gl_vbo: &VertexBufferBinding<V>,
                                            gl_ibo: &IndexBufferBinding<I>,
                                            mode: DrawingMode)
        where V: VertexData, I: IndexDatum
    {
        unsafe {
            self.draw_n_elements_buffered(&gl_vbo.gl_buffer,
                                          &gl_ibo.gl_buffer,
                                          mode,
                                          *gl_ibo.count,
                                          I::index_datum_type());
        }
    }

    pub fn draw_n_elements_vbo<V, I>(&mut self,
                                     gl_vbo: &VertexBufferBinding<'a, V>,
                                     mode: DrawingMode,
                                     count: usize,
                                     indices: &[I])
        where V: VertexData, I: IndexDatum, [I]: IndexData
    {
        unsafe {
            self.draw_n_elements(&gl_vbo.gl_buffer, mode, count, indices);
        }
    }

    pub fn draw_elements_vbo<V, I>(&mut self,
                                   gl_vbo: &VertexBufferBinding<'a, V>,
                                   mode: DrawingMode,
                                   indices: &[I])
        where V: VertexData, I: IndexDatum, [I]: IndexData
    {
        unsafe {
            self.draw_elements(&gl_vbo.gl_buffer, mode, indices);
        }
    }
}

impl<B, F, P, R, T> ContextOf<B, F, P, R, T> {
    pub fn new_vertex_buffer<V: VertexData>(&self) -> VertexBuffer<V> {
        VertexBuffer {
            attrib_binder: None,
            buffer: self.gen_buffer(),
            count: 0,
            phantom: PhantomData
        }
    }
}

pub trait VertexBufferContext {
    type Rest;

    fn bind_vertex_buffer<'a, V>(self, vbo: &'a mut VertexBuffer<V>)
        -> (VertexBufferBinding<V>, Self::Rest)
        where V: VertexData;
}

// TODO: Refactor to work with less restrictive generic params
impl<BA, BE, F, P, R, T> VertexBufferContext
    for ContextOf<BufferBinderOf<BA, BE>, F, P, R, T>
    where BA: BorrowMut<ArrayBufferBinder>
{
    type Rest = ContextOf<BufferBinderOf<(), BE>, F, P, R, T>;

    fn bind_vertex_buffer<'a, V>(mut self, vbo: &'a mut VertexBuffer<V>)
        -> (VertexBufferBinding<V>, Self::Rest)
        where V: VertexData
    {
        {
            let gl = self.borrowed_mut().map_buffers(|b| b.borrowed_mut());

            vbo.bind(gl).unwrap();
        }
        let (mut ba_binder, rest) = self.split_array_buffer();
        let mut ba_binder = ba_binder.borrow_mut();
        let gl_array_buffer = ba_binder.bind(&mut vbo.buffer);
        (
            VertexBufferBinding {
                gl_buffer: gl_array_buffer,
                count: &mut vbo.count,
                _phantom: PhantomData
            },
            rest
        )
    }
}

pub trait IndexBufferContext {
    type Rest;

    fn bind_index_buffer<'a, I>(self, ibo: &'a mut IndexBuffer<I>)
        -> (IndexBufferBinding<I>, Self::Rest)
        where I: IndexDatum;
}

impl<C> IndexBufferContext for C
    where C: ElementArrayBufferContext
{
    type Rest = C::Rest;

    fn bind_index_buffer<'a, I>(self, ibo: &'a mut IndexBuffer<I>)
        -> (IndexBufferBinding<I>, Self::Rest)
        where I: IndexDatum
    {
        let (gl_be, rest) = self.bind_element_array_buffer(&mut ibo.buffer);
        (
            IndexBufferBinding {
                gl_buffer: gl_be,
                count: &mut ibo.count,
                _phantom: PhantomData
            },
            rest
        )
    }
}

pub struct IndexBuffer<T: IndexDatum> {
    buffer: Buffer,
    count: usize,
    phantom: PhantomData<*const T>
}

impl<T: IndexDatum> IndexBuffer<T> {
    pub fn buffer(&self) -> &Buffer {
        &self.buffer
    }

    pub fn buffer_mut(&mut self) -> &mut Buffer {
        &mut self.buffer
    }
}

pub struct IndexBufferBinding<'a, T: IndexDatum + 'a> {
    gl_buffer: ElementArrayBufferBinding<'a>,
    count: &'a mut usize,
    _phantom: PhantomData<*const IndexBuffer<T>>
}

impl<'a, T: IndexDatum + 'a> IndexBufferBinding<'a, T> {
    pub fn buffer_data(&mut self, data: &[T], usage: super::BufferDataUsage)
        where [T]: IndexData
    {
        *self.count = data.len();
        self.gl_buffer.buffer_bytes(data.index_bytes(), usage);
    }
}

impl<B, F, P, R, T> ContextOf<B, F, P, R, T> {
    pub fn new_index_buffer<I: IndexDatum>(&self) -> IndexBuffer<I> {
        IndexBuffer {
            buffer: self.gen_buffer(),
            count: 0,
            phantom: PhantomData
        }
    }
}

#[macro_export]
macro_rules! attrib_pointers {
    ($gl:expr, $vbo:expr, {
        $($field_name:ident => $field_attrib:expr),*
    }) => {
        {
            let mut binder = $crate::AttribBinder::new();
            $(binder.add(stringify!($field_name), $field_attrib).unwrap());*;
            binder
        }
    }
}

#[macro_export]
macro_rules! bind_attrib_pointers {
    ($gl:expr, $vbo:expr, {
        $($field_name:ident => $field_attrib:expr),*
    }) => {
        {
            let vbo = $vbo;
            let binder = {
                attrib_pointers!($gl, vbo, {
                    $($field_name => $field_attrib),*
                })
            };
            vbo.bind_attrib_pointers(binder);
        }
    }
}

#[macro_export]
macro_rules! bind_vertex_buffer {
    ($gl:expr, $vbo:expr) => {
        {
            use std::borrow::BorrowMut;
            let mut vbo = $vbo;

            let (gl_vbo, _) = $gl.bind_vertex_buffer(vbo.borrow_mut());
            gl_vbo
        }
    }
}

#[macro_export]
macro_rules! bind_index_buffer {
    ($gl:expr, $ibo:expr) => {
        {
            use std::borrow::BorrowMut;
            let mut ibo = $ibo;

            let (gl_ibo, _) = $gl.bind_index_buffer(ibo.borrow_mut());
            gl_ibo
        }
    }
}
<|endoftext|>


<commit_before>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

use crate::DiskType;

use crate::sys::disk::{new_disk, Disk};
use crate::sys::processor::{self, Processor, Query};

use std::collections::HashMap;
use std::ffi::OsStr;
use std::mem::{size_of, zeroed};

use winapi::{ctypes::c_void, um::winbase::DRIVE_REMOVABLE};

use winapi::shared::minwindef::{BYTE, DWORD, MAX_PATH, TRUE};
use winapi::um::fileapi::{
    CreateFileW, GetDriveTypeW, GetLogicalDrives, GetVolumeInformationW, OPEN_EXISTING,
};
use winapi::um::handleapi::CloseHandle;
use winapi::um::handleapi::INVALID_HANDLE_VALUE;
use winapi::um::ioapiset::DeviceIoControl;
use winapi::um::sysinfoapi::{GetSystemInfo, SYSTEM_INFO};
use winapi::um::winbase::DRIVE_FIXED;
use winapi::um::winioctl::{
    DISK_GEOMETRY, IOCTL_DISK_GET_DRIVE_GEOMETRY, IOCTL_STORAGE_QUERY_PROPERTY,
};
use winapi::um::winnt::{BOOLEAN, FILE_SHARE_READ, FILE_SHARE_WRITE, HANDLE};

pub struct KeyHandler {
    pub unique_id: String,
    pub win_key: Vec<u16>,
}

impl KeyHandler {
    pub fn new(unique_id: String, win_key: Vec<u16>) -> KeyHandler {
        KeyHandler { unique_id, win_key }
    }
}

pub fn init_processors() -> (Vec<Processor>, String, String) {
    unsafe {
        let mut sys_info: SYSTEM_INFO = zeroed();
        GetSystemInfo(&mut sys_info);
        let (vendor_id, brand) = processor::get_vendor_id_and_brand(&sys_info);
        let frequencies = processor::get_frequencies(sys_info.dwNumberOfProcessors as usize);
        let mut ret = Vec::with_capacity(sys_info.dwNumberOfProcessors as usize + 1);
        for nb in 0..sys_info.dwNumberOfProcessors {
            ret.push(Processor::new_with_values(
                &format!("CPU {}", nb + 1),
                vendor_id.clone(),
                brand.clone(),
                frequencies[nb as usize],
            ));
        }
        (ret, vendor_id, brand)
    }
}

pub unsafe fn open_drive(drive_name: &[u16], open_rights: DWORD) -> HANDLE {
    CreateFileW(
        drive_name.as_ptr(),
        open_rights,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        std::ptr::null_mut(),
        OPEN_EXISTING,
        0,
        std::ptr::null_mut(),
    )
}

pub unsafe fn get_drive_size(handle: HANDLE) -> u64 {
    let mut pdg: DISK_GEOMETRY = std::mem::zeroed();
    let mut junk = 0;
    let result = DeviceIoControl(
        handle,
        IOCTL_DISK_GET_DRIVE_GEOMETRY,
        std::ptr::null_mut(),
        0,
        &mut pdg as *mut DISK_GEOMETRY as *mut c_void,
        size_of::<DISK_GEOMETRY>() as DWORD,
        &mut junk,
        std::ptr::null_mut(),
    );
    if result == TRUE {
        *pdg.Cylinders.QuadPart() as u64
            * pdg.TracksPerCylinder as u64
            * pdg.SectorsPerTrack as u64
            * pdg.BytesPerSector as u64
    } else {
        0
    }
}

pub unsafe fn get_disks() -> Vec<Disk> {
    let drives = GetLogicalDrives();
    if drives == 0 {
        return Vec::new();
    }

    #[cfg(feature = "multithread")]
    use rayon::iter::ParallelIterator;

    crate::utils::into_iter(0..size_of::<DWORD>() * 8)
        .filter_map(|x| {
            if (drives >> x) & 1 == 0 {
                return None;
            }
            let mount_point = [b'A' as u16 + x as u16, b':' as u16, b'\\' as u16, 0];

            let drive_type = GetDriveTypeW(mount_point.as_ptr());

            if drive_type != DRIVE_FIXED && drive_type != DRIVE_REMOVABLE {
                return None;
            }
            let mut name = [0u16; MAX_PATH + 1];
            let mut file_system = [0u16; 32];
            if GetVolumeInformationW(
                mount_point.as_ptr(),
                name.as_mut_ptr(),
                name.len() as DWORD,
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                file_system.as_mut_ptr(),
                file_system.len() as DWORD,
            ) == 0
            {
                return None;
            }
            let mut pos = 0;
            for x in name.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let name = String::from_utf16_lossy(&name[..pos]);
            let name = OsStr::new(&name);

            pos = 0;
            for x in file_system.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let file_system: Vec<u8> = file_system[..pos].iter().map(|x| *x as u8).collect();

            let drive_name = [
                b'\\' as u16,
                b'\\' as u16,
                b'.' as u16,
                b'\\' as u16,
                b'A' as u16 + x as u16,
                b':' as u16,
                0,
            ];
            let handle = open_drive(&drive_name, 0);
            if handle == INVALID_HANDLE_VALUE {
                CloseHandle(handle);
                return new_disk(name, &mount_point, &file_system, DiskType::Unknown(-1), 0);
            }
            let disk_size = get_drive_size(handle);
            /*let mut spq_trim: ffi::STORAGE_PROPERTY_QUERY = std::mem::zeroed();
            spq_trim.PropertyId = ffi::StorageDeviceTrimProperty;
            spq_trim.QueryType = ffi::PropertyStandardQuery;
            let mut dtd: ffi::DEVICE_TRIM_DESCRIPTOR = std::mem::zeroed();*/
            #[allow(non_snake_case)]
            #[repr(C)]
            struct STORAGE_PROPERTY_QUERY {
                PropertyId: i32,
                QueryType: i32,
                AdditionalParameters: [BYTE; 1],
            }
            #[allow(non_snake_case)]
            #[repr(C)]
            struct DEVICE_TRIM_DESCRIPTOR {
                Version: DWORD,
                Size: DWORD,
                TrimEnabled: BOOLEAN,
            }
            let mut spq_trim = STORAGE_PROPERTY_QUERY {
                PropertyId: 8i32,
                QueryType: 0i32,
                AdditionalParameters: [0],
            };
            let mut dtd: DEVICE_TRIM_DESCRIPTOR = std::mem::zeroed();

            let mut dw_size = 0;
            if DeviceIoControl(
                handle,
                IOCTL_STORAGE_QUERY_PROPERTY,
                &mut spq_trim as *mut STORAGE_PROPERTY_QUERY as *mut c_void,
                size_of::<STORAGE_PROPERTY_QUERY>() as DWORD,
                &mut dtd as *mut DEVICE_TRIM_DESCRIPTOR as *mut c_void,
                size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD,
                &mut dw_size,
                std::ptr::null_mut(),
            ) == 0
                || dw_size != size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD
            {
                CloseHandle(handle);
                return new_disk(
                    name,
                    &mount_point,
                    &file_system,
                    DiskType::Unknown(-1),
                    disk_size,
                );
            }
            let is_ssd = dtd.TrimEnabled != 0;
            CloseHandle(handle);
            new_disk(
                name,
                &mount_point,
                &file_system,
                if is_ssd { DiskType::SSD } else { DiskType::HDD },
                disk_size,
            )
        })
        .collect::<Vec<_>>()
}

#[allow(non_snake_case)]
pub unsafe fn load_symbols() -> HashMap<String, u32> {
    use winapi::um::winreg::{RegQueryValueExA, HKEY_PERFORMANCE_DATA};

    let mut cbCounters = 0;
    let mut dwType = 0;
    let mut ret = HashMap::new();

    let _dwStatus = RegQueryValueExA(
        HKEY_PERFORMANCE_DATA,
        b"Counter 009\0".as_ptr() as *const _,
        std::ptr::null_mut(),
        &mut dwType as *mut i32 as *mut _,
        std::ptr::null_mut(),
        &mut cbCounters as *mut i32 as *mut _,
    );

    let mut lpmszCounters = Vec::with_capacity(cbCounters as usize);
    lpmszCounters.set_len(cbCounters as usize);
    let _dwStatus = RegQueryValueExA(
        HKEY_PERFORMANCE_DATA,
        b"Counter 009\0".as_ptr() as *const _,
        std::ptr::null_mut(),
        &mut dwType as *mut i32 as *mut _,
        lpmszCounters.as_mut_ptr(),
        &mut cbCounters as *mut i32 as *mut _,
    );
    for (pos, s) in lpmszCounters
        .split(|x| *x == 0)
        .filter(|x| !x.is_empty())
        .collect::<Vec<_>>()
        .chunks(2)
        .filter(|&x| x.len() == 2)
        .filter_map(
            |x| match (std::str::from_utf8(x[0]), String::from_utf8(x[1].to_vec())) {
                (Ok(n), Ok(s)) => {
                    if let Ok(n) = u32::from_str_radix(n, 10) {
                        Some((n, s))
                    } else {
                        None
                    }
                }
                _ => None,
            },
        )
    {
        ret.insert(s, pos as u32);
    }
    ret
}

#[allow(clippy::ptr_arg)]
pub fn get_translation(s: &String, map: &HashMap<String, u32>) -> Option<String> {
    use winapi::um::pdh::PdhLookupPerfNameByIndexW;

    if let Some(index) = map.get(s) {
        let mut size: usize = 0;
        unsafe {
            let _res = PdhLookupPerfNameByIndexW(
                std::ptr::null(),
                *index,
                std::ptr::null_mut(),
                &mut size as *mut usize as *mut _,
            );
            if size == 0 {
                return Some(String::new());
            } else {
                let mut v = Vec::with_capacity(size);
                v.set_len(size);
                let _res = PdhLookupPerfNameByIndexW(
                    std::ptr::null(),
                    *index,
                    v.as_mut_ptr() as *mut _,
                    &mut size as *mut usize as *mut _,
                );
                return Some(String::from_utf16(&v[..size - 1]).expect("invalid utf16"));
            }
        }
    }
    None
}

pub fn add_counter(
    s: String,
    query: &mut Query,
    keys: &mut Option<KeyHandler>,
    counter_name: String,
) {
    let mut full = s.encode_utf16().collect::<Vec<_>>();
    full.push(0);
    if query.add_counter(&counter_name, full.clone()) {
        *keys = Some(KeyHandler::new(counter_name, full));
    }
}
<commit_msg>Also expose type<commit_after>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

use crate::DiskType;

use crate::sys::disk::{new_disk, Disk};
use crate::sys::processor::{self, Processor, Query};

use std::collections::HashMap;
use std::ffi::OsStr;
use std::mem::{size_of, zeroed};

use winapi::{ctypes::c_void, um::winbase::DRIVE_REMOVABLE};

use winapi::shared::minwindef::{BYTE, DWORD, MAX_PATH, TRUE};
use winapi::um::fileapi::{
    CreateFileW, GetDriveTypeW, GetLogicalDrives, GetVolumeInformationW, OPEN_EXISTING,
};
use winapi::um::handleapi::CloseHandle;
use winapi::um::handleapi::INVALID_HANDLE_VALUE;
use winapi::um::ioapiset::DeviceIoControl;
use winapi::um::sysinfoapi::{GetSystemInfo, SYSTEM_INFO};
use winapi::um::winbase::DRIVE_FIXED;
use winapi::um::winioctl::{
    DISK_GEOMETRY, IOCTL_DISK_GET_DRIVE_GEOMETRY, IOCTL_STORAGE_QUERY_PROPERTY,
};
use winapi::um::winnt::{BOOLEAN, FILE_SHARE_READ, FILE_SHARE_WRITE, HANDLE};

pub struct KeyHandler {
    pub unique_id: String,
    pub win_key: Vec<u16>,
}

impl KeyHandler {
    pub fn new(unique_id: String, win_key: Vec<u16>) -> KeyHandler {
        KeyHandler { unique_id, win_key }
    }
}

pub fn init_processors() -> (Vec<Processor>, String, String) {
    unsafe {
        let mut sys_info: SYSTEM_INFO = zeroed();
        GetSystemInfo(&mut sys_info);
        let (vendor_id, brand) = processor::get_vendor_id_and_brand(&sys_info);
        let frequencies = processor::get_frequencies(sys_info.dwNumberOfProcessors as usize);
        let mut ret = Vec::with_capacity(sys_info.dwNumberOfProcessors as usize + 1);
        for nb in 0..sys_info.dwNumberOfProcessors {
            ret.push(Processor::new_with_values(
                &format!("CPU {}", nb + 1),
                vendor_id.clone(),
                brand.clone(),
                frequencies[nb as usize],
            ));
        }
        (ret, vendor_id, brand)
    }
}

pub unsafe fn open_drive(drive_name: &[u16], open_rights: DWORD) -> HANDLE {
    CreateFileW(
        drive_name.as_ptr(),
        open_rights,
        FILE_SHARE_READ | FILE_SHARE_WRITE,
        std::ptr::null_mut(),
        OPEN_EXISTING,
        0,
        std::ptr::null_mut(),
    )
}

pub unsafe fn get_drive_size(handle: HANDLE) -> u64 {
    let mut pdg: DISK_GEOMETRY = std::mem::zeroed();
    let mut junk = 0;
    let result = DeviceIoControl(
        handle,
        IOCTL_DISK_GET_DRIVE_GEOMETRY,
        std::ptr::null_mut(),
        0,
        &mut pdg as *mut DISK_GEOMETRY as *mut c_void,
        size_of::<DISK_GEOMETRY>() as DWORD,
        &mut junk,
        std::ptr::null_mut(),
    );
    if result == TRUE {
        *pdg.Cylinders.QuadPart() as u64
            * pdg.TracksPerCylinder as u64
            * pdg.SectorsPerTrack as u64
            * pdg.BytesPerSector as u64
    } else {
        0
    }
}

pub unsafe fn get_disks() -> Vec<Disk> {
    let drives = GetLogicalDrives();
    if drives == 0 {
        return Vec::new();
    }

    #[cfg(feature = "multithread")]
    use rayon::iter::ParallelIterator;

    crate::utils::into_iter(0..size_of::<DWORD>() * 8)
        .filter_map(|x| {
            if (drives >> x) & 1 == 0 {
                return None;
            }
            let mount_point = [b'A' as u16 + x as u16, b':' as u16, b'\\' as u16, 0];

            let drive_type = GetDriveTypeW(mount_point.as_ptr());

            let is_removable = drive_type == DRIVE_REMOVABLE;

            if drive_type != DRIVE_FIXED && drive_type != DRIVE_REMOVABLE {
                return None;
            }
            let mut name = [0u16; MAX_PATH + 1];
            let mut file_system = [0u16; 32];
            if GetVolumeInformationW(
                mount_point.as_ptr(),
                name.as_mut_ptr(),
                name.len() as DWORD,
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                std::ptr::null_mut(),
                file_system.as_mut_ptr(),
                file_system.len() as DWORD,
            ) == 0
            {
                return None;
            }
            let mut pos = 0;
            for x in name.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let name = String::from_utf16_lossy(&name[..pos]);
            let name = OsStr::new(&name);

            pos = 0;
            for x in file_system.iter() {
                if *x == 0 {
                    break;
                }
                pos += 1;
            }
            let file_system: Vec<u8> = file_system[..pos].iter().map(|x| *x as u8).collect();

            let drive_name = [
                b'\\' as u16,
                b'\\' as u16,
                b'.' as u16,
                b'\\' as u16,
                b'A' as u16 + x as u16,
                b':' as u16,
                0,
            ];
            let handle = open_drive(&drive_name, 0);
            if handle == INVALID_HANDLE_VALUE {
                CloseHandle(handle);
                return new_disk(
                    name,
                    &mount_point,
                    &file_system,
                    if is_removable {
                        DiskType::Removable
                    } else {
                        DiskType::Unknown(-1)
                    },
                    0,
                );
            }
            let disk_size = get_drive_size(handle);
            /*let mut spq_trim: ffi::STORAGE_PROPERTY_QUERY = std::mem::zeroed();
            spq_trim.PropertyId = ffi::StorageDeviceTrimProperty;
            spq_trim.QueryType = ffi::PropertyStandardQuery;
            let mut dtd: ffi::DEVICE_TRIM_DESCRIPTOR = std::mem::zeroed();*/
            #[allow(non_snake_case)]
            #[repr(C)]
            struct STORAGE_PROPERTY_QUERY {
                PropertyId: i32,
                QueryType: i32,
                AdditionalParameters: [BYTE; 1],
            }
            #[allow(non_snake_case)]
            #[repr(C)]
            struct DEVICE_TRIM_DESCRIPTOR {
                Version: DWORD,
                Size: DWORD,
                TrimEnabled: BOOLEAN,
            }
            let mut spq_trim = STORAGE_PROPERTY_QUERY {
                PropertyId: 8i32,
                QueryType: 0i32,
                AdditionalParameters: [0],
            };
            let mut dtd: DEVICE_TRIM_DESCRIPTOR = std::mem::zeroed();

            let mut dw_size = 0;
            if DeviceIoControl(
                handle,
                IOCTL_STORAGE_QUERY_PROPERTY,
                &mut spq_trim as *mut STORAGE_PROPERTY_QUERY as *mut c_void,
                size_of::<STORAGE_PROPERTY_QUERY>() as DWORD,
                &mut dtd as *mut DEVICE_TRIM_DESCRIPTOR as *mut c_void,
                size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD,
                &mut dw_size,
                std::ptr::null_mut(),
            ) == 0
                || dw_size != size_of::<DEVICE_TRIM_DESCRIPTOR>() as DWORD
            {
                CloseHandle(handle);
                return new_disk(
                    name,
                    &mount_point,
                    &file_system,
                    if is_removable {
                        DiskType::Removable
                    } else {
                        DiskType::Unknown(-1)
                    },
                    disk_size,
                );
            }
            let is_ssd = dtd.TrimEnabled != 0;
            CloseHandle(handle);
            new_disk(
                name,
                &mount_point,
                &file_system,
                match (is_removable, is_ssd) {
                    (true, _) => DiskType::Removable,
                    (false, true) => DiskType::SSD,
                    (false, false) => DiskType::HDD,
                },
                disk_size,
            )
        })
        .collect::<Vec<_>>()
}

#[allow(non_snake_case)]
pub unsafe fn load_symbols() -> HashMap<String, u32> {
    use winapi::um::winreg::{RegQueryValueExA, HKEY_PERFORMANCE_DATA};

    let mut cbCounters = 0;
    let mut dwType = 0;
    let mut ret = HashMap::new();

    let _dwStatus = RegQueryValueExA(
        HKEY_PERFORMANCE_DATA,
        b"Counter 009\0".as_ptr() as *const _,
        std::ptr::null_mut(),
        &mut dwType as *mut i32 as *mut _,
        std::ptr::null_mut(),
        &mut cbCounters as *mut i32 as *mut _,
    );

    let mut lpmszCounters = Vec::with_capacity(cbCounters as usize);
    lpmszCounters.set_len(cbCounters as usize);
    let _dwStatus = RegQueryValueExA(
        HKEY_PERFORMANCE_DATA,
        b"Counter 009\0".as_ptr() as *const _,
        std::ptr::null_mut(),
        &mut dwType as *mut i32 as *mut _,
        lpmszCounters.as_mut_ptr(),
        &mut cbCounters as *mut i32 as *mut _,
    );
    for (pos, s) in lpmszCounters
        .split(|x| *x == 0)
        .filter(|x| !x.is_empty())
        .collect::<Vec<_>>()
        .chunks(2)
        .filter(|&x| x.len() == 2)
        .filter_map(
            |x| match (std::str::from_utf8(x[0]), String::from_utf8(x[1].to_vec())) {
                (Ok(n), Ok(s)) => {
                    if let Ok(n) = u32::from_str_radix(n, 10) {
                        Some((n, s))
                    } else {
                        None
                    }
                }
                _ => None,
            },
        )
    {
        ret.insert(s, pos as u32);
    }
    ret
}

#[allow(clippy::ptr_arg)]
pub fn get_translation(s: &String, map: &HashMap<String, u32>) -> Option<String> {
    use winapi::um::pdh::PdhLookupPerfNameByIndexW;

    if let Some(index) = map.get(s) {
        let mut size: usize = 0;
        unsafe {
            let _res = PdhLookupPerfNameByIndexW(
                std::ptr::null(),
                *index,
                std::ptr::null_mut(),
                &mut size as *mut usize as *mut _,
            );
            if size == 0 {
                return Some(String::new());
            } else {
                let mut v = Vec::with_capacity(size);
                v.set_len(size);
                let _res = PdhLookupPerfNameByIndexW(
                    std::ptr::null(),
                    *index,
                    v.as_mut_ptr() as *mut _,
                    &mut size as *mut usize as *mut _,
                );
                return Some(String::from_utf16(&v[..size - 1]).expect("invalid utf16"));
            }
        }
    }
    None
}

pub fn add_counter(
    s: String,
    query: &mut Query,
    keys: &mut Option<KeyHandler>,
    counter_name: String,
) {
    let mut full = s.encode_utf16().collect::<Vec<_>>();
    full.push(0);
    if query.add_counter(&counter_name, full.clone()) {
        *keys = Some(KeyHandler::new(counter_name, full));
    }
}
<|endoftext|>


<commit_before>use protobuf::descriptor::field_descriptor_proto::Type;
use protobuf::reflect::FieldDescriptor;
use protobuf::reflect::MessageDescriptor;
use protobuf_codegen::Codegen;
use protobuf_codegen::Customize;
use protobuf_codegen::CustomizeCallback;

fn main() {
    struct GenSerde;

    impl CustomizeCallback for GenSerde {
        fn message(&self, _message: &MessageDescriptor) -> Customize {
            Customize::default().before("#[derive(::serde::Serialize, ::serde::Deserialize)]")
        }

        fn field(&self, field: &FieldDescriptor) -> Customize {
            if field.get_proto().get_field_type() == Type::TYPE_ENUM {
                /// `EnumOrUnknown` is not a part of rust-protobuf, so external serializer is needed.
                Customize::default().before(
                    "#[serde(serialize_with = \"crate::serialize_enum_or_unknown\", deserialize_with = \"crate::deserialize_enum_or_unknown\")]")
            } else {
                Customize::default()
            }
        }

        fn special_field(&self, _message: &MessageDescriptor, _field: &str) -> Customize {
            Customize::default().before("#[serde(skip)]")
        }
    }

    Codegen::new()
        .cargo_out_dir("protos")
        .include("src")
        .inputs(&["src/customize_example.proto"])
        .customize_callback(GenSerde)
        .run_from_script();
}
<commit_msg>Fix doc syntax<commit_after>use protobuf::descriptor::field_descriptor_proto::Type;
use protobuf::reflect::FieldDescriptor;
use protobuf::reflect::MessageDescriptor;
use protobuf_codegen::Codegen;
use protobuf_codegen::Customize;
use protobuf_codegen::CustomizeCallback;

fn main() {
    struct GenSerde;

    impl CustomizeCallback for GenSerde {
        fn message(&self, _message: &MessageDescriptor) -> Customize {
            Customize::default().before("#[derive(::serde::Serialize, ::serde::Deserialize)]")
        }

        fn field(&self, field: &FieldDescriptor) -> Customize {
            if field.get_proto().get_field_type() == Type::TYPE_ENUM {
                // `EnumOrUnknown` is not a part of rust-protobuf, so external serializer is needed.
                Customize::default().before(
                    "#[serde(serialize_with = \"crate::serialize_enum_or_unknown\", deserialize_with = \"crate::deserialize_enum_or_unknown\")]")
            } else {
                Customize::default()
            }
        }

        fn special_field(&self, _message: &MessageDescriptor, _field: &str) -> Customize {
            Customize::default().before("#[serde(skip)]")
        }
    }

    Codegen::new()
        .cargo_out_dir("protos")
        .include("src")
        .inputs(&["src/customize_example.proto"])
        .customize_callback(GenSerde)
        .run_from_script();
}
<|endoftext|>