<commit_before>//! Support types for generated encoding tables.
//!
//! This module contains types and functions for working with the encoding tables generated by
//! `lib/codegen/meta-python/gen_encoding.py`.

use constant_hash::{probe, Table};
use ir::{Function, InstructionData, Opcode, Type};
use isa::{Encoding, Legalize};
use settings::PredicateView;
use std::ops::Range;

/// A recipe predicate.
///
/// This is a predicate function capable of testing ISA and instruction predicates simultaneously.
///
/// A None predicate is always satisfied.
pub type RecipePredicate = Option<fn(PredicateView, &InstructionData) -> bool>;

/// An instruction predicate.
///
/// This is a predicate function that needs to be tested in addition to the recipe predicate. It
/// can't depend on ISA settings.
pub type InstPredicate = fn(&Function, &InstructionData) -> bool;

/// Legalization action to perform when no encoding can be found for an instruction.
///
/// This is an index into an ISA-specific table of legalization actions.
pub type LegalizeCode = u8;

/// Level 1 hash table entry.
///
/// One level 1 hash table is generated per CPU mode. This table is keyed by the controlling type
/// variable, using `INVALID` for non-polymorphic instructions.
///
/// The hash table values are references to level 2 hash tables, encoded as an offset in `LEVEL2`
/// where the table begins, and the binary logarithm of its length. All the level 2 hash tables
/// have a power-of-two size.
///
/// Entries are generic over the offset type. It will typically be `u32` or `u16`, depending on the
/// size of the `LEVEL2` table.
///
/// Empty entries are encoded with a `!0` value for `log2len` which will always be out of range.
/// Entries that have a `legalize` value but no level 2 table have an `offset` field that is out f
/// bounds.
pub struct Level1Entry<OffT: Into<u32> + Copy> {
    pub ty: Type,
    pub log2len: u8,
    pub legalize: LegalizeCode,
    pub offset: OffT,
}

impl<OffT: Into<u32> + Copy> Level1Entry<OffT> {
    /// Get the level 2 table range indicated by this entry.
    fn range(&self) -> Range<usize> {
        let b = self.offset.into() as usize;
        b..b + (1 << self.log2len)
    }
}

impl<OffT: Into<u32> + Copy> Table<Type> for [Level1Entry<OffT>] {
    fn len(&self) -> usize {
        self.len()
    }

    fn key(&self, idx: usize) -> Option<Type> {
        if self[idx].log2len != !0 {
            Some(self[idx].ty)
        } else {
            None
        }
    }
}

/// Level 2 hash table entry.
///
/// The second level hash tables are keyed by `Opcode`, and contain an offset into the `ENCLISTS`
/// table where the encoding recipes for the instruction are stored.
///
/// Entries are generic over the offset type which depends on the size of `ENCLISTS`. A `u16`
/// offset allows the entries to be only 32 bits each. There is no benefit to dropping down to `u8`
/// for tiny ISAs. The entries won't shrink below 32 bits since the opcode is expected to be 16
/// bits.
///
/// Empty entries are encoded with a `NotAnOpcode` `opcode` field.
pub struct Level2Entry<OffT: Into<u32> + Copy> {
    pub opcode: Option<Opcode>,
    pub offset: OffT,
}

impl<OffT: Into<u32> + Copy> Table<Opcode> for [Level2Entry<OffT>] {
    fn len(&self) -> usize {
        self.len()
    }

    fn key(&self, idx: usize) -> Option<Opcode> {
        self[idx].opcode
    }
}

/// Two-level hash table lookup and iterator construction.
///
/// Given the controlling type variable and instruction opcode, find the corresponding encoding
/// list.
///
/// Returns an iterator that produces legal encodings for `inst`.
pub fn lookup_enclist<'a, OffT1, OffT2>(
    ctrl_typevar: Type,
    inst: &'a InstructionData,
    func: &'a Function,
    level1_table: &'static [Level1Entry<OffT1>],
    level2_table: &'static [Level2Entry<OffT2>],
    enclist: &'static [EncListEntry],
    legalize_actions: &'static [Legalize],
    recipe_preds: &'static [RecipePredicate],
    inst_preds: &'static [InstPredicate],
    isa_preds: PredicateView<'a>,
) -> Encodings<'a>
where
    OffT1: Into<u32> + Copy,
    OffT2: Into<u32> + Copy,
{
    let (offset, legalize) = match probe(level1_table, ctrl_typevar, ctrl_typevar.index()) {
        Err(l1idx) => {
            // No level 1 entry found for the type.
            // We have a sentinel entry with the default legalization code.
            (!0, level1_table[l1idx].legalize)
        }
        Ok(l1idx) => {
            // We have a valid level 1 entry for this type.
            let l1ent = &level1_table[l1idx];
            let offset = match level2_table.get(l1ent.range()) {
                Some(l2tab) => {
                    let opcode = inst.opcode();
                    match probe(l2tab, opcode, opcode as usize) {
                        Ok(l2idx) => l2tab[l2idx].offset.into() as usize,
                        Err(_) => !0,
                    }
                }
                // The l1ent range is invalid. This means that we just have a customized
                // legalization code for this type. The level 2 table is empty.
                None => !0,
            };
            (offset, l1ent.legalize)
        }
    };

    // Now we have an offset into `enclist` that is `!0` when no encoding list could be found.
    // The default legalization code is always valid.
    Encodings::new(
        offset,
        legalize,
        inst,
        func,
        enclist,
        legalize_actions,
        recipe_preds,
        inst_preds,
        isa_preds,
    )
}

/// Encoding list entry.
///
/// Encoding lists are represented as sequences of u16 words.
pub type EncListEntry = u16;

/// Number of bits used to represent a predicate. c.f. `meta-python/gen_encoding.py`.
const PRED_BITS: u8 = 12;
const PRED_MASK: usize = (1 << PRED_BITS) - 1;
/// First code word representing a predicate check. c.f. `meta-python/gen_encoding.py`.
const PRED_START: usize = 0x1000;

/// An iterator over legal encodings for the instruction.
pub struct Encodings<'a> {
    // Current offset into `enclist`, or out of bounds after we've reached the end.
    offset: usize,
    // Legalization code to use of no encoding is found.
    legalize: LegalizeCode,
    inst: &'a InstructionData,
    func: &'a Function,
    enclist: &'static [EncListEntry],
    legalize_actions: &'static [Legalize],
    recipe_preds: &'static [RecipePredicate],
    inst_preds: &'static [InstPredicate],
    isa_preds: PredicateView<'a>,
}

impl<'a> Encodings<'a> {
    /// Creates a new instance of `Encodings`.
    ///
    /// This iterator provides search for encodings that applies to the given instruction. The
    /// encoding lists are laid out such that first call to `next` returns valid entry in the list
    /// or `None`.
    pub fn new(
        offset: usize,
        legalize: LegalizeCode,
        inst: &'a InstructionData,
        func: &'a Function,
        enclist: &'static [EncListEntry],
        legalize_actions: &'static [Legalize],
        recipe_preds: &'static [RecipePredicate],
        inst_preds: &'static [InstPredicate],
        isa_preds: PredicateView<'a>,
    ) -> Self {
        Encodings {
            offset,
            inst,
            func,
            legalize,
            isa_preds,
            recipe_preds,
            inst_preds,
            enclist,
            legalize_actions,
        }
    }

    /// Get the legalization action that caused the enumeration of encodings to stop.
    /// This can be the default legalization action for the type or a custom code for the
    /// instruction.
    ///
    /// This method must only be called after the iterator returns `None`.
    pub fn legalize(&self) -> Legalize {
        debug_assert_eq!(self.offset, !0, "Premature Encodings::legalize()");
        self.legalize_actions[self.legalize as usize]
    }

    /// Check if the `rpred` recipe predicate is satisfied.
    fn check_recipe(&self, rpred: RecipePredicate) -> bool {
        match rpred {
            Some(p) => p(self.isa_preds, self.inst),
            None => true,
        }
    }

    /// Check an instruction or isa predicate.
    fn check_pred(&self, pred: usize) -> bool {
        if let Some(&p) = self.inst_preds.get(pred) {
            p(self.func, self.inst)
        } else {
            let pred = pred - self.inst_preds.len();
            self.isa_preds.test(pred)
        }
    }
}

impl<'a> Iterator for Encodings<'a> {
    type Item = Encoding;

    fn next(&mut self) -> Option<Encoding> {
        while let Some(entryref) = self.enclist.get(self.offset) {
            let entry = *entryref as usize;

            // Check for "recipe+bits".
            let recipe = entry >> 1;
            if let Some(&rpred) = self.recipe_preds.get(recipe) {
                let bits = self.offset + 1;
                if entry & 1 == 0 {
                    self.offset += 2; // Next entry.
                } else {
                    self.offset = !0; // Stop.
                }
                if self.check_recipe(rpred) {
                    return Some(Encoding::new(recipe as u16, self.enclist[bits]));
                }
                continue;
            }

            // Check for "stop with legalize".
            if entry < PRED_START {
                self.legalize = (entry - 2 * self.recipe_preds.len()) as LegalizeCode;
                self.offset = !0; // Stop.
                return None;
            }

            // Finally, this must be a predicate entry.
            let pred_entry = entry - PRED_START;
            let skip = pred_entry >> PRED_BITS;
            let pred = pred_entry & PRED_MASK;

            if self.check_pred(pred) {
                self.offset += 1;
            } else if skip == 0 {
                self.offset = !0; // Stop.
                return None;
            } else {
                self.offset += 1 + skip;
            }
        }
        None
    }
}
<commit_msg>Fix a typo in a comment.<commit_after>//! Support types for generated encoding tables.
//!
//! This module contains types and functions for working with the encoding tables generated by
//! `lib/codegen/meta-python/gen_encoding.py`.

