<commit_before>// Copyright 2016 Alexander Reece
//
// 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.

use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
use std::io;
use inflections::Inflect;

use WriteRust;
use common::Spec;

pub struct MethodModuleWriter<'a> {
    specs: &'a [Spec]
}

impl<'a> MethodModuleWriter<'a> {
    pub fn new(specs: &'a [Spec]) -> Self {
        MethodModuleWriter {
            specs: specs
        }
    }

    fn class_names(&self) -> HashSet<&'a str> {
        self.specs.iter()
            .flat_map(|spec| spec.classes())
            .map(|class| class.name())
            .collect()
    }

    fn class_methods(&'a self) -> BTreeMap<&'a str, BTreeSet<&'a str>> {
        self.class_names().into_iter()
            .filter_map(|class_name| {
                let method_names = self.specs.iter()
                    .filter_map(|spec| spec.class(class_name))
                    .flat_map(|class| class.methods())
                    .map(|method| method.name())
                    .collect::<BTreeSet<_>>();

                if method_names.is_empty() {
                    None
                } else {
                    Some((class_name, method_names))
                }
            })
            .collect()
    }

    fn class_method_fields(&'a self) -> BTreeMap<&'a str, BTreeMap<&'a str, BTreeSet<&'a str>>> {
        self.class_methods().into_iter()
            .map(|(class_name, method_names)| {
                // class_name
                let class_methods = method_names.into_iter()
                    .map(|method_name| {
                        // method_name
                        let method_fields = self.specs.iter()
                            .filter_map(|spec| spec.class(class_name))
                            .filter_map(|class| class.method(method_name))
                            .flat_map(|method| method.fields())
                            .filter(|field| !field.is_reserved())
                            .map(|field| field.var_name())
                            .collect();
                        (method_name, method_fields)
                        // end method_name
                    })
                    .collect();
                (class_name, class_methods)
                // end class name
            })
            .collect()
    }
}

impl<'a> WriteRust for MethodModuleWriter<'a> {
    fn write_rust_to<W>(&self, writer: &mut W) -> io::Result<()>
        where W: io::Write
    {
        for (class_name, class_methods) in self.class_method_fields() {
            let class_snake = class_name.to_snake_case();
            try!(writeln!(writer, "\npub mod {} {{", class_snake));

            for (method_name, field_names) in class_methods {
                let method_pascal = method_name.to_pascal_case();
                if field_names.is_empty() {
                    try!(writeln!(writer, "pub trait {}Method {{}}", method_pascal));
                    continue;
                }

                let has_lifetimes = self.specs.iter()
                    .filter_map(|spec| spec.class(class_name))
                    .filter_map(|class| class.method(method_name))
                    .any(|method| method.has_lifetimes());

                let lifetimes = if has_lifetimes { "<'a>" } else { "" };
                let section = format!("pub trait {}Method{}", method_pascal, lifetimes);

                try!(writeln!(writer, "{} {{", section));
                for field_name in field_names {
                    let tys = self.specs.iter()
                        .filter_map(|spec| spec.class(class_name))
                        .filter_map(|class| class.method(method_name))
                        .flat_map(|method| method.fields())
                        .filter(|field| field.var_name() == field_name)
                        .map(|field| (field.ty()))
                        .map(|ty| (ty.owned_type(), ty))
                        .collect::<HashMap<_, _>>();

                    if tys.len() > 1 {
                        panic!(
                            "Conflicting types for {}::{}::{}",
                            class_name,
                            method_name,
                            field_name
                        );
                    }

                    let var_name = field_name.to_snake_case();
                    let ty = tys.values().next().unwrap();
                    if ty.is_copy() {
                        try!(writeln!(
                                writer,
                            "fn set_{0}(_: {1}) {{}}",
                            var_name,
                            ty.owned_type()
                        ));
                    }
                    else {
                        try!(writeln!(
                            writer,
                            "fn set_{0}<V>(_: V) where V: Into<{1}> {{}}",
                            var_name,
                            ty.cow_definition("a")
                        ));
                    }
                }
                try!(writeln!(writer, "}} // {}\n", section));
            }
            try!(writeln!(writer, "}} // mod {}", class_snake));
        }

        Ok(())
    }
}<commit_msg>[primitives] Rename Method traits to SetMethodFields<commit_after>// Copyright 2016 Alexander Reece
//
// 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.

use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
use std::io;
use inflections::Inflect;

use WriteRust;
use common::Spec;

pub struct MethodModuleWriter<'a> {
    specs: &'a [Spec]
}

impl<'a> MethodModuleWriter<'a> {
    pub fn new(specs: &'a [Spec]) -> Self {
        MethodModuleWriter {
            specs: specs
        }
    }

    fn class_names(&self) -> HashSet<&'a str> {
        self.specs.iter()
            .flat_map(|spec| spec.classes())
            .map(|class| class.name())
            .collect()
    }

    fn class_methods(&'a self) -> BTreeMap<&'a str, BTreeSet<&'a str>> {
        self.class_names().into_iter()
            .filter_map(|class_name| {
                let method_names = self.specs.iter()
                    .filter_map(|spec| spec.class(class_name))
                    .flat_map(|class| class.methods())
                    .map(|method| method.name())
                    .collect::<BTreeSet<_>>();

                if method_names.is_empty() {
                    None
                } else {
                    Some((class_name, method_names))
                }
            })
            .collect()
    }

    fn class_method_fields(&'a self) -> BTreeMap<&'a str, BTreeMap<&'a str, BTreeSet<&'a str>>> {
        self.class_methods().into_iter()
            .map(|(class_name, method_names)| {
                // class_name
                let class_methods = method_names.into_iter()
                    .map(|method_name| {
                        // method_name
                        let method_fields = self.specs.iter()
                            .filter_map(|spec| spec.class(class_name))
                            .filter_map(|class| class.method(method_name))
                            .flat_map(|method| method.fields())
                            .filter(|field| !field.is_reserved())
                            .map(|field| field.var_name())
                            .collect();
                        (method_name, method_fields)
                        // end method_name
                    })
                    .collect();
                (class_name, class_methods)
                // end class name
            })
            .collect()
    }
}

impl<'a> WriteRust for MethodModuleWriter<'a> {
    fn write_rust_to<W>(&self, writer: &mut W) -> io::Result<()>
        where W: io::Write
    {
        for (class_name, class_methods) in self.class_method_fields() {
            let class_snake = class_name.to_snake_case();
            try!(writeln!(writer, "\npub mod {} {{", class_snake));

            for (method_name, field_names) in class_methods {
                let method_pascal = method_name.to_pascal_case();
                if field_names.is_empty() {
                    try!(writeln!(writer, "pub trait Set{}MethodFields {{}}", method_pascal));
                    continue;
                }

                let has_lifetimes = self.specs.iter()
                    .filter_map(|spec| spec.class(class_name))
                    .filter_map(|class| class.method(method_name))
                    .any(|method| method.has_lifetimes());

                let lifetimes = if has_lifetimes { "<'a>" } else { "" };
                let section = format!("pub trait Set{}MethodFields{}", method_pascal, lifetimes);

                try!(writeln!(writer, "{} {{", section));
                for field_name in field_names {
                    let tys = self.specs.iter()
                        .filter_map(|spec| spec.class(class_name))
                        .filter_map(|class| class.method(method_name))
                        .flat_map(|method| method.fields())
                        .filter(|field| field.var_name() == field_name)
                        .map(|field| (field.ty()))
                        .map(|ty| (ty.owned_type(), ty))
                        .collect::<HashMap<_, _>>();

                    if tys.len() > 1 {
                        panic!(
                            "Conflicting types for {}::{}::{}",
                            class_name,
                            method_name,
                            field_name
                        );
                    }

                    let var_name = field_name.to_snake_case();
                    let ty = tys.values().next().unwrap();
                    if ty.is_copy() {
                        try!(writeln!(
                                writer,
                            "fn set_{0}(_: {1}) {{}}",
                            var_name,
                            ty.owned_type()
                        ));
                    }
                    else {
                        try!(writeln!(
                            writer,
                            "fn set_{0}<V>(_: V) where V: Into<{1}> {{}}",
                            var_name,
                            ty.cow_definition("a")
                        ));
                    }
                }
                try!(writeln!(writer, "}} // {}\n", section));
            }
            try!(writeln!(writer, "}} // mod {}", class_snake));
        }

        Ok(())
    }
}<|endoftext|>


<commit_before>
use std::str;
use parser::{Node, Children, Data};
use self::unescape::unescapeString;

#[macro_escape]
mod macros;
mod unescape;

// TODO: figure out how to write this as a function on Node; the borrow checker
// was extra painful the last time you tried.
macro_rules! node_children( ( $node:expr ) => ( {
  match $node.contents {
    Children( ref nodes ) => nodes,
    _ => fail!( "No children in node." )
  } } ) )

pub fn codeForNode( node: &Node ) -> ~str {
  match node.name {
    "Definition" => wrapChildrenOutput( ~"rule!( ", node, ~" )\n" ),
    "Expression" => expressionOutput( node ),
    "Sequence" => sequenceOutput( node ),
    "Literal" => literalOutput( node ),
    "Class" => classOutput( node ),
    "Suffix" => suffixOutput( node ),
    "Prefix" => prefixOutput( node ),
    "Primary" => primaryOutput( node ),
    "DOT" => ~"base::Dot",
    "LEFTARROW" => ~" <- ",
    "SLASH" => ~", ",
    "Spacing" => ~"",
    "EndOfLine" => ~"",
    "OPEN" => ~"",
    "CLOSE" => ~"",
    _ => codeForNodeContents( node )
  }
}

