<commit_before>use slack_api as slack;
use tokio_tungstenite as tungstenite;

use futures_util::{SinkExt, StreamExt};
use futures_timer::Delay;

use tokio::sync::mpsc;

use atomic_counter::RelaxedCounter;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use atomic_counter::AtomicCounter;

use std::sync::Arc;
use std::env;
use std::path::Path;
use std::result::Result;
use std::thread;
use std::sync::RwLock;

use chrono::prelude::{Utc, DateTime};
use std::time::Duration;
use std::time::Instant;

use karmator_rust::database;
use karmator_rust::event;
use karmator_rust::user_event;
use karmator_rust::cache;
use karmator_rust::signal;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let filename = env::var("SQLITE_FILE").map_err(|_| "SQLITE_FILE env var must be set")?;
    let token = env::var("SLACK_API_TOKEN").map_err(|_| "SLACK_API_TOKEN env var must be set")?;
    let client = slack::default_client().map_err(|e| format!("Could not get default_client, {:?}", e))?;

    // Shared integer counter for message ids
    let msg_id = Arc::new(RelaxedCounter::new(0));

    // Atomic boolean for ensuring send can't happen till the hello is recieved from slack
    let can_send = Arc::new(AtomicBool::new(false));

    // Atomic boolean for exiting and re-establishing the connection
    let reconnect = Arc::new(AtomicBool::new(false));
    let reconnect_count = Arc::new(RelaxedCounter::new(0));

    // Shutdown v2
    let (sql_shutdown_tx, mut signal) = signal::Signal::new();

    // Monotonical clock for heartbeat/ping management
    let last_message_recieved = Arc::new(RwLock::new(Instant::now()));
    let last_ping_sent = Arc::new(RwLock::new(Instant::now()));

    // Uptime of program start
    let start_time: DateTime<Utc> = Utc::now();

    // System Cache manager
    let cache = cache::Cache::new(&token, client.clone());

    // Message Tx/Rx from the message loop to the downstream handlers
    let (tx, mut rx) = mpsc::channel(32);

    // Sql request/reply channels for the downstream handlers to talk to the
    // sqlite worker thread
    let (sql_tx, sql_rx) = mpsc::channel(32);

    // Launch the sqlite worker thread
    let sql_worker = thread::Builder::new().name("sqlite_worker".into()).spawn(move || {
        println!("INFO [Sql Worker]: Launching");

        let res = database::process_queries(Path::new(&filename), sql_rx);
        eprintln!("ERROR [Sql Worker]: {:?}", res);

        // Worker died, signal the shutdown signal
        let _ = sql_shutdown_tx.broadcast(true);

        println!("INFO [Sql Worker]: Exiting");
    })?;

    // Main server loop, exit if the database dies
    while !signal.should_shutdown() {
        // Work to establish the WS connection
        let response = slack::rtm::connect(&client, &token).await.map_err(|e| format!("Control - {:?}", e))?;
        let ws_url = response.url.ok_or(format!("Control - \tNo Ws url"))?;
        println!("SYSTEM [Slack RTM]: Websocket Url: {:?}", ws_url);

        let (ws_stream, _) = tungstenite::connect_async(ws_url).await.map_err(|e| format!("Control - \t\t{:?}", e))?;
        println!("SYSTEM [Slack RTM]: Websocket connection established");

        // Split the stream
        let (mut ws_write, mut ws_read) = ws_stream.split();

        // Set the can-send to false till its set true by the reciever
        can_send.store(false, Ordering::Relaxed);

        // Set the reconnect to till the system control sets it to true
        reconnect.store(false, Ordering::Relaxed);

        while !reconnect.load(Ordering::Relaxed) && !signal.should_shutdown() {
            tokio::select! {
                // Listens for shutdown signals from the system or the database
                _ = signal.shutdown() => (),

                ws_msg = ws_read.next() => {
                    match ws_msg {
                        Some(Ok(ws_msg)) => {
                            // TODO: we may have some ordering funnyness if we spawn a processor per
                            // inbound message, ponder doing one processor per channel so its consistent
                            // order per channel, if want to be intelligent about it, have a timeout, if a
                            // channel has been quiet long enough, go ahead and shut down the channel processor.
                            //
                            // TODO: this could explode if the outstream or database get backed up it will just
                            // spawn more and more inbound message, not good.
                            let ue = event::process_control_message(
                                tx.clone(),
                                can_send.clone(),
                                reconnect.clone(),
                                reconnect_count.clone(),
                                last_message_recieved.clone(),
                                ws_msg,
                            ).await;

                            // If there's an user event returned spawn off the user event processor
                            if let Ok(Some(event)) = ue {
                                let msg_id = msg_id.clone();
                                let tx2 = tx.clone();
                                let sql_tx2 = sql_tx.clone();
                                let cache = cache.clone();

                                tokio::spawn(async move {
                                    // TODO: check result
                                    let _ = user_event::process_user_message(
                                        msg_id,
                                        event,
                                        tx2,
                                        sql_tx2,
                                        start_time,
                                        cache,
                                    ).await;
                                });
                            }
                        },
                        Some(Err(e)) => {
                            eprintln!("SYSTEM [Slack RTM]: Connection error: {:?}", e);
                            reconnect.store(true, Ordering::Relaxed);
                            can_send.store(false, Ordering::Relaxed);
                        },

                        // Stream is done/closed
                        None => {
                            println!("SYSTEM [Slack RTM]: Stream closed, reconnecting");
                            reconnect.store(true, Ordering::Relaxed);
                            can_send.store(false, Ordering::Relaxed);
                        },
                    }
                },

                // The send may not happen right away, the select! checks the (if x)
                // part before it checks the queue/etc
                // TODO: this would be better if instead of exiting the select loop we
                // gracefully drain the rx queue before exiting
                Some(message) = rx.recv(), if can_send.load(Ordering::Relaxed) => {
                    // TODO: present some way to do plain vs fancy message, and if
                    // fancy do a webapi post, otherwise dump into the WS
                    //
                    // TODO: look into tracking the sent message with a confirmation
                    // that it was sent (via msg id) and if it fails, resend with
                    // backoff
                    match ws_write.send(message).await {
                        Ok(_) => (),
                        Err(e) => {
                            eprintln!("SYSTEM [Slack RTM]: Connection error: {:?}", e);
                            reconnect.store(true, Ordering::Relaxed);
                            can_send.store(false, Ordering::Relaxed);
                        },
                    }
                },

                // This is woken up peroidically to force a heartbeat check
                _ = Delay::new(Duration::from_secs(1)) => {
                    let now = Instant::now();

                    let last_message_delta = {
                        let timer = last_message_recieved.read().unwrap();
                        now.checked_duration_since(*timer)
                    };

                    let last_ping_delta = {
                        let timer = last_ping_sent.read().unwrap();
                        now.checked_duration_since(*timer)
                    };

                    // Check the delta, If either or both are None, that means
                    // their timer are *ahead* of the 'now' which is fine, stop checking.
                    match (last_message_delta, last_ping_delta) {
                        (Some(lmd), Some(lpd)) => {
                            // Check if more than 30s has past since the last slack message recieved
                            //  - [Yes] Check if more than 30s has past since the send of the last ping
                            //    - [Yes] Send Ping
                            //    - [No] Do nothing
                            //  - [No] Do nothing
                            //
                            // Check if more than 2m has past since the last slack message recieved
                            //  - [Yes] Reconnect
                            //  - [No] Do nothing
                            if lmd.as_secs() > 30 {
                                if lpd.as_secs() > 30 {
                                    println!("SYSTEM [Slack RTM]: Last message: {:?}s, Last ping: {:?}s",
                                        lmd.as_secs(),
                                        lpd.as_secs(),
                                    );
                                    let _ = event::send_slack_ping(
                                        msg_id.clone(),
                                        &mut tx.clone(),
                                        last_ping_sent.clone(),
                                    ).await;
                                }
                            } else if lmd.as_secs() > 120 {
                                eprintln!("SYSTEM [Slack RTM]: Last message: {:?}s, reconnecting", lmd.as_secs());
                                reconnect.store(true, Ordering::Relaxed);
                                can_send.store(false, Ordering::Relaxed);
                            }
                        },
                        _ => (),
                    }
                },
            }
        }

        // We exited the inner loop, wait here and do an exp backoff before trying to reconnect
        // wait for 20s, then increment the reconnect_count by 1, and if it exceeds 10, set the
        // shutdown flag and exit
        if !signal.should_shutdown() {
            let count = reconnect_count.clone().inc();

            if count <= 10 {
                println!("SYSTEM [Slack RTM]: Reconnecting, try: {:?}", count);
                Delay::new(Duration::from_secs(20)).await;
            } else {
                eprintln!("ERROR [Slack RTM]: Exceeded 10 retries, shutting down");
                signal.shutdown_now();
                break;
            }
        }
    }

    // Force drop the sender (since all other sender clone should be dropped by now)
    drop(sql_tx);
    let res = sql_worker.join();
    println!("SYSTEM [Sql Worker]: Thread Join: {:?}", res);

    Ok(())
}
<commit_msg>todos<commit_after>use slack_api as slack;
use tokio_tungstenite as tungstenite;

