<commit_before>//! Rustdoc's JSON backend
//!
//! This module contains the logic for rendering a crate as JSON rather than the normal static HTML
//! output. See [the RFC](https://github.com/rust-lang/rfcs/pull/2963) and the [`types`] module
//! docs for usage and details.

mod conversions;

use std::cell::RefCell;
use std::fs::File;
use std::path::PathBuf;
use std::rc::Rc;

use rustc_data_structures::fx::FxHashMap;
use rustc_middle::ty::TyCtxt;
use rustc_session::Session;
use rustc_span::edition::Edition;

use json_types as types;

use crate::clean;
use crate::config::{RenderInfo, RenderOptions};
use crate::error::Error;
use crate::formats::cache::Cache;
use crate::formats::FormatRenderer;
use crate::html::render::cache::ExternalLocation;
use crate::json::conversions::from_def_id;

#[derive(Clone)]
crate struct JsonRenderer<'tcx> {
    tcx: TyCtxt<'tcx>,
    /// A mapping of IDs that contains all local items for this crate which gets output as a top
    /// level field of the JSON blob.
    index: Rc<RefCell<FxHashMap<types::Id, types::Item>>>,
    /// The directory where the blob will be written to.
    out_path: PathBuf,
    cache: Rc<Cache>,
}

impl JsonRenderer<'_> {
    fn sess(&self) -> &Session {
        self.tcx.sess
    }

    fn get_trait_implementors(&mut self, id: rustc_span::def_id::DefId) -> Vec<types::Id> {
        Rc::clone(&self.cache)
            .implementors
            .get(&id)
            .map(|implementors| {
                implementors
                    .iter()
                    .map(|i| {
                        let item = &i.impl_item;
                        self.item(item.clone()).unwrap();
                        from_def_id(item.def_id)
                    })
                    .collect()
            })
            .unwrap_or_default()
    }

    fn get_impls(&mut self, id: rustc_span::def_id::DefId) -> Vec<types::Id> {
        Rc::clone(&self.cache)
            .impls
            .get(&id)
            .map(|impls| {
                impls
                    .iter()
                    .filter_map(|i| {
                        let item = &i.impl_item;
                        if item.def_id.is_local() {
                            self.item(item.clone()).unwrap();
                            Some(from_def_id(item.def_id))
                        } else {
                            None
                        }
                    })
                    .collect()
            })
            .unwrap_or_default()
    }

    fn get_trait_items(&mut self) -> Vec<(types::Id, types::Item)> {
        Rc::clone(&self.cache)
            .traits
            .iter()
            .filter_map(|(&id, trait_item)| {
                // only need to synthesize items for external traits
                if !id.is_local() {
                    trait_item.items.clone().into_iter().for_each(|i| self.item(i).unwrap());
                    Some((
                        from_def_id(id),
                        types::Item {
                            id: from_def_id(id),
                            crate_id: id.krate.as_u32(),
                            name: self
                                .cache
                                .paths
                                .get(&id)
                                .unwrap_or_else(|| {
                                    self.cache
                                        .external_paths
                                        .get(&id)
                                        .expect("Trait should either be in local or external paths")
                                })
                                .0
                                .last()
                                .map(Clone::clone),
                            visibility: types::Visibility::Public,
                            kind: types::ItemKind::Trait,
                            inner: types::ItemEnum::TraitItem(trait_item.clone().into()),
                            source: None,
                            docs: Default::default(),
                            links: Default::default(),
                            attrs: Default::default(),
                            deprecation: Default::default(),
                        },
                    ))
                } else {
                    None
                }
            })
            .collect()
    }
}

impl<'tcx> FormatRenderer<'tcx> for JsonRenderer<'tcx> {
    fn descr() -> &'static str {
        "json"
    }

    fn init(
        krate: clean::Crate,
        options: RenderOptions,
        _render_info: RenderInfo,
        _edition: Edition,
        cache: Cache,
        tcx: TyCtxt<'tcx>,
    ) -> Result<(Self, clean::Crate), Error> {
        debug!("Initializing json renderer");
        Ok((
            JsonRenderer {
                tcx,
                index: Rc::new(RefCell::new(FxHashMap::default())),
                out_path: options.output,
                cache: Rc::new(cache),
            },
            krate,
        ))
    }

    /// Inserts an item into the index. This should be used rather than directly calling insert on
    /// the hashmap because certain items (traits and types) need to have their mappings for trait
    /// implementations filled out before they're inserted.
    fn item(&mut self, item: clean::Item) -> Result<(), Error> {
        // Flatten items that recursively store other items
        item.kind.inner_items().for_each(|i| self.item(i.clone()).unwrap());

        let id = item.def_id;
        if let Some(mut new_item) = self.convert_item(item) {
            if let types::ItemEnum::TraitItem(ref mut t) = new_item.inner {
                t.implementors = self.get_trait_implementors(id)
            } else if let types::ItemEnum::StructItem(ref mut s) = new_item.inner {
                s.impls = self.get_impls(id)
            } else if let types::ItemEnum::EnumItem(ref mut e) = new_item.inner {
                e.impls = self.get_impls(id)
            }
            let removed = self.index.borrow_mut().insert(from_def_id(id), new_item.clone());
            // FIXME(adotinthevoid): Currently, the index is duplicated. This is a sanity check
            // to make sure the items are unique.
            if let Some(old_item) = removed {
                assert_eq!(old_item, new_item);
            }
        }

        Ok(())
    }

    fn mod_item_in(&mut self, item: &clean::Item, _module_name: &str) -> Result<(), Error> {
        use clean::types::ItemKind::*;
        if let ModuleItem(m) = &*item.kind {
            for item in &m.items {
                match &*item.kind {
                    // These don't have names so they don't get added to the output by default
                    ImportItem(_) => self.item(item.clone()).unwrap(),
                    ExternCrateItem(_, _) => self.item(item.clone()).unwrap(),
                    ImplItem(i) => i.items.iter().for_each(|i| self.item(i.clone()).unwrap()),
                    _ => {}
                }
            }
        }
        self.item(item.clone()).unwrap();
        Ok(())
    }

    fn mod_item_out(&mut self, _item_name: &str) -> Result<(), Error> {
        Ok(())
    }

    fn after_krate(
        &mut self,
        krate: &clean::Crate,
        _diag: &rustc_errors::Handler,
    ) -> Result<(), Error> {
        debug!("Done with crate");
        let mut index = (*self.index).clone().into_inner();
        index.extend(self.get_trait_items());
        let len = index.len();
        // This needs to be the default HashMap for compatibility with the public interface for
        // rustdoc-json
        #[allow(rustc::default_hash_types)]
        let output = types::Crate {
            root: types::Id(String::from("0:0")),
            crate_version: krate.version.clone(),
            includes_private: self.cache.document_private,
            index: index.into_iter().fold(
                std::collections::HashMap::with_capacity(len),
                |mut acc, (key, val)| {
                    acc.insert(key, val);
                    acc
                },
            ),
            paths: self
                .cache
                .paths
                .clone()
                .into_iter()
                .chain(self.cache.external_paths.clone().into_iter())
                .map(|(k, (path, kind))| {
                    (
                        from_def_id(k),
                        types::ItemSummary { crate_id: k.krate.as_u32(), path, kind: kind.into() },
                    )
                })
                .collect(),
            external_crates: self
                .cache
                .extern_locations
                .iter()
                .map(|(k, v)| {
                    (
                        k.as_u32(),
                        types::ExternalCrate {
                            name: v.0.to_string(),
                            html_root_url: match &v.2 {
                                ExternalLocation::Remote(s) => Some(s.clone()),
                                _ => None,
                            },
                        },
                    )
                })
                .collect(),
            format_version: 2,
        };
        let mut p = self.out_path.clone();
        p.push(output.index.get(&output.root).unwrap().name.clone().unwrap());
        p.set_extension("json");
        let file = File::create(&p).map_err(|error| Error { error: error.to_string(), file: p })?;
        serde_json::ser::to_writer(&file, &output).unwrap();
        Ok(())
    }

    fn cache(&self) -> &Cache {
        &self.cache
    }
}
<commit_msg>Simplify conversion<commit_after>//! Rustdoc's JSON backend
//!
//! This module contains the logic for rendering a crate as JSON rather than the normal static HTML
//! output. See [the RFC](https://github.com/rust-lang/rfcs/pull/2963) and the [`types`] module
//! docs for usage and details.

mod conversions;

use std::cell::RefCell;
use std::fs::File;
use std::path::PathBuf;
use std::rc::Rc;

use rustc_data_structures::fx::FxHashMap;
use rustc_middle::ty::TyCtxt;
use rustc_session::Session;
use rustc_span::edition::Edition;

use json_types as types;

use crate::clean;
use crate::config::{RenderInfo, RenderOptions};
use crate::error::Error;
use crate::formats::cache::Cache;
use crate::formats::FormatRenderer;
use crate::html::render::cache::ExternalLocation;
use crate::json::conversions::from_def_id;

#[derive(Clone)]
crate struct JsonRenderer<'tcx> {
    tcx: TyCtxt<'tcx>,
    /// A mapping of IDs that contains all local items for this crate which gets output as a top
    /// level field of the JSON blob.
    index: Rc<RefCell<FxHashMap<types::Id, types::Item>>>,
    /// The directory where the blob will be written to.
    out_path: PathBuf,
    cache: Rc<Cache>,
}