fn codeForNodeContents( node: &Node ) -> ~str {
  match node.contents {
    Children( ref children ) => {
      // TODO: use StrBuf here somehow when it's merged in
      children.iter().map( codeForNode ).collect::<Vec<~str>>().concat()
    }
    Data( data ) => str::from_utf8( data ).unwrap().to_owned(),
  }
}

fn wrapChildrenOutput( before: ~str, node: &Node, after: ~str ) -> ~str {
  [ before, codeForNodeContents( node ), after ].concat()
}

fn wrapNodeOutput( before: ~str, node: &Node, after: ~str ) -> ~str {
  [ before, codeForNode( node ), after ].concat()
}

fn expressionOutput( node: &Node ) -> ~str {
  let children = node_children!( node );
  if children.len() > 1 {
    wrapChildrenOutput( ~"or!( ", node , ~" )" )
  } else {
    codeForNodeContents( node )
  }
}

fn sequenceOutput( node: &Node ) -> ~str {
  let children = node_children!( node );
  if children.len() > 1 {
    let mut output = StrBuf::from_str( "seq!( " );
    for i in range( 0, children.len() ) {
      output.push_str( codeForNode( children.get( i ) ) );
      if i != children.len() -1 {
        output.push_str( ", " );
      }
    }
    output.push_str( " )" );
    output.into_owned()
  } else {
    codeForNodeContents( node )
  }
}

fn suffixOutput( node: &Node ) -> ~str {
  let children = node_children!( node );
  if children.len() == 2 {
    let macro_name = match children.get( 1 ).name {
      "QUESTION" => ~"opt",
      "STAR" => ~"star",
      "PLUS" => ~"plus",
      _ => fail!( "Bad second child." )
    };

    [ macro_name, ~"!( ", codeForNode( children.get( 0 ) ), ~" )" ].concat()
  } else {
    codeForNodeContents( node )
  }
}

fn prefixOutput( node: &Node ) -> ~str {
  let children = node_children!( node );
  if children.len() == 2 {
    let macro_name = children.get( 0 ).name.chars()
      .map( |x| x.to_lowercase() ).collect();
    [ macro_name, ~"!( ", codeForNode( children.get( 1 ) ), ~" )" ].concat()
  } else {
    codeForNodeContents( node )
  }
}

fn primaryOutput( node: &Node ) -> ~str {
  let children = node_children!( node );
  if children.len() == 1 && children.get( 0 ).name == "Identifier" {
    wrapNodeOutput( ~"ex!( ", children.get( 0 ), ~" )" )
  } else {
    codeForNodeContents( node )
  }
}

fn literalOutput( node: &Node ) -> ~str {
  // TODO: write escape func
  stringBasedRule( node, "lit" )
    .replace( "\n", r"\n" )
    .replace( "\t", r"\t" )
    .replace( "\r", r"\r" )
}

fn classOutput( node: &Node ) -> ~str {
  stringBasedRule( node, "class" )
    .replace( r"\]", r"]" )
    .replace( r"\[", r"[" )
}

fn stringBasedRule( node: &Node, rule_name: &str ) -> ~str {
  let content = unescapeString( codeForNodeContents( node ).trim() );

  [ rule_name.to_owned(),
    ~"!( r\"",
    content.slice_chars( 1, content.len() - 1 ).replace( r#"""#, r#"\""# ),
    ~"\" )" ].concat()
}
<commit_msg>Fully working generated code! (only macros)<commit_after>
use std::str;
use parser::{Node, Children, Data};
use self::unescape::unescapeString;

#[macro_escape]
mod macros;
mod unescape;

// TODO: figure out how to write this as a function on Node; the borrow checker
// was extra painful the last time you tried.
macro_rules! node_children( ( $node:expr ) => ( {
  match $node.contents {
    Children( ref nodes ) => nodes,
    _ => fail!( "No children in node." )
  } } ) )

pub fn codeForNode( node: &Node ) -> ~str {
  match node.name {
    "Definition" => wrapChildrenOutput( ~"rule!( ", node, ~" )\n" ),
    "Expression" => expressionOutput( node ),
    "Sequence" => sequenceOutput( node ),
    "Literal" => literalOutput( node ),
    "Class" => classOutput( node ),
    "Suffix" => suffixOutput( node ),
    "Prefix" => prefixOutput( node ),
    "Primary" => primaryOutput( node ),
    "DOT" => ~"base::Dot",
    "LEFTARROW" => ~" <- ",
    "SLASH" => ~", ",
    "Spacing" => ~"",
    "EndOfLine" => ~"",
    "OPEN" => ~"",
    "CLOSE" => ~"",
    _ => codeForNodeContents( node )
  }
}

fn codeForNodeContents( node: &Node ) -> ~str {
  match node.contents {
    Children( ref children ) => {
      // TODO: use StrBuf here somehow when it's merged in
      children.iter().map( codeForNode ).collect::<Vec<~str>>().concat()
    }
    Data( data ) => str::from_utf8( data ).unwrap().to_owned(),
  }
}

fn wrapChildrenOutput( before: ~str, node: &Node, after: ~str ) -> ~str {
  [ before, codeForNodeContents( node ), after ].concat()
}

fn wrapNodeOutput( before: ~str, node: &Node, after: ~str ) -> ~str {
  [ before, codeForNode( node ), after ].concat()
}

fn expressionOutput( node: &Node ) -> ~str {
  let children = node_children!( node );
  if children.len() > 1 {
    wrapChildrenOutput( ~"or!( ", node , ~" )" )
  } else {
    codeForNodeContents( node )
  }
}

fn sequenceOutput( node: &Node ) -> ~str {
  let children = node_children!( node );
  if children.len() > 1 {
    let mut output = StrBuf::from_str( "seq!( " );
    for i in range( 0, children.len() ) {
      output.push_str( codeForNode( children.get( i ) ) );
      if i != children.len() -1 {
        output.push_str( ", " );
      }
    }
    output.push_str( " )" );
    output.into_owned()
  } else {
    codeForNodeContents( node )
  }
}

fn suffixOutput( node: &Node ) -> ~str {
  let children = node_children!( node );
  if children.len() == 2 {
    let macro_name = match children.get( 1 ).name {
      "QUESTION" => ~"opt",
      "STAR" => ~"star",
      "PLUS" => ~"plus",
      _ => fail!( "Bad second child." )
    };

    [ macro_name, ~"!( ", codeForNode( children.get( 0 ) ), ~" )" ].concat()
  } else {
    codeForNodeContents( node )
  }
}

fn prefixOutput( node: &Node ) -> ~str {
  let children = node_children!( node );
  if children.len() == 2 {
    let macro_name = children.get( 0 ).name.chars()
      .map( |x| x.to_lowercase() ).collect();
    [ macro_name, ~"!( ", codeForNode( children.get( 1 ) ), ~" )" ].concat()
  } else {
    codeForNodeContents( node )
  }
}

fn primaryOutput( node: &Node ) -> ~str {
  let children = node_children!( node );
  if children.len() == 1 && children.get( 0 ).name == "Identifier" {
    wrapNodeOutput( ~"ex!( ", children.get( 0 ), ~" )" )
  } else {
    codeForNodeContents( node )
  }
}

fn literalOutput( node: &Node ) -> ~str {
  stringBasedRule( node, "lit" )
}

fn classOutput( node: &Node ) -> ~str {
  stringBasedRule( node, "class" )
    .replace( r"\\]", r"]" )
    .replace( r"\\[", r"[" )
}

fn stringBasedRule( node: &Node, rule_name: &str ) -> ~str {
  let full = codeForNodeContents( node );
  // TODO: write escape func
  let content = unescapeString( full.slice_chars( 1, full.len() - 1 ) )
    .replace( r"\", r"\\" )
    .replace( "\n", r"\n" )
    .replace( "\t", r"\t" )
    .replace( "\r", r"\r" )
    .replace( "\"", r#"\""# );
  format!( "{}!( \"{}\" )", rule_name, content )
}
<|endoftext|>


<commit_before>use std::cell::{Ref, RefCell};
use std::cmp;
use std::io;
use std::net::{Ipv4Addr, SocketAddr};
use std::num::Wrapping;
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_source::PacketSource;
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,
    packet_for_client_length: Option<u16>,
    closed: bool,
    tcb: TCB,
}

struct TCB {
    state: TCPState,
    syn_sequence_number: u32,
    sequence_number: Wrapping<u32>,
    acknowledgement_number: Wrapping<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: Wrapping(0),
            acknowledgement_number: Wrapping(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,
                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);
    }

    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) => {
                cx_error!(target: TAG, self.id, "Cannot write: {}", 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)) => {
                match Self::send_to_client(&self.client, selector, &ipv4_packet) {
                    Ok(_) => self.packet_for_client_length = None, // packet consumed
                    Err(_) => {
                        // ask to the client to pull when its buffer is not full
                        let client_rc = self.client.upgrade().expect("Expected client not found");
                        let mut client = client_rc.borrow_mut();
                        let self_rc = self.self_weak.upgrade().unwrap();
                        client.register_pending_packet_source(self_rc);
                    }
                }
                Ok(Some(()))
            },
            Ok(None) => Ok(None),
            Err(err) => Err(err),
        };
        match non_lexical_lifetime_workaround {
            Ok(None) => self.eof(selector),
            Err(err) => {
                cx_error!(target: TAG, self.id, "Cannot read: {}", 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_empty_packet_to_client(&mut self, selector: &mut Selector, flags: u16) {
        let ipv4_packet = Self::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_SYN);
        if let Err(err) = Self::send_to_client(&self.client, selector, &ipv4_packet) {
            // losing such an empty packet will not break the TCP connection
            cx_warn!(target: TAG, self.id, "Cannot send packet to client: {}", err);
        }
    }

    fn eof(&mut self, selector: &mut Selector) {
        self.tcb.remote_closed = true;
        if self.tcb.state == TCPState::CloseWait {
            self.send_empty_packet_to_client(selector, tcp_header::FLAG_FIN);
            self.tcb.sequence_number += Wrapping(1); // FIN counts for 1 byte
        }
    }

    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.0);
            tcp_header.set_acknowledgement_number(tcb.acknowledgement_number.0);
            tcp_header.set_flags(flags);
        } else {
            panic!("Not a TCP header");
        }
    }

    fn handle_packet(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        if let Some(TransportHeader::TCP(tcp_header)) = ipv4_packet.transport_header() {
            if self.tcb.state == TCPState::Init {
                self.handle_first_packet(selector, ipv4_packet);
                return;
            }

            if tcp_header.is_syn() {
                self.handle_duplicate_syn(selector, ipv4_packet);
                return;
            }

            if tcp_header.sequence_number() != self.tcb.acknowledgement_number.0 {
                // ignore packet already received or out-of-order, retransmission is already
                // managed by both sides
                cx_warn!(target: TAG, self.id, "Ignoring packet {}; expecting {}; flags={}", tcp_header.sequence_number(), tcp_header.acknowledgement_number(), tcp_header.flags());
                self.send_empty_packet_to_client(selector, tcp_header::FLAG_ACK); // re-ack
                return;
            }

            self.tcb.client_window = tcp_header.window();
            self.tcb.their_acknowledgement_number = tcp_header.acknowledgement_number();

            cx_debug!(target: TAG, self.id, "Receiving expected packet {} (flags={})", tcp_header.sequence_number(), tcp_header.flags());

            if tcp_header.is_rst() {
                self.close(selector);
                return;
            }

            if tcp_header.is_ack() {
                cx_debug!(target: TAG, self.id, "Client acked {}", tcp_header.acknowledgement_number());
            }

            if tcp_header.is_fin() {
                self.handle_fin(selector, ipv4_packet);
            } else if tcp_header.is_ack() {
                self.handle_ack(selector, ipv4_packet);
            }
        } else {
            panic!("Not a TCP packet");
        }
    }

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

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

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

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

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

    fn reset_connection(&mut self, selector: &mut Selector) {
        self.send_empty_packet_to_client(selector, tcp_header::FLAG_RST);
        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) {
        self.handle_packet(selector, ipv4_packet);
        self.update_interests(selector);
    }

    fn disconnect(&mut self, selector: &mut Selector) {
        cx_info!(target: TAG, self.id, "Close");
        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);
        }
    }
}