use futures_util::{SinkExt, StreamExt};
use futures_timer::Delay;

use tokio::sync::mpsc;

use atomic_counter::RelaxedCounter;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use atomic_counter::AtomicCounter;

use std::sync::Arc;
use std::env;
use std::path::Path;
use std::result::Result;
use std::thread;
use std::sync::RwLock;

use chrono::prelude::{Utc, DateTime};
use std::time::Duration;
use std::time::Instant;

use karmator_rust::database;
use karmator_rust::event;
use karmator_rust::user_event;
use karmator_rust::cache;
use karmator_rust::signal;

// TODO:
// Fix voting for - @luisp++ case it should -> luisp++ for the database
// Fix command parsing for <@user> -> user for commands

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let filename = env::var("SQLITE_FILE").map_err(|_| "SQLITE_FILE env var must be set")?;
    let token = env::var("SLACK_API_TOKEN").map_err(|_| "SLACK_API_TOKEN env var must be set")?;
    let client = slack::default_client().map_err(|e| format!("Could not get default_client, {:?}", e))?;

    // Shared integer counter for message ids
    let msg_id = Arc::new(RelaxedCounter::new(0));

    // Atomic boolean for ensuring send can't happen till the hello is recieved from slack
    let can_send = Arc::new(AtomicBool::new(false));

    // Atomic boolean for exiting and re-establishing the connection
    let reconnect = Arc::new(AtomicBool::new(false));
    let reconnect_count = Arc::new(RelaxedCounter::new(0));

    // Shutdown v2
    let (sql_shutdown_tx, mut signal) = signal::Signal::new();

    // Monotonical clock for heartbeat/ping management
    let last_message_recieved = Arc::new(RwLock::new(Instant::now()));
    let last_ping_sent = Arc::new(RwLock::new(Instant::now()));

    // Uptime of program start
    let start_time: DateTime<Utc> = Utc::now();

    // System Cache manager
    let cache = cache::Cache::new(&token, client.clone());

    // Message Tx/Rx from the message loop to the downstream handlers
    let (tx, mut rx) = mpsc::channel(32);

    // Sql request/reply channels for the downstream handlers to talk to the
    // sqlite worker thread
    let (sql_tx, sql_rx) = mpsc::channel(32);

    // Launch the sqlite worker thread
    let sql_worker = thread::Builder::new().name("sqlite_worker".into()).spawn(move || {
        println!("INFO [Sql Worker]: Launching");

        let res = database::process_queries(Path::new(&filename), sql_rx);
        eprintln!("ERROR [Sql Worker]: {:?}", res);

        // Worker died, signal the shutdown signal
        let _ = sql_shutdown_tx.broadcast(true);

        println!("INFO [Sql Worker]: Exiting");
    })?;

    // Main server loop, exit if the database dies
    while !signal.should_shutdown() {
        // Work to establish the WS connection
        let response = slack::rtm::connect(&client, &token).await.map_err(|e| format!("Control - {:?}", e))?;
        let ws_url = response.url.ok_or(format!("Control - \tNo Ws url"))?;
        println!("SYSTEM [Slack RTM]: Websocket Url: {:?}", ws_url);

        let (ws_stream, _) = tungstenite::connect_async(ws_url).await.map_err(|e| format!("Control - \t\t{:?}", e))?;
        println!("SYSTEM [Slack RTM]: Websocket connection established");

        // Split the stream
        let (mut ws_write, mut ws_read) = ws_stream.split();

        // Set the can-send to false till its set true by the reciever
        can_send.store(false, Ordering::Relaxed);

        // Set the reconnect to till the system control sets it to true
        reconnect.store(false, Ordering::Relaxed);

        while !reconnect.load(Ordering::Relaxed) && !signal.should_shutdown() {
            tokio::select! {
                // Listens for shutdown signals from the system or the database
                _ = signal.shutdown() => (),

                ws_msg = ws_read.next() => {
                    match ws_msg {
                        Some(Ok(ws_msg)) => {
                            // TODO: we may have some ordering funnyness if we spawn a processor per
                            // inbound message, ponder doing one processor per channel so its consistent
                            // order per channel, if want to be intelligent about it, have a timeout, if a
                            // channel has been quiet long enough, go ahead and shut down the channel processor.
                            //
                            // TODO: this could explode if the outstream or database get backed up it will just
                            // spawn more and more inbound message, not good.
                            let ue = event::process_control_message(
                                tx.clone(),
                                can_send.clone(),
                                reconnect.clone(),
                                reconnect_count.clone(),
                                last_message_recieved.clone(),
                                ws_msg,
                            ).await;

                            // If there's an user event returned spawn off the user event processor
                            if let Ok(Some(event)) = ue {
                                let msg_id = msg_id.clone();
                                let tx2 = tx.clone();
                                let sql_tx2 = sql_tx.clone();
                                let cache = cache.clone();

                                tokio::spawn(async move {
                                    // TODO: check result
                                    let _ = user_event::process_user_message(
                                        msg_id,
                                        event,
                                        tx2,
                                        sql_tx2,
                                        start_time,
                                        cache,
                                    ).await;
                                });
                            }
                        },
                        Some(Err(e)) => {
                            eprintln!("SYSTEM [Slack RTM]: Connection error: {:?}", e);
                            reconnect.store(true, Ordering::Relaxed);
                            can_send.store(false, Ordering::Relaxed);
                        },

                        // Stream is done/closed
                        None => {
                            println!("SYSTEM [Slack RTM]: Stream closed, reconnecting");
                            reconnect.store(true, Ordering::Relaxed);
                            can_send.store(false, Ordering::Relaxed);
                        },
                    }
                },

                // The send may not happen right away, the select! checks the (if x)
                // part before it checks the queue/etc
                // TODO: this would be better if instead of exiting the select loop we
                // gracefully drain the rx queue before exiting
                Some(message) = rx.recv(), if can_send.load(Ordering::Relaxed) => {
                    // TODO: present some way to do plain vs fancy message, and if
                    // fancy do a webapi post, otherwise dump into the WS
                    //
                    // TODO: look into tracking the sent message with a confirmation
                    // that it was sent (via msg id) and if it fails, resend with
                    // backoff
                    match ws_write.send(message).await {
                        Ok(_) => (),
                        Err(e) => {
                            eprintln!("SYSTEM [Slack RTM]: Connection error: {:?}", e);
                            reconnect.store(true, Ordering::Relaxed);
                            can_send.store(false, Ordering::Relaxed);
                        },
                    }
                },

                // This is woken up peroidically to force a heartbeat check
                _ = Delay::new(Duration::from_secs(1)) => {
                    let now = Instant::now();

                    let last_message_delta = {
                        let timer = last_message_recieved.read().unwrap();
                        now.checked_duration_since(*timer)
                    };

                    let last_ping_delta = {
                        let timer = last_ping_sent.read().unwrap();
                        now.checked_duration_since(*timer)
                    };

                    // Check the delta, If either or both are None, that means
                    // their timer are *ahead* of the 'now' which is fine, stop checking.
                    match (last_message_delta, last_ping_delta) {
                        (Some(lmd), Some(lpd)) => {
                            // Check if more than 30s has past since the last slack message recieved
                            //  - [Yes] Check if more than 30s has past since the send of the last ping
                            //    - [Yes] Send Ping
                            //    - [No] Do nothing
                            //  - [No] Do nothing
                            //
                            // Check if more than 2m has past since the last slack message recieved
                            //  - [Yes] Reconnect
                            //  - [No] Do nothing
                            if lmd.as_secs() > 30 {
                                if lpd.as_secs() > 30 {
                                    println!("SYSTEM [Slack RTM]: Last message: {:?}s, Last ping: {:?}s",
                                        lmd.as_secs(),
                                        lpd.as_secs(),
                                    );
                                    let _ = event::send_slack_ping(
                                        msg_id.clone(),
                                        &mut tx.clone(),
                                        last_ping_sent.clone(),
                                    ).await;
                                }
                            } else if lmd.as_secs() > 120 {
                                eprintln!("SYSTEM [Slack RTM]: Last message: {:?}s, reconnecting", lmd.as_secs());
                                reconnect.store(true, Ordering::Relaxed);
                                can_send.store(false, Ordering::Relaxed);
                            }
                        },
                        _ => (),
                    }
                },
            }
        }

        // We exited the inner loop, wait here and do an exp backoff before trying to reconnect
        // wait for 20s, then increment the reconnect_count by 1, and if it exceeds 10, set the
        // shutdown flag and exit
        if !signal.should_shutdown() {
            let count = reconnect_count.clone().inc();

            if count <= 10 {
                println!("SYSTEM [Slack RTM]: Reconnecting, try: {:?}", count);
                Delay::new(Duration::from_secs(20)).await;
            } else {
                eprintln!("ERROR [Slack RTM]: Exceeded 10 retries, shutting down");
                signal.shutdown_now();
                break;
            }
        }
    }

    // Force drop the sender (since all other sender clone should be dropped by now)
    drop(sql_tx);
    let res = sql_worker.join();
    println!("SYSTEM [Sql Worker]: Thread Join: {:?}", res);

    Ok(())
}
<|endoftext|>