impl JsonRenderer<'_> {
    fn sess(&self) -> &Session {
        self.tcx.sess
    }

    fn get_trait_implementors(&mut self, id: rustc_span::def_id::DefId) -> Vec<types::Id> {
        Rc::clone(&self.cache)
            .implementors
            .get(&id)
            .map(|implementors| {
                implementors
                    .iter()
                    .map(|i| {
                        let item = &i.impl_item;
                        self.item(item.clone()).unwrap();
                        from_def_id(item.def_id)
                    })
                    .collect()
            })
            .unwrap_or_default()
    }

    fn get_impls(&mut self, id: rustc_span::def_id::DefId) -> Vec<types::Id> {
        Rc::clone(&self.cache)
            .impls
            .get(&id)
            .map(|impls| {
                impls
                    .iter()
                    .filter_map(|i| {
                        let item = &i.impl_item;
                        if item.def_id.is_local() {
                            self.item(item.clone()).unwrap();
                            Some(from_def_id(item.def_id))
                        } else {
                            None
                        }
                    })
                    .collect()
            })
            .unwrap_or_default()
    }

    fn get_trait_items(&mut self) -> Vec<(types::Id, types::Item)> {
        Rc::clone(&self.cache)
            .traits
            .iter()
            .filter_map(|(&id, trait_item)| {
                // only need to synthesize items for external traits
                if !id.is_local() {
                    trait_item.items.clone().into_iter().for_each(|i| self.item(i).unwrap());
                    Some((
                        from_def_id(id),
                        types::Item {
                            id: from_def_id(id),
                            crate_id: id.krate.as_u32(),
                            name: self
                                .cache
                                .paths
                                .get(&id)
                                .unwrap_or_else(|| {
                                    self.cache
                                        .external_paths
                                        .get(&id)
                                        .expect("Trait should either be in local or external paths")
                                })
                                .0
                                .last()
                                .map(Clone::clone),
                            visibility: types::Visibility::Public,
                            kind: types::ItemKind::Trait,
                            inner: types::ItemEnum::TraitItem(trait_item.clone().into()),
                            source: None,
                            docs: Default::default(),
                            links: Default::default(),
                            attrs: Default::default(),
                            deprecation: Default::default(),
                        },
                    ))
                } else {
                    None
                }
            })
            .collect()
    }
}

impl<'tcx> FormatRenderer<'tcx> for JsonRenderer<'tcx> {
    fn descr() -> &'static str {
        "json"
    }

    fn init(
        krate: clean::Crate,
        options: RenderOptions,
        _render_info: RenderInfo,
        _edition: Edition,
        cache: Cache,
        tcx: TyCtxt<'tcx>,
    ) -> Result<(Self, clean::Crate), Error> {
        debug!("Initializing json renderer");
        Ok((
            JsonRenderer {
                tcx,
                index: Rc::new(RefCell::new(FxHashMap::default())),
                out_path: options.output,
                cache: Rc::new(cache),
            },
            krate,
        ))
    }

    /// Inserts an item into the index. This should be used rather than directly calling insert on
    /// the hashmap because certain items (traits and types) need to have their mappings for trait
    /// implementations filled out before they're inserted.
    fn item(&mut self, item: clean::Item) -> Result<(), Error> {
        // Flatten items that recursively store other items
        item.kind.inner_items().for_each(|i| self.item(i.clone()).unwrap());

        let id = item.def_id;
        if let Some(mut new_item) = self.convert_item(item) {
            if let types::ItemEnum::TraitItem(ref mut t) = new_item.inner {
                t.implementors = self.get_trait_implementors(id)
            } else if let types::ItemEnum::StructItem(ref mut s) = new_item.inner {
                s.impls = self.get_impls(id)
            } else if let types::ItemEnum::EnumItem(ref mut e) = new_item.inner {
                e.impls = self.get_impls(id)
            }
            let removed = self.index.borrow_mut().insert(from_def_id(id), new_item.clone());
            // FIXME(adotinthevoid): Currently, the index is duplicated. This is a sanity check
            // to make sure the items are unique.
            if let Some(old_item) = removed {
                assert_eq!(old_item, new_item);
            }
        }

        Ok(())
    }

    fn mod_item_in(&mut self, item: &clean::Item, _module_name: &str) -> Result<(), Error> {
        use clean::types::ItemKind::*;
        if let ModuleItem(m) = &*item.kind {
            for item in &m.items {
                match &*item.kind {
                    // These don't have names so they don't get added to the output by default
                    ImportItem(_) => self.item(item.clone()).unwrap(),
                    ExternCrateItem(_, _) => self.item(item.clone()).unwrap(),
                    ImplItem(i) => i.items.iter().for_each(|i| self.item(i.clone()).unwrap()),
                    _ => {}
                }
            }
        }
        self.item(item.clone()).unwrap();
        Ok(())
    }

    fn mod_item_out(&mut self, _item_name: &str) -> Result<(), Error> {
        Ok(())
    }

    fn after_krate(
        &mut self,
        krate: &clean::Crate,
        _diag: &rustc_errors::Handler,
    ) -> Result<(), Error> {
        debug!("Done with crate");
        let mut index = (*self.index).clone().into_inner();
        index.extend(self.get_trait_items());
        // This needs to be the default HashMap for compatibility with the public interface for
        // rustdoc-json
        #[allow(rustc::default_hash_types)]
        let output = types::Crate {
            root: types::Id(String::from("0:0")),
            crate_version: krate.version.clone(),
            includes_private: self.cache.document_private,
            index: index.into_iter().collect(),
            paths: self
                .cache
                .paths
                .clone()
                .into_iter()
                .chain(self.cache.external_paths.clone().into_iter())
                .map(|(k, (path, kind))| {
                    (
                        from_def_id(k),
                        types::ItemSummary { crate_id: k.krate.as_u32(), path, kind: kind.into() },
                    )
                })
                .collect(),
            external_crates: self
                .cache
                .extern_locations
                .iter()
                .map(|(k, v)| {
                    (
                        k.as_u32(),
                        types::ExternalCrate {
                            name: v.0.to_string(),
                            html_root_url: match &v.2 {
                                ExternalLocation::Remote(s) => Some(s.clone()),
                                _ => None,
                            },
                        },
                    )
                })
                .collect(),
            format_version: 2,
        };
        let mut p = self.out_path.clone();
        p.push(output.index.get(&output.root).unwrap().name.clone().unwrap());
        p.set_extension("json");
        let file = File::create(&p).map_err(|error| Error { error: error.to_string(), file: p })?;
        serde_json::ser::to_writer(&file, &output).unwrap();
        Ok(())
    }

    fn cache(&self) -> &Cache {
        &self.cache
    }
}
<|endoftext|>


<commit_before>use crate::ptr;
use crate::ops::{Deref, DerefMut};

/// A wrapper to inhibit compiler from automatically calling `T`’s destructor.
///
/// This wrapper is 0-cost.
///
/// `ManuallyDrop<T>` is subject to the same layout optimizations as `T`.
/// As a consequence, it has *no effect* on the assumptions that the compiler makes
/// about all values being initialized at their type.  In particular, initializing
/// a `ManuallyDrop<&mut T>` with [`mem::zeroed`] is undefined behavior.
/// If you need to handle uninitialized data, use [`MaybeUninit<T>`] instead.
///
/// # Examples
///
/// This wrapper helps with explicitly documenting the drop order dependencies between fields of
/// the type:
///
/// ```rust
/// use std::mem::ManuallyDrop;
/// struct Peach;
/// struct Banana;
/// struct Melon;
/// struct FruitBox {
///     // Immediately clear there’s something non-trivial going on with these fields.
///     peach: ManuallyDrop<Peach>,
///     melon: Melon, // Field that’s independent of the other two.
///     banana: ManuallyDrop<Banana>,
/// }
///
/// impl Drop for FruitBox {
///     fn drop(&mut self) {
///         unsafe {
///             // Explicit ordering in which field destructors are run specified in the intuitive
///             // location – the destructor of the structure containing the fields.
///             // Moreover, one can now reorder fields within the struct however much they want.
///             ManuallyDrop::drop(&mut self.peach);
///             ManuallyDrop::drop(&mut self.banana);
///         }
///         // After destructor for `FruitBox` runs (this function), the destructor for Melon gets
///         // invoked in the usual manner, as it is not wrapped in `ManuallyDrop`.
///     }
/// }
/// ```
///
/// [`mem::zeroed`]: fn.zeroed.html
/// [`MaybeUninit<T>`]: union.MaybeUninit.html
#[stable(feature = "manually_drop", since = "1.20.0")]
#[lang = "manually_drop"]
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct ManuallyDrop<T: ?Sized> {
    value: T,
}

impl<T> ManuallyDrop<T> {
    /// Wrap a value to be manually dropped.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use std::mem::ManuallyDrop;
    /// ManuallyDrop::new(Box::new(()));
    /// ```
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline(always)]
    pub const fn new(value: T) -> ManuallyDrop<T> {
        ManuallyDrop { value }
    }

    /// Extracts the value from the `ManuallyDrop` container.
    ///
    /// This allows the value to be dropped again.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use std::mem::ManuallyDrop;
    /// let x = ManuallyDrop::new(Box::new(()));
    /// let _: Box<()> = ManuallyDrop::into_inner(x); // This drops the `Box`.
    /// ```
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline(always)]
    pub const fn into_inner(slot: ManuallyDrop<T>) -> T {
        slot.value
    }

    /// Takes the contained value out.
    ///
    /// This method is primarily intended for moving out values in drop.
    /// Instead of using [`ManuallyDrop::drop`] to manually drop the value,
    /// you can use this method to take the value and use it however desired.
    /// `Drop` will be invoked on the returned value following normal end-of-scope rules.
    ///
    /// If you have ownership of the container, you can use [`ManuallyDrop::into_inner`] instead.
    ///
    /// # Safety
    ///
    /// This function semantically moves out the contained value without preventing further usage.
    /// It is up to the user of this method to ensure that this container is not used again.
    ///
    /// [`ManuallyDrop::drop`]: #method.drop
    /// [`ManuallyDrop::into_inner`]: #method.into_inner
    #[must_use = "if you don't need the value, you can use `ManuallyDrop::drop` instead"]
    #[unstable(feature = "manually_drop_take", issue = "55422")]
    #[inline]
    pub unsafe fn take(slot: &mut ManuallyDrop<T>) -> T {
        ManuallyDrop::into_inner(ptr::read(slot))
    }
}

