<commit_before>//! Various utilities for working with Cairo image surfaces.
use std::ops::DerefMut;

use cairo;

pub mod iterators;
pub mod shared_surface;

/// A pixel consisting of R, G, B and A values.
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub struct Pixel {
    pub r: u8,
    pub g: u8,
    pub b: u8,
    pub a: u8,
}

/// Modes which specify how the values of out of bounds pixels are computed.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum EdgeMode {
    /// The nearest inbounds pixel value is returned.
    Duplicate,
    /// The image is extended by taking the color values from the opposite of the image.
    ///
    /// Imagine the image being tiled infinitely, with the original image at the origin.
    Wrap,
    /// Zero RGBA values are returned.
    None,
}

/// Extension methods for `cairo::ImageSurfaceData`.
pub trait ImageSurfaceDataExt: DerefMut<Target = [u8]> {
    /// Sets the pixel at the given coordinates. Assumes the `ARgb32` format.
    #[inline]
    fn set_pixel(&mut self, stride: usize, pixel: Pixel, x: u32, y: u32) {
        let value = pixel.to_u32();
        unsafe {
            *(&mut self[y as usize * stride + x as usize * 4] as *mut u8 as *mut u32) = value;
        }
    }
}

impl Pixel {
    /// Returns an unpremultiplied value of this pixel.
    #[inline]
    pub fn unpremultiply(self) -> Self {
        if self.a == 0 {
            self
        } else {
            let alpha = f64::from(self.a) / 255.0;
            let unpremultiply = |x| ((f64::from(x) / alpha) + 0.5) as u8;

            Self {
                r: unpremultiply(self.r),
                g: unpremultiply(self.g),
                b: unpremultiply(self.b),
                a: self.a,
            }
        }
    }

    /// Returns a premultiplied value of this pixel.
    #[inline]
    pub fn premultiply(self) -> Self {
        let alpha = f64::from(self.a) / 255.0;
        let premultiply = |x| ((f64::from(x) * alpha) + 0.5) as u8;

        Self {
            r: premultiply(self.r),
            g: premultiply(self.g),
            b: premultiply(self.b),
            a: self.a,
        }
    }

    /// Returns the pixel value as a `u32`, in the same format as `cairo::Format::ARgb32`.
    #[inline]
    pub fn to_u32(self) -> u32 {
        (u32::from(self.a) << 24)
            | (u32::from(self.r) << 16)
            | (u32::from(self.g) << 8)
            | u32::from(self.b)
    }

    /// Converts a `u32` in the same format as `cairo::Format::ARgb32` into a `Pixel`.
    #[inline]
    pub fn from_u32(x: u32) -> Self {
        Self {
            r: ((x >> 16) & 0xFF) as u8,
            g: ((x >> 8) & 0xFF) as u8,
            b: (x & 0xFF) as u8,
            a: ((x >> 24) & 0xFF) as u8,
        }
    }

    /// Returns a 'mask' pixel with only the alpha channel
    ///
    /// Assuming, the pixel is linear RGB (not sRGB)
    /// y = luminance
    /// Y = 0.2126 R + 0.7152 G + 0.0722 B
    /// 1.0 opacity = 255
    ///
    /// When Y = 1.0, pixel for mask should be 0xFFFFFFFF
    /// (you get 1.0 luminance from 255 from R, G and B)
    ///
    /// r_mult = 0xFFFFFFFF / (255.0 * 255.0) * .2126 = 14042.45  ~= 14042
    /// g_mult = 0xFFFFFFFF / (255.0 * 255.0) * .7152 = 47239.69  ~= 47240
    /// b_mult = 0xFFFFFFFF / (255.0 * 255.0) * .0722 =  4768.88  ~= 4769
    ///
    /// This allows for the following expected behaviour:
    ///    (we only care about the most sig byte)
    /// if pixel = 0x00FFFFFF, pixel' = 0xFF......
    /// if pixel = 0x00020202, pixel' = 0x02......
    /// if pixel = 0x00000000, pixel' = 0x00......
    pub fn to_mask(self, opacity: u8) -> Self {
        let r = self.r as u32;
        let g = self.g as u32;
        let b = self.b as u32;
        let o = opacity as u32;

        Self {
            r: 0,
            g: 0,
            b: 0,
            a: (((r * 14042 + g * 47240 + b * 4769) * o) >> 24) as u8,
        }
    }

    #[inline]
    pub fn diff(self, p: &Pixel) -> Pixel {
        let a_r = self.r as i32;
        let a_g = self.g as i32;
        let a_b = self.b as i32;
        let a_a = self.a as i32;

        let b_r = p.r as i32;
        let b_g = p.g as i32;
        let b_b = p.b as i32;
        let b_a = p.a as i32;

        let r = (a_r - b_r).abs() as u8;
        let g = (a_g - b_g).abs() as u8;
        let b = (a_b - b_b).abs() as u8;
        let a = (a_a - b_a).abs() as u8;

        Pixel { r, g, b, a }
    }
}

impl<'a> ImageSurfaceDataExt for cairo::ImageSurfaceData<'a> {}
impl<'a> ImageSurfaceDataExt for &'a mut [u8] {}
<commit_msg>surface_utils: prefer explicit conversion<commit_after>//! Various utilities for working with Cairo image surfaces.
use std::ops::DerefMut;

use cairo;

pub mod iterators;
pub mod shared_surface;

/// A pixel consisting of R, G, B and A values.
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub struct Pixel {
    pub r: u8,
    pub g: u8,
    pub b: u8,
    pub a: u8,
}

/// Modes which specify how the values of out of bounds pixels are computed.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum EdgeMode {
    /// The nearest inbounds pixel value is returned.
    Duplicate,
    /// The image is extended by taking the color values from the opposite of the image.
    ///
    /// Imagine the image being tiled infinitely, with the original image at the origin.
    Wrap,
    /// Zero RGBA values are returned.
    None,
}

/// Extension methods for `cairo::ImageSurfaceData`.
pub trait ImageSurfaceDataExt: DerefMut<Target = [u8]> {
    /// Sets the pixel at the given coordinates. Assumes the `ARgb32` format.
    #[inline]
    fn set_pixel(&mut self, stride: usize, pixel: Pixel, x: u32, y: u32) {
        let value = pixel.to_u32();
        unsafe {
            *(&mut self[y as usize * stride + x as usize * 4] as *mut u8 as *mut u32) = value;
        }
    }
}

impl Pixel {
    /// Returns an unpremultiplied value of this pixel.
    #[inline]
    pub fn unpremultiply(self) -> Self {
        if self.a == 0 {
            self
        } else {
            let alpha = f64::from(self.a) / 255.0;
            let unpremultiply = |x| ((f64::from(x) / alpha) + 0.5) as u8;

            Self {
                r: unpremultiply(self.r),
                g: unpremultiply(self.g),
                b: unpremultiply(self.b),
                a: self.a,
            }
        }
    }

    /// Returns a premultiplied value of this pixel.
    #[inline]
    pub fn premultiply(self) -> Self {
        let alpha = f64::from(self.a) / 255.0;
        let premultiply = |x| ((f64::from(x) * alpha) + 0.5) as u8;

        Self {
            r: premultiply(self.r),
            g: premultiply(self.g),
            b: premultiply(self.b),
            a: self.a,
        }
    }

    /// Returns the pixel value as a `u32`, in the same format as `cairo::Format::ARgb32`.
    #[inline]
    pub fn to_u32(self) -> u32 {
        (u32::from(self.a) << 24)
            | (u32::from(self.r) << 16)
            | (u32::from(self.g) << 8)
            | u32::from(self.b)
    }

    /// Converts a `u32` in the same format as `cairo::Format::ARgb32` into a `Pixel`.
    #[inline]
    pub fn from_u32(x: u32) -> Self {
        Self {
            r: ((x >> 16) & 0xFF) as u8,
            g: ((x >> 8) & 0xFF) as u8,
            b: (x & 0xFF) as u8,
            a: ((x >> 24) & 0xFF) as u8,
        }
    }

    /// Returns a 'mask' pixel with only the alpha channel
    ///
    /// Assuming, the pixel is linear RGB (not sRGB)
    /// y = luminance
    /// Y = 0.2126 R + 0.7152 G + 0.0722 B
    /// 1.0 opacity = 255
    ///
    /// When Y = 1.0, pixel for mask should be 0xFFFFFFFF
    /// (you get 1.0 luminance from 255 from R, G and B)
    ///
    /// r_mult = 0xFFFFFFFF / (255.0 * 255.0) * .2126 = 14042.45  ~= 14042
    /// g_mult = 0xFFFFFFFF / (255.0 * 255.0) * .7152 = 47239.69  ~= 47240
    /// b_mult = 0xFFFFFFFF / (255.0 * 255.0) * .0722 =  4768.88  ~= 4769
    ///
    /// This allows for the following expected behaviour:
    ///    (we only care about the most sig byte)
    /// if pixel = 0x00FFFFFF, pixel' = 0xFF......
    /// if pixel = 0x00020202, pixel' = 0x02......
    /// if pixel = 0x00000000, pixel' = 0x00......
    pub fn to_mask(self, opacity: u8) -> Self {
        let r = u32::from(self.r);
        let g = u32::from(self.g);
        let b = u32::from(self.b);
        let o = u32::from(opacity);

        Self {
            r: 0,
            g: 0,
            b: 0,
            a: (((r * 14042 + g * 47240 + b * 4769) * o) >> 24) as u8,
        }
    }

    #[inline]
    pub fn diff(self, pixel: &Pixel) -> Pixel {
        let a_r = i32::from(self.r);
        let a_g = i32::from(self.g);
        let a_b = i32::from(self.b);
        let a_a = i32::from(self.a);

        let b_r = i32::from(pixel.r);
        let b_g = i32::from(pixel.g);
        let b_b = i32::from(pixel.b);
        let b_a = i32::from(pixel.a);

        let r = (a_r - b_r).abs() as u8;
        let g = (a_g - b_g).abs() as u8;
        let b = (a_b - b_b).abs() as u8;
        let a = (a_a - b_a).abs() as u8;

        Pixel { r, g, b, a }
    }
}

impl<'a> ImageSurfaceDataExt for cairo::ImageSurfaceData<'a> {}
impl<'a> ImageSurfaceDataExt for &'a mut [u8] {}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

extern crate string_cache_codegen;

use std::env;
use std::fs::File;
use std::io::{BufReader, BufRead};
use std::path::Path;

fn main() {
    let static_atoms = Path::new(&env::var_os("CARGO_MANIFEST_DIR").unwrap()).join("static_atoms.txt");
    let static_atoms = BufReader::new(File::open(&static_atoms).unwrap());
    let mut atom_type = string_cache_codegen::AtomType::new("Atom", "atom!");

    macro_rules! predefined {
        ($($name: expr,)+) => {
            {
                $(
                    atom_type.atom($name);
                )+
            }
        }
    }
    include!("../style/counter_style/predefined.rs");

    atom_type
        .atoms(static_atoms.lines().map(Result::unwrap))
        .write_to_file(&Path::new(&env::var_os("OUT_DIR").unwrap()).join("atom.rs"))
        .unwrap();
}
<commit_msg>Auto merge of #21440 - kingdido999:rustfmt/atoms, r=jdm<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

extern crate string_cache_codegen;

use std::env;
use std::fs::File;
use std::io::{BufReader, BufRead};
use std::path::Path;

