<commit_before>use bindings::{
    jit_context_t,
    jit_context_create,
    jit_context_destroy,
    jit_context_build_start,
    jit_context_build_end,
    jit_function_t,
    jit_function_next
};
use function::Function;
use std::kinds::marker::ContravariantLifetime;
use util::NativeRef;
/// Holds all of the functions you have built and compiled. There can be
/// multiple, but normally there is only one.
native_ref!(Context, _context, jit_context_t, ContravariantLifetime)
impl<'a> Context<'a> {
    #[inline(always)]
    /// Create a new JIT Context
    pub fn new() -> Context<'a> {
        unsafe {
            NativeRef::from_ptr(jit_context_create())
        }
    }
    #[inline(always)]
    /// Run a closure that can generate IR
    pub fn build<R>(&self, cb: || -> R) -> R {
        unsafe {
            jit_context_build_start(self.as_ptr());
            let rv = cb();
            jit_context_build_end(self.as_ptr());
            rv
        }
    }
    /// Iterate through all the functions in this context
    #[inline(always)]
    pub fn funcs(&self) -> Functions<'a> {
        Functions::new(self)
    }
}
#[unsafe_destructor]
impl<'a> Drop for Context<'a> {
    #[inline]
    fn drop(&mut self) {
        unsafe {
            jit_context_destroy(self.as_ptr());
        }
    }
}
/// Any JIT object which is in a context
pub trait InContext<'a> {
    /// Get the context for this object
    fn get_context(&self) -> Context<'a>;
}

/// An iterator over a context's functions
pub struct Functions<'a> {
    ctx: jit_context_t,
    last: jit_function_t,
    marker: ContravariantLifetime<'a>
}
impl<'a> Functions<'a> {
    fn new(ctx:&Context<'a>) -> Functions<'a> {
        unsafe {
            Functions {
                ctx: ctx.as_ptr(),
                last: RawPtr::null(),
                marker: ContravariantLifetime::<'a>
            }
        }
    }
}
impl<'a> Iterator<Function<'a>> for Functions<'a> {
    fn next(&mut self) -> Option<Function> {
        unsafe {
            let native_next = jit_function_next(self.ctx, self.last);
            self.last = native_next;
            NativeRef::from_ptr(native_next)
        }
    }
    fn size_hint(&self) -> (uint, Option<uint>) {
        let mut size : uint = 0;
        let mut last = self.last;
        loop {
            last = unsafe {
                jit_function_next(self.ctx, last)
            };
            if last.is_null() {
                break;
            } else {
                size += 1;
            }
        }
        (size, Some(size))
    }
}<commit_msg>Always inline drop for context<commit_after>use bindings::{
    jit_context_t,
    jit_context_create,
    jit_context_destroy,
    jit_context_build_start,
    jit_context_build_end,
    jit_function_t,
    jit_function_next
};
use function::Function;
use std::kinds::marker::ContravariantLifetime;
use util::NativeRef;
/// Holds all of the functions you have built and compiled. There can be
/// multiple, but normally there is only one.
native_ref!(Context, _context, jit_context_t, ContravariantLifetime)
impl<'a> Context<'a> {
    #[inline(always)]
    /// Create a new JIT Context
    pub fn new() -> Context<'a> {
        unsafe {
            NativeRef::from_ptr(jit_context_create())
        }
    }
    #[inline(always)]
    /// Run a closure that can generate IR
    pub fn build<R>(&self, cb: || -> R) -> R {
        unsafe {
            jit_context_build_start(self.as_ptr());
            let rv = cb();
            jit_context_build_end(self.as_ptr());
            rv
        }
    }
    /// Iterate through all the functions in this context
    #[inline(always)]
    pub fn funcs(&self) -> Functions<'a> {
        Functions::new(self)
    }
}
#[unsafe_destructor]
impl<'a> Drop for Context<'a> {
    #[inline(always)]
    fn drop(&mut self) {
        unsafe {
            jit_context_destroy(self.as_ptr());
        }
    }
}
/// Any JIT object which is in a context
pub trait InContext<'a> {
    /// Get the context for this object
    fn get_context(&self) -> Context<'a>;
}

/// An iterator over a context's functions
pub struct Functions<'a> {
    ctx: jit_context_t,
    last: jit_function_t,
    marker: ContravariantLifetime<'a>
}
impl<'a> Functions<'a> {
    fn new(ctx:&Context<'a>) -> Functions<'a> {
        unsafe {
            Functions {
                ctx: ctx.as_ptr(),
                last: RawPtr::null(),
                marker: ContravariantLifetime::<'a>
            }
        }
    }
}
impl<'a> Iterator<Function<'a>> for Functions<'a> {
    fn next(&mut self) -> Option<Function> {
        unsafe {
            let native_next = jit_function_next(self.ctx, self.last);
            self.last = native_next;
            NativeRef::from_ptr(native_next)
        }
    }
    fn size_hint(&self) -> (uint, Option<uint>) {
        let mut size : uint = 0;
        let mut last = self.last;
        loop {
            last = unsafe {
                jit_function_next(self.ctx, last)
            };
            if last.is_null() {
                break;
            } else {
                size += 1;
            }
        }
        (size, Some(size))
    }
}<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A wrapper around any Reader to treat it as an RNG.

use Rng;

/// An RNG that reads random bytes straight from a `Reader`. This will
/// work best with an infinite reader, but this is not required.
///
/// It will fail if it there is insufficient data to fulfill a request.
///
/// # Example
///
/// ```rust
/// use rand::{reader, Rng};
/// use std::io::MemReader;
///
/// let mut rng = reader::ReaderRng::new(MemReader::new(~[1,2,3,4,5,6,7,8]));
/// println!("{:x}", rng.gen::<uint>());
/// ```
pub struct ReaderRng<R> {
    priv reader: R
}

impl<R: Reader> ReaderRng<R> {
    /// Create a new `ReaderRng` from a `Reader`.
    pub fn new(r: R) -> ReaderRng<R> {
        ReaderRng {
            reader: r
        }
    }
}

impl<R: Reader> Rng for ReaderRng<R> {
    fn next_u32(&mut self) -> u32 {
        // This is designed for speed: reading a LE integer on a LE
        // platform just involves blitting the bytes into the memory
        // of the u32, similarly for BE on BE; avoiding byteswapping.
        if cfg!(target_endian="little") {
            self.reader.read_le_u32().unwrap()
        } else {
            self.reader.read_be_u32().unwrap()
        }
    }
    fn next_u64(&mut self) -> u64 {
        // see above for explanation.
        if cfg!(target_endian="little") {
            self.reader.read_le_u64().unwrap()
        } else {
            self.reader.read_be_u64().unwrap()
        }
    }
    fn fill_bytes(&mut self, v: &mut [u8]) {
        if v.len() == 0 { return }
        match self.reader.read(v) {
            Ok(n) if n == v.len() => return,
            Ok(n) => fail!("ReaderRng.fill_bytes could not fill buffer: \
                            read {} out of {} bytes.", n, v.len()),
            Err(e) => fail!("ReaderRng.fill_bytes error: {}", e)
        }
    }
}

#[cfg(test)]
mod test {
    use super::ReaderRng;
    use std::io::MemReader;
    use std::cast;
    use Rng;

    #[test]
    fn test_reader_rng_u64() {
        // transmute from the target to avoid endianness concerns.
        let v = ~[1u64, 2u64, 3u64];
        let bytes: ~[u8] = unsafe {cast::transmute(v)};
        let mut rng = ReaderRng::new(MemReader::new(bytes));

        assert_eq!(rng.next_u64(), 1);
        assert_eq!(rng.next_u64(), 2);
        assert_eq!(rng.next_u64(), 3);
    }
    #[test]
    fn test_reader_rng_u32() {
        // transmute from the target to avoid endianness concerns.
        let v = ~[1u32, 2u32, 3u32];
        let bytes: ~[u8] = unsafe {cast::transmute(v)};
        let mut rng = ReaderRng::new(MemReader::new(bytes));

        assert_eq!(rng.next_u32(), 1);
        assert_eq!(rng.next_u32(), 2);
        assert_eq!(rng.next_u32(), 3);
    }
    #[test]
    fn test_reader_rng_fill_bytes() {
        let v = [1u8, 2, 3, 4, 5, 6, 7, 8];
        let mut w = [0u8, .. 8];

        let mut rng = ReaderRng::new(MemReader::new(v.to_owned()));
        rng.fill_bytes(w);

        assert!(v == w);
    }

    #[test]
    #[should_fail]
    fn test_reader_rng_insufficient_bytes() {
        let mut rng = ReaderRng::new(MemReader::new(~[]));
        let mut v = [0u8, .. 3];
        rng.fill_bytes(v);
    }
}
<commit_msg>rand: Use fill() instead of read()<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! A wrapper around any Reader to treat it as an RNG.

use Rng;