impl<T: ?Sized> ManuallyDrop<T> {
    /// Manually drops the contained value.
    ///
    /// If you have ownership of the value, you can use [`ManuallyDrop::into_inner`] instead.
    ///
    /// # Safety
    ///
    /// This function runs the destructor of the contained value and thus the wrapped value
    /// now represents uninitialized data. It is up to the user of this method to ensure the
    /// uninitialized data is not actually used.
    ///
    /// [`ManuallyDrop::into_inner`]: #method.into_inner
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline]
    pub unsafe fn drop(slot: &mut ManuallyDrop<T>) {
        ptr::drop_in_place(&mut slot.value)
    }
}

#[stable(feature = "manually_drop", since = "1.20.0")]
impl<T: ?Sized> Deref for ManuallyDrop<T> {
    type Target = T;
    #[inline(always)]
    fn deref(&self) -> &T {
        &self.value
    }
}

#[stable(feature = "manually_drop", since = "1.20.0")]
impl<T: ?Sized> DerefMut for ManuallyDrop<T> {
    #[inline(always)]
    fn deref_mut(&mut self) -> &mut T {
        &mut self.value
    }
}
<commit_msg>Document that ManuallyDrop::drop should not called more than once<commit_after>use crate::ptr;
use crate::ops::{Deref, DerefMut};

/// A wrapper to inhibit compiler from automatically calling `T`’s destructor.
///
/// This wrapper is 0-cost.
///
/// `ManuallyDrop<T>` is subject to the same layout optimizations as `T`.
/// As a consequence, it has *no effect* on the assumptions that the compiler makes
/// about all values being initialized at their type.  In particular, initializing
/// a `ManuallyDrop<&mut T>` with [`mem::zeroed`] is undefined behavior.
/// If you need to handle uninitialized data, use [`MaybeUninit<T>`] instead.
///
/// # Examples
///
/// This wrapper helps with explicitly documenting the drop order dependencies between fields of
/// the type:
///
/// ```rust
/// use std::mem::ManuallyDrop;
/// struct Peach;
/// struct Banana;
/// struct Melon;
/// struct FruitBox {
///     // Immediately clear there’s something non-trivial going on with these fields.
///     peach: ManuallyDrop<Peach>,
///     melon: Melon, // Field that’s independent of the other two.
///     banana: ManuallyDrop<Banana>,
/// }
///
/// impl Drop for FruitBox {
///     fn drop(&mut self) {
///         unsafe {
///             // Explicit ordering in which field destructors are run specified in the intuitive
///             // location – the destructor of the structure containing the fields.
///             // Moreover, one can now reorder fields within the struct however much they want.
///             ManuallyDrop::drop(&mut self.peach);
///             ManuallyDrop::drop(&mut self.banana);
///         }
///         // After destructor for `FruitBox` runs (this function), the destructor for Melon gets
///         // invoked in the usual manner, as it is not wrapped in `ManuallyDrop`.
///     }
/// }
/// ```
///
/// [`mem::zeroed`]: fn.zeroed.html
/// [`MaybeUninit<T>`]: union.MaybeUninit.html
#[stable(feature = "manually_drop", since = "1.20.0")]
#[lang = "manually_drop"]
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct ManuallyDrop<T: ?Sized> {
    value: T,
}

impl<T> ManuallyDrop<T> {
    /// Wrap a value to be manually dropped.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use std::mem::ManuallyDrop;
    /// ManuallyDrop::new(Box::new(()));
    /// ```
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline(always)]
    pub const fn new(value: T) -> ManuallyDrop<T> {
        ManuallyDrop { value }
    }

    /// Extracts the value from the `ManuallyDrop` container.
    ///
    /// This allows the value to be dropped again.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use std::mem::ManuallyDrop;
    /// let x = ManuallyDrop::new(Box::new(()));
    /// let _: Box<()> = ManuallyDrop::into_inner(x); // This drops the `Box`.
    /// ```
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline(always)]
    pub const fn into_inner(slot: ManuallyDrop<T>) -> T {
        slot.value
    }

    /// Takes the contained value out.
    ///
    /// This method is primarily intended for moving out values in drop.
    /// Instead of using [`ManuallyDrop::drop`] to manually drop the value,
    /// you can use this method to take the value and use it however desired.
    /// `Drop` will be invoked on the returned value following normal end-of-scope rules.
    ///
    /// If you have ownership of the container, you can use [`ManuallyDrop::into_inner`] instead.
    ///
    /// # Safety
    ///
    /// This function semantically moves out the contained value without preventing further usage.
    /// It is up to the user of this method to ensure that this container is not used again.
    ///
    /// [`ManuallyDrop::drop`]: #method.drop
    /// [`ManuallyDrop::into_inner`]: #method.into_inner
    #[must_use = "if you don't need the value, you can use `ManuallyDrop::drop` instead"]
    #[unstable(feature = "manually_drop_take", issue = "55422")]
    #[inline]
    pub unsafe fn take(slot: &mut ManuallyDrop<T>) -> T {
        ManuallyDrop::into_inner(ptr::read(slot))
    }
}

impl<T: ?Sized> ManuallyDrop<T> {
    /// Manually drops the contained value.
    ///
    /// If you have ownership of the value, you can use [`ManuallyDrop::into_inner`] instead.
    ///
    /// # Safety
    ///
    /// This function runs the destructor of the contained value and thus the wrapped value
    /// now represents uninitialized data. It is up to the user of this method to ensure the
    /// uninitialized data is not actually used, and that this function is called at most once
    /// for a given instance of ManuallyDrop
    ///
    /// [`ManuallyDrop::into_inner`]: #method.into_inner
    #[stable(feature = "manually_drop", since = "1.20.0")]
    #[inline]
    pub unsafe fn drop(slot: &mut ManuallyDrop<T>) {
        ptr::drop_in_place(&mut slot.value)
    }
}

#[stable(feature = "manually_drop", since = "1.20.0")]
impl<T: ?Sized> Deref for ManuallyDrop<T> {
    type Target = T;
    #[inline(always)]
    fn deref(&self) -> &T {
        &self.value
    }
}

#[stable(feature = "manually_drop", since = "1.20.0")]
impl<T: ?Sized> DerefMut for ManuallyDrop<T> {
    #[inline(always)]
    fn deref_mut(&mut self) -> &mut T {
        &mut self.value
    }
}
<|endoftext|>


<commit_before>use proc_macro::TokenStream;
use syn::{
    Token, Ident, Type, Attribute, ReturnType, Expr, Block, Error,
    braced, parenthesized, parse_macro_input,
};
use syn::spanned::Spanned;
use syn::parse::{Result, Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn;
use quote::quote;
use itertools::Itertools;

#[allow(non_camel_case_types)]
mod kw {
    syn::custom_keyword!(query);
}

/// Ident or a wildcard `_`.
struct IdentOrWild(Ident);

impl Parse for IdentOrWild {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        Ok(if input.peek(Token![_]) {
            let underscore = input.parse::<Token![_]>()?;
            IdentOrWild(Ident::new("_", underscore.span()))
        } else {
            IdentOrWild(input.parse()?)
        })
    }
}

/// A modifier for a query
enum QueryModifier {
    /// The description of the query.
    Desc(Option<Ident>, Punctuated<Expr, Token![,]>),

    /// Cache the query to disk if the `Expr` returns true.
    Cache(Option<Ident>, Expr),

    /// Custom code to load the query from disk.
    LoadCached(Ident, Ident, Block),

    /// A cycle error for this query aborting the compilation with a fatal error.
    FatalCycle,

    /// Don't hash the result, instead just mark a query red if it runs
    NoHash,

    /// Don't force the query
    NoForce,

    /// Generate a dep node based on the dependencies of the query
    Anon,

    // Always evaluate the query, ignoring its depdendencies
    EvalAlways,
}

impl Parse for QueryModifier {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let modifier: Ident = input.parse()?;
        if modifier == "desc" {
            // Parse a description modifier like:
            // `desc { |tcx| "foo {}", tcx.item_path(key) }`
            let attr_content;
            braced!(attr_content in input);
            let tcx = if attr_content.peek(Token![|]) {
                attr_content.parse::<Token![|]>()?;
                let tcx = attr_content.parse()?;
                attr_content.parse::<Token![|]>()?;
                Some(tcx)
            } else {
                None
            };
            let desc = attr_content.parse_terminated(Expr::parse)?;
            Ok(QueryModifier::Desc(tcx, desc))
        } else if modifier == "cache" {
            // Parse a cache modifier like:
            // `cache { |tcx| key.is_local() }`
            let attr_content;
            braced!(attr_content in input);
            let tcx = if attr_content.peek(Token![|]) {
                attr_content.parse::<Token![|]>()?;
                let tcx = attr_content.parse()?;
                attr_content.parse::<Token![|]>()?;
                Some(tcx)
            } else {
                None
            };
            let expr = attr_content.parse()?;
            Ok(QueryModifier::Cache(tcx, expr))
        } else if modifier == "load_cached" {
            // Parse a load_cached modifier like:
            // `load_cached(tcx, id) { tcx.queries.on_disk_cache.try_load_query_result(tcx, id) }`
            let args;
            parenthesized!(args in input);
            let tcx = args.parse()?;
            args.parse::<Token![,]>()?;
            let id = args.parse()?;
            let block = input.parse()?;
            Ok(QueryModifier::LoadCached(tcx, id, block))
        } else if modifier == "fatal_cycle" {
            Ok(QueryModifier::FatalCycle)
        } else if modifier == "no_hash" {
            Ok(QueryModifier::NoHash)
        } else if modifier == "no_force" {
            Ok(QueryModifier::NoForce)
        } else if modifier == "anon" {
            Ok(QueryModifier::Anon)
        } else if modifier == "eval_always" {
            Ok(QueryModifier::EvalAlways)
        } else {
            Err(Error::new(modifier.span(), "unknown query modifier"))
        }
    }
}

