<commit_before>use XY;
use view::View;
use view::SizeCache;
use vec::Vec2;
use printer::Printer;
use orientation::Orientation;
use event::{Event, EventResult, Key};

use std::cmp::min;

/// Arranges its children linearly according to its orientation.
pub struct LinearLayout {
    children: Vec<Child>,
    orientation: Orientation,
    focus: usize,

    cache: Option<XY<SizeCache>>,
}

struct Child {
    view: Box<View>,
    size: Vec2,
    weight: usize,
}

impl Child {
    fn get_min_size(&mut self, req: Vec2) -> Vec2 {
        self.size = self.view.get_min_size(req);
        self.size
    }

    fn as_ref(&self) -> &View {
        &*self.view
    }

    fn as_mut(&mut self) -> &mut View {
        &mut *self.view
    }
}

impl LinearLayout {
    /// Creates a new layout with the given orientation.
    pub fn new(orientation: Orientation) -> Self {
        LinearLayout {
            children: Vec::new(),
            orientation: orientation,
            focus: 0,
            cache: None,
        }
    }

    /// Modifies the weight of the last child added.
    ///
    /// It is an error to call this before adding a child (and it will panic).
    pub fn weight(mut self, weight: usize) -> Self {
        self.children.last_mut().unwrap().weight = weight;

        self
    }

    /// Adds a child to the layout.
    pub fn child<V: View + 'static>(mut self, view: V) -> Self {
        self.children.push(Child {
            view: Box::new(view),
            size: Vec2::zero(),
            weight: 0,
        });
        self.invalidate();

        self
    }

    // Invalidate the view, to request a layout next time
    fn invalidate(&mut self) {
        self.cache = None;
    }

    /// Creates a new vertical layout.
    pub fn vertical() -> Self {
        LinearLayout::new(Orientation::Vertical)
    }

    /// Creates a new horizontal layout.
    pub fn horizontal() -> Self {
        LinearLayout::new(Orientation::Horizontal)
    }

    // If the cache can be used, return the cached size.
    // Otherwise, return None.
    fn get_cache(&self, req: Vec2) -> Option<Vec2> {
        match self.cache {
            None => None,
            Some(ref cache) => {
                // Is our cache even valid?
                // Also, is any child invalidating the layout?
                if cache.zip_map(req, SizeCache::accept).both() &&
                   self.children_are_sleeping() {
                    Some(cache.map(|s| s.value))
                } else {
                    None
                }
            }
        }
    }

    fn children_are_sleeping(&self) -> bool {
        !self.children
            .iter()
            .map(Child::as_ref)
            .any(View::needs_relayout)
    }

    fn focus_prev(&mut self) -> EventResult {
        if let Some(i) = self.children[..self.focus]
            .iter_mut()
            .rev()
            .map(Child::as_mut)
            .position(View::take_focus) {

            self.focus = i;
            EventResult::Consumed(None)
        } else {
            EventResult::Ignored
        }
    }

    fn focus_next(&mut self) -> EventResult {
        if let Some(i) = self.children[(self.focus + 1)..]
            .iter_mut()
            .rev()
            .map(Child::as_mut)
            .position(View::take_focus) {
            self.focus = i;
            EventResult::Consumed(None)
        } else {
            EventResult::Ignored
        }
    }
}

impl View for LinearLayout {
    fn draw(&mut self, printer: &Printer) {
        // Use pre-computed sizes
        let mut offset = Vec2::zero();
        for (i, child) in self.children.iter_mut().enumerate() {
            let printer =
                &printer.sub_printer(offset, child.size, i == self.focus);
            child.view.draw(printer);

            // On the axis given by the orientation,
            // add the child size to the offset.
            *self.orientation.get_ref(&mut offset) += self.orientation
                .get(&child.size);
        }
    }

    fn needs_relayout(&self) -> bool {
        if self.cache.is_none() {
            return true;
        }

        !self.children_are_sleeping()
    }

    fn layout(&mut self, size: Vec2) {
        // If we can get away without breaking a sweat, you can bet we will.
        if self.get_cache(size).is_none() {
            self.get_min_size(size);
        }

        for child in &mut self.children {
            child.view.layout(child.size);
        }
    }

    fn get_min_size(&mut self, req: Vec2) -> Vec2 {
        // Did anything change since last time?
        if let Some(size) = self.get_cache(req) {
            return size;
        }

        // First, make a naive scenario: everything will work fine.
        let sizes: Vec<Vec2> = self.children
            .iter_mut()
            .map(|c| c.get_min_size(req))
            .collect();
        // println_stderr!("Ideal sizes: {:?}", sizes);
        let ideal = self.orientation.stack(sizes.iter());
        // println_stderr!("Ideal result: {:?}", ideal);

        // Does it fit?
        if ideal.fits_in(req) {
            // Champagne!
            self.cache = Some(SizeCache::build(ideal, req));
            return ideal;
        }

        // Ok, so maybe it didn't.
        // Budget cuts, everyone.
        let budget_req = req.with(self.orientation, 1);
        // println_stderr!("Budget req: {:?}", budget_req);

        // See how they like it that way
        let min_sizes: Vec<Vec2> = self.children
            .iter_mut()
            .map(|c| c.get_min_size(budget_req))
            .collect();
        let desperate = self.orientation.stack(min_sizes.iter());
        // println_stderr!("Min sizes: {:?}", min_sizes);
        // println_stderr!("Desperate: {:?}", desperate);

        // This is the lowest we'll ever go. It better fit at least.
        if !desperate.fits_in(req) {
            // Just give up...
            // TODO: print some error message or something
            // println_stderr!("Seriously? {:?} > {:?}???", desperate, req);
            self.cache = Some(SizeCache::build(desperate, req));
            return desperate;
        }

        // This here is how much we're generously offered
        let mut available = self.orientation.get(&(req - desperate));
        // println_stderr!("Available: {:?}", available);

        // Here, we have to make a compromise between the ideal
        // and the desperate solutions.
        let mut overweight: Vec<(usize, usize)> = sizes.iter()
            .map(|v| self.orientation.get(v))
            .zip(min_sizes.iter().map(|v| self.orientation.get(v)))
            .map(|(a, b)| a - b)
            .enumerate()
            .collect();
        // println_stderr!("Overweight: {:?}", overweight);

        // So... distribute `available` to reduce the overweight...
        // TODO: use child weight in the distribution...

        // We'll give everyone his share of what we have left,
        // starting with those who ask the least.
        overweight.sort_by_key(|&(_, weight)| weight);
        let mut allocations = vec![0; overweight.len()];

        for (i, &(j, weight)) in overweight.iter().enumerate() {
            // This is the number of people we still have to feed.
            let remaining = overweight.len() - i;
            // How much we can spare on each one
            let budget = available / remaining;
            // Maybe he doesn't even need that much?
            let spent = min(budget, weight);
            allocations[j] = spent;
            available -= spent;
        }
        // println_stderr!("Allocations: {:?}", allocations);

        // Final lengths are the minimum ones + generous allocations
        let final_lengths: Vec<Vec2> = min_sizes.iter()
            .map(|v| self.orientation.get(v))
            .zip(allocations.iter())
            .map(|(a, b)| a + b)
            .map(|l| req.with(self.orientation, l))
            .collect();
        // println_stderr!("Final sizes: {:?}", final_lengths);

        // Let's ask everyone one last time. Everyone should be happy.
        // (But they may ask more on the other axis.)
        let final_sizes: Vec<Vec2> = self.children
            .iter_mut()
            .enumerate()
            .map(|(i, c)| c.get_min_size(final_lengths[i]))
            .collect();
        // println_stderr!("Final sizes2: {:?}", final_sizes);

        // Let's stack everything to see what it looks like.
        let compromise = self.orientation.stack(final_sizes.iter());

        // Phew, that was a lot of work! I'm not doing it again.
        self.cache = Some(SizeCache::build(compromise, req));

        compromise
    }

    fn take_focus(&mut self) -> bool {
        if let Some(i) = self.children
            .iter_mut()
            .map(Child::as_mut)
            .position(View::take_focus) {
            self.focus = i;
            true
        } else {
            false
        }
    }

    fn on_event(&mut self, event: Event) -> EventResult {
        match self.children[self.focus].view.on_event(event) {
            EventResult::Ignored => {
                match event {
                    Event::Shift(Key::Tab) if self.focus > 0 => {
                        self.focus_prev()
                    }
                    Event::Key(Key::Tab) if self.focus + 1 <
                                            self.children.len() => {
                        self.focus_next()
                    }
                    Event::Key(Key::Left) if self.orientation ==
                                             Orientation::Horizontal &&
                                             self.focus > 0 => {
                        self.focus_prev()
                    }
                    Event::Key(Key::Up) if self.orientation ==
                                           Orientation::Vertical &&
                                           self.focus > 0 => self.focus_prev(),
                    Event::Key(Key::Right) if self.orientation ==
                                              Orientation::Horizontal &&
                                              self.focus + 1 <
                                              self.children.len() => {
                        self.focus_next()
                    }
                    Event::Key(Key::Down) if self.orientation ==
                                             Orientation::Vertical &&
                                             self.focus + 1 <
                                             self.children.len() => {
                        self.focus_next()
                    }
                    _ => EventResult::Ignored,
                }
            }
            res => res,
        }
    }
}
<commit_msg>Fix linear layout focus switch<commit_after>use XY;
use view::View;
use view::SizeCache;
use vec::Vec2;
use printer::Printer;
use orientation::Orientation;
use event::{Event, EventResult, Key};

use std::cmp::min;

/// Arranges its children linearly according to its orientation.
pub struct LinearLayout {
    children: Vec<Child>,
    orientation: Orientation,
    focus: usize,

    cache: Option<XY<SizeCache>>,
}

struct Child {
    view: Box<View>,
    size: Vec2,
    weight: usize,
}

impl Child {
    fn get_min_size(&mut self, req: Vec2) -> Vec2 {
        self.size = self.view.get_min_size(req);
        self.size
    }

    fn as_ref(&self) -> &View {
        &*self.view
    }

    fn as_mut(&mut self) -> &mut View {
        &mut *self.view
    }
}

impl LinearLayout {
    /// Creates a new layout with the given orientation.
    pub fn new(orientation: Orientation) -> Self {
        LinearLayout {
            children: Vec::new(),
            orientation: orientation,
            focus: 0,
            cache: None,
        }
    }

    /// Modifies the weight of the last child added.
    ///
    /// It is an error to call this before adding a child (and it will panic).
    pub fn weight(mut self, weight: usize) -> Self {
        self.children.last_mut().unwrap().weight = weight;

        self
    }

    /// Adds a child to the layout.
    pub fn child<V: View + 'static>(mut self, view: V) -> Self {
        self.children.push(Child {
            view: Box::new(view),
            size: Vec2::zero(),
            weight: 0,
        });
        self.invalidate();