fn main() {
    let static_atoms =
        Path::new(&env::var_os("CARGO_MANIFEST_DIR").unwrap()).join("static_atoms.txt");
    let static_atoms = BufReader::new(File::open(&static_atoms).unwrap());
    let mut atom_type = string_cache_codegen::AtomType::new("Atom", "atom!");

    macro_rules! predefined {
        ($($name: expr,)+) => {
            {
                $(
                    atom_type.atom($name);
                )+
            }
        }
    }
    include!("../style/counter_style/predefined.rs");

    atom_type
        .atoms(static_atoms.lines().map(Result::unwrap))
        .write_to_file(&Path::new(&env::var_os("OUT_DIR").unwrap()).join("atom.rs"))
        .unwrap();
}
<|endoftext|>


<commit_before>use crate::Resolver;
use rustc_ast::token::{self, Token};
use rustc_ast::visit::{self, FnKind};
use rustc_ast::walk_list;
use rustc_ast::*;
use rustc_ast_lowering::ResolverAstLowering;
use rustc_expand::expand::AstFragment;
use rustc_hir::def_id::LocalDefId;
use rustc_hir::definitions::*;
use rustc_span::hygiene::ExpnId;
use rustc_span::symbol::{kw, sym};
use rustc_span::Span;
use tracing::debug;

crate fn collect_definitions(
    resolver: &mut Resolver<'_>,
    fragment: &AstFragment,
    expansion: ExpnId,
) {
    let parent_def = resolver.invocation_parents[&expansion];
    fragment.visit_with(&mut DefCollector { resolver, parent_def, expansion });
}

/// Creates `DefId`s for nodes in the AST.
struct DefCollector<'a, 'b> {
    resolver: &'a mut Resolver<'b>,
    parent_def: LocalDefId,
    expansion: ExpnId,
}

impl<'a, 'b> DefCollector<'a, 'b> {
    fn create_def(&mut self, node_id: NodeId, data: DefPathData, span: Span) -> LocalDefId {
        let parent_def = self.parent_def;
        debug!("create_def(node_id={:?}, data={:?}, parent_def={:?})", node_id, data, parent_def);
        self.resolver.create_def(parent_def, node_id, data, self.expansion, span)
    }

    fn with_parent<F: FnOnce(&mut Self)>(&mut self, parent_def: LocalDefId, f: F) {
        let orig_parent_def = std::mem::replace(&mut self.parent_def, parent_def);
        f(self);
        self.parent_def = orig_parent_def;
    }

    fn collect_field(&mut self, field: &'a StructField, index: Option<usize>) {
        let index = |this: &Self| {
            index.unwrap_or_else(|| {
                let node_id = NodeId::placeholder_from_expn_id(this.expansion);
                this.resolver.placeholder_field_indices[&node_id]
            })
        };

        if field.is_placeholder {
            let old_index = self.resolver.placeholder_field_indices.insert(field.id, index(self));
            assert!(old_index.is_none(), "placeholder field index is reset for a node ID");
            self.visit_macro_invoc(field.id);
        } else {
            let name = field.ident.map_or_else(|| sym::integer(index(self)), |ident| ident.name);
            let def = self.create_def(field.id, DefPathData::ValueNs(name), field.span);
            self.with_parent(def, |this| visit::walk_struct_field(this, field));
        }
    }

    fn visit_macro_invoc(&mut self, id: NodeId) {
        let old_parent =
            self.resolver.invocation_parents.insert(id.placeholder_to_expn_id(), self.parent_def);
        assert!(old_parent.is_none(), "parent `LocalDefId` is reset for an invocation");
    }
}

impl<'a, 'b> visit::Visitor<'a> for DefCollector<'a, 'b> {
    fn visit_item(&mut self, i: &'a Item) {
        debug!("visit_item: {:?}", i);

        // Pick the def data. This need not be unique, but the more
        // information we encapsulate into, the better
        let def_data = match &i.kind {
            ItemKind::Impl { .. } => DefPathData::Impl,
            ItemKind::Mod(..) if i.ident.name == kw::Invalid => {
                return visit::walk_item(self, i);
            }
            ItemKind::Mod(..)
            | ItemKind::Trait(..)
            | ItemKind::TraitAlias(..)
            | ItemKind::Enum(..)
            | ItemKind::Struct(..)
            | ItemKind::Union(..)
            | ItemKind::ExternCrate(..)
            | ItemKind::ForeignMod(..)
            | ItemKind::TyAlias(..) => DefPathData::TypeNs(i.ident.name),
            ItemKind::Static(..) | ItemKind::Const(..) | ItemKind::Fn(..) => {
                DefPathData::ValueNs(i.ident.name)
            }
            ItemKind::MacroDef(..) => DefPathData::MacroNs(i.ident.name),
            ItemKind::MacCall(..) => return self.visit_macro_invoc(i.id),
            ItemKind::GlobalAsm(..) => DefPathData::Misc,
            ItemKind::Use(..) => {
                return visit::walk_item(self, i);
            }
        };
        let def = self.create_def(i.id, def_data, i.span);

        self.with_parent(def, |this| {
            match i.kind {
                ItemKind::Struct(ref struct_def, _) | ItemKind::Union(ref struct_def, _) => {
                    // If this is a unit or tuple-like struct, register the constructor.
                    if let Some(ctor_hir_id) = struct_def.ctor_id() {
                        this.create_def(ctor_hir_id, DefPathData::Ctor, i.span);
                    }
                }
                _ => {}
            }
            visit::walk_item(this, i);
        });
    }

    fn visit_fn(&mut self, fn_kind: FnKind<'a>, span: Span, _: NodeId) {
        if let FnKind::Fn(_, _, sig, _, body) = fn_kind {
            if let Async::Yes { closure_id, return_impl_trait_id, .. } = sig.header.asyncness {
                self.create_def(return_impl_trait_id, DefPathData::ImplTrait, span);

                // For async functions, we need to create their inner defs inside of a
                // closure to match their desugared representation. Besides that,
                // we must mirror everything that `visit::walk_fn` below does.
                self.visit_fn_header(&sig.header);
                visit::walk_fn_decl(self, &sig.decl);
                let closure_def = self.create_def(closure_id, DefPathData::ClosureExpr, span);
                self.with_parent(closure_def, |this| walk_list!(this, visit_block, body));
                return;
            }
        }

        visit::walk_fn(self, fn_kind, span);
    }

    fn visit_use_tree(&mut self, use_tree: &'a UseTree, id: NodeId, _nested: bool) {
        self.create_def(id, DefPathData::Misc, use_tree.span);
        visit::walk_use_tree(self, use_tree, id);
    }

    fn visit_foreign_item(&mut self, foreign_item: &'a ForeignItem) {
        if let ForeignItemKind::MacCall(_) = foreign_item.kind {
            return self.visit_macro_invoc(foreign_item.id);
        }

        let def = self.create_def(
            foreign_item.id,
            DefPathData::ValueNs(foreign_item.ident.name),
            foreign_item.span,
        );

        self.with_parent(def, |this| {
            visit::walk_foreign_item(this, foreign_item);
        });
    }

    fn visit_variant(&mut self, v: &'a Variant) {
        if v.is_placeholder {
            return self.visit_macro_invoc(v.id);
        }
        let def = self.create_def(v.id, DefPathData::TypeNs(v.ident.name), v.span);
        self.with_parent(def, |this| {
            if let Some(ctor_hir_id) = v.data.ctor_id() {
                this.create_def(ctor_hir_id, DefPathData::Ctor, v.span);
            }
            visit::walk_variant(this, v)
        });
    }

    fn visit_variant_data(&mut self, data: &'a VariantData) {
        // The assumption here is that non-`cfg` macro expansion cannot change field indices.
        // It currently holds because only inert attributes are accepted on fields,
        // and every such attribute expands into a single field after it's resolved.
        for (index, field) in data.fields().iter().enumerate() {
            self.collect_field(field, Some(index));
        }
    }

    fn visit_generic_param(&mut self, param: &'a GenericParam) {
        if param.is_placeholder {
            self.visit_macro_invoc(param.id);
            return;
        }
        let name = param.ident.name;
        let def_path_data = match param.kind {
            GenericParamKind::Lifetime { .. } => DefPathData::LifetimeNs(name),
            GenericParamKind::Type { .. } => DefPathData::TypeNs(name),
            GenericParamKind::Const { .. } => DefPathData::ValueNs(name),
        };
        self.create_def(param.id, def_path_data, param.ident.span);

        visit::walk_generic_param(self, param);
    }

    fn visit_assoc_item(&mut self, i: &'a AssocItem, ctxt: visit::AssocCtxt) {
        let def_data = match &i.kind {
            AssocItemKind::Fn(..) | AssocItemKind::Const(..) => DefPathData::ValueNs(i.ident.name),
            AssocItemKind::TyAlias(..) => DefPathData::TypeNs(i.ident.name),
            AssocItemKind::MacCall(..) => return self.visit_macro_invoc(i.id),
        };

        let def = self.create_def(i.id, def_data, i.span);
        self.with_parent(def, |this| visit::walk_assoc_item(this, i, ctxt));
    }

    fn visit_pat(&mut self, pat: &'a Pat) {
        match pat.kind {
            PatKind::MacCall(..) => self.visit_macro_invoc(pat.id),
            _ => visit::walk_pat(self, pat),
        }
    }

    fn visit_anon_const(&mut self, constant: &'a AnonConst) {
        let def = self.create_def(constant.id, DefPathData::AnonConst, constant.value.span);
        self.with_parent(def, |this| visit::walk_anon_const(this, constant));
    }

    fn visit_expr(&mut self, expr: &'a Expr) {
        let parent_def = match expr.kind {
            ExprKind::MacCall(..) => return self.visit_macro_invoc(expr.id),
            ExprKind::Closure(_, asyncness, ..) => {
                // Async closures desugar to closures inside of closures, so
                // we must create two defs.
                let closure_def = self.create_def(expr.id, DefPathData::ClosureExpr, expr.span);
                match asyncness {
                    Async::Yes { closure_id, .. } => {
                        self.create_def(closure_id, DefPathData::ClosureExpr, expr.span)
                    }
                    Async::No => closure_def,
                }
            }
            ExprKind::Async(_, async_id, _) => {
                self.create_def(async_id, DefPathData::ClosureExpr, expr.span)
            }
            _ => self.parent_def,
        };

        self.with_parent(parent_def, |this| visit::walk_expr(this, expr));
    }

    fn visit_ty(&mut self, ty: &'a Ty) {
        match ty.kind {
            TyKind::MacCall(..) => return self.visit_macro_invoc(ty.id),
            TyKind::ImplTrait(node_id, _) => {
                self.create_def(node_id, DefPathData::ImplTrait, ty.span);
            }
            _ => {}
        }
        visit::walk_ty(self, ty);
    }

    fn visit_stmt(&mut self, stmt: &'a Stmt) {
        match stmt.kind {
            StmtKind::MacCall(..) => self.visit_macro_invoc(stmt.id),
            _ => visit::walk_stmt(self, stmt),
        }
    }

    fn visit_token(&mut self, t: Token) {
        if let token::Interpolated(nt) = t.kind {
            if let token::NtExpr(ref expr) = *nt {
                if let ExprKind::MacCall(..) = expr.kind {
                    self.visit_macro_invoc(expr.id);
                }
            }
        }
    }

    fn visit_arm(&mut self, arm: &'a Arm) {
        if arm.is_placeholder { self.visit_macro_invoc(arm.id) } else { visit::walk_arm(self, arm) }
    }

    fn visit_field(&mut self, f: &'a Field) {
        if f.is_placeholder { self.visit_macro_invoc(f.id) } else { visit::walk_field(self, f) }
    }

    fn visit_field_pattern(&mut self, fp: &'a FieldPat) {
        if fp.is_placeholder {
            self.visit_macro_invoc(fp.id)
        } else {
            visit::walk_field_pattern(self, fp)
        }
    }

    fn visit_param(&mut self, p: &'a Param) {
        if p.is_placeholder { self.visit_macro_invoc(p.id) } else { visit::walk_param(self, p) }
    }

    // This method is called only when we are visiting an individual field
    // after expanding an attribute on it.
    fn visit_struct_field(&mut self, field: &'a StructField) {
        self.collect_field(field, None);
    }
}
<commit_msg>fix def collector for impl trait<commit_after>use crate::Resolver;
use rustc_ast::token::{self, Token};
use rustc_ast::visit::{self, FnKind};
use rustc_ast::walk_list;
use rustc_ast::*;
use rustc_ast_lowering::ResolverAstLowering;
use rustc_expand::expand::AstFragment;
use rustc_hir::def_id::LocalDefId;
use rustc_hir::definitions::*;
use rustc_span::hygiene::ExpnId;
use rustc_span::symbol::{kw, sym};
use rustc_span::Span;
use tracing::debug;

