<commit_before>use std::cell::{Ref, RefCell};
use std::cmp;
use std::io;
use std::net::{Ipv4Addr, SocketAddr};
use std::rc::{Rc, Weak};
use log::LogLevel;
use mio::{Event, PollOpt, Ready, Token};
use mio::net::TcpStream;

use super::binary;
use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::ipv4_header::IPv4Header;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packet_storage::PacketStorage;
use super::packetizer::Packetizer;
use super::selector::{EventHandler, Selector};
use super::stream_buffer::StreamBuffer;
use super::tcp_header::{self, TCPHeader};
use super::transport_header::{TransportHeader, TransportHeaderMut};

const TAG: &'static str = "TCPConnection";

const MAX_PAYLOAD_LENGTH: u16 = 1400;

pub struct TCPConnection {
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    stream: TcpStream,
    token: Token,
    client_to_network: StreamBuffer,
    network_to_client: Packetizer,
    pending_packet_for_client: PacketStorage,
    packet_for_client_length: Option<u16>,
    closed: bool,
    tcb: TCB,
}

struct TCB {
    state: TCPState,
    syn_sequence_number: u32,
    sequence_number: u32,
    acknowledgement_number: u32,
    their_acknowledgement_number: u32,
    client_window: u16,
    remote_closed: bool,
}

#[derive(PartialEq, Eq)]
enum TCPState {
    Init,
    SynSent,
//    SynReceived,
    Established,
    CloseWait,
    LastAck,
}

impl TCB {
    fn new() -> Self {
        Self {
            state: TCPState::Init,
            syn_sequence_number: 0,
            sequence_number: 0,
            acknowledgement_number: 0,
            their_acknowledgement_number: 0,
            client_window: 0,
            remote_closed: false,
        }
    }

    fn numbers(&self) -> String {
        format!("(seq={}, ack={})", self.sequence_number, self.acknowledgement_number)
    }
}

impl TCPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, ipv4_header: &IPv4Header, transport_header: &TransportHeader) -> io::Result<Rc<RefCell<Self>>> {
        let stream = Self::create_stream(&id)?;

        if let TransportHeader::TCP(ref tcp_header) = *transport_header {
            // shrink the TCP options to pass a minimal refrence header to the packetizer
            let mut shrinked_tcp_header_raw = [0u8; 20];
            shrinked_tcp_header_raw.copy_from_slice(&tcp_header.raw()[..20]);
            let mut shrinked_tcp_header_data = tcp_header.data().clone();
            {
                let mut shrinked_tcp_header = shrinked_tcp_header_data.bind_mut(&mut shrinked_tcp_header_raw);
                shrinked_tcp_header.shrink_options();
                assert_eq!(20, shrinked_tcp_header.header_length());
            }

            let shrinked_transport_header = shrinked_tcp_header_data.bind(&shrinked_tcp_header_raw).into();

            let packetizer = Packetizer::new(&ipv4_header, &shrinked_transport_header);

            let rc = Rc::new(RefCell::new(Self {
                id: id,
                client: client,
                stream: stream,
                token: Token(0), // default value, will be set afterwards
                client_to_network: StreamBuffer::new(4 * MAX_PACKET_LENGTH),
                network_to_client: packetizer,
                pending_packet_for_client: PacketStorage::new(),
                packet_for_client_length: None,
                closed: false,
                tcb: TCB::new(),
            }));

            {
                let mut self_ref = rc.borrow_mut();
                // rc is an EventHandler, register() expects a Box<EventHandler>
                let handler = Box::new(rc.clone());
                let token = selector.register(&self_ref.stream, handler, Ready::readable(), PollOpt::level())?;
                self_ref.token = token;
            }
            Ok(rc)
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_stream(id: &ConnectionId) -> io::Result<TcpStream> {
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        TcpStream::connect(&rewritten_destination)
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        self.disconnect(selector);

        // route is embedded in router which is embedded in client: the client necessarily exists
        let client_rc = self.client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.router().remove(&self.id);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        match self.client_to_network.write_to(&mut self.stream) {
            Ok(w) => if w == 0 {
                self.close(selector);
            },
            Err(err) => {
                error!(target: TAG, "{} Cannot write: {}", self.id, err);
                self.reset_connection(selector);
            },
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        assert!(self.packet_for_client_length.is_none(), "A pending packet was not sent");
        let remaining_client_window = self.get_remaining_client_window();
        assert!(remaining_client_window > 0, "process_received() must not be called when window == 0");
        let max_payload_length = cmp::min(remaining_client_window, MAX_PAYLOAD_LENGTH) as usize;
        Self::update_headers(&mut self.network_to_client, &self.tcb, tcp_header::FLAG_ACK | tcp_header::FLAG_PSH);
        // the packet is bound to the lifetime of self, so we cannot borrow self to call methods
        // defer the other branches in a separate match-block
        let non_lexical_lifetime_workaround = match self.network_to_client.packetize_read(&mut self.stream, Some(max_payload_length)) {
            Ok(Some(ipv4_packet)) => {
                // TODO packet source pull
                Self::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
                Ok(Some(()))
            },
            Ok(None) => Ok(None),
            Err(err) => Err(err),
        };
        match non_lexical_lifetime_workaround {
            Ok(None) => self.eof(selector),
            Err(err) => {
                error!(target: TAG, "{} Cannot read: {}", self.id, err);
                self.reset_connection(selector);
            },
            Ok(Some(_)) => (), // already handled
        }
    }

    fn send_to_client(client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        let client_rc = client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.send_to_client(selector, &ipv4_packet)
    }

    fn send_to_client_ignore(id: &ConnectionId, client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        if let Err(err) = Self::send_to_client(client, selector, ipv4_packet) {
            warn!(target: TAG, "{} Cannot send packet to client: {}", id, err);
        }
    }

    fn eof(&mut self, selector: &mut Selector) {
        self.tcb.remote_closed = true;
        if self.tcb.state == TCPState::CloseWait {
            let ipv4_packet = Self::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_FIN);
            self.tcb.sequence_number += 1; // FIN counts for 1 byte
            Self::send_to_client_ignore(&self. id, &self.client, selector, &ipv4_packet);
        }
    }

    fn update_headers(packetizer: &mut Packetizer, tcb: &TCB, flags: u16) {
        if let TransportHeaderMut::TCP(ref mut tcp_header) = packetizer.transport_header_mut() {
            tcp_header.set_sequence_number(tcb.sequence_number);
            tcp_header.set_acknowledgement_number(tcb.acknowledgement_number);
            tcp_header.set_flags(flags);
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_empty_response_packet<'a>(id: &ConnectionId, packetizer: &'a mut Packetizer, tcb: &TCB, flags: u16) -> IPv4Packet<'a> {
        Self::update_headers(packetizer, tcb, flags);
        debug!(target: TAG, "{} Forging empty response (flags={}) {}", id, flags, tcb.numbers());
        if (flags & tcp_header::FLAG_ACK) != 0 {
            debug!(target: TAG, "{} Acking {}", id, tcb.numbers());
        }
        let ipv4_packet = packetizer.packetize_empty_payload();
        if log_enabled!(target: TAG, LogLevel::Trace) {
            binary::to_string(ipv4_packet.raw());
        }
        ipv4_packet
    }

    fn reset_connection(&mut self, selector: &mut Selector) {
        {
            let ipv4_packet = Self::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_RST);
            Self::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
        }
        self.close(selector);
    }

    fn update_interests(&mut self, selector: &mut Selector) {
        if !self.closed {
            let mut ready = Ready::empty();
            if self.may_read() {
                ready = ready | Ready::readable()
            }
            if self.may_write() {
                ready = ready | Ready::writable()
            }
            selector.reregister(&self.stream, self.token, ready, PollOpt::level()).expect("Cannot register on poll");
        }
    }

    fn may_read(&self) -> bool {
        !self.tcb.remote_closed &&
                self.packet_for_client_length.is_some() &&
                self.get_remaining_client_window() > 0
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn get_remaining_client_window(&self) -> u16 {
        // TODO
        42
    }
}

impl Connection for TCPConnection {
    fn id(&self) -> &ConnectionId {
        &self.id
    }

    fn send_to_network(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        // TODO
    }

    fn disconnect(&mut self, selector: &mut Selector) {
        info!(target: TAG, "{} Close", self.id);
        selector.deregister(&self.stream, self.token).unwrap();
        // socket will be closed by RAII
    }

    fn is_expired(&self) -> bool {
        // no external timeout expiration
        false
    }
}

impl EventHandler for TCPConnection {
    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }
}
<commit_msg>Add self-weak to TCPConnection<commit_after>use std::cell::{Ref, RefCell};
use std::cmp;
use std::io;
use std::net::{Ipv4Addr, SocketAddr};
use std::rc::{Rc, Weak};
use log::LogLevel;
use mio::{Event, PollOpt, Ready, Token};
use mio::net::TcpStream;

use super::binary;
use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::ipv4_header::IPv4Header;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packet_storage::PacketStorage;
use super::packetizer::Packetizer;
use super::selector::{EventHandler, Selector};
use super::stream_buffer::StreamBuffer;
use super::tcp_header::{self, TCPHeader};
use super::transport_header::{TransportHeader, TransportHeaderMut};

const TAG: &'static str = "TCPConnection";

const MAX_PAYLOAD_LENGTH: u16 = 1400;

pub struct TCPConnection {
    self_weak: Weak<RefCell<TCPConnection>>,
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    stream: TcpStream,
    token: Token,
    client_to_network: StreamBuffer,
    network_to_client: Packetizer,
    pending_packet_for_client: PacketStorage,
    packet_for_client_length: Option<u16>,
    closed: bool,
    tcb: TCB,
}