        self
    }

    // Invalidate the view, to request a layout next time
    fn invalidate(&mut self) {
        self.cache = None;
    }

    /// Creates a new vertical layout.
    pub fn vertical() -> Self {
        LinearLayout::new(Orientation::Vertical)
    }

    /// Creates a new horizontal layout.
    pub fn horizontal() -> Self {
        LinearLayout::new(Orientation::Horizontal)
    }

    // If the cache can be used, return the cached size.
    // Otherwise, return None.
    fn get_cache(&self, req: Vec2) -> Option<Vec2> {
        match self.cache {
            None => None,
            Some(ref cache) => {
                // Is our cache even valid?
                // Also, is any child invalidating the layout?
                if cache.zip_map(req, SizeCache::accept).both() &&
                   self.children_are_sleeping() {
                    Some(cache.map(|s| s.value))
                } else {
                    None
                }
            }
        }
    }

    fn children_are_sleeping(&self) -> bool {
        !self.children
            .iter()
            .map(Child::as_ref)
            .any(View::needs_relayout)
    }

    fn focus_prev(&mut self) -> EventResult {
        if let Some(i) = self.children[..self.focus]
            .iter_mut()
            .rev()
            .map(Child::as_mut)
            .position(View::take_focus) {

            // We're looking at the list in reverse
            self.focus -= i+1;
            EventResult::Consumed(None)
        } else {
            EventResult::Ignored
        }
    }

    fn focus_next(&mut self) -> EventResult {
        if let Some(i) = self.children[(self.focus + 1)..]
            .iter_mut()
            .rev()
            .map(Child::as_mut)
            .position(View::take_focus) {
            // Our slice doesn't start at 0
            self.focus += i + 1;
            EventResult::Consumed(None)
        } else {
            EventResult::Ignored
        }
    }
}

impl View for LinearLayout {
    fn draw(&mut self, printer: &Printer) {
        // Use pre-computed sizes
        let mut offset = Vec2::zero();
        for (i, child) in self.children.iter_mut().enumerate() {
            let printer =
                &printer.sub_printer(offset, child.size, i == self.focus);
            child.view.draw(printer);

            // On the axis given by the orientation,
            // add the child size to the offset.
            *self.orientation.get_ref(&mut offset) += self.orientation
                .get(&child.size);
        }
    }

    fn needs_relayout(&self) -> bool {
        if self.cache.is_none() {
            return true;
        }

        !self.children_are_sleeping()
    }

    fn layout(&mut self, size: Vec2) {
        // If we can get away without breaking a sweat, you can bet we will.
        if self.get_cache(size).is_none() {
            self.get_min_size(size);
        }

        for child in &mut self.children {
            child.view.layout(child.size);
        }
    }

    fn get_min_size(&mut self, req: Vec2) -> Vec2 {
        // Did anything change since last time?
        if let Some(size) = self.get_cache(req) {
            return size;
        }

        // First, make a naive scenario: everything will work fine.
        let sizes: Vec<Vec2> = self.children
            .iter_mut()
            .map(|c| c.get_min_size(req))
            .collect();
        // println_stderr!("Ideal sizes: {:?}", sizes);
        let ideal = self.orientation.stack(sizes.iter());
        // println_stderr!("Ideal result: {:?}", ideal);

        // Does it fit?
        if ideal.fits_in(req) {
            // Champagne!
            self.cache = Some(SizeCache::build(ideal, req));
            return ideal;
        }

        // Ok, so maybe it didn't.
        // Budget cuts, everyone.
        let budget_req = req.with(self.orientation, 1);
        // println_stderr!("Budget req: {:?}", budget_req);

        // See how they like it that way
        let min_sizes: Vec<Vec2> = self.children
            .iter_mut()
            .map(|c| c.get_min_size(budget_req))
            .collect();
        let desperate = self.orientation.stack(min_sizes.iter());
        // println_stderr!("Min sizes: {:?}", min_sizes);
        // println_stderr!("Desperate: {:?}", desperate);

        // This is the lowest we'll ever go. It better fit at least.
        if !desperate.fits_in(req) {
            // Just give up...
            // TODO: print some error message or something
            // println_stderr!("Seriously? {:?} > {:?}???", desperate, req);
            self.cache = Some(SizeCache::build(desperate, req));
            return desperate;
        }

        // This here is how much we're generously offered
        let mut available = self.orientation.get(&(req - desperate));
        // println_stderr!("Available: {:?}", available);

        // Here, we have to make a compromise between the ideal
        // and the desperate solutions.
        let mut overweight: Vec<(usize, usize)> = sizes.iter()
            .map(|v| self.orientation.get(v))
            .zip(min_sizes.iter().map(|v| self.orientation.get(v)))
            .map(|(a, b)| a - b)
            .enumerate()
            .collect();
        // println_stderr!("Overweight: {:?}", overweight);

        // So... distribute `available` to reduce the overweight...
        // TODO: use child weight in the distribution...

        // We'll give everyone his share of what we have left,
        // starting with those who ask the least.
        overweight.sort_by_key(|&(_, weight)| weight);
        let mut allocations = vec![0; overweight.len()];

        for (i, &(j, weight)) in overweight.iter().enumerate() {
            // This is the number of people we still have to feed.
            let remaining = overweight.len() - i;
            // How much we can spare on each one
            let budget = available / remaining;
            // Maybe he doesn't even need that much?
            let spent = min(budget, weight);
            allocations[j] = spent;
            available -= spent;
        }
        // println_stderr!("Allocations: {:?}", allocations);

        // Final lengths are the minimum ones + generous allocations
        let final_lengths: Vec<Vec2> = min_sizes.iter()
            .map(|v| self.orientation.get(v))
            .zip(allocations.iter())
            .map(|(a, b)| a + b)
            .map(|l| req.with(self.orientation, l))
            .collect();
        // println_stderr!("Final sizes: {:?}", final_lengths);

        // Let's ask everyone one last time. Everyone should be happy.
        // (But they may ask more on the other axis.)
        let final_sizes: Vec<Vec2> = self.children
            .iter_mut()
            .enumerate()
            .map(|(i, c)| c.get_min_size(final_lengths[i]))
            .collect();
        // println_stderr!("Final sizes2: {:?}", final_sizes);

        // Let's stack everything to see what it looks like.
        let compromise = self.orientation.stack(final_sizes.iter());

        // Phew, that was a lot of work! I'm not doing it again.
        self.cache = Some(SizeCache::build(compromise, req));

        compromise
    }

    fn take_focus(&mut self) -> bool {
        if let Some(i) = self.children
            .iter_mut()
            .map(Child::as_mut)
            .position(View::take_focus) {
            self.focus = i;
            true
        } else {
            false
        }
    }

    fn on_event(&mut self, event: Event) -> EventResult {
        match self.children[self.focus].view.on_event(event) {
            EventResult::Ignored => {
                match event {
                    Event::Shift(Key::Tab) if self.focus > 0 => {
                        self.focus_prev()
                    }
                    Event::Key(Key::Tab) if self.focus + 1 <
                                            self.children.len() => {
                        self.focus_next()
                    }
                    Event::Key(Key::Left) if self.orientation ==
                                             Orientation::Horizontal &&
                                             self.focus > 0 => {
                        self.focus_prev()
                    }
                    Event::Key(Key::Up) if self.orientation ==
                                           Orientation::Vertical &&
                                           self.focus > 0 => self.focus_prev(),
                    Event::Key(Key::Right) if self.orientation ==
                                              Orientation::Horizontal &&
                                              self.focus + 1 <
                                              self.children.len() => {
                        self.focus_next()
                    }
                    Event::Key(Key::Down) if self.orientation ==
                                             Orientation::Vertical &&
                                             self.focus + 1 <
                                             self.children.len() => {
                        self.focus_next()
                    }
                    _ => EventResult::Ignored,
                }
            }
            res => res,
        }
    }
}
<|endoftext|>


<commit_before>//! Name resolution algorithm. The end result of the algorithm is `ItemMap`: a
//! map with maps each module to it's scope: the set of items, visible in the
//! module. That is, we only resolve imports here, name resolution of item
//! bodies will be done in a separate step.
//!
//! Like Rustc, we use an interative per-crate algorithm: we start with scopes
//! containing only directly defined items, and then iteratively resolve
//! imports.
//!
//! To make this work nicely in the IDE scenarios, we place `InputModuleItems`
//! in between raw syntax and name resolution. `InputModuleItems` are computed
//! using only the module's syntax, and it is all directly defined items plus
//! imports. The plain is to make `InputModuleItems` independent of local
//! modifications (that is, typing inside a function shold not change IMIs),
//! such that the results of name resolution can be preserved unless the module
//! structure itself is modified.
use std::{
    sync::Arc,
};

use rustc_hash::FxHashMap;
use ra_syntax::{
    TextRange,
    SmolStr, SyntaxKind::{self, *},
    ast::{self, AstNode}
};
use ra_db::SourceRootId;

use crate::{
    Cancelable, FileId,
    DefId, DefLoc, DefKind,
    SourceItemId, SourceFileItemId, SourceFileItems,
    Path, PathKind,
    HirDatabase,
    module::{ModuleId, ModuleTree},
};

/// Item map is the result of the name resolution. Item map contains, for each
/// module, the set of visible items.
#[derive(Default, Debug, PartialEq, Eq)]
pub struct ItemMap {
    pub per_module: FxHashMap<ModuleId, ModuleScope>,
}

#[derive(Debug, Default, PartialEq, Eq, Clone)]
pub struct ModuleScope {
    items: FxHashMap<SmolStr, Resolution>,
}

impl ModuleScope {
    pub fn entries<'a>(&'a self) -> impl Iterator<Item = (&'a SmolStr, &Resolution)> + 'a {
        self.items.iter()
    }
    pub fn get(&self, name: &SmolStr) -> Option<&Resolution> {
        self.items.get(name)
    }
}

/// A set of items and imports declared inside a module, without relation to
/// other modules.
///
/// This stands in-between raw syntax and name resolution and alow us to avoid
/// recomputing name res: if `InputModuleItems` are the same, we can avoid
/// running name resolution.
#[derive(Debug, Default, PartialEq, Eq)]
pub struct InputModuleItems {
    items: Vec<ModuleItem>,
    imports: Vec<Import>,
}

#[derive(Debug, PartialEq, Eq)]
struct ModuleItem {
    id: SourceFileItemId,
    name: SmolStr,
    kind: SyntaxKind,
    vis: Vis,
}

#[derive(Debug, PartialEq, Eq)]
enum Vis {
    // Priv,
    Other,
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct Import {
    path: Path,
    kind: ImportKind,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct NamedImport {
    pub file_item_id: SourceFileItemId,
    pub relative_range: TextRange,
}

impl NamedImport {
    pub fn range(&self, db: &impl HirDatabase, file_id: FileId) -> TextRange {
        let source_item_id = SourceItemId {
            file_id,
            item_id: self.file_item_id,
        };
        let syntax = db.file_item(source_item_id);
        let offset = syntax.borrowed().range().start();
        self.relative_range + offset
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum ImportKind {
    Glob,
    Named(NamedImport),
}

/// Resolution is basically `DefId` atm, but it should account for stuff like
/// multiple namespaces, ambiguity and errors.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Resolution {
    /// None for unresolved
    pub def_id: Option<DefId>,
    /// ident by whitch this is imported into local scope.
    pub import: Option<NamedImport>,
}

// #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
// enum Namespace {
//     Types,
//     Values,
// }

// #[derive(Debug)]
// struct PerNs<T> {
//     types: Option<T>,
//     values: Option<T>,
// }

impl InputModuleItems {
    pub(crate) fn new<'a>(
        file_items: &SourceFileItems,
        items: impl Iterator<Item = ast::ModuleItem<'a>>,
    ) -> InputModuleItems {
        let mut res = InputModuleItems::default();
        for item in items {
            res.add_item(file_items, item);
        }
        res
    }

    fn add_item(&mut self, file_items: &SourceFileItems, item: ast::ModuleItem) -> Option<()> {
        match item {
            ast::ModuleItem::StructDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::EnumDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::FnDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::TraitDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::TypeDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::ImplItem(_) => {
                // impls don't define items
            }
            ast::ModuleItem::UseItem(it) => self.add_use_item(file_items, it),
            ast::ModuleItem::ExternCrateItem(_) => {
                // TODO
            }
            ast::ModuleItem::ConstDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::StaticDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::Module(it) => self.items.push(ModuleItem::new(file_items, it)?),
        }
        Some(())
    }

    fn add_use_item(&mut self, file_items: &SourceFileItems, item: ast::UseItem) {
        let file_item_id = file_items.id_of_unchecked(item.syntax());
        let start_offset = item.syntax().range().start();
        Path::expand_use_item(item, |path, range| {
            let kind = match range {
                None => ImportKind::Glob,
                Some(range) => ImportKind::Named(NamedImport {
                    file_item_id,
                    relative_range: range - start_offset,
                }),
            };
            self.imports.push(Import { kind, path })
        })
    }
}

impl ModuleItem {
    fn new<'a>(file_items: &SourceFileItems, item: impl ast::NameOwner<'a>) -> Option<ModuleItem> {
        let name = item.name()?.text();
        let kind = item.syntax().kind();
        let vis = Vis::Other;
        let id = file_items.id_of_unchecked(item.syntax());
        let res = ModuleItem {
            id,
            name,
            kind,
            vis,
        };
        Some(res)
    }
}

pub(crate) struct Resolver<'a, DB> {
    db: &'a DB,
    input: &'a FxHashMap<ModuleId, Arc<InputModuleItems>>,
    source_root: SourceRootId,
    module_tree: Arc<ModuleTree>,
    result: ItemMap,
}