crate fn collect_definitions(
    resolver: &mut Resolver<'_>,
    fragment: &AstFragment,
    expansion: ExpnId,
) {
    let parent_def = resolver.invocation_parents[&expansion];
    fragment.visit_with(&mut DefCollector { resolver, parent_def, expansion });
}

/// Creates `DefId`s for nodes in the AST.
struct DefCollector<'a, 'b> {
    resolver: &'a mut Resolver<'b>,
    parent_def: LocalDefId,
    expansion: ExpnId,
}

impl<'a, 'b> DefCollector<'a, 'b> {
    fn create_def(&mut self, node_id: NodeId, data: DefPathData, span: Span) -> LocalDefId {
        let parent_def = self.parent_def;
        debug!("create_def(node_id={:?}, data={:?}, parent_def={:?})", node_id, data, parent_def);
        self.resolver.create_def(parent_def, node_id, data, self.expansion, span)
    }

    fn with_parent<F: FnOnce(&mut Self)>(&mut self, parent_def: LocalDefId, f: F) {
        let orig_parent_def = std::mem::replace(&mut self.parent_def, parent_def);
        f(self);
        self.parent_def = orig_parent_def;
    }

    fn collect_field(&mut self, field: &'a StructField, index: Option<usize>) {
        let index = |this: &Self| {
            index.unwrap_or_else(|| {
                let node_id = NodeId::placeholder_from_expn_id(this.expansion);
                this.resolver.placeholder_field_indices[&node_id]
            })
        };

        if field.is_placeholder {
            let old_index = self.resolver.placeholder_field_indices.insert(field.id, index(self));
            assert!(old_index.is_none(), "placeholder field index is reset for a node ID");
            self.visit_macro_invoc(field.id);
        } else {
            let name = field.ident.map_or_else(|| sym::integer(index(self)), |ident| ident.name);
            let def = self.create_def(field.id, DefPathData::ValueNs(name), field.span);
            self.with_parent(def, |this| visit::walk_struct_field(this, field));
        }
    }

    fn visit_macro_invoc(&mut self, id: NodeId) {
        let old_parent =
            self.resolver.invocation_parents.insert(id.placeholder_to_expn_id(), self.parent_def);
        assert!(old_parent.is_none(), "parent `LocalDefId` is reset for an invocation");
    }
}

impl<'a, 'b> visit::Visitor<'a> for DefCollector<'a, 'b> {
    fn visit_item(&mut self, i: &'a Item) {
        debug!("visit_item: {:?}", i);

        // Pick the def data. This need not be unique, but the more
        // information we encapsulate into, the better
        let def_data = match &i.kind {
            ItemKind::Impl { .. } => DefPathData::Impl,
            ItemKind::Mod(..) if i.ident.name == kw::Invalid => {
                return visit::walk_item(self, i);
            }
            ItemKind::Mod(..)
            | ItemKind::Trait(..)
            | ItemKind::TraitAlias(..)
            | ItemKind::Enum(..)
            | ItemKind::Struct(..)
            | ItemKind::Union(..)
            | ItemKind::ExternCrate(..)
            | ItemKind::ForeignMod(..)
            | ItemKind::TyAlias(..) => DefPathData::TypeNs(i.ident.name),
            ItemKind::Static(..) | ItemKind::Const(..) | ItemKind::Fn(..) => {
                DefPathData::ValueNs(i.ident.name)
            }
            ItemKind::MacroDef(..) => DefPathData::MacroNs(i.ident.name),
            ItemKind::MacCall(..) => return self.visit_macro_invoc(i.id),
            ItemKind::GlobalAsm(..) => DefPathData::Misc,
            ItemKind::Use(..) => {
                return visit::walk_item(self, i);
            }
        };
        let def = self.create_def(i.id, def_data, i.span);

        self.with_parent(def, |this| {
            match i.kind {
                ItemKind::Struct(ref struct_def, _) | ItemKind::Union(ref struct_def, _) => {
                    // If this is a unit or tuple-like struct, register the constructor.
                    if let Some(ctor_hir_id) = struct_def.ctor_id() {
                        this.create_def(ctor_hir_id, DefPathData::Ctor, i.span);
                    }
                }
                _ => {}
            }
            visit::walk_item(this, i);
        });
    }

    fn visit_fn(&mut self, fn_kind: FnKind<'a>, span: Span, _: NodeId) {
        if let FnKind::Fn(_, _, sig, _, body) = fn_kind {
            if let Async::Yes { closure_id, return_impl_trait_id, .. } = sig.header.asyncness {
                self.create_def(return_impl_trait_id, DefPathData::ImplTrait, span);

                // For async functions, we need to create their inner defs inside of a
                // closure to match their desugared representation. Besides that,
                // we must mirror everything that `visit::walk_fn` below does.
                self.visit_fn_header(&sig.header);
                visit::walk_fn_decl(self, &sig.decl);
                let closure_def = self.create_def(closure_id, DefPathData::ClosureExpr, span);
                self.with_parent(closure_def, |this| walk_list!(this, visit_block, body));
                return;
            }
        }

        visit::walk_fn(self, fn_kind, span);
    }

    fn visit_use_tree(&mut self, use_tree: &'a UseTree, id: NodeId, _nested: bool) {
        self.create_def(id, DefPathData::Misc, use_tree.span);
        visit::walk_use_tree(self, use_tree, id);
    }

    fn visit_foreign_item(&mut self, foreign_item: &'a ForeignItem) {
        if let ForeignItemKind::MacCall(_) = foreign_item.kind {
            return self.visit_macro_invoc(foreign_item.id);
        }

        let def = self.create_def(
            foreign_item.id,
            DefPathData::ValueNs(foreign_item.ident.name),
            foreign_item.span,
        );

        self.with_parent(def, |this| {
            visit::walk_foreign_item(this, foreign_item);
        });
    }

    fn visit_variant(&mut self, v: &'a Variant) {
        if v.is_placeholder {
            return self.visit_macro_invoc(v.id);
        }
        let def = self.create_def(v.id, DefPathData::TypeNs(v.ident.name), v.span);
        self.with_parent(def, |this| {
            if let Some(ctor_hir_id) = v.data.ctor_id() {
                this.create_def(ctor_hir_id, DefPathData::Ctor, v.span);
            }
            visit::walk_variant(this, v)
        });
    }

    fn visit_variant_data(&mut self, data: &'a VariantData) {
        // The assumption here is that non-`cfg` macro expansion cannot change field indices.
        // It currently holds because only inert attributes are accepted on fields,
        // and every such attribute expands into a single field after it's resolved.
        for (index, field) in data.fields().iter().enumerate() {
            self.collect_field(field, Some(index));
        }
    }

    fn visit_generic_param(&mut self, param: &'a GenericParam) {
        if param.is_placeholder {
            self.visit_macro_invoc(param.id);
            return;
        }
        let name = param.ident.name;
        let def_path_data = match param.kind {
            GenericParamKind::Lifetime { .. } => DefPathData::LifetimeNs(name),
            GenericParamKind::Type { .. } => DefPathData::TypeNs(name),
            GenericParamKind::Const { .. } => DefPathData::ValueNs(name),
        };
        self.create_def(param.id, def_path_data, param.ident.span);

        visit::walk_generic_param(self, param);
    }

    fn visit_assoc_item(&mut self, i: &'a AssocItem, ctxt: visit::AssocCtxt) {
        let def_data = match &i.kind {
            AssocItemKind::Fn(..) | AssocItemKind::Const(..) => DefPathData::ValueNs(i.ident.name),
            AssocItemKind::TyAlias(..) => DefPathData::TypeNs(i.ident.name),
            AssocItemKind::MacCall(..) => return self.visit_macro_invoc(i.id),
        };

        let def = self.create_def(i.id, def_data, i.span);
        self.with_parent(def, |this| visit::walk_assoc_item(this, i, ctxt));
    }

    fn visit_pat(&mut self, pat: &'a Pat) {
        match pat.kind {
            PatKind::MacCall(..) => self.visit_macro_invoc(pat.id),
            _ => visit::walk_pat(self, pat),
        }
    }

    fn visit_anon_const(&mut self, constant: &'a AnonConst) {
        let def = self.create_def(constant.id, DefPathData::AnonConst, constant.value.span);
        self.with_parent(def, |this| visit::walk_anon_const(this, constant));
    }

    fn visit_expr(&mut self, expr: &'a Expr) {
        let parent_def = match expr.kind {
            ExprKind::MacCall(..) => return self.visit_macro_invoc(expr.id),
            ExprKind::Closure(_, asyncness, ..) => {
                // Async closures desugar to closures inside of closures, so
                // we must create two defs.
                let closure_def = self.create_def(expr.id, DefPathData::ClosureExpr, expr.span);
                match asyncness {
                    Async::Yes { closure_id, .. } => {
                        self.create_def(closure_id, DefPathData::ClosureExpr, expr.span)
                    }
                    Async::No => closure_def,
                }
            }
            ExprKind::Async(_, async_id, _) => {
                self.create_def(async_id, DefPathData::ClosureExpr, expr.span)
            }
            _ => self.parent_def,
        };

        self.with_parent(parent_def, |this| visit::walk_expr(this, expr));
    }

    fn visit_ty(&mut self, ty: &'a Ty) {
        match ty.kind {
            TyKind::MacCall(..) => self.visit_macro_invoc(ty.id),
            TyKind::ImplTrait(node_id, _) => {
                let parent_def = self.create_def(node_id, DefPathData::ImplTrait, ty.span);
                self.with_parent(parent_def, |this| visit::walk_ty(this, ty));
            }
            _ => visit::walk_ty(self, ty),
        }
    }

    fn visit_stmt(&mut self, stmt: &'a Stmt) {
        match stmt.kind {
            StmtKind::MacCall(..) => self.visit_macro_invoc(stmt.id),
            _ => visit::walk_stmt(self, stmt),
        }
    }

    fn visit_token(&mut self, t: Token) {
        if let token::Interpolated(nt) = t.kind {
            if let token::NtExpr(ref expr) = *nt {
                if let ExprKind::MacCall(..) = expr.kind {
                    self.visit_macro_invoc(expr.id);
                }
            }
        }
    }

    fn visit_arm(&mut self, arm: &'a Arm) {
        if arm.is_placeholder { self.visit_macro_invoc(arm.id) } else { visit::walk_arm(self, arm) }
    }

    fn visit_field(&mut self, f: &'a Field) {
        if f.is_placeholder { self.visit_macro_invoc(f.id) } else { visit::walk_field(self, f) }
    }

    fn visit_field_pattern(&mut self, fp: &'a FieldPat) {
        if fp.is_placeholder {
            self.visit_macro_invoc(fp.id)
        } else {
            visit::walk_field_pattern(self, fp)
        }
    }

    fn visit_param(&mut self, p: &'a Param) {
        if p.is_placeholder { self.visit_macro_invoc(p.id) } else { visit::walk_param(self, p) }
    }

    // This method is called only when we are visiting an individual field
    // after expanding an attribute on it.
    fn visit_struct_field(&mut self, field: &'a StructField) {
        self.collect_field(field, None);
    }
}
<|endoftext|>


