<commit_before>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! A definition of `StorageKey` trait and implementations for common types.

use byteorder::{ByteOrder, BigEndian};
use crypto::{Hash, PublicKey, HASH_SIZE, PUBLIC_KEY_LENGTH};

/// A type that can be (de)serialized as a key in the blockchain storage.
///
/// Since internally the keys are sorted in a serialized form, the big-endian encoding is used.
///
/// # Examples
///
/// ```
/// # extern crate exonum;
/// # extern crate byteorder;
/// use std::mem;
/// use exonum::storage::StorageKey;
/// use byteorder::{LittleEndian, ByteOrder};
///
/// struct Key {
///     a: i16,
///     b: u32,
/// }
///
/// impl StorageKey for Key {
///     fn size(&self) -> usize {
///         mem::size_of_val(&self.a) + mem::size_of_val(&self.b)
///     }
///
///     fn write(&self, buffer: &mut [u8]) {
///         LittleEndian::write_i16(&mut buffer[0..2], self.a);
///         LittleEndian::write_u32(&mut buffer[2..6], self.b);
///     }
///
///     fn read(buffer: &[u8]) -> Self {
///         let a = LittleEndian::read_i16(&buffer[0..2]);
///         let b = LittleEndian::read_u32(&buffer[2..6]);
///         Key { a, b }
///     }
/// }
/// # fn main() {}
/// ```
pub trait StorageKey {
    /// Returns the size of the serialized key in bytes.
    fn size(&self) -> usize;

    /// Serializes the key into the specified buffer of bytes.
    ///
    /// The caller must guarantee that the size of the buffer is equal to the precalculated size
    /// of the serialized key.
    // TODO: should be unsafe (ECR-174)?
    fn write(&self, buffer: &mut [u8]);

    /// Deserializes the key from the specified buffer of bytes.
    // TODO: should be unsafe (ECR-174)?
    fn read(buffer: &[u8]) -> Self;
}

/// No-op implementation.
impl StorageKey for () {
    fn size(&self) -> usize {
        0
    }

    fn write(&self, _buffer: &mut [u8]) {
        // no-op
    }

    fn read(_buffer: &[u8]) -> Self {
        ()
    }
}

impl StorageKey for u8 {
    fn size(&self) -> usize {
        1
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer[0] = *self
    }

    fn read(buffer: &[u8]) -> Self {
        buffer[0]
    }
}

/// Uses big-endian encoding.
impl StorageKey for u16 {
    fn size(&self) -> usize {
        2
    }

    fn write(&self, buffer: &mut [u8]) {
        BigEndian::write_u16(buffer, *self)
    }

    fn read(buffer: &[u8]) -> Self {
        BigEndian::read_u16(buffer)
    }
}

/// Uses big-endian encoding.
impl StorageKey for u32 {
    fn size(&self) -> usize {
        4
    }

    fn write(&self, buffer: &mut [u8]) {
        BigEndian::write_u32(buffer, *self)
    }

    fn read(buffer: &[u8]) -> Self {
        BigEndian::read_u32(buffer)
    }
}

/// Uses big-endian encoding.
impl StorageKey for u64 {
    fn size(&self) -> usize {
        8
    }

    fn write(&self, buffer: &mut [u8]) {
        BigEndian::write_u64(buffer, *self)
    }

    fn read(buffer: &[u8]) -> Self {
        BigEndian::read_u64(buffer)
    }
}

/// **Not sorted in the natural order.**
impl StorageKey for i8 {
    fn size(&self) -> usize {
        1
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer[0] = *self as u8
    }

    fn read(buffer: &[u8]) -> Self {
        buffer[0] as i8
    }
}

/// Uses big-endian encoding. **Not sorted in the natural order.**
impl StorageKey for i16 {
    fn size(&self) -> usize {
        2
    }

    fn write(&self, buffer: &mut [u8]) {
        BigEndian::write_i16(buffer, *self)
    }

    fn read(buffer: &[u8]) -> Self {
        BigEndian::read_i16(buffer)
    }
}

/// Uses big-endian encoding. **Not sorted in the natural order.**
impl StorageKey for i32 {
    fn size(&self) -> usize {
        4
    }

    fn write(&self, buffer: &mut [u8]) {
        BigEndian::write_i32(buffer, *self)
    }

    fn read(buffer: &[u8]) -> Self {
        BigEndian::read_i32(buffer)
    }
}

/// Uses big-endian encoding. **Not sorted in the natural order.**
impl StorageKey for i64 {
    fn size(&self) -> usize {
        8
    }

    fn write(&self, buffer: &mut [u8]) {
        BigEndian::write_i64(buffer, *self)
    }

    fn read(buffer: &[u8]) -> Self {
        BigEndian::read_i64(buffer)
    }
}

impl StorageKey for Hash {
    fn size(&self) -> usize {
        HASH_SIZE
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer.copy_from_slice(self.as_ref())
    }

    fn read(buffer: &[u8]) -> Self {
        Hash::from_slice(buffer).unwrap()
    }
}

impl StorageKey for PublicKey {
    fn size(&self) -> usize {
        PUBLIC_KEY_LENGTH
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer.copy_from_slice(self.as_ref())
    }

    fn read(buffer: &[u8]) -> Self {
        PublicKey::from_slice(buffer).unwrap()
    }
}

impl StorageKey for Vec<u8> {
    fn size(&self) -> usize {
        self.len()
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer.copy_from_slice(self)
    }

    fn read(buffer: &[u8]) -> Self {
        buffer.to_vec()
    }
}

/// Uses UTF-8 string serialization.
impl StorageKey for String {
    fn size(&self) -> usize {
        self.len()
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer.copy_from_slice(self.as_bytes())
    }

    fn read(buffer: &[u8]) -> Self {
        unsafe { ::std::str::from_utf8_unchecked(buffer).to_string() }
    }
}
<commit_msg>Rewrite `StorageKey` example<commit_after>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! A definition of `StorageKey` trait and implementations for common types.

use byteorder::{ByteOrder, BigEndian};
use crypto::{Hash, PublicKey, HASH_SIZE, PUBLIC_KEY_LENGTH};

/// A type that can be (de)serialized as a key in the blockchain storage.
///
/// Since keys are sorted in a serialized form, the big-endian encoding should be used
/// with unsigned integer types. Note however that the big-endian encoding
/// will **not** sort signed integer types in the natural order; a possible solution is
/// mapping the type to a corresponding unsigned one as shown in the example below.
///
/// # Examples
///
/// ```
/// # extern crate exonum;
/// # extern crate byteorder;
/// use std::mem;
/// use exonum::storage::StorageKey;
/// use byteorder::{BigEndian, ByteOrder};
///
/// struct Key {
///     a: i16,
///     b: u32,
/// }
///
/// impl StorageKey for Key {
///     fn size(&self) -> usize {
///         mem::size_of_val(&self.a) + mem::size_of_val(&self.b)
///     }
///
///     fn write(&self, buffer: &mut [u8]) {
///         // Maps `a` to the `u16` range in the natural order:
///         // -32768 -> 0, -32767 -> 1, ..., 32767 -> 65535
///         let mapped_a = self.a.wrapping_add(i16::min_value()) as u16;
///         BigEndian::write_u16(&mut buffer[0..2], mapped_a);
///         BigEndian::write_u32(&mut buffer[2..6], self.b);
///     }
///
///     fn read(buffer: &[u8]) -> Self {
///         let mapped_a = BigEndian::read_u16(&buffer[0..2]);
///         let a = mapped_a.wrapping_add(i16::min_value() as u16) as i16;
///         let b = BigEndian::read_u32(&buffer[2..6]);
///         Key { a, b }
///     }
/// }
/// # fn main() {
/// # // Check the natural ordering of keys
/// # let (mut x, mut y) = (vec![0u8; 6], vec![0u8; 6]);
/// # Key { a: -1, b: 2 }.write(&mut x);
/// # Key { a: 1, b: 513 }.write(&mut y);
/// # assert!(x < y);
/// # // Check the roundtrip
/// # let key = Key::read(&x);
/// # assert_eq!(key.a, -1);
/// # assert_eq!(key.b, 2);
/// # }
/// ```
pub trait StorageKey {
    /// Returns the size of the serialized key in bytes.
    fn size(&self) -> usize;

    /// Serializes the key into the specified buffer of bytes.
    ///
    /// The caller must guarantee that the size of the buffer is equal to the precalculated size
    /// of the serialized key.
    // TODO: should be unsafe (ECR-174)?
    fn write(&self, buffer: &mut [u8]);

    /// Deserializes the key from the specified buffer of bytes.
    // TODO: should be unsafe (ECR-174)?
    fn read(buffer: &[u8]) -> Self;
}

/// No-op implementation.
impl StorageKey for () {
    fn size(&self) -> usize {
        0
    }

    fn write(&self, _buffer: &mut [u8]) {
        // no-op
    }

    fn read(_buffer: &[u8]) -> Self {
        ()
    }
}

impl StorageKey for u8 {
    fn size(&self) -> usize {
        1
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer[0] = *self
    }

    fn read(buffer: &[u8]) -> Self {
        buffer[0]
    }
}

/// Uses big-endian encoding.
impl StorageKey for u16 {
    fn size(&self) -> usize {
        2
    }

    fn write(&self, buffer: &mut [u8]) {
        BigEndian::write_u16(buffer, *self)
    }

    fn read(buffer: &[u8]) -> Self {
        BigEndian::read_u16(buffer)
    }
}

/// Uses big-endian encoding.
impl StorageKey for u32 {
    fn size(&self) -> usize {
        4
    }

    fn write(&self, buffer: &mut [u8]) {
        BigEndian::write_u32(buffer, *self)
    }

    fn read(buffer: &[u8]) -> Self {
        BigEndian::read_u32(buffer)
    }
}

/// Uses big-endian encoding.
impl StorageKey for u64 {
    fn size(&self) -> usize {
        8
    }

    fn write(&self, buffer: &mut [u8]) {
        BigEndian::write_u64(buffer, *self)
    }

    fn read(buffer: &[u8]) -> Self {
        BigEndian::read_u64(buffer)
    }
}

/// **Not sorted in the natural order.**
impl StorageKey for i8 {
    fn size(&self) -> usize {
        1
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer[0] = *self as u8
    }

    fn read(buffer: &[u8]) -> Self {
        buffer[0] as i8
    }
}

/// Uses big-endian encoding. **Not sorted in the natural order.**
impl StorageKey for i16 {
    fn size(&self) -> usize {
        2
    }

    fn write(&self, buffer: &mut [u8]) {
        BigEndian::write_i16(buffer, *self)
    }

    fn read(buffer: &[u8]) -> Self {
        BigEndian::read_i16(buffer)
    }
}

/// Uses big-endian encoding. **Not sorted in the natural order.**
impl StorageKey for i32 {
    fn size(&self) -> usize {
        4
    }

    fn write(&self, buffer: &mut [u8]) {
        BigEndian::write_i32(buffer, *self)
    }

    fn read(buffer: &[u8]) -> Self {
        BigEndian::read_i32(buffer)
    }
}