impl<'a, DB> Resolver<'a, DB>
where
    DB: HirDatabase,
{
    pub(crate) fn new(
        db: &'a DB,
        input: &'a FxHashMap<ModuleId, Arc<InputModuleItems>>,
        source_root: SourceRootId,
        module_tree: Arc<ModuleTree>,
    ) -> Resolver<'a, DB> {
        Resolver {
            db: db,
            input: &input,
            source_root,
            module_tree,
            result: ItemMap::default(),
        }
    }

    pub(crate) fn resolve(mut self) -> Cancelable<ItemMap> {
        for (&module_id, items) in self.input.iter() {
            self.populate_module(module_id, items)
        }

        for &module_id in self.input.keys() {
            self.db.check_canceled()?;
            self.resolve_imports(module_id);
        }
        Ok(self.result)
    }

    fn populate_module(&mut self, module_id: ModuleId, input: &InputModuleItems) {
        let file_id = module_id.source(&self.module_tree).file_id();

        let mut module_items = ModuleScope::default();

        for import in input.imports.iter() {
            if let Some(name) = import.path.segments.iter().last() {
                if let ImportKind::Named(import) = import.kind {
                    module_items.items.insert(
                        name.clone(),
                        Resolution {
                            def_id: None,
                            import: Some(import),
                        },
                    );
                }
            }
        }

        for item in input.items.iter() {
            if item.kind == MODULE {
                // handle submodules separatelly
                continue;
            }
            let def_loc = DefLoc {
                kind: DefKind::Item,
                source_root_id: self.source_root,
                module_id,
                source_item_id: SourceItemId {
                    file_id,
                    item_id: item.id,
                },
            };
            let def_id = def_loc.id(self.db);
            let resolution = Resolution {
                def_id: Some(def_id),
                import: None,
            };
            module_items.items.insert(item.name.clone(), resolution);
        }

        for (name, module_id) in module_id.children(&self.module_tree) {
            let def_loc = DefLoc {
                kind: DefKind::Module,
                source_root_id: self.source_root,
                module_id,
                source_item_id: module_id.source(&self.module_tree).0,
            };
            let def_id = def_loc.id(self.db);
            let resolution = Resolution {
                def_id: Some(def_id),
                import: None,
            };
            module_items.items.insert(name, resolution);
        }

        self.result.per_module.insert(module_id, module_items);
    }

    fn resolve_imports(&mut self, module_id: ModuleId) {
        for import in self.input[&module_id].imports.iter() {
            self.resolve_import(module_id, import);
        }
    }

    fn resolve_import(&mut self, module_id: ModuleId, import: &Import) {
        let ptr = match import.kind {
            ImportKind::Glob => return,
            ImportKind::Named(ptr) => ptr,
        };

        let mut curr = match import.path.kind {
            // TODO: handle extern crates
            PathKind::Plain => return,
            PathKind::Self_ => module_id,
            PathKind::Super => {
                match module_id.parent(&self.module_tree) {
                    Some(it) => it,
                    // TODO: error
                    None => return,
                }
            }
            PathKind::Crate => module_id.crate_root(&self.module_tree),
        };

        for (i, name) in import.path.segments.iter().enumerate() {
            let is_last = i == import.path.segments.len() - 1;

            let def_id = match self.result.per_module[&curr].items.get(name) {
                None => return,
                Some(res) => match res.def_id {
                    Some(it) => it,
                    None => return,
                },
            };

            if !is_last {
                curr = match def_id.loc(self.db) {
                    DefLoc {
                        kind: DefKind::Module,
                        module_id,
                        ..
                    } => module_id,
                    _ => return,
                }
            } else {
                self.update(module_id, |items| {
                    let res = Resolution {
                        def_id: Some(def_id),
                        import: Some(ptr),
                    };
                    items.items.insert(name.clone(), res);
                })
            }
        }
    }

    fn update(&mut self, module_id: ModuleId, f: impl FnOnce(&mut ModuleScope)) {
        let module_items = self.result.per_module.get_mut(&module_id).unwrap();
        f(module_items)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use salsa::Database;
    use ra_db::FilesDatabase;
    use ra_syntax::SmolStr;

    use crate::{
        self as hir,
        db::HirDatabase,
        mock::MockDatabase,
};

    fn item_map(fixture: &str) -> (Arc<hir::ItemMap>, hir::ModuleId) {
        let (db, pos) = MockDatabase::with_position(fixture);
        let source_root = db.file_source_root(pos.file_id);
        let module = hir::source_binder::module_from_position(&db, pos)
            .unwrap()
            .unwrap();
        let module_id = module.module_id;
        (db.item_map(source_root).unwrap(), module_id)
    }

    #[test]
    fn test_item_map() {
        let (item_map, module_id) = item_map(
            "
            //- /lib.rs
            mod foo;

            use crate::foo::bar::Baz;
            <|>

            //- /foo/mod.rs
            pub mod bar;

            //- /foo/bar.rs
            pub struct Baz;
        ",
        );
        let name = SmolStr::from("Baz");
        let resolution = &item_map.per_module[&module_id].items[&name];
        assert!(resolution.def_id.is_some());
    }

    #[test]
    fn typing_inside_a_function_should_not_invalidate_item_map() {
        let (mut db, pos) = MockDatabase::with_position(
            "
            //- /lib.rs
            mod foo;<|>

            use crate::foo::bar::Baz;

            fn foo() -> i32 {
                1 + 1
            }
            //- /foo/mod.rs
            pub mod bar;

            //- /foo/bar.rs
            pub struct Baz;
        ",
        );
        let source_root = db.file_source_root(pos.file_id);
        {
            let events = db.log_executed(|| {
                db.item_map(source_root).unwrap();
            });
            assert!(format!("{:?}", events).contains("item_map"))
        }

        let new_text = "
            mod foo;

            use crate::foo::bar::Baz;

            fn foo() -> i32 { 92 }
        "
        .to_string();

        db.query_mut(ra_db::FileTextQuery)
            .set(pos.file_id, Arc::new(new_text));

        {
            let events = db.log_executed(|| {
                db.item_map(source_root).unwrap();
            });
            assert!(
                !format!("{:?}", events).contains("_item_map"),
                "{:#?}",
                events
            )
        }
    }
}
<commit_msg>WIP: resolve across crates<commit_after>//! Name resolution algorithm. The end result of the algorithm is `ItemMap`: a
//! map with maps each module to it's scope: the set of items, visible in the
//! module. That is, we only resolve imports here, name resolution of item
//! bodies will be done in a separate step.
//!
//! Like Rustc, we use an interative per-crate algorithm: we start with scopes
//! containing only directly defined items, and then iteratively resolve
//! imports.
//!
//! To make this work nicely in the IDE scenarios, we place `InputModuleItems`
//! in between raw syntax and name resolution. `InputModuleItems` are computed
//! using only the module's syntax, and it is all directly defined items plus
//! imports. The plain is to make `InputModuleItems` independent of local
//! modifications (that is, typing inside a function shold not change IMIs),
//! such that the results of name resolution can be preserved unless the module
//! structure itself is modified.
use std::{
    sync::Arc,
};

use rustc_hash::FxHashMap;
use ra_syntax::{
    TextRange,
    SmolStr, SyntaxKind::{self, *},
    ast::{self, AstNode}
};
use ra_db::SourceRootId;

use crate::{
    Cancelable, FileId,
    DefId, DefLoc, DefKind,
    SourceItemId, SourceFileItemId, SourceFileItems,
    Path, PathKind,
    HirDatabase, Crate,
    module::{ModuleId, ModuleTree},
};

/// Item map is the result of the name resolution. Item map contains, for each
/// module, the set of visible items.
#[derive(Default, Debug, PartialEq, Eq)]
pub struct ItemMap {
    pub per_module: FxHashMap<ModuleId, ModuleScope>,
}

#[derive(Debug, Default, PartialEq, Eq, Clone)]
pub struct ModuleScope {
    items: FxHashMap<SmolStr, Resolution>,
}

impl ModuleScope {
    pub fn entries<'a>(&'a self) -> impl Iterator<Item = (&'a SmolStr, &Resolution)> + 'a {
        self.items.iter()
    }
    pub fn get(&self, name: &SmolStr) -> Option<&Resolution> {
        self.items.get(name)
    }
}

/// A set of items and imports declared inside a module, without relation to
/// other modules.
///
/// This stands in-between raw syntax and name resolution and alow us to avoid
/// recomputing name res: if `InputModuleItems` are the same, we can avoid
/// running name resolution.
#[derive(Debug, Default, PartialEq, Eq)]
pub struct InputModuleItems {
    items: Vec<ModuleItem>,
    imports: Vec<Import>,
}

#[derive(Debug, PartialEq, Eq)]
struct ModuleItem {
    id: SourceFileItemId,
    name: SmolStr,
    kind: SyntaxKind,
    vis: Vis,
}

