<commit_before>//! Simple hierarchical profiler
use std::{
    cell::RefCell,
    collections::{BTreeMap, HashSet},
    io::{stderr, Write},
    sync::{
        atomic::{AtomicBool, Ordering},
        RwLock,
    },
    time::{Duration, Instant},
};

use once_cell::sync::Lazy;

/// Filtering syntax
/// env RA_PROFILE=*             // dump everything
/// env RA_PROFILE=foo|bar|baz   // enabled only selected entries
/// env RA_PROFILE=*@3>10        // dump everything, up to depth 3, if it takes more than 10 ms
pub fn init() {
    let spec = std::env::var("RA_PROFILE").unwrap_or_default();
    init_from(&spec);
}

pub fn init_from(spec: &str) {
    let filter = if spec.is_empty() { Filter::disabled() } else { Filter::from_spec(spec) };
    filter.install();
}

pub type Label = &'static str;

/// This function starts a profiling scope in the current execution stack with a given description.
/// It returns a Profile structure and measure elapsed time between this method invocation and Profile structure drop.
/// It supports nested profiling scopes in case when this function invoked multiple times at the execution stack. In this case the profiling information will be nested at the output.
/// Profiling information is being printed in the stderr.
///
/// # Example
/// ```
/// use ra_prof::{profile, set_filter, Filter};
///
/// let f = Filter::from_spec("profile1|profile2@2");
/// set_filter(f);
/// profiling_function1();
///
/// fn profiling_function1() {
///     let _p = profile("profile1");
///     profiling_function2();
/// }
///
/// fn profiling_function2() {
///     let _p = profile("profile2");
/// }
/// ```
/// This will print in the stderr the following:
/// ```text
///  0ms - profile
///      0ms - profile2
/// ```
pub fn profile(label: Label) -> Profiler {
    assert!(!label.is_empty());
    let enabled = PROFILING_ENABLED.load(Ordering::Relaxed)
        && PROFILE_STACK.with(|stack| stack.borrow_mut().push(label));
    let label = if enabled { Some(label) } else { None };
    Profiler { label, detail: None }
}

pub struct Profiler {
    label: Option<Label>,
    detail: Option<String>,
}

impl Profiler {
    pub fn detail(mut self, detail: impl FnOnce() -> String) -> Profiler {
        if self.label.is_some() {
            self.detail = Some(detail())
        }
        self
    }
}

static PROFILING_ENABLED: AtomicBool = AtomicBool::new(false);
static FILTER: Lazy<RwLock<Filter>> = Lazy::new(Default::default);
thread_local!(static PROFILE_STACK: RefCell<ProfileStack> = RefCell::new(ProfileStack::new()));

#[derive(Default, Clone, Debug)]
struct Filter {
    depth: usize,
    allowed: HashSet<String>,
    longer_than: Duration,
    version: usize,
}

impl Filter {
    fn new(depth: usize, allowed: HashSet<String>, longer_than: Duration) -> Filter {
        Filter { depth, allowed, longer_than, version: 0 }
    }

    fn disabled() -> Filter {
        Filter::default()
    }

    fn from_spec(mut spec: &str) -> Filter {
        let longer_than = if let Some(idx) = spec.rfind('>') {
            let longer_than = spec[idx + 1..].parse().expect("invalid profile longer_than");
            spec = &spec[..idx];
            Duration::from_millis(longer_than)
        } else {
            Duration::new(0, 0)
        };

        let depth = if let Some(idx) = spec.rfind('@') {
            let depth: usize = spec[idx + 1..].parse().expect("invalid profile depth");
            spec = &spec[..idx];
            depth
        } else {
            999
        };
        let allowed =
            if spec == "*" { HashSet::new() } else { spec.split('|').map(String::from).collect() };
        Filter::new(depth, allowed, longer_than)
    }

    fn install(mut self) {
        PROFILING_ENABLED.store(self.depth > 0, Ordering::SeqCst);
        let mut old = FILTER.write().unwrap();
        self.version = old.version + 1;
        *old = self;
    }
}

struct ProfileStack {
    starts: Vec<Instant>,
    messages: Vec<Message>,
    filter: Filter,
}

struct Message {
    level: usize,
    duration: Duration,
    label: Label,
    detail: Option<String>,
}

impl ProfileStack {
    fn new() -> ProfileStack {
        ProfileStack { starts: Vec::new(), messages: Vec::new(), filter: Default::default() }
    }

    fn push(&mut self, label: Label) -> bool {
        if self.starts.is_empty() {
            if let Ok(f) = FILTER.try_read() {
                if f.version > self.filter.version {
                    self.filter = f.clone();
                }
            };
        }
        if self.starts.len() > self.filter.depth {
            return false;
        }
        let allowed = &self.filter.allowed;
        if self.starts.is_empty() && !allowed.is_empty() && !allowed.contains(label) {
            return false;
        }

        self.starts.push(Instant::now());
        true
    }

    pub fn pop(&mut self, label: Label, detail: Option<String>) {
        let start = self.starts.pop().unwrap();
        let duration = start.elapsed();
        let level = self.starts.len();
        self.messages.push(Message { level, duration, label, detail });
        if level == 0 {
            let stdout = stderr();
            let longer_than = self.filter.longer_than;
            // Convert to millis for comparison to avoid problems with rounding
            // (otherwise we could print `0ms` despite user's `>0` filter when
            // `duration` is just a few nanos).
            if duration.as_millis() > longer_than.as_millis() {
                print(&self.messages, longer_than, &mut stdout.lock());
            }
            self.messages.clear();
        }
    }
}

impl Drop for Profiler {
    fn drop(&mut self) {
        match self {
            Profiler { label: Some(label), detail } => {
                PROFILE_STACK.with(|stack| {
                    stack.borrow_mut().pop(label, detail.take());
                });
            }
            Profiler { label: None, .. } => (),
        }
    }
}

fn print(msgs: &[Message], longer_than: Duration, out: &mut impl Write) {
    if msgs.is_empty() {
        return;
    }
    let children_map = idx_to_children(msgs);
    let root_idx = msgs.len() - 1;
    print_for_idx(root_idx, &children_map, msgs, longer_than, out);
}

fn print_for_idx(
    current_idx: usize,
    children_map: &[Vec<usize>],
    msgs: &[Message],
    longer_than: Duration,
    out: &mut impl Write,
) {
    let current = &msgs[current_idx];
    let current_indent = "    ".repeat(current.level);
    let detail = current.detail.as_ref().map(|it| format!(" @ {}", it)).unwrap_or_default();
    writeln!(
        out,
        "{}{:5}ms - {}{}",
        current_indent,
        current.duration.as_millis(),
        current.label,
        detail,
    )
    .expect("printing profiling info");

    let longer_than_millis = longer_than.as_millis();
    let children_indices = &children_map[current_idx];
    let mut accounted_for = Duration::default();
    let mut short_children = BTreeMap::new(); // Use `BTreeMap` to get deterministic output.

    for child_idx in children_indices.iter() {
        let child = &msgs[*child_idx];
        if child.duration.as_millis() > longer_than_millis {
            print_for_idx(*child_idx, children_map, msgs, longer_than, out);
        } else {
            let pair = short_children.entry(child.label).or_insert((Duration::default(), 0));
            pair.0 += child.duration;
            pair.1 += 1;
        }
        accounted_for += child.duration;
    }

    for (child_msg, (duration, count)) in short_children.iter() {
        let millis = duration.as_millis();
        writeln!(out, "    {}{:5}ms - {} ({} calls)", current_indent, millis, child_msg, count)
            .expect("printing profiling info");
    }

    let unaccounted_millis = (current.duration - accounted_for).as_millis();
    if !children_indices.is_empty()
        && unaccounted_millis > 0
        && unaccounted_millis > longer_than_millis
    {
        writeln!(out, "    {}{:5}ms - ???", current_indent, unaccounted_millis)
            .expect("printing profiling info");
    }
}