struct TCB {
    state: TCPState,
    syn_sequence_number: u32,
    sequence_number: u32,
    acknowledgement_number: u32,
    their_acknowledgement_number: u32,
    client_window: u16,
    remote_closed: bool,
}

#[derive(PartialEq, Eq)]
enum TCPState {
    Init,
    SynSent,
//    SynReceived,
    Established,
    CloseWait,
    LastAck,
}

impl TCB {
    fn new() -> Self {
        Self {
            state: TCPState::Init,
            syn_sequence_number: 0,
            sequence_number: 0,
            acknowledgement_number: 0,
            their_acknowledgement_number: 0,
            client_window: 0,
            remote_closed: false,
        }
    }

    fn numbers(&self) -> String {
        format!("(seq={}, ack={})", self.sequence_number, self.acknowledgement_number)
    }
}

impl TCPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, ipv4_header: &IPv4Header, transport_header: &TransportHeader) -> io::Result<Rc<RefCell<Self>>> {
        let stream = Self::create_stream(&id)?;

        if let TransportHeader::TCP(ref tcp_header) = *transport_header {
            // shrink the TCP options to pass a minimal refrence header to the packetizer
            let mut shrinked_tcp_header_raw = [0u8; 20];
            shrinked_tcp_header_raw.copy_from_slice(&tcp_header.raw()[..20]);
            let mut shrinked_tcp_header_data = tcp_header.data().clone();
            {
                let mut shrinked_tcp_header = shrinked_tcp_header_data.bind_mut(&mut shrinked_tcp_header_raw);
                shrinked_tcp_header.shrink_options();
                assert_eq!(20, shrinked_tcp_header.header_length());
            }

            let shrinked_transport_header = shrinked_tcp_header_data.bind(&shrinked_tcp_header_raw).into();

            let packetizer = Packetizer::new(&ipv4_header, &shrinked_transport_header);

            let rc = Rc::new(RefCell::new(Self {
                self_weak: Weak::new(),
                id: id,
                client: client,
                stream: stream,
                token: Token(0), // default value, will be set afterwards
                client_to_network: StreamBuffer::new(4 * MAX_PACKET_LENGTH),
                network_to_client: packetizer,
                pending_packet_for_client: PacketStorage::new(),
                packet_for_client_length: None,
                closed: false,
                tcb: TCB::new(),
            }));

            {
                let mut self_ref = rc.borrow_mut();

                // keep a shared reference to this
                self_ref.self_weak = Rc::downgrade(&rc);

                // rc is an EventHandler, register() expects a Box<EventHandler>
                let handler = Box::new(rc.clone());
                let token = selector.register(&self_ref.stream, handler, Ready::readable(), PollOpt::level())?;
                self_ref.token = token;
            }
            Ok(rc)
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_stream(id: &ConnectionId) -> io::Result<TcpStream> {
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        TcpStream::connect(&rewritten_destination)
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        self.disconnect(selector);

        // route is embedded in router which is embedded in client: the client necessarily exists
        let client_rc = self.client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.router().remove(&self.id);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        match self.client_to_network.write_to(&mut self.stream) {
            Ok(w) => if w == 0 {
                self.close(selector);
            },
            Err(err) => {
                error!(target: TAG, "{} Cannot write: {}", self.id, err);
                self.reset_connection(selector);
            },
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        assert!(self.packet_for_client_length.is_none(), "A pending packet was not sent");
        let remaining_client_window = self.get_remaining_client_window();
        assert!(remaining_client_window > 0, "process_received() must not be called when window == 0");
        let max_payload_length = cmp::min(remaining_client_window, MAX_PAYLOAD_LENGTH) as usize;
        Self::update_headers(&mut self.network_to_client, &self.tcb, tcp_header::FLAG_ACK | tcp_header::FLAG_PSH);
        // the packet is bound to the lifetime of self, so we cannot borrow self to call methods
        // defer the other branches in a separate match-block
        let non_lexical_lifetime_workaround = match self.network_to_client.packetize_read(&mut self.stream, Some(max_payload_length)) {
            Ok(Some(ipv4_packet)) => {
                // TODO packet source pull
                Self::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
                Ok(Some(()))
            },
            Ok(None) => Ok(None),
            Err(err) => Err(err),
        };
        match non_lexical_lifetime_workaround {
            Ok(None) => self.eof(selector),
            Err(err) => {
                error!(target: TAG, "{} Cannot read: {}", self.id, err);
                self.reset_connection(selector);
            },
            Ok(Some(_)) => (), // already handled
        }
    }

    fn send_to_client(client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        let client_rc = client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.send_to_client(selector, &ipv4_packet)
    }

    fn send_to_client_ignore(id: &ConnectionId, client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        if let Err(err) = Self::send_to_client(client, selector, ipv4_packet) {
            warn!(target: TAG, "{} Cannot send packet to client: {}", id, err);
        }
    }

    fn eof(&mut self, selector: &mut Selector) {
        self.tcb.remote_closed = true;
        if self.tcb.state == TCPState::CloseWait {
            let ipv4_packet = Self::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_FIN);
            self.tcb.sequence_number += 1; // FIN counts for 1 byte
            Self::send_to_client_ignore(&self. id, &self.client, selector, &ipv4_packet);
        }
    }

    fn update_headers(packetizer: &mut Packetizer, tcb: &TCB, flags: u16) {
        if let TransportHeaderMut::TCP(ref mut tcp_header) = packetizer.transport_header_mut() {
            tcp_header.set_sequence_number(tcb.sequence_number);
            tcp_header.set_acknowledgement_number(tcb.acknowledgement_number);
            tcp_header.set_flags(flags);
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_empty_response_packet<'a>(id: &ConnectionId, packetizer: &'a mut Packetizer, tcb: &TCB, flags: u16) -> IPv4Packet<'a> {
        Self::update_headers(packetizer, tcb, flags);
        debug!(target: TAG, "{} Forging empty response (flags={}) {}", id, flags, tcb.numbers());
        if (flags & tcp_header::FLAG_ACK) != 0 {
            debug!(target: TAG, "{} Acking {}", id, tcb.numbers());
        }
        let ipv4_packet = packetizer.packetize_empty_payload();
        if log_enabled!(target: TAG, LogLevel::Trace) {
            binary::to_string(ipv4_packet.raw());
        }
        ipv4_packet
    }

    fn reset_connection(&mut self, selector: &mut Selector) {
        {
            let ipv4_packet = Self::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_RST);
            Self::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
        }
        self.close(selector);
    }

    fn update_interests(&mut self, selector: &mut Selector) {
        if !self.closed {
            let mut ready = Ready::empty();
            if self.may_read() {
                ready = ready | Ready::readable()
            }
            if self.may_write() {
                ready = ready | Ready::writable()
            }
            selector.reregister(&self.stream, self.token, ready, PollOpt::level()).expect("Cannot register on poll");
        }
    }

    fn may_read(&self) -> bool {
        !self.tcb.remote_closed &&
                self.packet_for_client_length.is_some() &&
                self.get_remaining_client_window() > 0
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn get_remaining_client_window(&self) -> u16 {
        // TODO
        42
    }
}

impl Connection for TCPConnection {
    fn id(&self) -> &ConnectionId {
        &self.id
    }

    fn send_to_network(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        // TODO
    }

    fn disconnect(&mut self, selector: &mut Selector) {
        info!(target: TAG, "{} Close", self.id);
        selector.deregister(&self.stream, self.token).unwrap();
        // socket will be closed by RAII
    }

    fn is_expired(&self) -> bool {
        // no external timeout expiration
        false
    }
}

impl EventHandler for TCPConnection {
    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }
}
<|endoftext|>


<commit_before>#![allow(clippy::option_if_let_else)]

use std::env;
use std::fs;
use std::path::Path;
use std::process::{Command, ExitStatus, Stdio};
use std::str;

#[cfg(all(feature = "backtrace", not(feature = "std")))]
compile_error! {
    "`backtrace` feature without `std` feature is not supported"
}

// This code exercises the surface area that we expect of the std Backtrace
// type. If the current toolchain is able to compile it, we go ahead and use
// backtrace in anyhow.
const PROBE: &str = r#"
    #![feature(backtrace)]
    #![allow(dead_code)]

    use std::backtrace::{Backtrace, BacktraceStatus};
    use std::error::Error;
    use std::fmt::{self, Display};

    #[derive(Debug)]
    struct E;

    impl Display for E {
        fn fmt(&self, _formatter: &mut fmt::Formatter) -> fmt::Result {
            unimplemented!()
        }
    }

    impl Error for E {
        fn backtrace(&self) -> Option<&Backtrace> {
            let backtrace = Backtrace::capture();
            match backtrace.status() {
                BacktraceStatus::Captured | BacktraceStatus::Disabled | _ => {}
            }
            unimplemented!()
        }
    }
"#;

fn main() {
    if cfg!(feature = "std") {
        match compile_probe() {
            Some(status) if status.success() => println!("cargo:rustc-cfg=backtrace"),
            _ => {}
        }
    }

    let rustc = match rustc_minor_version() {
        Some(rustc) => rustc,
        None => return,
    };

    if rustc < 51 {
        println!("cargo:rustc-cfg=anyhow_no_ptr_addr_of");
    }

    if rustc < 52 {
        println!("cargo:rustc-cfg=anyhow_no_fmt_arguments_as_str");
    }
}