impl PacketSource for TCPConnection {
    fn get(&mut self) -> Option<IPv4Packet> {
        if let Some(len) = self.packet_for_client_length {
            Some(self.network_to_client.inflate(len))
        } else {
            None
        }
    }

    fn next(&mut self, selector: &mut Selector) {
        let len = self.packet_for_client_length.expect("next() called on empty packet source");
        self.tcb.sequence_number += Wrapping(len as u32);
        self.packet_for_client_length = None;
        self.update_interests(selector);
    }
}
<commit_msg>Extract tcp_header without "if let"<commit_after>use std::cell::{Ref, RefCell};
use std::cmp;
use std::io;
use std::net::{Ipv4Addr, SocketAddr};
use std::num::Wrapping;
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_source::PacketSource;
use super::packetizer::Packetizer;
use super::selector::{EventHandler, Selector};
use super::stream_buffer::StreamBuffer;
use super::tcp_header::{self, TCPHeader, TCPHeaderMut};
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,
    packet_for_client_length: Option<u16>,
    closed: bool,
    tcb: TCB,
}

struct TCB {
    state: TCPState,
    syn_sequence_number: u32,
    sequence_number: Wrapping<u32>,
    acknowledgement_number: Wrapping<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: Wrapping(0),
            acknowledgement_number: Wrapping(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,
                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);
    }

    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) => {
                cx_error!(target: TAG, self.id, "Cannot write: {}", 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)) => {
                match Self::send_to_client(&self.client, selector, &ipv4_packet) {
                    Ok(_) => self.packet_for_client_length = None, // packet consumed
                    Err(_) => {
                        // ask to the client to pull when its buffer is not full
                        let client_rc = self.client.upgrade().expect("Expected client not found");
                        let mut client = client_rc.borrow_mut();
                        let self_rc = self.self_weak.upgrade().unwrap();
                        client.register_pending_packet_source(self_rc);
                    }
                }
                Ok(Some(()))
            },
            Ok(None) => Ok(None),
            Err(err) => Err(err),
        };
        match non_lexical_lifetime_workaround {
            Ok(None) => self.eof(selector),
            Err(err) => {
                cx_error!(target: TAG, self.id, "Cannot read: {}", 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_empty_packet_to_client(&mut self, selector: &mut Selector, flags: u16) {
        let ipv4_packet = Self::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_SYN);
        if let Err(err) = Self::send_to_client(&self.client, selector, &ipv4_packet) {
            // losing such an empty packet will not break the TCP connection
            cx_warn!(target: TAG, self.id, "Cannot send packet to client: {}", err);
        }
    }

    fn eof(&mut self, selector: &mut Selector) {
        self.tcb.remote_closed = true;
        if self.tcb.state == TCPState::CloseWait {
            self.send_empty_packet_to_client(selector, tcp_header::FLAG_FIN);
            self.tcb.sequence_number += Wrapping(1); // FIN counts for 1 byte
        }
    }

    #[inline]
    fn tcp_header_of_transport_mut<'a>(transport_header: TransportHeaderMut<'a>) -> TCPHeaderMut<'a> {
        if let TransportHeaderMut::TCP(tcp_header) = transport_header {
            tcp_header
        } else {
            panic!("Not a TCP header");
        }
    }

    #[inline]
    fn tcp_header_of_packet<'a>(ipv4_packet: &'a IPv4Packet) -> TCPHeader<'a> {
        if let Some(TransportHeader::TCP(tcp_header)) = ipv4_packet.transport_header() {
            tcp_header
        } else {
            panic!("Not a TCP packet");
        }
    }

    fn update_headers(packetizer: &mut Packetizer, tcb: &TCB, flags: u16) {
        let mut tcp_header = Self::tcp_header_of_transport_mut(packetizer.transport_header_mut());
        tcp_header.set_sequence_number(tcb.sequence_number.0);
        tcp_header.set_acknowledgement_number(tcb.acknowledgement_number.0);
        tcp_header.set_flags(flags);
    }

    fn handle_packet(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        let tcp_header = Self::tcp_header_of_packet(ipv4_packet);
        if self.tcb.state == TCPState::Init {
            self.handle_first_packet(selector, ipv4_packet);
            return;
        }

        if tcp_header.is_syn() {
            self.handle_duplicate_syn(selector, ipv4_packet);
            return;
        }

        if tcp_header.sequence_number() != self.tcb.acknowledgement_number.0 {
            // ignore packet already received or out-of-order, retransmission is already
            // managed by both sides
            cx_warn!(target: TAG, self.id, "Ignoring packet {}; expecting {}; flags={}", tcp_header.sequence_number(), tcp_header.acknowledgement_number(), tcp_header.flags());
            self.send_empty_packet_to_client(selector, tcp_header::FLAG_ACK); // re-ack
            return;
        }

        self.tcb.client_window = tcp_header.window();
        self.tcb.their_acknowledgement_number = tcp_header.acknowledgement_number();

        cx_debug!(target: TAG, self.id, "Receiving expected packet {} (flags={})", tcp_header.sequence_number(), tcp_header.flags());

        if tcp_header.is_rst() {
            self.close(selector);
            return;
        }

        if tcp_header.is_ack() {
            cx_debug!(target: TAG, self.id, "Client acked {}", tcp_header.acknowledgement_number());
        }

        if tcp_header.is_fin() {
            self.handle_fin(selector, ipv4_packet);
        } else if tcp_header.is_ack() {
            self.handle_ack(selector, ipv4_packet);
        }
    }

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

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

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

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

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

    fn reset_connection(&mut self, selector: &mut Selector) {
        self.send_empty_packet_to_client(selector, tcp_header::FLAG_RST);
        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) {
        self.handle_packet(selector, ipv4_packet);
        self.update_interests(selector);
    }

    fn disconnect(&mut self, selector: &mut Selector) {
        cx_info!(target: TAG, self.id, "Close");
        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);
        }
    }
}

