<commit_before>#![stable(feature = "core_hint", since = "1.27.0")]

//! Hints to compiler that affects how code should be emitted or optimized.
//! Hints may be compile time or runtime.

use crate::intrinsics;

/// Informs the compiler that this point in the code is not reachable, enabling
/// further optimizations.
///
/// # Safety
///
/// Reaching this function is completely *undefined behavior* (UB). In
/// particular, the compiler assumes that all UB must never happen, and
/// therefore will eliminate all branches that reach to a call to
/// `unreachable_unchecked()`.
///
/// Like all instances of UB, if this assumption turns out to be wrong, i.e., the
/// `unreachable_unchecked()` call is actually reachable among all possible
/// control flow, the compiler will apply the wrong optimization strategy, and
/// may sometimes even corrupt seemingly unrelated code, causing
/// difficult-to-debug problems.
///
/// Use this function only when you can prove that the code will never call it.
/// Otherwise, consider using the [`unreachable!`] macro, which does not allow
/// optimizations but will panic when executed.
///
/// # Example
///
/// ```
/// fn div_1(a: u32, b: u32) -> u32 {
///     use std::hint::unreachable_unchecked;
///
///     // `b.saturating_add(1)` is always positive (not zero),
///     // hence `checked_div` will never return `None`.
///     // Therefore, the else branch is unreachable.
///     a.checked_div(b.saturating_add(1))
///         .unwrap_or_else(|| unsafe { unreachable_unchecked() })
/// }
///
/// assert_eq!(div_1(7, 0), 7);
/// assert_eq!(div_1(9, 1), 4);
/// assert_eq!(div_1(11, u32::MAX), 0);
/// ```
#[inline]
#[stable(feature = "unreachable", since = "1.27.0")]
#[rustc_const_unstable(feature = "const_unreachable_unchecked", issue = "53188")]
pub const unsafe fn unreachable_unchecked() -> ! {
    // SAFETY: the safety contract for `intrinsics::unreachable` must
    // be upheld by the caller.
    unsafe { intrinsics::unreachable() }
}

/// Emits a machine instruction to signal the processor that it is running in
/// a busy-wait spin-loop ("spin lock").
///
/// Upon receiving the spin-loop signal the processor can optimize its behavior by,
/// for example, saving power or switching hyper-threads.
///
/// This function is different from [`thread::yield_now`] which directly
/// yields to the system's scheduler, whereas `spin_loop` does not interact
/// with the operating system.
///
/// A common use case for `spin_loop` is implementing bounded optimistic
/// spinning in a CAS loop in synchronization primitives. To avoid problems
/// like priority inversion, it is strongly recommended that the spin loop is
/// terminated after a finite amount of iterations and an appropriate blocking
/// syscall is made.
///
/// **Note**: On platforms that do not support receiving spin-loop hints this
/// function does not do anything at all.
///
/// # Examples
///
/// ```
/// use std::sync::atomic::{AtomicBool, Ordering};
/// use std::sync::Arc;
/// use std::{hint, thread};
///
/// // A shared atomic value that threads will use to coordinate
/// let live = Arc::new(AtomicBool::new(false));
///
/// // In a background thread we'll eventually set the value
/// let bg_work = {
///     let live = live.clone();
///     thread::spawn(move || {
///         // Do some work, then make the value live
///         do_some_work();
///         live.store(true, Ordering::Release);
///     })
/// };
///
/// // Back on our current thread, we wait for the value to be set
/// while live.load(Ordering::Acquire) {
///     // The spin loop is a hint to the CPU that we're waiting, but probably
///     // not for very long
///     hint::spin_loop();
/// }
///
/// // The value is now set
/// # fn do_some_work() {}
/// do_some_work();
/// bg_work.join()?;
/// # Ok::<(), Box<dyn core::any::Any + Send + 'static>>(())
/// ```
///
/// [`thread::yield_now`]: ../../std/thread/fn.yield_now.html
#[inline]
#[stable(feature = "renamed_spin_loop", since = "1.49.0")]
pub fn spin_loop() {
    #[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "sse2"))]
    {
        #[cfg(target_arch = "x86")]
        {
            // SAFETY: the `cfg` attr ensures that we only execute this on x86 targets.
            unsafe { crate::arch::x86::_mm_pause() };
        }

        #[cfg(target_arch = "x86_64")]
        {
            // SAFETY: the `cfg` attr ensures that we only execute this on x86_64 targets.
            unsafe { crate::arch::x86_64::_mm_pause() };
        }
    }

    #[cfg(any(target_arch = "aarch64", all(target_arch = "arm", target_feature = "v6")))]
    {
        #[cfg(target_arch = "aarch64")]
        {
            // SAFETY: the `cfg` attr ensures that we only execute this on aarch64 targets.
            unsafe { crate::arch::aarch64::__yield() };
        }
        #[cfg(target_arch = "arm")]
        {
            // SAFETY: the `cfg` attr ensures that we only execute this on arm targets
            // with support for the v6 feature.
            unsafe { crate::arch::arm::__yield() };
        }
    }
}

/// An identity function that *__hints__* to the compiler to be maximally pessimistic about what
/// `black_box` could do.
///
/// Unlike [`std::convert::identity`], a Rust compiler is encouraged to assume that `black_box` can
/// use `dummy` in any possible valid way that Rust code is allowed to without introducing undefined
/// behavior in the calling code. This property makes `black_box` useful for writing code in which
/// certain optimizations are not desired, such as benchmarks.
///
/// Note however, that `black_box` is only (and can only be) provided on a "best-effort" basis. The
/// extent to which it can block optimisations may vary depending upon the platform and code-gen
/// backend used. Programs cannot rely on `black_box` for *correctness* in any way.
///
/// [`std::convert::identity`]: crate::convert::identity
#[cfg_attr(not(miri), inline)]
#[cfg_attr(miri, inline(never))]
#[unstable(feature = "test", issue = "50297")]
#[cfg_attr(miri, allow(unused_mut))]
pub fn black_box<T>(mut dummy: T) -> T {
    // We need to "use" the argument in some way LLVM can't introspect, and on
    // targets that support it we can typically leverage inline assembly to do
    // this. LLVM's interpretation of inline assembly is that it's, well, a black
    // box. This isn't the greatest implementation since it probably deoptimizes
    // more than we want, but it's so far good enough.

    #[cfg(not(miri))] // This is just a hint, so it is fine to skip in Miri.
    // SAFETY: the inline assembly is a no-op.
    unsafe {
        // FIXME: Cannot use `asm!` because it doesn't support MIPS and other architectures.
        llvm_asm!("" : : "r"(&mut dummy) : "memory" : "volatile");
    }

    dummy
}
<commit_msg>Rollup merge of #80656 - booleancoercion:master, r=sfackler<commit_after>#![stable(feature = "core_hint", since = "1.27.0")]

//! Hints to compiler that affects how code should be emitted or optimized.
//! Hints may be compile time or runtime.

use crate::intrinsics;

/// Informs the compiler that this point in the code is not reachable, enabling
/// further optimizations.
///
/// # Safety
///
/// Reaching this function is completely *undefined behavior* (UB). In
/// particular, the compiler assumes that all UB must never happen, and
/// therefore will eliminate all branches that reach to a call to
/// `unreachable_unchecked()`.
///
/// Like all instances of UB, if this assumption turns out to be wrong, i.e., the
/// `unreachable_unchecked()` call is actually reachable among all possible
/// control flow, the compiler will apply the wrong optimization strategy, and
/// may sometimes even corrupt seemingly unrelated code, causing
/// difficult-to-debug problems.
///
/// Use this function only when you can prove that the code will never call it.
/// Otherwise, consider using the [`unreachable!`] macro, which does not allow
/// optimizations but will panic when executed.
///
/// # Example
///
/// ```
/// fn div_1(a: u32, b: u32) -> u32 {
///     use std::hint::unreachable_unchecked;
///
///     // `b.saturating_add(1)` is always positive (not zero),
///     // hence `checked_div` will never return `None`.
///     // Therefore, the else branch is unreachable.
///     a.checked_div(b.saturating_add(1))
///         .unwrap_or_else(|| unsafe { unreachable_unchecked() })
/// }
///
/// assert_eq!(div_1(7, 0), 7);
/// assert_eq!(div_1(9, 1), 4);
/// assert_eq!(div_1(11, u32::MAX), 0);
/// ```
#[inline]
#[stable(feature = "unreachable", since = "1.27.0")]
#[rustc_const_unstable(feature = "const_unreachable_unchecked", issue = "53188")]
pub const unsafe fn unreachable_unchecked() -> ! {
    // SAFETY: the safety contract for `intrinsics::unreachable` must
    // be upheld by the caller.
    unsafe { intrinsics::unreachable() }
}

/// Emits a machine instruction to signal the processor that it is running in
/// a busy-wait spin-loop ("spin lock").
///
/// Upon receiving the spin-loop signal the processor can optimize its behavior by,
/// for example, saving power or switching hyper-threads.
///
/// This function is different from [`thread::yield_now`] which directly
/// yields to the system's scheduler, whereas `spin_loop` does not interact
/// with the operating system.
///
/// A common use case for `spin_loop` is implementing bounded optimistic
/// spinning in a CAS loop in synchronization primitives. To avoid problems
/// like priority inversion, it is strongly recommended that the spin loop is
/// terminated after a finite amount of iterations and an appropriate blocking
/// syscall is made.
///
/// **Note**: On platforms that do not support receiving spin-loop hints this
/// function does not do anything at all.
///
/// # Examples
///
/// ```
/// use std::sync::atomic::{AtomicBool, Ordering};
/// use std::sync::Arc;
/// use std::{hint, thread};
///
/// // A shared atomic value that threads will use to coordinate
/// let live = Arc::new(AtomicBool::new(false));
///
/// // In a background thread we'll eventually set the value
/// let bg_work = {
///     let live = live.clone();
///     thread::spawn(move || {
///         // Do some work, then make the value live
///         do_some_work();
///         live.store(true, Ordering::Release);
///     })
/// };
///
/// // Back on our current thread, we wait for the value to be set
/// while !live.load(Ordering::Acquire) {
///     // The spin loop is a hint to the CPU that we're waiting, but probably
///     // not for very long
///     hint::spin_loop();
/// }
///
/// // The value is now set
/// # fn do_some_work() {}
/// do_some_work();
/// bg_work.join()?;
/// # Ok::<(), Box<dyn core::any::Any + Send + 'static>>(())
/// ```
///
/// [`thread::yield_now`]: ../../std/thread/fn.yield_now.html
#[inline]
#[stable(feature = "renamed_spin_loop", since = "1.49.0")]
pub fn spin_loop() {
    #[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "sse2"))]
    {
        #[cfg(target_arch = "x86")]
        {
            // SAFETY: the `cfg` attr ensures that we only execute this on x86 targets.
            unsafe { crate::arch::x86::_mm_pause() };
        }

        #[cfg(target_arch = "x86_64")]
        {
            // SAFETY: the `cfg` attr ensures that we only execute this on x86_64 targets.
            unsafe { crate::arch::x86_64::_mm_pause() };
        }
    }

    #[cfg(any(target_arch = "aarch64", all(target_arch = "arm", target_feature = "v6")))]
    {
        #[cfg(target_arch = "aarch64")]
        {
            // SAFETY: the `cfg` attr ensures that we only execute this on aarch64 targets.
            unsafe { crate::arch::aarch64::__yield() };
        }
        #[cfg(target_arch = "arm")]
        {
            // SAFETY: the `cfg` attr ensures that we only execute this on arm targets
            // with support for the v6 feature.
            unsafe { crate::arch::arm::__yield() };
        }
    }
}