fn compile_probe() -> Option<ExitStatus> {
    let rustc = env::var_os("RUSTC")?;
    let out_dir = env::var_os("OUT_DIR")?;
    let probefile = Path::new(&out_dir).join("probe.rs");
    fs::write(&probefile, PROBE).ok()?;

    // Make sure to pick up Cargo rustc configuration.
    let mut cmd = if let Some(wrapper) = env::var_os("RUSTC_WRAPPER") {
        let mut cmd = Command::new(wrapper);
        // The wrapper's first argument is supposed to be the path to rustc.
        cmd.arg(rustc);
        cmd
    } else {
        Command::new(rustc)
    };

    cmd.stderr(Stdio::null())
        .arg("--edition=2018")
        .arg("--crate-name=anyhow_build")
        .arg("--crate-type=lib")
        .arg("--emit=metadata")
        .arg("--out-dir")
        .arg(out_dir)
        .arg(probefile);

    if let Some(target) = env::var_os("TARGET") {
        cmd.arg("--target").arg(target);
    }

    // If Cargo wants to set RUSTFLAGS, use that.
    if let Ok(rustflags) = env::var("CARGO_ENCODED_RUSTFLAGS") {
        if !rustflags.is_empty() {
            for arg in rustflags.split('\x1f') {
                cmd.arg(arg);
            }
        }
    }

    cmd.status().ok()
}

fn rustc_minor_version() -> Option<u32> {
    let rustc = env::var_os("RUSTC")?;
    let output = Command::new(rustc).arg("--version").output().ok()?;
    let version = str::from_utf8(&output.stdout).ok()?;
    let mut pieces = version.split('.');
    if pieces.next() != Some("rustc 1") {
        return None;
    }
    pieces.next()?.parse().ok()
}
<commit_msg>Work around RA's RUSTC_WRAPPER poisoning all subsequent command-line builds<commit_after>#![allow(clippy::option_if_let_else)]

use std::env;
use std::fs;
use std::path::Path;
use std::process::{Command, ExitStatus, Stdio};
use std::str;

#[cfg(all(feature = "backtrace", not(feature = "std")))]
compile_error! {
    "`backtrace` feature without `std` feature is not supported"
}

// This code exercises the surface area that we expect of the std Backtrace
// type. If the current toolchain is able to compile it, we go ahead and use
// backtrace in anyhow.
const PROBE: &str = r#"
    #![feature(backtrace)]
    #![allow(dead_code)]

    use std::backtrace::{Backtrace, BacktraceStatus};
    use std::error::Error;
    use std::fmt::{self, Display};

    #[derive(Debug)]
    struct E;

    impl Display for E {
        fn fmt(&self, _formatter: &mut fmt::Formatter) -> fmt::Result {
            unimplemented!()
        }
    }

    impl Error for E {
        fn backtrace(&self) -> Option<&Backtrace> {
            let backtrace = Backtrace::capture();
            match backtrace.status() {
                BacktraceStatus::Captured | BacktraceStatus::Disabled | _ => {}
            }
            unimplemented!()
        }
    }
"#;

fn main() {
    if cfg!(feature = "std") {
        match compile_probe() {
            Some(status) if status.success() => println!("cargo:rustc-cfg=backtrace"),
            _ => {}
        }
    }

    let rustc = match rustc_minor_version() {
        Some(rustc) => rustc,
        None => return,
    };

    if rustc < 51 {
        println!("cargo:rustc-cfg=anyhow_no_ptr_addr_of");
    }

    if rustc < 52 {
        println!("cargo:rustc-cfg=anyhow_no_fmt_arguments_as_str");
    }
}

fn compile_probe() -> Option<ExitStatus> {
    println!("cargo:rerun-if-env-changed=RUSTC_WRAPPER");

    let rustc = env::var_os("RUSTC")?;
    let out_dir = env::var_os("OUT_DIR")?;
    let probefile = Path::new(&out_dir).join("probe.rs");
    fs::write(&probefile, PROBE).ok()?;

    // Make sure to pick up Cargo rustc configuration.
    let mut cmd = if let Some(wrapper) = env::var_os("RUSTC_WRAPPER") {
        let mut cmd = Command::new(wrapper);
        // The wrapper's first argument is supposed to be the path to rustc.
        cmd.arg(rustc);
        cmd
    } else {
        Command::new(rustc)
    };

    cmd.stderr(Stdio::null())
        .arg("--edition=2018")
        .arg("--crate-name=anyhow_build")
        .arg("--crate-type=lib")
        .arg("--emit=metadata")
        .arg("--out-dir")
        .arg(out_dir)
        .arg(probefile);

    if let Some(target) = env::var_os("TARGET") {
        cmd.arg("--target").arg(target);
    }

    // If Cargo wants to set RUSTFLAGS, use that.
    if let Ok(rustflags) = env::var("CARGO_ENCODED_RUSTFLAGS") {
        if !rustflags.is_empty() {
            for arg in rustflags.split('\x1f') {
                cmd.arg(arg);
            }
        }
    }

    cmd.status().ok()
}

fn rustc_minor_version() -> Option<u32> {
    let rustc = env::var_os("RUSTC")?;
    let output = Command::new(rustc).arg("--version").output().ok()?;
    let version = str::from_utf8(&output.stdout).ok()?;
    let mut pieces = version.split('.');
    if pieces.next() != Some("rustc 1") {
        return None;
    }
    pieces.next()?.parse().ok()
}
<|endoftext|>


<commit_before>use std::env;
use std::path::{Path};
use std::process::{Command};

const SEARCH_LINUX: &'static [&'static str] = &[
    "/usr/lib",
    "/usr/lib/llvm",
    "/usr/lib/llvm-3.8/lib",
    "/usr/lib/llvm-3.7/lib",
    "/usr/lib/llvm-3.6/lib",
    "/usr/lib/llvm-3.5/lib",
    "/usr/lib/llvm-3.4/lib",
    "/usr/lib64/llvm",
    "/usr/lib/x86_64-linux-gnu",
    "/usr/local/lib",
];

const SEARCH_OSX: &'static [&'static str] = &[
    "/usr/local/opt/llvm/lib",
    "/Library/Developer/CommandLineTools/usr/lib",
    "/Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/lib",
];

const SEARCH_WINDOWS: &'static [&'static str] = &[
    "C:\\Program Files\\LLVM\\bin",
    "C:\\Program Files\\LLVM\\lib",
];

fn error(prefix: &str) -> ! {
    panic!("{:?}, set the LIBCLANG_PATH environment variable to your system's libclang directory", prefix);
}

fn run(command: &str, arguments: &[&str]) -> Option<String> {
    Command::new(command).args(arguments).output().map(|o| {
        String::from_utf8_lossy(&o.stdout).into_owned()
    }).ok()
}

fn find_libclang() -> Option<(String, String)> {
    let search = if let Ok(directory) = env::var("LIBCLANG_PATH") {
        vec![directory]
    } else {
        run("llvm-config", &["--libdir"]).map(|d| vec![d]).unwrap_or_else(|| {
            if cfg!(any(target_os="freebsd", target_os="linux")) {
                SEARCH_LINUX
            } else if cfg!(target_os="osx") {
                SEARCH_OSX
            } else if cfg!(target_os="windows") {
                SEARCH_WINDOWS
            } else {
                error("unsupported operating system");
            }.into_iter().map(|s| s.to_string()).collect()
        })
    };

    let library = format!("{}clang{}", env::consts::DLL_PREFIX, env::consts::DLL_SUFFIX);
    let directory = search.into_iter().find(|d| Path::new(&d).join(&library).exists());

    if directory.is_none() && cfg!(target_os="linux") {
        if let Some(output) = run("llvm-config", &["--libdir"]) {
            output.lines().next().map(|l| (l.into(), library))
        } else {
            None
        }
    } else {
        directory.map(|d| (d, library))
    }
}

const LIBRARIES: &'static [&'static str] = &[
    "LLVMAnalysis",
    "LLVMBitReader",
    "LLVMCore",
    "LLVMLTO",
    "LLVMLinker",
    "LLVMMC",
    "LLVMMCParser",
    "LLVMObjCARCOpts",
    "LLVMObject",
    "LLVMOption",
    "LLVMScalarOpts",
    "LLVMSupport",
    "LLVMTarget",
    "LLVMTransformUtils",
    "LLVMVectorize",
    "LLVMipa",
    "LLVMipo",
    "clang",
    "clangARCMigrate",
    "clangAST",
    "clangASTMatchers",
    "clangAnalysis",
    "clangBasic",
    "clangDriver",
    "clangEdit",
    "clangFormat",
    "clangFrontend",
    "clangIndex",
    "clangLex",
    "clangParse",
    "clangRewrite",
    "clangRewriteFrontend",
    "clangSema",
    "clangSerialization",
    "clangStaticAnalyzerCheckers",
    "clangStaticAnalyzerCore",
    "clangStaticAnalyzerFrontend",
    "clangTooling",
];

fn get_libraries() -> Vec<String> {
    run("llvm-config", &["--libs"]).map(|o| {
        o.split_whitespace().filter_map(|p| {
            Path::new(p).file_stem().map(|l| l.to_string_lossy()[2..].into())
        }).collect()
    }).unwrap_or_else(|| LIBRARIES.iter().map(|l| (*l).into()).collect())
}