impl PacketSource for TCPConnection {
    fn get(&mut self) -> Option<IPv4Packet> {
        if let Some(len) = self.packet_for_client_length {
            Some(self.network_to_client.inflate(len))
        } else {
            None
        }
    }

    fn next(&mut self, selector: &mut Selector) {
        let len = self.packet_for_client_length.expect("next() called on empty packet source");
        self.tcb.sequence_number += Wrapping(len as u32);
        self.packet_for_client_length = None;
        self.update_interests(selector);
    }
}
<|endoftext|>


<commit_before>use linked_hash_map::LinkedHashMap;
use std::collections::HashSet;

use crate::options::{ImportStyle, Options, StringTransform};
use crate::shape::{self, Shape};
use crate::to_singular::to_singular;
use crate::util::{kebab_case, lower_camel_case, snake_case, type_case};

#[derive(PartialEq, PartialOrd, Ord, Eq, Hash, Clone, Copy)]
enum Import {
    Any,
    Optional,
    BaseModel,
    Field,
}

impl Import {
    fn pair(&self) -> (&'static str, &'static str) {
        match self {
            Import::Any => ("typing", "Any"),
            Import::Optional => ("typing", "Optional"),
            Import::BaseModel => ("pydantic", "BaseModel"),
            Import::Field => ("pydantic", "Field"),
        }
    }
    fn module(&self) -> &'static str {
        self.pair().0
    }
    fn identifier(&self) -> &'static str {
        self.pair().1
    }
    fn qualified(&self) -> String {
        let (module, identifier) = self.pair();
        format!("{}.{}", module, identifier)
    }
}

struct Ctxt {
    options: Options,
    type_names: HashSet<String>,
    imports: HashSet<Import>,
    created_classes: Vec<(Shape, Ident)>,
}

pub type Ident = String;
pub type Code = String;

pub fn python_types(name: &str, shape: &Shape, options: Options) -> Code {
    let mut ctxt = Ctxt {
        options,
        type_names: HashSet::new(),
        imports: HashSet::new(),
        created_classes: Vec::new(),
    };

    if !matches!(shape, Shape::Struct { .. }) {
        // reserve the requested name
        ctxt.type_names.insert(name.to_string());
    }

    let (ident, code) = type_from_shape(&mut ctxt, name, shape);
    let mut code = code.unwrap_or_default();

    if !ctxt.imports.is_empty() {
        let mut imports: Vec<_> = ctxt.imports.drain().collect();
        imports.sort();
        let mut import_code = String::new();
        match ctxt.options.import_style {
            ImportStyle::AssumeExisting => {}
            ImportStyle::AddImports => {
                for import in imports {
                    let (module, identifier) = import.pair();
                    import_code += &format!("from {} import {}\n", module, identifier);
                }
            }
            ImportStyle::QualifiedPaths => {
                let mut seen = HashSet::new();
                for import in imports {
                    let module = import.module();
                    if seen.insert(module) {
                        import_code += &format!("import {}\n", module);
                    }
                }
            }
        }
        import_code += "\n\n";
        code = import_code + &code;
    }

    if ident != name {
        code += &format!("\n\n{} = {}", name, ident);
    }
    code
}

fn type_from_shape(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    use crate::shape::Shape::*;
    match shape {
        Null | Any | Bottom => (import(ctxt, Import::Any), None),
        Bool => ("bool".into(), None),
        StringT => ("str".into(), None),
        Integer => ("int".into(), None),
        Floating => ("float".into(), None),
        Tuple(shapes, _n) => {
            let folded = shape::fold_shapes(shapes.clone());
            if folded == Any && shapes.iter().any(|s| s != &Any) {
                generate_tuple_type(ctxt, path, shapes)
            } else {
                generate_vec_type(ctxt, path, &folded)
            }
        }
        VecT { elem_type: e } => generate_vec_type(ctxt, path, e),
        Struct { fields } => generate_data_class(ctxt, path, fields, shape),
        MapT { val_type: v } => generate_map_type(ctxt, path, v),
        Opaque(t) => (t.clone(), None),
        Optional(e) => {
            let (inner, defs) = type_from_shape(ctxt, path, e);
            if ctxt.options.use_default_for_missing_fields {
                (inner, defs)
            } else {
                let optional = import(ctxt, Import::Optional);
                (format!("{}[{}]", optional, inner), defs)
            }
        }
    }
}

fn generate_vec_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("list[{}]", inner), defs)
}

fn generate_map_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("dict[str, {}]", inner), defs)
}

fn generate_tuple_type(ctxt: &mut Ctxt, path: &str, shapes: &[Shape]) -> (Ident, Option<Code>) {
    let mut types = Vec::new();
    let mut defs = Vec::new();

    for shape in shapes {
        let (typ, def) = type_from_shape(ctxt, path, shape);
        types.push(typ);
        if let Some(code) = def {
            defs.push(code)
        }
    }

    (
        format!("tuple[{}]", types.join(", ")),
        Some(defs.join("\n\n")),
    )
}

fn field_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "field", snake_case)
}

fn type_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "GeneratedType", type_case)
}

// https://docs.python.org/3/reference/lexical_analysis.html#keywords
#[rustfmt::skip]
const PYTHON_KEYWORDS: &[&str] = &[
    "False", "None", "True",
    "and", "as", "assert", "async", "await", "break", "class", "continue",
    "def", "del", "elif", "else", "except", "finally", "for", "from", "global",
    "if", "import", "in", "is", "lambda", "nonlocal", "not", "or", "pass",
    "raise", "return", "try", "while", "with", "yield",
];

fn type_or_field_name(
    name: &str,
    used_names: &HashSet<String>,
    default_name: &str,
    case_fn: fn(&str) -> String,
) -> Ident {
    let name = name.trim();
    let mut output_name = case_fn(name);
    if PYTHON_KEYWORDS.contains(&&*output_name) {
        output_name.push_str("_field");
    }
    if output_name.is_empty() {
        output_name.push_str(default_name);
    }
    if let Some(c) = output_name.chars().next() {
        if c.is_ascii_digit() {
            output_name = String::from("n") + &output_name;
        }
    }
    if !used_names.contains(&output_name) {
        return output_name;
    }
    for n in 2.. {
        let temp = format!("{}{}", output_name, n);
        if !used_names.contains(&temp) {
            return temp;
        }
    }
    unreachable!()
}

fn import(ctxt: &mut Ctxt, import: Import) -> String {
    ctxt.imports.insert(import);
    match ctxt.options.import_style {
        ImportStyle::QualifiedPaths => import.qualified(),
        _ => import.identifier().into(),
    }
}

fn generate_data_class(
    ctxt: &mut Ctxt,
    path: &str,
    field_shapes: &LinkedHashMap<String, Shape>,
    containing_shape: &Shape,
) -> (Ident, Option<Code>) {
    for (created_for_shape, ident) in ctxt.created_classes.iter() {
        if created_for_shape.is_acceptable_substitution_for(containing_shape) {
            return (ident.into(), None);
        }
    }

    let type_name = type_name(path, &ctxt.type_names);
    ctxt.type_names.insert(type_name.clone());
    ctxt.created_classes
        .push((containing_shape.clone(), type_name.clone()));

    let mut field_names = HashSet::new();
    let mut defs = Vec::new();

    let fields: Vec<Code> = field_shapes
        .iter()
        .map(|(name, typ)| {
            let field_name = field_name(name, &field_names);
            field_names.insert(field_name.clone());

            let (field_type, child_defs) = type_from_shape(ctxt, name, typ);

            if let Some(code) = child_defs {
                defs.push(code);
            }

            let mut field_code = String::new();
            let transformed = apply_transform(ctxt, &field_name, name);
            if transformed != field_name {
                field_code += &format!(" = {}(alias = \"{}\")", import(ctxt, Import::Field), transformed)
            }

            format!("    {}: {}{}", field_name, field_type, field_code)
        })
        .collect();

    let mut code = String::new();

    code += &format!(
        "class {}({}):\n",
        type_name,
        import(ctxt, Import::BaseModel)
    );

    if fields.is_empty() {
        code += "    pass\n";
    } else {
        code += &fields.join("\n");
        code += "\n";
    }

    if !defs.is_empty() {
        let mut d = defs.join("\n\n");
        d += "\n\n";
        d += &code;
        code = d;
    }

    (type_name, Some(code))
}

fn apply_transform(ctxt: &Ctxt, field_name: &str, name: &str) -> String {
    match ctxt.options.property_name_format {
        Some(StringTransform::LowerCase) => field_name.to_ascii_lowercase(),
        Some(StringTransform::PascalCase) => type_case(field_name),
        Some(StringTransform::SnakeCase) => snake_case(field_name),
        Some(StringTransform::KebabCase) => kebab_case(field_name),
        Some(StringTransform::UpperCase) => field_name.to_ascii_uppercase(),
        Some(StringTransform::CamelCase) => lower_camel_case(field_name),
        Some(StringTransform::ScreamingSnakeCase) => snake_case(field_name).to_ascii_uppercase(),
        Some(StringTransform::ScreamingKebabCase) => kebab_case(field_name).to_ascii_uppercase(),
        None => name.to_string(),
    }
}

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

    #[test]
    fn field_names_test() {
        fn field_name_test(from: &str, to: &str) {
            assert_eq!(
                field_name(from, &HashSet::new()),
                to.to_string(),
                r#"From "{}" to "{}""#,
                from,
                to
            );
        }

        field_name_test("valid", "valid");
        field_name_test("1", "n1");
        field_name_test("+1", "n1");
        field_name_test("", "field");
        field_name_test("def", "def_field");
    }
}
<commit_msg>shared: Prevent redundant new lines for python.<commit_after>use linked_hash_map::LinkedHashMap;
use std::collections::HashSet;