/// Ensures only doc comment attributes are used
fn check_attributes(attrs: Vec<Attribute>) -> Result<()> {
    for attr in attrs {
        if !attr.path.is_ident("doc") {
            return Err(Error::new(attr.span(), "attributes not supported on queries"));
        }
    }
    Ok(())
}

/// A compiler query. `query ... { ... }`
struct Query {
    modifiers: List<QueryModifier>,
    name: Ident,
    key: IdentOrWild,
    arg: Type,
    result: ReturnType,
}

impl Parse for Query {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        check_attributes(input.call(Attribute::parse_outer)?)?;

        // Parse the query declaration. Like `query type_of(key: DefId) -> Ty<'tcx>`
        input.parse::<kw::query>()?;
        let name: Ident = input.parse()?;
        let arg_content;
        parenthesized!(arg_content in input);
        let key = arg_content.parse()?;
        arg_content.parse::<Token![:]>()?;
        let arg = arg_content.parse()?;
        let result = input.parse()?;

        // Parse the query modifiers
        let content;
        braced!(content in input);
        let modifiers = content.parse()?;

        Ok(Query {
            modifiers,
            name,
            key,
            arg,
            result,
        })
    }
}

/// A type used to greedily parse another type until the input is empty.
struct List<T>(Vec<T>);

impl<T: Parse> Parse for List<T> {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let mut list = Vec::new();
        while !input.is_empty() {
            list.push(input.parse()?);
        }
        Ok(List(list))
    }
}

/// A named group containing queries.
struct Group {
    name: Ident,
    queries: List<Query>,
}

impl Parse for Group {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let name: Ident = input.parse()?;
        let content;
        braced!(content in input);
        Ok(Group {
            name,
            queries: content.parse()?,
        })
    }
}

struct QueryModifiers {
    /// The description of the query.
    desc: Option<(Option<Ident>, Punctuated<Expr, Token![,]>)>,

    /// Cache the query to disk if the `Expr` returns true.
    cache: Option<(Option<Ident>, Expr)>,

    /// Custom code to load the query from disk.
    load_cached: Option<(Ident, Ident, Block)>,

    /// A cycle error for this query aborting the compilation with a fatal error.
    fatal_cycle: bool,

    /// Don't hash the result, instead just mark a query red if it runs
    no_hash: bool,

    /// Don't force the query
    no_force: bool,

    /// Generate a dep node based on the dependencies of the query
    anon: bool,

    // Always evaluate the query, ignoring its depdendencies
    eval_always: bool,
}

/// Process query modifiers into a struct, erroring on duplicates
fn process_modifiers(query: &mut Query) -> QueryModifiers {
    let mut load_cached = None;
    let mut cache = None;
    let mut desc = None;
    let mut fatal_cycle = false;
    let mut no_hash = false;
    let mut no_force = false;
    let mut anon = false;
    let mut eval_always = false;
    for modifier in query.modifiers.0.drain(..) {
        match modifier {
            QueryModifier::LoadCached(tcx, id, block) => {
                if load_cached.is_some() {
                    panic!("duplicate modifier `load_cached` for query `{}`", query.name);
                }
                load_cached = Some((tcx, id, block));
            }
            QueryModifier::Cache(tcx, expr) => {
                if cache.is_some() {
                    panic!("duplicate modifier `cache` for query `{}`", query.name);
                }
                cache = Some((tcx, expr));
            }
            QueryModifier::Desc(tcx, list) => {
                if desc.is_some() {
                    panic!("duplicate modifier `desc` for query `{}`", query.name);
                }
                desc = Some((tcx, list));
            }
            QueryModifier::FatalCycle => {
                if fatal_cycle {
                    panic!("duplicate modifier `fatal_cycle` for query `{}`", query.name);
                }
                fatal_cycle = true;
            }
            QueryModifier::NoHash => {
                if no_hash {
                    panic!("duplicate modifier `no_hash` for query `{}`", query.name);
                }
                no_hash = true;
            }
            QueryModifier::NoForce => {
                if no_force {
                    panic!("duplicate modifier `no_force` for query `{}`", query.name);
                }
                no_force = true;
            }
            QueryModifier::Anon => {
                if anon {
                    panic!("duplicate modifier `anon` for query `{}`", query.name);
                }
                anon = true;
            }
            QueryModifier::EvalAlways => {
                if eval_always {
                    panic!("duplicate modifier `eval_always` for query `{}`", query.name);
                }
                eval_always = true;
            }
        }
    }
    QueryModifiers {
        load_cached,
        cache,
        desc,
        fatal_cycle,
        no_hash,
        no_force,
        anon,
        eval_always,
    }
}

/// Add the impl of QueryDescription for the query to `impls` if one is requested
fn add_query_description_impl(
    query: &Query,
    modifiers: QueryModifiers,
    impls: &mut proc_macro2::TokenStream
) {
    let name = &query.name;
    let arg = &query.arg;
    let key = &query.key.0;

    // Find out if we should cache the query on disk
    let cache = modifiers.cache.as_ref().map(|(tcx, expr)| {
        let try_load_from_disk = if let Some((tcx, id, block)) = modifiers.load_cached.as_ref() {
            // Use custom code to load the query from disk
            quote! {
                #[inline]
                fn try_load_from_disk(
                    #tcx: TyCtxt<'_, 'tcx, 'tcx>,
                    #id: SerializedDepNodeIndex
                ) -> Option<Self::Value> {
                    #block
                }
            }
        } else {
            // Use the default code to load the query from disk
            quote! {
                #[inline]
                fn try_load_from_disk(
                    tcx: TyCtxt<'_, 'tcx, 'tcx>,
                    id: SerializedDepNodeIndex
                ) -> Option<Self::Value> {
                    tcx.queries.on_disk_cache.try_load_query_result(tcx, id)
                }
            }
        };

        let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
        quote! {
            #[inline]
            fn cache_on_disk(#tcx: TyCtxt<'_, 'tcx, 'tcx>, #key: Self::Key) -> bool {
                #expr
            }

            #try_load_from_disk
        }
    });

    if cache.is_none() && modifiers.load_cached.is_some() {
        panic!("load_cached modifier on query `{}` without a cache modifier", name);
    }

    let desc = modifiers.desc.as_ref().map(|(tcx, desc)| {
        let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
        quote! {
            fn describe(
                #tcx: TyCtxt<'_, '_, '_>,
                #key: #arg,
            ) -> Cow<'static, str> {
                format!(#desc).into()
            }
        }
    });

    if desc.is_some() || cache.is_some() {
        let cache = cache.unwrap_or(quote! {});
        let desc = desc.unwrap_or(quote! {});

        impls.extend(quote! {
            impl<'tcx> QueryDescription<'tcx> for queries::#name<'tcx> {
                #desc
                #cache
            }
        });
    }
}