fn main() {
    let (directory, file) = match find_libclang() {
        Some((directory, file)) => (directory, file),
        _ => error("unable to find libclang"),
    };

    if cfg!(feature="static") {
        print!("cargo:rustc-flags=");
        if let Ok(directory) = env::var("LIBCLANG_STATIC_PATH") {
            print!("-L {} ", directory);
        }
        for library in get_libraries() {
            print!("-l static={} ", library)
        }
        println!("-L {} -l ncursesw -l z -l stdc++", directory);
    } else {
        println!("cargo:rustc-link-search={}", directory);
        println!("cargo:rustc-link-lib=dylib=:{}", file);
    }
}
<commit_msg>Improve find_libclang (#11)<commit_after>use std::env;
use std::path::{Path};
use std::process::{Command};

const SEARCH_LINUX: &'static [&'static str] = &[
    "/usr/lib",
    "/usr/lib/llvm",
    "/usr/lib/llvm-3.8/lib",
    "/usr/lib/llvm-3.7/lib",
    "/usr/lib/llvm-3.6/lib",
    "/usr/lib/llvm-3.5/lib",
    "/usr/lib/llvm-3.4/lib",
    "/usr/lib64/llvm",
    "/usr/lib/x86_64-linux-gnu",
    "/usr/local/lib",
];

const SEARCH_OSX: &'static [&'static str] = &[
    "/usr/local/opt/llvm/lib",
    "/Library/Developer/CommandLineTools/usr/lib",
    "/Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/lib",
];

const SEARCH_WINDOWS: &'static [&'static str] = &[
    "C:\\Program Files\\LLVM\\bin",
    "C:\\Program Files\\LLVM\\lib",
];

fn error(prefix: &str) -> ! {
    panic!("{:?}, set the LIBCLANG_PATH environment variable to your system's libclang directory", prefix);
}

fn run(command: &str, arguments: &[&str]) -> Option<String> {
    Command::new(command).args(arguments).output().map(|o| {
        String::from_utf8_lossy(&o.stdout).into_owned()
    }).ok()
}

fn find_libclang() -> Option<(String, String)> {
    let search = if let Ok(directory) = env::var("LIBCLANG_PATH") {
        vec![directory]
    } else if let Some(output) = run("llvm-config", &["--libdir"]) {
        vec![output.lines()
                   .map(|s| s.to_string())
                   .next()
                   .unwrap()]
    } else {
        if cfg!(any(target_os="freebsd", target_os="linux")) {
            SEARCH_LINUX
        } else if cfg!(target_os="osx") {
            SEARCH_OSX
        } else if cfg!(target_os="windows") {
            SEARCH_WINDOWS
        } else {
            error("unsupported operating system");
        }.into_iter().map(|s| s.to_string()).collect()
    };

    let library = format!("{}clang{}",
                          env::consts::DLL_PREFIX,
                          env::consts::DLL_SUFFIX);

    let directory = search.into_iter()
                          .find(|d| Path::new(&d)
                                         .join(&library)
                                         .exists());

    directory.map(|d| (d, library))
}

const LIBRARIES: &'static [&'static str] = &[
    "LLVMAnalysis",
    "LLVMBitReader",
    "LLVMCore",
    "LLVMLTO",
    "LLVMLinker",
    "LLVMMC",
    "LLVMMCParser",
    "LLVMObjCARCOpts",
    "LLVMObject",
    "LLVMOption",
    "LLVMScalarOpts",
    "LLVMSupport",
    "LLVMTarget",
    "LLVMTransformUtils",
    "LLVMVectorize",
    "LLVMipa",
    "LLVMipo",
    "clang",
    "clangARCMigrate",
    "clangAST",
    "clangASTMatchers",
    "clangAnalysis",
    "clangBasic",
    "clangDriver",
    "clangEdit",
    "clangFormat",
    "clangFrontend",
    "clangIndex",
    "clangLex",
    "clangParse",
    "clangRewrite",
    "clangRewriteFrontend",
    "clangSema",
    "clangSerialization",
    "clangStaticAnalyzerCheckers",
    "clangStaticAnalyzerCore",
    "clangStaticAnalyzerFrontend",
    "clangTooling",
];

fn get_libraries() -> Vec<String> {
    run("llvm-config", &["--libs"]).map(|o| {
        o.split_whitespace().filter_map(|p| {
            Path::new(p).file_stem().map(|l| l.to_string_lossy()[2..].into())
        }).collect()
    }).unwrap_or_else(|| LIBRARIES.iter().map(|l| (*l).into()).collect())
}

fn main() {
    let (directory, file) = match find_libclang() {
        Some((directory, file)) => (directory, file),
        _ => error("unable to find libclang"),
    };

    if cfg!(feature="static") {
        print!("cargo:rustc-flags=");
        if let Ok(directory) = env::var("LIBCLANG_STATIC_PATH") {
            print!("-L {} ", directory);
        }
        for library in get_libraries() {
            print!("-l static={} ", library)
        }
        println!("-L {} -l ncursesw -l z -l stdc++", directory);
    } else {
        println!("cargo:rustc-link-search={}", directory);
        println!("cargo:rustc-link-lib=dylib=:{}", file);
    }
}
<|endoftext|>


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

use kite::{Document, Term, TermRef};
use kite::schema::FieldRef;
use byteorder::{BigEndian, WriteBytesExt};
use roaring::RoaringBitmap;

use key_builder::KeyBuilder;

#[derive(Debug)]
pub struct SegmentBuilder {
    current_doc: u16,
    pub term_dictionary: HashMap<Term, TermRef>,
    current_term_ref: u32,
    pub term_directories: HashMap<(FieldRef, TermRef), RoaringBitmap>,
    pub statistics: HashMap<Vec<u8>, i64>,
    pub stored_field_values: HashMap<(FieldRef, u16, Vec<u8>), Vec<u8>>,
}

#[derive(Debug)]
pub enum DocumentInsertError {
    /// Segment couldn't hold any more docs
    SegmentFull,
}

impl SegmentBuilder {
    pub fn new() -> SegmentBuilder {
        SegmentBuilder {
            current_doc: 0,
            term_dictionary: HashMap::new(),
            current_term_ref: 0,
            term_directories: HashMap::new(),
            statistics: HashMap::new(),
            stored_field_values: HashMap::new(),
        }
    }

    fn get_term_ref(&mut self, term: &Term) -> TermRef {
        if let Some(term_ref) = self.term_dictionary.get(term) {
            return *term_ref;
        }

        // Add the term to the dictionary
        let term_ref = TermRef::new(self.current_term_ref);
        self.current_term_ref += 1;
        self.term_dictionary.insert(term.clone(), term_ref);

        term_ref
    }

    // TODO: Need to translate field names to field refs and terms to term refs
    pub fn add_document(&mut self, doc: &Document) -> Result<u16, DocumentInsertError> {
        // Get document ord
        let doc_id = self.current_doc;
        self.current_doc += 1;
        try!(self.current_doc.checked_add(1).ok_or(DocumentInsertError::SegmentFull));

        // Insert indexed fields
        let mut term_frequencies = HashMap::new();
        for (field, tokens) in doc.indexed_fields.iter() {
            let mut field_token_count = 0;

            for token in tokens.iter() {
                field_token_count += 1;

                // Get term ref
                let term_ref = self.get_term_ref(&token.term);

                // Term frequency
                let mut term_frequency = term_frequencies.entry(term_ref).or_insert(0);
                *term_frequency += 1;

                // Write directory list
                self.term_directories.entry((*field, term_ref)).or_insert_with(RoaringBitmap::new).insert(doc_id as u32);
            }

            // Term frequencies
            for (term_ref, frequency) in term_frequencies.drain() {
                // Write term frequency
                // 1 is by far the most common frequency. At search time, we interpret a missing
                // key as meaning there is a term frequency of 1
                if frequency != 1 {
                    let mut value_type = vec![b't', b'f'];
                    value_type.extend(term_ref.ord().to_string().as_bytes());

                    let mut frequency_bytes: Vec<u8> = Vec::new();
                    frequency_bytes.write_i64::<BigEndian>(frequency).unwrap();

                    self.stored_field_values.insert((*field, doc_id, value_type), frequency_bytes);
                }

                // Increment term document frequency
                let stat_name = KeyBuilder::segment_stat_term_doc_frequency_stat_name(field.ord(), term_ref.ord());
                let mut stat = self.statistics.entry(stat_name).or_insert(0);
                *stat += 1;
            }

            // Field length
            // Used by the BM25 similarity model
            let length = ((field_token_count as f64).sqrt() - 1.0) * 3.0;
            let length = if length > 255.0 { 255.0 } else { length } as u8;
            if length != 0 {
                self.stored_field_values.insert((*field, doc_id, b"len".to_vec()), vec![length]);
            }

            // Increment total field docs
            {
                let stat_name = KeyBuilder::segment_stat_total_field_docs_stat_name(field.ord());
                let mut stat = self.statistics.entry(stat_name).or_insert(0);
                *stat += 1;
            }

            // Increment total field tokens
            {
                let stat_name = KeyBuilder::segment_stat_total_field_tokens_stat_name(field.ord());
                let mut stat = self.statistics.entry(stat_name).or_insert(0);
                *stat += field_token_count;
            }
        }

        // Insert stored fields
        for (field, value) in doc.stored_fields.iter() {
            self.stored_field_values.insert((*field, doc_id, b"val".to_vec()), value.to_bytes());
        }

        // Increment total docs
        {
            let mut stat = self.statistics.entry(b"total_docs".to_vec()).or_insert(0);
            *stat += 1;
        }

        Ok(doc_id)
    }
}
<commit_msg>Implement Segment for SegmentBuilder<commit_after>use std::collections::HashMap;