/// An RNG that reads random bytes straight from a `Reader`. This will
/// work best with an infinite reader, but this is not required.
///
/// It will fail if it there is insufficient data to fulfill a request.
///
/// # Example
///
/// ```rust
/// use rand::{reader, Rng};
/// use std::io::MemReader;
///
/// let mut rng = reader::ReaderRng::new(MemReader::new(~[1,2,3,4,5,6,7,8]));
/// println!("{:x}", rng.gen::<uint>());
/// ```
pub struct ReaderRng<R> {
    priv reader: R
}

impl<R: Reader> ReaderRng<R> {
    /// Create a new `ReaderRng` from a `Reader`.
    pub fn new(r: R) -> ReaderRng<R> {
        ReaderRng {
            reader: r
        }
    }
}

impl<R: Reader> Rng for ReaderRng<R> {
    fn next_u32(&mut self) -> u32 {
        // This is designed for speed: reading a LE integer on a LE
        // platform just involves blitting the bytes into the memory
        // of the u32, similarly for BE on BE; avoiding byteswapping.
        if cfg!(target_endian="little") {
            self.reader.read_le_u32().unwrap()
        } else {
            self.reader.read_be_u32().unwrap()
        }
    }
    fn next_u64(&mut self) -> u64 {
        // see above for explanation.
        if cfg!(target_endian="little") {
            self.reader.read_le_u64().unwrap()
        } else {
            self.reader.read_be_u64().unwrap()
        }
    }
    fn fill_bytes(&mut self, v: &mut [u8]) {
        if v.len() == 0 { return }
        match self.reader.fill(v) {
            Ok(()) => {}
            Err(e) => fail!("ReaderRng.fill_bytes error: {}", e)
        }
    }
}

#[cfg(test)]
mod test {
    use super::ReaderRng;
    use std::io::MemReader;
    use std::cast;
    use Rng;

    #[test]
    fn test_reader_rng_u64() {
        // transmute from the target to avoid endianness concerns.
        let v = ~[1u64, 2u64, 3u64];
        let bytes: ~[u8] = unsafe {cast::transmute(v)};
        let mut rng = ReaderRng::new(MemReader::new(bytes));

        assert_eq!(rng.next_u64(), 1);
        assert_eq!(rng.next_u64(), 2);
        assert_eq!(rng.next_u64(), 3);
    }
    #[test]
    fn test_reader_rng_u32() {
        // transmute from the target to avoid endianness concerns.
        let v = ~[1u32, 2u32, 3u32];
        let bytes: ~[u8] = unsafe {cast::transmute(v)};
        let mut rng = ReaderRng::new(MemReader::new(bytes));

        assert_eq!(rng.next_u32(), 1);
        assert_eq!(rng.next_u32(), 2);
        assert_eq!(rng.next_u32(), 3);
    }
    #[test]
    fn test_reader_rng_fill_bytes() {
        let v = [1u8, 2, 3, 4, 5, 6, 7, 8];
        let mut w = [0u8, .. 8];

        let mut rng = ReaderRng::new(MemReader::new(v.to_owned()));
        rng.fill_bytes(w);

        assert!(v == w);
    }

    #[test]
    #[should_fail]
    fn test_reader_rng_insufficient_bytes() {
        let mut rng = ReaderRng::new(MemReader::new(~[]));
        let mut v = [0u8, .. 3];
        rng.fill_bytes(v);
    }
}
<|endoftext|>


<commit_before>// Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(all(unix, not(target_os = "ios")))]
mod imp {
    extern crate libc;

    use self::OsRngInner::*;