use crate::options::{ImportStyle, Options, StringTransform};
use crate::shape::{self, Shape};
use crate::to_singular::to_singular;
use crate::util::{kebab_case, lower_camel_case, snake_case, type_case};

#[derive(PartialEq, PartialOrd, Ord, Eq, Hash, Clone, Copy)]
enum Import {
    Any,
    Optional,
    BaseModel,
    Field,
}

impl Import {
    fn pair(&self) -> (&'static str, &'static str) {
        match self {
            Import::Any => ("typing", "Any"),
            Import::Optional => ("typing", "Optional"),
            Import::BaseModel => ("pydantic", "BaseModel"),
            Import::Field => ("pydantic", "Field"),
        }
    }
    fn module(&self) -> &'static str {
        self.pair().0
    }
    fn identifier(&self) -> &'static str {
        self.pair().1
    }
    fn qualified(&self) -> String {
        let (module, identifier) = self.pair();
        format!("{}.{}", module, identifier)
    }
}

struct Ctxt {
    options: Options,
    type_names: HashSet<String>,
    imports: HashSet<Import>,
    created_classes: Vec<(Shape, Ident)>,
}

pub type Ident = String;
pub type Code = String;

pub fn python_types(name: &str, shape: &Shape, options: Options) -> Code {
    let mut ctxt = Ctxt {
        options,
        type_names: HashSet::new(),
        imports: HashSet::new(),
        created_classes: Vec::new(),
    };

    if !matches!(shape, Shape::Struct { .. }) {
        // reserve the requested name
        ctxt.type_names.insert(name.to_string());
    }

    let (ident, code) = type_from_shape(&mut ctxt, name, shape);
    let mut code = code.unwrap_or_default();

    if !ctxt.imports.is_empty() {
        let mut imports: Vec<_> = ctxt.imports.drain().collect();
        imports.sort();
        let mut import_code = String::new();
        match ctxt.options.import_style {
            ImportStyle::AssumeExisting => {}
            ImportStyle::AddImports => {
                for import in imports {
                    let (module, identifier) = import.pair();
                    import_code += &format!("from {} import {}\n", module, identifier);
                }
            }
            ImportStyle::QualifiedPaths => {
                let mut seen = HashSet::new();
                for import in imports {
                    let module = import.module();
                    if seen.insert(module) {
                        import_code += &format!("import {}\n", module);
                    }
                }
            }
        }
        if !import_code.is_empty(){
            import_code += "\n\n";
            code = import_code + &code;
        }
    }

    if ident != name {
        code += &format!("\n\n{} = {}", name, ident);
    }
    code
}

fn type_from_shape(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    use crate::shape::Shape::*;
    match shape {
        Null | Any | Bottom => (import(ctxt, Import::Any), None),
        Bool => ("bool".into(), None),
        StringT => ("str".into(), None),
        Integer => ("int".into(), None),
        Floating => ("float".into(), None),
        Tuple(shapes, _n) => {
            let folded = shape::fold_shapes(shapes.clone());
            if folded == Any && shapes.iter().any(|s| s != &Any) {
                generate_tuple_type(ctxt, path, shapes)
            } else {
                generate_vec_type(ctxt, path, &folded)
            }
        }
        VecT { elem_type: e } => generate_vec_type(ctxt, path, e),
        Struct { fields } => generate_data_class(ctxt, path, fields, shape),
        MapT { val_type: v } => generate_map_type(ctxt, path, v),
        Opaque(t) => (t.clone(), None),
        Optional(e) => {
            let (inner, defs) = type_from_shape(ctxt, path, e);
            if ctxt.options.use_default_for_missing_fields {
                (inner, defs)
            } else {
                let optional = import(ctxt, Import::Optional);
                (format!("{}[{}]", optional, inner), defs)
            }
        }
    }
}

fn generate_vec_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("list[{}]", inner), defs)
}

fn generate_map_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("dict[str, {}]", inner), defs)
}

fn generate_tuple_type(ctxt: &mut Ctxt, path: &str, shapes: &[Shape]) -> (Ident, Option<Code>) {
    let mut types = Vec::new();
    let mut defs = Vec::new();

    for shape in shapes {
        let (typ, def) = type_from_shape(ctxt, path, shape);
        types.push(typ);
        if let Some(code) = def {
            if !code.is_empty() {
                defs.push(code)
            }
        }
    }

    (
        format!("tuple[{}]", types.join(", ")),
        Some(defs.join("\n\n")),
    )
}

fn field_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "field", snake_case)
}

fn type_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "GeneratedType", type_case)
}

// https://docs.python.org/3/reference/lexical_analysis.html#keywords
#[rustfmt::skip]
const PYTHON_KEYWORDS: &[&str] = &[
    "False", "None", "True",
    "and", "as", "assert", "async", "await", "break", "class", "continue",
    "def", "del", "elif", "else", "except", "finally", "for", "from", "global",
    "if", "import", "in", "is", "lambda", "nonlocal", "not", "or", "pass",
    "raise", "return", "try", "while", "with", "yield",
];

fn type_or_field_name(
    name: &str,
    used_names: &HashSet<String>,
    default_name: &str,
    case_fn: fn(&str) -> String,
) -> Ident {
    let name = name.trim();
    let mut output_name = case_fn(name);
    if PYTHON_KEYWORDS.contains(&&*output_name) {
        output_name.push_str("_field");
    }
    if output_name.is_empty() {
        output_name.push_str(default_name);
    }
    if let Some(c) = output_name.chars().next() {
        if c.is_ascii_digit() {
            output_name = String::from("n") + &output_name;
        }
    }
    if !used_names.contains(&output_name) {
        return output_name;
    }
    for n in 2.. {
        let temp = format!("{}{}", output_name, n);
        if !used_names.contains(&temp) {
            return temp;
        }
    }
    unreachable!()
}

fn import(ctxt: &mut Ctxt, import: Import) -> String {
    ctxt.imports.insert(import);
    match ctxt.options.import_style {
        ImportStyle::QualifiedPaths => import.qualified(),
        _ => import.identifier().into(),
    }
}

fn generate_data_class(
    ctxt: &mut Ctxt,
    path: &str,
    field_shapes: &LinkedHashMap<String, Shape>,
    containing_shape: &Shape,
) -> (Ident, Option<Code>) {
    for (created_for_shape, ident) in ctxt.created_classes.iter() {
        if created_for_shape.is_acceptable_substitution_for(containing_shape) {
            return (ident.into(), None);
        }
    }

    let type_name = type_name(path, &ctxt.type_names);
    ctxt.type_names.insert(type_name.clone());
    ctxt.created_classes
        .push((containing_shape.clone(), type_name.clone()));

    let mut field_names = HashSet::new();
    let mut defs = Vec::new();

    let fields: Vec<Code> = field_shapes
        .iter()
        .map(|(name, typ)| {
            let field_name = field_name(name, &field_names);
            field_names.insert(field_name.clone());

            let (field_type, child_defs) = type_from_shape(ctxt, name, typ);

            if let Some(code) = child_defs {
                if !code.is_empty() {
                    defs.push(code);
                }
            }

            let mut field_code = String::new();
            let transformed = apply_transform(ctxt, &field_name, name);
            if transformed != field_name {
                field_code += &format!(" = {}(alias = \"{}\")", import(ctxt, Import::Field), transformed)
            }

            format!("    {}: {}{}", field_name, field_type, field_code)
        })
        .collect();

    let mut code = String::new();

    code += &format!(
        "class {}({}):\n",
        type_name,
        import(ctxt, Import::BaseModel)
    );

    if fields.is_empty() {
        code += "    pass\n";
    } else {
        code += &fields.join("\n");
        code += "\n";
    }

    if !defs.is_empty() {
        let mut d = defs.join("\n\n");
        d += "\n\n";
        d += &code;
        code = d;
    }

    (type_name, Some(code))
}

fn apply_transform(ctxt: &Ctxt, field_name: &str, name: &str) -> String {
    match ctxt.options.property_name_format {
        Some(StringTransform::LowerCase) => field_name.to_ascii_lowercase(),
        Some(StringTransform::PascalCase) => type_case(field_name),
        Some(StringTransform::SnakeCase) => snake_case(field_name),
        Some(StringTransform::KebabCase) => kebab_case(field_name),
        Some(StringTransform::UpperCase) => field_name.to_ascii_uppercase(),
        Some(StringTransform::CamelCase) => lower_camel_case(field_name),
        Some(StringTransform::ScreamingSnakeCase) => snake_case(field_name).to_ascii_uppercase(),
        Some(StringTransform::ScreamingKebabCase) => kebab_case(field_name).to_ascii_uppercase(),
        None => name.to_string(),
    }
}

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

    #[test]
    fn field_names_test() {
        fn field_name_test(from: &str, to: &str) {
            assert_eq!(
                field_name(from, &HashSet::new()),
                to.to_string(),
                r#"From "{}" to "{}""#,
                from,
                to
            );
        }

        field_name_test("valid", "valid");
        field_name_test("1", "n1");
        field_name_test("+1", "n1");
        field_name_test("", "field");
        field_name_test("def", "def_field");
    }
}
<|endoftext|>


