<commit_before>use std::{
    cell::RefCell,
    time::{Duration, Instant},
    mem,
    io::{stderr, Write},
    iter::repeat,
    collections::HashSet,
    sync::{RwLock, atomic::{AtomicBool, Ordering}},
};

use once_cell::sync::Lazy;
use itertools::Itertools;

/// Set profiling filter. It specifies descriptions allowed to profile.
/// This is helpful when call stack has too many nested profiling scopes.
/// Additionally filter can specify maximum depth of profiling scopes nesting.
///
/// #Example
/// ```
/// use ra_prof::{set_filter, Filter};
/// let f = Filter::from_spec("profile1|profile2@2");
/// set_filter(f);
/// ```
pub fn set_filter(f: Filter) {
    PROFILING_ENABLED.store(f.depth > 0, Ordering::SeqCst);
    let set: HashSet<_> = f.allowed.iter().cloned().collect();
    let mut old = FILTER.write().unwrap();
    let filter_data = FilterData {
        depth: f.depth,
        allowed: set,
        longer_than: f.longer_than,
        version: old.version + 1,
    };
    *old = filter_data;
}

/// 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(desc: &str) -> Profiler {
    assert!(!desc.is_empty());
    if !PROFILING_ENABLED.load(Ordering::Relaxed) {
        return Profiler { desc: None };
    }

    PROFILE_STACK.with(|stack| {
        let mut stack = stack.borrow_mut();
        if stack.starts.is_empty() {
            if let Ok(f) = FILTER.try_read() {
                if f.version > stack.filter_data.version {
                    stack.filter_data = f.clone();
                }
            };
        }

        if stack.starts.len() > stack.filter_data.depth {
            return Profiler { desc: None };
        }

        if stack.filter_data.allowed.is_empty() || stack.filter_data.allowed.contains(desc) {
            stack.starts.push(Instant::now());
            Profiler { desc: Some(desc.to_string()) }
        } else {
            Profiler { desc: None }
        }
    })
}

pub struct Profiler {
    desc: Option<String>,
}

pub struct Filter {
    depth: usize,
    allowed: Vec<String>,
    longer_than: Duration,
}

impl Filter {
    // 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 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 == "*" { Vec::new() } else { spec.split('|').map(String::from).collect() };
        Filter::new(depth, allowed, longer_than)
    }

    pub fn disabled() -> Filter {
        Filter::new(0, Vec::new(), Duration::new(0, 0))
    }

    pub fn new(depth: usize, allowed: Vec<String>, longer_than: Duration) -> Filter {
        Filter { depth, allowed, longer_than }
    }
}

struct ProfileStack {
    starts: Vec<Instant>,
    messages: Vec<Message>,
    filter_data: FilterData,
}

struct Message {
    level: usize,
    duration: Duration,
    message: String,
}

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

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

static PROFILING_ENABLED: AtomicBool = AtomicBool::new(false);

static FILTER: Lazy<RwLock<FilterData>> = Lazy::new(Default::default);

thread_local!(static PROFILE_STACK: RefCell<ProfileStack> = RefCell::new(ProfileStack::new()));

impl Drop for Profiler {
    fn drop(&mut self) {
        match self {
            Profiler { desc: Some(desc) } => {
                PROFILE_STACK.with(|stack| {
                    let mut stack = stack.borrow_mut();
                    let start = stack.starts.pop().unwrap();
                    let duration = start.elapsed();
                    let level = stack.starts.len();
                    let message = mem::replace(desc, String::new());
                    stack.messages.push(Message { level, duration, message });
                    if level == 0 {
                        let stdout = stderr();
                        let longer_than = stack.filter_data.longer_than;
                        if duration >= longer_than {
                            print(0, &stack.messages, &mut stdout.lock(), longer_than);
                        }
                        stack.messages.clear();
                    }
                });
            }
            Profiler { desc: None } => (),
        }
    }
}

fn print(lvl: usize, msgs: &[Message], out: &mut impl Write, longer_than: Duration) {
    let mut last = 0;
    let indent = repeat("    ").take(lvl + 1).collect::<String>();
    // We output hierarchy for long calls, but sum up all short calls
    let mut short = Vec::new();
    for (i, &Message { level, duration, message: ref msg }) in msgs.iter().enumerate() {
        if level != lvl {
            continue;
        }
        if duration >= longer_than {
            writeln!(out, "{} {:6}ms - {}", indent, duration.as_millis(), msg)
                .expect("printing profiling info to stdout");

            print(lvl + 1, &msgs[last..i], out, longer_than);
        } else {
            short.push((msg, duration))
        }

        last = i;
    }
    short.sort_by_key(|(msg, _time)| *msg);
    for (msg, entires) in short.iter().group_by(|(msg, _time)| msg).into_iter() {
        let mut count = 0;
        let mut total_duration = Duration::default();
        entires.for_each(|(_msg, time)| {
            count += 1;
            total_duration += *time;
        });
        writeln!(out, "{} {:6}ms - {} ({} calls)", indent, total_duration.as_millis(), msg, count)
            .expect("printing profiling info to stdout");
    }
}

#[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");
    }
}
<commit_msg>apply profile filter to top-level entries only<commit_after>use std::{
    cell::RefCell,
    time::{Duration, Instant},
    mem,
    io::{stderr, Write},
    iter::repeat,
    collections::HashSet,
    sync::{RwLock, atomic::{AtomicBool, Ordering}},
};

use once_cell::sync::Lazy;
use itertools::Itertools;

/// Set profiling filter. It specifies descriptions allowed to profile.
/// This is helpful when call stack has too many nested profiling scopes.
/// Additionally filter can specify maximum depth of profiling scopes nesting.
///
/// #Example
/// ```
/// use ra_prof::{set_filter, Filter};
/// let f = Filter::from_spec("profile1|profile2@2");
/// set_filter(f);
/// ```
pub fn set_filter(f: Filter) {
    PROFILING_ENABLED.store(f.depth > 0, Ordering::SeqCst);
    let set: HashSet<_> = f.allowed.iter().cloned().collect();
    let mut old = FILTER.write().unwrap();
    let filter_data = FilterData {
        depth: f.depth,
        allowed: set,
        longer_than: f.longer_than,
        version: old.version + 1,
    };
    *old = filter_data;
}

/// 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(desc: &str) -> Profiler {
    assert!(!desc.is_empty());
    if !PROFILING_ENABLED.load(Ordering::Relaxed) {
        return Profiler { desc: None };
    }

    PROFILE_STACK.with(|stack| {
        let mut stack = stack.borrow_mut();
        if stack.starts.is_empty() {
            if let Ok(f) = FILTER.try_read() {
                if f.version > stack.filter_data.version {
                    stack.filter_data = f.clone();
                }
            };
        }
        if stack.starts.len() > stack.filter_data.depth {
            return Profiler { desc: None };
        }
        let allowed = &stack.filter_data.allowed;
        if stack.starts.is_empty() && !allowed.is_empty() && !allowed.contains(desc) {
            return Profiler { desc: None };
        }

        stack.starts.push(Instant::now());
        Profiler { desc: Some(desc.to_string()) }
    })
}

pub struct Profiler {
    desc: Option<String>,
}

pub struct Filter {
    depth: usize,
    allowed: Vec<String>,
    longer_than: Duration,
}

impl Filter {
    // 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 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 == "*" { Vec::new() } else { spec.split('|').map(String::from).collect() };
        Filter::new(depth, allowed, longer_than)
    }

    pub fn disabled() -> Filter {
        Filter::new(0, Vec::new(), Duration::new(0, 0))
    }

    pub fn new(depth: usize, allowed: Vec<String>, longer_than: Duration) -> Filter {
        Filter { depth, allowed, longer_than }
    }
}

struct ProfileStack {
    starts: Vec<Instant>,
    messages: Vec<Message>,
    filter_data: FilterData,
}

struct Message {
    level: usize,
    duration: Duration,
    message: String,
}

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

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

static PROFILING_ENABLED: AtomicBool = AtomicBool::new(false);

static FILTER: Lazy<RwLock<FilterData>> = Lazy::new(Default::default);

thread_local!(static PROFILE_STACK: RefCell<ProfileStack> = RefCell::new(ProfileStack::new()));

impl Drop for Profiler {
    fn drop(&mut self) {
        match self {
            Profiler { desc: Some(desc) } => {
                PROFILE_STACK.with(|stack| {
                    let mut stack = stack.borrow_mut();
                    let start = stack.starts.pop().unwrap();
                    let duration = start.elapsed();
                    let level = stack.starts.len();
                    let message = mem::replace(desc, String::new());
                    stack.messages.push(Message { level, duration, message });
                    if level == 0 {
                        let stdout = stderr();
                        let longer_than = stack.filter_data.longer_than;
                        if duration >= longer_than {
                            print(0, &stack.messages, &mut stdout.lock(), longer_than);
                        }
                        stack.messages.clear();
                    }
                });
            }
            Profiler { desc: None } => (),
        }
    }
}

fn print(lvl: usize, msgs: &[Message], out: &mut impl Write, longer_than: Duration) {
    let mut last = 0;
    let indent = repeat("    ").take(lvl + 1).collect::<String>();
    // We output hierarchy for long calls, but sum up all short calls
    let mut short = Vec::new();
    for (i, &Message { level, duration, message: ref msg }) in msgs.iter().enumerate() {
        if level != lvl {
            continue;
        }
        if duration >= longer_than {
            writeln!(out, "{} {:6}ms - {}", indent, duration.as_millis(), msg)
                .expect("printing profiling info to stdout");

            print(lvl + 1, &msgs[last..i], out, longer_than);
        } else {
            short.push((msg, duration))
        }

        last = i;
    }
    short.sort_by_key(|(msg, _time)| *msg);
    for (msg, entires) in short.iter().group_by(|(msg, _time)| msg).into_iter() {
        let mut count = 0;
        let mut total_duration = Duration::default();
        entires.for_each(|(_msg, time)| {
            count += 1;
            total_duration += *time;
        });
        writeln!(out, "{} {:6}ms - {} ({} calls)", indent, total_duration.as_millis(), msg, count)
            .expect("printing profiling info to stdout");
    }
}

#[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");
    }
}
<|endoftext|>