/// Returns a mapping from an index in the `msgs` to the vector with the indices of its children.
///
/// This assumes that the entries in `msgs` are in the order of when the calls to `profile` finish.
/// In other words, a postorder of the call graph. In particular, the root is the last element of
/// `msgs`.
fn idx_to_children(msgs: &[Message]) -> Vec<Vec<usize>> {
    // Initialize with the index of the root; `msgs` and `ancestors` should be never empty.
    assert!(!msgs.is_empty());
    let mut ancestors = vec![msgs.len() - 1];
    let mut result: Vec<Vec<usize>> = vec![vec![]; msgs.len()];
    for (idx, msg) in msgs[..msgs.len() - 1].iter().enumerate().rev() {
        // We need to find the parent of the current message, i.e., the last ancestor that has a
        // level lower than the current message.
        while msgs[*ancestors.last().unwrap()].level >= msg.level {
            ancestors.pop();
        }
        result[*ancestors.last().unwrap()].push(idx);
        ancestors.push(idx);
    }
    // Note that above we visited all children from the last to the first one. Let's reverse vectors
    // to get the more natural order where the first element is the first child.
    for vec in result.iter_mut() {
        vec.reverse();
    }
    result
}

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

    #[test]
    fn test_basic_profile() {
        let s = vec!["profile1".to_string(), "profile2".to_string()];
        let f = Filter::new(2, s, Duration::new(0, 0));
        set_filter(f);
        profiling_function1();
    }

    fn profiling_function1() {
        let _p = profile("profile1");
        profiling_function2();
    }

    fn profiling_function2() {
        let _p = profile("profile2");
    }

    #[test]
    fn test_longer_than() {
        let mut result = vec![];
        let msgs = vec![
            Message { level: 1, duration: Duration::from_nanos(3), label: "bar", detail: None },
            Message { level: 1, duration: Duration::from_nanos(2), label: "bar", detail: None },
            Message { level: 0, duration: Duration::from_millis(1), label: "foo", detail: None },
        ];
        print(&msgs, Duration::from_millis(0), &mut result);
        // The calls to `bar` are so short that they'll be rounded to 0ms and should get collapsed
        // when printing.
        assert_eq!(
            std::str::from_utf8(&result).unwrap(),
            "    1ms - foo\n        0ms - bar (2 calls)\n"
        );
    }

    #[test]
    fn test_unaccounted_for_topmost() {
        let mut result = vec![];
        let msgs = vec![
            Message { level: 1, duration: Duration::from_millis(2), label: "bar", detail: None },
            Message { level: 0, duration: Duration::from_millis(5), label: "foo", detail: None },
        ];
        print(&msgs, Duration::from_millis(0), &mut result);
        assert_eq!(
            std::str::from_utf8(&result).unwrap().lines().collect::<Vec<_>>(),
            vec![
                "    5ms - foo",
                "        2ms - bar",
                "        3ms - ???",
                // Dummy comment to improve formatting
            ]
        );
    }

    #[test]
    fn test_unaccounted_for_multiple_levels() {
        let mut result = vec![];
        let msgs = vec![
            Message { level: 2, duration: Duration::from_millis(3), label: "baz", detail: None },
            Message { level: 1, duration: Duration::from_millis(5), label: "bar", detail: None },
            Message { level: 2, duration: Duration::from_millis(2), label: "baz", detail: None },
            Message { level: 1, duration: Duration::from_millis(4), label: "bar", detail: None },
            Message { level: 0, duration: Duration::from_millis(9), label: "foo", detail: None },
        ];
        print(&msgs, Duration::from_millis(0), &mut result);
        assert_eq!(
            std::str::from_utf8(&result).unwrap().lines().collect::<Vec<_>>(),
            vec![
                "    9ms - foo",
                "        5ms - bar",
                "            3ms - baz",
                "            2ms - ???",
                "        4ms - bar",
                "            2ms - baz",
                "            2ms - ???",
            ]
        );
    }
}
<commit_msg>Simplify<commit_after>//! Simple hierarchical profiler
use std::{
    cell::RefCell,
    collections::{BTreeMap, HashSet},
    io::{stderr, Write},
    sync::{
        atomic::{AtomicBool, Ordering},
        RwLock,
    },
    time::{Duration, Instant},
};

use once_cell::sync::Lazy;

/// Filtering syntax
/// env RA_PROFILE=*             // dump everything
/// env RA_PROFILE=foo|bar|baz   // enabled only selected entries
/// env RA_PROFILE=*@3>10        // dump everything, up to depth 3, if it takes more than 10 ms
pub fn init() {
    let spec = std::env::var("RA_PROFILE").unwrap_or_default();
    init_from(&spec);
}

pub fn init_from(spec: &str) {
    let filter = if spec.is_empty() { Filter::disabled() } else { Filter::from_spec(spec) };
    filter.install();
}

pub type Label = &'static str;

/// This function starts a profiling scope in the current execution stack with a given description.
/// It returns a Profile structure and measure elapsed time between this method invocation and Profile structure drop.
/// It supports nested profiling scopes in case when this function invoked multiple times at the execution stack. In this case the profiling information will be nested at the output.
/// Profiling information is being printed in the stderr.
///
/// # Example
/// ```
/// use ra_prof::{profile, set_filter, Filter};
///
/// let f = Filter::from_spec("profile1|profile2@2");
/// set_filter(f);
/// profiling_function1();
///
/// fn profiling_function1() {
///     let _p = profile("profile1");
///     profiling_function2();
/// }
///
/// fn profiling_function2() {
///     let _p = profile("profile2");
/// }
/// ```
/// This will print in the stderr the following:
/// ```text
///  0ms - profile
///      0ms - profile2
/// ```
pub fn profile(label: Label) -> Profiler {
    assert!(!label.is_empty());
    let enabled = PROFILING_ENABLED.load(Ordering::Relaxed)
        && PROFILE_STACK.with(|stack| stack.borrow_mut().push(label));
    let label = if enabled { Some(label) } else { None };
    Profiler { label, detail: None }
}

pub struct Profiler {
    label: Option<Label>,
    detail: Option<String>,
}

impl Profiler {
    pub fn detail(mut self, detail: impl FnOnce() -> String) -> Profiler {
        if self.label.is_some() {
            self.detail = Some(detail())
        }
        self
    }
}

impl Drop for Profiler {
    fn drop(&mut self) {
        match self {
            Profiler { label: Some(label), detail } => {
                PROFILE_STACK.with(|stack| {
                    stack.borrow_mut().pop(label, detail.take());
                });
            }
            Profiler { label: None, .. } => (),
        }
    }
}

static PROFILING_ENABLED: AtomicBool = AtomicBool::new(false);
static FILTER: Lazy<RwLock<Filter>> = Lazy::new(Default::default);
thread_local!(static PROFILE_STACK: RefCell<ProfileStack> = RefCell::new(ProfileStack::new()));

#[derive(Default, Clone, Debug)]
struct Filter {
    depth: usize,
    allowed: HashSet<String>,
    longer_than: Duration,
    version: usize,
}

impl Filter {
    fn disabled() -> Filter {
        Filter::default()
    }

    fn from_spec(mut spec: &str) -> Filter {
        let longer_than = if let Some(idx) = spec.rfind('>') {
            let longer_than = spec[idx + 1..].parse().expect("invalid profile longer_than");
            spec = &spec[..idx];
            Duration::from_millis(longer_than)
        } else {
            Duration::new(0, 0)
        };

        let depth = if let Some(idx) = spec.rfind('@') {
            let depth: usize = spec[idx + 1..].parse().expect("invalid profile depth");
            spec = &spec[..idx];
            depth
        } else {
            999
        };
        let allowed =
            if spec == "*" { HashSet::new() } else { spec.split('|').map(String::from).collect() };
        Filter { depth, allowed, longer_than, version: 0 }
    }

    fn install(mut self) {
        PROFILING_ENABLED.store(self.depth > 0, Ordering::SeqCst);
        let mut old = FILTER.write().unwrap();
        self.version = old.version + 1;
        *old = self;
    }
}

struct ProfileStack {
    starts: Vec<Instant>,
    messages: Vec<Message>,
    filter: Filter,
}

struct Message {
    level: usize,
    duration: Duration,
    label: Label,
    detail: Option<String>,
}

impl ProfileStack {
    fn new() -> ProfileStack {
        ProfileStack { starts: Vec::new(), messages: Vec::new(), filter: Default::default() }
    }

    fn push(&mut self, label: Label) -> bool {
        if self.starts.is_empty() {
            if let Ok(f) = FILTER.try_read() {
                if f.version > self.filter.version {
                    self.filter = f.clone();
                }
            };
        }
        if self.starts.len() > self.filter.depth {
            return false;
        }
        let allowed = &self.filter.allowed;
        if self.starts.is_empty() && !allowed.is_empty() && !allowed.contains(label) {
            return false;
        }

        self.starts.push(Instant::now());
        true
    }

    pub fn pop(&mut self, label: Label, detail: Option<String>) {
        let start = self.starts.pop().unwrap();
        let duration = start.elapsed();
        let level = self.starts.len();
        self.messages.push(Message { level, duration, label, detail });
        if level == 0 {
            let longer_than = self.filter.longer_than;
            // Convert to millis for comparison to avoid problems with rounding
            // (otherwise we could print `0ms` despite user's `>0` filter when
            // `duration` is just a few nanos).
            if duration.as_millis() > longer_than.as_millis() {
                let stderr = stderr();
                print(&self.messages, longer_than, &mut stderr.lock());
            }
            self.messages.clear();
            assert!(self.starts.is_empty())
        }
    }
}

fn print(msgs: &[Message], longer_than: Duration, out: &mut impl Write) {
    if msgs.is_empty() {
        return;
    }
    let children_map = idx_to_children(msgs);
    let root_idx = msgs.len() - 1;
    print_for_idx(root_idx, &children_map, msgs, longer_than, out);
}

fn print_for_idx(
    current_idx: usize,
    children_map: &[Vec<usize>],
    msgs: &[Message],
    longer_than: Duration,
    out: &mut impl Write,
) {
    let current = &msgs[current_idx];
    let current_indent = "    ".repeat(current.level);
    let detail = current.detail.as_ref().map(|it| format!(" @ {}", it)).unwrap_or_default();
    writeln!(
        out,
        "{}{:5}ms - {}{}",
        current_indent,
        current.duration.as_millis(),
        current.label,
        detail,
    )
    .expect("printing profiling info");

    let longer_than_millis = longer_than.as_millis();
    let children_indices = &children_map[current_idx];
    let mut accounted_for = Duration::default();
    let mut short_children = BTreeMap::new(); // Use `BTreeMap` to get deterministic output.

    for child_idx in children_indices.iter() {
        let child = &msgs[*child_idx];
        if child.duration.as_millis() > longer_than_millis {
            print_for_idx(*child_idx, children_map, msgs, longer_than, out);
        } else {
            let pair = short_children.entry(child.label).or_insert((Duration::default(), 0));
            pair.0 += child.duration;
            pair.1 += 1;
        }
        accounted_for += child.duration;
    }

    for (child_msg, (duration, count)) in short_children.iter() {
        let millis = duration.as_millis();
        writeln!(out, "    {}{:5}ms - {} ({} calls)", current_indent, millis, child_msg, count)
            .expect("printing profiling info");
    }

    let unaccounted_millis = (current.duration - accounted_for).as_millis();
    if !children_indices.is_empty()
        && unaccounted_millis > 0
        && unaccounted_millis > longer_than_millis
    {
        writeln!(out, "    {}{:5}ms - ???", current_indent, unaccounted_millis)
            .expect("printing profiling info");
    }
}