<commit_before>//! Intel Instruction Set Architectures.

pub mod settings;
mod abi;
mod binemit;
mod enc_tables;
mod registers;

use binemit::{CodeSink, MemoryCodeSink, emit_function};
use super::super::settings as shared_settings;
use isa::enc_tables::{self as shared_enc_tables, lookup_enclist, Encodings};
use isa::Builder as IsaBuilder;
use isa::{TargetIsa, RegInfo, RegClass, EncInfo, RegUnit};
use self::registers::RU;
use ir;
use regalloc;
use result;
use ir::{InstBuilder, InstructionData, Opcode};
use stack_layout::layout_stack;
use cursor::{Cursor, EncCursor};

#[allow(dead_code)]
struct Isa {
    shared_flags: shared_settings::Flags,
    isa_flags: settings::Flags,
    cpumode: &'static [shared_enc_tables::Level1Entry<u16>],
}

/// Get an ISA builder for creating Intel targets.
pub fn isa_builder() -> IsaBuilder {
    IsaBuilder {
        setup: settings::builder(),
        constructor: isa_constructor,
    }
}

fn isa_constructor(
    shared_flags: shared_settings::Flags,
    builder: &shared_settings::Builder,
) -> Box<TargetIsa> {
    let level1 = if shared_flags.is_64bit() {
        &enc_tables::LEVEL1_I64[..]
    } else {
        &enc_tables::LEVEL1_I32[..]
    };
    Box::new(Isa {
        isa_flags: settings::Flags::new(&shared_flags, builder),
        shared_flags,
        cpumode: level1,
    })
}

impl TargetIsa for Isa {
    fn name(&self) -> &'static str {
        "intel"
    }

    fn flags(&self) -> &shared_settings::Flags {
        &self.shared_flags
    }

    fn register_info(&self) -> RegInfo {
        registers::INFO.clone()
    }

    fn encoding_info(&self) -> EncInfo {
        enc_tables::INFO.clone()
    }

    fn legal_encodings<'a>(
        &'a self,
        dfg: &'a ir::DataFlowGraph,
        inst: &'a ir::InstructionData,
        ctrl_typevar: ir::Type,
    ) -> Encodings<'a> {
        lookup_enclist(
            ctrl_typevar,
            inst,
            dfg,
            self.cpumode,
            &enc_tables::LEVEL2[..],
            &enc_tables::ENCLISTS[..],
            &enc_tables::LEGALIZE_ACTIONS[..],
            &enc_tables::RECIPE_PREDICATES[..],
            &enc_tables::INST_PREDICATES[..],
            self.isa_flags.predicate_view(),
        )
    }

    fn legalize_signature(&self, sig: &mut ir::Signature, current: bool) {
        abi::legalize_signature(sig, &self.shared_flags, current)
    }

    fn regclass_for_abi_type(&self, ty: ir::Type) -> RegClass {
        abi::regclass_for_abi_type(ty)
    }

    fn allocatable_registers(&self, func: &ir::Function) -> regalloc::AllocatableSet {
        abi::allocatable_registers(func, &self.shared_flags)
    }

    fn emit_inst(
        &self,
        func: &ir::Function,
        inst: ir::Inst,
        divert: &mut regalloc::RegDiversions,
        sink: &mut CodeSink,
    ) {
        binemit::emit_inst(func, inst, divert, sink)
    }

    fn emit_function(&self, func: &ir::Function, sink: &mut MemoryCodeSink) {
        emit_function(func, binemit::emit_inst, sink)
    }

    fn reloc_names(&self) -> &'static [&'static str] {
        &binemit::RELOC_NAMES
    }

    fn prologue_epilogue(&self, func: &mut ir::Function) -> result::CtonResult {
        let word_size = if self.flags().is_64bit() { 8 } else { 4 };

        let mut csr_stack_size = word_size; // Size of RBP to start with
        for _reg in abi::CSR_GPRS.iter() {
            csr_stack_size += word_size;
        }

        let stack_offset = -(csr_stack_size as i32);
        let slot = ir::StackSlotData {
            kind: ir::StackSlotKind::IncomingArg,
            size: csr_stack_size,
            offset: stack_offset,
        };
        func.create_stack_slot(slot);

        let total_stack_size = layout_stack(&mut func.stack_slots, word_size)?;
        let local_stack_size = total_stack_size - csr_stack_size;

        // Append frame pointer to function signature
        let rbp_arg = ir::AbiParam::special_reg(
            ir::types::I64,
            ir::ArgumentPurpose::FramePointer,
            RU::rbp as RegUnit,
        );
        func.signature.params.push(rbp_arg);
        func.signature.returns.push(rbp_arg);

        for reg in abi::CSR_GPRS.iter() {
            let csr_arg = ir::AbiParam::special_reg(
                ir::types::I64,
                ir::ArgumentPurpose::CalleeSaved,
                *reg as RegUnit,
            );
            func.signature.params.push(csr_arg);
            func.signature.returns.push(csr_arg);
        }

        // Append param to entry EBB
        let entry_ebb = func.layout.entry_block().expect("missing entry block");
        func.dfg.append_ebb_param(entry_ebb, ir::types::I64);

        // Find our frame pointer parameter Value
        let fp = func.special_param(ir::ArgumentPurpose::FramePointer)
            .expect("missing frame pointer");

        // Assign it a location
        func.locations[fp] = ir::ValueLoc::Reg(RU::rbp as RegUnit);

        let mut csr_vals = Vec::new();
        for reg in abi::CSR_GPRS.iter() {
            // Append param to entry EBB
            func.dfg.append_ebb_param(entry_ebb, ir::types::I64);

            let csr_arg = func.dfg.ebb_params(entry_ebb).last().expect(
                "no last argument",
            );

            // Assign it a location
            func.locations[*csr_arg] = ir::ValueLoc::Reg(*reg as RegUnit);

            // Remember it so we can push it momentarily
            csr_vals.push(*csr_arg);
        }

        // Insert prologue
        {
            let mut pos = EncCursor::new(func, self).at_first_insertion_point(entry_ebb);
            pos.ins().x86_push(fp);
            pos.ins().copy_special(
                RU::rsp as RegUnit,
                RU::rbp as RegUnit,
            );
            if local_stack_size > 0 {
                pos.ins().adjust_sp_imm(-(local_stack_size as i32));
            }

            for csr_arg in csr_vals {
                pos.ins().x86_push(csr_arg);
            }
        }

        // Find all 'return' instructions
        let mut return_insts = Vec::new();
        for ebb in func.layout.ebbs() {
            for inst in func.layout.ebb_insts(ebb) {
                if let InstructionData::MultiAry { opcode, .. } = func.dfg[inst] {
                    if opcode == Opcode::Return {
                        return_insts.push(inst);
                    }
                }

            }
        }

        // Insert an epilogue directly before every 'return'
        for inst in return_insts {
            let fp_ret = self.insert_epilogue(inst, local_stack_size as i32, func);
            func.locations[fp_ret] = ir::ValueLoc::Reg(RU::rbp as RegUnit);
            func.dfg.append_inst_arg(inst, fp_ret);
        }

        Ok(())
    }
}