<commit_before>//! This module defines `AssistCtx` -- the API surface that is exposed to assists.
use hir::{db::HirDatabase, InFile, SourceAnalyzer};
use ra_db::FileRange;
use ra_fmt::{leading_indent, reindent};
use ra_syntax::{
    algo::{self, find_covering_element, find_node_at_offset},
    AstNode, SourceFile, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxToken, TextRange, TextUnit,
    TokenAtOffset,
};
use ra_text_edit::TextEditBuilder;

use crate::{AssistAction, AssistId, AssistLabel};

#[derive(Clone, Debug)]
pub(crate) enum Assist {
    Unresolved { label: AssistLabel },
    Resolved { label: AssistLabel, action: AssistAction },
}

/// `AssistCtx` allows to apply an assist or check if it could be applied.
///
/// Assists use a somewhat over-engineered approach, given the current needs. The
/// assists workflow consists of two phases. In the first phase, a user asks for
/// the list of available assists. In the second phase, the user picks a
/// particular assist and it gets applied.
///
/// There are two peculiarities here:
///
/// * first, we ideally avoid computing more things then necessary to answer
///   "is assist applicable" in the first phase.
/// * second, when we are applying assist, we don't have a guarantee that there
///   weren't any changes between the point when user asked for assists and when
///   they applied a particular assist. So, when applying assist, we need to do
///   all the checks from scratch.
///
/// To avoid repeating the same code twice for both "check" and "apply"
/// functions, we use an approach reminiscent of that of Django's function based
/// views dealing with forms. Each assist receives a runtime parameter,
/// `should_compute_edit`. It first check if an edit is applicable (potentially
/// computing info required to compute the actual edit). If it is applicable,
/// and `should_compute_edit` is `true`, it then computes the actual edit.
///
/// So, to implement the original assists workflow, we can first apply each edit
/// with `should_compute_edit = false`, and then applying the selected edit
/// again, with `should_compute_edit = true` this time.
///
/// Note, however, that we don't actually use such two-phase logic at the
/// moment, because the LSP API is pretty awkward in this place, and it's much
/// easier to just compute the edit eagerly :-)
#[derive(Debug)]
pub(crate) struct AssistCtx<'a, DB> {
    pub(crate) db: &'a DB,
    pub(crate) frange: FileRange,
    source_file: SourceFile,
    should_compute_edit: bool,
}

impl<'a, DB> Clone for AssistCtx<'a, DB> {
    fn clone(&self) -> Self {
        AssistCtx {
            db: self.db,
            frange: self.frange,
            source_file: self.source_file.clone(),
            should_compute_edit: self.should_compute_edit,
        }
    }
}

impl<'a, DB: HirDatabase> AssistCtx<'a, DB> {
    pub(crate) fn with_ctx<F, T>(db: &DB, frange: FileRange, should_compute_edit: bool, f: F) -> T
    where
        F: FnOnce(AssistCtx<DB>) -> T,
    {
        let parse = db.parse(frange.file_id);

        let ctx = AssistCtx { db, frange, source_file: parse.tree(), should_compute_edit };
        f(ctx)
    }

    pub(crate) fn add_assist(
        self,
        id: AssistId,
        label: impl Into<String>,
        f: impl FnOnce(&mut AssistBuilder),
    ) -> Option<Assist> {
        let label = AssistLabel { label: label.into(), id };
        assert_eq!(
            label.label.chars().nth(0).and_then(|c| Some(c.is_uppercase())),
            Some(true),
            "First character should be uppercase"
        );

        let assist = if self.should_compute_edit {
            let action = {
                let mut edit = AssistBuilder::default();
                f(&mut edit);
                edit.build()
            };
            Assist::Resolved { label, action }
        } else {
            Assist::Unresolved { label }
        };

        Some(assist)
    }

    pub(crate) fn token_at_offset(&self) -> TokenAtOffset<SyntaxToken> {
        self.source_file.syntax().token_at_offset(self.frange.range.start())
    }

    pub(crate) fn find_token_at_offset(&self, kind: SyntaxKind) -> Option<SyntaxToken> {
        self.token_at_offset().find(|it| it.kind() == kind)
    }

    pub(crate) fn find_node_at_offset<N: AstNode>(&self) -> Option<N> {
        find_node_at_offset(self.source_file.syntax(), self.frange.range.start())
    }
    pub(crate) fn covering_element(&self) -> SyntaxElement {
        find_covering_element(self.source_file.syntax(), self.frange.range)
    }
    pub(crate) fn source_analyzer(
        &self,
        node: &SyntaxNode,
        offset: Option<TextUnit>,
    ) -> SourceAnalyzer {
        SourceAnalyzer::new(self.db, InFile::new(self.frange.file_id.into(), node), offset)
    }

    pub(crate) fn covering_node_for_range(&self, range: TextRange) -> SyntaxElement {
        find_covering_element(self.source_file.syntax(), range)
    }
}

#[derive(Default)]
pub(crate) struct AssistBuilder {
    edit: TextEditBuilder,
    cursor_position: Option<TextUnit>,
    target: Option<TextRange>,
}

impl AssistBuilder {
    /// Replaces specified `range` of text with a given string.
    pub(crate) fn replace(&mut self, range: TextRange, replace_with: impl Into<String>) {
        self.edit.replace(range, replace_with.into())
    }

    /// Replaces specified `node` of text with a given string, reindenting the
    /// string to maintain `node`'s existing indent.
    // FIXME: remove in favor of ra_syntax::edit::IndentLevel::increase_indent
    pub(crate) fn replace_node_and_indent(
        &mut self,
        node: &SyntaxNode,
        replace_with: impl Into<String>,
    ) {
        let mut replace_with = replace_with.into();
        if let Some(indent) = leading_indent(node) {
            replace_with = reindent(&replace_with, &indent)
        }
        self.replace(node.text_range(), replace_with)
    }

    /// Remove specified `range` of text.
    #[allow(unused)]
    pub(crate) fn delete(&mut self, range: TextRange) {
        self.edit.delete(range)
    }

    /// Append specified `text` at the given `offset`
    pub(crate) fn insert(&mut self, offset: TextUnit, text: impl Into<String>) {
        self.edit.insert(offset, text.into())
    }

    /// Specify desired position of the cursor after the assist is applied.
    pub(crate) fn set_cursor(&mut self, offset: TextUnit) {
        self.cursor_position = Some(offset)
    }

    /// Specify that the assist should be active withing the `target` range.
    ///
    /// Target ranges are used to sort assists: the smaller the target range,
    /// the more specific assist is, and so it should be sorted first.
    pub(crate) fn target(&mut self, target: TextRange) {
        self.target = Some(target)
    }

    /// Get access to the raw `TextEditBuilder`.
    pub(crate) fn text_edit_builder(&mut self) -> &mut TextEditBuilder {
        &mut self.edit
    }

    pub(crate) fn replace_ast<N: AstNode>(&mut self, old: N, new: N) {
        algo::diff(old.syntax(), new.syntax()).into_text_edit(&mut self.edit)
    }

    fn build(self) -> AssistAction {
        AssistAction {
            edit: self.edit.finish(),
            cursor_position: self.cursor_position,
            target: self.target,
        }
    }
}
<commit_msg>unwrap<commit_after>//! This module defines `AssistCtx` -- the API surface that is exposed to assists.
use hir::{db::HirDatabase, InFile, SourceAnalyzer};
use ra_db::FileRange;
use ra_fmt::{leading_indent, reindent};
use ra_syntax::{
    algo::{self, find_covering_element, find_node_at_offset},
    AstNode, SourceFile, SyntaxElement, SyntaxKind, SyntaxNode, SyntaxToken, TextRange, TextUnit,
    TokenAtOffset,
};
use ra_text_edit::TextEditBuilder;

use crate::{AssistAction, AssistId, AssistLabel};

#[derive(Clone, Debug)]
pub(crate) enum Assist {
    Unresolved { label: AssistLabel },
    Resolved { label: AssistLabel, action: AssistAction },
}

/// `AssistCtx` allows to apply an assist or check if it could be applied.
///
/// Assists use a somewhat over-engineered approach, given the current needs. The
/// assists workflow consists of two phases. In the first phase, a user asks for
/// the list of available assists. In the second phase, the user picks a
/// particular assist and it gets applied.
///
/// There are two peculiarities here:
///
/// * first, we ideally avoid computing more things then necessary to answer
///   "is assist applicable" in the first phase.
/// * second, when we are applying assist, we don't have a guarantee that there
///   weren't any changes between the point when user asked for assists and when
///   they applied a particular assist. So, when applying assist, we need to do
///   all the checks from scratch.
///
/// To avoid repeating the same code twice for both "check" and "apply"
/// functions, we use an approach reminiscent of that of Django's function based
/// views dealing with forms. Each assist receives a runtime parameter,
/// `should_compute_edit`. It first check if an edit is applicable (potentially
/// computing info required to compute the actual edit). If it is applicable,
/// and `should_compute_edit` is `true`, it then computes the actual edit.
///
/// So, to implement the original assists workflow, we can first apply each edit
/// with `should_compute_edit = false`, and then applying the selected edit
/// again, with `should_compute_edit = true` this time.
///
/// Note, however, that we don't actually use such two-phase logic at the
/// moment, because the LSP API is pretty awkward in this place, and it's much
/// easier to just compute the edit eagerly :-)
#[derive(Debug)]
pub(crate) struct AssistCtx<'a, DB> {
    pub(crate) db: &'a DB,
    pub(crate) frange: FileRange,
    source_file: SourceFile,
    should_compute_edit: bool,
}

impl<'a, DB> Clone for AssistCtx<'a, DB> {
    fn clone(&self) -> Self {
        AssistCtx {
            db: self.db,
            frange: self.frange,
            source_file: self.source_file.clone(),
            should_compute_edit: self.should_compute_edit,
        }
    }
}

impl<'a, DB: HirDatabase> AssistCtx<'a, DB> {
    pub(crate) fn with_ctx<F, T>(db: &DB, frange: FileRange, should_compute_edit: bool, f: F) -> T
    where
        F: FnOnce(AssistCtx<DB>) -> T,
    {
        let parse = db.parse(frange.file_id);

        let ctx = AssistCtx { db, frange, source_file: parse.tree(), should_compute_edit };
        f(ctx)
    }

