<commit_before>// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

use std::sync::{atomic::AtomicUsize, Arc};

use super::{ActivateResult, Queue};
use crate::virtio::{AsAny, VIRTIO_MMIO_INT_CONFIG, VIRTIO_MMIO_INT_VRING};
use logger::{error, warn};
use std::sync::atomic::Ordering;
use utils::eventfd::EventFd;
use vm_memory::GuestMemoryMmap;

/// Enum that indicates if a VirtioDevice is inactive or has been activated
/// and memory attached to it.
pub enum DeviceState {
    Inactive,
    Activated(GuestMemoryMmap),
}

impl DeviceState {
    /// Checks if the device is activated.
    pub fn is_activated(&self) -> bool {
        match self {
            DeviceState::Inactive => false,
            DeviceState::Activated(_) => true,
        }
    }

    /// Gets the memory attached to the device if it is activated.
    pub fn mem(&self) -> Option<&GuestMemoryMmap> {
        match self {
            DeviceState::Activated(ref mem) => Some(mem),
            DeviceState::Inactive => None,
        }
    }
}

pub enum IrqType {
    Config,
    Vring,
}

/// Helper struct that is responsible for triggering guest IRQs
pub struct IrqTrigger {
    pub(crate) irq_status: Arc<AtomicUsize>,
    pub(crate) irq_evt: EventFd,
}

impl IrqTrigger {
    pub fn new() -> std::io::Result<Self> {
        Ok(Self {
            irq_status: Arc::new(AtomicUsize::new(0)),
            irq_evt: EventFd::new(libc::EFD_NONBLOCK)?,
        })
    }

    pub fn trigger_irq(&self, irq_type: IrqType) -> std::result::Result<(), std::io::Error> {
        let irq = match irq_type {
            IrqType::Config => VIRTIO_MMIO_INT_CONFIG,
            IrqType::Vring => VIRTIO_MMIO_INT_VRING,
        };
        self.irq_status.fetch_or(irq as usize, Ordering::SeqCst);

        self.irq_evt.write(1).map_err(|e| {
            error!("Failed to send irq to the guest: {:?}", e);
            e
        })?;

        Ok(())
    }
}

/// Trait for virtio devices to be driven by a virtio transport.
///
/// The lifecycle of a virtio device is to be moved to a virtio transport, which will then query the
/// device. The virtio devices needs to create queues, events and event fds for interrupts and expose
/// them to the transport via get_queues/get_queue_events/get_interrupt/get_interrupt_status fns.
pub trait VirtioDevice: AsAny + Send {
    /// Get the available features offered by device.
    fn avail_features(&self) -> u64;

    /// Get acknowledged features of the driver.
    fn acked_features(&self) -> u64;

    /// Set acknowledged features of the driver.
    /// This function must maintain the following invariant:
    /// - self.avail_features() & self.acked_features() = self.get_acked_features()
    fn set_acked_features(&mut self, acked_features: u64);

    fn has_feature(&self, feature: u64) -> bool {
        (self.acked_features() & 1 << feature) != 0
    }

    /// The virtio device type.
    fn device_type(&self) -> u32;

    /// Returns the device queues.
    fn queues(&self) -> &[Queue];

    /// Returns a mutable reference to the device queues.
    fn queues_mut(&mut self) -> &mut [Queue];

    /// Returns the device queues event fds.
    fn queue_events(&self) -> &[EventFd];

    /// Returns the device interrupt eventfd.
    fn interrupt_evt(&self) -> &EventFd;

    /// Returns the current device interrupt status.
    fn interrupt_status(&self) -> Arc<AtomicUsize>;

    /// The set of feature bits shifted by `page * 32`.
    fn avail_features_by_page(&self, page: u32) -> u32 {
        let avail_features = self.avail_features();
        match page {
            // Get the lower 32-bits of the features bitfield.
            0 => avail_features as u32,
            // Get the upper 32-bits of the features bitfield.
            1 => (avail_features >> 32) as u32,
            _ => {
                warn!("Received request for unknown features page.");
                0u32
            }
        }
    }

    /// Acknowledges that this set of features should be enabled.
    fn ack_features_by_page(&mut self, page: u32, value: u32) {
        let mut v = match page {
            0 => u64::from(value),
            1 => u64::from(value) << 32,
            _ => {
                warn!("Cannot acknowledge unknown features page: {}", page);
                0u64
            }
        };

        // Check if the guest is ACK'ing a feature that we didn't claim to have.
        let avail_features = self.avail_features();
        let unrequested_features = v & !avail_features;
        if unrequested_features != 0 {
            warn!("Received acknowledge request for unknown feature: {:x}", v);
            // Don't count these features as acked.
            v &= !unrequested_features;
        }
        self.set_acked_features(self.acked_features() | v);
    }

    /// Reads this device configuration space at `offset`.
    fn read_config(&self, offset: u64, data: &mut [u8]);

    /// Writes to this device configuration space at `offset`.
    fn write_config(&mut self, offset: u64, data: &[u8]);

    /// Performs the formal activation for a device, which can be verified also with `is_activated`.
    fn activate(&mut self, mem: GuestMemoryMmap) -> ActivateResult;

    /// Checks if the resources of this device are activated.
    fn is_activated(&self) -> bool;

    /// Optionally deactivates this device and returns ownership of the guest memory map, interrupt
    /// event, and queue events.
    fn reset(&mut self) -> Option<(EventFd, Vec<EventFd>)> {
        None
    }
}

impl std::fmt::Debug for dyn VirtioDevice {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "VirtioDevice type {}", self.device_type())
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;

    impl IrqTrigger {
        pub fn has_pending_irq(&self, irq_type: IrqType) -> bool {
            if let Ok(1) = self.irq_evt.read() {
                let irq_status = self.irq_status.load(Ordering::SeqCst) as u32;
                return matches!(
                    (irq_status, irq_type),
                    (VIRTIO_MMIO_INT_CONFIG, IrqType::Config)
                        | (VIRTIO_MMIO_INT_VRING, IrqType::Vring)
                );
            }

            false
        }
    }

    #[test]
    fn irq_trigger() {
        let irq_trigger = IrqTrigger::new().unwrap();
        assert_eq!(irq_trigger.irq_status.load(Ordering::SeqCst), 0);

        // Check that there are no pending irqs.
        assert!(!irq_trigger.has_pending_irq(IrqType::Config));
        assert!(!irq_trigger.has_pending_irq(IrqType::Vring));

        // Check that trigger_irq() correctly generates irqs.
        irq_trigger.trigger_irq(IrqType::Config).unwrap();
        assert!(irq_trigger.has_pending_irq(IrqType::Config));
        irq_trigger.irq_status.store(0, Ordering::SeqCst);
        irq_trigger.trigger_irq(IrqType::Vring).unwrap();
        assert!(irq_trigger.has_pending_irq(IrqType::Vring));

        // Check trigger_irq() failure case (irq_evt is full).
        irq_trigger.irq_evt.write(u64::MAX - 1).unwrap();
        assert!(irq_trigger.trigger_irq(IrqType::Config).is_err());
        assert!(irq_trigger.trigger_irq(IrqType::Vring).is_err());
    }

    struct MockVirtioDevice {
        acked_features: u64,
    }

    impl VirtioDevice for MockVirtioDevice {
        fn avail_features(&self) -> u64 {
            todo!()
        }

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

        fn set_acked_features(&mut self, _acked_features: u64) {
            todo!()
        }

        fn device_type(&self) -> u32 {
            todo!()
        }

        fn queues(&self) -> &[Queue] {
            todo!()
        }

        fn queues_mut(&mut self) -> &mut [Queue] {
            todo!()
        }

        fn queue_events(&self) -> &[EventFd] {
            todo!()
        }

        fn interrupt_evt(&self) -> &EventFd {
            todo!()
        }

        fn interrupt_status(&self) -> Arc<AtomicUsize> {
            todo!()
        }

        fn read_config(&self, _offset: u64, _data: &mut [u8]) {
            todo!()
        }

        fn write_config(&mut self, _offset: u64, _data: &[u8]) {
            todo!()
        }

        fn activate(&mut self, _mem: GuestMemoryMmap) -> ActivateResult {
            todo!()
        }

        fn is_activated(&self) -> bool {
            todo!()
        }
    }

    #[test]
    fn test_has_feature() {
        let mut device = MockVirtioDevice { acked_features: 0 };

        let mock_feature_1 = 1u64;
        assert!(!device.has_feature(mock_feature_1));
        device.acked_features = 1 << mock_feature_1;
        assert!(device.has_feature(mock_feature_1));

        let mock_feature_2 = 2u64;
        assert!(!device.has_feature(mock_feature_2));
        device.acked_features = (1 << mock_feature_1) | (1 << mock_feature_2);
        assert!(device.has_feature(mock_feature_1));
        assert!(device.has_feature(mock_feature_2));
    }
}
<commit_msg>[virtio][unit tests] Fix has_pending_irq()<commit_after>// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

use std::sync::{atomic::AtomicUsize, Arc};

use super::{ActivateResult, Queue};
use crate::virtio::{AsAny, VIRTIO_MMIO_INT_CONFIG, VIRTIO_MMIO_INT_VRING};
use logger::{error, warn};
use std::sync::atomic::Ordering;
use utils::eventfd::EventFd;
use vm_memory::GuestMemoryMmap;

/// Enum that indicates if a VirtioDevice is inactive or has been activated
/// and memory attached to it.
pub enum DeviceState {
    Inactive,
    Activated(GuestMemoryMmap),
}

impl DeviceState {
    /// Checks if the device is activated.
    pub fn is_activated(&self) -> bool {
        match self {
            DeviceState::Inactive => false,
            DeviceState::Activated(_) => true,
        }
    }

    /// Gets the memory attached to the device if it is activated.
    pub fn mem(&self) -> Option<&GuestMemoryMmap> {
        match self {
            DeviceState::Activated(ref mem) => Some(mem),
            DeviceState::Inactive => None,
        }
    }
}

pub enum IrqType {
    Config,
    Vring,
}

/// Helper struct that is responsible for triggering guest IRQs
pub struct IrqTrigger {
    pub(crate) irq_status: Arc<AtomicUsize>,
    pub(crate) irq_evt: EventFd,
}

impl IrqTrigger {
    pub fn new() -> std::io::Result<Self> {
        Ok(Self {
            irq_status: Arc::new(AtomicUsize::new(0)),
            irq_evt: EventFd::new(libc::EFD_NONBLOCK)?,
        })
    }