use kite::{Document, Term, TermRef};
use kite::schema::FieldRef;
use kite::segment::Segment;
use byteorder::{BigEndian, WriteBytesExt};
use roaring::RoaringBitmap;

use key_builder::KeyBuilder;

#[derive(Debug)]
pub struct SegmentBuilder {
    current_doc: u16,
    pub term_dictionary: HashMap<Term, TermRef>,
    current_term_ref: u32,
    pub term_directories: HashMap<(FieldRef, TermRef), RoaringBitmap>,
    pub statistics: HashMap<Vec<u8>, i64>,
    pub stored_field_values: HashMap<(FieldRef, u16, Vec<u8>), Vec<u8>>,
}

#[derive(Debug)]
pub enum DocumentInsertError {
    /// Segment couldn't hold any more docs
    SegmentFull,
}

impl SegmentBuilder {
    pub fn new() -> SegmentBuilder {
        SegmentBuilder {
            current_doc: 0,
            term_dictionary: HashMap::new(),
            current_term_ref: 0,
            term_directories: HashMap::new(),
            statistics: HashMap::new(),
            stored_field_values: HashMap::new(),
        }
    }

    fn get_term_ref(&mut self, term: &Term) -> TermRef {
        if let Some(term_ref) = self.term_dictionary.get(term) {
            return *term_ref;
        }

        // Add the term to the dictionary
        let term_ref = TermRef::new(self.current_term_ref);
        self.current_term_ref += 1;
        self.term_dictionary.insert(term.clone(), term_ref);

        term_ref
    }

    // TODO: Need to translate field names to field refs and terms to term refs
    pub fn add_document(&mut self, doc: &Document) -> Result<u16, DocumentInsertError> {
        // Get document ord
        let doc_id = self.current_doc;
        self.current_doc += 1;
        try!(self.current_doc.checked_add(1).ok_or(DocumentInsertError::SegmentFull));

        // Insert indexed fields
        let mut term_frequencies = HashMap::new();
        for (field, tokens) in doc.indexed_fields.iter() {
            let mut field_token_count = 0;

            for token in tokens.iter() {
                field_token_count += 1;

                // Get term ref
                let term_ref = self.get_term_ref(&token.term);

                // Term frequency
                let mut term_frequency = term_frequencies.entry(term_ref).or_insert(0);
                *term_frequency += 1;

                // Write directory list
                self.term_directories.entry((*field, term_ref)).or_insert_with(RoaringBitmap::new).insert(doc_id as u32);
            }

            // Term frequencies
            for (term_ref, frequency) in term_frequencies.drain() {
                // Write term frequency
                // 1 is by far the most common frequency. At search time, we interpret a missing
                // key as meaning there is a term frequency of 1
                if frequency != 1 {
                    let mut value_type = vec![b't', b'f'];
                    value_type.extend(term_ref.ord().to_string().as_bytes());

                    let mut frequency_bytes: Vec<u8> = Vec::new();
                    frequency_bytes.write_i64::<BigEndian>(frequency).unwrap();

                    self.stored_field_values.insert((*field, doc_id, value_type), frequency_bytes);
                }

                // Increment term document frequency
                let stat_name = KeyBuilder::segment_stat_term_doc_frequency_stat_name(field.ord(), term_ref.ord());
                let mut stat = self.statistics.entry(stat_name).or_insert(0);
                *stat += 1;
            }

            // Field length
            // Used by the BM25 similarity model
            let length = ((field_token_count as f64).sqrt() - 1.0) * 3.0;
            let length = if length > 255.0 { 255.0 } else { length } as u8;
            if length != 0 {
                self.stored_field_values.insert((*field, doc_id, b"len".to_vec()), vec![length]);
            }

            // Increment total field docs
            {
                let stat_name = KeyBuilder::segment_stat_total_field_docs_stat_name(field.ord());
                let mut stat = self.statistics.entry(stat_name).or_insert(0);
                *stat += 1;
            }

            // Increment total field tokens
            {
                let stat_name = KeyBuilder::segment_stat_total_field_tokens_stat_name(field.ord());
                let mut stat = self.statistics.entry(stat_name).or_insert(0);
                *stat += field_token_count;
            }
        }

        // Insert stored fields
        for (field, value) in doc.stored_fields.iter() {
            self.stored_field_values.insert((*field, doc_id, b"val".to_vec()), value.to_bytes());
        }

        // Increment total docs
        {
            let mut stat = self.statistics.entry(b"total_docs".to_vec()).or_insert(0);
            *stat += 1;
        }

        Ok(doc_id)
    }
}

impl Segment for SegmentBuilder {
    fn id(&self) -> u32 {
        0
    }

    fn load_statistic(&self, stat_name: &[u8]) -> Result<Option<i64>, String> {
        Ok(self.statistics.get(stat_name).cloned())
    }

    fn load_stored_field_value_raw(&self, doc_ord: u16, field_ref: FieldRef, value_type: &[u8]) -> Result<Option<Vec<u8>>, String> {
        Ok(self.stored_field_values.get(&(field_ref, doc_ord, value_type.to_vec())).cloned())
    }

    fn load_term_directory(&self, field_ref: FieldRef, term_ref: TermRef) -> Result<Option<RoaringBitmap>, String> {
        Ok(self.term_directories.get(&(field_ref, term_ref)).cloned())
    }