    pub(crate) fn add_assist(
        self,
        id: AssistId,
        label: impl Into<String>,
        f: impl FnOnce(&mut AssistBuilder),
    ) -> Option<Assist> {
        let label = AssistLabel { label: label.into(), id };
        assert_eq!(
            label.label.chars().nth(0).and_then(|c| Some(c.is_uppercase())).unwrap(),
            true,
            "First character should be uppercase"
        );

        let assist = if self.should_compute_edit {
            let action = {
                let mut edit = AssistBuilder::default();
                f(&mut edit);
                edit.build()
            };
            Assist::Resolved { label, action }
        } else {
            Assist::Unresolved { label }
        };

        Some(assist)
    }

    pub(crate) fn token_at_offset(&self) -> TokenAtOffset<SyntaxToken> {
        self.source_file.syntax().token_at_offset(self.frange.range.start())
    }

    pub(crate) fn find_token_at_offset(&self, kind: SyntaxKind) -> Option<SyntaxToken> {
        self.token_at_offset().find(|it| it.kind() == kind)
    }

    pub(crate) fn find_node_at_offset<N: AstNode>(&self) -> Option<N> {
        find_node_at_offset(self.source_file.syntax(), self.frange.range.start())
    }
    pub(crate) fn covering_element(&self) -> SyntaxElement {
        find_covering_element(self.source_file.syntax(), self.frange.range)
    }
    pub(crate) fn source_analyzer(
        &self,
        node: &SyntaxNode,
        offset: Option<TextUnit>,
    ) -> SourceAnalyzer {
        SourceAnalyzer::new(self.db, InFile::new(self.frange.file_id.into(), node), offset)
    }

    pub(crate) fn covering_node_for_range(&self, range: TextRange) -> SyntaxElement {
        find_covering_element(self.source_file.syntax(), range)
    }
}

#[derive(Default)]
pub(crate) struct AssistBuilder {
    edit: TextEditBuilder,
    cursor_position: Option<TextUnit>,
    target: Option<TextRange>,
}

impl AssistBuilder {
    /// Replaces specified `range` of text with a given string.
    pub(crate) fn replace(&mut self, range: TextRange, replace_with: impl Into<String>) {
        self.edit.replace(range, replace_with.into())
    }

    /// Replaces specified `node` of text with a given string, reindenting the
    /// string to maintain `node`'s existing indent.
    // FIXME: remove in favor of ra_syntax::edit::IndentLevel::increase_indent
    pub(crate) fn replace_node_and_indent(
        &mut self,
        node: &SyntaxNode,
        replace_with: impl Into<String>,
    ) {
        let mut replace_with = replace_with.into();
        if let Some(indent) = leading_indent(node) {
            replace_with = reindent(&replace_with, &indent)
        }
        self.replace(node.text_range(), replace_with)
    }

    /// Remove specified `range` of text.
    #[allow(unused)]
    pub(crate) fn delete(&mut self, range: TextRange) {
        self.edit.delete(range)
    }

    /// Append specified `text` at the given `offset`
    pub(crate) fn insert(&mut self, offset: TextUnit, text: impl Into<String>) {
        self.edit.insert(offset, text.into())
    }

    /// Specify desired position of the cursor after the assist is applied.
    pub(crate) fn set_cursor(&mut self, offset: TextUnit) {
        self.cursor_position = Some(offset)
    }

    /// Specify that the assist should be active withing the `target` range.
    ///
    /// Target ranges are used to sort assists: the smaller the target range,
    /// the more specific assist is, and so it should be sorted first.
    pub(crate) fn target(&mut self, target: TextRange) {
        self.target = Some(target)
    }

    /// Get access to the raw `TextEditBuilder`.
    pub(crate) fn text_edit_builder(&mut self) -> &mut TextEditBuilder {
        &mut self.edit
    }

    pub(crate) fn replace_ast<N: AstNode>(&mut self, old: N, new: N) {
        algo::diff(old.syntax(), new.syntax()).into_text_edit(&mut self.edit)
    }

    fn build(self) -> AssistAction {
        AssistAction {
            edit: self.edit.finish(),
            cursor_position: self.cursor_position,
            target: self.target,
        }
    }
}
<|endoftext|>


<commit_before>use std::{iter, ops::RangeInclusive};

use arrayvec::ArrayVec;
use ra_text_edit::TextEditBuilder;
use ra_syntax::{AstNode, TreeArc, ast, SyntaxKind::*, SyntaxElement, SourceFile, InsertPosition, Direction};
use ra_fmt::leading_indent;

pub struct AstEditor<N: AstNode> {
    original_ast: TreeArc<N>,
    ast: TreeArc<N>,
}

impl<N: AstNode> AstEditor<N> {
    pub fn new(node: &N) -> AstEditor<N> {
        AstEditor { original_ast: node.to_owned(), ast: node.to_owned() }
    }

    pub fn into_text_edit(self, builder: &mut TextEditBuilder) {
        // FIXME: compute a more fine-grained diff here.
        // If *you* know a nice algorithm to compute diff between two syntax
        // tree, tell me about it!
        builder.replace(self.original_ast.syntax().range(), self.ast().syntax().text().to_string());
    }

    pub fn ast(&self) -> &N {
        &*self.ast
    }

    #[must_use]
    fn insert_children<'a>(
        &self,
        position: InsertPosition<SyntaxElement<'_>>,
        to_insert: impl Iterator<Item = SyntaxElement<'a>>,
    ) -> TreeArc<N> {
        let new_syntax = self.ast().syntax().insert_children(position, to_insert);
        N::cast(&new_syntax).unwrap().to_owned()
    }

    #[must_use]
    fn replace_children<'a>(
        &self,
        to_delete: RangeInclusive<SyntaxElement<'_>>,
        to_insert: impl Iterator<Item = SyntaxElement<'a>>,
    ) -> TreeArc<N> {
        let new_syntax = self.ast().syntax().replace_children(to_delete, to_insert);
        N::cast(&new_syntax).unwrap().to_owned()
    }

    fn do_make_multiline(&mut self) {
        let l_curly =
            match self.ast().syntax().children_with_tokens().find(|it| it.kind() == L_CURLY) {
                Some(it) => it,
                None => return,
            };
        let sibling = match l_curly.next_sibling_or_token() {
            Some(it) => it,
            None => return,
        };
        let existing_ws = match sibling.as_token() {
            None => None,
            Some(tok) if tok.kind() != WHITESPACE => None,
            Some(ws) => {
                if ws.text().contains('\n') {
                    return;
                }
                Some(ws)
            }
        };

        let indent = leading_indent(self.ast().syntax()).unwrap_or("");
        let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
        let to_insert = iter::once(ws.ws().into());
        self.ast = match existing_ws {
            None => self.insert_children(InsertPosition::After(l_curly), to_insert),
            Some(ws) => self.replace_children(RangeInclusive::new(ws.into(), ws.into()), to_insert),
        };
    }
}

impl AstEditor<ast::NamedFieldList> {
    pub fn append_field(&mut self, field: &ast::NamedField) {
        self.insert_field(InsertPosition::Last, field)
    }

    pub fn make_multiline(&mut self) {
        self.do_make_multiline()
    }

    pub fn insert_field(
        &mut self,
        position: InsertPosition<&'_ ast::NamedField>,
        field: &ast::NamedField,
    ) {
        let is_multiline = self.ast().syntax().text().contains('\n');
        let ws;
        let space = if is_multiline {
            ws = tokens::WsBuilder::new(&format!(
                "\n{}    ",
                leading_indent(self.ast().syntax()).unwrap_or("")
            ));
            ws.ws()
        } else {
            tokens::single_space()
        };

        let mut to_insert: ArrayVec<[SyntaxElement; 4]> = ArrayVec::new();
        to_insert.push(space.into());
        to_insert.push(field.syntax().into());
        to_insert.push(tokens::comma().into());

        macro_rules! after_l_curly {
            () => {{
                let anchor = match self.l_curly() {
                    Some(it) => it,
                    None => return,
                };
                InsertPosition::After(anchor)
            }};
        }

        macro_rules! after_field {
            ($anchor:expr) => {
                if let Some(comma) = $anchor
                    .syntax()
                    .siblings_with_tokens(Direction::Next)
                    .find(|it| it.kind() == COMMA)
                {
                    InsertPosition::After(comma)
                } else {
                    to_insert.insert(0, tokens::comma().into());
                    InsertPosition::After($anchor.syntax().into())
                }
            };
        };

        let position = match position {
            InsertPosition::First => after_l_curly!(),
            InsertPosition::Last => {
                if !is_multiline {
                    // don't insert comma before curly
                    to_insert.pop();
                }
                match self.ast().fields().last() {
                    Some(it) => after_field!(it),
                    None => after_l_curly!(),
                }
            }
            InsertPosition::Before(anchor) => InsertPosition::Before(anchor.syntax().into()),
            InsertPosition::After(anchor) => after_field!(anchor),
        };

        self.ast = self.insert_children(position, to_insert.iter().cloned());
    }

    fn l_curly(&self) -> Option<SyntaxElement> {
        self.ast().syntax().children_with_tokens().find(|it| it.kind() == L_CURLY)
    }
}

impl AstEditor<ast::ItemList> {
    pub fn make_multiline(&mut self) {
        self.do_make_multiline()
    }

    pub fn append_functions<'a>(&mut self, fns: impl Iterator<Item = &'a ast::FnDef>) {
        fns.for_each(|it| self.append_function(it))
    }

    pub fn append_function(&mut self, fn_def: &ast::FnDef) {
        let (indent, position) = match self.ast().impl_items().last() {
            Some(it) => (
                leading_indent(it.syntax()).unwrap_or("").to_string(),
                InsertPosition::After(it.syntax().into()),
            ),
            None => match self.l_curly() {
                Some(it) => (
                    "    ".to_string() + leading_indent(self.ast().syntax()).unwrap_or(""),
                    InsertPosition::After(it),
                ),
                None => return,
            },
        };
        let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
        let to_insert: ArrayVec<[SyntaxElement; 2]> =
            [ws.ws().into(), fn_def.syntax().into()].into();
        self.ast = self.insert_children(position, to_insert.into_iter());
    }

    fn l_curly(&self) -> Option<SyntaxElement> {
        self.ast().syntax().children_with_tokens().find(|it| it.kind() == L_CURLY)
    }
}

