<commit_before>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/// This example shows a basic packet logger using libpnet

extern crate pnet;

use pnet::datalink::{self, NetworkInterface};

use pnet::packet::Packet;
use pnet::packet::arp::ArpPacket;
use pnet::packet::ethernet::{EtherTypes, EthernetPacket};
use pnet::packet::icmpv6::Icmpv6Packet;
use pnet::packet::icmp::{echo_reply, echo_request, IcmpPacket, IcmpTypes};
use pnet::packet::ip::{IpNextHeaderProtocol, IpNextHeaderProtocols};
use pnet::packet::ipv4::Ipv4Packet;
use pnet::packet::ipv6::Ipv6Packet;
use pnet::packet::tcp::TcpPacket;
use pnet::packet::udp::UdpPacket;
use std::env;
use std::io::{self, Write};
use std::process;
use std::net::IpAddr;

fn handle_udp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let udp = UdpPacket::new(packet);

    if let Some(udp) = udp {
        println!("[{}]: UDP Packet: {}:{} > {}:{}; length: {}",
                 interface_name,
                 source,
                 udp.get_source(),
                 destination,
                 udp.get_destination(),
                 udp.get_length());
    } else {
        println!("[{}]: Malformed UDP Packet", interface_name);
    }
}

fn handle_icmp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let icmp_packet = IcmpPacket::new(packet);
    if let Some(icmp_packet) = icmp_packet {
        match icmp_packet.get_icmp_type() {
            IcmpTypes::EchoReply => {
                let echo_reply_packet = echo_reply::EchoReplyPacket::new(packet).unwrap();
                println!("[{}]: ICMP echo reply {} -> {} (seq={:?}, id={:?})",
                         interface_name,
                         source,
                         destination,
                         echo_reply_packet.get_sequence_number(),
                         echo_reply_packet.get_identifier());
            }
            IcmpTypes::EchoRequest => {
                let echo_request_packet = echo_request::EchoRequestPacket::new(packet).unwrap();
                println!("[{}]: ICMP echo request {} -> {} (seq={:?}, id={:?})",
                         interface_name,
                         source,
                         destination,
                         echo_request_packet.get_sequence_number(),
                         echo_request_packet.get_identifier());
            }
            _ => {
                println!("[{}]: ICMP packet {} -> {} (type={:?})",
                         interface_name,
                         source,
                         destination,
                         icmp_packet.get_icmp_type())
            }
        }
    } else {
        println!("[{}]: Malformed ICMP Packet", interface_name);
    }
}

fn handle_icmpv6_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let icmpv6_packet = Icmpv6Packet::new(packet);
    if let Some(icmpv6_packet) = icmpv6_packet {
        println!(
            "[{}]: ICMPv6 packet {} -> {} (type={:?})",
            interface_name,
            source,
            destination,
            icmpv6_packet.get_icmpv6_type()
        )
    } else {
        println!("[{}]: Malformed ICMPv6 Packet", interface_name);
    }
}

fn handle_tcp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let tcp = TcpPacket::new(packet);
    if let Some(tcp) = tcp {
        println!("[{}]: TCP Packet: {}:{} > {}:{}; length: {}",
                 interface_name,
                 source,
                 tcp.get_source(),
                 destination,
                 tcp.get_destination(),
                 packet.len());
    } else {
        println!("[{}]: Malformed TCP Packet", interface_name);
    }
}

fn handle_transport_protocol(interface_name: &str,
                             source: IpAddr,
                             destination: IpAddr,
                             protocol: IpNextHeaderProtocol,
                             packet: &[u8]) {
    match protocol {
        IpNextHeaderProtocols::Udp => {
            handle_udp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Tcp => {
            handle_tcp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Icmp => {
            handle_icmp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Icmpv6 => {
            handle_icmpv6_packet(interface_name, source, destination, packet)
        }
        _ => {
            println!("[{}]: Unknown {} packet: {} > {}; protocol: {:?} length: {}",
                     interface_name,
                     match source {
                         IpAddr::V4(..) => "IPv4",
                         _ => "IPv6",
                     },
                     source,
                     destination,
                     protocol,
                     packet.len())
        }

    }
}

fn handle_ipv4_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = Ipv4Packet::new(ethernet.payload());
    if let Some(header) = header {
        handle_transport_protocol(interface_name,
                                  IpAddr::V4(header.get_source()),
                                  IpAddr::V4(header.get_destination()),
                                  header.get_next_level_protocol(),
                                  header.payload());
    } else {
        println!("[{}]: Malformed IPv4 Packet", interface_name);
    }
}

fn handle_ipv6_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = Ipv6Packet::new(ethernet.payload());
    if let Some(header) = header {
        handle_transport_protocol(interface_name,
                                  IpAddr::V6(header.get_source()),
                                  IpAddr::V6(header.get_destination()),
                                  header.get_next_header(),
                                  header.payload());
    } else {
        println!("[{}]: Malformed IPv6 Packet", interface_name);
    }
}

fn handle_arp_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = ArpPacket::new(ethernet.payload());
    if let Some(header) = header {
        println!("[{}]: ARP packet: {}({}) > {}({}); operation: {:?}",
                 interface_name,
                 ethernet.get_source(),
                 header.get_sender_proto_addr(),
                 ethernet.get_destination(),
                 header.get_target_proto_addr(),
                 header.get_operation());
    } else {
        println!("[{}]: Malformed ARP Packet", interface_name);
    }
}

fn handle_packet(interface_name: &str, ethernet: &EthernetPacket) {
    match ethernet.get_ethertype() {
        EtherTypes::Ipv4 => handle_ipv4_packet(interface_name, ethernet),
        EtherTypes::Ipv6 => handle_ipv6_packet(interface_name, ethernet),
        EtherTypes::Arp => handle_arp_packet(interface_name, ethernet),
        _ => {
            println!("[{}]: Unknown packet: {} > {}; ethertype: {:?} length: {}",
                     interface_name,
                     ethernet.get_source(),
                     ethernet.get_destination(),
                     ethernet.get_ethertype(),
                     ethernet.packet().len())
        }
    }
}

fn main() {
    use pnet::datalink::Channel::Ethernet;

    let iface_name = match env::args().nth(1) {
        Some(n) => n,
        None => {
            writeln!(io::stderr(), "USAGE: packetdump <NETWORK INTERFACE>").unwrap();
            process::exit(1);
        },
    };
    let interface_names_match = |iface: &NetworkInterface| iface.name == iface_name;

    // Find the network interface with the provided name
    let interfaces = datalink::interfaces();
    let interface = interfaces.into_iter().filter(interface_names_match).next().unwrap();

    // Create a channel to receive on
    let (_, mut rx) = match datalink::channel(&interface, Default::default()) {
        Ok(Ethernet(tx, rx)) => (tx, rx),
        Ok(_) => panic!("packetdump: unhandled channel type: {}"),
        Err(e) => panic!("packetdump: unable to create channel: {}", e),
    };

    loop {
        match rx.next() {
            Ok(packet) => handle_packet(&interface.name[..], &EthernetPacket::new(packet).unwrap()),
            Err(e) => panic!("packetdump: unable to receive packet: {}", e),
        }
    }
}
<commit_msg>Support `TUN` interface packet dump for Mac OS X platform<commit_after>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/// This example shows a basic packet logger using libpnet

extern crate pnet;

use pnet::datalink::{self, NetworkInterface};

use pnet::packet::Packet;
use pnet::packet::arp::ArpPacket;
use pnet::packet::ethernet::{EtherTypes, EthernetPacket, MutableEthernetPacket};
use pnet::packet::icmpv6::Icmpv6Packet;
use pnet::packet::icmp::{echo_reply, echo_request, IcmpPacket, IcmpTypes};
use pnet::packet::ip::{IpNextHeaderProtocol, IpNextHeaderProtocols};
use pnet::packet::ipv4::Ipv4Packet;
use pnet::packet::ipv6::Ipv6Packet;
use pnet::packet::tcp::TcpPacket;
use pnet::packet::udp::UdpPacket;
use pnet::util::MacAddr;

use std::env;
use std::io::{self, Write};
use std::process;
use std::net::IpAddr;

fn handle_udp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let udp = UdpPacket::new(packet);

    if let Some(udp) = udp {
        println!("[{}]: UDP Packet: {}:{} > {}:{}; length: {}",
                 interface_name,
                 source,
                 udp.get_source(),
                 destination,
                 udp.get_destination(),
                 udp.get_length());
    } else {
        println!("[{}]: Malformed UDP Packet", interface_name);
    }
}

fn handle_icmp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let icmp_packet = IcmpPacket::new(packet);
    if let Some(icmp_packet) = icmp_packet {
        match icmp_packet.get_icmp_type() {
            IcmpTypes::EchoReply => {
                let echo_reply_packet = echo_reply::EchoReplyPacket::new(packet).unwrap();
                println!("[{}]: ICMP echo reply {} -> {} (seq={:?}, id={:?})",
                         interface_name,
                         source,
                         destination,
                         echo_reply_packet.get_sequence_number(),
                         echo_reply_packet.get_identifier());
            }
            IcmpTypes::EchoRequest => {
                let echo_request_packet = echo_request::EchoRequestPacket::new(packet).unwrap();
                println!("[{}]: ICMP echo request {} -> {} (seq={:?}, id={:?})",
                         interface_name,
                         source,
                         destination,
                         echo_request_packet.get_sequence_number(),
                         echo_request_packet.get_identifier());
            }
            _ => {
                println!("[{}]: ICMP packet {} -> {} (type={:?})",
                         interface_name,
                         source,
                         destination,
                         icmp_packet.get_icmp_type())
            }
        }
    } else {
        println!("[{}]: Malformed ICMP Packet", interface_name);
    }
}

fn handle_icmpv6_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let icmpv6_packet = Icmpv6Packet::new(packet);
    if let Some(icmpv6_packet) = icmpv6_packet {
        println!(
            "[{}]: ICMPv6 packet {} -> {} (type={:?})",
            interface_name,
            source,
            destination,
            icmpv6_packet.get_icmpv6_type()
        )
    } else {
        println!("[{}]: Malformed ICMPv6 Packet", interface_name);
    }
}

fn handle_tcp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let tcp = TcpPacket::new(packet);
    if let Some(tcp) = tcp {
        println!("[{}]: TCP Packet: {}:{} > {}:{}; length: {}",
                 interface_name,
                 source,
                 tcp.get_source(),
                 destination,
                 tcp.get_destination(),
                 packet.len());
    } else {
        println!("[{}]: Malformed TCP Packet", interface_name);
    }
}