/// An identity function that *__hints__* to the compiler to be maximally pessimistic about what
/// `black_box` could do.
///
/// Unlike [`std::convert::identity`], a Rust compiler is encouraged to assume that `black_box` can
/// use `dummy` in any possible valid way that Rust code is allowed to without introducing undefined
/// behavior in the calling code. This property makes `black_box` useful for writing code in which
/// certain optimizations are not desired, such as benchmarks.
///
/// Note however, that `black_box` is only (and can only be) provided on a "best-effort" basis. The
/// extent to which it can block optimisations may vary depending upon the platform and code-gen
/// backend used. Programs cannot rely on `black_box` for *correctness* in any way.
///
/// [`std::convert::identity`]: crate::convert::identity
#[cfg_attr(not(miri), inline)]
#[cfg_attr(miri, inline(never))]
#[unstable(feature = "test", issue = "50297")]
#[cfg_attr(miri, allow(unused_mut))]
pub fn black_box<T>(mut dummy: T) -> T {
    // We need to "use" the argument in some way LLVM can't introspect, and on
    // targets that support it we can typically leverage inline assembly to do
    // this. LLVM's interpretation of inline assembly is that it's, well, a black
    // box. This isn't the greatest implementation since it probably deoptimizes
    // more than we want, but it's so far good enough.

    #[cfg(not(miri))] // This is just a hint, so it is fine to skip in Miri.
    // SAFETY: the inline assembly is a no-op.
    unsafe {
        // FIXME: Cannot use `asm!` because it doesn't support MIPS and other architectures.
        llvm_asm!("" : : "r"(&mut dummy) : "memory" : "volatile");
    }

    dummy
}
<|endoftext|>


<commit_before>use memory::Memory;

static NMI_VECTOR: u16 = 0xfffa;
static RESET_VECTOR: u16 = 0xfffc;
static IRQ_VECTOR: u16 = 0xfffe;

struct Registers {
	pc: u16,				// program counter
	ac: u8,					// accumulator
	x: u8,					// x register
	y: u8,					// y register
	sr: u8,					// status register (NV-BDIZC: Negative, oVerflow, Break, Decimal, Interrupt, Zero, Carry)
	sp: u8,					// stack pointer
}

enum Operand {
	Implied,							// OPC				// operand implied
	Immediate(u8),						// OPC #$BB			// operand is byte (BB)
	Accumulator,						// OPC A			// operand is AC
	Relative(i8),						// OPC $BB			// branch target is PC + offset (BB), bit 7 signifies negative offset
	Absolute(u16),						// OPC $HHLL		// operand is address $HHLL
	AbsoluteIndexedWithX(u16),			// OPC $HHLL,X		// operand is address $HHLL incremented by X with carry
	AbsoluteIndexedWithY(u16),			// OPC $HHLL,Y		// operand is address $HHLL incremented by Y with carry
	Indirect(u16),						// OPC ($HHLL)		// operand is effective address; effective address is value of address
	ZeroPage(u8),						// OPC $LL			// operand is address $00LL
	ZeroPageIndexedWithX(u8),			// OPC $LL,X		// operand is address $00LL incremented by X, no page transition
	ZeroPageIndexedWithY(u8),			// OPC $LL,Y		// operand is address $00LL incremented by Y, no page transition
	ZeroPageIndexedWithXIndirect(u8),	// OPC ($LL,X)		// operand is effective address; effective address is $00LL incremented by X without carry
	ZeroPageIndirectIndexedWithY(u8),	// OPC ($LL),Y		// operand is effective address incremented by Y with carry; effective address is word at $00LL
}

impl Operand {
	fn addr (&self, cpu: &Mos6502) -> u16 {
		match * self {
			Implied								=> fail!("mos6510: Implied operand is never targetted to an address"),
			Immediate(_)						=> fail!("mos6510: Immediade operand is never targetted to an address"),
			Accumulator							=> fail!("mos6510: Accumulator operand is never targetted to an address"),
			Relative(offset)					=> cpu.reg.pc + offset as u16,
			Absolute(addr)						=> addr,
			AbsoluteIndexedWithX(addr)			=> addr + cpu.reg.x as u16,
			AbsoluteIndexedWithY(addr)			=> addr + cpu.reg.y as u16,
			Indirect(addr)						=> cpu.mem.get_le(addr),
			ZeroPage(addr)						=> addr as u16,
			ZeroPageIndexedWithX(addr)			=> (addr + cpu.reg.x) as u16,					// no page transition
			ZeroPageIndexedWithY(addr)			=> (addr + cpu.reg.y) as u16,					// no page transition
			ZeroPageIndexedWithXIndirect(addr)	=> cpu.mem.get_le((addr + cpu.reg.x) as u16),	// no page transition
			ZeroPageIndirectIndexedWithY(addr)	=> { let iaddr: u16 = cpu.mem.get_le(addr as u16); iaddr + cpu.reg.y as u16 },
		}
	}

	fn get (&self, cpu: &Mos6502) -> u8 {
		match *self {
			Implied								=> fail!("mos6510: Implied operand never has a value"),
			Immediate(val)						=> val,
			Accumulator							=> cpu.reg.ac,
			Relative(_target)					=> fail!("mos6510: Relative operand never has a value"),
			op									=> { let addr = op.addr(cpu); cpu.mem.get(addr) },
		}
	}

	fn set (&self, cpu: &mut Mos6502, val: u8) {
		match *self {
			Implied								=> fail!("mos6510: Implied operand never sets a value"),
			Immediate(_)						=> fail!("mos6510: Immediate operand never sets a value"),
			Accumulator							=> cpu.reg.ac = val,
			Relative(_target)					=> fail!("mos6510: Relative operand never sets a value"),
			op									=> { let addr = op.addr(cpu); cpu.mem.set(addr, val); },
		}
	}
}

pub struct Mos6502<'self> {
	priv reg: Registers,					// internal CPU registers
	priv mem: &'self Memory<'self, u16>,	// memory as accessible to the CPU
}

impl<'self> Mos6502<'self> {
	pub fn new (mem: &'self Memory<u16>) -> Mos6502<'self> {
		Mos6502 {
			reg: Registers { pc: 0, ac: 0, x: 0, y: 0, sr: 0, sp: 0 },
			mem: mem,
		}
	}
}

pub struct Mos6510<'self> {
	priv cpu: Mos6502<'self>,				// Core CPU
	priv port_ddr: u8,						// CPU port data direction register
	priv port_dat: u8,						// CPU port data register
}

impl<'self> Mos6510<'self> {
	pub fn new (mem: &'self Memory<u16>) -> Mos6510<'self> {
		// TODO: addresses $0000 (data direction) and $0001 (data) are hardwired for the processor I/O port
		Mos6510 { cpu: Mos6502::new(mem), port_ddr: 0, port_dat: 0 }
	}
}
<commit_msg>Use type parameters to make CPU accept any addressable memory object<commit_after>use addressable::Addressable;
use addressable::AddressableUtil;

static NMI_VECTOR: u16 = 0xfffa;
static RESET_VECTOR: u16 = 0xfffc;
static IRQ_VECTOR: u16 = 0xfffe;

struct Registers {
	pc: u16,				// program counter
	ac: u8,					// accumulator
	x: u8,					// x register
	y: u8,					// y register
	sr: u8,					// status register (NV-BDIZC: Negative, oVerflow, Break, Decimal, Interrupt, Zero, Carry)
	sp: u8,					// stack pointer
}

enum Operand {
	Implied,							// OPC				// operand implied
	Immediate(u8),						// OPC #$BB			// operand is byte (BB)
	Accumulator,						// OPC A			// operand is AC
	Relative(i8),						// OPC $BB			// branch target is PC + offset (BB), bit 7 signifies negative offset
	Absolute(u16),						// OPC $HHLL		// operand is address $HHLL
	AbsoluteIndexedWithX(u16),			// OPC $HHLL,X		// operand is address $HHLL incremented by X with carry
	AbsoluteIndexedWithY(u16),			// OPC $HHLL,Y		// operand is address $HHLL incremented by Y with carry
	Indirect(u16),						// OPC ($HHLL)		// operand is effective address; effective address is value of address
	ZeroPage(u8),						// OPC $LL			// operand is address $00LL
	ZeroPageIndexedWithX(u8),			// OPC $LL,X		// operand is address $00LL incremented by X, no page transition
	ZeroPageIndexedWithY(u8),			// OPC $LL,Y		// operand is address $00LL incremented by Y, no page transition
	ZeroPageIndexedWithXIndirect(u8),	// OPC ($LL,X)		// operand is effective address; effective address is $00LL incremented by X without carry
	ZeroPageIndirectIndexedWithY(u8),	// OPC ($LL),Y		// operand is effective address incremented by Y with carry; effective address is word at $00LL
}

impl Operand {
	fn addr<M: Addressable<u16>> (&self, cpu: &Mos6502<M>) -> u16 {
		match *self {
			Implied								=> fail!("mos6510: Implied operand is never targetted to an address"),
			Immediate(_)						=> fail!("mos6510: Immediade operand is never targetted to an address"),
			Accumulator							=> fail!("mos6510: Accumulator operand is never targetted to an address"),
			Relative(offset)					=> cpu.reg.pc + offset as u16,
			Absolute(addr)						=> addr,
			AbsoluteIndexedWithX(addr)			=> addr + cpu.reg.x as u16,
			AbsoluteIndexedWithY(addr)			=> addr + cpu.reg.y as u16,
			Indirect(addr)						=> cpu.mem.get_le(addr),
			ZeroPage(addr)						=> addr as u16,
			ZeroPageIndexedWithX(addr)			=> (addr + cpu.reg.x) as u16,					// no page transition
			ZeroPageIndexedWithY(addr)			=> (addr + cpu.reg.y) as u16,					// no page transition
			ZeroPageIndexedWithXIndirect(addr)	=> cpu.mem.get_le((addr + cpu.reg.x) as u16),	// no page transition
			ZeroPageIndirectIndexedWithY(addr)	=> { let iaddr: u16 = cpu.mem.get_le(addr as u16); iaddr + cpu.reg.y as u16 },
		}
	}

	fn get<M: Addressable<u16>> (&self, cpu: &Mos6502<M>) -> u8 {
		match *self {
			Implied								=> fail!("mos6510: Implied operand never has a value"),
			Immediate(val)						=> val,
			Accumulator							=> cpu.reg.ac,
			Relative(_target)					=> fail!("mos6510: Relative operand never has a value"),
			op									=> { let addr = op.addr(cpu); cpu.mem.get(addr) },
		}
	}