impl AstEditor<ast::FnDef> {
    pub fn set_body(&mut self, body: &ast::Block) {
        let mut to_insert: ArrayVec<[SyntaxElement; 2]> = ArrayVec::new();
        let old_body_or_semi: SyntaxElement = if let Some(old_body) = self.ast().body() {
            old_body.syntax().into()
        } else if let Some(semi) = self.ast().semicolon_token() {
            to_insert.push(tokens::single_space().into());
            semi.into()
        } else {
            to_insert.push(tokens::single_space().into());
            to_insert.push(body.syntax().into());
            self.ast = self.insert_children(InsertPosition::Last, to_insert.into_iter());
            return;
        };
        to_insert.push(body.syntax().into());
        let replace_range = RangeInclusive::new(old_body_or_semi, old_body_or_semi);
        self.ast = self.replace_children(replace_range, to_insert.into_iter())
    }

    pub fn strip_attrs_and_docs(&mut self) {
        loop {
            if let Some(start) = self
                .ast()
                .syntax()
                .children_with_tokens()
                .find(|it| it.kind() == ATTR || it.kind() == COMMENT)
            {
                let end = match start.next_sibling_or_token() {
                    Some(el) if el.kind() == WHITESPACE => el,
                    Some(_) | None => start,
                };
                self.ast = self.replace_children(RangeInclusive::new(start, end), iter::empty());
            } else {
                break;
            }
        }
    }
}

pub struct AstBuilder<N: AstNode> {
    _phantom: std::marker::PhantomData<N>,
}

impl AstBuilder<ast::NamedField> {
    fn from_text(text: &str) -> TreeArc<ast::NamedField> {
        ast_node_from_file_text(&format!("fn f() {{ S {{ {}, }} }}", text))
    }

    pub fn from_pieces(name: &ast::NameRef, expr: Option<&ast::Expr>) -> TreeArc<ast::NamedField> {
        match expr {
            Some(expr) => Self::from_text(&format!("{}: {}", name.syntax(), expr.syntax())),
            None => Self::from_text(&name.syntax().to_string()),
        }
    }
}

impl AstBuilder<ast::Block> {
    fn from_text(text: &str) -> TreeArc<ast::Block> {
        ast_node_from_file_text(&format!("fn f() {}", text))
    }

    pub fn single_expr(e: &ast::Expr) -> TreeArc<ast::Block> {
        Self::from_text(&format!("{{ {} }}", e.syntax()))
    }
}

impl AstBuilder<ast::Expr> {
    fn from_text(text: &str) -> TreeArc<ast::Expr> {
        ast_node_from_file_text(&format!("fn f() {{ {}; }}", text))
    }

    pub fn unit() -> TreeArc<ast::Expr> {
        Self::from_text("()")
    }

    pub fn unimplemented() -> TreeArc<ast::Expr> {
        Self::from_text("unimplemented!()")
    }
}

impl AstBuilder<ast::NameRef> {
    pub fn new(text: &str) -> TreeArc<ast::NameRef> {
        ast_node_from_file_text(&format!("fn f() {{ {}; }}", text))
    }
}

fn ast_node_from_file_text<N: AstNode>(text: &str) -> TreeArc<N> {
    let file = SourceFile::parse(text);
    let res = file.syntax().descendants().find_map(N::cast).unwrap().to_owned();
    res
}

mod tokens {
    use lazy_static::lazy_static;
    use ra_syntax::{AstNode, SourceFile, TreeArc, SyntaxToken, SyntaxKind::*};

    lazy_static! {
        static ref SOURCE_FILE: TreeArc<SourceFile> = SourceFile::parse(",\n; ;");
    }

    pub(crate) fn comma() -> SyntaxToken<'static> {
        SOURCE_FILE
            .syntax()
            .descendants_with_tokens()
            .filter_map(|it| it.as_token())
            .find(|it| it.kind() == COMMA)
            .unwrap()
    }

    pub(crate) fn single_space() -> SyntaxToken<'static> {
        SOURCE_FILE
            .syntax()
            .descendants_with_tokens()
            .filter_map(|it| it.as_token())
            .find(|it| it.kind() == WHITESPACE && it.text().as_str() == " ")
            .unwrap()
    }

    #[allow(unused)]
    pub(crate) fn single_newline() -> SyntaxToken<'static> {
        SOURCE_FILE
            .syntax()
            .descendants_with_tokens()
            .filter_map(|it| it.as_token())
            .find(|it| it.kind() == WHITESPACE && it.text().as_str() == "\n")
            .unwrap()
    }

    pub(crate) struct WsBuilder(TreeArc<SourceFile>);

    impl WsBuilder {
        pub(crate) fn new(text: &str) -> WsBuilder {
            WsBuilder(SourceFile::parse(text))
        }
        pub(crate) fn ws(&self) -> SyntaxToken<'_> {
            self.0.syntax().first_child_or_token().unwrap().as_token().unwrap()
        }
    }

}

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

    use ra_syntax::SourceFile;

    #[test]
    fn structure_editing() {
        let file = SourceFile::parse(
            "\
fn foo() {
    let s = S {
        original: 92,
    }
}
",
        );
        let field_list = file.syntax().descendants().find_map(ast::NamedFieldList::cast).unwrap();
        let mut editor = AstEditor::new(field_list);

        let field = AstBuilder::<ast::NamedField>::from_text("first_inserted: 1");
        editor.append_field(&field);
        let field = AstBuilder::<ast::NamedField>::from_text("second_inserted: 2");
        editor.append_field(&field);
        eprintln!("{}", editor.ast().syntax());
    }
}
<commit_msg>drop useless test<commit_after>use std::{iter, ops::RangeInclusive};

use arrayvec::ArrayVec;
use ra_text_edit::TextEditBuilder;
use ra_syntax::{AstNode, TreeArc, ast, SyntaxKind::*, SyntaxElement, SourceFile, InsertPosition, Direction};
use ra_fmt::leading_indent;

pub struct AstEditor<N: AstNode> {
    original_ast: TreeArc<N>,
    ast: TreeArc<N>,
}

impl<N: AstNode> AstEditor<N> {
    pub fn new(node: &N) -> AstEditor<N> {
        AstEditor { original_ast: node.to_owned(), ast: node.to_owned() }
    }

    pub fn into_text_edit(self, builder: &mut TextEditBuilder) {
        // FIXME: compute a more fine-grained diff here.
        // If *you* know a nice algorithm to compute diff between two syntax
        // tree, tell me about it!
        builder.replace(self.original_ast.syntax().range(), self.ast().syntax().text().to_string());
    }

    pub fn ast(&self) -> &N {
        &*self.ast
    }

    #[must_use]
    fn insert_children<'a>(
        &self,
        position: InsertPosition<SyntaxElement<'_>>,
        to_insert: impl Iterator<Item = SyntaxElement<'a>>,
    ) -> TreeArc<N> {
        let new_syntax = self.ast().syntax().insert_children(position, to_insert);
        N::cast(&new_syntax).unwrap().to_owned()
    }

    #[must_use]
    fn replace_children<'a>(
        &self,
        to_delete: RangeInclusive<SyntaxElement<'_>>,
        to_insert: impl Iterator<Item = SyntaxElement<'a>>,
    ) -> TreeArc<N> {
        let new_syntax = self.ast().syntax().replace_children(to_delete, to_insert);
        N::cast(&new_syntax).unwrap().to_owned()
    }

    fn do_make_multiline(&mut self) {
        let l_curly =
            match self.ast().syntax().children_with_tokens().find(|it| it.kind() == L_CURLY) {
                Some(it) => it,
                None => return,
            };
        let sibling = match l_curly.next_sibling_or_token() {
            Some(it) => it,
            None => return,
        };
        let existing_ws = match sibling.as_token() {
            None => None,
            Some(tok) if tok.kind() != WHITESPACE => None,
            Some(ws) => {
                if ws.text().contains('\n') {
                    return;
                }
                Some(ws)
            }
        };

        let indent = leading_indent(self.ast().syntax()).unwrap_or("");
        let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
        let to_insert = iter::once(ws.ws().into());
        self.ast = match existing_ws {
            None => self.insert_children(InsertPosition::After(l_curly), to_insert),
            Some(ws) => self.replace_children(RangeInclusive::new(ws.into(), ws.into()), to_insert),
        };
    }
}

impl AstEditor<ast::NamedFieldList> {
    pub fn append_field(&mut self, field: &ast::NamedField) {
        self.insert_field(InsertPosition::Last, field)
    }

    pub fn make_multiline(&mut self) {
        self.do_make_multiline()
    }

    pub fn insert_field(
        &mut self,
        position: InsertPosition<&'_ ast::NamedField>,
        field: &ast::NamedField,
    ) {
        let is_multiline = self.ast().syntax().text().contains('\n');
        let ws;
        let space = if is_multiline {
            ws = tokens::WsBuilder::new(&format!(
                "\n{}    ",
                leading_indent(self.ast().syntax()).unwrap_or("")
            ));
            ws.ws()
        } else {
            tokens::single_space()
        };

        let mut to_insert: ArrayVec<[SyntaxElement; 4]> = ArrayVec::new();
        to_insert.push(space.into());
        to_insert.push(field.syntax().into());
        to_insert.push(tokens::comma().into());

        macro_rules! after_l_curly {
            () => {{
                let anchor = match self.l_curly() {
                    Some(it) => it,
                    None => return,
                };
                InsertPosition::After(anchor)
            }};
        }

        macro_rules! after_field {
            ($anchor:expr) => {
                if let Some(comma) = $anchor
                    .syntax()
                    .siblings_with_tokens(Direction::Next)
                    .find(|it| it.kind() == COMMA)
                {
                    InsertPosition::After(comma)
                } else {
                    to_insert.insert(0, tokens::comma().into());
                    InsertPosition::After($anchor.syntax().into())
                }
            };
        };

        let position = match position {
            InsertPosition::First => after_l_curly!(),
            InsertPosition::Last => {
                if !is_multiline {
                    // don't insert comma before curly
                    to_insert.pop();
                }
                match self.ast().fields().last() {
                    Some(it) => after_field!(it),
                    None => after_l_curly!(),
                }
            }
            InsertPosition::Before(anchor) => InsertPosition::Before(anchor.syntax().into()),
            InsertPosition::After(anchor) => after_field!(anchor),
        };

        self.ast = self.insert_children(position, to_insert.iter().cloned());
    }

    fn l_curly(&self) -> Option<SyntaxElement> {
        self.ast().syntax().children_with_tokens().find(|it| it.kind() == L_CURLY)
    }
}

impl AstEditor<ast::ItemList> {
    pub fn make_multiline(&mut self) {
        self.do_make_multiline()
    }