fn handle_transport_protocol(interface_name: &str,
                             source: IpAddr,
                             destination: IpAddr,
                             protocol: IpNextHeaderProtocol,
                             packet: &[u8]) {
    match protocol {
        IpNextHeaderProtocols::Udp => {
            handle_udp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Tcp => {
            handle_tcp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Icmp => {
            handle_icmp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Icmpv6 => {
            handle_icmpv6_packet(interface_name, source, destination, packet)
        }
        _ => {
            println!("[{}]: Unknown {} packet: {} > {}; protocol: {:?} length: {}",
                     interface_name,
                     match source {
                         IpAddr::V4(..) => "IPv4",
                         _ => "IPv6",
                     },
                     source,
                     destination,
                     protocol,
                     packet.len())
        }

    }
}

fn handle_ipv4_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = Ipv4Packet::new(ethernet.payload());
    if let Some(header) = header {
        handle_transport_protocol(interface_name,
                                  IpAddr::V4(header.get_source()),
                                  IpAddr::V4(header.get_destination()),
                                  header.get_next_level_protocol(),
                                  header.payload());
    } else {
        println!("[{}]: Malformed IPv4 Packet", interface_name);
    }
}

fn handle_ipv6_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = Ipv6Packet::new(ethernet.payload());
    if let Some(header) = header {
        handle_transport_protocol(interface_name,
                                  IpAddr::V6(header.get_source()),
                                  IpAddr::V6(header.get_destination()),
                                  header.get_next_header(),
                                  header.payload());
    } else {
        println!("[{}]: Malformed IPv6 Packet", interface_name);
    }
}

fn handle_arp_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = ArpPacket::new(ethernet.payload());
    if let Some(header) = header {
        println!("[{}]: ARP packet: {}({}) > {}({}); operation: {:?}",
                 interface_name,
                 ethernet.get_source(),
                 header.get_sender_proto_addr(),
                 ethernet.get_destination(),
                 header.get_target_proto_addr(),
                 header.get_operation());
    } else {
        println!("[{}]: Malformed ARP Packet", interface_name);
    }
}

fn handle_ethernet_frame(interface: &NetworkInterface, ethernet: &EthernetPacket) {
    let interface_name = &interface.name[..];
    match ethernet.get_ethertype() {
        EtherTypes::Ipv4 => handle_ipv4_packet(interface_name, ethernet),
        EtherTypes::Ipv6 => handle_ipv6_packet(interface_name, ethernet),
        EtherTypes::Arp => handle_arp_packet(interface_name, ethernet),
        _ => {
            println!(
                "[{}]: Unknown packet: {} > {}; ethertype: {:?} length: {}",
                interface_name,
                ethernet.get_source(),
                ethernet.get_destination(),
                ethernet.get_ethertype(),
                ethernet.packet().len()
            )
        }
    }
}

fn main() {
    use pnet::datalink::Channel::Ethernet;

    let iface_name = match env::args().nth(1) {
        Some(n) => n,
        None => {
            writeln!(io::stderr(), "USAGE: packetdump <NETWORK INTERFACE>").unwrap();
            process::exit(1);
        },
    };
    let interface_names_match = |iface: &NetworkInterface| iface.name == iface_name;

    // Find the network interface with the provided name
    let interfaces = datalink::interfaces();
    let interface = interfaces.into_iter().filter(interface_names_match).next().unwrap();

    // Create a channel to receive on
    let (_, mut rx) = match datalink::channel(&interface, Default::default()) {
        Ok(Ethernet(tx, rx)) => (tx, rx),
        Ok(_) => panic!("packetdump: unhandled channel type: {}"),
        Err(e) => panic!("packetdump: unable to create channel: {}", e),
    };

    loop {
        let mut buf: [u8; 1600] = [0u8; 1600];
        let mut fake_ethernet_frame = MutableEthernetPacket::new(&mut buf[..]).unwrap();
        match rx.next() {
            Ok(packet) => {
                if cfg!(target_os = "macos")
                    && interface.is_up()
                    && !interface.is_broadcast()
                    && !interface.is_loopback()
                    && interface.is_point_to_point() {
                        // Maybe is TUN interface
                        let version = Ipv4Packet::new(&packet).unwrap().get_version();
                        if version == 4 {
                            fake_ethernet_frame.set_destination(MacAddr(0,0,0,0,0,0));
                            fake_ethernet_frame.set_source(MacAddr(0,0,0,0,0,0));
                            fake_ethernet_frame.set_ethertype(EtherTypes::Ipv4);
                            fake_ethernet_frame.set_payload(&packet);
                            handle_ethernet_frame(&interface, &fake_ethernet_frame.to_immutable());
                            continue;
                        } else if version == 6 {
                            fake_ethernet_frame.set_destination(MacAddr(0,0,0,0,0,0));
                            fake_ethernet_frame.set_source(MacAddr(0,0,0,0,0,0));
                            fake_ethernet_frame.set_ethertype(EtherTypes::Ipv6);
                            fake_ethernet_frame.set_payload(&packet);
                            handle_ethernet_frame(&interface, &fake_ethernet_frame.to_immutable());
                            continue;
                        }
                }
                handle_ethernet_frame(&interface, &EthernetPacket::new(packet).unwrap());
            },
            Err(e) => panic!("packetdump: unable to receive packet: {}", e),
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/// This example shows a basic packet logger using libpnet
extern crate pnet;

use pnet::datalink::{self, NetworkInterface};

use pnet::packet::arp::ArpPacket;
use pnet::packet::ethernet::{EtherTypes, EthernetPacket, MutableEthernetPacket};
use pnet::packet::icmp::{echo_reply, echo_request, IcmpPacket, IcmpTypes};
use pnet::packet::icmpv6::Icmpv6Packet;
use pnet::packet::ip::{IpNextHeaderProtocol, IpNextHeaderProtocols};
use pnet::packet::ipv4::Ipv4Packet;
use pnet::packet::ipv6::Ipv6Packet;
use pnet::packet::tcp::TcpPacket;
use pnet::packet::udp::UdpPacket;
use pnet::packet::Packet;
use pnet::util::MacAddr;

use std::env;
use std::io::{self, Write};
use std::net::IpAddr;
use std::process;

fn handle_udp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let udp = UdpPacket::new(packet);

    if let Some(udp) = udp {
        println!(
            "[{}]: UDP Packet: {}:{} > {}:{}; length: {}",
            interface_name,
            source,
            udp.get_source(),
            destination,
            udp.get_destination(),
            udp.get_length()
        );
    } else {
        println!("[{}]: Malformed UDP Packet", interface_name);
    }
}

fn handle_icmp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let icmp_packet = IcmpPacket::new(packet);
    if let Some(icmp_packet) = icmp_packet {
        match icmp_packet.get_icmp_type() {
            IcmpTypes::EchoReply => {
                let echo_reply_packet = echo_reply::EchoReplyPacket::new(packet).unwrap();
                println!(
                    "[{}]: ICMP echo reply {} -> {} (seq={:?}, id={:?})",
                    interface_name,
                    source,
                    destination,
                    echo_reply_packet.get_sequence_number(),
                    echo_reply_packet.get_identifier()
                );
            }
            IcmpTypes::EchoRequest => {
                let echo_request_packet = echo_request::EchoRequestPacket::new(packet).unwrap();
                println!(
                    "[{}]: ICMP echo request {} -> {} (seq={:?}, id={:?})",
                    interface_name,
                    source,
                    destination,
                    echo_request_packet.get_sequence_number(),
                    echo_request_packet.get_identifier()
                );
            }
            _ => println!(
                "[{}]: ICMP packet {} -> {} (type={:?})",
                interface_name,
                source,
                destination,
                icmp_packet.get_icmp_type()
            ),
        }
    } else {
        println!("[{}]: Malformed ICMP Packet", interface_name);
    }
}

fn handle_icmpv6_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let icmpv6_packet = Icmpv6Packet::new(packet);
    if let Some(icmpv6_packet) = icmpv6_packet {
        println!(
            "[{}]: ICMPv6 packet {} -> {} (type={:?})",
            interface_name,
            source,
            destination,
            icmpv6_packet.get_icmpv6_type()
        )
    } else {
        println!("[{}]: Malformed ICMPv6 Packet", interface_name);
    }
}

fn handle_tcp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let tcp = TcpPacket::new(packet);
    if let Some(tcp) = tcp {
        println!(
            "[{}]: TCP Packet: {}:{} > {}:{}; length: {}",
            interface_name,
            source,
            tcp.get_source(),
            destination,
            tcp.get_destination(),
            packet.len()
        );
    } else {
        println!("[{}]: Malformed TCP Packet", interface_name);
    }
}

fn handle_transport_protocol(
    interface_name: &str,
    source: IpAddr,
    destination: IpAddr,
    protocol: IpNextHeaderProtocol,
    packet: &[u8],
) {
    match protocol {
        IpNextHeaderProtocols::Udp => {
            handle_udp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Tcp => {
            handle_tcp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Icmp => {
            handle_icmp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Icmpv6 => {
            handle_icmpv6_packet(interface_name, source, destination, packet)
        }
        _ => println!(
            "[{}]: Unknown {} packet: {} > {}; protocol: {:?} length: {}",
            interface_name,
            match source {
                IpAddr::V4(..) => "IPv4",
                _ => "IPv6",
            },
            source,
            destination,
            protocol,
            packet.len()
        ),
    }
}

fn handle_ipv4_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = Ipv4Packet::new(ethernet.payload());
    if let Some(header) = header {
        handle_transport_protocol(
            interface_name,
            IpAddr::V4(header.get_source()),
            IpAddr::V4(header.get_destination()),
            header.get_next_level_protocol(),
            header.payload(),
        );
    } else {
        println!("[{}]: Malformed IPv4 Packet", interface_name);
    }
}

fn handle_ipv6_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = Ipv6Packet::new(ethernet.payload());
    if let Some(header) = header {
        handle_transport_protocol(
            interface_name,
            IpAddr::V6(header.get_source()),
            IpAddr::V6(header.get_destination()),
            header.get_next_header(),
            header.payload(),
        );
    } else {
        println!("[{}]: Malformed IPv6 Packet", interface_name);
    }
}

fn handle_arp_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = ArpPacket::new(ethernet.payload());
    if let Some(header) = header {
        println!(
            "[{}]: ARP packet: {}({}) > {}({}); operation: {:?}",
            interface_name,
            ethernet.get_source(),
            header.get_sender_proto_addr(),
            ethernet.get_destination(),
            header.get_target_proto_addr(),
            header.get_operation()
        );
    } else {
        println!("[{}]: Malformed ARP Packet", interface_name);
    }
}

fn handle_ethernet_frame(interface: &NetworkInterface, ethernet: &EthernetPacket) {
    let interface_name = &interface.name[..];
    match ethernet.get_ethertype() {
        EtherTypes::Ipv4 => handle_ipv4_packet(interface_name, ethernet),
        EtherTypes::Ipv6 => handle_ipv6_packet(interface_name, ethernet),
        EtherTypes::Arp => handle_arp_packet(interface_name, ethernet),
        _ => println!(
            "[{}]: Unknown packet: {} > {}; ethertype: {:?} length: {}",
            interface_name,
            ethernet.get_source(),
            ethernet.get_destination(),
            ethernet.get_ethertype(),
            ethernet.packet().len()
        ),
    }
}