<commit_before>//! A number of passes which remove various redundancies in the CFG.
//!
//! The `SimplifyCfg` pass gets rid of unnecessary blocks in the CFG, whereas the `SimplifyLocals`
//! gets rid of all the unnecessary local variable declarations.
//!
//! The `SimplifyLocals` pass is kinda expensive and therefore not very suitable to be run often.
//! Most of the passes should not care or be impacted in meaningful ways due to extra locals
//! either, so running the pass once, right before codegen, should suffice.
//!
//! On the other side of the spectrum, the `SimplifyCfg` pass is considerably cheap to run, thus
//! one should run it after every pass which may modify CFG in significant ways. This pass must
//! also be run before any analysis passes because it removes dead blocks, and some of these can be
//! ill-typed.
//!
//! The cause of this typing issue is typeck allowing most blocks whose end is not reachable have
//! an arbitrary return type, rather than having the usual () return type (as a note, typeck's
//! notion of reachability is in fact slightly weaker than MIR CFG reachability - see #31617). A
//! standard example of the situation is:
//!
//! ```rust
//!   fn example() {
//!       let _a: char = { return; };
//!   }
//! ```
//!
//! Here the block (`{ return; }`) has the return type `char`, rather than `()`, but the MIR we
//! naively generate still contains the `_a = ()` write in the unreachable block "after" the
//! return.

use crate::transform::{MirPass, MirSource};
use rustc::mir::visit::{MutVisitor, MutatingUseContext, PlaceContext, Visitor};
use rustc::mir::*;
use rustc::ty::{self, TyCtxt};
use rustc_index::bit_set::BitSet;
use rustc_index::vec::{Idx, IndexVec};
use std::borrow::Cow;

pub struct SimplifyCfg {
    label: String,
}

impl SimplifyCfg {
    pub fn new(label: &str) -> Self {
        SimplifyCfg { label: format!("SimplifyCfg-{}", label) }
    }
}

pub fn simplify_cfg(body: &mut BodyAndCache<'_>) {
    CfgSimplifier::new(body).simplify();
    remove_dead_blocks(body);

    // FIXME: Should probably be moved into some kind of pass manager
    body.basic_blocks_mut().raw.shrink_to_fit();
}

impl<'tcx> MirPass<'tcx> for SimplifyCfg {
    fn name(&self) -> Cow<'_, str> {
        Cow::Borrowed(&self.label)
    }

    fn run_pass(&self, _tcx: TyCtxt<'tcx>, _src: MirSource<'tcx>, body: &mut BodyAndCache<'tcx>) {
        debug!("SimplifyCfg({:?}) - simplifying {:?}", self.label, body);
        simplify_cfg(body);
    }
}

pub struct CfgSimplifier<'a, 'tcx> {
    basic_blocks: &'a mut IndexVec<BasicBlock, BasicBlockData<'tcx>>,
    pred_count: IndexVec<BasicBlock, u32>,
}

impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
    pub fn new(body: &'a mut BodyAndCache<'tcx>) -> Self {
        let mut pred_count = IndexVec::from_elem(0u32, body.basic_blocks());

        // we can't use mir.predecessors() here because that counts
        // dead blocks, which we don't want to.
        pred_count[START_BLOCK] = 1;

        for (_, data) in traversal::preorder(body) {
            if let Some(ref term) = data.terminator {
                for &tgt in term.successors() {
                    pred_count[tgt] += 1;
                }
            }
        }

        let basic_blocks = body.basic_blocks_mut();

        CfgSimplifier { basic_blocks, pred_count }
    }

    pub fn simplify(mut self) {
        self.strip_nops();

        let mut start = START_BLOCK;

        loop {
            let mut changed = false;

            self.collapse_goto_chain(&mut start, &mut changed);

            for bb in self.basic_blocks.indices() {
                if self.pred_count[bb] == 0 {
                    continue;
                }

                debug!("simplifying {:?}", bb);

                let mut terminator =
                    self.basic_blocks[bb].terminator.take().expect("invalid terminator state");

                for successor in terminator.successors_mut() {
                    self.collapse_goto_chain(successor, &mut changed);
                }

                let mut new_stmts = vec![];
                let mut inner_changed = true;
                while inner_changed {
                    inner_changed = false;
                    inner_changed |= self.simplify_branch(&mut terminator);
                    inner_changed |= self.merge_successor(&mut new_stmts, &mut terminator);
                    changed |= inner_changed;
                }

                let data = &mut self.basic_blocks[bb];
                data.statements.extend(new_stmts);
                data.terminator = Some(terminator);

                changed |= inner_changed;
            }

            if !changed {
                break;
            }
        }

        if start != START_BLOCK {
            debug_assert!(self.pred_count[START_BLOCK] == 0);
            self.basic_blocks.swap(START_BLOCK, start);
            self.pred_count.swap(START_BLOCK, start);

            // pred_count == 1 if the start block has no predecessor _blocks_.
            if self.pred_count[START_BLOCK] > 1 {
                for (bb, data) in self.basic_blocks.iter_enumerated_mut() {
                    if self.pred_count[bb] == 0 {
                        continue;
                    }

                    for target in data.terminator_mut().successors_mut() {
                        if *target == start {
                            *target = START_BLOCK;
                        }
                    }
                }
            }
        }
    }

    // Collapse a goto chain starting from `start`
    fn collapse_goto_chain(&mut self, start: &mut BasicBlock, changed: &mut bool) {
        let mut terminator = match self.basic_blocks[*start] {
            BasicBlockData {
                ref statements,
                terminator:
                    ref mut terminator @ Some(Terminator { kind: TerminatorKind::Goto { .. }, .. }),
                ..
            } if statements.is_empty() => terminator.take(),
            // if `terminator` is None, this means we are in a loop. In that
            // case, let all the loop collapse to its entry.
            _ => return,
        };

        let target = match terminator {
            Some(Terminator { kind: TerminatorKind::Goto { ref mut target }, .. }) => {
                self.collapse_goto_chain(target, changed);
                *target
            }
            _ => unreachable!(),
        };
        self.basic_blocks[*start].terminator = terminator;

        debug!("collapsing goto chain from {:?} to {:?}", *start, target);

        *changed |= *start != target;

        if self.pred_count[*start] == 1 {
            // This is the last reference to *start, so the pred-count to
            // to target is moved into the current block.
            self.pred_count[*start] = 0;
        } else {
            self.pred_count[target] += 1;
            self.pred_count[*start] -= 1;
        }

        *start = target;
    }

    // merge a block with 1 `goto` predecessor to its parent
    fn merge_successor(
        &mut self,
        new_stmts: &mut Vec<Statement<'tcx>>,
        terminator: &mut Terminator<'tcx>,
    ) -> bool {
        let target = match terminator.kind {
            TerminatorKind::Goto { target } if self.pred_count[target] == 1 => target,
            _ => return false,
        };

        debug!("merging block {:?} into {:?}", target, terminator);
        *terminator = match self.basic_blocks[target].terminator.take() {
            Some(terminator) => terminator,
            None => {
                // unreachable loop - this should not be possible, as we
                // don't strand blocks, but handle it correctly.
                return false;
            }
        };
        new_stmts.extend(self.basic_blocks[target].statements.drain(..));
        self.pred_count[target] = 0;

        true
    }

    // turn a branch with all successors identical to a goto
    fn simplify_branch(&mut self, terminator: &mut Terminator<'tcx>) -> bool {
        match terminator.kind {
            TerminatorKind::SwitchInt { .. } => {}
            _ => return false,
        };

        let first_succ = {
            if let Some(&first_succ) = terminator.successors().next() {
                if terminator.successors().all(|s| *s == first_succ) {
                    let count = terminator.successors().count();
                    self.pred_count[first_succ] -= (count - 1) as u32;
                    first_succ
                } else {
                    return false;
                }
            } else {
                return false;
            }
        };

        debug!("simplifying branch {:?}", terminator);
        terminator.kind = TerminatorKind::Goto { target: first_succ };
        true
    }

    fn strip_nops(&mut self) {
        for blk in self.basic_blocks.iter_mut() {
            blk.statements
                .retain(|stmt| if let StatementKind::Nop = stmt.kind { false } else { true })
        }
    }
}

pub fn remove_dead_blocks(body: &mut BodyAndCache<'_>) {
    let mut seen = BitSet::new_empty(body.basic_blocks().len());
    for (bb, _) in traversal::preorder(body) {
        seen.insert(bb.index());
    }

    let basic_blocks = body.basic_blocks_mut();

    let num_blocks = basic_blocks.len();
    let mut replacements: Vec<_> = (0..num_blocks).map(BasicBlock::new).collect();
    let mut used_blocks = 0;
    for alive_index in seen.iter() {
        replacements[alive_index] = BasicBlock::new(used_blocks);
        if alive_index != used_blocks {
            // Swap the next alive block data with the current available slot. Since
            // alive_index is non-decreasing this is a valid operation.
            basic_blocks.raw.swap(alive_index, used_blocks);
        }
        used_blocks += 1;
    }
    basic_blocks.raw.truncate(used_blocks);

    for block in basic_blocks {
        for target in block.terminator_mut().successors_mut() {
            *target = replacements[target.index()];
        }
    }
}