    pub fn trigger_irq(&self, irq_type: IrqType) -> std::result::Result<(), std::io::Error> {
        let irq = match irq_type {
            IrqType::Config => VIRTIO_MMIO_INT_CONFIG,
            IrqType::Vring => VIRTIO_MMIO_INT_VRING,
        };
        self.irq_status.fetch_or(irq as usize, Ordering::SeqCst);

        self.irq_evt.write(1).map_err(|e| {
            error!("Failed to send irq to the guest: {:?}", e);
            e
        })?;

        Ok(())
    }
}

/// Trait for virtio devices to be driven by a virtio transport.
///
/// The lifecycle of a virtio device is to be moved to a virtio transport, which will then query the
/// device. The virtio devices needs to create queues, events and event fds for interrupts and expose
/// them to the transport via get_queues/get_queue_events/get_interrupt/get_interrupt_status fns.
pub trait VirtioDevice: AsAny + Send {
    /// Get the available features offered by device.
    fn avail_features(&self) -> u64;

    /// Get acknowledged features of the driver.
    fn acked_features(&self) -> u64;

    /// Set acknowledged features of the driver.
    /// This function must maintain the following invariant:
    /// - self.avail_features() & self.acked_features() = self.get_acked_features()
    fn set_acked_features(&mut self, acked_features: u64);

    fn has_feature(&self, feature: u64) -> bool {
        (self.acked_features() & 1 << feature) != 0
    }

    /// The virtio device type.
    fn device_type(&self) -> u32;

    /// Returns the device queues.
    fn queues(&self) -> &[Queue];

    /// Returns a mutable reference to the device queues.
    fn queues_mut(&mut self) -> &mut [Queue];

    /// Returns the device queues event fds.
    fn queue_events(&self) -> &[EventFd];

    /// Returns the device interrupt eventfd.
    fn interrupt_evt(&self) -> &EventFd;

    /// Returns the current device interrupt status.
    fn interrupt_status(&self) -> Arc<AtomicUsize>;

    /// The set of feature bits shifted by `page * 32`.
    fn avail_features_by_page(&self, page: u32) -> u32 {
        let avail_features = self.avail_features();
        match page {
            // Get the lower 32-bits of the features bitfield.
            0 => avail_features as u32,
            // Get the upper 32-bits of the features bitfield.
            1 => (avail_features >> 32) as u32,
            _ => {
                warn!("Received request for unknown features page.");
                0u32
            }
        }
    }

    /// Acknowledges that this set of features should be enabled.
    fn ack_features_by_page(&mut self, page: u32, value: u32) {
        let mut v = match page {
            0 => u64::from(value),
            1 => u64::from(value) << 32,
            _ => {
                warn!("Cannot acknowledge unknown features page: {}", page);
                0u64
            }
        };

        // Check if the guest is ACK'ing a feature that we didn't claim to have.
        let avail_features = self.avail_features();
        let unrequested_features = v & !avail_features;
        if unrequested_features != 0 {
            warn!("Received acknowledge request for unknown feature: {:x}", v);
            // Don't count these features as acked.
            v &= !unrequested_features;
        }
        self.set_acked_features(self.acked_features() | v);
    }

    /// Reads this device configuration space at `offset`.
    fn read_config(&self, offset: u64, data: &mut [u8]);

    /// Writes to this device configuration space at `offset`.
    fn write_config(&mut self, offset: u64, data: &[u8]);

    /// Performs the formal activation for a device, which can be verified also with `is_activated`.
    fn activate(&mut self, mem: GuestMemoryMmap) -> ActivateResult;

    /// Checks if the resources of this device are activated.
    fn is_activated(&self) -> bool;

    /// Optionally deactivates this device and returns ownership of the guest memory map, interrupt
    /// event, and queue events.
    fn reset(&mut self) -> Option<(EventFd, Vec<EventFd>)> {
        None
    }
}

impl std::fmt::Debug for dyn VirtioDevice {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "VirtioDevice type {}", self.device_type())
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;

    impl IrqTrigger {
        pub fn has_pending_irq(&self, irq_type: IrqType) -> bool {
            if let Ok(num_irqs) = self.irq_evt.read() {
                if num_irqs == 0 {
                    return false;
                }

                let irq_status = self.irq_status.load(Ordering::SeqCst) as u32;
                return matches!(
                    (irq_status, irq_type),
                    (VIRTIO_MMIO_INT_CONFIG, IrqType::Config)
                        | (VIRTIO_MMIO_INT_VRING, IrqType::Vring)
                );
            }

            false
        }
    }

    #[test]
    fn irq_trigger() {
        let irq_trigger = IrqTrigger::new().unwrap();
        assert_eq!(irq_trigger.irq_status.load(Ordering::SeqCst), 0);

        // Check that there are no pending irqs.
        assert!(!irq_trigger.has_pending_irq(IrqType::Config));
        assert!(!irq_trigger.has_pending_irq(IrqType::Vring));

        // Check that trigger_irq() correctly generates irqs.
        irq_trigger.trigger_irq(IrqType::Config).unwrap();
        assert!(irq_trigger.has_pending_irq(IrqType::Config));
        irq_trigger.irq_status.store(0, Ordering::SeqCst);
        irq_trigger.trigger_irq(IrqType::Vring).unwrap();
        assert!(irq_trigger.has_pending_irq(IrqType::Vring));

        // Check trigger_irq() failure case (irq_evt is full).
        irq_trigger.irq_evt.write(u64::MAX - 1).unwrap();
        assert!(irq_trigger.trigger_irq(IrqType::Config).is_err());
        assert!(irq_trigger.trigger_irq(IrqType::Vring).is_err());
    }

    struct MockVirtioDevice {
        acked_features: u64,
    }

    impl VirtioDevice for MockVirtioDevice {
        fn avail_features(&self) -> u64 {
            todo!()
        }

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

        fn set_acked_features(&mut self, _acked_features: u64) {
            todo!()
        }

        fn device_type(&self) -> u32 {
            todo!()
        }

        fn queues(&self) -> &[Queue] {
            todo!()
        }

        fn queues_mut(&mut self) -> &mut [Queue] {
            todo!()
        }

        fn queue_events(&self) -> &[EventFd] {
            todo!()
        }

        fn interrupt_evt(&self) -> &EventFd {
            todo!()
        }

        fn interrupt_status(&self) -> Arc<AtomicUsize> {
            todo!()
        }

        fn read_config(&self, _offset: u64, _data: &mut [u8]) {
            todo!()
        }

        fn write_config(&mut self, _offset: u64, _data: &[u8]) {
            todo!()
        }

        fn activate(&mut self, _mem: GuestMemoryMmap) -> ActivateResult {
            todo!()
        }

        fn is_activated(&self) -> bool {
            todo!()
        }
    }

    #[test]
    fn test_has_feature() {
        let mut device = MockVirtioDevice { acked_features: 0 };

        let mock_feature_1 = 1u64;
        assert!(!device.has_feature(mock_feature_1));
        device.acked_features = 1 << mock_feature_1;
        assert!(device.has_feature(mock_feature_1));

        let mock_feature_2 = 2u64;
        assert!(!device.has_feature(mock_feature_2));
        device.acked_features = (1 << mock_feature_1) | (1 << mock_feature_2);
        assert!(device.has_feature(mock_feature_1));
        assert!(device.has_feature(mock_feature_2));
    }
}
<|endoftext|>


<commit_before>mod reader;
mod writer;

pub use self::reader::MultiValuedFastFieldReader;
pub use self::writer::MultiValuedFastFieldWriter;

#[cfg(test)]
mod tests {

    use crate::collector::TopDocs;
    use crate::query::QueryParser;
    use crate::schema::Cardinality;
    use crate::schema::Facet;
    use crate::schema::IntOptions;
    use crate::schema::Schema;
    use crate::schema::INDEXED;
    use crate::Index;
    use chrono::Duration;

    #[test]
    fn test_multivalued_u64() {
        let mut schema_builder = Schema::builder();
        let field = schema_builder.add_u64_field(
            "multifield",
            IntOptions::default().set_fast(Cardinality::MultiValues),
        );
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        let mut index_writer = index.writer_for_tests().unwrap();
        index_writer.add_document(doc!(field=>1u64, field=>3u64));
        index_writer.add_document(doc!());
        index_writer.add_document(doc!(field=>4u64));
        index_writer.add_document(doc!(field=>5u64, field=>20u64,field=>1u64));
        assert!(index_writer.commit().is_ok());

        let searcher = index.reader().unwrap().searcher();
        let segment_reader = searcher.segment_reader(0);
        let mut vals = Vec::new();
        let multi_value_reader = segment_reader.fast_fields().u64s(field).unwrap();
        {
            multi_value_reader.get_vals(2, &mut vals);
            assert_eq!(&vals, &[4u64]);
        }
        {
            multi_value_reader.get_vals(0, &mut vals);
            assert_eq!(&vals, &[1u64, 3u64]);
        }
        {
            multi_value_reader.get_vals(1, &mut vals);
            assert!(vals.is_empty());
        }
    }