fn main() {
    use pnet::datalink::Channel::Ethernet;

    let iface_name = match env::args().nth(1) {
        Some(n) => n,
        None => {
            writeln!(io::stderr(), "USAGE: packetdump <NETWORK INTERFACE>").unwrap();
            process::exit(1);
        }
    };
    let interface_names_match = |iface: &NetworkInterface| iface.name == iface_name;

    // Find the network interface with the provided name
    let interfaces = datalink::interfaces();
    let interface = interfaces
        .into_iter()
        .filter(interface_names_match)
        .next()
        .unwrap();

    // Create a channel to receive on
    let (_, mut rx) = match datalink::channel(&interface, Default::default()) {
        Ok(Ethernet(tx, rx)) => (tx, rx),
        Ok(_) => panic!("packetdump: unhandled channel type: {}"),
        Err(e) => panic!("packetdump: unable to create channel: {}", e),
    };

    loop {
        let mut buf: [u8; 1600] = [0u8; 1600];
        let mut fake_ethernet_frame = MutableEthernetPacket::new(&mut buf[..]).unwrap();
        match rx.next() {
            Ok(packet) => {
                if cfg!(target_os = "macos") && interface.is_up() && !interface.is_broadcast()
                    && !interface.is_loopback() && interface.is_point_to_point()
                {
                    // Maybe is TUN interface
                    let version = Ipv4Packet::new(&packet).unwrap().get_version();
                    if version == 4 {
                        fake_ethernet_frame.set_destination(MacAddr(0, 0, 0, 0, 0, 0));
                        fake_ethernet_frame.set_source(MacAddr(0, 0, 0, 0, 0, 0));
                        fake_ethernet_frame.set_ethertype(EtherTypes::Ipv4);
                        fake_ethernet_frame.set_payload(&packet);
                        handle_ethernet_frame(&interface, &fake_ethernet_frame.to_immutable());
                        continue;
                    } else if version == 6 {
                        fake_ethernet_frame.set_destination(MacAddr(0, 0, 0, 0, 0, 0));
                        fake_ethernet_frame.set_source(MacAddr(0, 0, 0, 0, 0, 0));
                        fake_ethernet_frame.set_ethertype(EtherTypes::Ipv6);
                        fake_ethernet_frame.set_payload(&packet);
                        handle_ethernet_frame(&interface, &fake_ethernet_frame.to_immutable());
                        continue;
                    }
                }
                handle_ethernet_frame(&interface, &EthernetPacket::new(packet).unwrap());
            }
            Err(e) => panic!("packetdump: unable to receive packet: {}", e),
        }
    }
}
<commit_msg>Fix packet dump example on macos loopback interface.<commit_after>// Copyright (c) 2014, 2015 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/// This example shows a basic packet logger using libpnet
extern crate pnet;

use pnet::datalink::{self, NetworkInterface};

use pnet::packet::arp::ArpPacket;
use pnet::packet::ethernet::{EtherTypes, EthernetPacket, MutableEthernetPacket};
use pnet::packet::icmp::{echo_reply, echo_request, IcmpPacket, IcmpTypes};
use pnet::packet::icmpv6::Icmpv6Packet;
use pnet::packet::ip::{IpNextHeaderProtocol, IpNextHeaderProtocols};
use pnet::packet::ipv4::Ipv4Packet;
use pnet::packet::ipv6::Ipv6Packet;
use pnet::packet::tcp::TcpPacket;
use pnet::packet::udp::UdpPacket;
use pnet::packet::Packet;
use pnet::util::MacAddr;

use std::env;
use std::io::{self, Write};
use std::net::IpAddr;
use std::process;

fn handle_udp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let udp = UdpPacket::new(packet);

    if let Some(udp) = udp {
        println!(
            "[{}]: UDP Packet: {}:{} > {}:{}; length: {}",
            interface_name,
            source,
            udp.get_source(),
            destination,
            udp.get_destination(),
            udp.get_length()
        );
    } else {
        println!("[{}]: Malformed UDP Packet", interface_name);
    }
}

fn handle_icmp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let icmp_packet = IcmpPacket::new(packet);
    if let Some(icmp_packet) = icmp_packet {
        match icmp_packet.get_icmp_type() {
            IcmpTypes::EchoReply => {
                let echo_reply_packet = echo_reply::EchoReplyPacket::new(packet).unwrap();
                println!(
                    "[{}]: ICMP echo reply {} -> {} (seq={:?}, id={:?})",
                    interface_name,
                    source,
                    destination,
                    echo_reply_packet.get_sequence_number(),
                    echo_reply_packet.get_identifier()
                );
            }
            IcmpTypes::EchoRequest => {
                let echo_request_packet = echo_request::EchoRequestPacket::new(packet).unwrap();
                println!(
                    "[{}]: ICMP echo request {} -> {} (seq={:?}, id={:?})",
                    interface_name,
                    source,
                    destination,
                    echo_request_packet.get_sequence_number(),
                    echo_request_packet.get_identifier()
                );
            }
            _ => println!(
                "[{}]: ICMP packet {} -> {} (type={:?})",
                interface_name,
                source,
                destination,
                icmp_packet.get_icmp_type()
            ),
        }
    } else {
        println!("[{}]: Malformed ICMP Packet", interface_name);
    }
}

fn handle_icmpv6_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let icmpv6_packet = Icmpv6Packet::new(packet);
    if let Some(icmpv6_packet) = icmpv6_packet {
        println!(
            "[{}]: ICMPv6 packet {} -> {} (type={:?})",
            interface_name,
            source,
            destination,
            icmpv6_packet.get_icmpv6_type()
        )
    } else {
        println!("[{}]: Malformed ICMPv6 Packet", interface_name);
    }
}

fn handle_tcp_packet(interface_name: &str, source: IpAddr, destination: IpAddr, packet: &[u8]) {
    let tcp = TcpPacket::new(packet);
    if let Some(tcp) = tcp {
        println!(
            "[{}]: TCP Packet: {}:{} > {}:{}; length: {}",
            interface_name,
            source,
            tcp.get_source(),
            destination,
            tcp.get_destination(),
            packet.len()
        );
    } else {
        println!("[{}]: Malformed TCP Packet", interface_name);
    }
}

fn handle_transport_protocol(
    interface_name: &str,
    source: IpAddr,
    destination: IpAddr,
    protocol: IpNextHeaderProtocol,
    packet: &[u8],
) {
    match protocol {
        IpNextHeaderProtocols::Udp => {
            handle_udp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Tcp => {
            handle_tcp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Icmp => {
            handle_icmp_packet(interface_name, source, destination, packet)
        }
        IpNextHeaderProtocols::Icmpv6 => {
            handle_icmpv6_packet(interface_name, source, destination, packet)
        }
        _ => println!(
            "[{}]: Unknown {} packet: {} > {}; protocol: {:?} length: {}",
            interface_name,
            match source {
                IpAddr::V4(..) => "IPv4",
                _ => "IPv6",
            },
            source,
            destination,
            protocol,
            packet.len()
        ),
    }
}

fn handle_ipv4_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = Ipv4Packet::new(ethernet.payload());
    if let Some(header) = header {
        handle_transport_protocol(
            interface_name,
            IpAddr::V4(header.get_source()),
            IpAddr::V4(header.get_destination()),
            header.get_next_level_protocol(),
            header.payload(),
        );
    } else {
        println!("[{}]: Malformed IPv4 Packet", interface_name);
    }
}

fn handle_ipv6_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = Ipv6Packet::new(ethernet.payload());
    if let Some(header) = header {
        handle_transport_protocol(
            interface_name,
            IpAddr::V6(header.get_source()),
            IpAddr::V6(header.get_destination()),
            header.get_next_header(),
            header.payload(),
        );
    } else {
        println!("[{}]: Malformed IPv6 Packet", interface_name);
    }
}

fn handle_arp_packet(interface_name: &str, ethernet: &EthernetPacket) {
    let header = ArpPacket::new(ethernet.payload());
    if let Some(header) = header {
        println!(
            "[{}]: ARP packet: {}({}) > {}({}); operation: {:?}",
            interface_name,
            ethernet.get_source(),
            header.get_sender_proto_addr(),
            ethernet.get_destination(),
            header.get_target_proto_addr(),
            header.get_operation()
        );
    } else {
        println!("[{}]: Malformed ARP Packet", interface_name);
    }
}

fn handle_ethernet_frame(interface: &NetworkInterface, ethernet: &EthernetPacket) {
    let interface_name = &interface.name[..];
    match ethernet.get_ethertype() {
        EtherTypes::Ipv4 => handle_ipv4_packet(interface_name, ethernet),
        EtherTypes::Ipv6 => handle_ipv6_packet(interface_name, ethernet),
        EtherTypes::Arp => handle_arp_packet(interface_name, ethernet),
        _ => println!(
            "[{}]: Unknown packet: {} > {}; ethertype: {:?} length: {}",
            interface_name,
            ethernet.get_source(),
            ethernet.get_destination(),
            ethernet.get_ethertype(),
            ethernet.packet().len()
        ),
    }
}

fn main() {
    use pnet::datalink::Channel::Ethernet;

    let iface_name = match env::args().nth(1) {
        Some(n) => n,
        None => {
            writeln!(io::stderr(), "USAGE: packetdump <NETWORK INTERFACE>").unwrap();
            process::exit(1);
        }
    };
    let interface_names_match = |iface: &NetworkInterface| iface.name == iface_name;

    // Find the network interface with the provided name
    let interfaces = datalink::interfaces();
    let interface = interfaces
        .into_iter()
        .filter(interface_names_match)
        .next()
        .unwrap();

    // Create a channel to receive on
    let (_, mut rx) = match datalink::channel(&interface, Default::default()) {
        Ok(Ethernet(tx, rx)) => (tx, rx),
        Ok(_) => panic!("packetdump: unhandled channel type: {}"),
        Err(e) => panic!("packetdump: unable to create channel: {}", e),
    };

    loop {
        let mut buf: [u8; 1600] = [0u8; 1600];
        let mut fake_ethernet_frame = MutableEthernetPacket::new(&mut buf[..]).unwrap();
        match rx.next() {
            Ok(packet) => {
                let payload_offset;
                if cfg!(target_os = "macos") && interface.is_up() && !interface.is_broadcast()
                    && ((!interface.is_loopback() && interface.is_point_to_point())
                         || interface.is_loopback())
                {
                    if interface.is_loopback()
                    {
                        // The pnet code for BPF loopback adds a zero'd out Ethernet header
                        payload_offset = 14;
                    }
                    else
                    {
                        // Maybe is TUN interface
                        payload_offset = 0;
                    }
                    if packet.len() > payload_offset
                    {
                        let version = Ipv4Packet::new(&packet[payload_offset..]).unwrap().get_version();
                        if version == 4 {
                            fake_ethernet_frame.set_destination(MacAddr(0, 0, 0, 0, 0, 0));
                            fake_ethernet_frame.set_source(MacAddr(0, 0, 0, 0, 0, 0));
                            fake_ethernet_frame.set_ethertype(EtherTypes::Ipv4);
                            fake_ethernet_frame.set_payload(&packet[payload_offset..]);
                            handle_ethernet_frame(&interface, &fake_ethernet_frame.to_immutable());
                            continue;
                        } else if version == 6 {
                            fake_ethernet_frame.set_destination(MacAddr(0, 0, 0, 0, 0, 0));
                            fake_ethernet_frame.set_source(MacAddr(0, 0, 0, 0, 0, 0));
                            fake_ethernet_frame.set_ethertype(EtherTypes::Ipv6);
                            fake_ethernet_frame.set_payload(&packet[payload_offset..]);
                            handle_ethernet_frame(&interface, &fake_ethernet_frame.to_immutable());
                            continue;
                        }
                    }
                }
                handle_ethernet_frame(&interface, &EthernetPacket::new(packet).unwrap());
            }
            Err(e) => panic!("packetdump: unable to receive packet: {}", e),
        }
    }
}
<|endoftext|>