	fn set<M: Addressable<u16>> (&self, cpu: &mut Mos6502<M>, val: u8) {
		match *self {
			Implied								=> fail!("mos6510: Implied operand never sets a value"),
			Immediate(_)						=> fail!("mos6510: Immediate operand never sets a value"),
			Accumulator							=> cpu.reg.ac = val,
			Relative(_target)					=> fail!("mos6510: Relative operand never sets a value"),
			op									=> { let addr = op.addr(cpu); cpu.mem.set(addr, val); },
		}
	}
}

pub struct Mos6502<M> {
	priv reg: Registers,					// internal CPU registers
	priv mem: M,							// memory as accessible to the CPU
}

impl<M: Addressable<u16>> Mos6502<M> {
	pub fn new (mem: M) -> Mos6502<M> {
		Mos6502 {
			reg: Registers { pc: 0, ac: 0, x: 0, y: 0, sr: 0, sp: 0 },
			mem: mem,
		}
	}
}

pub struct Mos6510<'self> {
	priv cpu: Mos6502<'self>,				// Core CPU
	priv port_ddr: u8,						// CPU port data direction register
	priv port_dat: u8,						// CPU port data register
}

impl<'self> Mos6510<'self> {
	pub fn new (mem: &'self Memory<u16>) -> Mos6510<'self> {
		// TODO: addresses $0000 (data direction) and $0001 (data) are hardwired for the processor I/O port
		Mos6510 { cpu: Mos6502::new(mem), port_ddr: 0, port_dat: 0 }
	}
}
<|endoftext|>


<commit_before>use std::collections::HashSet;
use std::sync::Arc;
use super::structs::AttestationRecord;
use super::AttesterMap;
use super::attestation_parent_hashes::{
    attestation_parent_hashes,
    ParentHashesError,
};
use super::db::{
    ClientDB,
    DBError
};
use super::db::stores::{
    BlockStore,
    ValidatorStore,
};
use super::utils::types::{
    Hash256,
};
use super::message_generation::generate_signed_message;
use super::signature_verification::{
    verify_aggregate_signature_for_indices,
    SignatureVerificationError,
};

#[derive(Debug,PartialEq)]
pub enum AttestationValidationError {
    SlotTooHigh,
    SlotTooLow,
    JustifiedSlotTooHigh,
    UnknownJustifiedBlock,
    TooManyObliqueHashes,
    BadCurrentHashes,
    BadObliqueHashes,
    BadAttesterMap,
    IntWrapping,
    PublicKeyCorrupt,
    NoPublicKeyForValidator,
    BadBitfieldLength,
    InvalidBitfield,
    InvalidBitfieldEndBits,
    NoSignatures,
    NonZeroTrailingBits,
    AggregateSignatureFail,
    DBError(String),
}

fn bytes_for_bits(bits: usize) -> usize {
    (bits.saturating_sub(1) / 8) + 1
}

fn any_of_last_n_bits_are_set(byte: u8, n: usize) -> bool {
    let shift = 8_u16.saturating_sub(n as u16);
    ((!0 << shift) & u16::from(byte)) > 0
}

pub fn validate_attestation<T>(a: &AttestationRecord,
                               block_slot: u64,
                               cycle_length: u8,
                               known_last_justified_slot: u64,
                               known_parent_hashes: &Arc<Vec<Hash256>>,
                               block_store: &Arc<BlockStore<T>>,
                               validator_store: &Arc<ValidatorStore<T>>,
                               attester_map: &Arc<AttesterMap>)
    -> Result<Option<HashSet<usize>>, AttestationValidationError>
    where T: ClientDB + Sized
{
    /*
     * The attesation slot must not be higher than the block that contained it.
     */
    if a.slot > block_slot {
        return Err(AttestationValidationError::SlotTooHigh);
    }

    /*
     * The slot of this attestation must not be more than cycle_length + 1 distance
     * from the block that contained it.
     *
     * The below code stays overflow-safe as long as cycle length is a < 64 bit integer.
     */
    if a.slot < block_slot.saturating_sub(u64::from(cycle_length) + 1) {
        return Err(AttestationValidationError::SlotTooLow);
    }

    /*
     * The attestation must indicate that its last justified slot is the same as the last justified
     * slot known to us.
     */
    if a.justified_slot > known_last_justified_slot {
        return Err(AttestationValidationError::JustifiedSlotTooHigh);
    }

    /*
     * There is no need to include more oblique parents hashes than there are blocks
     * in a cycle.
     */
    if a.oblique_parent_hashes.len() > usize::from(cycle_length) {
        return Err(AttestationValidationError::TooManyObliqueHashes);
    }

    /*
     * Retrieve the set of attestation indices for this slot and shard id.
     *
     * This is an array mapping the order that validators will appear in the bitfield to the
     * canonincal index of a validator.
     */
    let attestation_indices = attester_map.get(&(a.slot, a.shard_id))
        .ok_or(AttestationValidationError::BadAttesterMap)?;

    /*
     * The bitfield must be no longer than the minimum required to represent each validator in the
     * attestation indicies for this slot and shard id.
     */
    if a.attester_bitfield.num_bytes() !=
        bytes_for_bits(attestation_indices.len())
    {
        return Err(AttestationValidationError::BadBitfieldLength);
    }

    /*
     * If there are excess bits in the bitfield because the number of a validators in not a
     * multiple of 8, reject this attestation record.
     *
     * Allow extra set bits would permit mutliple different byte layouts (and therefore hashes) to
     * refer to the same AttesationRecord.
     */
    let last_byte =
        a.attester_bitfield.get_byte(a.attester_bitfield.num_bytes() - 1)
        .ok_or(AttestationValidationError::InvalidBitfield)?;
    if any_of_last_n_bits_are_set(*last_byte, a.attester_bitfield.len() % 8) {
        return Err(AttestationValidationError::InvalidBitfieldEndBits)
    }

    /*
     * The specified justified block hash must be known to us
     */
    if !block_store.block_exists(&a.justified_block_hash)? {
        return Err(AttestationValidationError::UnknownJustifiedBlock)
    }

    let signed_message = {
        let parent_hashes = attestation_parent_hashes(
            cycle_length,
            block_slot,
            a.slot,
            &known_parent_hashes,
            &a.oblique_parent_hashes)?;
        generate_signed_message(
            a.slot,
            &parent_hashes,
            a.shard_id,
            &a.shard_block_hash,
            a.justified_slot)
    };

    let voted_hashmap =
        verify_aggregate_signature_for_indices(
            &signed_message,
            &a.aggregate_sig,
            &attestation_indices,
            &a.attester_bitfield,
            &validator_store)?;

    Ok(voted_hashmap)
}

impl From<ParentHashesError> for AttestationValidationError {
    fn from(e: ParentHashesError) -> Self {
        match e {
            ParentHashesError::BadCurrentHashes
                => AttestationValidationError::BadCurrentHashes,
            ParentHashesError::BadObliqueHashes
                => AttestationValidationError::BadObliqueHashes,
            ParentHashesError::SlotTooLow
                => AttestationValidationError::SlotTooLow,
            ParentHashesError::SlotTooHigh
                => AttestationValidationError::SlotTooHigh,
            ParentHashesError::IntWrapping
                => AttestationValidationError::IntWrapping
        }
    }
}

impl From<DBError> for AttestationValidationError {
    fn from(e: DBError) -> Self {
        AttestationValidationError::DBError(e.message)
    }
}

impl From<SignatureVerificationError> for AttestationValidationError {
    fn from(e: SignatureVerificationError) -> Self {
        match e {
            SignatureVerificationError::BadValidatorIndex
                => AttestationValidationError::BadAttesterMap,
            SignatureVerificationError::PublicKeyCorrupt
                => AttestationValidationError::PublicKeyCorrupt,
            SignatureVerificationError::NoPublicKeyForValidator
                => AttestationValidationError::NoPublicKeyForValidator,
            SignatureVerificationError::DBError(s)
                => AttestationValidationError::DBError(s),
        }
    }
}
<commit_msg>Fix bug with bitfield validation<commit_after>use std::collections::HashSet;
use std::sync::Arc;
use super::structs::AttestationRecord;
use super::AttesterMap;
use super::attestation_parent_hashes::{
    attestation_parent_hashes,
    ParentHashesError,
};
use super::db::{
    ClientDB,
    DBError
};
use super::db::stores::{
    BlockStore,
    ValidatorStore,
};
use super::utils::types::{
    Hash256,
};
use super::message_generation::generate_signed_message;
use super::signature_verification::{
    verify_aggregate_signature_for_indices,
    SignatureVerificationError,
};

#[derive(Debug,PartialEq)]
pub enum AttestationValidationError {
    SlotTooHigh,
    SlotTooLow,
    JustifiedSlotTooHigh,
    UnknownJustifiedBlock,
    TooManyObliqueHashes,
    BadCurrentHashes,
    BadObliqueHashes,
    BadAttesterMap,
    IntWrapping,
    PublicKeyCorrupt,
    NoPublicKeyForValidator,
    BadBitfieldLength,
    InvalidBitfield,
    InvalidBitfieldEndBits,
    NoSignatures,
    NonZeroTrailingBits,
    AggregateSignatureFail,
    DBError(String),
}

pub fn validate_attestation<T>(a: &AttestationRecord,
                               block_slot: u64,
                               cycle_length: u8,
                               known_last_justified_slot: u64,
                               known_parent_hashes: &Arc<Vec<Hash256>>,
                               block_store: &Arc<BlockStore<T>>,
                               validator_store: &Arc<ValidatorStore<T>>,
                               attester_map: &Arc<AttesterMap>)
    -> Result<Option<HashSet<usize>>, AttestationValidationError>
    where T: ClientDB + Sized
{
    /*
     * The attesation slot must not be higher than the block that contained it.
     */
    if a.slot > block_slot {
        return Err(AttestationValidationError::SlotTooHigh);
    }

    /*
     * The slot of this attestation must not be more than cycle_length + 1 distance
     * from the block that contained it.
     *
     * The below code stays overflow-safe as long as cycle length is a < 64 bit integer.
     */
    if a.slot < block_slot.saturating_sub(u64::from(cycle_length) + 1) {
        return Err(AttestationValidationError::SlotTooLow);
    }

    /*
     * The attestation must indicate that its last justified slot is the same as the last justified
     * slot known to us.
     */
    if a.justified_slot > known_last_justified_slot {
        return Err(AttestationValidationError::JustifiedSlotTooHigh);
    }

    /*
     * There is no need to include more oblique parents hashes than there are blocks
     * in a cycle.
     */
    if a.oblique_parent_hashes.len() > usize::from(cycle_length) {
        return Err(AttestationValidationError::TooManyObliqueHashes);
    }

    /*
     * Retrieve the set of attestation indices for this slot and shard id.
     *
     * This is an array mapping the order that validators will appear in the bitfield to the
     * canonincal index of a validator.
     */
    let attestation_indices = attester_map.get(&(a.slot, a.shard_id))
        .ok_or(AttestationValidationError::BadAttesterMap)?;

    /*
     * The bitfield must be no longer than the minimum required to represent each validator in the
     * attestation indicies for this slot and shard id.
     */
    if a.attester_bitfield.num_bytes() !=
        bytes_for_bits(attestation_indices.len())
    {
        return Err(AttestationValidationError::BadBitfieldLength);
    }

    /*
     * If there are excess bits in the bitfield because the number of a validators in not a
     * multiple of 8, reject this attestation record.
     *
     * Allow extra set bits would permit mutliple different byte layouts (and therefore hashes) to
     * refer to the same AttesationRecord.
     */
    let last_byte =
        a.attester_bitfield.get_byte(a.attester_bitfield.num_bytes() - 1)
        .ok_or(AttestationValidationError::InvalidBitfield)?;
    if any_of_last_n_bits_are_set(*last_byte, a.attester_bitfield.len() % 8) {
        return Err(AttestationValidationError::InvalidBitfieldEndBits)
    }

    /*
     * The specified justified block hash must be known to us
     */
    if !block_store.block_exists(&a.justified_block_hash)? {
        return Err(AttestationValidationError::UnknownJustifiedBlock)
    }

    let signed_message = {
        let parent_hashes = attestation_parent_hashes(
            cycle_length,
            block_slot,
            a.slot,
            &known_parent_hashes,
            &a.oblique_parent_hashes)?;
        generate_signed_message(
            a.slot,
            &parent_hashes,
            a.shard_id,
            &a.shard_block_hash,
            a.justified_slot)
    };

    let voted_hashmap =
        verify_aggregate_signature_for_indices(
            &signed_message,
            &a.aggregate_sig,
            &attestation_indices,
            &a.attester_bitfield,
            &validator_store)?;

    Ok(voted_hashmap)
}

fn bytes_for_bits(bits: usize) -> usize {
    (bits.saturating_sub(1) / 8) + 1
}

fn any_of_last_n_bits_are_set(byte: u8, n: usize) -> bool {
    for i in 0..n {
        let mask = 0_u8 >> 8_usize - i as usize;
        if byte & mask > 0 {
            return true
        }
    }
    false
}

impl From<ParentHashesError> for AttestationValidationError {
    fn from(e: ParentHashesError) -> Self {
        match e {
            ParentHashesError::BadCurrentHashes
                => AttestationValidationError::BadCurrentHashes,
            ParentHashesError::BadObliqueHashes
                => AttestationValidationError::BadObliqueHashes,
            ParentHashesError::SlotTooLow
                => AttestationValidationError::SlotTooLow,
            ParentHashesError::SlotTooHigh
                => AttestationValidationError::SlotTooHigh,
            ParentHashesError::IntWrapping
                => AttestationValidationError::IntWrapping
        }
    }
}

impl From<DBError> for AttestationValidationError {
    fn from(e: DBError) -> Self {
        AttestationValidationError::DBError(e.message)
    }
}

impl From<SignatureVerificationError> for AttestationValidationError {
    fn from(e: SignatureVerificationError) -> Self {
        match e {
            SignatureVerificationError::BadValidatorIndex
                => AttestationValidationError::BadAttesterMap,
            SignatureVerificationError::PublicKeyCorrupt
                => AttestationValidationError::PublicKeyCorrupt,
            SignatureVerificationError::NoPublicKeyForValidator
                => AttestationValidationError::NoPublicKeyForValidator,
            SignatureVerificationError::DBError(s)
                => AttestationValidationError::DBError(s),
        }
    }
}
<|endoftext|>


<commit_before>// "Tifflin" Kernel
// - By John Hodge (thePowersGang)
//
// Core/main.rs
// - Kernel main
#![crate_name="kernel"]
#![crate_type="lib"]
#![feature(no_std)]
#![feature(asm)]	// Enables the asm! syntax extension
#![feature(box_syntax)]	// Enables 'box' syntax
#![feature(thread_local)]	// Allows use of thread_local
#![feature(lang_items)]	// Allow definition of lang_items
#![feature(core)]	// silences warnings about write!
#![feature(optin_builtin_traits)]	// Negative impls
#![feature(unique)]	// Unique
#![feature(slice_patterns)]	// Slice (array) destructuring patterns, used by multiboot code
#![feature(step_by)]	// Range::step_by
#![feature(linkage)]	// allows using #[linkage="external"]
#![feature(const_fn)]	// Allows defining `const fn`
#![no_std]

#![feature(plugin)]
#![feature(custom_attribute)]
#![plugin(tag_safe)]
use prelude::*;

#[macro_use]
extern crate core;

pub use arch::memory::PAGE_SIZE;

#[doc(hidden)]
#[macro_use] pub mod logmacros;
#[doc(hidden)]
#[macro_use] pub mod macros;
#[doc(hidden)]
#[macro_use] #[cfg(arch__amd64)] #[path="arch/amd64/mod-macros.rs"] pub mod arch_macros;

// Evil Hack: For some reason, write! (and friends) will expand pointing to std instead of core
#[doc(hidden)]
mod std {
	pub use core::option;
	pub use core::{default,fmt,cmp};
	pub use core::marker;	// needed for derive(Copy)
	pub use core::iter;	// needed for 'for'
}

/// Kernel's version of 'std::prelude'
pub mod prelude;

/// Library datatypes (Vec, Queue, ...)
#[macro_use]
pub mod lib;	// Clone of libstd

/// Heavy synchronisation primitives (Mutex, Semaphore, RWLock, ...)
#[macro_use]
pub mod sync;

/// Asynchrnous wait support
pub mod async;

/// Logging framework
pub mod logging;
/// Memory management (physical, virtual, heap)
pub mod memory;
/// Thread management
#[macro_use]
pub mod threads;
/// Timekeeping (timers and wall time)
pub mod time;

// Module/Executable loading (and symbol lookup)
pub mod loading;
/// Module management (loading and initialisation of kernel modules)
pub mod modules;

/// Meta devices (the Hardware Abstraction Layer)
pub mod metadevs;
/// Device to driver mapping manager
///
/// Starts driver instances for the devices it sees
pub mod device_manager;

/// User output, via a kernel-provided compositing "WM"
pub mod gui;

// Public for driver modules
pub mod vfs;

mod config;

/// Stack unwinding (panic) handling
pub mod unwind;

pub mod irqs;

/// Built-in device drivers
mod hw;

/// Achitecture-specific code - AMD64 (aka x86-64)
#[macro_use]
#[cfg(arch__amd64)] #[path="arch/amd64/mod.rs"] pub mod arch;	// Needs to be pub for exports to be avaliable

/// Kernel entrypoint
#[no_mangle]
pub extern "C" fn kmain()
{
	log_notice!("Tifflin Kernel v{} build {} starting", env!("TK_VERSION"), env!("TK_BUILD"));
	log_notice!("> Git state : {}", env!("TK_GITSPEC"));
	log_notice!("> Built with {}", env!("RUST_VERSION"));

	// Initialise core services before attempting modules
	::memory::phys::init();
	::memory::virt::init();
	::memory::heap::init();
	::threads::init();

	log_log!("Command line = '{}'", ::arch::boot::get_boot_string());
	// TODO: Parse command line into config

	// Dump active video mode
	let vidmode = ::arch::boot::get_video_mode();
	match vidmode {
	Some(m) => {
		log_debug!("Video mode : {}x{} @ {:#x}", m.width, m.height, m.base);
		::metadevs::video::set_boot_mode(m);
		},
	None => log_debug!("No video mode present")
	}

	// Modules (dependency tree included)
	// - Requests that the GUI be started as soon as possible
	::modules::init(&["GUI"]);

	// Yield to allow init threads to run
	::threads::yield_time();

	// Run system init
	sysinit();

	// Thread 0 idle loop
	log_info!("Entering idle");
	loop
	{
		log_trace!("TID0 napping");
		::threads::yield_time();
	}
}

// Initialise the system once drivers are up
fn sysinit()
{
	use metadevs::storage::VolumeHandle;
	use vfs::{mount,handle};
	use vfs::Path;

	// 1. Mount /system to the specified volume
	let sysdisk = ::config::get_string(::config::Value::SysDisk);
	match VolumeHandle::open_named(sysdisk)
	{
	Err(e) => {
		log_error!("Unable to open /system volume {}: {}", sysdisk, e);
		return ;
		},
	Ok(vh) => match mount::mount("/system".as_ref(), vh, "", &[])
		{
		Ok(_) => {},
		Err(e) => {
			log_error!("Unable to mount /system from {}: {:?}", sysdisk, e);
			return ;
			},
		},
	}

	// 2. Symbolic link /sysroot to the specified folder
	let sysroot = ::config::get_string(::config::Value::SysRoot);
	handle::Dir::open(Path::new("/")).unwrap()
		.symlink("sysroot", Path::new(sysroot)).unwrap();

	// *. Testing: open a file known to exist on the testing disk
	{
		match handle::File::open( Path::new("/system/1.TXT"), handle::FileOpenMode::SharedRO )
		{
		Err(e) => log_warning!("VFS test file can't be opened: {:?}", e),
		Ok(h) => {
			log_debug!("VFS open test = {:?}", h);
			let mut buf = [0; 16];
			let sz = h.read(0, &mut buf).unwrap();
			log_debug!("- Contents: {:?}", ::lib::RawString(&buf[..sz]));
			},
		}

		// - Iterate root dir
		match handle::Dir::open( Path::new("/") )
		{
		Err(e) => log_warning!("VFS root cannot be opened: {:?}", e),
		Ok(h) =>
			for name in h.iter() {
				log_log!("{}: {:?}", 0, name);
			},
		}
		// - Iterate /system
		match handle::Dir::open( Path::new("/system") )
		{
		Err(e) => log_warning!("System volume cannot be opened: {:?}", e),
		Ok(h) =>
			for name in h.iter() {
				log_log!("{:?}", name);
			},
		}
	}

	// *. TEST Automount
	// - Probably shouldn't be included in the final version, but works for testing filesystem and storage drivers
	handle::Dir::open( Path::new("/") ).and_then(|h| h.mkdir("mount")).unwrap();
	for (_,v) in ::metadevs::storage::enum_lvs()
	{
		let vh = match VolumeHandle::open_named(&v)
			{
			Err(e) => {
				log_log!("Unable to open '{}': {}", v, e);
				continue;
				},
			Ok(v) => v,
			};
		handle::Dir::open( Path::new("/mount") ).and_then(|h| h.mkdir(&v)).unwrap();
		let mountpt = format!("/mount/{}",v);
		match mount::mount( mountpt.as_ref(), vh, "", &[] )
		{
		Ok(_) => log_log!("Auto-mounted to {}", mountpt),
		Err(e) => log_notice!("Unable to automount '{}': {:?}", v, e),
		}

	}
}

// vim: ft=rust

<commit_msg>(minor) Kernel/main - Silence some warnings<commit_after>// "Tifflin" Kernel
// - By John Hodge (thePowersGang)
//
// Core/main.rs
// - Kernel main
#![crate_name="kernel"]
#![crate_type="lib"]
#![feature(no_std)]
#![feature(asm)]	// Enables the asm! syntax extension
#![feature(box_syntax)]	// Enables 'box' syntax
#![feature(thread_local)]	// Allows use of thread_local
#![feature(lang_items)]	// Allow definition of lang_items
#![feature(core)]	// silences warnings about write!
#![feature(optin_builtin_traits)]	// Negative impls
#![feature(unique)]	// Unique
#![feature(slice_patterns)]	// Slice (array) destructuring patterns, used by multiboot code
#![feature(step_by)]	// Range::step_by
#![feature(linkage)]	// allows using #[linkage="external"]
#![feature(const_fn)]	// Allows defining `const fn`
#![no_std]

#![feature(plugin)]
#![feature(custom_attribute)]
#![plugin(tag_safe)]
use prelude::*;

#[macro_use]
extern crate core;

pub use arch::memory::PAGE_SIZE;

#[doc(hidden)]
#[macro_use] pub mod logmacros;
#[doc(hidden)]
#[macro_use] pub mod macros;
#[doc(hidden)]
#[macro_use] #[cfg(arch__amd64)] #[path="arch/amd64/mod-macros.rs"] pub mod arch_macros;

// Evil Hack: For some reason, write! (and friends) will expand pointing to std instead of core
#[doc(hidden)]
mod std {
	pub use core::option;
	pub use core::{default,fmt,cmp};
	pub use core::marker;	// needed for derive(Copy)
	pub use core::iter;	// needed for 'for'
}

/// Kernel's version of 'std::prelude'
pub mod prelude;

/// Library datatypes (Vec, Queue, ...)
#[macro_use]
pub mod lib;	// Clone of libstd

/// Heavy synchronisation primitives (Mutex, Semaphore, RWLock, ...)
#[macro_use]
pub mod sync;

/// Asynchrnous wait support
pub mod async;

/// Logging framework
pub mod logging;
/// Memory management (physical, virtual, heap)
pub mod memory;
/// Thread management
#[macro_use]
pub mod threads;
/// Timekeeping (timers and wall time)
pub mod time;

// Module/Executable loading (and symbol lookup)
pub mod loading;
/// Module management (loading and initialisation of kernel modules)
pub mod modules;

/// Meta devices (the Hardware Abstraction Layer)
pub mod metadevs;
/// Device to driver mapping manager
///
/// Starts driver instances for the devices it sees
pub mod device_manager;

/// User output, via a kernel-provided compositing "WM"
pub mod gui;

// Public for driver modules
pub mod vfs;

mod config;

/// Stack unwinding (panic) handling
pub mod unwind;

pub mod irqs;

/// Built-in device drivers
mod hw;

/// Achitecture-specific code - AMD64 (aka x86-64)
#[macro_use]
#[cfg(arch__amd64)] #[path="arch/amd64/mod.rs"] pub mod arch;	// Needs to be pub for exports to be avaliable

/// Kernel entrypoint
#[no_mangle]
pub extern "C" fn kmain()
{
	log_notice!("Tifflin Kernel v{} build {} starting", env!("TK_VERSION"), env!("TK_BUILD"));
	log_notice!("> Git state : {}", env!("TK_GITSPEC"));
	log_notice!("> Built with {}", env!("RUST_VERSION"));

	// Initialise core services before attempting modules
	::memory::phys::init();
	::memory::virt::init();
	::memory::heap::init();
	::threads::init();

	log_log!("Command line = '{}'", ::arch::boot::get_boot_string());
	// TODO: Parse command line into config

	// Dump active video mode
	let vidmode = ::arch::boot::get_video_mode();
	match vidmode {
	Some(m) => {
		log_debug!("Video mode : {}x{} @ {:#x}", m.width, m.height, m.base);
		::metadevs::video::set_boot_mode(m);
		},
	None => log_debug!("No video mode present")
	}

	// Modules (dependency tree included)
	// - Requests that the GUI be started as soon as possible
	::modules::init(&["GUI"]);

	// Yield to allow init threads to run
	::threads::yield_time();

	// Run system init
	sysinit();

	// Thread 0 idle loop
	log_info!("Entering idle");
	loop
	{
		log_trace!("TID0 napping");
		::threads::yield_time();
	}
}

// Initialise the system once drivers are up
fn sysinit()
{
	use metadevs::storage::VolumeHandle;
	use vfs::{mount,handle};
	use vfs::Path;

	// 1. Mount /system to the specified volume
	let sysdisk = ::config::get_string(::config::Value::SysDisk);
	match VolumeHandle::open_named(sysdisk)
	{
	Err(e) => {
		log_error!("Unable to open /system volume {}: {}", sysdisk, e);
		return ;
		},
	Ok(vh) => match mount::mount("/system".as_ref(), vh, "", &[])
		{
		Ok(_) => {},
		Err(e) => {
			log_error!("Unable to mount /system from {}: {:?}", sysdisk, e);
			return ;
			},
		},
	}

	// 2. Symbolic link /sysroot to the specified folder
	let sysroot = ::config::get_string(::config::Value::SysRoot);
	handle::Dir::open(Path::new("/")).unwrap()
		.symlink("sysroot", Path::new(sysroot)).unwrap();

	// *. Testing: open a file known to exist on the testing disk
	{
		match handle::File::open( Path::new("/system/1.TXT"), handle::FileOpenMode::SharedRO )
		{
		Err(e) => log_warning!("VFS test file can't be opened: {:?}", e),
		Ok(h) => {
			log_debug!("VFS open test = {:?}", h);
			let mut buf = [0; 16];
			let sz = h.read(0, &mut buf).unwrap();
			log_debug!("- Contents: {:?}", ::lib::RawString(&buf[..sz]));
			},
		}

		// - Iterate root dir
		match handle::Dir::open( Path::new("/") )
		{
		Err(e) => log_warning!("VFS root cannot be opened: {:?}", e),
		Ok(h) =>
			for name in h.iter() {
				log_log!("{}: {:?}", 0, name);
			},
		}
		// - Iterate /system
		match handle::Dir::open( Path::new("/system") )
		{
		Err(e) => log_warning!("System volume cannot be opened: {:?}", e),
		Ok(h) =>
			for name in h.iter() {
				log_log!("{:?}", name);
			},
		}
	}

	// *. TEST Automount
	// - Probably shouldn't be included in the final version, but works for testing filesystem and storage drivers
	let mountdir = handle::Dir::open( Path::new("/") ).and_then(|h| h.mkdir("mount")).unwrap();
	for (_,v) in ::metadevs::storage::enum_lvs()
	{
		let vh = match VolumeHandle::open_named(&v)
			{
			Err(e) => {
				log_log!("Unable to open '{}': {}", v, e);
				continue;
				},
			Ok(v) => v,
			};
		mountdir.mkdir(&v).unwrap();
		let mountpt = format!("/mount/{}",v);
		match mount::mount( mountpt.as_ref(), vh, "", &[] )
		{
		Ok(_) => log_log!("Auto-mounted to {}", mountpt),
		Err(e) => log_notice!("Unable to automount '{}': {:?}", v, e),
		}

	}
}

// vim: ft=rust

<|endoftext|>


<commit_before>extern crate crossbeam;
#[macro_use] extern crate version;
extern crate walkdir;

use std::env;
use std::process::exit;
use std::process::Command;
use std::process::Stdio;
use std::io::ErrorKind;

use walkdir::WalkDir;
use crossbeam::*;

fn help(cmds: Vec<String>) {
    println!(r#"

     _
    (_)_ __ ___   __ _  __ _
    | | '_ \` _ \/ _\`|/ _\`|
    | | | | | | | (_| | (_| |
    |_|_| |_| |_|\__,_|\__, |
                       |___/
    -------------------------

    Usage: imag [--version | --versions | -h | --help] <command> <args...>

    imag - the personal information management suite for the commandline

    imag is a PIM suite for the commandline. It consists of several commands,
    called "modules". Each module implements one PIM aspect and all of these
    modules can be used independently.

    Available commands:
    "#);

    for cmd in cmds.iter() {
        println!("\t{}", cmd);
    }

    println!(r#"

    Call a command with 'imag <command> <args>'
    Each command can be called with "--help" to get the respective helptext.

    Please visit https://github.com/matthiasbeyer/imag to view the source code,
    follow the development of imag or maybe even contribute to imag.

    imag is free software. It is released under the terms of LGPLv2.1

    (c) 2016 Matthias Beyer and contributors"#);
}

fn get_commands() -> Vec<String> {
    let path = env::var("PATH");
    if path.is_err() {
        println!("PATH error: {:?}", path);
        exit(1);
    }
    let pathelements = path.unwrap();
    let pathelements = pathelements.split(":");

    let joinhandles : Vec<ScopedJoinHandle<Vec<String>>> = pathelements
        .map(|elem| {
            crossbeam::scope(|scope| {
                scope.spawn(|| {
                    WalkDir::new(elem)
                        .max_depth(1)
                        .into_iter()
                        .filter(|path| {
                            match path {
                                &Ok(ref p) => p.file_name()
                                    .to_str()
                                    .map_or(false, |filename| filename.starts_with("imag-")),
                                &Err(_)   => false,
                            }
                        })
                        .filter_map(|x| x.ok())
                        .filter_map(|path| {
                           path.file_name()
                               .to_str()
                               .map(String::from)
                        })
                        .collect()
                })
            })
        })
        .collect();

    let mut execs = vec![];
    for joinhandle in joinhandles.into_iter() {
        let mut v = joinhandle.join();
        execs.append(&mut v);
    }

    execs
}

fn find_command() -> Option<String> {
    env::args().skip(1).filter(|x| !x.starts_with("-")).next()
}

fn find_flag() -> Option<String> {
    env::args().skip(1).filter(|x| x.starts_with("-")).next()
}

fn find_args(command: &str) -> Vec<String> {
    env::args()
        .skip(1)
        .position(|e| e == command)
        .map(|pos| env::args().skip(pos + 2).collect::<Vec<String>>())
        .unwrap_or(vec![])
}

fn main() {
    let commands  = get_commands();
    let mut args  = env::args();
    let _         = args.next();
    let first_arg = match find_command() {
        Some(s) => s,
        None    => match find_flag() {
            Some(s) => s,
            None => {
                help(commands);
                exit(0);
            },
        },
    };
    let is_debug = env::args().skip(1).filter(|x| x == "--debug").next().is_some();

    match &first_arg[..] {
        "--help" | "-h" => {
            help(commands);
            exit(0);
        },

        "--version"  => println!("imag {}", &version!()[..]),

        "--versions" => {
            let mut result = vec![];
            for command in commands.iter() {
                result.push(crossbeam::scope(|scope| {
                    scope.spawn(|| {
                        let v = Command::new(command).arg("--version").output();
                        match v {
                            Ok(v) => match String::from_utf8(v.stdout) {
                                Ok(s) => format!("{} -> {}", command, s),
                                Err(e) => format!("Failed calling {} -> {:?}", command, e),
                            },
                            Err(e) => format!("Failed calling {} -> {:?}", command, e),
                        }
                    })
                }))
            }

            for versionstring in result.into_iter().map(|handle| handle.join()) {
                println!("{}", versionstring);
            }
        },

        s => {
            let mut subcommand_args = find_args(s);
            if is_debug && subcommand_args.filter(|x| x == "debug").count() == 0 {
                subcommand_args.push(String::from("--debug"));
            }
            match Command::new(format!("imag-{}", s))
                .stdin(Stdio::inherit())
                .stdout(Stdio::inherit())
                .stderr(Stdio::inherit())
                .args(&subcommand_args[..])
                .spawn()
                .and_then(|mut handle| handle.wait())
            {
                Ok(exit_status) => {
                    if !exit_status.success() {
                        println!("{} exited with non-zero exit code", s);
                        exit(exit_status.code().unwrap_or(42));
                    }
                },

                Err(e) => {
                    match e.kind() {
                        ErrorKind::NotFound => {
                            println!("No such command: 'imag-{}'", s);
                            exit(2);
                        },
                        ErrorKind::PermissionDenied => {
                            println!("No permission to execute: 'imag-{}'", s);
                            exit(1);
                        },
                        _ => {
                            println!("Error spawning: {:?}", e);
                            exit(1337);
                        }
                    }
                }
            }

        },
    }
}
<commit_msg>Fixed compile errors<commit_after>extern crate crossbeam;
#[macro_use] extern crate version;
extern crate walkdir;

use std::env;
use std::process::exit;
use std::process::Command;
use std::process::Stdio;
use std::io::ErrorKind;

use walkdir::WalkDir;
use crossbeam::*;

fn help(cmds: Vec<String>) {
    println!(r#"

     _
    (_)_ __ ___   __ _  __ _
    | | '_ \` _ \/ _\`|/ _\`|
    | | | | | | | (_| | (_| |
    |_|_| |_| |_|\__,_|\__, |
                       |___/
    -------------------------

    Usage: imag [--version | --versions | -h | --help] <command> <args...>

    imag - the personal information management suite for the commandline

    imag is a PIM suite for the commandline. It consists of several commands,
    called "modules". Each module implements one PIM aspect and all of these
    modules can be used independently.

    Available commands:
    "#);

    for cmd in cmds.iter() {
        println!("\t{}", cmd);
    }

    println!(r#"

    Call a command with 'imag <command> <args>'
    Each command can be called with "--help" to get the respective helptext.

    Please visit https://github.com/matthiasbeyer/imag to view the source code,
    follow the development of imag or maybe even contribute to imag.

    imag is free software. It is released under the terms of LGPLv2.1

    (c) 2016 Matthias Beyer and contributors"#);
}

fn get_commands() -> Vec<String> {
    let path = env::var("PATH");
    if path.is_err() {
        println!("PATH error: {:?}", path);
        exit(1);
    }
    let pathelements = path.unwrap();
    let pathelements = pathelements.split(":");

    let joinhandles : Vec<ScopedJoinHandle<Vec<String>>> = pathelements
        .map(|elem| {
            crossbeam::scope(|scope| {
                scope.spawn(|| {
                    WalkDir::new(elem)
                        .max_depth(1)
                        .into_iter()
                        .filter(|path| {
                            match path {
                                &Ok(ref p) => p.file_name()
                                    .to_str()
                                    .map_or(false, |filename| filename.starts_with("imag-")),
                                &Err(_)   => false,
                            }
                        })
                        .filter_map(|x| x.ok())
                        .filter_map(|path| {
                           path.file_name()
                               .to_str()
                               .map(String::from)
                        })
                        .collect()
                })
            })
        })
        .collect();

    let mut execs = vec![];
    for joinhandle in joinhandles.into_iter() {
        let mut v = joinhandle.join();
        execs.append(&mut v);
    }

    execs
}

fn find_command() -> Option<String> {
    env::args().skip(1).filter(|x| !x.starts_with("-")).next()
}

fn find_flag() -> Option<String> {
    env::args().skip(1).filter(|x| x.starts_with("-")).next()
}

fn find_args(command: &str) -> Vec<String> {
    env::args()
        .skip(1)
        .position(|e| e == command)
        .map(|pos| env::args().skip(pos + 2).collect::<Vec<String>>())
        .unwrap_or(vec![])
}

fn main() {
    let commands  = get_commands();
    let mut args  = env::args();
    let _         = args.next();
    let first_arg = match find_command() {
        Some(s) => s,
        None    => match find_flag() {
            Some(s) => s,
            None => {
                help(commands);
                exit(0);
            },
        },
    };
    let is_debug = env::args().skip(1).filter(|x| x == "--debug").next().is_some();

    match &first_arg[..] {
        "--help" | "-h" => {
            help(commands);
            exit(0);
        },

        "--version"  => println!("imag {}", &version!()[..]),

        "--versions" => {
            let mut result = vec![];
            for command in commands.iter() {
                result.push(crossbeam::scope(|scope| {
                    scope.spawn(|| {
                        let v = Command::new(command).arg("--version").output();
                        match v {
                            Ok(v) => match String::from_utf8(v.stdout) {
                                Ok(s) => format!("{} -> {}", command, s),
                                Err(e) => format!("Failed calling {} -> {:?}", command, e),
                            },
                            Err(e) => format!("Failed calling {} -> {:?}", command, e),
                        }
                    })
                }))
            }

            for versionstring in result.into_iter().map(|handle| handle.join()) {
                println!("{}", versionstring);
            }
        },

        s => {
            let mut subcommand_args = find_args(s);
            if is_debug && subcommand_args.iter().find(|x| *x == "debug").is_none() {
                subcommand_args.push(String::from("--debug"));
            }
            match Command::new(format!("imag-{}", s))
                .stdin(Stdio::inherit())
                .stdout(Stdio::inherit())
                .stderr(Stdio::inherit())
                .args(&subcommand_args[..])
                .spawn()
                .and_then(|mut handle| handle.wait())
            {
                Ok(exit_status) => {
                    if !exit_status.success() {
                        println!("{} exited with non-zero exit code", s);
                        exit(exit_status.code().unwrap_or(42));
                    }
                },

                Err(e) => {
                    match e.kind() {
                        ErrorKind::NotFound => {
                            println!("No such command: 'imag-{}'", s);
                            exit(2);
                        },
                        ErrorKind::PermissionDenied => {
                            println!("No permission to execute: 'imag-{}'", s);
                            exit(1);
                        },
                        _ => {
                            println!("Error spawning: {:?}", e);
                            exit(1337);
                        }
                    }
                }
            }

        },
    }
}
<|endoftext|>


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

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

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

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

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

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

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

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

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

        Ok((thread, srv))
    }

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

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

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

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

    pub(crate) fn send_task<R>(&self, req: Request) -> Result<R, tt::ExpansionError>
    where
        R: TryFrom<Response, Error = &'static str>,
    {
        let (result_tx, result_rx) = bounded(0);
        let sender = match self.inner.upgrade() {
            None => {
                return Err(tt::ExpansionError::Unknown("Proc macro process is closed.".into()))
            }
            Some(it) => it,
        };
        sender.send(Task { req, result_tx }).unwrap();
        let res = result_rx
            .recv()
            .map_err(|_| tt::ExpansionError::Unknown("Proc macro thread is closed.".into()))?;

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

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

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

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

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

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

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

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

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

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

        Some((stdin, read))
    }
}

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

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

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

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

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

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

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

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

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

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

        Ok((thread, srv))
    }

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

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

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

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

    pub(crate) fn send_task<R>(&self, req: Request) -> Result<R, tt::ExpansionError>
    where
        R: TryFrom<Response, Error = &'static str>,
    {
        let (result_tx, result_rx) = bounded(0);
        let sender = match self.inner.upgrade() {
            None => {
                return Err(tt::ExpansionError::Unknown("Proc macro process is closed.".into()))
            }
            Some(it) => it,
        };
        sender.send(Task { req, result_tx }).unwrap();
        let res = result_rx
            .recv()
            .map_err(|_| tt::ExpansionError::Unknown("Proc macro thread is closed.".into()))?;

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

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

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

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

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

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

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

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

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

        Some((stdin, read))
    }
}

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

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


<commit_before>//! This module handles fuzzy-searching of functions, structs and other symbols
//! by name across the whole workspace and dependencies.
//!
//! It works by building an incrementally-updated text-search index of all
//! symbols. The backbone of the index is the **awesome** `fst` crate by
//! @BurntSushi.
//!
//! In a nutshell, you give a set of strings to `fst`, and it builds a
//! finite state machine describing this set of strings. The strings which
//! could fuzzy-match a pattern can also be described by a finite state machine.
//! What is freaking cool is that you can now traverse both state machines in
//! lock-step to enumerate the strings which are both in the input set and
//! fuzz-match the query. Or, more formally, given two languages described by
//! FSTs, one can build a product FST which describes the intersection of the
//! languages.
//!
//! `fst` does not support cheap updating of the index, but it supports unioning
//! of state machines. So, to account for changing source code, we build an FST
//! for each library (which is assumed to never change) and an FST for each Rust
//! file in the current workspace, and run a query against the union of all
//! those FSTs.

use std::{
    cmp::Ordering,
    fmt,
    hash::{Hash, Hasher},
    mem,
    sync::Arc,
};

use fst::{self, Streamer};
use ra_db::{
    salsa::{self, ParallelDatabase},
    FileId, SourceDatabaseExt, SourceRootId,
};
use ra_syntax::{
    ast::{self, NameOwner},
    match_ast, AstNode, Parse, SmolStr, SourceFile,
    SyntaxKind::{self, *},
    SyntaxNode, SyntaxNodePtr, TextRange, WalkEvent,
};
#[cfg(not(feature = "wasm"))]
use rayon::prelude::*;

use crate::RootDatabase;

#[derive(Debug)]
pub struct Query {
    query: String,
    lowercased: String,
    only_types: bool,
    libs: bool,
    exact: bool,
    limit: usize,
}

impl Query {
    pub fn new(query: String) -> Query {
        let lowercased = query.to_lowercase();
        Query {
            query,
            lowercased,
            only_types: false,
            libs: false,
            exact: false,
            limit: usize::max_value(),
        }
    }

    pub fn only_types(&mut self) {
        self.only_types = true;
    }

    pub fn libs(&mut self) {
        self.libs = true;
    }

    pub fn exact(&mut self) {
        self.exact = true;
    }

    pub fn limit(&mut self, limit: usize) {
        self.limit = limit
    }
}

#[salsa::query_group(SymbolsDatabaseStorage)]
pub trait SymbolsDatabase: hir::db::HirDatabase {
    fn file_symbols(&self, file_id: FileId) -> Arc<SymbolIndex>;
    #[salsa::input]
    fn library_symbols(&self, id: SourceRootId) -> Arc<SymbolIndex>;
    /// The set of "local" (that is, from the current workspace) roots.
    /// Files in local roots are assumed to change frequently.
    #[salsa::input]
    fn local_roots(&self) -> Arc<Vec<SourceRootId>>;
    /// The set of roots for crates.io libraries.
    /// Files in libraries are assumed to never change.
    #[salsa::input]
    fn library_roots(&self) -> Arc<Vec<SourceRootId>>;
}

fn file_symbols(db: &impl SymbolsDatabase, file_id: FileId) -> Arc<SymbolIndex> {
    db.check_canceled();
    let parse = db.parse(file_id);

    let symbols = source_file_to_file_symbols(&parse.tree(), file_id);

    // FIXME: add macros here

    Arc::new(SymbolIndex::new(symbols))
}

// Feature: Workspace Symbol
//
// Uses fuzzy-search to find types, modules and functions by name across your
// project and dependencies. This is **the** most useful feature, which improves code
// navigation tremendously. It mostly works on top of the built-in LSP
// functionality, however `#` and `*` symbols can be used to narrow down the
// search. Specifically,
//
// - `Foo` searches for `Foo` type in the current workspace
// - `foo#` searches for `foo` function in the current workspace
// - `Foo*` searches for `Foo` type among dependencies, including `stdlib`
// - `foo#*` searches for `foo` function among dependencies
//
// That is, `#` switches from "types" to all symbols, `*` switches from the current
// workspace to dependencies.
//
// |===
// | Editor  | Shortcut
//
// | VS Code | kbd:[Ctrl+T]
// |===
pub fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
    /// Need to wrap Snapshot to provide `Clone` impl for `map_with`
    struct Snap(salsa::Snapshot<RootDatabase>);
    impl Clone for Snap {
        fn clone(&self) -> Snap {
            Snap(self.0.snapshot())
        }
    }

    let buf: Vec<Arc<SymbolIndex>> = if query.libs {
        let snap = Snap(db.snapshot());
        #[cfg(not(feature = "wasm"))]
        let buf = db
            .library_roots()
            .par_iter()
            .map_with(snap, |db, &lib_id| db.0.library_symbols(lib_id))
            .collect();

        #[cfg(feature = "wasm")]
        let buf = db.library_roots().iter().map(|&lib_id| snap.0.library_symbols(lib_id)).collect();

        buf
    } else {
        let mut files = Vec::new();
        for &root in db.local_roots().iter() {
            let sr = db.source_root(root);
            files.extend(sr.walk())
        }

        let snap = Snap(db.snapshot());
        #[cfg(not(feature = "wasm"))]
        let buf =
            files.par_iter().map_with(snap, |db, &file_id| db.0.file_symbols(file_id)).collect();

        #[cfg(feature = "wasm")]
        let buf = files.iter().map(|&file_id| snap.0.file_symbols(file_id)).collect();

        buf
    };
    query.search(&buf)
}

pub fn index_resolve(db: &RootDatabase, name_ref: &ast::NameRef) -> Vec<FileSymbol> {
    let name = name_ref.text();
    let mut query = Query::new(name.to_string());
    query.exact();
    query.limit(4);
    world_symbols(db, query)
}

#[derive(Default)]
pub struct SymbolIndex {
    symbols: Vec<FileSymbol>,
    map: fst::Map<Vec<u8>>,
}

impl fmt::Debug for SymbolIndex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("SymbolIndex").field("n_symbols", &self.symbols.len()).finish()
    }
}

impl PartialEq for SymbolIndex {
    fn eq(&self, other: &SymbolIndex) -> bool {
        self.symbols == other.symbols
    }
}

impl Eq for SymbolIndex {}

impl Hash for SymbolIndex {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.symbols.hash(hasher)
    }
}

impl SymbolIndex {
    fn new(mut symbols: Vec<FileSymbol>) -> SymbolIndex {
        fn cmp(lhs: &FileSymbol, rhs: &FileSymbol) -> Ordering {
            let lhs_chars = lhs.name.chars().map(|c| c.to_ascii_lowercase());
            let rhs_chars = rhs.name.chars().map(|c| c.to_ascii_lowercase());
            lhs_chars.cmp(rhs_chars)
        }

        #[cfg(not(feature = "wasm"))]
        symbols.par_sort_by(cmp);

        #[cfg(feature = "wasm")]
        symbols.sort_by(cmp);

        let mut builder = fst::MapBuilder::memory();

        let mut last_batch_start = 0;

        for idx in 0..symbols.len() {
            if let Some(next_symbol) = symbols.get(idx + 1) {
                if cmp(&symbols[last_batch_start], next_symbol) == Ordering::Equal {
                    continue;
                }
            }

            let start = last_batch_start;
            let end = idx + 1;
            last_batch_start = end;

            let key = symbols[start].name.as_str().to_ascii_lowercase();
            let value = SymbolIndex::range_to_map_value(start, end);

            builder.insert(key, value).unwrap();
        }

        let map = fst::Map::new(builder.into_inner().unwrap()).unwrap();
        SymbolIndex { symbols, map }
    }