/// Returns a mapping from an index in the `msgs` to the vector with the indices of its children.
///
/// This assumes that the entries in `msgs` are in the order of when the calls to `profile` finish.
/// In other words, a postorder of the call graph. In particular, the root is the last element of
/// `msgs`.
fn idx_to_children(msgs: &[Message]) -> Vec<Vec<usize>> {
    // Initialize with the index of the root; `msgs` and `ancestors` should be never empty.
    assert!(!msgs.is_empty());
    let mut ancestors = vec![msgs.len() - 1];
    let mut result: Vec<Vec<usize>> = vec![vec![]; msgs.len()];
    for (idx, msg) in msgs[..msgs.len() - 1].iter().enumerate().rev() {
        // We need to find the parent of the current message, i.e., the last ancestor that has a
        // level lower than the current message.
        while msgs[*ancestors.last().unwrap()].level >= msg.level {
            ancestors.pop();
        }
        result[*ancestors.last().unwrap()].push(idx);
        ancestors.push(idx);
    }
    // Note that above we visited all children from the last to the first one. Let's reverse vectors
    // to get the more natural order where the first element is the first child.
    for vec in result.iter_mut() {
        vec.reverse();
    }
    result
}

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

    #[test]
    fn test_basic_profile() {
        let s = vec!["profile1".to_string(), "profile2".to_string()];
        let f = Filter::new(2, s, Duration::new(0, 0));
        set_filter(f);
        profiling_function1();
    }

    fn profiling_function1() {
        let _p = profile("profile1");
        profiling_function2();
    }

    fn profiling_function2() {
        let _p = profile("profile2");
    }

    #[test]
    fn test_longer_than() {
        let mut result = vec![];
        let msgs = vec![
            Message { level: 1, duration: Duration::from_nanos(3), label: "bar", detail: None },
            Message { level: 1, duration: Duration::from_nanos(2), label: "bar", detail: None },
            Message { level: 0, duration: Duration::from_millis(1), label: "foo", detail: None },
        ];
        print(&msgs, Duration::from_millis(0), &mut result);
        // The calls to `bar` are so short that they'll be rounded to 0ms and should get collapsed
        // when printing.
        assert_eq!(
            std::str::from_utf8(&result).unwrap(),
            "    1ms - foo\n        0ms - bar (2 calls)\n"
        );
    }

    #[test]
    fn test_unaccounted_for_topmost() {
        let mut result = vec![];
        let msgs = vec![
            Message { level: 1, duration: Duration::from_millis(2), label: "bar", detail: None },
            Message { level: 0, duration: Duration::from_millis(5), label: "foo", detail: None },
        ];
        print(&msgs, Duration::from_millis(0), &mut result);
        assert_eq!(
            std::str::from_utf8(&result).unwrap().lines().collect::<Vec<_>>(),
            vec![
                "    5ms - foo",
                "        2ms - bar",
                "        3ms - ???",
                // Dummy comment to improve formatting
            ]
        );
    }

    #[test]
    fn test_unaccounted_for_multiple_levels() {
        let mut result = vec![];
        let msgs = vec![
            Message { level: 2, duration: Duration::from_millis(3), label: "baz", detail: None },
            Message { level: 1, duration: Duration::from_millis(5), label: "bar", detail: None },
            Message { level: 2, duration: Duration::from_millis(2), label: "baz", detail: None },
            Message { level: 1, duration: Duration::from_millis(4), label: "bar", detail: None },
            Message { level: 0, duration: Duration::from_millis(9), label: "foo", detail: None },
        ];
        print(&msgs, Duration::from_millis(0), &mut result);
        assert_eq!(
            std::str::from_utf8(&result).unwrap().lines().collect::<Vec<_>>(),
            vec![
                "    9ms - foo",
                "        5ms - bar",
                "            3ms - baz",
                "            2ms - ???",
                "        4ms - bar",
                "            2ms - baz",
                "            2ms - ???",
            ]
        );
    }
}
<|endoftext|>


<commit_before>use fallible_iterator::FallibleIterator;
use futures::sink;
use futures::sync::mpsc;
use futures::{try_ready, Future, Poll, Sink, Stream};
use postgres_protocol::authentication;
use postgres_protocol::authentication::sasl::{self, ScramSha256};
use postgres_protocol::message::backend::Message;
use postgres_protocol::message::frontend;
use state_machine_future::{transition, RentToOwn, StateMachineFuture};
use std::collections::HashMap;
use tokio_codec::Framed;
use tokio_io::{AsyncRead, AsyncWrite};

use crate::proto::{Client, Connection, MaybeTlsStream, PostgresCodec, TlsFuture, FrontendMessage};
use crate::tls::ChannelBinding;
use crate::{Config, Error, TlsConnect};

#[derive(StateMachineFuture)]
pub enum ConnectRaw<S, T>
where
    S: AsyncRead + AsyncWrite,
    T: TlsConnect<S>,
{
    #[state_machine_future(start, transitions(SendingStartup))]
    Start {
        future: TlsFuture<S, T>,
        config: Config,
        idx: Option<usize>,
    },
    #[state_machine_future(transitions(ReadingAuth))]
    SendingStartup {
        future: sink::Send<Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>>,
        config: Config,
        idx: Option<usize>,
        channel_binding: ChannelBinding,
    },
    #[state_machine_future(transitions(ReadingInfo, SendingPassword, SendingSasl))]
    ReadingAuth {
        stream: Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>,
        config: Config,
        idx: Option<usize>,
        channel_binding: ChannelBinding,
    },
    #[state_machine_future(transitions(ReadingAuthCompletion))]
    SendingPassword {
        future: sink::Send<Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>>,
        config: Config,
        idx: Option<usize>,
    },
    #[state_machine_future(transitions(ReadingSasl))]
    SendingSasl {
        future: sink::Send<Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>>,
        scram: ScramSha256,
        config: Config,
        idx: Option<usize>,
    },
    #[state_machine_future(transitions(SendingSasl, ReadingAuthCompletion))]
    ReadingSasl {
        stream: Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>,
        scram: ScramSha256,
        config: Config,
        idx: Option<usize>,
    },
    #[state_machine_future(transitions(ReadingInfo))]
    ReadingAuthCompletion {
        stream: Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>,
        config: Config,
        idx: Option<usize>,
    },
    #[state_machine_future(transitions(Finished))]
    ReadingInfo {
        stream: Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>,
        process_id: i32,
        secret_key: i32,
        parameters: HashMap<String, String>,
        config: Config,
        idx: Option<usize>,
    },
    #[state_machine_future(ready)]
    Finished((Client, Connection<MaybeTlsStream<S, T::Stream>>)),
    #[state_machine_future(error)]
    Failed(Error),
}