    pub fn append_functions<'a>(&mut self, fns: impl Iterator<Item = &'a ast::FnDef>) {
        fns.for_each(|it| self.append_function(it))
    }

    pub fn append_function(&mut self, fn_def: &ast::FnDef) {
        let (indent, position) = match self.ast().impl_items().last() {
            Some(it) => (
                leading_indent(it.syntax()).unwrap_or("").to_string(),
                InsertPosition::After(it.syntax().into()),
            ),
            None => match self.l_curly() {
                Some(it) => (
                    "    ".to_string() + leading_indent(self.ast().syntax()).unwrap_or(""),
                    InsertPosition::After(it),
                ),
                None => return,
            },
        };
        let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
        let to_insert: ArrayVec<[SyntaxElement; 2]> =
            [ws.ws().into(), fn_def.syntax().into()].into();
        self.ast = self.insert_children(position, to_insert.into_iter());
    }

    fn l_curly(&self) -> Option<SyntaxElement> {
        self.ast().syntax().children_with_tokens().find(|it| it.kind() == L_CURLY)
    }
}

impl AstEditor<ast::FnDef> {
    pub fn set_body(&mut self, body: &ast::Block) {
        let mut to_insert: ArrayVec<[SyntaxElement; 2]> = ArrayVec::new();
        let old_body_or_semi: SyntaxElement = if let Some(old_body) = self.ast().body() {
            old_body.syntax().into()
        } else if let Some(semi) = self.ast().semicolon_token() {
            to_insert.push(tokens::single_space().into());
            semi.into()
        } else {
            to_insert.push(tokens::single_space().into());
            to_insert.push(body.syntax().into());
            self.ast = self.insert_children(InsertPosition::Last, to_insert.into_iter());
            return;
        };
        to_insert.push(body.syntax().into());
        let replace_range = RangeInclusive::new(old_body_or_semi, old_body_or_semi);
        self.ast = self.replace_children(replace_range, to_insert.into_iter())
    }

    pub fn strip_attrs_and_docs(&mut self) {
        loop {
            if let Some(start) = self
                .ast()
                .syntax()
                .children_with_tokens()
                .find(|it| it.kind() == ATTR || it.kind() == COMMENT)
            {
                let end = match start.next_sibling_or_token() {
                    Some(el) if el.kind() == WHITESPACE => el,
                    Some(_) | None => start,
                };
                self.ast = self.replace_children(RangeInclusive::new(start, end), iter::empty());
            } else {
                break;
            }
        }
    }
}

pub struct AstBuilder<N: AstNode> {
    _phantom: std::marker::PhantomData<N>,
}

impl AstBuilder<ast::NamedField> {
    fn from_text(text: &str) -> TreeArc<ast::NamedField> {
        ast_node_from_file_text(&format!("fn f() {{ S {{ {}, }} }}", text))
    }

    pub fn from_pieces(name: &ast::NameRef, expr: Option<&ast::Expr>) -> TreeArc<ast::NamedField> {
        match expr {
            Some(expr) => Self::from_text(&format!("{}: {}", name.syntax(), expr.syntax())),
            None => Self::from_text(&name.syntax().to_string()),
        }
    }
}

impl AstBuilder<ast::Block> {
    fn from_text(text: &str) -> TreeArc<ast::Block> {
        ast_node_from_file_text(&format!("fn f() {}", text))
    }

    pub fn single_expr(e: &ast::Expr) -> TreeArc<ast::Block> {
        Self::from_text(&format!("{{ {} }}", e.syntax()))
    }
}

impl AstBuilder<ast::Expr> {
    fn from_text(text: &str) -> TreeArc<ast::Expr> {
        ast_node_from_file_text(&format!("fn f() {{ {}; }}", text))
    }

    pub fn unit() -> TreeArc<ast::Expr> {
        Self::from_text("()")
    }

    pub fn unimplemented() -> TreeArc<ast::Expr> {
        Self::from_text("unimplemented!()")
    }
}

impl AstBuilder<ast::NameRef> {
    pub fn new(text: &str) -> TreeArc<ast::NameRef> {
        ast_node_from_file_text(&format!("fn f() {{ {}; }}", text))
    }
}

fn ast_node_from_file_text<N: AstNode>(text: &str) -> TreeArc<N> {
    let file = SourceFile::parse(text);
    let res = file.syntax().descendants().find_map(N::cast).unwrap().to_owned();
    res
}

mod tokens {
    use lazy_static::lazy_static;
    use ra_syntax::{AstNode, SourceFile, TreeArc, SyntaxToken, SyntaxKind::*};

    lazy_static! {
        static ref SOURCE_FILE: TreeArc<SourceFile> = SourceFile::parse(",\n; ;");
    }

    pub(crate) fn comma() -> SyntaxToken<'static> {
        SOURCE_FILE
            .syntax()
            .descendants_with_tokens()
            .filter_map(|it| it.as_token())
            .find(|it| it.kind() == COMMA)
            .unwrap()
    }

    pub(crate) fn single_space() -> SyntaxToken<'static> {
        SOURCE_FILE
            .syntax()
            .descendants_with_tokens()
            .filter_map(|it| it.as_token())
            .find(|it| it.kind() == WHITESPACE && it.text().as_str() == " ")
            .unwrap()
    }

    #[allow(unused)]
    pub(crate) fn single_newline() -> SyntaxToken<'static> {
        SOURCE_FILE
            .syntax()
            .descendants_with_tokens()
            .filter_map(|it| it.as_token())
            .find(|it| it.kind() == WHITESPACE && it.text().as_str() == "\n")
            .unwrap()
    }

    pub(crate) struct WsBuilder(TreeArc<SourceFile>);

    impl WsBuilder {
        pub(crate) fn new(text: &str) -> WsBuilder {
            WsBuilder(SourceFile::parse(text))
        }
        pub(crate) fn ws(&self) -> SyntaxToken<'_> {
            self.0.syntax().first_child_or_token().unwrap().as_token().unwrap()
        }
    }

}
<|endoftext|>


<commit_before>//! Set of allocatable registers as a bit vector of register units.
//!
//! While allocating registers, we need to keep track of which registers are available and which
//! registers are in use. Since registers can alias in different ways, we track this via the
//! "register unit" abstraction. Every register contains one or more register units. Registers that
//! share a register unit can't be in use at the same time.

use std::mem::size_of_val;
use isa::registers::{RegUnit, RegUnitMask, RegClass};

/// Set of registers available for allocation.
#[derive(Clone)]
pub struct AllocatableSet {
    avail: RegUnitMask,
}

// Given a register class and a register unit in the class, compute a word index and a bit mask of
// register units representing that register.
//
// Note that a register is not allowed to straddle words.
fn bitmask(rc: RegClass, reg: RegUnit) -> (usize, u32) {
    // Bit mask representing the register. It is `rc.width` consecutive units.
    let width_bits = (1 << rc.width) - 1;
    // Index into avail[] of the word containing `reg`.
    let word_index = (reg / 32) as usize;
    // The actual bits in the word that cover `reg`.
    let reg_bits = width_bits << (reg % 32);

    (word_index, reg_bits)
}

impl AllocatableSet {
    /// Create a new register set with all registers available.
    ///
    /// Note that this includes *all* registers. Query the `TargetIsa` object to get a set of
    /// allocatable registers where reserved registers have been filtered out.
    pub fn new() -> AllocatableSet {
        AllocatableSet { avail: [!0; 3] }
    }

    /// Returns `true` if the specified register is available.
    pub fn is_avail(&self, rc: RegClass, reg: RegUnit) -> bool {
        let (idx, bits) = bitmask(rc, reg);
        (self.avail[idx] & bits) == bits
    }

    /// Allocate `reg` from `rc` so it is no longer available.
    ///
    /// It is an error to take a register that doesn't have all of its register units available.
    pub fn take(&mut self, rc: RegClass, reg: RegUnit) {
        let (idx, bits) = bitmask(rc, reg);
        debug_assert!((self.avail[idx] & bits) == bits, "Not available");
        self.avail[idx] &= !bits;
    }

    /// Make `reg` available for allocation again.
    pub fn free(&mut self, rc: RegClass, reg: RegUnit) {
        let (idx, bits) = bitmask(rc, reg);
        debug_assert!((self.avail[idx] & bits) == 0, "Not allocated");
        self.avail[idx] |= bits;
    }

    /// Return an iterator over all available registers belonging to the register class `rc`.
    ///
    /// This doesn't allocate anything from the set; use `take()` for that.
    pub fn iter(&self, rc: RegClass) -> RegSetIter {
        // Start by copying the RC mask. It is a single set bit for each register in the class.
        let mut rsi = RegSetIter { regs: rc.mask };

        // Mask out the unavailable units.
        for idx in 0..self.avail.len() {
            // If a single unit in a register is unavailable, the whole register can't be used.
            // If a register straddles a word boundary, it will be marked as unavailable.
            // There's an assertion in `cdsl/registers.py` to check for that.
            for i in 0..rc.width {
                rsi.regs[idx] &= self.avail[idx] >> i;
            }
        }
        rsi
    }

    /// Check if any register units allocated out of this set interferes with units allocated out
    /// of `other`.
    ///
    /// This assumes that unused bits are 1.
    pub fn interferes_with(&self, other: &AllocatableSet) -> bool {
        self.avail
            .iter()
            .zip(&other.avail)
            .any(|(&x, &y)| (x | y) != !0)
    }

    /// Intersect this set of allocatable registers with `other`. This has the effect of removing
    /// any register units from this set that are not in `other`.
    pub fn intersect(&mut self, other: &AllocatableSet) {
        for (x, &y) in self.avail.iter_mut().zip(&other.avail) {
            *x &= y;
        }
    }
}

/// Iterator over available registers in a register class.
pub struct RegSetIter {
    regs: RegUnitMask,
}

impl Iterator for RegSetIter {
    type Item = RegUnit;