    pub fn len(&self) -> usize {
        self.symbols.len()
    }

    pub fn memory_size(&self) -> usize {
        self.map.as_fst().size() + self.symbols.len() * mem::size_of::<FileSymbol>()
    }

    #[cfg(not(feature = "wasm"))]
    pub(crate) fn for_files(
        files: impl ParallelIterator<Item = (FileId, Parse<ast::SourceFile>)>,
    ) -> SymbolIndex {
        let symbols = files
            .flat_map(|(file_id, file)| source_file_to_file_symbols(&file.tree(), file_id))
            .collect::<Vec<_>>();
        SymbolIndex::new(symbols)
    }

    #[cfg(feature = "wasm")]
    pub(crate) fn for_files(
        files: impl Iterator<Item = (FileId, Parse<ast::SourceFile>)>,
    ) -> SymbolIndex {
        let symbols = files
            .flat_map(|(file_id, file)| source_file_to_file_symbols(&file.tree(), file_id))
            .collect::<Vec<_>>();
        SymbolIndex::new(symbols)
    }

    fn range_to_map_value(start: usize, end: usize) -> u64 {
        debug_assert![start <= (std::u32::MAX as usize)];
        debug_assert![end <= (std::u32::MAX as usize)];

        ((start as u64) << 32) | end as u64
    }