impl Isa {
    fn insert_epilogue(
        &self,
        inst: ir::Inst,
        stack_size: i32,
        func: &mut ir::Function,
    ) -> ir::Value {
        let mut pos = EncCursor::new(func, self).at_inst(inst);
        if stack_size > 0 {
            pos.ins().adjust_sp_imm(stack_size);
        }
        pos.ins().x86_pop(ir::types::I64)
    }
}
<commit_msg>Pop the callee-saved registers in the function epilogue(s).<commit_after>//! Intel Instruction Set Architectures.

pub mod settings;
mod abi;
mod binemit;
mod enc_tables;
mod registers;

use binemit::{CodeSink, MemoryCodeSink, emit_function};
use super::super::settings as shared_settings;
use isa::enc_tables::{self as shared_enc_tables, lookup_enclist, Encodings};
use isa::Builder as IsaBuilder;
use isa::{TargetIsa, RegInfo, RegClass, EncInfo, RegUnit};
use self::registers::RU;
use ir;
use regalloc;
use result;
use ir::{InstBuilder, InstructionData, Opcode};
use stack_layout::layout_stack;
use cursor::{Cursor, EncCursor};

#[allow(dead_code)]
struct Isa {
    shared_flags: shared_settings::Flags,
    isa_flags: settings::Flags,
    cpumode: &'static [shared_enc_tables::Level1Entry<u16>],
}

/// Get an ISA builder for creating Intel targets.
pub fn isa_builder() -> IsaBuilder {
    IsaBuilder {
        setup: settings::builder(),
        constructor: isa_constructor,
    }
}

fn isa_constructor(
    shared_flags: shared_settings::Flags,
    builder: &shared_settings::Builder,
) -> Box<TargetIsa> {
    let level1 = if shared_flags.is_64bit() {
        &enc_tables::LEVEL1_I64[..]
    } else {
        &enc_tables::LEVEL1_I32[..]
    };
    Box::new(Isa {
        isa_flags: settings::Flags::new(&shared_flags, builder),
        shared_flags,
        cpumode: level1,
    })
}

impl TargetIsa for Isa {
    fn name(&self) -> &'static str {
        "intel"
    }

    fn flags(&self) -> &shared_settings::Flags {
        &self.shared_flags
    }

    fn register_info(&self) -> RegInfo {
        registers::INFO.clone()
    }

    fn encoding_info(&self) -> EncInfo {
        enc_tables::INFO.clone()
    }

    fn legal_encodings<'a>(
        &'a self,
        dfg: &'a ir::DataFlowGraph,
        inst: &'a ir::InstructionData,
        ctrl_typevar: ir::Type,
    ) -> Encodings<'a> {
        lookup_enclist(
            ctrl_typevar,
            inst,
            dfg,
            self.cpumode,
            &enc_tables::LEVEL2[..],
            &enc_tables::ENCLISTS[..],
            &enc_tables::LEGALIZE_ACTIONS[..],
            &enc_tables::RECIPE_PREDICATES[..],
            &enc_tables::INST_PREDICATES[..],
            self.isa_flags.predicate_view(),
        )
    }

    fn legalize_signature(&self, sig: &mut ir::Signature, current: bool) {
        abi::legalize_signature(sig, &self.shared_flags, current)
    }

    fn regclass_for_abi_type(&self, ty: ir::Type) -> RegClass {
        abi::regclass_for_abi_type(ty)
    }

    fn allocatable_registers(&self, func: &ir::Function) -> regalloc::AllocatableSet {
        abi::allocatable_registers(func, &self.shared_flags)
    }

    fn emit_inst(
        &self,
        func: &ir::Function,
        inst: ir::Inst,
        divert: &mut regalloc::RegDiversions,
        sink: &mut CodeSink,
    ) {
        binemit::emit_inst(func, inst, divert, sink)
    }

    fn emit_function(&self, func: &ir::Function, sink: &mut MemoryCodeSink) {
        emit_function(func, binemit::emit_inst, sink)
    }

    fn reloc_names(&self) -> &'static [&'static str] {
        &binemit::RELOC_NAMES
    }

    fn prologue_epilogue(&self, func: &mut ir::Function) -> result::CtonResult {
        let word_size = if self.flags().is_64bit() { 8 } else { 4 };

        let mut csr_stack_size = word_size; // Size of RBP to start with
        for _reg in abi::CSR_GPRS.iter() {
            csr_stack_size += word_size;
        }

        let stack_offset = -(csr_stack_size as i32);
        let slot = ir::StackSlotData {
            kind: ir::StackSlotKind::IncomingArg,
            size: csr_stack_size,
            offset: stack_offset,
        };
        func.create_stack_slot(slot);

        let total_stack_size = layout_stack(&mut func.stack_slots, word_size)?;
        let local_stack_size = total_stack_size - csr_stack_size;

        // Append frame pointer to function signature
        let rbp_arg = ir::AbiParam::special_reg(
            ir::types::I64,
            ir::ArgumentPurpose::FramePointer,
            RU::rbp as RegUnit,
        );
        func.signature.params.push(rbp_arg);
        func.signature.returns.push(rbp_arg);

        for reg in abi::CSR_GPRS.iter() {
            let csr_arg = ir::AbiParam::special_reg(
                ir::types::I64,
                ir::ArgumentPurpose::CalleeSaved,
                *reg as RegUnit,
            );
            func.signature.params.push(csr_arg);
            func.signature.returns.push(csr_arg);
        }

        // Append param to entry EBB
        let entry_ebb = func.layout.entry_block().expect("missing entry block");
        func.dfg.append_ebb_param(entry_ebb, ir::types::I64);

        // Find our frame pointer parameter Value
        let fp = func.special_param(ir::ArgumentPurpose::FramePointer)
            .expect("missing frame pointer");

        // Assign it a location
        func.locations[fp] = ir::ValueLoc::Reg(RU::rbp as RegUnit);

        let mut csr_vals = Vec::new();
        for reg in abi::CSR_GPRS.iter() {
            // Append param to entry EBB
            func.dfg.append_ebb_param(entry_ebb, ir::types::I64);

            let csr_arg = func.dfg.ebb_params(entry_ebb).last().expect(
                "no last argument",
            );

            // Assign it a location
            func.locations[*csr_arg] = ir::ValueLoc::Reg(*reg as RegUnit);

            // Remember it so we can push it momentarily
            csr_vals.push(*csr_arg);
        }

        // Insert prologue
        {
            let mut pos = EncCursor::new(func, self).at_first_insertion_point(entry_ebb);
            pos.ins().x86_push(fp);
            pos.ins().copy_special(
                RU::rsp as RegUnit,
                RU::rbp as RegUnit,
            );
            if local_stack_size > 0 {
                pos.ins().adjust_sp_imm(-(local_stack_size as i32));
            }

            for csr_arg in csr_vals {
                pos.ins().x86_push(csr_arg);
            }
        }

        // Find all 'return' instructions
        let mut return_insts = Vec::new();
        for ebb in func.layout.ebbs() {
            for inst in func.layout.ebb_insts(ebb) {
                if let InstructionData::MultiAry { opcode, .. } = func.dfg[inst] {
                    if opcode == Opcode::Return {
                        return_insts.push(inst);
                    }
                }

            }
        }

        // Insert an epilogue directly before every 'return'
        for inst in return_insts {
            self.insert_epilogue(inst, local_stack_size as i32, func);
        }

        Ok(())
    }
}