    fn next(&mut self) -> Option<RegUnit> {
        let mut unit_offset = 0;

        // Find the first set bit in `self.regs`.
        for word in &mut self.regs {
            if *word != 0 {
                // Compute the register unit number from the lowest set bit in the word.
                let unit = unit_offset + word.trailing_zeros() as RegUnit;

                // Clear that lowest bit so we won't find it again.
                *word = *word & (*word - 1);

                return Some(unit);
            }
            // How many register units was there in the word? This is a constant 32 for `u32` etc.
            unit_offset += 8 * size_of_val(word) as RegUnit;
        }

        // All of `self.regs` is 0.
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use isa::registers::{RegClass, RegClassData};

    // Register classes for testing.
    const GPR: RegClass = &RegClassData {
        name: "GPR",
        index: 0,
        width: 1,
        bank: 0,
        toprc: 0,
        first: 28,
        subclasses: 0,
        mask: [0xf0000000, 0x0000000f, 0],
    };
    const DPR: RegClass = &RegClassData {
        name: "DPR",
        index: 0,
        width: 2,
        bank: 0,
        toprc: 0,
        first: 28,
        subclasses: 0,
        mask: [0x50000000, 0x0000000a, 0],
    };

    #[test]
    fn put_and_take() {
        let mut regs = AllocatableSet::new();

        // `GPR` has units 28-36.
        assert_eq!(regs.iter(GPR).count(), 8);
        assert_eq!(regs.iter(DPR).collect::<Vec<_>>(), [28, 30, 33, 35]);

        assert!(regs.is_avail(GPR, 29));
        regs.take(&GPR, 29);
        assert!(!regs.is_avail(GPR, 29));

        assert_eq!(regs.iter(GPR).count(), 7);
        assert_eq!(regs.iter(DPR).collect::<Vec<_>>(), [30, 33, 35]);

        assert!(regs.is_avail(GPR, 30));
        regs.take(&GPR, 30);
        assert!(!regs.is_avail(GPR, 30));

        assert_eq!(regs.iter(GPR).count(), 6);
        assert_eq!(regs.iter(DPR).collect::<Vec<_>>(), [33, 35]);

        assert!(regs.is_avail(GPR, 32));
        regs.take(&GPR, 32);
        assert!(!regs.is_avail(GPR, 32));

        assert_eq!(regs.iter(GPR).count(), 5);
        assert_eq!(regs.iter(DPR).collect::<Vec<_>>(), [33, 35]);

        regs.free(&GPR, 30);
        assert!(regs.is_avail(GPR, 30));
        assert!(!regs.is_avail(GPR, 29));
        assert!(!regs.is_avail(GPR, 32));

        assert_eq!(regs.iter(GPR).count(), 6);
        assert_eq!(regs.iter(DPR).collect::<Vec<_>>(), [30, 33, 35]);

        regs.free(&GPR, 32);
        assert!(regs.is_avail(GPR, 31));
        assert!(!regs.is_avail(GPR, 29));
        assert!(regs.is_avail(GPR, 32));

        assert_eq!(regs.iter(GPR).count(), 7);
        assert_eq!(regs.iter(DPR).collect::<Vec<_>>(), [30, 33, 35]);
    }

    #[test]
    fn interference() {
        let mut regs1 = AllocatableSet::new();
        let mut regs2 = AllocatableSet::new();

        assert!(!regs1.interferes_with(&regs2));
        regs1.take(&GPR, 32);
        assert!(!regs1.interferes_with(&regs2));
        regs2.take(&GPR, 31);
        assert!(!regs1.interferes_with(&regs2));
        regs1.intersect(&regs2);
        assert!(regs1.interferes_with(&regs2));
    }
}
<commit_msg>Implement ExactSizeIterator for RegSetIter.<commit_after>//! Set of allocatable registers as a bit vector of register units.
//!
//! While allocating registers, we need to keep track of which registers are available and which
//! registers are in use. Since registers can alias in different ways, we track this via the
//! "register unit" abstraction. Every register contains one or more register units. Registers that
//! share a register unit can't be in use at the same time.

use isa::registers::{RegUnit, RegUnitMask, RegClass};
use std::iter::ExactSizeIterator;
use std::mem::size_of_val;

/// Set of registers available for allocation.
#[derive(Clone)]
pub struct AllocatableSet {
    avail: RegUnitMask,
}

// Given a register class and a register unit in the class, compute a word index and a bit mask of
// register units representing that register.
//
// Note that a register is not allowed to straddle words.
fn bitmask(rc: RegClass, reg: RegUnit) -> (usize, u32) {
    // Bit mask representing the register. It is `rc.width` consecutive units.
    let width_bits = (1 << rc.width) - 1;
    // Index into avail[] of the word containing `reg`.
    let word_index = (reg / 32) as usize;
    // The actual bits in the word that cover `reg`.
    let reg_bits = width_bits << (reg % 32);

    (word_index, reg_bits)
}

impl AllocatableSet {
    /// Create a new register set with all registers available.
    ///
    /// Note that this includes *all* registers. Query the `TargetIsa` object to get a set of
    /// allocatable registers where reserved registers have been filtered out.
    pub fn new() -> AllocatableSet {
        AllocatableSet { avail: [!0; 3] }
    }

    /// Returns `true` if the specified register is available.
    pub fn is_avail(&self, rc: RegClass, reg: RegUnit) -> bool {
        let (idx, bits) = bitmask(rc, reg);
        (self.avail[idx] & bits) == bits
    }

    /// Allocate `reg` from `rc` so it is no longer available.
    ///
    /// It is an error to take a register that doesn't have all of its register units available.
    pub fn take(&mut self, rc: RegClass, reg: RegUnit) {
        let (idx, bits) = bitmask(rc, reg);
        debug_assert!((self.avail[idx] & bits) == bits, "Not available");
        self.avail[idx] &= !bits;
    }

    /// Make `reg` available for allocation again.
    pub fn free(&mut self, rc: RegClass, reg: RegUnit) {
        let (idx, bits) = bitmask(rc, reg);
        debug_assert!((self.avail[idx] & bits) == 0, "Not allocated");
        self.avail[idx] |= bits;
    }

    /// Return an iterator over all available registers belonging to the register class `rc`.
    ///
    /// This doesn't allocate anything from the set; use `take()` for that.
    pub fn iter(&self, rc: RegClass) -> RegSetIter {
        // Start by copying the RC mask. It is a single set bit for each register in the class.
        let mut rsi = RegSetIter { regs: rc.mask };

        // Mask out the unavailable units.
        for idx in 0..self.avail.len() {
            // If a single unit in a register is unavailable, the whole register can't be used.
            // If a register straddles a word boundary, it will be marked as unavailable.
            // There's an assertion in `cdsl/registers.py` to check for that.
            for i in 0..rc.width {
                rsi.regs[idx] &= self.avail[idx] >> i;
            }
        }
        rsi
    }

    /// Check if any register units allocated out of this set interferes with units allocated out
    /// of `other`.
    ///
    /// This assumes that unused bits are 1.
    pub fn interferes_with(&self, other: &AllocatableSet) -> bool {
        self.avail
            .iter()
            .zip(&other.avail)
            .any(|(&x, &y)| (x | y) != !0)
    }

    /// Intersect this set of allocatable registers with `other`. This has the effect of removing
    /// any register units from this set that are not in `other`.
    pub fn intersect(&mut self, other: &AllocatableSet) {
        for (x, &y) in self.avail.iter_mut().zip(&other.avail) {
            *x &= y;
        }
    }
}

/// Iterator over available registers in a register class.
pub struct RegSetIter {
    regs: RegUnitMask,
}

impl Iterator for RegSetIter {
    type Item = RegUnit;

    fn next(&mut self) -> Option<RegUnit> {
        let mut unit_offset = 0;

        // Find the first set bit in `self.regs`.
        for word in &mut self.regs {
            if *word != 0 {
                // Compute the register unit number from the lowest set bit in the word.
                let unit = unit_offset + word.trailing_zeros() as RegUnit;

                // Clear that lowest bit so we won't find it again.
                *word = *word & (*word - 1);

                return Some(unit);
            }
            // How many register units was there in the word? This is a constant 32 for `u32` etc.
            unit_offset += 8 * size_of_val(word) as RegUnit;
        }

        // All of `self.regs` is 0.
        None
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let bits = self.regs.iter().map(|&w| w.count_ones() as usize).sum();
        (bits, Some(bits))
    }
}

impl ExactSizeIterator for RegSetIter {}

#[cfg(test)]
mod tests {
    use super::*;
    use isa::registers::{RegClass, RegClassData};

    // Register classes for testing.
    const GPR: RegClass = &RegClassData {
        name: "GPR",
        index: 0,
        width: 1,
        bank: 0,
        toprc: 0,
        first: 28,
        subclasses: 0,
        mask: [0xf0000000, 0x0000000f, 0],
    };
    const DPR: RegClass = &RegClassData {
        name: "DPR",
        index: 0,
        width: 2,
        bank: 0,
        toprc: 0,
        first: 28,
        subclasses: 0,
        mask: [0x50000000, 0x0000000a, 0],
    };

    #[test]
    fn put_and_take() {
        let mut regs = AllocatableSet::new();

        // `GPR` has units 28-36.
        assert_eq!(regs.iter(GPR).len(), 8);
        assert_eq!(regs.iter(GPR).count(), 8);
        assert_eq!(regs.iter(DPR).collect::<Vec<_>>(), [28, 30, 33, 35]);

        assert!(regs.is_avail(GPR, 29));
        regs.take(&GPR, 29);
        assert!(!regs.is_avail(GPR, 29));

        assert_eq!(regs.iter(GPR).count(), 7);
        assert_eq!(regs.iter(DPR).collect::<Vec<_>>(), [30, 33, 35]);

        assert!(regs.is_avail(GPR, 30));
        regs.take(&GPR, 30);
        assert!(!regs.is_avail(GPR, 30));

        assert_eq!(regs.iter(GPR).count(), 6);
        assert_eq!(regs.iter(DPR).collect::<Vec<_>>(), [33, 35]);

        assert!(regs.is_avail(GPR, 32));
        regs.take(&GPR, 32);
        assert!(!regs.is_avail(GPR, 32));

        assert_eq!(regs.iter(GPR).count(), 5);
        assert_eq!(regs.iter(DPR).collect::<Vec<_>>(), [33, 35]);

        regs.free(&GPR, 30);
        assert!(regs.is_avail(GPR, 30));
        assert!(!regs.is_avail(GPR, 29));
        assert!(!regs.is_avail(GPR, 32));

        assert_eq!(regs.iter(GPR).count(), 6);
        assert_eq!(regs.iter(DPR).collect::<Vec<_>>(), [30, 33, 35]);

        regs.free(&GPR, 32);
        assert!(regs.is_avail(GPR, 31));
        assert!(!regs.is_avail(GPR, 29));
        assert!(regs.is_avail(GPR, 32));

        assert_eq!(regs.iter(GPR).count(), 7);
        assert_eq!(regs.iter(DPR).collect::<Vec<_>>(), [30, 33, 35]);
    }