    fn map_value_to_range(value: u64) -> (usize, usize) {
        let end = value as u32 as usize;
        let start = (value >> 32) as usize;
        (start, end)
    }
}

impl Query {
    pub(crate) fn search(self, indices: &[Arc<SymbolIndex>]) -> Vec<FileSymbol> {
        let mut op = fst::map::OpBuilder::new();
        for file_symbols in indices.iter() {
            let automaton = fst::automaton::Subsequence::new(&self.lowercased);
            op = op.add(file_symbols.map.search(automaton))
        }
        let mut stream = op.union();
        let mut res = Vec::new();
        while let Some((_, indexed_values)) = stream.next() {
            if res.len() >= self.limit {
                break;
            }
            for indexed_value in indexed_values {
                let symbol_index = &indices[indexed_value.index];
                let (start, end) = SymbolIndex::map_value_to_range(indexed_value.value);

                for symbol in &symbol_index.symbols[start..end] {
                    if self.only_types && !is_type(symbol.kind) {
                        continue;
                    }
                    if self.exact && symbol.name != self.query {
                        continue;
                    }
                    res.push(symbol.clone());
                }
            }
        }
        res
    }
}

fn is_type(kind: SyntaxKind) -> bool {
    match kind {
        STRUCT_DEF | ENUM_DEF | TRAIT_DEF | TYPE_ALIAS_DEF => true,
        _ => false,
    }
}

/// The actual data that is stored in the index. It should be as compact as
/// possible.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct FileSymbol {
    pub file_id: FileId,
    pub name: SmolStr,
    pub kind: SyntaxKind,
    pub range: TextRange,
    pub ptr: SyntaxNodePtr,
    pub name_range: Option<TextRange>,
    pub container_name: Option<SmolStr>,
}