    #[test]
    fn test_multivalued_date() {
        let mut schema_builder = Schema::builder();
        let date_field = schema_builder.add_date_field(
            "multi_date_field",
            IntOptions::default()
                .set_fast(Cardinality::MultiValues)
                .set_indexed()
                .set_stored(),
        );
        let time_i =
            schema_builder.add_i64_field("time_stamp_i", IntOptions::default().set_stored());
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        let mut index_writer = index.writer_for_tests().unwrap();
        let first_time_stamp = chrono::Utc::now();
        index_writer.add_document(
            doc!(date_field=>first_time_stamp, date_field=>first_time_stamp, time_i=>1i64),
        );
        index_writer.add_document(doc!(time_i=>0i64));
        // add one second
        index_writer
            .add_document(doc!(date_field=>first_time_stamp + Duration::seconds(1), time_i=>2i64));
        // add another second
        let two_secs_ahead = first_time_stamp + Duration::seconds(2);
        index_writer.add_document(doc!(date_field=>two_secs_ahead, date_field=>two_secs_ahead,date_field=>two_secs_ahead, time_i=>3i64));
        assert!(index_writer.commit().is_ok());

        let reader = index.reader().unwrap();
        let searcher = reader.searcher();
        let reader = searcher.segment_reader(0);
        assert_eq!(reader.num_docs(), 4);

        {
            let parser = QueryParser::for_index(&index, vec![date_field]);
            let query = parser
                .parse_query(&format!("\"{}\"", first_time_stamp.to_rfc3339()).to_string())
                .expect("could not parse query");
            let results = searcher
                .search(&query, &TopDocs::with_limit(5))
                .expect("could not query index");

            assert_eq!(results.len(), 1);
            for (_score, doc_address) in results {
                let retrieved_doc = searcher.doc(doc_address).expect("cannot fetch doc");
                assert_eq!(
                    retrieved_doc
                        .get_first(date_field)
                        .expect("cannot find value")
                        .date_value()
                        .unwrap()
                        .timestamp(),
                    first_time_stamp.timestamp()
                );
                assert_eq!(
                    retrieved_doc
                        .get_first(time_i)
                        .expect("cannot find value")
                        .i64_value(),
                    Some(1i64)
                );
            }
        }

        {
            let parser = QueryParser::for_index(&index, vec![date_field]);
            let query = parser
                .parse_query(&format!("\"{}\"", two_secs_ahead.to_rfc3339()).to_string())
                .expect("could not parse query");
            let results = searcher
                .search(&query, &TopDocs::with_limit(5))
                .expect("could not query index");

            assert_eq!(results.len(), 1);

            for (_score, doc_address) in results {
                let retrieved_doc = searcher.doc(doc_address).expect("cannot fetch doc");
                assert_eq!(
                    retrieved_doc
                        .get_first(date_field)
                        .expect("cannot find value")
                        .date_value()
                        .unwrap()
                        .timestamp(),
                    two_secs_ahead.timestamp()
                );
                assert_eq!(
                    retrieved_doc
                        .get_first(time_i)
                        .expect("cannot find value")
                        .i64_value(),
                    Some(3i64)
                );
            }
        }

        {
            let parser = QueryParser::for_index(&index, vec![date_field]);
            let range_q = format!(
                "[{} TO {}]",
                (first_time_stamp + Duration::seconds(1)).to_rfc3339(),
                (first_time_stamp + Duration::seconds(3)).to_rfc3339()
            );
            let query = parser.parse_query(&range_q).expect("could not parse query");
            let results = searcher
                .search(&query, &TopDocs::with_limit(5))
                .expect("could not query index");

            assert_eq!(results.len(), 2);
            for (i, doc_pair) in results.iter().enumerate() {
                let retrieved_doc = searcher.doc(doc_pair.1).expect("cannot fetch doc");
                let offset_sec = match i {
                    0 => 1,
                    1 => 2,
                    _ => panic!("should not have more than 2 docs"),
                };
                let time_i_val = match i {
                    0 => 2,
                    1 => 3,
                    _ => panic!("should not have more than 2 docs"),
                };
                assert_eq!(
                    retrieved_doc
                        .get_first(date_field)
                        .expect("cannot find value")
                        .date_value()
                        .expect("value not of Date type")
                        .timestamp(),
                    (first_time_stamp + Duration::seconds(offset_sec)).timestamp()
                );
                assert_eq!(
                    retrieved_doc
                        .get_first(time_i)
                        .expect("cannot find value")
                        .i64_value()
                        .expect("value not of i64 type"),
                    time_i_val
                );
            }
        }
    }

    #[test]
    fn test_multivalued_i64() {
        let mut schema_builder = Schema::builder();
        let field = schema_builder.add_i64_field(
            "multifield",
            IntOptions::default().set_fast(Cardinality::MultiValues),
        );
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        let mut index_writer = index.writer_for_tests().unwrap();
        index_writer.add_document(doc!(field=> 1i64, field => 3i64));
        index_writer.add_document(doc!());
        index_writer.add_document(doc!(field=> -4i64));
        index_writer.add_document(doc!(field=> -5i64, field => -20i64, field=>1i64));
        assert!(index_writer.commit().is_ok());

        let searcher = index.reader().unwrap().searcher();
        let segment_reader = searcher.segment_reader(0);
        let mut vals = Vec::new();
        let multi_value_reader = segment_reader.fast_fields().i64s(field).unwrap();
        multi_value_reader.get_vals(2, &mut vals);
        assert_eq!(&vals, &[-4i64]);
        multi_value_reader.get_vals(0, &mut vals);
        assert_eq!(&vals, &[1i64, 3i64]);
        multi_value_reader.get_vals(1, &mut vals);
        assert!(vals.is_empty());
        multi_value_reader.get_vals(3, &mut vals);
        assert_eq!(&vals, &[-5i64, -20i64, 1i64]);
    }
    #[test]
    #[ignore]
    fn test_many_facets() {
        let mut schema_builder = Schema::builder();
        let field = schema_builder.add_facet_field("facetfield", INDEXED);
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        let mut index_writer = index.writer_for_tests().unwrap();
        for i in 0..100_000 {
            index_writer.add_document(doc!(field=> Facet::from(format!("/lang/{}", i).as_str())));
        }
        assert!(index_writer.commit().is_ok());
    }
}
<commit_msg>Change to an exlusive bound in the date range test<commit_after>mod reader;
mod writer;

pub use self::reader::MultiValuedFastFieldReader;
pub use self::writer::MultiValuedFastFieldWriter;

#[cfg(test)]
mod tests {

    use crate::collector::TopDocs;
    use crate::query::QueryParser;
    use crate::schema::Cardinality;
    use crate::schema::Facet;
    use crate::schema::IntOptions;
    use crate::schema::Schema;
    use crate::schema::INDEXED;
    use crate::Index;
    use chrono::Duration;

    #[test]
    fn test_multivalued_u64() {
        let mut schema_builder = Schema::builder();
        let field = schema_builder.add_u64_field(
            "multifield",
            IntOptions::default().set_fast(Cardinality::MultiValues),
        );
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        let mut index_writer = index.writer_for_tests().unwrap();
        index_writer.add_document(doc!(field=>1u64, field=>3u64));
        index_writer.add_document(doc!());
        index_writer.add_document(doc!(field=>4u64));
        index_writer.add_document(doc!(field=>5u64, field=>20u64,field=>1u64));
        assert!(index_writer.commit().is_ok());

        let searcher = index.reader().unwrap().searcher();
        let segment_reader = searcher.segment_reader(0);
        let mut vals = Vec::new();
        let multi_value_reader = segment_reader.fast_fields().u64s(field).unwrap();
        {
            multi_value_reader.get_vals(2, &mut vals);
            assert_eq!(&vals, &[4u64]);
        }
        {
            multi_value_reader.get_vals(0, &mut vals);
            assert_eq!(&vals, &[1u64, 3u64]);
        }
        {
            multi_value_reader.get_vals(1, &mut vals);
            assert!(vals.is_empty());
        }
    }

    #[test]
    fn test_multivalued_date() {
        let mut schema_builder = Schema::builder();
        let date_field = schema_builder.add_date_field(
            "multi_date_field",
            IntOptions::default()
                .set_fast(Cardinality::MultiValues)
                .set_indexed()
                .set_stored(),
        );
        let time_i =
            schema_builder.add_i64_field("time_stamp_i", IntOptions::default().set_stored());
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        let mut index_writer = index.writer_for_tests().unwrap();
        let first_time_stamp = chrono::Utc::now();
        index_writer.add_document(
            doc!(date_field=>first_time_stamp, date_field=>first_time_stamp, time_i=>1i64),
        );
        index_writer.add_document(doc!(time_i=>0i64));
        // add one second
        index_writer
            .add_document(doc!(date_field=>first_time_stamp + Duration::seconds(1), time_i=>2i64));
        // add another second
        let two_secs_ahead = first_time_stamp + Duration::seconds(2);
        index_writer.add_document(doc!(date_field=>two_secs_ahead, date_field=>two_secs_ahead,date_field=>two_secs_ahead, time_i=>3i64));
        // add three seconds
        index_writer.add_document(doc!(date_field=>first_time_stamp + Duration::seconds(3), time_i=>4i64));
        assert!(index_writer.commit().is_ok());

        let reader = index.reader().unwrap();
        let searcher = reader.searcher();
        let reader = searcher.segment_reader(0);
        assert_eq!(reader.num_docs(), 5);

        {
            let parser = QueryParser::for_index(&index, vec![date_field]);
            let query = parser
                .parse_query(&format!("\"{}\"", first_time_stamp.to_rfc3339()).to_string())
                .expect("could not parse query");
            let results = searcher
                .search(&query, &TopDocs::with_limit(5))
                .expect("could not query index");

            assert_eq!(results.len(), 1);
            for (_score, doc_address) in results {
                let retrieved_doc = searcher.doc(doc_address).expect("cannot fetch doc");
                assert_eq!(
                    retrieved_doc
                        .get_first(date_field)
                        .expect("cannot find value")
                        .date_value()
                        .unwrap()
                        .timestamp(),
                    first_time_stamp.timestamp()
                );
                assert_eq!(
                    retrieved_doc
                        .get_first(time_i)
                        .expect("cannot find value")
                        .i64_value(),
                    Some(1i64)
                );
            }
        }

        {
            let parser = QueryParser::for_index(&index, vec![date_field]);
            let query = parser
                .parse_query(&format!("\"{}\"", two_secs_ahead.to_rfc3339()).to_string())
                .expect("could not parse query");
            let results = searcher
                .search(&query, &TopDocs::with_limit(5))
                .expect("could not query index");

            assert_eq!(results.len(), 1);

            for (_score, doc_address) in results {
                let retrieved_doc = searcher.doc(doc_address).expect("cannot fetch doc");
                assert_eq!(
                    retrieved_doc
                        .get_first(date_field)
                        .expect("cannot find value")
                        .date_value()
                        .unwrap()
                        .timestamp(),
                    two_secs_ahead.timestamp()
                );
                assert_eq!(
                    retrieved_doc
                        .get_first(time_i)
                        .expect("cannot find value")
                        .i64_value(),
                    Some(3i64)
                );
            }
        }

        {
            let parser = QueryParser::for_index(&index, vec![date_field]);
            let range_q = format!(
                "[{} TO {}}}",
                (first_time_stamp + Duration::seconds(1)).to_rfc3339(),
                (first_time_stamp + Duration::seconds(3)).to_rfc3339()
            );
            let query = parser.parse_query(&range_q).expect("could not parse query");
            let results = searcher
                .search(&query, &TopDocs::with_limit(5))
                .expect("could not query index");

            assert_eq!(results.len(), 2);
            for (i, doc_pair) in results.iter().enumerate() {
                let retrieved_doc = searcher.doc(doc_pair.1).expect("cannot fetch doc");
                let offset_sec = match i {
                    0 => 1,
                    1 => 2,
                    _ => panic!("should not have more than 2 docs"),
                };
                let time_i_val = match i {
                    0 => 2,
                    1 => 3,
                    _ => panic!("should not have more than 2 docs"),
                };
                assert_eq!(
                    retrieved_doc
                        .get_first(date_field)
                        .expect("cannot find value")
                        .date_value()
                        .expect("value not of Date type")
                        .timestamp(),
                    (first_time_stamp + Duration::seconds(offset_sec)).timestamp()
                );
                assert_eq!(
                    retrieved_doc
                        .get_first(time_i)
                        .expect("cannot find value")
                        .i64_value()
                        .expect("value not of i64 type"),
                    time_i_val
                );
            }
        }
    }