    use io::{IoResult, File};
    use path::Path;
    use rand::Rng;
    use rand::reader::ReaderRng;
    use result::Result::{Ok, Err};
    use slice::SlicePrelude;
    use mem;
    use os::errno;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        extern "C" {
            fn syscall(number: libc::c_long, ...) -> libc::c_long;
        }

        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(target_arch = "arm")]
        const NR_GETRANDOM: libc::c_long = 384;

        unsafe {
            syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0u)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        let len = v.len();
        while read < len {
            let result = getrandom(v[mut read..]);
            if result == -1 {
                let err = errno() as libc::c_int;
                if err == libc::EINTR {
                    continue;
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as uint;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8, ..4] = [0u8, ..4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8, ..8] = [0u8, ..8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn is_getrandom_available() -> bool {
        use sync::atomic::{AtomicBool, INIT_ATOMIC_BOOL, Relaxed};

        static GETRANDOM_CHECKED: AtomicBool = INIT_ATOMIC_BOOL;
        static GETRANDOM_AVAILABLE: AtomicBool = INIT_ATOMIC_BOOL;

        if !GETRANDOM_CHECKED.load(Relaxed) {
            let mut buf: [u8, ..0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = errno() as libc::c_int;
                err != libc::ENOSYS
            } else {
                true
            };
            GETRANDOM_AVAILABLE.store(available, Relaxed);
            GETRANDOM_CHECKED.store(true, Relaxed);
            available
        } else {
            GETRANDOM_AVAILABLE.load(Relaxed)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn is_getrandom_available() -> bool { false }

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed
    /// This does not block.
    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReaderRng(ReaderRng<File>),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OsRng { inner: OsReaderRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReaderRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReaderRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReaderRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    extern crate libc;

    use io::{IoResult};
    use kinds::marker;
    use mem;
    use os;
    use rand::Rng;
    use result::Result::{Ok};
    use self::libc::{c_int, size_t};
    use slice::{SlicePrelude};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed
    /// This does not block.
    pub struct OsRng {
        marker: marker::NoCopy
    }

    #[repr(C)]
    struct SecRandom;

    #[allow(non_upper_case_globals)]
    static kSecRandomDefault: *const SecRandom = 0 as *const SecRandom;

    #[link(name = "Security", kind = "framework")]
    extern "C" {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            Ok(OsRng {marker: marker::NoCopy} )
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t, v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate libc;

    use io::{IoResult, IoError};
    use mem;
    use ops::Drop;
    use os;
    use rand::Rng;
    use result::Result::{Ok, Err};
    use self::libc::{DWORD, BYTE, LPCSTR, BOOL};
    use self::libc::types::os::arch::extra::{LONG_PTR};
    use slice::{SlicePrelude};

    type HCRYPTPROV = LONG_PTR;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed
    /// This does not block.
    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;

    #[allow(non_snake_case)]
    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

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

    use super::OsRng;
    use rand::Rng;
    use task;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0u, 20) {
            let (tx, rx) = channel();
            txs.push(tx);
            task::spawn(move|| {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0u, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(&mut v);
                    task::deschedule();
                }
            })
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<commit_msg>Fix Markdown syntax in docs for OsRng<commit_after>// Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(all(unix, not(target_os = "ios")))]
mod imp {
    extern crate libc;

    use self::OsRngInner::*;

    use io::{IoResult, File};
    use path::Path;
    use rand::Rng;
    use rand::reader::ReaderRng;
    use result::Result::{Ok, Err};
    use slice::SlicePrelude;
    use mem;
    use os::errno;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        extern "C" {
            fn syscall(number: libc::c_long, ...) -> libc::c_long;
        }

        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(target_arch = "arm")]
        const NR_GETRANDOM: libc::c_long = 384;

        unsafe {
            syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0u)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        let len = v.len();
        while read < len {
            let result = getrandom(v[mut read..]);
            if result == -1 {
                let err = errno() as libc::c_int;
                if err == libc::EINTR {
                    continue;
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as uint;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8, ..4] = [0u8, ..4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8, ..8] = [0u8, ..8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8, ..8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm")))]
    fn is_getrandom_available() -> bool {
        use sync::atomic::{AtomicBool, INIT_ATOMIC_BOOL, Relaxed};

        static GETRANDOM_CHECKED: AtomicBool = INIT_ATOMIC_BOOL;
        static GETRANDOM_AVAILABLE: AtomicBool = INIT_ATOMIC_BOOL;

        if !GETRANDOM_CHECKED.load(Relaxed) {
            let mut buf: [u8, ..0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = errno() as libc::c_int;
                err != libc::ENOSYS
            } else {
                true
            };
            GETRANDOM_AVAILABLE.store(available, Relaxed);
            GETRANDOM_CHECKED.store(true, Relaxed);
            available
        } else {
            GETRANDOM_AVAILABLE.load(Relaxed)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64", target_arch = "x86", target_arch = "arm"))))]
    fn is_getrandom_available() -> bool { false }

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReaderRng(ReaderRng<File>),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OsRng { inner: OsReaderRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReaderRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReaderRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReaderRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    extern crate libc;

    use io::{IoResult};
    use kinds::marker;
    use mem;
    use os;
    use rand::Rng;
    use result::Result::{Ok};
    use self::libc::{c_int, size_t};
    use slice::{SlicePrelude};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        marker: marker::NoCopy
    }

    #[repr(C)]
    struct SecRandom;

    #[allow(non_upper_case_globals)]
    static kSecRandomDefault: *const SecRandom = 0 as *const SecRandom;

    #[link(name = "Security", kind = "framework")]
    extern "C" {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            Ok(OsRng {marker: marker::NoCopy} )
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t, v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate libc;

    use io::{IoResult, IoError};
    use mem;
    use ops::Drop;
    use os;
    use rand::Rng;
    use result::Result::{Ok, Err};
    use self::libc::{DWORD, BYTE, LPCSTR, BOOL};
    use self::libc::types::os::arch::extra::{LONG_PTR};
    use slice::{SlicePrelude};

    type HCRYPTPROV = LONG_PTR;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;

    #[allow(non_snake_case)]
    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                panic!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

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

    use super::OsRng;
    use rand::Rng;
    use task;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0u, 20) {
            let (tx, rx) = channel();
            txs.push(tx);
            task::spawn(move|| {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0u, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(&mut v);
                    task::deschedule();
                }
            })
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! # The Unicode Library
//!
//! Unicode-intensive functions for `char` and `str` types.
//!
//! This crate provides a collection of Unicode-related functionality,
//! including decompositions, conversions, etc., and provides traits
//! implementing these functions for the `char` and `str` types.
//!
//! The functionality included here is only that which is necessary to
//! provide for basic string-related manipulations. This crate does not
//! (yet) aim to provide a full set of Unicode tables.

#![crate_name = "unicode"]
#![unstable(feature = "unicode")]
#![feature(staged_api)]
#![staged_api]
#![crate_type = "rlib"]
#![doc(html_logo_url = "http://www.rust-lang.org/logos/rust-logo-128x128-blk-v2.png",
       html_favicon_url = "http://www.rust-lang.org/favicon.ico",
       html_root_url = "http://doc.rust-lang.org/nightly/",
       html_playground_url = "http://play.rust-lang.org/")]
#![feature(no_std)]
#![no_std]
#![feature(slicing_syntax)]
#![feature(int_uint)]
#![feature(core)]

extern crate core;

// regex module
pub use tables::regex;

mod normalize;
mod tables;
mod u_char;
mod u_str;

// re-export char so that std et al see it correctly
/// Character manipulation (`char` type, Unicode Scalar Value)
///
/// This module provides the `CharExt` trait, as well as its
/// implementation for the primitive `char` type, in order to allow
/// basic character manipulation.
///
/// A `char` actually represents a
/// *[Unicode Scalar Value](http://www.unicode.org/glossary/#unicode_scalar_value)*,
/// as it can contain any Unicode code point except high-surrogate and
/// low-surrogate code points.
///
/// As such, only values in the ranges \[0x0,0xD7FF\] and \[0xE000,0x10FFFF\]
/// (inclusive) are allowed. A `char` can always be safely cast to a `u32`;
/// however the converse is not always true due to the above range limits
/// and, as such, should be performed via the `from_u32` function..
#[stable(feature = "rust1", since = "1.0.0")]
pub mod char {
    pub use core::char::{MAX, from_u32, from_digit};

    pub use normalize::{decompose_canonical, decompose_compatible, compose};

    pub use tables::normalization::canonical_combining_class;
    pub use tables::UNICODE_VERSION;

    pub use u_char::CharExt;
}

pub mod str {
    pub use u_str::{UnicodeStr, Words, Graphemes, GraphemeIndices};
    pub use u_str::{utf8_char_width, is_utf16, Utf16Items, Utf16Item};
    pub use u_str::{utf16_items, Utf16Encoder};
}

// NOTE: remove after next snapshot
// this lets us use #[derive(..)]
#[cfg(stage0)]
mod std {
    pub use core::clone;
    pub use core::cmp;
    pub use core::fmt;
    pub use core::marker;
    // for-loops
    pub use core::iter;
    pub use core::option;
}
<commit_msg>Auto merge of #22217 - mmatyas:char_typofix, r=huonw<commit_after>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! # The Unicode Library
//!
//! Unicode-intensive functions for `char` and `str` types.
//!
//! This crate provides a collection of Unicode-related functionality,
//! including decompositions, conversions, etc., and provides traits
//! implementing these functions for the `char` and `str` types.
//!
//! The functionality included here is only that which is necessary to
//! provide for basic string-related manipulations. This crate does not
//! (yet) aim to provide a full set of Unicode tables.

#![crate_name = "unicode"]
#![unstable(feature = "unicode")]
#![feature(staged_api)]
#![staged_api]
#![crate_type = "rlib"]
#![doc(html_logo_url = "http://www.rust-lang.org/logos/rust-logo-128x128-blk-v2.png",
       html_favicon_url = "http://www.rust-lang.org/favicon.ico",
       html_root_url = "http://doc.rust-lang.org/nightly/",
       html_playground_url = "http://play.rust-lang.org/")]
#![feature(no_std)]
#![no_std]
#![feature(slicing_syntax)]
#![feature(int_uint)]
#![feature(core)]

extern crate core;

// regex module
pub use tables::regex;

mod normalize;
mod tables;
mod u_char;
mod u_str;

// re-export char so that std et al see it correctly
/// Character manipulation (`char` type, Unicode Scalar Value)
///
/// This module provides the `CharExt` trait, as well as its
/// implementation for the primitive `char` type, in order to allow
/// basic character manipulation.
///
/// A `char` actually represents a
/// *[Unicode Scalar Value](http://www.unicode.org/glossary/#unicode_scalar_value)*,
/// as it can contain any Unicode code point except high-surrogate and
/// low-surrogate code points.
///
/// As such, only values in the ranges \[0x0,0xD7FF\] and \[0xE000,0x10FFFF\]
/// (inclusive) are allowed. A `char` can always be safely cast to a `u32`;
/// however the converse is not always true due to the above range limits
/// and, as such, should be performed via the `from_u32` function.
#[stable(feature = "rust1", since = "1.0.0")]
pub mod char {
    pub use core::char::{MAX, from_u32, from_digit};

    pub use normalize::{decompose_canonical, decompose_compatible, compose};

    pub use tables::normalization::canonical_combining_class;
    pub use tables::UNICODE_VERSION;

    pub use u_char::CharExt;
}

pub mod str {
    pub use u_str::{UnicodeStr, Words, Graphemes, GraphemeIndices};
    pub use u_str::{utf8_char_width, is_utf16, Utf16Items, Utf16Item};
    pub use u_str::{utf16_items, Utf16Encoder};
}

// NOTE: remove after next snapshot
// this lets us use #[derive(..)]
#[cfg(stage0)]
mod std {
    pub use core::clone;
    pub use core::cmp;
    pub use core::fmt;
    pub use core::marker;
    // for-loops
    pub use core::iter;
    pub use core::option;
}
<|endoftext|>


<commit_before>use std;
use base::*;
use super::{LoudsTrie, NodeID, LoudsPos, LinkID};

struct History<'a> {
    trie_: &'a LoudsTrie,
    node_id_: NodeID,
    louds_pos_: LoudsPos,
    link_id_: LinkID,
    key_pos_: u32,
    //key_id_: u32,
}

impl<'a> History<'a> {
    fn new(trie: &'a LoudsTrie, node_id: NodeID, louds_pos: LoudsPos,
           link_id: LinkID, key_pos: u32) -> History<'a> {
        History { trie_: trie, node_id_: node_id, louds_pos_: louds_pos,
                  link_id_: link_id, key_pos_: key_pos }
    }
}

struct State<'a> {
    history_: Vec<History<'a> >,
    key_buf_: Vec<u8>,
}

impl<'a> State<'a> {
    fn new() -> State<'a> {
        State { history_: Vec::new(), key_buf_: Vec::new() }
    }

    fn push<'b>(&mut self, key: &'b[u8], trie: &'a LoudsTrie, node_id: NodeID,
                louds_pos: LoudsPos, link_id: LinkID) {
        self.key_buf_.extend(key);
        assert!(self.key_buf_.len() <= std::u32::MAX as usize);
        self.history_.push(History::new(trie, node_id, louds_pos, link_id,
                                        self.key_buf_.len() as u32));
    }

    fn get_key(&self) -> &[u8] {
        &self.key_buf_[..]
    }

    fn pop(&mut self) -> Option<History<'a> > {
        self.history_.pop()
    }

    fn get_node_id(&self) -> NodeID {
        self.history_.last().unwrap().node_id_
    }

    fn get_link_id(&self) -> LinkID {
        self.history_.last().unwrap().link_id_
    }
}

pub struct Nav<'a> {
    state_: State<'a>,
    trie_: &'a LoudsTrie,
}

// For lookups, marisa does caching based on the input character.
// We can't do that here. May want to remove or rethink the cache
// implementation in light of this.