fn source_file_to_file_symbols(source_file: &SourceFile, file_id: FileId) -> Vec<FileSymbol> {
    let mut symbols = Vec::new();
    let mut stack = Vec::new();

    for event in source_file.syntax().preorder() {
        match event {
            WalkEvent::Enter(node) => {
                if let Some(mut symbol) = to_file_symbol(&node, file_id) {
                    symbol.container_name = stack.last().cloned();

                    stack.push(symbol.name.clone());
                    symbols.push(symbol);
                }
            }

            WalkEvent::Leave(node) => {
                if to_symbol(&node).is_some() {
                    stack.pop();
                }
            }
        }
    }

    symbols
}

fn to_symbol(node: &SyntaxNode) -> Option<(SmolStr, SyntaxNodePtr, TextRange)> {
    fn decl<N: NameOwner>(node: N) -> Option<(SmolStr, SyntaxNodePtr, TextRange)> {
        let name = node.name()?;
        let name_range = name.syntax().text_range();
        let name = name.text().clone();
        let ptr = SyntaxNodePtr::new(node.syntax());

        Some((name, ptr, name_range))
    }
    match_ast! {
        match node {
            ast::FnDef(it) => decl(it),
            ast::StructDef(it) => decl(it),
            ast::EnumDef(it) => decl(it),
            ast::TraitDef(it) => decl(it),
            ast::Module(it) => decl(it),
            ast::TypeAliasDef(it) => decl(it),
            ast::ConstDef(it) => decl(it),
            ast::StaticDef(it) => decl(it),
            ast::MacroCall(it) => {
                if it.is_macro_rules().is_some() {
                    decl(it)
                } else {
                    None
                }
            },
            _ => None,
        }
    }
}