    #[test]
    fn interference() {
        let mut regs1 = AllocatableSet::new();
        let mut regs2 = AllocatableSet::new();

        assert!(!regs1.interferes_with(&regs2));
        regs1.take(&GPR, 32);
        assert!(!regs1.interferes_with(&regs2));
        regs2.take(&GPR, 31);
        assert!(!regs1.interferes_with(&regs2));
        regs1.intersect(&regs2);
        assert!(regs1.interferes_with(&regs2));
    }
}
<|endoftext|>


<commit_before>
extern crate amqp;

use amqp::{ConsumerCallBackFn, Session, Table, Basic, Channel, Options};
use amqp::protocol;
use std::default::Default;

fn consumer_function(channel: &mut Channel, deliver: protocol::basic::Deliver, headers: protocol::basic::BasicProperties, body: Vec<u8>){
    println!("Got a delivery:");
    println!("Deliver info: {:?}", deliver);
    println!("Content headers: {:?}", headers);
    println!("Content body: {:?}", body);
    channel.basic_ack(deliver.delivery_tag, false);
}

fn main() {
    let mut session = Session::new(Options{vhost: "/", .. Default::default()}).ok().expect("Can't create session");
    println!("Openned session:");
    let mut channel = session.open_channel(1).expect("channel");
    println!("Openned channel: {}", channel.id);

    let queue_name = "test_queue";
    //queue: &str, passive: bool, durable: bool, exclusive: bool, auto_delete: bool, nowait: bool, arguments: Table
    let queue_declare = channel.queue_declare(queue_name, false, true, false, false, false, Table::new());
    println!("Queue declare: {:?}", queue_declare);

    //queue: &str, consumer_tag: &str, no_local: bool, no_ack: bool, exclusive: bool, nowait: bool, arguments: Table
    println!("Declaring consumer...");
    let consumer_name = channel.basic_consume(consumer_function as ConsumerCallBackFn, queue_name, "", false, false, false, false, Table::new());
    println!("Starting consumer {:?}", consumer_name);

    let props = protocol::basic::BasicProperties{ content_type: Some("text".to_string()), ..Default::default()};
    let res = channel.basic_publish("", queue_name, true, false, props, (b"Hello from rust!").to_vec());

    println!("Published");

    channel.start_consuming();

    let closed = channel.close(200, "Bye".to_string());
    session.close(200, "Good Bye".to_string());
}
<commit_msg>Split into different functions<commit_after>
extern crate amqp;

use amqp::{ConsumerCallBackFn, Session, Table, Basic, Channel, Options};
use amqp::protocol;
use std::default::Default;

// for debugging
fn listen_discovery(channel: &mut Channel) {
    fn participant_discovered(channel: &mut Channel, deliver: protocol::basic::Deliver,
                            headers: protocol::basic::BasicProperties, body: Vec<u8>){
        println!("Discovered MsgFlo participant !!:");
        //println!("Deliver info: {:?}", deliver);
        //println!("Content headers: {:?}", headers);
        //println!("Content body: {:?}", body);
        channel.basic_ack(deliver.delivery_tag, false);
    }

    let queue_name = "fbp";
    let consumer_name = channel.basic_consume(participant_discovered as ConsumerCallBackFn, queue_name,
                                                "", false, false, false, false, Table::new());
}

fn send_discovery(channel: &mut Channel) {
    let queue_name = "fbp"; // TODO: use an exchange istead, requires protocol change in msgflo

    let queue_declare = channel.queue_declare(queue_name, false, true, false, false, false, Table::new());
    let content_type = Some("application/json".to_string());
    let props = protocol::basic::BasicProperties { content_type: content_type, ..Default::default() };

    let payload = b"Hello from rust!";
    let res = channel.basic_publish("", queue_name, true, false, props, payload.to_vec());
}

// for debugging
fn selftest() {
    let mut session = Session::new(Options{vhost: "/", .. Default::default()}).ok().expect("Can't create session");
    let mut channel = session.open_channel(1).expect("channel");

    send_discovery(&mut channel);
    listen_discovery(&mut channel);

    channel.start_consuming();

    let closed = channel.close(200, "Bye".to_string());
    session.close(200, "Good Bye".to_string());
}

fn main() {
    selftest();
}
<|endoftext|>


<commit_before>
use std::iter;
use std::ascii::AsciiExt;

static STANDARD_CHARS: &'static[u8] = b"abcdefghijklmnopqrstuvwxyz234567";

pub fn encode(b: &[u8]) -> String {
    let mut out = Vec::new();
    let mut n: u64 = 0;

    for i in 0 .. b.len() {
        // Calculate the shift (32, 24, 16, 8, 0)
        let shift = 8 * (4 - (i % 5));
        n |= (b[i] as u64) << shift;

        // If the shift was zero, or this is the last character of input, extract
        // each 5-bit group from n.
        if shift == 0 || (i == b.len()-1) {
            // Determine number of output groups, based on the shift
            let groups = match shift {
                32 => 2,
                24 => 4,
                16 => 5,
                8  => 7,
                _  => 8
            };

            // Pull out each 5-bit index
            for j in iter::range_step(35, (35 - (groups * 5)), -5) {
                let ndx = (n >> j as u32) & 31;
                out.push(STANDARD_CHARS[ndx as usize]);
            }

            // Reset n
            n = 0;
        }
    }

    let result = match String::from_utf8(out) {
        Ok(s) => s,
        Err(e) => panic!(e)
    };
    result
}

pub fn decode(b: &[u8]) -> Vec<u8> {
    let mut out: Vec<u8> = Vec::new();
    let mut n: u64 = 0;

    for i in range(0, b.len()) {
        // Use the character to generate an index
        let c = b[i];
        let idx = match c as char {
            'A' ... 'Z' => c - 65, // A..Z => 65..90
            'a' ... 'z' => c - 97, // a..z => 97..122
            _ => 26 + (c - 50)     // 2..7 => 50..55
        };

        // Calculate the shift for this 5-bit group (35, 25, ..., 0)
        let shift = 35 - (5 * (i % 8));
        n |= (idx as u64 & 31) << shift;

        // If the shift is zero or this is the last character of input, extract
        // each 8-bit group from n.
        if shift == 0 || (i == b.len()-1) {
            // Determine number of output groups, based on shift
            let groups = match shift {
                35 => 0,
                30 => 1,
                25 => 2,
                20 => 2,
                15 => 3,
                10 => 4,
                5  => 4,
                _  => 5
            };

            // Pull out each 8-bit group
            for j in iter::range_step(32, (32 - (groups * 8)), -8) {
                let v = ((n >> j as u8) & 255) as u8;
                out.push(v);
            }

            // Reset n
            n = 0;
        }
    }

    out
}

macro_rules! check_encode(
    ($input:expr, $expected:expr) => ( {
        let x = encode($input.as_bytes());
        assert!(x.eq_ignore_ascii_case($expected));
    });
    );

macro_rules! check_decode(
    ($expected:expr, $input:expr) => (
        assert_eq!($expected.as_bytes(), decode($input.as_bytes()));
        );
);

#[test]
fn encode_tests() {
    check_encode!("", "");
    check_encode!("f", "MY");
    check_encode!("foo", "MZXw6");
    check_encode!("foob", "MZXW6YQ");
    check_encode!("fooba", "MZXW6YTB");
    check_encode!("foobar", "MZXW6YTBOI");
}

#[test]
fn decode_tests() {
    check_decode!("", "");
    check_decode!("f", "MY");
    check_decode!("foo", "MZXW6");
    check_decode!("foob", "MZxw6YQ");
    check_decode!("fooba", "MZXW6YTB");
    check_decode!("foobar", "MZXW6YTBOI");
}

<commit_msg>Fix iteration for 1.0.0-beta<commit_after>
use std::iter;
use std::ascii::AsciiExt;

static STANDARD_CHARS: &'static[u8] = b"abcdefghijklmnopqrstuvwxyz234567";

pub fn encode(b: &[u8]) -> String {
    let mut out = Vec::new();
    let mut n: u64 = 0;

    for i in 0 .. b.len() {
        // Calculate the shift (32, 24, 16, 8, 0)
        let shift = 8 * (4 - (i % 5));
        n |= (b[i] as u64) << shift;

        // If the shift was zero, or this is the last character of input, extract
        // each 5-bit group from n.
        if shift == 0 || (i == b.len()-1) {
            // Determine number of output groups, based on the shift
            let groups = match shift {
                32 => 2,
                24 => 4,
                16 => 5,
                8  => 7,
                _  => 8
            };

            // Pull out each 5-bit index
            for j in iter::range_step_inclusive(35, (35 - (groups * 5)), -5) {
                let ndx = (n >> j as u32) & 31;
                out.push(STANDARD_CHARS[ndx as usize]);
            }

            // Reset n
            n = 0;
        }
    }

    let result = match String::from_utf8(out) {
        Ok(s) => s,
        Err(e) => panic!(e)
    };
    result
}

pub fn decode(b: &[u8]) -> Vec<u8> {
    let mut out: Vec<u8> = Vec::new();
    let mut n: u64 = 0;

    for i in iter::range_inclusive(0, b.len()) {
        // Use the character to generate an index
        let c = b[i];
        let idx = match c as char {
            'A' ... 'Z' => c - 65, // A..Z => 65..90
            'a' ... 'z' => c - 97, // a..z => 97..122
            _ => 26 + (c - 50)     // 2..7 => 50..55
        };

        // Calculate the shift for this 5-bit group (35, 25, ..., 0)
        let shift = 35 - (5 * (i % 8));
        n |= (idx as u64 & 31) << shift;

        // If the shift is zero or this is the last character of input, extract
        // each 8-bit group from n.
        if shift == 0 || (i == b.len()-1) {
            // Determine number of output groups, based on shift
            let groups = match shift {
                35 => 0,
                30 => 1,
                25 => 2,
                20 => 2,
                15 => 3,
                10 => 4,
                5  => 4,
                _  => 5
            };

            // Pull out each 8-bit group
            for j in iter::range_step_inclusive(32, (32 - (groups * 8)), -8) {
                let v = ((n >> j as u8) & 255) as u8;
                out.push(v);
            }

            // Reset n
            n = 0;
        }
    }

    out
}

macro_rules! check_encode(
    ($input:expr, $expected:expr) => ( {
        let x = encode($input.as_bytes());
        assert!(x.eq_ignore_ascii_case($expected));
    });
    );

macro_rules! check_decode(
    ($expected:expr, $input:expr) => (
        assert_eq!($expected, decode($input.as_bytes()));
        );
);

#[test]
fn encode_tests() {
    check_encode!("", "");
    check_encode!("f", "MY");
    check_encode!("foo", "MZXw6");
    check_encode!("foob", "MZXW6YQ");
    check_encode!("fooba", "MZXW6YTB");
    check_encode!("foobar", "MZXW6YTBOI");
}

#[test]
fn decode_tests() {
    check_decode!("", "");
    check_decode!("f", "MY");
    check_decode!("foo", "MZXW6");
    check_decode!("foob", "MZxw6YQ");
    check_decode!("fooba", "MZXW6YTB");
    check_decode!("foobar", "MZXW6YTBOI");
}

<|endoftext|>


<commit_before>use std::{mem, ptr, u32};

/// This is a pretty naive implementation of a BigUint abstracting all
/// math out to a vector of `u32` chunks.
///
/// It can only do 3 things:
/// - Be instantiated from an arbitrary byte slice
/// - Do a division by `u32`, mutating self and returning the remainder.
/// - Check if it's zero.
///
/// Turns out those are all the operations you need to encode base58,
/// or anything else, really.
pub struct BigUint {
    pub chunks: Vec<u32>
}

impl BigUint {
    /// Divide self by `divider`, return the remainder of the operation.
    #[inline(always)]
    pub fn rem_div(&mut self, divider: u32) -> u32 {
        let mut carry = 0u64;

        let mut pop = 0;

        for chunk in self.chunks.iter_mut() {
            carry = (carry << 32) | *chunk as u64;
            *chunk = (carry / divider as u64) as u32;
            carry %= divider as u64;

            if *chunk == 0 {
                pop += 1;
            }
        }

        // TODO: Do this without a loop (might need unsafe)
        for _ in 0..pop {
            self.chunks.remove(0);
        }

        carry as u32
    }