impl<'a> Nav<'a> {
    pub fn new(trie: &'a LoudsTrie) -> Nav<'a> {
        Nav { state_: State::new(), trie_: trie }
    }

    pub fn has_child(&mut self) -> bool {
        self.trie_.has_child(self.state_.get_node_id())
    }
    pub fn go_to_child(&mut self) -> bool {
        let init_node_id = self.state_.get_node_id();
        let init_link_id = self.state_.get_link_id();
        if let Some((node_id, louds_pos)) = self.trie_.child_pos(init_node_id) {
            if self.trie_.link_flags_.at(node_id.0 as usize) {
                let link_id = self.trie_.update_link_id(init_link_id.0 as usize,
                                                        node_id.0 as usize);

                let link = self.trie_.get_link_2(node_id.0 as usize, link_id);
                // Proceed either to next trie or tail
                match &self.trie_.next_trie_ {
                    &Some(ref trie) => {
                        // TODO: this.
                    },
                    &None => {
                        // FIXME: Shouldn't need this temporary vector.
                        //        'restore' should return an iterator, and
                        //        state.push should consume it.
                        let mut v = Vec::new();
                        self.trie_.tail_.restore(link, &mut v);

                        // Not sure if these values are correct/useful.
                        // If some stuff is only needed for some nodes... should
                        // reflect that in the types we use
                        self.state_.push(&v, self.trie_, node_id, louds_pos,
                                         LinkID(link_id as u32));
                    }
                }
            } else {
                let node_char = [ self.trie_.bases_[node_id.0 as usize] ];
                self.state_.push(&node_char, self.trie_, node_id, louds_pos,
                                 init_link_id);
            }
            true
        } else {
            false
        }
    }
    pub fn has_sibling(&self) -> bool {
        panic!("not implemented")
    }
    pub fn go_to_sibling(&mut self) -> bool {
        panic!("not implemented")
    }
    pub fn has_parent(&self) -> bool {
        panic!("not implemented")
    }
    pub fn go_to_parent(&self) -> bool {
        panic!("not implemented")
    }
    pub fn is_terminal(&self) -> bool {
        panic!("not implemented")
    }
    pub fn get_string(&self) -> &str {
        panic!("not implemented")
    }
    pub fn is_end(&self) -> bool {
        panic!("not implemented")
    }

}

/*

bool LoudsTrie::lookup(Agent &agent) const {
  MARISA_DEBUG_IF(!agent.has_state(), MARISA_STATE_ERROR);

  State &state = agent.state();
  state.lookup_init();
  while (state.query_pos() < agent.query().length()) {
    if (!find_child(agent)) {
      return false;
    }
  }
  if (!terminal_flags_[state.node_id()]) {
    return false;
  }
  agent.set_key(agent.query().ptr(), agent.query().length());
  agent.set_key(terminal_flags_.rank1(state.node_id()));
  return true;
}

bool LoudsTrie::find_child(Agent &agent) const {
  MARISA_DEBUG_IF(agent.state().query_pos() >= agent.query().length(),
      MARISA_BOUND_ERROR);

  State &state = agent.state();
  const std::size_t cache_id = get_cache_id(state.node_id(),
      agent.query()[state.query_pos()]);
  if (state.node_id() == cache_[cache_id].parent()) {
    if (cache_[cache_id].extra() != MARISA_INVALID_EXTRA) {
      if (!match(agent, cache_[cache_id].link())) {
        return false;
      }
    } else {
      state.set_query_pos(state.query_pos() + 1);
    }
    state.set_node_id(cache_[cache_id].child());
    return true;
  }

  std::size_t louds_pos = louds_.select0(state.node_id()) + 1;
  if (!louds_[louds_pos]) {
    return false;
  }
  state.set_node_id(louds_pos - state.node_id() - 1);
  std::size_t link_id = MARISA_INVALID_LINK_ID;
  do {
    if (link_flags_[state.node_id()]) {
      link_id = update_link_id(link_id, state.node_id());
      const std::size_t prev_query_pos = state.query_pos();
      if (match(agent, get_link(state.node_id(), link_id))) {
        return true;
      } else if (state.query_pos() != prev_query_pos) {
        return false;
      }
    } else if (bases_[state.node_id()] ==
        (UInt8)agent.query()[state.query_pos()]) {
      state.set_query_pos(state.query_pos() + 1);
      return true;
    }
    state.set_node_id(state.node_id() + 1);
    ++louds_pos;
  } while (louds_[louds_pos]);
  return false;
}

std::size_t LoudsTrie::get_cache_id(std::size_t node_id, char label) const {
  return (node_id ^ (node_id << 5) ^ (UInt8)label) & cache_mask_;
}

std::size_t LoudsTrie::get_cache_id(std::size_t node_id) const {
  return node_id & cache_mask_;
}

bool LoudsTrie::match(Agent &agent, std::size_t link) const {
  if (next_trie_.get() != NULL) {
    return next_trie_->match_(agent, link);
  } else {
    return tail_.match(agent, link);
  }
}

bool LoudsTrie::match_(Agent &agent, std::size_t node_id) const {
  MARISA_DEBUG_IF(agent.state().query_pos() >= agent.query().length(),
      MARISA_BOUND_ERROR);
  MARISA_DEBUG_IF(node_id == 0, MARISA_RANGE_ERROR);

  State &state = agent.state();
  for ( ; ; ) {
    const std::size_t cache_id = get_cache_id(node_id);
    if (node_id == cache_[cache_id].child()) {
      if (cache_[cache_id].extra() != MARISA_INVALID_EXTRA) {
        if (!match(agent, cache_[cache_id].link())) {
          return false;
        }
      } else if (cache_[cache_id].label() ==
          agent.query()[state.query_pos()]) {
        state.set_query_pos(state.query_pos() + 1);
      } else {
        return false;
      }

      node_id = cache_[cache_id].parent();
      if (node_id == 0) {
        return true;
      } else if (state.query_pos() >= agent.query().length()) {
        return false;
      }
      continue;
    }

    if (link_flags_[node_id]) {
      if (next_trie_.get() != NULL) {
        if (!match(agent, get_link(node_id))) {
          return false;
        }
      } else if (!tail_.match(agent, get_link(node_id))) {
        return false;
      }
    } else if (bases_[node_id] == (UInt8)agent.query()[state.query_pos()]) {
      state.set_query_pos(state.query_pos() + 1);
    } else {
      return false;
    }

    if (node_id <= num_l1_nodes_) {
      return true;
    } else if (state.query_pos() >= agent.query().length()) {
      return false;
    }
    node_id = louds_.select1(node_id) - node_id - 1;
  }
}

bool Tail::match(Agent &agent, std::size_t offset) const {
  MARISA_DEBUG_IF(buf_.empty(), MARISA_STATE_ERROR);
  MARISA_DEBUG_IF(agent.state().query_pos() >= agent.query().length(),
      MARISA_BOUND_ERROR);

  State &state = agent.state();
  if (end_flags_.empty()) {
    const char * const ptr = &buf_[offset] - state.query_pos();
    do {
      if (ptr[state.query_pos()] != agent.query()[state.query_pos()]) {
        return false;
      }
      state.set_query_pos(state.query_pos() + 1);
      if (ptr[state.query_pos()] == '\0') {
        return true;
      }
    } while (state.query_pos() < agent.query().length());
    return false;
  } else {
    do {
      if (buf_[offset] != agent.query()[state.query_pos()]) {
        return false;
      }
      state.set_query_pos(state.query_pos() + 1);
      if (end_flags_[offset++]) {
        return true;
      }
    } while (state.query_pos() < agent.query().length());
    return false;
  }
}

*/
<commit_msg>stack variable for trie<commit_after>use std;
use base::*;
use super::{LoudsTrie, NodeID, LoudsPos, LinkID};

struct History<'a> {
    trie_: &'a LoudsTrie,
    node_id_: NodeID,
    louds_pos_: LoudsPos,
    link_id_: LinkID,
    key_pos_: u32,
    //key_id_: u32,
}

impl<'a> History<'a> {
    fn new(trie: &'a LoudsTrie, node_id: NodeID, louds_pos: LoudsPos,
           link_id: LinkID, key_pos: u32) -> History<'a> {
        History { trie_: trie, node_id_: node_id, louds_pos_: louds_pos,
                  link_id_: link_id, key_pos_: key_pos }
    }
}

struct State<'a> {
    history_: Vec<History<'a> >,
    key_buf_: Vec<u8>,
}

impl<'a> State<'a> {
    fn new() -> State<'a> {
        State { history_: Vec::new(), key_buf_: Vec::new() }
    }

    fn push<'b>(&mut self, key: &'b[u8], trie: &'a LoudsTrie, node_id: NodeID,
                louds_pos: LoudsPos, link_id: LinkID) {
        self.key_buf_.extend(key);
        assert!(self.key_buf_.len() <= std::u32::MAX as usize);
        self.history_.push(History::new(trie, node_id, louds_pos, link_id,
                                        self.key_buf_.len() as u32));
    }

    fn get_key(&self) -> &[u8] {
        &self.key_buf_[..]
    }

    fn pop(&mut self) -> Option<History<'a> > {
        self.history_.pop()
    }

    fn get_node_id(&self) -> NodeID {
        self.history_.last().unwrap().node_id_
    }

    fn get_link_id(&self) -> LinkID {
        self.history_.last().unwrap().link_id_
    }
}