#[derive(Debug, PartialEq, Eq)]
enum Vis {
    // Priv,
    Other,
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct Import {
    path: Path,
    kind: ImportKind,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct NamedImport {
    pub file_item_id: SourceFileItemId,
    pub relative_range: TextRange,
}

impl NamedImport {
    pub fn range(&self, db: &impl HirDatabase, file_id: FileId) -> TextRange {
        let source_item_id = SourceItemId {
            file_id,
            item_id: self.file_item_id,
        };
        let syntax = db.file_item(source_item_id);
        let offset = syntax.borrowed().range().start();
        self.relative_range + offset
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum ImportKind {
    Glob,
    Named(NamedImport),
}

/// Resolution is basically `DefId` atm, but it should account for stuff like
/// multiple namespaces, ambiguity and errors.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Resolution {
    /// None for unresolved
    pub def_id: Option<DefId>,
    /// ident by whitch this is imported into local scope.
    pub import: Option<NamedImport>,
}

// #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
// enum Namespace {
//     Types,
//     Values,
// }

// #[derive(Debug)]
// struct PerNs<T> {
//     types: Option<T>,
//     values: Option<T>,
// }

impl InputModuleItems {
    pub(crate) fn new<'a>(
        file_items: &SourceFileItems,
        items: impl Iterator<Item = ast::ModuleItem<'a>>,
    ) -> InputModuleItems {
        let mut res = InputModuleItems::default();
        for item in items {
            res.add_item(file_items, item);
        }
        res
    }

    fn add_item(&mut self, file_items: &SourceFileItems, item: ast::ModuleItem) -> Option<()> {
        match item {
            ast::ModuleItem::StructDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::EnumDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::FnDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::TraitDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::TypeDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::ImplItem(_) => {
                // impls don't define items
            }
            ast::ModuleItem::UseItem(it) => self.add_use_item(file_items, it),
            ast::ModuleItem::ExternCrateItem(_) => {
                // TODO
            }
            ast::ModuleItem::ConstDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::StaticDef(it) => self.items.push(ModuleItem::new(file_items, it)?),
            ast::ModuleItem::Module(it) => self.items.push(ModuleItem::new(file_items, it)?),
        }
        Some(())
    }

    fn add_use_item(&mut self, file_items: &SourceFileItems, item: ast::UseItem) {
        let file_item_id = file_items.id_of_unchecked(item.syntax());
        let start_offset = item.syntax().range().start();
        Path::expand_use_item(item, |path, range| {
            let kind = match range {
                None => ImportKind::Glob,
                Some(range) => ImportKind::Named(NamedImport {
                    file_item_id,
                    relative_range: range - start_offset,
                }),
            };
            self.imports.push(Import { kind, path })
        })
    }
}

impl ModuleItem {
    fn new<'a>(file_items: &SourceFileItems, item: impl ast::NameOwner<'a>) -> Option<ModuleItem> {
        let name = item.name()?.text();
        let kind = item.syntax().kind();
        let vis = Vis::Other;
        let id = file_items.id_of_unchecked(item.syntax());
        let res = ModuleItem {
            id,
            name,
            kind,
            vis,
        };
        Some(res)
    }
}

pub(crate) struct Resolver<'a, DB> {
    db: &'a DB,
    input: &'a FxHashMap<ModuleId, Arc<InputModuleItems>>,
    source_root: SourceRootId,
    module_tree: Arc<ModuleTree>,
    result: ItemMap,
}

impl<'a, DB> Resolver<'a, DB>
where
    DB: HirDatabase,
{
    pub(crate) fn new(
        db: &'a DB,
        input: &'a FxHashMap<ModuleId, Arc<InputModuleItems>>,
        source_root: SourceRootId,
        module_tree: Arc<ModuleTree>,
    ) -> Resolver<'a, DB> {
        Resolver {
            db,
            input,
            source_root,
            module_tree,
            result: ItemMap::default(),
        }
    }

    pub(crate) fn resolve(mut self) -> Cancelable<ItemMap> {
        for (&module_id, items) in self.input.iter() {
            self.populate_module(module_id, items)
        }

        for &module_id in self.input.keys() {
            self.db.check_canceled()?;
            self.resolve_imports(module_id)?;
        }
        Ok(self.result)
    }

    fn populate_module(&mut self, module_id: ModuleId, input: &InputModuleItems) {
        let file_id = module_id.source(&self.module_tree).file_id();

        let mut module_items = ModuleScope::default();

        for import in input.imports.iter() {
            if let Some(name) = import.path.segments.iter().last() {
                if let ImportKind::Named(import) = import.kind {
                    module_items.items.insert(
                        name.clone(),
                        Resolution {
                            def_id: None,
                            import: Some(import),
                        },
                    );
                }
            }
        }

        for item in input.items.iter() {
            if item.kind == MODULE {
                // handle submodules separatelly
                continue;
            }
            let def_loc = DefLoc {
                kind: DefKind::Item,
                source_root_id: self.source_root,
                module_id,
                source_item_id: SourceItemId {
                    file_id,
                    item_id: item.id,
                },
            };
            let def_id = def_loc.id(self.db);
            let resolution = Resolution {
                def_id: Some(def_id),
                import: None,
            };
            module_items.items.insert(item.name.clone(), resolution);
        }

        for (name, module_id) in module_id.children(&self.module_tree) {
            let def_loc = DefLoc {
                kind: DefKind::Module,
                source_root_id: self.source_root,
                module_id,
                source_item_id: module_id.source(&self.module_tree).0,
            };
            let def_id = def_loc.id(self.db);
            let resolution = Resolution {
                def_id: Some(def_id),
                import: None,
            };
            module_items.items.insert(name, resolution);
        }

        self.result.per_module.insert(module_id, module_items);
    }

    fn resolve_imports(&mut self, module_id: ModuleId) -> Cancelable<()> {
        for import in self.input[&module_id].imports.iter() {
            self.resolve_import(module_id, import)?;
        }
        Ok(())
    }

    fn resolve_import(&mut self, module_id: ModuleId, import: &Import) -> Cancelable<()> {
        let ptr = match import.kind {
            ImportKind::Glob => return Ok(()),
            ImportKind::Named(ptr) => ptr,
        };
        let mut segments = import.path.segments.iter().enumerate();

        let mut curr = match import.path.kind {
            // TODO: handle extern crates
            PathKind::Plain => {
                let root_id = module_id.crate_root(&self.module_tree);
                let file_id = root_id.source(&self.module_tree).file_id();
                let crate_graph = self.db.crate_graph();
                let crate_id = match crate_graph.crate_id_for_crate_root(file_id) {
                    None => return Ok(()),
                    Some(it) => it,
                };
                let krate = Crate::new(crate_id);
                let crate_name = match segments.next() {
                    None => return Ok(()),
                    Some((_, it)) => it,
                };
                match krate
                    .dependencies(self.db)
                    .into_iter()
                    .find(|it| &it.name == crate_name)
                {
                    None => return Ok(()),
                    Some(dep) => match dep.krate.root_module(self.db)? {
                        None => return Ok(()),
                        Some(it) => it.module_id,
                    },
                }
            }
            PathKind::Self_ => module_id,
            PathKind::Super => {
                match module_id.parent(&self.module_tree) {
                    Some(it) => it,
                    // TODO: error
                    None => return Ok(()),
                }
            }
            PathKind::Crate => module_id.crate_root(&self.module_tree),
        };

        for (i, name) in segments {
            let is_last = i == import.path.segments.len() - 1;

            let def_id = match self.result.per_module[&curr].items.get(name) {
                None => return Ok(()),
                Some(res) => match res.def_id {
                    Some(it) => it,
                    None => return Ok(()),
                },
            };

            if !is_last {
                curr = match def_id.loc(self.db) {
                    DefLoc {
                        kind: DefKind::Module,
                        module_id,
                        ..
                    } => module_id,
                    _ => return Ok(()),
                }
            } else {
                self.update(module_id, |items| {
                    let res = Resolution {
                        def_id: Some(def_id),
                        import: Some(ptr),
                    };
                    items.items.insert(name.clone(), res);
                })
            }
        }
        Ok(())
    }

    fn update(&mut self, module_id: ModuleId, f: impl FnOnce(&mut ModuleScope)) {
        let module_items = self.result.per_module.get_mut(&module_id).unwrap();
        f(module_items)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use salsa::Database;
    use ra_db::FilesDatabase;
    use ra_syntax::SmolStr;

    use crate::{
        self as hir,
        db::HirDatabase,
        mock::MockDatabase,
};

    fn item_map(fixture: &str) -> (Arc<hir::ItemMap>, hir::ModuleId) {
        let (db, pos) = MockDatabase::with_position(fixture);
        let source_root = db.file_source_root(pos.file_id);
        let module = hir::source_binder::module_from_position(&db, pos)
            .unwrap()
            .unwrap();
        let module_id = module.module_id;
        (db.item_map(source_root).unwrap(), module_id)
    }

    #[test]
    fn test_item_map() {
        let (item_map, module_id) = item_map(
            "
            //- /lib.rs
            mod foo;

            use crate::foo::bar::Baz;
            <|>

            //- /foo/mod.rs
            pub mod bar;

            //- /foo/bar.rs
            pub struct Baz;
        ",
        );
        let name = SmolStr::from("Baz");
        let resolution = &item_map.per_module[&module_id].items[&name];
        assert!(resolution.def_id.is_some());
    }

    #[test]
    fn typing_inside_a_function_should_not_invalidate_item_map() {
        let (mut db, pos) = MockDatabase::with_position(
            "
            //- /lib.rs
            mod foo;<|>

            use crate::foo::bar::Baz;

            fn foo() -> i32 {
                1 + 1
            }
            //- /foo/mod.rs
            pub mod bar;

            //- /foo/bar.rs
            pub struct Baz;
        ",
        );
        let source_root = db.file_source_root(pos.file_id);
        {
            let events = db.log_executed(|| {
                db.item_map(source_root).unwrap();
            });
            assert!(format!("{:?}", events).contains("item_map"))
        }

        let new_text = "
            mod foo;

            use crate::foo::bar::Baz;

            fn foo() -> i32 { 92 }
        "
        .to_string();

        db.query_mut(ra_db::FileTextQuery)
            .set(pos.file_id, Arc::new(new_text));

        {
            let events = db.log_executed(|| {
                db.item_map(source_root).unwrap();
            });
            assert!(
                !format!("{:?}", events).contains("_item_map"),
                "{:#?}",
                events
            )
        }
    }
}
<|endoftext|>


<commit_before>//! Handle process life-time and message passing for proc-macro client

use crossbeam_channel::{bounded, Receiver, Sender};
use ra_tt::Subtree;

use crate::msg::{ErrorCode, Message, Request, Response, ResponseError};
use crate::rpc::{ExpansionResult, ExpansionTask, ListMacrosResult, ListMacrosTask, ProcMacroKind};

use io::{BufRead, BufReader};
use std::{
    convert::{TryFrom, TryInto},
    ffi::{OsStr, OsString},
    io::{self, Write},
    path::{Path, PathBuf},
    process::{Child, Command, Stdio},
    sync::{Arc, Weak},
};

#[derive(Debug, Default)]
pub(crate) struct ProcMacroProcessSrv {
    inner: Option<Weak<Sender<Task>>>,
}

#[derive(Debug)]
pub(crate) struct ProcMacroProcessThread {
    // XXX: drop order is significant
    sender: Arc<Sender<Task>>,
    handle: jod_thread::JoinHandle<()>,
}

impl ProcMacroProcessSrv {
    pub fn run(
        process_path: PathBuf,
        args: impl IntoIterator<Item = impl AsRef<OsStr>>,
    ) -> io::Result<(ProcMacroProcessThread, ProcMacroProcessSrv)> {
        let process = Process::run(process_path, args)?;

        let (task_tx, task_rx) = bounded(0);
        let handle = jod_thread::spawn(move || {
            client_loop(task_rx, process);
        });

        let task_tx = Arc::new(task_tx);
        let srv = ProcMacroProcessSrv { inner: Some(Arc::downgrade(&task_tx)) };
        let thread = ProcMacroProcessThread { handle, sender: task_tx };

        Ok((thread, srv))
    }

    pub fn find_proc_macros(
        &self,
        dylib_path: &Path,
    ) -> Result<Vec<(String, ProcMacroKind)>, ra_tt::ExpansionError> {
        let task = ListMacrosTask { lib: dylib_path.to_path_buf() };

        let result: ListMacrosResult = self.send_task(Request::ListMacro(task))?;
        Ok(result.macros)
    }