    fn load_deletion_list(&self) -> Result<Option<RoaringBitmap>, String> {
        Ok(None)
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::{
    connections::ConnectionInterface,
    handle_fields::{build_handle_field_directive, HandleFieldDirectiveValues},
};
use common::{Diagnostic, DiagnosticsResult, NamedItem};
use graphql_ir::{
    Field, FragmentDefinition, LinkedField, Program, ScalarField, Selection, Transformed,
    Transformer,
};
use intern::string_key::{Intern, StringKey};
use lazy_static::lazy_static;
use schema::{Schema, Type};
use std::sync::Arc;
use thiserror::Error;

pub fn transform_declarative_connection(
    program: &Program,
    connection_interface: &ConnectionInterface,
) -> DiagnosticsResult<Program> {
    let mut transform = DeclarativeConnectionMutationTransform::new(program, connection_interface);
    let next_program = transform
        .transform_program(program)
        .replace_or_else(|| program.clone());

    if transform.errors.is_empty() {
        Ok(next_program)
    } else {
        Err(transform.errors)
    }
}

lazy_static! {
    static ref APPEND_EDGE: StringKey = "appendEdge".intern();
    static ref APPEND_NODE: StringKey = "appendNode".intern();
    static ref CONNECTIONS_ARG_NAME: StringKey = "connections".intern();
    static ref DELETE_RECORD: StringKey = "deleteRecord".intern();
    static ref DELETE_EDGE: StringKey = "deleteEdge".intern();
    static ref PREPEND_EDGE: StringKey = "prependEdge".intern();
    static ref PREPEND_NODE: StringKey = "prependNode".intern();
    static ref EDGE_TYPENAME_ARG: StringKey = "edgeTypeName".intern();
}

struct DeclarativeConnectionMutationTransform<'s, 'c> {
    program: &'s Program,
    errors: Vec<Diagnostic>,
    connection_interface: &'c ConnectionInterface,
}

impl<'s, 'c> DeclarativeConnectionMutationTransform<'s, 'c> {
    fn new(program: &'s Program, connection_interface: &'c ConnectionInterface) -> Self {
        Self {
            program,
            connection_interface,
            errors: vec![],
        }
    }
}

impl<'s, 'c> Transformer for DeclarativeConnectionMutationTransform<'s, 'c> {
    const NAME: &'static str = "DeclarativeConnectionMutationTransform";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_fragment(&mut self, _: &FragmentDefinition) -> Transformed<FragmentDefinition> {
        Transformed::Keep
    }

    fn transform_scalar_field(&mut self, field: &ScalarField) -> Transformed<Selection> {
        let linked_field_directive = field.directives.iter().find(|directive| {
            directive.name.item == *APPEND_EDGE
                || directive.name.item == *PREPEND_EDGE
                || directive.name.item == *APPEND_NODE
                || directive.name.item == *PREPEND_NODE
        });
        if let Some(linked_field_directive) = linked_field_directive {
            self.errors.push(Diagnostic::error(
                ValidationMessage::ConnectionMutationDirectiveOnScalarField {
                    directive_name: linked_field_directive.name.item,
                    field_name: field.alias_or_name(&self.program.schema),
                },
                field.definition.location,
            ));
        }
        let delete_directive = field.directives.iter().find(|directive| {
            directive.name.item == *DELETE_RECORD || directive.name.item == *DELETE_EDGE
        });
        let field_definition = self.program.schema.field(field.definition.item);
        match delete_directive {
            None => Transformed::Keep,
            Some(delete_directive) => {
                let is_id = self.program.schema.is_id(field_definition.type_.inner());
                if !is_id {
                    self.errors.push(Diagnostic::error(
                        ValidationMessage::DeleteRecordDirectiveOnUnsupportedType {
                            directive_name: delete_directive.name.item,
                            field_name: field.alias_or_name(&self.program.schema),
                            current_type: self
                                .program
                                .schema
                                .get_type_string(&field_definition.type_),
                        },
                        delete_directive.name.location,
                    ));
                    Transformed::Keep
                } else {
                    let connections_arg = delete_directive.arguments.named(*CONNECTIONS_ARG_NAME);
                    let handle_directive =
                        build_handle_field_directive(HandleFieldDirectiveValues {
                            handle: delete_directive.name.item,
                            key: "".intern(),
                            dynamic_key: None,
                            filters: None,
                            handle_args: connections_arg
                                .map(|connections_arg| vec![connections_arg.clone()]),
                        });
                    let mut next_directives: Vec<_> = field
                        .directives
                        .iter()
                        .filter(|directive| directive.name.item != *DELETE_RECORD)
                        .cloned()
                        .collect();
                    next_directives.push(handle_directive);
                    Transformed::Replace(Selection::ScalarField(Arc::new(ScalarField {
                        directives: next_directives,
                        ..field.clone()
                    })))
                }
            }
        }
    }

    fn transform_linked_field(&mut self, field: &LinkedField) -> Transformed<Selection> {
        let transformed_field = self.default_transform_linked_field(field);
        let delete_directive = field.directives.named(*DELETE_RECORD);
        if let Some(delete_directive) = delete_directive {
            self.errors.push(Diagnostic::error(
                ValidationMessage::DeleteRecordDirectiveOnLinkedField {
                    directive_name: delete_directive.name.item,
                    field_name: field.alias_or_name(&self.program.schema),
                },
                field.definition.location,
            ));
        }
        let edge_directive = field.directives.iter().find(|directive| {
            directive.name.item == *APPEND_EDGE || directive.name.item == *PREPEND_EDGE
        });
        let node_directive = field.directives.iter().find(|directive| {
            directive.name.item == *APPEND_NODE || directive.name.item == *PREPEND_NODE
        });
        match (edge_directive, node_directive) {
            (Some(edge_directive), Some(node_directive)) => {
                self.errors.push(Diagnostic::error(
                    ValidationMessage::ConflictingEdgeAndNodeDirectives {
                        edge_directive_name: edge_directive.name.item,
                        node_directive_name: node_directive.name.item,
                        field_name: field.alias_or_name(&self.program.schema),
                    },
                    edge_directive.name.location,
                ));
                transformed_field
            }
            (None, None) => transformed_field,
            (Some(edge_directive), None) => {
                let connections_arg = edge_directive.arguments.named(*CONNECTIONS_ARG_NAME);
                match connections_arg {
                    None => {
                        self.errors.push(Diagnostic::error(
                            ValidationMessage::ConnectionsArgumentRequired {
                                directive_name: edge_directive.name.item,
                            },
                            edge_directive.name.location,
                        ));
                        transformed_field
                    }
                    Some(connections_arg) => {
                        let field_definition = self.program.schema.field(field.definition.item);
                        let mut has_cursor_field = false;
                        let mut has_node_field = false;
                        if let Type::Object(id) = field_definition.type_.inner() {
                            let object = self.program.schema.object(id);
                            for field_id in &object.fields {
                                let current_field_name =
                                    self.program.schema.field(*field_id).name.item;
                                if current_field_name == self.connection_interface.cursor {
                                    has_cursor_field = true;
                                } else if current_field_name == self.connection_interface.node {
                                    has_node_field = true;
                                }
                            }
                        }
                        if has_cursor_field && has_node_field {
                            let handle_directive =
                                build_handle_field_directive(HandleFieldDirectiveValues {
                                    handle: edge_directive.name.item,
                                    key: "".intern(),
                                    dynamic_key: None,
                                    filters: None,
                                    handle_args: Some(vec![connections_arg.clone()]),
                                });
                            let mut next_field = match transformed_field {
                                Transformed::Replace(Selection::LinkedField(linked_field)) => {
                                    (*linked_field).clone()
                                }
                                Transformed::Keep => field.clone(),
                                _ => panic!(
                                    "DeclarativeConnection got unexpected transform result: `{:?}`.",
                                    transformed_field
                                ),
                            };
                            let index = next_field
                                .directives
                                .iter()
                                .position(|directive| {
                                    directive.name.item == edge_directive.name.item
                                })
                                .expect("Expected the edge directive to exist.");
                            next_field.directives[index] = handle_directive;
                            Transformed::Replace(Selection::LinkedField(Arc::new(next_field)))
                        } else {
                            self.errors.push(Diagnostic::error(
                                ValidationMessage::EdgeDirectiveOnUnsupportedType {
                                    directive_name: edge_directive.name.item,
                                    field_name: field.alias_or_name(&self.program.schema),
                                },
                                edge_directive.name.location,
                            ));
                            Transformed::Keep
                        }
                    }
                }
            }
            (None, Some(node_directive)) => {
                let connections_arg = node_directive.arguments.named(*CONNECTIONS_ARG_NAME);
                match connections_arg {
                    None => {
                        self.errors.push(Diagnostic::error(
                            ValidationMessage::ConnectionsArgumentRequired {
                                directive_name: node_directive.name.item,
                            },
                            node_directive.name.location,
                        ));
                        transformed_field
                    }
                    Some(connections_arg) => {
                        let edge_typename_arg = node_directive.arguments.named(*EDGE_TYPENAME_ARG);
                        if let Some(edge_typename_arg) = edge_typename_arg {
                            let field_definition = self.program.schema.field(field.definition.item);
                            match field_definition.type_.inner() {
                                Type::Object(_) | Type::Interface(_) | Type::Union(_) => {
                                    let handle_directive =
                                        build_handle_field_directive(HandleFieldDirectiveValues {
                                            handle: node_directive.name.item,
                                            key: "".intern(),
                                            dynamic_key: None,
                                            filters: None,
                                            handle_args: Some(vec![
                                                connections_arg.clone(),
                                                edge_typename_arg.clone(),
                                            ]),
                                        });
                                    let mut next_field = match transformed_field {
                                        Transformed::Replace(Selection::LinkedField(
                                            linked_field,
                                        )) => (*linked_field).clone(),
                                        Transformed::Keep => field.clone(),
                                        _ => panic!(
                                            "DeclarativeConnection got unexpected transform result: `{:?}`.",
                                            transformed_field
                                        ),
                                    };
                                    let index = next_field
                                        .directives
                                        .iter()
                                        .position(|directive| {
                                            directive.name.item == node_directive.name.item
                                        })
                                        .expect("Expected the edge directive to exist.");
                                    next_field.directives[index] = handle_directive;
                                    Transformed::Replace(Selection::LinkedField(Arc::new(
                                        next_field,
                                    )))
                                }
                                _ => {
                                    self.errors.push(Diagnostic::error(
                                        ValidationMessage::NodeDirectiveOnUnsupportedType {
                                            directive_name: node_directive.name.item,
                                            field_name: field.alias_or_name(&self.program.schema),
                                            current_type: self
                                                .program
                                                .schema
                                                .get_type_string(&field_definition.type_),
                                        },
                                        node_directive.name.location,
                                    ));
                                    Transformed::Keep
                                }
                            }
                        } else {
                            self.errors.push(Diagnostic::error(
                                ValidationMessage::NodeDirectiveMissesRequiredEdgeTypeName {
                                    directive_name: node_directive.name.item,
                                    field_name: field.alias_or_name(&self.program.schema),
                                },
                                node_directive.name.location,
                            ));
                            Transformed::Keep
                        }
                    }
                }
            }
        }
    }
}

#[derive(Debug, Error)]
enum ValidationMessage {
    #[error(
        "Unsupported use of @{directive_name} on field '${field_name}', 'edgeTypeName' argument must be provided."
    )]
    NodeDirectiveMissesRequiredEdgeTypeName {
        directive_name: StringKey,
        field_name: StringKey,
    },

    #[error("Invalid use of @{directive_name} on scalar field '{field_name}'.")]
    ConnectionMutationDirectiveOnScalarField {
        directive_name: StringKey,
        field_name: StringKey,
    },

    #[error(
        "Invalid use of @{directive_name} on field '{field_name}'. Expected field type 'ID', got '{current_type}'."
    )]
    DeleteRecordDirectiveOnUnsupportedType {
        directive_name: StringKey,
        field_name: StringKey,
        current_type: String,
    },

    #[error("Invalid use of @{directive_name} on linked field '{field_name}'.")]
    DeleteRecordDirectiveOnLinkedField {
        directive_name: StringKey,
        field_name: StringKey,
    },

    #[error(
        "Invalid use of @{edge_directive_name} and @{node_directive_name} on field '{field_name}' - these directives cannot be used together."
    )]
    ConflictingEdgeAndNodeDirectives {
        edge_directive_name: StringKey,
        node_directive_name: StringKey,
        field_name: StringKey,
    },

    #[error("Expected the 'connections' argument to be defined on @{directive_name}.")]
    ConnectionsArgumentRequired { directive_name: StringKey },

    #[error(
        "Unsupported use of @{directive_name} on field '{field_name}', expected an edge field (a field with 'cursor' and 'node' selection)."
    )]
    EdgeDirectiveOnUnsupportedType {
        directive_name: StringKey,
        field_name: StringKey,
    },

    #[error(
        "Unsupported use of @{directive_name} on field '{field_name}'. Expected an object, union or interface, but got '{current_type}'."
    )]
    NodeDirectiveOnUnsupportedType {
        directive_name: StringKey,
        field_name: StringKey,
        current_type: String,
    },
}
<commit_msg>declarative_connection: simplify lifetimes<commit_after>/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::{
    connections::ConnectionInterface,
    handle_fields::{build_handle_field_directive, HandleFieldDirectiveValues},
};
use common::{Diagnostic, DiagnosticsResult, NamedItem};
use graphql_ir::{
    Field, FragmentDefinition, LinkedField, Program, ScalarField, Selection, Transformed,
    Transformer,
};
use intern::string_key::{Intern, StringKey};
use lazy_static::lazy_static;
use schema::{SDLSchema, Schema, Type};
use std::sync::Arc;
use thiserror::Error;