/// Uses big-endian encoding. **Not sorted in the natural order.**
impl StorageKey for i64 {
    fn size(&self) -> usize {
        8
    }

    fn write(&self, buffer: &mut [u8]) {
        BigEndian::write_i64(buffer, *self)
    }

    fn read(buffer: &[u8]) -> Self {
        BigEndian::read_i64(buffer)
    }
}

impl StorageKey for Hash {
    fn size(&self) -> usize {
        HASH_SIZE
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer.copy_from_slice(self.as_ref())
    }

    fn read(buffer: &[u8]) -> Self {
        Hash::from_slice(buffer).unwrap()
    }
}

impl StorageKey for PublicKey {
    fn size(&self) -> usize {
        PUBLIC_KEY_LENGTH
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer.copy_from_slice(self.as_ref())
    }

    fn read(buffer: &[u8]) -> Self {
        PublicKey::from_slice(buffer).unwrap()
    }
}

impl StorageKey for Vec<u8> {
    fn size(&self) -> usize {
        self.len()
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer.copy_from_slice(self)
    }

    fn read(buffer: &[u8]) -> Self {
        buffer.to_vec()
    }
}

/// Uses UTF-8 string serialization.
impl StorageKey for String {
    fn size(&self) -> usize {
        self.len()
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer.copy_from_slice(self.as_bytes())
    }

    fn read(buffer: &[u8]) -> Self {
        unsafe { ::std::str::from_utf8_unchecked(buffer).to_string() }
    }
}
<|endoftext|>


<commit_before>//! A pass that simplifies branches when their condition is known.

use crate::transform::MirPass;
use rustc_middle::mir::*;
use rustc_middle::ty::TyCtxt;

use std::borrow::Cow;

pub struct SimplifyBranches {
    label: String,
}

impl SimplifyBranches {
    pub fn new(label: &str) -> Self {
        SimplifyBranches { label: format!("SimplifyBranches-{}", label) }
    }
}

impl<'tcx> MirPass<'tcx> for SimplifyBranches {
    fn name(&self) -> Cow<'_, str> {
        Cow::Borrowed(&self.label)
    }

    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        let param_env = tcx.param_env(body.source.def_id());
        for block in body.basic_blocks_mut() {
            let terminator = block.terminator_mut();
            terminator.kind = match terminator.kind {
                TerminatorKind::SwitchInt {
                    discr: Operand::Constant(ref c),
                    switch_ty,
                    ref targets,
                    ..
                } => {
                    let constant = c.literal.try_eval_bits(tcx, param_env, switch_ty);
                    if let Some(constant) = constant {
                        let otherwise = targets.otherwise();
                        let mut ret = TerminatorKind::Goto { target: otherwise };
                        for (v, t) in targets.iter() {
                            if v == constant {
                                ret = TerminatorKind::Goto { target: t };
                                break;
                            }
                        }
                        ret
                    } else {
                        continue;
                    }
                }
                TerminatorKind::Assert {
                    target, cond: Operand::Constant(ref c), expected, ..
                } if (c.literal.try_eval_bool(tcx, param_env) == Some(true)) == expected => {
                    TerminatorKind::Goto { target }
                }
                TerminatorKind::FalseEdge { real_target, .. } => {
                    TerminatorKind::Goto { target: real_target }
                }
                TerminatorKind::FalseUnwind { real_target, .. } => {
                    TerminatorKind::Goto { target: real_target }
                }
                _ => continue,
            };
        }
    }
}
<commit_msg>Rollup merge of #78198 - tmiasko:assert, r=davidtwco<commit_after>//! A pass that simplifies branches when their condition is known.

use crate::transform::MirPass;
use rustc_middle::mir::*;
use rustc_middle::ty::TyCtxt;

use std::borrow::Cow;

pub struct SimplifyBranches {
    label: String,
}

impl SimplifyBranches {
    pub fn new(label: &str) -> Self {
        SimplifyBranches { label: format!("SimplifyBranches-{}", label) }
    }
}

impl<'tcx> MirPass<'tcx> for SimplifyBranches {
    fn name(&self) -> Cow<'_, str> {
        Cow::Borrowed(&self.label)
    }

    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        let param_env = tcx.param_env(body.source.def_id());
        for block in body.basic_blocks_mut() {
            let terminator = block.terminator_mut();
            terminator.kind = match terminator.kind {
                TerminatorKind::SwitchInt {
                    discr: Operand::Constant(ref c),
                    switch_ty,
                    ref targets,
                    ..
                } => {
                    let constant = c.literal.try_eval_bits(tcx, param_env, switch_ty);
                    if let Some(constant) = constant {
                        let otherwise = targets.otherwise();
                        let mut ret = TerminatorKind::Goto { target: otherwise };
                        for (v, t) in targets.iter() {
                            if v == constant {
                                ret = TerminatorKind::Goto { target: t };
                                break;
                            }
                        }
                        ret
                    } else {
                        continue;
                    }
                }
                TerminatorKind::Assert {
                    target, cond: Operand::Constant(ref c), expected, ..
                } => match c.literal.try_eval_bool(tcx, param_env) {
                    Some(v) if v == expected => TerminatorKind::Goto { target },
                    _ => continue,
                },
                TerminatorKind::FalseEdge { real_target, .. } => {
                    TerminatorKind::Goto { target: real_target }
                }
                TerminatorKind::FalseUnwind { real_target, .. } => {
                    TerminatorKind::Goto { target: real_target }
                }
                _ => continue,
            };
        }
    }
}
<|endoftext|>


<commit_before>use core::slice;
use core::str;
use collections::range::RangeArgument;
use core::cmp::{max, min};

/// Bounded slice abstraction
///
/// # Code Migration
///
/// `foo[a..b]` => `foo.get_slice(a..b)`
///
/// `foo[a..]` => `foo.get_slice(a..)`
///
/// `foo[..b]` => `foo.get_slice(..b)`
///
pub trait GetSlice {
    fn get_slice<T: RangeArgument<usize>>(&self, a: T) -> &Self;
}

impl GetSlice for str {
    fn get_slice<T: RangeArgument<usize>>(&self, a: T) -> &Self {
        let start = min(a.start().map(|&x| x).unwrap_or(self.len() - 1), self.len() - 1);
        let end = min(a.end().map(|&x| x).unwrap_or(self.len() - 1), self.len() - 1);

        if start <= end {
            &self[start..end + 1]
        } else {
            ""
        }
    }
}

impl<T> GetSlice for [T] {
    fn get_slice<U: RangeArgument<usize>>(&self, a: U) -> &Self {
        let start = min(a.start().map(|&x| x).unwrap_or(self.len() - 1), self.len() - 1);
        let end = min(a.end().map(|&x| x).unwrap_or(self.len() - 1), self.len() - 1);

        if start <= end {
            &self[start..end + 1]
        } else {
            &[]
        }
    }
}
<commit_msg>aaaargh! bug bugs everywhere<commit_after>use core::slice;
use core::str;
use collections::range::RangeArgument;
use core::ops::Range;
use core::cmp::{max, min};

/// Bounded slice abstraction
///
/// # Code Migration
///
/// `foo[a..b]` => `foo.get_slice(a..b)`
///
/// `foo[a..]` => `foo.get_slice(a..)`
///
/// `foo[..b]` => `foo.get_slice(..b)`
///
pub trait GetSlice {
    fn get_slice<T: RangeArgument<usize>>(&self, a: T) -> &Self;
    fn get_slice_mut<T: RangeArgument<usize>>(&mut self, a: T) -> &mut Self;
}

fn bound<T: RangeArgument<usize>>(len: usize, a: T) -> Range<usize> {
    let start = min(a.start().map(|&x| x).unwrap_or(0), len - 1);
    let end = min(a.end().map(|&x| x).unwrap_or(len), len);

    if start <= end {
        start..end
    } else {
        0..0
    }
}

impl GetSlice for str {
    fn get_slice<T: RangeArgument<usize>>(&self, a: T) -> &Self {
        &self[bound(self.len(), a)]
    }

    fn get_slice_mut<T: RangeArgument<usize>>(&mut self, a: T) -> &mut Self {
        let len = self.len();
        &mut self[bound(len, a)]
    }
}

impl<T> GetSlice for [T] {
    fn get_slice<U: RangeArgument<usize>>(&self, a: U) -> &Self {
        &self[bound(self.len(), a)]
    }