pub struct SimplifyLocals;

impl<'tcx> MirPass<'tcx> for SimplifyLocals {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut BodyAndCache<'tcx>) {
        trace!("running SimplifyLocals on {:?}", source);
        let locals = {
            let read_only_cache = read_only!(body);
            let mut marker = DeclMarker { locals: BitSet::new_empty(body.local_decls.len()), body };
            marker.visit_body(read_only_cache);
            // Return pointer and arguments are always live
            marker.locals.insert(RETURN_PLACE);
            for arg in body.args_iter() {
                marker.locals.insert(arg);
            }

            marker.locals
        };

        let map = make_local_map(&mut body.local_decls, locals);
        // Update references to all vars and tmps now
        LocalUpdater { map, tcx }.visit_body(body);
        body.local_decls.shrink_to_fit();
    }
}

/// Construct the mapping while swapping out unused stuff out from the `vec`.
fn make_local_map<V>(
    vec: &mut IndexVec<Local, V>,
    mask: BitSet<Local>,
) -> IndexVec<Local, Option<Local>> {
    let mut map: IndexVec<Local, Option<Local>> = IndexVec::from_elem(None, &*vec);
    let mut used = Local::new(0);
    for alive_index in mask.iter() {
        map[alive_index] = Some(used);
        if alive_index != used {
            vec.swap(alive_index, used);
        }
        used.increment_by(1);
    }
    vec.truncate(used.index());
    map
}

struct DeclMarker<'a, 'tcx> {
    pub locals: BitSet<Local>,
    pub body: &'a Body<'tcx>,
}

impl<'a, 'tcx> Visitor<'tcx> for DeclMarker<'a, 'tcx> {
    fn visit_local(&mut self, local: &Local, ctx: PlaceContext, location: Location) {
        // Ignore storage markers altogether, they get removed along with their otherwise unused
        // decls.
        // FIXME: Extend this to all non-uses.
        if ctx.is_storage_marker() {
            return;
        }

        // Ignore stores of constants because `ConstProp` and `CopyProp` can remove uses of many
        // of these locals. However, if the local is still needed, then it will be referenced in
        // another place and we'll mark it as being used there.
        if ctx == PlaceContext::MutatingUse(MutatingUseContext::Store)
            || ctx == PlaceContext::MutatingUse(MutatingUseContext::Projection)
        {
            let block = &self.body.basic_blocks()[location.block];
            if location.statement_index != block.statements.len() {
                let stmt = &block.statements[location.statement_index];

                if let StatementKind::Assign(box (p, Rvalue::Use(Operand::Constant(c)))) =
                    &stmt.kind
                {
                    match c.literal.val {
                        // Keep assignments from unevaluated constants around, since the evaluation
                        // may report errors, even if the use of the constant is dead code.
                        ty::ConstKind::Unevaluated(..) => {}
                        _ => {
                            if !p.is_indirect() {
                                trace!("skipping store of const value {:?} to {:?}", c, p);
                                return;
                            }
                        }
                    }
                }
            }
        }

        self.locals.insert(*local);
    }
}

struct LocalUpdater<'tcx> {
    map: IndexVec<Local, Option<Local>>,
    tcx: TyCtxt<'tcx>,
}

impl<'tcx> MutVisitor<'tcx> for LocalUpdater<'tcx> {
    fn tcx(&self) -> TyCtxt<'tcx> {
        self.tcx
    }

    fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) {
        // Remove unnecessary StorageLive and StorageDead annotations.
        data.statements.retain(|stmt| match &stmt.kind {
            StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => self.map[*l].is_some(),
            StatementKind::Assign(box (place, _)) => self.map[place.local].is_some(),
            _ => true,
        });
        self.super_basic_block_data(block, data);
    }

    fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) {
        *l = self.map[*l].unwrap();
    }

    fn process_projection_elem(&mut self, elem: &PlaceElem<'tcx>) -> Option<PlaceElem<'tcx>> {
        match elem {
            PlaceElem::Index(local) => Some(PlaceElem::Index(self.map[*local].unwrap())),
            _ => None,
        }
    }
}
<commit_msg>Auto merge of #68551 - Marwes:allocations_mir, r=ecstatic-morse<commit_after>//! A number of passes which remove various redundancies in the CFG.
//!
//! The `SimplifyCfg` pass gets rid of unnecessary blocks in the CFG, whereas the `SimplifyLocals`
//! gets rid of all the unnecessary local variable declarations.
//!
//! The `SimplifyLocals` pass is kinda expensive and therefore not very suitable to be run often.
//! Most of the passes should not care or be impacted in meaningful ways due to extra locals
//! either, so running the pass once, right before codegen, should suffice.
//!
//! On the other side of the spectrum, the `SimplifyCfg` pass is considerably cheap to run, thus
//! one should run it after every pass which may modify CFG in significant ways. This pass must
//! also be run before any analysis passes because it removes dead blocks, and some of these can be
//! ill-typed.
//!
//! The cause of this typing issue is typeck allowing most blocks whose end is not reachable have
//! an arbitrary return type, rather than having the usual () return type (as a note, typeck's
//! notion of reachability is in fact slightly weaker than MIR CFG reachability - see #31617). A
//! standard example of the situation is:
//!
//! ```rust
//!   fn example() {
//!       let _a: char = { return; };
//!   }
//! ```
//!
//! Here the block (`{ return; }`) has the return type `char`, rather than `()`, but the MIR we
//! naively generate still contains the `_a = ()` write in the unreachable block "after" the
//! return.

use crate::transform::{MirPass, MirSource};
use rustc::mir::visit::{MutVisitor, MutatingUseContext, PlaceContext, Visitor};
use rustc::mir::*;
use rustc::ty::{self, TyCtxt};
use rustc_index::bit_set::BitSet;
use rustc_index::vec::{Idx, IndexVec};
use std::borrow::Cow;

pub struct SimplifyCfg {
    label: String,
}

impl SimplifyCfg {
    pub fn new(label: &str) -> Self {
        SimplifyCfg { label: format!("SimplifyCfg-{}", label) }
    }
}

pub fn simplify_cfg(body: &mut BodyAndCache<'_>) {
    CfgSimplifier::new(body).simplify();
    remove_dead_blocks(body);

    // FIXME: Should probably be moved into some kind of pass manager
    body.basic_blocks_mut().raw.shrink_to_fit();
}

impl<'tcx> MirPass<'tcx> for SimplifyCfg {
    fn name(&self) -> Cow<'_, str> {
        Cow::Borrowed(&self.label)
    }

    fn run_pass(&self, _tcx: TyCtxt<'tcx>, _src: MirSource<'tcx>, body: &mut BodyAndCache<'tcx>) {
        debug!("SimplifyCfg({:?}) - simplifying {:?}", self.label, body);
        simplify_cfg(body);
    }
}

pub struct CfgSimplifier<'a, 'tcx> {
    basic_blocks: &'a mut IndexVec<BasicBlock, BasicBlockData<'tcx>>,
    pred_count: IndexVec<BasicBlock, u32>,
}

impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
    pub fn new(body: &'a mut BodyAndCache<'tcx>) -> Self {
        let mut pred_count = IndexVec::from_elem(0u32, body.basic_blocks());

        // we can't use mir.predecessors() here because that counts
        // dead blocks, which we don't want to.
        pred_count[START_BLOCK] = 1;

        for (_, data) in traversal::preorder(body) {
            if let Some(ref term) = data.terminator {
                for &tgt in term.successors() {
                    pred_count[tgt] += 1;
                }
            }
        }

        let basic_blocks = body.basic_blocks_mut();

        CfgSimplifier { basic_blocks, pred_count }
    }

    pub fn simplify(mut self) {
        self.strip_nops();

        let mut start = START_BLOCK;

        // Vec of the blocks that should be merged. We store the indices here, instead of the
        // statements itself to avoid moving the (relatively) large statements twice.
        // We do not push the statements directly into the target block (`bb`) as that is slower
        // due to additional reallocations
        let mut merged_blocks = Vec::new();
        loop {
            let mut changed = false;

            self.collapse_goto_chain(&mut start, &mut changed);

            for bb in self.basic_blocks.indices() {
                if self.pred_count[bb] == 0 {
                    continue;
                }

                debug!("simplifying {:?}", bb);

                let mut terminator =
                    self.basic_blocks[bb].terminator.take().expect("invalid terminator state");

                for successor in terminator.successors_mut() {
                    self.collapse_goto_chain(successor, &mut changed);
                }

                let mut inner_changed = true;
                merged_blocks.clear();
                while inner_changed {
                    inner_changed = false;
                    inner_changed |= self.simplify_branch(&mut terminator);
                    inner_changed |= self.merge_successor(&mut merged_blocks, &mut terminator);
                    changed |= inner_changed;
                }

                let statements_to_merge =
                    merged_blocks.iter().map(|&i| self.basic_blocks[i].statements.len()).sum();

                if statements_to_merge > 0 {
                    let mut statements = std::mem::take(&mut self.basic_blocks[bb].statements);
                    statements.reserve(statements_to_merge);
                    for &from in &merged_blocks {
                        statements.append(&mut self.basic_blocks[from].statements);
                    }
                    self.basic_blocks[bb].statements = statements;
                }

                self.basic_blocks[bb].terminator = Some(terminator);

                changed |= inner_changed;
            }

            if !changed {
                break;
            }
        }

        if start != START_BLOCK {
            debug_assert!(self.pred_count[START_BLOCK] == 0);
            self.basic_blocks.swap(START_BLOCK, start);
            self.pred_count.swap(START_BLOCK, start);

            // pred_count == 1 if the start block has no predecessor _blocks_.
            if self.pred_count[START_BLOCK] > 1 {
                for (bb, data) in self.basic_blocks.iter_enumerated_mut() {
                    if self.pred_count[bb] == 0 {
                        continue;
                    }

                    for target in data.terminator_mut().successors_mut() {
                        if *target == start {
                            *target = START_BLOCK;
                        }
                    }
                }
            }
        }
    }

    // Collapse a goto chain starting from `start`
    fn collapse_goto_chain(&mut self, start: &mut BasicBlock, changed: &mut bool) {
        let mut terminator = match self.basic_blocks[*start] {
            BasicBlockData {
                ref statements,
                terminator:
                    ref mut terminator @ Some(Terminator { kind: TerminatorKind::Goto { .. }, .. }),
                ..
            } if statements.is_empty() => terminator.take(),
            // if `terminator` is None, this means we are in a loop. In that
            // case, let all the loop collapse to its entry.
            _ => return,
        };

        let target = match terminator {
            Some(Terminator { kind: TerminatorKind::Goto { ref mut target }, .. }) => {
                self.collapse_goto_chain(target, changed);
                *target
            }
            _ => unreachable!(),
        };
        self.basic_blocks[*start].terminator = terminator;

        debug!("collapsing goto chain from {:?} to {:?}", *start, target);

        *changed |= *start != target;

        if self.pred_count[*start] == 1 {
            // This is the last reference to *start, so the pred-count to
            // to target is moved into the current block.
            self.pred_count[*start] = 0;
        } else {
            self.pred_count[target] += 1;
            self.pred_count[*start] -= 1;
        }

        *start = target;
    }

    // merge a block with 1 `goto` predecessor to its parent
    fn merge_successor(
        &mut self,
        merged_blocks: &mut Vec<BasicBlock>,
        terminator: &mut Terminator<'tcx>,
    ) -> bool {
        let target = match terminator.kind {
            TerminatorKind::Goto { target } if self.pred_count[target] == 1 => target,
            _ => return false,
        };

        debug!("merging block {:?} into {:?}", target, terminator);
        *terminator = match self.basic_blocks[target].terminator.take() {
            Some(terminator) => terminator,
            None => {
                // unreachable loop - this should not be possible, as we
                // don't strand blocks, but handle it correctly.
                return false;
            }
        };

        merged_blocks.push(target);
        self.pred_count[target] = 0;

        true
    }

    // turn a branch with all successors identical to a goto
    fn simplify_branch(&mut self, terminator: &mut Terminator<'tcx>) -> bool {
        match terminator.kind {
            TerminatorKind::SwitchInt { .. } => {}
            _ => return false,
        };

        let first_succ = {
            if let Some(&first_succ) = terminator.successors().next() {
                if terminator.successors().all(|s| *s == first_succ) {
                    let count = terminator.successors().count();
                    self.pred_count[first_succ] -= (count - 1) as u32;
                    first_succ
                } else {
                    return false;
                }
            } else {
                return false;
            }
        };

        debug!("simplifying branch {:?}", terminator);
        terminator.kind = TerminatorKind::Goto { target: first_succ };
        true
    }

    fn strip_nops(&mut self) {
        for blk in self.basic_blocks.iter_mut() {
            blk.statements
                .retain(|stmt| if let StatementKind::Nop = stmt.kind { false } else { true })
        }
    }
}

pub fn remove_dead_blocks(body: &mut BodyAndCache<'_>) {
    let mut seen = BitSet::new_empty(body.basic_blocks().len());
    for (bb, _) in traversal::preorder(body) {
        seen.insert(bb.index());
    }

    let basic_blocks = body.basic_blocks_mut();

    let num_blocks = basic_blocks.len();
    let mut replacements: Vec<_> = (0..num_blocks).map(BasicBlock::new).collect();
    let mut used_blocks = 0;
    for alive_index in seen.iter() {
        replacements[alive_index] = BasicBlock::new(used_blocks);
        if alive_index != used_blocks {
            // Swap the next alive block data with the current available slot. Since
            // alive_index is non-decreasing this is a valid operation.
            basic_blocks.raw.swap(alive_index, used_blocks);
        }
        used_blocks += 1;
    }
    basic_blocks.raw.truncate(used_blocks);

    for block in basic_blocks {
        for target in block.terminator_mut().successors_mut() {
            *target = replacements[target.index()];
        }
    }
}

pub struct SimplifyLocals;

impl<'tcx> MirPass<'tcx> for SimplifyLocals {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut BodyAndCache<'tcx>) {
        trace!("running SimplifyLocals on {:?}", source);
        let locals = {
            let read_only_cache = read_only!(body);
            let mut marker = DeclMarker { locals: BitSet::new_empty(body.local_decls.len()), body };
            marker.visit_body(read_only_cache);
            // Return pointer and arguments are always live
            marker.locals.insert(RETURN_PLACE);
            for arg in body.args_iter() {
                marker.locals.insert(arg);
            }

            marker.locals
        };

        let map = make_local_map(&mut body.local_decls, locals);
        // Update references to all vars and tmps now
        LocalUpdater { map, tcx }.visit_body(body);
        body.local_decls.shrink_to_fit();
    }
}

/// Construct the mapping while swapping out unused stuff out from the `vec`.
fn make_local_map<V>(
    vec: &mut IndexVec<Local, V>,
    mask: BitSet<Local>,
) -> IndexVec<Local, Option<Local>> {
    let mut map: IndexVec<Local, Option<Local>> = IndexVec::from_elem(None, &*vec);
    let mut used = Local::new(0);
    for alive_index in mask.iter() {
        map[alive_index] = Some(used);
        if alive_index != used {
            vec.swap(alive_index, used);
        }
        used.increment_by(1);
    }
    vec.truncate(used.index());
    map
}

struct DeclMarker<'a, 'tcx> {
    pub locals: BitSet<Local>,
    pub body: &'a Body<'tcx>,
}

impl<'a, 'tcx> Visitor<'tcx> for DeclMarker<'a, 'tcx> {
    fn visit_local(&mut self, local: &Local, ctx: PlaceContext, location: Location) {
        // Ignore storage markers altogether, they get removed along with their otherwise unused
        // decls.
        // FIXME: Extend this to all non-uses.
        if ctx.is_storage_marker() {
            return;
        }

        // Ignore stores of constants because `ConstProp` and `CopyProp` can remove uses of many
        // of these locals. However, if the local is still needed, then it will be referenced in
        // another place and we'll mark it as being used there.
        if ctx == PlaceContext::MutatingUse(MutatingUseContext::Store)
            || ctx == PlaceContext::MutatingUse(MutatingUseContext::Projection)
        {
            let block = &self.body.basic_blocks()[location.block];
            if location.statement_index != block.statements.len() {
                let stmt = &block.statements[location.statement_index];

                if let StatementKind::Assign(box (p, Rvalue::Use(Operand::Constant(c)))) =
                    &stmt.kind
                {
                    match c.literal.val {
                        // Keep assignments from unevaluated constants around, since the evaluation
                        // may report errors, even if the use of the constant is dead code.
                        ty::ConstKind::Unevaluated(..) => {}
                        _ => {
                            if !p.is_indirect() {
                                trace!("skipping store of const value {:?} to {:?}", c, p);
                                return;
                            }
                        }
                    }
                }
            }
        }

        self.locals.insert(*local);
    }
}

struct LocalUpdater<'tcx> {
    map: IndexVec<Local, Option<Local>>,
    tcx: TyCtxt<'tcx>,
}

impl<'tcx> MutVisitor<'tcx> for LocalUpdater<'tcx> {
    fn tcx(&self) -> TyCtxt<'tcx> {
        self.tcx
    }

    fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) {
        // Remove unnecessary StorageLive and StorageDead annotations.
        data.statements.retain(|stmt| match &stmt.kind {
            StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => self.map[*l].is_some(),
            StatementKind::Assign(box (place, _)) => self.map[place.local].is_some(),
            _ => true,
        });
        self.super_basic_block_data(block, data);
    }

    fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) {
        *l = self.map[*l].unwrap();
    }

    fn process_projection_elem(&mut self, elem: &PlaceElem<'tcx>) -> Option<PlaceElem<'tcx>> {
        match elem {
            PlaceElem::Index(local) => Some(PlaceElem::Index(self.map[*local].unwrap())),
            _ => None,
        }
    }
}
<|endoftext|>