pub struct Nav<'a> {
    state_: State<'a>,
    trie_: &'a LoudsTrie,
}

// For lookups, marisa does caching based on the input character.
// We can't do that here. May want to remove or rethink the cache
// implementation in light of this.

impl<'a> Nav<'a> {
    pub fn new(trie: &'a LoudsTrie) -> Nav<'a> {
        Nav { state_: State::new(), trie_: trie }
    }

    pub fn has_child(&mut self) -> bool {
        self.trie_.has_child(self.state_.get_node_id())
    }
    pub fn go_to_child(&mut self) -> bool {
        let init_node_id = self.state_.get_node_id();
        let init_link_id = self.state_.get_link_id();

        let mut trie = &self.trie_;

        if let Some((node_id, louds_pos)) = trie.child_pos(init_node_id) {
            if trie.link_flags_.at(node_id.0 as usize) {
                let link_id = trie.update_link_id(init_link_id.0 as usize,
                                                  node_id.0 as usize);

                let link = trie.get_link_2(node_id.0 as usize, link_id);
                // Proceed either to next trie or tail
                match &trie.next_trie_ {
                    &Some(ref trie) => {
                        // TODO: this.
                    },
                    &None => {
                        // FIXME: Shouldn't need this temporary vector.
                        //        'restore' should return an iterator, and
                        //        state.push should consume it.
                        let mut v = Vec::new();
                        trie.tail_.restore(link, &mut v);

                        // Not sure if these values are correct/useful.
                        // If some stuff is only needed for some nodes... should
                        // reflect that in the types we use
                        self.state_.push(&v, trie, node_id, louds_pos,
                                         LinkID(link_id as u32));
                    }
                }
            } else {
                let node_char = [ trie.bases_[node_id.0 as usize] ];
                self.state_.push(&node_char, trie, node_id, louds_pos,
                                 init_link_id);
            }
            true
        } else {
            false
        }
    }
    pub fn has_sibling(&self) -> bool {
        panic!("not implemented")
    }
    pub fn go_to_sibling(&mut self) -> bool {
        panic!("not implemented")
    }
    pub fn has_parent(&self) -> bool {
        panic!("not implemented")
    }
    pub fn go_to_parent(&self) -> bool {
        panic!("not implemented")
    }
    pub fn is_terminal(&self) -> bool {
        panic!("not implemented")
    }
    pub fn get_string(&self) -> &str {
        panic!("not implemented")
    }
    pub fn is_end(&self) -> bool {
        panic!("not implemented")
    }

}

/*

bool LoudsTrie::lookup(Agent &agent) const {
  MARISA_DEBUG_IF(!agent.has_state(), MARISA_STATE_ERROR);

  State &state = agent.state();
  state.lookup_init();
  while (state.query_pos() < agent.query().length()) {
    if (!find_child(agent)) {
      return false;
    }
  }
  if (!terminal_flags_[state.node_id()]) {
    return false;
  }
  agent.set_key(agent.query().ptr(), agent.query().length());
  agent.set_key(terminal_flags_.rank1(state.node_id()));
  return true;
}

bool LoudsTrie::find_child(Agent &agent) const {
  MARISA_DEBUG_IF(agent.state().query_pos() >= agent.query().length(),
      MARISA_BOUND_ERROR);

  State &state = agent.state();
  const std::size_t cache_id = get_cache_id(state.node_id(),
      agent.query()[state.query_pos()]);
  if (state.node_id() == cache_[cache_id].parent()) {
    if (cache_[cache_id].extra() != MARISA_INVALID_EXTRA) {
      if (!match(agent, cache_[cache_id].link())) {
        return false;
      }
    } else {
      state.set_query_pos(state.query_pos() + 1);
    }
    state.set_node_id(cache_[cache_id].child());
    return true;
  }

  std::size_t louds_pos = louds_.select0(state.node_id()) + 1;
  if (!louds_[louds_pos]) {
    return false;
  }
  state.set_node_id(louds_pos - state.node_id() - 1);
  std::size_t link_id = MARISA_INVALID_LINK_ID;
  do {
    if (link_flags_[state.node_id()]) {
      link_id = update_link_id(link_id, state.node_id());
      const std::size_t prev_query_pos = state.query_pos();
      if (match(agent, get_link(state.node_id(), link_id))) {
        return true;
      } else if (state.query_pos() != prev_query_pos) {
        return false;
      }
    } else if (bases_[state.node_id()] ==
        (UInt8)agent.query()[state.query_pos()]) {
      state.set_query_pos(state.query_pos() + 1);
      return true;
    }
    state.set_node_id(state.node_id() + 1);
    ++louds_pos;
  } while (louds_[louds_pos]);
  return false;
}

std::size_t LoudsTrie::get_cache_id(std::size_t node_id, char label) const {
  return (node_id ^ (node_id << 5) ^ (UInt8)label) & cache_mask_;
}

std::size_t LoudsTrie::get_cache_id(std::size_t node_id) const {
  return node_id & cache_mask_;
}

bool LoudsTrie::match(Agent &agent, std::size_t link) const {
  if (next_trie_.get() != NULL) {
    return next_trie_->match_(agent, link);
  } else {
    return tail_.match(agent, link);
  }
}

bool LoudsTrie::match_(Agent &agent, std::size_t node_id) const {
  MARISA_DEBUG_IF(agent.state().query_pos() >= agent.query().length(),
      MARISA_BOUND_ERROR);
  MARISA_DEBUG_IF(node_id == 0, MARISA_RANGE_ERROR);

  State &state = agent.state();
  for ( ; ; ) {
    const std::size_t cache_id = get_cache_id(node_id);
    if (node_id == cache_[cache_id].child()) {
      if (cache_[cache_id].extra() != MARISA_INVALID_EXTRA) {
        if (!match(agent, cache_[cache_id].link())) {
          return false;
        }
      } else if (cache_[cache_id].label() ==
          agent.query()[state.query_pos()]) {
        state.set_query_pos(state.query_pos() + 1);
      } else {
        return false;
      }

      node_id = cache_[cache_id].parent();
      if (node_id == 0) {
        return true;
      } else if (state.query_pos() >= agent.query().length()) {
        return false;
      }
      continue;
    }

    if (link_flags_[node_id]) {
      if (next_trie_.get() != NULL) {
        if (!match(agent, get_link(node_id))) {
          return false;
        }
      } else if (!tail_.match(agent, get_link(node_id))) {
        return false;
      }
    } else if (bases_[node_id] == (UInt8)agent.query()[state.query_pos()]) {
      state.set_query_pos(state.query_pos() + 1);
    } else {
      return false;
    }

    if (node_id <= num_l1_nodes_) {
      return true;
    } else if (state.query_pos() >= agent.query().length()) {
      return false;
    }
    node_id = louds_.select1(node_id) - node_id - 1;
  }
}

bool Tail::match(Agent &agent, std::size_t offset) const {
  MARISA_DEBUG_IF(buf_.empty(), MARISA_STATE_ERROR);
  MARISA_DEBUG_IF(agent.state().query_pos() >= agent.query().length(),
      MARISA_BOUND_ERROR);

  State &state = agent.state();
  if (end_flags_.empty()) {
    const char * const ptr = &buf_[offset] - state.query_pos();
    do {
      if (ptr[state.query_pos()] != agent.query()[state.query_pos()]) {
        return false;
      }
      state.set_query_pos(state.query_pos() + 1);
      if (ptr[state.query_pos()] == '\0') {
        return true;
      }
    } while (state.query_pos() < agent.query().length());
    return false;
  } else {
    do {
      if (buf_[offset] != agent.query()[state.query_pos()]) {
        return false;
      }
      state.set_query_pos(state.query_pos() + 1);
      if (end_flags_[offset++]) {
        return true;
      }
    } while (state.query_pos() < agent.query().length());
    return false;
  }
}

*/
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::ast::{Ast, AstBuilder, AstKey, ObjectEntry, Primitive, RequestParameters};
use crate::build_ast::{
    build_fragment, build_operation, build_request, build_request_params,
    build_request_params_ast_key,
};
use crate::constants::CODEGEN_CONSTANTS;
use crate::indentation::print_indentation;
use crate::utils::escape;

use graphql_ir::{FragmentDefinition, OperationDefinition};
use schema::Schema;

use fnv::{FnvBuildHasher, FnvHashSet};
use indexmap::IndexMap;
use interner::StringKey;
use std::fmt::{Result as FmtResult, Write};

pub struct Printer {
    builder: AstBuilder,
}

type VariableDefinitions = IndexMap<AstKey, String, FnvBuildHasher>;
struct DedupedJSONPrinter<'b> {
    variable_definitions: VariableDefinitions,
    duplicates: FnvHashSet<AstKey>,
    builder: &'b AstBuilder,
    root_key: AstKey,
}