    fn get_slice_mut<U: RangeArgument<usize>>(&mut self, a: U) -> &mut Self {
        let len = self.len();
        &mut self[bound(len, a)]
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use bytecode_verifier::{verify_module, DependencyChecker};
use log::LevelFilter;
use move_lang::{compiled_unit::CompiledUnit, move_compile, shared::Address};
use sha2::{Digest, Sha256};
use std::{
    collections::BTreeMap,
    fs::File,
    io::{Read, Write},
    path::{Path, PathBuf},
};
use vm::CompiledModule;

pub const STD_LIB_DIR: &str = "modules";
pub const MOVE_EXTENSION: &str = "move";

pub const TRANSACTION_SCRIPTS: &str = "transaction_scripts";
/// The output path under which compiled files will be put
pub const COMPILED_OUTPUT_PATH: &str = "compiled";
/// The file name for the compiled stdlib
pub const COMPILED_STDLIB_DIR: &str = "stdlib";
/// The extension for compiled files
pub const COMPILED_EXTENSION: &str = "mv";
/// The file name of the debug module
pub const DEBUG_MODULE_FILE_NAME: &str = "debug.move";

/// The output path for stdlib documentation.
pub const STD_LIB_DOC_DIR: &str = "modules/doc";
/// The output path for transaction script documentation.
pub const TRANSACTION_SCRIPTS_DOC_DIR: &str = "transaction_scripts/doc";

/// The output path under which compiled script files can be found
pub const COMPILED_TRANSACTION_SCRIPTS_DIR: &str = "compiled/transaction_scripts";
/// The output path for transaction script ABIs.
pub const COMPILED_TRANSACTION_SCRIPTS_ABI_DIR: &str = "compiled/transaction_scripts/abi";

/// Where to write generated transaction builders.
pub const TRANSACTION_BUILDERS_GENERATED_SOURCE_PATH: &str =
    "../transaction-builder-generated/src/stdlib.rs";

pub fn filter_move_files(dir_iter: impl Iterator<Item = PathBuf>) -> impl Iterator<Item = PathBuf> {
    dir_iter.flat_map(|path| {
        if path.extension()?.to_str()? == MOVE_EXTENSION {
            Some(path)
        } else {
            None
        }
    })
}

pub fn stdlib_files() -> Vec<String> {
    let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    path.push(STD_LIB_DIR);
    let dirfiles = datatest_stable::utils::iterate_directory(&path);
    filter_move_files(dirfiles)
        .flat_map(|path| path.into_os_string().into_string())
        .collect()
}

pub fn build_stdlib() -> BTreeMap<String, CompiledModule> {
    let (_, compiled_units) =
        move_compile(&stdlib_files(), &[], Some(Address::LIBRA_CORE)).unwrap();
    let mut modules = BTreeMap::new();
    for (i, compiled_unit) in compiled_units.into_iter().enumerate() {
        let name = compiled_unit.name();
        match compiled_unit {
            CompiledUnit::Module { module, .. } => {
                verify_module(&module).expect("stdlib module failed to verify");
                DependencyChecker::verify_module(&module, modules.values())
                    .expect("stdlib module dependency failed to verify");
                // Tag each module with its index in the module dependency order. Needed for
                // when they are deserialized and verified later on.
                modules.insert(format!("{}_{}", i, name), module);
            }
            CompiledUnit::Script { .. } => panic!("Unexpected Script in stdlib"),
        }
    }
    modules
}

pub fn compile_script(source_file_str: String) -> Vec<u8> {
    let (_, mut compiled_program) = move_compile(
        &[source_file_str],
        &stdlib_files(),
        Some(Address::LIBRA_CORE),
    )
    .unwrap();
    let mut script_bytes = vec![];
    assert!(compiled_program.len() == 1);
    match compiled_program.pop().unwrap() {
        CompiledUnit::Module { .. } => panic!("Unexpected module when compiling script"),
        CompiledUnit::Script { script, .. } => script.serialize(&mut script_bytes).unwrap(),
    };
    script_bytes
}

pub fn save_binary(path: &Path, binary: &[u8]) -> bool {
    if path.exists() {
        let mut bytes = vec![];
        File::open(path).unwrap().read_to_end(&mut bytes).unwrap();
        if Sha256::digest(binary) == Sha256::digest(&bytes) {
            return false;
        }
    }
    File::create(path).unwrap().write_all(binary).unwrap();
    true
}

pub fn build_stdlib_doc() {
    build_doc(STD_LIB_DOC_DIR, "", stdlib_files().as_slice(), "")
}

pub fn build_transaction_script_doc(script_file_str: String) {
    build_doc(
        TRANSACTION_SCRIPTS_DOC_DIR,
        STD_LIB_DOC_DIR,
        &[script_file_str],
        STD_LIB_DIR,
    )
}

pub fn build_transaction_script_abi(script_file_str: String) {
    build_abi(
        COMPILED_TRANSACTION_SCRIPTS_ABI_DIR,
        &[script_file_str],
        STD_LIB_DIR,
        COMPILED_TRANSACTION_SCRIPTS_DIR,
    )
}

fn build_doc(output_path: &str, doc_path: &str, sources: &[String], dep_path: &str) {
    let mut options = move_prover::cli::Options::default();
    options.move_sources = sources.to_vec();
    if !dep_path.is_empty() {
        options.move_deps = vec![dep_path.to_string()]
    }
    options.verbosity_level = LevelFilter::Warn;
    options.run_docgen = true;
    options.docgen.include_impl = true;
    options.docgen.include_private_fun = true;
    options.docgen.specs_inlined = false;
    if !doc_path.is_empty() {
        options.docgen.doc_path = vec![doc_path.to_string()];
    }
    options.docgen.output_directory = output_path.to_string();
    options.setup_logging_for_test();
    move_prover::run_move_prover_errors_to_stderr(options).unwrap();
}

fn build_abi(output_path: &str, sources: &[String], dep_path: &str, compiled_script_path: &str) {
    let mut options = move_prover::cli::Options::default();
    options.move_sources = sources.to_vec();
    if !dep_path.is_empty() {
        options.move_deps = vec![dep_path.to_string()]
    }
    options.verbosity_level = LevelFilter::Warn;
    options.run_abigen = true;
    options.abigen.output_directory = output_path.to_string();
    options.abigen.compiled_script_directory = compiled_script_path.to_string();
    options.setup_logging_for_test();
    move_prover::run_move_prover_errors_to_stderr(options).unwrap();
}

pub fn generate_rust_transaction_builders() {
    let abis = transaction_builder_generator::read_abis(COMPILED_TRANSACTION_SCRIPTS_ABI_DIR)
        .expect("Failed to read generated ABIs");
    let mut f = std::fs::File::create(TRANSACTION_BUILDERS_GENERATED_SOURCE_PATH)
        .expect("Failed to open file for Rust script build generation");
    transaction_builder_generator::rust::output(&mut f, &abis, /* local types */ true)
        .expect("Failed to generate Rust builders for Libra");

    std::process::Command::new("rustfmt")
        .arg("--config")
        .arg("merge_imports=true")
        .arg(TRANSACTION_BUILDERS_GENERATED_SOURCE_PATH)
        .status()
        .expect("Failed to run rustfmt on generated code");
}
<commit_msg>[Move stdlib] fix last minute typo with #5319<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use bytecode_verifier::{verify_module, DependencyChecker};
use log::LevelFilter;
use move_lang::{compiled_unit::CompiledUnit, move_compile, shared::Address};
use sha2::{Digest, Sha256};
use std::{
    collections::BTreeMap,
    fs::File,
    io::{Read, Write},
    path::{Path, PathBuf},
};
use vm::CompiledModule;

pub const STD_LIB_DIR: &str = "modules";
pub const MOVE_EXTENSION: &str = "move";

pub const TRANSACTION_SCRIPTS: &str = "transaction_scripts";
/// The output path under which compiled files will be put
pub const COMPILED_OUTPUT_PATH: &str = "compiled";
/// The file name for the compiled stdlib
pub const COMPILED_STDLIB_DIR: &str = "stdlib";
/// The extension for compiled files
pub const COMPILED_EXTENSION: &str = "mv";
/// The file name of the debug module
pub const DEBUG_MODULE_FILE_NAME: &str = "debug.move";

/// The output path for stdlib documentation.
pub const STD_LIB_DOC_DIR: &str = "modules/doc";
/// The output path for transaction script documentation.
pub const TRANSACTION_SCRIPTS_DOC_DIR: &str = "transaction_scripts/doc";

/// The output path under which compiled script files can be found
pub const COMPILED_TRANSACTION_SCRIPTS_DIR: &str = "compiled/transaction_scripts";
/// The output path for transaction script ABIs.
pub const COMPILED_TRANSACTION_SCRIPTS_ABI_DIR: &str = "compiled/transaction_scripts/abi";

/// Where to write generated transaction builders.
pub const TRANSACTION_BUILDERS_GENERATED_SOURCE_PATH: &str =
    "../transaction-builder/generated/src/stdlib.rs";

pub fn filter_move_files(dir_iter: impl Iterator<Item = PathBuf>) -> impl Iterator<Item = PathBuf> {
    dir_iter.flat_map(|path| {
        if path.extension()?.to_str()? == MOVE_EXTENSION {
            Some(path)
        } else {
            None
        }
    })
}

pub fn stdlib_files() -> Vec<String> {
    let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    path.push(STD_LIB_DIR);
    let dirfiles = datatest_stable::utils::iterate_directory(&path);
    filter_move_files(dirfiles)
        .flat_map(|path| path.into_os_string().into_string())
        .collect()
}

pub fn build_stdlib() -> BTreeMap<String, CompiledModule> {
    let (_, compiled_units) =
        move_compile(&stdlib_files(), &[], Some(Address::LIBRA_CORE)).unwrap();
    let mut modules = BTreeMap::new();
    for (i, compiled_unit) in compiled_units.into_iter().enumerate() {
        let name = compiled_unit.name();
        match compiled_unit {
            CompiledUnit::Module { module, .. } => {
                verify_module(&module).expect("stdlib module failed to verify");
                DependencyChecker::verify_module(&module, modules.values())
                    .expect("stdlib module dependency failed to verify");
                // Tag each module with its index in the module dependency order. Needed for
                // when they are deserialized and verified later on.
                modules.insert(format!("{}_{}", i, name), module);
            }
            CompiledUnit::Script { .. } => panic!("Unexpected Script in stdlib"),
        }
    }
    modules
}

pub fn compile_script(source_file_str: String) -> Vec<u8> {
    let (_, mut compiled_program) = move_compile(
        &[source_file_str],
        &stdlib_files(),
        Some(Address::LIBRA_CORE),
    )
    .unwrap();
    let mut script_bytes = vec![];
    assert!(compiled_program.len() == 1);
    match compiled_program.pop().unwrap() {
        CompiledUnit::Module { .. } => panic!("Unexpected module when compiling script"),
        CompiledUnit::Script { script, .. } => script.serialize(&mut script_bytes).unwrap(),
    };
    script_bytes
}

pub fn save_binary(path: &Path, binary: &[u8]) -> bool {
    if path.exists() {
        let mut bytes = vec![];
        File::open(path).unwrap().read_to_end(&mut bytes).unwrap();
        if Sha256::digest(binary) == Sha256::digest(&bytes) {
            return false;
        }
    }
    File::create(path).unwrap().write_all(binary).unwrap();
    true
}

pub fn build_stdlib_doc() {
    build_doc(STD_LIB_DOC_DIR, "", stdlib_files().as_slice(), "")
}

pub fn build_transaction_script_doc(script_file_str: String) {
    build_doc(
        TRANSACTION_SCRIPTS_DOC_DIR,
        STD_LIB_DOC_DIR,
        &[script_file_str],
        STD_LIB_DIR,
    )
}

pub fn build_transaction_script_abi(script_file_str: String) {
    build_abi(
        COMPILED_TRANSACTION_SCRIPTS_ABI_DIR,
        &[script_file_str],
        STD_LIB_DIR,
        COMPILED_TRANSACTION_SCRIPTS_DIR,
    )
}

fn build_doc(output_path: &str, doc_path: &str, sources: &[String], dep_path: &str) {
    let mut options = move_prover::cli::Options::default();
    options.move_sources = sources.to_vec();
    if !dep_path.is_empty() {
        options.move_deps = vec![dep_path.to_string()]
    }
    options.verbosity_level = LevelFilter::Warn;
    options.run_docgen = true;
    options.docgen.include_impl = true;
    options.docgen.include_private_fun = true;
    options.docgen.specs_inlined = false;
    if !doc_path.is_empty() {
        options.docgen.doc_path = vec![doc_path.to_string()];
    }
    options.docgen.output_directory = output_path.to_string();
    options.setup_logging_for_test();
    move_prover::run_move_prover_errors_to_stderr(options).unwrap();
}

fn build_abi(output_path: &str, sources: &[String], dep_path: &str, compiled_script_path: &str) {
    let mut options = move_prover::cli::Options::default();
    options.move_sources = sources.to_vec();
    if !dep_path.is_empty() {
        options.move_deps = vec![dep_path.to_string()]
    }
    options.verbosity_level = LevelFilter::Warn;
    options.run_abigen = true;
    options.abigen.output_directory = output_path.to_string();
    options.abigen.compiled_script_directory = compiled_script_path.to_string();
    options.setup_logging_for_test();
    move_prover::run_move_prover_errors_to_stderr(options).unwrap();
}

pub fn generate_rust_transaction_builders() {
    let abis = transaction_builder_generator::read_abis(COMPILED_TRANSACTION_SCRIPTS_ABI_DIR)
        .expect("Failed to read generated ABIs");
    let mut f = std::fs::File::create(TRANSACTION_BUILDERS_GENERATED_SOURCE_PATH)
        .expect("Failed to open file for Rust script build generation");
    transaction_builder_generator::rust::output(&mut f, &abis, /* local types */ true)
        .expect("Failed to generate Rust builders for Libra");

    std::process::Command::new("rustfmt")
        .arg("--config")
        .arg("merge_imports=true")
        .arg(TRANSACTION_BUILDERS_GENERATED_SOURCE_PATH)
        .status()
        .expect("Failed to run rustfmt on generated code");
}
<|endoftext|>


<commit_before>//! The data that we will serialize and deserialize.

use super::query::DepGraphQuery;
use super::{DepKind, DepNode, DepNodeIndex};
use rustc_data_structures::fingerprint::Fingerprint;
use rustc_data_structures::fx::FxHashMap;
use rustc_data_structures::sync::Lock;
use rustc_index::vec::{Idx, IndexVec};
use rustc_serialize::opaque::{self, FileEncodeResult, FileEncoder, IntEncodedWithFixedSize};
use rustc_serialize::{Decodable, Decoder, Encodable};
use smallvec::SmallVec;
use std::convert::TryInto;

// The maximum value of `SerializedDepNodeIndex` leaves the upper two bits
// unused so that we can store multiple index types in `CompressedHybridIndex`,
// and use those bits to encode which index type it contains.
rustc_index::newtype_index! {
    pub struct SerializedDepNodeIndex {
        MAX = 0x7FFF_FFFF
    }
}

/// Data for use when recompiling the **current crate**.
#[derive(Debug)]
pub struct SerializedDepGraph<K: DepKind> {
    /// The set of all DepNodes in the graph
    pub nodes: IndexVec<SerializedDepNodeIndex, DepNode<K>>,
    /// The set of all Fingerprints in the graph. Each Fingerprint corresponds to
    /// the DepNode at the same index in the nodes vector.
    pub fingerprints: IndexVec<SerializedDepNodeIndex, Fingerprint>,
    /// For each DepNode, stores the list of edges originating from that
    /// DepNode. Encoded as a [start, end) pair indexing into edge_list_data,
    /// which holds the actual DepNodeIndices of the target nodes.
    pub edge_list_indices: IndexVec<SerializedDepNodeIndex, (u32, u32)>,
    /// A flattened list of all edge targets in the graph. Edge sources are
    /// implicit in edge_list_indices.
    pub edge_list_data: Vec<SerializedDepNodeIndex>,
}

impl<K: DepKind> Default for SerializedDepGraph<K> {
    fn default() -> Self {
        SerializedDepGraph {
            nodes: Default::default(),
            fingerprints: Default::default(),
            edge_list_indices: Default::default(),
            edge_list_data: Default::default(),
        }
    }
}

impl<K: DepKind> SerializedDepGraph<K> {
    #[inline]
    pub fn edge_targets_from(&self, source: SerializedDepNodeIndex) -> &[SerializedDepNodeIndex] {
        let targets = self.edge_list_indices[source];
        &self.edge_list_data[targets.0 as usize..targets.1 as usize]
    }
}

impl<'a, K: DepKind + Decodable<opaque::Decoder<'a>>> Decodable<opaque::Decoder<'a>>
    for SerializedDepGraph<K>
{
    #[instrument(skip(d))]
    fn decode(d: &mut opaque::Decoder<'a>) -> Result<SerializedDepGraph<K>, String> {
        let position = d.position();

        // The last 16 bytes are the node count and edge count.
        debug!("position: {:?}", d.position());
        d.set_position(d.data.len() - 2 * IntEncodedWithFixedSize::ENCODED_SIZE);
        debug!("position: {:?}", d.position());

        let node_count = IntEncodedWithFixedSize::decode(d)?.0 as usize;
        let edge_count = IntEncodedWithFixedSize::decode(d)?.0 as usize;
        debug!(?node_count, ?edge_count);

        debug!("position: {:?}", d.position());
        d.set_position(position);
        debug!("position: {:?}", d.position());

        let mut nodes = IndexVec::with_capacity(node_count);
        let mut fingerprints = IndexVec::with_capacity(node_count);
        let mut edge_list_indices = IndexVec::with_capacity(node_count);
        let mut edge_list_data = Vec::with_capacity(edge_count);

        for _index in 0..node_count {
            d.read_struct("NodeInfo", 3, |d| {
                let dep_node: DepNode<K> = d.read_struct_field("node", 0, Decodable::decode)?;
                let _i: SerializedDepNodeIndex = nodes.push(dep_node);
                debug_assert_eq!(_i.index(), _index);

                let fingerprint: Fingerprint =
                    d.read_struct_field("fingerprint", 1, Decodable::decode)?;
                let _i: SerializedDepNodeIndex = fingerprints.push(fingerprint);
                debug_assert_eq!(_i.index(), _index);

                d.read_struct_field("edges", 2, |d| {
                    d.read_seq(|d, len| {
                        let start = edge_list_data.len().try_into().unwrap();
                        for e in 0..len {
                            let edge = d.read_seq_elt(e, Decodable::decode)?;
                            edge_list_data.push(edge);
                        }
                        let end = edge_list_data.len().try_into().unwrap();
                        let _i: SerializedDepNodeIndex = edge_list_indices.push((start, end));
                        debug_assert_eq!(_i.index(), _index);
                        Ok(())
                    })
                })
            })?;
        }

        Ok(SerializedDepGraph { nodes, fingerprints, edge_list_indices, edge_list_data })
    }
}

#[derive(Debug, Encodable, Decodable)]
pub struct NodeInfo<K: DepKind> {
    node: DepNode<K>,
    fingerprint: Fingerprint,
    edges: SmallVec<[DepNodeIndex; 8]>,
}

struct Stat<K: DepKind> {
    kind: K,
    node_counter: u64,
    edge_counter: u64,
}

struct EncodingStatus<K: DepKind> {
    encoder: FileEncoder,
    total_node_count: usize,
    total_edge_count: usize,
    result: FileEncodeResult,
    stats: Option<FxHashMap<K, Stat<K>>>,
}

impl<K: DepKind> EncodingStatus<K> {
    fn new(encoder: FileEncoder, record_stats: bool) -> Self {
        Self {
            encoder,
            total_edge_count: 0,
            total_node_count: 0,
            result: Ok(()),
            stats: if record_stats { Some(FxHashMap::default()) } else { None },
        }
    }

    #[instrument(skip(self, _record_graph))]
    fn encode_node(
        &mut self,
        node: &NodeInfo<K>,
        _record_graph: &Option<Lock<DepGraphQuery<K>>>,
    ) -> DepNodeIndex {
        let index = DepNodeIndex::new(self.total_node_count);
        self.total_node_count += 1;

        let edge_count = node.edges.len();
        self.total_edge_count += edge_count;

        #[cfg(debug_assertions)]
        if let Some(record_graph) = &_record_graph {
            // Do not ICE when a query is called from within `with_query`.
            if let Some(record_graph) = &mut record_graph.try_lock() {
                record_graph.push(index, node.node, &node.edges);
            }
        }

        if let Some(stats) = &mut self.stats {
            let kind = node.node.kind;

            let stat = stats.entry(kind).or_insert(Stat { kind, node_counter: 0, edge_counter: 0 });
            stat.node_counter += 1;
            stat.edge_counter += edge_count as u64;
        }

        debug!(?index, ?node);
        let encoder = &mut self.encoder;
        self.result =
            std::mem::replace(&mut self.result, Ok(())).and_then(|()| node.encode(encoder));
        index
    }

    fn finish(self) -> FileEncodeResult {
        let Self { mut encoder, total_node_count, total_edge_count, result, stats: _ } = self;
        let () = result?;

        let node_count = total_node_count.try_into().unwrap();
        let edge_count = total_edge_count.try_into().unwrap();

        debug!(?node_count, ?edge_count);
        debug!("position: {:?}", encoder.position());
        IntEncodedWithFixedSize(node_count).encode(&mut encoder)?;
        IntEncodedWithFixedSize(edge_count).encode(&mut encoder)?;
        debug!("position: {:?}", encoder.position());
        // Drop the encoder so that nothing is written after the counts.
        encoder.flush()
    }
}

pub struct GraphEncoder<K: DepKind> {
    status: Lock<EncodingStatus<K>>,
    record_graph: Option<Lock<DepGraphQuery<K>>>,
}

impl<K: DepKind + Encodable<FileEncoder>> GraphEncoder<K> {
    pub fn new(
        encoder: FileEncoder,
        prev_node_count: usize,
        record_graph: bool,
        record_stats: bool,
    ) -> Self {
        let record_graph = if cfg!(debug_assertions) && record_graph {
            Some(Lock::new(DepGraphQuery::new(prev_node_count)))
        } else {
            None
        };
        let status = Lock::new(EncodingStatus::new(encoder, record_stats));
        GraphEncoder { status, record_graph }
    }

    pub(crate) fn with_query(&self, f: impl Fn(&DepGraphQuery<K>)) {
        if let Some(record_graph) = &self.record_graph {
            f(&record_graph.lock())
        }
    }

    pub(crate) fn print_incremental_info(
        &self,
        total_read_count: u64,
        total_duplicate_read_count: u64,
    ) {
        let status = self.status.lock();
        if let Some(record_stats) = &status.stats {
            let mut stats: Vec<_> = record_stats.values().collect();
            stats.sort_by_key(|s| -(s.node_counter as i64));

            const SEPARATOR: &str = "[incremental] --------------------------------\
                                     ----------------------------------------------\
                                     ------------";

            eprintln!("[incremental]");
            eprintln!("[incremental] DepGraph Statistics");
            eprintln!("{}", SEPARATOR);
            eprintln!("[incremental]");
            eprintln!("[incremental] Total Node Count: {}", status.total_node_count);
            eprintln!("[incremental] Total Edge Count: {}", status.total_edge_count);

            if cfg!(debug_assertions) {
                eprintln!("[incremental] Total Edge Reads: {}", total_read_count);
                eprintln!(
                    "[incremental] Total Duplicate Edge Reads: {}",
                    total_duplicate_read_count
                );
            }

            eprintln!("[incremental]");
            eprintln!(
                "[incremental]  {:<36}| {:<17}| {:<12}| {:<17}|",
                "Node Kind", "Node Frequency", "Node Count", "Avg. Edge Count"
            );
            eprintln!("{}", SEPARATOR);

            for stat in stats {
                let node_kind_ratio =
                    (100.0 * (stat.node_counter as f64)) / (status.total_node_count as f64);
                let node_kind_avg_edges = (stat.edge_counter as f64) / (stat.node_counter as f64);

                eprintln!(
                    "[incremental]  {:<36}|{:>16.1}% |{:>12} |{:>17.1} |",
                    format!("{:?}", stat.kind),
                    node_kind_ratio,
                    stat.node_counter,
                    node_kind_avg_edges,
                );
            }

            eprintln!("{}", SEPARATOR);
            eprintln!("[incremental]");
        }
    }

    pub(crate) fn send(
        &self,
        node: DepNode<K>,
        fingerprint: Fingerprint,
        edges: SmallVec<[DepNodeIndex; 8]>,
    ) -> DepNodeIndex {
        let node = NodeInfo { node, fingerprint, edges };
        self.status.lock().encode_node(&node, &self.record_graph)
    }

    pub fn finish(self) -> FileEncodeResult {
        self.status.into_inner().finish()
    }
}
<commit_msg>Add documentation.<commit_after>//! The data that we will serialize and deserialize.
//!
//! The dep-graph is serialized as a sequence of NodeInfo, with the dependencies
//! specified inline.  The total number of nodes and edges are stored as the last
//! 16 bytes of the file, so we can find them easily at decoding time.
//!
//! The serialisation is performed on-demand when each node is emitted. Using this
//! scheme, we do not need to keep the current graph in memory.
//!
//! The deserisalisation is performed manually, in order to convert from the stored
//! sequence of NodeInfos to the different arrays in SerializedDepGraph.  Since the
//! node and edge count are stored at the end of the file, all the arrays can be
//! pre-allocated with the right length.

use super::query::DepGraphQuery;
use super::{DepKind, DepNode, DepNodeIndex};
use rustc_data_structures::fingerprint::Fingerprint;
use rustc_data_structures::fx::FxHashMap;
use rustc_data_structures::sync::Lock;
use rustc_index::vec::{Idx, IndexVec};
use rustc_serialize::opaque::{self, FileEncodeResult, FileEncoder, IntEncodedWithFixedSize};
use rustc_serialize::{Decodable, Decoder, Encodable};
use smallvec::SmallVec;
use std::convert::TryInto;

// The maximum value of `SerializedDepNodeIndex` leaves the upper two bits
// unused so that we can store multiple index types in `CompressedHybridIndex`,
// and use those bits to encode which index type it contains.
rustc_index::newtype_index! {
    pub struct SerializedDepNodeIndex {
        MAX = 0x7FFF_FFFF
    }
}

/// Data for use when recompiling the **current crate**.
#[derive(Debug)]
pub struct SerializedDepGraph<K: DepKind> {
    /// The set of all DepNodes in the graph
    pub nodes: IndexVec<SerializedDepNodeIndex, DepNode<K>>,
    /// The set of all Fingerprints in the graph. Each Fingerprint corresponds to
    /// the DepNode at the same index in the nodes vector.
    pub fingerprints: IndexVec<SerializedDepNodeIndex, Fingerprint>,
    /// For each DepNode, stores the list of edges originating from that
    /// DepNode. Encoded as a [start, end) pair indexing into edge_list_data,
    /// which holds the actual DepNodeIndices of the target nodes.
    pub edge_list_indices: IndexVec<SerializedDepNodeIndex, (u32, u32)>,
    /// A flattened list of all edge targets in the graph. Edge sources are
    /// implicit in edge_list_indices.
    pub edge_list_data: Vec<SerializedDepNodeIndex>,
}

impl<K: DepKind> Default for SerializedDepGraph<K> {
    fn default() -> Self {
        SerializedDepGraph {
            nodes: Default::default(),
            fingerprints: Default::default(),
            edge_list_indices: Default::default(),
            edge_list_data: Default::default(),
        }
    }
}

impl<K: DepKind> SerializedDepGraph<K> {
    #[inline]
    pub fn edge_targets_from(&self, source: SerializedDepNodeIndex) -> &[SerializedDepNodeIndex] {
        let targets = self.edge_list_indices[source];
        &self.edge_list_data[targets.0 as usize..targets.1 as usize]
    }
}

impl<'a, K: DepKind + Decodable<opaque::Decoder<'a>>> Decodable<opaque::Decoder<'a>>
    for SerializedDepGraph<K>
{
    #[instrument(skip(d))]
    fn decode(d: &mut opaque::Decoder<'a>) -> Result<SerializedDepGraph<K>, String> {
        let position = d.position();

        // The last 16 bytes are the node count and edge count.
        debug!("position: {:?}", d.position());
        d.set_position(d.data.len() - 2 * IntEncodedWithFixedSize::ENCODED_SIZE);
        debug!("position: {:?}", d.position());

        let node_count = IntEncodedWithFixedSize::decode(d)?.0 as usize;
        let edge_count = IntEncodedWithFixedSize::decode(d)?.0 as usize;
        debug!(?node_count, ?edge_count);

        debug!("position: {:?}", d.position());
        d.set_position(position);
        debug!("position: {:?}", d.position());

        let mut nodes = IndexVec::with_capacity(node_count);
        let mut fingerprints = IndexVec::with_capacity(node_count);
        let mut edge_list_indices = IndexVec::with_capacity(node_count);
        let mut edge_list_data = Vec::with_capacity(edge_count);

        for _index in 0..node_count {
            d.read_struct("NodeInfo", 3, |d| {
                let dep_node: DepNode<K> = d.read_struct_field("node", 0, Decodable::decode)?;
                let _i: SerializedDepNodeIndex = nodes.push(dep_node);
                debug_assert_eq!(_i.index(), _index);

                let fingerprint: Fingerprint =
                    d.read_struct_field("fingerprint", 1, Decodable::decode)?;
                let _i: SerializedDepNodeIndex = fingerprints.push(fingerprint);
                debug_assert_eq!(_i.index(), _index);

                d.read_struct_field("edges", 2, |d| {
                    d.read_seq(|d, len| {
                        let start = edge_list_data.len().try_into().unwrap();
                        for e in 0..len {
                            let edge = d.read_seq_elt(e, Decodable::decode)?;
                            edge_list_data.push(edge);
                        }
                        let end = edge_list_data.len().try_into().unwrap();
                        let _i: SerializedDepNodeIndex = edge_list_indices.push((start, end));
                        debug_assert_eq!(_i.index(), _index);
                        Ok(())
                    })
                })
            })?;
        }

        Ok(SerializedDepGraph { nodes, fingerprints, edge_list_indices, edge_list_data })
    }
}

#[derive(Debug, Encodable, Decodable)]
pub struct NodeInfo<K: DepKind> {
    node: DepNode<K>,
    fingerprint: Fingerprint,
    edges: SmallVec<[DepNodeIndex; 8]>,
}

struct Stat<K: DepKind> {
    kind: K,
    node_counter: u64,
    edge_counter: u64,
}

struct EncodingStatus<K: DepKind> {
    encoder: FileEncoder,
    total_node_count: usize,
    total_edge_count: usize,
    result: FileEncodeResult,
    stats: Option<FxHashMap<K, Stat<K>>>,
}

impl<K: DepKind> EncodingStatus<K> {
    fn new(encoder: FileEncoder, record_stats: bool) -> Self {
        Self {
            encoder,
            total_edge_count: 0,
            total_node_count: 0,
            result: Ok(()),
            stats: if record_stats { Some(FxHashMap::default()) } else { None },
        }
    }

    #[instrument(skip(self, _record_graph))]
    fn encode_node(
        &mut self,
        node: &NodeInfo<K>,
        _record_graph: &Option<Lock<DepGraphQuery<K>>>,
    ) -> DepNodeIndex {
        let index = DepNodeIndex::new(self.total_node_count);
        self.total_node_count += 1;

        let edge_count = node.edges.len();
        self.total_edge_count += edge_count;

        #[cfg(debug_assertions)]
        if let Some(record_graph) = &_record_graph {
            // Do not ICE when a query is called from within `with_query`.
            if let Some(record_graph) = &mut record_graph.try_lock() {
                record_graph.push(index, node.node, &node.edges);
            }
        }

        if let Some(stats) = &mut self.stats {
            let kind = node.node.kind;

            let stat = stats.entry(kind).or_insert(Stat { kind, node_counter: 0, edge_counter: 0 });
            stat.node_counter += 1;
            stat.edge_counter += edge_count as u64;
        }

        debug!(?index, ?node);
        let encoder = &mut self.encoder;
        self.result =
            std::mem::replace(&mut self.result, Ok(())).and_then(|()| node.encode(encoder));
        index
    }

    fn finish(self) -> FileEncodeResult {
        let Self { mut encoder, total_node_count, total_edge_count, result, stats: _ } = self;
        let () = result?;

        let node_count = total_node_count.try_into().unwrap();
        let edge_count = total_edge_count.try_into().unwrap();

        debug!(?node_count, ?edge_count);
        debug!("position: {:?}", encoder.position());
        IntEncodedWithFixedSize(node_count).encode(&mut encoder)?;
        IntEncodedWithFixedSize(edge_count).encode(&mut encoder)?;
        debug!("position: {:?}", encoder.position());
        // Drop the encoder so that nothing is written after the counts.
        encoder.flush()
    }
}

pub struct GraphEncoder<K: DepKind> {
    status: Lock<EncodingStatus<K>>,
    record_graph: Option<Lock<DepGraphQuery<K>>>,
}

impl<K: DepKind + Encodable<FileEncoder>> GraphEncoder<K> {
    pub fn new(
        encoder: FileEncoder,
        prev_node_count: usize,
        record_graph: bool,
        record_stats: bool,
    ) -> Self {
        let record_graph = if cfg!(debug_assertions) && record_graph {
            Some(Lock::new(DepGraphQuery::new(prev_node_count)))
        } else {
            None
        };
        let status = Lock::new(EncodingStatus::new(encoder, record_stats));
        GraphEncoder { status, record_graph }
    }

    pub(crate) fn with_query(&self, f: impl Fn(&DepGraphQuery<K>)) {
        if let Some(record_graph) = &self.record_graph {
            f(&record_graph.lock())
        }
    }

    pub(crate) fn print_incremental_info(
        &self,
        total_read_count: u64,
        total_duplicate_read_count: u64,
    ) {
        let status = self.status.lock();
        if let Some(record_stats) = &status.stats {
            let mut stats: Vec<_> = record_stats.values().collect();
            stats.sort_by_key(|s| -(s.node_counter as i64));

            const SEPARATOR: &str = "[incremental] --------------------------------\
                                     ----------------------------------------------\
                                     ------------";

            eprintln!("[incremental]");
            eprintln!("[incremental] DepGraph Statistics");
            eprintln!("{}", SEPARATOR);
            eprintln!("[incremental]");
            eprintln!("[incremental] Total Node Count: {}", status.total_node_count);
            eprintln!("[incremental] Total Edge Count: {}", status.total_edge_count);

            if cfg!(debug_assertions) {
                eprintln!("[incremental] Total Edge Reads: {}", total_read_count);
                eprintln!(
                    "[incremental] Total Duplicate Edge Reads: {}",
                    total_duplicate_read_count
                );
            }

            eprintln!("[incremental]");
            eprintln!(
                "[incremental]  {:<36}| {:<17}| {:<12}| {:<17}|",
                "Node Kind", "Node Frequency", "Node Count", "Avg. Edge Count"
            );
            eprintln!("{}", SEPARATOR);

            for stat in stats {
                let node_kind_ratio =
                    (100.0 * (stat.node_counter as f64)) / (status.total_node_count as f64);
                let node_kind_avg_edges = (stat.edge_counter as f64) / (stat.node_counter as f64);

                eprintln!(
                    "[incremental]  {:<36}|{:>16.1}% |{:>12} |{:>17.1} |",
                    format!("{:?}", stat.kind),
                    node_kind_ratio,
                    stat.node_counter,
                    node_kind_avg_edges,
                );
            }

            eprintln!("{}", SEPARATOR);
            eprintln!("[incremental]");
        }
    }

    pub(crate) fn send(
        &self,
        node: DepNode<K>,
        fingerprint: Fingerprint,
        edges: SmallVec<[DepNodeIndex; 8]>,
    ) -> DepNodeIndex {
        let node = NodeInfo { node, fingerprint, edges };
        self.status.lock().encode_node(&node, &self.record_graph)
    }

    pub fn finish(self) -> FileEncodeResult {
        self.status.into_inner().finish()
    }
}
<|endoftext|>


<commit_before>// Feature: Format String Completion.
//
// `"Result {result} is {2 + 2}"` is expanded to the `"Result {} is {}", result, 2 + 2`.
//
// The following postfix snippets are available:
//
// - `format` -> `format!(...)`
// - `panic` -> `panic!(...)`
// - `println` -> `println!(...)`
// - `log`:
//   + `logd` -> `log::debug!(...)`
//   + `logt` -> `log::trace!(...)`
//   + `logi` -> `log::info!(...)`
//   + `logw` -> `log::warn!(...)`
//   + `loge` -> `log::error!(...)`

use crate::{
    completions::postfix::postfix_snippet, config::SnippetCap, context::CompletionContext,
    Completions,
};
use syntax::ast::{self, AstToken};

/// Mapping ("postfix completion item" => "macro to use")
static KINDS: &[(&str, &str)] = &[
    ("fmt", "format!"),
    ("panic", "panic!"),
    ("println", "println!"),
    ("eprintln", "eprintln!"),
    ("logd", "log::debug!"),
    ("logt", "log::trace!"),
    ("logi", "log::info!"),
    ("logw", "log::warn!"),
    ("loge", "log::error!"),
];

pub(crate) fn add_format_like_completions(
    acc: &mut Completions,
    ctx: &CompletionContext,
    dot_receiver: &ast::Expr,
    cap: SnippetCap,
    receiver_text: &ast::String,
) {
    let input = match string_literal_contents(receiver_text) {
        // It's not a string literal, do not parse input.
        Some(input) => input,
        None => return,
    };

    let mut parser = FormatStrParser::new(input);

    if parser.parse().is_ok() {
        for (label, macro_name) in KINDS {
            let snippet = parser.into_suggestion(macro_name);

            postfix_snippet(ctx, cap, &dot_receiver, label, macro_name, &snippet).add_to(acc);
        }
    }
}

/// Checks whether provided item is a string literal.
fn string_literal_contents(item: &ast::String) -> Option<String> {
    let item = item.text();
    if item.len() >= 2 && item.starts_with("\"") && item.ends_with("\"") {
        return Some(item[1..item.len() - 1].to_owned());
    }

    None
}

/// Parser for a format-like string. It is more allowing in terms of string contents,
/// as we expect variable placeholders to be filled with expressions.
#[derive(Debug)]
pub(crate) struct FormatStrParser {
    input: String,
    output: String,
    extracted_expressions: Vec<String>,
    state: State,
    parsed: bool,
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum State {
    NotExpr,
    MaybeExpr,
    Expr,
    MaybeIncorrect,
    FormatOpts,
}

impl FormatStrParser {
    pub(crate) fn new(input: String) -> Self {
        Self {
            input: input.into(),
            output: String::new(),
            extracted_expressions: Vec::new(),
            state: State::NotExpr,
            parsed: false,
        }
    }

    pub(crate) fn parse(&mut self) -> Result<(), ()> {
        let mut current_expr = String::new();

        let mut placeholder_id = 1;

        // Count of open braces inside of an expression.
        // We assume that user knows what they're doing, thus we treat it like a correct pattern, e.g.
        // "{MyStruct { val_a: 0, val_b: 1 }}".
        let mut inexpr_open_count = 0;

        for chr in self.input.chars() {
            match (self.state, chr) {
                (State::NotExpr, '{') => {
                    self.output.push(chr);
                    self.state = State::MaybeExpr;
                }
                (State::NotExpr, '}') => {
                    self.output.push(chr);
                    self.state = State::MaybeIncorrect;
                }
                (State::NotExpr, _) => {
                    self.output.push(chr);
                }
                (State::MaybeIncorrect, '}') => {
                    // It's okay, we met "}}".
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::MaybeIncorrect, _) => {
                    // Error in the string.
                    return Err(());
                }
                (State::MaybeExpr, '{') => {
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::MaybeExpr, '}') => {
                    // This is an empty sequence '{}'. Replace it with placeholder.
                    self.output.push(chr);
                    self.extracted_expressions.push(format!("${}", placeholder_id));
                    placeholder_id += 1;
                    self.state = State::NotExpr;
                }
                (State::MaybeExpr, _) => {
                    current_expr.push(chr);
                    self.state = State::Expr;
                }
                (State::Expr, '}') => {
                    if inexpr_open_count == 0 {
                        self.output.push(chr);
                        self.extracted_expressions.push(current_expr.trim().into());
                        current_expr = String::new();
                        self.state = State::NotExpr;
                    } else {
                        // We're closing one brace met before inside of the expression.
                        current_expr.push(chr);
                        inexpr_open_count -= 1;
                    }
                }
                (State::Expr, ':') => {
                    if inexpr_open_count == 0 {
                        // We're outside of braces, thus assume that it's a specifier, like "{Some(value):?}"
                        self.output.push(chr);
                        self.extracted_expressions.push(current_expr.trim().into());
                        current_expr = String::new();
                        self.state = State::FormatOpts;
                    } else {
                        // We're inside of braced expression, assume that it's a struct field name/value delimeter.
                        current_expr.push(chr);
                    }
                }
                (State::Expr, '{') => {
                    current_expr.push(chr);
                    inexpr_open_count += 1;
                }
                (State::Expr, _) => {
                    current_expr.push(chr);
                }
                (State::FormatOpts, '}') => {
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::FormatOpts, _) => {
                    self.output.push(chr);
                }
            }
        }

        if self.state != State::NotExpr {
            return Err(());
        }

        self.parsed = true;
        Ok(())
    }

    pub(crate) fn into_suggestion(&self, macro_name: &str) -> String {
        assert!(self.parsed, "Attempt to get a suggestion from not parsed expression");

        let expressions_as_string = self.extracted_expressions.join(", ");
        format!(r#"{}("{}", {})"#, macro_name, self.output, expressions_as_string)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use expect_test::{expect, Expect};

    fn check(input: &str, expect: &Expect) {
        let mut parser = FormatStrParser::new((*input).to_owned());
        let outcome_repr = if parser.parse().is_ok() {
            // Parsing should be OK, expected repr is "string; expr_1, expr_2".
            if parser.extracted_expressions.is_empty() {
                parser.output
            } else {
                format!("{}; {}", parser.output, parser.extracted_expressions.join(", "))
            }
        } else {
            // Parsing should fail, expected repr is "-".
            "-".to_owned()
        };

        expect.assert_eq(&outcome_repr);
    }

    #[test]
    fn format_str_parser() {
        let test_vector = &[
            ("no expressions", expect![["no expressions"]]),
            ("{expr} is {2 + 2}", expect![["{} is {}; expr, 2 + 2"]]),
            ("{expr:?}", expect![["{:?}; expr"]]),
            ("{malformed", expect![["-"]]),
            ("malformed}", expect![["-"]]),
            ("{{correct", expect![["{{correct"]]),
            ("correct}}", expect![["correct}}"]]),
            ("{correct}}}", expect![["{}}}; correct"]]),
            ("{correct}}}}}", expect![["{}}}}}; correct"]]),
            ("{incorrect}}", expect![["-"]]),
            ("placeholders {} {}", expect![["placeholders {} {}; $1, $2"]]),
            ("mixed {} {2 + 2} {}", expect![["mixed {} {} {}; $1, 2 + 2, $2"]]),
            (
                "{SomeStruct { val_a: 0, val_b: 1 }}",
                expect![["{}; SomeStruct { val_a: 0, val_b: 1 }"]],
            ),
            ("{expr:?} is {2.32f64:.5}", expect![["{:?} is {:.5}; expr, 2.32f64"]]),
            (
                "{SomeStruct { val_a: 0, val_b: 1 }:?}",
                expect![["{:?}; SomeStruct { val_a: 0, val_b: 1 }"]],
            ),
            ("{     2 + 2        }", expect![["{}; 2 + 2"]]),
        ];

        for (input, output) in test_vector {
            check(input, output)
        }
    }

    #[test]
    fn test_into_suggestion() {
        let test_vector = &[
            ("println!", "{}", r#"println!("{}", $1)"#),
            ("eprintln!", "{}", r#"eprintln!("{}", $1)"#),
            (
                "log::info!",
                "{} {expr} {} {2 + 2}",
                r#"log::info!("{} {} {} {}", $1, expr, $2, 2 + 2)"#,
            ),
            ("format!", "{expr:?}", r#"format!("{:?}", expr)"#),
        ];

        for (kind, input, output) in test_vector {
            let mut parser = FormatStrParser::new((*input).to_owned());
            parser.parse().expect("Parsing must succeed");

            assert_eq!(&parser.into_suggestion(*kind), output);
        }
    }
}
<commit_msg>Properly parse path separators in format-like postfix<commit_after>// Feature: Format String Completion.
//
// `"Result {result} is {2 + 2}"` is expanded to the `"Result {} is {}", result, 2 + 2`.
//
// The following postfix snippets are available:
//
// - `format` -> `format!(...)`
// - `panic` -> `panic!(...)`
// - `println` -> `println!(...)`
// - `log`:
//   + `logd` -> `log::debug!(...)`
//   + `logt` -> `log::trace!(...)`
//   + `logi` -> `log::info!(...)`
//   + `logw` -> `log::warn!(...)`
//   + `loge` -> `log::error!(...)`

use crate::{
    completions::postfix::postfix_snippet, config::SnippetCap, context::CompletionContext,
    Completions,
};
use syntax::ast::{self, AstToken};

/// Mapping ("postfix completion item" => "macro to use")
static KINDS: &[(&str, &str)] = &[
    ("fmt", "format!"),
    ("panic", "panic!"),
    ("println", "println!"),
    ("eprintln", "eprintln!"),
    ("logd", "log::debug!"),
    ("logt", "log::trace!"),
    ("logi", "log::info!"),
    ("logw", "log::warn!"),
    ("loge", "log::error!"),
];

pub(crate) fn add_format_like_completions(
    acc: &mut Completions,
    ctx: &CompletionContext,
    dot_receiver: &ast::Expr,
    cap: SnippetCap,
    receiver_text: &ast::String,
) {
    let input = match string_literal_contents(receiver_text) {
        // It's not a string literal, do not parse input.
        Some(input) => input,
        None => return,
    };

    let mut parser = FormatStrParser::new(input);

    if parser.parse().is_ok() {
        for (label, macro_name) in KINDS {
            let snippet = parser.into_suggestion(macro_name);

            postfix_snippet(ctx, cap, &dot_receiver, label, macro_name, &snippet).add_to(acc);
        }
    }
}

/// Checks whether provided item is a string literal.
fn string_literal_contents(item: &ast::String) -> Option<String> {
    let item = item.text();
    if item.len() >= 2 && item.starts_with("\"") && item.ends_with("\"") {
        return Some(item[1..item.len() - 1].to_owned());
    }

    None
}

/// Parser for a format-like string. It is more allowing in terms of string contents,
/// as we expect variable placeholders to be filled with expressions.
#[derive(Debug)]
pub(crate) struct FormatStrParser {
    input: String,
    output: String,
    extracted_expressions: Vec<String>,
    state: State,
    parsed: bool,
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum State {
    NotExpr,
    MaybeExpr,
    Expr,
    MaybeIncorrect,
    FormatOpts,
}

impl FormatStrParser {
    pub(crate) fn new(input: String) -> Self {
        Self {
            input: input.into(),
            output: String::new(),
            extracted_expressions: Vec::new(),
            state: State::NotExpr,
            parsed: false,
        }
    }

    pub(crate) fn parse(&mut self) -> Result<(), ()> {
        let mut current_expr = String::new();

        let mut placeholder_id = 1;

        // Count of open braces inside of an expression.
        // We assume that user knows what they're doing, thus we treat it like a correct pattern, e.g.
        // "{MyStruct { val_a: 0, val_b: 1 }}".
        let mut inexpr_open_count = 0;

        let mut chars = self.input.chars().peekable();
        while let Some(chr) = chars.next() {
            match (self.state, chr) {
                (State::NotExpr, '{') => {
                    self.output.push(chr);
                    self.state = State::MaybeExpr;
                }
                (State::NotExpr, '}') => {
                    self.output.push(chr);
                    self.state = State::MaybeIncorrect;
                }
                (State::NotExpr, _) => {
                    self.output.push(chr);
                }
                (State::MaybeIncorrect, '}') => {
                    // It's okay, we met "}}".
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::MaybeIncorrect, _) => {
                    // Error in the string.
                    return Err(());
                }
                (State::MaybeExpr, '{') => {
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::MaybeExpr, '}') => {
                    // This is an empty sequence '{}'. Replace it with placeholder.
                    self.output.push(chr);
                    self.extracted_expressions.push(format!("${}", placeholder_id));
                    placeholder_id += 1;
                    self.state = State::NotExpr;
                }
                (State::MaybeExpr, _) => {
                    current_expr.push(chr);
                    self.state = State::Expr;
                }
                (State::Expr, '}') => {
                    if inexpr_open_count == 0 {
                        self.output.push(chr);
                        self.extracted_expressions.push(current_expr.trim().into());
                        current_expr = String::new();
                        self.state = State::NotExpr;
                    } else {
                        // We're closing one brace met before inside of the expression.
                        current_expr.push(chr);
                        inexpr_open_count -= 1;
                    }
                }
                (State::Expr, ':') if chars.peek().copied() == Some(':') => {
                    // path seperator
                    current_expr.push_str("::");
                    chars.next();
                }
                (State::Expr, ':') => {
                    if inexpr_open_count == 0 {
                        // We're outside of braces, thus assume that it's a specifier, like "{Some(value):?}"
                        self.output.push(chr);
                        self.extracted_expressions.push(current_expr.trim().into());
                        current_expr = String::new();
                        self.state = State::FormatOpts;
                    } else {
                        // We're inside of braced expression, assume that it's a struct field name/value delimeter.
                        current_expr.push(chr);
                    }
                }
                (State::Expr, '{') => {
                    current_expr.push(chr);
                    inexpr_open_count += 1;
                }
                (State::Expr, _) => {
                    current_expr.push(chr);
                }
                (State::FormatOpts, '}') => {
                    self.output.push(chr);
                    self.state = State::NotExpr;
                }
                (State::FormatOpts, _) => {
                    self.output.push(chr);
                }
            }
        }

        if self.state != State::NotExpr {
            return Err(());
        }

        self.parsed = true;
        Ok(())
    }

    pub(crate) fn into_suggestion(&self, macro_name: &str) -> String {
        assert!(self.parsed, "Attempt to get a suggestion from not parsed expression");

        let expressions_as_string = self.extracted_expressions.join(", ");
        format!(r#"{}("{}", {})"#, macro_name, self.output, expressions_as_string)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use expect_test::{expect, Expect};

    fn check(input: &str, expect: &Expect) {
        let mut parser = FormatStrParser::new((*input).to_owned());
        let outcome_repr = if parser.parse().is_ok() {
            // Parsing should be OK, expected repr is "string; expr_1, expr_2".
            if parser.extracted_expressions.is_empty() {
                parser.output
            } else {
                format!("{}; {}", parser.output, parser.extracted_expressions.join(", "))
            }
        } else {
            // Parsing should fail, expected repr is "-".
            "-".to_owned()
        };

        expect.assert_eq(&outcome_repr);
    }

    #[test]
    fn format_str_parser() {
        let test_vector = &[
            ("no expressions", expect![["no expressions"]]),
            ("{expr} is {2 + 2}", expect![["{} is {}; expr, 2 + 2"]]),
            ("{expr:?}", expect![["{:?}; expr"]]),
            ("{malformed", expect![["-"]]),
            ("malformed}", expect![["-"]]),
            ("{{correct", expect![["{{correct"]]),
            ("correct}}", expect![["correct}}"]]),
            ("{correct}}}", expect![["{}}}; correct"]]),
            ("{correct}}}}}", expect![["{}}}}}; correct"]]),
            ("{incorrect}}", expect![["-"]]),
            ("placeholders {} {}", expect![["placeholders {} {}; $1, $2"]]),
            ("mixed {} {2 + 2} {}", expect![["mixed {} {} {}; $1, 2 + 2, $2"]]),
            (
                "{SomeStruct { val_a: 0, val_b: 1 }}",
                expect![["{}; SomeStruct { val_a: 0, val_b: 1 }"]],
            ),
            ("{expr:?} is {2.32f64:.5}", expect![["{:?} is {:.5}; expr, 2.32f64"]]),
            (
                "{SomeStruct { val_a: 0, val_b: 1 }:?}",
                expect![["{:?}; SomeStruct { val_a: 0, val_b: 1 }"]],
            ),
            ("{     2 + 2        }", expect![["{}; 2 + 2"]]),
            ("{strsim::jaro_winkle(a)}", expect![["{}; strsim::jaro_winkle(a)"]]),
            ("{foo::bar::baz()}", expect![["{}; foo::bar::baz()"]]),
            ("{foo::bar():?}", expect![["{:?}; foo::bar()"]]),
        ];

        for (input, output) in test_vector {
            check(input, output)
        }
    }

    #[test]
    fn test_into_suggestion() {
        let test_vector = &[
            ("println!", "{}", r#"println!("{}", $1)"#),
            ("eprintln!", "{}", r#"eprintln!("{}", $1)"#),
            (
                "log::info!",
                "{} {expr} {} {2 + 2}",
                r#"log::info!("{} {} {} {}", $1, expr, $2, 2 + 2)"#,
            ),
            ("format!", "{expr:?}", r#"format!("{:?}", expr)"#),
        ];

        for (kind, input, output) in test_vector {
            let mut parser = FormatStrParser::new((*input).to_owned());
            parser.parse().expect("Parsing must succeed");

            assert_eq!(&parser.into_suggestion(*kind), output);
        }
    }
}
<|endoftext|>


<commit_before>/// Easy to use Let's Encrypt client to issue and renew TLS certs

extern crate acme_client;
extern crate clap;
extern crate env_logger;

use std::io::{self, Write};
use acme_client::Directory;
use acme_client::error::Result;
use clap::{Arg, App, SubCommand, ArgMatches};

fn main() {
    let matches = App::new(env!("CARGO_PKG_NAME"))
        .version(env!("CARGO_PKG_VERSION"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .usage("letsencrypt-rs sign -D example.org -P /var/www -k domain.key -o domain.crt\
                \n    letsencrypt-rs revoke -K user_or_domain.key -C signed.crt")
        .subcommand(SubCommand::with_name("sign")
            .about("Signs a certificate")
            .display_order(1)
            .arg(Arg::with_name("USER_KEY_PATH")
                .help("User private key path to use it in account registration.")
                .long("user-key")
                .short("U")
                .takes_value(true))
            .arg(Arg::with_name("DOMAIN_KEY_PATH")
                .help("Domain private key path to use it in CSR generation.")
                .short("K")
                .long("domain-key")
                .takes_value(true))
            .arg(Arg::with_name("DOMAIN")
                .help("Domain name to obtain certificate. You can use more than one domain name.")
                .short("D")
                .long("domain")
                .multiple(true)
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("PUBLIC_DIR")
                .help("Directory to save ACME simple http challenge. This option is required.")
                .short("P")
                .long("public-dir")
                .takes_value(true))
            .arg(Arg::with_name("EMAIL")
                .help("Contact email address (optional).")
                .short("E")
                .long("email")
                .takes_value(true))
            .arg(Arg::with_name("DOMAIN_CSR")
                .help("Path to domain certificate signing request.")
                .short("C")
                .long("domain-csr")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_USER_KEY")
                .help("Path to save private user key.")
                .long("save-user-key")
                .short("u")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_DOMAIN_KEY")
                .help("Path to save domain private key.")
                .short("k")
                .long("save-domain-key")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_DOMAIN_CSR")
                .help("Path to save domain certificate signing request.")
                .long("save-csr")
                .short("S")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_SIGNED_CERTIFICATE")
                .help("Path to save signed certificate. Default is STDOUT.")
                .short("o")
                .long("save-crt")
                .takes_value(true))
            .arg(Arg::with_name("CHAIN")
                .help("Chains the signed certificate with Let's Encrypt Authority X3 \
                       (IdenTrust cross-signed) intermediate certificate.")
                .short("c")
                .long("chain")
                .takes_value(false))
            .arg(Arg::with_name("DNS_CHALLENGE")
                 .help("Use DNS challenge instead of HTTP. This option requires user \
                        to generate a TXT record for domain")
                 .short("d")
                 .long("dns")
                 .takes_value(false)))
        .subcommand(SubCommand::with_name("revoke")
            .about("Revokes a signed certificate")
            .display_order(2)
            .arg(Arg::with_name("USER_KEY")
                .help("User or domain private key path.")
                .long("user-key")
                .short("K")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("SIGNED_CRT")
                .help("Path to signed domain certificate to revoke.")
                .long("signed-crt")
                .short("C")
                .required(true)
                .takes_value(true)))
        .arg(Arg::with_name("verbose")
             .help("Show verbose output")
             .short("v")
             .multiple(true))
        .get_matches();

    init_logger(matches.occurrences_of("verbose"));

    let res = if let Some(matches) = matches.subcommand_matches("sign") {
        // TODO: remove unwrap
        sign_certificate(matches)
    } else if let Some(matches) = matches.subcommand_matches("revoke") {
        revoke_certificate(matches)
    } else {
        Err(matches.usage().into())
    };

    if let Err(e) = res {
        writeln!(io::stderr(), "{}", e).expect("Failed to write stderr");
    }
}

fn sign_certificate(matches: &ArgMatches) -> Result<()> {
    let directory = Directory::lets_encrypt()?;

    let mut account_registration = directory.account_registration();

    if let Some(email) = matches.value_of("EMAIL") {
        account_registration = account_registration.email(email);
    }

    if let Some(user_key_path) = matches.value_of("USER_KEY_PATH") {
        account_registration = account_registration.pkey_from_file(user_key_path)?;
    }

    let account = account_registration.register()?;
    let domains: Vec<_> = matches.values_of("DOMAIN")
        .ok_or("You need to provide at least one domain name")?
        .collect();
    for domain in &domains {
        let authorization = account.authorization(domain)?;
        if !matches.is_present("DNS_CHALLENGE") {
            let challenge = authorization.get_http_challenge().ok_or("HTTP challenge not found")?;
            challenge.save_key_authorization(matches.value_of("PUBLIC_DIR")
                                                 .ok_or("--public-dir not defined. \
                                                            You need to define a public \
                                                            directory to use http challenge \
                                                            verification")?)?;
            challenge.validate()?;
        } else {
            let challenge = authorization.get_dns_challenge().ok_or("DNS challenge not found")?;
            println!("Please create a TXT record for _acme-challenge.{}: {}\n\
                      Press enter to continue",
                     domain,
                     challenge.signature()?);
            io::stdin().read_line(&mut String::new()).unwrap();
            challenge.validate()?;
        }
    }

    let mut certificate_signer = account.certificate_signer(&domains);

    if let Some(domain_key_path) = matches.value_of("DOMAIN_KEY_PATH") {
        if let Some(csr_path) = matches.value_of("DOMAIN_CSR") {
            certificate_signer = certificate_signer.csr_from_file(domain_key_path, csr_path)?;
        } else {
            certificate_signer = certificate_signer.pkey_from_file(domain_key_path)?;
        }
    }

    let certificate = certificate_signer.sign_certificate()?;
    let signed_certificate_path = matches.value_of("SAVE_SIGNED_CERTIFICATE")
        .ok_or("You need to save signed certificate")?;
    if matches.is_present("CHAIN") {
        certificate.save_signed_certificate_and_chain(None, signed_certificate_path)?;
    } else {
        certificate.save_signed_certificate(signed_certificate_path)?;
    }

    if let Some(path) = matches.value_of("SAVE_DOMAIN_KEY") {
        certificate.save_private_key(path)?;
    }
    if let Some(path) = matches.value_of("SAVE_DOMAIN_CSR") {
        certificate.save_csr(path)?;
    }
    if let Some(path) = matches.value_of("SAVE_USER_KEY") {
        account.save_private_key(path)?;
    }

    Ok(())
}

fn revoke_certificate(matches: &ArgMatches) -> Result<()> {
    let directory = Directory::lets_encrypt()?;
    let account = directory.account_registration()
        .pkey_from_file(matches.value_of("USER_KEY_PATH")
                            .ok_or("You need to provide user \
                                   or domain private key used \
                                   to sign certificate.")?)?
        .register()?;
    account.revoke_certificate_from_file(matches.value_of("SIGNED_CRT")
                                             .ok_or("You need to provide \
                                                    a signed certificate to \
                                                    revoke.")?)?;
    Ok(())
}

fn init_logger(level: u64) {
    let level = match level {
        0 => "",
        1 => "acme_client=info",
        _ => "acme_client=debug",
    };
    let mut builder = env_logger::LogBuilder::new();
    builder.parse(&::std::env::var("RUST_LOG").unwrap_or(level.to_owned()));
    let _ = builder.init();
}
<commit_msg>Print error description and exit with a code<commit_after>/// Easy to use Let's Encrypt client to issue and renew TLS certs

extern crate acme_client;
extern crate clap;
extern crate env_logger;

use std::io::{self, Write};
use acme_client::Directory;
use acme_client::error::Result;
use clap::{Arg, App, SubCommand, ArgMatches};

fn main() {
    let matches = App::new(env!("CARGO_PKG_NAME"))
        .version(env!("CARGO_PKG_VERSION"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .usage("letsencrypt-rs sign -D example.org -P /var/www -k domain.key -o domain.crt\
                \n    letsencrypt-rs revoke -K user_or_domain.key -C signed.crt")
        .subcommand(SubCommand::with_name("sign")
            .about("Signs a certificate")
            .display_order(1)
            .arg(Arg::with_name("USER_KEY_PATH")
                .help("User private key path to use it in account registration.")
                .long("user-key")
                .short("U")
                .takes_value(true))
            .arg(Arg::with_name("DOMAIN_KEY_PATH")
                .help("Domain private key path to use it in CSR generation.")
                .short("K")
                .long("domain-key")
                .takes_value(true))
            .arg(Arg::with_name("DOMAIN")
                .help("Domain name to obtain certificate. You can use more than one domain name.")
                .short("D")
                .long("domain")
                .multiple(true)
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("PUBLIC_DIR")
                .help("Directory to save ACME simple http challenge. This option is required.")
                .short("P")
                .long("public-dir")
                .takes_value(true))
            .arg(Arg::with_name("EMAIL")
                .help("Contact email address (optional).")
                .short("E")
                .long("email")
                .takes_value(true))
            .arg(Arg::with_name("DOMAIN_CSR")
                .help("Path to domain certificate signing request.")
                .short("C")
                .long("domain-csr")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_USER_KEY")
                .help("Path to save private user key.")
                .long("save-user-key")
                .short("u")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_DOMAIN_KEY")
                .help("Path to save domain private key.")
                .short("k")
                .long("save-domain-key")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_DOMAIN_CSR")
                .help("Path to save domain certificate signing request.")
                .long("save-csr")
                .short("S")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_SIGNED_CERTIFICATE")
                .help("Path to save signed certificate. Default is STDOUT.")
                .short("o")
                .long("save-crt")
                .takes_value(true))
            .arg(Arg::with_name("CHAIN")
                .help("Chains the signed certificate with Let's Encrypt Authority X3 \
                       (IdenTrust cross-signed) intermediate certificate.")
                .short("c")
                .long("chain")
                .takes_value(false))
            .arg(Arg::with_name("DNS_CHALLENGE")
                 .help("Use DNS challenge instead of HTTP. This option requires user \
                        to generate a TXT record for domain")
                 .short("d")
                 .long("dns")
                 .takes_value(false)))
        .subcommand(SubCommand::with_name("revoke")
            .about("Revokes a signed certificate")
            .display_order(2)
            .arg(Arg::with_name("USER_KEY")
                .help("User or domain private key path.")
                .long("user-key")
                .short("K")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("SIGNED_CRT")
                .help("Path to signed domain certificate to revoke.")
                .long("signed-crt")
                .short("C")
                .required(true)
                .takes_value(true)))
        .arg(Arg::with_name("verbose")
             .help("Show verbose output")
             .short("v")
             .multiple(true))
        .get_matches();

    init_logger(matches.occurrences_of("verbose"));

    let res = if let Some(matches) = matches.subcommand_matches("sign") {
        sign_certificate(matches)
    } else if let Some(matches) = matches.subcommand_matches("revoke") {
        revoke_certificate(matches)
    } else {
        println!("{}", matches.usage());
        Ok(())
    };

    if let Err(e) = res {
        writeln!(io::stderr(), "{}", e.description()).expect("Failed to write stderr");
        ::std::process::exit(1);
    }
}

fn sign_certificate(matches: &ArgMatches) -> Result<()> {
    let directory = Directory::lets_encrypt()?;

    let mut account_registration = directory.account_registration();

    if let Some(email) = matches.value_of("EMAIL") {
        account_registration = account_registration.email(email);
    }

    if let Some(user_key_path) = matches.value_of("USER_KEY_PATH") {
        account_registration = account_registration.pkey_from_file(user_key_path)?;
    }

    let account = account_registration.register()?;
    let domains: Vec<_> = matches.values_of("DOMAIN")
        .ok_or("You need to provide at least one domain name")?
        .collect();
    for domain in &domains {
        let authorization = account.authorization(domain)?;
        if !matches.is_present("DNS_CHALLENGE") {
            let challenge = authorization.get_http_challenge().ok_or("HTTP challenge not found")?;
            challenge.save_key_authorization(matches.value_of("PUBLIC_DIR")
                                                 .ok_or("--public-dir not defined. \
                                                            You need to define a public \
                                                            directory to use http challenge \
                                                            verification")?)?;
            challenge.validate()?;
        } else {
            let challenge = authorization.get_dns_challenge().ok_or("DNS challenge not found")?;
            println!("Please create a TXT record for _acme-challenge.{}: {}\n\
                      Press enter to continue",
                     domain,
                     challenge.signature()?);
            io::stdin().read_line(&mut String::new()).unwrap();
            challenge.validate()?;
        }
    }

    let mut certificate_signer = account.certificate_signer(&domains);

    if let Some(domain_key_path) = matches.value_of("DOMAIN_KEY_PATH") {
        if let Some(csr_path) = matches.value_of("DOMAIN_CSR") {
            certificate_signer = certificate_signer.csr_from_file(domain_key_path, csr_path)?;
        } else {
            certificate_signer = certificate_signer.pkey_from_file(domain_key_path)?;
        }
    }

    let certificate = certificate_signer.sign_certificate()?;
    let signed_certificate_path = matches.value_of("SAVE_SIGNED_CERTIFICATE")
        .ok_or("You need to save signed certificate")?;
    if matches.is_present("CHAIN") {
        certificate.save_signed_certificate_and_chain(None, signed_certificate_path)?;
    } else {
        certificate.save_signed_certificate(signed_certificate_path)?;
    }

    if let Some(path) = matches.value_of("SAVE_DOMAIN_KEY") {
        certificate.save_private_key(path)?;
    }
    if let Some(path) = matches.value_of("SAVE_DOMAIN_CSR") {
        certificate.save_csr(path)?;
    }
    if let Some(path) = matches.value_of("SAVE_USER_KEY") {
        account.save_private_key(path)?;
    }

    Ok(())
}

fn revoke_certificate(matches: &ArgMatches) -> Result<()> {
    let directory = Directory::lets_encrypt()?;
    let account = directory.account_registration()
        .pkey_from_file(matches.value_of("USER_KEY_PATH")
                            .ok_or("You need to provide user \
                                   or domain private key used \
                                   to sign certificate.")?)?
        .register()?;
    account.revoke_certificate_from_file(matches.value_of("SIGNED_CRT")
                                             .ok_or("You need to provide \
                                                    a signed certificate to \
                                                    revoke.")?)?;
    Ok(())
}

fn init_logger(level: u64) {
    let level = match level {
        0 => "",
        1 => "acme_client=info",
        _ => "acme_client=debug",
    };
    let mut builder = env_logger::LogBuilder::new();
    builder.parse(&::std::env::var("RUST_LOG").unwrap_or(level.to_owned()));
    let _ = builder.init();
}
<|endoftext|>