pub fn transform_declarative_connection(
    program: &Program,
    connection_interface: &ConnectionInterface,
) -> DiagnosticsResult<Program> {
    let mut transform = DeclarativeConnectionMutationTransform::new(program, connection_interface);
    let next_program = transform
        .transform_program(program)
        .replace_or_else(|| program.clone());

    if transform.errors.is_empty() {
        Ok(next_program)
    } else {
        Err(transform.errors)
    }
}

lazy_static! {
    static ref APPEND_EDGE: StringKey = "appendEdge".intern();
    static ref APPEND_NODE: StringKey = "appendNode".intern();
    static ref CONNECTIONS_ARG_NAME: StringKey = "connections".intern();
    static ref DELETE_RECORD: StringKey = "deleteRecord".intern();
    static ref DELETE_EDGE: StringKey = "deleteEdge".intern();
    static ref PREPEND_EDGE: StringKey = "prependEdge".intern();
    static ref PREPEND_NODE: StringKey = "prependNode".intern();
    static ref EDGE_TYPENAME_ARG: StringKey = "edgeTypeName".intern();
}

struct DeclarativeConnectionMutationTransform<'a> {
    schema: &'a SDLSchema,
    errors: Vec<Diagnostic>,
    connection_interface: &'a ConnectionInterface,
}

impl<'a> DeclarativeConnectionMutationTransform<'a> {
    fn new(program: &'a Program, connection_interface: &'a ConnectionInterface) -> Self {
        Self {
            schema: &program.schema,
            connection_interface,
            errors: vec![],
        }
    }
}

impl Transformer for DeclarativeConnectionMutationTransform<'_> {
    const NAME: &'static str = "DeclarativeConnectionMutationTransform";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_fragment(&mut self, _: &FragmentDefinition) -> Transformed<FragmentDefinition> {
        Transformed::Keep
    }

    fn transform_scalar_field(&mut self, field: &ScalarField) -> Transformed<Selection> {
        let linked_field_directive = field.directives.iter().find(|directive| {
            directive.name.item == *APPEND_EDGE
                || directive.name.item == *PREPEND_EDGE
                || directive.name.item == *APPEND_NODE
                || directive.name.item == *PREPEND_NODE
        });
        if let Some(linked_field_directive) = linked_field_directive {
            self.errors.push(Diagnostic::error(
                ValidationMessage::ConnectionMutationDirectiveOnScalarField {
                    directive_name: linked_field_directive.name.item,
                    field_name: field.alias_or_name(self.schema),
                },
                field.definition.location,
            ));
        }
        let delete_directive = field.directives.iter().find(|directive| {
            directive.name.item == *DELETE_RECORD || directive.name.item == *DELETE_EDGE
        });
        let field_definition = self.schema.field(field.definition.item);
        match delete_directive {
            None => Transformed::Keep,
            Some(delete_directive) => {
                let is_id = self.schema.is_id(field_definition.type_.inner());
                if !is_id {
                    self.errors.push(Diagnostic::error(
                        ValidationMessage::DeleteRecordDirectiveOnUnsupportedType {
                            directive_name: delete_directive.name.item,
                            field_name: field.alias_or_name(self.schema),
                            current_type: self.schema.get_type_string(&field_definition.type_),
                        },
                        delete_directive.name.location,
                    ));
                    Transformed::Keep
                } else {
                    let connections_arg = delete_directive.arguments.named(*CONNECTIONS_ARG_NAME);
                    let handle_directive =
                        build_handle_field_directive(HandleFieldDirectiveValues {
                            handle: delete_directive.name.item,
                            key: "".intern(),
                            dynamic_key: None,
                            filters: None,
                            handle_args: connections_arg
                                .map(|connections_arg| vec![connections_arg.clone()]),
                        });
                    let mut next_directives: Vec<_> = field
                        .directives
                        .iter()
                        .filter(|directive| directive.name.item != *DELETE_RECORD)
                        .cloned()
                        .collect();
                    next_directives.push(handle_directive);
                    Transformed::Replace(Selection::ScalarField(Arc::new(ScalarField {
                        directives: next_directives,
                        ..field.clone()
                    })))
                }
            }
        }
    }

    fn transform_linked_field(&mut self, field: &LinkedField) -> Transformed<Selection> {
        let transformed_field = self.default_transform_linked_field(field);
        let delete_directive = field.directives.named(*DELETE_RECORD);
        if let Some(delete_directive) = delete_directive {
            self.errors.push(Diagnostic::error(
                ValidationMessage::DeleteRecordDirectiveOnLinkedField {
                    directive_name: delete_directive.name.item,
                    field_name: field.alias_or_name(self.schema),
                },
                field.definition.location,
            ));
        }
        let edge_directive = field.directives.iter().find(|directive| {
            directive.name.item == *APPEND_EDGE || directive.name.item == *PREPEND_EDGE
        });
        let node_directive = field.directives.iter().find(|directive| {
            directive.name.item == *APPEND_NODE || directive.name.item == *PREPEND_NODE
        });
        match (edge_directive, node_directive) {
            (Some(edge_directive), Some(node_directive)) => {
                self.errors.push(Diagnostic::error(
                    ValidationMessage::ConflictingEdgeAndNodeDirectives {
                        edge_directive_name: edge_directive.name.item,
                        node_directive_name: node_directive.name.item,
                        field_name: field.alias_or_name(self.schema),
                    },
                    edge_directive.name.location,
                ));
                transformed_field
            }
            (None, None) => transformed_field,
            (Some(edge_directive), None) => {
                let connections_arg = edge_directive.arguments.named(*CONNECTIONS_ARG_NAME);
                match connections_arg {
                    None => {
                        self.errors.push(Diagnostic::error(
                            ValidationMessage::ConnectionsArgumentRequired {
                                directive_name: edge_directive.name.item,
                            },
                            edge_directive.name.location,
                        ));
                        transformed_field
                    }
                    Some(connections_arg) => {
                        let field_definition = self.schema.field(field.definition.item);
                        let mut has_cursor_field = false;
                        let mut has_node_field = false;
                        if let Type::Object(id) = field_definition.type_.inner() {
                            let object = self.schema.object(id);
                            for field_id in &object.fields {
                                let current_field_name = self.schema.field(*field_id).name.item;
                                if current_field_name == self.connection_interface.cursor {
                                    has_cursor_field = true;
                                } else if current_field_name == self.connection_interface.node {
                                    has_node_field = true;
                                }
                            }
                        }
                        if has_cursor_field && has_node_field {
                            let handle_directive =
                                build_handle_field_directive(HandleFieldDirectiveValues {
                                    handle: edge_directive.name.item,
                                    key: "".intern(),
                                    dynamic_key: None,
                                    filters: None,
                                    handle_args: Some(vec![connections_arg.clone()]),
                                });
                            let mut next_field = match transformed_field {
                                Transformed::Replace(Selection::LinkedField(linked_field)) => {
                                    (*linked_field).clone()
                                }
                                Transformed::Keep => field.clone(),
                                _ => panic!(
                                    "DeclarativeConnection got unexpected transform result: `{:?}`.",
                                    transformed_field
                                ),
                            };
                            let index = next_field
                                .directives
                                .iter()
                                .position(|directive| {
                                    directive.name.item == edge_directive.name.item
                                })
                                .expect("Expected the edge directive to exist.");
                            next_field.directives[index] = handle_directive;
                            Transformed::Replace(Selection::LinkedField(Arc::new(next_field)))
                        } else {
                            self.errors.push(Diagnostic::error(
                                ValidationMessage::EdgeDirectiveOnUnsupportedType {
                                    directive_name: edge_directive.name.item,
                                    field_name: field.alias_or_name(self.schema),
                                },
                                edge_directive.name.location,
                            ));
                            Transformed::Keep
                        }
                    }
                }
            }
            (None, Some(node_directive)) => {
                let connections_arg = node_directive.arguments.named(*CONNECTIONS_ARG_NAME);
                match connections_arg {
                    None => {
                        self.errors.push(Diagnostic::error(
                            ValidationMessage::ConnectionsArgumentRequired {
                                directive_name: node_directive.name.item,
                            },
                            node_directive.name.location,
                        ));
                        transformed_field
                    }
                    Some(connections_arg) => {
                        let edge_typename_arg = node_directive.arguments.named(*EDGE_TYPENAME_ARG);
                        if let Some(edge_typename_arg) = edge_typename_arg {
                            let field_definition = self.schema.field(field.definition.item);
                            match field_definition.type_.inner() {
                                Type::Object(_) | Type::Interface(_) | Type::Union(_) => {
                                    let handle_directive =
                                        build_handle_field_directive(HandleFieldDirectiveValues {
                                            handle: node_directive.name.item,
                                            key: "".intern(),
                                            dynamic_key: None,
                                            filters: None,
                                            handle_args: Some(vec![
                                                connections_arg.clone(),
                                                edge_typename_arg.clone(),
                                            ]),
                                        });
                                    let mut next_field = match transformed_field {
                                        Transformed::Replace(Selection::LinkedField(
                                            linked_field,
                                        )) => (*linked_field).clone(),
                                        Transformed::Keep => field.clone(),
                                        _ => panic!(
                                            "DeclarativeConnection got unexpected transform result: `{:?}`.",
                                            transformed_field
                                        ),
                                    };
                                    let index = next_field
                                        .directives
                                        .iter()
                                        .position(|directive| {
                                            directive.name.item == node_directive.name.item
                                        })
                                        .expect("Expected the edge directive to exist.");
                                    next_field.directives[index] = handle_directive;
                                    Transformed::Replace(Selection::LinkedField(Arc::new(
                                        next_field,
                                    )))
                                }
                                _ => {
                                    self.errors.push(Diagnostic::error(
                                        ValidationMessage::NodeDirectiveOnUnsupportedType {
                                            directive_name: node_directive.name.item,
                                            field_name: field.alias_or_name(self.schema),
                                            current_type: self
                                                .schema
                                                .get_type_string(&field_definition.type_),
                                        },
                                        node_directive.name.location,
                                    ));
                                    Transformed::Keep
                                }
                            }
                        } else {
                            self.errors.push(Diagnostic::error(
                                ValidationMessage::NodeDirectiveMissesRequiredEdgeTypeName {
                                    directive_name: node_directive.name.item,
                                    field_name: field.alias_or_name(self.schema),
                                },
                                node_directive.name.location,
                            ));
                            Transformed::Keep
                        }
                    }
                }
            }
        }
    }
}