impl<'b> DedupedJSONPrinter<'b> {
    fn new(builder: &'b AstBuilder, root_key: AstKey) -> Self {
        let mut visited = Default::default();
        let mut duplicates = Default::default();
        Self::collect_value_duplicates(&mut visited, &mut duplicates, builder, root_key);
        Self {
            variable_definitions: Default::default(),
            duplicates,
            builder,
            root_key,
        }
    }

    fn new_without_dedupe(builder: &'b AstBuilder, root_key: AstKey) -> Self {
        Self {
            variable_definitions: Default::default(),
            duplicates: Default::default(),
            builder,
            root_key,
        }
    }

    /// We don't dedupe in an already deduped AST unless the duplicate
    /// also appears on a subtree that's not a duplicate
    /// Input:
    /// [
    ///     [{a: 1}, {b:2}],
    ///     [{a: 1}, {b:2}],
    ///     {b: 2}
    /// ]
    /// Output:
    /// v0 = {b: 2};
    /// v1 = [{a: 1}, v0];
    fn collect_value_duplicates(
        visited: &mut FnvHashSet<AstKey>,
        duplicates: &mut FnvHashSet<AstKey>,
        builder: &AstBuilder,
        key: AstKey,
    ) {
        match builder.lookup(key) {
            Ast::Array(array) => {
                if array.is_empty() {
                    return;
                }
                if !visited.insert(key) {
                    duplicates.insert(key);
                    return;
                }
                for val in array {
                    if let Primitive::Key(key) = val {
                        Self::collect_value_duplicates(visited, duplicates, builder, *key);
                    }
                }
            }
            Ast::Object(object) => {
                if object.is_empty() {
                    return;
                }
                if !visited.insert(key) {
                    duplicates.insert(key);
                    return;
                }
                for (_, val) in object {
                    if let Primitive::Key(key) = val {
                        Self::collect_value_duplicates(visited, duplicates, builder, *key);
                    }
                }
            }
        }
    }

    fn print(mut self) -> String {
        let mut result = String::new();
        self.print_ast(&mut result, self.root_key, 0, false);
        if self.variable_definitions.is_empty() {
            result
        } else {
            let mut with_variables = String::new();
            with_variables.push_str("(function(){\nvar ");
            let last = self.variable_definitions.len() - 1;
            for (i, (_, value)) in self.variable_definitions.drain(..).enumerate() {
                writeln!(
                    &mut with_variables,
                    "v{} = {}{}",
                    i,
                    value,
                    if i == last { ";" } else { "," }
                )
                .unwrap();
            }
            write!(&mut with_variables, "return {};\n}})()", result).unwrap();
            with_variables
        }
    }

    fn print_ast(&mut self, f: &mut String, key: AstKey, indent: usize, is_dedupe_var: bool) {
        // Only use variable references at depth beyond the top level.
        if indent > 0 && self.duplicates.contains(&key) {
            let v = if self.variable_definitions.contains_key(&key) {
                self.variable_definitions.get_full(&key).unwrap().0
            } else {
                let mut variable = String::new();
                self.print_ast(&mut variable, key, 0, true);
                let v = self.variable_definitions.len();
                self.variable_definitions.insert(key, variable);
                v
            };
            return write!(f, "(v{}/*: any*/)", v).unwrap();
        }

        let ast = self.builder.lookup(key);
        match ast {
            Ast::Object(object) => {
                if object.is_empty() {
                    f.push_str("{}");
                } else {
                    let next_indent = indent + 1;
                    f.push('{');
                    for (key, value) in object {
                        f.push('\n');
                        print_indentation(f, next_indent);
                        write!(f, "\"{}\": ", key.lookup()).unwrap();
                        self.print_primitive(f, value, next_indent, is_dedupe_var)
                            .unwrap();
                        f.push(',');
                    }
                    f.pop();
                    f.push('\n');
                    print_indentation(f, indent);
                    f.push('}');
                }
            }
            Ast::Array(array) => {
                if array.is_empty() {
                    if is_dedupe_var {
                        // Empty arrays can only have one inferred flow type and then conflict if
                        // used in different places, this is unsound if we would write to them but
                        // this whole module is based on the idea of a read only JSON tree.
                        f.push_str("([]/*: any*/)");
                    } else {
                        f.push_str("[]");
                    }
                } else {
                    f.push('[');
                    let next_indent = indent + 1;
                    for value in array {
                        f.push('\n');
                        print_indentation(f, next_indent);
                        self.print_primitive(f, value, next_indent, is_dedupe_var)
                            .unwrap();
                        f.push(',');
                    }
                    f.pop();
                    f.push('\n');
                    print_indentation(f, indent);
                    f.push(']');
                }
            }
        }
    }

    fn print_primitive(
        &mut self,
        f: &mut String,
        primitive: &Primitive,
        indent: usize,
        is_dedupe_var: bool,
    ) -> FmtResult {
        match primitive {
            Primitive::Null => write!(f, "null"),
            Primitive::Bool(b) => write!(f, "{}", if *b { "true" } else { "false" }),
            Primitive::RawString(str) => {
                f.push('\"');
                escape(str, f);
                f.push('\"');
                Ok(())
            }
            Primitive::String(key) => write!(f, "\"{}\"", key),
            Primitive::Float(value) => write!(f, "{}", value.as_float()),
            Primitive::Int(value) => write!(f, "{}", value),
            Primitive::Key(key) => {
                self.print_ast(f, *key, indent, is_dedupe_var);
                Ok(())
            }
            Primitive::StorageKey(field_name, key) => {
                print_static_storage_key(f, &self.builder, *field_name, *key)
            }
            Primitive::ModuleDependency(key) => write!(f, "require('{}.graphql')", key),
        }
    }
}

impl Default for Printer {
    fn default() -> Self {
        Self {
            builder: AstBuilder::default(),
        }
    }
}

pub fn print_operation(schema: &Schema, operation: &OperationDefinition) -> String {
    let mut builder = AstBuilder::default();
    let key = build_operation(schema, &mut builder, operation);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_fragment(schema: &Schema, fragment: &FragmentDefinition) -> String {
    let mut builder = AstBuilder::default();
    let key = build_fragment(schema, &mut builder, fragment);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_request(
    schema: &Schema,
    operation: &OperationDefinition,
    fragment: &FragmentDefinition,
    request_parameters: RequestParameters,
) -> String {
    let mut builder = AstBuilder::default();
    let request_parameters =
        build_request_params_ast_key(schema, request_parameters, &mut builder, operation);
    let key = build_request(
        schema,
        &mut builder,
        operation,
        fragment,
        request_parameters,
    );
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_request_params(
    schema: &Schema,
    operation: &OperationDefinition,
    query_id: Option<String>,
) -> String {
    let mut request_parameters = build_request_params(operation);
    request_parameters.id = query_id;

    let mut builder = AstBuilder::default();
    let request_parameters_ast_key =
        build_request_params_ast_key(schema, request_parameters, &mut builder, operation);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, request_parameters_ast_key);

    printer.print()
}

impl Printer {
    pub fn print_request_deduped(
        &mut self,
        schema: &Schema,
        operation: &OperationDefinition,
        fragment: &FragmentDefinition,
        request_parameters: RequestParameters,
    ) -> String {
        let request_parameters =
            build_request_params_ast_key(schema, request_parameters, &mut self.builder, operation);
        let key = build_request(
            schema,
            &mut self.builder,
            operation,
            fragment,
            request_parameters,
        );
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }

    pub fn print_operation_deduped(
        &mut self,
        schema: &Schema,
        operation: &OperationDefinition,
    ) -> String {
        let key = build_operation(schema, &mut self.builder, operation);
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }

    pub fn print_fragment_deduped(
        &mut self,
        schema: &Schema,
        fragment: &FragmentDefinition,
    ) -> String {
        let key = build_fragment(schema, &mut self.builder, fragment);
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }
}

/// Pre-computes storage key if possible and advantageous. Storage keys are
/// generated for fields with supplied arguments that are all statically known
/// (ie. literals, no variables) at build time.
fn print_static_storage_key(
    f: &mut String,
    builder: &AstBuilder,
    field_name: StringKey,
    args_key: AstKey,
) -> FmtResult {
    // TODO (T64585375): JS compiler has an option to force a storageKey.
    let args = builder.lookup(args_key).assert_array();
    write_static_storage_key(f, builder, field_name, args)
}

fn write_static_storage_key(
    f: &mut String,
    builder: &AstBuilder,
    field_name: StringKey,
    args: &[Primitive],
) -> FmtResult {
    write!(f, "\"{}(", field_name)?;
    for arg_key in args {
        let arg = builder.lookup(arg_key.assert_key()).assert_object();
        let (_, name) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.name)
            .expect("Expected `name` to exist");
        let name = name.assert_string();
        write!(f, "{}:", name)?;
        write_argument_value(f, builder, arg)?;
        f.push(',');
    }
    f.pop(); // args won't be empty
    f.push_str(")\"");
    Ok(())
}

fn write_argument_value(f: &mut String, builder: &AstBuilder, arg: &[ObjectEntry]) -> FmtResult {
    let (_, key) = arg
        .iter()
        .find(|(key, _)| *key == CODEGEN_CONSTANTS.kind)
        .expect("Expected `kind` to exist");
    let key = key.assert_string();
    // match doesn't allow `CODEGEN_CONSTANTS.<>` on the match arm, falling back to if statements
    if key == CODEGEN_CONSTANTS.literal {
        let (_, literal) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.value)
            .expect("Expected `name` to exist");
        write_constant_value(f, builder, literal)?;
    } else if key == CODEGEN_CONSTANTS.list_value {
        let (_, items) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.items)
            .expect("Expected `items` to exist");
        let array = builder.lookup(items.assert_key()).assert_array();

        f.push('[');
        let mut after_first = false;
        for key_or_null in array {
            match key_or_null {
                Primitive::Null => {}
                Primitive::Key(key) => {
                    if after_first {
                        f.push(',');
                    } else {
                        after_first = true;
                    }
                    let object = builder.lookup(*key).assert_object();
                    write_argument_value(f, builder, object)?;
                }
                _ => panic!("Expected an object key or null"),
            }
        }
        f.push(']');
    } else {
        // We filtered out Variables, here it should only be ObjectValue
        let (_, value) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.fields)
            .expect("Expected `fields` to exist");
        let array = builder.lookup(value.assert_key()).assert_array();

        f.push('{');
        for key in array {
            let (_, name) = arg
                .iter()
                .find(|(key, _)| *key == CODEGEN_CONSTANTS.name)
                .expect("Expected `name` to exist");
            let name = name.assert_string();
            write!(f, "\\\"{}\\\":", name)?;
            let object = builder.lookup(key.assert_key()).assert_object();
            write_argument_value(f, builder, object)?;
            f.push(',');
        }
        if !array.is_empty() {
            f.pop();
        }
        f.push('}');
    }
    Ok(())
}