    pub fn custom_derive(
        &self,
        dylib_path: &Path,
        subtree: &Subtree,
        derive_name: &str,
    ) -> Result<Subtree, ra_tt::ExpansionError> {
        let task = ExpansionTask {
            macro_body: subtree.clone(),
            macro_name: derive_name.to_string(),
            attributes: None,
            lib: dylib_path.to_path_buf(),
        };

        let result: ExpansionResult = self.send_task(Request::ExpansionMacro(task))?;
        Ok(result.expansion)
    }

    pub fn send_task<R>(&self, req: Request) -> Result<R, ra_tt::ExpansionError>
    where
        R: TryFrom<Response, Error = &'static str>,
    {
        let sender = match &self.inner {
            None => return Err(ra_tt::ExpansionError::Unknown("No sender is found.".to_string())),
            Some(it) => it,
        };

        let (result_tx, result_rx) = bounded(0);
        let sender = match sender.upgrade() {
            None => {
                return Err(ra_tt::ExpansionError::Unknown("Proc macro process is closed.".into()))
            }
            Some(it) => it,
        };
        sender.send(Task { req: req.into(), result_tx }).unwrap();
        let res = result_rx
            .recv()
            .map_err(|_| ra_tt::ExpansionError::Unknown("Proc macro thread is closed.".into()))?;

        match res {
            Some(Response::Error(err)) => {
                return Err(ra_tt::ExpansionError::ExpansionError(err.message));
            }
            Some(res) => Ok(res.try_into().map_err(|err| {
                ra_tt::ExpansionError::Unknown(format!(
                    "Fail to get response, reason : {:#?} ",
                    err
                ))
            })?),
            None => Err(ra_tt::ExpansionError::Unknown("Empty result".into())),
        }
    }
}

fn client_loop(task_rx: Receiver<Task>, mut process: Process) {
    let (mut stdin, mut stdout) = match process.stdio() {
        None => return,
        Some(it) => it,
    };

    for task in task_rx {
        let Task { req, result_tx } = task;

        match send_request(&mut stdin, &mut stdout, req) {
            Ok(res) => result_tx.send(res).unwrap(),
            Err(_err) => {
                let res = Response::Error(ResponseError {
                    code: ErrorCode::ServerErrorEnd,
                    message: "Server closed".into(),
                });
                result_tx.send(res.into()).unwrap();
                // Restart the process
                if process.restart().is_err() {
                    break;
                }
                let stdio = match process.stdio() {
                    None => break,
                    Some(it) => it,
                };
                stdin = stdio.0;
                stdout = stdio.1;
            }
        }
    }
}

struct Task {
    req: Request,
    result_tx: Sender<Option<Response>>,
}

struct Process {
    path: PathBuf,
    args: Vec<OsString>,
    child: Child,
}

impl Drop for Process {
    fn drop(&mut self) {
        let _ = self.child.kill();
    }
}

impl Process {
    fn run(
        path: PathBuf,
        args: impl IntoIterator<Item = impl AsRef<OsStr>>,
    ) -> io::Result<Process> {
        let args = args.into_iter().map(|s| s.as_ref().into()).collect();
        let child = mk_child(&path, &args)?;
        Ok(Process { path, args, child })
    }

    fn restart(&mut self) -> io::Result<()> {
        let _ = self.child.kill();
        self.child = mk_child(&self.path, &self.args)?;
        Ok(())
    }

    fn stdio(&mut self) -> Option<(impl Write, impl BufRead)> {
        let stdin = self.child.stdin.take()?;
        let stdout = self.child.stdout.take()?;
        let read = BufReader::new(stdout);

        Some((stdin, read))
    }
}

fn mk_child(path: &Path, args: impl IntoIterator<Item = impl AsRef<OsStr>>) -> io::Result<Child> {
    Command::new(&path)
        .args(args)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .spawn()
}

fn send_request(
    mut writer: &mut impl Write,
    mut reader: &mut impl BufRead,
    req: Request,
) -> io::Result<Option<Response>> {
    req.write(&mut writer)?;
    Ok(Response::read(&mut reader)?)
}
<commit_msg>proc_macro: add ability to log to stderr and view output in vscode<commit_after>//! Handle process life-time and message passing for proc-macro client

use crossbeam_channel::{bounded, Receiver, Sender};
use ra_tt::Subtree;

use crate::msg::{ErrorCode, Message, Request, Response, ResponseError};
use crate::rpc::{ExpansionResult, ExpansionTask, ListMacrosResult, ListMacrosTask, ProcMacroKind};

use io::{BufRead, BufReader};
use std::{
    convert::{TryFrom, TryInto},
    ffi::{OsStr, OsString},
    io::{self, Write},
    path::{Path, PathBuf},
    process::{Child, Command, Stdio},
    sync::{Arc, Weak},
};

#[derive(Debug, Default)]
pub(crate) struct ProcMacroProcessSrv {
    inner: Option<Weak<Sender<Task>>>,
}

#[derive(Debug)]
pub(crate) struct ProcMacroProcessThread {
    // XXX: drop order is significant
    sender: Arc<Sender<Task>>,
    handle: jod_thread::JoinHandle<()>,
}

impl ProcMacroProcessSrv {
    pub fn run(
        process_path: PathBuf,
        args: impl IntoIterator<Item = impl AsRef<OsStr>>,
    ) -> io::Result<(ProcMacroProcessThread, ProcMacroProcessSrv)> {
        let process = Process::run(process_path, args)?;

        let (task_tx, task_rx) = bounded(0);
        let handle = jod_thread::spawn(move || {
            client_loop(task_rx, process);
        });

        let task_tx = Arc::new(task_tx);
        let srv = ProcMacroProcessSrv { inner: Some(Arc::downgrade(&task_tx)) };
        let thread = ProcMacroProcessThread { handle, sender: task_tx };

        Ok((thread, srv))
    }

    pub fn find_proc_macros(
        &self,
        dylib_path: &Path,
    ) -> Result<Vec<(String, ProcMacroKind)>, ra_tt::ExpansionError> {
        let task = ListMacrosTask { lib: dylib_path.to_path_buf() };

        let result: ListMacrosResult = self.send_task(Request::ListMacro(task))?;
        Ok(result.macros)
    }

    pub fn custom_derive(
        &self,
        dylib_path: &Path,
        subtree: &Subtree,
        derive_name: &str,
    ) -> Result<Subtree, ra_tt::ExpansionError> {
        let task = ExpansionTask {
            macro_body: subtree.clone(),
            macro_name: derive_name.to_string(),
            attributes: None,
            lib: dylib_path.to_path_buf(),
        };

        let result: ExpansionResult = self.send_task(Request::ExpansionMacro(task))?;
        Ok(result.expansion)
    }

    pub fn send_task<R>(&self, req: Request) -> Result<R, ra_tt::ExpansionError>
    where
        R: TryFrom<Response, Error = &'static str>,
    {
        let sender = match &self.inner {
            None => return Err(ra_tt::ExpansionError::Unknown("No sender is found.".to_string())),
            Some(it) => it,
        };

        let (result_tx, result_rx) = bounded(0);
        let sender = match sender.upgrade() {
            None => {
                return Err(ra_tt::ExpansionError::Unknown("Proc macro process is closed.".into()))
            }
            Some(it) => it,
        };
        sender.send(Task { req: req.into(), result_tx }).unwrap();
        let res = result_rx
            .recv()
            .map_err(|_| ra_tt::ExpansionError::Unknown("Proc macro thread is closed.".into()))?;

        match res {
            Some(Response::Error(err)) => {
                return Err(ra_tt::ExpansionError::ExpansionError(err.message));
            }
            Some(res) => Ok(res.try_into().map_err(|err| {
                ra_tt::ExpansionError::Unknown(format!(
                    "Fail to get response, reason : {:#?} ",
                    err
                ))
            })?),
            None => Err(ra_tt::ExpansionError::Unknown("Empty result".into())),
        }
    }
}

fn client_loop(task_rx: Receiver<Task>, mut process: Process) {
    let (mut stdin, mut stdout) = match process.stdio() {
        None => return,
        Some(it) => it,
    };

    for task in task_rx {
        let Task { req, result_tx } = task;

        match send_request(&mut stdin, &mut stdout, req) {
            Ok(res) => result_tx.send(res).unwrap(),
            Err(_err) => {
                let res = Response::Error(ResponseError {
                    code: ErrorCode::ServerErrorEnd,
                    message: "Server closed".into(),
                });
                result_tx.send(res.into()).unwrap();
                // Restart the process
                if process.restart().is_err() {
                    break;
                }
                let stdio = match process.stdio() {
                    None => break,
                    Some(it) => it,
                };
                stdin = stdio.0;
                stdout = stdio.1;
            }
        }
    }
}

struct Task {
    req: Request,
    result_tx: Sender<Option<Response>>,
}

struct Process {
    path: PathBuf,
    args: Vec<OsString>,
    child: Child,
}

impl Drop for Process {
    fn drop(&mut self) {
        let _ = self.child.kill();
    }
}

impl Process {
    fn run(
        path: PathBuf,
        args: impl IntoIterator<Item = impl AsRef<OsStr>>,
    ) -> io::Result<Process> {
        let args = args.into_iter().map(|s| s.as_ref().into()).collect();
        let child = mk_child(&path, &args)?;
        Ok(Process { path, args, child })
    }

    fn restart(&mut self) -> io::Result<()> {
        let _ = self.child.kill();
        self.child = mk_child(&self.path, &self.args)?;
        Ok(())
    }

    fn stdio(&mut self) -> Option<(impl Write, impl BufRead)> {
        let stdin = self.child.stdin.take()?;
        let stdout = self.child.stdout.take()?;
        let read = BufReader::new(stdout);

        Some((stdin, read))
    }
}

fn mk_child(path: &Path, args: impl IntoIterator<Item = impl AsRef<OsStr>>) -> io::Result<Child> {
    Command::new(&path)
        .args(args)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::inherit())
        .spawn()
}

fn send_request(
    mut writer: &mut impl Write,
    mut reader: &mut impl BufRead,
    req: Request,
) -> io::Result<Option<Response>> {
    req.write(&mut writer)?;
    Ok(Response::read(&mut reader)?)
}
<|endoftext|>


<commit_before>use ruma_identifiers::{EventId, RoomVersionId, ServerNameBox};
use thiserror::Error;