    #[test]
    fn test_multivalued_i64() {
        let mut schema_builder = Schema::builder();
        let field = schema_builder.add_i64_field(
            "multifield",
            IntOptions::default().set_fast(Cardinality::MultiValues),
        );
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        let mut index_writer = index.writer_for_tests().unwrap();
        index_writer.add_document(doc!(field=> 1i64, field => 3i64));
        index_writer.add_document(doc!());
        index_writer.add_document(doc!(field=> -4i64));
        index_writer.add_document(doc!(field=> -5i64, field => -20i64, field=>1i64));
        assert!(index_writer.commit().is_ok());

        let searcher = index.reader().unwrap().searcher();
        let segment_reader = searcher.segment_reader(0);
        let mut vals = Vec::new();
        let multi_value_reader = segment_reader.fast_fields().i64s(field).unwrap();
        multi_value_reader.get_vals(2, &mut vals);
        assert_eq!(&vals, &[-4i64]);
        multi_value_reader.get_vals(0, &mut vals);
        assert_eq!(&vals, &[1i64, 3i64]);
        multi_value_reader.get_vals(1, &mut vals);
        assert!(vals.is_empty());
        multi_value_reader.get_vals(3, &mut vals);
        assert_eq!(&vals, &[-5i64, -20i64, 1i64]);
    }
    #[test]
    #[ignore]
    fn test_many_facets() {
        let mut schema_builder = Schema::builder();
        let field = schema_builder.add_facet_field("facetfield", INDEXED);
        let schema = schema_builder.build();
        let index = Index::create_in_ram(schema);
        let mut index_writer = index.writer_for_tests().unwrap();
        for i in 0..100_000 {
            index_writer.add_document(doc!(field=> Facet::from(format!("/lang/{}", i).as_str())));
        }
        assert!(index_writer.commit().is_ok());
    }
}
<|endoftext|>


<commit_before>use std::collections::{HashMap};
use grit;

pub fn optimize(prog: grit::ProgDef) -> grit::ProgDef {
  let mut prog = prog;
  let prog_info = collect_prog_info(&prog);
  prog.fun_defs = prog.fun_defs.map_in_place(|mut fun_def| {
    let fun_info = prog_info.get(&fun_def.name).unwrap();
    fun_def.blocks = fun_def.blocks.map_in_place(|mut block| {
      block.ops = block.ops.map_in_place(|op|
        update_op(&prog_info, fun_info, op));
      block.jump = update_jump(&prog_info, fun_info, block.jump);
      block
    });
    fun_def
  });
  prog
}

type ProgInfo = HashMap<grit::FunName, FunInfo>;

#[derive(Debug, PartialEq, Clone)]
struct FunInfo {
  arg_count: usize,
  capture_infos: Vec<Info>,
  var_infos: Vec<Info>,
  return_info: Info,
}

#[derive(Debug, PartialEq, Clone)]
enum Info {
  Any,
  Fun(grit::FunName),
  False,
  True,
  Int(i32),
  Obj(grit::ObjName),
  AnyTrue,
  No,
}

#[derive(Debug, PartialEq, Copy, Clone)]
enum BoolInfo {
  Any,
  True,
  False,
  No,
}

fn update_op(prog_info: &ProgInfo, fun_info: &FunInfo, op: grit::Op) -> grit::Op {
  match op {
    grit::Op::Call(target_var, callee, args) =>
      grit::Op::Call(target_var,
        update_callee(prog_info, fun_info, callee, args.len()),
        args.map_in_place(|a| update_val(prog_info, fun_info, a))),
    grit::Op::ExternCall(target_var, ext_name, args) =>
      grit::Op::ExternCall(target_var, ext_name,
        args.map_in_place(|a| update_val(prog_info, fun_info, a))),
    grit::Op::AllocClos(closs) =>
      grit::Op::AllocClos(closs.map_in_place(|(var, fun_name, captures)| {
        (var, fun_name, captures.map_in_place(|c| update_val(prog_info, fun_info, c)))
      })),
    grit::Op::Assign(var_vals) =>
      grit::Op::Assign(var_vals.map_in_place(|(var, val)| {
        (var, update_val(prog_info, fun_info, val))
      })),
  }
}

fn update_jump(prog_info: &ProgInfo, fun_info: &FunInfo, jump: grit::Jump) -> grit::Jump {
  match jump {
    grit::Jump::Goto(label) =>
      grit::Jump::Goto(label),
    grit::Jump::Return(val) =>
      grit::Jump::Return(update_val(prog_info, fun_info, val)),
    grit::Jump::TailCall(callee, args) =>
      grit::Jump::TailCall(update_callee(prog_info, fun_info, callee, args.len()),
        args.map_in_place(|arg| update_val(prog_info, fun_info, arg))),
    grit::Jump::Branch(boolval, then_label, else_label) =>
      match boolval_info(fun_info, &boolval) {
        BoolInfo::True => grit::Jump::Goto(then_label),
        BoolInfo::False => grit::Jump::Goto(else_label),
        BoolInfo::Any =>
          grit::Jump::Branch(grit::Boolval::IsTrue(grit::Val::Undefined),
            then_label, else_label),
        BoolInfo::No =>
          grit::Jump::Branch(update_boolval(prog_info, fun_info, boolval),
            then_label, else_label),
      },
  }
}

fn update_callee(prog_info: &ProgInfo, fun_info: &FunInfo,
  callee: grit::Callee, arg_count: usize) -> grit::Callee
{
  match callee {
    grit::Callee::Unknown(val) =>
      match val_info(fun_info, &val) {
        Info::Fun(callee_name) => {
          let callee_info = prog_info.get(&callee_name).unwrap();
          if callee_info.arg_count == arg_count {
            if callee_info.capture_infos.len() == 0 {
              grit::Callee::Combinator(callee_name)
            } else {
              grit::Callee::KnownClosure(callee_name, val)
            }
          } else {
            grit::Callee::Unknown(val)
          }
        },
        _ =>
          grit::Callee::Unknown(val)
      },
    callee => callee,
  }
}

fn update_val(prog_info: &ProgInfo, fun_info: &FunInfo, val: grit::Val) -> grit::Val {
  match val_info(fun_info, &val) {
    Info::Any => grit::Val::Int(666),
    Info::Fun(fun_name) =>
      if prog_info.get(&fun_name).unwrap().capture_infos.len() == 0 {
        grit::Val::Combinator(fun_name)
      } else {
        val
      },
    Info::False => grit::Val::False,
    Info::True => grit::Val::True,
    Info::Int(num) => grit::Val::Int(num),
    Info::Obj(name) => grit::Val::Obj(name),
    Info::AnyTrue | Info::No => val,
  }
}

fn update_boolval(prog_info: &ProgInfo, fun_info: &FunInfo,
  boolval: grit::Boolval) -> grit::Boolval
{
  match boolval {
    grit::Boolval::IsTrue(val) =>
      grit::Boolval::IsTrue(update_val(prog_info, fun_info, val)),
    grit::Boolval::IsFalse(val) =>
      grit::Boolval::IsFalse(update_val(prog_info, fun_info, val)),
  }
}

fn collect_prog_info(prog: &grit::ProgDef) -> ProgInfo {
  let mut prog_info: ProgInfo = prog.fun_defs.iter().map(|fun_def| {
    (fun_def.name.clone(), FunInfo {
      arg_count: fun_def.arg_count,
      capture_infos: (0..fun_def.capture_count).map(|_| Info::Any).collect(),
      var_infos: (0..fun_def.var_count).map(|_| Info::Any).collect(),
      return_info: Info::Any,
    })
  }).collect();

  loop {
    let prev_info = prog_info.clone();

    for fun_def in prog.fun_defs.iter() {
      for block in fun_def.blocks.iter() {
        for op in block.ops.iter() {
          op_info_update(&mut prog_info, &fun_def.name, op);
        }
        jump_info_update(&mut prog_info, &fun_def.name, &block.jump);
      }
    }

    if prev_info == prog_info {
      break
    }
  }

  prog_info
}

fn op_info_update(prog_info: &mut ProgInfo, fun_name: &grit::FunName, op: &grit::Op) {
  match *op {
    grit::Op::Call(ref result_var, ref callee, _) =>
      prog_info.get_mut(fun_name).unwrap().var_infos[result_var.0] = join_info(
        &prog_info.get(fun_name).unwrap().var_infos[result_var.0],
        &return_info(prog_info, callee)),
    grit::Op::ExternCall(ref result_var, _, _) =>
      prog_info.get_mut(fun_name).unwrap().var_infos[result_var.0] = Info::No,
    grit::Op::AllocClos(ref closs) =>
      for &(ref clos_var, ref clos_name, ref captures) in closs.iter() {
        prog_info.get_mut(fun_name).unwrap().var_infos[clos_var.0] = join_info(
          &prog_info.get(fun_name).unwrap().var_infos[clos_var.0],
          &Info::Fun(clos_name.clone()));
        for (idx, capture_val) in captures.iter().enumerate() {
          prog_info.get_mut(clos_name).unwrap().capture_infos[idx] = join_info(
            &prog_info.get(clos_name).unwrap().capture_infos[idx],
            &val_info(prog_info.get(fun_name).unwrap(), capture_val));
        }
      },
    grit::Op::Assign(ref var_vals) =>
      for &(ref var, ref val) in var_vals.iter() {
        prog_info.get_mut(fun_name).unwrap().var_infos[var.0] = join_info(
          &prog_info.get(fun_name).unwrap().var_infos[var.0],
          &val_info(prog_info.get(fun_name).unwrap(), val));
      },
  }
}

fn jump_info_update(prog_info: &mut ProgInfo,
  fun_name: &grit::FunName, jump: &grit::Jump)
{
  match *jump {
    grit::Jump::Goto(_) => { },
    grit::Jump::Return(ref val) =>
      prog_info.get_mut(fun_name).unwrap().return_info = join_info(
        &prog_info.get(fun_name).unwrap().return_info,
        &val_info(prog_info.get(fun_name).unwrap(), val)),
    grit::Jump::TailCall(ref callee, _) =>
      prog_info.get_mut(fun_name).unwrap().return_info = join_info(
        &prog_info.get(fun_name).unwrap().return_info,
        &return_info(prog_info, callee)),
    grit::Jump::Branch(_, _, _) => { },
  }
}