<commit_before>use hir::db::HirDatabase;
use ra_syntax::{
    ast::{self, AstNode},
    SyntaxKind::{
        WHITESPACE, MATCH_ARM, LAMBDA_EXPR, PATH_EXPR, BREAK_EXPR, LOOP_EXPR, RETURN_EXPR, COMMENT
    }, SyntaxNode, TextUnit,
};

use crate::{AssistCtx, Assist};

pub(crate) fn introduce_variable(ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let node = ctx.covering_node();
    if !valid_covering_node(node) {
        return None;
    }
    let expr = node.ancestors().filter_map(valid_target_expr).next()?;
    let (anchor_stmt, wrap_in_block) = anchor_stmt(expr)?;
    let indent = anchor_stmt.prev_sibling()?;
    if indent.kind() != WHITESPACE {
        return None;
    }
    ctx.build("introduce variable", move |edit| {
        let mut buf = String::new();

        let cursor_offset = if wrap_in_block {
            buf.push_str("{ let var_name = ");
            TextUnit::of_str("{ let ")
        } else {
            buf.push_str("let var_name = ");
            TextUnit::of_str("let ")
        };

        expr.syntax().text().push_to(&mut buf);
        let full_stmt = ast::ExprStmt::cast(anchor_stmt);
        let is_full_stmt = if let Some(expr_stmt) = full_stmt {
            Some(expr.syntax()) == expr_stmt.expr().map(|e| e.syntax())
        } else {
            false
        };
        if is_full_stmt {
            if !full_stmt.unwrap().has_semi() {
                buf.push_str(";");
            }
            edit.replace(expr.syntax().range(), buf);
        } else {
            buf.push_str(";");
            indent.text().push_to(&mut buf);
            edit.target(expr.syntax().range());
            edit.replace(expr.syntax().range(), "var_name".to_string());
            edit.insert(anchor_stmt.range().start(), buf);
            if wrap_in_block {
                edit.insert(anchor_stmt.range().end(), " }");
            }
        }
        edit.set_cursor(anchor_stmt.range().start() + cursor_offset);
    })
}

fn valid_covering_node(node: &SyntaxNode) -> bool {
    node.kind() != COMMENT
}
/// Check whether the node is a valid expression which can be extracted to a variable.
/// In general that's true for any expression, but in some cases that would produce invalid code.
fn valid_target_expr(node: &SyntaxNode) -> Option<&ast::Expr> {
    match node.kind() {
        PATH_EXPR => None,
        BREAK_EXPR => ast::BreakExpr::cast(node).and_then(|e| e.expr()),
        RETURN_EXPR => ast::ReturnExpr::cast(node).and_then(|e| e.expr()),
        LOOP_EXPR => ast::ReturnExpr::cast(node).and_then(|e| e.expr()),
        _ => ast::Expr::cast(node),
    }
}