<commit_before>extern crate hal;
extern crate docopt;
extern crate arrayfire as af;

use af::{Array, Dim4};
use hal::{Model, Layer};
use hal::optimizer::{Optimizer, SGD};
use hal::error::HALError;
use hal::model::{Sequential};
use hal::layer::{Dense};

fn build_optimizer(name: &'static str) -> Result<Box<Optimizer>, HALError> {
  match name{
    "SGD" => Ok(Box::new(SGD::default())),
    _     => Err(HALError::UNKNOWN),
  }
}

fn generate_sin_wave(input_dims: u64, num_rows: u64) -> Vec<Array> {
  let dims = Dim4::new(&[input_dims, 1, 1, 1]);
  let mut range = af::range(dims, 0, af::Aftype::F32).unwrap();
  let mut waves = Vec::<Array>::new();
  for _ in (0..num_rows) {
    waves.push(af::sin(&range).unwrap());
    range = af::add(&(1.0 as f32), &range).unwrap();
  }
  waves
}

fn main() {
  // First we need to parameterize our network
  let input_dims = 8;
  let hidden_dims = 4;
  let output_dims = 8;
  let num_train_samples = 1024;
  let iter = 200;
  let batch_size = 128;
  let optimizer_type = "SGD";

  // Now, let's build a model with an optimizer and a loss function
  let mut model = Box::new(Sequential::new(build_optimizer(optimizer_type).unwrap(), "mse"));

  // Let's add a few layers why don't we?
  model.add(Box::new(Dense::new(input_dims, hidden_dims, "tanh", "normal", "ones")));
  model.add(Box::new(Dense::new(hidden_dims, output_dims, "tanh", "normal", "ones")));

  // Get some nice information about our model
  model.info();

  // Test with learning to predict sin wave
  let data = generate_sin_wave(input_dims, num_train_samples);

  // iterate our model
  let (loss, prediction) = model.fit(&data, &data, batch_size, iter, true);
  println!("prediction: "); af::print(&prediction);
  println!("loss: {:?}", loss);
}
<commit_msg>cleanup data generator<commit_after>extern crate hal;
extern crate docopt;
extern crate arrayfire as af;

use af::{Array, Dim4};
use hal::{Model, Layer};
use hal::optimizer::{Optimizer, SGD};
use hal::error::HALError;
use hal::model::{Sequential};
use hal::layer::{Dense};

fn build_optimizer(name: &'static str) -> Result<Box<Optimizer>, HALError> {
  match name{
    "SGD" => Ok(Box::new(SGD::default())),
    _     => Err(HALError::UNKNOWN),
  }
}

fn generate_sin_wave(input_dims: u64, num_rows: u64) -> Vec<Array> {
  let dims = Dim4::new(&[input_dims, 1, 1, 1]);
  let delta = 2.0*3.1415/(input_dims as f32 - 1.0);

  let index = af::range(dims, 0, af::Aftype::F32).unwrap();
  let mut range = af::mul(&index, &delta).unwrap();
  let mut waves = Vec::<Array>::new();

  for i in (input_dims..num_rows+input_dims) {
    waves.push(af::sin(&range).unwrap());
    range = af::mul(&af::add(&index, &(1.0 as f32)).unwrap(), &delta).unwrap();
  }
  waves
}

fn main() {
  // First we need to parameterize our network
  let input_dims = 8;
  let hidden_dims = 4;
  let output_dims = 8;
  let num_train_samples = 1024;
  let iter = 200;
  let batch_size = 128;
  let optimizer_type = "SGD";

  // Now, let's build a model with an optimizer and a loss function
  let mut model = Box::new(Sequential::new(build_optimizer(optimizer_type).unwrap(), "mse"));

  // Let's add a few layers why don't we?
  model.add(Box::new(Dense::new(input_dims, hidden_dims, "tanh", "normal", "ones")));
  model.add(Box::new(Dense::new(hidden_dims, output_dims, "tanh", "normal", "ones")));

  // Get some nice information about our model
  model.info();

  // Test with learning to predict sin wave
  let data = generate_sin_wave(input_dims, num_train_samples);

  // iterate our model
  let (loss, prediction) = model.fit(&data, &data, batch_size, iter, true);
  println!("prediction: "); af::print(&prediction);
  println!("loss: {:?}", loss);
}
<|endoftext|>


<commit_before>//! Example of using urlencoded
extern crate iron;
extern crate urlencoded;

use std::io::net::ip::Ipv4Addr;

use iron::{Iron, ServerT, Request, Response, Alloy, Chain};

use urlencoded::{UrlEncoded, Encoded};

// urlencoded returns a Hashmap
// Here we create a function to log the hashmap we are storing in Alloy.
// Alloy is where your middleware can store data and we access it through
// the `find` API exposed by alloy.
fn log_hashmap( _ : &mut Request, _ : &mut Response, alloy: &mut Alloy) {
    let hashmap = alloy.find::<Encoded>();
    match hashmap {
        Some(&Encoded(ref encoded)) => println!("Url Encoded:\n {}", encoded),
        None => ()
    }
}

// test out the server with `curl -i "127.0.0.1:3000/?name=franklin&name=trevor"`
fn main() {
    let mut server: ServerT = Iron::new();
    server.chain.link(UrlEncoded::new());
    server.chain.link(log_hashmap);
    server.listen(Ipv4Addr(127, 0, 0, 1), 3000);
}
<commit_msg>(refactor) Updated to use FromFn<commit_after>//! Example of using urlencoded
extern crate iron;
extern crate urlencoded;

use std::io::net::ip::Ipv4Addr;

use iron::{Iron, ServerT, Request, Response, Alloy, Chain};
use iron::middleware::{Status, Continue, FromFn};

use urlencoded::{UrlEncoded, Encoded};

// urlencoded returns a Hashmap
// Here we create a function to log the hashmap we are storing in Alloy.
// Alloy is where your middleware can store data and we access it through
// the `find` API exposed by alloy.
fn log_hashmap( _ : &mut Request, _ : &mut Response, alloy: &mut Alloy) -> Status {
    let hashmap = alloy.find::<Encoded>();
    match hashmap {
        Some(&Encoded(ref encoded)) => println!("Url Encoded:\n {}", encoded),
        None => ()
    }
    Continue
}

// test out the server with `curl -i "127.0.0.1:3000/?name=franklin&name=trevor"`
fn main() {
    let mut server: ServerT = Iron::new();
    server.chain.link(UrlEncoded::new());
    server.chain.link(FromFn::new(log_hashmap));
    server.listen(Ipv4Addr(127, 0, 0, 1), 3000);
}
<|endoftext|>


<commit_before>use crate::runtime::{blocking, context, io, time, Spawner};
use std::{error, fmt};

cfg_blocking! {
    use crate::runtime::task;
    use crate::runtime::blocking::task::BlockingTask;
}

cfg_rt_core! {
    use crate::task::JoinHandle;

    use std::future::Future;
}

/// Handle to the runtime.
///
/// The handle is internally reference-counted and can be freely cloned. A handle can be
/// obtained using the [`Runtime::handle`] method.
///
/// [`Runtime::handle`]: crate::runtime::Runtime::handle()
#[derive(Debug, Clone)]
pub struct Handle {
    pub(super) spawner: Spawner,

    /// Handles to the I/O drivers
    pub(super) io_handle: io::Handle,

    /// Handles to the time drivers
    pub(super) time_handle: time::Handle,

    /// Source of `Instant::now()`
    pub(super) clock: time::Clock,

    /// Blocking pool spawner
    pub(super) blocking_spawner: blocking::Spawner,
}

impl Handle {
    /// Enter the runtime context. This allows you to construct types that must
    /// have an executor available on creation such as [`Delay`] or [`TcpStream`].
    /// It will also allow you to call methods such as [`tokio::spawn`].
    ///
    /// This function is also available as [`Runtime::enter`].
    ///
    /// [`Delay`]: struct@crate::time::Delay
    /// [`TcpStream`]: struct@crate::net::TcpStream
    /// [`Runtime::enter`]: fn@crate::runtime::Runtime::enter
    /// [`tokio::spawn`]: fn@crate::spawn
    ///
    /// # Example
    ///
    /// ```
    /// use tokio::runtime::Runtime;
    ///
    /// fn function_that_spawns(msg: String) {
    ///     // Had we not used `handle.enter` below, this would panic.
    ///     tokio::spawn(async move {
    ///         println!("{}", msg);
    ///     });
    /// }
    ///
    /// fn main() {
    ///     let rt = Runtime::new().unwrap();
    ///     let handle = rt.handle().clone();
    ///
    ///     let s = "Hello World!".to_string();
    ///
    ///     // By entering the context, we tie `tokio::spawn` to this executor.
    ///     handle.enter(|| function_that_spawns(s));
    /// }
    /// ```
    pub fn enter<F, R>(&self, f: F) -> R
    where
        F: FnOnce() -> R,
    {
        context::enter(self.clone(), f)
    }

    /// Returns a Handle view over the currently running Runtime
    ///
    /// # Panic
    ///
    /// This will panic if called outside the context of a Tokio runtime.
    ///
    /// # Examples
    ///
    /// This can be used to obtain the handle of the surrounding runtime from an async
    /// block or function running on that runtime.
    ///
    /// ```
    /// # use tokio::runtime::Runtime;
    /// # fn dox() {
    /// # let rt = Runtime::new().unwrap();
    /// # rt.spawn(async {
    /// use tokio::runtime::Handle;
    ///
    /// // Inside an async block or function.
    /// let handle = Handle::current();
    /// handle.spawn(async {
    ///     println!("now running in the existing Runtime");
    /// })
    /// # });
    /// # }
    /// ```
    pub fn current() -> Self {
        context::current().expect("not currently running on the Tokio runtime.")
    }

    /// Returns a Handle view over the currently running Runtime
    ///
    /// Returns an error if no Runtime has been started
    ///
    /// Contrary to `current`, this never panics
    pub fn try_current() -> Result<Self, TryCurrentError> {
        context::current().ok_or(TryCurrentError(()))
    }
}