impl<S, T> PollConnectRaw<S, T> for ConnectRaw<S, T>
where
    S: AsyncRead + AsyncWrite,
    T: TlsConnect<S>,
{
    fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start<S, T>>) -> Poll<AfterStart<S, T>, Error> {
        let (stream, channel_binding) = try_ready!(state.future.poll());
        let state = state.take();

        let mut params = vec![("client_encoding", "UTF8"), ("timezone", "GMT")];
        if let Some(user) = &state.config.0.user {
            params.push(("user", &**user));
        }
        if let Some(dbname) = &state.config.0.dbname {
            params.push(("database", &**dbname));
        }
        if let Some(options) = &state.config.0.options {
            params.push(("options", &**options));
        }
        if let Some(application_name) = &state.config.0.application_name {
            params.push(("application_name", &**application_name));
        }

        let mut buf = vec![];
        frontend::startup_message(params, &mut buf).map_err(Error::encode)?;

        let stream = Framed::new(stream, PostgresCodec);

        transition!(SendingStartup {
            future: stream.send(FrontendMessage::Raw(buf)),
            config: state.config,
            idx: state.idx,
            channel_binding,
        })
    }

    fn poll_sending_startup<'a>(
        state: &'a mut RentToOwn<'a, SendingStartup<S, T>>,
    ) -> Poll<AfterSendingStartup<S, T>, Error> {
        let stream = try_ready!(state.future.poll().map_err(Error::io));
        let state = state.take();
        transition!(ReadingAuth {
            stream,
            config: state.config,
            idx: state.idx,
            channel_binding: state.channel_binding,
        })
    }

    fn poll_reading_auth<'a>(
        state: &'a mut RentToOwn<'a, ReadingAuth<S, T>>,
    ) -> Poll<AfterReadingAuth<S, T>, Error> {
        let message = try_ready!(state.stream.poll().map_err(Error::io));
        let state = state.take();

        match message {
            Some(Message::AuthenticationOk) => transition!(ReadingInfo {
                stream: state.stream,
                process_id: 0,
                secret_key: 0,
                parameters: HashMap::new(),
                config: state.config,
                idx: state.idx,
            }),
            Some(Message::AuthenticationCleartextPassword) => {
                let pass = state
                    .config
                    .0
                    .password
                    .as_ref()
                    .ok_or_else(|| Error::config("password missing".into()))?;
                let mut buf = vec![];
                frontend::password_message(pass, &mut buf).map_err(Error::encode)?;
                transition!(SendingPassword {
                    future: state.stream.send(FrontendMessage::Raw(buf)),
                    config: state.config,
                    idx: state.idx,
                })
            }
            Some(Message::AuthenticationMd5Password(body)) => {
                let user = state
                    .config
                    .0
                    .user
                    .as_ref()
                    .ok_or_else(|| Error::config("user missing".into()))?;
                let pass = state
                    .config
                    .0
                    .password
                    .as_ref()
                    .ok_or_else(|| Error::config("password missing".into()))?;
                let output = authentication::md5_hash(user.as_bytes(), pass, body.salt());
                let mut buf = vec![];
                frontend::password_message(output.as_bytes(), &mut buf).map_err(Error::encode)?;
                transition!(SendingPassword {
                    future: state.stream.send(FrontendMessage::Raw(buf)),
                    config: state.config,
                    idx: state.idx,
                })
            }
            Some(Message::AuthenticationSasl(body)) => {
                let pass = state
                    .config
                    .0
                    .password
                    .as_ref()
                    .ok_or_else(|| Error::config("password missing".into()))?;

                let mut has_scram = false;
                let mut has_scram_plus = false;
                let mut mechanisms = body.mechanisms();
                while let Some(mechanism) = mechanisms.next().map_err(Error::parse)? {
                    match mechanism {
                        sasl::SCRAM_SHA_256 => has_scram = true,
                        sasl::SCRAM_SHA_256_PLUS => has_scram_plus = true,
                        _ => {}
                    }
                }

                let channel_binding = if let Some(tls_server_end_point) =
                    state.channel_binding.tls_server_end_point
                {
                    Some(sasl::ChannelBinding::tls_server_end_point(
                        tls_server_end_point,
                    ))
                } else {
                    None
                };

                let (channel_binding, mechanism) = if has_scram_plus {
                    match channel_binding {
                        Some(channel_binding) => (channel_binding, sasl::SCRAM_SHA_256_PLUS),
                        None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
                    }
                } else if has_scram {
                    match channel_binding {
                        Some(_) => (sasl::ChannelBinding::unrequested(), sasl::SCRAM_SHA_256),
                        None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
                    }
                } else {
                    return Err(Error::authentication(
                        "unsupported authentication method".into(),
                    ));
                };

                let scram = ScramSha256::new(pass, channel_binding);

                let mut buf = vec![];
                frontend::sasl_initial_response(mechanism, scram.message(), &mut buf)
                    .map_err(Error::encode)?;

                transition!(SendingSasl {
                    future: state.stream.send(FrontendMessage::Raw(buf)),
                    scram,
                    config: state.config,
                    idx: state.idx,
                })
            }
            Some(Message::AuthenticationKerberosV5)
            | Some(Message::AuthenticationScmCredential)
            | Some(Message::AuthenticationGss)
            | Some(Message::AuthenticationSspi) => Err(Error::authentication(
                "unsupported authentication method".into(),
            )),
            Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
            Some(_) => Err(Error::unexpected_message()),
            None => Err(Error::closed()),
        }
    }

    fn poll_sending_password<'a>(
        state: &'a mut RentToOwn<'a, SendingPassword<S, T>>,
    ) -> Poll<AfterSendingPassword<S, T>, Error> {
        let stream = try_ready!(state.future.poll().map_err(Error::io));
        let state = state.take();
        transition!(ReadingAuthCompletion {
            stream,
            config: state.config,
            idx: state.idx,
        })
    }

    fn poll_sending_sasl<'a>(
        state: &'a mut RentToOwn<'a, SendingSasl<S, T>>,
    ) -> Poll<AfterSendingSasl<S, T>, Error> {
        let stream = try_ready!(state.future.poll().map_err(Error::io));
        let state = state.take();
        transition!(ReadingSasl {
            stream,
            scram: state.scram,
            config: state.config,
            idx: state.idx,
        })
    }

    fn poll_reading_sasl<'a>(
        state: &'a mut RentToOwn<'a, ReadingSasl<S, T>>,
    ) -> Poll<AfterReadingSasl<S, T>, Error> {
        let message = try_ready!(state.stream.poll().map_err(Error::io));
        let mut state = state.take();

        match message {
            Some(Message::AuthenticationSaslContinue(body)) => {
                state
                    .scram
                    .update(body.data())
                    .map_err(|e| Error::authentication(Box::new(e)))?;
                let mut buf = vec![];
                frontend::sasl_response(state.scram.message(), &mut buf).map_err(Error::encode)?;
                transition!(SendingSasl {
                    future: state.stream.send(FrontendMessage::Raw(buf)),
                    scram: state.scram,
                    config: state.config,
                    idx: state.idx,
                })
            }
            Some(Message::AuthenticationSaslFinal(body)) => {
                state
                    .scram
                    .finish(body.data())
                    .map_err(|e| Error::authentication(Box::new(e)))?;
                transition!(ReadingAuthCompletion {
                    stream: state.stream,
                    config: state.config,
                    idx: state.idx,
                })
            }
            Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
            Some(_) => Err(Error::unexpected_message()),
            None => Err(Error::closed()),
        }
    }

    fn poll_reading_auth_completion<'a>(
        state: &'a mut RentToOwn<'a, ReadingAuthCompletion<S, T>>,
    ) -> Poll<AfterReadingAuthCompletion<S, T>, Error> {
        let message = try_ready!(state.stream.poll().map_err(Error::io));
        let state = state.take();

        match message {
            Some(Message::AuthenticationOk) => transition!(ReadingInfo {
                stream: state.stream,
                process_id: 0,
                secret_key: 0,
                parameters: HashMap::new(),
                config: state.config,
                idx: state.idx,
            }),
            Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
            Some(_) => Err(Error::unexpected_message()),
            None => Err(Error::closed()),
        }
    }

    fn poll_reading_info<'a>(
        state: &'a mut RentToOwn<'a, ReadingInfo<S, T>>,
    ) -> Poll<AfterReadingInfo<S, T>, Error> {
        loop {
            let message = try_ready!(state.stream.poll().map_err(Error::io));
            match message {
                Some(Message::BackendKeyData(body)) => {
                    state.process_id = body.process_id();
                    state.secret_key = body.secret_key();
                }
                Some(Message::ParameterStatus(body)) => {
                    state.parameters.insert(
                        body.name().map_err(Error::parse)?.to_string(),
                        body.value().map_err(Error::parse)?.to_string(),
                    );
                }
                Some(Message::ReadyForQuery(_)) => {
                    let state = state.take();
                    let (sender, receiver) = mpsc::unbounded();
                    let client = Client::new(
                        sender,
                        state.process_id,
                        state.secret_key,
                        state.config,
                        state.idx,
                    );
                    let connection = Connection::new(state.stream, state.parameters, receiver);
                    transition!(Finished((client, connection)))
                }
                Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
                Some(Message::NoticeResponse(_)) => {}
                Some(_) => return Err(Error::unexpected_message()),
                None => return Err(Error::closed()),
            }
        }
    }
}

impl<S, T> ConnectRawFuture<S, T>
where
    S: AsyncRead + AsyncWrite,
    T: TlsConnect<S>,
{
    pub fn new(stream: S, tls: T, config: Config, idx: Option<usize>) -> ConnectRawFuture<S, T> {
        ConnectRaw::start(TlsFuture::new(stream, config.0.ssl_mode, tls), config, idx)
    }
}
<commit_msg>rustfmt<commit_after>use fallible_iterator::FallibleIterator;
use futures::sink;
use futures::sync::mpsc;
use futures::{try_ready, Future, Poll, Sink, Stream};
use postgres_protocol::authentication;
use postgres_protocol::authentication::sasl::{self, ScramSha256};
use postgres_protocol::message::backend::Message;
use postgres_protocol::message::frontend;
use state_machine_future::{transition, RentToOwn, StateMachineFuture};
use std::collections::HashMap;
use tokio_codec::Framed;
use tokio_io::{AsyncRead, AsyncWrite};

use crate::proto::{Client, Connection, FrontendMessage, MaybeTlsStream, PostgresCodec, TlsFuture};
use crate::tls::ChannelBinding;
use crate::{Config, Error, TlsConnect};