/// Returns the syntax node which will follow the freshly introduced var
/// and a boolean indicating whether we have to wrap it within a { } block
/// to produce correct code.
/// It can be a statement, the last in a block expression or a wanna be block
/// expression like a lambda or match arm.
fn anchor_stmt(expr: &ast::Expr) -> Option<(&SyntaxNode, bool)> {
    expr.syntax().ancestors().find_map(|node| {
        if ast::Stmt::cast(node).is_some() {
            return Some((node, false));
        }

        if let Some(expr) = node.parent().and_then(ast::Block::cast).and_then(|it| it.expr()) {
            if expr.syntax() == node {
                return Some((node, false));
            }
        }

        if let Some(parent) = node.parent() {
            if parent.kind() == MATCH_ARM || parent.kind() == LAMBDA_EXPR {
                return Some((node, true));
            }
        }

        None
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_range, check_assist_target, check_assist_range_target};

    #[test]
    fn test_introduce_var_simple() {
        check_assist_range(
            introduce_variable,
            "
fn foo() {
    foo(<|>1 + 1<|>);
}",
            "
fn foo() {
    let <|>var_name = 1 + 1;
    foo(var_name);
}",
        );
    }

    #[test]
    fn test_introduce_var_expr_stmt() {
        check_assist_range(
            introduce_variable,
            "
fn foo() {
    <|>1 + 1<|>;
}",
            "
fn foo() {
    let <|>var_name = 1 + 1;
}",
        );
    }

    #[test]
    fn test_introduce_var_part_of_expr_stmt() {
        check_assist_range(
            introduce_variable,
            "
fn foo() {
    <|>1<|> + 1;
}",
            "
fn foo() {
    let <|>var_name = 1;
    var_name + 1;
}",
        );
    }

    #[test]
    fn test_introduce_var_last_expr() {
        check_assist_range(
            introduce_variable,
            "
fn foo() {
    bar(<|>1 + 1<|>)
}",
            "
fn foo() {
    let <|>var_name = 1 + 1;
    bar(var_name)
}",
        );
    }

    #[test]
    fn test_introduce_var_last_full_expr() {
        check_assist_range(
            introduce_variable,
            "
fn foo() {
    <|>bar(1 + 1)<|>
}",
            "
fn foo() {
    let <|>var_name = bar(1 + 1);
    var_name
}",
        );
    }

    #[test]
    fn test_introduce_var_block_expr_second_to_last() {
        check_assist_range(
            introduce_variable,
            "
fn foo() {
    <|>{ let x = 0; x }<|>
    something_else();
}",
            "
fn foo() {
    let <|>var_name = { let x = 0; x };
    something_else();
}",
        );
    }

    #[test]
    fn test_introduce_var_in_match_arm_no_block() {
        check_assist_range(
            introduce_variable,
            "
fn main() {
    let x = true;
    let tuple = match x {
        true => (<|>2 + 2<|>, true)
        _ => (0, false)
    };
}
",
            "
fn main() {
    let x = true;
    let tuple = match x {
        true => { let <|>var_name = 2 + 2; (var_name, true) }
        _ => (0, false)
    };
}
",
        );
    }

    #[test]
    fn test_introduce_var_in_match_arm_with_block() {
        check_assist_range(
            introduce_variable,
            "
fn main() {
    let x = true;
    let tuple = match x {
        true => {
            let y = 1;
            (<|>2 + y<|>, true)
        }
        _ => (0, false)
    };
}
",
            "
fn main() {
    let x = true;
    let tuple = match x {
        true => {
            let y = 1;
            let <|>var_name = 2 + y;
            (var_name, true)
        }
        _ => (0, false)
    };
}
",
        );
    }

    #[test]
    fn test_introduce_var_in_closure_no_block() {
        check_assist_range(
            introduce_variable,
            "
fn main() {
    let lambda = |x: u32| <|>x * 2<|>;
}
",
            "
fn main() {
    let lambda = |x: u32| { let <|>var_name = x * 2; var_name };
}
",
        );
    }

    #[test]
    fn test_introduce_var_in_closure_with_block() {
        check_assist_range(
            introduce_variable,
            "
fn main() {
    let lambda = |x: u32| { <|>x * 2<|> };
}
",
            "
fn main() {
    let lambda = |x: u32| { let <|>var_name = x * 2; var_name };
}
",
        );
    }

    #[test]
    fn test_introduce_var_path_simple() {
        check_assist(
            introduce_variable,
            "
fn main() {
    let o = S<|>ome(true);
}
",
            "
fn main() {
    let <|>var_name = Some(true);
    let o = var_name;
}
",
        );
    }

    #[test]
    fn test_introduce_var_path_method() {
        check_assist(
            introduce_variable,
            "
fn main() {
    let v = b<|>ar.foo();
}
",
            "
fn main() {
    let <|>var_name = bar.foo();
    let v = var_name;
}
",
        );
    }

    #[test]
    fn test_introduce_var_return() {
        check_assist(
            introduce_variable,
            "
fn foo() -> u32 {
    r<|>eturn 2 + 2;
}
",
            "
fn foo() -> u32 {
    let <|>var_name = 2 + 2;
    return var_name;
}
",
        );
    }

    #[test]
    fn test_introduce_var_break() {
        check_assist(
            introduce_variable,
            "
fn main() {
    let result = loop {
        b<|>reak 2 + 2;
    };
}
",
            "
fn main() {
    let result = loop {
        let <|>var_name = 2 + 2;
        break var_name;
    };
}
",
        );
    }

    #[test]
    fn test_introduce_var_for_cast() {
        check_assist(
            introduce_variable,
            "
fn main() {
    let v = 0f32 a<|>s u32;
}
",
            "
fn main() {
    let <|>var_name = 0f32 as u32;
    let v = var_name;
}
",
        );
    }

    #[test]
    fn test_introduce_var_for_return_not_applicable() {
        check_assist_not_applicable(
            introduce_variable,
            "
fn foo() {
    r<|>eturn;
}
",
        );
    }

    #[test]
    fn test_introduce_var_for_break_not_applicable() {
        check_assist_not_applicable(
            introduce_variable,
            "
fn main() {
    loop {
        b<|>reak;
    };
}
",
        );
    }

    #[test]
    fn test_introduce_var_in_comment_not_applicable() {
        check_assist_not_applicable(
            introduce_variable,
            "
fn main() {
    let x = true;
    let tuple = match x {
        // c<|>omment
        true => (2 + 2, true)
        _ => (0, false)
    };
}
",
        );
    }

    // FIXME: This is not quite correct, but good enough(tm) for the sorting heuristic
    #[test]
    fn introduce_var_target() {
        check_assist_target(
            introduce_variable,
            "
fn foo() -> u32 {
    r<|>eturn 2 + 2;
}
",
            "2 + 2",
        );

        check_assist_range_target(
            introduce_variable,
            "
fn main() {
    let x = true;
    let tuple = match x {
        true => (<|>2 + 2<|>, true)
        _ => (0, false)
    };
}
",
            "2 + 2",
        );
    }
}
<commit_msg>Fix introduce var duplicating newlines<commit_after>use hir::db::HirDatabase;
use ra_syntax::{
    ast::{self, AstNode},
    SyntaxKind::{
        WHITESPACE, MATCH_ARM, LAMBDA_EXPR, PATH_EXPR, BREAK_EXPR, LOOP_EXPR, RETURN_EXPR, COMMENT
    }, SyntaxNode, TextUnit,
};

use crate::{AssistCtx, Assist};

pub(crate) fn introduce_variable(ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let node = ctx.covering_node();
    if !valid_covering_node(node) {
        return None;
    }
    let expr = node.ancestors().filter_map(valid_target_expr).next()?;
    let (anchor_stmt, wrap_in_block) = anchor_stmt(expr)?;
    let indent = anchor_stmt.prev_sibling()?;
    if indent.kind() != WHITESPACE {
        return None;
    }
    ctx.build("introduce variable", move |edit| {
        let mut buf = String::new();

        let cursor_offset = if wrap_in_block {
            buf.push_str("{ let var_name = ");
            TextUnit::of_str("{ let ")
        } else {
            buf.push_str("let var_name = ");
            TextUnit::of_str("let ")
        };

        expr.syntax().text().push_to(&mut buf);
        let full_stmt = ast::ExprStmt::cast(anchor_stmt);
        let is_full_stmt = if let Some(expr_stmt) = full_stmt {
            Some(expr.syntax()) == expr_stmt.expr().map(|e| e.syntax())
        } else {
            false
        };
        if is_full_stmt {
            if !full_stmt.unwrap().has_semi() {
                buf.push_str(";");
            }
            edit.replace(expr.syntax().range(), buf);
        } else {
            buf.push_str(";");

            // We want to maintain the indent level,
            // but we do not want to duplicate possible
            // extra newlines in the indent block
            for chunk in indent.text().chunks() {
                if chunk.starts_with("\r\n") {
                    buf.push_str("\r\n");
                    buf.push_str(chunk.trim_start_matches("\r\n"));
                } else if chunk.starts_with("\n") {
                    buf.push_str("\n");
                    buf.push_str(chunk.trim_start_matches("\n"));
                } else {
                    buf.push_str(chunk);
                }
            }

            edit.target(expr.syntax().range());
            edit.replace(expr.syntax().range(), "var_name".to_string());
            edit.insert(anchor_stmt.range().start(), buf);
            if wrap_in_block {
                edit.insert(anchor_stmt.range().end(), " }");
            }
        }
        edit.set_cursor(anchor_stmt.range().start() + cursor_offset);
    })
}

fn valid_covering_node(node: &SyntaxNode) -> bool {
    node.kind() != COMMENT
}
/// Check whether the node is a valid expression which can be extracted to a variable.
/// In general that's true for any expression, but in some cases that would produce invalid code.
fn valid_target_expr(node: &SyntaxNode) -> Option<&ast::Expr> {
    match node.kind() {
        PATH_EXPR => None,
        BREAK_EXPR => ast::BreakExpr::cast(node).and_then(|e| e.expr()),
        RETURN_EXPR => ast::ReturnExpr::cast(node).and_then(|e| e.expr()),
        LOOP_EXPR => ast::ReturnExpr::cast(node).and_then(|e| e.expr()),
        _ => ast::Expr::cast(node),
    }
}

/// Returns the syntax node which will follow the freshly introduced var
/// and a boolean indicating whether we have to wrap it within a { } block
/// to produce correct code.
/// It can be a statement, the last in a block expression or a wanna be block
/// expression like a lambda or match arm.
fn anchor_stmt(expr: &ast::Expr) -> Option<(&SyntaxNode, bool)> {
    expr.syntax().ancestors().find_map(|node| {
        if ast::Stmt::cast(node).is_some() {
            return Some((node, false));
        }

        if let Some(expr) = node.parent().and_then(ast::Block::cast).and_then(|it| it.expr()) {
            if expr.syntax() == node {
                return Some((node, false));
            }
        }

        if let Some(parent) = node.parent() {
            if parent.kind() == MATCH_ARM || parent.kind() == LAMBDA_EXPR {
                return Some((node, true));
            }
        }

        None
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable, check_assist_range, check_assist_target, check_assist_range_target};

    #[test]
    fn test_introduce_var_simple() {
        check_assist_range(
            introduce_variable,
            "
fn foo() {
    foo(<|>1 + 1<|>);
}",
            "
fn foo() {
    let <|>var_name = 1 + 1;
    foo(var_name);
}",
        );
    }

    #[test]
    fn test_introduce_var_expr_stmt() {
        check_assist_range(
            introduce_variable,
            "
fn foo() {
    <|>1 + 1<|>;
}",
            "
fn foo() {
    let <|>var_name = 1 + 1;
}",
        );
    }

    #[test]
    fn test_introduce_var_part_of_expr_stmt() {
        check_assist_range(
            introduce_variable,
            "
fn foo() {
    <|>1<|> + 1;
}",
            "
fn foo() {
    let <|>var_name = 1;
    var_name + 1;
}",
        );
    }

    #[test]
    fn test_introduce_var_last_expr() {
        check_assist_range(
            introduce_variable,
            "
fn foo() {
    bar(<|>1 + 1<|>)
}",
            "
fn foo() {
    let <|>var_name = 1 + 1;
    bar(var_name)
}",
        );
    }

    #[test]
    fn test_introduce_var_last_full_expr() {
        check_assist_range(
            introduce_variable,
            "
fn foo() {
    <|>bar(1 + 1)<|>
}",
            "
fn foo() {
    let <|>var_name = bar(1 + 1);
    var_name
}",
        );
    }

    #[test]
    fn test_introduce_var_block_expr_second_to_last() {
        check_assist_range(
            introduce_variable,
            "
fn foo() {
    <|>{ let x = 0; x }<|>
    something_else();
}",
            "
fn foo() {
    let <|>var_name = { let x = 0; x };
    something_else();
}",
        );
    }

    #[test]
    fn test_introduce_var_in_match_arm_no_block() {
        check_assist_range(
            introduce_variable,
            "
fn main() {
    let x = true;
    let tuple = match x {
        true => (<|>2 + 2<|>, true)
        _ => (0, false)
    };
}
",
            "
fn main() {
    let x = true;
    let tuple = match x {
        true => { let <|>var_name = 2 + 2; (var_name, true) }
        _ => (0, false)
    };
}
",
        );
    }

    #[test]
    fn test_introduce_var_in_match_arm_with_block() {
        check_assist_range(
            introduce_variable,
            "
fn main() {
    let x = true;
    let tuple = match x {
        true => {
            let y = 1;
            (<|>2 + y<|>, true)
        }
        _ => (0, false)
    };
}
",
            "
fn main() {
    let x = true;
    let tuple = match x {
        true => {
            let y = 1;
            let <|>var_name = 2 + y;
            (var_name, true)
        }
        _ => (0, false)
    };
}
",
        );
    }

    #[test]
    fn test_introduce_var_in_closure_no_block() {
        check_assist_range(
            introduce_variable,
            "
fn main() {
    let lambda = |x: u32| <|>x * 2<|>;
}
",
            "
fn main() {
    let lambda = |x: u32| { let <|>var_name = x * 2; var_name };
}
",
        );
    }

    #[test]
    fn test_introduce_var_in_closure_with_block() {
        check_assist_range(
            introduce_variable,
            "
fn main() {
    let lambda = |x: u32| { <|>x * 2<|> };
}
",
            "
fn main() {
    let lambda = |x: u32| { let <|>var_name = x * 2; var_name };
}
",
        );
    }

    #[test]
    fn test_introduce_var_path_simple() {
        check_assist(
            introduce_variable,
            "
fn main() {
    let o = S<|>ome(true);
}
",
            "
fn main() {
    let <|>var_name = Some(true);
    let o = var_name;
}
",
        );
    }

    #[test]
    fn test_introduce_var_path_method() {
        check_assist(
            introduce_variable,
            "
fn main() {
    let v = b<|>ar.foo();
}
",
            "
fn main() {
    let <|>var_name = bar.foo();
    let v = var_name;
}
",
        );
    }

    #[test]
    fn test_introduce_var_return() {
        check_assist(
            introduce_variable,
            "
fn foo() -> u32 {
    r<|>eturn 2 + 2;
}
",
            "
fn foo() -> u32 {
    let <|>var_name = 2 + 2;
    return var_name;
}
",
        );
    }

    #[test]
    fn test_introduce_var_does_not_add_extra_whitespace() {
        check_assist(
            introduce_variable,
            "
fn foo() -> u32 {

    r<|>eturn 2 + 2;
}
",
            "
fn foo() -> u32 {

    let <|>var_name = 2 + 2;
    return var_name;
}
",
        );

        check_assist(
            introduce_variable,
            "
fn foo() -> u32 {

        r<|>eturn 2 + 2;
}
",
            "
fn foo() -> u32 {

        let <|>var_name = 2 + 2;
        return var_name;
}
",
        );

        check_assist(
            introduce_variable,
            "
fn foo() -> u32 {
    let foo = 1;

    // bar

    r<|>eturn 2 + 2;
}
",
            "
fn foo() -> u32 {
    let foo = 1;

    // bar

    let <|>var_name = 2 + 2;
    return var_name;
}
",
        );
    }

    #[test]
    fn test_introduce_var_break() {
        check_assist(
            introduce_variable,
            "
fn main() {
    let result = loop {
        b<|>reak 2 + 2;
    };
}
",
            "
fn main() {
    let result = loop {
        let <|>var_name = 2 + 2;
        break var_name;
    };
}
",
        );
    }

    #[test]
    fn test_introduce_var_for_cast() {
        check_assist(
            introduce_variable,
            "
fn main() {
    let v = 0f32 a<|>s u32;
}
",
            "
fn main() {
    let <|>var_name = 0f32 as u32;
    let v = var_name;
}
",
        );
    }

    #[test]
    fn test_introduce_var_for_return_not_applicable() {
        check_assist_not_applicable(
            introduce_variable,
            "
fn foo() {
    r<|>eturn;
}
",
        );
    }

    #[test]
    fn test_introduce_var_for_break_not_applicable() {
        check_assist_not_applicable(
            introduce_variable,
            "
fn main() {
    loop {
        b<|>reak;
    };
}
",
        );
    }

    #[test]
    fn test_introduce_var_in_comment_not_applicable() {
        check_assist_not_applicable(
            introduce_variable,
            "
fn main() {
    let x = true;
    let tuple = match x {
        // c<|>omment
        true => (2 + 2, true)
        _ => (0, false)
    };
}
",
        );
    }

    // FIXME: This is not quite correct, but good enough(tm) for the sorting heuristic
    #[test]
    fn introduce_var_target() {
        check_assist_target(
            introduce_variable,
            "
fn foo() -> u32 {
    r<|>eturn 2 + 2;
}
",
            "2 + 2",
        );

        check_assist_range_target(
            introduce_variable,
            "
fn main() {
    let x = true;
    let tuple = match x {
        true => (<|>2 + 2<|>, true)
        _ => (0, false)
    };
}
",
            "2 + 2",
        );
    }
}
<|endoftext|>


<commit_before>use rand::Rng;
use std::fmt;

/// Simple implementation of [Reservoir Sampling](https://en.wikipedia.org/wiki/Reservoir_sampling)
/// with [fast approximation](https://erikerlandson.github.io/blog/2015/11/20/very-fast-reservoir-sampling/)
#[derive(Clone)]
pub struct ReservoirSampling<T, R>
where
    R: Rng,
{
    k: usize,
    rng: R,
    reservoir: Vec<T>,
    i: usize,
    skip_until: usize,
}

impl<T, R> ReservoirSampling<T, R>
where
    R: Rng,
{
    /// Create new reservoir sampler that keeps `k` samples and uses `rng` for its random
    /// decisions.
    pub fn new(k: usize, rng: R) -> ReservoirSampling<T, R> {
        ReservoirSampling {
            k: k,
            rng: rng,
            reservoir: vec![],
            i: 0,
            skip_until: 0,
        }
    }

    /// Number of samples that should be kept.
    pub fn k(&self) -> usize {
        self.k
    }

    /// Read-only copy of the reservoir. Contains at most `k` entries.
    pub fn reservoir(&self) -> &Vec<T> {
        &self.reservoir
    }

    /// Number of data points seen.
    pub fn i(&self) -> usize {
        self.i
    }

    /// Observe new data point.
    pub fn add(&mut self, obj: T) {
        let t = self.k * 4; // TODO: make this a parameter

        if self.i < self.k {
            // initial fill-up
            self.reservoir.push(obj)
        } else if self.i < t {
            // normal reservoir sampling
            let j = self.rng.gen_range::<usize>(0, self.i);
            if j < self.k {
                self.reservoir[j] = obj;
            }
        } else if self.i >= self.skip_until {
            // fast skipping approximation
            let j = self.rng.gen_range::<usize>(0, self.k);
            self.reservoir[j] = obj;

            // calculate next skip
            let p = (self.k as f64) / ((self.i + 1) as f64);
            let u = 1. - self.rng.gen_range::<f64>(0., 1.); // (0.0, 1.0]
            let g = (u.ln() / (1. - p).ln()).floor() as usize;
            self.skip_until = self.i + g;
        }

        self.i += 1;
    }

    /// Checks if reservoir is empty (i.e. no data points where observed)
    pub fn is_empty(&self) -> bool {
        self.i == 0
    }

    /// Clear sampler state. It now behaves like a fresh one (i.e. no data points seen so far).
    pub fn clear(&mut self) {
        self.i = 0;
        self.skip_until = 0;
        self.reservoir.clear();
    }
}

impl<T, R> fmt::Debug for ReservoirSampling<T, R>
where
    R: Rng,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "ReservoirSampling {{ k: {} }}", self.k)
    }
}