cfg_rt_core! {
    impl Handle {
        /// Spawns a future onto the Tokio runtime.
        ///
        /// This spawns the given future onto the runtime's executor, usually a
        /// thread pool. The thread pool is then responsible for polling the future
        /// until it completes.
        ///
        /// See [module level][mod] documentation for more details.
        ///
        /// [mod]: index.html
        ///
        /// # Examples
        ///
        /// ```
        /// use tokio::runtime::Runtime;
        ///
        /// # fn dox() {
        /// // Create the runtime
        /// let rt = Runtime::new().unwrap();
        /// let handle = rt.handle();
        ///
        /// // Spawn a future onto the runtime
        /// handle.spawn(async {
        ///     println!("now running on a worker thread");
        /// });
        /// # }
        /// ```
        ///
        /// # Panics
        ///
        /// This function will not panic unless task execution is disabled on the
        /// executor. This can only happen if the runtime was built using
        /// [`Builder`] without picking either [`basic_scheduler`] or
        /// [`threaded_scheduler`].
        ///
        /// [`Builder`]: struct@crate::runtime::Builder
        /// [`threaded_scheduler`]: fn@crate::runtime::Builder::threaded_scheduler
        /// [`basic_scheduler`]: fn@crate::runtime::Builder::basic_scheduler
        pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
        where
            F: Future + Send + 'static,
            F::Output: Send + 'static,
        {
            self.spawner.spawn(future)
        }

        /// Run a future to completion on the Tokio runtime from a synchronous
        /// context.
        ///
        /// This runs the given future on the runtime, blocking until it is
        /// complete, and yielding its resolved result. Any tasks or timers which
        /// the future spawns internally will be executed on the runtime.
        ///
        /// If the provided executor currently has no active core thread, this
        /// function might hang until a core thread is added. This is not a
        /// concern when using the [threaded scheduler], as it always has active
        /// core threads, but if you use the [basic scheduler], some other
        /// thread must currently be inside a call to [`Runtime::block_on`].
        /// See also [the module level documentation][1], which has a section on
        /// scheduler types.
        ///
        /// This method may not be called from an asynchronous context.
        ///
        /// [threaded scheduler]: fn@crate::runtime::Builder::threaded_scheduler
        /// [basic scheduler]: fn@crate::runtime::Builder::basic_scheduler
        /// [`Runtime::block_on`]: fn@crate::runtime::Runtime::block_on
        /// [1]: index.html#runtime-configurations
        ///
        /// # Panics
        ///
        /// This function panics if the provided future panics, or if called
        /// within an asynchronous execution context.
        ///
        /// # Examples
        ///
        /// Using `block_on` with the [threaded scheduler].
        ///
        /// ```
        /// use tokio::runtime::Runtime;
        /// use std::thread;
        ///
        /// // Create the runtime.
        /// //
        /// // If the rt-threaded feature is enabled, this creates a threaded
        /// // scheduler by default.
        /// let rt = Runtime::new().unwrap();
        /// let handle = rt.handle().clone();
        ///
        /// // Use the runtime from another thread.
        /// let th = thread::spawn(move || {
        ///     // Execute the future, blocking the current thread until completion.
        ///     //
        ///     // This example uses the threaded scheduler, so no concurrent call to
        ///     // `rt.block_on` is required.
        ///     handle.block_on(async {
        ///         println!("hello");
        ///     });
        /// });
        ///
        /// th.join().unwrap();
        /// ```
        ///
        /// Using the [basic scheduler] requires a concurrent call to
        /// [`Runtime::block_on`]:
        ///
        /// [threaded scheduler]: fn@crate::runtime::Builder::threaded_scheduler
        /// [basic scheduler]: fn@crate::runtime::Builder::basic_scheduler
        /// [`Runtime::block_on`]: fn@crate::runtime::Runtime::block_on
        ///
        /// ```
        /// use tokio::runtime::Builder;
        /// use tokio::sync::oneshot;
        /// use std::thread;
        ///
        /// // Create the runtime.
        /// let mut rt = Builder::new()
        ///     .enable_all()
        ///     .basic_scheduler()
        ///     .build()
        ///     .unwrap();
        ///
        /// let handle = rt.handle().clone();
        ///
        /// // Signal main thread when task has finished.
        /// let (send, recv) = oneshot::channel();
        ///
        /// // Use the runtime from another thread.
        /// let th = thread::spawn(move || {
        ///     // Execute the future, blocking the current thread until completion.
        ///     handle.block_on(async {
        ///         send.send("done").unwrap();
        ///     });
        /// });
        ///
        /// // The basic scheduler is used, so the thread above might hang if we
        /// // didn't call block_on on the rt too.
        /// rt.block_on(async {
        ///     assert_eq!(recv.await.unwrap(), "done");
        /// });
        /// # th.join().unwrap();
        /// ```
        ///
        pub fn block_on<F: Future>(&self, future: F) -> F::Output {
            self.enter(|| {
                let mut enter = crate::runtime::enter(true);
                enter.block_on(future).expect("failed to park thread")
            })
        }
    }
}

cfg_blocking! {
    impl Handle {
        /// Runs the provided closure on a thread where blocking is acceptable.
        ///
        /// In general, issuing a blocking call or performing a lot of compute in a
        /// future without yielding is not okay, as it may prevent the executor from
        /// driving other futures forward. This function runs the provided closure
        /// on a thread dedicated to blocking operations. See the [CPU-bound tasks
        /// and blocking code][blocking] section for more information.
        ///
        /// Tokio will spawn more blocking threads when they are requested through
        /// this function until the upper limit configured on the [`Builder`] is
        /// reached.  This limit is very large by default, because `spawn_blocking` is
        /// often used for various kinds of IO operations that cannot be performed
        /// asynchronously. When you run CPU-bound code using `spawn_blocking`, you
        /// should keep this large upper limit in mind; to run your CPU-bound
        /// computations on only a few threads, you should use a separate thread
        /// pool such as [rayon] rather than configuring the number of blocking
        /// threads.
        ///
        /// This function is intended for non-async operations that eventually
        /// finish on their own. If you want to spawn an ordinary thread, you should
        /// use [`thread::spawn`] instead.
        ///
        /// Closures spawned using `spawn_blocking` cannot be cancelled. When you
        /// shut down the executor, it will wait indefinitely for all blocking
        /// operations to finish. You can use [`shutdown_timeout`] to stop waiting
        /// for them after a certain timeout. Be aware that this will still not
        /// cancel the tasks — they are simply allowed to keep running after the
        /// method returns.
        ///
        /// Note that if you are using the [basic scheduler], this function will
        /// still spawn additional threads for blocking operations. The basic
        /// scheduler's single thread is only used for asynchronous code.
        ///
        /// [`Builder`]: struct@crate::runtime::Builder
        /// [blocking]: ../index.html#cpu-bound-tasks-and-blocking-code
        /// [rayon]: https://docs.rs/rayon
        /// [basic scheduler]: fn@crate::runtime::Builder::basic_scheduler
        /// [`thread::spawn`]: fn@std::thread::spawn
        /// [`shutdown_timeout`]: fn@crate::runtime::Runtime::shutdown_timeout
        ///
        /// # Examples
        ///
        /// ```
        /// use tokio::runtime::Runtime;
        ///
        /// # async fn docs() -> Result<(), Box<dyn std::error::Error>>{
        /// // Create the runtime
        /// let rt = Runtime::new().unwrap();
        /// let handle = rt.handle();
        ///
        /// let res = handle.spawn_blocking(move || {
        ///     // do some compute-heavy work or call synchronous code
        ///     "done computing"
        /// }).await?;
        ///
        /// assert_eq!(res, "done computing");
        /// # Ok(())
        /// # }
        /// ```
        pub fn spawn_blocking<F, R>(&self, f: F) -> JoinHandle<R>
        where
            F: FnOnce() -> R + Send + 'static,
            R: Send + 'static,
        {
            let (task, handle) = task::joinable(BlockingTask::new(f));
            let _ = self.blocking_spawner.spawn(task, self);
            handle
        }
    }
}

/// Error returned by `try_current` when no Runtime has been started
pub struct TryCurrentError(());

impl fmt::Debug for TryCurrentError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("TryCurrentError").finish()
    }
}

impl fmt::Display for TryCurrentError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("no tokio Runtime has been initialized")
    }
}