#[derive(StateMachineFuture)]
pub enum ConnectRaw<S, T>
where
    S: AsyncRead + AsyncWrite,
    T: TlsConnect<S>,
{
    #[state_machine_future(start, transitions(SendingStartup))]
    Start {
        future: TlsFuture<S, T>,
        config: Config,
        idx: Option<usize>,
    },
    #[state_machine_future(transitions(ReadingAuth))]
    SendingStartup {
        future: sink::Send<Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>>,
        config: Config,
        idx: Option<usize>,
        channel_binding: ChannelBinding,
    },
    #[state_machine_future(transitions(ReadingInfo, SendingPassword, SendingSasl))]
    ReadingAuth {
        stream: Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>,
        config: Config,
        idx: Option<usize>,
        channel_binding: ChannelBinding,
    },
    #[state_machine_future(transitions(ReadingAuthCompletion))]
    SendingPassword {
        future: sink::Send<Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>>,
        config: Config,
        idx: Option<usize>,
    },
    #[state_machine_future(transitions(ReadingSasl))]
    SendingSasl {
        future: sink::Send<Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>>,
        scram: ScramSha256,
        config: Config,
        idx: Option<usize>,
    },
    #[state_machine_future(transitions(SendingSasl, ReadingAuthCompletion))]
    ReadingSasl {
        stream: Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>,
        scram: ScramSha256,
        config: Config,
        idx: Option<usize>,
    },
    #[state_machine_future(transitions(ReadingInfo))]
    ReadingAuthCompletion {
        stream: Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>,
        config: Config,
        idx: Option<usize>,
    },
    #[state_machine_future(transitions(Finished))]
    ReadingInfo {
        stream: Framed<MaybeTlsStream<S, T::Stream>, PostgresCodec>,
        process_id: i32,
        secret_key: i32,
        parameters: HashMap<String, String>,
        config: Config,
        idx: Option<usize>,
    },
    #[state_machine_future(ready)]
    Finished((Client, Connection<MaybeTlsStream<S, T::Stream>>)),
    #[state_machine_future(error)]
    Failed(Error),
}

impl<S, T> PollConnectRaw<S, T> for ConnectRaw<S, T>
where
    S: AsyncRead + AsyncWrite,
    T: TlsConnect<S>,
{
    fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start<S, T>>) -> Poll<AfterStart<S, T>, Error> {
        let (stream, channel_binding) = try_ready!(state.future.poll());
        let state = state.take();

        let mut params = vec![("client_encoding", "UTF8"), ("timezone", "GMT")];
        if let Some(user) = &state.config.0.user {
            params.push(("user", &**user));
        }
        if let Some(dbname) = &state.config.0.dbname {
            params.push(("database", &**dbname));
        }
        if let Some(options) = &state.config.0.options {
            params.push(("options", &**options));
        }
        if let Some(application_name) = &state.config.0.application_name {
            params.push(("application_name", &**application_name));
        }

        let mut buf = vec![];
        frontend::startup_message(params, &mut buf).map_err(Error::encode)?;

        let stream = Framed::new(stream, PostgresCodec);

        transition!(SendingStartup {
            future: stream.send(FrontendMessage::Raw(buf)),
            config: state.config,
            idx: state.idx,
            channel_binding,
        })
    }

    fn poll_sending_startup<'a>(
        state: &'a mut RentToOwn<'a, SendingStartup<S, T>>,
    ) -> Poll<AfterSendingStartup<S, T>, Error> {
        let stream = try_ready!(state.future.poll().map_err(Error::io));
        let state = state.take();
        transition!(ReadingAuth {
            stream,
            config: state.config,
            idx: state.idx,
            channel_binding: state.channel_binding,
        })
    }

    fn poll_reading_auth<'a>(
        state: &'a mut RentToOwn<'a, ReadingAuth<S, T>>,
    ) -> Poll<AfterReadingAuth<S, T>, Error> {
        let message = try_ready!(state.stream.poll().map_err(Error::io));
        let state = state.take();

        match message {
            Some(Message::AuthenticationOk) => transition!(ReadingInfo {
                stream: state.stream,
                process_id: 0,
                secret_key: 0,
                parameters: HashMap::new(),
                config: state.config,
                idx: state.idx,
            }),
            Some(Message::AuthenticationCleartextPassword) => {
                let pass = state
                    .config
                    .0
                    .password
                    .as_ref()
                    .ok_or_else(|| Error::config("password missing".into()))?;
                let mut buf = vec![];
                frontend::password_message(pass, &mut buf).map_err(Error::encode)?;
                transition!(SendingPassword {
                    future: state.stream.send(FrontendMessage::Raw(buf)),
                    config: state.config,
                    idx: state.idx,
                })
            }
            Some(Message::AuthenticationMd5Password(body)) => {
                let user = state
                    .config
                    .0
                    .user
                    .as_ref()
                    .ok_or_else(|| Error::config("user missing".into()))?;
                let pass = state
                    .config
                    .0
                    .password
                    .as_ref()
                    .ok_or_else(|| Error::config("password missing".into()))?;
                let output = authentication::md5_hash(user.as_bytes(), pass, body.salt());
                let mut buf = vec![];
                frontend::password_message(output.as_bytes(), &mut buf).map_err(Error::encode)?;
                transition!(SendingPassword {
                    future: state.stream.send(FrontendMessage::Raw(buf)),
                    config: state.config,
                    idx: state.idx,
                })
            }
            Some(Message::AuthenticationSasl(body)) => {
                let pass = state
                    .config
                    .0
                    .password
                    .as_ref()
                    .ok_or_else(|| Error::config("password missing".into()))?;

                let mut has_scram = false;
                let mut has_scram_plus = false;
                let mut mechanisms = body.mechanisms();
                while let Some(mechanism) = mechanisms.next().map_err(Error::parse)? {
                    match mechanism {
                        sasl::SCRAM_SHA_256 => has_scram = true,
                        sasl::SCRAM_SHA_256_PLUS => has_scram_plus = true,
                        _ => {}
                    }
                }

                let channel_binding = if let Some(tls_server_end_point) =
                    state.channel_binding.tls_server_end_point
                {
                    Some(sasl::ChannelBinding::tls_server_end_point(
                        tls_server_end_point,
                    ))
                } else {
                    None
                };

                let (channel_binding, mechanism) = if has_scram_plus {
                    match channel_binding {
                        Some(channel_binding) => (channel_binding, sasl::SCRAM_SHA_256_PLUS),
                        None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
                    }
                } else if has_scram {
                    match channel_binding {
                        Some(_) => (sasl::ChannelBinding::unrequested(), sasl::SCRAM_SHA_256),
                        None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
                    }
                } else {
                    return Err(Error::authentication(
                        "unsupported authentication method".into(),
                    ));
                };

                let scram = ScramSha256::new(pass, channel_binding);

                let mut buf = vec![];
                frontend::sasl_initial_response(mechanism, scram.message(), &mut buf)
                    .map_err(Error::encode)?;

                transition!(SendingSasl {
                    future: state.stream.send(FrontendMessage::Raw(buf)),
                    scram,
                    config: state.config,
                    idx: state.idx,
                })
            }
            Some(Message::AuthenticationKerberosV5)
            | Some(Message::AuthenticationScmCredential)
            | Some(Message::AuthenticationGss)
            | Some(Message::AuthenticationSspi) => Err(Error::authentication(
                "unsupported authentication method".into(),
            )),
            Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
            Some(_) => Err(Error::unexpected_message()),
            None => Err(Error::closed()),
        }
    }

    fn poll_sending_password<'a>(
        state: &'a mut RentToOwn<'a, SendingPassword<S, T>>,
    ) -> Poll<AfterSendingPassword<S, T>, Error> {
        let stream = try_ready!(state.future.poll().map_err(Error::io));
        let state = state.take();
        transition!(ReadingAuthCompletion {
            stream,
            config: state.config,
            idx: state.idx,
        })
    }

    fn poll_sending_sasl<'a>(
        state: &'a mut RentToOwn<'a, SendingSasl<S, T>>,
    ) -> Poll<AfterSendingSasl<S, T>, Error> {
        let stream = try_ready!(state.future.poll().map_err(Error::io));
        let state = state.take();
        transition!(ReadingSasl {
            stream,
            scram: state.scram,
            config: state.config,
            idx: state.idx,
        })
    }

    fn poll_reading_sasl<'a>(
        state: &'a mut RentToOwn<'a, ReadingSasl<S, T>>,
    ) -> Poll<AfterReadingSasl<S, T>, Error> {
        let message = try_ready!(state.stream.poll().map_err(Error::io));
        let mut state = state.take();

        match message {
            Some(Message::AuthenticationSaslContinue(body)) => {
                state
                    .scram
                    .update(body.data())
                    .map_err(|e| Error::authentication(Box::new(e)))?;
                let mut buf = vec![];
                frontend::sasl_response(state.scram.message(), &mut buf).map_err(Error::encode)?;
                transition!(SendingSasl {
                    future: state.stream.send(FrontendMessage::Raw(buf)),
                    scram: state.scram,
                    config: state.config,
                    idx: state.idx,
                })
            }
            Some(Message::AuthenticationSaslFinal(body)) => {
                state
                    .scram
                    .finish(body.data())
                    .map_err(|e| Error::authentication(Box::new(e)))?;
                transition!(ReadingAuthCompletion {
                    stream: state.stream,
                    config: state.config,
                    idx: state.idx,
                })
            }
            Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
            Some(_) => Err(Error::unexpected_message()),
            None => Err(Error::closed()),
        }
    }

    fn poll_reading_auth_completion<'a>(
        state: &'a mut RentToOwn<'a, ReadingAuthCompletion<S, T>>,
    ) -> Poll<AfterReadingAuthCompletion<S, T>, Error> {
        let message = try_ready!(state.stream.poll().map_err(Error::io));
        let state = state.take();

        match message {
            Some(Message::AuthenticationOk) => transition!(ReadingInfo {
                stream: state.stream,
                process_id: 0,
                secret_key: 0,
                parameters: HashMap::new(),
                config: state.config,
                idx: state.idx,
            }),
            Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
            Some(_) => Err(Error::unexpected_message()),
            None => Err(Error::closed()),
        }
    }

    fn poll_reading_info<'a>(
        state: &'a mut RentToOwn<'a, ReadingInfo<S, T>>,
    ) -> Poll<AfterReadingInfo<S, T>, Error> {
        loop {
            let message = try_ready!(state.stream.poll().map_err(Error::io));
            match message {
                Some(Message::BackendKeyData(body)) => {
                    state.process_id = body.process_id();
                    state.secret_key = body.secret_key();
                }
                Some(Message::ParameterStatus(body)) => {
                    state.parameters.insert(
                        body.name().map_err(Error::parse)?.to_string(),
                        body.value().map_err(Error::parse)?.to_string(),
                    );
                }
                Some(Message::ReadyForQuery(_)) => {
                    let state = state.take();
                    let (sender, receiver) = mpsc::unbounded();
                    let client = Client::new(
                        sender,
                        state.process_id,
                        state.secret_key,
                        state.config,
                        state.idx,
                    );
                    let connection = Connection::new(state.stream, state.parameters, receiver);
                    transition!(Finished((client, connection)))
                }
                Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
                Some(Message::NoticeResponse(_)) => {}
                Some(_) => return Err(Error::unexpected_message()),
                None => return Err(Error::closed()),
            }
        }
    }
}

