<commit_before>import result::{ok, err};
import std::io;
import std::json;
import std::map;
import std::map::map;
import zmq::{context, socket, error};

import json = std::json::json;

iface transport {
    fn head(path: str) -> response::t;
    fn get(path: str) -> response::t;
    fn put(path: str, doc: map<str, json>) -> response::t;
    fn post(path: str, doc: map<str, json>) -> response::t;
    fn delete(path: str) -> response::t;
}

type client = { transport: transport };

fn mk_client(transport: transport) -> client {
    { transport: transport }
}

impl client for client {
    fn get(index: str, typ: str, id: str) -> response::t {
        let path = index + "/" + typ + "/" + id;
        self.transport.get(path)
    }
}

enum consistency { CONSISTENCY_DEFAULT, ONE, QUORUM, ALL }
enum replication { REPLICATION_DEFAULT, SYNC, ASYNC }
enum op_type { CREATE, INDEX }
enum version_type { INTERNAL, EXTERNAL }

type index_builder = {
    index: str,
    typ: str,
    mut id: option<str>,
    mut routing: option<str>,
    mut parent: option<str>,
    mut timestamp: option<str>,
    mut ttl: option<str>,
    mut op_type: op_type,
    mut refresh: bool,
    mut version: uint,
    mut version_type: version_type,
    mut percolate: option<str>,
    mut consistency: consistency,
    mut replication: replication,
    mut source: option<map<str, json>>,
};

fn mk_index_builder(index: str, typ: str) -> index_builder {
    {
        index: index,
        typ: typ,
        mut id: none,
        mut routing: none,
        mut parent: none,
        mut timestamp: none,
        mut ttl: none,
        mut op_type: INDEX,
        mut refresh: false,
        mut version: 0u,
        mut version_type: INTERNAL,
        mut percolate: none,
        mut consistency: CONSISTENCY_DEFAULT,
        mut replication: REPLICATION_DEFAULT,
        mut source: none,
    }
}

impl index_builder for index_builder {
    fn set_id(id: str) -> index_builder { self.id = some(id); self }
    fn set_routing(routing: str) -> index_builder {
        self.routing = some(routing);
        self
    }
    fn set_parent(parent: str) -> index_builder {
        self.parent = some(parent);
        self
    }
    fn set_timestamp(timestamp: str) -> index_builder {
        self.timestamp = some(timestamp);
        self
    }
    fn set_ttl(ttl: str) -> index_builder {
        self.ttl = some(ttl);
        self
    }
    fn set_op_type(op_type: op_type) -> index_builder {
        self.op_type = op_type;
        self
    }
    fn set_refresh() -> index_builder { self.refresh = true; self }
    fn set_version(version: uint) -> index_builder {
        self.version = version;
        self
    }
    fn set_version_type(version_type: version_type) -> index_builder {
        self.version_type = version_type;
        self
    }
    fn set_percolate(percolate: str) -> index_builder {
        self.percolate = some(percolate);
        self
    }
    fn set_consistency(consistency: consistency) -> index_builder {
        self.consistency = consistency;
        self
    }
    fn set_replication(replication: replication) -> index_builder {
        self.replication = replication;
        self
    }
    fn set_source(source: map<str, json>) -> index_builder {
        self.source = some(source);
        self
    }
    fn execute(client: client) -> response::t {
        let path = [self.index, self.typ];
        alt self.id {
          none {}
          some(id) { vec::push(path, id); }
        }

        let path = str::connect(path, "/");
        let params = [];

        alt self.routing {
          none {}
          some(routing) { vec::push(params, "routing=" + routing); }
        }

        alt self.parent {
          none {}
          some(parent) { vec::push(params, "parent=" + parent); }
        }

        alt self.timestamp {
          none {}
          some(timestamp) { vec::push(params, "timestamp=" + timestamp); }
        }

        alt self.ttl {
          none {}
          some(ttl) { vec::push(params, "ttl=" + ttl); }
        }

        alt self.op_type {
          CREATE { vec::push(params, "op_type=create"); }
          INDEX { vec::push(params, "op_type=index"); }
        }

        if self.refresh {
            vec::push(params, "refresh=true");
        }

        alt self.percolate {
          none {}
          some(percolate) { vec::push(params, "percolate=" + percolate); }
        }

        alt self.consistency {
          CONSISTENCY_DEFAULT {}
          ONE { vec::push(params, "consistency=one"); }
          QUORUM { vec::push(params, "consistency=quorum"); }
          ALL { vec::push(params, "consistency=all"); }
        }

        alt self.replication {
          REPLICATION_DEFAULT {}
          SYNC { vec::push(params, "replication=sync"); }
          ASYNC { vec::push(params, "replication=async"); }
        }

        if vec::is_not_empty(params) {
            path += "?" + str::connect(params, "&");
        }

        let source = alt self.source {
          none { map::new_str_hash() }
          some(source) { source }
        };

        alt self.id {
          none { client.transport.put(path, source) }
          some(_) { client.transport.post(path, source) }
        }
    }
}

type json_dict_builder = { dict: map<str, json> };

fn mk_json_dict_builder() -> json_dict_builder {
    { dict: map::new_str_hash() }
}

impl json_dict_builder for json_dict_builder {
    fn field_float(key: str, value: float) -> json_dict_builder {
        self.dict.insert(key, json::num(value));
        self
    }
    fn field_str(key: str, value: str) -> json_dict_builder {
        self.dict.insert(key, json::string(value));
        self
    }
    fn field_bool(key: str, value: bool) -> json_dict_builder {
        self.dict.insert(key, json::boolean(value));
        self
    }
    fn field_null(key: str) -> json_dict_builder {
        self.dict.insert(key, json::null);
        self
    }
    fn field_dict(key: str, f: fn(json_dict_builder))
      -> json_dict_builder {
        let builder = mk_json_dict_builder();
        f(builder);
        self.dict.insert(key, json::dict(builder.dict));
        self
    }
    fn field_list(key: str, f: fn(json_list_builder)) -> json_dict_builder {
        let builder = mk_json_list_builder();
        f(builder);
        self.dict.insert(key, json::list(builder.list));
        self
    }
}

type json_list_builder = { mut list: [json] };

fn mk_json_list_builder() -> json_list_builder {
    { mut list: [] }
}

impl json_list_builder for json_list_builder {
    fn add_float(value: float) -> json_list_builder {
        vec::push(self.list, json::num(value));
        self
    }
    fn add_str(value: str) -> json_list_builder {
        vec::push(self.list, json::string(value));
        self
    }
    fn add_bool(value: bool) -> json_list_builder {
        vec::push(self.list, json::boolean(value));
        self
    }
    fn add_null() -> json_list_builder {
        vec::push(self.list, json::null);
        self
    }
    fn add_dict(f: fn(json_dict_builder)) -> json_list_builder {
        let builder = mk_json_dict_builder();
        f(builder);
        vec::push(self.list, json::dict(builder.dict));
        self
    }
    fn add_list(f: fn(json_list_builder)) -> json_list_builder {
        let builder = mk_json_list_builder();
        f(builder);
        vec::push(self.list, json::list(builder.list));
        self
    }
}

#[doc = "Transport to talk to Elasticsearch with zeromq"]
type zmq_transport = { socket: zmq::socket };

#[doc = "Zeromq transport implementation"]
impl of transport for zmq_transport {
    fn head(path: str) -> response::t { self.send("HEAD|" + path) }
    fn get(path: str) -> response::t { self.send("GET|" + path) }
    fn put(path: str, doc: map<str, json>) -> response::t {
        self.send("PUT|" + path + "|" + json::to_str(json::dict(doc)))
    }
    fn post(path: str, doc: map<str, json>) -> response::t {
        self.send("POST|" + path + "|" + json::to_str(json::dict(doc)))
    }
    fn delete(path: str) -> response::t {
        self.send("DELETE|" + path)
    }

    fn send(request: str) -> response::t {
        str::as_bytes(request) { |bytes|
            alt self.socket.send_between(bytes, 0u, str::len(request), 0) {
              ok(()) {}
              err(e) { fail e.to_str(); }
            }
        }

        alt self.socket.recv(0) {
          ok(msg) { response::parse(msg) }
          err(e) { fail e.to_str(); }
        }
    }
}

#[doc = "Create a zeromq transport to Elasticsearch"]
fn mk_zmq_transport(ctx: zmq::context, addr: str) -> transport {
    let socket = alt ctx.socket(zmq::REQ) {
      ok(socket) { socket }
      err(e) { fail e.to_str() }
    };

    str::as_bytes(addr) { |bytes|
        alt socket.connect(bytes) {
          ok(()) {}
          err(e) { fail e.to_str(); }
        }
    }

    { socket: socket } as transport
}

#[doc = "Helper function to creating a client with zeromq"]
fn connect_with_zmq(ctx: zmq::context, addr: str) -> client {
    let transport = mk_zmq_transport(ctx, addr);
    mk_client(transport)
}

mod response {
    type t = {
        code: uint,
        status: str,
        doc: json,
    };