impl error::Error for TryCurrentError {}
<commit_msg>Clarifying that Handle::current must be called on a thread managed by tokio (#2493)<commit_after>use crate::runtime::{blocking, context, io, time, Spawner};
use std::{error, fmt};

cfg_blocking! {
    use crate::runtime::task;
    use crate::runtime::blocking::task::BlockingTask;
}

cfg_rt_core! {
    use crate::task::JoinHandle;

    use std::future::Future;
}

/// Handle to the runtime.
///
/// The handle is internally reference-counted and can be freely cloned. A handle can be
/// obtained using the [`Runtime::handle`] method.
///
/// [`Runtime::handle`]: crate::runtime::Runtime::handle()
#[derive(Debug, Clone)]
pub struct Handle {
    pub(super) spawner: Spawner,

    /// Handles to the I/O drivers
    pub(super) io_handle: io::Handle,

    /// Handles to the time drivers
    pub(super) time_handle: time::Handle,

    /// Source of `Instant::now()`
    pub(super) clock: time::Clock,

    /// Blocking pool spawner
    pub(super) blocking_spawner: blocking::Spawner,
}

impl Handle {
    /// Enter the runtime context. This allows you to construct types that must
    /// have an executor available on creation such as [`Delay`] or [`TcpStream`].
    /// It will also allow you to call methods such as [`tokio::spawn`].
    ///
    /// This function is also available as [`Runtime::enter`].
    ///
    /// [`Delay`]: struct@crate::time::Delay
    /// [`TcpStream`]: struct@crate::net::TcpStream
    /// [`Runtime::enter`]: fn@crate::runtime::Runtime::enter
    /// [`tokio::spawn`]: fn@crate::spawn
    ///
    /// # Example
    ///
    /// ```
    /// use tokio::runtime::Runtime;
    ///
    /// fn function_that_spawns(msg: String) {
    ///     // Had we not used `handle.enter` below, this would panic.
    ///     tokio::spawn(async move {
    ///         println!("{}", msg);
    ///     });
    /// }
    ///
    /// fn main() {
    ///     let rt = Runtime::new().unwrap();
    ///     let handle = rt.handle().clone();
    ///
    ///     let s = "Hello World!".to_string();
    ///
    ///     // By entering the context, we tie `tokio::spawn` to this executor.
    ///     handle.enter(|| function_that_spawns(s));
    /// }
    /// ```
    pub fn enter<F, R>(&self, f: F) -> R
    where
        F: FnOnce() -> R,
    {
        context::enter(self.clone(), f)
    }

    /// Returns a `Handle` view over the currently running `Runtime`
    ///
    /// # Panic
    ///
    /// This will panic if called outside the context of a Tokio runtime. That means that you must
    /// call this on one of the threads **being run by the runtime**. Calling this from within a
    /// thread created by `std::thread::spawn` (for example) will cause a panic.
    ///
    /// # Examples
    ///
    /// This can be used to obtain the handle of the surrounding runtime from an async
    /// block or function running on that runtime.
    ///
    /// ```
    /// # use std::thread;
    /// # use tokio::runtime::Runtime;
    /// # fn dox() {
    /// # let rt = Runtime::new().unwrap();
    /// # rt.spawn(async {
    /// use tokio::runtime::Handle;
    ///
    /// // Inside an async block or function.
    /// let handle = Handle::current();
    /// handle.spawn(async {
    ///     println!("now running in the existing Runtime");
    /// });
    ///
    /// # let handle =
    /// thread::spawn(move || {
    ///     // Notice that the handle is created outside of this thread and then moved in
    ///     handle.block_on(async { /* ... */ })
    ///     // This next line would cause a panic
    ///     // let handle2 = Handle::current();
    /// });
    /// # handle.join().unwrap();
    /// # });
    /// # }
    /// ```
    pub fn current() -> Self {
        context::current().expect("not currently running on the Tokio runtime.")
    }

    /// Returns a Handle view over the currently running Runtime
    ///
    /// Returns an error if no Runtime has been started
    ///
    /// Contrary to `current`, this never panics
    pub fn try_current() -> Result<Self, TryCurrentError> {
        context::current().ok_or(TryCurrentError(()))
    }
}

cfg_rt_core! {
    impl Handle {
        /// Spawns a future onto the Tokio runtime.
        ///
        /// This spawns the given future onto the runtime's executor, usually a
        /// thread pool. The thread pool is then responsible for polling the future
        /// until it completes.
        ///
        /// See [module level][mod] documentation for more details.
        ///
        /// [mod]: index.html
        ///
        /// # Examples
        ///
        /// ```
        /// use tokio::runtime::Runtime;
        ///
        /// # fn dox() {
        /// // Create the runtime
        /// let rt = Runtime::new().unwrap();
        /// let handle = rt.handle();
        ///
        /// // Spawn a future onto the runtime
        /// handle.spawn(async {
        ///     println!("now running on a worker thread");
        /// });
        /// # }
        /// ```
        ///
        /// # Panics
        ///
        /// This function will not panic unless task execution is disabled on the
        /// executor. This can only happen if the runtime was built using
        /// [`Builder`] without picking either [`basic_scheduler`] or
        /// [`threaded_scheduler`].
        ///
        /// [`Builder`]: struct@crate::runtime::Builder
        /// [`threaded_scheduler`]: fn@crate::runtime::Builder::threaded_scheduler
        /// [`basic_scheduler`]: fn@crate::runtime::Builder::basic_scheduler
        pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
        where
            F: Future + Send + 'static,
            F::Output: Send + 'static,
        {
            self.spawner.spawn(future)
        }

        /// Run a future to completion on the Tokio runtime from a synchronous
        /// context.
        ///
        /// This runs the given future on the runtime, blocking until it is
        /// complete, and yielding its resolved result. Any tasks or timers which
        /// the future spawns internally will be executed on the runtime.
        ///
        /// If the provided executor currently has no active core thread, this
        /// function might hang until a core thread is added. This is not a
        /// concern when using the [threaded scheduler], as it always has active
        /// core threads, but if you use the [basic scheduler], some other
        /// thread must currently be inside a call to [`Runtime::block_on`].
        /// See also [the module level documentation][1], which has a section on
        /// scheduler types.
        ///
        /// This method may not be called from an asynchronous context.
        ///
        /// [threaded scheduler]: fn@crate::runtime::Builder::threaded_scheduler
        /// [basic scheduler]: fn@crate::runtime::Builder::basic_scheduler
        /// [`Runtime::block_on`]: fn@crate::runtime::Runtime::block_on
        /// [1]: index.html#runtime-configurations
        ///
        /// # Panics
        ///
        /// This function panics if the provided future panics, or if called
        /// within an asynchronous execution context.
        ///
        /// # Examples
        ///
        /// Using `block_on` with the [threaded scheduler].
        ///
        /// ```
        /// use tokio::runtime::Runtime;
        /// use std::thread;
        ///
        /// // Create the runtime.
        /// //
        /// // If the rt-threaded feature is enabled, this creates a threaded
        /// // scheduler by default.
        /// let rt = Runtime::new().unwrap();
        /// let handle = rt.handle().clone();
        ///
        /// // Use the runtime from another thread.
        /// let th = thread::spawn(move || {
        ///     // Execute the future, blocking the current thread until completion.
        ///     //
        ///     // This example uses the threaded scheduler, so no concurrent call to
        ///     // `rt.block_on` is required.
        ///     handle.block_on(async {
        ///         println!("hello");
        ///     });
        /// });
        ///
        /// th.join().unwrap();
        /// ```
        ///
        /// Using the [basic scheduler] requires a concurrent call to
        /// [`Runtime::block_on`]:
        ///
        /// [threaded scheduler]: fn@crate::runtime::Builder::threaded_scheduler
        /// [basic scheduler]: fn@crate::runtime::Builder::basic_scheduler
        /// [`Runtime::block_on`]: fn@crate::runtime::Runtime::block_on
        ///
        /// ```
        /// use tokio::runtime::Builder;
        /// use tokio::sync::oneshot;
        /// use std::thread;
        ///
        /// // Create the runtime.
        /// let mut rt = Builder::new()
        ///     .enable_all()
        ///     .basic_scheduler()
        ///     .build()
        ///     .unwrap();
        ///
        /// let handle = rt.handle().clone();
        ///
        /// // Signal main thread when task has finished.
        /// let (send, recv) = oneshot::channel();
        ///
        /// // Use the runtime from another thread.
        /// let th = thread::spawn(move || {
        ///     // Execute the future, blocking the current thread until completion.
        ///     handle.block_on(async {
        ///         send.send("done").unwrap();
        ///     });
        /// });
        ///
        /// // The basic scheduler is used, so the thread above might hang if we
        /// // didn't call block_on on the rt too.
        /// rt.block_on(async {
        ///     assert_eq!(recv.await.unwrap(), "done");
        /// });
        /// # th.join().unwrap();
        /// ```
        ///
        pub fn block_on<F: Future>(&self, future: F) -> F::Output {
            self.enter(|| {
                let mut enter = crate::runtime::enter(true);
                enter.block_on(future).expect("failed to park thread")
            })
        }
    }
}

cfg_blocking! {
    impl Handle {
        /// Runs the provided closure on a thread where blocking is acceptable.
        ///
        /// In general, issuing a blocking call or performing a lot of compute in a
        /// future without yielding is not okay, as it may prevent the executor from
        /// driving other futures forward. This function runs the provided closure
        /// on a thread dedicated to blocking operations. See the [CPU-bound tasks
        /// and blocking code][blocking] section for more information.
        ///
        /// Tokio will spawn more blocking threads when they are requested through
        /// this function until the upper limit configured on the [`Builder`] is
        /// reached.  This limit is very large by default, because `spawn_blocking` is
        /// often used for various kinds of IO operations that cannot be performed
        /// asynchronously. When you run CPU-bound code using `spawn_blocking`, you
        /// should keep this large upper limit in mind; to run your CPU-bound
        /// computations on only a few threads, you should use a separate thread
        /// pool such as [rayon] rather than configuring the number of blocking
        /// threads.
        ///
        /// This function is intended for non-async operations that eventually
        /// finish on their own. If you want to spawn an ordinary thread, you should
        /// use [`thread::spawn`] instead.
        ///
        /// Closures spawned using `spawn_blocking` cannot be cancelled. When you
        /// shut down the executor, it will wait indefinitely for all blocking
        /// operations to finish. You can use [`shutdown_timeout`] to stop waiting
        /// for them after a certain timeout. Be aware that this will still not
        /// cancel the tasks — they are simply allowed to keep running after the
        /// method returns.
        ///
        /// Note that if you are using the [basic scheduler], this function will
        /// still spawn additional threads for blocking operations. The basic
        /// scheduler's single thread is only used for asynchronous code.
        ///
        /// [`Builder`]: struct@crate::runtime::Builder
        /// [blocking]: ../index.html#cpu-bound-tasks-and-blocking-code
        /// [rayon]: https://docs.rs/rayon
        /// [basic scheduler]: fn@crate::runtime::Builder::basic_scheduler
        /// [`thread::spawn`]: fn@std::thread::spawn
        /// [`shutdown_timeout`]: fn@crate::runtime::Runtime::shutdown_timeout
        ///
        /// # Examples
        ///
        /// ```
        /// use tokio::runtime::Runtime;
        ///
        /// # async fn docs() -> Result<(), Box<dyn std::error::Error>>{
        /// // Create the runtime
        /// let rt = Runtime::new().unwrap();
        /// let handle = rt.handle();
        ///
        /// let res = handle.spawn_blocking(move || {
        ///     // do some compute-heavy work or call synchronous code
        ///     "done computing"
        /// }).await?;
        ///
        /// assert_eq!(res, "done computing");
        /// # Ok(())
        /// # }
        /// ```
        pub fn spawn_blocking<F, R>(&self, f: F) -> JoinHandle<R>
        where
            F: FnOnce() -> R + Send + 'static,
            R: Send + 'static,
        {
            let (task, handle) = task::joinable(BlockingTask::new(f));
            let _ = self.blocking_spawner.spawn(task, self);
            handle
        }
    }
}

/// Error returned by `try_current` when no Runtime has been started
pub struct TryCurrentError(());

impl fmt::Debug for TryCurrentError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("TryCurrentError").finish()
    }
}

impl fmt::Display for TryCurrentError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("no tokio Runtime has been initialized")
    }
}

impl error::Error for TryCurrentError {}
<|endoftext|>


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

//! This module provides an abstraction for positioning a node in a binary tree,
//! A `Position` uniquely identifies the location of a node
//!
//! In this implementation, `Position` is represented by the in-order-traversal sequence number
//! of the node.
//! The process of locating a node and jumping between nodes is done through position calculation,
//! which comes from treebits.
//!
//! For example
//! ```text
//!      3
//!     /  \
//!    /    \
//!   1      5 <-[Node index, a.k.a, Position]
//!  / \    / \
//! 0   2  4   6
//!
//! 0   1  2   3 <[Leaf index]
//! ```
//! Note1: The in-order-traversal counts from 0
//! Note2: The level of tree counts from leaf level, start from 0
//! Note3: The leaf index starting from left-most leaf, starts from 0

use super::treebits;

#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct Position(u64);

impl Position {
    pub fn from_inorder_index(index: u64) -> Self {
        Position(index)
    }

    pub fn to_inorder_index(self) -> u64 {
        self.0
    }

    pub fn from_postorder_index(index: u64) -> Self {
        Self::from_inorder_index(treebits::postorder_to_inorder(index))
    }

    pub fn to_postorder_index(self) -> u64 {
        treebits::inorder_to_postorder(self.to_inorder_index())
    }

    pub fn get_parent(self) -> Self {
        Self::from_inorder_index(treebits::parent(self.0))
    }

    // Note: if self is root, the sibling will overflow
    pub fn get_sibling(self) -> Self {
        Self::from_inorder_index(treebits::sibling(self.0))
    }

    // Requirement: self can not be leaf.
    pub fn get_left_child(self) -> Position {
        Self::from_inorder_index(treebits::left_child(self.0))
    }

    // Requirement: self can not be leaf.
    pub fn get_right_child(self) -> Position {
        Self::from_inorder_index(treebits::right_child(self.0))
    }

    // Note: if self is root, the direction will overflow (and will always be left)
    pub fn get_direction_for_self(self) -> treebits::NodeDirection {
        treebits::direction_from_parent(self.0)
    }

    // The level start from 0 counting from the leaf level
    pub fn get_level(self) -> u32 {
        treebits::level(self.0)
    }

    // Given the position, return the leaf index counting from the left
    pub fn to_leaf_index(self) -> u64 {
        treebits::pos_counting_from_left(self.0)
    }

    // Opposite of get_left_node_count_from_position.
    pub fn from_leaf_index(leaf_index: u64) -> Position {
        Self::from_inorder_index(treebits::node_from_level_and_pos(0, leaf_index))
    }

    /// Given a position, returns the position next to it on the right on the same level. For
    /// example, given input 5 this function should return 9.
    ///
    /// ```text
    ///       3
    ///    /     \
    ///   1       5       9
    ///  / \     / \     / \
    /// 0   2   4   6   8   10
    /// ```
    pub fn get_next_sibling(self) -> Position {
        let level = self.get_level();
        let pos = treebits::pos_counting_from_left(self.0);
        Position(treebits::node_from_level_and_pos(level, pos + 1))
    }