impl<S, T> ConnectRawFuture<S, T>
where
    S: AsyncRead + AsyncWrite,
    T: TlsConnect<S>,
{
    pub fn new(stream: S, tls: T, config: Config, idx: Option<usize>) -> ConnectRawFuture<S, T> {
        ConnectRaw::start(TlsFuture::new(stream, config.0.ssl_mode, tls), config, idx)
    }
}
<|endoftext|>


<commit_before>extern crate rustyline;

extern crate gluon_completion as completion;

use std::error::Error as StdError;
use std::fmt;
use std::sync::Mutex;

use futures::Future;

use self::rustyline::error::ReadlineError;

use base::ast::{Expr, Pattern, SpannedPattern, Typed};
use base::error::InFile;
use base::kind::Kind;
use base::pos;
use base::symbol::{Symbol, SymbolModule};
use base::types::ArcType;
use parser::parse_partial_let_or_expr;
use vm::Error as VMError;
use vm::api::{Function, Userdata, VmType, WithVM, IO};
use vm::gc::{Gc, Traverseable};
use vm::internal::ValuePrinter;
use vm::thread::{RootStr, RootedValue, Thread, ThreadInternal};

use gluon::{new_vm, Compiler, Result as GluonResult, RootedThread};
use gluon::compiler_pipeline::{run_io, Executable};

fn type_of_expr(args: WithVM<RootStr>) -> IO<Result<String, String>> {
    let WithVM { vm, value: args } = args;
    let mut compiler = Compiler::new();
    IO::Value(match compiler.typecheck_str(vm, "<repl>", &args, None) {
        Ok((expr, _)) => {
            let env = vm.get_env();
            Ok(format!("{}", expr.env_type_of(&*env)))
        }
        Err(msg) => Err(format!("{}", msg)),
    })
}

fn find_kind(args: WithVM<RootStr>) -> IO<Result<String, String>> {
    let vm = args.vm;
    let args = args.value.trim();
    IO::Value(match vm.find_type_info(args) {
        Ok(ref alias) => {
            let kind = alias.args.iter().rev().fold(
                Kind::typ(),
                |acc, arg| Kind::function(arg.kind.clone(), acc),
            );
            Ok(format!("{}", kind))
        }
        Err(err) => Err(format!("{}", err)),
    })
}

fn find_info(args: WithVM<RootStr>) -> IO<Result<String, String>> {
    use std::fmt::Write;
    let vm = args.vm;
    let args = args.value.trim();
    let env = vm.get_env();
    let mut buffer = String::new();
    match env.find_type_info(args) {
        Ok(alias) => {
            // Found a type alias
            let mut fmt = || -> Result<(), ::std::fmt::Error> {
                write!(&mut buffer, "type {}", args)?;
                for g in &alias.args {
                    write!(&mut buffer, " {}", g.id)?;
                }
                write!(&mut buffer, " = {}", alias.unresolved_type())
            };
            fmt().unwrap();
        }
        Err(err) => {
            // Try to find a value at `args` to print its type and documentation comment (if any)
            match env.get_binding(args) {
                Ok((_, typ)) => {
                    write!(&mut buffer, "{}: {}", args, typ).unwrap();
                }
                Err(_) => return IO::Value(Err(format!("{}", err))),
            }
        }
    }
    let maybe_comment = env.get_metadata(args)
        .ok()
        .and_then(|metadata| metadata.comment.as_ref());
    if let Some(comment) = maybe_comment {
        for line in comment.lines() {
            write!(&mut buffer, "\n/// {}", line).unwrap();
        }
    }
    IO::Value(Ok(buffer))
}

fn complete(thread: &Thread, name: &str, fileinput: &str, pos: usize) -> GluonResult<Vec<String>> {
    use base::pos::BytePos;
    use gluon::compiler_pipeline::*;

    let mut compiler = Compiler::new();

    // The parser may find parse errors but still produce an expression
    // For that case still typecheck the expression but return the parse error afterwards
    let (mut expr, _parse_result): (_, GluonResult<()>) =
        match compiler.parse_partial_expr(thread.global_env().type_cache(), &name, fileinput) {
            Ok(expr) => (expr, Ok(())),
            Err((None, err)) => return Err(err.into()),
            Err((Some(expr), err)) => (expr, Err(err.into())),
        };

    expr.expand_macro(&mut compiler, thread, &name)?;

    // Only need the typechecker to fill infer the types as best it can regardless of errors
    let _ = compiler.typecheck_expr(thread, &name, fileinput, &mut expr);
    let suggestions = completion::suggest(&*thread.get_env(), &expr, BytePos::from(pos));
    Ok(
        suggestions
            .into_iter()
            .map(|ident| {
                let s: &str = ident.name.as_ref();
                s.to_string()
            })
            .collect(),
    )
}

struct Completer(RootedThread);

impl rustyline::completion::Completer for Completer {
    fn complete(&self, line: &str, pos: usize) -> rustyline::Result<(usize, Vec<String>)> {
        let result = complete(&self.0, "<repl>", line, pos);

        // Get the start of the completed identifier
        let ident_start = line[..pos]
            .rfind(|c: char| c.is_whitespace() || c == '.')
            .map_or(0, |i| i + 1);
        Ok((ident_start, result.unwrap_or(Vec::new())))
    }
}

struct Editor(Mutex<rustyline::Editor<Completer>>);

impl Userdata for Editor {}

impl fmt::Debug for Editor {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Editor(..)")
    }
}

impl VmType for Editor {
    type Type = Editor;
}

impl Traverseable for Editor {
    fn traverse(&self, _: &mut Gc) {}
}

fn new_editor(vm: WithVM<()>) -> Editor {
    let mut editor = rustyline::Editor::new();
    editor.set_completer(Some(Completer(vm.vm.root_thread())));
    Editor(Mutex::new(editor))
}

fn readline(editor: &Editor, prompt: &str) -> IO<Option<String>> {
    let mut editor = editor.0.lock().unwrap();
    let input = match editor.readline(prompt) {
        Ok(input) => input,
        Err(ReadlineError::Eof) | Err(ReadlineError::Interrupted) => return IO::Value(None),
        Err(err) => return IO::Exception(format!("{}", err)),
    };
    if !input.trim().is_empty() {
        editor.add_history_entry(&input);
    }

    IO::Value(Some(input))
}

fn eval_line(WithVM { vm, value: line }: WithVM<&str>) -> IO<String> {
    IO::Value(match eval_line_(vm, line) {
        Ok(x) => x,
        Err(x) => x.to_string(),
    })
}

fn eval_line_(vm: &Thread, line: &str) -> GluonResult<String> {
    let mut compiler = Compiler::new();
    let let_or_expr = {
        let mut module = SymbolModule::new("<line>".into(), compiler.mut_symbols());
        match parse_partial_let_or_expr(&mut module, line) {
            Ok(x) => x,
            Err((_, err)) => return Err(InFile::new("<line>", line, err).into()),
        }
    };
    let mut eval_expr;
    let (value, typ) = match let_or_expr {
        Ok(expr) => {
            eval_expr = expr;
            eval_expr
                .run_expr(&mut compiler, vm, "<line>", line, None)
                .and_then(|v| run_io(vm, v))
                .wait()?
        }
        Err(let_binding) => {
            let unpack_pattern = let_binding.name.clone();
            eval_expr = match unpack_pattern.value {
                Pattern::Ident(ref id) if !let_binding.args.is_empty() => {
                    // We can't compile function bindings by only looking at `let_binding.expr`
                    // so rewrite `let f x y = <expr>` into `let f x y = <expr> in f`
                    let id = pos::spanned2(0.into(), 0.into(), Expr::Ident(id.clone()));
                    let expr = Expr::LetBindings(vec![let_binding], Box::new(id));
                    pos::spanned2(0.into(), 0.into(), expr)
                }
                _ => let_binding.expr,
            };
            let value = eval_expr
                .run_expr(&mut compiler, vm, "<line>", line, None)
                .wait()?;
            set_globals(vm, &unpack_pattern, &value.typ, &value.value)?;
            (value.value, value.typ)
        }
    };

    let env = vm.global_env().get_env();
    Ok(ValuePrinter::new(&*env, &typ, *value).width(80).to_string())
}