fn return_info(prog_info: &ProgInfo, callee: &grit::Callee) -> Info {
  match *callee {
    grit::Callee::Combinator(ref callee_name) |
    grit::Callee::KnownClosure(ref callee_name, _) =>
      prog_info.get(callee_name).unwrap().return_info.clone(),
    grit::Callee::Unknown(_) =>
      Info::No,
  }
}

fn val_info(fun_info: &FunInfo, val: &grit::Val) -> Info {
  match *val {
    grit::Val::Var(ref var) => fun_info.var_infos[var.0].clone(),
    grit::Val::Arg(_) => Info::No,
    grit::Val::Capture(idx) => fun_info.capture_infos[idx].clone(),
    grit::Val::Combinator(ref fun_name) => Info::Fun(fun_name.clone()),
    grit::Val::Obj(ref obj_name) => Info::Obj(obj_name.clone()),
    grit::Val::Int(num) => Info::Int(num),
    grit::Val::True => Info::True,
    grit::Val::False => Info::False,
    grit::Val::Undefined => Info::Any,
  }
}

fn boolval_info(fun_info: &FunInfo, boolval: &grit::Boolval) -> BoolInfo {
  match *boolval {
    grit::Boolval::IsTrue(ref val) =>
      info_to_bool(&val_info(fun_info, val)),
    grit::Boolval::IsFalse(ref val) =>
      negate_boolinfo(info_to_bool(&val_info(fun_info, val))),
  }
}

fn negate_boolinfo(boolinfo: BoolInfo) -> BoolInfo {
  match boolinfo {
    BoolInfo::Any => BoolInfo::Any,
    BoolInfo::True => BoolInfo::False,
    BoolInfo::False => BoolInfo::True,
    BoolInfo::No => BoolInfo::No,
  }
}

fn info_to_bool(info: &Info) -> BoolInfo {
  match *info {
    Info::Any => BoolInfo::Any,
    Info::No => BoolInfo::No,
    Info::Fun(_) | Info::Int(_) | Info::Obj(_) => BoolInfo::True,
    Info::True | Info::AnyTrue => BoolInfo::True,
    Info::False => BoolInfo::False,
  }
}

fn join_info(info_1: &Info, info_2: &Info) -> Info {
  match (info_1, info_2) {
    (&Info::Any, info_2) => info_2.clone(),
    (info_1, &Info::Any) => info_1.clone(),
    (&Info::No, _) => Info::No,
    (_, &Info::No) => Info::No,
    (&Info::Fun(ref name_1), &Info::Fun(ref name_2)) =>
      if name_1 == name_2 {
        Info::Fun(name_1.clone())
      } else {
        Info::AnyTrue
      },
    (&Info::Int(num_1), &Info::Int(num_2)) =>
      if num_1 == num_2 {
        Info::Int(num_1)
      } else {
        Info::AnyTrue
      },
    (&Info::Obj(ref name_1), &Info::Obj(ref name_2)) =>
      if name_1 == name_2 {
        Info::Obj(name_1.clone())
      } else {
        Info::AnyTrue
      },
    (&Info::False, &Info::False) => Info::False,
    (&Info::True, &Info::True) => Info::True,
    (&Info::False, _) => Info::Any,
    (_, &Info::False) => Info::Any,
    (_, _) => Info::AnyTrue,
  }
}
<commit_msg>Fix a broken optimization<commit_after>use std::collections::{HashMap};
use grit;

pub fn optimize(prog: grit::ProgDef) -> grit::ProgDef {
  let mut prog = prog;
  let prog_info = collect_prog_info(&prog);
  prog.fun_defs = prog.fun_defs.map_in_place(|mut fun_def| {
    let fun_info = prog_info.get(&fun_def.name).unwrap();
    fun_def.blocks = fun_def.blocks.map_in_place(|mut block| {
      block.ops = block.ops.map_in_place(|op|
        update_op(&prog_info, fun_info, op));
      block.jump = update_jump(&prog_info, fun_info, block.jump);
      block
    });
    fun_def
  });
  prog
}

type ProgInfo = HashMap<grit::FunName, FunInfo>;

#[derive(Debug, PartialEq, Clone)]
struct FunInfo {
  arg_count: usize,
  capture_infos: Vec<Info>,
  var_infos: Vec<Info>,
  return_info: Info,
}

#[derive(Debug, PartialEq, Clone)]
enum Info {
  Any,
  Fun(grit::FunName),
  False,
  True,
  Int(i32),
  Obj(grit::ObjName),
  AnyTrue,
  No,
}

#[derive(Debug, PartialEq, Copy, Clone)]
enum BoolInfo {
  Any,
  True,
  False,
  No,
}

fn update_op(prog_info: &ProgInfo, fun_info: &FunInfo, op: grit::Op) -> grit::Op {
  match op {
    grit::Op::Call(target_var, callee, args) =>
      grit::Op::Call(target_var,
        update_callee(prog_info, fun_info, callee, args.len()),
        args.map_in_place(|a| update_val(prog_info, fun_info, a))),
    grit::Op::ExternCall(target_var, ext_name, args) =>
      grit::Op::ExternCall(target_var, ext_name,
        args.map_in_place(|a| update_val(prog_info, fun_info, a))),
    grit::Op::AllocClos(closs) =>
      grit::Op::AllocClos(closs.map_in_place(|(var, fun_name, captures)| {
        (var, fun_name, captures.map_in_place(|c| update_val(prog_info, fun_info, c)))
      })),
    grit::Op::Assign(var_vals) =>
      grit::Op::Assign(var_vals.map_in_place(|(var, val)| {
        (var, update_val(prog_info, fun_info, val))
      })),
  }
}

fn update_jump(prog_info: &ProgInfo, fun_info: &FunInfo, jump: grit::Jump) -> grit::Jump {
  match jump {
    grit::Jump::Goto(label) =>
      grit::Jump::Goto(label),
    grit::Jump::Return(val) =>
      grit::Jump::Return(update_val(prog_info, fun_info, val)),
    grit::Jump::TailCall(callee, args) =>
      grit::Jump::TailCall(update_callee(prog_info, fun_info, callee, args.len()),
        args.map_in_place(|arg| update_val(prog_info, fun_info, arg))),
    grit::Jump::Branch(boolval, then_label, else_label) =>
      match boolval_info(fun_info, &boolval) {
        BoolInfo::True => grit::Jump::Goto(then_label),
        BoolInfo::False => grit::Jump::Goto(else_label),
        BoolInfo::Any =>
          grit::Jump::Branch(grit::Boolval::IsTrue(grit::Val::Undefined),
            then_label, else_label),
        BoolInfo::No =>
          grit::Jump::Branch(update_boolval(prog_info, fun_info, boolval),
            then_label, else_label),
      },
  }
}

fn update_callee(prog_info: &ProgInfo, fun_info: &FunInfo,
  callee: grit::Callee, arg_count: usize) -> grit::Callee
{
  match callee {
    grit::Callee::Unknown(val) =>
      match val_info(fun_info, &val) {
        Info::Fun(callee_name) => {
          let callee_info = prog_info.get(&callee_name).unwrap();
          if callee_info.arg_count == arg_count {
            if callee_info.capture_infos.len() == 0 {
              grit::Callee::Combinator(callee_name)
            } else {
              grit::Callee::KnownClosure(callee_name, val)
            }
          } else {
            grit::Callee::Unknown(val)
          }
        },
        _ =>
          grit::Callee::Unknown(val)
      },
    callee => callee,
  }
}

fn update_val(prog_info: &ProgInfo, fun_info: &FunInfo, val: grit::Val) -> grit::Val {
  match val_info(fun_info, &val) {
    Info::Any => grit::Val::Undefined,
    Info::Fun(fun_name) =>
      if prog_info.get(&fun_name).unwrap().capture_infos.len() == 0 {
        grit::Val::Combinator(fun_name)
      } else {
        val
      },
    Info::False => grit::Val::False,
    Info::True => grit::Val::True,
    Info::Int(num) => grit::Val::Int(num),
    Info::Obj(name) => grit::Val::Obj(name),
    Info::AnyTrue | Info::No => val,
  }
}

fn update_boolval(prog_info: &ProgInfo, fun_info: &FunInfo,
  boolval: grit::Boolval) -> grit::Boolval
{
  match boolval {
    grit::Boolval::IsTrue(val) =>
      grit::Boolval::IsTrue(update_val(prog_info, fun_info, val)),
    grit::Boolval::IsFalse(val) =>
      grit::Boolval::IsFalse(update_val(prog_info, fun_info, val)),
  }
}

fn collect_prog_info(prog: &grit::ProgDef) -> ProgInfo {
  let mut prog_info: ProgInfo = prog.fun_defs.iter().map(|fun_def| {
    (fun_def.name.clone(), FunInfo {
      arg_count: fun_def.arg_count,
      capture_infos: (0..fun_def.capture_count).map(|_| Info::Any).collect(),
      var_infos: (0..fun_def.var_count).map(|_| Info::Any).collect(),
      return_info: Info::Any,
    })
  }).collect();

  loop {
    let prev_info = prog_info.clone();

    for fun_def in prog.fun_defs.iter() {
      for block in fun_def.blocks.iter() {
        for op in block.ops.iter() {
          op_info_update(&mut prog_info, &fun_def.name, op);
        }
        jump_info_update(&mut prog_info, &fun_def.name, &block.jump);
      }
    }

    if prev_info == prog_info {
      break
    }
  }

  prog_info
}

fn op_info_update(prog_info: &mut ProgInfo, fun_name: &grit::FunName, op: &grit::Op) {
  match *op {
    grit::Op::Call(ref result_var, ref callee, _) =>
      prog_info.get_mut(fun_name).unwrap().var_infos[result_var.0] = join_info(
        &prog_info.get(fun_name).unwrap().var_infos[result_var.0],
        &return_info(prog_info, callee)),
    grit::Op::ExternCall(ref result_var, _, _) =>
      prog_info.get_mut(fun_name).unwrap().var_infos[result_var.0] = Info::No,
    grit::Op::AllocClos(ref closs) =>
      for &(ref clos_var, ref clos_name, ref captures) in closs.iter() {
        prog_info.get_mut(fun_name).unwrap().var_infos[clos_var.0] = join_info(
          &prog_info.get(fun_name).unwrap().var_infos[clos_var.0],
          &Info::Fun(clos_name.clone()));
        for (idx, capture_val) in captures.iter().enumerate() {
          prog_info.get_mut(clos_name).unwrap().capture_infos[idx] = join_info(
            &prog_info.get(clos_name).unwrap().capture_infos[idx],
            &val_info(prog_info.get(fun_name).unwrap(), capture_val));
        }
      },
    grit::Op::Assign(ref var_vals) =>
      for &(ref var, ref val) in var_vals.iter() {
        prog_info.get_mut(fun_name).unwrap().var_infos[var.0] = join_info(
          &prog_info.get(fun_name).unwrap().var_infos[var.0],
          &val_info(prog_info.get(fun_name).unwrap(), val));
      },
  }
}