fn write_constant_value(f: &mut String, builder: &AstBuilder, value: &Primitive) -> FmtResult {
    match value {
        Primitive::Bool(b) => write!(f, "{}", if *b { "true" } else { "false" }),
        Primitive::String(key) => write!(f, "\\\"{}\\\"", key),
        Primitive::Float(value) => write!(f, "{}", value.as_float()),
        Primitive::Int(value) => write!(f, "{}", value),
        Primitive::Key(key) => {
            let ast = builder.lookup(*key);
            match ast {
                Ast::Array(arr) => {
                    f.push('[');
                    for value in arr {
                        write_constant_value(f, builder, value)?;
                        f.push(',');
                    }
                    if !arr.is_empty() {
                        f.pop();
                    }
                    f.push(']');
                    Ok(())
                }
                Ast::Object(obj) => {
                    f.push('{');
                    for (name, value) in obj {
                        write!(f, "\\\"{}\\\":", name)?;
                        write_constant_value(f, builder, value)?;
                        f.push(',');
                    }
                    if !obj.is_empty() {
                        f.pop();
                    }
                    f.push('}');
                    Ok(())
                }
            }
        }
        Primitive::Null => {
            f.push_str("null");
            Ok(())
        }
        Primitive::StorageKey(_, _) => panic!("Unexpected StorageKey"),
        Primitive::RawString(_) => panic!("Unexpected RawString"),
        Primitive::ModuleDependency(_) => panic!("Unexpected ModuleDependency"),
    }
}
<commit_msg>Fix a bug in object storageKey printing<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::ast::{Ast, AstBuilder, AstKey, ObjectEntry, Primitive, RequestParameters};
use crate::build_ast::{
    build_fragment, build_operation, build_request, build_request_params,
    build_request_params_ast_key,
};
use crate::constants::CODEGEN_CONSTANTS;
use crate::indentation::print_indentation;
use crate::utils::escape;

use graphql_ir::{FragmentDefinition, OperationDefinition};
use schema::Schema;

use fnv::{FnvBuildHasher, FnvHashSet};
use indexmap::IndexMap;
use interner::StringKey;
use std::fmt::{Result as FmtResult, Write};

pub struct Printer {
    builder: AstBuilder,
}

type VariableDefinitions = IndexMap<AstKey, String, FnvBuildHasher>;
struct DedupedJSONPrinter<'b> {
    variable_definitions: VariableDefinitions,
    duplicates: FnvHashSet<AstKey>,
    builder: &'b AstBuilder,
    root_key: AstKey,
}

impl<'b> DedupedJSONPrinter<'b> {
    fn new(builder: &'b AstBuilder, root_key: AstKey) -> Self {
        let mut visited = Default::default();
        let mut duplicates = Default::default();
        Self::collect_value_duplicates(&mut visited, &mut duplicates, builder, root_key);
        Self {
            variable_definitions: Default::default(),
            duplicates,
            builder,
            root_key,
        }
    }

    fn new_without_dedupe(builder: &'b AstBuilder, root_key: AstKey) -> Self {
        Self {
            variable_definitions: Default::default(),
            duplicates: Default::default(),
            builder,
            root_key,
        }
    }

    /// We don't dedupe in an already deduped AST unless the duplicate
    /// also appears on a subtree that's not a duplicate
    /// Input:
    /// [
    ///     [{a: 1}, {b:2}],
    ///     [{a: 1}, {b:2}],
    ///     {b: 2}
    /// ]
    /// Output:
    /// v0 = {b: 2};
    /// v1 = [{a: 1}, v0];
    fn collect_value_duplicates(
        visited: &mut FnvHashSet<AstKey>,
        duplicates: &mut FnvHashSet<AstKey>,
        builder: &AstBuilder,
        key: AstKey,
    ) {
        match builder.lookup(key) {
            Ast::Array(array) => {
                if array.is_empty() {
                    return;
                }
                if !visited.insert(key) {
                    duplicates.insert(key);
                    return;
                }
                for val in array {
                    if let Primitive::Key(key) = val {
                        Self::collect_value_duplicates(visited, duplicates, builder, *key);
                    }
                }
            }
            Ast::Object(object) => {
                if object.is_empty() {
                    return;
                }
                if !visited.insert(key) {
                    duplicates.insert(key);
                    return;
                }
                for (_, val) in object {
                    if let Primitive::Key(key) = val {
                        Self::collect_value_duplicates(visited, duplicates, builder, *key);
                    }
                }
            }
        }
    }

    fn print(mut self) -> String {
        let mut result = String::new();
        self.print_ast(&mut result, self.root_key, 0, false);
        if self.variable_definitions.is_empty() {
            result
        } else {
            let mut with_variables = String::new();
            with_variables.push_str("(function(){\nvar ");
            let last = self.variable_definitions.len() - 1;
            for (i, (_, value)) in self.variable_definitions.drain(..).enumerate() {
                writeln!(
                    &mut with_variables,
                    "v{} = {}{}",
                    i,
                    value,
                    if i == last { ";" } else { "," }
                )
                .unwrap();
            }
            write!(&mut with_variables, "return {};\n}})()", result).unwrap();
            with_variables
        }
    }

    fn print_ast(&mut self, f: &mut String, key: AstKey, indent: usize, is_dedupe_var: bool) {
        // Only use variable references at depth beyond the top level.
        if indent > 0 && self.duplicates.contains(&key) {
            let v = if self.variable_definitions.contains_key(&key) {
                self.variable_definitions.get_full(&key).unwrap().0
            } else {
                let mut variable = String::new();
                self.print_ast(&mut variable, key, 0, true);
                let v = self.variable_definitions.len();
                self.variable_definitions.insert(key, variable);
                v
            };
            return write!(f, "(v{}/*: any*/)", v).unwrap();
        }

        let ast = self.builder.lookup(key);
        match ast {
            Ast::Object(object) => {
                if object.is_empty() {
                    f.push_str("{}");
                } else {
                    let next_indent = indent + 1;
                    f.push('{');
                    for (key, value) in object {
                        f.push('\n');
                        print_indentation(f, next_indent);
                        write!(f, "\"{}\": ", key.lookup()).unwrap();
                        self.print_primitive(f, value, next_indent, is_dedupe_var)
                            .unwrap();
                        f.push(',');
                    }
                    f.pop();
                    f.push('\n');
                    print_indentation(f, indent);
                    f.push('}');
                }
            }
            Ast::Array(array) => {
                if array.is_empty() {
                    if is_dedupe_var {
                        // Empty arrays can only have one inferred flow type and then conflict if
                        // used in different places, this is unsound if we would write to them but
                        // this whole module is based on the idea of a read only JSON tree.
                        f.push_str("([]/*: any*/)");
                    } else {
                        f.push_str("[]");
                    }
                } else {
                    f.push('[');
                    let next_indent = indent + 1;
                    for value in array {
                        f.push('\n');
                        print_indentation(f, next_indent);
                        self.print_primitive(f, value, next_indent, is_dedupe_var)
                            .unwrap();
                        f.push(',');
                    }
                    f.pop();
                    f.push('\n');
                    print_indentation(f, indent);
                    f.push(']');
                }
            }
        }
    }