use constant_hash::{probe, Table};
use ir::{Function, InstructionData, Opcode, Type};
use isa::{Encoding, Legalize};
use settings::PredicateView;
use std::ops::Range;

/// A recipe predicate.
///
/// This is a predicate function capable of testing ISA and instruction predicates simultaneously.
///
/// A None predicate is always satisfied.
pub type RecipePredicate = Option<fn(PredicateView, &InstructionData) -> bool>;

/// An instruction predicate.
///
/// This is a predicate function that needs to be tested in addition to the recipe predicate. It
/// can't depend on ISA settings.
pub type InstPredicate = fn(&Function, &InstructionData) -> bool;

/// Legalization action to perform when no encoding can be found for an instruction.
///
/// This is an index into an ISA-specific table of legalization actions.
pub type LegalizeCode = u8;

/// Level 1 hash table entry.
///
/// One level 1 hash table is generated per CPU mode. This table is keyed by the controlling type
/// variable, using `INVALID` for non-polymorphic instructions.
///
/// The hash table values are references to level 2 hash tables, encoded as an offset in `LEVEL2`
/// where the table begins, and the binary logarithm of its length. All the level 2 hash tables
/// have a power-of-two size.
///
/// Entries are generic over the offset type. It will typically be `u32` or `u16`, depending on the
/// size of the `LEVEL2` table.
///
/// Empty entries are encoded with a `!0` value for `log2len` which will always be out of range.
/// Entries that have a `legalize` value but no level 2 table have an `offset` field that is out of
/// bounds.
pub struct Level1Entry<OffT: Into<u32> + Copy> {
    pub ty: Type,
    pub log2len: u8,
    pub legalize: LegalizeCode,
    pub offset: OffT,
}

impl<OffT: Into<u32> + Copy> Level1Entry<OffT> {
    /// Get the level 2 table range indicated by this entry.
    fn range(&self) -> Range<usize> {
        let b = self.offset.into() as usize;
        b..b + (1 << self.log2len)
    }
}

impl<OffT: Into<u32> + Copy> Table<Type> for [Level1Entry<OffT>] {
    fn len(&self) -> usize {
        self.len()
    }

    fn key(&self, idx: usize) -> Option<Type> {
        if self[idx].log2len != !0 {
            Some(self[idx].ty)
        } else {
            None
        }
    }
}

/// Level 2 hash table entry.
///
/// The second level hash tables are keyed by `Opcode`, and contain an offset into the `ENCLISTS`
/// table where the encoding recipes for the instruction are stored.
///
/// Entries are generic over the offset type which depends on the size of `ENCLISTS`. A `u16`
/// offset allows the entries to be only 32 bits each. There is no benefit to dropping down to `u8`
/// for tiny ISAs. The entries won't shrink below 32 bits since the opcode is expected to be 16
/// bits.
///
/// Empty entries are encoded with a `NotAnOpcode` `opcode` field.
pub struct Level2Entry<OffT: Into<u32> + Copy> {
    pub opcode: Option<Opcode>,
    pub offset: OffT,
}

impl<OffT: Into<u32> + Copy> Table<Opcode> for [Level2Entry<OffT>] {
    fn len(&self) -> usize {
        self.len()
    }

    fn key(&self, idx: usize) -> Option<Opcode> {
        self[idx].opcode
    }
}

/// Two-level hash table lookup and iterator construction.
///
/// Given the controlling type variable and instruction opcode, find the corresponding encoding
/// list.
///
/// Returns an iterator that produces legal encodings for `inst`.
pub fn lookup_enclist<'a, OffT1, OffT2>(
    ctrl_typevar: Type,
    inst: &'a InstructionData,
    func: &'a Function,
    level1_table: &'static [Level1Entry<OffT1>],
    level2_table: &'static [Level2Entry<OffT2>],
    enclist: &'static [EncListEntry],
    legalize_actions: &'static [Legalize],
    recipe_preds: &'static [RecipePredicate],
    inst_preds: &'static [InstPredicate],
    isa_preds: PredicateView<'a>,
) -> Encodings<'a>
where
    OffT1: Into<u32> + Copy,
    OffT2: Into<u32> + Copy,
{
    let (offset, legalize) = match probe(level1_table, ctrl_typevar, ctrl_typevar.index()) {
        Err(l1idx) => {
            // No level 1 entry found for the type.
            // We have a sentinel entry with the default legalization code.
            (!0, level1_table[l1idx].legalize)
        }
        Ok(l1idx) => {
            // We have a valid level 1 entry for this type.
            let l1ent = &level1_table[l1idx];
            let offset = match level2_table.get(l1ent.range()) {
                Some(l2tab) => {
                    let opcode = inst.opcode();
                    match probe(l2tab, opcode, opcode as usize) {
                        Ok(l2idx) => l2tab[l2idx].offset.into() as usize,
                        Err(_) => !0,
                    }
                }
                // The l1ent range is invalid. This means that we just have a customized
                // legalization code for this type. The level 2 table is empty.
                None => !0,
            };
            (offset, l1ent.legalize)
        }
    };

    // Now we have an offset into `enclist` that is `!0` when no encoding list could be found.
    // The default legalization code is always valid.
    Encodings::new(
        offset,
        legalize,
        inst,
        func,
        enclist,
        legalize_actions,
        recipe_preds,
        inst_preds,
        isa_preds,
    )
}

/// Encoding list entry.
///
/// Encoding lists are represented as sequences of u16 words.
pub type EncListEntry = u16;

/// Number of bits used to represent a predicate. c.f. `meta-python/gen_encoding.py`.
const PRED_BITS: u8 = 12;
const PRED_MASK: usize = (1 << PRED_BITS) - 1;
/// First code word representing a predicate check. c.f. `meta-python/gen_encoding.py`.
const PRED_START: usize = 0x1000;

/// An iterator over legal encodings for the instruction.
pub struct Encodings<'a> {
    // Current offset into `enclist`, or out of bounds after we've reached the end.
    offset: usize,
    // Legalization code to use of no encoding is found.
    legalize: LegalizeCode,
    inst: &'a InstructionData,
    func: &'a Function,
    enclist: &'static [EncListEntry],
    legalize_actions: &'static [Legalize],
    recipe_preds: &'static [RecipePredicate],
    inst_preds: &'static [InstPredicate],
    isa_preds: PredicateView<'a>,
}

impl<'a> Encodings<'a> {
    /// Creates a new instance of `Encodings`.
    ///
    /// This iterator provides search for encodings that applies to the given instruction. The
    /// encoding lists are laid out such that first call to `next` returns valid entry in the list
    /// or `None`.
    pub fn new(
        offset: usize,
        legalize: LegalizeCode,
        inst: &'a InstructionData,
        func: &'a Function,
        enclist: &'static [EncListEntry],
        legalize_actions: &'static [Legalize],
        recipe_preds: &'static [RecipePredicate],
        inst_preds: &'static [InstPredicate],
        isa_preds: PredicateView<'a>,
    ) -> Self {
        Encodings {
            offset,
            inst,
            func,
            legalize,
            isa_preds,
            recipe_preds,
            inst_preds,
            enclist,
            legalize_actions,
        }
    }

    /// Get the legalization action that caused the enumeration of encodings to stop.
    /// This can be the default legalization action for the type or a custom code for the
    /// instruction.
    ///
    /// This method must only be called after the iterator returns `None`.
    pub fn legalize(&self) -> Legalize {
        debug_assert_eq!(self.offset, !0, "Premature Encodings::legalize()");
        self.legalize_actions[self.legalize as usize]
    }

    /// Check if the `rpred` recipe predicate is satisfied.
    fn check_recipe(&self, rpred: RecipePredicate) -> bool {
        match rpred {
            Some(p) => p(self.isa_preds, self.inst),
            None => true,
        }
    }

    /// Check an instruction or isa predicate.
    fn check_pred(&self, pred: usize) -> bool {
        if let Some(&p) = self.inst_preds.get(pred) {
            p(self.func, self.inst)
        } else {
            let pred = pred - self.inst_preds.len();
            self.isa_preds.test(pred)
        }
    }
}

impl<'a> Iterator for Encodings<'a> {
    type Item = Encoding;