    // Given a leaf index, calculate the position of a minimum root which contains this leaf
    pub fn get_root_position(leaf_index: u64) -> Position {
        let leaf = Self::from_leaf_index(leaf_index);
        Self::from_inorder_index(treebits::get_root(leaf.0))
    }

    // Given index of right most leaf, calculate if a position is the root
    // of a perfect subtree that does not contains placeholder nodes.
    pub fn is_freezable(self, leaf_index: u64) -> bool {
        let leaf = Self::from_leaf_index(leaf_index);
        treebits::is_freezable(self.0, leaf.0)
    }

    // Given index of right most leaf, calculate if a position should be a placeholder node at this
    // moment
    pub fn is_placeholder(self, leaf_index: u64) -> bool {
        let leaf = Self::from_leaf_index(leaf_index);
        treebits::is_placeholder(self.0, leaf.0)
    }

    /// Creates an `AncestorIterator` using this position.
    pub fn iter_ancestor(self) -> AncestorIterator {
        AncestorIterator { position: self }
    }

    /// Creates an `AncestorSiblingIterator` using this position.
    pub fn iter_ancestor_sibling(self) -> AncestorSiblingIterator {
        AncestorSiblingIterator { position: self }
    }
}

/// `AncestorSiblingIterator` generates current sibling position and moves itself to its parent
/// position for each iteration.
#[derive(Debug)]
pub struct AncestorSiblingIterator {
    position: Position,
}

impl Iterator for AncestorSiblingIterator {
    type Item = Position;

    fn next(&mut self) -> Option<Position> {
        let current_sibling_position = self.position.get_sibling();
        self.position = self.position.get_parent();
        Some(current_sibling_position)
    }
}

/// `AncestorIterator` generates current position and moves itself to its parent position for each
/// iteration.
#[derive(Debug)]
pub struct AncestorIterator {
    position: Position,
}

impl Iterator for AncestorIterator {
    type Item = Position;

    fn next(&mut self) -> Option<Position> {
        let current_position = self.position;
        self.position = self.position.get_parent();
        Some(current_position)
    }
}

/// Traverse leaves from left to right in groups that forms full subtrees, yielding root positions
/// of such subtrees.
/// Note that each 1-bit in num_leaves corresponds to a full subtree.
/// For example, in the below tree of 5=0b101 leaves, the two 1-bits corresponds to Fzn2 and L4
/// accordingly.
///
/// ```text
///            Non-fzn
///           /       \
///          /         \
///         /           \
///       Fzn2         Non-fzn
///      /   \           /   \
///     /     \         /     \
///    Fzn1    Fzn3  Non-fzn  [Placeholder]
///   /  \    /  \    /    \
///  L0  L1  L2  L3 L4   [Placeholder]
/// ```
pub struct FrozenSubTreeIterator {
    bitmap: u64,
    seen_leaves: u64,
}

impl FrozenSubTreeIterator {
    pub fn new(num_leaves: u64) -> Self {
        Self {
            bitmap: num_leaves,
            seen_leaves: 0,
        }
    }
}

impl Iterator for FrozenSubTreeIterator {
    type Item = Position;

    fn next(&mut self) -> Option<Position> {
        if self.bitmap == 0 {
            return None;
        }

        // Find the remaining biggest full subtree.
        // The MSB of the bitmap represents it. For example for a tree of 0x1010=10 leaves, the
        // biggest and leftmost full subtree has 0x1000=8 leaves, which can be got by smearing all
        // bits after MSB with 1-bits (got 0x1111), right shift once (got 0x0111) and add 1 (got
        // 0x1000=8). At the same time, we also observe that the in-order numbering of a full
        // subtree root is (num_leaves - 1) greater than that of the leftmost leaf, and also
        // (num_leaves - 1) less than that of the rightmost leaf.
        let root_offset = treebits::smear_ones_for_u64(self.bitmap) >> 1;
        let num_leaves = root_offset + 1;
        let leftmost_leaf = Position::from_leaf_index(self.seen_leaves);
        let root = Position::from_inorder_index(leftmost_leaf.to_inorder_index() + root_offset);

        // Mark it consumed.
        self.bitmap &= !num_leaves;
        self.seen_leaves += num_leaves;

        Some(root)
    }
}
<commit_msg>[Proof] Replace incorrect 0x with 0b in FrozenSubTreeIterator<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! This module provides an abstraction for positioning a node in a binary tree,
//! A `Position` uniquely identifies the location of a node
//!
//! In this implementation, `Position` is represented by the in-order-traversal sequence number
//! of the node.
//! The process of locating a node and jumping between nodes is done through position calculation,
//! which comes from treebits.
//!
//! For example
//! ```text
//!      3
//!     /  \
//!    /    \
//!   1      5 <-[Node index, a.k.a, Position]
//!  / \    / \
//! 0   2  4   6
//!
//! 0   1  2   3 <[Leaf index]
//! ```
//! Note1: The in-order-traversal counts from 0
//! Note2: The level of tree counts from leaf level, start from 0
//! Note3: The leaf index starting from left-most leaf, starts from 0

use super::treebits;

#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct Position(u64);

impl Position {
    pub fn from_inorder_index(index: u64) -> Self {
        Position(index)
    }

    pub fn to_inorder_index(self) -> u64 {
        self.0
    }

    pub fn from_postorder_index(index: u64) -> Self {
        Self::from_inorder_index(treebits::postorder_to_inorder(index))
    }

    pub fn to_postorder_index(self) -> u64 {
        treebits::inorder_to_postorder(self.to_inorder_index())
    }

    pub fn get_parent(self) -> Self {
        Self::from_inorder_index(treebits::parent(self.0))
    }

    // Note: if self is root, the sibling will overflow
    pub fn get_sibling(self) -> Self {
        Self::from_inorder_index(treebits::sibling(self.0))
    }

    // Requirement: self can not be leaf.
    pub fn get_left_child(self) -> Position {
        Self::from_inorder_index(treebits::left_child(self.0))
    }

    // Requirement: self can not be leaf.
    pub fn get_right_child(self) -> Position {
        Self::from_inorder_index(treebits::right_child(self.0))
    }

    // Note: if self is root, the direction will overflow (and will always be left)
    pub fn get_direction_for_self(self) -> treebits::NodeDirection {
        treebits::direction_from_parent(self.0)
    }

    // The level start from 0 counting from the leaf level
    pub fn get_level(self) -> u32 {
        treebits::level(self.0)
    }

    // Given the position, return the leaf index counting from the left
    pub fn to_leaf_index(self) -> u64 {
        treebits::pos_counting_from_left(self.0)
    }

    // Opposite of get_left_node_count_from_position.
    pub fn from_leaf_index(leaf_index: u64) -> Position {
        Self::from_inorder_index(treebits::node_from_level_and_pos(0, leaf_index))
    }

    /// Given a position, returns the position next to it on the right on the same level. For
    /// example, given input 5 this function should return 9.
    ///
    /// ```text
    ///       3
    ///    /     \
    ///   1       5       9
    ///  / \     / \     / \
    /// 0   2   4   6   8   10
    /// ```
    pub fn get_next_sibling(self) -> Position {
        let level = self.get_level();
        let pos = treebits::pos_counting_from_left(self.0);
        Position(treebits::node_from_level_and_pos(level, pos + 1))
    }

    // Given a leaf index, calculate the position of a minimum root which contains this leaf
    pub fn get_root_position(leaf_index: u64) -> Position {
        let leaf = Self::from_leaf_index(leaf_index);
        Self::from_inorder_index(treebits::get_root(leaf.0))
    }

    // Given index of right most leaf, calculate if a position is the root
    // of a perfect subtree that does not contains placeholder nodes.
    pub fn is_freezable(self, leaf_index: u64) -> bool {
        let leaf = Self::from_leaf_index(leaf_index);
        treebits::is_freezable(self.0, leaf.0)
    }

    // Given index of right most leaf, calculate if a position should be a placeholder node at this
    // moment
    pub fn is_placeholder(self, leaf_index: u64) -> bool {
        let leaf = Self::from_leaf_index(leaf_index);
        treebits::is_placeholder(self.0, leaf.0)
    }

    /// Creates an `AncestorIterator` using this position.
    pub fn iter_ancestor(self) -> AncestorIterator {
        AncestorIterator { position: self }
    }

    /// Creates an `AncestorSiblingIterator` using this position.
    pub fn iter_ancestor_sibling(self) -> AncestorSiblingIterator {
        AncestorSiblingIterator { position: self }
    }
}

/// `AncestorSiblingIterator` generates current sibling position and moves itself to its parent
/// position for each iteration.
#[derive(Debug)]
pub struct AncestorSiblingIterator {
    position: Position,
}

impl Iterator for AncestorSiblingIterator {
    type Item = Position;

    fn next(&mut self) -> Option<Position> {
        let current_sibling_position = self.position.get_sibling();
        self.position = self.position.get_parent();
        Some(current_sibling_position)
    }
}

/// `AncestorIterator` generates current position and moves itself to its parent position for each
/// iteration.
#[derive(Debug)]
pub struct AncestorIterator {
    position: Position,
}

impl Iterator for AncestorIterator {
    type Item = Position;

    fn next(&mut self) -> Option<Position> {
        let current_position = self.position;
        self.position = self.position.get_parent();
        Some(current_position)
    }
}

/// Traverse leaves from left to right in groups that forms full subtrees, yielding root positions
/// of such subtrees.
/// Note that each 1-bit in num_leaves corresponds to a full subtree.
/// For example, in the below tree of 5=0b101 leaves, the two 1-bits corresponds to Fzn2 and L4
/// accordingly.
///
/// ```text
///            Non-fzn
///           /       \
///          /         \
///         /           \
///       Fzn2         Non-fzn
///      /   \           /   \
///     /     \         /     \
///    Fzn1    Fzn3  Non-fzn  [Placeholder]
///   /  \    /  \    /    \
///  L0  L1  L2  L3 L4   [Placeholder]
/// ```
pub struct FrozenSubTreeIterator {
    bitmap: u64,
    seen_leaves: u64,
}

impl FrozenSubTreeIterator {
    pub fn new(num_leaves: u64) -> Self {
        Self {
            bitmap: num_leaves,
            seen_leaves: 0,
        }
    }
}

impl Iterator for FrozenSubTreeIterator {
    type Item = Position;

    fn next(&mut self) -> Option<Position> {
        if self.bitmap == 0 {
            return None;
        }

        // Find the remaining biggest full subtree.
        // The MSB of the bitmap represents it. For example for a tree of 0b1010=10 leaves, the
        // biggest and leftmost full subtree has 0b1000=8 leaves, which can be got by smearing all
        // bits after MSB with 1-bits (got 0b1111), right shift once (got 0b0111) and add 1 (got
        // 0b1000=8). At the same time, we also observe that the in-order numbering of a full
        // subtree root is (num_leaves - 1) greater than that of the leftmost leaf, and also
        // (num_leaves - 1) less than that of the rightmost leaf.
        let root_offset = treebits::smear_ones_for_u64(self.bitmap) >> 1;
        let num_leaves = root_offset + 1;
        let leftmost_leaf = Position::from_leaf_index(self.seen_leaves);
        let root = Position::from_inorder_index(leftmost_leaf.to_inorder_index() + root_offset);

        // Mark it consumed.
        self.bitmap &= !num_leaves;
        self.seen_leaves += num_leaves;

        Some(root)
    }
}
<|endoftext|>