    fn print_primitive(
        &mut self,
        f: &mut String,
        primitive: &Primitive,
        indent: usize,
        is_dedupe_var: bool,
    ) -> FmtResult {
        match primitive {
            Primitive::Null => write!(f, "null"),
            Primitive::Bool(b) => write!(f, "{}", if *b { "true" } else { "false" }),
            Primitive::RawString(str) => {
                f.push('\"');
                escape(str, f);
                f.push('\"');
                Ok(())
            }
            Primitive::String(key) => write!(f, "\"{}\"", key),
            Primitive::Float(value) => write!(f, "{}", value.as_float()),
            Primitive::Int(value) => write!(f, "{}", value),
            Primitive::Key(key) => {
                self.print_ast(f, *key, indent, is_dedupe_var);
                Ok(())
            }
            Primitive::StorageKey(field_name, key) => {
                print_static_storage_key(f, &self.builder, *field_name, *key)
            }
            Primitive::ModuleDependency(key) => write!(f, "require('{}.graphql')", key),
        }
    }
}

impl Default for Printer {
    fn default() -> Self {
        Self {
            builder: AstBuilder::default(),
        }
    }
}

pub fn print_operation(schema: &Schema, operation: &OperationDefinition) -> String {
    let mut builder = AstBuilder::default();
    let key = build_operation(schema, &mut builder, operation);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_fragment(schema: &Schema, fragment: &FragmentDefinition) -> String {
    let mut builder = AstBuilder::default();
    let key = build_fragment(schema, &mut builder, fragment);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_request(
    schema: &Schema,
    operation: &OperationDefinition,
    fragment: &FragmentDefinition,
    request_parameters: RequestParameters,
) -> String {
    let mut builder = AstBuilder::default();
    let request_parameters =
        build_request_params_ast_key(schema, request_parameters, &mut builder, operation);
    let key = build_request(
        schema,
        &mut builder,
        operation,
        fragment,
        request_parameters,
    );
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, key);
    printer.print()
}

pub fn print_request_params(
    schema: &Schema,
    operation: &OperationDefinition,
    query_id: Option<String>,
) -> String {
    let mut request_parameters = build_request_params(operation);
    request_parameters.id = query_id;

    let mut builder = AstBuilder::default();
    let request_parameters_ast_key =
        build_request_params_ast_key(schema, request_parameters, &mut builder, operation);
    let printer = DedupedJSONPrinter::new_without_dedupe(&builder, request_parameters_ast_key);

    printer.print()
}

impl Printer {
    pub fn print_request_deduped(
        &mut self,
        schema: &Schema,
        operation: &OperationDefinition,
        fragment: &FragmentDefinition,
        request_parameters: RequestParameters,
    ) -> String {
        let request_parameters =
            build_request_params_ast_key(schema, request_parameters, &mut self.builder, operation);
        let key = build_request(
            schema,
            &mut self.builder,
            operation,
            fragment,
            request_parameters,
        );
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }

    pub fn print_operation_deduped(
        &mut self,
        schema: &Schema,
        operation: &OperationDefinition,
    ) -> String {
        let key = build_operation(schema, &mut self.builder, operation);
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }

    pub fn print_fragment_deduped(
        &mut self,
        schema: &Schema,
        fragment: &FragmentDefinition,
    ) -> String {
        let key = build_fragment(schema, &mut self.builder, fragment);
        let deduped_printer = DedupedJSONPrinter::new(&self.builder, key);
        deduped_printer.print()
    }
}

/// Pre-computes storage key if possible and advantageous. Storage keys are
/// generated for fields with supplied arguments that are all statically known
/// (ie. literals, no variables) at build time.
fn print_static_storage_key(
    f: &mut String,
    builder: &AstBuilder,
    field_name: StringKey,
    args_key: AstKey,
) -> FmtResult {
    // TODO (T64585375): JS compiler has an option to force a storageKey.
    let args = builder.lookup(args_key).assert_array();
    write_static_storage_key(f, builder, field_name, args)
}

fn write_static_storage_key(
    f: &mut String,
    builder: &AstBuilder,
    field_name: StringKey,
    args: &[Primitive],
) -> FmtResult {
    write!(f, "\"{}(", field_name)?;
    for arg_key in args {
        let arg = builder.lookup(arg_key.assert_key()).assert_object();
        let (_, name) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.name)
            .expect("Expected `name` to exist");
        let name = name.assert_string();
        write!(f, "{}:", name)?;
        write_argument_value(f, builder, arg)?;
        f.push(',');
    }
    f.pop(); // args won't be empty
    f.push_str(")\"");
    Ok(())
}

fn write_argument_value(f: &mut String, builder: &AstBuilder, arg: &[ObjectEntry]) -> FmtResult {
    let (_, key) = arg
        .iter()
        .find(|(key, _)| *key == CODEGEN_CONSTANTS.kind)
        .expect("Expected `kind` to exist");
    let key = key.assert_string();
    // match doesn't allow `CODEGEN_CONSTANTS.<>` on the match arm, falling back to if statements
    if key == CODEGEN_CONSTANTS.literal {
        let (_, literal) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.value)
            .expect("Expected `name` to exist");
        write_constant_value(f, builder, literal)?;
    } else if key == CODEGEN_CONSTANTS.list_value {
        let (_, items) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.items)
            .expect("Expected `items` to exist");
        let array = builder.lookup(items.assert_key()).assert_array();

        f.push('[');
        let mut after_first = false;
        for key_or_null in array {
            match key_or_null {
                Primitive::Null => {}
                Primitive::Key(key) => {
                    if after_first {
                        f.push(',');
                    } else {
                        after_first = true;
                    }
                    let object = builder.lookup(*key).assert_object();
                    write_argument_value(f, builder, object)?;
                }
                _ => panic!("Expected an object key or null"),
            }
        }
        f.push(']');
    } else {
        // We filtered out Variables, here it should only be ObjectValue
        let (_, fields) = arg
            .iter()
            .find(|(key, _)| *key == CODEGEN_CONSTANTS.fields)
            .expect("Expected `fields` to exist");
        let fields = builder.lookup(fields.assert_key()).assert_array();

        f.push('{');
        for field in fields {
            let field = builder.lookup(field.assert_key()).assert_object();
            let (_, name) = field
                .iter()
                .find(|(key, _)| *key == CODEGEN_CONSTANTS.name)
                .expect("Expected `name` to exist");
            let name = name.assert_string();
            write!(f, "\\\"{}\\\":", name)?;
            write_argument_value(f, builder, field)?;
            f.push(',');
        }
        if !fields.is_empty() {
            f.pop();
        }
        f.push('}');
    }
    Ok(())
}

fn write_constant_value(f: &mut String, builder: &AstBuilder, value: &Primitive) -> FmtResult {
    match value {
        Primitive::Bool(b) => write!(f, "{}", if *b { "true" } else { "false" }),
        Primitive::String(key) => write!(f, "\\\"{}\\\"", key),
        Primitive::Float(value) => write!(f, "{}", value.as_float()),
        Primitive::Int(value) => write!(f, "{}", value),
        Primitive::Key(key) => {
            let ast = builder.lookup(*key);
            match ast {
                Ast::Array(arr) => {
                    f.push('[');
                    for value in arr {
                        write_constant_value(f, builder, value)?;
                        f.push(',');
                    }
                    if !arr.is_empty() {
                        f.pop();
                    }
                    f.push(']');
                    Ok(())
                }
                Ast::Object(obj) => {
                    f.push('{');
                    for (name, value) in obj {
                        write!(f, "\\\"{}\\\":", name)?;
                        write_constant_value(f, builder, value)?;
                        f.push(',');
                    }
                    if !obj.is_empty() {
                        f.pop();
                    }
                    f.push('}');
                    Ok(())
                }
            }
        }
        Primitive::Null => {
            f.push_str("null");
            Ok(())
        }
        Primitive::StorageKey(_, _) => panic!("Unexpected StorageKey"),
        Primitive::RawString(_) => panic!("Unexpected RawString"),
        Primitive::ModuleDependency(_) => panic!("Unexpected ModuleDependency"),
    }
}
<|endoftext|>