pub fn rustc_queries(input: TokenStream) -> TokenStream {
    let groups = parse_macro_input!(input as List<Group>);

    let mut query_stream = quote! {};
    let mut query_description_stream = quote! {};
    let mut dep_node_def_stream = quote! {};
    let mut dep_node_force_stream = quote! {};
    let mut no_force_queries = Vec::new();

    for group in groups.0 {
        let mut group_stream = quote! {};
        for mut query in group.queries.0 {
            let modifiers = process_modifiers(&mut query);
            let name = &query.name;
            let arg = &query.arg;
            let result_full = &query.result;
            let result = match query.result {
                ReturnType::Default => quote! { -> () },
                _ => quote! { #result_full },
            };

            let mut attributes = Vec::new();

            // Pass on the fatal_cycle modifier
            if modifiers.fatal_cycle {
                attributes.push(quote! { fatal_cycle });
            };
            // Pass on the no_hash modifier
            if modifiers.no_hash {
                attributes.push(quote! { no_hash });
            };

            let mut attribute_stream = quote! {};

            for e in attributes.into_iter().intersperse(quote! {,}) {
                attribute_stream.extend(e);
            }

            // Add the query to the group
            group_stream.extend(quote! {
                [#attribute_stream] fn #name: #name(#arg) #result,
            });

            let mut attributes = Vec::new();

            // Pass on the anon modifier
            if modifiers.anon {
                attributes.push(quote! { anon });
            };
            // Pass on the eval_always modifier
            if modifiers.eval_always {
                attributes.push(quote! { eval_always });
            };

            let mut attribute_stream = quote! {};
            for e in attributes.into_iter().intersperse(quote! {,}) {
                attribute_stream.extend(e);
            }
            // Create a dep node for the query
            dep_node_def_stream.extend(quote! {
                [#attribute_stream] #name(#arg),
            });

            if modifiers.no_force {
                no_force_queries.push(name.clone());
            } else {
                // Add a match arm to force the query given the dep node
                dep_node_force_stream.extend(quote! {
                    DepKind::#name => {
                        if let Some(key) = RecoverKey::recover($tcx, $dep_node) {
                            force_ex!($tcx, #name, key);
                        } else {
                            return false;
                        }
                    }
                });
            }

            add_query_description_impl(&query, modifiers, &mut query_description_stream);
        }
        let name = &group.name;
        query_stream.extend(quote! {
            #name { #group_stream },
        });
    }

    // Add an arm for the no force queries to panic when trying to force them
    for query in no_force_queries {
        dep_node_force_stream.extend(quote! {
            DepKind::#query |
        });
    }
    dep_node_force_stream.extend(quote! {
        DepKind::Null => {
            bug!("Cannot force dep node: {:?}", $dep_node)
        }
    });

    TokenStream::from(quote! {
        macro_rules! rustc_query_append {
            ([$($macro:tt)*][$($other:tt)*]) => {
                $($macro)* {
                    $($other)*

                    #query_stream

                }
            }
        }
        macro_rules! rustc_dep_node_append {
            ([$($macro:tt)*][$($other:tt)*]) => {
                $($macro)*(
                    $($other)*

                    #dep_node_def_stream
                );
            }
        }
        macro_rules! rustc_dep_node_force {
            ([$dep_node:expr, $tcx:expr] $($other:tt)*) => {
                match $dep_node.kind {
                    $($other)*

                    #dep_node_force_stream
                }
            }
        }
        #query_description_stream
    })
}
<commit_msg>Add cycle_delay_bug to proc macro<commit_after>use proc_macro::TokenStream;
use syn::{
    Token, Ident, Type, Attribute, ReturnType, Expr, Block, Error,
    braced, parenthesized, parse_macro_input,
};
use syn::spanned::Spanned;
use syn::parse::{Result, Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn;
use quote::quote;
use itertools::Itertools;

#[allow(non_camel_case_types)]
mod kw {
    syn::custom_keyword!(query);
}

/// Ident or a wildcard `_`.
struct IdentOrWild(Ident);

impl Parse for IdentOrWild {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        Ok(if input.peek(Token![_]) {
            let underscore = input.parse::<Token![_]>()?;
            IdentOrWild(Ident::new("_", underscore.span()))
        } else {
            IdentOrWild(input.parse()?)
        })
    }
}

/// A modifier for a query
enum QueryModifier {
    /// The description of the query.
    Desc(Option<Ident>, Punctuated<Expr, Token![,]>),

    /// Cache the query to disk if the `Expr` returns true.
    Cache(Option<Ident>, Expr),

    /// Custom code to load the query from disk.
    LoadCached(Ident, Ident, Block),

    /// A cycle error for this query aborting the compilation with a fatal error.
    FatalCycle,

    /// A cycle error results in a delay_bug call
    CycleDelayBug,

    /// Don't hash the result, instead just mark a query red if it runs
    NoHash,

    /// Don't force the query
    NoForce,

    /// Generate a dep node based on the dependencies of the query
    Anon,

    // Always evaluate the query, ignoring its depdendencies
    EvalAlways,
}

impl Parse for QueryModifier {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let modifier: Ident = input.parse()?;
        if modifier == "desc" {
            // Parse a description modifier like:
            // `desc { |tcx| "foo {}", tcx.item_path(key) }`
            let attr_content;
            braced!(attr_content in input);
            let tcx = if attr_content.peek(Token![|]) {
                attr_content.parse::<Token![|]>()?;
                let tcx = attr_content.parse()?;
                attr_content.parse::<Token![|]>()?;
                Some(tcx)
            } else {
                None
            };
            let desc = attr_content.parse_terminated(Expr::parse)?;
            Ok(QueryModifier::Desc(tcx, desc))
        } else if modifier == "cache" {
            // Parse a cache modifier like:
            // `cache { |tcx| key.is_local() }`
            let attr_content;
            braced!(attr_content in input);
            let tcx = if attr_content.peek(Token![|]) {
                attr_content.parse::<Token![|]>()?;
                let tcx = attr_content.parse()?;
                attr_content.parse::<Token![|]>()?;
                Some(tcx)
            } else {
                None
            };
            let expr = attr_content.parse()?;
            Ok(QueryModifier::Cache(tcx, expr))
        } else if modifier == "load_cached" {
            // Parse a load_cached modifier like:
            // `load_cached(tcx, id) { tcx.queries.on_disk_cache.try_load_query_result(tcx, id) }`
            let args;
            parenthesized!(args in input);
            let tcx = args.parse()?;
            args.parse::<Token![,]>()?;
            let id = args.parse()?;
            let block = input.parse()?;
            Ok(QueryModifier::LoadCached(tcx, id, block))
        } else if modifier == "fatal_cycle" {
            Ok(QueryModifier::FatalCycle)
        } else if modifier == "cycle_delay_bug" {
            Ok(QueryModifier::CycleDelayBug)
        } else if modifier == "no_hash" {
            Ok(QueryModifier::NoHash)
        } else if modifier == "no_force" {
            Ok(QueryModifier::NoForce)
        } else if modifier == "anon" {
            Ok(QueryModifier::Anon)
        } else if modifier == "eval_always" {
            Ok(QueryModifier::EvalAlways)
        } else {
            Err(Error::new(modifier.span(), "unknown query modifier"))
        }
    }
}

/// Ensures only doc comment attributes are used
fn check_attributes(attrs: Vec<Attribute>) -> Result<()> {
    for attr in attrs {
        if !attr.path.is_ident("doc") {
            return Err(Error::new(attr.span(), "attributes not supported on queries"));
        }
    }
    Ok(())
}

/// A compiler query. `query ... { ... }`
struct Query {
    modifiers: List<QueryModifier>,
    name: Ident,
    key: IdentOrWild,
    arg: Type,
    result: ReturnType,
}

impl Parse for Query {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        check_attributes(input.call(Attribute::parse_outer)?)?;

        // Parse the query declaration. Like `query type_of(key: DefId) -> Ty<'tcx>`
        input.parse::<kw::query>()?;
        let name: Ident = input.parse()?;
        let arg_content;
        parenthesized!(arg_content in input);
        let key = arg_content.parse()?;
        arg_content.parse::<Token![:]>()?;
        let arg = arg_content.parse()?;
        let result = input.parse()?;

        // Parse the query modifiers
        let content;
        braced!(content in input);
        let modifiers = content.parse()?;

        Ok(Query {
            modifiers,
            name,
            key,
            arg,
            result,
        })
    }
}

/// A type used to greedily parse another type until the input is empty.
struct List<T>(Vec<T>);

impl<T: Parse> Parse for List<T> {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let mut list = Vec::new();
        while !input.is_empty() {
            list.push(input.parse()?);
        }
        Ok(List(list))
    }
}

/// A named group containing queries.
struct Group {
    name: Ident,
    queries: List<Query>,
}

impl Parse for Group {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let name: Ident = input.parse()?;
        let content;
        braced!(content in input);
        Ok(Group {
            name,
            queries: content.parse()?,
        })
    }
}

struct QueryModifiers {
    /// The description of the query.
    desc: Option<(Option<Ident>, Punctuated<Expr, Token![,]>)>,

    /// Cache the query to disk if the `Expr` returns true.
    cache: Option<(Option<Ident>, Expr)>,

    /// Custom code to load the query from disk.
    load_cached: Option<(Ident, Ident, Block)>,

    /// A cycle error for this query aborting the compilation with a fatal error.
    fatal_cycle: bool,

    /// A cycle error results in a delay_bug call
    cycle_delay_bug: bool,

    /// Don't hash the result, instead just mark a query red if it runs
    no_hash: bool,

    /// Don't force the query
    no_force: bool,

    /// Generate a dep node based on the dependencies of the query
    anon: bool,

    // Always evaluate the query, ignoring its depdendencies
    eval_always: bool,
}

/// Process query modifiers into a struct, erroring on duplicates
fn process_modifiers(query: &mut Query) -> QueryModifiers {
    let mut load_cached = None;
    let mut cache = None;
    let mut desc = None;
    let mut fatal_cycle = false;
    let mut cycle_delay_bug = false;
    let mut no_hash = false;
    let mut no_force = false;
    let mut anon = false;
    let mut eval_always = false;
    for modifier in query.modifiers.0.drain(..) {
        match modifier {
            QueryModifier::LoadCached(tcx, id, block) => {
                if load_cached.is_some() {
                    panic!("duplicate modifier `load_cached` for query `{}`", query.name);
                }
                load_cached = Some((tcx, id, block));
            }
            QueryModifier::Cache(tcx, expr) => {
                if cache.is_some() {
                    panic!("duplicate modifier `cache` for query `{}`", query.name);
                }
                cache = Some((tcx, expr));
            }
            QueryModifier::Desc(tcx, list) => {
                if desc.is_some() {
                    panic!("duplicate modifier `desc` for query `{}`", query.name);
                }
                desc = Some((tcx, list));
            }
            QueryModifier::FatalCycle => {
                if fatal_cycle {
                    panic!("duplicate modifier `fatal_cycle` for query `{}`", query.name);
                }
                fatal_cycle = true;
            }
            QueryModifier::CycleDelayBug => {
                if cycle_delay_bug {
                    panic!("duplicate modifier `cycle_delay_bug` for query `{}`", query.name);
                }
                cycle_delay_bug = true;
            }
            QueryModifier::NoHash => {
                if no_hash {
                    panic!("duplicate modifier `no_hash` for query `{}`", query.name);
                }
                no_hash = true;
            }
            QueryModifier::NoForce => {
                if no_force {
                    panic!("duplicate modifier `no_force` for query `{}`", query.name);
                }
                no_force = true;
            }
            QueryModifier::Anon => {
                if anon {
                    panic!("duplicate modifier `anon` for query `{}`", query.name);
                }
                anon = true;
            }
            QueryModifier::EvalAlways => {
                if eval_always {
                    panic!("duplicate modifier `eval_always` for query `{}`", query.name);
                }
                eval_always = true;
            }
        }
    }
    QueryModifiers {
        load_cached,
        cache,
        desc,
        fatal_cycle,
        cycle_delay_bug,
        no_hash,
        no_force,
        anon,
        eval_always,
    }
}