fn set_globals(
    vm: &Thread,
    pattern: &SpannedPattern<Symbol>,
    typ: &ArcType,
    value: &RootedValue<&Thread>,
) -> GluonResult<()> {
    match pattern.value {
        Pattern::Ident(ref id) => {
            vm.set_global(id.name.clone(), typ.clone(), Default::default(), **value)?;
            Ok(())
        }
        Pattern::Tuple { ref elems, .. } => {
            let iter = elems.iter().zip(::vm::dynamic::field_iter(&value, typ, vm));
            for (elem_pattern, (elem_value, elem_type)) in iter {
                set_globals(vm, elem_pattern, &elem_type, &elem_value)?;
            }
            Ok(())
        }
        Pattern::Record { ref fields, .. } => {
            let iter = fields
                .iter()
                .zip(::vm::dynamic::field_iter(&value, typ, vm));
            for (field, (field_value, field_type)) in iter {
                match field.value {
                    Some(ref field_pattern) => {
                        set_globals(vm, field_pattern, &field_type, &field_value)?
                    }
                    None => vm.set_global(
                        field.name.value.clone(),
                        field_type,
                        Default::default(),
                        *field_value,
                    )?,
                }
            }
            Ok(())
        }
        _ => {
            Err(VMError::Message("The repl cannot bind variables from this pattern".into()).into())
        }
    }
}

fn compile_repl(vm: &Thread) -> Result<(), Box<StdError + Send + Sync>> {
    vm.register_type::<Editor>("Editor", &[])?;

    vm.define_global(
        "rustyline",
        record!(
        new_editor => primitive!(1 new_editor),
        readline => primitive!(2 readline)
    ),
    )?;
    vm.define_global(
        "repl_prim",
        record!(
        type_of_expr => primitive!(1 type_of_expr),
        find_info => primitive!(1 find_info),
        find_kind => primitive!(1 find_kind),
        eval_line => primitive!(1 eval_line)
    ),
    )?;

    const REPL_SOURCE: &'static str =
        include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/src/repl.glu"));

    let mut compiler = Compiler::new();
    compiler
        .load_script_async(vm, "std.repl", REPL_SOURCE)
        .sync_or_error()?;
    Ok(())
}

#[allow(dead_code)]
pub fn run() -> Result<(), Box<StdError + Send + Sync>> {
    let vm = new_vm();
    compile_repl(&vm)?;
    let mut repl: Function<&Thread, fn(()) -> IO<()>> = vm.get_global("std.repl")?;
    debug!("Starting repl");
    repl.call(())?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::compile_repl;
    use vm::api::{FunctionRef, IO};
    use gluon::{self, RootedThread};
    use gluon::import::Import;

    fn new_vm() -> RootedThread {
        let vm = gluon::new_vm();
        let import = vm.get_macros().get("import");
        import
            .as_ref()
            .and_then(|import| import.downcast_ref::<Import>())
            .expect("Import macro")
            .add_path("..");
        vm
    }

    #[test]
    fn compile_repl_test() {
        let _ = ::env_logger::init();
        let vm = new_vm();
        compile_repl(&vm).unwrap_or_else(|err| ice!("{}", err));
        let repl: Result<FunctionRef<fn(()) -> IO<()>>, _> = vm.get_global("std.repl");
        assert!(repl.is_ok(), "{}", repl.err().unwrap());
    }

    type QueryFn = fn(&'static str) -> IO<Result<String, String>>;

    #[test]
    fn type_of_expr() {
        let _ = ::env_logger::init();
        let vm = new_vm();
        compile_repl(&vm).unwrap_or_else(|err| ice!("{}", err));
        let mut type_of: FunctionRef<QueryFn> = vm.get_global("repl_prim.type_of_expr").unwrap();
        assert_eq!(type_of.call("123"), Ok(IO::Value(Ok("Int".into()))));
    }

    #[test]
    fn find_kind() {
        let _ = ::env_logger::init();
        let vm = new_vm();
        compile_repl(&vm).unwrap_or_else(|err| ice!("{}", err));
        let mut find_kind: FunctionRef<QueryFn> = vm.get_global("repl_prim.find_kind").unwrap();
        assert_eq!(
            find_kind.call("std.prelude.Semigroup"),
            Ok(IO::Value(Ok("Type -> Type".into())))
        );
    }

    #[test]
    fn find_info() {
        let _ = ::env_logger::init();
        let vm = new_vm();
        compile_repl(&vm).unwrap_or_else(|err| ice!("{}", err));
        let mut find_info: FunctionRef<QueryFn> = vm.get_global("repl_prim.find_info").unwrap();
        match find_info.call("std.prelude.Semigroup") {
            Ok(IO::Value(Ok(_))) => (),
            x => assert!(false, "{:?}", x),
        }
        match find_info.call("std.prelude.make_Semigroup") {
            Ok(IO::Value(Ok(_))) => (),
            x => assert!(false, "{:?}", x),
        }
        match find_info.call("float_prim") {
            Ok(IO::Value(Ok(_))) => (),
            x => assert!(false, "{:?}", x),
        }
    }
}
<commit_msg>Revert tests to panic! macro.<commit_after>extern crate rustyline;

extern crate gluon_completion as completion;

use std::error::Error as StdError;
use std::fmt;
use std::sync::Mutex;

use futures::Future;

use self::rustyline::error::ReadlineError;

use base::ast::{Expr, Pattern, SpannedPattern, Typed};
use base::error::InFile;
use base::kind::Kind;
use base::pos;
use base::symbol::{Symbol, SymbolModule};
use base::types::ArcType;
use parser::parse_partial_let_or_expr;
use vm::Error as VMError;
use vm::api::{Function, Userdata, VmType, WithVM, IO};
use vm::gc::{Gc, Traverseable};
use vm::internal::ValuePrinter;
use vm::thread::{RootStr, RootedValue, Thread, ThreadInternal};

use gluon::{new_vm, Compiler, Result as GluonResult, RootedThread};
use gluon::compiler_pipeline::{run_io, Executable};

fn type_of_expr(args: WithVM<RootStr>) -> IO<Result<String, String>> {
    let WithVM { vm, value: args } = args;
    let mut compiler = Compiler::new();
    IO::Value(match compiler.typecheck_str(vm, "<repl>", &args, None) {
        Ok((expr, _)) => {
            let env = vm.get_env();
            Ok(format!("{}", expr.env_type_of(&*env)))
        }
        Err(msg) => Err(format!("{}", msg)),
    })
}

fn find_kind(args: WithVM<RootStr>) -> IO<Result<String, String>> {
    let vm = args.vm;
    let args = args.value.trim();
    IO::Value(match vm.find_type_info(args) {
        Ok(ref alias) => {
            let kind = alias.args.iter().rev().fold(
                Kind::typ(),
                |acc, arg| Kind::function(arg.kind.clone(), acc),
            );
            Ok(format!("{}", kind))
        }
        Err(err) => Err(format!("{}", err)),
    })
}

fn find_info(args: WithVM<RootStr>) -> IO<Result<String, String>> {
    use std::fmt::Write;
    let vm = args.vm;
    let args = args.value.trim();
    let env = vm.get_env();
    let mut buffer = String::new();
    match env.find_type_info(args) {
        Ok(alias) => {
            // Found a type alias
            let mut fmt = || -> Result<(), ::std::fmt::Error> {
                write!(&mut buffer, "type {}", args)?;
                for g in &alias.args {
                    write!(&mut buffer, " {}", g.id)?;
                }
                write!(&mut buffer, " = {}", alias.unresolved_type())
            };
            fmt().unwrap();
        }
        Err(err) => {
            // Try to find a value at `args` to print its type and documentation comment (if any)
            match env.get_binding(args) {
                Ok((_, typ)) => {
                    write!(&mut buffer, "{}: {}", args, typ).unwrap();
                }
                Err(_) => return IO::Value(Err(format!("{}", err))),
            }
        }
    }
    let maybe_comment = env.get_metadata(args)
        .ok()
        .and_then(|metadata| metadata.comment.as_ref());
    if let Some(comment) = maybe_comment {
        for line in comment.lines() {
            write!(&mut buffer, "\n/// {}", line).unwrap();
        }
    }
    IO::Value(Ok(buffer))
}

fn complete(thread: &Thread, name: &str, fileinput: &str, pos: usize) -> GluonResult<Vec<String>> {
    use base::pos::BytePos;
    use gluon::compiler_pipeline::*;

    let mut compiler = Compiler::new();

    // The parser may find parse errors but still produce an expression
    // For that case still typecheck the expression but return the parse error afterwards
    let (mut expr, _parse_result): (_, GluonResult<()>) =
        match compiler.parse_partial_expr(thread.global_env().type_cache(), &name, fileinput) {
            Ok(expr) => (expr, Ok(())),
            Err((None, err)) => return Err(err.into()),
            Err((Some(expr), err)) => (expr, Err(err.into())),
        };

    expr.expand_macro(&mut compiler, thread, &name)?;

    // Only need the typechecker to fill infer the types as best it can regardless of errors
    let _ = compiler.typecheck_expr(thread, &name, fileinput, &mut expr);
    let suggestions = completion::suggest(&*thread.get_env(), &expr, BytePos::from(pos));
    Ok(
        suggestions
            .into_iter()
            .map(|ident| {
                let s: &str = ident.name.as_ref();
                s.to_string()
            })
            .collect(),
    )
}

struct Completer(RootedThread);

impl rustyline::completion::Completer for Completer {
    fn complete(&self, line: &str, pos: usize) -> rustyline::Result<(usize, Vec<String>)> {
        let result = complete(&self.0, "<repl>", line, pos);

        // Get the start of the completed identifier
        let ident_start = line[..pos]
            .rfind(|c: char| c.is_whitespace() || c == '.')
            .map_or(0, |i| i + 1);
        Ok((ident_start, result.unwrap_or(Vec::new())))
    }
}

struct Editor(Mutex<rustyline::Editor<Completer>>);

impl Userdata for Editor {}

impl fmt::Debug for Editor {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Editor(..)")
    }
}