fn to_file_symbol(node: &SyntaxNode, file_id: FileId) -> Option<FileSymbol> {
    to_symbol(node).map(move |(name, ptr, name_range)| FileSymbol {
        name,
        kind: node.kind(),
        range: node.text_range(),
        ptr,
        file_id,
        name_range: Some(name_range),
        container_name: None,
    })
}
<commit_msg>Profile `world_symbols`<commit_after>//! This module handles fuzzy-searching of functions, structs and other symbols
//! by name across the whole workspace and dependencies.
//!
//! It works by building an incrementally-updated text-search index of all
//! symbols. The backbone of the index is the **awesome** `fst` crate by
//! @BurntSushi.
//!
//! In a nutshell, you give a set of strings to `fst`, and it builds a
//! finite state machine describing this set of strings. The strings which
//! could fuzzy-match a pattern can also be described by a finite state machine.
//! What is freaking cool is that you can now traverse both state machines in
//! lock-step to enumerate the strings which are both in the input set and
//! fuzz-match the query. Or, more formally, given two languages described by
//! FSTs, one can build a product FST which describes the intersection of the
//! languages.
//!
//! `fst` does not support cheap updating of the index, but it supports unioning
//! of state machines. So, to account for changing source code, we build an FST
//! for each library (which is assumed to never change) and an FST for each Rust
//! file in the current workspace, and run a query against the union of all
//! those FSTs.

use std::{
    cmp::Ordering,
    fmt,
    hash::{Hash, Hasher},
    mem,
    sync::Arc,
};

use fst::{self, Streamer};
use ra_db::{
    salsa::{self, ParallelDatabase},
    FileId, SourceDatabaseExt, SourceRootId,
};
use ra_syntax::{
    ast::{self, NameOwner},
    match_ast, AstNode, Parse, SmolStr, SourceFile,
    SyntaxKind::{self, *},
    SyntaxNode, SyntaxNodePtr, TextRange, WalkEvent,
};
#[cfg(not(feature = "wasm"))]
use rayon::prelude::*;

use crate::RootDatabase;

#[derive(Debug)]
pub struct Query {
    query: String,
    lowercased: String,
    only_types: bool,
    libs: bool,
    exact: bool,
    limit: usize,
}