fn jump_info_update(prog_info: &mut ProgInfo,
  fun_name: &grit::FunName, jump: &grit::Jump)
{
  match *jump {
    grit::Jump::Goto(_) => { },
    grit::Jump::Return(ref val) =>
      prog_info.get_mut(fun_name).unwrap().return_info = join_info(
        &prog_info.get(fun_name).unwrap().return_info,
        &val_info(prog_info.get(fun_name).unwrap(), val)),
    grit::Jump::TailCall(ref callee, _) =>
      prog_info.get_mut(fun_name).unwrap().return_info = join_info(
        &prog_info.get(fun_name).unwrap().return_info,
        &return_info(prog_info, callee)),
    grit::Jump::Branch(_, _, _) => { },
  }
}

fn return_info(prog_info: &ProgInfo, callee: &grit::Callee) -> Info {
  match *callee {
    grit::Callee::Combinator(ref callee_name) |
    grit::Callee::KnownClosure(ref callee_name, _) =>
      prog_info.get(callee_name).unwrap().return_info.clone(),
    grit::Callee::Unknown(_) =>
      Info::No,
  }
}

fn val_info(fun_info: &FunInfo, val: &grit::Val) -> Info {
  match *val {
    grit::Val::Var(ref var) => fun_info.var_infos[var.0].clone(),
    grit::Val::Arg(_) => Info::No,
    grit::Val::Capture(idx) => fun_info.capture_infos[idx].clone(),
    grit::Val::Combinator(ref fun_name) => Info::Fun(fun_name.clone()),
    grit::Val::Obj(ref obj_name) => Info::Obj(obj_name.clone()),
    grit::Val::Int(num) => Info::Int(num),
    grit::Val::True => Info::True,
    grit::Val::False => Info::False,
    grit::Val::Undefined => Info::Any,
  }
}

fn boolval_info(fun_info: &FunInfo, boolval: &grit::Boolval) -> BoolInfo {
  match *boolval {
    grit::Boolval::IsTrue(ref val) =>
      info_to_bool(&val_info(fun_info, val)),
    grit::Boolval::IsFalse(ref val) =>
      negate_boolinfo(info_to_bool(&val_info(fun_info, val))),
  }
}

fn negate_boolinfo(boolinfo: BoolInfo) -> BoolInfo {
  match boolinfo {
    BoolInfo::Any => BoolInfo::Any,
    BoolInfo::True => BoolInfo::False,
    BoolInfo::False => BoolInfo::True,
    BoolInfo::No => BoolInfo::No,
  }
}

fn info_to_bool(info: &Info) -> BoolInfo {
  match *info {
    Info::Any => BoolInfo::Any,
    Info::No => BoolInfo::No,
    Info::Fun(_) | Info::Int(_) | Info::Obj(_) => BoolInfo::True,
    Info::True | Info::AnyTrue => BoolInfo::True,
    Info::False => BoolInfo::False,
  }
}

fn join_info(info_1: &Info, info_2: &Info) -> Info {
  match (info_1, info_2) {
    (&Info::Any, info_2) => info_2.clone(),
    (info_1, &Info::Any) => info_1.clone(),
    (&Info::No, _) => Info::No,
    (_, &Info::No) => Info::No,
    (&Info::Fun(ref name_1), &Info::Fun(ref name_2)) =>
      if name_1 == name_2 {
        Info::Fun(name_1.clone())
      } else {
        Info::AnyTrue
      },
    (&Info::Int(num_1), &Info::Int(num_2)) =>
      if num_1 == num_2 {
        Info::Int(num_1)
      } else {
        Info::AnyTrue
      },
    (&Info::Obj(ref name_1), &Info::Obj(ref name_2)) =>
      if name_1 == name_2 {
        Info::Obj(name_1.clone())
      } else {
        Info::AnyTrue
      },
    (&Info::False, &Info::False) => Info::False,
    (&Info::True, &Info::True) => Info::True,
    (&Info::False, _) => Info::No,
    (_, &Info::False) => Info::No,
    (_, _) => Info::AnyTrue,
  }
}
<|endoftext|>


<commit_before>use super::Status;
use crate as ion_shell;
use builtins_proc::builtin;

macro_rules! string_function {
    (#[$outer:meta], $method:tt) => {
        #[$outer]
        pub fn $method(args: &[small::String], _shell: &mut crate::Shell<'_>) -> Status {
            if args.len() <= 2 {
                return Status::bad_argument(concat!(
                    "ion: ",
                    stringify!($method),
                    ": two arguments must be supplied",
                ));
            }
            args[2..].iter().any(|arg| args[1].$method(arg.as_str())).into()
        }
    };
}

string_function!(
#[builtin(
    desc = "check if a given string starts with another one",
    man = "
SYNOPSIS
    starts_with <PATTERN> tests...

DESCRIPTION
    Returns 0 if any argument starts_with contains the first argument, else returns 0"
)], starts_with);
string_function!(
#[builtin(
    desc = "check if a given string starts with another one",
    man = "
SYNOPSIS
    starts_with <PATTERN> tests...

DESCRIPTION
    Returns 0 if any argument starts_with contains the first argument, else returns 0"
)], ends_with);
string_function!(
#[builtin(
    desc = "check if a given string starts with another one",
    man = "
SYNOPSIS
    starts_with <PATTERN> tests...

DESCRIPTION
    Returns 0 if any argument starts_with contains the first argument, else returns 0"
)], contains);
<commit_msg>Change starts-with, ends-with and contains builtin descriptions<commit_after>use super::Status;
use crate as ion_shell;
use builtins_proc::builtin;

macro_rules! string_function {
    (#[$outer:meta], $method:tt) => {
        #[$outer]
        pub fn $method(args: &[small::String], _shell: &mut crate::Shell<'_>) -> Status {
            if args.len() <= 2 {
                return Status::bad_argument(concat!(
                    "ion: ",
                    stringify!($method),
                    ": two arguments must be supplied",
                ));
            }
            args[2..].iter().any(|arg| args[1].$method(arg.as_str())).into()
        }
    };
}

string_function!(
#[builtin(
    desc = "check if a given string starts with another one",
    man = "
SYNOPSIS
    starts-with <PATTERN> tests...

DESCRIPTION
    Returns 0 if the first argument starts with any other argument, else returns 0"
)], starts_with);
string_function!(
#[builtin(
    desc = "check if a given string ends with another one",
    man = "
SYNOPSIS
    ends-with <PATTERN> tests...

DESCRIPTION
    Returns 0 if the first argument ends with any other argument, else returns 0"
)], ends_with);
string_function!(
#[builtin(
    desc = "check if a given string contains another one",
    man = "
SYNOPSIS
    contains <PATTERN> tests...

DESCRIPTION
    Returns 0 if the first argument contains any other argument, else returns 0"
)], contains);
<|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.

//! Sampling from random distributions.
//!
//! This is a generalization of `Rand` to allow parameters to control the
//! exact properties of the generated values, e.g. the mean and standard
//! deviation of a normal distribution. The `Sample` trait is the most
//! general, and allows for generating values that change some state
//! internally. The `IndependentSample` trait is for generating values
//! that do not need to record state.

#![experimental]

use core::prelude::*;
use core::num::{Float, Int};

use {Rng, Rand};

pub use self::range::Range;
pub use self::gamma::{Gamma, ChiSquared, FisherF, StudentT};
pub use self::normal::{Normal, LogNormal};
pub use self::exponential::Exp;

pub mod range;
pub mod gamma;
pub mod normal;
pub mod exponential;

/// Types that can be used to create a random instance of `Support`.
pub trait Sample<Support> {
    /// Generate a random value of `Support`, using `rng` as the
    /// source of randomness.
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Support;
}

/// `Sample`s that do not require keeping track of state.
///
/// Since no state is recorded, each sample is (statistically)
/// independent of all others, assuming the `Rng` used has this
/// property.
// FIXME maybe having this separate is overkill (the only reason is to
// take &self rather than &mut self)? or maybe this should be the
// trait called `Sample` and the other should be `DependentSample`.
pub trait IndependentSample<Support>: Sample<Support> {
    /// Generate a random value.
    fn ind_sample<R: Rng>(&self, &mut R) -> Support;
}

/// A wrapper for generating types that implement `Rand` via the
/// `Sample` & `IndependentSample` traits.
pub struct RandSample<Sup>;

impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Sup { self.ind_sample(rng) }
}

impl<Sup: Rand> IndependentSample<Sup> for RandSample<Sup> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {
        rng.gen()
    }
}

/// A value with a particular weight for use with `WeightedChoice`.
pub struct Weighted<T> {
    /// The numerical weight of this item
    pub weight: uint,
    /// The actual item which is being weighted
    pub item: T,
}

/// A distribution that selects from a finite collection of weighted items.
///
/// Each item has an associated weight that influences how likely it
/// is to be chosen: higher weight is more likely.
///
/// The `Clone` restriction is a limitation of the `Sample` and
/// `IndependentSample` traits. Note that `&T` is (cheaply) `Clone` for
/// all `T`, as is `uint`, so one can store references or indices into
/// another vector.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{Weighted, WeightedChoice, IndependentSample};
///
/// let mut items = vec!(Weighted { weight: 2, item: 'a' },
///                      Weighted { weight: 4, item: 'b' },
///                      Weighted { weight: 1, item: 'c' });
/// let wc = WeightedChoice::new(items.as_mut_slice());
/// let mut rng = rand::task_rng();
/// for _ in range(0u, 16) {
///      // on average prints 'a' 4 times, 'b' 8 and 'c' twice.
///      println!("{}", wc.ind_sample(&mut rng));
/// }
/// ```
pub struct WeightedChoice<'a, T:'a> {
    items: &'a mut [Weighted<T>],
    weight_range: Range<uint>
}