impl VmType for Editor {
    type Type = Editor;
}

impl Traverseable for Editor {
    fn traverse(&self, _: &mut Gc) {}
}

fn new_editor(vm: WithVM<()>) -> Editor {
    let mut editor = rustyline::Editor::new();
    editor.set_completer(Some(Completer(vm.vm.root_thread())));
    Editor(Mutex::new(editor))
}

fn readline(editor: &Editor, prompt: &str) -> IO<Option<String>> {
    let mut editor = editor.0.lock().unwrap();
    let input = match editor.readline(prompt) {
        Ok(input) => input,
        Err(ReadlineError::Eof) | Err(ReadlineError::Interrupted) => return IO::Value(None),
        Err(err) => return IO::Exception(format!("{}", err)),
    };
    if !input.trim().is_empty() {
        editor.add_history_entry(&input);
    }

    IO::Value(Some(input))
}

fn eval_line(WithVM { vm, value: line }: WithVM<&str>) -> IO<String> {
    IO::Value(match eval_line_(vm, line) {
        Ok(x) => x,
        Err(x) => x.to_string(),
    })
}

fn eval_line_(vm: &Thread, line: &str) -> GluonResult<String> {
    let mut compiler = Compiler::new();
    let let_or_expr = {
        let mut module = SymbolModule::new("<line>".into(), compiler.mut_symbols());
        match parse_partial_let_or_expr(&mut module, line) {
            Ok(x) => x,
            Err((_, err)) => return Err(InFile::new("<line>", line, err).into()),
        }
    };
    let mut eval_expr;
    let (value, typ) = match let_or_expr {
        Ok(expr) => {
            eval_expr = expr;
            eval_expr
                .run_expr(&mut compiler, vm, "<line>", line, None)
                .and_then(|v| run_io(vm, v))
                .wait()?
        }
        Err(let_binding) => {
            let unpack_pattern = let_binding.name.clone();
            eval_expr = match unpack_pattern.value {
                Pattern::Ident(ref id) if !let_binding.args.is_empty() => {
                    // We can't compile function bindings by only looking at `let_binding.expr`
                    // so rewrite `let f x y = <expr>` into `let f x y = <expr> in f`
                    let id = pos::spanned2(0.into(), 0.into(), Expr::Ident(id.clone()));
                    let expr = Expr::LetBindings(vec![let_binding], Box::new(id));
                    pos::spanned2(0.into(), 0.into(), expr)
                }
                _ => let_binding.expr,
            };
            let value = eval_expr
                .run_expr(&mut compiler, vm, "<line>", line, None)
                .wait()?;
            set_globals(vm, &unpack_pattern, &value.typ, &value.value)?;
            (value.value, value.typ)
        }
    };

    let env = vm.global_env().get_env();
    Ok(ValuePrinter::new(&*env, &typ, *value).width(80).to_string())
}

fn set_globals(
    vm: &Thread,
    pattern: &SpannedPattern<Symbol>,
    typ: &ArcType,
    value: &RootedValue<&Thread>,
) -> GluonResult<()> {
    match pattern.value {
        Pattern::Ident(ref id) => {
            vm.set_global(id.name.clone(), typ.clone(), Default::default(), **value)?;
            Ok(())
        }
        Pattern::Tuple { ref elems, .. } => {
            let iter = elems.iter().zip(::vm::dynamic::field_iter(&value, typ, vm));
            for (elem_pattern, (elem_value, elem_type)) in iter {
                set_globals(vm, elem_pattern, &elem_type, &elem_value)?;
            }
            Ok(())
        }
        Pattern::Record { ref fields, .. } => {
            let iter = fields
                .iter()
                .zip(::vm::dynamic::field_iter(&value, typ, vm));
            for (field, (field_value, field_type)) in iter {
                match field.value {
                    Some(ref field_pattern) => {
                        set_globals(vm, field_pattern, &field_type, &field_value)?
                    }
                    None => vm.set_global(
                        field.name.value.clone(),
                        field_type,
                        Default::default(),
                        *field_value,
                    )?,
                }
            }
            Ok(())
        }
        _ => {
            Err(VMError::Message("The repl cannot bind variables from this pattern".into()).into())
        }
    }
}

fn compile_repl(vm: &Thread) -> Result<(), Box<StdError + Send + Sync>> {
    vm.register_type::<Editor>("Editor", &[])?;

    vm.define_global(
        "rustyline",
        record!(
        new_editor => primitive!(1 new_editor),
        readline => primitive!(2 readline)
    ),
    )?;
    vm.define_global(
        "repl_prim",
        record!(
        type_of_expr => primitive!(1 type_of_expr),
        find_info => primitive!(1 find_info),
        find_kind => primitive!(1 find_kind),
        eval_line => primitive!(1 eval_line)
    ),
    )?;

    const REPL_SOURCE: &'static str =
        include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/src/repl.glu"));

    let mut compiler = Compiler::new();
    compiler
        .load_script_async(vm, "std.repl", REPL_SOURCE)
        .sync_or_error()?;
    Ok(())
}

#[allow(dead_code)]
pub fn run() -> Result<(), Box<StdError + Send + Sync>> {
    let vm = new_vm();
    compile_repl(&vm)?;
    let mut repl: Function<&Thread, fn(()) -> IO<()>> = vm.get_global("std.repl")?;
    debug!("Starting repl");
    repl.call(())?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::compile_repl;
    use vm::api::{FunctionRef, IO};
    use gluon::{self, RootedThread};
    use gluon::import::Import;

    fn new_vm() -> RootedThread {
        let vm = gluon::new_vm();
        let import = vm.get_macros().get("import");
        import
            .as_ref()
            .and_then(|import| import.downcast_ref::<Import>())
            .expect("Import macro")
            .add_path("..");
        vm
    }

    #[test]
    fn compile_repl_test() {
        let _ = ::env_logger::init();
        let vm = new_vm();
        compile_repl(&vm).unwrap_or_else(|err| panic!("{}", err));
        let repl: Result<FunctionRef<fn(()) -> IO<()>>, _> = vm.get_global("std.repl");
        assert!(repl.is_ok(), "{}", repl.err().unwrap());
    }

    type QueryFn = fn(&'static str) -> IO<Result<String, String>>;

    #[test]
    fn type_of_expr() {
        let _ = ::env_logger::init();
        let vm = new_vm();
        compile_repl(&vm).unwrap_or_else(|err| panic!("{}", err));
        let mut type_of: FunctionRef<QueryFn> = vm.get_global("repl_prim.type_of_expr").unwrap();
        assert_eq!(type_of.call("123"), Ok(IO::Value(Ok("Int".into()))));
    }

    #[test]
    fn find_kind() {
        let _ = ::env_logger::init();
        let vm = new_vm();
        compile_repl(&vm).unwrap_or_else(|err| panic!("{}", err));
        let mut find_kind: FunctionRef<QueryFn> = vm.get_global("repl_prim.find_kind").unwrap();
        assert_eq!(
            find_kind.call("std.prelude.Semigroup"),
            Ok(IO::Value(Ok("Type -> Type".into())))
        );
    }

    #[test]
    fn find_info() {
        let _ = ::env_logger::init();
        let vm = new_vm();
        compile_repl(&vm).unwrap_or_else(|err| panic!("{}", err));
        let mut find_info: FunctionRef<QueryFn> = vm.get_global("repl_prim.find_info").unwrap();
        match find_info.call("std.prelude.Semigroup") {
            Ok(IO::Value(Ok(_))) => (),
            x => assert!(false, "{:?}", x),
        }
        match find_info.call("std.prelude.make_Semigroup") {
            Ok(IO::Value(Ok(_))) => (),
            x => assert!(false, "{:?}", x),
        }
        match find_info.call("float_prim") {
            Ok(IO::Value(Ok(_))) => (),
            x => assert!(false, "{:?}", x),
        }
    }
}
<|endoftext|>


<commit_before>fn sum_vec(nums: &Vec<i32>) -> i32 {
    let mut result = 0;
    for i in nums {
        result += *i;
    }
    result
}

fn check_sub_array_sum(nums: Vec<i32>, k: i32) -> bool {
    let length = nums.len();
    if length <= 1 {
        return false;
    }

    if sum_vec(&nums) == 0 {
        return true;
    }

    if k == 0 {
        return false;
    }

    if length >= 2 && k > 0 {
        return true;
    }
}

fn main() {}
<commit_msg>not finish<commit_after>fn sum_vec(nums: &Vec<i32>) -> i32 {
    let mut result = 0;
    for i in nums {
        result += *i;
    }
    result
}

fn check_sub_array_sum(nums: Vec<i32>, k: i32) -> bool {
    let length = nums.len();
    if length <= 1 {
        return false;
    }

    if sum_vec(&nums) == 0 {
        return true;
    }

    if k == 0 {
        return false;
    }

    if length >= 2 && k > 0 {
        return true;
    }

    let inner_k = k.abs();
    //not finish
}

fn main() {}
<|endoftext|>