impl Query {
    pub fn new(query: String) -> Query {
        let lowercased = query.to_lowercase();
        Query {
            query,
            lowercased,
            only_types: false,
            libs: false,
            exact: false,
            limit: usize::max_value(),
        }
    }

    pub fn only_types(&mut self) {
        self.only_types = true;
    }

    pub fn libs(&mut self) {
        self.libs = true;
    }

    pub fn exact(&mut self) {
        self.exact = true;
    }

    pub fn limit(&mut self, limit: usize) {
        self.limit = limit
    }
}

#[salsa::query_group(SymbolsDatabaseStorage)]
pub trait SymbolsDatabase: hir::db::HirDatabase {
    fn file_symbols(&self, file_id: FileId) -> Arc<SymbolIndex>;
    #[salsa::input]
    fn library_symbols(&self, id: SourceRootId) -> Arc<SymbolIndex>;
    /// The set of "local" (that is, from the current workspace) roots.
    /// Files in local roots are assumed to change frequently.
    #[salsa::input]
    fn local_roots(&self) -> Arc<Vec<SourceRootId>>;
    /// The set of roots for crates.io libraries.
    /// Files in libraries are assumed to never change.
    #[salsa::input]
    fn library_roots(&self) -> Arc<Vec<SourceRootId>>;
}

fn file_symbols(db: &impl SymbolsDatabase, file_id: FileId) -> Arc<SymbolIndex> {
    db.check_canceled();
    let parse = db.parse(file_id);

    let symbols = source_file_to_file_symbols(&parse.tree(), file_id);

    // FIXME: add macros here

    Arc::new(SymbolIndex::new(symbols))
}

// Feature: Workspace Symbol
//
// Uses fuzzy-search to find types, modules and functions by name across your
// project and dependencies. This is **the** most useful feature, which improves code
// navigation tremendously. It mostly works on top of the built-in LSP
// functionality, however `#` and `*` symbols can be used to narrow down the
// search. Specifically,
//
// - `Foo` searches for `Foo` type in the current workspace
// - `foo#` searches for `foo` function in the current workspace
// - `Foo*` searches for `Foo` type among dependencies, including `stdlib`
// - `foo#*` searches for `foo` function among dependencies
//
// That is, `#` switches from "types" to all symbols, `*` switches from the current
// workspace to dependencies.
//
// |===
// | Editor  | Shortcut
//
// | VS Code | kbd:[Ctrl+T]
// |===
pub fn world_symbols(db: &RootDatabase, query: Query) -> Vec<FileSymbol> {
    let _p = ra_prof::profile("world_symbols").detail(|| query.query.clone());

    /// Need to wrap Snapshot to provide `Clone` impl for `map_with`
    struct Snap(salsa::Snapshot<RootDatabase>);
    impl Clone for Snap {
        fn clone(&self) -> Snap {
            Snap(self.0.snapshot())
        }
    }

    let buf: Vec<Arc<SymbolIndex>> = if query.libs {
        let snap = Snap(db.snapshot());
        #[cfg(not(feature = "wasm"))]
        let buf = db
            .library_roots()
            .par_iter()
            .map_with(snap, |db, &lib_id| db.0.library_symbols(lib_id))
            .collect();

        #[cfg(feature = "wasm")]
        let buf = db.library_roots().iter().map(|&lib_id| snap.0.library_symbols(lib_id)).collect();

        buf
    } else {
        let mut files = Vec::new();
        for &root in db.local_roots().iter() {
            let sr = db.source_root(root);
            files.extend(sr.walk())
        }

        let snap = Snap(db.snapshot());
        #[cfg(not(feature = "wasm"))]
        let buf =
            files.par_iter().map_with(snap, |db, &file_id| db.0.file_symbols(file_id)).collect();

        #[cfg(feature = "wasm")]
        let buf = files.iter().map(|&file_id| snap.0.file_symbols(file_id)).collect();

        buf
    };
    query.search(&buf)
}

pub fn index_resolve(db: &RootDatabase, name_ref: &ast::NameRef) -> Vec<FileSymbol> {
    let name = name_ref.text();
    let mut query = Query::new(name.to_string());
    query.exact();
    query.limit(4);
    world_symbols(db, query)
}

#[derive(Default)]
pub struct SymbolIndex {
    symbols: Vec<FileSymbol>,
    map: fst::Map<Vec<u8>>,
}

impl fmt::Debug for SymbolIndex {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("SymbolIndex").field("n_symbols", &self.symbols.len()).finish()
    }
}

impl PartialEq for SymbolIndex {
    fn eq(&self, other: &SymbolIndex) -> bool {
        self.symbols == other.symbols
    }
}

impl Eq for SymbolIndex {}

impl Hash for SymbolIndex {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        self.symbols.hash(hasher)
    }
}

impl SymbolIndex {
    fn new(mut symbols: Vec<FileSymbol>) -> SymbolIndex {
        fn cmp(lhs: &FileSymbol, rhs: &FileSymbol) -> Ordering {
            let lhs_chars = lhs.name.chars().map(|c| c.to_ascii_lowercase());
            let rhs_chars = rhs.name.chars().map(|c| c.to_ascii_lowercase());
            lhs_chars.cmp(rhs_chars)
        }

        #[cfg(not(feature = "wasm"))]
        symbols.par_sort_by(cmp);

        #[cfg(feature = "wasm")]
        symbols.sort_by(cmp);

        let mut builder = fst::MapBuilder::memory();

        let mut last_batch_start = 0;

        for idx in 0..symbols.len() {
            if let Some(next_symbol) = symbols.get(idx + 1) {
                if cmp(&symbols[last_batch_start], next_symbol) == Ordering::Equal {
                    continue;
                }
            }

            let start = last_batch_start;
            let end = idx + 1;
            last_batch_start = end;

            let key = symbols[start].name.as_str().to_ascii_lowercase();
            let value = SymbolIndex::range_to_map_value(start, end);

            builder.insert(key, value).unwrap();
        }

        let map = fst::Map::new(builder.into_inner().unwrap()).unwrap();
        SymbolIndex { symbols, map }
    }

    pub fn len(&self) -> usize {
        self.symbols.len()
    }

    pub fn memory_size(&self) -> usize {
        self.map.as_fst().size() + self.symbols.len() * mem::size_of::<FileSymbol>()
    }

    #[cfg(not(feature = "wasm"))]
    pub(crate) fn for_files(
        files: impl ParallelIterator<Item = (FileId, Parse<ast::SourceFile>)>,
    ) -> SymbolIndex {
        let symbols = files
            .flat_map(|(file_id, file)| source_file_to_file_symbols(&file.tree(), file_id))
            .collect::<Vec<_>>();
        SymbolIndex::new(symbols)
    }

    #[cfg(feature = "wasm")]
    pub(crate) fn for_files(
        files: impl Iterator<Item = (FileId, Parse<ast::SourceFile>)>,
    ) -> SymbolIndex {
        let symbols = files
            .flat_map(|(file_id, file)| source_file_to_file_symbols(&file.tree(), file_id))
            .collect::<Vec<_>>();
        SymbolIndex::new(symbols)
    }

    fn range_to_map_value(start: usize, end: usize) -> u64 {
        debug_assert![start <= (std::u32::MAX as usize)];
        debug_assert![end <= (std::u32::MAX as usize)];

        ((start as u64) << 32) | end as u64
    }

    fn map_value_to_range(value: u64) -> (usize, usize) {
        let end = value as u32 as usize;
        let start = (value >> 32) as usize;
        (start, end)
    }
}

impl Query {
    pub(crate) fn search(self, indices: &[Arc<SymbolIndex>]) -> Vec<FileSymbol> {
        let mut op = fst::map::OpBuilder::new();
        for file_symbols in indices.iter() {
            let automaton = fst::automaton::Subsequence::new(&self.lowercased);
            op = op.add(file_symbols.map.search(automaton))
        }
        let mut stream = op.union();
        let mut res = Vec::new();
        while let Some((_, indexed_values)) = stream.next() {
            if res.len() >= self.limit {
                break;
            }
            for indexed_value in indexed_values {
                let symbol_index = &indices[indexed_value.index];
                let (start, end) = SymbolIndex::map_value_to_range(indexed_value.value);

                for symbol in &symbol_index.symbols[start..end] {
                    if self.only_types && !is_type(symbol.kind) {
                        continue;
                    }
                    if self.exact && symbol.name != self.query {
                        continue;
                    }
                    res.push(symbol.clone());
                }
            }
        }
        res
    }
}

fn is_type(kind: SyntaxKind) -> bool {
    match kind {
        STRUCT_DEF | ENUM_DEF | TRAIT_DEF | TYPE_ALIAS_DEF => true,
        _ => false,
    }
}

/// The actual data that is stored in the index. It should be as compact as
/// possible.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct FileSymbol {
    pub file_id: FileId,
    pub name: SmolStr,
    pub kind: SyntaxKind,
    pub range: TextRange,
    pub ptr: SyntaxNodePtr,
    pub name_range: Option<TextRange>,
    pub container_name: Option<SmolStr>,
}

fn source_file_to_file_symbols(source_file: &SourceFile, file_id: FileId) -> Vec<FileSymbol> {
    let mut symbols = Vec::new();
    let mut stack = Vec::new();

    for event in source_file.syntax().preorder() {
        match event {
            WalkEvent::Enter(node) => {
                if let Some(mut symbol) = to_file_symbol(&node, file_id) {
                    symbol.container_name = stack.last().cloned();

                    stack.push(symbol.name.clone());
                    symbols.push(symbol);
                }
            }

            WalkEvent::Leave(node) => {
                if to_symbol(&node).is_some() {
                    stack.pop();
                }
            }
        }
    }

    symbols
}

fn to_symbol(node: &SyntaxNode) -> Option<(SmolStr, SyntaxNodePtr, TextRange)> {
    fn decl<N: NameOwner>(node: N) -> Option<(SmolStr, SyntaxNodePtr, TextRange)> {
        let name = node.name()?;
        let name_range = name.syntax().text_range();
        let name = name.text().clone();
        let ptr = SyntaxNodePtr::new(node.syntax());

        Some((name, ptr, name_range))
    }
    match_ast! {
        match node {
            ast::FnDef(it) => decl(it),
            ast::StructDef(it) => decl(it),
            ast::EnumDef(it) => decl(it),
            ast::TraitDef(it) => decl(it),
            ast::Module(it) => decl(it),
            ast::TypeAliasDef(it) => decl(it),
            ast::ConstDef(it) => decl(it),
            ast::StaticDef(it) => decl(it),
            ast::MacroCall(it) => {
                if it.is_macro_rules().is_some() {
                    decl(it)
                } else {
                    None
                }
            },
            _ => None,
        }
    }
}

fn to_file_symbol(node: &SyntaxNode, file_id: FileId) -> Option<FileSymbol> {
    to_symbol(node).map(move |(name, ptr, name_range)| FileSymbol {
        name,
        kind: node.kind(),
        range: node.text_range(),
        ptr,
        file_id,
        name_range: Some(name_range),
        container_name: None,
    })
}
<|endoftext|>