/// Add the impl of QueryDescription for the query to `impls` if one is requested
fn add_query_description_impl(
    query: &Query,
    modifiers: QueryModifiers,
    impls: &mut proc_macro2::TokenStream
) {
    let name = &query.name;
    let arg = &query.arg;
    let key = &query.key.0;

    // Find out if we should cache the query on disk
    let cache = modifiers.cache.as_ref().map(|(tcx, expr)| {
        let try_load_from_disk = if let Some((tcx, id, block)) = modifiers.load_cached.as_ref() {
            // Use custom code to load the query from disk
            quote! {
                #[inline]
                fn try_load_from_disk(
                    #tcx: TyCtxt<'_, 'tcx, 'tcx>,
                    #id: SerializedDepNodeIndex
                ) -> Option<Self::Value> {
                    #block
                }
            }
        } else {
            // Use the default code to load the query from disk
            quote! {
                #[inline]
                fn try_load_from_disk(
                    tcx: TyCtxt<'_, 'tcx, 'tcx>,
                    id: SerializedDepNodeIndex
                ) -> Option<Self::Value> {
                    tcx.queries.on_disk_cache.try_load_query_result(tcx, id)
                }
            }
        };

        let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
        quote! {
            #[inline]
            fn cache_on_disk(#tcx: TyCtxt<'_, 'tcx, 'tcx>, #key: Self::Key) -> bool {
                #expr
            }

            #try_load_from_disk
        }
    });

    if cache.is_none() && modifiers.load_cached.is_some() {
        panic!("load_cached modifier on query `{}` without a cache modifier", name);
    }

    let desc = modifiers.desc.as_ref().map(|(tcx, desc)| {
        let tcx = tcx.as_ref().map(|t| quote! { #t }).unwrap_or(quote! { _ });
        quote! {
            fn describe(
                #tcx: TyCtxt<'_, '_, '_>,
                #key: #arg,
            ) -> Cow<'static, str> {
                format!(#desc).into()
            }
        }
    });

    if desc.is_some() || cache.is_some() {
        let cache = cache.unwrap_or(quote! {});
        let desc = desc.unwrap_or(quote! {});

        impls.extend(quote! {
            impl<'tcx> QueryDescription<'tcx> for queries::#name<'tcx> {
                #desc
                #cache
            }
        });
    }
}

pub fn rustc_queries(input: TokenStream) -> TokenStream {
    let groups = parse_macro_input!(input as List<Group>);

    let mut query_stream = quote! {};
    let mut query_description_stream = quote! {};
    let mut dep_node_def_stream = quote! {};
    let mut dep_node_force_stream = quote! {};
    let mut no_force_queries = Vec::new();

    for group in groups.0 {
        let mut group_stream = quote! {};
        for mut query in group.queries.0 {
            let modifiers = process_modifiers(&mut query);
            let name = &query.name;
            let arg = &query.arg;
            let result_full = &query.result;
            let result = match query.result {
                ReturnType::Default => quote! { -> () },
                _ => quote! { #result_full },
            };

            let mut attributes = Vec::new();

            // Pass on the fatal_cycle modifier
            if modifiers.fatal_cycle {
                attributes.push(quote! { fatal_cycle });
            };
            // Pass on the cycle_delay_bug modifier
            if modifiers.cycle_delay_bug {
                attributes.push(quote! { cycle_delay_bug });
            };
            // Pass on the no_hash modifier
            if modifiers.no_hash {
                attributes.push(quote! { no_hash });
            };

            let mut attribute_stream = quote! {};

            for e in attributes.into_iter().intersperse(quote! {,}) {
                attribute_stream.extend(e);
            }

            // Add the query to the group
            group_stream.extend(quote! {
                [#attribute_stream] fn #name: #name(#arg) #result,
            });

            let mut attributes = Vec::new();

            // Pass on the anon modifier
            if modifiers.anon {
                attributes.push(quote! { anon });
            };
            // Pass on the eval_always modifier
            if modifiers.eval_always {
                attributes.push(quote! { eval_always });
            };

            let mut attribute_stream = quote! {};
            for e in attributes.into_iter().intersperse(quote! {,}) {
                attribute_stream.extend(e);
            }
            // Create a dep node for the query
            dep_node_def_stream.extend(quote! {
                [#attribute_stream] #name(#arg),
            });

            if modifiers.no_force {
                no_force_queries.push(name.clone());
            } else {
                // Add a match arm to force the query given the dep node
                dep_node_force_stream.extend(quote! {
                    DepKind::#name => {
                        if let Some(key) = RecoverKey::recover($tcx, $dep_node) {
                            force_ex!($tcx, #name, key);
                        } else {
                            return false;
                        }
                    }
                });
            }

            add_query_description_impl(&query, modifiers, &mut query_description_stream);
        }
        let name = &group.name;
        query_stream.extend(quote! {
            #name { #group_stream },
        });
    }

    // Add an arm for the no force queries to panic when trying to force them
    for query in no_force_queries {
        dep_node_force_stream.extend(quote! {
            DepKind::#query |
        });
    }
    dep_node_force_stream.extend(quote! {
        DepKind::Null => {
            bug!("Cannot force dep node: {:?}", $dep_node)
        }
    });

    TokenStream::from(quote! {
        macro_rules! rustc_query_append {
            ([$($macro:tt)*][$($other:tt)*]) => {
                $($macro)* {
                    $($other)*

                    #query_stream

                }
            }
        }
        macro_rules! rustc_dep_node_append {
            ([$($macro:tt)*][$($other:tt)*]) => {
                $($macro)*(
                    $($other)*

                    #dep_node_def_stream
                );
            }
        }
        macro_rules! rustc_dep_node_force {
            ([$dep_node:expr, $tcx:expr] $($other:tt)*) => {
                match $dep_node.kind {
                    $($other)*

                    #dep_node_force_stream
                }
            }
        }
        #query_description_stream
    })
}
<|endoftext|>


<commit_before>#[allow(unused_imports, dead_code)]
mod test {
    use std::str;
    use std::io::{Cursor, Write};

    use regex::Regex;

    use clap::{App, Arg, SubCommand, ArgGroup};

    fn compare<S, S2>(l: S, r: S2) -> bool
        where S: AsRef<str>,
              S2: AsRef<str>
    {
        let re = Regex::new("\x1b[^m]*m").unwrap();
        // Strip out any mismatching \r character on windows that might sneak in on either side
        let ls = l.as_ref().trim().replace("\r", "");
        let rs = r.as_ref().trim().replace("\r", "");
        let left = re.replace_all(&*ls, "");
        let right = re.replace_all(&*rs, "");
        let b = left == right;
        if !b {
            println!("");
            println!("--> left");
            println!("{}", left);
            println!("--> right");
            println!("{}", right);
            println!("--")
        }
        b
    }

    pub fn compare_output(l: App, args: &str, right: &str, stderr: bool) -> bool {
        let mut buf = Cursor::new(Vec::with_capacity(50));
        let res = l.get_matches_from_safe(args.split(' ').collect::<Vec<_>>());
        let err = res.unwrap_err();
        err.write_to(&mut buf).unwrap();
        let content = buf.into_inner();
        let left = String::from_utf8(content).unwrap();
        assert_eq!(stderr, err.use_stderr());
        compare(left, right)
    }

    // Legacy tests from the pyhton script days

    pub fn complex_app() -> App<'static, 'static> {
        let args = "-o --option=[opt]... 'tests options'
                    [positional] 'tests positionals'";
        let opt3_vals = ["fast", "slow"];
        let pos3_vals = ["vi", "emacs"];
        App::new("clap-test")
            .version("v1.4.8")
            .about("tests clap library")
            .author("Kevin K. <kbknapp@gmail.com>")
            .args_from_usage(args)
            .arg(Arg::from_usage("-f --flag... 'tests flags'")
                .global(true))
            .args(&[
                Arg::from_usage("[flag2] -F 'tests flags with exclusions'").conflicts_with("flag").requires("long-option-2"),
                Arg::from_usage("--long-option-2 [option2] 'tests long options with exclusions'").conflicts_with("option").requires("positional2"),
                Arg::from_usage("[positional2] 'tests positionals with exclusions'"),
                Arg::from_usage("-O --Option [option3] 'specific vals'").possible_values(&opt3_vals),
                Arg::from_usage("[positional3]... 'tests specific values'").possible_values(&pos3_vals),
                Arg::from_usage("--multvals [one] [two] 'Tests mutliple values, not mult occs'"),
                Arg::from_usage("--multvalsmo... [one] [two] 'Tests mutliple values, and mult occs'"),
                Arg::from_usage("--minvals2 [minvals]... 'Tests 2 min vals'").min_values(2),
                Arg::from_usage("--maxvals3 [maxvals]... 'Tests 3 max vals'").max_values(3)
            ])
            .subcommand(SubCommand::with_name("subcmd")
                                    .about("tests subcommands")
                                    .version("0.1")
                                    .author("Kevin K. <kbknapp@gmail.com>")
                                    .arg_from_usage("-o --option [scoption]... 'tests options'")
                                    .arg_from_usage("-s --subcmdarg [subcmdarg] 'tests other args'")
                                    .arg_from_usage("[scpositional] 'tests positionals'"))
    }
}
<commit_msg>tests: adds a check_output2 function to compare output to one of two possible solutions<commit_after>#[allow(unused_imports, dead_code)]
mod test {
    use std::str;
    use std::io::{Cursor, Write};

    use regex::Regex;

    use clap::{App, Arg, SubCommand, ArgGroup};

    fn compare<S, S2>(l: S, r: S2) -> bool
        where S: AsRef<str>,
              S2: AsRef<str>
    {
        let re = Regex::new("\x1b[^m]*m").unwrap();
        // Strip out any mismatching \r character on windows that might sneak in on either side
        let ls = l.as_ref().trim().replace("\r", "");
        let rs = r.as_ref().trim().replace("\r", "");
        let left = re.replace_all(&*ls, "");
        let right = re.replace_all(&*rs, "");
        let b = left == right;
        if !b {
            println!("");
            println!("--> left");
            println!("{}", left);
            println!("--> right");
            println!("{}", right);
            println!("--")
        }
        b
    }

    pub fn compare_output(l: App, args: &str, right: &str, stderr: bool) -> bool {
        let mut buf = Cursor::new(Vec::with_capacity(50));
        let res = l.get_matches_from_safe(args.split(' ').collect::<Vec<_>>());
        let err = res.unwrap_err();
        err.write_to(&mut buf).unwrap();
        let content = buf.into_inner();
        let left = String::from_utf8(content).unwrap();
        assert_eq!(stderr, err.use_stderr());
        compare(left, right)
    }
    pub fn compare_output2(l: App, args: &str, right1: &str, right2: &str, stderr: bool) -> bool {
        let mut buf = Cursor::new(Vec::with_capacity(50));
        let res = l.get_matches_from_safe(args.split(' ').collect::<Vec<_>>());
        let err = res.unwrap_err();
        err.write_to(&mut buf).unwrap();
        let content = buf.into_inner();
        let left = String::from_utf8(content).unwrap();
        assert_eq!(stderr, err.use_stderr());
        compare(&*left, right1) || compare(&*left, right2)
    }

    // Legacy tests from the pyhton script days

    pub fn complex_app() -> App<'static, 'static> {
        let args = "-o --option=[opt]... 'tests options'
                    [positional] 'tests positionals'";
        let opt3_vals = ["fast", "slow"];
        let pos3_vals = ["vi", "emacs"];
        App::new("clap-test")
            .version("v1.4.8")
            .about("tests clap library")
            .author("Kevin K. <kbknapp@gmail.com>")
            .args_from_usage(args)
            .arg(Arg::from_usage("-f --flag... 'tests flags'")
                .global(true))
            .args(&[
                Arg::from_usage("[flag2] -F 'tests flags with exclusions'").conflicts_with("flag").requires("long-option-2"),
                Arg::from_usage("--long-option-2 [option2] 'tests long options with exclusions'").conflicts_with("option").requires("positional2"),
                Arg::from_usage("[positional2] 'tests positionals with exclusions'"),
                Arg::from_usage("-O --Option [option3] 'specific vals'").possible_values(&opt3_vals),
                Arg::from_usage("[positional3]... 'tests specific values'").possible_values(&pos3_vals),
                Arg::from_usage("--multvals [one] [two] 'Tests mutliple values, not mult occs'"),
                Arg::from_usage("--multvalsmo... [one] [two] 'Tests mutliple values, and mult occs'"),
                Arg::from_usage("--minvals2 [minvals]... 'Tests 2 min vals'").min_values(2),
                Arg::from_usage("--maxvals3 [maxvals]... 'Tests 3 max vals'").max_values(3)
            ])
            .subcommand(SubCommand::with_name("subcmd")
                                    .about("tests subcommands")
                                    .version("0.1")
                                    .author("Kevin K. <kbknapp@gmail.com>")
                                    .arg_from_usage("-o --option [scoption]... 'tests options'")
                                    .arg_from_usage("-s --subcmdarg [subcmdarg] 'tests other args'")
                                    .arg_from_usage("[scpositional] 'tests positionals'"))
    }
}
<|endoftext|>


<commit_before>mod static_string_indexing;
mod instructions_template;

use static_string_indexing::index_static_strings;
use instructions_template::generate_instructions_rs;

use std::env;
use std::fs;
use std::fs::File;
use std::io::Write;
use std::path::Path;
use std::process::Command;

fn find_prolog_files(libraries: &mut File, prefix: &str, current_dir: &Path) {
    let entries = match current_dir.read_dir() {
        Ok(entries) => entries,
        Err(_) => return,
    };

    for entry in entries.filter_map(Result::ok).map(|e| e.path()) {
        if entry.is_dir() {
            if let Some(file_name) = entry.file_name() {
                let new_prefix = prefix.to_owned() + file_name.to_str().unwrap() + "/";
                find_prolog_files(libraries, &new_prefix, &entry);
            }
        } else if entry.is_file() {
            let ext = std::ffi::OsStr::new("pl");
            if entry.extension() == Some(ext) {
                let contain = String::from_utf8(fs::read(&entry).unwrap()).unwrap();
                let name = entry.file_stem().unwrap().to_str().unwrap();

                let line = format!(
                    "        m.insert(\"{}\",\n{:?});\n",
                    prefix.to_owned() + name,
                    contain
                );

                libraries.write_all(line.as_bytes()).unwrap();
            }
        }
    }
}

fn main() {
    let out_dir = env::var("OUT_DIR").unwrap();
    let dest_path = Path::new(&out_dir).join("libraries.rs");

    let mut libraries = File::create(&dest_path).unwrap();
    let lib_path = Path::new("src/lib");

    libraries
        .write_all(
            b"ref_thread_local::ref_thread_local! {
    pub(crate) static managed LIBRARIES: IndexMap<&'static str, &'static str> = {
        let mut m = IndexMap::new();\n",
        )
        .unwrap();

    find_prolog_files(&mut libraries, "", &lib_path);
    libraries.write_all(b"\n        m\n    };\n}\n").unwrap();

    let instructions_path = Path::new(&out_dir).join("instructions.rs");
    let mut instructions_file = File::create(&instructions_path).unwrap();

    let quoted_output = generate_instructions_rs();

    instructions_file
        .write_all(quoted_output.to_string().as_bytes())
        .unwrap();

    Command::new("rustfmt")
        .arg(instructions_path.as_os_str())
        .spawn().unwrap()
        .wait().unwrap();

    let static_atoms_path = Path::new(&out_dir).join("static_atoms.rs");
    let mut static_atoms_file = File::create(&static_atoms_path).unwrap();

    let quoted_output = index_static_strings(&instructions_path);

    static_atoms_file
        .write_all(quoted_output.to_string().as_bytes())
        .unwrap();

    Command::new("rustfmt")
        .arg(static_atoms_path.as_os_str())
        .spawn().unwrap()
        .wait().unwrap();

    println!("cargo:rerun-if-changed=src/");
}
<commit_msg>Improve rustfmt missing error during build (#1379)<commit_after>mod static_string_indexing;
mod instructions_template;

use static_string_indexing::index_static_strings;
use instructions_template::generate_instructions_rs;

use std::env;
use std::fs;
use std::fs::File;
use std::io::Write;
use std::path::Path;
use std::process::{self, Command};

fn find_prolog_files(libraries: &mut File, prefix: &str, current_dir: &Path) {
    let entries = match current_dir.read_dir() {
        Ok(entries) => entries,
        Err(_) => return,
    };

    for entry in entries.filter_map(Result::ok).map(|e| e.path()) {
        if entry.is_dir() {
            if let Some(file_name) = entry.file_name() {
                let new_prefix = prefix.to_owned() + file_name.to_str().unwrap() + "/";
                find_prolog_files(libraries, &new_prefix, &entry);
            }
        } else if entry.is_file() {
            let ext = std::ffi::OsStr::new("pl");
            if entry.extension() == Some(ext) {
                let contain = String::from_utf8(fs::read(&entry).unwrap()).unwrap();
                let name = entry.file_stem().unwrap().to_str().unwrap();

                let line = format!(
                    "        m.insert(\"{}\",\n{:?});\n",
                    prefix.to_owned() + name,
                    contain
                );

                libraries.write_all(line.as_bytes()).unwrap();
            }
        }
    }
}

fn main() {
    let out_dir = env::var("OUT_DIR").unwrap();
    let dest_path = Path::new(&out_dir).join("libraries.rs");

    let mut libraries = File::create(&dest_path).unwrap();
    let lib_path = Path::new("src/lib");

    libraries
        .write_all(
            b"ref_thread_local::ref_thread_local! {
    pub(crate) static managed LIBRARIES: IndexMap<&'static str, &'static str> = {
        let mut m = IndexMap::new();\n",
        )
        .unwrap();

    find_prolog_files(&mut libraries, "", &lib_path);
    libraries.write_all(b"\n        m\n    };\n}\n").unwrap();

    let instructions_path = Path::new(&out_dir).join("instructions.rs");
    let mut instructions_file = File::create(&instructions_path).unwrap();

    let quoted_output = generate_instructions_rs();

    instructions_file
        .write_all(quoted_output.to_string().as_bytes())
        .unwrap();

    format_generated_file(instructions_path.as_path());

    let static_atoms_path = Path::new(&out_dir).join("static_atoms.rs");
    let mut static_atoms_file = File::create(&static_atoms_path).unwrap();

    let quoted_output = index_static_strings(&instructions_path);

    static_atoms_file
        .write_all(quoted_output.to_string().as_bytes())
        .unwrap();

    format_generated_file(static_atoms_path.as_path());

    println!("cargo:rerun-if-changed=src/");
}

fn format_generated_file(path: &Path) {
    Command::new("rustfmt")
        .arg(path.as_os_str())
        .spawn()
        .unwrap_or_else(|err| {
            eprintln!(
                "{}: rustfmt (https://github.com/rust-lang/rustfmt) \
                is required build time dependency and must be available \
                in the environment's PATH variable.",
                err
            );
            process::exit(1);
        })
        .wait()
        .unwrap();
}
<|endoftext|>


<commit_before>#![allow(non_snake_case, non_upper_case_globals)]
// #![cfg_attr(test, feature(test))]
extern crate time;

use time::precise_time_s;
use std::mem;

///////////////////////////////////////////////////////////////////////////////
// 1. Calibration
///////////////////////////////////////////////////////////////////////////////
const aalpha:f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const bbeta:f64 = 0.95;           // Discount factor

// Productivity values
const vProductivity:[f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
const mTransition:[[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                   [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                   [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                   [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                   [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const nGridCapital: usize = 17820;
const nGridProductivity: usize = 5;

fn solve(print: bool) -> f64 {
  /////////////////////////////////////////////////////////////////////////////
  // 2. Steady State
  /////////////////////////////////////////////////////////////////////////////

    let capitalSteadyState: f64 = (aalpha * bbeta).powf(1_f64 / (1_f64 - aalpha));
    let outputSteadyState: f64 = capitalSteadyState.powf(aalpha);
    let consumptionSteadyState: f64 = outputSteadyState - capitalSteadyState;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 outputSteadyState,
                 capitalSteadyState,
                 consumptionSteadyState);
    }

    let mut vGridCapital = [0f64; nGridCapital];

    for (i, val) in vGridCapital.iter_mut().enumerate() {
        *val = 0.5 * capitalSteadyState + 0.00001 * (i as f64)
    }

    // 3. Required matrices and vectors
    let mut mValueFunction = vec![[0f64; nGridProductivity]; nGridCapital];
    let mut mValueFunctionNew = mValueFunction.clone();
    let mut mPolicyFunction = mValueFunction.clone();
    let mut expectedValueFunction = mValueFunction.clone();

    // 4. We pre-build output for each point in the grid
    let mOutput : Vec<[f64; nGridProductivity]> = (0..nGridCapital).map(|nCapital| {
        let mut arr = [0.0; nGridProductivity];

        for (nProductivity, slot) in arr.iter_mut().enumerate() {
            *slot = vProductivity[nProductivity] * vGridCapital[nCapital].powf(aalpha)
        }

        arr
    }).collect();

  // 5. Main iteration
  // TODO: one could implement a macro for the multiple declarations
    let mut maxDifference = 10_f64;
    const tolerance: f64 = 0.0000001; // compiler warn: variable does not need to be mutable
    let mut capitalChoice: f64;

    let mut iteration = 0;

    while maxDifference > tolerance {
        for nProductivity in 0..nGridProductivity {
            let transitions = &mTransition[nProductivity];

            for nCapital in 0..nGridCapital {
                let value_fns = &mValueFunction[nCapital];

                let expected_value = transitions.iter()
                                                .zip(value_fns.iter())
                                                .fold(0.0f64, |acc, (transition, value_fn)| {
                    acc + (transition * value_fn)
                });

                expectedValueFunction[nCapital][nProductivity] = expected_value;
            }
        }

        for nProductivity in 0..nGridProductivity {

        // We start from previous choice (monotonicity of policy function)
            let mut gridCapitalNextPeriod = 0;

            for nCapital in 0..nGridCapital {

                let mut valueHighSoFar = -100000.0;

                for nCapitalNextPeriod in gridCapitalNextPeriod..nGridCapital {

                    let consumption = mOutput[nCapital][nProductivity] -
                                  vGridCapital[nCapitalNextPeriod];
                    let valueProvisional = (1_f64 - bbeta) * (consumption.ln()) +
                                       bbeta *
                                       expectedValueFunction[nCapitalNextPeriod][nProductivity];

                    if valueProvisional > valueHighSoFar {
                        valueHighSoFar = valueProvisional;
                        capitalChoice = vGridCapital[nCapitalNextPeriod];
                        gridCapitalNextPeriod = nCapitalNextPeriod;
                    } else {
                        break; // We break when we have achieved the max
                    }

                    mValueFunctionNew[nCapital][nProductivity] = valueHighSoFar;
                    mPolicyFunction[nCapital][nProductivity] = capitalChoice;
                }

            }

        }

        {
            let old_vals = mValueFunction.iter().flat_map(|arr| arr.iter());
            let new_vals = mValueFunctionNew.iter().flat_map(|arr| arr.iter());

            maxDifference = -100000.0;
            for diff in old_vals.zip(new_vals).map(|(old, new)| (old - new).abs()) {
                if diff > maxDifference {
                    maxDifference = diff
                }
            }
        }

        // swap buffers after the loop
        mem::swap(&mut mValueFunction, &mut mValueFunctionNew);

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1) {
            println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
        }
    }

    if print {
        println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
    }

    mPolicyFunction[999][2]
}

fn main() {
    let sample = false;

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}

// FIXME: Can't use bench as the loop is too slow (takes 8mins to get samples)
// #[cfg(test)]
// mod test {
//     extern crate test;
//     use self::test::Bencher;
//     use super::solve;

//     #[bench]
//     fn bench(b: &mut Bencher) {
//         b.iter(|| {
//             // use `test::black_box` to prevent compiler optimizations from disregarding
//             // unused values
//             test::black_box(solve(false));
//         });
//     }
// }
<commit_msg>Reduce memory access in main loop<commit_after>#![allow(non_snake_case, non_upper_case_globals)]
// #![cfg_attr(test, feature(test))]
extern crate time;

use time::precise_time_s;
use std::mem;

///////////////////////////////////////////////////////////////////////////////
// 1. Calibration
///////////////////////////////////////////////////////////////////////////////
const aalpha:f64 = 0.33333333333;  // Elasticity of output w.r.t. capital
const bbeta:f64 = 0.95;           // Discount factor

// Productivity values
const vProductivity:[f64; 5] = [0.9792, 0.9896, 1.0000, 1.0106, 1.0212];

// Transition matrix
const mTransition:[[f64; 5]; 5] = [[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],
                                   [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],
                                   [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],
                                   [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],
                                   [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]];

// Dimensions to generate the grid of capital
const nGridCapital: usize = 17820;
const nGridProductivity: usize = 5;

fn solve(print: bool) -> f64 {
  /////////////////////////////////////////////////////////////////////////////
  // 2. Steady State
  /////////////////////////////////////////////////////////////////////////////

    let capitalSteadyState: f64 = (aalpha * bbeta).powf(1_f64 / (1_f64 - aalpha));
    let outputSteadyState: f64 = capitalSteadyState.powf(aalpha);
    let consumptionSteadyState: f64 = outputSteadyState - capitalSteadyState;

    if print {
        println!("Output = {}, Capital = {}, Consumption = {}",
                 outputSteadyState,
                 capitalSteadyState,
                 consumptionSteadyState);
    }

    let mut vGridCapital = [0f64; nGridCapital];

    for (i, val) in vGridCapital.iter_mut().enumerate() {
        *val = 0.5 * capitalSteadyState + 0.00001 * (i as f64)
    }

    // 3. Required matrices and vectors
    let mut mValueFunction = vec![[0f64; nGridProductivity]; nGridCapital];
    let mut mValueFunctionNew = mValueFunction.clone();
    let mut mPolicyFunction = mValueFunction.clone();
    let mut expectedValueFunction = mValueFunction.clone();

    // 4. We pre-build output for each point in the grid
    let mOutput : Vec<[f64; nGridProductivity]> = (0..nGridCapital).map(|nCapital| {
        let mut arr = [0.0; nGridProductivity];

        for (nProductivity, slot) in arr.iter_mut().enumerate() {
            *slot = vProductivity[nProductivity] * vGridCapital[nCapital].powf(aalpha)
        }

        arr
    }).collect();

  // 5. Main iteration
  // TODO: one could implement a macro for the multiple declarations
    let mut maxDifference = 10_f64;
    const tolerance: f64 = 0.0000001; // compiler warn: variable does not need to be mutable

    let mut iteration = 0;

    while maxDifference > tolerance {
        for nProductivity in 0..nGridProductivity {
            let transitions = &mTransition[nProductivity];

            for nCapital in 0..nGridCapital {
                let value_fns = &mValueFunction[nCapital];

                let expected_value = transitions.iter()
                                                .zip(value_fns.iter())
                                                .fold(0.0f64, |acc, (transition, value_fn)| {
                    acc + (transition * value_fn)
                });

                expectedValueFunction[nCapital][nProductivity] = expected_value;
            }
        }

        for nProductivity in 0..nGridProductivity {
            // We start from previous choice (monotonicity of policy function)
            let mut gridCapitalNextPeriod = 0;

            for nCapital in 0..nGridCapital {
                let mut valueHighSoFar = -100000.0;
                let mut capitalChoice = vGridCapital[0];
                let mOutput_cache = mOutput[nCapital][nProductivity];

                for nCapitalNextPeriod in gridCapitalNextPeriod..nGridCapital {
                    let consumption = mOutput_cache - vGridCapital[nCapitalNextPeriod];
                    let valueProvisional = (1_f64 - bbeta) * (consumption.ln()) +
                                       bbeta *
                                       expectedValueFunction[nCapitalNextPeriod][nProductivity];

                    if valueProvisional > valueHighSoFar {
                        valueHighSoFar = valueProvisional;
                        capitalChoice = vGridCapital[nCapitalNextPeriod];
                        gridCapitalNextPeriod = nCapitalNextPeriod;
                    } else {
                        break; // We break when we have achieved the max
                    }
                }

                mValueFunctionNew[nCapital][nProductivity] = valueHighSoFar;
                mPolicyFunction[nCapital][nProductivity] = capitalChoice;
            }
        }

        {
            let old_vals = mValueFunction.iter().flat_map(|arr| arr.iter());
            let new_vals = mValueFunctionNew.iter().flat_map(|arr| arr.iter());

            maxDifference = -100000.0;
            for diff in old_vals.zip(new_vals).map(|(old, new)| (old - new).abs()) {
                if diff > maxDifference {
                    maxDifference = diff
                }
            }
        }

        // swap buffers after the loop
        mem::swap(&mut mValueFunction, &mut mValueFunctionNew);

        iteration += 1;
        if print && (iteration % 10 == 0 || iteration == 1) {
            println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
        }
    }

    if print {
        println!("Iteration = {}, Sup Diff = {}", iteration, maxDifference);
    }

    mPolicyFunction[999][2]
}

fn main() {
    let sample = false;

    if sample {
        let mut samples = (0..5).map(|i| {
            let cpu0 = precise_time_s();
            let result = solve(false);
            let cpu1 = precise_time_s();

            assert_eq!(result, 0.1465491436956954);

            let diff = cpu1 - cpu0;
            println!("Sample #{}, Time: {}s", i + 1, diff);
            diff
        }).collect::<Vec<f64>>();

        samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
        println!("\nMedian time is = {}", samples[2]);
    } else {
        let cpu0 = precise_time_s();
        let result = solve(true);
        let cpu1 = precise_time_s();

        println!("My check = {}\n", result);
        assert_eq!(result, 0.1465491436956954);
        println!("Elapsed time is = {}", cpu1 - cpu0);
    }
}

// FIXME: Can't use bench as the loop is too slow (takes 8mins to get samples)
// #[cfg(test)]
// mod test {
//     extern crate test;
//     use self::test::Bencher;
//     use super::solve;

//     #[bench]
//     fn bench(b: &mut Bencher) {
//         b.iter(|| {
//             // use `test::black_box` to prevent compiler optimizations from disregarding
//             // unused values
//             test::black_box(solve(false));
//         });
//     }
// }
<|endoftext|>