impl<'a, T: Clone> WeightedChoice<'a, T> {
    /// Create a new `WeightedChoice`.
    ///
    /// Panics if:
    /// - `v` is empty
    /// - the total weight is 0
    /// - the total weight is larger than a `uint` can contain.
    pub fn new<'a>(items: &'a mut [Weighted<T>]) -> WeightedChoice<'a, T> {
        // strictly speaking, this is subsumed by the total weight == 0 case
        assert!(!items.is_empty(), "WeightedChoice::new called with no items");

        let mut running_total = 0u;

        // we convert the list from individual weights to cumulative
        // weights so we can binary search. This *could* drop elements
        // with weight == 0 as an optimisation.
        for item in items.iter_mut() {
            running_total = match running_total.checked_add(item.weight) {
                Some(n) => n,
                None => panic!("WeightedChoice::new called with a total weight \
                               larger than a uint can contain")
            };

            item.weight = running_total;
        }
        assert!(running_total != 0, "WeightedChoice::new called with a total weight of 0");

        WeightedChoice {
            items: items,
            // we're likely to be generating numbers in this range
            // relatively often, so might as well cache it
            weight_range: Range::new(0, running_total)
        }
    }
}

impl<'a, T: Clone> Sample<T> for WeightedChoice<'a, T> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> T { self.ind_sample(rng) }
}

impl<'a, T: Clone> IndependentSample<T> for WeightedChoice<'a, T> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> T {
        // we want to find the first element that has cumulative
        // weight > sample_weight, which we do by binary since the
        // cumulative weights of self.items are sorted.

        // choose a weight in [0, total_weight)
        let sample_weight = self.weight_range.ind_sample(rng);

        // short circuit when it's the first item
        if sample_weight < self.items[0].weight {
            return self.items[0].item.clone();
        }

        let mut idx = 0;
        let mut modifier = self.items.len();

        // now we know that every possibility has an element to the
        // left, so we can just search for the last element that has
        // cumulative weight <= sample_weight, then the next one will
        // be "it". (Note that this greatest element will never be the
        // last element of the vector, since sample_weight is chosen
        // in [0, total_weight) and the cumulative weight of the last
        // one is exactly the total weight.)
        while modifier > 1 {
            let i = idx + modifier / 2;
            if self.items[i].weight <= sample_weight {
                // we're small, so look to the right, but allow this
                // exact element still.
                idx = i;
                // we need the `/ 2` to round up otherwise we'll drop
                // the trailing elements when `modifier` is odd.
                modifier += 1;
            } else {
                // otherwise we're too big, so go left. (i.e. do
                // nothing)
            }
            modifier /= 2;
        }
        return self.items[idx + 1].item.clone();
    }
}

mod ziggurat_tables;

/// Sample a random number using the Ziggurat method (specifically the
/// ZIGNOR variant from Doornik 2005). Most of the arguments are
/// directly from the paper:
///
/// * `rng`: source of randomness
/// * `symmetric`: whether this is a symmetric distribution, or one-sided with P(x < 0) = 0.
/// * `X`: the $x_i$ abscissae.
/// * `F`: precomputed values of the PDF at the $x_i$, (i.e. $f(x_i)$)
/// * `F_DIFF`: precomputed values of $f(x_i) - f(x_{i+1})$
/// * `pdf`: the probability density function
/// * `zero_case`: manual sampling from the tail when we chose the
///    bottom box (i.e. i == 0)

// the perf improvement (25-50%) is definitely worth the extra code
// size from force-inlining.
#[inline(always)]
fn ziggurat<R: Rng, P, Z>(
            rng: &mut R,
            symmetric: bool,
            x_tab: ziggurat_tables::ZigTable,
            f_tab: ziggurat_tables::ZigTable,
            mut pdf: P,
            mut zero_case: Z)
            -> f64 where P: FnMut(f64) -> f64, Z: FnMut(&mut R, f64) -> f64 {
    static SCALE: f64 = (1u64 << 53) as f64;
    loop {
        // reimplement the f64 generation as an optimisation suggested
        // by the Doornik paper: we have a lot of precision-space
        // (i.e. there are 11 bits of the 64 of a u64 to use after
        // creating a f64), so we might as well reuse some to save
        // generating a whole extra random number. (Seems to be 15%
        // faster.)
        //
        // This unfortunately misses out on the benefits of direct
        // floating point generation if an RNG like dSMFT is
        // used. (That is, such RNGs create floats directly, highly
        // efficiently and overload next_f32/f64, so by not calling it
        // this may be slower than it would be otherwise.)
        // FIXME: investigate/optimise for the above.
        let bits: u64 = rng.gen();
        let i = (bits & 0xff) as uint;
        let f = (bits >> 11) as f64 / SCALE;

        // u is either U(-1, 1) or U(0, 1) depending on if this is a
        // symmetric distribution or not.
        let u = if symmetric {2.0 * f - 1.0} else {f};
        let x = u * x_tab[i];

        let test_x = if symmetric { x.abs() } else {x};

        // algebraically equivalent to |u| < x_tab[i+1]/x_tab[i] (or u < x_tab[i+1]/x_tab[i])
        if test_x < x_tab[i + 1] {
            return x;
        }
        if i == 0 {
            return zero_case(rng, u);
        }
        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
        if f_tab[i + 1] + (f_tab[i] - f_tab[i + 1]) * rng.gen() < pdf(x) {
            return x;
        }
    }
}

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

    use {Rng, Rand};
    use super::{RandSample, WeightedChoice, Weighted, Sample, IndependentSample};

    #[deriving(PartialEq, Show)]
    struct ConstRand(uint);
    impl Rand for ConstRand {
        fn rand<R: Rng>(_: &mut R) -> ConstRand {
            ConstRand(0)
        }
    }

    // 0, 1, 2, 3, ...
    struct CountingRng { i: u32 }
    impl Rng for CountingRng {
        fn next_u32(&mut self) -> u32 {
            self.i += 1;
            self.i - 1
        }
        fn next_u64(&mut self) -> u64 {
            self.next_u32() as u64
        }
    }

    #[test]
    fn test_rand_sample() {
        let mut rand_sample = RandSample::<ConstRand>;

        assert_eq!(rand_sample.sample(&mut ::test::rng()), ConstRand(0));
        assert_eq!(rand_sample.ind_sample(&mut ::test::rng()), ConstRand(0));
    }
    #[test]
    fn test_weighted_choice() {
        // this makes assumptions about the internal implementation of
        // WeightedChoice, specifically: it doesn't reorder the items,
        // it doesn't do weird things to the RNG (so 0 maps to 0, 1 to
        // 1, internally; modulo a modulo operation).

        macro_rules! t (
            ($items:expr, $expected:expr) => {{
                let mut items = $items;
                let wc = WeightedChoice::new(items.as_mut_slice());
                let expected = $expected;

                let mut rng = CountingRng { i: 0 };

                for &val in expected.iter() {
                    assert_eq!(wc.ind_sample(&mut rng), val)
                }
            }}
        );

        t!(vec!(Weighted { weight: 1, item: 10i}), [10]);

        // skip some
        t!(vec!(Weighted { weight: 0, item: 20i},
                Weighted { weight: 2, item: 21i},
                Weighted { weight: 0, item: 22i},
                Weighted { weight: 1, item: 23i}),
           [21,21, 23]);

        // different weights
        t!(vec!(Weighted { weight: 4, item: 30i},
                Weighted { weight: 3, item: 31i}),
           [30,30,30,30, 31,31,31]);

        // check that we're binary searching
        // correctly with some vectors of odd
        // length.
        t!(vec!(Weighted { weight: 1, item: 40i},
                Weighted { weight: 1, item: 41i},
                Weighted { weight: 1, item: 42i},
                Weighted { weight: 1, item: 43i},
                Weighted { weight: 1, item: 44i}),
           [40, 41, 42, 43, 44]);
        t!(vec!(Weighted { weight: 1, item: 50i},
                Weighted { weight: 1, item: 51i},
                Weighted { weight: 1, item: 52i},
                Weighted { weight: 1, item: 53i},
                Weighted { weight: 1, item: 54i},
                Weighted { weight: 1, item: 55i},
                Weighted { weight: 1, item: 56i}),
           [50, 51, 52, 53, 54, 55, 56]);
    }

    #[test] #[should_fail]
    fn test_weighted_choice_no_items() {
        WeightedChoice::<int>::new(&mut []);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_zero_weight() {
        WeightedChoice::new(&mut [Weighted { weight: 0, item: 0i},
                                  Weighted { weight: 0, item: 1i}]);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_weight_overflows() {
        let x = (-1) as uint / 2; // x + x + 2 is the overflow
        WeightedChoice::new(&mut [Weighted { weight: x, item: 0i },
                                  Weighted { weight: 1, item: 1i },
                                  Weighted { weight: x, item: 2i },
                                  Weighted { weight: 1, item: 3i }]);
    }
}
<commit_msg>auto merge of #19777 : nikomatsakis/rust/warn-on-shadowing, r=acrichto<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.

//! Sampling from random distributions.
//!
//! This is a generalization of `Rand` to allow parameters to control the
//! exact properties of the generated values, e.g. the mean and standard
//! deviation of a normal distribution. The `Sample` trait is the most
//! general, and allows for generating values that change some state
//! internally. The `IndependentSample` trait is for generating values
//! that do not need to record state.

#![experimental]

use core::prelude::*;
use core::num::{Float, Int};

use {Rng, Rand};

pub use self::range::Range;
pub use self::gamma::{Gamma, ChiSquared, FisherF, StudentT};
pub use self::normal::{Normal, LogNormal};
pub use self::exponential::Exp;

pub mod range;
pub mod gamma;
pub mod normal;
pub mod exponential;

/// Types that can be used to create a random instance of `Support`.
pub trait Sample<Support> {
    /// Generate a random value of `Support`, using `rng` as the
    /// source of randomness.
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Support;
}

/// `Sample`s that do not require keeping track of state.
///
/// Since no state is recorded, each sample is (statistically)
/// independent of all others, assuming the `Rng` used has this
/// property.
// FIXME maybe having this separate is overkill (the only reason is to
// take &self rather than &mut self)? or maybe this should be the
// trait called `Sample` and the other should be `DependentSample`.
pub trait IndependentSample<Support>: Sample<Support> {
    /// Generate a random value.
    fn ind_sample<R: Rng>(&self, &mut R) -> Support;
}

/// A wrapper for generating types that implement `Rand` via the
/// `Sample` & `IndependentSample` traits.
pub struct RandSample<Sup>;

impl<Sup: Rand> Sample<Sup> for RandSample<Sup> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Sup { self.ind_sample(rng) }
}