impl<T, R> Extend<T> for ReservoirSampling<T, R>
where
    R: Rng,
{
    fn extend<S: IntoIterator<Item = T>>(&mut self, iter: S) {
        for elem in iter {
            self.add(elem);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::ReservoirSampling;
    use rand::{ChaChaRng, SeedableRng};

    #[test]
    fn getter() {
        let rs = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        assert_eq!(rs.k(), 10);
        assert_eq!(rs.i(), 0);
        assert_eq!(rs.reservoir(), &vec![]);
    }

    #[test]
    fn empty() {
        let rs = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        assert!(rs.is_empty());
        assert!(rs.reservoir().is_empty());
    }

    #[test]
    fn add_k() {
        let mut rs = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        for i in 0..10 {
            rs.add(i);
        }
        assert_eq!(rs.i(), 10);
        assert_eq!(rs.reservoir(), &vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
    }

    #[test]
    fn add_parts() {
        let mut rs = ReservoirSampling::<u64, ChaChaRng>::new(100, ChaChaRng::from_seed([0; 32]));
        for _ in 0..1500 {
            rs.add(0);
        }
        for _ in 0..7500 {
            rs.add(1);
        }
        for _ in 0..1000 {
            rs.add(0);
        }
        assert!((rs.reservoir().iter().sum::<u64>() as i64 - 75).abs() < 5);
    }

    #[test]
    fn debug() {
        let rs = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        assert_eq!(format!("{:?}", rs), "ReservoirSampling { k: 10 }");
    }

    #[test]
    fn clear() {
        let mut rs = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        for i in 0..10 {
            rs.add(i);
        }
        rs.clear();
        for i in 10..20 {
            rs.add(i);
        }
        assert_eq!(rs.i(), 10);
        assert_eq!(
            rs.reservoir(),
            &vec![10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
        );
    }

    #[test]
    fn extend() {
        let mut rs = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        rs.extend(0..10);
        assert_eq!(rs.i(), 10);
        assert_eq!(rs.reservoir(), &vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
    }

    #[test]
    fn clone() {
        let mut rs1 = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        for i in 0..10 {
            rs1.add(i);
        }
        let mut rs2 = rs1.clone();
        for _ in 0..1000 {
            rs2.add(99);
        }
        assert_eq!(rs1.i(), 10);
        assert_eq!(rs2.i(), 1010);
        assert_eq!(rs1.reservoir(), &vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
        assert_eq!(
            rs2.reservoir(),
            &vec![99, 99, 99, 99, 99, 99, 99, 99, 99, 99]
        );
    }
}
<commit_msg>reservoidsampling: test asserts<commit_after>use rand::Rng;
use std::fmt;

/// Simple implementation of [Reservoir Sampling](https://en.wikipedia.org/wiki/Reservoir_sampling)
/// with [fast approximation](https://erikerlandson.github.io/blog/2015/11/20/very-fast-reservoir-sampling/)
#[derive(Clone)]
pub struct ReservoirSampling<T, R>
where
    R: Rng,
{
    k: usize,
    rng: R,
    reservoir: Vec<T>,
    i: usize,
    skip_until: usize,
}

impl<T, R> ReservoirSampling<T, R>
where
    R: Rng,
{
    /// Create new reservoir sampler that keeps `k` samples and uses `rng` for its random
    /// decisions.
    pub fn new(k: usize, rng: R) -> ReservoirSampling<T, R> {
        assert!(k > 0, "k must be greater than 0");

        ReservoirSampling {
            k: k,
            rng: rng,
            reservoir: vec![],
            i: 0,
            skip_until: 0,
        }
    }

    /// Number of samples that should be kept.
    pub fn k(&self) -> usize {
        self.k
    }

    /// Read-only copy of the reservoir. Contains at most `k` entries.
    pub fn reservoir(&self) -> &Vec<T> {
        &self.reservoir
    }

    /// Number of data points seen.
    pub fn i(&self) -> usize {
        self.i
    }

    /// Observe new data point.
    pub fn add(&mut self, obj: T) {
        let t = self.k * 4; // TODO: make this a parameter

        if self.i < self.k {
            // initial fill-up
            self.reservoir.push(obj)
        } else if self.i < t {
            // normal reservoir sampling
            let j = self.rng.gen_range::<usize>(0, self.i);
            if j < self.k {
                self.reservoir[j] = obj;
            }
        } else if self.i >= self.skip_until {
            // fast skipping approximation
            let j = self.rng.gen_range::<usize>(0, self.k);
            self.reservoir[j] = obj;

            // calculate next skip
            let p = (self.k as f64) / ((self.i + 1) as f64);
            let u = 1. - self.rng.gen_range::<f64>(0., 1.); // (0.0, 1.0]
            let g = (u.ln() / (1. - p).ln()).floor() as usize;
            self.skip_until = self.i + g;
        }

        self.i += 1;
    }

    /// Checks if reservoir is empty (i.e. no data points where observed)
    pub fn is_empty(&self) -> bool {
        self.i == 0
    }

    /// Clear sampler state. It now behaves like a fresh one (i.e. no data points seen so far).
    pub fn clear(&mut self) {
        self.i = 0;
        self.skip_until = 0;
        self.reservoir.clear();
    }
}

impl<T, R> fmt::Debug for ReservoirSampling<T, R>
where
    R: Rng,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "ReservoirSampling {{ k: {} }}", self.k)
    }
}

impl<T, R> Extend<T> for ReservoirSampling<T, R>
where
    R: Rng,
{
    fn extend<S: IntoIterator<Item = T>>(&mut self, iter: S) {
        for elem in iter {
            self.add(elem);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::ReservoirSampling;
    use rand::{ChaChaRng, SeedableRng};

    #[test]
    #[should_panic(expected = "k must be greater than 0")]
    fn new_panics_k0() {
        ReservoirSampling::<u64, ChaChaRng>::new(0, ChaChaRng::from_seed([0; 32]));
    }

    #[test]
    fn getter() {
        let rs = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        assert_eq!(rs.k(), 10);
        assert_eq!(rs.i(), 0);
        assert_eq!(rs.reservoir(), &vec![]);
    }

    #[test]
    fn empty() {
        let rs = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        assert!(rs.is_empty());
        assert!(rs.reservoir().is_empty());
    }

    #[test]
    fn add_k() {
        let mut rs = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        for i in 0..10 {
            rs.add(i);
        }
        assert_eq!(rs.i(), 10);
        assert_eq!(rs.reservoir(), &vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
    }

    #[test]
    fn add_parts() {
        let mut rs = ReservoirSampling::<u64, ChaChaRng>::new(100, ChaChaRng::from_seed([0; 32]));
        for _ in 0..1500 {
            rs.add(0);
        }
        for _ in 0..7500 {
            rs.add(1);
        }
        for _ in 0..1000 {
            rs.add(0);
        }
        assert!((rs.reservoir().iter().sum::<u64>() as i64 - 75).abs() < 5);
    }

    #[test]
    fn debug() {
        let rs = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        assert_eq!(format!("{:?}", rs), "ReservoirSampling { k: 10 }");
    }

    #[test]
    fn clear() {
        let mut rs = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        for i in 0..10 {
            rs.add(i);
        }
        rs.clear();
        for i in 10..20 {
            rs.add(i);
        }
        assert_eq!(rs.i(), 10);
        assert_eq!(
            rs.reservoir(),
            &vec![10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
        );
    }

    #[test]
    fn extend() {
        let mut rs = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        rs.extend(0..10);
        assert_eq!(rs.i(), 10);
        assert_eq!(rs.reservoir(), &vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
    }

    #[test]
    fn clone() {
        let mut rs1 = ReservoirSampling::<u64, ChaChaRng>::new(10, ChaChaRng::from_seed([0; 32]));
        for i in 0..10 {
            rs1.add(i);
        }
        let mut rs2 = rs1.clone();
        for _ in 0..1000 {
            rs2.add(99);
        }
        assert_eq!(rs1.i(), 10);
        assert_eq!(rs2.i(), 1010);
        assert_eq!(rs1.reservoir(), &vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
        assert_eq!(
            rs2.reservoir(),
            &vec![99, 99, 99, 99, 99, 99, 99, 99, 99, 99]
        );
    }

    #[test]
    fn k1() {
        let mut rs = ReservoirSampling::<u64, ChaChaRng>::new(1, ChaChaRng::from_seed([0; 32]));
        for _ in 0..10 {
            rs.add(1);
        }
        assert_eq!(rs.reservoir(), &vec![1]);
    }
}
<|endoftext|>


<commit_before>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::rc::Rc;
use std::sync::Arc;

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Context;
use argh::FromArgs;
use base::error;
use base::warn;
use base::Event;
use cros_async::sync::Mutex as AsyncMutex;
use cros_async::EventAsync;
use cros_async::Executor;
use data_model::DataInit;
use futures::channel::mpsc;
use futures::future::AbortHandle;
use futures::future::Abortable;
use hypervisor::ProtectionType;
use once_cell::sync::OnceCell;
use sync::Mutex;
use vm_memory::GuestMemory;
use vmm_vhost::message::VhostUserProtocolFeatures;
use vmm_vhost::message::VhostUserVirtioFeatures;

use crate::virtio;
use crate::virtio::copy_config;
use crate::virtio::device_constants::snd::virtio_snd_config;
use crate::virtio::snd::common_backend::async_funcs::handle_ctrl_queue;
use crate::virtio::snd::common_backend::async_funcs::handle_pcm_queue;
use crate::virtio::snd::common_backend::async_funcs::send_pcm_response_worker;
use crate::virtio::snd::common_backend::create_stream_source_generators;
use crate::virtio::snd::common_backend::hardcoded_snd_data;
use crate::virtio::snd::common_backend::hardcoded_virtio_snd_config;
use crate::virtio::snd::common_backend::PcmResponse;
use crate::virtio::snd::common_backend::SndData;
use crate::virtio::snd::common_backend::StreamInfo;
use crate::virtio::snd::common_backend::MAX_QUEUE_NUM;
use crate::virtio::snd::common_backend::MAX_VRING_LEN;
use crate::virtio::snd::parameters::Parameters;
#[cfg(feature = "audio_cras")]
use crate::virtio::snd::parameters::StreamSourceBackend as Backend;
#[cfg(feature = "audio_cras")]
use crate::virtio::snd::sys::StreamSourceBackend as SysBackend;
use crate::virtio::vhost::user::device::handler::sys::Doorbell;
use crate::virtio::vhost::user::device::handler::VhostUserBackend;
use crate::virtio::vhost::user::device::listener::sys::VhostUserListener;
use crate::virtio::vhost::user::device::listener::VhostUserListenerTrait;

static SND_EXECUTOR: OnceCell<Executor> = OnceCell::new();

// Async workers:
// 0 - ctrl
// 1 - event
// 2 - tx
// 3 - rx
const PCM_RESPONSE_WORKER_IDX_OFFSET: usize = 2;
struct SndBackend {
    cfg: virtio_snd_config,
    avail_features: u64,
    acked_features: u64,
    acked_protocol_features: VhostUserProtocolFeatures,
    workers: [Option<AbortHandle>; MAX_QUEUE_NUM],
    response_workers: [Option<AbortHandle>; 2], // tx and rx
    snd_data: Rc<SndData>,
    streams: Rc<AsyncMutex<Vec<AsyncMutex<StreamInfo>>>>,
    tx_send: mpsc::UnboundedSender<PcmResponse>,
    rx_send: mpsc::UnboundedSender<PcmResponse>,
    tx_recv: Option<mpsc::UnboundedReceiver<PcmResponse>>,
    rx_recv: Option<mpsc::UnboundedReceiver<PcmResponse>>,
}

impl SndBackend {
    pub fn new(params: Parameters) -> anyhow::Result<Self> {
        let cfg = hardcoded_virtio_snd_config(&params);
        let avail_features = virtio::base_features(ProtectionType::Unprotected)
            | VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits();

        let snd_data = hardcoded_snd_data(&params);
        let generators = create_stream_source_generators(&params, &snd_data);

        if snd_data.pcm_info_len() != generators.len() {
            error!(
                "snd: expected {} stream source generators, got {}",
                snd_data.pcm_info_len(),
                generators.len(),
            )
        }

        let streams = generators
            .into_iter()
            .map(|generator| AsyncMutex::new(StreamInfo::new(generator)))
            .collect();
        let streams = Rc::new(AsyncMutex::new(streams));

        let (tx_send, tx_recv) = mpsc::unbounded();
        let (rx_send, rx_recv) = mpsc::unbounded();

        Ok(SndBackend {
            cfg,
            avail_features,
            acked_features: 0,
            acked_protocol_features: VhostUserProtocolFeatures::empty(),
            workers: Default::default(),
            response_workers: Default::default(),
            snd_data: Rc::new(snd_data),
            streams,
            tx_send,
            rx_send,
            tx_recv: Some(tx_recv),
            rx_recv: Some(rx_recv),
        })
    }
}

impl VhostUserBackend for SndBackend {
    fn max_queue_num(&self) -> usize {
        return MAX_QUEUE_NUM;
    }

    fn max_vring_len(&self) -> u16 {
        return MAX_VRING_LEN;
    }

    fn features(&self) -> u64 {
        self.avail_features
    }

    fn ack_features(&mut self, value: u64) -> anyhow::Result<()> {
        let unrequested_features = value & !self.avail_features;
        if unrequested_features != 0 {
            bail!("invalid features are given: {:#x}", unrequested_features);
        }

        self.acked_features |= value;

        Ok(())
    }

    fn acked_features(&self) -> u64 {
        self.acked_features
    }

    fn protocol_features(&self) -> VhostUserProtocolFeatures {
        VhostUserProtocolFeatures::CONFIG | VhostUserProtocolFeatures::MQ
    }

    fn ack_protocol_features(&mut self, features: u64) -> anyhow::Result<()> {
        let features = VhostUserProtocolFeatures::from_bits(features)
            .ok_or_else(|| anyhow!("invalid protocol features are given: {:#x}", features))?;
        let supported = self.protocol_features();
        self.acked_protocol_features = features & supported;
        Ok(())
    }

    fn acked_protocol_features(&self) -> u64 {
        self.acked_protocol_features.bits()
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        copy_config(data, 0, self.cfg.as_slice(), offset)
    }

    fn reset(&mut self) {
        for handle in self.workers.iter_mut().filter_map(Option::take) {
            handle.abort();
        }
    }

    fn start_queue(
        &mut self,
        idx: usize,
        mut queue: virtio::Queue,
        mem: GuestMemory,
        doorbell: Arc<Mutex<Doorbell>>,
        kick_evt: Event,
    ) -> anyhow::Result<()> {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            warn!("Starting new queue handler without stopping old handler");
            handle.abort();
        }

        // Safe because the executor is initialized in main() below.
        let ex = SND_EXECUTOR.get().expect("Executor not initialized");

        // Enable any virtqueue features that were negotiated (like VIRTIO_RING_F_EVENT_IDX).
        queue.ack_features(self.acked_features);

        let kick_evt =
            EventAsync::new(kick_evt, ex).context("failed to create EventAsync for kick_evt")?;
        let (handle, registration) = AbortHandle::new_pair();
        match idx {
            0 => {
                // ctrl queue
                let streams = self.streams.clone();
                let snd_data = self.snd_data.clone();
                let tx_send = self.tx_send.clone();
                let rx_send = self.rx_send.clone();
                ex.spawn_local(Abortable::new(
                    async move {
                        handle_ctrl_queue(
                            ex, &mem, &streams, &*snd_data, queue, kick_evt, &doorbell, tx_send,
                            rx_send,
                        )
                        .await
                    },
                    registration,
                ))
                .detach();
            }
            1 => {} // TODO(woodychow): Add event queue support
            2 | 3 => {
                let (send, recv) = if idx == 2 {
                    (self.tx_send.clone(), self.tx_recv.take())
                } else {
                    (self.rx_send.clone(), self.rx_recv.take())
                };
                let mut recv = recv.ok_or_else(|| anyhow!("queue restart is not supported"))?;
                let queue = Rc::new(AsyncMutex::new(queue));
                let queue2 = Rc::clone(&queue);
                let mem = Rc::new(mem);
                let mem2 = Rc::clone(&mem);
                let streams = Rc::clone(&self.streams);
                ex.spawn_local(Abortable::new(
                    async move { handle_pcm_queue(&*mem, &streams, send, &queue, kick_evt).await },
                    registration,
                ))
                .detach();

                let (handle2, registration2) = AbortHandle::new_pair();

                ex.spawn_local(Abortable::new(
                    async move {
                        send_pcm_response_worker(&*mem2, &queue2, &doorbell, &mut recv).await
                    },
                    registration2,
                ))
                .detach();

                self.response_workers[idx - PCM_RESPONSE_WORKER_IDX_OFFSET] = Some(handle2);
            }
            _ => bail!("attempted to start unknown queue: {}", idx),
        }

        self.workers[idx] = Some(handle);
        Ok(())
    }

    fn stop_queue(&mut self, idx: usize) {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            handle.abort();
        }
        if idx == 2 || idx == 3 {
            if let Some(handle) = self
                .response_workers
                .get_mut(idx - PCM_RESPONSE_WORKER_IDX_OFFSET)
                .and_then(Option::take)
            {
                handle.abort();
            }
        }
    }
}

// TODO(b/241489181): Remove once cras-snd calls are changed to snd.
#[cfg(feature = "audio_cras")]
#[derive(FromArgs)]
#[argh(subcommand, name = "cras-snd")]
/// CRAS snd device
pub struct CrasOptions {
    #[argh(option, arg_name = "PATH")]
    /// path to a socket
    socket: String,
    #[argh(option, arg_name = "CONFIG")]
    /// comma separated key=value pairs for setting up cras snd devices.
    /// Possible key values:
    /// capture - Enable audio capture. Default to false.
    /// client_type - Set specific client type for cras backend.
    /// num_output_devices - Set number of output PCM devices.
    /// num_input_devices - Set number of input PCM devices.
    /// num_output_streams - Set number of output PCM streams per device.
    /// num_input_streams - Set number of input PCM streams per device.
    /// Example: [capture=true,client=crosvm,socket=unified,
    /// num_output_devices=1,num_input_devices=1,num_output_streams=1,num_input_streams=1]
    config: Option<String>,
}

#[derive(FromArgs)]
#[argh(subcommand, name = "snd")]
/// Snd device
pub struct Options {
    #[argh(option, arg_name = "PATH")]
    /// path to a socket
    socket: String,
    #[argh(option, arg_name = "CONFIG")]
    /// comma separated key=value pairs for setting up cras snd devices.
    /// Possible key values:
    /// capture - Enable audio capture. Default to false.
    /// backend(null,[cras]) - Which backend to use for vhost-snd.
    /// client_type - Set specific client type for cras backend.
    /// socket_type - Set socket type for cras backend.
    /// num_output_devices - Set number of output PCM devices.
    /// num_input_devices - Set number of input PCM devices.
    /// num_output_streams - Set number of output PCM streams per device.
    /// num_input_streams - Set number of input PCM streams per device.
    /// Example: [capture=true,backend=BACKEND,
    /// num_output_devices=1,num_input_devices=1,num_output_streams=1,num_input_streams=1]
    config: Option<String>,
}

// TODO(b/241489181): Remove once cras-snd calls are changed to snd.
#[cfg(feature = "audio_cras")]
pub fn run_cras_snd_device(opts: CrasOptions) -> anyhow::Result<()> {
    let params = opts
        .config
        .unwrap_or("".to_string())
        .parse::<Parameters>()?;
    let params = Parameters {
        backend: Backend::Sys(SysBackend::CRAS),
        ..params
    };

    run(params, opts.socket)
}

/// Starts a vhost-user snd device.
/// Returns an error if the given `args` is invalid or the device fails to run.
pub fn run_snd_device(opts: Options) -> anyhow::Result<()> {
    let params = opts
        .config
        .unwrap_or("".to_string())
        .parse::<Parameters>()?;

    run(params, opts.socket)
}

fn run(params: Parameters, socket_path: String) -> anyhow::Result<()> {
    let snd_device = Box::new(SndBackend::new(params)?);

    // Child, we can continue by spawning the executor and set up the device
    let ex = Executor::new().context("Failed to create executor")?;

    let _ = SND_EXECUTOR.set(ex.clone());

    let listener = VhostUserListener::new_socket(&socket_path, None)?;
    // run_until() returns an Result<Result<..>> which the ? operator lets us flatten.
    ex.run_until(listener.run_backend(snd_device, &ex))?
}
<commit_msg>devices: vhost: user: snd:: Implement VVU snd<commit_after>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::rc::Rc;
use std::sync::Arc;

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Context;
use argh::FromArgs;
use base::error;
use base::warn;
use base::Event;
use cros_async::sync::Mutex as AsyncMutex;
use cros_async::EventAsync;
use cros_async::Executor;
use data_model::DataInit;
use futures::channel::mpsc;
use futures::future::AbortHandle;
use futures::future::Abortable;
use hypervisor::ProtectionType;
use once_cell::sync::OnceCell;
use sync::Mutex;
use vm_memory::GuestMemory;
use vmm_vhost::message::VhostUserProtocolFeatures;
use vmm_vhost::message::VhostUserVirtioFeatures;

use crate::virtio;
use crate::virtio::copy_config;
use crate::virtio::device_constants::snd::virtio_snd_config;
use crate::virtio::snd::common_backend::async_funcs::handle_ctrl_queue;
use crate::virtio::snd::common_backend::async_funcs::handle_pcm_queue;
use crate::virtio::snd::common_backend::async_funcs::send_pcm_response_worker;
use crate::virtio::snd::common_backend::create_stream_source_generators;
use crate::virtio::snd::common_backend::hardcoded_snd_data;
use crate::virtio::snd::common_backend::hardcoded_virtio_snd_config;
use crate::virtio::snd::common_backend::PcmResponse;
use crate::virtio::snd::common_backend::SndData;
use crate::virtio::snd::common_backend::StreamInfo;
use crate::virtio::snd::common_backend::MAX_QUEUE_NUM;
use crate::virtio::snd::common_backend::MAX_VRING_LEN;
use crate::virtio::snd::parameters::Parameters;
#[cfg(feature = "audio_cras")]
use crate::virtio::snd::parameters::StreamSourceBackend as Backend;
#[cfg(feature = "audio_cras")]
use crate::virtio::snd::sys::StreamSourceBackend as SysBackend;
use crate::virtio::vhost::user::device::handler::sys::Doorbell;
use crate::virtio::vhost::user::device::handler::VhostUserBackend;
use crate::virtio::vhost::user::device::listener::sys::VhostUserListener;
use crate::virtio::vhost::user::device::listener::VhostUserListenerTrait;

static SND_EXECUTOR: OnceCell<Executor> = OnceCell::new();

// Async workers:
// 0 - ctrl
// 1 - event
// 2 - tx
// 3 - rx
const PCM_RESPONSE_WORKER_IDX_OFFSET: usize = 2;
struct SndBackend {
    cfg: virtio_snd_config,
    avail_features: u64,
    acked_features: u64,
    acked_protocol_features: VhostUserProtocolFeatures,
    workers: [Option<AbortHandle>; MAX_QUEUE_NUM],
    response_workers: [Option<AbortHandle>; 2], // tx and rx
    snd_data: Rc<SndData>,
    streams: Rc<AsyncMutex<Vec<AsyncMutex<StreamInfo>>>>,
    tx_send: mpsc::UnboundedSender<PcmResponse>,
    rx_send: mpsc::UnboundedSender<PcmResponse>,
    tx_recv: Option<mpsc::UnboundedReceiver<PcmResponse>>,
    rx_recv: Option<mpsc::UnboundedReceiver<PcmResponse>>,
}

impl SndBackend {
    pub fn new(params: Parameters) -> anyhow::Result<Self> {
        let cfg = hardcoded_virtio_snd_config(&params);
        let avail_features = virtio::base_features(ProtectionType::Unprotected)
            | VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits();

        let snd_data = hardcoded_snd_data(&params);
        let generators = create_stream_source_generators(&params, &snd_data);

        if snd_data.pcm_info_len() != generators.len() {
            error!(
                "snd: expected {} stream source generators, got {}",
                snd_data.pcm_info_len(),
                generators.len(),
            )
        }

        let streams = generators
            .into_iter()
            .map(|generator| AsyncMutex::new(StreamInfo::new(generator)))
            .collect();
        let streams = Rc::new(AsyncMutex::new(streams));

        let (tx_send, tx_recv) = mpsc::unbounded();
        let (rx_send, rx_recv) = mpsc::unbounded();

        Ok(SndBackend {
            cfg,
            avail_features,
            acked_features: 0,
            acked_protocol_features: VhostUserProtocolFeatures::empty(),
            workers: Default::default(),
            response_workers: Default::default(),
            snd_data: Rc::new(snd_data),
            streams,
            tx_send,
            rx_send,
            tx_recv: Some(tx_recv),
            rx_recv: Some(rx_recv),
        })
    }
}

impl VhostUserBackend for SndBackend {
    fn max_queue_num(&self) -> usize {
        return MAX_QUEUE_NUM;
    }

    fn max_vring_len(&self) -> u16 {
        return MAX_VRING_LEN;
    }

    fn features(&self) -> u64 {
        self.avail_features
    }

    fn ack_features(&mut self, value: u64) -> anyhow::Result<()> {
        let unrequested_features = value & !self.avail_features;
        if unrequested_features != 0 {
            bail!("invalid features are given: {:#x}", unrequested_features);
        }

        self.acked_features |= value;

        Ok(())
    }

    fn acked_features(&self) -> u64 {
        self.acked_features
    }

    fn protocol_features(&self) -> VhostUserProtocolFeatures {
        VhostUserProtocolFeatures::CONFIG | VhostUserProtocolFeatures::MQ
    }

    fn ack_protocol_features(&mut self, features: u64) -> anyhow::Result<()> {
        let features = VhostUserProtocolFeatures::from_bits(features)
            .ok_or_else(|| anyhow!("invalid protocol features are given: {:#x}", features))?;
        let supported = self.protocol_features();
        self.acked_protocol_features = features & supported;
        Ok(())
    }

    fn acked_protocol_features(&self) -> u64 {
        self.acked_protocol_features.bits()
    }

    fn read_config(&self, offset: u64, data: &mut [u8]) {
        copy_config(data, 0, self.cfg.as_slice(), offset)
    }

    fn reset(&mut self) {
        for handle in self.workers.iter_mut().filter_map(Option::take) {
            handle.abort();
        }
    }

    fn start_queue(
        &mut self,
        idx: usize,
        mut queue: virtio::Queue,
        mem: GuestMemory,
        doorbell: Arc<Mutex<Doorbell>>,
        kick_evt: Event,
    ) -> anyhow::Result<()> {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            warn!("Starting new queue handler without stopping old handler");
            handle.abort();
        }

        // Safe because the executor is initialized in main() below.
        let ex = SND_EXECUTOR.get().expect("Executor not initialized");

        // Enable any virtqueue features that were negotiated (like VIRTIO_RING_F_EVENT_IDX).
        queue.ack_features(self.acked_features);

        let kick_evt =
            EventAsync::new(kick_evt, ex).context("failed to create EventAsync for kick_evt")?;
        let (handle, registration) = AbortHandle::new_pair();
        match idx {
            0 => {
                // ctrl queue
                let streams = self.streams.clone();
                let snd_data = self.snd_data.clone();
                let tx_send = self.tx_send.clone();
                let rx_send = self.rx_send.clone();
                ex.spawn_local(Abortable::new(
                    async move {
                        handle_ctrl_queue(
                            ex, &mem, &streams, &*snd_data, queue, kick_evt, &doorbell, tx_send,
                            rx_send,
                        )
                        .await
                    },
                    registration,
                ))
                .detach();
            }
            1 => {} // TODO(woodychow): Add event queue support
            2 | 3 => {
                let (send, recv) = if idx == 2 {
                    (self.tx_send.clone(), self.tx_recv.take())
                } else {
                    (self.rx_send.clone(), self.rx_recv.take())
                };
                let mut recv = recv.ok_or_else(|| anyhow!("queue restart is not supported"))?;
                let queue = Rc::new(AsyncMutex::new(queue));
                let queue2 = Rc::clone(&queue);
                let mem = Rc::new(mem);
                let mem2 = Rc::clone(&mem);
                let streams = Rc::clone(&self.streams);
                ex.spawn_local(Abortable::new(
                    async move { handle_pcm_queue(&*mem, &streams, send, &queue, kick_evt).await },
                    registration,
                ))
                .detach();

                let (handle2, registration2) = AbortHandle::new_pair();

                ex.spawn_local(Abortable::new(
                    async move {
                        send_pcm_response_worker(&*mem2, &queue2, &doorbell, &mut recv).await
                    },
                    registration2,
                ))
                .detach();

                self.response_workers[idx - PCM_RESPONSE_WORKER_IDX_OFFSET] = Some(handle2);
            }
            _ => bail!("attempted to start unknown queue: {}", idx),
        }

        self.workers[idx] = Some(handle);
        Ok(())
    }

    fn stop_queue(&mut self, idx: usize) {
        if let Some(handle) = self.workers.get_mut(idx).and_then(Option::take) {
            handle.abort();
        }
        if idx == 2 || idx == 3 {
            if let Some(handle) = self
                .response_workers
                .get_mut(idx - PCM_RESPONSE_WORKER_IDX_OFFSET)
                .and_then(Option::take)
            {
                handle.abort();
            }
        }
    }
}

// TODO(b/241489181): Remove once cras-snd calls are changed to snd.
#[cfg(feature = "audio_cras")]
#[derive(FromArgs)]
#[argh(subcommand, name = "cras-snd")]
/// CRAS snd device
pub struct CrasOptions {
    #[argh(option, arg_name = "PATH")]
    /// path to a socket
    socket: Option<String>,
    #[argh(option, arg_name = "CONFIG")]
    /// VFIO-PCI device name (e.g. '0000:00:07.0')
    vfio: Option<String>,
    #[argh(option, arg_name = "CONFIG")]
    /// comma separated key=value pairs for setting up cras snd devices.
    /// Possible key values:
    /// capture - Enable audio capture. Default to false.
    /// client_type - Set specific client type for cras backend.
    /// num_output_devices - Set number of output PCM devices.
    /// num_input_devices - Set number of input PCM devices.
    /// num_output_streams - Set number of output PCM streams per device.
    /// num_input_streams - Set number of input PCM streams per device.
    /// Example: [capture=true,client=crosvm,socket=unified,
    /// num_output_devices=1,num_input_devices=1,num_output_streams=1,num_input_streams=1]
    config: Option<String>,
}

#[derive(FromArgs)]
#[argh(subcommand, name = "snd")]
/// Snd device
pub struct Options {
    #[argh(option, arg_name = "PATH")]
    /// path to bind a listening vhost-user socket
    socket: Option<String>,
    #[argh(option, arg_name = "STRING")]
    /// VFIO-PCI device name (e.g. '0000:00:07.0')
    vfio: Option<String>,
    #[argh(option, arg_name = "CONFIG")]
    /// comma separated key=value pairs for setting up cras snd devices.
    /// Possible key values:
    /// capture - Enable audio capture. Default to false.
    /// backend(null,[cras]) - Which backend to use for vhost-snd.
    /// client_type - Set specific client type for cras backend.
    /// socket_type - Set socket type for cras backend.
    /// num_output_devices - Set number of output PCM devices.
    /// num_input_devices - Set number of input PCM devices.
    /// num_output_streams - Set number of output PCM streams per device.
    /// num_input_streams - Set number of input PCM streams per device.
    /// Example: [capture=true,backend=BACKEND,
    /// num_output_devices=1,num_input_devices=1,num_output_streams=1,num_input_streams=1]
    config: Option<String>,
}

// TODO(b/241489181): Remove once cras-snd calls are changed to snd.
#[cfg(feature = "audio_cras")]
pub fn run_cras_snd_device(opts: CrasOptions) -> anyhow::Result<()> {
    let params = opts
        .config
        .unwrap_or("".to_string())
        .parse::<Parameters>()?;
    let params = Parameters {
        backend: Backend::Sys(SysBackend::CRAS),
        ..params
    };

    run(params, opts.socket, opts.vfio)
}

/// Starts a vhost-user snd device.
/// Returns an error if the given `args` is invalid or the device fails to run.
pub fn run_snd_device(opts: Options) -> anyhow::Result<()> {
    let params = opts
        .config
        .unwrap_or("".to_string())
        .parse::<Parameters>()?;

    run(params, opts.socket, opts.vfio)
}

fn run(
    params: Parameters,
    socket_path: Option<String>,
    vfio: Option<String>,
) -> anyhow::Result<()> {
    let snd_device = Box::new(SndBackend::new(params)?);

    // Child, we can continue by spawning the executor and set up the device
    let ex = Executor::new().context("Failed to create executor")?;

    let _ = SND_EXECUTOR.set(ex.clone());

    let listener = VhostUserListener::new_from_socket_or_vfio(
        &socket_path,
        &vfio,
        snd_device.max_queue_num(),
        None,
    )?;
    // run_until() returns an Result<Result<..>> which the ? operator lets us flatten.
    ex.run_until(listener.run_backend(snd_device, &ex))?
}
<|endoftext|>