    fn next(&mut self) -> Option<Encoding> {
        while let Some(entryref) = self.enclist.get(self.offset) {
            let entry = *entryref as usize;

            // Check for "recipe+bits".
            let recipe = entry >> 1;
            if let Some(&rpred) = self.recipe_preds.get(recipe) {
                let bits = self.offset + 1;
                if entry & 1 == 0 {
                    self.offset += 2; // Next entry.
                } else {
                    self.offset = !0; // Stop.
                }
                if self.check_recipe(rpred) {
                    return Some(Encoding::new(recipe as u16, self.enclist[bits]));
                }
                continue;
            }

            // Check for "stop with legalize".
            if entry < PRED_START {
                self.legalize = (entry - 2 * self.recipe_preds.len()) as LegalizeCode;
                self.offset = !0; // Stop.
                return None;
            }

            // Finally, this must be a predicate entry.
            let pred_entry = entry - PRED_START;
            let skip = pred_entry >> PRED_BITS;
            let pred = pred_entry & PRED_MASK;

            if self.check_pred(pred) {
                self.offset += 1;
            } else if skip == 0 {
                self.offset = !0; // Stop.
                return None;
            } else {
                self.offset += 1 + skip;
            }
        }
        None
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str, ptr};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

extern "C" {
    fn register_writer_tx(tx: *mut Sender<Vec<u8>>);
    fn register_stop_tx(tx: *mut Sender<()>);
}

/// Connects to the provided address, (eg "123.123.123.123:3000") and registers
/// the on data received handler
#[no_mangle]
pub extern "C" fn connect(address: *const c_char,
    handler: extern fn(*const c_char, c_int),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    println!("Rust - connect");

    let mut r_address;
    unsafe {
        r_address = CStr::from_ptr(address);
    }
    let s_address = r_address.to_bytes();
    let host_address = match str::from_utf8(s_address) {
        Ok(safe_str) => safe_str,
        Err(_) => {
            println!("Invalid host address");
            return -1 as c_int;
        }
    };

    // Create and register a way to kill this client
    let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    let kill_tx = k_tx.clone();
    let mut k_tx_ptr = Box::new(k_tx);
    unsafe {
        register_stop_tx(&mut *k_tx_ptr);
    }

    // Writer thread's channel
    let (w_tx, w_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    let mut w_tx_ptr = Box::new(w_tx);
    unsafe {
        register_writer_tx(&mut *w_tx_ptr);
    }

    println!("Attempting connect to: {}", host_address);

    let result = TcpStream::connect(host_address);
    if result.is_err() {
        println!("Error connecting to {} - {}", host_address, result.unwrap_err());
        return -1 as c_int;
    }
    println!("Connected");
    on_connect_handler();

    let stream = result.unwrap();
    let client = Bstream::new(stream);

    let r_client = client.clone();
    let r_kill_tx = kill_tx.clone();

    let w_client = client.clone();
    let w_kill_tx = kill_tx.clone();

    // Start the reader thread
    thread::Builder::new()
        .name("ReaderThread".to_string())
        .spawn(move||{
            reader_thread(r_client, handler, r_kill_tx)
        }).unwrap();

    // Start the writer thread
    thread::Builder::new()
        .name("WriterThread".to_string())
        .spawn(move||{
            writer_thread(w_rx, w_client, w_kill_tx)
        }).unwrap();

    // Wait for the kill signal
    match kill_rx.recv() {
        Ok(_) => { }
        Err(e) => {
            println!("Error on kill channel: {}", e);
            return -1 as c_int;
        }
    };
    on_disconnect_handler();

    // Exit out in standard C fashion
    0 as c_int
}

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn send_to_writer(w_tx: *mut Sender<Vec<u8>>,
                                 buffer: *const c_char,
                                 count: c_int,
                                 k_tx: *mut Sender<()>) -> c_int {
    if count < 1 {
        println!("Error - count must be greater than zero");
        return -1 as c_int;
    }

    let num_elts = count as usize;
    let mut n_buffer = Vec::<u8>::with_capacity(num_elts);
    for x in 0..num_elts as isize {
        unsafe {
            n_buffer.push(ptr::read(buffer.offset(x)) as u8);
        }
    }

    unsafe {
        match (*w_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                println!("Error sending buffer: {}", e);
                let _ = (*k_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream,
                 event_handler: extern fn(*const c_char, c_int),
                 kill_tx: Sender<()>) {

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        event_handler(c_buffer.as_ptr(), buffer.len() as c_int);
                    }).unwrap();
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Reader thread finished");
    let _ = kill_tx.send(());
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream, kill_tx: Sender<()>) {
    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        println!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Writer thread finished");
    let _ = kill_tx.send(());
}
<commit_msg>Debugging<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str, ptr};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

extern "C" {
    fn register_writer_tx(tx: *mut Sender<Vec<u8>>);
    fn register_stop_tx(tx: *mut Sender<()>);
}

/// Connects to the provided address, (eg "123.123.123.123:3000") and registers
/// the on data received handler
#[no_mangle]
pub extern "C" fn connect(address: *const c_char,
    handler: extern fn(*const c_char, c_int),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");
    println!("Rust - connect");

    let mut r_address;
    unsafe {
        r_address = CStr::from_ptr(address);
    }
    let s_address = r_address.to_bytes();
    let host_address = match str::from_utf8(s_address) {
        Ok(safe_str) => safe_str,
        Err(_) => {
            println!("Invalid host address");
            return -1 as c_int;
        }
    };

    // Create and register a way to kill this client
    let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    let kill_tx = k_tx.clone();
    let mut k_tx_ptr = Box::new(k_tx);
    unsafe {
        register_stop_tx(&mut *k_tx_ptr);
    }

    // Writer thread's channel
    let (w_tx, w_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    let mut w_tx_ptr = Box::new(w_tx);
    unsafe {
        register_writer_tx(&mut *w_tx_ptr);
    }

    println!("Attempting connect to: {}", host_address);

    let result = TcpStream::connect(host_address);
    if result.is_err() {
        println!("Error connecting to {} - {}", host_address, result.unwrap_err());
        return -1 as c_int;
    }
    println!("Connected");
    on_connect_handler();

    let stream = result.unwrap();
    let client = Bstream::new(stream);

    let r_client = client.clone();
    let r_kill_tx = kill_tx.clone();

    let w_client = client.clone();
    let w_kill_tx = kill_tx.clone();

    // Start the reader thread
    thread::Builder::new()
        .name("ReaderThread".to_string())
        .spawn(move||{
            reader_thread(r_client, handler, r_kill_tx)
        }).unwrap();

    // Start the writer thread
    thread::Builder::new()
        .name("WriterThread".to_string())
        .spawn(move||{
            writer_thread(w_rx, w_client, w_kill_tx)
        }).unwrap();

    // Wait for the kill signal
    match kill_rx.recv() {
        Ok(_) => { }
        Err(e) => {
            println!("Error on kill channel: {}", e);
            return -1 as c_int;
        }
    };
    on_disconnect_handler();

    // Exit out in standard C fashion
    0 as c_int
}

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn send_to_writer(w_tx: *mut Sender<Vec<u8>>,
                                 buffer: *const c_char,
                                 count: c_int,
                                 k_tx: *mut Sender<()>) -> c_int {
    if count < 1 {
        println!("Error - count must be greater than zero");
        return -1 as c_int;
    }

    let num_elts = count as usize;
    let mut n_buffer = Vec::<u8>::with_capacity(num_elts);
    for x in 0..num_elts as isize {
        unsafe {
            n_buffer.push(ptr::read(buffer.offset(x)) as u8);
        }
    }

    unsafe {
        match (*w_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                println!("Error sending buffer: {}", e);
                let _ = (*k_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream,
                 event_handler: extern fn(*const c_char, c_int),
                 kill_tx: Sender<()>) {

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        event_handler(c_buffer.as_ptr(), buffer.len() as c_int);
                    }).unwrap();
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Reader thread finished");
    let _ = kill_tx.send(());
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream, kill_tx: Sender<()>) {
    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        println!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Writer thread finished");
    let _ = kill_tx.send(());
}
<|endoftext|>


<commit_before>// Copyright 2019-2020 Twitter, Inc.
// Licensed under the Apache License, Version 2.0
// http://www.apache.org/licenses/LICENSE-2.0

use std::collections::VecDeque;
use std::io::BufRead;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Instant;

use mio::{Events, Poll, Token};
use rand::prelude::SliceRandom;
use rand::rngs::ThreadRng;
use rand::thread_rng;
use rustcommon_timer::Wheel;
use rustls::ClientConfig;
use slab::Slab;

use crate::codec::*;
use crate::config::Protocol;
use crate::session::{Session, State};
use crate::stats::*;
use crate::*;

pub struct Client {
    codec: Box<dyn Codec>,
    sessions: Slab<Session>,
    config: Arc<Config>,
    ready_queue: VecDeque<usize>,
    connect_queue: VecDeque<SocketAddr>,
    tls_config: Option<Arc<ClientConfig>>,
    metrics: Arc<Metrics>,
    timers: Wheel<usize>,
    last_timeout: Instant,
    events: Option<Events>,
    poll: Poll,
    id: usize,
    connect: Option<Arc<Ratelimiter>>,
    request: Option<Arc<Ratelimiter>>,
    close: Option<Arc<Ratelimiter>>,
}

impl Client {
    pub fn new(
        id: usize,
        config: Arc<Config>,
        connect: Option<Arc<Ratelimiter>>,
        request: Option<Arc<Ratelimiter>>,
        close: Option<Arc<Ratelimiter>>,
        metrics: Arc<Metrics>,
    ) -> Self {
        let codec: Box<dyn Codec> = match config.protocol() {
            Protocol::Echo => Box::new(crate::codec::Echo::new()),
            Protocol::Memcache => Box::new(crate::codec::Memcache::new()),
            Protocol::ThriftCache => Box::new(crate::codec::ThriftCache::new()),
            Protocol::PelikanRds => Box::new(crate::codec::PelikanRds::new()),
            Protocol::Ping => Box::new(crate::codec::Ping::new()),
            Protocol::RedisResp => {
                Box::new(crate::codec::Redis::new(crate::codec::RedisMode::Resp))
            }
            Protocol::RedisInline => {
                Box::new(crate::codec::Redis::new(crate::codec::RedisMode::Inline))
            }
        };

        let tls_config = load_tls_config(&config);

        Self {
            codec,
            sessions: Slab::new(),
            config,
            ready_queue: VecDeque::new(),
            connect_queue: VecDeque::new(),
            metrics,
            tls_config,
            timers: Wheel::<usize>::new(SECOND / MICROSECOND),
            last_timeout: Instant::now(),
            events: None,
            poll: Poll::new().expect("failed to create mio::Poll"),
            id,
            connect,
            request,
            close,
        }
    }

    pub fn add_endpoint(&mut self, addr: &SocketAddr) {
        debug!("client({}) adding endpoint: {}", self.id, addr);
        for _ in 0..self.config.poolsize() {
            self.connect_queue.push_back(*addr);
        }
        self.connect_shuffle();
    }

    fn connect_shuffle(&mut self) {
        let mut tmp: Vec<SocketAddr> = self.connect_queue.drain(0..).collect();
        let mut rng = thread_rng();
        tmp.shuffle(&mut rng);
        for addr in tmp {
            self.connect_queue.push_back(addr);
        }
    }

    fn do_timeouts(&mut self) {
        let last = self.last_timeout;
        let now = Instant::now();
        let ticks =
            (now - last).as_secs() as usize * 1000000 + (now - last).subsec_nanos() as usize / 1000;

        let timeouts = self.timers.tick(ticks);

        for token in timeouts {
            if let Some(state) = self.sessions.get(token).map(|v| v.state()) {
                match state {
                    State::Connecting => {
                        // timeout while connecting
                        self.stat_increment(Stat::ConnectionsTimeout);
                    }
                    State::Reading => {
                        // timeout while reading
                        self.stat_increment(Stat::RequestsTimeout);
                    }
                    _ => {
                        // ignore other timeouts
                        continue;
                    }
                }
                if !self.config.soft_timeout() {
                    let session = self.sessions.remove(token);
                    self.connect_queue.push_back(session.addr());
                }
            }
        }

        self.last_timeout = now;
    }

    fn server_closed(&mut self, token: usize) {
        // server has closed connection
        trace!("server closed: {}", token);
        self.metrics.increment(&Stat::ConnectionsClosed);
        self.metrics.increment(&Stat::ConnectionsServerClosed);
        let mut session = self.sessions.remove(token);
        session.deregister(&self.poll);
        self.connect_queue.push_back(session.addr());
    }

    fn hangup(&mut self, token: usize) {
        if self.sessions.contains(token) {
            self.metrics.increment(&Stat::ConnectionsClosed);
            self.metrics.increment(&Stat::ConnectionsClientClosed);
            let mut session = self.sessions.remove(token);
            session.deregister(&self.poll);
            self.connect_queue.push_back(session.addr());
        }
    }

    fn do_events(&mut self) {
        let mut events = self
            .events
            .take()
            .unwrap_or_else(|| Events::with_capacity(1024));
        self.poll
            .poll(&mut events, Some(Duration::from_millis(1)))
            .unwrap();
        for event in events.iter() {
            let token = event.token();
            if let Some(session) = self.sessions.get_mut(token.0) {
                let read_status = if event.is_readable() {
                    trace!("handle read for: {}", token.0);
                    session.do_read()
                } else {
                    Ok(None)
                };

                let write_status = if event.is_writable() {
                    trace!("handle write for: {}", token.0);
                    session.do_write()
                } else {
                    Ok(None)
                };

                match read_status {
                    Ok(Some(0)) => {
                        self.server_closed(token.0);
                        continue;
                    }
                    Ok(Some(bytes)) => {
                        let start = session.timestamp();
                        // parse response
                        trace!("read {} bytes: {}", bytes, token.0);
                        if let Ok(content) = session.buffer.fill_buf() {
                            trace!("read: {:?}", content);
                            match self.codec.decode(content) {
                                Ok(response) => {
                                    let stop = Instant::now();

                                    self.metrics.heatmap_increment(start, stop);
                                    self.metrics.time_interval(
                                        &Stat::ResponsesLatency,
                                        start,
                                        stop,
                                    );
                                    self.metrics.increment(&Stat::ResponsesTotal);
                                    self.metrics.increment(&Stat::ResponsesOk);

                                    self.ready_queue.push_back(token.0);
                                    session.set_state(State::Writing);

                                    match response {
                                        Response::Hit => {
                                            self.metrics.increment(&Stat::ResponsesHit);
                                        }
                                        Response::Miss => {
                                            self.metrics.increment(&Stat::ResponsesMiss);
                                        }
                                        _ => {}
                                    }
                                }
                                Err(error) => {
                                    if error != Error::Incomplete {
                                        self.metrics.increment(&Stat::ResponsesTotal);
                                        self.metrics.increment(&Stat::ResponsesError);
                                    }

                                    match error {
                                        Error::ChecksumMismatch(a, b) => {
                                            let stop = Instant::now();
                                            let start = session.timestamp();
                                            self.metrics.heatmap_increment(start, stop);
                                            self.metrics.time_interval(
                                                &Stat::ResponsesLatency,
                                                start,
                                                stop,
                                            );
                                            warn!("Response checksum mismatch!");
                                            warn!("Expected: {:?}", a);
                                            warn!("Got: {:?}", b);
                                            self.ready_queue.push_back(token.0);
                                            session.set_state(State::Writing);
                                        }
                                        _ => {
                                            self.hangup(token.0);
                                            continue;
                                        }
                                    }
                                }
                            }
                        }
                        session.buffer.consume(bytes);
                    }
                    Ok(None) => {
                        // wasn't ready
                        trace!("spurious read: {}", token.0);
                    }
                    Err(_) => {
                        // got some error, close connection
                        self.metrics.increment(&Stat::ResponsesTotal);
                        self.metrics.increment(&Stat::ResponsesError);
                        self.hangup(token.0);
                        continue;
                    }
                }

                match write_status {
                    Ok(Some(bytes)) => {
                        trace!("wrote: {} bytes: {}", bytes, token.0);
                        if session.tx_pending() > 0 {
                            // incomplete write
                            println!("have: {} bytes pending: {}", session.tx_pending(), token.0);
                        } else {
                            // completed write
                            self.metrics.increment(&Stat::RequestsDequeued);
                            session.set_state(State::Reading);
                        }
                    }
                    Ok(None) => {
                        trace!("spurious write: {}", token.0);
                    }
                    Err(_) => {
                        // got some error, close connection
                        let session = self.sessions.remove(token.0);
                        self.connect_queue.push_back(session.addr());
                        continue;
                    }
                }

                if session.state() == State::Connecting && !session.is_handshaking() {
                    // increment time interval
                    let stop = Instant::now();
                    let start = session.timestamp();
                    self.metrics
                        .time_interval(&Stat::ConnectionsLatency, start, stop);
                    self.metrics.increment(&Stat::ConnectionsOpened);

                    // finished connecting
                    session.set_state(State::Connected);
                    self.ready_queue.push_back(token.0);
                }
                session.reregister(&self.poll);
            } else {
                // ignore event for unknown session?
            }
        }

        self.events = Some(events);
    }

    fn send_request(&mut self, rng: &mut ThreadRng, token: usize) {
        if let Some(session) = self.sessions.get_mut(token) {
            trace!("send request: {}", token);
            session.set_timestamp(Instant::now());
            self.metrics.increment(&Stat::RequestsEnqueued);
            self.codec.encode(&mut session.buffer, rng);
            session.set_state(State::Writing);
            session.reregister(&self.poll);
        }
    }

    fn do_requests(&mut self, rng: &mut ThreadRng) {
        loop {
            if let Some(token) = self.ready_queue.pop_front() {
                if let Some(ref mut request) = self.request {
                    if request.try_wait().is_ok() {
                        self.send_request(rng, token);
                    } else {
                        self.ready_queue.push_front(token);
                    }
                } else {
                    self.send_request(rng, token);
                }
            } else {
                break;
            }
        }
    }

    fn do_hangups(&mut self) {
        if self.close.is_some() {
            loop {
                if let Some(token) = self.ready_queue.pop_front() {
                    trace!("hangup: {}", token);
                    if self.sessions.contains(token) {
                        if self.close.as_ref().unwrap().try_wait().is_ok() {
                            self.hangup(token);
                        } else {
                            self.ready_queue.push_front(token);
                            break;
                        }
                    }
                } else {
                    break;
                }
            }
        }
    }

    fn connect(&mut self, addr: SocketAddr) {
        let session = self.sessions.vacant_entry();
        let tls = if let Some(ref mut tls_config) = self.tls_config {
            Some(rustls::ClientSession::new(
                &tls_config,
                webpki::DNSNameRef::try_from_ascii_str("localhost").expect("invalid dns name"),
            ))
        } else {
            None
        };
        let start = Instant::now();
        if let Ok(mut s) = Session::new(addr, Token(session.key()), tls) {
            s.set_nodelay(self.config.tcp_nodelay());
            self.metrics.increment(&Stat::ConnectionsTotal);
            if self.tls_config.is_some() {
                s.register(&self.poll);
            } else {
                self.metrics
                    .time_interval(&Stat::ConnectionsLatency, start, Instant::now());
                self.metrics.increment(&Stat::ConnectionsOpened);
                s.register(&self.poll);
                self.ready_queue.push_back(session.key());
            }
            session.insert(s);
        } else {
            self.metrics.increment(&Stat::ConnectionsError);
            self.connect_queue.push_back(addr);
        }
    }

    fn do_connects(&mut self) {
        while let Some(addr) = self.connect_queue.pop_front() {
            trace!("connect: {}", addr);
            if let Some(ref mut connect) = self.connect {
                if connect.try_wait().is_ok() {
                    self.connect(addr);
                } else {
                    self.connect_queue.push_back(addr);
                    break;
                }
            } else {
                self.connect(addr);
            }
        }
    }

    pub fn run(&mut self, rng: &mut ThreadRng) {
        self.do_timeouts();
        self.do_events();
        self.do_connects();
        if self.close.is_some() {
            self.do_hangups();
        }
        self.do_requests(rng);
    }

    fn stat_increment(&self, label: Stat) {
        self.metrics.increment(&label)
    }
}

fn load_tls_config(config: &Arc<Config>) -> Option<Arc<rustls::ClientConfig>> {
    let cert_chain = config.tls_ca();
    let cert = config.tls_cert();
    let key = config.tls_key();

    if cert_chain.is_some() && cert.is_some() && key.is_some() {
        let mut config = rustls::ClientConfig::new();

        let certificate_chain =
            std::fs::File::open(cert_chain.unwrap()).expect("failed to open cert chain");
        config
            .root_store
            .add_pem_file(&mut std::io::BufReader::new(certificate_chain))
            .expect("failed to load cert chain");

        config
            .dangerous()
            .set_certificate_verifier(Arc::new(NoCertificateVerification {}));

        let cert = std::fs::File::open(cert.unwrap()).expect("failed to open cert");
        let cert = rustls::internal::pemfile::certs(&mut std::io::BufReader::new(cert)).unwrap();

        let key = std::fs::File::open(key.unwrap()).expect("failed to open private key");
        let keys = rustls::internal::pemfile::pkcs8_private_keys(&mut std::io::BufReader::new(key))
            .unwrap();
        assert_eq!(keys.len(), 1);
        let key = keys[0].clone();

        config
            .set_single_client_cert(cert, key)
            .expect("invalid cert or key");

        Some(Arc::new(config))
    } else if cert_chain.is_none() && cert.is_none() && key.is_none() {
        None
    } else {
        fatal!("Invalid TLS configuration");
    }
}

pub struct NoCertificateVerification {}

impl rustls::ServerCertVerifier for NoCertificateVerification {
    fn verify_server_cert(
        &self,
        _roots: &rustls::RootCertStore,
        _presented_certs: &[rustls::Certificate],
        _dns_name: webpki::DNSNameRef<'_>,
        _ocsp: &[u8],
    ) -> Result<rustls::ServerCertVerified, rustls::TLSError> {
        Ok(rustls::ServerCertVerified::assertion())
    }
}
<commit_msg>fix: incorrect latency measurements when ratelimited (#310)<commit_after>// Copyright 2019-2020 Twitter, Inc.
// Licensed under the Apache License, Version 2.0
// http://www.apache.org/licenses/LICENSE-2.0

use std::collections::VecDeque;
use std::io::BufRead;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Instant;

use mio::{Events, Poll, Token};
use rand::prelude::SliceRandom;
use rand::rngs::ThreadRng;
use rand::thread_rng;
use rustcommon_timer::Wheel;
use rustls::ClientConfig;
use slab::Slab;

use crate::codec::*;
use crate::config::Protocol;
use crate::session::{Session, State};
use crate::stats::*;
use crate::*;

pub struct Client {
    codec: Box<dyn Codec>,
    sessions: Slab<Session>,
    config: Arc<Config>,
    ready_queue: VecDeque<usize>,
    connect_queue: VecDeque<SocketAddr>,
    tls_config: Option<Arc<ClientConfig>>,
    metrics: Arc<Metrics>,
    timers: Wheel<usize>,
    last_timeout: Instant,
    events: Option<Events>,
    poll: Poll,
    id: usize,
    connect: Option<Arc<Ratelimiter>>,
    request: Option<Arc<Ratelimiter>>,
    close: Option<Arc<Ratelimiter>>,
}

impl Client {
    pub fn new(
        id: usize,
        config: Arc<Config>,
        connect: Option<Arc<Ratelimiter>>,
        request: Option<Arc<Ratelimiter>>,
        close: Option<Arc<Ratelimiter>>,
        metrics: Arc<Metrics>,
    ) -> Self {
        let codec: Box<dyn Codec> = match config.protocol() {
            Protocol::Echo => Box::new(crate::codec::Echo::new()),
            Protocol::Memcache => Box::new(crate::codec::Memcache::new()),
            Protocol::ThriftCache => Box::new(crate::codec::ThriftCache::new()),
            Protocol::PelikanRds => Box::new(crate::codec::PelikanRds::new()),
            Protocol::Ping => Box::new(crate::codec::Ping::new()),
            Protocol::RedisResp => {
                Box::new(crate::codec::Redis::new(crate::codec::RedisMode::Resp))
            }
            Protocol::RedisInline => {
                Box::new(crate::codec::Redis::new(crate::codec::RedisMode::Inline))
            }
        };

        let tls_config = load_tls_config(&config);

        Self {
            codec,
            sessions: Slab::new(),
            config,
            ready_queue: VecDeque::new(),
            connect_queue: VecDeque::new(),
            metrics,
            tls_config,
            timers: Wheel::<usize>::new(SECOND / MICROSECOND),
            last_timeout: Instant::now(),
            events: None,
            poll: Poll::new().expect("failed to create mio::Poll"),
            id,
            connect,
            request,
            close,
        }
    }

    pub fn add_endpoint(&mut self, addr: &SocketAddr) {
        debug!("client({}) adding endpoint: {}", self.id, addr);
        for _ in 0..self.config.poolsize() {
            self.connect_queue.push_back(*addr);
        }
        self.connect_shuffle();
    }

    fn connect_shuffle(&mut self) {
        let mut tmp: Vec<SocketAddr> = self.connect_queue.drain(0..).collect();
        let mut rng = thread_rng();
        tmp.shuffle(&mut rng);
        for addr in tmp {
            self.connect_queue.push_back(addr);
        }
    }

    fn do_timeouts(&mut self) {
        let last = self.last_timeout;
        let now = Instant::now();
        let ticks =
            (now - last).as_secs() as usize * 1000000 + (now - last).subsec_nanos() as usize / 1000;

        let timeouts = self.timers.tick(ticks);

        for token in timeouts {
            if let Some(state) = self.sessions.get(token).map(|v| v.state()) {
                match state {
                    State::Connecting => {
                        // timeout while connecting
                        self.stat_increment(Stat::ConnectionsTimeout);
                    }
                    State::Reading => {
                        // timeout while reading
                        self.stat_increment(Stat::RequestsTimeout);
                    }
                    _ => {
                        // ignore other timeouts
                        continue;
                    }
                }
                if !self.config.soft_timeout() {
                    let session = self.sessions.remove(token);
                    self.connect_queue.push_back(session.addr());
                }
            }
        }

        self.last_timeout = now;
    }

    fn server_closed(&mut self, token: usize) {
        // server has closed connection
        trace!("server closed: {}", token);
        self.metrics.increment(&Stat::ConnectionsClosed);
        self.metrics.increment(&Stat::ConnectionsServerClosed);
        let mut session = self.sessions.remove(token);
        session.deregister(&self.poll);
        self.connect_queue.push_back(session.addr());
    }

    fn hangup(&mut self, token: usize) {
        if self.sessions.contains(token) {
            self.metrics.increment(&Stat::ConnectionsClosed);
            self.metrics.increment(&Stat::ConnectionsClientClosed);
            let mut session = self.sessions.remove(token);
            session.deregister(&self.poll);
            self.connect_queue.push_back(session.addr());
        }
    }

    fn do_events(&mut self) {
        let mut events = self
            .events
            .take()
            .unwrap_or_else(|| Events::with_capacity(1024));
        self.poll
            .poll(&mut events, Some(Duration::from_millis(1)))
            .unwrap();
        for event in events.iter() {
            let token = event.token();
            if let Some(session) = self.sessions.get_mut(token.0) {
                let read_status = if event.is_readable() {
                    trace!("handle read for: {}", token.0);
                    session.do_read()
                } else {
                    Ok(None)
                };

                let write_status = if event.is_writable() {
                    trace!("handle write for: {}", token.0);
                    session.set_timestamp(Instant::now());
                    session.do_write()
                } else {
                    Ok(None)
                };

                match read_status {
                    Ok(Some(0)) => {
                        self.server_closed(token.0);
                        continue;
                    }
                    Ok(Some(bytes)) => {
                        let start = session.timestamp();
                        // parse response
                        trace!("read {} bytes: {}", bytes, token.0);
                        if let Ok(content) = session.buffer.fill_buf() {
                            trace!("read: {:?}", content);
                            match self.codec.decode(content) {
                                Ok(response) => {
                                    let stop = Instant::now();

                                    self.metrics.heatmap_increment(start, stop);
                                    self.metrics.time_interval(
                                        &Stat::ResponsesLatency,
                                        start,
                                        stop,
                                    );
                                    self.metrics.increment(&Stat::ResponsesTotal);
                                    self.metrics.increment(&Stat::ResponsesOk);

                                    self.ready_queue.push_back(token.0);
                                    session.set_state(State::Writing);

                                    match response {
                                        Response::Hit => {
                                            self.metrics.increment(&Stat::ResponsesHit);
                                        }
                                        Response::Miss => {
                                            self.metrics.increment(&Stat::ResponsesMiss);
                                        }
                                        _ => {}
                                    }
                                }
                                Err(error) => {
                                    if error != Error::Incomplete {
                                        self.metrics.increment(&Stat::ResponsesTotal);
                                        self.metrics.increment(&Stat::ResponsesError);
                                    }

                                    match error {
                                        Error::ChecksumMismatch(a, b) => {
                                            let stop = Instant::now();
                                            let start = session.timestamp();
                                            self.metrics.heatmap_increment(start, stop);
                                            self.metrics.time_interval(
                                                &Stat::ResponsesLatency,
                                                start,
                                                stop,
                                            );
                                            warn!("Response checksum mismatch!");
                                            warn!("Expected: {:?}", a);
                                            warn!("Got: {:?}", b);
                                            self.ready_queue.push_back(token.0);
                                            session.set_state(State::Writing);
                                        }
                                        _ => {
                                            self.hangup(token.0);
                                            continue;
                                        }
                                    }
                                }
                            }
                        }
                        session.buffer.consume(bytes);
                    }
                    Ok(None) => {
                        // wasn't ready
                        trace!("spurious read: {}", token.0);
                    }
                    Err(_) => {
                        // got some error, close connection
                        self.metrics.increment(&Stat::ResponsesTotal);
                        self.metrics.increment(&Stat::ResponsesError);
                        self.hangup(token.0);
                        continue;
                    }
                }

                match write_status {
                    Ok(Some(bytes)) => {
                        trace!("wrote: {} bytes: {}", bytes, token.0);
                        if session.tx_pending() > 0 {
                            // incomplete write
                            println!("have: {} bytes pending: {}", session.tx_pending(), token.0);
                        } else if bytes > 0 {
                            // completed write
                            self.metrics.increment(&Stat::RequestsDequeued);
                            session.set_state(State::Reading);
                        }
                    }
                    Ok(None) => {
                        trace!("spurious write: {}", token.0);
                    }
                    Err(_) => {
                        // got some error, close connection
                        let session = self.sessions.remove(token.0);
                        self.connect_queue.push_back(session.addr());
                        continue;
                    }
                }

                if session.state() == State::Connecting && !session.is_handshaking() {
                    // increment time interval
                    let stop = Instant::now();
                    let start = session.timestamp();
                    self.metrics
                        .time_interval(&Stat::ConnectionsLatency, start, stop);
                    self.metrics.increment(&Stat::ConnectionsOpened);

                    // finished connecting
                    session.set_state(State::Connected);
                    self.ready_queue.push_back(token.0);
                }
                session.reregister(&self.poll);
            } else {
                // ignore event for unknown session?
            }
        }

        self.events = Some(events);
    }

    fn send_request(&mut self, rng: &mut ThreadRng, token: usize) {
        if let Some(session) = self.sessions.get_mut(token) {
            trace!("send request: {}", token);
            session.set_timestamp(Instant::now());
            self.metrics.increment(&Stat::RequestsEnqueued);
            self.codec.encode(&mut session.buffer, rng);
            session.set_state(State::Writing);
            session.reregister(&self.poll);
        }
    }

    fn do_requests(&mut self, rng: &mut ThreadRng) {
        loop {
            if let Some(token) = self.ready_queue.pop_front() {
                if let Some(ref mut request) = self.request {
                    if request.try_wait().is_ok() {
                        self.send_request(rng, token);
                    } else {
                        self.ready_queue.push_front(token);
                        break;
                    }
                } else {
                    self.send_request(rng, token);
                }
            } else {
                break;
            }
        }
    }

    fn do_hangups(&mut self) {
        if self.close.is_some() {
            loop {
                if let Some(token) = self.ready_queue.pop_front() {
                    trace!("hangup: {}", token);
                    if self.sessions.contains(token) {
                        if self.close.as_ref().unwrap().try_wait().is_ok() {
                            self.hangup(token);
                        } else {
                            self.ready_queue.push_front(token);
                            break;
                        }
                    }
                } else {
                    break;
                }
            }
        }
    }

    fn connect(&mut self, addr: SocketAddr) {
        let session = self.sessions.vacant_entry();
        let tls = if let Some(ref mut tls_config) = self.tls_config {
            Some(rustls::ClientSession::new(
                &tls_config,
                webpki::DNSNameRef::try_from_ascii_str("localhost").expect("invalid dns name"),
            ))
        } else {
            None
        };
        let start = Instant::now();
        if let Ok(mut s) = Session::new(addr, Token(session.key()), tls) {
            s.set_nodelay(self.config.tcp_nodelay());
            self.metrics.increment(&Stat::ConnectionsTotal);
            if self.tls_config.is_some() {
                s.register(&self.poll);
            } else {
                self.metrics
                    .time_interval(&Stat::ConnectionsLatency, start, Instant::now());
                self.metrics.increment(&Stat::ConnectionsOpened);
                s.register(&self.poll);
                self.ready_queue.push_back(session.key());
            }
            session.insert(s);
        } else {
            self.metrics.increment(&Stat::ConnectionsError);
            self.connect_queue.push_back(addr);
        }
    }

    fn do_connects(&mut self) {
        while let Some(addr) = self.connect_queue.pop_front() {
            trace!("connect: {}", addr);
            if let Some(ref mut connect) = self.connect {
                if connect.try_wait().is_ok() {
                    self.connect(addr);
                } else {
                    self.connect_queue.push_back(addr);
                    break;
                }
            } else {
                self.connect(addr);
            }
        }
    }

    pub fn run(&mut self, rng: &mut ThreadRng) {
        self.do_timeouts();
        self.do_events();
        self.do_connects();
        if self.close.is_some() {
            self.do_hangups();
        }
        self.do_requests(rng);
    }

    fn stat_increment(&self, label: Stat) {
        self.metrics.increment(&label)
    }
}

fn load_tls_config(config: &Arc<Config>) -> Option<Arc<rustls::ClientConfig>> {
    let cert_chain = config.tls_ca();
    let cert = config.tls_cert();
    let key = config.tls_key();

    if cert_chain.is_some() && cert.is_some() && key.is_some() {
        let mut config = rustls::ClientConfig::new();

        let certificate_chain =
            std::fs::File::open(cert_chain.unwrap()).expect("failed to open cert chain");
        config
            .root_store
            .add_pem_file(&mut std::io::BufReader::new(certificate_chain))
            .expect("failed to load cert chain");

        config
            .dangerous()
            .set_certificate_verifier(Arc::new(NoCertificateVerification {}));

        let cert = std::fs::File::open(cert.unwrap()).expect("failed to open cert");
        let cert = rustls::internal::pemfile::certs(&mut std::io::BufReader::new(cert)).unwrap();

        let key = std::fs::File::open(key.unwrap()).expect("failed to open private key");
        let keys = rustls::internal::pemfile::pkcs8_private_keys(&mut std::io::BufReader::new(key))
            .unwrap();
        assert_eq!(keys.len(), 1);
        let key = keys[0].clone();

        config
            .set_single_client_cert(cert, key)
            .expect("invalid cert or key");

        Some(Arc::new(config))
    } else if cert_chain.is_none() && cert.is_none() && key.is_none() {
        None
    } else {
        fatal!("Invalid TLS configuration");
    }
}

pub struct NoCertificateVerification {}

impl rustls::ServerCertVerifier for NoCertificateVerification {
    fn verify_server_cert(
        &self,
        _roots: &rustls::RootCertStore,
        _presented_certs: &[rustls::Certificate],
        _dns_name: webpki::DNSNameRef<'_>,
        _ocsp: &[u8],
    ) -> Result<rustls::ServerCertVerified, rustls::TLSError> {
        Ok(rustls::ServerCertVerified::assertion())
    }
}
<|endoftext|>


<commit_before>extern crate docopt;
extern crate regex;
extern crate regex_syntax as syntax;
extern crate serde;
#[macro_use]
extern crate serde_derive;

use std::error;
use std::io::{self, Write};
use std::process;
use std::result;

use docopt::Docopt;
use regex::internal::{Compiler, LiteralSearcher};
use syntax::hir::literal::Literals;
use syntax::hir::Hir;

const USAGE: &'static str = "
Usage:
    regex-debug [options] ast <pattern>
    regex-debug [options] hir <pattern>
    regex-debug [options] prefixes <patterns> ...
    regex-debug [options] suffixes <patterns> ...
    regex-debug [options] anchors <pattern>
    regex-debug [options] captures <pattern>
    regex-debug [options] compile <patterns> ...
    regex-debug [options] utf8-ranges <class>
    regex-debug --help

Options:
    --help               Show this usage message.
    --size-limit ARG     An approximate size limit on the total size (in bytes)
                         of a compiled regular expression program.
                         [default: 10485760]
    --bytes              Show the instruction codes for byte oriented programs.
                         (As opposed to Unicode oriented programs.)
    --dfa                Show the instruction codes for a DFA.
    --dfa-reverse        Show the instruction codes for a reverse DFA.
                         This implies --dfa.
    -a, --all-literals   Shows all literals extracted.
                         By default, only unambiguous literals are shown.
    --literal-limit ARG  An approximate limit on the total size (in bytes)
                         of all literals extracted. [default: 250]
    --class-limit ARG    A limit on the size of character classes used to
                         extract literals. [default: 10]
    --literal-bytes      Show raw literal bytes instead of Unicode chars.
    --lcp                Show the longest common prefix of all the literals
                         extracted.
    --lcs                Show the longest common suffix of all the literals
                         extracted.
    --searcher           Show the debug output for the literal searcher
                         constructed by the literals found.
";

#[derive(Deserialize)]
struct Args {
    cmd_ast: bool,
    cmd_hir: bool,
    cmd_prefixes: bool,
    cmd_suffixes: bool,
    cmd_anchors: bool,
    cmd_captures: bool,
    cmd_compile: bool,
    cmd_utf8_ranges: bool,

    arg_pattern: String,
    arg_patterns: Vec<String>,
    arg_class: String,

    flag_size_limit: usize,
    flag_bytes: bool,
    flag_dfa: bool,
    flag_dfa_reverse: bool,
    flag_all_literals: bool,
    flag_literal_limit: usize,
    flag_class_limit: usize,
    flag_literal_bytes: bool,
    flag_lcp: bool,
    flag_lcs: bool,
    flag_searcher: bool,
}

type Result<T> = result::Result<T, Box<error::Error + Send + Sync>>;

fn main() {
    let mut args: Args = Docopt::new(USAGE)
        .and_then(|d| d.deserialize())
        .unwrap_or_else(|e| e.exit());
    if args.flag_dfa_reverse {
        args.flag_dfa = true;
    }
    match run(&args) {
        Ok(_) => process::exit(0),
        Err(err) => {
            let _ = writeln!(&mut io::stderr(), "{}", err);
            process::exit(1)
        }
    }
}

fn run(args: &Args) -> Result<()> {
    if args.cmd_ast {
        cmd_ast(args)
    } else if args.cmd_hir {
        cmd_hir(args)
    } else if args.cmd_prefixes {
        cmd_literals(args)
    } else if args.cmd_suffixes {
        cmd_literals(args)
    } else if args.cmd_anchors {
        cmd_anchors(args)
    } else if args.cmd_captures {
        cmd_captures(args)
    } else if args.cmd_compile {
        cmd_compile(args)
    } else if args.cmd_utf8_ranges {
        cmd_utf8_ranges(args)
    } else {
        unreachable!()
    }
}

fn cmd_ast(args: &Args) -> Result<()> {
    use syntax::ast::parse::Parser;

    let mut parser = Parser::new();
    let ast = parser.parse(&args.arg_pattern)?;
    println!("{:#?}", ast);
    Ok(())
}

fn cmd_hir(args: &Args) -> Result<()> {
    use syntax::ParserBuilder;

    let mut parser = ParserBuilder::new().allow_invalid_utf8(false).build();
    let hir = parser.parse(&args.arg_pattern)?;
    println!("{:#?}", hir);
    Ok(())
}

fn cmd_literals(args: &Args) -> Result<()> {
    let exprs = args.parse_many()?;
    let mut lits = if args.cmd_prefixes {
        args.literals(&exprs, |lits, e| lits.union_prefixes(e))
    } else {
        args.literals(&exprs, |lits, e| lits.union_suffixes(e))
    };
    if !args.flag_all_literals {
        if args.cmd_prefixes {
            lits = lits.unambiguous_prefixes();
        } else {
            lits = lits.unambiguous_suffixes();
        }
    }
    if args.flag_searcher {
        if args.cmd_prefixes {
            println!("{:?}", LiteralSearcher::prefixes(lits))
        } else {
            println!("{:?}", LiteralSearcher::suffixes(lits))
        }
    } else if args.flag_lcp {
        println!("{}", escape_unicode(lits.longest_common_prefix()));
    } else if args.flag_lcs {
        println!("{}", escape_unicode(lits.longest_common_suffix()));
    } else {
        for lit in lits.literals() {
            if args.flag_literal_bytes {
                if lit.is_cut() {
                    println!("Cut({})", escape_bytes(lit));
                } else {
                    println!("Complete({})", escape_bytes(lit));
                }
            } else {
                println!("{:?}", lit);
            }
        }
    }
    Ok(())
}

fn cmd_anchors(args: &Args) -> Result<()> {
    let expr = args.parse_one()?;
    if expr.is_anchored_start() {
        println!("start");
    }
    if expr.is_anchored_end() {
        println!("end");
    }
    Ok(())
}

fn cmd_captures(args: &Args) -> Result<()> {
    let expr = args.parse_one()?;
    let prog = args.compiler().only_utf8(false).compile(&[expr])?;
    for (i, name) in prog.captures.iter().enumerate() {
        match *name {
            None => println!("{}", i),
            Some(ref name) => println!("{}:{}", i, name),
        }
    }
    Ok(())
}

fn cmd_compile(args: &Args) -> Result<()> {
    let exprs = args.parse_many()?;
    let compiler = args
        .compiler()
        .bytes(args.flag_bytes)
        .only_utf8(!args.flag_bytes)
        .dfa(args.flag_dfa)
        .reverse(args.flag_dfa_reverse);
    let prog = compiler.compile(&exprs)?;
    print!("{:?}", prog);
    Ok(())
}

fn cmd_utf8_ranges(args: &Args) -> Result<()> {
    use syntax::hir::{self, HirKind};
    use syntax::utf8::Utf8Sequences;
    use syntax::ParserBuilder;

    let hir = ParserBuilder::new()
        .build()
        .parse(&format!("[{}]", args.arg_class))?;
    let cls = match hir.into_kind() {
        HirKind::Class(hir::Class::Unicode(cls)) => cls,
        _ => {
            return Err(
                format!("unexpected HIR, expected Unicode class").into()
            )
        }
    };
    let mut char_count = 0;
    for (i, range) in cls.iter().enumerate() {
        if i > 0 {
            println!("----------------------------");
        }
        char_count += (range.end() as u32) - (range.start() as u32) + 1;
        for seq in Utf8Sequences::new(range.start(), range.end()) {
            for utf8_range in seq.into_iter() {
                print!("[{:02X}-{:02X}]", utf8_range.start, utf8_range.end);
            }
            println!();
        }
    }
    println!("codepoint count: {}", char_count);
    Ok(())
}

impl Args {
    fn parse_one(&self) -> Result<Hir> {
        parse(&self.arg_pattern)
    }

    fn parse_many(&self) -> Result<Vec<Hir>> {
        self.arg_patterns.iter().map(|s| parse(s)).collect()
    }

    fn literals<F: Fn(&mut Literals, &Hir) -> bool>(
        &self,
        exprs: &[Hir],
        get_literals: F,
    ) -> Literals {
        let mut lits = Some(self.empty_literals());
        for e in exprs {
            lits = lits.and_then(|mut lits| {
                if !get_literals(&mut lits, e) {
                    None
                } else {
                    Some(lits)
                }
            });
        }
        lits.unwrap_or(self.empty_literals())
    }

    fn empty_literals(&self) -> Literals {
        let mut lits = Literals::empty();
        lits.set_limit_size(self.flag_literal_limit);
        lits.set_limit_class(self.flag_class_limit);
        lits
    }

    fn compiler(&self) -> Compiler {
        Compiler::new().size_limit(self.flag_size_limit)
    }
}

fn parse(re: &str) -> Result<Hir> {
    use syntax::ParserBuilder;
    ParserBuilder::new()
        .allow_invalid_utf8(true)
        .build()
        .parse(re)
        .map_err(From::from)
}

fn escape_unicode(bytes: &[u8]) -> String {
    let show = match ::std::str::from_utf8(bytes) {
        Ok(v) => v.to_string(),
        Err(_) => escape_bytes(bytes),
    };
    let mut space_escaped = String::new();
    for c in show.chars() {
        if c.is_whitespace() {
            let escaped = if c as u32 <= 0x7F {
                escape_byte(c as u8)
            } else {
                if c as u32 <= 0xFFFF {
                    format!(r"\u{{{:04x}}}", c as u32)
                } else {
                    format!(r"\U{{{:08x}}}", c as u32)
                }
            };
            space_escaped.push_str(&escaped);
        } else {
            space_escaped.push(c);
        }
    }
    space_escaped
}

fn escape_bytes(bytes: &[u8]) -> String {
    let mut s = String::new();
    for &b in bytes {
        s.push_str(&escape_byte(b));
    }
    s
}

fn escape_byte(byte: u8) -> String {
    use std::ascii::escape_default;

    let escaped: Vec<u8> = escape_default(byte).collect();
    String::from_utf8_lossy(&escaped).into_owned()
}
<commit_msg>debug: add utf8-ranges-rev<commit_after>extern crate docopt;
extern crate regex;
extern crate regex_syntax as syntax;
extern crate serde;
#[macro_use]
extern crate serde_derive;

use std::error;
use std::io::{self, Write};
use std::process;
use std::result;

use docopt::Docopt;
use regex::internal::{Compiler, LiteralSearcher};
use syntax::hir::literal::Literals;
use syntax::hir::Hir;

const USAGE: &'static str = "
Usage:
    regex-debug [options] ast <pattern>
    regex-debug [options] hir <pattern>
    regex-debug [options] prefixes <patterns> ...
    regex-debug [options] suffixes <patterns> ...
    regex-debug [options] anchors <pattern>
    regex-debug [options] captures <pattern>
    regex-debug [options] compile <patterns> ...
    regex-debug [options] utf8-ranges <class>
    regex-debug [options] utf8-ranges-rev <class>
    regex-debug --help

Options:
    --help               Show this usage message.
    --size-limit ARG     An approximate size limit on the total size (in bytes)
                         of a compiled regular expression program.
                         [default: 10485760]
    --bytes              Show the instruction codes for byte oriented programs.
                         (As opposed to Unicode oriented programs.)
    --dfa                Show the instruction codes for a DFA.
    --dfa-reverse        Show the instruction codes for a reverse DFA.
                         This implies --dfa.
    -a, --all-literals   Shows all literals extracted.
                         By default, only unambiguous literals are shown.
    --literal-limit ARG  An approximate limit on the total size (in bytes)
                         of all literals extracted. [default: 250]
    --class-limit ARG    A limit on the size of character classes used to
                         extract literals. [default: 10]
    --literal-bytes      Show raw literal bytes instead of Unicode chars.
    --lcp                Show the longest common prefix of all the literals
                         extracted.
    --lcs                Show the longest common suffix of all the literals
                         extracted.
    --searcher           Show the debug output for the literal searcher
                         constructed by the literals found.
";

#[derive(Deserialize)]
struct Args {
    cmd_ast: bool,
    cmd_hir: bool,
    cmd_prefixes: bool,
    cmd_suffixes: bool,
    cmd_anchors: bool,
    cmd_captures: bool,
    cmd_compile: bool,
    cmd_utf8_ranges: bool,
    cmd_utf8_ranges_rev: bool,

    arg_pattern: String,
    arg_patterns: Vec<String>,
    arg_class: String,

    flag_size_limit: usize,
    flag_bytes: bool,
    flag_dfa: bool,
    flag_dfa_reverse: bool,
    flag_all_literals: bool,
    flag_literal_limit: usize,
    flag_class_limit: usize,
    flag_literal_bytes: bool,
    flag_lcp: bool,
    flag_lcs: bool,
    flag_searcher: bool,
}

type Result<T> = result::Result<T, Box<dyn error::Error + Send + Sync>>;

fn main() {
    let mut args: Args = Docopt::new(USAGE)
        .and_then(|d| d.deserialize())
        .unwrap_or_else(|e| e.exit());
    if args.flag_dfa_reverse {
        args.flag_dfa = true;
    }
    match run(&args) {
        Ok(_) => process::exit(0),
        Err(err) => {
            let _ = writeln!(&mut io::stderr(), "{}", err);
            process::exit(1)
        }
    }
}

fn run(args: &Args) -> Result<()> {
    if args.cmd_ast {
        cmd_ast(args)
    } else if args.cmd_hir {
        cmd_hir(args)
    } else if args.cmd_prefixes {
        cmd_literals(args)
    } else if args.cmd_suffixes {
        cmd_literals(args)
    } else if args.cmd_anchors {
        cmd_anchors(args)
    } else if args.cmd_captures {
        cmd_captures(args)
    } else if args.cmd_compile {
        cmd_compile(args)
    } else if args.cmd_utf8_ranges {
        cmd_utf8_ranges(args)
    } else if args.cmd_utf8_ranges_rev {
        cmd_utf8_ranges_rev(args)
    } else {
        unreachable!()
    }
}

fn cmd_ast(args: &Args) -> Result<()> {
    use syntax::ast::parse::Parser;

    let mut parser = Parser::new();
    let ast = parser.parse(&args.arg_pattern)?;
    println!("{:#?}", ast);
    Ok(())
}

fn cmd_hir(args: &Args) -> Result<()> {
    use syntax::ParserBuilder;

    let mut parser = ParserBuilder::new().allow_invalid_utf8(false).build();
    let hir = parser.parse(&args.arg_pattern)?;
    println!("{:#?}", hir);
    Ok(())
}

fn cmd_literals(args: &Args) -> Result<()> {
    let exprs = args.parse_many()?;
    let mut lits = if args.cmd_prefixes {
        args.literals(&exprs, |lits, e| lits.union_prefixes(e))
    } else {
        args.literals(&exprs, |lits, e| lits.union_suffixes(e))
    };
    if !args.flag_all_literals {
        if args.cmd_prefixes {
            lits = lits.unambiguous_prefixes();
        } else {
            lits = lits.unambiguous_suffixes();
        }
    }
    if args.flag_searcher {
        if args.cmd_prefixes {
            println!("{:?}", LiteralSearcher::prefixes(lits))
        } else {
            println!("{:?}", LiteralSearcher::suffixes(lits))
        }
    } else if args.flag_lcp {
        println!("{}", escape_unicode(lits.longest_common_prefix()));
    } else if args.flag_lcs {
        println!("{}", escape_unicode(lits.longest_common_suffix()));
    } else {
        for lit in lits.literals() {
            if args.flag_literal_bytes {
                if lit.is_cut() {
                    println!("Cut({})", escape_bytes(lit));
                } else {
                    println!("Complete({})", escape_bytes(lit));
                }
            } else {
                println!("{:?}", lit);
            }
        }
    }
    Ok(())
}

fn cmd_anchors(args: &Args) -> Result<()> {
    let expr = args.parse_one()?;
    if expr.is_anchored_start() {
        println!("start");
    }
    if expr.is_anchored_end() {
        println!("end");
    }
    Ok(())
}

fn cmd_captures(args: &Args) -> Result<()> {
    let expr = args.parse_one()?;
    let prog = args.compiler().only_utf8(false).compile(&[expr])?;
    for (i, name) in prog.captures.iter().enumerate() {
        match *name {
            None => println!("{}", i),
            Some(ref name) => println!("{}:{}", i, name),
        }
    }
    Ok(())
}

fn cmd_compile(args: &Args) -> Result<()> {
    let exprs = args.parse_many()?;
    let compiler = args
        .compiler()
        .bytes(args.flag_bytes)
        .only_utf8(!args.flag_bytes)
        .dfa(args.flag_dfa)
        .reverse(args.flag_dfa_reverse);
    let prog = compiler.compile(&exprs)?;
    print!("{:?}", prog);
    Ok(())
}

fn cmd_utf8_ranges(args: &Args) -> Result<()> {
    use syntax::hir::{self, HirKind};
    use syntax::utf8::Utf8Sequences;
    use syntax::ParserBuilder;

    let hir = ParserBuilder::new()
        .build()
        .parse(&format!("[{}]", args.arg_class))?;
    let cls = match hir.into_kind() {
        HirKind::Class(hir::Class::Unicode(cls)) => cls,
        _ => {
            return Err(
                format!("unexpected HIR, expected Unicode class").into()
            )
        }
    };
    let mut char_count = 0;
    for (i, range) in cls.iter().enumerate() {
        if i > 0 {
            println!("----------------------------");
        }
        char_count += (range.end() as u32) - (range.start() as u32) + 1;
        for seq in Utf8Sequences::new(range.start(), range.end()) {
            for utf8_range in seq.into_iter() {
                print!("[{:02X}-{:02X}]", utf8_range.start, utf8_range.end);
            }
            println!();
        }
    }
    println!("codepoint count: {}", char_count);
    Ok(())
}

fn cmd_utf8_ranges_rev(args: &Args) -> Result<()> {
    use syntax::hir::{self, HirKind};
    use syntax::utf8::Utf8Sequences;
    use syntax::ParserBuilder;

    let hir = ParserBuilder::new()
        .build()
        .parse(&format!("[{}]", args.arg_class))?;
    let cls = match hir.into_kind() {
        HirKind::Class(hir::Class::Unicode(cls)) => cls,
        _ => {
            return Err(
                format!("unexpected HIR, expected Unicode class").into()
            )
        }
    };
    let mut char_count = 0;
    let mut seqs = vec![];
    for (_, range) in cls.iter().enumerate() {
        char_count += (range.end() as u32) - (range.start() as u32) + 1;
        for seq in Utf8Sequences::new(range.start(), range.end()) {
            let mut seq = seq.as_slice().to_vec();
            seq.reverse();
            seqs.push(seq);
        }
    }
    seqs.sort();
    for seq in seqs {
        for utf8_range in seq.into_iter() {
            print!("[{:02X}-{:02X}]", utf8_range.start, utf8_range.end);
        }
        println!();
    }
    println!("codepoint count: {}", char_count);
    Ok(())
}

impl Args {
    fn parse_one(&self) -> Result<Hir> {
        parse(&self.arg_pattern)
    }

    fn parse_many(&self) -> Result<Vec<Hir>> {
        self.arg_patterns.iter().map(|s| parse(s)).collect()
    }

    fn literals<F: Fn(&mut Literals, &Hir) -> bool>(
        &self,
        exprs: &[Hir],
        get_literals: F,
    ) -> Literals {
        let mut lits = Some(self.empty_literals());
        for e in exprs {
            lits = lits.and_then(|mut lits| {
                if !get_literals(&mut lits, e) {
                    None
                } else {
                    Some(lits)
                }
            });
        }
        lits.unwrap_or(self.empty_literals())
    }

    fn empty_literals(&self) -> Literals {
        let mut lits = Literals::empty();
        lits.set_limit_size(self.flag_literal_limit);
        lits.set_limit_class(self.flag_class_limit);
        lits
    }

    fn compiler(&self) -> Compiler {
        Compiler::new().size_limit(self.flag_size_limit)
    }
}

fn parse(re: &str) -> Result<Hir> {
    use syntax::ParserBuilder;
    ParserBuilder::new()
        .allow_invalid_utf8(true)
        .build()
        .parse(re)
        .map_err(From::from)
}

fn escape_unicode(bytes: &[u8]) -> String {
    let show = match ::std::str::from_utf8(bytes) {
        Ok(v) => v.to_string(),
        Err(_) => escape_bytes(bytes),
    };
    let mut space_escaped = String::new();
    for c in show.chars() {
        if c.is_whitespace() {
            let escaped = if c as u32 <= 0x7F {
                escape_byte(c as u8)
            } else {
                if c as u32 <= 0xFFFF {
                    format!(r"\u{{{:04x}}}", c as u32)
                } else {
                    format!(r"\U{{{:08x}}}", c as u32)
                }
            };
            space_escaped.push_str(&escaped);
        } else {
            space_escaped.push(c);
        }
    }
    space_escaped
}

fn escape_bytes(bytes: &[u8]) -> String {
    let mut s = String::new();
    for &b in bytes {
        s.push_str(&escape_byte(b));
    }
    s
}

fn escape_byte(byte: u8) -> String {
    use std::ascii::escape_default;

    let escaped: Vec<u8> = escape_default(byte).collect();
    String::from_utf8_lossy(&escaped).into_owned()
}
<|endoftext|>