impl<Sup: Rand> IndependentSample<Sup> for RandSample<Sup> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Sup {
        rng.gen()
    }
}

/// A value with a particular weight for use with `WeightedChoice`.
pub struct Weighted<T> {
    /// The numerical weight of this item
    pub weight: uint,
    /// The actual item which is being weighted
    pub item: T,
}

/// A distribution that selects from a finite collection of weighted items.
///
/// Each item has an associated weight that influences how likely it
/// is to be chosen: higher weight is more likely.
///
/// The `Clone` restriction is a limitation of the `Sample` and
/// `IndependentSample` traits. Note that `&T` is (cheaply) `Clone` for
/// all `T`, as is `uint`, so one can store references or indices into
/// another vector.
///
/// # Example
///
/// ```rust
/// use std::rand;
/// use std::rand::distributions::{Weighted, WeightedChoice, IndependentSample};
///
/// let mut items = vec!(Weighted { weight: 2, item: 'a' },
///                      Weighted { weight: 4, item: 'b' },
///                      Weighted { weight: 1, item: 'c' });
/// let wc = WeightedChoice::new(items.as_mut_slice());
/// let mut rng = rand::task_rng();
/// for _ in range(0u, 16) {
///      // on average prints 'a' 4 times, 'b' 8 and 'c' twice.
///      println!("{}", wc.ind_sample(&mut rng));
/// }
/// ```
pub struct WeightedChoice<'a, T:'a> {
    items: &'a mut [Weighted<T>],
    weight_range: Range<uint>
}

impl<'a, T: Clone> WeightedChoice<'a, T> {
    /// Create a new `WeightedChoice`.
    ///
    /// Panics if:
    /// - `v` is empty
    /// - the total weight is 0
    /// - the total weight is larger than a `uint` can contain.
    pub fn new(items: &'a mut [Weighted<T>]) -> WeightedChoice<'a, T> {
        // strictly speaking, this is subsumed by the total weight == 0 case
        assert!(!items.is_empty(), "WeightedChoice::new called with no items");

        let mut running_total = 0u;

        // we convert the list from individual weights to cumulative
        // weights so we can binary search. This *could* drop elements
        // with weight == 0 as an optimisation.
        for item in items.iter_mut() {
            running_total = match running_total.checked_add(item.weight) {
                Some(n) => n,
                None => panic!("WeightedChoice::new called with a total weight \
                               larger than a uint can contain")
            };

            item.weight = running_total;
        }
        assert!(running_total != 0, "WeightedChoice::new called with a total weight of 0");

        WeightedChoice {
            items: items,
            // we're likely to be generating numbers in this range
            // relatively often, so might as well cache it
            weight_range: Range::new(0, running_total)
        }
    }
}

impl<'a, T: Clone> Sample<T> for WeightedChoice<'a, T> {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> T { self.ind_sample(rng) }
}

impl<'a, T: Clone> IndependentSample<T> for WeightedChoice<'a, T> {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> T {
        // we want to find the first element that has cumulative
        // weight > sample_weight, which we do by binary since the
        // cumulative weights of self.items are sorted.

        // choose a weight in [0, total_weight)
        let sample_weight = self.weight_range.ind_sample(rng);

        // short circuit when it's the first item
        if sample_weight < self.items[0].weight {
            return self.items[0].item.clone();
        }

        let mut idx = 0;
        let mut modifier = self.items.len();

        // now we know that every possibility has an element to the
        // left, so we can just search for the last element that has
        // cumulative weight <= sample_weight, then the next one will
        // be "it". (Note that this greatest element will never be the
        // last element of the vector, since sample_weight is chosen
        // in [0, total_weight) and the cumulative weight of the last
        // one is exactly the total weight.)
        while modifier > 1 {
            let i = idx + modifier / 2;
            if self.items[i].weight <= sample_weight {
                // we're small, so look to the right, but allow this
                // exact element still.
                idx = i;
                // we need the `/ 2` to round up otherwise we'll drop
                // the trailing elements when `modifier` is odd.
                modifier += 1;
            } else {
                // otherwise we're too big, so go left. (i.e. do
                // nothing)
            }
            modifier /= 2;
        }
        return self.items[idx + 1].item.clone();
    }
}

mod ziggurat_tables;

/// Sample a random number using the Ziggurat method (specifically the
/// ZIGNOR variant from Doornik 2005). Most of the arguments are
/// directly from the paper:
///
/// * `rng`: source of randomness
/// * `symmetric`: whether this is a symmetric distribution, or one-sided with P(x < 0) = 0.
/// * `X`: the $x_i$ abscissae.
/// * `F`: precomputed values of the PDF at the $x_i$, (i.e. $f(x_i)$)
/// * `F_DIFF`: precomputed values of $f(x_i) - f(x_{i+1})$
/// * `pdf`: the probability density function
/// * `zero_case`: manual sampling from the tail when we chose the
///    bottom box (i.e. i == 0)

// the perf improvement (25-50%) is definitely worth the extra code
// size from force-inlining.
#[inline(always)]
fn ziggurat<R: Rng, P, Z>(
            rng: &mut R,
            symmetric: bool,
            x_tab: ziggurat_tables::ZigTable,
            f_tab: ziggurat_tables::ZigTable,
            mut pdf: P,
            mut zero_case: Z)
            -> f64 where P: FnMut(f64) -> f64, Z: FnMut(&mut R, f64) -> f64 {
    static SCALE: f64 = (1u64 << 53) as f64;
    loop {
        // reimplement the f64 generation as an optimisation suggested
        // by the Doornik paper: we have a lot of precision-space
        // (i.e. there are 11 bits of the 64 of a u64 to use after
        // creating a f64), so we might as well reuse some to save
        // generating a whole extra random number. (Seems to be 15%
        // faster.)
        //
        // This unfortunately misses out on the benefits of direct
        // floating point generation if an RNG like dSMFT is
        // used. (That is, such RNGs create floats directly, highly
        // efficiently and overload next_f32/f64, so by not calling it
        // this may be slower than it would be otherwise.)
        // FIXME: investigate/optimise for the above.
        let bits: u64 = rng.gen();
        let i = (bits & 0xff) as uint;
        let f = (bits >> 11) as f64 / SCALE;

        // u is either U(-1, 1) or U(0, 1) depending on if this is a
        // symmetric distribution or not.
        let u = if symmetric {2.0 * f - 1.0} else {f};
        let x = u * x_tab[i];

        let test_x = if symmetric { x.abs() } else {x};

        // algebraically equivalent to |u| < x_tab[i+1]/x_tab[i] (or u < x_tab[i+1]/x_tab[i])
        if test_x < x_tab[i + 1] {
            return x;
        }
        if i == 0 {
            return zero_case(rng, u);
        }
        // algebraically equivalent to f1 + DRanU()*(f0 - f1) < 1
        if f_tab[i + 1] + (f_tab[i] - f_tab[i + 1]) * rng.gen() < pdf(x) {
            return x;
        }
    }
}

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

    use {Rng, Rand};
    use super::{RandSample, WeightedChoice, Weighted, Sample, IndependentSample};

    #[deriving(PartialEq, Show)]
    struct ConstRand(uint);
    impl Rand for ConstRand {
        fn rand<R: Rng>(_: &mut R) -> ConstRand {
            ConstRand(0)
        }
    }

    // 0, 1, 2, 3, ...
    struct CountingRng { i: u32 }
    impl Rng for CountingRng {
        fn next_u32(&mut self) -> u32 {
            self.i += 1;
            self.i - 1
        }
        fn next_u64(&mut self) -> u64 {
            self.next_u32() as u64
        }
    }

    #[test]
    fn test_rand_sample() {
        let mut rand_sample = RandSample::<ConstRand>;

        assert_eq!(rand_sample.sample(&mut ::test::rng()), ConstRand(0));
        assert_eq!(rand_sample.ind_sample(&mut ::test::rng()), ConstRand(0));
    }
    #[test]
    fn test_weighted_choice() {
        // this makes assumptions about the internal implementation of
        // WeightedChoice, specifically: it doesn't reorder the items,
        // it doesn't do weird things to the RNG (so 0 maps to 0, 1 to
        // 1, internally; modulo a modulo operation).

        macro_rules! t (
            ($items:expr, $expected:expr) => {{
                let mut items = $items;
                let wc = WeightedChoice::new(items.as_mut_slice());
                let expected = $expected;

                let mut rng = CountingRng { i: 0 };

                for &val in expected.iter() {
                    assert_eq!(wc.ind_sample(&mut rng), val)
                }
            }}
        );

        t!(vec!(Weighted { weight: 1, item: 10i}), [10]);

        // skip some
        t!(vec!(Weighted { weight: 0, item: 20i},
                Weighted { weight: 2, item: 21i},
                Weighted { weight: 0, item: 22i},
                Weighted { weight: 1, item: 23i}),
           [21,21, 23]);

        // different weights
        t!(vec!(Weighted { weight: 4, item: 30i},
                Weighted { weight: 3, item: 31i}),
           [30,30,30,30, 31,31,31]);

        // check that we're binary searching
        // correctly with some vectors of odd
        // length.
        t!(vec!(Weighted { weight: 1, item: 40i},
                Weighted { weight: 1, item: 41i},
                Weighted { weight: 1, item: 42i},
                Weighted { weight: 1, item: 43i},
                Weighted { weight: 1, item: 44i}),
           [40, 41, 42, 43, 44]);
        t!(vec!(Weighted { weight: 1, item: 50i},
                Weighted { weight: 1, item: 51i},
                Weighted { weight: 1, item: 52i},
                Weighted { weight: 1, item: 53i},
                Weighted { weight: 1, item: 54i},
                Weighted { weight: 1, item: 55i},
                Weighted { weight: 1, item: 56i}),
           [50, 51, 52, 53, 54, 55, 56]);
    }

    #[test] #[should_fail]
    fn test_weighted_choice_no_items() {
        WeightedChoice::<int>::new(&mut []);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_zero_weight() {
        WeightedChoice::new(&mut [Weighted { weight: 0, item: 0i},
                                  Weighted { weight: 0, item: 1i}]);
    }
    #[test] #[should_fail]
    fn test_weighted_choice_weight_overflows() {
        let x = (-1) as uint / 2; // x + x + 2 is the overflow
        WeightedChoice::new(&mut [Weighted { weight: x, item: 0i },
                                  Weighted { weight: 1, item: 1i },
                                  Weighted { weight: x, item: 2i },
                                  Weighted { weight: 1, item: 3i }]);
    }
}
<|endoftext|>