<commit_before>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// Based off of libcss's examples/example1.c
mod example1 {

    use CssResult;
    use CssProperty;
    use types::*;
    use hint::*;
    use select::*;
    use util::VoidPtrLike;
    use wapcaplet::LwcString;

    struct MyDomNode {
        name: @LwcString
    }

    impl VoidPtrLike for MyDomNode {
        fn from_void_ptr(node: *libc::c_void) -> MyDomNode {
            assert!(node.is_not_null());
            MyDomNode {
                name: unsafe {
                    let box = cast::reinterpret_cast(&node);
                    cast::bump_box_refcount(box);
                    box
                }
            }
        }

        fn to_void_ptr(&self) -> *libc::c_void {
            unsafe { cast::reinterpret_cast(&self.name) }
        }
    }

    #[test]
    fn run() {
        use super::super::stylesheet::{CssUrlResolutionFn, CssStylesheetParams, CssStylesheet,
                                       css_stylesheet_create, CssStylesheetParamsVersion1};
        use super::super::computed::CssComputedStyle;
        use super::super::values::{CssColorColor, CssColorInherit};
        use super::super::ll::types::{CSS_ORIGIN_AUTHOR, CSS_MEDIA_ALL, CSS_MEDIA_SCREEN};
        use super::super::conversions::ToLl;
        use wapcaplet::{LwcString, from_rust_string};

        let data = "h1 { color: red; }\
                    h4 { color: #321; }\
                    h4, h5 { color: #123456; }";

        let resolve: CssUrlResolutionFn = |a,b| resolve_url(a, b);

        let params: CssStylesheetParams = CssStylesheetParams {
            params_version: CssStylesheetParamsVersion1,
            level: CssLevel21,
            charset: ~"UTF-8",
            url: ~"foo",
            title: ~"foo",
            allow_quirks: false,
            inline_style: false,
            resolve: Some(resolve),
            import: None,
            color: None,
            font: None,
        };

        let mut sheet: CssStylesheet = css_stylesheet_create(&params);
        debug!("stylesheet: %?", sheet);
        debug!("stylesheet size: %?", sheet.size());

        sheet.append_data(str::to_bytes(data));
        sheet.data_done();
        debug!("stylesheet size: %?", sheet.size());

        let mut select_ctx: CssSelectCtx = css_select_ctx_create();
        assert!(select_ctx.count_sheets() == 0);
        select_ctx.append_sheet(sheet, CSS_ORIGIN_AUTHOR, CSS_MEDIA_ALL);
        debug!("count sheets: %?", select_ctx.count_sheets());
        assert!(select_ctx.count_sheets() == 1);

        for uint::range(1, 7) |hh| {
            let element = fmt!("h%u", hh);
            let element_name: @LwcString = @from_rust_string(element);
            let node = MyDomNode { name: element_name };
            let select_handler = SelectHandler { bogus: () };
            let style: CssSelectResults = select_ctx.select_style(&node,
                                                                  CSS_MEDIA_SCREEN,
                                                                  None,
                                                                  &select_handler);

            let computed: CssComputedStyle = style.computed_style(CssPseudoElementNone);

            match computed.color() {
                CssColorInherit => {
                    debug!("color of h%u is 'inherit'", hh);
                }
                CssColorColor(color) => {
                    debug!("color of h%u is %x", hh, color.to_ll() as uint);
                }
            }
        }
    }

    fn resolve_url(_base: &str, _rel: &LwcString) -> CssResult<LwcString> {
        fail!(~"resolving url");
    }

    struct SelectHandler {
        bogus: ()
    }

    impl CssSelectHandler<MyDomNode> for SelectHandler {
        fn node_name(&self, node: &MyDomNode) -> CssQName {
            debug!("HL node_name!");
            debug!("SS %?", (*node.name).to_str());

            CssQName {
                ns: None,
                name: node.name.clone()
            }
        }

        fn node_id(&self, _node: &MyDomNode) -> Option<LwcString> { None }

        fn named_parent_node(&self, _node: &MyDomNode, _qname: &CssQName) -> Option<MyDomNode> {
            None
        }

        fn parent_node(&self, _node: &MyDomNode) -> Option<MyDomNode> {
            None
        }

        fn node_has_id(&self, _node: &MyDomNode, _name: LwcString) -> bool { false }

        fn named_ancestor_node(&self, _node: &MyDomNode, _qname: &CssQName) -> Option<MyDomNode> {
            None
        }

        fn node_is_root(&self, _node: &MyDomNode) -> bool { false }

        fn node_is_link(&self, _node: &MyDomNode) -> bool { false }

        fn ua_default_for_property(&self, property: CssProperty) -> CssHint {
            match property {
                _ => CssHintDefault
            }
        }
    }
}

#[test]
fn test_arc() {
    use std::arc::ARC;
    use stylesheet::*;
    use types::CssLevel21;
    use wapcaplet::LwcString;
    use super::CssResult;

    let resolve: CssUrlResolutionFn = |a,b| resolve_url(a, b);
    let params: CssStylesheetParams = CssStylesheetParams {
        params_version: CssStylesheetParamsVersion1,
        level: CssLevel21,
        charset: ~"UTF-8",
        url: ~"foo",
        title: ~"foo",
        allow_quirks: false,
        inline_style: false,
        resolve: Some(resolve),
        import: None,
        color: None,
        font: None,
    };

    let sheet: CssStylesheet = css_stylesheet_create(&params);
    let _arc = ARC(sheet);

    fn resolve_url(_base: &str, _rel: &LwcString) -> CssResult<LwcString> {
        fail!(~"resolving url");
    }
}
<commit_msg>Fix incorrect clone.<commit_after>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// Based off of libcss's examples/example1.c
mod example1 {

    use CssResult;
    use CssProperty;
    use types::*;
    use hint::*;
    use select::*;
    use util::VoidPtrLike;
    use wapcaplet::LwcString;

    struct MyDomNode {
        name: @LwcString
    }

    impl VoidPtrLike for MyDomNode {
        fn from_void_ptr(node: *libc::c_void) -> MyDomNode {
            assert!(node.is_not_null());
            MyDomNode {
                name: unsafe {
                    let box = cast::reinterpret_cast(&node);
                    cast::bump_box_refcount(box);
                    box
                }
            }
        }

        fn to_void_ptr(&self) -> *libc::c_void {
            unsafe { cast::reinterpret_cast(&self.name) }
        }
    }

    #[test]
    fn run() {
        use super::super::stylesheet::{CssUrlResolutionFn, CssStylesheetParams, CssStylesheet,
                                       css_stylesheet_create, CssStylesheetParamsVersion1};
        use super::super::computed::CssComputedStyle;
        use super::super::values::{CssColorColor, CssColorInherit};
        use super::super::ll::types::{CSS_ORIGIN_AUTHOR, CSS_MEDIA_ALL, CSS_MEDIA_SCREEN};
        use super::super::conversions::ToLl;
        use wapcaplet::{LwcString, from_rust_string};

        let data = "h1 { color: red; }\
                    h4 { color: #321; }\
                    h4, h5 { color: #123456; }";

        let resolve: CssUrlResolutionFn = |a,b| resolve_url(a, b);

        let params: CssStylesheetParams = CssStylesheetParams {
            params_version: CssStylesheetParamsVersion1,
            level: CssLevel21,
            charset: ~"UTF-8",
            url: ~"foo",
            title: ~"foo",
            allow_quirks: false,
            inline_style: false,
            resolve: Some(resolve),
            import: None,
            color: None,
            font: None,
        };

        let mut sheet: CssStylesheet = css_stylesheet_create(&params);
        debug!("stylesheet: %?", sheet);
        debug!("stylesheet size: %?", sheet.size());

        sheet.append_data(str::to_bytes(data));
        sheet.data_done();
        debug!("stylesheet size: %?", sheet.size());

        let mut select_ctx: CssSelectCtx = css_select_ctx_create();
        assert!(select_ctx.count_sheets() == 0);
        select_ctx.append_sheet(sheet, CSS_ORIGIN_AUTHOR, CSS_MEDIA_ALL);
        debug!("count sheets: %?", select_ctx.count_sheets());
        assert!(select_ctx.count_sheets() == 1);

        for uint::range(1, 7) |hh| {
            let element = fmt!("h%u", hh);
            let element_name: @LwcString = @from_rust_string(element);
            let node = MyDomNode { name: element_name };
            let select_handler = SelectHandler { bogus: () };
            let style: CssSelectResults = select_ctx.select_style(&node,
                                                                  CSS_MEDIA_SCREEN,
                                                                  None,
                                                                  &select_handler);

            let computed: CssComputedStyle = style.computed_style(CssPseudoElementNone);

            match computed.color() {
                CssColorInherit => {
                    debug!("color of h%u is 'inherit'", hh);
                }
                CssColorColor(color) => {
                    debug!("color of h%u is %x", hh, color.to_ll() as uint);
                }
            }
        }
    }

    fn resolve_url(_base: &str, _rel: &LwcString) -> CssResult<LwcString> {
        fail!(~"resolving url");
    }

    struct SelectHandler {
        bogus: ()
    }

    impl CssSelectHandler<MyDomNode> for SelectHandler {
        fn node_name(&self, node: &MyDomNode) -> CssQName {
            debug!("HL node_name!");
            debug!("SS %?", (*node.name).to_str());

            CssQName {
                ns: None,
                name: (*node.name).clone()
            }
        }

        fn node_id(&self, _node: &MyDomNode) -> Option<LwcString> { None }

        fn named_parent_node(&self, _node: &MyDomNode, _qname: &CssQName) -> Option<MyDomNode> {
            None
        }

        fn parent_node(&self, _node: &MyDomNode) -> Option<MyDomNode> {
            None
        }

        fn node_has_id(&self, _node: &MyDomNode, _name: LwcString) -> bool { false }

        fn named_ancestor_node(&self, _node: &MyDomNode, _qname: &CssQName) -> Option<MyDomNode> {
            None
        }

        fn node_is_root(&self, _node: &MyDomNode) -> bool { false }

        fn node_is_link(&self, _node: &MyDomNode) -> bool { false }

        fn ua_default_for_property(&self, property: CssProperty) -> CssHint {
            match property {
                _ => CssHintDefault
            }
        }
    }
}

#[test]
fn test_arc() {
    use std::arc::ARC;
    use stylesheet::*;
    use types::CssLevel21;
    use wapcaplet::LwcString;
    use super::CssResult;

    let resolve: CssUrlResolutionFn = |a,b| resolve_url(a, b);
    let params: CssStylesheetParams = CssStylesheetParams {
        params_version: CssStylesheetParamsVersion1,
        level: CssLevel21,
        charset: ~"UTF-8",
        url: ~"foo",
        title: ~"foo",
        allow_quirks: false,
        inline_style: false,
        resolve: Some(resolve),
        import: None,
        color: None,
        font: None,
    };

    let sheet: CssStylesheet = css_stylesheet_create(&params);
    let _arc = ARC(sheet);

    fn resolve_url(_base: &str, _rel: &LwcString) -> CssResult<LwcString> {
        fail!(~"resolving url");
    }
}
<|endoftext|>