#[derive(Debug, Error)]
enum ValidationMessage {
    #[error(
        "Unsupported use of @{directive_name} on field '${field_name}', 'edgeTypeName' argument must be provided."
    )]
    NodeDirectiveMissesRequiredEdgeTypeName {
        directive_name: StringKey,
        field_name: StringKey,
    },

    #[error("Invalid use of @{directive_name} on scalar field '{field_name}'.")]
    ConnectionMutationDirectiveOnScalarField {
        directive_name: StringKey,
        field_name: StringKey,
    },

    #[error(
        "Invalid use of @{directive_name} on field '{field_name}'. Expected field type 'ID', got '{current_type}'."
    )]
    DeleteRecordDirectiveOnUnsupportedType {
        directive_name: StringKey,
        field_name: StringKey,
        current_type: String,
    },

    #[error("Invalid use of @{directive_name} on linked field '{field_name}'.")]
    DeleteRecordDirectiveOnLinkedField {
        directive_name: StringKey,
        field_name: StringKey,
    },

    #[error(
        "Invalid use of @{edge_directive_name} and @{node_directive_name} on field '{field_name}' - these directives cannot be used together."
    )]
    ConflictingEdgeAndNodeDirectives {
        edge_directive_name: StringKey,
        node_directive_name: StringKey,
        field_name: StringKey,
    },

    #[error("Expected the 'connections' argument to be defined on @{directive_name}.")]
    ConnectionsArgumentRequired { directive_name: StringKey },

    #[error(
        "Unsupported use of @{directive_name} on field '{field_name}', expected an edge field (a field with 'cursor' and 'node' selection)."
    )]
    EdgeDirectiveOnUnsupportedType {
        directive_name: StringKey,
        field_name: StringKey,
    },

    #[error(
        "Unsupported use of @{directive_name} on field '{field_name}'. Expected an object, union or interface, but got '{current_type}'."
    )]
    NodeDirectiveOnUnsupportedType {
        directive_name: StringKey,
        field_name: StringKey,
        current_type: String,
    },
}
<|endoftext|>


<commit_before>use std::prelude::v1::*;

use std::any::TypeId;
use std::hash::{BuildHasherDefault, Hasher};
use std::collections::HashMap;

use task;

/// A macro to create a `static` of type `LocalKey`
///
/// This macro is intentionally similar to the `thread_local!`, and creates a
/// `static` which has a `get_mut` method to access the data on a task.
///
/// The data associated with each task local is per-task, so different tasks
/// will contain different values.
#[macro_export]
macro_rules! task_local {
    (static $NAME:ident: $t:ty = $e:expr) => (
        static $NAME: $crate::task::LocalKey<$t> = {
            fn __init() -> $t { $e }
            fn __key() -> ::std::any::TypeId {
                struct __A;
                ::std::any::TypeId::of::<__A>()
            }
            $crate::task::LocalKey {
                __init: __init,
                __key: __key,
            }
        };
    )
}

pub struct LocalMap(HashMap<TypeId, Box<Opaque>, BuildHasherDefault<IdHasher>>);

pub fn local_map() -> LocalMap {
    LocalMap(HashMap::default())
}

pub trait Opaque: Send {}
impl<T: Send> Opaque for T {}

/// A key for task-local data stored in a future's task.
///
/// This type is generated by the `task_local!` macro and performs very
/// similarly to the `thread_local!` macro and `std::thread::LocalKey` types.
/// Data associated with a `LocalKey<T>` is stored inside of a future's task,
/// and the data is destroyed when the future is completed and the task is
/// destroyed.
///
/// Task-local data can migrate between threads and hence requires a `Send`
/// bound. Additionally, task-local data also requires the `'static` bound to
/// ensure it lives long enough. When a key is accessed for the first time the
/// task's data is initialized with the provided initialization expression to
/// the macro.
#[derive(Debug)]
pub struct LocalKey<T> {
    // "private" fields which have to be public to get around macro hygiene, not
    // included in the stability story for this type. Can change at any time.
    #[doc(hidden)]
    pub __key: fn() -> TypeId,
    #[doc(hidden)]
    pub __init: fn() -> T,
}

pub struct IdHasher {
    id: u64,
}

impl Default for IdHasher {
    fn default() -> IdHasher {
        IdHasher { id: 0 }
    }
}

impl Hasher for IdHasher {
    fn write(&mut self, _bytes: &[u8]) {
        // TODO: need to do something sensible
        panic!("can only hash u64");
    }

    fn write_u64(&mut self, u: u64) {
        self.id = u;
    }

    fn finish(&self) -> u64 {
        self.id
    }
}

impl<T: Send + 'static> LocalKey<T> {
    /// Access this task-local key.
    ///
    /// This function will access this task-local key to retrieve the data
    /// associated with the current task and this key. If this is the first time
    /// this key has been accessed on this task, then the key will be
    /// initialized with the initialization expression provided at the time the
    /// `task_local!` macro was called.
    pub fn get_mut(&'static self, cx: &mut task::Context) -> &mut T {
        let key = (self.__key)();
        let data = &mut cx.map.inner.0;
        let entry: &mut Box<Opaque> = data.entry(key).or_insert_with(|| {
            Box::new((self.__init)())
        });
        unsafe { &mut *(&mut **entry as *mut Opaque as *mut T) }
    }
}
<commit_msg>Unelide the signature of task-local get_mut<commit_after>use std::prelude::v1::*;

use std::any::TypeId;
use std::hash::{BuildHasherDefault, Hasher};
use std::collections::HashMap;

use task;

/// A macro to create a `static` of type `LocalKey`
///
/// This macro is intentionally similar to the `thread_local!`, and creates a
/// `static` which has a `get_mut` method to access the data on a task.
///
/// The data associated with each task local is per-task, so different tasks
/// will contain different values.
#[macro_export]
macro_rules! task_local {
    (static $NAME:ident: $t:ty = $e:expr) => (
        static $NAME: $crate::task::LocalKey<$t> = {
            fn __init() -> $t { $e }
            fn __key() -> ::std::any::TypeId {
                struct __A;
                ::std::any::TypeId::of::<__A>()
            }
            $crate::task::LocalKey {
                __init: __init,
                __key: __key,
            }
        };
    )
}

pub struct LocalMap(HashMap<TypeId, Box<Opaque>, BuildHasherDefault<IdHasher>>);

pub fn local_map() -> LocalMap {
    LocalMap(HashMap::default())
}

pub trait Opaque: Send {}
impl<T: Send> Opaque for T {}

/// A key for task-local data stored in a future's task.
///
/// This type is generated by the `task_local!` macro and performs very
/// similarly to the `thread_local!` macro and `std::thread::LocalKey` types.
/// Data associated with a `LocalKey<T>` is stored inside of a future's task,
/// and the data is destroyed when the future is completed and the task is
/// destroyed.
///
/// Task-local data can migrate between threads and hence requires a `Send`
/// bound. Additionally, task-local data also requires the `'static` bound to
/// ensure it lives long enough. When a key is accessed for the first time the
/// task's data is initialized with the provided initialization expression to
/// the macro.
#[derive(Debug)]
pub struct LocalKey<T> {
    // "private" fields which have to be public to get around macro hygiene, not
    // included in the stability story for this type. Can change at any time.
    #[doc(hidden)]
    pub __key: fn() -> TypeId,
    #[doc(hidden)]
    pub __init: fn() -> T,
}

pub struct IdHasher {
    id: u64,
}

impl Default for IdHasher {
    fn default() -> IdHasher {
        IdHasher { id: 0 }
    }
}

impl Hasher for IdHasher {
    fn write(&mut self, _bytes: &[u8]) {
        // TODO: need to do something sensible
        panic!("can only hash u64");
    }

    fn write_u64(&mut self, u: u64) {
        self.id = u;
    }

    fn finish(&self) -> u64 {
        self.id
    }
}

impl<T: Send + 'static> LocalKey<T> {
    /// Access this task-local key.
    ///
    /// This function will access this task-local key to retrieve the data
    /// associated with the current task and this key. If this is the first time
    /// this key has been accessed on this task, then the key will be
    /// initialized with the initialization expression provided at the time the
    /// `task_local!` macro was called.
    pub fn get_mut<'a>(&'static self, cx: &'a mut task::Context) -> &'a mut T {
        let key = (self.__key)();
        let data = &mut cx.map.inner.0;
        let entry: &mut Box<Opaque> = data.entry(key).or_insert_with(|| {
            Box::new((self.__init)())
        });
        unsafe { &mut *(&mut **entry as *mut Opaque as *mut T) }
    }
}
<|endoftext|>