impl Isa {
    fn insert_epilogue(&self, inst: ir::Inst, stack_size: i32, func: &mut ir::Function) {
        let mut return_values = Vec::new();

        let mut pos = EncCursor::new(func, self).at_inst(inst);
        if stack_size > 0 {
            pos.ins().adjust_sp_imm(stack_size);
        }
        for reg in abi::CSR_GPRS.iter().rev() {
            let csr_ret = pos.ins().x86_pop(ir::types::I64);
            return_values.push((csr_ret, *reg));
        }
        let fp_ret = pos.ins().x86_pop(ir::types::I64);
        return_values.push((fp_ret, RU::rbp));

        let func = pos.func;
        for (val, reg) in return_values {
            func.locations[val] = ir::ValueLoc::Reg(reg as RegUnit);
            func.dfg.append_inst_arg(inst, val);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2020 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate proc_macro;

use quote::{quote, ToTokens};
use std::borrow::Cow;
use synstructure::{decl_derive, Structure, VariantInfo};

macro_rules! syn_throw {
    ($err:expr) => {
        return syn::Error::to_compile_error(&$err);
    };
}

macro_rules! syn_try {
    ($expr:expr) => {
        match $expr {
            Ok(expr) => expr,
            Err(err) => syn_throw!(err),
        }
    };
}

fn discriminant<'v>(v: &VariantInfo<'v>) -> syn::Result<Option<Cow<'v, syn::Expr>>> {
    v.ast()
        .discriminant
        .iter()
        .map(|(_, discriminant)| Ok(Cow::Borrowed(discriminant)))
        .chain(v.ast().attrs.iter().filter_map(|attr| match attr {
            syn::Attribute {
                style: syn::AttrStyle::Outer,
                path,
                ..
            } if path.is_ident("wasmbin") => {
                syn::custom_keyword!(discriminant);

                Some(
                    attr.parse_args_with(|parser: syn::parse::ParseStream| {
                        parser.parse::<discriminant>()?;
                        parser.parse::<syn::Token![=]>()?;
                        parser.parse()
                    })
                    .map(Cow::Owned),
                )
            }
            _ => None,
        }))
        .try_fold(None, |prev, discriminant| {
            let discriminant = discriminant?;
            if let Some(prev) = prev {
                let mut err = syn::Error::new_spanned(
                    discriminant,
                    "#[derive(Wasmbin)]: duplicate discriminant",
                );
                err.combine(syn::Error::new_spanned(
                    prev,
                    "#[derive(Wasmbin)]: previous discriminant here",
                ));
                return Err(err);
            }
            Ok(Some(discriminant))
        })
}

fn gen_encode_discriminant(repr: &syn::Type, discriminant: &syn::Expr) -> proc_macro2::TokenStream {
    quote!(<#repr as Encode>::encode(&#discriminant, w)?)
}

fn is_newtype_like(v: &VariantInfo) -> bool {
    matches!(v.ast().fields, fields @ syn::Fields::Unnamed(_) if fields.len() == 1)
}

fn track_err_in_field(
    mut res: proc_macro2::TokenStream,
    v: &VariantInfo,
    field: &syn::Field,
    index: usize,
) -> proc_macro2::TokenStream {
    if !is_newtype_like(v) {
        let field_name = match &field.ident {
            Some(ident) => ident.to_string(),
            None => index.to_string(),
        };
        res = quote!(#res.map_err(|err| err.in_path(PathItem::Name(#field_name))));
    }
    res
}

fn track_err_in_variant(
    res: proc_macro2::TokenStream,
    v: &VariantInfo,
) -> proc_macro2::TokenStream {
    use std::fmt::Write;

    let mut variant_name = String::new();
    if let Some(prefix) = v.prefix {
        write!(variant_name, "{}::", prefix).unwrap();
    }
    write!(variant_name, "{}", v.ast().ident).unwrap();

    quote!(#res.map_err(|err| err.in_path(PathItem::Variant(#variant_name))))
}

fn catch_expr(
    res: proc_macro2::TokenStream,
    err: proc_macro2::TokenStream,
) -> proc_macro2::TokenStream {
    quote!(
        (move || -> Result<_, #err> {
            Ok({ #res })
        })()
    )
}

fn gen_decode(v: &VariantInfo) -> proc_macro2::TokenStream {
    let mut res = v.construct(|field, index| {
        let res = track_err_in_field(quote!(Decode::decode(r)), v, field, index);
        quote!(#res?)
    });
    res = catch_expr(res, quote!(DecodeError));
    res = track_err_in_variant(res, v);
    res
}

fn parse_repr(s: &Structure) -> syn::Result<syn::Type> {
    s.ast()
        .attrs
        .iter()
        .find(|attr| attr.path.is_ident("repr"))
        .ok_or_else(|| {
            syn::Error::new_spanned(
                &s.ast().ident,
                "Wasmbin enums must have a #[repr(type)] attribute",
            )
        })?
        .parse_args()
}

fn wasmbin_derive(s: Structure) -> proc_macro2::TokenStream {
    let (encode_discriminant, decode) = match s.ast().data {
        syn::Data::Enum(_) => {
            let repr = syn_try!(parse_repr(&s));

            let mut encode_discriminant = quote!();

            let mut decoders = quote!();
            let mut decode_other = quote!({ return Ok(None) });

            for v in s.variants() {
                let discriminant = syn_try!(discriminant(v));

                match discriminant {
                    Some(discriminant) => {
                        let pat = v.pat();

                        let encode = gen_encode_discriminant(&repr, &discriminant);
                        (quote!(#pat => #encode,)).to_tokens(&mut encode_discriminant);

                        let decode = gen_decode(v);
                        (quote!(
                            #discriminant => #decode?,
                        ))
                        .to_tokens(&mut decoders);
                    }
                    None => {
                        let fields = v.ast().fields;
                        if fields.len() != 1 {
                            syn_throw!(syn::Error::new_spanned(
                                fields,
                                "Catch-all variants without discriminant must have a single field."
                            ));
                        }
                        let field = fields.iter().next().unwrap();
                        let construct = match &field.ident {
                            Some(ident) => quote!({ #ident: res }),
                            None => quote!((res)),
                        };
                        let variant_name = v.ast().ident;
                        decode_other = quote! {
                            if let Some(res) = DecodeWithDiscriminant::maybe_decode_with_discriminant(discriminant, r)? {
                                Self::#variant_name #construct
                            } else #decode_other
                        };
                    }
                }
            }

            let name = s.ast().ident.to_string();

            (
                quote! {
                    match *self {
                        #encode_discriminant
                        _ => {}
                    }
                },
                quote! {
                    gen impl DecodeWithDiscriminant for @Self {
                        const NAME: &'static str = #name;
                        type Discriminant = #repr;

                        fn maybe_decode_with_discriminant(discriminant: #repr, r: &mut impl std::io::Read) -> Result<Option<Self>, DecodeError> {
                            Ok(Some(match discriminant {
                                #decoders
                                _ => #decode_other
                            }))
                        }
                    }

                    gen impl Decode for @Self {
                        fn decode(r: &mut impl std::io::Read) -> Result<Self, DecodeError> {
                            DecodeWithDiscriminant::decode_without_discriminant(r)
                        }
                    }
                },
            )
        }
        _ => {
            let variants = s.variants();
            assert_eq!(variants.len(), 1);
            let v = &variants[0];
            let decode = gen_decode(v);
            match syn_try!(discriminant(v)) {
                Some(discriminant) => {
                    let name = s.ast().ident.to_string();
                    (
                        gen_encode_discriminant(&syn::parse_quote!(u8), &discriminant),
                        quote! {
                            gen impl DecodeWithDiscriminant for @Self {
                                const NAME: &'static str = #name;
                                type Discriminant = u8;

                                fn maybe_decode_with_discriminant(discriminant: u8, r: &mut impl std::io::Read) -> Result<Option<Self>, DecodeError> {
                                    match discriminant {
                                        #discriminant => #decode.map(Some),
                                        _ => Ok(None),
                                    }
                                }
                            }

                            gen impl Decode for @Self {
                                fn decode(r: &mut impl std::io::Read) -> Result<Self, DecodeError> {
                                    DecodeWithDiscriminant::decode_without_discriminant(r)
                                }
                            }
                        },
                    )
                }
                None => (
                    quote! {},
                    quote! {
                        gen impl Decode for @Self {
                            fn decode(r: &mut impl std::io::Read) -> Result<Self, DecodeError> {
                                #decode
                            }
                        }
                    },
                ),
            }
        }
    };

    let encode_body = s.each(|bi| {
        quote! {
            Encode::encode(#bi, w)?
        }
    });

    s.gen_impl(quote! {
        use crate::io::{Encode, Decode, DecodeWithDiscriminant, DecodeError, PathItem};

        gen impl Encode for @Self {
            fn encode(&self, w: &mut impl std::io::Write) -> std::io::Result<()> {
                #encode_discriminant;
                match *self { #encode_body }
                Ok(())
            }
        }

        #decode
    })
}

fn wasmbin_countable_derive(s: Structure) -> proc_macro2::TokenStream {
    s.gen_impl(quote! {
        gen impl crate::builtins::WasmbinCountable for @Self {}
    })
}

fn wasmbin_visit_derive(mut s: Structure) -> proc_macro2::TokenStream {
    s.bind_with(|_| synstructure::BindStyle::Move);

    fn generate_visit_body(s: &Structure, method: &'static str) -> proc_macro2::TokenStream {
        let method = proc_macro2::Ident::new(method, proc_macro2::Span::call_site());

        let body = s.each_variant(|v| {
            let res = v.bindings().iter().enumerate().map(|(i, bi)| {
                let res = quote!(Visit::#method(#bi, f));
                track_err_in_field(res, v, bi.ast(), i)
            });
            let mut res = quote!(#(#res?;)*);
            res = catch_expr(res, quote!(VisitError<VisitE>));
            res = track_err_in_variant(res, v);
            quote!(#res?)
        });
        quote!(
            match self { #body }
            Ok(())
        )
    }

    let visit_children_body = generate_visit_body(&s, "visit_child");

    let visit_children_mut_body = generate_visit_body(&s, "visit_child_mut");

    s.gen_impl(quote! {
        use crate::visit::{Visit, VisitError};
        use crate::io::PathItem;

        gen impl Visit for @Self where Self: 'static {
            fn visit_children<'a, VisitT: 'static, VisitE, VisitF: FnMut(&'a VisitT) -> Result<(), VisitE>>(&'a self, f: &mut VisitF) -> Result<(), VisitError<VisitE>> {
                #visit_children_body
            }

            fn visit_children_mut<VisitT: 'static, VisitE, VisitF: FnMut(&mut VisitT) -> Result<(), VisitE>>(&mut self, f: &mut VisitF) -> Result<(), VisitError<VisitE>> {
                #visit_children_mut_body
            }
        }
    })
}

#[proc_macro_attribute]
pub fn wasmbin_discriminants(
    _attr: proc_macro::TokenStream,
    input: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    let mut input: syn::DeriveInput = syn::parse(input).unwrap();
    let e = match &mut input.data {
        syn::Data::Enum(e) => e,
        _ => panic!("This attribute can only be used on enums"),
    };
    let mut seen_non_units = false;
    for v in &mut e.variants {
        match v.fields {
            syn::Fields::Unit => {}
            _ => seen_non_units = true,
        }
        #[cfg(not(feature = "nightly"))]
        {
            if let Some((_, discriminant)) = v.discriminant.take() {
                v.attrs
                    .push(syn::parse_quote!(#[wasmbin(discriminant = #discriminant)]));
            }
        }
    }
    assert!(
        seen_non_units,
        "Attribute shouldn't be used on C-like enums"
    );
    input.into_token_stream().into()
}

decl_derive!([Wasmbin, attributes(wasmbin)] => wasmbin_derive);
decl_derive!([WasmbinCountable] => wasmbin_countable_derive);
decl_derive!([Visit] => wasmbin_visit_derive);
<commit_msg>Minor simplification<commit_after>// Copyright 2020 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

extern crate proc_macro;

use quote::{quote, ToTokens};
use std::borrow::Cow;
use synstructure::{decl_derive, Structure, VariantInfo};

macro_rules! syn_throw {
    ($err:expr) => {
        return syn::Error::to_compile_error(&$err);
    };
}

macro_rules! syn_try {
    ($expr:expr) => {
        match $expr {
            Ok(expr) => expr,
            Err(err) => syn_throw!(err),
        }
    };
}

fn discriminant<'v>(v: &VariantInfo<'v>) -> syn::Result<Option<Cow<'v, syn::Expr>>> {
    v.ast()
        .discriminant
        .iter()
        .map(|(_, discriminant)| Ok(Cow::Borrowed(discriminant)))
        .chain(v.ast().attrs.iter().filter_map(|attr| match attr {
            syn::Attribute {
                style: syn::AttrStyle::Outer,
                path,
                ..
            } if path.is_ident("wasmbin") => {
                syn::custom_keyword!(discriminant);

                Some(
                    attr.parse_args_with(|parser: syn::parse::ParseStream| {
                        parser.parse::<discriminant>()?;
                        parser.parse::<syn::Token![=]>()?;
                        parser.parse()
                    })
                    .map(Cow::Owned),
                )
            }
            _ => None,
        }))
        .try_fold(None, |prev, discriminant| {
            let discriminant = discriminant?;
            if let Some(prev) = prev {
                let mut err = syn::Error::new_spanned(
                    discriminant,
                    "#[derive(Wasmbin)]: duplicate discriminant",
                );
                err.combine(syn::Error::new_spanned(
                    prev,
                    "#[derive(Wasmbin)]: previous discriminant here",
                ));
                return Err(err);
            }
            Ok(Some(discriminant))
        })
}

fn gen_encode_discriminant(repr: &syn::Type, discriminant: &syn::Expr) -> proc_macro2::TokenStream {
    quote!(<#repr as Encode>::encode(&#discriminant, w)?)
}

fn is_newtype_like(v: &VariantInfo) -> bool {
    matches!(v.ast().fields, fields @ syn::Fields::Unnamed(_) if fields.len() == 1)
}

fn track_err_in_field(
    mut res: proc_macro2::TokenStream,
    v: &VariantInfo,
    field: &syn::Field,
    index: usize,
) -> proc_macro2::TokenStream {
    if !is_newtype_like(v) {
        let field_name = match &field.ident {
            Some(ident) => ident.to_string(),
            None => index.to_string(),
        };
        res = quote!(#res.map_err(|err| err.in_path(PathItem::Name(#field_name))));
    }
    res
}

fn track_err_in_variant(
    res: proc_macro2::TokenStream,
    v: &VariantInfo,
) -> proc_macro2::TokenStream {
    use std::fmt::Write;

    let mut variant_name = String::new();
    if let Some(prefix) = v.prefix {
        write!(variant_name, "{}::", prefix).unwrap();
    }
    write!(variant_name, "{}", v.ast().ident).unwrap();

    quote!(#res.map_err(|err| err.in_path(PathItem::Variant(#variant_name))))
}

fn catch_expr(
    res: proc_macro2::TokenStream,
    err: proc_macro2::TokenStream,
) -> proc_macro2::TokenStream {
    quote!(
        (move || -> Result<_, #err> {
            Ok({ #res })
        })()
    )
}

fn gen_decode(v: &VariantInfo) -> proc_macro2::TokenStream {
    let mut res = v.construct(|field, index| {
        let res = track_err_in_field(quote!(Decode::decode(r)), v, field, index);
        quote!(#res?)
    });
    res = catch_expr(res, quote!(DecodeError));
    res = track_err_in_variant(res, v);
    res
}

fn parse_repr(s: &Structure) -> syn::Result<syn::Type> {
    s.ast()
        .attrs
        .iter()
        .find(|attr| attr.path.is_ident("repr"))
        .ok_or_else(|| {
            syn::Error::new_spanned(
                &s.ast().ident,
                "Wasmbin enums must have a #[repr(type)] attribute",
            )
        })?
        .parse_args()
}

fn wasmbin_derive(s: Structure) -> proc_macro2::TokenStream {
    let (encode_discriminant, decode) = match s.ast().data {
        syn::Data::Enum(_) => {
            let repr = syn_try!(parse_repr(&s));

            let mut encode_discriminant = quote!();

            let mut decoders = quote!();
            let mut decode_other = quote!({ return Ok(None) });

            for v in s.variants() {
                let discriminant = syn_try!(discriminant(v));

                match discriminant {
                    Some(discriminant) => {
                        let pat = v.pat();

                        let encode = gen_encode_discriminant(&repr, &discriminant);
                        (quote!(#pat => #encode,)).to_tokens(&mut encode_discriminant);

                        let decode = gen_decode(v);
                        (quote!(
                            #discriminant => #decode?,
                        ))
                        .to_tokens(&mut decoders);
                    }
                    None => {
                        let fields = v.ast().fields;
                        if fields.len() != 1 {
                            syn_throw!(syn::Error::new_spanned(
                                fields,
                                "Catch-all variants without discriminant must have a single field."
                            ));
                        }
                        let field = fields.iter().next().unwrap();
                        let construct = match &field.ident {
                            Some(ident) => quote!({ #ident: res }),
                            None => quote!((res)),
                        };
                        let variant_name = v.ast().ident;
                        decode_other = quote! {
                            if let Some(res) = DecodeWithDiscriminant::maybe_decode_with_discriminant(discriminant, r)? {
                                Self::#variant_name #construct
                            } else #decode_other
                        };
                    }
                }
            }

            let name = s.ast().ident.to_string();

            (
                quote! {
                    match *self {
                        #encode_discriminant
                        _ => {}
                    }
                },
                quote! {
                    gen impl DecodeWithDiscriminant for @Self {
                        const NAME: &'static str = #name;
                        type Discriminant = #repr;

                        fn maybe_decode_with_discriminant(discriminant: #repr, r: &mut impl std::io::Read) -> Result<Option<Self>, DecodeError> {
                            Ok(Some(match discriminant {
                                #decoders
                                _ => #decode_other
                            }))
                        }
                    }

                    gen impl Decode for @Self {
                        fn decode(r: &mut impl std::io::Read) -> Result<Self, DecodeError> {
                            DecodeWithDiscriminant::decode_without_discriminant(r)
                        }
                    }
                },
            )
        }
        _ => {
            let variants = s.variants();
            assert_eq!(variants.len(), 1);
            let v = &variants[0];
            let decode = gen_decode(v);
            match syn_try!(discriminant(v)) {
                Some(discriminant) => {
                    let name = s.ast().ident.to_string();
                    (
                        gen_encode_discriminant(&syn::parse_quote!(u8), &discriminant),
                        quote! {
                            gen impl DecodeWithDiscriminant for @Self {
                                const NAME: &'static str = #name;
                                type Discriminant = u8;

                                fn maybe_decode_with_discriminant(discriminant: u8, r: &mut impl std::io::Read) -> Result<Option<Self>, DecodeError> {
                                    match discriminant {
                                        #discriminant => #decode.map(Some),
                                        _ => Ok(None),
                                    }
                                }
                            }

                            gen impl Decode for @Self {
                                fn decode(r: &mut impl std::io::Read) -> Result<Self, DecodeError> {
                                    DecodeWithDiscriminant::decode_without_discriminant(r)
                                }
                            }
                        },
                    )
                }
                None => (
                    quote! {},
                    quote! {
                        gen impl Decode for @Self {
                            fn decode(r: &mut impl std::io::Read) -> Result<Self, DecodeError> {
                                #decode
                            }
                        }
                    },
                ),
            }
        }
    };

    let encode_body = s.each(|bi| {
        quote! {
            Encode::encode(#bi, w)?
        }
    });

    s.gen_impl(quote! {
        use crate::io::{Encode, Decode, DecodeWithDiscriminant, DecodeError, PathItem};

        gen impl Encode for @Self {
            fn encode(&self, w: &mut impl std::io::Write) -> std::io::Result<()> {
                #encode_discriminant;
                match *self { #encode_body }
                Ok(())
            }
        }

        #decode
    })
}

fn wasmbin_countable_derive(s: Structure) -> proc_macro2::TokenStream {
    s.gen_impl(quote! {
        gen impl crate::builtins::WasmbinCountable for @Self {}
    })
}

fn wasmbin_visit_derive(mut s: Structure) -> proc_macro2::TokenStream {
    s.bind_with(|_| synstructure::BindStyle::Move);

    fn generate_visit_body(s: &Structure, method: proc_macro2::TokenStream) -> proc_macro2::TokenStream {
        let body = s.each_variant(|v| {
            let res = v.bindings().iter().enumerate().map(|(i, bi)| {
                let res = quote!(Visit::#method(#bi, f));
                track_err_in_field(res, v, bi.ast(), i)
            });
            let mut res = quote!(#(#res?;)*);
            res = catch_expr(res, quote!(VisitError<VisitE>));
            res = track_err_in_variant(res, v);
            quote!(#res?)
        });
        quote!(
            match self { #body }
            Ok(())
        )
    }

    let visit_children_body = generate_visit_body(&s, quote!(visit_child));

    let visit_children_mut_body = generate_visit_body(&s, quote!(visit_child_mut));

    s.gen_impl(quote! {
        use crate::visit::{Visit, VisitError};
        use crate::io::PathItem;

        gen impl Visit for @Self where Self: 'static {
            fn visit_children<'a, VisitT: 'static, VisitE, VisitF: FnMut(&'a VisitT) -> Result<(), VisitE>>(&'a self, f: &mut VisitF) -> Result<(), VisitError<VisitE>> {
                #visit_children_body
            }

            fn visit_children_mut<VisitT: 'static, VisitE, VisitF: FnMut(&mut VisitT) -> Result<(), VisitE>>(&mut self, f: &mut VisitF) -> Result<(), VisitError<VisitE>> {
                #visit_children_mut_body
            }
        }
    })
}

#[proc_macro_attribute]
pub fn wasmbin_discriminants(
    _attr: proc_macro::TokenStream,
    input: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    let mut input: syn::DeriveInput = syn::parse(input).unwrap();
    let e = match &mut input.data {
        syn::Data::Enum(e) => e,
        _ => panic!("This attribute can only be used on enums"),
    };
    let mut seen_non_units = false;
    for v in &mut e.variants {
        match v.fields {
            syn::Fields::Unit => {}
            _ => seen_non_units = true,
        }
        #[cfg(not(feature = "nightly"))]
        {
            if let Some((_, discriminant)) = v.discriminant.take() {
                v.attrs
                    .push(syn::parse_quote!(#[wasmbin(discriminant = #discriminant)]));
            }
        }
    }
    assert!(
        seen_non_units,
        "Attribute shouldn't be used on C-like enums"
    );
    input.into_token_stream().into()
}

decl_derive!([Wasmbin, attributes(wasmbin)] => wasmbin_derive);
decl_derive!([WasmbinCountable] => wasmbin_countable_derive);
decl_derive!([Visit] => wasmbin_visit_derive);
<|endoftext|>


<commit_before>//! Buffering wrappers for I/O traits

mod bufreader;
mod bufwriter;
mod linewriter;
mod linewritershim;

#[cfg(test)]
mod tests;

use crate::error;
use crate::fmt;
use crate::io::Error;

pub use bufreader::BufReader;
pub use bufwriter::BufWriter;
pub use linewriter::LineWriter;
use linewritershim::LineWriterShim;

/// An error returned by [`BufWriter::into_inner`] which combines an error that
/// happened while writing out the buffer, and the buffered writer object
/// which may be used to recover from the condition.
///
/// # Examples
///
/// ```no_run
/// use std::io::BufWriter;
/// use std::net::TcpStream;
///
/// let mut stream = BufWriter::new(TcpStream::connect("127.0.0.1:34254").unwrap());
///
/// // do stuff with the stream
///
/// // we want to get our `TcpStream` back, so let's try:
///
/// let stream = match stream.into_inner() {
///     Ok(s) => s,
///     Err(e) => {
///         // Here, e is an IntoInnerError
///         panic!("An error occurred");
///     }
/// };
/// ```
#[derive(Debug)]
#[stable(feature = "rust1", since = "1.0.0")]
pub struct IntoInnerError<W>(W, Error);

impl<W> IntoInnerError<W> {
    /// Construct a new IntoInnerError
    fn new(writer: W, error: Error) -> Self {
        Self(writer, error)
    }

    /// Helper to construct a new IntoInnerError; intended to help with
    /// adapters that wrap other adapters
    fn new_wrapped<W2>(self, f: impl FnOnce(W) -> W2) -> IntoInnerError<W2> {
        let Self(writer, error) = self;
        IntoInnerError::new(f(writer), error)
    }

    /// Returns the error which caused the call to [`BufWriter::into_inner()`]
    /// to fail.
    ///
    /// This error was returned when attempting to write the internal buffer.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::io::BufWriter;
    /// use std::net::TcpStream;
    ///
    /// let mut stream = BufWriter::new(TcpStream::connect("127.0.0.1:34254").unwrap());
    ///
    /// // do stuff with the stream
    ///
    /// // we want to get our `TcpStream` back, so let's try:
    ///
    /// let stream = match stream.into_inner() {
    ///     Ok(s) => s,
    ///     Err(e) => {
    ///         // Here, e is an IntoInnerError, let's log the inner error.
    ///         //
    ///         // We'll just 'log' to stdout for this example.
    ///         println!("{}", e.error());
    ///
    ///         panic!("An unexpected error occurred.");
    ///     }
    /// };
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn error(&self) -> &Error {
        &self.1
    }

    /// Returns the buffered writer instance which generated the error.
    ///
    /// The returned object can be used for error recovery, such as
    /// re-inspecting the buffer.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::io::BufWriter;
    /// use std::net::TcpStream;
    ///
    /// let mut stream = BufWriter::new(TcpStream::connect("127.0.0.1:34254").unwrap());
    ///
    /// // do stuff with the stream
    ///
    /// // we want to get our `TcpStream` back, so let's try:
    ///
    /// let stream = match stream.into_inner() {
    ///     Ok(s) => s,
    ///     Err(e) => {
    ///         // Here, e is an IntoInnerError, let's re-examine the buffer:
    ///         let buffer = e.into_inner();
    ///
    ///         // do stuff to try to recover
    ///
    ///         // afterwards, let's just return the stream
    ///         buffer.into_inner().unwrap()
    ///     }
    /// };
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn into_inner(self) -> W {
        self.0
    }

    /// Consumes the [`IntoInnerError`] and returns the error which caused the call to
    /// [`BufWriter::into_inner()`] to fail.  Unlike `error`, this can be used to
    /// obtain ownership of the underlying error.
    ///
    /// # Example
    /// ```
    /// #![feature(io_into_inner_error_parts)]
    /// use std::io::{BufWriter, ErrorKind, Write};
    ///
    /// let mut not_enough_space = [0u8; 10];
    /// let mut stream = BufWriter::new(not_enough_space.as_mut());
    /// write!(stream, "this cannot be actually written").unwrap();
    /// let into_inner_err = stream.into_inner().expect_err("now we discover it's too small");
    /// let err = into_inner_err.into_error();
    /// assert_eq!(err.kind(), ErrorKind::WriteZero);
    /// ```
    #[unstable(feature = "io_into_inner_error_parts", issue = "79704")]
    pub fn into_error(self) -> Error {
        self.1
    }

    /// Consumes the [`IntoInnerError`] and returns the error which caused the call to
    /// [`BufWriter::into_inner()`] to fail, and the underlying writer.
    ///
    /// This can be used to simply obtain ownership of the underlying error; it can also be used for
    /// advanced error recovery.
    ///
    /// # Example
    /// ```
    /// #![feature(io_into_inner_error_parts)]
    /// use std::io::{BufWriter, ErrorKind, Write};
    ///
    /// let mut not_enough_space = [0u8; 10];
    /// let mut stream = BufWriter::new(not_enough_space.as_mut());
    /// write!(stream, "this cannot be actually written").unwrap();
    /// let into_inner_err = stream.into_inner().expect_err("now we discover it's too small");
    /// let (err, recovered_writer) = into_inner_err.into_parts();
    /// assert_eq!(err.kind(), ErrorKind::WriteZero);
    /// assert_eq!(recovered_writer.buffer(), b"t be actually written");
    /// ```
    #[unstable(feature = "io_into_inner_error_parts", issue = "79704")]
    pub fn into_parts(self) -> (Error, W) {
        (self.1, self.0)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<W> From<IntoInnerError<W>> for Error {
    fn from(iie: IntoInnerError<W>) -> Error {
        iie.1
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<W: Send + fmt::Debug> error::Error for IntoInnerError<W> {
    #[allow(deprecated, deprecated_in_future)]
    fn description(&self) -> &str {
        error::Error::description(self.error())
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<W> fmt::Display for IntoInnerError<W> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.error().fmt(f)
    }
}
<commit_msg>Rollup merge of #87175 - inquisitivecrystal:inner-error, r=kennytm<commit_after>//! Buffering wrappers for I/O traits

mod bufreader;
mod bufwriter;
mod linewriter;
mod linewritershim;

#[cfg(test)]
mod tests;

use crate::error;
use crate::fmt;
use crate::io::Error;

pub use bufreader::BufReader;
pub use bufwriter::BufWriter;
pub use linewriter::LineWriter;
use linewritershim::LineWriterShim;

/// An error returned by [`BufWriter::into_inner`] which combines an error that
/// happened while writing out the buffer, and the buffered writer object
/// which may be used to recover from the condition.
///
/// # Examples
///
/// ```no_run
/// use std::io::BufWriter;
/// use std::net::TcpStream;
///
/// let mut stream = BufWriter::new(TcpStream::connect("127.0.0.1:34254").unwrap());
///
/// // do stuff with the stream
///
/// // we want to get our `TcpStream` back, so let's try:
///
/// let stream = match stream.into_inner() {
///     Ok(s) => s,
///     Err(e) => {
///         // Here, e is an IntoInnerError
///         panic!("An error occurred");
///     }
/// };
/// ```
#[derive(Debug)]
#[stable(feature = "rust1", since = "1.0.0")]
pub struct IntoInnerError<W>(W, Error);

impl<W> IntoInnerError<W> {
    /// Construct a new IntoInnerError
    fn new(writer: W, error: Error) -> Self {
        Self(writer, error)
    }

    /// Helper to construct a new IntoInnerError; intended to help with
    /// adapters that wrap other adapters
    fn new_wrapped<W2>(self, f: impl FnOnce(W) -> W2) -> IntoInnerError<W2> {
        let Self(writer, error) = self;
        IntoInnerError::new(f(writer), error)
    }

    /// Returns the error which caused the call to [`BufWriter::into_inner()`]
    /// to fail.
    ///
    /// This error was returned when attempting to write the internal buffer.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::io::BufWriter;
    /// use std::net::TcpStream;
    ///
    /// let mut stream = BufWriter::new(TcpStream::connect("127.0.0.1:34254").unwrap());
    ///
    /// // do stuff with the stream
    ///
    /// // we want to get our `TcpStream` back, so let's try:
    ///
    /// let stream = match stream.into_inner() {
    ///     Ok(s) => s,
    ///     Err(e) => {
    ///         // Here, e is an IntoInnerError, let's log the inner error.
    ///         //
    ///         // We'll just 'log' to stdout for this example.
    ///         println!("{}", e.error());
    ///
    ///         panic!("An unexpected error occurred.");
    ///     }
    /// };
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn error(&self) -> &Error {
        &self.1
    }

    /// Returns the buffered writer instance which generated the error.
    ///
    /// The returned object can be used for error recovery, such as
    /// re-inspecting the buffer.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::io::BufWriter;
    /// use std::net::TcpStream;
    ///
    /// let mut stream = BufWriter::new(TcpStream::connect("127.0.0.1:34254").unwrap());
    ///
    /// // do stuff with the stream
    ///
    /// // we want to get our `TcpStream` back, so let's try:
    ///
    /// let stream = match stream.into_inner() {
    ///     Ok(s) => s,
    ///     Err(e) => {
    ///         // Here, e is an IntoInnerError, let's re-examine the buffer:
    ///         let buffer = e.into_inner();
    ///
    ///         // do stuff to try to recover
    ///
    ///         // afterwards, let's just return the stream
    ///         buffer.into_inner().unwrap()
    ///     }
    /// };
    /// ```
    #[stable(feature = "rust1", since = "1.0.0")]
    pub fn into_inner(self) -> W {
        self.0
    }

    /// Consumes the [`IntoInnerError`] and returns the error which caused the call to
    /// [`BufWriter::into_inner()`] to fail.  Unlike `error`, this can be used to
    /// obtain ownership of the underlying error.
    ///
    /// # Example
    /// ```
    /// use std::io::{BufWriter, ErrorKind, Write};
    ///
    /// let mut not_enough_space = [0u8; 10];
    /// let mut stream = BufWriter::new(not_enough_space.as_mut());
    /// write!(stream, "this cannot be actually written").unwrap();
    /// let into_inner_err = stream.into_inner().expect_err("now we discover it's too small");
    /// let err = into_inner_err.into_error();
    /// assert_eq!(err.kind(), ErrorKind::WriteZero);
    /// ```
    #[stable(feature = "io_into_inner_error_parts", since = "1.55.0")]
    pub fn into_error(self) -> Error {
        self.1
    }

    /// Consumes the [`IntoInnerError`] and returns the error which caused the call to
    /// [`BufWriter::into_inner()`] to fail, and the underlying writer.
    ///
    /// This can be used to simply obtain ownership of the underlying error; it can also be used for
    /// advanced error recovery.
    ///
    /// # Example
    /// ```
    /// use std::io::{BufWriter, ErrorKind, Write};
    ///
    /// let mut not_enough_space = [0u8; 10];
    /// let mut stream = BufWriter::new(not_enough_space.as_mut());
    /// write!(stream, "this cannot be actually written").unwrap();
    /// let into_inner_err = stream.into_inner().expect_err("now we discover it's too small");
    /// let (err, recovered_writer) = into_inner_err.into_parts();
    /// assert_eq!(err.kind(), ErrorKind::WriteZero);
    /// assert_eq!(recovered_writer.buffer(), b"t be actually written");
    /// ```
    #[stable(feature = "io_into_inner_error_parts", since = "1.55.0")]
    pub fn into_parts(self) -> (Error, W) {
        (self.1, self.0)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<W> From<IntoInnerError<W>> for Error {
    fn from(iie: IntoInnerError<W>) -> Error {
        iie.1
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<W: Send + fmt::Debug> error::Error for IntoInnerError<W> {
    #[allow(deprecated, deprecated_in_future)]
    fn description(&self) -> &str {
        error::Error::description(self.error())
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<W> fmt::Display for IntoInnerError<W> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.error().fmt(f)
    }
}
<|endoftext|>