    fn parse(msg: [u8]) -> t {
        io::println(#fmt("parse: %?", str::from_bytes(msg)));

        let end = vec::len(msg);

        let (start, code) = parse_code(msg, end);
        let (start, status) = parse_status(msg, start, end);
        let doc = parse_doc(msg, start, end);

        { code: code, status: status, doc: doc }
    }

    fn parse_code(msg: [u8], end: uint) -> (uint, uint) {
        alt vec::position_from(msg, 0u, end) { |c| c == '|' as u8 } {
          none { fail "invalid response" }
          some(i) {
            alt uint::parse_buf(vec::slice(msg, 0u, i), 10u) {
              some(code) { (i + 1u, code) }
              none { fail "invalid status code" }
            }
          }
        }
    }

    fn parse_status(msg: [u8], start: uint, end: uint) -> (uint, str) {
        alt vec::position_from(msg, start, end) { |c| c == '|' as u8 } {
          none { fail "invalid response" }
          some(i) { (i + 1u, str::from_bytes(vec::slice(msg, start, i))) }
        }
    }

    fn parse_doc(msg: [u8], start: uint, end: uint) -> json {
        if start == end { ret json::null; }

        io::with_bytes_reader_between(msg, start, end) { |rdr|
            alt json::from_reader(rdr) {
              ok(json) { json }
              err({line, col, msg}) { fail #fmt("%u:%u: %s", line, col, msg); }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    import std::map;

    #[test]
    fn test() {
        let ctx = alt zmq::init(1) {
          ok(ctx) { ctx }
          err(e) { fail e.to_str(); }
        };

        let client = connect_with_zmq(ctx, "tcp://localhost:9700");
        io::println(#fmt("%?\n", client.transport.head("/")));
        io::println(#fmt("%?\n", client.transport.get("/")));

        io::println(#fmt("%?\n", client.transport.get("/test/test/1")));

        mk_index_builder("test", "test")
          .set_id("1")
          .set_version(2u)
          .set_source(mk_json_dict_builder()
              .field_float("foo", 5.0)
              .field_str("bar", "wee")
              .field_dict("baz") { |bld|
                  bld.field_float("a", 2.0);
              }
              .field_list("boo") { |bld|
                  bld.add_float(1.0).add_str("zzz");
              }
              .dict
          )
          .execute(client);

        //io::println(#fmt("%?\n", client.get("test", "test", "1")));
        //io::println(#fmt("%?\n", client.transport.put("/test/test/1", doc)));
        io::println(#fmt("%?\n", client.transport.get("/test/test/1")));

        io::println(#fmt("%?\n", client.transport.delete("/test/test/1")));
    }
}
<commit_msg>add a client.delete function<commit_after>import result::{ok, err};
import std::io;
import std::json;
import std::map;
import std::map::map;
import zmq::{context, socket, error};

import json = std::json::json;

iface transport {
    fn head(path: str) -> response::t;
    fn get(path: str) -> response::t;
    fn put(path: str, doc: map<str, json>) -> response::t;
    fn post(path: str, doc: map<str, json>) -> response::t;
    fn delete(path: str) -> response::t;
}

type client = { transport: transport };

fn mk_client(transport: transport) -> client {
    { transport: transport }
}

impl client for client {
    fn get(index: str, typ: str, id: str) -> response::t {
        let path = index + "/" + typ + "/" + id;
        self.transport.get(path)
    }
    fn delete(index: str, typ: str, id: str) -> response::t {
        let path = index + "/" + typ + "/" + id;
        self.transport.delete(path)
    }
}

enum consistency { CONSISTENCY_DEFAULT, ONE, QUORUM, ALL }
enum replication { REPLICATION_DEFAULT, SYNC, ASYNC }
enum op_type { CREATE, INDEX }
enum version_type { INTERNAL, EXTERNAL }

type index_builder = {
    index: str,
    typ: str,
    mut id: option<str>,
    mut routing: option<str>,
    mut parent: option<str>,
    mut timestamp: option<str>,
    mut ttl: option<str>,
    mut op_type: op_type,
    mut refresh: bool,
    mut version: uint,
    mut version_type: version_type,
    mut percolate: option<str>,
    mut consistency: consistency,
    mut replication: replication,
    mut source: option<map<str, json>>,
};

fn mk_index_builder(index: str, typ: str) -> index_builder {
    {
        index: index,
        typ: typ,
        mut id: none,
        mut routing: none,
        mut parent: none,
        mut timestamp: none,
        mut ttl: none,
        mut op_type: INDEX,
        mut refresh: false,
        mut version: 0u,
        mut version_type: INTERNAL,
        mut percolate: none,
        mut consistency: CONSISTENCY_DEFAULT,
        mut replication: REPLICATION_DEFAULT,
        mut source: none,
    }
}

impl index_builder for index_builder {
    fn set_id(id: str) -> index_builder { self.id = some(id); self }
    fn set_routing(routing: str) -> index_builder {
        self.routing = some(routing);
        self
    }
    fn set_parent(parent: str) -> index_builder {
        self.parent = some(parent);
        self
    }
    fn set_timestamp(timestamp: str) -> index_builder {
        self.timestamp = some(timestamp);
        self
    }
    fn set_ttl(ttl: str) -> index_builder {
        self.ttl = some(ttl);
        self
    }
    fn set_op_type(op_type: op_type) -> index_builder {
        self.op_type = op_type;
        self
    }
    fn set_refresh() -> index_builder { self.refresh = true; self }
    fn set_version(version: uint) -> index_builder {
        self.version = version;
        self
    }
    fn set_version_type(version_type: version_type) -> index_builder {
        self.version_type = version_type;
        self
    }
    fn set_percolate(percolate: str) -> index_builder {
        self.percolate = some(percolate);
        self
    }
    fn set_consistency(consistency: consistency) -> index_builder {
        self.consistency = consistency;
        self
    }
    fn set_replication(replication: replication) -> index_builder {
        self.replication = replication;
        self
    }
    fn set_source(source: map<str, json>) -> index_builder {
        self.source = some(source);
        self
    }
    fn execute(client: client) -> response::t {
        let path = [self.index, self.typ];
        alt self.id {
          none {}
          some(id) { vec::push(path, id); }
        }

        let path = str::connect(path, "/");
        let params = [];

        alt self.routing {
          none {}
          some(routing) { vec::push(params, "routing=" + routing); }
        }

        alt self.parent {
          none {}
          some(parent) { vec::push(params, "parent=" + parent); }
        }

        alt self.timestamp {
          none {}
          some(timestamp) { vec::push(params, "timestamp=" + timestamp); }
        }

        alt self.ttl {
          none {}
          some(ttl) { vec::push(params, "ttl=" + ttl); }
        }

        alt self.op_type {
          CREATE { vec::push(params, "op_type=create"); }
          INDEX { vec::push(params, "op_type=index"); }
        }

        if self.refresh {
            vec::push(params, "refresh=true");
        }

        alt self.percolate {
          none {}
          some(percolate) { vec::push(params, "percolate=" + percolate); }
        }

        alt self.consistency {
          CONSISTENCY_DEFAULT {}
          ONE { vec::push(params, "consistency=one"); }
          QUORUM { vec::push(params, "consistency=quorum"); }
          ALL { vec::push(params, "consistency=all"); }
        }

        alt self.replication {
          REPLICATION_DEFAULT {}
          SYNC { vec::push(params, "replication=sync"); }
          ASYNC { vec::push(params, "replication=async"); }
        }

        if vec::is_not_empty(params) {
            path += "?" + str::connect(params, "&");
        }

        let source = alt self.source {
          none { map::new_str_hash() }
          some(source) { source }
        };

        alt self.id {
          none { client.transport.put(path, source) }
          some(_) { client.transport.post(path, source) }
        }
    }
}

type json_dict_builder = { dict: map<str, json> };

fn mk_json_dict_builder() -> json_dict_builder {
    { dict: map::new_str_hash() }
}

impl json_dict_builder for json_dict_builder {
    fn field_float(key: str, value: float) -> json_dict_builder {
        self.dict.insert(key, json::num(value));
        self
    }
    fn field_str(key: str, value: str) -> json_dict_builder {
        self.dict.insert(key, json::string(value));
        self
    }
    fn field_bool(key: str, value: bool) -> json_dict_builder {
        self.dict.insert(key, json::boolean(value));
        self
    }
    fn field_null(key: str) -> json_dict_builder {
        self.dict.insert(key, json::null);
        self
    }
    fn field_dict(key: str, f: fn(json_dict_builder))
      -> json_dict_builder {
        let builder = mk_json_dict_builder();
        f(builder);
        self.dict.insert(key, json::dict(builder.dict));
        self
    }
    fn field_list(key: str, f: fn(json_list_builder)) -> json_dict_builder {
        let builder = mk_json_list_builder();
        f(builder);
        self.dict.insert(key, json::list(builder.list));
        self
    }
}

type json_list_builder = { mut list: [json] };

fn mk_json_list_builder() -> json_list_builder {
    { mut list: [] }
}

impl json_list_builder for json_list_builder {
    fn add_float(value: float) -> json_list_builder {
        vec::push(self.list, json::num(value));
        self
    }
    fn add_str(value: str) -> json_list_builder {
        vec::push(self.list, json::string(value));
        self
    }
    fn add_bool(value: bool) -> json_list_builder {
        vec::push(self.list, json::boolean(value));
        self
    }
    fn add_null() -> json_list_builder {
        vec::push(self.list, json::null);
        self
    }
    fn add_dict(f: fn(json_dict_builder)) -> json_list_builder {
        let builder = mk_json_dict_builder();
        f(builder);
        vec::push(self.list, json::dict(builder.dict));
        self
    }
    fn add_list(f: fn(json_list_builder)) -> json_list_builder {
        let builder = mk_json_list_builder();
        f(builder);
        vec::push(self.list, json::list(builder.list));
        self
    }
}

#[doc = "Transport to talk to Elasticsearch with zeromq"]
type zmq_transport = { socket: zmq::socket };

#[doc = "Zeromq transport implementation"]
impl of transport for zmq_transport {
    fn head(path: str) -> response::t { self.send("HEAD|" + path) }
    fn get(path: str) -> response::t { self.send("GET|" + path) }
    fn put(path: str, doc: map<str, json>) -> response::t {
        self.send("PUT|" + path + "|" + json::to_str(json::dict(doc)))
    }
    fn post(path: str, doc: map<str, json>) -> response::t {
        self.send("POST|" + path + "|" + json::to_str(json::dict(doc)))
    }
    fn delete(path: str) -> response::t {
        self.send("DELETE|" + path)
    }

    fn send(request: str) -> response::t {
        str::as_bytes(request) { |bytes|
            alt self.socket.send_between(bytes, 0u, str::len(request), 0) {
              ok(()) {}
              err(e) { fail e.to_str(); }
            }
        }

        alt self.socket.recv(0) {
          ok(msg) { response::parse(msg) }
          err(e) { fail e.to_str(); }
        }
    }
}

#[doc = "Create a zeromq transport to Elasticsearch"]
fn mk_zmq_transport(ctx: zmq::context, addr: str) -> transport {
    let socket = alt ctx.socket(zmq::REQ) {
      ok(socket) { socket }
      err(e) { fail e.to_str() }
    };

    str::as_bytes(addr) { |bytes|
        alt socket.connect(bytes) {
          ok(()) {}
          err(e) { fail e.to_str(); }
        }
    }

    { socket: socket } as transport
}

#[doc = "Helper function to creating a client with zeromq"]
fn connect_with_zmq(ctx: zmq::context, addr: str) -> client {
    let transport = mk_zmq_transport(ctx, addr);
    mk_client(transport)
}

mod response {
    type t = {
        code: uint,
        status: str,
        doc: json,
    };

    fn parse(msg: [u8]) -> t {
        io::println(#fmt("parse: %?", str::from_bytes(msg)));

        let end = vec::len(msg);

        let (start, code) = parse_code(msg, end);
        let (start, status) = parse_status(msg, start, end);
        let doc = parse_doc(msg, start, end);

        { code: code, status: status, doc: doc }
    }

    fn parse_code(msg: [u8], end: uint) -> (uint, uint) {
        alt vec::position_from(msg, 0u, end) { |c| c == '|' as u8 } {
          none { fail "invalid response" }
          some(i) {
            alt uint::parse_buf(vec::slice(msg, 0u, i), 10u) {
              some(code) { (i + 1u, code) }
              none { fail "invalid status code" }
            }
          }
        }
    }

    fn parse_status(msg: [u8], start: uint, end: uint) -> (uint, str) {
        alt vec::position_from(msg, start, end) { |c| c == '|' as u8 } {
          none { fail "invalid response" }
          some(i) { (i + 1u, str::from_bytes(vec::slice(msg, start, i))) }
        }
    }

    fn parse_doc(msg: [u8], start: uint, end: uint) -> json {
        if start == end { ret json::null; }

        io::with_bytes_reader_between(msg, start, end) { |rdr|
            alt json::from_reader(rdr) {
              ok(json) { json }
              err({line, col, msg}) { fail #fmt("%u:%u: %s", line, col, msg); }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    import std::map;

    #[test]
    fn test() {
        let ctx = alt zmq::init(1) {
          ok(ctx) { ctx }
          err(e) { fail e.to_str(); }
        };

        let client = connect_with_zmq(ctx, "tcp://localhost:9700");
        io::println(#fmt("%?\n", client.transport.head("/")));
        io::println(#fmt("%?\n", client.transport.get("/")));

        io::println(#fmt("%?\n", client.transport.get("/test/test/1")));

        mk_index_builder("test", "test")
          .set_id("1")
          .set_version(2u)
          .set_source(mk_json_dict_builder()
              .field_float("foo", 5.0)
              .field_str("bar", "wee")
              .field_dict("baz") { |bld|
                  bld.field_float("a", 2.0);
              }
              .field_list("boo") { |bld|
                  bld.add_float(1.0).add_str("zzz");
              }
              .dict
          )
          .execute(client);

        //io::println(#fmt("%?\n", client.get("test", "test", "1")));
        //io::println(#fmt("%?\n", client.transport.put("/test/test/1", doc)));
        io::println(#fmt("%?\n", client.transport.get("/test/test/1")));

        io::println(#fmt("%?\n", client.delete("test", "test", "1")));
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    json_rpc::JsonRpcClientWrapper,
    validator_set::{validator_set_full_node_addresses, validator_set_validator_addresses},
};
use diem_config::{
    config::{RoleType, HANDSHAKE_VERSION},
    network_id::{NetworkContext, NetworkId},
};
use diem_crypto::{x25519, x25519::PRIVATE_KEY_SIZE};
use diem_management::error::Error;
use diem_network_address_encryption::Encryptor;
use diem_secure_storage::{InMemoryStorage, Storage};
use diem_types::{
    account_address,
    chain_id::ChainId,
    network_address::{
        encrypted::{Key, KeyVersion, KEY_LEN},
        NetworkAddress,
    },
    PeerId,
};
use fallible::copy_from_slice::copy_slice_to_vec;
use netcore::transport::tcp::{resolve_and_connect, TcpSocket};
use network::{
    noise::{HandshakeAuthMode, NoiseUpgrader},
    protocols::wire::handshake::v1::SupportedProtocols,
    transport::{upgrade_outbound, UpgradeContext, SUPPORTED_MESSAGING_PROTOCOL},
    ProtocolId,
};
use std::{collections::BTreeMap, sync::Arc};
use structopt::StructOpt;
use tokio::{runtime::Runtime, time::Duration};

const DEFAULT_TIMEOUT_SECONDS: u64 = 5;

#[derive(Debug, StructOpt)]
pub struct CheckEndpoint {
    /// `NetworkAddress` of remote server interface
    #[structopt(long)]
    address: NetworkAddress,
    /// `ChainId` of remote server
    #[structopt(long)]
    chain_id: ChainId,
    /// `NetworkId` of remote server interface
    #[structopt(long)]
    network_id: NetworkId,
    /// `PrivateKey` to connect to remote server
    #[structopt(long, parse(try_from_str = parse_private_key_hex))]
    private_key: x25519::PrivateKey,
    /// Optional number of seconds to timeout attempting to connect to endpoint
    #[structopt(long)]
    timeout_seconds: Option<u64>,
}

fn parse_private_key_hex(src: &str) -> Result<x25519::PrivateKey, Error> {
    let input = src.trim();
    if input.len() != 64 {
        return Err(Error::CommandArgumentError(
            "Invalid private key length, must be 64 hex characters".to_string(),
        ));
    }

    let value_slice = hex::decode(src.trim())
        .map_err(|_| Error::CommandArgumentError(format!("Not a valid private key: {}", src)))?;

    let mut value = [0; PRIVATE_KEY_SIZE];
    copy_slice_to_vec(&value_slice, &mut value)
        .map_err(|e| Error::CommandArgumentError(format!("{}", e)))?;

    Ok(x25519::PrivateKey::from(value))
}

impl CheckEndpoint {
    pub fn execute(self) -> Result<String, Error> {
        validate_address(&self.address)?;
        let (peer_id, public_key) = private_key_to_public_info(&self.private_key);
        check_endpoint(
            build_upgrade_context(self.chain_id, self.network_id, peer_id, self.private_key),
            self.address,
            peer_id,
            public_key,
            timeout_duration(self.timeout_seconds),
        )
    }
}

#[derive(Debug, StructOpt)]
pub struct CheckValidatorSetEndpoints {
    /// JSON-RPC Endpoint (e.g. http://localhost:8080)
    #[structopt(long, required_unless = "config")]
    json_server: String,
    /// Specifies whether or not to evaluate validators or fullnodes
    #[structopt(long)]
    role: RoleType,
    /// The expected on-chain key, only required for validator checks
    #[structopt(long, required_if("role", "validator"), parse(try_from_str = parse_validator_key_hex))]
    address_encryption_key: Option<Key>,
    /// The expected on-chain key version, only required for validator checks
    #[structopt(long, required_if("role", "validator"))]
    version: Option<KeyVersion>,
    /// `ChainId` of remote server
    #[structopt(long)]
    chain_id: ChainId,
    /// Private key to connect to remote server
    #[structopt(long, parse(try_from_str = parse_private_key_hex))]
    private_key: x25519::PrivateKey,
    /// Optional number of seconds to timeout attempting to connect to endpoint
    #[structopt(long)]
    timeout_seconds: Option<u64>,
}

fn parse_validator_key_hex(src: &str) -> Result<Key, Error> {
    let potential_err_msg = format!("Not a valid encryption key: {}", src);
    let value_slice =
        hex::decode(src.trim()).map_err(|_| Error::CommandArgumentError(potential_err_msg))?;

    let mut value = [0; KEY_LEN];
    copy_slice_to_vec(&value_slice, &mut value)
        .map_err(|e| Error::CommandArgumentError(format!("{}", e)))?;
    Ok(value)
}

impl CheckValidatorSetEndpoints {
    pub fn execute(self) -> Result<String, Error> {
        let is_validator = self.role.is_validator();
        let client = JsonRpcClientWrapper::new(self.json_server);

        let nodes = if is_validator {
            // Following unwraps shouldn't fail as it is in memory
            let mut encryptor = Encryptor::new(Storage::InMemoryStorage(InMemoryStorage::new()));
            encryptor.initialize().unwrap();
            encryptor
                .add_key(self.version.unwrap(), self.address_encryption_key.unwrap())
                .unwrap();
            encryptor
                .set_current_version(self.version.unwrap())
                .unwrap();

            validator_set_validator_addresses(client, &encryptor, None)?
        } else {
            validator_set_full_node_addresses(client, None)?
        };

        // Build a single upgrade context to run all the checks
        let network_id = if is_validator {
            NetworkId::Validator
        } else {
            NetworkId::Public
        };

        let (peer_id, public_key) = private_key_to_public_info(&self.private_key);
        let upgrade_context =
            build_upgrade_context(self.chain_id, network_id, peer_id, self.private_key);

        let timeout = timeout_duration(self.timeout_seconds);

        // Check all the addresses accordingly
        for (name, peer_id, addrs) in nodes {
            for addr in addrs {
                match check_endpoint(upgrade_context.clone(), addr, peer_id, public_key, timeout) {
                    Ok(_) => println!("{} -- good", name),
                    Err(err) => println!("{} -- bad -- {}", name, err),
                };
            }
        }

        Ok("Complete!".to_string())
    }
}

/// Builds a listener free noise connector
fn build_upgrade_context(
    chain_id: ChainId,
    network_id: NetworkId,
    peer_id: PeerId,
    private_key: x25519::PrivateKey,
) -> Arc<UpgradeContext> {
    // RoleType doesn't matter, but the `NetworkId` and `PeerId` are used in handshakes
    let network_context = Arc::new(NetworkContext::new(
        RoleType::FullNode,
        network_id.clone(),
        peer_id,
    ));

    // Let's make sure some protocol can be connected.  In the future we may want to allow for specifics
    let mut supported_protocols = BTreeMap::new();
    supported_protocols.insert(
        SUPPORTED_MESSAGING_PROTOCOL,
        SupportedProtocols::from(ProtocolId::all().iter()),
    );

    // Build the noise and network handshake, without running a full Noise server with listener
    Arc::new(UpgradeContext::new(
        NoiseUpgrader::new(
            network_context,
            private_key,
            // If we had an incoming message, auth mode would matter
            HandshakeAuthMode::server_only(),
        ),
        HANDSHAKE_VERSION,
        supported_protocols,
        chain_id,
        network_id,
    ))
}

fn timeout_duration(maybe_secs: Option<u64>) -> Duration {
    Duration::from_secs(if let Some(secs) = maybe_secs {
        secs
    } else {
        DEFAULT_TIMEOUT_SECONDS
    })
}

fn validate_address(address: &NetworkAddress) -> Result<(), Error> {
    if !address.is_diemnet_addr() {
        Err(Error::CommandArgumentError(
            "Address must have ip, tcp, noise key, and handshake".to_string(),
        ))
    } else {
        Ok(())
    }
}

/// Wrapper for `check_endpoint_inner` to handle runtime
fn check_endpoint(
    upgrade_context: Arc<UpgradeContext>,
    address: NetworkAddress,
    peer_id: PeerId,
    public_key: x25519::PublicKey,
    timeout: Duration,
) -> Result<String, Error> {
    let runtime = Runtime::new().unwrap();
    let remote_pubkey = address.find_noise_proto().unwrap();
    println!(
        "Connecting with peer_id {} and pubkey {} to {} with timeout: {:?}",
        peer_id, public_key, address, timeout
    );

    let connect_future = async {
        tokio::time::timeout(
            timeout,
            check_endpoint_inner(upgrade_context, address.clone(), remote_pubkey),
        )
        .await
        .map_err(|_| Error::Timeout("CheckEndpoint", address.to_string()))
    };

    runtime.block_on(connect_future)?
}

/// Connects via Noise, and then drops the connection
async fn check_endpoint_inner(
    upgrade_context: Arc<UpgradeContext>,
    address: NetworkAddress,
    remote_pubkey: x25519::PublicKey,
) -> Result<String, Error> {
    // Connect to the address, this should handle DNS resolution
    let fut_socket = async {
        resolve_and_connect(address.clone())
            .await
            .map(TcpSocket::new)
    };

    // The peer id doesn't matter because we don't validate it
    let remote_peer_id = account_address::from_identity_public_key(remote_pubkey);
    match upgrade_outbound(
        upgrade_context,
        fut_socket,
        address.clone(),
        remote_peer_id,
        remote_pubkey,
    )
    .await
    {
        Ok(conn) => {
            let msg = format!("Successfully connected to {}", conn.metadata.addr);

            // Disconnect
            drop(conn);
            Ok(msg)
        }
        Err(error) => Err(Error::UnexpectedError(format!(
            "Failed to connect to {} due to {}",
            address, error
        ))),
    }
}

/// Derive the peer id that we're using.  This is a convenience to only have to provide private key
fn private_key_to_public_info(private_key: &x25519::PrivateKey) -> (PeerId, x25519::PublicKey) {
    let public_key = private_key.public_key();
    let peer_id = account_address::from_identity_public_key(public_key);
    (peer_id, public_key)
}
<commit_msg>[operational-tool] Simplify Network Checker<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    json_rpc::JsonRpcClientWrapper,
    validator_set::{validator_set_full_node_addresses, validator_set_validator_addresses},
};
use diem_config::{
    config::{RoleType, HANDSHAKE_VERSION},
    network_id::{NetworkContext, NetworkId},
};
use diem_crypto::{x25519, x25519::PRIVATE_KEY_SIZE};
use diem_management::error::Error;
use diem_network_address_encryption::Encryptor;
use diem_secure_storage::{InMemoryStorage, Storage};
use diem_types::{
    account_address,
    chain_id::ChainId,
    network_address::{
        encrypted::{Key, KeyVersion, KEY_LEN},
        NetworkAddress,
    },
    PeerId,
};
use fallible::copy_from_slice::copy_slice_to_vec;
use netcore::transport::tcp::{resolve_and_connect, TcpSocket};
use network::{
    noise::{HandshakeAuthMode, NoiseUpgrader},
    protocols::wire::handshake::v1::SupportedProtocols,
    transport::{upgrade_outbound, UpgradeContext, SUPPORTED_MESSAGING_PROTOCOL},
    ProtocolId,
};
use std::{collections::BTreeMap, sync::Arc};
use structopt::StructOpt;
use tokio::{runtime::Runtime, time::Duration};

const DEFAULT_TIMEOUT_SECONDS: u64 = 5;

#[derive(Debug, StructOpt)]
pub struct CheckEndpoint {
    /// `NetworkAddress` of remote server interface
    #[structopt(long)]
    address: NetworkAddress,
    /// `ChainId` of remote server
    #[structopt(long)]
    chain_id: ChainId,
    /// `NetworkId` of remote server interface
    #[structopt(long)]
    network_id: NetworkId,
    /// `PrivateKey` to connect to remote server
    #[structopt(long, parse(try_from_str = parse_private_key_hex))]
    private_key: Option<x25519::PrivateKey>,
    /// Optional number of seconds to timeout attempting to connect to endpoint
    #[structopt(long)]
    timeout_seconds: Option<u64>,
}

fn parse_private_key_hex(src: &str) -> Result<x25519::PrivateKey, Error> {
    let input = src.trim();
    if input.len() != 64 {
        return Err(Error::CommandArgumentError(
            "Invalid private key length, must be 64 hex characters".to_string(),
        ));
    }

    let value_slice = hex::decode(src.trim())
        .map_err(|_| Error::CommandArgumentError(format!("Not a valid private key: {}", src)))?;

    let mut value = [0; PRIVATE_KEY_SIZE];
    copy_slice_to_vec(&value_slice, &mut value)
        .map_err(|e| Error::CommandArgumentError(format!("{}", e)))?;

    Ok(x25519::PrivateKey::from(value))
}

impl CheckEndpoint {
    pub fn execute(self) -> Result<String, Error> {
        validate_address(&self.address)?;
        let private_key = self.private_key.unwrap_or_else(|| {
            let dummy = [0; PRIVATE_KEY_SIZE];
            x25519::PrivateKey::from(dummy)
        });
        let (peer_id, public_key) = private_key_to_public_info(&private_key);
        let timeout = timeout_duration(self.timeout_seconds);
        println!(
            "Connecting with peer_id {} and pubkey {} to {} with timeout: {:?}",
            peer_id, public_key, self.address, timeout
        );
        check_endpoint(
            build_upgrade_context(self.chain_id, self.network_id, peer_id, private_key),
            self.address,
            timeout,
        )
    }
}

#[derive(Debug, StructOpt)]
pub struct CheckValidatorSetEndpoints {
    /// JSON-RPC Endpoint (e.g. http://localhost:8080)
    #[structopt(long, required_unless = "config")]
    json_server: String,
    /// Specifies whether or not to evaluate validators or fullnodes
    #[structopt(long)]
    role: RoleType,
    /// The expected on-chain key, only required for validator checks
    #[structopt(long, required_if("role", "validator"), parse(try_from_str = parse_validator_key_hex))]
    address_encryption_key: Option<Key>,
    /// The expected on-chain key version, only required for validator checks
    #[structopt(long, required_if("role", "validator"))]
    version: Option<KeyVersion>,
    /// `ChainId` of remote server
    #[structopt(long)]
    chain_id: ChainId,
    /// Private key to connect to remote server
    #[structopt(long, parse(try_from_str = parse_private_key_hex))]
    private_key: Option<x25519::PrivateKey>,
    /// Optional number of seconds to timeout attempting to connect to endpoint
    #[structopt(long)]
    timeout_seconds: Option<u64>,
}

fn parse_validator_key_hex(src: &str) -> Result<Key, Error> {
    let potential_err_msg = format!("Not a valid encryption key: {}", src);
    let value_slice =
        hex::decode(src.trim()).map_err(|_| Error::CommandArgumentError(potential_err_msg))?;

    let mut value = [0; KEY_LEN];
    copy_slice_to_vec(&value_slice, &mut value)
        .map_err(|e| Error::CommandArgumentError(format!("{}", e)))?;
    Ok(value)
}

impl CheckValidatorSetEndpoints {
    pub fn execute(self) -> Result<String, Error> {
        let is_validator = self.role.is_validator();
        let client = JsonRpcClientWrapper::new(self.json_server);
        let private_key = if let Some(private_key) = self.private_key {
            private_key
        } else if is_validator {
            return Err(Error::CommandArgumentError(
                "Must provide a private key for validators".into(),
            ));
        } else {
            let dummy = [0; PRIVATE_KEY_SIZE];
            x25519::PrivateKey::from(dummy)
        };

        let nodes = if is_validator {
            let address_encryption_key = self.address_encryption_key.ok_or_else(|| {
                Error::CommandArgumentError(
                    "Must provide address encryption key for validators".into(),
                )
            })?;
            let version = self.version.ok_or_else(|| {
                Error::CommandArgumentError("Must provide version for validators".into())
            })?;

            // Following unwraps shouldn't fail as it is in memory
            let mut encryptor = Encryptor::new(Storage::InMemoryStorage(InMemoryStorage::new()));
            encryptor.initialize().unwrap();
            encryptor.add_key(version, address_encryption_key).unwrap();
            encryptor.set_current_version(version).unwrap();

            validator_set_validator_addresses(client, &encryptor, None)?
        } else {
            validator_set_full_node_addresses(client, None)?
        };

        // Build a single upgrade context to run all the checks
        let network_id = if is_validator {
            NetworkId::Validator
        } else {
            NetworkId::Public
        };

        let (peer_id, public_key) = private_key_to_public_info(&private_key);
        let upgrade_context =
            build_upgrade_context(self.chain_id, network_id, peer_id, private_key);

        let timeout = timeout_duration(self.timeout_seconds);
        println!(
            "Checking nodes with peer_id {} and public_key {}, timeout {:?}",
            peer_id, public_key, timeout
        );

        // Check all the addresses accordingly
        for (name, peer_id, addrs) in nodes {
            for addr in addrs {
                match check_endpoint(upgrade_context.clone(), addr, timeout) {
                    Ok(_) => println!("{} -- good", name),
                    Err(err) => println!("{} : {} -- bad -- {}", name, peer_id, err),
                };
            }
        }

        Ok("Complete!".to_string())
    }
}

/// Builds a listener free noise connector
fn build_upgrade_context(
    chain_id: ChainId,
    network_id: NetworkId,
    peer_id: PeerId,
    private_key: x25519::PrivateKey,
) -> Arc<UpgradeContext> {
    // RoleType doesn't matter, but the `NetworkId` and `PeerId` are used in handshakes
    let network_context = Arc::new(NetworkContext::new(
        RoleType::FullNode,
        network_id.clone(),
        peer_id,
    ));

    // Let's make sure some protocol can be connected.  In the future we may want to allow for specifics
    let mut supported_protocols = BTreeMap::new();
    supported_protocols.insert(
        SUPPORTED_MESSAGING_PROTOCOL,
        SupportedProtocols::from(ProtocolId::all().iter()),
    );

    // Build the noise and network handshake, without running a full Noise server with listener
    Arc::new(UpgradeContext::new(
        NoiseUpgrader::new(
            network_context,
            private_key,
            // If we had an incoming message, auth mode would matter
            HandshakeAuthMode::server_only(),
        ),
        HANDSHAKE_VERSION,
        supported_protocols,
        chain_id,
        network_id,
    ))
}

fn timeout_duration(maybe_secs: Option<u64>) -> Duration {
    Duration::from_secs(if let Some(secs) = maybe_secs {
        secs
    } else {
        DEFAULT_TIMEOUT_SECONDS
    })
}

fn validate_address(address: &NetworkAddress) -> Result<(), Error> {
    if !address.is_diemnet_addr() {
        Err(Error::CommandArgumentError(
            "Address must have ip, tcp, noise key, and handshake".to_string(),
        ))
    } else {
        Ok(())
    }
}

/// Wrapper for `check_endpoint_inner` to handle runtime
fn check_endpoint(
    upgrade_context: Arc<UpgradeContext>,
    address: NetworkAddress,
    timeout: Duration,
) -> Result<String, Error> {
    let runtime = Runtime::new().unwrap();
    let remote_pubkey = address.find_noise_proto().unwrap();

    let connect_future = async {
        tokio::time::timeout(
            timeout,
            check_endpoint_inner(upgrade_context, address.clone(), remote_pubkey),
        )
        .await
        .map_err(|_| Error::Timeout("CheckEndpoint", address.to_string()))
    };

    runtime.block_on(connect_future)?
}

/// Connects via Noise, and then drops the connection
async fn check_endpoint_inner(
    upgrade_context: Arc<UpgradeContext>,
    address: NetworkAddress,
    remote_pubkey: x25519::PublicKey,
) -> Result<String, Error> {
    // Connect to the address, this should handle DNS resolution
    let fut_socket = async {
        resolve_and_connect(address.clone())
            .await
            .map(TcpSocket::new)
    };

    // The peer id doesn't matter because we don't validate it
    let remote_peer_id = account_address::from_identity_public_key(remote_pubkey);
    match upgrade_outbound(
        upgrade_context,
        fut_socket,
        address.clone(),
        remote_peer_id,
        remote_pubkey,
    )
    .await
    {
        Ok(conn) => {
            let msg = format!("Successfully connected to {}", conn.metadata.addr);

            // Disconnect
            drop(conn);
            Ok(msg)
        }
        Err(error) => Err(Error::UnexpectedError(format!(
            "Failed to connect to {} due to {}",
            address, error
        ))),
    }
}

/// Derive the peer id that we're using.  This is a convenience to only have to provide private key
fn private_key_to_public_info(private_key: &x25519::PrivateKey) -> (PeerId, x25519::PublicKey) {
    let public_key = private_key.public_key();
    let peer_id = account_address::from_identity_public_key(public_key);
    (peer_id, public_key)
}
<|endoftext|>


<commit_before>use addressable::Addressable;
use ram::Ram;
use rom::Rom;

// C64 memory map overview: http://www.c64-wiki.com/index.php/Memory_Map
// Details about the PLA: http://www.c64-wiki.de/index.php/PLA_(C64-Chip)
// Even more PLA details: http://skoe.de/docs/c64-dissected/pla/c64_pla_dissected_r1.1_a4ss.pdf

// 64K RAM (not every part is always accessible)
// 8K BASIC ROM at $A000
// 4K character ROM at $D000
// 8K kernal ROM at $E000

struct C64Memory {
	ram: Ram<u16>,
	kernal: Rom<u16>,
	basic: Rom<u16>,
	characters: Rom<u16>,
}

impl C64Memory {
	fn new () -> C64Memory {
		C64Memory {
			ram: Ram::new(),
			kernal: Rom::new(&Path("kernal.rom")),
			basic: Rom::new(&Path("basic.rom")),
			characters: Rom::new(&Path("characters.rom")),
		}
	}
}

impl Addressable<u16> for C64Memory {
	fn get (&self, addr: u16) -> u8 {
		match addr {
			// TODO: Switch memory access based on LORAM/HIRAM/CHAREN
			0xa000..0xbfff => self.basic.get(addr - 0xa000),
			0xd000..0xdfff => self.characters.get(addr - 0xd000),
			0xe000..0xffff => self.kernal.get(addr - 0xe000),
			_              => self.ram.get(addr),
		}
	}

	fn set (&mut self, addr: u16, data: u8) {
		self.ram.set(addr, data);
	}
}

pub struct C64;

impl C64 {
	pub fn new () -> C64 {
		C64
	}

	pub fn run (&mut self) {
		// TODO
	}
}
<commit_msg>Added CPU and memory to C64<commit_after>use addressable::Addressable;
use mos65xx::Mos6510;
use ram::Ram;
use rom::Rom;

// C64 memory map overview: http://www.c64-wiki.com/index.php/Memory_Map
// Details about the PLA: http://www.c64-wiki.de/index.php/PLA_(C64-Chip)
// Even more PLA details: http://skoe.de/docs/c64-dissected/pla/c64_pla_dissected_r1.1_a4ss.pdf

// 64K RAM (not every part is always accessible)
// 8K BASIC ROM at $A000
// 4K character ROM at $D000
// 8K kernal ROM at $E000

struct C64Memory {
	ram: Ram<u16>,
	kernal: Rom<u16>,
	basic: Rom<u16>,
	characters: Rom<u16>,
}

impl C64Memory {
	fn new () -> C64Memory {
		C64Memory {
			ram: Ram::new(),
			kernal: Rom::new(&Path("kernal.rom")),
			basic: Rom::new(&Path("basic.rom")),
			characters: Rom::new(&Path("characters.rom")),
		}
	}
}

impl Addressable<u16> for C64Memory {
	fn get (&self, addr: u16) -> u8 {
		match addr {
			// TODO: Switch memory access based on LORAM/HIRAM/CHAREN
			0xa000..0xbfff => self.basic.get(addr - 0xa000),
			0xd000..0xdfff => self.characters.get(addr - 0xd000),
			0xe000..0xffff => self.kernal.get(addr - 0xe000),
			_              => self.ram.get(addr),
		}
	}

	fn set (&mut self, addr: u16, data: u8) {
		self.ram.set(addr, data);
	}
}

pub struct C64 {
	priv cpu: Mos6510,
	priv mem: C64Memory,
}

impl C64 {
	pub fn new () -> C64 {
		C64 {
			cpu: Mos6510::new(),
			mem: C64Memory::new(),
		}
	}

	pub fn run (&mut self) {
		// TODO
	}
}
<|endoftext|>


<commit_before>use cgmath::prelude::*;

use crate::bvh::BVHNode;
use crate::color::Color;
use crate::config::*;
use crate::float::*;
use crate::intersect::{Interaction, Ray};
use crate::light::Light;
use crate::scene::Scene;

fn sample_light(
    isect: &Interaction,
    scene: &Scene,
    flash: &dyn Light,
    config: &RenderConfig,
) -> (Color, Ray, Float) {
    let (light, pdf) = match config.light_mode {
        LightMode::Scene => scene.sample_light().unwrap_or((flash, 1.0)),
        LightMode::Camera => (flash, 1.0),
    };
    let (li, ray, lpdf) = light.sample_li(isect);
    (li, ray, pdf * lpdf)
}

pub fn path_trace<'a>(
    mut ray: Ray,
    scene: &'a Scene,
    flash: &dyn Light,
    config: &RenderConfig,
    node_stack: &mut Vec<(&'a BVHNode, Float)>,
) -> Color {
    let mut c = Color::black();
    let mut beta = Color::white();
    let mut bounce = 0;
    let mut specular_bounce = false;
    while let Some(hit) = scene.intersect(&mut ray, node_stack) {
        let isect = hit.interaction(&config);
        if bounce == 0 || specular_bounce {
            c += beta * isect.le(&ray);
        }
        let (le, mut shadow_ray, light_pdf) = sample_light(&isect, scene, flash, config);
        if scene.intersect(&mut shadow_ray, node_stack).is_none() {
            let cos_t = isect.ns.dot(shadow_ray.dir).abs();
            c += beta * le * isect.bsdf(&ray, &shadow_ray) * cos_t / light_pdf;
        }
        let mut pdf = 1.0;
        let terminate = if bounce < config.bounces {
            false
        } else if config.russian_roulette {
            // Survival probability
            let prob = beta.luma().min(0.95);
            pdf *= prob;
            rand::random::<Float>() > prob
        } else {
            true
        };
        if !terminate {
            if let Some((brdf, new_ray, brdf_pdf)) = isect.sample_bsdf(&ray) {
                ray = new_ray;
                pdf *= brdf_pdf;
                beta *= isect.ns.dot(ray.dir).abs() * brdf / pdf;
                bounce += 1;
                specular_bounce = isect.is_specular();
            } else {
                break;
            }
        } else {
            break;
        }
    }
    c
}
<commit_msg>Intersect shadow ray only if bsdf is not black<commit_after>use cgmath::prelude::*;

use crate::bvh::BVHNode;
use crate::color::Color;
use crate::config::*;
use crate::float::*;
use crate::intersect::{Interaction, Ray};
use crate::light::Light;
use crate::scene::Scene;

fn sample_light(
    isect: &Interaction,
    scene: &Scene,
    flash: &dyn Light,
    config: &RenderConfig,
) -> (Color, Ray, Float) {
    let (light, pdf) = match config.light_mode {
        LightMode::Scene => scene.sample_light().unwrap_or((flash, 1.0)),
        LightMode::Camera => (flash, 1.0),
    };
    let (li, ray, lpdf) = light.sample_li(isect);
    (li, ray, pdf * lpdf)
}

pub fn path_trace<'a>(
    mut ray: Ray,
    scene: &'a Scene,
    flash: &dyn Light,
    config: &RenderConfig,
    node_stack: &mut Vec<(&'a BVHNode, Float)>,
) -> Color {
    let mut c = Color::black();
    let mut beta = Color::white();
    let mut bounce = 0;
    let mut specular_bounce = false;
    while let Some(hit) = scene.intersect(&mut ray, node_stack) {
        let isect = hit.interaction(&config);
        if bounce == 0 || specular_bounce {
            c += beta * isect.le(&ray);
        }
        let (le, mut shadow_ray, light_pdf) = sample_light(&isect, scene, flash, config);
        let bsdf = isect.bsdf(&ray, &shadow_ray);
        if !bsdf.is_black() && scene.intersect(&mut shadow_ray, node_stack).is_none() {
            let cos_t = isect.ns.dot(shadow_ray.dir).abs();
            c += beta * le * bsdf * cos_t / light_pdf;
        }
        let mut pdf = 1.0;
        let terminate = if bounce < config.bounces {
            false
        } else if config.russian_roulette {
            // Survival probability
            let prob = beta.luma().min(0.95);
            pdf *= prob;
            rand::random::<Float>() > prob
        } else {
            true
        };
        if !terminate {
            if let Some((brdf, new_ray, brdf_pdf)) = isect.sample_bsdf(&ray) {
                ray = new_ray;
                pdf *= brdf_pdf;
                beta *= isect.ns.dot(ray.dir).abs() * brdf / pdf;
                bounce += 1;
                specular_bounce = isect.is_specular();
            } else {
                break;
            }
        } else {
            break;
        }
    }
    c
}
<|endoftext|>


<commit_before>use {
    BackEnd,
    Color,
    Field,
    Line,
    Matrix2d,
    Stroke,
};
use triangulation::{
    with_round_border_line_tri_list_xy_f32_rgba_f32
};

/// A line context with round border information.
pub struct RoundBorderLineColorContext<'a> {
    /// Base/original transform.
    pub base: Field<'a, Matrix2d>,
    /// Current transform.
    pub transform: Field<'a, Matrix2d>,
    /// Current line.
    pub line: Field<'a, Line>,
    /// Current color.
    pub color: Field<'a, Color>,
    /// Current round border.
    pub round_border_radius: Field<'a, f64>,
}

impl<'a> Stroke<'a> for RoundBorderLineColorContext<'a> {
    fn stroke<B: BackEnd>(&'a self, back_end: &mut B) {
        if back_end.supports_tri_list_xy_f32_rgba_f32() {
            let line = self.line.get();
            let round_border_radius = self.round_border_radius.get();
            let color = self.color.get();
            // Complete transparency does not need to be rendered.
            if color[3] == 0.0 { return; }
            // Turn on alpha blending if not completely opaque.
            let needs_alpha = color[3] != 1.0;
            if needs_alpha { back_end.enable_alpha_blend(); }
            with_round_border_line_tri_list_xy_f32_rgba_f32(
                64,
                self.transform.get(),
                line,
                round_border_radius,
                color,
                |vertices, colors| {
                    back_end.tri_list_xy_f32_rgba_f32(vertices, colors)
                }
            );
            if needs_alpha { back_end.disable_alpha_blend(); }
        } else {
            unimplemented!();
        }
    }
}
<commit_msg>Implemented `View` for `RoundBorderLineColorContext`<commit_after>
use {
    BackEnd,
    Borrowed,
    Color,
    Field,
    Line,
    Matrix2d,
    Stroke,
    Value,
    View,
};
use triangulation::{
    with_round_border_line_tri_list_xy_f32_rgba_f32
};
use vecmath::{
    identity,
};

/// A line context with round border information.
pub struct RoundBorderLineColorContext<'a> {
    /// Base/original transform.
    pub base: Field<'a, Matrix2d>,
    /// Current transform.
    pub transform: Field<'a, Matrix2d>,
    /// Current line.
    pub line: Field<'a, Line>,
    /// Current color.
    pub color: Field<'a, Color>,
    /// Current round border.
    pub round_border_radius: Field<'a, f64>,
}

impl<'a> Stroke<'a> for RoundBorderLineColorContext<'a> {
    fn stroke<B: BackEnd>(&'a self, back_end: &mut B) {
        if back_end.supports_tri_list_xy_f32_rgba_f32() {
            let line = self.line.get();
            let round_border_radius = self.round_border_radius.get();
            let color = self.color.get();
            // Complete transparency does not need to be rendered.
            if color[3] == 0.0 { return; }
            // Turn on alpha blending if not completely opaque.
            let needs_alpha = color[3] != 1.0;
            if needs_alpha { back_end.enable_alpha_blend(); }
            with_round_border_line_tri_list_xy_f32_rgba_f32(
                64,
                self.transform.get(),
                line,
                round_border_radius,
                color,
                |vertices, colors| {
                    back_end.tri_list_xy_f32_rgba_f32(vertices, colors)
                }
            );
            if needs_alpha { back_end.disable_alpha_blend(); }
        } else {
            unimplemented!();
        }
    }
}

impl<'a> View<'a> for RoundBorderLineColorContext<'a> {
    #[inline(always)]
    fn view(&'a self) -> RoundBorderLineColorContext<'a> {
        RoundBorderLineColorContext {
            base: Borrowed(self.base.get()),
            transform: Borrowed(self.base.get()),
            line: Borrowed(self.line.get()),
            color: Borrowed(self.color.get()),
            round_border_radius: Borrowed(self.round_border_radius.get()),
        }
    }

    #[inline(always)]
    fn reset(&'a self) -> RoundBorderLineColorContext<'a> {
        RoundBorderLineColorContext {
            base: Borrowed(self.base.get()),
            transform: Value(identity()),
            line: Borrowed(self.line.get()),
            color: Borrowed(self.color.get()),
            round_border_radius: Borrowed(self.round_border_radius.get()),
        }
    }

    #[inline(always)]
    fn store_view(&'a self) -> RoundBorderLineColorContext<'a> {
        RoundBorderLineColorContext {
            base: Borrowed(self.transform.get()),
            transform: Borrowed(self.transform.get()),
            line: Borrowed(self.line.get()),
            color: Borrowed(self.color.get()),
            round_border_radius: Borrowed(self.round_border_radius.get()),
        }
    }
}

<|endoftext|>


<commit_before>use services::crypto::wrappers::bn::BigNumber;
use services::crypto::wrappers::pair::{GroupOrderElement, PointG1, Pair};
use errors::crypto::CryptoError;
use services::crypto::anoncreds::helpers::AppendBigNumArray;
use std::collections::{HashMap, HashSet};
use std::cell::RefCell;

pub enum ByteOrder {
    Big,
    Little
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct SchemaKey {
    pub name: String,
    pub version: String,
    pub issue_id: String
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Schema {
    pub name: String,
    pub version: String,
    pub attribute_names: HashSet<String>
}

#[derive(Debug)]
pub struct PublicKey {
    pub n: BigNumber,
    pub s: BigNumber,
    pub rms: BigNumber,
    pub r: HashMap<String, BigNumber>,
    pub rctxt: BigNumber,
    pub z: BigNumber
}

pub struct RevocationPublicKey {
    pub g: PointG1,
    pub h: PointG1,
    pub h0: PointG1,
    pub h1: PointG1,
    pub h2: PointG1,
    pub htilde: PointG1,
    pub u: PointG1,
    pub pk: PointG1,
    pub y: PointG1,
    pub x: GroupOrderElement
}

pub struct RevocationSecretKey {
    pub x: GroupOrderElement,
    pub sk: GroupOrderElement
}

#[derive(Debug)]
pub struct SecretKey {
    pub p: BigNumber,
    pub q: BigNumber
}

pub struct AccumulatorPublicKey {
    pub z: Pair
}

pub struct AccumulatorSecretKey {
    pub gamma: GroupOrderElement
}

pub struct Accumulator {
    pub accumulator_id: i32,
    pub acc: PointG1,
    pub v: HashSet<i32>,
    pub max_claim_num: i32,
    pub current_i: i32
}

impl Accumulator {
    pub fn is_full(&self) -> bool {
        self.current_i > self.max_claim_num
    }
}

pub struct Witness {
    pub sigma_i: PointG1,
    pub u_i: PointG1,
    pub g_i: PointG1,
    pub omega: PointG1,
    pub v: HashSet<i32>
}

pub struct ClaimRequest {
    pub user_id: String,
    pub u: BigNumber,
    pub ur: Option<PointG1>
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Predicate {
    pub attr_name: String,
    pub p_type: String,
    pub value: i32
}

#[derive(Debug)]
pub struct Attribute {
    pub name: String,
    pub value: String,
    pub encode: bool
}

pub struct ClaimInitData {
    pub u: BigNumber,
    pub v_prime: BigNumber
}

pub struct Claims {
    pub primary_claim: PrimaryClaim,
    pub non_revocation_claim: Option<RefCell<NonRevocationClaim>>
}

#[derive(Debug)]
pub struct PrimaryClaim {
    pub encoded_attributes: HashMap<String, BigNumber>,
    pub m2: BigNumber,
    pub a: BigNumber,
    pub e: BigNumber,
    pub v_prime: BigNumber
}

pub struct NonRevocationClaim {
    pub accumulator_id: i32,
    pub sigma: PointG1,
    pub c: GroupOrderElement,
    pub vr_prime_prime: GroupOrderElement,
    pub witness: Witness,
    pub g_i: PointG1,
    pub i: i32,
    pub m2: GroupOrderElement
}

pub struct NonRevocProofXList {
    pub rho: GroupOrderElement,
    pub r: GroupOrderElement,
    pub r_prime: GroupOrderElement,
    pub r_prime_prime: GroupOrderElement,
    pub r_prime_prime_prime: GroupOrderElement,
    pub o: GroupOrderElement,
    pub o_prime: GroupOrderElement,
    pub m: GroupOrderElement,
    pub m_prime: GroupOrderElement,
    pub t: GroupOrderElement,
    pub t_prime: GroupOrderElement,
    pub m2: GroupOrderElement,
    pub s: GroupOrderElement,
    pub c: GroupOrderElement
}

pub struct NonRevocProofTauList {
    pub t1: PointG1,
    pub t2: PointG1,
    pub t3: Pair,
    pub t4: Pair,
    pub t5: PointG1,
    pub t6: PointG1,
    pub t7: Pair,
    pub t8: Pair
}

#[derive(Clone)]
pub struct NonRevocProofCList {
    pub e: PointG1,
    pub d: PointG1,
    pub a: PointG1,
    pub g: PointG1,
    pub w: PointG1,
    pub s: PointG1,
    pub u: PointG1
}

pub struct ProofInput {
    pub revealed_attrs: HashSet<String>,
    pub predicates: Vec<Predicate>,
    pub ts: Option<String>,
    pub pubseq_no: Option<String>
}

pub struct ProofClaims {
    pub claims: Claims,
    pub revealed_attrs: HashSet<String>,
    pub predicates: Vec<Predicate>
}

pub struct FullProof {
    pub c_hash: BigNumber,
    pub schema_keys: Vec<SchemaKey>,
    pub proofs: Vec<Proof>,
    pub c_list: Vec<Vec<u8>>
}

pub struct Proof {
    pub primary_proof: PrimaryProof,
    pub non_revoc_proof: Option<NonRevocProof>
}

pub struct InitProof {
    pub primary_init_proof: PrimaryInitProof,
    pub non_revoc_init_proof: Option<NonRevocInitProof>
}

pub struct PrimaryInitProof {
    pub eq_proof: PrimaryEqualInitProof,
    pub ge_proofs: Vec<PrimaryPrecicateGEInitProof>
}

pub struct PrimaryProof {
    pub eq_proof: PrimaryEqualProof,
    pub ge_proofs: Vec<PrimaryPredicateGEProof>
}

pub struct PrimaryEqualInitProof {
    pub a_prime: BigNumber,
    pub t: BigNumber,
    pub etilde: BigNumber,
    pub eprime: BigNumber,
    pub vtilde: BigNumber,
    pub vprime: BigNumber,
    pub mtilde: HashMap<String, BigNumber>,
    pub m1_tilde: BigNumber,
    pub m2_tilde: BigNumber,
    pub unrevealed_attrs: HashSet<String>,
    pub revealed_attrs: HashSet<String>,
    pub encoded_attributes: HashMap<String, BigNumber>,
    pub m2: BigNumber
}

pub struct PrimaryPrecicateGEInitProof {
    pub c_list: Vec<BigNumber>,
    pub tau_list: Vec<BigNumber>,
    pub u: HashMap<String, BigNumber>,
    pub u_tilde: HashMap<String, BigNumber>,
    pub r: HashMap<String, BigNumber>,
    pub r_tilde: HashMap<String, BigNumber>,
    pub alpha_tilde: BigNumber,
    pub predicate: Predicate,
    pub t: HashMap<String, BigNumber>
}

pub struct PrimaryEqualProof {
    pub revealed_attr_names: HashSet<String>,
    pub a_prime: BigNumber,
    pub e: BigNumber,
    pub v: BigNumber,
    pub m: HashMap<String, BigNumber>,
    pub m1: BigNumber,
    pub m2: BigNumber
}

pub struct PrimaryPredicateGEProof {
    pub u: HashMap<String, BigNumber>,
    pub r: HashMap<String, BigNumber>,
    pub mj: BigNumber,
    pub alpha: BigNumber,
    pub t: HashMap<String, BigNumber>,
    pub predicate: Predicate
}

pub struct RevocationClaimInitData {
    pub u: PointG1,
    pub v_prime: GroupOrderElement
}

pub struct NonRevocInitProof {
    pub c_list_params: NonRevocProofXList,
    pub tau_list_params: NonRevocProofXList,
    pub c_list: NonRevocProofCList,
    pub tau_list: NonRevocProofTauList
}

pub struct NonRevocProof {
    pub x_list: NonRevocProofXList,
    pub c_list: NonRevocProofCList
}

//impl block

impl Schema {
    pub fn new(name: String, version: String, attributes_names: HashSet<String>) -> Schema {
        Schema {
            name: name,
            version: version,
            attribute_names: attributes_names
        }
    }
}

impl SchemaKey {
    pub fn new(name: String, version: String, issuer_id: String) -> SchemaKey {
        SchemaKey {
            name: name,
            version: version,
            issue_id: issuer_id
        }
    }
}

impl PublicKey {
    pub fn new(n: BigNumber, s: BigNumber, rms: BigNumber, r: HashMap<String, BigNumber>,
               rctxt: BigNumber, z: BigNumber) -> PublicKey {
        PublicKey {
            n: n,
            s: s,
            rms: rms,
            r: r,
            rctxt: rctxt,
            z: z
        }
    }
}

impl RevocationPublicKey {
    pub fn new(g: PointG1, h: PointG1, h0: PointG1, h1: PointG1, h2: PointG1, htilde: PointG1,
               u: PointG1, pk: PointG1, y: PointG1, x: GroupOrderElement) -> RevocationPublicKey {
        RevocationPublicKey {
            g: g,
            h: h,
            h0: h0,
            h1: h1,
            h2: h2,
            htilde: htilde,
            u: u,
            pk: pk,
            y: y,
            x: x
        }
    }
}

impl RevocationSecretKey {
    pub fn new(x: GroupOrderElement, sk: GroupOrderElement) -> RevocationSecretKey {
        RevocationSecretKey {
            x: x,
            sk: sk
        }
    }
}

impl SecretKey {
    pub fn new(p: BigNumber, q: BigNumber) -> SecretKey {
        SecretKey {
            p: p,
            q: q
        }
    }
}

impl AccumulatorPublicKey {
    pub fn new(z: Pair) -> AccumulatorPublicKey {
        AccumulatorPublicKey {
            z: z
        }
    }
}

impl AccumulatorSecretKey {
    pub fn new(gamma: GroupOrderElement) -> AccumulatorSecretKey {
        AccumulatorSecretKey {
            gamma: gamma
        }
    }
}

impl Accumulator {
    pub fn new(accumulator_id: i32, acc: PointG1, v: HashSet<i32>, max_claim_num: i32,
               current_i: i32) -> Accumulator {
        Accumulator {
            accumulator_id: accumulator_id,
            acc: acc,
            v: v,
            max_claim_num: max_claim_num,
            current_i: current_i
        }
    }
}

impl NonRevocProof {
    pub fn new(x_list: NonRevocProofXList, c_list: NonRevocProofCList) -> NonRevocProof {
        NonRevocProof {
            x_list: x_list,
            c_list: c_list
        }
    }
}

impl PrimaryEqualInitProof {
    pub fn as_list(&self) -> Result<Vec<BigNumber>, CryptoError> {
        Ok(vec![self.a_prime.clone()?])
    }

    pub fn as_tau_list(&self) -> Result<Vec<BigNumber>, CryptoError> {
        Ok(vec![self.t.clone()?])
    }
}

impl PrimaryPrecicateGEInitProof {
    pub fn as_list(&self) -> Result<&Vec<BigNumber>, CryptoError> {
        Ok(&self.c_list)
    }

    pub fn as_tau_list(&self) -> Result<&Vec<BigNumber>, CryptoError> {
        Ok(&self.tau_list)
    }
}

impl PrimaryInitProof {
    pub fn as_c_list(&self) -> Result<Vec<BigNumber>, CryptoError> {
        let mut c_list: Vec<BigNumber> = self.eq_proof.as_list()?;
        for ge_proof in self.ge_proofs.iter() {
            c_list.append_vec(ge_proof.as_list()?)?;
        }
        Ok(c_list)
    }

    pub fn as_tau_list(&self) -> Result<Vec<BigNumber>, CryptoError> {
        let mut tau_list: Vec<BigNumber> = self.eq_proof.as_tau_list()?;
        for ge_proof in self.ge_proofs.iter() {
            tau_list.append_vec(ge_proof.as_tau_list()?)?;
        }
        Ok(tau_list)
    }
}

impl NonRevocProofTauList {
    pub fn as_slice(&self) -> Result<Vec<Vec<u8>>, CryptoError> {
        Ok(vec![self.t1.to_bytes()?, self.t2.to_bytes()?, self.t3.to_bytes()?, self.t4.to_bytes()?,
                self.t5.to_bytes()?, self.t6.to_bytes()?, self.t7.to_bytes()?, self.t8.to_bytes()?])
    }
}

impl NonRevocProofCList {
    pub fn as_list(&self) -> Result<Vec<PointG1>, CryptoError> {
        Ok(vec![self.e, self.d, self.a, self.g, self.w, self.s, self.u])
    }
}

impl NonRevocInitProof {
    pub fn as_c_list(&self) -> Result<Vec<PointG1>, CryptoError> {
        let vec = self.c_list.as_list()?;
        Ok(vec)
    }

    pub fn as_tau_list(&self) -> Result<Vec<Vec<u8>>, CryptoError> {
        let vec = self.tau_list.as_slice()?;
        Ok(vec)
    }
}

impl NonRevocProofXList {
    pub fn new(rho: GroupOrderElement, r: GroupOrderElement, r_prime: GroupOrderElement,
               r_prime_prime: GroupOrderElement, r_prime_prime_prime: GroupOrderElement,
               o: GroupOrderElement, o_prime: GroupOrderElement, m: GroupOrderElement,
               m_prime: GroupOrderElement, t: GroupOrderElement, t_prime: GroupOrderElement,
               m2: GroupOrderElement, s: GroupOrderElement,
               c: GroupOrderElement) -> NonRevocProofXList {
        NonRevocProofXList {
            rho: rho,
            r: r,
            r_prime: r_prime,
            r_prime_prime: r_prime_prime,
            r_prime_prime_prime: r_prime_prime_prime,
            o: o,
            o_prime: o_prime,
            m: m,
            m_prime: m_prime,
            t: t,
            t_prime: t_prime,
            m2: m2,
            s: s,
            c: c
        }
    }

    pub fn as_list(&self) -> Result<Vec<GroupOrderElement>, CryptoError> {
        Ok(vec![self.rho, self.o, self.c, self.o_prime, self.m, self.m_prime, self.t, self.t_prime,
                self.m2, self.s, self.r, self.r_prime, self.r_prime_prime, self.r_prime_prime_prime])
    }

    pub fn from_list(seq: Vec<GroupOrderElement>) -> NonRevocProofXList {
        NonRevocProofXList::new(seq[0], seq[1], seq[2], seq[3], seq[4], seq[5], seq[6], seq[7],
                                seq[8], seq[9], seq[10], seq[11], seq[12], seq[13])
    }
}<commit_msg>refactoring<commit_after>use services::crypto::wrappers::bn::BigNumber;
use services::crypto::wrappers::pair::{GroupOrderElement, PointG1, Pair};
use errors::crypto::CryptoError;
use services::crypto::anoncreds::helpers::AppendBigNumArray;
use std::collections::{HashMap, HashSet};
use std::cell::RefCell;

pub enum ByteOrder {
    Big,
    Little
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct SchemaKey {
    pub name: String,
    pub version: String,
    pub issue_id: String
}

impl SchemaKey {
    pub fn new(name: String, version: String, issuer_id: String) -> SchemaKey {
        SchemaKey {
            name: name,
            version: version,
            issue_id: issuer_id
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Schema {
    pub name: String,
    pub version: String,
    pub attribute_names: HashSet<String>
}

impl Schema {
    pub fn new(name: String, version: String, attributes_names: HashSet<String>) -> Schema {
        Schema {
            name: name,
            version: version,
            attribute_names: attributes_names
        }
    }
}

#[derive(Debug)]
pub struct PublicKey {
    pub n: BigNumber,
    pub s: BigNumber,
    pub rms: BigNumber,
    pub r: HashMap<String, BigNumber>,
    pub rctxt: BigNumber,
    pub z: BigNumber
}

impl PublicKey {
    pub fn new(n: BigNumber, s: BigNumber, rms: BigNumber, r: HashMap<String, BigNumber>,
               rctxt: BigNumber, z: BigNumber) -> PublicKey {
        PublicKey {
            n: n,
            s: s,
            rms: rms,
            r: r,
            rctxt: rctxt,
            z: z
        }
    }
}

pub struct RevocationPublicKey {
    pub g: PointG1,
    pub h: PointG1,
    pub h0: PointG1,
    pub h1: PointG1,
    pub h2: PointG1,
    pub htilde: PointG1,
    pub u: PointG1,
    pub pk: PointG1,
    pub y: PointG1,
    pub x: GroupOrderElement
}

impl RevocationPublicKey {
    pub fn new(g: PointG1, h: PointG1, h0: PointG1, h1: PointG1, h2: PointG1, htilde: PointG1,
               u: PointG1, pk: PointG1, y: PointG1, x: GroupOrderElement) -> RevocationPublicKey {
        RevocationPublicKey {
            g: g,
            h: h,
            h0: h0,
            h1: h1,
            h2: h2,
            htilde: htilde,
            u: u,
            pk: pk,
            y: y,
            x: x
        }
    }
}

pub struct RevocationSecretKey {
    pub x: GroupOrderElement,
    pub sk: GroupOrderElement
}

impl RevocationSecretKey {
    pub fn new(x: GroupOrderElement, sk: GroupOrderElement) -> RevocationSecretKey {
        RevocationSecretKey {
            x: x,
            sk: sk
        }
    }
}

#[derive(Debug)]
pub struct SecretKey {
    pub p: BigNumber,
    pub q: BigNumber
}

impl SecretKey {
    pub fn new(p: BigNumber, q: BigNumber) -> SecretKey {
        SecretKey {
            p: p,
            q: q
        }
    }
}

pub struct AccumulatorPublicKey {
    pub z: Pair
}

impl AccumulatorPublicKey {
    pub fn new(z: Pair) -> AccumulatorPublicKey {
        AccumulatorPublicKey {
            z: z
        }
    }
}

pub struct AccumulatorSecretKey {
    pub gamma: GroupOrderElement
}

impl AccumulatorSecretKey {
    pub fn new(gamma: GroupOrderElement) -> AccumulatorSecretKey {
        AccumulatorSecretKey {
            gamma: gamma
        }
    }
}

pub struct Accumulator {
    pub accumulator_id: i32,
    pub acc: PointG1,
    pub v: HashSet<i32>,
    pub max_claim_num: i32,
    pub current_i: i32
}

impl Accumulator {
    pub fn new(accumulator_id: i32, acc: PointG1, v: HashSet<i32>, max_claim_num: i32,
               current_i: i32) -> Accumulator {
        Accumulator {
            accumulator_id: accumulator_id,
            acc: acc,
            v: v,
            max_claim_num: max_claim_num,
            current_i: current_i
        }
    }

    pub fn is_full(&self) -> bool {
        self.current_i > self.max_claim_num
    }
}

pub struct Witness {
    pub sigma_i: PointG1,
    pub u_i: PointG1,
    pub g_i: PointG1,
    pub omega: PointG1,
    pub v: HashSet<i32>
}

pub struct ClaimRequest {
    pub user_id: String,
    pub u: BigNumber,
    pub ur: Option<PointG1>
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Predicate {
    pub attr_name: String,
    pub p_type: String,
    pub value: i32
}

#[derive(Debug)]
pub struct Attribute {
    pub name: String,
    pub value: String,
    pub encode: bool
}

pub struct ClaimInitData {
    pub u: BigNumber,
    pub v_prime: BigNumber
}

pub struct Claims {
    pub primary_claim: PrimaryClaim,
    pub non_revocation_claim: Option<RefCell<NonRevocationClaim>>
}

#[derive(Debug)]
pub struct PrimaryClaim {
    pub encoded_attributes: HashMap<String, BigNumber>,
    pub m2: BigNumber,
    pub a: BigNumber,
    pub e: BigNumber,
    pub v_prime: BigNumber
}

pub struct NonRevocationClaim {
    pub accumulator_id: i32,
    pub sigma: PointG1,
    pub c: GroupOrderElement,
    pub vr_prime_prime: GroupOrderElement,
    pub witness: Witness,
    pub g_i: PointG1,
    pub i: i32,
    pub m2: GroupOrderElement
}

pub struct NonRevocProofXList {
    pub rho: GroupOrderElement,
    pub r: GroupOrderElement,
    pub r_prime: GroupOrderElement,
    pub r_prime_prime: GroupOrderElement,
    pub r_prime_prime_prime: GroupOrderElement,
    pub o: GroupOrderElement,
    pub o_prime: GroupOrderElement,
    pub m: GroupOrderElement,
    pub m_prime: GroupOrderElement,
    pub t: GroupOrderElement,
    pub t_prime: GroupOrderElement,
    pub m2: GroupOrderElement,
    pub s: GroupOrderElement,
    pub c: GroupOrderElement
}

impl NonRevocProofXList {
    pub fn new(rho: GroupOrderElement, r: GroupOrderElement, r_prime: GroupOrderElement,
               r_prime_prime: GroupOrderElement, r_prime_prime_prime: GroupOrderElement,
               o: GroupOrderElement, o_prime: GroupOrderElement, m: GroupOrderElement,
               m_prime: GroupOrderElement, t: GroupOrderElement, t_prime: GroupOrderElement,
               m2: GroupOrderElement, s: GroupOrderElement,
               c: GroupOrderElement) -> NonRevocProofXList {
        NonRevocProofXList {
            rho: rho,
            r: r,
            r_prime: r_prime,
            r_prime_prime: r_prime_prime,
            r_prime_prime_prime: r_prime_prime_prime,
            o: o,
            o_prime: o_prime,
            m: m,
            m_prime: m_prime,
            t: t,
            t_prime: t_prime,
            m2: m2,
            s: s,
            c: c
        }
    }

    pub fn as_list(&self) -> Result<Vec<GroupOrderElement>, CryptoError> {
        Ok(vec![self.rho, self.o, self.c, self.o_prime, self.m, self.m_prime, self.t, self.t_prime,
                self.m2, self.s, self.r, self.r_prime, self.r_prime_prime, self.r_prime_prime_prime])
    }

    pub fn from_list(seq: Vec<GroupOrderElement>) -> NonRevocProofXList {
        NonRevocProofXList::new(seq[0], seq[1], seq[2], seq[3], seq[4], seq[5], seq[6], seq[7],
                                seq[8], seq[9], seq[10], seq[11], seq[12], seq[13])
    }
}

pub struct NonRevocProofTauList {
    pub t1: PointG1,
    pub t2: PointG1,
    pub t3: Pair,
    pub t4: Pair,
    pub t5: PointG1,
    pub t6: PointG1,
    pub t7: Pair,
    pub t8: Pair
}

impl NonRevocProofTauList {
    pub fn as_slice(&self) -> Result<Vec<Vec<u8>>, CryptoError> {
        Ok(vec![self.t1.to_bytes()?, self.t2.to_bytes()?, self.t3.to_bytes()?, self.t4.to_bytes()?,
                self.t5.to_bytes()?, self.t6.to_bytes()?, self.t7.to_bytes()?, self.t8.to_bytes()?])
    }
}

#[derive(Clone)]
pub struct NonRevocProofCList {
    pub e: PointG1,
    pub d: PointG1,
    pub a: PointG1,
    pub g: PointG1,
    pub w: PointG1,
    pub s: PointG1,
    pub u: PointG1
}

impl NonRevocProofCList {
    pub fn as_list(&self) -> Result<Vec<PointG1>, CryptoError> {
        Ok(vec![self.e, self.d, self.a, self.g, self.w, self.s, self.u])
    }
}

pub struct ProofInput {
    pub revealed_attrs: HashSet<String>,
    pub predicates: Vec<Predicate>,
    pub ts: Option<String>,
    pub pubseq_no: Option<String>
}

pub struct ProofClaims {
    pub claims: Claims,
    pub revealed_attrs: HashSet<String>,
    pub predicates: Vec<Predicate>
}

pub struct FullProof {
    pub c_hash: BigNumber,
    pub schema_keys: Vec<SchemaKey>,
    pub proofs: Vec<Proof>,
    pub c_list: Vec<Vec<u8>>
}

pub struct Proof {
    pub primary_proof: PrimaryProof,
    pub non_revoc_proof: Option<NonRevocProof>
}

pub struct InitProof {
    pub primary_init_proof: PrimaryInitProof,
    pub non_revoc_init_proof: Option<NonRevocInitProof>
}

pub struct PrimaryInitProof {
    pub eq_proof: PrimaryEqualInitProof,
    pub ge_proofs: Vec<PrimaryPrecicateGEInitProof>
}

impl PrimaryInitProof {
    pub fn as_c_list(&self) -> Result<Vec<BigNumber>, CryptoError> {
        let mut c_list: Vec<BigNumber> = self.eq_proof.as_list()?;
        for ge_proof in self.ge_proofs.iter() {
            c_list.append_vec(ge_proof.as_list()?)?;
        }
        Ok(c_list)
    }

    pub fn as_tau_list(&self) -> Result<Vec<BigNumber>, CryptoError> {
        let mut tau_list: Vec<BigNumber> = self.eq_proof.as_tau_list()?;
        for ge_proof in self.ge_proofs.iter() {
            tau_list.append_vec(ge_proof.as_tau_list()?)?;
        }
        Ok(tau_list)
    }
}

pub struct PrimaryProof {
    pub eq_proof: PrimaryEqualProof,
    pub ge_proofs: Vec<PrimaryPredicateGEProof>
}

pub struct PrimaryEqualInitProof {
    pub a_prime: BigNumber,
    pub t: BigNumber,
    pub etilde: BigNumber,
    pub eprime: BigNumber,
    pub vtilde: BigNumber,
    pub vprime: BigNumber,
    pub mtilde: HashMap<String, BigNumber>,
    pub m1_tilde: BigNumber,
    pub m2_tilde: BigNumber,
    pub unrevealed_attrs: HashSet<String>,
    pub revealed_attrs: HashSet<String>,
    pub encoded_attributes: HashMap<String, BigNumber>,
    pub m2: BigNumber
}

impl PrimaryEqualInitProof {
    pub fn as_list(&self) -> Result<Vec<BigNumber>, CryptoError> {
        Ok(vec![self.a_prime.clone()?])
    }

    pub fn as_tau_list(&self) -> Result<Vec<BigNumber>, CryptoError> {
        Ok(vec![self.t.clone()?])
    }
}

pub struct PrimaryPrecicateGEInitProof {
    pub c_list: Vec<BigNumber>,
    pub tau_list: Vec<BigNumber>,
    pub u: HashMap<String, BigNumber>,
    pub u_tilde: HashMap<String, BigNumber>,
    pub r: HashMap<String, BigNumber>,
    pub r_tilde: HashMap<String, BigNumber>,
    pub alpha_tilde: BigNumber,
    pub predicate: Predicate,
    pub t: HashMap<String, BigNumber>
}

pub struct PrimaryEqualProof {
    pub revealed_attr_names: HashSet<String>,
    pub a_prime: BigNumber,
    pub e: BigNumber,
    pub v: BigNumber,
    pub m: HashMap<String, BigNumber>,
    pub m1: BigNumber,
    pub m2: BigNumber
}

pub struct PrimaryPredicateGEProof {
    pub u: HashMap<String, BigNumber>,
    pub r: HashMap<String, BigNumber>,
    pub mj: BigNumber,
    pub alpha: BigNumber,
    pub t: HashMap<String, BigNumber>,
    pub predicate: Predicate
}

impl PrimaryPrecicateGEInitProof {
    pub fn as_list(&self) -> Result<&Vec<BigNumber>, CryptoError> {
        Ok(&self.c_list)
    }

    pub fn as_tau_list(&self) -> Result<&Vec<BigNumber>, CryptoError> {
        Ok(&self.tau_list)
    }
}

pub struct RevocationClaimInitData {
    pub u: PointG1,
    pub v_prime: GroupOrderElement
}

pub struct NonRevocInitProof {
    pub c_list_params: NonRevocProofXList,
    pub tau_list_params: NonRevocProofXList,
    pub c_list: NonRevocProofCList,
    pub tau_list: NonRevocProofTauList
}

impl NonRevocInitProof {
    pub fn as_c_list(&self) -> Result<Vec<PointG1>, CryptoError> {
        let vec = self.c_list.as_list()?;
        Ok(vec)
    }

    pub fn as_tau_list(&self) -> Result<Vec<Vec<u8>>, CryptoError> {
        let vec = self.tau_list.as_slice()?;
        Ok(vec)
    }
}

pub struct NonRevocProof {
    pub x_list: NonRevocProofXList,
    pub c_list: NonRevocProofCList
}

impl NonRevocProof {
    pub fn new(x_list: NonRevocProofXList, c_list: NonRevocProofCList) -> NonRevocProof {
        NonRevocProof {
            x_list: x_list,
            c_list: c_list
        }
    }
}
<|endoftext|>


<commit_before>extern crate ring;
extern crate rustc_serialize;

use std::fmt::Display;
use std::cmp;

use self::ring::digest::{ Algorithm, Digest };
use self::rustc_serialize::{ Encodable, Encoder, Decodable, Decoder, json };

use services::ledger::merkletree::hashutils::{ Hashable, HashUtils };

pub use services::ledger::merkletree::proof::{
    Proof,
    Lemma,
    Positioned
};

/// Binary Tree where leaves hold a stand-alone value.
#[derive(Clone, Debug, PartialEq)]
pub enum Tree<T> {
    Empty {
        hash: Vec<u8>
    },

    Leaf {
        hash: Vec<u8>,
        value: T
    },

    Node {
        hash: Vec<u8>,
        left: Box<Tree<T>>,
        right: Box<Tree<T>>
    }
}

impl <T> Tree<T> {
    /// Create an empty tree
    pub fn empty(hash: Digest) -> Self {
        Tree::Empty {
            hash: hash.as_ref().into()
        }
    }

    /// Create a new tree
    pub fn new(hash: Digest, value: T) -> Self {
        Tree::Leaf {
            hash: hash.as_ref().into(),
            value: value
        }
    }

    /// Create a new leaf
    pub fn new_leaf(algo: &'static Algorithm, value: T) -> Tree<T>
            where T: Hashable {

        let hash = algo.hash_leaf(&value);
        Tree::new(hash, value)
    }

    /// Returns a hash from the tree.
    pub fn hash(&self) -> &Vec<u8> {
        match *self {
            Tree::Empty { ref hash }    => hash,
            Tree::Leaf { ref hash, .. } => hash,
            Tree::Node { ref hash, .. } => hash
        }
    }

    /// Returns a borrowing iterator over the leaves of the tree.
    pub fn iter(&self) -> LeavesIterator<T> {
        LeavesIterator::new(self)
    }

    pub fn get_height(&self) -> usize {
        match *self {
            Tree::Empty { .. } => { 0 },
            Tree::Node { ref left, ref right, .. } => {
                1 + cmp::max(left.get_height(),right.get_height())
            },
            Tree::Leaf { .. } => { 0 }
        }
    }

    pub fn get_count(&self) -> usize {
        match *self {
            Tree::Empty { .. } => { 0 },
            Tree::Node { ref left, ref right, .. } => {
                left.get_count() + right.get_count()
            },
            Tree::Leaf { .. } => { 1 }
        }
    }
}

impl<T: AsRef<[u8]> + Encodable> Encodable for Tree<T> {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        match *self {
            Tree::Empty { ref hash, .. } => {
                s.emit_struct("node", 4, |s| {
                    try!(s.emit_struct_field("type", 0, |s| {
                        s.emit_str("empty")
                    }));
                    try!(s.emit_struct_field("hash", 1, |s| {
                        hash.encode(s)
                    }));
                    try!(s.emit_struct_field("", 2, |s| {
                        s.emit_str("")
                    }));
                    try!(s.emit_struct_field("", 3, |s| {
                        s.emit_str("")
                    }));
                    Ok(())
                })
            }
            Tree::Node { ref hash, ref left, ref right, .. } => {
                s.emit_struct("node", 4, |s| {
                    try!(s.emit_struct_field("type", 0, |s| {
                        s.emit_str("node")
                    }));
                    try!(s.emit_struct_field("hash", 1, |s| {
                        hash.encode(s)
                    }));
                    try!(s.emit_struct_field("left", 2, |s| {
                        left.encode(s)
                    }));
                    try!(s.emit_struct_field("right", 3, |s| {
                        right.encode(s)
                    }));
                    Ok(())
                })
            }
            Tree::Leaf { ref hash, ref value, .. } => {
                s.emit_struct("node", 4, |s| {
                    try!(s.emit_struct_field("type", 0, |s| {
                        s.emit_str("leaf")
                    }));
                    try!(s.emit_struct_field("hash", 1, |s| {
                        hash.encode(s)
                    }));
                    try!(s.emit_struct_field("value", 2, |s| {
                        value.encode(s)
                    }));
                    try!(s.emit_struct_field("", 3, |s| {
                        s.emit_str("")
                    }));
                    Ok(())
                })
            }
        }
    }
}

impl<T: AsRef<[u8]> + Decodable + Display> Decodable for Tree<T> {
    fn decode<D: Decoder>(d: &mut D) -> Result<Tree<T>, D::Error> {
        d.read_struct("node", 4, |d| {
            let nodetype = try!(d.read_struct_field("type", 0, |d| {
                d.read_str()
            }));
            let hash = try!(d.read_struct_field("hash", 0, |d| {
                Vec::<u8>::decode(d)
            }));
            match nodetype.as_ref() {
                "empty" => {
                    Ok(Tree::Empty{
                        hash: hash
                    })
                }
                "node" => {
                    let left = try!(d.read_struct_field("left", 1, |d| {
                        Tree::<T>::decode(d)
                    }));
                    let right = try!(d.read_struct_field("right", 2, |d| {
                        Tree::<T>::decode(d)
                    }));
                    Ok(Tree::Node{
                        hash: hash,
                        left: Box::new(left),
                        right: Box::new(right)
                    })
                }
                "leaf" => {
                    let value = try!(d.read_struct_field("value", 1, |d| {
                        T::decode(d)
                    }));
                    Ok(Tree::Leaf{
                        hash: hash,
                        value: value
                    })
                }
                _ => {
                    Err(d.error("bad node type"))
                }
            }
        })
    }
}

/// An borrowing iterator over the leaves of a `Tree`.
/// Adapted from http://codereview.stackexchange.com/q/110283.
#[allow(missing_debug_implementations)]
pub struct LeavesIterator<'a, T> where T: 'a {
    current_value: Option<&'a T>,
    right_nodes: Vec<&'a Tree<T>>
}

impl <'a, T> LeavesIterator<'a, T> {

    fn new(root: &'a Tree<T>) -> Self {
        let mut iter = LeavesIterator {
            current_value: None,
            right_nodes: Vec::new()
        };

        iter.add_left(root);

        iter
    }

    fn add_left(&mut self, mut tree: &'a Tree<T>) {
        loop {
            match *tree {
                Tree::Empty { .. } => {
                    self.current_value = None;
                    break;
                },

                Tree::Node { ref left, ref right, .. } => {
                    self.right_nodes.push(right);
                    tree = left;
                },

                Tree::Leaf { ref value, .. } => {
                    self.current_value = Some(value);
                    break;
                }
            }
        }
    }

}

impl <'a, T> Iterator for LeavesIterator<'a, T> {

    type Item = &'a T;

    fn next(&mut self) -> Option<&'a T> {
        let result = self.current_value.take();

        if let Some(rest) = self.right_nodes.pop() {
            self.add_left(rest);
        }

        result
    }

}

/// An iterator over the leaves of a `Tree`.
#[allow(missing_debug_implementations)]
pub struct LeavesIntoIterator<T> {
    current_value: Option<T>,
    right_nodes: Vec<Tree<T>>
}

impl <T> LeavesIntoIterator<T> {

    fn new(root: Tree<T>) -> Self {
        let mut iter = LeavesIntoIterator {
            current_value: None,
            right_nodes: Vec::new()
        };

        iter.add_left(root);

        iter
    }

    fn add_left(&mut self, mut tree: Tree<T>) {
        loop {
            match tree {
                Tree::Empty { .. } => {
                    self.current_value = None;
                    break;
                },

                Tree::Node { left, right, .. } => {
                    self.right_nodes.push(*right);
                    tree = *left;
                },

                Tree::Leaf { value, .. } => {
                    self.current_value = Some(value);
                    break;
                }
            }
        }
    }

}

impl <T> Iterator for LeavesIntoIterator<T> {

    type Item = T;

    fn next(&mut self) -> Option<T> {
        let result = self.current_value.take();

        if let Some(rest) = self.right_nodes.pop() {
            self.add_left(rest);
        }

        result
    }

}

impl <T> IntoIterator for Tree<T> {

    type Item     = T;
    type IntoIter = LeavesIntoIterator<T>;

    fn into_iter(self) -> Self::IntoIter {
        LeavesIntoIterator::new(self)
    }

}
<commit_msg>replace try! with ?<commit_after>extern crate ring;
extern crate rustc_serialize;

use std::fmt::Display;
use std::cmp;

use self::ring::digest::{ Algorithm, Digest };
use self::rustc_serialize::{ Encodable, Encoder, Decodable, Decoder, json };

use services::ledger::merkletree::hashutils::{ Hashable, HashUtils };

pub use services::ledger::merkletree::proof::{
    Proof,
    Lemma,
    Positioned
};

/// Binary Tree where leaves hold a stand-alone value.
#[derive(Clone, Debug, PartialEq)]
pub enum Tree<T> {
    Empty {
        hash: Vec<u8>
    },

    Leaf {
        hash: Vec<u8>,
        value: T
    },

    Node {
        hash: Vec<u8>,
        left: Box<Tree<T>>,
        right: Box<Tree<T>>
    }
}

impl <T> Tree<T> {
    /// Create an empty tree
    pub fn empty(hash: Digest) -> Self {
        Tree::Empty {
            hash: hash.as_ref().into()
        }
    }

    /// Create a new tree
    pub fn new(hash: Digest, value: T) -> Self {
        Tree::Leaf {
            hash: hash.as_ref().into(),
            value: value
        }
    }

    /// Create a new leaf
    pub fn new_leaf(algo: &'static Algorithm, value: T) -> Tree<T>
            where T: Hashable {

        let hash = algo.hash_leaf(&value);
        Tree::new(hash, value)
    }

    /// Returns a hash from the tree.
    pub fn hash(&self) -> &Vec<u8> {
        match *self {
            Tree::Empty { ref hash }    => hash,
            Tree::Leaf { ref hash, .. } => hash,
            Tree::Node { ref hash, .. } => hash
        }
    }

    /// Returns a borrowing iterator over the leaves of the tree.
    pub fn iter(&self) -> LeavesIterator<T> {
        LeavesIterator::new(self)
    }

    pub fn get_height(&self) -> usize {
        match *self {
            Tree::Empty { .. } => { 0 },
            Tree::Node { ref left, ref right, .. } => {
                1 + cmp::max(left.get_height(),right.get_height())
            },
            Tree::Leaf { .. } => { 0 }
        }
    }

    pub fn get_count(&self) -> usize {
        match *self {
            Tree::Empty { .. } => { 0 },
            Tree::Node { ref left, ref right, .. } => {
                left.get_count() + right.get_count()
            },
            Tree::Leaf { .. } => { 1 }
        }
    }
}

impl<T: AsRef<[u8]> + Encodable> Encodable for Tree<T> {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        match *self {
            Tree::Empty { ref hash, .. } => {
                s.emit_struct("node", 4, |s| {
                    s.emit_struct_field("type", 0, |s| { s.emit_str("empty") })?;
                    s.emit_struct_field("hash", 1, |s| { hash.encode(s) })?;
                    s.emit_struct_field("", 2, |s| { s.emit_str("") })?;
                    s.emit_struct_field("", 3, |s| { s.emit_str("") })?;
                    Ok(())
                })
            }
            Tree::Node { ref hash, ref left, ref right, .. } => {
                s.emit_struct("node", 4, |s| {
                    s.emit_struct_field("type", 0, |s| { s.emit_str("node") })?;
                    s.emit_struct_field("hash", 1, |s| { hash.encode(s) })?;
                    s.emit_struct_field("left", 2, |s| { left.encode(s) })?;
                    s.emit_struct_field("right", 3, |s| { right.encode(s) })?;
                    Ok(())
                })
            }
            Tree::Leaf { ref hash, ref value, .. } => {
                s.emit_struct("node", 4, |s| {
                    s.emit_struct_field("type", 0, |s| { s.emit_str("leaf") })?;
                    s.emit_struct_field("hash", 1, |s| { hash.encode(s) })?;
                    s.emit_struct_field("value", 2, |s| { value.encode(s) })?;
                    s.emit_struct_field("", 3, |s| { s.emit_str("") })?;
                    Ok(())
                })
            }
        }
    }
}

impl<T: AsRef<[u8]> + Decodable + Display> Decodable for Tree<T> {
    fn decode<D: Decoder>(d: &mut D) -> Result<Tree<T>, D::Error> {
        d.read_struct("node", 4, |d| {
            let nodetype = d.read_struct_field("type", 0, |d| { d.read_str() })?;
            let hash = d.read_struct_field("hash", 0, |d| { Vec::<u8>::decode(d) })?;
            match nodetype.as_ref() {
                "empty" => {
                    Ok(Tree::Empty{
                        hash: hash
                    })
                }
                "node" => {
                    let left = d.read_struct_field("left", 1, |d| { Tree::<T>::decode(d) })?;
                    let right = d.read_struct_field("right", 2, |d| { Tree::<T>::decode(d) })?;
                    Ok(Tree::Node{
                        hash: hash,
                        left: Box::new(left),
                        right: Box::new(right)
                    })
                }
                "leaf" => {
                    let value = d.read_struct_field("value", 1, |d| { T::decode(d) })?;
                    Ok(Tree::Leaf{
                        hash: hash,
                        value: value
                    })
                }
                _ => {
                    Err(d.error("bad node type"))
                }
            }
        })
    }
}

/// An borrowing iterator over the leaves of a `Tree`.
/// Adapted from http://codereview.stackexchange.com/q/110283.
#[allow(missing_debug_implementations)]
pub struct LeavesIterator<'a, T> where T: 'a {
    current_value: Option<&'a T>,
    right_nodes: Vec<&'a Tree<T>>
}

impl <'a, T> LeavesIterator<'a, T> {

    fn new(root: &'a Tree<T>) -> Self {
        let mut iter = LeavesIterator {
            current_value: None,
            right_nodes: Vec::new()
        };

        iter.add_left(root);

        iter
    }

    fn add_left(&mut self, mut tree: &'a Tree<T>) {
        loop {
            match *tree {
                Tree::Empty { .. } => {
                    self.current_value = None;
                    break;
                },

                Tree::Node { ref left, ref right, .. } => {
                    self.right_nodes.push(right);
                    tree = left;
                },

                Tree::Leaf { ref value, .. } => {
                    self.current_value = Some(value);
                    break;
                }
            }
        }
    }

}

impl <'a, T> Iterator for LeavesIterator<'a, T> {

    type Item = &'a T;

    fn next(&mut self) -> Option<&'a T> {
        let result = self.current_value.take();

        if let Some(rest) = self.right_nodes.pop() {
            self.add_left(rest);
        }

        result
    }

}

/// An iterator over the leaves of a `Tree`.
#[allow(missing_debug_implementations)]
pub struct LeavesIntoIterator<T> {
    current_value: Option<T>,
    right_nodes: Vec<Tree<T>>
}

impl <T> LeavesIntoIterator<T> {

    fn new(root: Tree<T>) -> Self {
        let mut iter = LeavesIntoIterator {
            current_value: None,
            right_nodes: Vec::new()
        };

        iter.add_left(root);

        iter
    }

    fn add_left(&mut self, mut tree: Tree<T>) {
        loop {
            match tree {
                Tree::Empty { .. } => {
                    self.current_value = None;
                    break;
                },

                Tree::Node { left, right, .. } => {
                    self.right_nodes.push(*right);
                    tree = *left;
                },

                Tree::Leaf { value, .. } => {
                    self.current_value = Some(value);
                    break;
                }
            }
        }
    }

}

impl <T> Iterator for LeavesIntoIterator<T> {

    type Item = T;

    fn next(&mut self) -> Option<T> {
        let result = self.current_value.take();

        if let Some(rest) = self.right_nodes.pop() {
            self.add_left(rest);
        }

        result
    }

}

impl <T> IntoIterator for Tree<T> {

    type Item     = T;
    type IntoIter = LeavesIntoIterator<T>;

    fn into_iter(self) -> Self::IntoIter {
        LeavesIntoIterator::new(self)
    }

}
<|endoftext|>