/// `ruma-signature`'s error type, wraps a number of other error types.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum Error {
    /// [`JsonError`] wrapper.
    #[error("JSON error: {0}")]
    Json(#[from] JsonError),

    /// [`VerificationError`] wrapper.
    #[error("Verification error: {0}")]
    Verification(#[from] VerificationError),

    /// [`ParseError`] wrapper.
    #[error("Parse error: {0}")]
    Parse(#[from] ParseError),

    /// Wrapper for [`pkcs8::der::Error`].
    #[error("DER Parse error: {0}")]
    DerParse(pkcs8::der::Error),

    /// [`SplitError`] wrapper.
    #[error("Split error: {0}")]
    SplitError(#[from] SplitError),
}

/// All errors related to JSON validation/parsing.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum JsonError {
    /// Signals that `target` is not of type `of_type` ([`JsonType`]).
    #[error("Value in {target:?} must be a JSON {of_type:?}")]
    NotOfType {
        /// An arbitrary "target" that doesn't have the required type.
        target: String,
        /// The JSON type the target was expected to be.
        of_type: JsonType,
    },

    /// Like [`JsonError::NotOfType`], only called when the `target` is a multiple;
    /// array, set, etc.
    #[error("Values in {target:?} must be JSON {of_type:?}s")]
    NotMultiplesOfType {
        /// An arbitrary "target" where
        /// each or one of it's elements doesn't have the required type.
        target: String,
        /// The JSON type the element was expected to be.
        of_type: JsonType,
    },

    /// Signals that a specific field is missing from a JSON object.
    #[error("JSON object must contain the field {0:?}")]
    JsonFieldMissingFromObject(String),

    /// Signals a key missing from a JSON object.
    ///
    /// Note that this is different from [`JsonError::JsonFieldMissingFromObject`],
    /// this error talks about an expected identifying key (`"ed25519:abcd"`)
    /// missing from a target, where the key has a specific "type"/name.
    #[error("JSON object {for_target:?} does not have {type_of} key {with_key:?}")]
    JsonKeyMissing {
        /// The target from which the key is missing.
        for_target: String,
        /// The kind of thing the key indicates.
        type_of: String,
        /// The key that is missing.
        with_key: String,
    },

    /// A derivative error from [`ruma_serde::CanonicalJsonError`],
    /// captured here.
    #[error("Canonical JSON error: {0}")]
    CanonicalJson(#[from] ruma_serde::CanonicalJsonError),

    /// A more generic JSON error from [`serde_json`].
    #[error(transparent)]
    Serde(#[from] serde_json::Error),
}

// TODO: make macro for this
impl JsonError {
    pub(crate) fn not_of_type<T: Into<String>>(target: T, of_type: JsonType) -> Error {
        Self::NotOfType { target: target.into(), of_type }.into()
    }

    pub(crate) fn not_multiples_of_type<T: Into<String>>(target: T, of_type: JsonType) -> Error {
        Self::NotMultiplesOfType { target: target.into(), of_type }.into()
    }

    pub(crate) fn field_missing_from_object<T: Into<String>>(target: T) -> Error {
        Self::JsonFieldMissingFromObject(target.into()).into()
    }

    pub(crate) fn key_missing<T1: Into<String>, T2: Into<String>, T3: Into<String>>(
        for_target: T1,
        type_of: T2,
        with_key: T3,
    ) -> Error {
        Self::JsonKeyMissing {
            for_target: for_target.into(),
            type_of: type_of.into(),
            with_key: with_key.into(),
        }
        .into()
    }
}

/// A JSON type enum for [`JsonError`] variants.
#[derive(Debug)]
pub enum JsonType {
    /// A JSON Object.
    Object,
    /// A JSON String.
    String,
    /// A JSON Integer.
    Integer,
    /// A JSON Array.
    Array,
    /// A JSON Boolean.
    Boolean,
    /// JSON Null.
    Null,
}

/// Errors relating to verification of events and signatures.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum VerificationError {
    /// For when a signature cannot be found for a `target`.
    #[error("Could not find signatures for {0:?}")]
    SignatureNotFound(ServerNameBox),

    /// For when a public key cannot be found for a `target`.
    #[error("Could not find public key for {0:?}")]
    PublicKeyNotFound(ServerNameBox),

    /// For when no public key matches the signature given.
    #[error("Not signed with any of the given public keys")]
    UnknownPublicKeysForSignature,

    /// For when [`ed25519_dalek`] cannot verify a signature.
    #[error("Could not verify signature: {0}")]
    Signature(#[source] ed25519_dalek::SignatureError),
}

impl VerificationError {
    pub(crate) fn signature_not_found<T: Into<ServerNameBox>>(target: T) -> Error {
        Self::SignatureNotFound(target.into()).into()
    }

    pub(crate) fn public_key_not_found<T: Into<ServerNameBox>>(target: T) -> Error {
        Self::PublicKeyNotFound(target.into()).into()
    }
}

/// Errors relating to parsing of all sorts.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum ParseError {
    /// For user ID parsing errors.
    #[error("Could not parse User ID: {0}")]
    UserId(#[source] ruma_identifiers::Error),
    /// For event ID parsing errors.
    #[error("Could not parse Event ID: {0}")]
    EventId(#[source] ruma_identifiers::Error),

    /// For when an event ID, coupled with a specific room version, doesn't have a server name
    /// embedded.
    #[error("Event Id {0:?} should have a server name for the given room version {1:?}")]
    ServerNameFromEventIdByRoomVersion(EventId, RoomVersionId),

    /// For when the extracted/"parsed" public key from a PKCS#8 v2 document doesn't match the
    /// public key derived from it's private key.
    #[error("PKCS#8 Document public key does not match public key derived from private key; derived: {0:X?} (len {}), parsed: {1:X?} (len {})", .derived_key.len(), .parsed_key.len())]
    DerivedPublicKeyDoesNotMatchParsedKey {
        /// The parsed key.
        parsed_key: Vec<u8>,
        /// The derived key.
        derived_key: Vec<u8>,
    },

    /// For when the ASN.1 Object Identifier on a PKCS#8 document doesn't match the expected one.
    ///
    /// e.g. the document describes a RSA key, while an ed25519 key was expected.
    #[error("Algorithm OID does not match ed25519, expected {expected}, found {found}")]
    Oid {
        /// The expected OID.
        expected: pkcs8::ObjectIdentifier,
        /// The OID that was found instead.
        found: pkcs8::ObjectIdentifier,
    },

    /// For when [`ed25519_dalek`] cannot parse a secret/private key.
    #[error("Could not parse secret key: {0}")]
    SecretKey(#[source] ed25519_dalek::SignatureError),

    /// For when [`ed25519_dalek`] cannot parse a public key.
    #[error("Could not parse public key: {0}")]
    PublicKey(#[source] ed25519_dalek::SignatureError),

    /// For when [`ed25519_dalek`] cannot parse a signature.
    #[error("Could not parse signature: {0}")]
    Signature(#[source] ed25519_dalek::SignatureError),

    /// For when parsing base64 gives an error.
    #[error("Could not parse {of_type} base64 string {string:?}: {source}")]
    Base64 {
        /// The "type"/name of the base64 string
        of_type: String,
        /// The string itself.
        string: String,
        /// The originating error.
        #[source]
        source: base64::DecodeError,
    },
}

impl ParseError {
    pub(crate) fn from_event_id_by_room_version(
        event_id: &EventId,
        room_version: &RoomVersionId,
    ) -> Error {
        Self::ServerNameFromEventIdByRoomVersion(event_id.clone(), room_version.clone()).into()
    }

    pub(crate) fn derived_vs_parsed_mismatch<P: Into<Vec<u8>>, D: Into<Vec<u8>>>(
        parsed: P,
        derived: D,
    ) -> Error {
        Self::DerivedPublicKeyDoesNotMatchParsedKey {
            parsed_key: parsed.into(),
            derived_key: derived.into(),
        }
        .into()
    }

    pub(crate) fn base64<T1: Into<String>, T2: Into<String>>(
        of_type: T1,
        string: T2,
        source: base64::DecodeError,
    ) -> Error {
        Self::Base64 { of_type: of_type.into(), string: string.into(), source }.into()
    }
}

/// An error when trying to extract the algorithm and version from a key identifier.
#[derive(Error, Debug)]
pub enum SplitError {
    /// The signature's ID does not have exactly two components separated by a colon.
    #[error("malformed signature ID: expected exactly 2 segment separated by a colon, found {0}")]
    InvalidLength(usize),

    /// The signature's ID contains invalid characters in its version.
    #[error("malformed signature ID: expected version to contain only characters in the character set `[a-zA-Z0-9_]`, found `{0}`")]
    InvalidVersion(String),

    /// The signature uses an unknown algorithm.
    #[error("unknown algorithm: {0}")]
    UnknownAlgorithm(String),
}
<commit_msg>signatures: Add missing newlines<commit_after>use ruma_identifiers::{EventId, RoomVersionId, ServerNameBox};
use thiserror::Error;

/// `ruma-signature`'s error type, wraps a number of other error types.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum Error {
    /// [`JsonError`] wrapper.
    #[error("JSON error: {0}")]
    Json(#[from] JsonError),

    /// [`VerificationError`] wrapper.
    #[error("Verification error: {0}")]
    Verification(#[from] VerificationError),

    /// [`ParseError`] wrapper.
    #[error("Parse error: {0}")]
    Parse(#[from] ParseError),

    /// Wrapper for [`pkcs8::der::Error`].
    #[error("DER Parse error: {0}")]
    DerParse(pkcs8::der::Error),

    /// [`SplitError`] wrapper.
    #[error("Split error: {0}")]
    SplitError(#[from] SplitError),
}

/// All errors related to JSON validation/parsing.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum JsonError {
    /// Signals that `target` is not of type `of_type` ([`JsonType`]).
    #[error("Value in {target:?} must be a JSON {of_type:?}")]
    NotOfType {
        /// An arbitrary "target" that doesn't have the required type.
        target: String,
        /// The JSON type the target was expected to be.
        of_type: JsonType,
    },

    /// Like [`JsonError::NotOfType`], only called when the `target` is a multiple;
    /// array, set, etc.
    #[error("Values in {target:?} must be JSON {of_type:?}s")]
    NotMultiplesOfType {
        /// An arbitrary "target" where
        /// each or one of it's elements doesn't have the required type.
        target: String,
        /// The JSON type the element was expected to be.
        of_type: JsonType,
    },

    /// Signals that a specific field is missing from a JSON object.
    #[error("JSON object must contain the field {0:?}")]
    JsonFieldMissingFromObject(String),

    /// Signals a key missing from a JSON object.
    ///
    /// Note that this is different from [`JsonError::JsonFieldMissingFromObject`],
    /// this error talks about an expected identifying key (`"ed25519:abcd"`)
    /// missing from a target, where the key has a specific "type"/name.
    #[error("JSON object {for_target:?} does not have {type_of} key {with_key:?}")]
    JsonKeyMissing {
        /// The target from which the key is missing.
        for_target: String,
        /// The kind of thing the key indicates.
        type_of: String,
        /// The key that is missing.
        with_key: String,
    },

    /// A derivative error from [`ruma_serde::CanonicalJsonError`],
    /// captured here.
    #[error("Canonical JSON error: {0}")]
    CanonicalJson(#[from] ruma_serde::CanonicalJsonError),

    /// A more generic JSON error from [`serde_json`].
    #[error(transparent)]
    Serde(#[from] serde_json::Error),
}

// TODO: make macro for this
impl JsonError {
    pub(crate) fn not_of_type<T: Into<String>>(target: T, of_type: JsonType) -> Error {
        Self::NotOfType { target: target.into(), of_type }.into()
    }

    pub(crate) fn not_multiples_of_type<T: Into<String>>(target: T, of_type: JsonType) -> Error {
        Self::NotMultiplesOfType { target: target.into(), of_type }.into()
    }

    pub(crate) fn field_missing_from_object<T: Into<String>>(target: T) -> Error {
        Self::JsonFieldMissingFromObject(target.into()).into()
    }

    pub(crate) fn key_missing<T1: Into<String>, T2: Into<String>, T3: Into<String>>(
        for_target: T1,
        type_of: T2,
        with_key: T3,
    ) -> Error {
        Self::JsonKeyMissing {
            for_target: for_target.into(),
            type_of: type_of.into(),
            with_key: with_key.into(),
        }
        .into()
    }
}

/// A JSON type enum for [`JsonError`] variants.
#[derive(Debug)]
pub enum JsonType {
    /// A JSON Object.
    Object,

    /// A JSON String.
    String,

    /// A JSON Integer.
    Integer,

    /// A JSON Array.
    Array,

    /// A JSON Boolean.
    Boolean,

    /// JSON Null.
    Null,
}

/// Errors relating to verification of events and signatures.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum VerificationError {
    /// For when a signature cannot be found for a `target`.
    #[error("Could not find signatures for {0:?}")]
    SignatureNotFound(ServerNameBox),

    /// For when a public key cannot be found for a `target`.
    #[error("Could not find public key for {0:?}")]
    PublicKeyNotFound(ServerNameBox),

    /// For when no public key matches the signature given.
    #[error("Not signed with any of the given public keys")]
    UnknownPublicKeysForSignature,

    /// For when [`ed25519_dalek`] cannot verify a signature.
    #[error("Could not verify signature: {0}")]
    Signature(#[source] ed25519_dalek::SignatureError),
}

impl VerificationError {
    pub(crate) fn signature_not_found<T: Into<ServerNameBox>>(target: T) -> Error {
        Self::SignatureNotFound(target.into()).into()
    }

    pub(crate) fn public_key_not_found<T: Into<ServerNameBox>>(target: T) -> Error {
        Self::PublicKeyNotFound(target.into()).into()
    }
}

/// Errors relating to parsing of all sorts.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum ParseError {
    /// For user ID parsing errors.
    #[error("Could not parse User ID: {0}")]
    UserId(#[source] ruma_identifiers::Error),

    /// For event ID parsing errors.
    #[error("Could not parse Event ID: {0}")]
    EventId(#[source] ruma_identifiers::Error),

    /// For when an event ID, coupled with a specific room version, doesn't have a server name
    /// embedded.
    #[error("Event Id {0:?} should have a server name for the given room version {1:?}")]
    ServerNameFromEventIdByRoomVersion(EventId, RoomVersionId),

    /// For when the extracted/"parsed" public key from a PKCS#8 v2 document doesn't match the
    /// public key derived from it's private key.
    #[error("PKCS#8 Document public key does not match public key derived from private key; derived: {0:X?} (len {}), parsed: {1:X?} (len {})", .derived_key.len(), .parsed_key.len())]
    DerivedPublicKeyDoesNotMatchParsedKey {
        /// The parsed key.
        parsed_key: Vec<u8>,
        /// The derived key.
        derived_key: Vec<u8>,
    },

    /// For when the ASN.1 Object Identifier on a PKCS#8 document doesn't match the expected one.
    ///
    /// e.g. the document describes a RSA key, while an ed25519 key was expected.
    #[error("Algorithm OID does not match ed25519, expected {expected}, found {found}")]
    Oid {
        /// The expected OID.
        expected: pkcs8::ObjectIdentifier,
        /// The OID that was found instead.
        found: pkcs8::ObjectIdentifier,
    },

    /// For when [`ed25519_dalek`] cannot parse a secret/private key.
    #[error("Could not parse secret key: {0}")]
    SecretKey(#[source] ed25519_dalek::SignatureError),

    /// For when [`ed25519_dalek`] cannot parse a public key.
    #[error("Could not parse public key: {0}")]
    PublicKey(#[source] ed25519_dalek::SignatureError),

    /// For when [`ed25519_dalek`] cannot parse a signature.
    #[error("Could not parse signature: {0}")]
    Signature(#[source] ed25519_dalek::SignatureError),

    /// For when parsing base64 gives an error.
    #[error("Could not parse {of_type} base64 string {string:?}: {source}")]
    Base64 {
        /// The "type"/name of the base64 string
        of_type: String,
        /// The string itself.
        string: String,
        /// The originating error.
        #[source]
        source: base64::DecodeError,
    },
}

impl ParseError {
    pub(crate) fn from_event_id_by_room_version(
        event_id: &EventId,
        room_version: &RoomVersionId,
    ) -> Error {
        Self::ServerNameFromEventIdByRoomVersion(event_id.clone(), room_version.clone()).into()
    }

    pub(crate) fn derived_vs_parsed_mismatch<P: Into<Vec<u8>>, D: Into<Vec<u8>>>(
        parsed: P,
        derived: D,
    ) -> Error {
        Self::DerivedPublicKeyDoesNotMatchParsedKey {
            parsed_key: parsed.into(),
            derived_key: derived.into(),
        }
        .into()
    }

    pub(crate) fn base64<T1: Into<String>, T2: Into<String>>(
        of_type: T1,
        string: T2,
        source: base64::DecodeError,
    ) -> Error {
        Self::Base64 { of_type: of_type.into(), string: string.into(), source }.into()
    }
}

/// An error when trying to extract the algorithm and version from a key identifier.
#[derive(Error, Debug)]
pub enum SplitError {
    /// The signature's ID does not have exactly two components separated by a colon.
    #[error("malformed signature ID: expected exactly 2 segment separated by a colon, found {0}")]
    InvalidLength(usize),

    /// The signature's ID contains invalid characters in its version.
    #[error("malformed signature ID: expected version to contain only characters in the character set `[a-zA-Z0-9_]`, found `{0}`")]
    InvalidVersion(String),

    /// The signature uses an unknown algorithm.
    #[error("unknown algorithm: {0}")]
    UnknownAlgorithm(String),
}
<|endoftext|>


<commit_before>//! System operations instructions

use util::address::Address;
use util::bigint::{U256, M256};
use util::gas::Gas;
use vm::{Memory, Log, Context, Transaction};
use super::State;

use std::cmp::min;
use tiny_keccak::Keccak;
use vm::eval::util::copy_from_memory;

pub fn suicide<M: Memory + Default>(state: &mut State<M>) {
    pop!(state, address: Address);
    let balance = state.account_state.balance(state.context.address).unwrap();
    state.account_state.remove(state.context.address).unwrap();
    state.removed.push(state.context.address);
    state.account_state.increase_balance(address, balance);
}

pub fn log<M: Memory + Default>(state: &mut State<M>, topic_len: usize) {
    pop!(state, index, len);
    let data = copy_from_memory(&state.memory, index, len);
    let mut topics = Vec::new();
    for _ in 0..topic_len {
        topics.push(state.stack.pop().unwrap());
    }

    state.logs.push(Log {
        address: state.context.address,
        data: data,
        topics: topics,
    });
}

pub fn sha3<M: Memory + Default>(state: &mut State<M>) {
    pop!(state, from, len);
    let data = copy_from_memory(&state.memory, from, len);
    let mut sha3 = Keccak::new_keccak256();
    sha3.update(data.as_slice());
    let mut ret = [0u8; 32];
    sha3.finalize(&mut ret);
    push!(state, M256::from(ret.as_ref()));
}

pub fn create<M: Memory + Default>(state: &mut State<M>, after_gas: Gas) -> Option<Context> {
    pop!(state, value: U256);
    pop!(state, init_start, init_len);
    if state.account_state.balance(state.context.address).unwrap() < value {
        push!(state, M256::zero());
        return None;
    }

    let init = copy_from_memory(&state.memory, init_start, init_len);
    let transaction = Transaction::ContractCreation {
        caller: state.context.address,
        gas_price: state.context.gas_price,
        gas_limit: after_gas,
        value: value,
        init: init,
    };
    let context = transaction.into_context(
        Gas::zero(), Some(state.context.origin), &mut state.account_state, true
    ).unwrap();

    push!(state, context.address.into());
    Some(context)
}

pub fn call<M: Memory + Default>(state: &mut State<M>, stipend_gas: Gas, after_gas: Gas, as_self: bool) -> Option<(Context, (M256, M256))> {
    pop!(state, gas: Gas, to: Address, value: U256);
    pop!(state, in_start, in_len, out_start, out_len);
    if state.account_state.balance(state.context.address).unwrap() < value {
        push!(state, M256::zero());
        return None;
    }

    let input = copy_from_memory(&state.memory, in_start, in_len);
    let gas_limit = min(gas + stipend_gas, after_gas);

    let transaction = Transaction::MessageCall {
        address: to,
        caller: state.context.address,
        gas_price: state.context.gas_price,
        gas_limit: gas_limit,
        value: value,
        data: input,
    };

    let mut context = transaction.into_context(
        Gas::zero(), Some(state.context.origin), &mut state.account_state, true
    ).unwrap();
    if as_self {
        context.address = state.context.address;
    }

    push!(state, M256::from(1u64));
    Some((context, (out_start, out_len)))
}

pub fn delegate_call<M: Memory + Default>(state: &mut State<M>, after_gas: Gas) -> Option<(Context, (M256, M256))> {
    pop!(state, gas: Gas, to: Address);
    pop!(state, in_start, in_len, out_start, out_len);

    let input = copy_from_memory(&state.memory, in_start, in_len);
    let gas_limit = min(gas, after_gas);

    let transaction = Transaction::MessageCall {
        address: to,
        caller: state.context.caller,
        gas_price: state.context.gas_price,
        gas_limit: gas_limit,
        value: state.context.value,
        data: input,
    };

    let mut context = transaction.into_context(
        Gas::zero(), Some(state.context.origin), &mut state.account_state, true
    ).unwrap();
    context.value = U256::zero();
    context.address = state.context.address;

    push!(state, M256::from(1u64));
    Some((context, (out_start, out_len)))
}
<commit_msg>SUICIDE: wrong remove order<commit_after>//! System operations instructions

use util::address::Address;
use util::bigint::{U256, M256};
use util::gas::Gas;
use vm::{Memory, Log, Context, Transaction};
use super::State;

use std::cmp::min;
use tiny_keccak::Keccak;
use vm::eval::util::copy_from_memory;

pub fn suicide<M: Memory + Default>(state: &mut State<M>) {
    pop!(state, address: Address);
    let balance = state.account_state.balance(state.context.address).unwrap();
    state.removed.push(state.context.address);
    state.account_state.increase_balance(address, balance);
    state.account_state.remove(state.context.address).unwrap();
}

pub fn log<M: Memory + Default>(state: &mut State<M>, topic_len: usize) {
    pop!(state, index, len);
    let data = copy_from_memory(&state.memory, index, len);
    let mut topics = Vec::new();
    for _ in 0..topic_len {
        topics.push(state.stack.pop().unwrap());
    }

    state.logs.push(Log {
        address: state.context.address,
        data: data,
        topics: topics,
    });
}

pub fn sha3<M: Memory + Default>(state: &mut State<M>) {
    pop!(state, from, len);
    let data = copy_from_memory(&state.memory, from, len);
    let mut sha3 = Keccak::new_keccak256();
    sha3.update(data.as_slice());
    let mut ret = [0u8; 32];
    sha3.finalize(&mut ret);
    push!(state, M256::from(ret.as_ref()));
}

pub fn create<M: Memory + Default>(state: &mut State<M>, after_gas: Gas) -> Option<Context> {
    pop!(state, value: U256);
    pop!(state, init_start, init_len);
    if state.account_state.balance(state.context.address).unwrap() < value {
        push!(state, M256::zero());
        return None;
    }

    let init = copy_from_memory(&state.memory, init_start, init_len);
    let transaction = Transaction::ContractCreation {
        caller: state.context.address,
        gas_price: state.context.gas_price,
        gas_limit: after_gas,
        value: value,
        init: init,
    };
    let context = transaction.into_context(
        Gas::zero(), Some(state.context.origin), &mut state.account_state, true
    ).unwrap();

    push!(state, context.address.into());
    Some(context)
}

pub fn call<M: Memory + Default>(state: &mut State<M>, stipend_gas: Gas, after_gas: Gas, as_self: bool) -> Option<(Context, (M256, M256))> {
    pop!(state, gas: Gas, to: Address, value: U256);
    pop!(state, in_start, in_len, out_start, out_len);
    if state.account_state.balance(state.context.address).unwrap() < value {
        push!(state, M256::zero());
        return None;
    }

    let input = copy_from_memory(&state.memory, in_start, in_len);
    let gas_limit = min(gas + stipend_gas, after_gas);

    let transaction = Transaction::MessageCall {
        address: to,
        caller: state.context.address,
        gas_price: state.context.gas_price,
        gas_limit: gas_limit,
        value: value,
        data: input,
    };

    let mut context = transaction.into_context(
        Gas::zero(), Some(state.context.origin), &mut state.account_state, true
    ).unwrap();
    if as_self {
        context.address = state.context.address;
    }

    push!(state, M256::from(1u64));
    Some((context, (out_start, out_len)))
}

pub fn delegate_call<M: Memory + Default>(state: &mut State<M>, after_gas: Gas) -> Option<(Context, (M256, M256))> {
    pop!(state, gas: Gas, to: Address);
    pop!(state, in_start, in_len, out_start, out_len);

    let input = copy_from_memory(&state.memory, in_start, in_len);
    let gas_limit = min(gas, after_gas);

    let transaction = Transaction::MessageCall {
        address: to,
        caller: state.context.caller,
        gas_price: state.context.gas_price,
        gas_limit: gas_limit,
        value: state.context.value,
        data: input,
    };

    let mut context = transaction.into_context(
        Gas::zero(), Some(state.context.origin), &mut state.account_state, true
    ).unwrap();
    context.value = U256::zero();
    context.address = state.context.address;

    push!(state, M256::from(1u64));
    Some((context, (out_start, out_len)))
}
<|endoftext|>


<commit_before>use image::ImageResult;
use mli::{Forward, Graph, Train};
use mli_conv::Conv2;
use mli_ndarray::Map2One;
use mli_relu::Blu;
use ndarray::{array, s, Array, Array2};
use ndarray_image::{open_gray_image, save_gray_image};
use rand_core::SeedableRng;
use rand_distr::{Distribution, Normal};
use rand_pcg::Pcg64;
use std::path::{Path, PathBuf};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "image", about = "Loads and convolves an image as a test")]
struct Opt {
    /// Number of epochs
    #[structopt(short = "t", default_value = "2000")]
    epochs: usize,
    /// Number of epochs per output image
    #[structopt(short = "s", default_value = "10")]
    show_every: usize,
    /// Initial learning rate
    #[structopt(short = "i", default_value = "0.000000000001")]
    initial_learning_rate: f32,
    /// Learning rate multiplier per epoch
    #[structopt(short = "m", default_value = "1.003")]
    learning_rate_multiplier: f32,
    /// File to load
    #[structopt(parse(from_os_str))]
    file: PathBuf,
    /// Output directory
    #[structopt(parse(from_os_str))]
    output_dir: PathBuf,
}

fn sobel(image: &Array2<f32>) -> Array2<f32> {
    let down_filter = array![[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0],];
    let right_filter = down_filter.t().to_owned();
    let down = Conv2::new(down_filter).forward(&image).1;
    let right = Conv2::new(right_filter).forward(&image).1;
    (down.map(|f| f.powi(2)) + right.map(|f| f.powi(2))).map(|f| f.sqrt())
}

fn open_image(path: impl AsRef<Path>) -> ImageResult<Array2<f32>> {
    let image = open_gray_image(path)?;
    Ok(image.map(|&n| f32::from(n) / 255.0))
}

fn save_image(path: impl AsRef<Path>, image: &Array2<f32>) -> ImageResult<()> {
    let image = image.map(|&n| num::clamp(n * 255.0, 0.0, 255.0) as u8);
    save_gray_image(path, image.view())
}

fn main() -> ImageResult<()> {
    let opt = Opt::from_args();
    let mut prng = Pcg64::from_seed([0; 32]);
    let image = open_image(opt.file)?;
    let sobel_image = sobel(&image);
    let conv_layers = 4;
    let filter_radius = 1usize;
    let filter_len = (filter_radius * 2 + 1).pow(2);
    let padding = (conv_layers * filter_radius - 1) as i32;
    #[allow(clippy::deref_addrof)]
    let expected = sobel_image.slice(s![padding..-padding, padding..-padding]);
    save_image(opt.output_dir.join("actual.png"), &expected.to_owned())?;
    let mut random_filter = |mean: f32, variance: f32| -> Array2<f32> {
        // Xavier initialize by changing the variance to be 1/N where N is the number of neurons.
        Array::from_iter(
            Normal::new(mean / filter_len as f32, variance / filter_len as f32)
                .unwrap()
                .sample_iter(&mut prng)
                .take(filter_len),
        )
        .into_shape((filter_radius * 2 + 1, filter_radius * 2 + 1))
        .unwrap()
    };
    let mut train_filter = Conv2::new(random_filter(1.0, 9.0f32.powi(-1)))
        .chain(Map2One(Blu::new(0.1, 0.2)))
        .chain(Conv2::new(random_filter(0.0, 9.0f32.powi(0))))
        .chain(Map2One(Blu::new(-0.3, 0.7)))
        .chain(Conv2::new(random_filter(0.0, 9.0f32.powi(1))))
        .chain(Map2One(Blu::new(0.4, -0.2)))
        .chain(Conv2::new(random_filter(8.0, 9.0f32.powi(2))));
    let mut learn_rate = opt.initial_learning_rate;
    for i in 0..opt.epochs {
        let (internal, output) = train_filter.forward(&image.view());
        if i % opt.show_every == 0 {
            save_image(opt.output_dir.join(format!("epoch{:04}.png", i)), &output)?;
        }
        let loss = (output.clone() - expected.view()).map(|n| n.powi(2)).sum();
        eprintln!("epoch {:04} loss: {}, learn_rate: {}", i, loss, learn_rate);
        // The loss function is (n - t)^2, so 2*(n - t) is dE/df where
        // `E` is loss and `f` is output.
        let delta_loss = (output - expected).map(|n| 2.0 * n);
        let output_delta = -learn_rate * delta_loss;
        train_filter.propogate(&image.view(), &internal, &output_delta);
        learn_rate *= opt.learning_rate_multiplier;
    }
    Ok(())
}
<commit_msg>change deep train to have a seed<commit_after>use image::ImageResult;
use mli::{Forward, Graph, Train};
use mli_conv::Conv2;
use mli_ndarray::Map2One;
use mli_relu::Blu;
use ndarray::{array, s, Array, Array2, OwnedRepr};
use ndarray_image::{open_gray_image, save_gray_image};
use rand_core::SeedableRng;
use rand_distr::{Distribution, Normal};
use rand_pcg::Pcg64;
use std::path::{Path, PathBuf};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "image", about = "Loads and convolves an image as a test")]
struct Opt {
    /// Number of epochs
    #[structopt(short = "t", default_value = "2000")]
    epochs: usize,
    /// Number of epochs per output image
    #[structopt(short = "s", default_value = "1")]
    show_every: usize,
    /// Initial learning rate
    #[structopt(short = "i", default_value = "0.000000001")]
    initial_learning_rate: f32,
    /// Learning rate multiplier per epoch
    #[structopt(short = "m", default_value = "1.003")]
    learning_rate_multiplier: f32,
    /// Seed
    #[structopt(short = "z", default_value = "42")]
    seed: u32,
    /// File to load
    #[structopt(parse(from_os_str))]
    file: PathBuf,
    /// Output directory
    #[structopt(parse(from_os_str))]
    output_dir: PathBuf,
}

fn sobel(image: &Array2<f32>) -> Array2<f32> {
    let down_filter = array![[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0],];
    let right_filter = down_filter.t().to_owned();
    let down = Conv2::new(down_filter).forward(&image).1;
    let right = Conv2::new(right_filter).forward(&image).1;
    (down.map(|f| f.powi(2)) + right.map(|f| f.powi(2))).map(|f| f.sqrt())
}

fn open_image(path: impl AsRef<Path>) -> ImageResult<Array2<f32>> {
    let image = open_gray_image(path)?;
    Ok(image.map(|&n| f32::from(n) / 255.0))
}

fn save_image(path: impl AsRef<Path>, image: &Array2<f32>) -> ImageResult<()> {
    let image = image.map(|&n| num::clamp(n * 255.0, 0.0, 255.0) as u8);
    save_gray_image(path, image.view())
}

fn main() -> ImageResult<()> {
    let opt = Opt::from_args();
    let image = open_image(opt.file)?;
    let sobel_image = sobel(&image);
    let conv_layers = 5;
    let filter_radius = 1usize;
    let filter_len = (filter_radius * 2 + 1).pow(2);
    let padding = (conv_layers * filter_radius - 1) as i32;
    #[allow(clippy::deref_addrof)]
    let expected = sobel_image.slice(s![padding..-padding, padding..-padding]);
    save_image(opt.output_dir.join("actual.png"), &expected.to_owned())?;
    let mut prng = Pcg64::from_seed([
        (opt.seed & 0xFF) as u8,
        ((opt.seed >> 8) & 0xFF) as u8,
        ((opt.seed >> 16) & 0xFF) as u8,
        ((opt.seed >> 24) & 0xFF) as u8,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
    ]);
    let mut random_filter = |mean: f32, variance: f32| -> Conv2<OwnedRepr<f32>> {
        // Xavier initialize by changing the variance to be 1/N where N is the number of neurons.
        Conv2::new(
            Array::from_iter(
                Normal::new(mean / filter_len as f32, variance / filter_len as f32)
                    .unwrap()
                    .sample_iter(&mut prng)
                    .take(filter_len),
            )
            .into_shape((filter_radius * 2 + 1, filter_radius * 2 + 1))
            .unwrap(),
        )
    };
    let mut prng = Pcg64::from_seed([
        (opt.seed & 0xFF) as u8,
        ((opt.seed >> 8) & 0xFF) as u8,
        ((opt.seed >> 16) & 0xFF) as u8,
        ((opt.seed >> 24) & 0xFF) as u8,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
        0,
    ]);
    let mut random_blu = |mean: f32, variance: f32| -> Blu {
        // Xavier initialize by changing the variance to be 1/N where N is the number of neurons.
        let distr = Normal::new(mean, variance).unwrap();
        Blu::new(distr.sample(&mut prng), distr.sample(&mut prng))
    };
    let mut train_filter = random_filter(0.0, 10.0f32.powi(0))
        .chain(Map2One(random_blu(0.0, 0.05)))
        .chain(random_filter(0.0, 10.0f32.powi(0)))
        .chain(Map2One(random_blu(0.0, 0.05)))
        .chain(random_filter(0.0, 10.0f32.powi(0)))
        .chain(Map2One(random_blu(0.0, 0.05)))
        .chain(random_filter(0.0, 10.0f32.powi(0)))
        .chain(Map2One(random_blu(0.0, 0.05)))
        .chain(random_filter(64.0, 10.0f32.powi(3)));
    let mut learn_rate = opt.initial_learning_rate;
    for i in 0..opt.epochs {
        let (internal, output) = train_filter.forward(&image);
        if i % opt.show_every == 0 {
            save_image(opt.output_dir.join(format!("epoch{:04}.png", i)), &output)?;
        }
        let loss = (output.clone() - expected.view()).map(|n| n.powi(2)).sum();
        eprintln!("epoch {:04} loss: {}, learn_rate: {}", i, loss, learn_rate);
        // The loss function is (n - t)^2, so 2*(n - t) is dE/df where
        // `E` is loss and `f` is output.
        let delta_loss = (output - expected).map(|n| 2.0 * n);
        let output_delta = -learn_rate * delta_loss;
        train_filter.propogate(&image, &internal, &output_delta);
        learn_rate *= opt.learning_rate_multiplier;
    }
    Ok(())
}
<|endoftext|>