    /// Check if self is zero.
    #[inline(always)]
    pub fn is_zero(&self) -> bool {
        self.chunks.iter().all(|chunk| *chunk == 0)
    }
}

/// Helper function for transmuting 4 bytes into a chunk.
#[inline(always)]
fn bytes_to_u32(bytes: [u8; 4]) -> u32 {
    u32::from_be(unsafe { mem::transmute(bytes) })
}

/// Helper function for transmuting a slice into a chunk.
#[inline(always)]
fn slice_to_u32(slice: &[u8]) -> u32 {
    debug_assert!(slice.len() == 4);

    unsafe {
        let mut bytes: [u8; 4] = mem::uninitialized();

        ptr::copy_nonoverlapping(
            slice.as_ptr(),
            bytes.as_mut_ptr(),
            4
        );

        bytes_to_u32(bytes)
    }
}

impl<'a> From<&'a [u8]> for BigUint {
    #[inline]
    fn from(mut bytes: &[u8]) -> Self {
        let modulo = bytes.len() % 4;

        let mut chunks = Vec::new();

        if modulo > 0 {
            chunks.reserve(bytes.len() / 4 + 1);

            let mut first = [0u8; 4];

            for (r, w) in bytes[..modulo].iter().zip(first[4-modulo..].iter_mut()) {
                *w = *r;
            }

            bytes = &bytes[modulo..];

            chunks.push(bytes_to_u32(first));
        } else {
            chunks.reserve(bytes.len() / 4);
        }

        for slice in bytes.chunks(4) {
            chunks.push(slice_to_u32(slice))
        }

        BigUint {
            chunks: chunks
        }
    }
}

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

    #[test]
    fn big_uint_from_bytes() {
        let bytes: &[u8] = &[
                      0xDE,0xAD,0x00,0x00,0x00,0x13,
            0x37,0xAD,0x00,0x00,0x00,0x00,0xDE,0xAD,
        ];

        let big = BigUint::from(bytes);

        assert_eq!(big.chunks, vec![0x0000DEAD, 0x00000013, 0x37AD0000, 0x0000DEAD]);
    }

    #[test]
    fn big_uint_rem_div() {
        let mut big = BigUint {
            chunks: vec![0x136AD712,0x84322759]
        };

        let rem = big.rem_div(58);
        let merged = ((big.chunks[0] as u64) << 32) | big.chunks[1] as u64;

        assert_eq!(merged, 0x136AD71284322759 / 58);
        assert_eq!(rem as u64, 0x136AD71284322759 % 58);
    }
}
<commit_msg>Faster BigUint init<commit_after>use std::{ptr, u32};

/// This is a pretty naive implementation of a BigUint abstracting all
/// math out to a vector of `u32` chunks.
///
/// It can only do 3 things:
/// - Be instantiated from an arbitrary byte slice
/// - Do a division by `u32`, mutating self and returning the remainder.
/// - Check if it's zero.
///
/// Turns out those are all the operations you need to encode base58,
/// or anything else, really.
pub struct BigUint {
    pub chunks: Vec<u32>
}

impl BigUint {
    /// Divide self by `divider`, return the remainder of the operation.
    #[inline(always)]
    pub fn rem_div(&mut self, divider: u32) -> u32 {
        let mut carry = 0u64;

        for chunk in self.chunks.iter_mut() {
            carry = (carry << 32) | *chunk as u64;
            *chunk = (carry / divider as u64) as u32;
            carry %= divider as u64;
        }

        if self.chunks[0] == 0 {
            self.chunks.remove(0);
        }

        carry as u32
    }

    /// Check if self is zero.
    #[inline(always)]
    pub fn is_zero(&self) -> bool {
        self.chunks.iter().all(|chunk| *chunk == 0)
    }
}

impl<'a> From<&'a [u8]> for BigUint {
    #[inline]
    fn from(bytes: &[u8]) -> Self {
        let modulo = bytes.len() % 4;

        let len = bytes.len() / 4 + (modulo > 0) as usize;

        let mut chunks = Vec::with_capacity(len);

        unsafe {
            chunks.set_len(len);
            *chunks.get_unchecked_mut(0) = 0u32;

            ptr::copy_nonoverlapping(
                bytes.as_ptr(),
                (chunks.as_mut_ptr() as *mut u8).offset(modulo as isize),
                bytes.len()
            );
        }

        for chunk in chunks.iter_mut() {
            *chunk = u32::from_be(*chunk);
        }

        BigUint {
            chunks: chunks
        }
    }
}

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

    #[test]
    fn big_uint_from_bytes() {
        let bytes: &[u8] = &[
                      0xDE,0xAD,0x00,0x00,0x00,0x13,
            0x37,0xAD,0x00,0x00,0x00,0x00,0xDE,0xAD,
        ];

        let big = BigUint::from(bytes);

        assert_eq!(big.chunks, vec![0x0000DEAD, 0x00000013, 0x37AD0000, 0x0000DEAD]);
    }

    #[test]
    fn big_uint_rem_div() {
        let mut big = BigUint {
            chunks: vec![0x136AD712,0x84322759]
        };

        let rem = big.rem_div(58);
        let merged = ((big.chunks[0] as u64) << 32) | big.chunks[1] as u64;

        assert_eq!(merged, 0x136AD71284322759 / 58);
        assert_eq!(rem as u64, 0x136AD71284322759 % 58);
    }
}
<|endoftext|>


<commit_before>use crate::cell::Cell;

pub struct RWLock {
    mode: Cell<isize>,
}

unsafe impl Send for RWLock {}
unsafe impl Sync for RWLock {} // no threads on this platform

impl RWLock {
    pub const fn new() -> RWLock {
        RWLock { mode: Cell::new(0) }
    }

    #[inline]
    pub unsafe fn read(&self) {
        let m = self.mode.get();
        if m >= 0 {
            self.mode.set(m + 1);
        } else {
            rtabort!("rwlock locked for writing");
        }
    }

    #[inline]
    pub unsafe fn try_read(&self) -> bool {
        let m = self.mode.get();
        if m >= 0 {
            self.mode.set(m + 1);
            true
        } else {
            false
        }
    }

    #[inline]
    pub unsafe fn write(&self) {
        if self.mode.replace(-1) != 0 {
            rtabort!("rwlock locked for reading")
        }
    }

    #[inline]
    pub unsafe fn try_write(&self) -> bool {
        if self.mode.get() == 0 {
            self.mode.set(-1);
            true
        } else {
            false
        }
    }

    #[inline]
    pub unsafe fn read_unlock(&self) {
        self.mode.set(self.mode.get() - 1);
    }

    #[inline]
    pub unsafe fn write_unlock(&self) {
        self.mode.set(0);
    }

    #[inline]
    pub unsafe fn destroy(&self) {}
}
<commit_msg>Assert state in sys/unsupported's RwLock::write_unlock.<commit_after>use crate::cell::Cell;

pub struct RWLock {
    mode: Cell<isize>,
}

unsafe impl Send for RWLock {}
unsafe impl Sync for RWLock {} // no threads on this platform

impl RWLock {
    pub const fn new() -> RWLock {
        RWLock { mode: Cell::new(0) }
    }

    #[inline]
    pub unsafe fn read(&self) {
        let m = self.mode.get();
        if m >= 0 {
            self.mode.set(m + 1);
        } else {
            rtabort!("rwlock locked for writing");
        }
    }

    #[inline]
    pub unsafe fn try_read(&self) -> bool {
        let m = self.mode.get();
        if m >= 0 {
            self.mode.set(m + 1);
            true
        } else {
            false
        }
    }

    #[inline]
    pub unsafe fn write(&self) {
        if self.mode.replace(-1) != 0 {
            rtabort!("rwlock locked for reading")
        }
    }

    #[inline]
    pub unsafe fn try_write(&self) -> bool {
        if self.mode.get() == 0 {
            self.mode.set(-1);
            true
        } else {
            false
        }
    }

    #[inline]
    pub unsafe fn read_unlock(&self) {
        self.mode.set(self.mode.get() - 1);
    }

    #[inline]
    pub unsafe fn write_unlock(&self) {
        assert_eq!(self.mode.replace(0), -1);
    }

    #[inline]
    pub unsafe fn destroy(&self) {}
}
<|endoftext|>